Skip to content

Ptycho improvements #570

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ercius merged 12 commits into from
Feb 11, 2026
Merged

Ptycho improvements #570

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
2f05473
remove some logging. Reduce n_trials to 25
ercius Jan 30, 2026
86eeafe
use int for defocus and defocus search range. This allows comparison …
ercius Jan 30, 2026
a8e3c4e
remove unused parameters
ercius Jan 30, 2026
77de16a
add bool parameters to enable optimization of diffraction_angle, defo…
ercius Jan 31, 2026
07e3f38
fix json comma
ercius Jan 31, 2026
501aefc
fix typos. Fix error where probe step size was x10 twice
ercius Jan 31, 2026
14d3c53
always use parallax for hyperparameter optimization
ercius Feb 2, 2026
aac6738
use only needed parameters.
ercius Feb 2, 2026
fc30716
add max_batch_size parameter to better manage memory
ercius Feb 2, 2026
d02a30b
fix up all the parameters
ercius Feb 2, 2026
e229893
fix json commas
ercius Feb 2, 2026
5f0a225
remove naming typo
ercius Feb 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 42 additions & 42 deletions operators/quantem-direct-ptycho/operator.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,6 @@
}
],
"parameters": [
{
"name": "calculation_frequency",
"label": "Calculation Frequency",
"type": "int",
"default": "100",
"description": "Number of frames to accumulate before recalculating the center and emitting a BF image.",
"required": true
},
{
"name": "max_concurrent_scans",
"label": "Max Concurrent Scans",
"type": "int",
"default": "1",
"description": "Maximum number of scans to keep in memory simultaneously. Oldest scans are evicted when this limit is exceeded.",
"required": false
},
{
"name": "accelerating_voltage",
"label": "Accelerating voltage",
Expand Down Expand Up @@ -63,11 +47,27 @@
{
"name": "initial_defocus",
"label": "STEM defocus",
"type": "float",
"default": "0.0",
"type": "int",
"default": "0",
"description": "The STEM defocus in nm.",
"required": false
},
{
"name": "optimize_defocus",
"label": "Optimize the defocus",
"type": "bool",
"default": true,
"description": "Use the optimization routine to optimize the defocus.",
"required": false
},
{
"name": "defocus_search_range",
"label": "Defocus search range",
"type": "int",
"default": "50",
"description": "The defocus search range in nanometers. Set to zero to disable defocus optimization.",
"required": false
},
{
"name": "diffraction_rotation_angle",
"label": "Diffraction pattern rotation angle",
Expand All @@ -76,6 +76,30 @@
"description": "The rotation of the diffraction pattern on the detector in degrees.",
"required": false
},
{
"name": "optimize_rotation_angle",
"label": "Optimize the rotation angle",
"type": "bool",
"default": false,
"description": "Use the optimization routine to optimize the rotation angle.",
"required": false
},
{
"name": "maximum_C12_magnitude",
"label": "Maximum C12 magnitude",
"type": "int",
"default": "2",
"description": "The maximum C12 magnitude in nanometers.",
"required": false
},
{
"name": "optimize_C12",
"label": "Optimize the C12",
"type": "bool",
"default": true,
"description": "Use the optimization routine to optimize the C12.",
"required": false
},
{
"name": "crop_probes",
"label": "Crop probes on each side",
Expand All @@ -92,22 +116,6 @@
"description": "The upsampling factor for the final reconstruction.",
"required": false
},
{
"name": "defocus_search_range_nm",
"label": "Defocus search range",
"type": "float",
"default": "50.0",
"description": "The defocus search range in nanometers. Unused if defocus is input.",
"required": false
},
{
"name": "maximum_C12_magnitude_nm",
"label": "Maximum C12 magnitude",
"type": "int",
"default": "2",
"description": "The maximum C12 magnitude in nanometers.",
"required": false
},
{
"name": "deconvolution_kernel",
"label": "Deconvolution kernel",
Expand All @@ -116,14 +124,6 @@
"options": ["parallax", "ssb", "icom"],
"description": "The deconvolution kernel.",
"required": false
},
{
"name": "use_optimization",
"label": "Use optimization routine",
"type": "bool",
"default": true,
"description": "Use the optimization routine with initial parameters.",
"required": false
}
],
"parallel_config": {
Expand Down
99 changes: 47 additions & 52 deletions operators/quantem-direct-ptycho/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ def quantem_direct_ptycho(
scan_number = batch.header.scan_number

# --- 2. Get or Create FrameAccumulator ---
max_concurrent_scans = int(parameters.get("max_concurrent_scans", 1))
max_concurrent_scans = 1 # TODO: remove this

if scan_number not in accumulators:
# Check if we need to evict old accumulators before creating new one
Expand Down Expand Up @@ -114,56 +114,50 @@ def quantem_direct_ptycho(
return None

# --- 5. Perform Calculation ---
logger.info(
f"Scan {scan_number}: Triggering calculation after {accumulator.num_batches_added} messages."
)
logger.info(f"Accumulator finished: {accumulator.finished}")

logger.info(f"Scan {scan_number}: Calculating ptycho images.")

# Calculation parameters
probe_semiangle = parameters.get("probe_semiangle", 25.0)
energy = parameters.get("accelerating_voltage", 300e3)
probe_step_size = parameters.get(
"probe_step_size", 0.1
) # test data set: 0.14383155 nm
crop_probes = parameters.get("crop_probes", 0)
probe_semiangle = parameters.get("probe_semiangle", 25.0)
# test data set: 0.14383155 nm probe step size
probe_step_size_nm = parameters.get("probe_step_size", 0.1)
probe_step_size_A = probe_step_size_nm * 10
upsampling_factor = parameters.get("upsampling_factor", 2)

# Parameters for optimize_hyperparameters function
initial_defocus_nm = parameters.get(
"initial_defocus", None
) # in nanometers, can be None
if initial_defocus_nm is not None:
initial_defocus_nm = initial_defocus_nm
initial_defocus_A = initial_defocus_nm * 10 # convert to Angstroms
optimize_defocus = bool(parameters.get("optimize_defocus", True))
if optimize_defocus:
defocus_search_max_nm = parameters.get("defocus_search_range", 50)
defocus_search_max_A = defocus_search_max_nm * 10 # convert to Angstroms
else:
initial_defocus_A = None

diffraction_rotation_angle = parameters.get(
"diffraction_rotation_angle", None
) # in degrees, can be None
if diffraction_rotation_angle is not None:
diffraction_rotation_angle = diffraction_rotation_angle
rotation_angle = diffraction_rotation_angle * np.pi / 180 # convert to radians
defocus_search_max_A = 0

initial_defocus_nm = parameters.get("initial_defocus", 0)
initial_defocus_A = initial_defocus_nm * 10

# Only search in one direction from initial guess
if initial_defocus_A > 0:
defocus_search_range_A = (0, defocus_search_max_A)
elif initial_defocus_A < 0:
defocus_search_range_A = (-defocus_search_max_A, 0)
else:
rotation_angle = None
defocus_search_range_A = (-defocus_search_max_A, defocus_search_max_A)

defocus_search_range_nm = parameters.get(
"defocus_search_range", 50
) # in nanometers
defocus_search_range_A = defocus_search_range_nm * 10 # convert to Angstroms
# in degrees
diffraction_rotation_angle_deg = parameters.get("diffraction_rotation_angle", 0)
rotation_angle = diffraction_rotation_angle_deg * np.pi / 180 # convert to radians

maximum_C12_magnitude_nm = parameters.get(
"maximum_C12_magnitude", 10
) # in nanometers
maximum_C12_magnitude_A = maximum_C12_magnitude_nm * 10 # convert to Angstroms
optimize_angle = bool(parameters.get("optimize_rotation_angle", False))

deconvolution_kernel = parameters.get("deconvolution_kernel", "parallax")
optimize_C12 = bool(parameters.get("optimize_C12", False))
if optimize_C12:
maximum_C12_magnitude_nm = parameters.get("maximum_C12_magnitude", 10)
maximum_C12_magnitude_A = maximum_C12_magnitude_nm * 10 # convert to Angstroms
else:
maximum_C12_magnitude_A = None

# Determine whether to use optimization or manual settings
use_optimization = bool(parameters.get("use_optimization", True))
deconvolution_kernel = parameters.get("deconvolution_kernel", "parallax")

crop_probes = parameters.get("crop_probes", 0)
if crop_probes == 0:
logger.info(f"Scan {scan_number}: No cropping of probes applied.")
dense_data = accumulator[:, :-1, :, :].to_dense() ## remove the flyback column
Expand All @@ -173,6 +167,7 @@ def quantem_direct_ptycho(
crop_probes:-crop_probes, crop_probes : -crop_probes - 1, :, :
].to_dense() ## crop the edges if needed and remove the flyback column

# Convert SparseArray to Dataset4dstem
dset = em.datastructures.Dataset4dstem.from_array(array=dense_data)
logger.debug(f"dense shape = {dense_data.shape}")

Expand All @@ -184,10 +179,7 @@ def quantem_direct_ptycho(
dset.sampling[3] = probe_semiangle / probe_R
dset.units[2:] = ["mrad", "mrad"]

dset.sampling[0] = (
probe_step_size * 10
) ## convert to be Anggstrom for quantem. distiller will give nanometers.
dset.sampling[1] = probe_step_size * 10
dset.sampling[0:2] = probe_step_size_A
dset.units[0:2] = ["A", "A"]

logger.info(f"Scan {scan_number}: Start direct ptycho")
Expand All @@ -207,25 +199,28 @@ def quantem_direct_ptycho(
rotation_angle=rotation_angle, # need radians
)

if use_optimization:
if optimize_C12 or optimize_angle or optimize_defocus:
logger.info(f"Scan {scan_number}: Optimizing hyperparameters")

# Build optimization aberration coefficients
opt_aberration_coefs = {}
if initial_defocus_A is None:

if optimize_defocus:
opt_aberration_coefs["C10"] = OptimizationParameter(
defocus_search_range_A, defocus_search_range_A
defocus_search_range_A[0], defocus_search_range_A[1]
)
else:
opt_aberration_coefs["C10"] = -initial_defocus_A

opt_aberration_coefs["C12"] = OptimizationParameter(
0, maximum_C12_magnitude_A
)
opt_aberration_coefs["phi12"] = OptimizationParameter(-np.pi / 2, np.pi / 2)
if optimize_C12:
opt_aberration_coefs["C12"] = OptimizationParameter(
0, maximum_C12_magnitude_A
)
opt_aberration_coefs["phi12"] = OptimizationParameter(
-np.pi / 2, np.pi / 2
)

# Build rotation angle optimization
if rotation_angle is None:
if optimize_angle:
opt_rotation_angle = OptimizationParameter(0, np.pi)
else:
opt_rotation_angle = rotation_angle
Expand All @@ -234,7 +229,7 @@ def quantem_direct_ptycho(
aberration_coefs=opt_aberration_coefs,
rotation_angle=opt_rotation_angle,
deconvolution_kernel=deconvolution_kernel,
n_trials=50,
n_trials=25,
max_batch_size=10,
)
else:
Expand Down
Loading