ptyrad.analysis.extract#

Pure extraction functions for PtyRAD model.hdf5 outputs.

Every public extraction worker accepts either the full loaded dict (from ptyrad.io.load.load_ptyrad()) or pre-sliced sub-dicts (opt_tensors, model_attrs) so callers can use these helpers without instantiating ptyrad.analysis.Analyzer.

Geometry / FOV helpers (get_probe_center_positions, get_scanned_fov_bbox, apply_fov) are re-exported on purpose: they encode the single most error-prone convention in the saved file (crop_pos is the top-left of the probe window, not the center). The public position extractor exposes both conventions explicitly via target='probe' and target='crop'.

Functions

apply_fov(arr, bbox)

Crop the last two dims of arr to an inclusive bbox.

extract_keys(data[, delimiter])

Return a flat dotted-key listing of a loaded data dict.

extract_loss_curves(data)

Pull the loss-history fields from a loaded data dict.

extract_object(data_or_opt, *[, fov, ...])

Return the complex object obja * exp(1j * objp).

extract_object_amplitude(data_or_opt, *[, ...])

Return the saved object amplitude obja (no recomposition).

extract_object_phase(data_or_opt, *[, fov, ...])

Return the saved object phase objp (no recomposition).

extract_probe(data_or_opt, *[, space, ...])

Return the probe modes.

extract_probe_positions(data_or_opt[, ...])

Return positions, shape (N, 2) as [y, x], float32.

extract_provenance(path)

Read the provenance_json root attribute from an HDF5 file.

get_probe_center_positions(crop_pos, probe_shape)

Convert top-left crop_pos to probe-center positions.

get_scanned_fov_bbox(crop_pos, probe_shape)

Inclusive (y_min, y_max, x_min, x_max) of probe-center positions.

ptyrad.analysis.extract.apply_fov(arr, bbox)[source]#

Crop the last two dims of arr to an inclusive bbox.

bbox is (y_min, y_max, x_min, x_max) and both bounds are inclusive. Returns arr unchanged when bbox is None.

Parameters:
  • arr (ndarray)

  • bbox (tuple[int, int, int, int] | None)

Return type:

ndarray

ptyrad.analysis.extract.extract_keys(data, delimiter='.')[source]#

Return a flat dotted-key listing of a loaded data dict.

Thin wrapper over ptyrad.io.hierarchy.list_nested_keys() that accepts a nested dict (the form returned by load_ptyrad) and yields keys like 'optimizable_tensors.probe', 'model_attributes.crop_pos'.

Parameters:
  • data (Mapping[str, Any])

  • delimiter (str)

Return type:

list[str]

ptyrad.analysis.extract.extract_loss_curves(data)[source]#

Pull the loss-history fields from a loaded data dict.

Returns a dict with these keys (missing fields become None):

  • loss_iters(niter, 2) float64 ndarray of (iter_number, total_loss) pairs.

  • batch_lossesdict[str, list[float]] of per-loss batch values from the most recent iteration.

  • avg_lossesdict[str, float] of averaged batch losses.

  • niter — final iteration number (int).

Parameters:

data (Mapping[str, Any])

Return type:

dict[str, Any]

ptyrad.analysis.extract.extract_object(data_or_opt, *, fov='full', model_attrs=None, probe_shape=None, as_torch=False, device='cpu')[source]#

Return the complex object obja * exp(1j * objp).

Shape (omode, Nz, Ny, Nx), complex64. The 4D shape is preserved even for single-slice runs (Nz == 1) so downstream code can branch on Nz directly; squeeze yourself if you need 2D.

fov='crop' clips to the scanned-FOV bbox (see get_scanned_fov_bbox()); it requires crop_pos from model_attrs. probe_shape is auto-inferred from opt_tensors['probe'] and only needs to be passed when callers have stripped the probe out of opt_tensors.

Parameters:
  • data_or_opt (Mapping[str, Any])

  • fov (Literal['full', 'crop'])

  • model_attrs (Mapping[str, Any] | None)

  • probe_shape (tuple[int, int] | None)

  • as_torch (bool)

  • device (str)

ptyrad.analysis.extract.extract_object_amplitude(data_or_opt, *, fov='full', model_attrs=None, probe_shape=None, as_torch=False, device='cpu')[source]#

Return the saved object amplitude obja (no recomposition).

Shape (omode, Nz, Ny, Nx), float32. fov behaves the same as extract_object().

Parameters:
  • data_or_opt (Mapping[str, Any])

  • fov (Literal['full', 'crop'])

  • model_attrs (Mapping[str, Any] | None)

  • probe_shape (tuple[int, int] | None)

  • as_torch (bool)

  • device (str)

ptyrad.analysis.extract.extract_object_phase(data_or_opt, *, fov='full', model_attrs=None, probe_shape=None, as_torch=False, device='cpu')[source]#

Return the saved object phase objp (no recomposition).

Shape (omode, Nz, Ny, Nx), float32. fov behaves the same as extract_object(). Sign convention follows PtychoModel: the complex object is obja * exp(1j * objp); no sign flip is applied here.

Parameters:
  • data_or_opt (Mapping[str, Any])

  • fov (Literal['full', 'crop'])

  • model_attrs (Mapping[str, Any] | None)

  • probe_shape (tuple[int, int] | None)

  • as_torch (bool)

  • device (str)

