nm_model

class GaussianMixtureNMConfig(BaseModel):
    """Gaussian mixture noise model."""

    model_config = ConfigDict(
        protected_namespaces=(),
        validate_assignment=True,
        arbitrary_types_allowed=True,
        extra="allow",
    )
    # model type
    model_type: Literal["GaussianMixtureNoiseModel"]

    path: Optional[Union[Path, str]] = None
    """Path to the directory where the trained noise model (*.npz) is saved in the
    `train` method."""

    # TODO remove and use as parameters to the NM functions?
    signal: Optional[Union[str, Path, np.ndarray]] = Field(default=None, exclude=True)
    """Path to the file containing signal or respective numpy array."""

    # TODO remove and use as parameters to the NM functions?
    observation: Optional[Union[str, Path, np.ndarray]] = Field(
        default=None, exclude=True
    )
    """Path to the file containing observation or respective numpy array."""

    weight: Optional[Array] = None
    """A [3*n_gaussian, n_coeff] sized array containing the values of the weights
    describing the GMM noise model, with each row corresponding to one
    parameter of each gaussian, namely [mean, standard deviation and weight].
    Specifically, rows are organized as follows:
    - first n_gaussian rows correspond to the means
    - next n_gaussian rows correspond to the weights
    - last n_gaussian rows correspond to the standard deviations
    If `weight=None`, the weight array is initialized using the `min_signal`
    and `max_signal` parameters."""

    n_gaussian: int = Field(default=1, ge=1)
    """Number of gaussians used for the GMM."""

    n_coeff: int = Field(default=2, ge=2)
    """Number of coefficients to describe the functional relationship between gaussian
    parameters and the signal. 2 implies a linear relationship, 3 implies a quadratic
    relationship and so on."""

    min_signal: float = Field(default=0.0, ge=0.0)
    """Minimum signal intensity expected in the image."""

    max_signal: float = Field(default=1.0, ge=0.0)
    """Maximum signal intensity expected in the image."""

    min_sigma: float = Field(default=200.0, ge=0.0)  # TODO took from nb in pn2v
    """Minimum value of `standard deviation` allowed in the GMM.
    All values of `standard deviation` below this are clamped to this value."""

    tol: float = Field(default=1e-10)
    """Tolerance used in the computation of the noise model likelihood."""

    @model_validator(mode="after")
    def validate_path_to_pretrained_vs_training_data(self: Self) -> Self:
        """Validate paths provided in the config.

        Returns
        -------
        Self
            Returns itself.
        """
        if self.path and (self.signal is not None or self.observation is not None):
            raise ValueError(
                "Either only 'path' to pre-trained noise model should be"
                "provided or only signal and observation in form of paths"
                "or numpy arrays."
            )
        if not self.path and (self.signal is None or self.observation is None):
            raise ValueError(
                "Either only 'path' to pre-trained noise model should be"
                "provided or only signal and observation in form of paths"
                "or numpy arrays."
            )
        return self

class MultiChannelNMConfig(BaseModel):
    """Noise Model config aggregating noise models for single output channels."""

    # TODO: check that this model config is OK
    model_config = ConfigDict(
        validate_assignment=True, arbitrary_types_allowed=True, extra="allow"
    )
    noise_models: list[GaussianMixtureNMConfig]
    """List of noise models, one for each target channel."""