Utils

Utils

ptychi.utils.add_additional_opr_probe_modes_to_probe(probe: Tensor, n_opr_modes_to_add: int, normalize: bool = True) → Tensor[source]

Add additional OPR modes to the probe.

Parameters:

  • probe (Tensor) – A (n_opr_modes, n_modes, h, w) tensor of the probe.

  • n_opr_modes_to_add (int) – The number of OPR modes to add.

  • normalize (bool, optional) – Whether to normalize the OPR modes using mnorm so that the power of each mode is the number of pixels in a mode.

Returns:

A (n_opr_modes + n_opr_modes_to_add, n_modes, h, w) tensor of the probe with additional OPR modes.

Return type:

Tensor

ptychi.utils.auto_transfer_to_device(data: Tensor) → Tensor[source]

Automatically determine the device that the data should be placed on, and transfer the data to that device.

The logic of this function is as follows:

  1. If torch.get_default_device() is cuda, transfer the data to cuda.

  2. If torch.get_default_device() is cpu, it could either be GPU is unavailable or intentionally disabled, OR the current code is executed by DataParallel.

    1. If torch.cuda.device_count() is 0, we assume it is the former case, and we keep the data as is.

    2. If torch.cuda.device_count() is not 0, we assume it is the latter case, and we transfer the data to cuda.

ptychi.utils.calculate_data_size_gb(shape: Sequence[int], dtype: dtype | dtype) → float[source]

Calculate the size of the data in GB.

ptychi.utils.chunked_processing(func: Callable, common_kwargs: dict, chunkable_kwargs: dict, iterated_kwargs: dict, replicated_kwargs: dict = None, chunk_size: int = 96)[source]
Parameters:

  • func (callable) – The callable to be executed.

  • common_kwargs (dict) – A dictionary of arguments that should stay constant across chunks.

  • chunkable_kwargs (dict) – A dictionary of arguments that should be chunked.

  • iterated_kwargs (dict) – A dictionary of arguments that should be returned by func, then passed to func for the next chunk. The order of arguments should be the same as the returns of func.

  • replicated_kwargs (dict, optional) – A dictionary of arguments that should be replicated for each chunk along the first dimension to match the chunk size. Tensors given here should have a first dimension of size 1 intended as the batch dimension.

  • chunk_size (int, optional) – The size of each chunk. Default is 96.

Return type:

The returns of func as if it is executed for the entire data.

ptychi.utils.clear_memory(task: PtychographyTask | None = None)[source]

Clear the memory of the device used. If a Task object is provided, it will be deleted and the memory will be released.

Parameters:

task (PtychographyTask, optional) – The Task object to be deleted.

ptychi.utils.generate_gaussian_random_image(shape: tuple[int, ...], loc: float = 0.9, sigma: float = 0.1, smoothing: float = 3.0) → Tensor[source]
ptychi.utils.generate_initial_object(shape: tuple[int, ...], method: Literal['random'] = 'random') → Tensor[source]
ptychi.utils.generate_initial_opr_mode_weights(n_points: int, n_opr_modes: int, eigenmode_weight: float | None = None, probe: Tensor | None = None) → Tensor[source]

Generate initial weights for OPR modes, where the weights of the main OPR mode are set to 1, and the weights of eigenmodes are set to 0.

Parameters:

  • n_points (int) – number of scan points.

  • n_opr_modes (int) – number of OPR modes.

  • eigenmode_weight (float) – initial weight for eigenmodes.

  • probe (Tensor) – The probe. If provided, the weights will be normalized to match the power of the probe.

Returns:

weights – a (n_points, n_opr_modes) tensor of weights.

Return type:

Tensor

ptychi.utils.generate_secondary_probe_modes_hermite(probe: Tensor, m: int, n: int) → Tensor[source]

Generate secondary probe modes using Hermite polynomials.

Parameters:

  • probe (Tensor) – A (h, w) tensor of the primary mode of the probe.

  • m (int) – The orders of the Hermite polynomial.

  • n (int) – The orders of the Hermite polynomial.

Returns:

A ((m + 1) * (n + 1), h, w) tensor of the secondary probe modes.

