# 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
from pydantic import Field as PydanticField
from pydantic.dataclasses import dataclass
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]
@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
@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 = PydanticField(default=32, ge=1)
"""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.
"""
@dataclass
class AutodiffPtychographyObjectExperimentalOptions(base.Options):
deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)
[docs]
@dataclass
class AutodiffPtychographyObjectOptions(base.ObjectOptions):
experimental: AutodiffPtychographyObjectExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyObjectExperimentalOptions)
@dataclass
class AutodiffPtychographyProbeExperimentalOptions(base.Options):
deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)
[docs]
@dataclass
class AutodiffPtychographyProbeOptions(base.ProbeOptions):
experimental: AutodiffPtychographyProbeExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyProbeExperimentalOptions)
[docs]
@dataclass
class AutodiffPtychographyProbePositionOptions(base.ProbePositionOptions):
pass
[docs]
@dataclass
class AutodiffPtychographyOPRModeWeightsOptions(base.OPRModeWeightsOptions):
pass
[docs]
@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
)