ptyrad.analysis.extract.extract_probe(data_or_opt, *, space='real', as_torch=False, device='cpu')[source]#

Return the probe modes. Shape (pmode, Ny, Nx), complex64.

space='real' returns the stored real-space complex wavefunction (the form held in optimizable_tensors['probe'] after the view_as_complex post-process at save time).

space='fourier' returns fftshift(fft2(ifftshift(probe), norm='ortho')) along the last two axes. The fftshift sandwich matches ptyrad.plotting.plot_probe_modes() exactly so amplitudes and phases line up between getters and plotters; in particular the pre-fftshift to the corner avoids the checkerboard-phase artifact that a plain fft2 would produce.

Parameters:
  • data_or_opt (Mapping[str, Any])

  • space (Literal['real', 'fourier'])

  • as_torch (bool)

  • device (str)

ptyrad.analysis.extract.extract_probe_positions(data_or_opt, model_attrs=None, *, fov='full', target='probe', units='px', include_sub_px_shifts=True, as_torch=False, device='cpu')[source]#

Return positions, shape (N, 2) as [y, x], float32.

Parameters:
  • data_or_opt (Mapping[str, Any]) – Full data dict from ptyrad.io.load.load_ptyrad(), or the optimizable_tensors sub-dict. model_attrs must be supplied when the latter is passed.

  • model_attrs (Mapping[str, Any] | None) – Optional explicit model_attributes sub-dict. Must contain crop_pos (always) and dx (only when units='Ang').

  • fov (Literal['full', 'crop']) – 'full' returns positions in the full saved-object coordinate frame. 'crop' subtracts the probe-center scanned-FOV bbox top-left (see get_scanned_fov_bbox()) so that, regardless of target, the returned positions live in the same coordinate frame as the array returned by extract_object(fov='crop'). Because that cropped FOV is anchored on probe centers, target='crop' can produce negative local coordinates near the top/left edge.

  • target (Literal['probe', 'crop']) – 'probe' (default) returns probe-center positions (crop_pos + (Hp // 2, Wp // 2)). 'crop' returns the top-left crop-window positions used by PtychoModel.get_obj_patches(). The two differ by exactly (Hp // 2, Wp // 2) and that offset survives fov='crop' because the same bbox is subtracted in both cases.

  • units (Literal['px', 'pixel', 'Ang']) – 'px' / 'pixel' returns object-space pixels. 'Ang' multiplies by model_attrs['dx'] to return Ångströms.

  • include_sub_px_shifts (bool) – When True (default), add opt_tensors['probe_pos_shifts'] if present. These are the optimized sub-pixel offsets relative to the integer crop_pos grid. Older saves without the key silently fall back to zero shifts.

  • as_torch (bool)

  • device (str)

ptyrad.analysis.extract.extract_provenance(path)[source]#

Read the provenance_json root attribute from an HDF5 file.

The provenance JSON is stored as a root-level HDF5 attribute, not as a dataset, so ptyrad.io.load.load_ptyrad() silently drops it. This helper opens the file directly via ptyrad.io.provenance.load_provenance_from_h5() and returns the parsed dict (with keys like 'probe', 'obj', 'pos', 'tilt'). Returns None when the attribute is missing or the JSON parse fails.

Parameters:

path (str | PathLike)

Return type:

dict | None

ptyrad.analysis.extract.get_nested(d, key, delimiter='.', safe=False, default=None)[source]#

Get a value from a nested dictionary either safely (return default if not found) or stricly to fail early.

Parameters: - d (dict): The dictionary to traverse. - key (str, or list or tuple of string): A sequence of keys to access nested values. - delimiter (str): The string used to seperate different parts of the displayed key path - safe (boolean): The flag to switch between safe/strict mode of getting values from a nested dict. - default: The value to return if any key is missing or intermediate value is None.

Returns: - The nested value if found, otherwise default in safe mode or error in strict mode.

ptyrad.analysis.extract.get_probe_center_positions(crop_pos, probe_shape, probe_pos_shifts=None)[source]#

Convert top-left crop_pos to probe-center positions.

The crop_pos stored in model.hdf5 is the top-left corner of the probe window in object pixels — that is the form consumed by PtychoModel.get_obj_patches(), not the visual probe center. The probe center sits (Hp // 2, Wp // 2) further in.

Parameters:
  • crop_pos (array-like, shape (N, 2), int) – Top-left [y, x] of each probe window in object pixels.

  • probe_shape (tuple) – Probe shape; only the last two dims (Hp, Wp) are used.

  • probe_pos_shifts (array-like, shape (N, 2), optional) – Optimized sub-pixel offsets; added to the center positions when given.

Returns:

Shape (N, 2), float32, [y, x] order, in object pixels.

Return type:

np.ndarray

ptyrad.analysis.extract.get_scanned_fov_bbox(crop_pos, probe_shape)[source]#

Inclusive (y_min, y_max, x_min, x_max) of probe-center positions.

The bbox is computed from the integer probe-center positions (crop_pos + probe_shape // 2) without probe_pos_shifts — matching the bbox math in ptyrad.io.save.save_results() so the result lines up exactly with the saved *_crop TIFFs (sliced as [y_min:y_max + 1, x_min:x_max + 1]).

Parameters:
  • crop_pos (ndarray)

  • probe_shape (tuple[int, int])

Return type:

tuple[int, int, int, int]