Return type:

Tensor

ptychi.utils.get_default_complex_dtype()[source]

Get the default complex dtype.

Returns:

The default complex dtype.

Return type:

torch.dtype

ptychi.utils.get_max_batch_size(probe_shape: Sequence[int], object_shape: Sequence[int], double_precision: bool = False, data_saved_on_device: bool = False, all_data_shape: Sequence[int] = None, reconstructor_type: Literal['lsqml'] = 'lsqml', margin_factor: float = 0.2) → int[source]

Estimate the maximum batch size that fits in the available device memory.

We estimate the memory usage using an empirical formula: ` mem = x0 * n_p * batch_size + x1 * n_p + x2 * object_numel ` where n_p = n_modes * probe_size ** 2 and the coefficients x0, x1, x2 were fit from experimental data.

Parameters:

  • probe_shape (Sequence[int]) – The shape of the 4D probe, expected to be (n_opr_modes, n_modes, h, w).

  • object_shape (Sequence[int]) – The shape of the object, expected to be (n_slcies, h, w).

  • double_precision (bool, optional) – Whether to use double precision.

  • data_saved_on_device (bool, optional) – Whether the raw data is kept on device.

  • all_data_shape (Sequence[int], optional) – The shape of the data, expected to be (n_points, h, w).

  • reconstructor_type (Literal["lsqml"], optional) – The type of reconstructor. Currently only lsqml is supported.

  • margin_factor (float, optional) – The fraction of the device memory to be left free.

Returns:

The suggested batch size.

Return type:

int

ptychi.utils.get_probe_renormalization_factor(patterns: Tensor | ndarray) → float[source]

Calculate the renormalization factor that should be applied to the probe to match the maximum power of the diffraction patterns.

Parameters:

patterns (Tensor | ndarray) – A (n, h, w) buffer of diffraction patterns.

Returns:

The renormalization factor.

Return type:

float

ptychi.utils.get_suggested_object_size(positions_px, probe_shape, extra=0)[source]
ptychi.utils.jsonize(val)[source]

Convert a value to a JSON-serializable object.

ptychi.utils.move_nested_tensors_to_device(value, device)[source]

Recursively move tensors contained in lists/tuples/dicts to a device.

ptychi.utils.orthogonalize_initial_probe(probe: Tensor, secondary_mode_energy: float = 0.02, method: Literal['hermite'] = 'hermite') → Tensor[source]

Orthogonalize initial probe.

Parameters:

  • probe (Tensor) – A (n_opr_modes, n_modes, h, w) tensor of the probe. This function only generates incoherent modes; OPR modes are kept as they are. Only the first incoherent mode of the input probe is used. As such, the rest of the incoherent modes can be arbotrarily initialized, but the shape of the input probe should be indicate the number of incoherent modes intended to be generated.

  • secondary_mode_energy (float, optional) – The energy of the secondary mode relative to the principal mode, which is always 1.0.

  • method (Literal["hermite"], optional) – The method to use for orthogonalization.

Returns:

The orthogonalized probe.

Return type:

Tensor

ptychi.utils.rescale_probe(probe: ndarray | Tensor, patterns: ndarray | Tensor, weights: ndarray | Tensor | None = None) → None[source]

Scale probe so that the sum of intensity matches that of the diffraction patterns.

Parameters:

  • probe (Tensor) – A (n_modes, h, w) or (n_opr_modes, n_modes, h, w) tensor of the probe.

  • patterns (Tensor) – A (n, h, w) tensor of diffraction patterns.

  • weights (Tensor, optional) – A (n_points, n_opr_modes) tensor of weights for each OPR mode.

Returns:

scaled_probe – The scaled probe.

Return type:

Tensor

ptychi.utils.set_default_complex_dtype(dtype)[source]

Set the default complex dtype.

Parameters:

dtype (torch.dtype) – The default complex dtype.

ptychi.utils.to_numpy(data: ndarray | Tensor) → ndarray[source]
ptychi.utils.to_tensor(data: ndarray | Tensor, device=None, dtype=None) → Tensor[source]