careamics 0.0.19__py3-none-any.whl

careamics/config/algorithms/vae_algorithm_config.py

@@ -0,0 +1,178 @@
+"""VAE-based algorithm Pydantic model."""
+
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Literal, Self
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from careamics.config.architectures import LVAEConfig
+from careamics.config.lightning.optimizer_configs import (
+    LrSchedulerConfig,
+    OptimizerConfig,
+)
+from careamics.config.losses.loss_config import LVAELossConfig
+from careamics.config.noise_model.likelihood_config import (
+    GaussianLikelihoodConfig,
+    NMLikelihoodConfig,
+)
+from careamics.config.noise_model.noise_model_config import MultiChannelNMConfig
+from careamics.config.support import SupportedAlgorithm, SupportedLoss
+
+
+class VAEBasedAlgorithm(BaseModel):
+    """VAE-based algorithm configuration.
+
+    # TODO
+
+    Examples
+    --------
+    # TODO add once finalized
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+        extra="allow",
+    )
+
+    # Mandatory fields
+    # defined in SupportedAlgorithm
+    # TODO: Use supported Enum classes for typing?
+    #   - values can still be passed as strings and they will be cast to Enum
+    algorithm: Literal["hdn", "microsplit"]
+
+    # NOTE: these are all configs (pydantic models)
+    loss: LVAELossConfig
+    model: LVAEConfig
+    noise_model: MultiChannelNMConfig | None = None
+    noise_model_likelihood: NMLikelihoodConfig | None = None
+    gaussian_likelihood: GaussianLikelihoodConfig | None = None  # TODO change to str
+
+    mmse_count: int = 1
+    is_supervised: bool = False
+
+    # Optional fields
+    optimizer: OptimizerConfig = OptimizerConfig()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
+    lr_scheduler: LrSchedulerConfig = LrSchedulerConfig()
+
+    @model_validator(mode="after")
+    def algorithm_cross_validation(self: Self) -> Self:
+        """Validate the algorithm model based on `algorithm`.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # hdn
+        # TODO move to designated configurations
+        if self.algorithm == SupportedAlgorithm.HDN:
+            if self.loss.loss_type != SupportedLoss.HDN:
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `hdn`."
+                )
+            if self.model.multiscale_count > 1:
+                raise ValueError("Algorithm `hdn` does not support multiscale models.")
+        # musplit
+        if self.algorithm == SupportedAlgorithm.MICROSPLIT:
+            if self.loss.loss_type not in [
+                SupportedLoss.MUSPLIT,
+                SupportedLoss.DENOISPLIT,
+                SupportedLoss.DENOISPLIT_MUSPLIT,
+            ]:  # TODO Update losses configs, make loss just microsplit
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `microsplit`."
+                )  # TODO Update losses configs
+
+            if (
+                self.loss.loss_type == SupportedLoss.DENOISPLIT
+                and self.model.predict_logvar is not None
+            ):
+                raise ValueError(
+                    "Algorithm `denoisplit` with loss `denoisplit` only supports "
+                    "`predict_logvar` as `None`."
+                )
+            if (
+                self.loss.loss_type == SupportedLoss.DENOISPLIT
+                and self.noise_model is None
+            ):
+                raise ValueError("Algorithm `denoisplit` requires a noise model.")
+        # TODO: what if algorithm is not musplit or denoisplit
+        return self
+
+    @model_validator(mode="after")
+    def output_channels_validation(self: Self) -> Self:
+        """Validate the consistency between number of out channels and noise models.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        if self.noise_model is not None:
+            assert self.model.output_channels == len(self.noise_model.noise_models), (
+                f"Number of output channels ({self.model.output_channels}) must match "
+                f"the number of noise models ({len(self.noise_model.noise_models)})."
+            )
+
+        if self.algorithm == SupportedAlgorithm.HDN:
+            assert self.model.output_channels == 1, (
+                f"Number of output channels ({self.model.output_channels}) must be 1 "
+                "for algorithm `hdn`."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def predict_logvar_validation(self: Self) -> Self:
+        """Validate the consistency of `predict_logvar` throughout the model.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        if self.gaussian_likelihood is not None:
+            assert (
+                self.model.predict_logvar == self.gaussian_likelihood.predict_logvar
+            ), (
+                f"Model `predict_logvar` ({self.model.predict_logvar}) must match "
+                "Gaussian likelihood model `predict_logvar` "
+                f"({self.gaussian_likelihood.predict_logvar}).",
+            )
+        # if self.algorithm == SupportedAlgorithm.HDN:
+        #     assert (
+        #         self.model.predict_logvar is None
+        #     ), "Model `predict_logvar` must be `None` for algorithm `hdn`."
+        #     if self.gaussian_likelihood is not None:
+        #         assert self.gaussian_likelihood.predict_logvar is None, (
+        #             "Gaussian likelihood model `predict_logvar` must be `None` "
+        #             "for algorithm `hdn`."
+        #         )
+        # TODO check this
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    @classmethod
+    def get_compatible_algorithms(cls) -> list[str]:
+        """Get the list of compatible algorithms.
+
+        Returns
+        -------
+        list of str
+            List of compatible algorithms.
+        """
+        return ["hdn", "microsplit"]

