careamics 0.0.19__py3-none-any.whl

careamics/config/lightning/optimizer_configs.py

@@ -0,0 +1,186 @@
+"""Optimizers and schedulers Pydantic models."""
+
+from __future__ import annotations
+
+from typing import Literal, Self
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
+from torch import optim
+
+from careamics.utils.torch_utils import filter_parameters
+
+from ..support import SupportedOptimizer
+
+
+class OptimizerConfig(BaseModel):
+    """Torch optimizer Pydantic model.
+
+    Only parameters supported by the corresponding torch optimizer will be taken
+    into account. For more details, check:
+    https://pytorch.org/docs/stable/optim.html#algorithms
+
+    Note that mandatory parameters (see the specific Optimizer signature in the
+    link above) must be provided. For example, SGD requires `lr`.
+
+    Attributes
+    ----------
+    name : {"Adam", "SGD"}
+        Name of the optimizer.
+    parameters : dict
+        Parameters of the optimizer (see torch documentation).
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    # Mandatory field
+    name: Literal["Adam", "SGD", "Adamax"] = Field(
+        default="Adam", validate_default=True
+    )
+    """Name of the optimizer, supported optimizers are defined in SupportedOptimizer."""
+
+    # Optional parameters, empty dict default value to allow filtering dictionary
+    parameters: dict = Field(
+        default={},
+        validate_default=True,
+    )
+    """Parameters of the optimizer, see PyTorch documentation for more details."""
+
+    @field_validator("parameters")
+    @classmethod
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> dict:
+        """
+        Validate optimizer parameters.
+
+        This method filters out unknown parameters, given the optimizer name.
+
+        Parameters
+        ----------
+        user_params : dict
+            Parameters passed on to the torch optimizer.
+        values : ValidationInfo
+            Pydantic field validation info, used to get the optimizer name.
+
+        Returns
+        -------
+        dict
+            Filtered optimizer parameters.
+
+        Raises
+        ------
+        ValueError
+            If the optimizer name is not specified.
+        """
+        optimizer_name = values.data["name"]
+
+        # retrieve the corresponding optimizer class
+        optimizer_class = getattr(optim, optimizer_name)
+
+        # filter the user parameters according to the optimizer's signature
+        parameters = filter_parameters(optimizer_class, user_params)
+
+        return parameters
+
+    @model_validator(mode="after")
+    def sgd_lr_parameter(self) -> Self:
+        """
+        Check that SGD optimizer has the mandatory `lr` parameter specified.
+
+        This is specific for PyTorch < 2.2.
+
+        Returns
+        -------
+        Self
+            Validated optimizer.
+
+        Raises
+        ------
+        ValueError
+            If the optimizer is SGD and the lr parameter is not specified.
+        """
+        if self.name == SupportedOptimizer.SGD and "lr" not in self.parameters:
+            raise ValueError(
+                "SGD optimizer requires `lr` parameter, check that it has correctly "
+                "been specified in `parameters`."
+            )
+
+        return self
+
+
+class LrSchedulerConfig(BaseModel):
+    """Torch learning rate scheduler Pydantic model.
+
+    Only parameters supported by the corresponding torch lr scheduler will be taken
+    into account. For more details, check:
+    https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
+
+    Note that mandatory parameters (see the specific LrScheduler signature in the
+    link above) must be provided. For example, StepLR requires `step_size`.
+
+    Attributes
+    ----------
+    name : {"ReduceLROnPlateau", "StepLR"}
+        Name of the learning rate scheduler.
+    parameters : dict
+        Parameters of the learning rate scheduler (see torch documentation).
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    # Mandatory field
+    name: Literal["ReduceLROnPlateau", "StepLR"] = Field(default="ReduceLROnPlateau")
+    """Name of the learning rate scheduler, supported schedulers are defined in
+    SupportedScheduler."""
+
+    # Optional parameters
+    parameters: dict = Field(default={}, validate_default=True)
+    """Parameters of the learning rate scheduler, see PyTorch documentation for more
+    details."""
+
+    @field_validator("parameters")
+    @classmethod
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> dict:
+        """Filter parameters based on the learning rate scheduler's signature.
+
+        Parameters
+        ----------
+        user_params : dict
+            User parameters.
+        values : ValidationInfo
+            Pydantic field validation info, used to get the scheduler name.
+
+        Returns
+        -------
+        dict
+            Filtered scheduler parameters.
+
+        Raises
+        ------
+        ValueError
+            If the scheduler is StepLR and the step_size parameter is not specified.
+        """
+        # retrieve the corresponding scheduler class
+        scheduler_class = getattr(optim.lr_scheduler, values.data["name"])
+
+        # filter the user parameters according to the scheduler's signature
+        parameters = filter_parameters(scheduler_class, user_params)
+
+        if values.data["name"] == "StepLR" and "step_size" not in parameters:
+            raise ValueError(
+                "StepLR scheduler requires `step_size` parameter, check that it has "
+                "correctly been specified in `parameters`."
+            )
+
+        return parameters

