AD phase reconstruction

Automatic differentiation (AD) based methods of phase reconstruction.

AD phase reconstruction

class PyLorentz.phase.AD_phase.ADPhase(dd: DefocusedDataset, device: str | int, save_dir: PathLike | None = None, name: str | None = None, verbose: bool = 1, scope: Microscope | None = None, sample_params: dict = {}, rng_seed: int | None = None, LRs: dict = {}, scheduler_type: str | None = None, noise_frac: float = 0.01, gaussian_sigma: float = 1)[source]

Bases: BasePhaseReconstruction

ADPhase class for phase reconstruction using defocused datasets and DIPs.

__init__(dd: DefocusedDataset, device: str | int, save_dir: PathLike | None = None, name: str | None = None, verbose: bool = 1, scope: Microscope | None = None, sample_params: dict = {}, rng_seed: int | None = None, LRs: dict = {}, scheduler_type: str | None = None, noise_frac: float = 0.01, gaussian_sigma: float = 1)[source]

Initialize the ADPhase object.

Parameters:

  • dd (DefocusedDataset) – The defocused dataset.

  • device (Union[str, int]) – The device to use (CPU or GPU).

  • save_dir (Optional[os.PathLike], optional) – Directory to save results.

  • name (Optional[str], optional) – Name for the results.

  • verbose (bool, optional) – Verbosity level.

  • scope (Optional[Microscope], optional) – Microscope object.

  • sample_params (dict, optional) – Sample parameters.

  • rng_seed (Optional[int], optional) – Random seed.

  • LRs (dict, optional) – Learning rates for optimization.

  • scheduler_type (Optional[str], optional) – Type of learning rate scheduler.

  • noise_frac (float, optional) – Fraction of noise to add.

  • gaussian_sigma (float, optional) – Sigma value for Gaussian filter.

property recon_phase: ndarray | None

Returns the reconstructed phase after applying Gaussian filter.

Returns:

Reconstructed phase image.

Return type:

Optional[np.ndarray]

property best_phase: ndarray | None

Returns the best phase after applying Gaussian filter.

Returns:

Best phase image.

Return type:

Optional[np.ndarray]

set_best_phase(iter_ind: int = -1) → None[source]

Sets the best phase from the specified iteration index.

Parameters:

iter_ind (int, optional) – Index of the iteration to use for the best phase.

property best_amp: ndarray | None

Returns the best amplitude after applying Gaussian filter.

Returns:

Best amplitude image.

Return type:

Optional[np.ndarray]

property gaussian_sigma: float

Returns the Gaussian sigma value.

Returns:

Gaussian sigma value.

Return type:

float

property recon_amp: ndarray | None

Returns the reconstructed amplitude after applying Gaussian filter.

Returns:

Reconstructed amplitude image.

Return type:

Optional[np.ndarray]

property scope: Microscope | None

Returns the microscope object.

Returns:

Microscope object.

Return type:

Optional[Microscope]

property device: str

Returns the device used for computation.

Returns:

Device for computation.

Return type:

str

property model_input: torch.Tensor | None

Returns the model input tensor.

Returns:

Model input tensor.

Return type:

Optional[torch.Tensor]

property guess_phase: torch.Tensor | None

Returns the guess phase used to pre-train the DIP.

Returns:

Guess phase tensor.

Return type:

Optional[torch.Tensor]

property guess_amp: torch.Tensor | None

Returns the guess amplitude used to pre-train the DIP or if solve_amp is False.

Returns:

Guess amplitude tensor.

Return type:

Optional[torch.Tensor]

property input_DIP: torch.Tensor | None

Returns the noise used as input for one or both DIPs.

Returns:

Input noise tensor.

Return type:

Optional[torch.Tensor]

reconstruct(num_iter: int, model: torch.nn.Module | List[torch.nn.Module] | None = None, num_pretrain_iter: int = 0, solve_amp: bool = False, solve_amp_scale: bool = True, guess_amp: float | ndarray | None = None, LRs: dict = {}, scheduler_type: str | None = None, save: bool = False, name: str | None = None, save_dir: PathLike | None = None, noise_frac: float | None = None, guess_phase: str | ndarray | None = 'SITIE', reset: bool = True, print_every: int = -1, verbose: int = 1, store_iters_every: int = -1, qc: any | None = None, **kwargs) → None[source]

Performs the reconstruction process.