careamics/config/architectures/__init__.py

@@ -0,0 +1,7 @@
+"""Deep-learning model configurations."""
+
+__all__ = ["ArchitectureConfig", "LVAEConfig", "UNetConfig"]
+
+from .architecture_config import ArchitectureConfig
+from .lvae_config import LVAEConfig
+from .unet_config import UNetConfig

careamics/config/architectures/architecture_config.py

@@ -0,0 +1,37 @@
+"""Base model for the various CAREamics architectures."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class ArchitectureConfig(BaseModel):
+    """
+    Base Pydantic model for all model architectures.
+
+    The `model_dump` method allows removing the `architecture` key from the model.
+    """
+
+    architecture: str
+    """Name of the architecture."""
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        """
+        Dump the model as a dictionary, ignoring the architecture keyword.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments from Pydantic BaseModel model_dump method.
+
+        Returns
+        -------
+        {str: Any}
+            Model as a dictionary.
+        """
+        model_dict = super().model_dump(**kwargs)
+
+        # remove the architecture key
+        model_dict.pop("architecture")
+
+        return model_dict

careamics/config/architectures/lvae_config.py

@@ -0,0 +1,262 @@
+"""LVAE Pydantic model."""
+
+from typing import Literal, Self
+
+from pydantic import ConfigDict, Field, field_validator, model_validator
+
+from .architecture_config import ArchitectureConfig
+
+
+# TODO: it is quite confusing to call this LVAEModel, as it is basically a config
+class LVAEConfig(ArchitectureConfig):
+    """LVAE model."""
+
+    model_config = ConfigDict(validate_assignment=True, validate_default=True)
+
+    architecture: Literal["LVAE"]
+
+    input_shape: tuple[int, ...] = Field(default=(64, 64), validate_default=True)
+    """Shape of the input patch (Z, Y, X) or (Y, X) if the data is 2D."""
+    encoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
+
+    # TODO make this per hierarchy step ?
+    decoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
+    multiscale_count: int = Field(default=1)
+    # TODO there should be a check for multiscale_count in dataset !!
+
+    # 1 - off, len(z_dims) + 1 # TODO Consider starting from 0
+    z_dims: list = Field(default=[128, 128, 128, 128])
+    output_channels: int = Field(default=1, ge=1)
+    encoder_n_filters: int = Field(default=64, ge=8, le=1024)
+    decoder_n_filters: int = Field(default=64, ge=8, le=1024)
+    encoder_dropout: float = Field(default=0.1, ge=0.0, le=0.9)
+    decoder_dropout: float = Field(default=0.1, ge=0.0, le=0.9)
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ] = Field(
+        default="ELU",
+    )
+
+    predict_logvar: Literal[None, "pixelwise"] = "pixelwise"
+    analytical_kl: bool = Field(default=False)
+
+    @model_validator(mode="after")
+    def validate_conv_strides(self: Self) -> Self:
+        """
+        Validate the convolutional strides.
+
+        Returns
+        -------
+        list
+            Validated strides.
+
+        Raises
+        ------
+        ValueError
+            If the number of strides is not 2.
+        """
+        if len(self.encoder_conv_strides) < 2 or len(self.encoder_conv_strides) > 3:
+            raise ValueError(
+                f"Strides must be 2 or 3 (got {len(self.encoder_conv_strides)})."
+            )
+
+        if len(self.decoder_conv_strides) < 2 or len(self.decoder_conv_strides) > 3:
+            raise ValueError(
+                f"Strides must be 2 or 3 (got {len(self.decoder_conv_strides)})."
+            )
+
+        # adding 1 to encoder strides for the number of input channels
+        if len(self.input_shape) != len(self.encoder_conv_strides):
+            raise ValueError(
+                f"Input dimensions must be equal to the number of encoder conv strides"
+                f" (got {len(self.input_shape)} and {len(self.encoder_conv_strides)})."
+            )
+
+        if len(self.encoder_conv_strides) < len(self.decoder_conv_strides):
+            raise ValueError(
+                f"Decoder can't be 3D when encoder is 2D (got"
+                f" {len(self.encoder_conv_strides)} and"
+                f"{len(self.decoder_conv_strides)})."
+            )
+
+        if any(s < 1 for s in self.encoder_conv_strides) or any(
+            s < 1 for s in self.decoder_conv_strides
+        ):
+            raise ValueError(
+                f"All strides must be greater or equal to 1"
+                f"(got {self.encoder_conv_strides} and {self.decoder_conv_strides})."
+            )
+        # TODO: validate max stride size ?
+        return self
+
+    @field_validator("input_shape")
+    @classmethod
+    def validate_input_shape(cls, input_shape: list) -> list:
+        """
+        Validate the input shape.
+
+        Parameters
+        ----------
+        input_shape : list
+            Shape of the input patch.
+
+        Returns
+        -------
+        list
+            Validated input shape.
+
+        Raises
+        ------
+        ValueError
+            If the number of dimensions is not 3 or 4.
+        """
+        if len(input_shape) < 2 or len(input_shape) > 3:
+            raise ValueError(
+                f"Number of input dimensions must be 2 for 2D data 3 for 3D"
+                f"(got {len(input_shape)})."
+            )
+
+        if any(s < 1 for s in input_shape):
+            raise ValueError(
+                f"Input shape must be greater than 1 in all dimensions"
+                f"(got {input_shape})."
+            )
+
+        if any(s < 64 for s in input_shape[-2:]):
+            raise ValueError(
+                f"Input shape must be greater or equal to 64 in XY dimensions"
+                f"(got {input_shape})."
+            )
+
+        return input_shape
+
+    @field_validator("encoder_n_filters")
+    @classmethod
+    def validate_encoder_even(cls, encoder_n_filters: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        encoder_n_filters : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if encoder_n_filters % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {encoder_n_filters})."
+            )
+
+        return encoder_n_filters
+
+    @field_validator("decoder_n_filters")
+    @classmethod
+    def validate_decoder_even(cls, decoder_n_filters: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        decoder_n_filters : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if decoder_n_filters % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {decoder_n_filters})."
+            )
+
+        return decoder_n_filters
+
+    @field_validator("z_dims")
+    def validate_z_dims(cls, z_dims: tuple) -> tuple:
+        """
+        Validate the z_dims.
+
+        Parameters
+        ----------
+        z_dims : tuple
+            Tuple of z dimensions.
+
+        Returns
+        -------
+        tuple
+            Validated z dimensions.
+
+        Raises
+        ------
+        ValueError
+            If the number of z dimensions is not 4.
+        """
+        if len(z_dims) < 2:
+            raise ValueError(
+                f"Number of z dimensions must be at least 2 (got {len(z_dims)})."
+            )
+
+        return z_dims
+
+    @model_validator(mode="after")
+    def validate_multiscale_count(self: Self) -> Self:
+        """
+        Validate the multiscale count.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        if self.multiscale_count < 1 or self.multiscale_count > len(self.z_dims) + 1:
+            raise ValueError(
+                f"Multiscale count must be 1 for LC off or less or equal to the number"
+                f" of Z dims + 1 (got {self.multiscale_count} and {len(self.z_dims)})."
+            )
+        return self
+
+    def set_3D(self, is_3D: bool) -> None:
+        """
+        Set 3D model by setting the `conv_dims` parameters.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        """
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        return len(self.input_shape) == 3

