careamics 0.0.4.2__py3-none-any.whl → 0.0.5__py3-none-any.whl

careamics/config/architectures/lvae_model.py

@@ -15,9 +15,16 @@ class LVAEModel(ArchitectureModel):
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
     architecture: Literal["LVAE"]
-    input_shape: int = Field(default=64, ge=8, le=1024)
-    multiscale_count: int = Field(default=5)  # TODO clarify
-    # 0 - off, len(z_dims) + 1 # TODO can/should be le to z_dims len + 1
+    input_shape: list[int] = Field(default=(64, 64), validate_default=True)
+    """Shape of the input patch (C, Z, Y, X) or (C, 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)
@@ -31,10 +38,90 @@ class LVAEModel(ArchitectureModel):
     )
 
     predict_logvar: Literal[None, "pixelwise"] = None
+    analytical_kl: bool = Field(default=False)
 
-    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})."
+            )
+        return input_shape
 
     @field_validator("encoder_n_filters")
     @classmethod
@@ -124,27 +211,20 @@ class LVAEModel(ArchitectureModel):
         return z_dims
 
     @model_validator(mode="after")
-    def validate_multiscale_count(cls, self: Self) -> Self:
+    def validate_multiscale_count(self: Self) -> Self:
         """
         Validate the multiscale count.
 
-        Parameters
-        ----------
-        self : Self
-            The model.
-
         Returns
         -------
         Self
             The validated model.
         """
-        # if self.multiscale_count != 0:
-        #     if self.multiscale_count != len(self.z_dims) - 1:
-        #         raise ValueError(
-        #             f"Multiscale count must be 0 or equal to the number of Z "
-        #             f"dims - 1 (got {self.multiscale_count} and {len(self.z_dims)})."
-        #         )
-
+        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:
@@ -156,7 +236,10 @@ class LVAEModel(ArchitectureModel):
         is_3D : bool
             Whether the algorithm is 3D or not.
         """
-        raise NotImplementedError("VAE is not implemented yet.")
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
 
     def is_3D(self) -> bool:
         """
@@ -167,4 +250,4 @@ class LVAEModel(ArchitectureModel):
         bool
             Whether the model is 3D or not.
         """
-        raise NotImplementedError("VAE is not implemented yet.")
+        return self.conv_dims == 3

careamics/config/configuration_factory.py

