# Copyright © 2025 UChicago Argonne, LLC All right reserved
# Full license accessible at https://github.com//AdvancedPhotonSource/pty-chi/blob/main/LICENSE
from typing import Optional, Type, Union
import dataclasses
from dataclasses import field
import ptychi.api.options.base as base
import ptychi.api.options.task as task_options
import ptychi.api.enums as enums
import ptychi.api.options.ad_general as ad_general
import ptychi.forward_models as fm
[docs]
@dataclasses.dataclass
class AutodiffPtychographyReconstructorOptions(ad_general.AutodiffReconstructorOptions):
forward_model_class: Union["enums.ForwardModels", Type["fm.ForwardModel"]] = enums.ForwardModels.PLANAR_PTYCHOGRAPHY
[docs]
def get_reconstructor_type(self) -> enums.Reconstructors:
return enums.Reconstructors.AD_PTYCHO
@dataclasses.dataclass
class DeepImagePriorOptions(base.Options):
enabled: bool = False
"""Whether to use deep image prior."""
model: enums.DIPModels = enums.DIPModels.UNET
"""The model used to represent the reconstruction parameter."""
model_params: dict = dataclasses.field(default_factory=dict)
"""The parameters for the model's constructor."""
net_input_channels: int = 32
"""The number of channels in the input to the network."""
constrain_object_outside_network: bool = False
"""If True, magnitude and phase of the object generated by the model is
further constrained so that the magnitude is between 0 and 1 outside the
network model. Do not enable this option if your model already includes
such constraints.
"""
residual_generation: bool = False
"""If True, data generated by the DIP network is added to the initial guess
of the reconstruction parameter, i.e., the network learns a residue. When
using this feature, the network is usually supposed to start from a state
that maps everything to zeros, so that the overall DIP generation is an
identity function.
"""
@dataclasses.dataclass
class AutodiffPtychographyObjectExperimentalOptions(base.Options):
deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)
[docs]
@dataclasses.dataclass
class AutodiffPtychographyObjectOptions(base.ObjectOptions):
experimental: AutodiffPtychographyObjectExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyObjectExperimentalOptions)
@dataclasses.dataclass
class AutodiffPtychographyProbeExperimentalOptions(base.Options):
deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)
[docs]
@dataclasses.dataclass
class AutodiffPtychographyProbeOptions(base.ProbeOptions):
experimental: AutodiffPtychographyProbeExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyProbeExperimentalOptions)
[docs]
@dataclasses.dataclass
class AutodiffPtychographyProbePositionOptions(base.ProbePositionOptions):
pass
[docs]
@dataclasses.dataclass
class AutodiffPtychographyOPRModeWeightsOptions(base.OPRModeWeightsOptions):
pass
[docs]
@dataclasses.dataclass
class AutodiffPtychographyOptions(task_options.PtychographyTaskOptions):
reconstructor_options: AutodiffPtychographyReconstructorOptions = field(
default_factory=AutodiffPtychographyReconstructorOptions
)
object_options: AutodiffPtychographyObjectOptions = field(
default_factory=AutodiffPtychographyObjectOptions
)
probe_options: AutodiffPtychographyProbeOptions = field(
default_factory=AutodiffPtychographyProbeOptions
)
probe_position_options: AutodiffPtychographyProbePositionOptions = field(
default_factory=AutodiffPtychographyProbePositionOptions
)
opr_mode_weight_options: AutodiffPtychographyOPRModeWeightsOptions = field(
default_factory=AutodiffPtychographyOPRModeWeightsOptions
)