careamics/config/lightning/training_config.py

@@ -0,0 +1,70 @@
+"""Training configuration."""
+
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .callbacks.callback_config import CheckpointConfig, EarlyStoppingConfig
+
+
+class TrainingConfig(BaseModel):
+    """
+    Parameters related to the training.
+
+    Mandatory parameters are:
+        - num_epochs: number of epochs, greater than 0.
+        - batch_size: batch size, greater than 0.
+        - augmentation: whether to use data augmentation or not (True or False).
+
+    Attributes
+    ----------
+    num_epochs : int
+        Number of epochs, greater than 0.
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    lightning_trainer_config: dict | None = None
+    """Configuration for the PyTorch Lightning Trainer, following PyTorch Lightning
+    Trainer class"""
+
+    logger: Literal["wandb", "tensorboard"] | None = None
+    """Logger to use during training. If None, no logger will be used. Available
+    loggers are defined in SupportedLogger."""
+
+    # Only basic callbacks
+    checkpoint_callback: CheckpointConfig = CheckpointConfig()
+    """Checkpoint callback configuration, following PyTorch Lightning Checkpoint
+    callback."""
+
+    early_stopping_callback: EarlyStoppingConfig | None = Field(
+        default=None, validate_default=True
+    )
+    """Early stopping callback configuration, following PyTorch Lightning Checkpoint
+    callback."""
+
+    def __str__(self) -> str:
+        """Pretty string reprensenting the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def has_logger(self) -> bool:
+        """Check if the logger is defined.
+
+        Returns
+        -------
+        bool
+            Whether the logger is defined or not.
+        """
+        return self.logger is not None

careamics/config/losses/__init__.py

@@ -0,0 +1,8 @@
+"""Losses Pydantic configurations."""
+
+__all__ = [
+    "KLLossConfig",
+    "LVAELossConfig",
+]
+
+from .loss_config import KLLossConfig, LVAELossConfig

careamics/config/losses/loss_config.py