@@ -234,8 +234,8 @@ def _create_supervised_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = 1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -267,10 +267,10 @@ def _create_supervised_configuration(
         Whether to train all channels independently, by default False.
     loss : Literal["mae", "mse"], optional
         Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default 1.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_params : dict, optional
@@ -282,19 +282,29 @@ def _create_supervised_configuration(
     -------
     Configuration
         Configuration for training CARE or Noise2Noise.
+
+    Raises
+    ------
+    ValueError
+        If the number of channels is not specified when using channels.
+    ValueError
+        If the number of channels is specified but "C" is not in the axes.
     """
     # if there are channels, we need to specify their number
-    if "C" in axes and n_channels_in == 1:
-        raise ValueError(
-            f"Number of channels in must be specified when using channels "
-            f"(got {n_channels_in} channel)."
-        )
-    elif "C" not in axes and n_channels_in > 1:
+    if "C" in axes and n_channels_in is None:
+        raise ValueError("Number of channels in must be specified when using channels ")
+    elif "C" not in axes and (n_channels_in is not None and n_channels_in > 1):
         raise ValueError(
             f"C is not present in the axes, but number of channels is specified "
             f"(got {n_channels_in} channels)."
         )
 
+    if n_channels_in is None:
+        n_channels_in = 1
+
+    if n_channels_out is None:
+        n_channels_out = n_channels_in
+
     # augmentations
     transform_list = _list_augmentations(augmentations)
 
@@ -327,8 +337,8 @@ def create_care_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = -1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -374,16 +384,16 @@ def create_care_configuration(
         and XYRandomRotate90 (in XY) to the images.
     independent_channels : bool, optional
         Whether to train all channels independently, by default False.
-    loss : Literal["mae", "mse"], optional
-        Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default -1.
-    logger : Literal["wandb", "tensorboard", "none"], optional
-        Logger to use, by default "none".
-    model_params : dict, optional
-        UNetModel parameters, by default None.
+    loss : Literal["mae", "mse"], default="mae"
+        Loss function to use.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
+    logger : Literal["wandb", "tensorboard", "none"], default="none"
+        Logger to use.
+    model_params : dict, default=None
+        UNetModel parameters.
     dataloader_params : dict, optional
         Parameters for the dataloader, see PyTorch notes, by default None.
 
@@ -459,9 +469,6 @@ def create_care_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    if n_channels_out == -1:
-        n_channels_out = n_channels_in
-
     return _create_supervised_configuration(
         algorithm="care",
         experiment_name=experiment_name,
@@ -491,8 +498,8 @@ def create_n2n_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = -1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -540,10 +547,10 @@ def create_n2n_configuration(
         Whether to train all channels independently, by default False.
     loss : Literal["mae", "mse"], optional
         Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default -1.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_params : dict, optional
@@ -623,9 +630,6 @@ def create_n2n_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    if n_channels_out == -1:
-        n_channels_out = n_channels_in
-
     return _create_supervised_configuration(
         algorithm="n2n",
         experiment_name=experiment_name,
@@ -655,7 +659,7 @@ def create_n2v_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     use_n2v2: bool = False,
-    n_channels: int = 1,
+    n_channels: Optional[int] = None,
     roi_size: int = 11,
     masked_pixel_percentage: float = 0.2,
     struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
@@ -727,8 +731,8 @@ def create_n2v_configuration(
         Whether to train all channels together, by default True.
     use_n2v2 : bool, optional
         Whether to use N2V2, by default False.
-    n_channels : int, optional
-        Number of channels (in and out), by default 1.
+    n_channels : int or None, default=None
+        Number of channels (in and out).
     roi_size : int, optional
         N2V pixel manipulation area, by default 11.
     masked_pixel_percentage : float, optional
@@ -837,17 +841,17 @@ def create_n2v_configuration(
     ... )
     """
     # if there are channels, we need to specify their number
-    if "C" in axes and n_channels == 1:
-        raise ValueError(
-            f"Number of channels must be specified when using channels "
-            f"(got {n_channels} channel)."
-        )
-    elif "C" not in axes and n_channels > 1:
+    if "C" in axes and n_channels is None:
+        raise ValueError("Number of channels must be specified when using channels.")
+    elif "C" not in axes and (n_channels is not None and n_channels > 1):
         raise ValueError(
             f"C is not present in the axes, but number of channels is specified "
             f"(got {n_channels} channel)."
         )
 
+    if n_channels is None:
+        n_channels = 1
+
     # augmentations
     transform_list = _list_augmentations(augmentations)
 

careamics/config/configuration_model.py

@@ -565,8 +565,8 @@ def save_configuration(config: Configuration, path: Union[str, Path]) -> Path:
     config : Configuration
         Configuration to save.
     path : str or Path
-        Path to a existing folder in which to save the configuration or to an existing
-        configuration file.
+        Path to a existing folder in which to save the configuration, or to a valid
+        configuration file path (uses a .yml or .yaml extension).
 
     Returns
     -------

careamics/config/likelihood_model.py

@@ -4,7 +4,7 @@ from typing import Literal, Optional, Union
 
 import numpy as np
 import torch
-from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator
+from pydantic import BaseModel, ConfigDict, PlainSerializer, PlainValidator
 from typing_extensions import Annotated
 
 from careamics.models.lvae.noise_models import (
@@ -41,7 +41,12 @@ class GaussianLikelihoodConfig(BaseModel):
 
 
 class NMLikelihoodConfig(BaseModel):
-    """Noise model likelihood configuration."""
+    """Noise model likelihood configuration.
+
+    NOTE: we need to define the data mean and std here because the noise model
+    is trained on not-normalized data. Hence, we need to unnormalize the model
+    output to compute the noise model likelihood.
+    """
 
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
@@ -54,7 +59,3 @@ class NMLikelihoodConfig(BaseModel):
     data_std: Tensor = torch.ones(1)
     """The standard deviation of the data, used to unnormalize data for noise
     model evaluation. Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
-
-    # TODO: serialization/deserialization for this
-    noise_model: Optional[NoiseModel] = Field(default=None, exclude=True)
-    """The noise model instance used to compute the likelihood."""

careamics/config/loss_model.py

@@ -0,0 +1,56 @@
+"""Configuration classes for LVAE losses."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+
+class KLLossConfig(BaseModel):
+    """KL loss configuration."""
+
+    model_config = ConfigDict(validate_assignment=True, validate_default=True)
+
+    loss_type: Literal["kl", "kl_restricted"] = "kl"
+    """Type of KL divergence used as KL loss."""
+    rescaling: Literal["latent_dim", "image_dim"] = "latent_dim"
+    """Rescaling of the KL loss."""
+    aggregation: Literal["sum", "mean"] = "mean"
+    """Aggregation of the KL loss across different layers."""
+    free_bits_coeff: float = 0.0
+    """Free bits coefficient for the KL loss."""
+    annealing: bool = False
+    """Whether to apply KL loss annealing."""
+    start: int = -1
+    """Epoch at which KL loss annealing starts."""
+    annealtime: int = 10
+    """Number of epochs for which KL loss annealing is applied."""
+    current_epoch: int = 0
+    """Current epoch in the training loop."""
+
+
+class LVAELossConfig(BaseModel):
+    """LVAE loss configuration."""
+
+    model_config = ConfigDict(
+        validate_assignment=True, validate_default=True, arbitrary_types_allowed=True
+    )
+
+    loss_type: Literal["musplit", "denoisplit", "denoisplit_musplit"]
+    """Type of loss to use for LVAE."""
+
+    reconstruction_weight: float = 1.0
+    """Weight for the reconstruction loss in the total net loss
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    kl_weight: float = 1.0
+    """Weight for the KL loss in the total net loss.
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    musplit_weight: float = 0.1
+    """Weight for the muSplit loss (used in the muSplit-denoiSplit loss)."""
+    denoisplit_weight: float = 0.9
+    """Weight for the denoiSplit loss (used in the muSplit-deonoiSplit loss)."""
+    kl_params: KLLossConfig = KLLossConfig()
+    """KL loss configuration."""
+
+    # TODO: remove?
+    non_stochastic: bool = False
+    """Whether to sample latents and compute KL."""

careamics/config/nm_model.py

@@ -11,9 +11,8 @@ from pydantic import (
     Field,
     PlainSerializer,
     PlainValidator,
-    model_validator,
 )
-from typing_extensions import Annotated, Self
+from typing_extensions import Annotated
 
 from careamics.utils.serializers import _array_to_json, _to_numpy
 
@@ -90,28 +89,29 @@ class GaussianMixtureNMConfig(BaseModel):
     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
+    # @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
+    # TODO revisit validation
 
 
 # The noise model is given by a set of GMMs, one for each target

careamics/config/vae_algorithm_model.py

@@ -12,6 +12,7 @@ from careamics.config.support import SupportedAlgorithm, SupportedLoss
 
 from .architectures import CustomModel, LVAEModel
 from .likelihood_model import GaussianLikelihoodConfig, NMLikelihoodConfig
+from .loss_model import LVAELossConfig
 from .nm_model import MultiChannelNMConfig
 from .optimizer_models import LrSchedulerModel, OptimizerModel
 
@@ -38,13 +39,13 @@ class VAEAlgorithmConfig(BaseModel):
     # TODO: Use supported Enum classes for typing?
     #   - values can still be passed as strings and they will be cast to Enum
     algorithm: Literal["musplit", "denoisplit"]
-    loss: Literal["musplit", "denoisplit", "denoisplit_musplit"]
-    model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
 
-    # TODO: these are configs, change naming of attrs
+    # NOTE: these are all configs (pydantic models)
+    loss: LVAELossConfig
+    model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
     noise_model: Optional[MultiChannelNMConfig] = None
-    noise_model_likelihood_model: Optional[NMLikelihoodConfig] = None
-    gaussian_likelihood_model: Optional[GaussianLikelihoodConfig] = None
+    noise_model_likelihood: Optional[NMLikelihoodConfig] = None
+    gaussian_likelihood: Optional[GaussianLikelihoodConfig] = None
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
@@ -63,13 +64,13 @@ class VAEAlgorithmConfig(BaseModel):
         """
         # musplit
         if self.algorithm == SupportedAlgorithm.MUSPLIT:
-            if self.loss != SupportedLoss.MUSPLIT:
+            if self.loss.loss_type != SupportedLoss.MUSPLIT:
                 raise ValueError(
                     f"Algorithm {self.algorithm} only supports loss `musplit`."
                 )
 
         if self.algorithm == SupportedAlgorithm.DENOISPLIT:
-            if self.loss not in [
+            if self.loss.loss_type not in [
                 SupportedLoss.DENOISPLIT,
                 SupportedLoss.DENOISPLIT_MUSPLIT,
             ]:
@@ -78,16 +79,17 @@ class VAEAlgorithmConfig(BaseModel):
                     "or `denoisplit_musplit."
                 )
             if (
-                self.loss == SupportedLoss.DENOISPLIT
+                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.noise_model is None:
                 raise ValueError("Algorithm `denoisplit` requires a noise model.")
-        # TODO: what if algorithm is not musplit or denoisplit (HDN?)
+        # TODO: what if algorithm is not musplit or denoisplit
         return self
 
     @model_validator(mode="after")
@@ -115,14 +117,13 @@ class VAEAlgorithmConfig(BaseModel):
         Self
             The validated model.
         """
-        if self.gaussian_likelihood_model is not None:
+        if self.gaussian_likelihood is not None:
             assert (
-                self.model.predict_logvar
-                == self.gaussian_likelihood_model.predict_logvar
+                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_model.predict_logvar}).",
+                f"({self.gaussian_likelihood.predict_logvar}).",
             )
         return self