careamics 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

careamics/config/__init__.py

@@ -1,7 +1,8 @@
 """Configuration module."""
 
 __all__ = [
-    "AlgorithmConfig",
+    "FCNAlgorithmConfig",
+    "VAEAlgorithmConfig",
     "DataConfig",
     "Configuration",
     "CheckpointModel",
@@ -15,9 +16,9 @@ __all__ = [
     "register_model",
     "CustomModel",
     "clear_custom_models",
+    "GaussianMixtureNMConfig",
+    "MultiChannelNMConfig",
 ]
-
-from .algorithm_model import AlgorithmConfig
 from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
 from .configuration_factory import (
@@ -31,5 +32,8 @@ from .configuration_model import (
     save_configuration,
 )
 from .data_model import DataConfig
+from .fcn_algorithm_model import FCNAlgorithmConfig
 from .inference_model import InferenceConfig
+from .nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from .training_model import TrainingConfig
+from .vae_algorithm_model import VAEAlgorithmConfig

careamics/config/architectures/__init__.py

@@ -4,7 +4,7 @@ __all__ = [
     "ArchitectureModel",
     "CustomModel",
     "UNetModel",
-    "VAEModel",
+    "LVAEModel",
     "clear_custom_models",
     "get_custom_model",
     "register_model",
@@ -12,6 +12,6 @@ __all__ = [
 
 from .architecture_model import ArchitectureModel
 from .custom_model import CustomModel
+from .lvae_model import LVAEModel
 from .register_model import clear_custom_models, get_custom_model, register_model
 from .unet_model import UNetModel
-from .vae_model import VAEModel

careamics/config/architectures/architecture_model.py

@@ -27,7 +27,7 @@ class ArchitectureModel(BaseModel):
         Returns
         -------
         dict[str, Any]
-            Model as a dictionnary.
+            Model as a dictionary.
         """
         model_dict = super().model_dump(**kwargs)
 

careamics/config/architectures/custom_model.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 from pprint import pformat
 from typing import Any, Literal
 
@@ -23,12 +24,13 @@ class CustomModel(ArchitectureModel):
 
     Attributes
     ----------
-    architecture : Literal["Custom"]
-        Discriminator for the custom model, must be set to "Custom".
+    architecture : Literal["custom"]
+        Discriminator for the custom model, must be set to "custom".
     name : str
         Name of the custom model.
     parameters : CustomParametersModel
-        Parameters of the custom model.
+        All parameters, required for the initialization of the torch module have to be
+        passed here.
 
     Raises
     ------
@@ -57,7 +59,7 @@ class CustomModel(ArchitectureModel):
     ...
     >>> # Create a configuration
     >>> config_dict = {
-    ...     "architecture": "Custom",
+    ...     "architecture": "custom",
     ...     "name": "my_linear",
     ...     "in_features": 10,
     ...     "out_features": 5,
@@ -71,10 +73,9 @@ class CustomModel(ArchitectureModel):
     )
 
     # discriminator used for choosing the pydantic model in Model
-    architecture: Literal["Custom"]
+    architecture: Literal["custom"]
     """Name of the architecture."""
 
-    # name of the custom model
     name: str
     """Name of the custom model."""
 
@@ -120,10 +121,12 @@ class CustomModel(ArchitectureModel):
             get_custom_model(self.name)(**self.model_dump())
         except Exception as e:
             raise ValueError(
-                f"error while passing parameters to the model {e}. Verify that all "
+                f"while passing parameters to the model {e}. Verify that all "
                 f"mandatory parameters are provided, and that either the {e} accepts "
                 f"*args and **kwargs in its __init__() method, or that no additional"
-                f"parameter is provided."
+                f"parameter is provided. Trace: "
+                f"filename: {inspect.trace()[-1].filename}, function: "
+                f"{inspect.trace()[-1].function}, line: {inspect.trace()[-1].lineno}"
             ) from None
 
         return self

careamics/config/architectures/lvae_model.py

@@ -0,0 +1,170 @@
+"""LVAE Pydantic model."""
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, field_validator, model_validator
+from typing_extensions import Self
+
+from .architecture_model import ArchitectureModel
+
+
+# TODO: it is quite confusing to call this LVAEModel, as it is basically a config
+class LVAEModel(ArchitectureModel):
+    """LVAE model."""
+
+    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
+    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"] = None
+
+    analytical_kl: bool = Field(
+        default=False,
+    )
+
+    @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(cls, 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)})."
+        #         )
+
+        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.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")