@@ -0,0 +1,60 @@
+"""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[
+        "hdn", "microsplit", "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 revisit weights for the losses
+    # TODO: remove?
+    non_stochastic: bool = False
+    """Whether to sample latents and compute KL."""
+
+    # TODO what are the correct parameters for HDN ?

careamics/config/ng_configs/__init__.py

@@ -0,0 +1,5 @@
+"""Definitions of configurations for CAREamics, compatible with the NG dataset."""
+
+__all__ = ["N2VConfiguration"]
+
+from .n2v_configuration import N2VConfiguration

careamics/config/ng_configs/n2v_configuration.py

@@ -0,0 +1,64 @@
+"""Configuration for N2V."""
+
+from typing import Self
+
+import numpy as np
+from pydantic import model_validator
+
+from careamics.config.algorithms import N2VAlgorithm
+from careamics.config.data.patching_strategies import RandomPatchingConfig
+
+from .ng_configuration import NGConfiguration
+
+
+class N2VConfiguration(NGConfiguration):
+    """N2V-specific configuration."""
+
+    algorithm_config: N2VAlgorithm
+
+    @model_validator(mode="after")
+    def validate_n2v_mask_pixel_perc(self: Self) -> Self:
+        """
+        Validate that there will always be at least one blind-spot pixel in every patch.
+
+        The probability of creating a blind-spot pixel is a function of the chosen
+        masked pixel percentage and patch size.
+
+        Returns
+        -------
+        Self
+            Validated configuration.
+
+        Raises
+        ------
+        ValueError
+            If the probability of masking a pixel within a patch is less than 1 for the
+            chosen masked pixel percentage and patch size.
+        """
+        if self.data_config.mode == "training":
+            assert isinstance(self.data_config.patching, RandomPatchingConfig)
+
+            mask_pixel_perc = self.algorithm_config.n2v_config.masked_pixel_percentage
+            patch_size = self.data_config.patching.patch_size
+            expected_area_per_pixel = 1 / (mask_pixel_perc / 100)
+
+            n_dims = 3 if self.algorithm_config.model.is_3D() else 2
+            patch_size_lower_bound = int(
+                np.ceil(expected_area_per_pixel ** (1 / n_dims))
+            )
+            required_patch_size = tuple(
+                2 ** int(np.ceil(np.log2(patch_size_lower_bound)))
+                for _ in range(n_dims)
+            )
+            required_mask_pixel_perc = (1 / np.prod(patch_size)) * 100
+
+            if expected_area_per_pixel > np.prod(patch_size):
+                raise ValueError(
+                    "The probability of creating a blind-spot pixel within a patch is "
+                    f"below 1, for a patch size of {patch_size} with a masked pixel "
+                    f"percentage of {mask_pixel_perc}%. Either increase the patch size "
+                    f"to {required_patch_size} or increase the masked pixel percentage "
+                    f"to at least {required_mask_pixel_perc}%."
+                )
+
+        return self

careamics/config/ng_configs/ng_configuration.py

@@ -0,0 +1,256 @@
+"""CAREamics configuration compatible with the NG Dataset."""
+
+from __future__ import annotations
+
+import re
+from pprint import pformat
+from typing import Any, Literal, Self, Union
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+from careamics.config.algorithms import (
+    CAREAlgorithm,
+    N2NAlgorithm,
+    N2VAlgorithm,
+)
+from careamics.config.data import NGDataConfig
+from careamics.config.lightning.training_config import TrainingConfig
+
+ALGORITHMS = Union[
+    CAREAlgorithm,
+    N2NAlgorithm,
+    N2VAlgorithm,
+]
+
+
+class NGConfiguration(BaseModel):
+    """
+    CAREamics configuration.
+
+    The configuration defines all parameters used to build and train a CAREamics model.
+    These parameters are validated to ensure that they are compatible with each other.
+
+    It contains three sub-configurations:
+
+    - AlgorithmModel: configuration for the algorithm training, which includes the
+        architecture, loss function, optimizer, and other hyperparameters.
+    - DataModel: configuration for the dataloader, which includes the type of data,
+        transformations, mean/std and other parameters.
+    - TrainingModel: configuration for the training, which includes the number of
+        epochs or the callbacks.
+
+    Attributes
+    ----------
+    experiment_name : str
+        Name of the experiment, used when saving logs and checkpoints.
+    algorithm : AlgorithmModel
+        Algorithm configuration.
+    data : DataModel
+        Data configuration.
+    training : TrainingModel
+        Training configuration.
+
+    Methods
+    -------
+    set_3D(is_3D: bool, axes: str, patch_size: List[int]) -> None
+        Switch configuration between 2D and 3D.
+    model_dump(
+        exclude_defaults: bool = False, exclude_none: bool = True, **kwargs: Dict
+        ) -> Dict
+        Export configuration to a dictionary.
+
+    Raises
+    ------
+    ValueError
+        Configuration parameter type validation errors.
+    ValueError
+        If the experiment name contains invalid characters or is empty.
+    ValueError
+        If the algorithm is 3D but there is not "Z" in the data axes, or 2D algorithm
+        with "Z" in data axes.
+    ValueError
+        Algorithm, data or training validation errors.
+    """
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        arbitrary_types_allowed=True,
+    )
+
+    # version
+    version: Literal["0.1.0"] = "0.1.0"
+    """CAREamics configuration version."""
+
+    # required parameters
+    experiment_name: str
+    """Name of the experiment, used to name logs and checkpoints."""
+
+    # Sub-configurations
+    algorithm_config: ALGORITHMS = Field(discriminator="algorithm")
+    """Algorithm configuration, holding all parameters required to configure the
+    model."""
+
+    data_config: NGDataConfig
+    """Data configuration, holding all parameters required to configure the training
+    data loader."""
+
+    training_config: TrainingConfig
+    """Training configuration, holding all parameters required to configure the
+    training process."""
+
+    @field_validator("experiment_name")
+    @classmethod
+    def no_symbol(cls, name: str) -> str:
+        """
+        Validate experiment name.
+
+        A valid experiment name is a non-empty string with only contains letters,
+        numbers, underscores, dashes and spaces.
+
+        Parameters
+        ----------
+        name : str
+            Name to validate.
+
+        Returns
+        -------
+        str
+            Validated name.
+
+        Raises
+        ------
+        ValueError
+            If the name is empty or contains invalid characters.
+        """
+        if len(name) == 0 or name.isspace():
+            raise ValueError("Experiment name is empty.")
+
+        # Validate using a regex that it contains only letters, numbers, underscores,
+        # dashes and spaces
+        if not re.match(r"^[a-zA-Z0-9_\- ]*$", name):
+            raise ValueError(
+                f"Experiment name contains invalid characters (got {name}). "
+                f"Only letters, numbers, underscores, dashes and spaces are allowed."
+            )
+
+        return name
+
+    @model_validator(mode="after")
+    def validate_3D(self: Self) -> Self:
+        """
+        Validate algorithm dimensions to match data dimensions.
+
+        Returns
+        -------
+        Self
+            Validated configuration.
+        """
+        if self.data_config.is_3D() != self.algorithm_config.model.is_3D():
+            raise ValueError(
+                f"Mismatch between data ({'3D' if self.data_config.is_3D() else '2D'}) "
+                f"and algorithm ("
+                f"{'3D' if self.algorithm_config.model.is_3D() else '2D'}). Data "
+                f"dimensionality is determined by the axes ({self.data_config.axes}), "
+                f"as well as patch size (if applicable) and data type (if data type "
+                f"is 'czi', which uses 3D when 'T' axis is specified)."
+            )
+
+        return self
+
+    def __str__(self) -> str:
+        """
+        Pretty string reprensenting the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm name.
+
+        Returns
+        -------
+        str
+            Algorithm name.
+        """
+        return self.algorithm_config.get_algorithm_friendly_name()
+
+    def get_algorithm_description(self) -> str:
+        """
+        Return a description of the algorithm.
+
+        This method is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Description of the algorithm.
+        """
+        return self.algorithm_config.get_algorithm_description()
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return self.algorithm_config.get_algorithm_citations()
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return self.algorithm_config.get_algorithm_references()
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return self.algorithm_config.get_algorithm_keywords()
+
+    def model_dump(
+        self,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Override model_dump method in order to set default values.
+
+        As opposed to the parent model_dump method, this method sets exclude none by
+        default.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional arguments to pass to the parent model_dump method.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the model parameters.
+        """
+        if "exclude_none" not in kwargs:
+            kwargs["exclude_none"] = True
+
+        return super().model_dump(**kwargs)

careamics/config/ng_factories/__init__.py

@@ -0,0 +1,9 @@
+"""Convenience functions to create coherent configurations for CAREamics."""
+
+__all__ = [
+    "create_n2v_configuration",
+    "create_ng_data_configuration",
+]
+
+from .data_factory import create_ng_data_configuration
+from .n2v_factory import create_n2v_configuration