careamics/config/architectures/unet_config.py

@@ -0,0 +1,125 @@
+"""UNet Pydantic model."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, field_validator
+
+from .architecture_config import ArchitectureConfig
+
+
+# TODO tests activation <-> pydantic model, test the literals!
+# TODO annotations for the json schema?
+class UNetConfig(ArchitectureConfig):
+    """
+    Pydantic model for a N2V(2)-compatible UNet.
+
+    Attributes
+    ----------
+    depth : int
+        Depth of the model, between 1 and 10 (default 2).
+    num_channels_init : int
+        Number of filters of the first level of the network, should be even
+        and minimum 8 (default 96).
+    """
+
+    # pydantic model config
+    model_config = ConfigDict(validate_assignment=True)
+
+    # discriminator used for choosing the pydantic model in Model
+    architecture: Literal["UNet"]
+    """Name of the architecture."""
+
+    # parameters
+    # validate_defaults allow ignoring default values in the dump if they were not set
+    conv_dims: Literal[2, 3] = Field(default=2, validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
+    num_classes: int = Field(default=1, ge=1, validate_default=True)
+    """Number of classes or channels in the model output."""
+
+    in_channels: int = Field(default=1, ge=1, validate_default=True)
+    """Number of channels in the input to the model."""
+
+    depth: int = Field(default=2, ge=1, le=10, validate_default=True)
+    """Number of levels in the UNet."""
+
+    num_channels_init: int = Field(default=32, ge=8, le=1024, validate_default=True)
+    """Number of convolutional filters in the first layer of the UNet."""
+
+    # TODO we are not using this, so why make it a choice?
+    final_activation: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
+    ] = Field(default="None", validate_default=True)
+    """Final activation function."""
+
+    n2v2: bool = Field(default=False, validate_default=True)
+    """Whether to use N2V2 architecture modifications, with blur pool layers and fewer
+    skip connections.
+    """
+
+    independent_channels: bool = Field(default=True, validate_default=True)
+    """Whether information is processed independently in each channel, used to train
+    channels independently."""
+
+    use_batch_norm: bool = Field(default=True, validate_default=True)
+    """Whether to use batch normalization in the model."""
+
+    @field_validator("num_channels_init")
+    @classmethod
+    def validate_num_channels_init(cls, num_channels_init: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        num_channels_init : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if num_channels_init % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {num_channels_init})."
+            )
+
+        return num_channels_init
+
+    def set_3D(self, is_3D: bool) -> None:
+        """
+        Set 3D model by setting the `conv_dims` parameters.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        """
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        This method is used in the NG configuration validation to check that the model
+        dimensions match the data dimensions.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        return self.conv_dims == 3