Parameters:

  • num_iter (int) – Number of iterations for reconstruction.

  • model (Optional[Union[nn.Module, List[nn.Module]]], optional) – Model or list of models for DIP.

  • num_pretrain_iter (int, optional) – Number of pretraining iterations.

  • solve_amp (bool, optional) – Whether to solve for amplitude.

  • solve_amp_scale (bool, optional) – Whether to solve amplitude scale.

  • guess_amp (Optional[Union[float, np.ndarray]], optional) – Guess amplitude.

  • LRs (dict, optional) – Learning rates for optimization.

  • scheduler_type (Optional[str], optional) – Type of learning rate scheduler.

  • save (bool, optional) – Whether to save results.

  • name (Optional[str], optional) – Name for the saved results.

  • save_dir (Optional[os.PathLike], optional) – Directory to save results.

  • noise_frac (Optional[float], optional) – Fraction of noise to add.

  • guess_phase (Union[str, np.ndarray, None], optional) – Guess phase or method to obtain it.

  • reset (bool, optional) – Whether to reset the model.

  • print_every (int, optional) – Frequency of printing progress.

  • verbose (int, optional) – Verbosity level.

  • store_iters_every (int, optional) – Frequency of storing iterations.

  • qc (Optional[any], optional) – Quality control object.

  • **kwargs – Additional keyword arguments for scheduler parameters.

show_final(crop: int = 5, **kwargs) → None[source]

Show the phase and induction of the final iteration.

Parameters:

  • crop (int, optional) – Amount to crop off of induction maps before displaying; often

  • 5. (necessary in order to avoid edge artifacts. Defaults to) –

visualize(crop=5)[source]

Plot the best reconstructed phase and induction maps.

Parameters:

  • crop (int, optional) – Amount to crop off of induction maps before displaying; often

  • 5. (necessary in order to avoid edge artifacts. Defaults to) –

save_results(iter_ind: int = None, save_dir: PathLike | None = None, name: str | None = None, overwrite: bool = False)[source]

Save the recontructed phase, Bx, By, and color images. # TODO add saving amplitude and phase_E

Parameters:

  • iter_ind (int, optional) – Index to save. Defaults to None which saves best phase.

  • save_dir (os.PathLik], optional) – Directory to save in. Defaults to None

  • self.save_dir. (which saves in) –

  • name (str, optional) – Name to prepend saved files. Defaults to None.

  • overwrite (bool, optional) – Whether or not to overwrite files. Defaults to False.

Deep image prior model

class PyLorentz.phase.DIP_NN.DIP_NN(*args: Any, **kwargs: Any)[source]

Bases: Module

Autoencoder for reconstructing object wave and amplitude of LTEM images.

Parameters:

  • num_images – int, number of input channels, equal to number of images in TFS.

  • nb_filters – int, number of filters in the first convolutional block.

  • use_dropout – bool, whether to use dropout in the inner layers.

  • batch_norm – bool, whether to use batch normalization after each convolutional layer.

  • upsampling_mode – str, “bilinear” or “nearest” upsampling method.

__init__(num_images: int = 1, nb_filters: int = 16, use_dropout: bool = False, batch_norm: bool = False, upsampling_mode: str = 'nearest')[source]
forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass for the DIP network.

Parameters:

x – Input tensor.

Returns:

Output tensor after passing through the network.

class PyLorentz.phase.DIP_NN.conv2dblock(*args: Any, **kwargs: Any)[source]

Bases: Module

A block consisting of convolutional layers with optional batch normalization and dropout.

Parameters:

  • nb_layers – int, number of convolutional layers.

  • input_channels – int, number of input channels.

  • output_channels – int, number of output channels.

  • kernel_size – int, size of the convolutional kernel.

  • stride – int, stride of the convolution.

  • padding – int, padding for the convolution.

  • use_batchnorm – bool, whether to use batch normalization.

  • lrelu_a – float, negative slope for the Leaky ReLU activation.

  • dropout – float, dropout rate.

  • last_sigmoid – bool, whether to use a sigmoid activation on the last layer.

  • last_tanh – bool, whether to use a tanh activation on the last layer.

  • last_skipReLU – bool, whether to skip ReLU activation on the last layer.

__init__(nb_layers: int, input_channels: int, output_channels: int, kernel_size: int = 3, stride: int = 1, padding: int = 1, use_batchnorm: bool = False, lrelu_a: float = 0.01, dropout_: float = 0, last_sigmoid: bool = False, last_tanh: bool = False, last_skipReLU: bool = False)[source]
forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass for the conv2dblock.

Parameters:

x – Input tensor.

Returns:

Output tensor after passing through the block.

class PyLorentz.phase.DIP_NN.upsample_block(*args: Any, **kwargs: Any)[source]

Bases: Module

Upsampling block using interpolation followed by a convolution.

Parameters:

  • input_channels – int, number of input channels.

  • output_channels – int, number of output channels.

  • scale_factor – int, factor by which to scale the input.

  • mode – str, interpolation mode, either “bilinear” or “nearest”.

__init__(input_channels: int, output_channels: int, scale_factor: int = 2, mode: str = 'bilinear')[source]
forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass for the upsample_block.

Parameters:

x – Input tensor.

Returns:

Output tensor after upsampling and convolution.

PyLorentz.phase.DIP_NN.rng_seed(seed: int) → None[source]

Set the random seed for reproducibility.

Parameters:

seed – The seed value to use.

Returns:

None

PyLorentz.phase.DIP_NN.weight_reset(m: torch.nn.Module) → None[source]

Reset the weights of a given module.

Parameters:

m – The neural network module to reset.

Returns:

None