careamics 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

careamics/config/configuration_factories.py

@@ -5,9 +5,20 @@ from typing import Annotated, Any, Literal, Union
 
 from pydantic import Field, TypeAdapter
 
-from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
-from careamics.config.architectures import UNetModel
+from careamics.config.algorithms import (
+    CAREAlgorithm,
+    MicroSplitAlgorithm,
+    N2NAlgorithm,
+    N2VAlgorithm,
+)
+from careamics.config.architectures import LVAEModel, UNetModel
 from careamics.config.data import DataConfig, NGDataConfig
+from careamics.config.likelihood_model import (
+    GaussianLikelihoodConfig,
+    NMLikelihoodConfig,
+)
+from careamics.config.loss_model import LVAELossConfig
+from careamics.config.nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from careamics.config.support import (
     SupportedArchitecture,
     SupportedPixelManipulation,
@@ -20,6 +31,7 @@ from careamics.config.transformations import (
     XYFlipModel,
     XYRandomRotate90Model,
 )
+from careamics.lvae_training.dataset.config import MicroSplitDataConfig
 
 from .configuration import Configuration
 
@@ -224,7 +236,7 @@ def _create_algorithm_configuration(
 def _create_data_configuration(
     data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
-    patch_size: list[int],
+    patch_size: Sequence[int],
     batch_size: int,
     augmentations: Union[list[SPATIAL_TRANSFORMS_UNION]],
     train_dataloader_params: dict[str, Any] | None = None,
@@ -277,6 +289,70 @@ def _create_data_configuration(
     return DataConfig(**data)
 
 
+def _create_microsplit_data_configuration(
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: Sequence[int],
+    grid_size: int,
+    multiscale_count: int,
+    batch_size: int,
+    augmentations: Union[list[SPATIAL_TRANSFORMS_UNION]],
+    train_dataloader_params: dict[str, Any] | None = None,
+    val_dataloader_params: dict[str, Any] | None = None,
+) -> DataConfig:
+    """
+    Create a dictionary with the parameters of the data model.
+
+    Parameters
+    ----------
+    data_type : {"array", "tiff", "czi", "custom"}
+        Type of the data.
+    axes : str
+        Axes of the data.
+    patch_size : list of int
+        Size of the patches along the spatial dimensions.
+    grid_size : int
+        Grid size for patch extraction.
+    multiscale_count : int
+        Number of LC scales.
+    batch_size : int
+        Batch size.
+    augmentations : list of transforms
+        List of transforms to apply.
+    train_dataloader_params : dict
+        Parameters for the training dataloader, see PyTorch notes, by default None.
+    val_dataloader_params : dict
+        Parameters for the validation dataloader, see PyTorch notes, by default None.
+
+    Returns
+    -------
+    DataConfig
+        Data model with the specified parameters.
+    """
+    # data model
+    data = {
+        "data_type": data_type,
+        "axes": axes,
+        "image_size": patch_size,
+        "grid_size": grid_size,
+        "multiscale_lowres_count": multiscale_count,
+        "batch_size": batch_size,
+        "transforms": augmentations,
+    }
+    # Don't override defaults set in DataConfig class
+    if train_dataloader_params is not None:
+        # DataConfig enforces the presence of `shuffle` key in the dataloader parameters
+        if "shuffle" not in train_dataloader_params:
+            train_dataloader_params["shuffle"] = True
+
+        data["train_dataloader_params"] = train_dataloader_params
+
+    if val_dataloader_params is not None:
+        data["val_dataloader_params"] = val_dataloader_params
+
+    return MicroSplitDataConfig(**data)
+
+
 def _create_ng_data_configuration(
     data_type: Literal["array", "tiff", "custom"],
     axes: str,
@@ -357,7 +433,7 @@ def _create_ng_data_configuration(
 
 
 def _create_training_configuration(
-    num_epochs: int,
+    trainer_params: dict,
     logger: Literal["wandb", "tensorboard", "none"],
     checkpoint_params: dict[str, Any] | None = None,
 ) -> TrainingConfig:
@@ -366,8 +442,8 @@ def _create_training_configuration(
 
     Parameters
     ----------
-    num_epochs : int
-        Number of epochs.
+    trainer_params : dict
+        Parameters for Lightning Trainer class, see PyTorch Lightning documentation.
     logger : {"wandb", "tensorboard", "none"}
         Logger to use.
     checkpoint_params : dict, default=None
@@ -380,7 +456,7 @@ def _create_training_configuration(
         Training model with the specified parameters.
     """
     return TrainingConfig(
-        num_epochs=num_epochs,
+        lightning_trainer_config=trainer_params,
         logger=None if logger == "none" else logger,
         checkpoint_callback={} if checkpoint_params is None else checkpoint_params,
     )
@@ -392,9 +468,9 @@ def _create_supervised_config_dict(
     experiment_name: str,
     data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
-    patch_size: list[int],
+    patch_size: Sequence[int],
     batch_size: int,
-    num_epochs: int,
+    trainer_params: dict | None = None,
     augmentations: list[SPATIAL_TRANSFORMS_UNION] | None = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
@@ -409,6 +485,8 @@ def _create_supervised_config_dict(
     train_dataloader_params: dict[str, Any] | None = None,
     val_dataloader_params: dict[str, Any] | None = None,
     checkpoint_params: dict[str, Any] | None = None,
+    num_epochs: int | None = None,
+    num_steps: int | None = None,
 ) -> dict:
     """
     Create a configuration for training CARE or Noise2Noise.
@@ -427,8 +505,8 @@ def _create_supervised_config_dict(
         Size of the patches along the spatial dimensions (e.g. [64, 64]).
     batch_size : int
         Batch size.
-    num_epochs : int
-        Number of epochs.
+    trainer_params : dict
+        Parameters for the training configuration.
     augmentations : list of transforms, default=None
         List of transforms to apply, either both or one of XYFlipModel and
         XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
@@ -461,6 +539,13 @@ def _create_supervised_config_dict(
     checkpoint_params : dict, default=None
         Parameters for the checkpoint callback, see PyTorch Lightning documentation
         (`ModelCheckpoint`) for the list of available parameters.
+    num_epochs : int or None, default=None
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int or None, default=None
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
 
     Returns
     -------
@@ -518,9 +603,18 @@ def _create_supervised_config_dict(
         val_dataloader_params=val_dataloader_params,
     )
 
+    # Handle trainer parameters with num_epochs and num_steps
+    final_trainer_params = {} if trainer_params is None else trainer_params.copy()
+
+    # Add num_epochs and num_steps if provided
+    if num_epochs is not None:
+        final_trainer_params["max_epochs"] = num_epochs
+    if num_steps is not None:
+        final_trainer_params["limit_train_batches"] = num_steps
+
     # training
     training_params = _create_training_configuration(
-        num_epochs=num_epochs,
+        trainer_params=final_trainer_params,
         logger=logger,
         checkpoint_params=checkpoint_params,
     )
@@ -537,15 +631,17 @@ def create_care_configuration(
     experiment_name: str,
     data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
-    patch_size: list[int],
+    patch_size: Sequence[int],
     batch_size: int,
-    num_epochs: int,
+    num_epochs: int = 100,
+    num_steps: int | None = None,
     augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
     n_channels_in: int | None = None,
     n_channels_out: int | None = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
+    trainer_params: dict | None = None,
     model_params: dict | None = None,
     optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
     optimizer_params: dict[str, Any] | None = None,
@@ -588,8 +684,13 @@ def create_care_configuration(
         Size of the patches along the spatial dimensions (e.g. [64, 64]).
     batch_size : int
         Batch size.
-    num_epochs : int
-        Number of epochs.
+    num_epochs : int, default=100
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int, optional
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
     augmentations : list of transforms, default=None
         List of transforms to apply, either both or one of XYFlipModel and
         XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
@@ -604,6 +705,8 @@ def create_care_configuration(
         Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], default="none"
         Logger to use.
+    trainer_params : dict, optional
+        Parameters for the trainer class, see PyTorch Lightning documentation.
     model_params : dict, default=None
         UNetModel parameters.
     optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
@@ -644,6 +747,16 @@ def create_care_configuration(
     ...     num_epochs=100
     ... )
 
+    You can also limit the number of batches per epoch:
+    >>> config = create_care_configuration(
+    ...     experiment_name="care_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_steps=100  # limit to 100 batches per epoch
+    ... )
+
     To disable transforms, simply set `augmentations` to an empty list:
     >>> config = create_care_configuration(
     ...     experiment_name="care_experiment",
@@ -730,13 +843,13 @@ def create_care_configuration(
             axes=axes,
             patch_size=patch_size,
             batch_size=batch_size,
-            num_epochs=num_epochs,
             augmentations=augmentations,
             independent_channels=independent_channels,
             loss=loss,
             n_channels_in=n_channels_in,
             n_channels_out=n_channels_out,
             logger=logger,
+            trainer_params=trainer_params,
             model_params=model_params,
             optimizer=optimizer,
             optimizer_params=optimizer_params,
@@ -745,6 +858,8 @@ def create_care_configuration(
             train_dataloader_params=train_dataloader_params,
             val_dataloader_params=val_dataloader_params,
             checkpoint_params=checkpoint_params,
+            num_epochs=num_epochs,
+            num_steps=num_steps,
         )
     )
 
@@ -753,15 +868,17 @@ def create_n2n_configuration(
     experiment_name: str,
     data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
-    patch_size: list[int],
+    patch_size: Sequence[int],
     batch_size: int,
-    num_epochs: int,
+    num_epochs: int = 100,
+    num_steps: int | None = None,
     augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
     n_channels_in: int | None = None,
     n_channels_out: int | None = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
+    trainer_params: dict | None = None,
     model_params: dict | None = None,
     optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
     optimizer_params: dict[str, Any] | None = None,
@@ -804,8 +921,13 @@ def create_n2n_configuration(
         Size of the patches along the spatial dimensions (e.g. [64, 64]).
     batch_size : int
         Batch size.
-    num_epochs : int
-        Number of epochs.
+    num_epochs : int, default=100
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int, optional
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
     augmentations : list of transforms, default=None
         List of transforms to apply, either both or one of XYFlipModel and
         XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
@@ -820,6 +942,8 @@ def create_n2n_configuration(
         Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
+    trainer_params : dict, optional
+        Parameters for the trainer class, see PyTorch Lightning documentation.
     model_params : dict, default=None
         UNetModel parameters.
     optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
@@ -860,6 +984,16 @@ def create_n2n_configuration(
     ...     num_epochs=100
     ... )
 
+    You can also limit the number of batches per epoch:
+    >>> config = create_n2n_configuration(
+    ...     experiment_name="n2n_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_steps=100  # limit to 100 batches per epoch
+    ... )
+
     To disable transforms, simply set `augmentations` to an empty list:
     >>> config = create_n2n_configuration(
     ...     experiment_name="n2n_experiment",
@@ -871,8 +1005,7 @@ def create_n2n_configuration(
     ...     augmentations=[]
     ... )
 
-    A list of transforms can be passed to the `augmentations` parameter to replace the
-    default augmentations:
+    A list of transforms can be passed to the `augmentations` parameter:
     >>> from careamics.config.transformations import XYFlipModel
     >>> config = create_n2n_configuration(
     ...     experiment_name="n2n_experiment",
@@ -946,7 +1079,7 @@ def create_n2n_configuration(
             axes=axes,
             patch_size=patch_size,
             batch_size=batch_size,
-            num_epochs=num_epochs,
+            trainer_params=trainer_params,
             augmentations=augmentations,
             independent_channels=independent_channels,
             loss=loss,
@@ -961,6 +1094,8 @@ def create_n2n_configuration(
             train_dataloader_params=train_dataloader_params,
             val_dataloader_params=val_dataloader_params,
             checkpoint_params=checkpoint_params,
+            num_epochs=num_epochs,
+            num_steps=num_steps,
         )
     )
 
@@ -969,9 +1104,10 @@ def create_n2v_configuration(
     experiment_name: str,
     data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
-    patch_size: list[int],
+    patch_size: Sequence[int],
     batch_size: int,
-    num_epochs: int,
+    num_epochs: int = 100,
+    num_steps: int | None = None,
     augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
     independent_channels: bool = True,
     use_n2v2: bool = False,
@@ -980,6 +1116,7 @@ def create_n2v_configuration(
     masked_pixel_percentage: float = 0.2,
     struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
     struct_n2v_span: int = 5,
+    trainer_params: dict | None = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: dict | None = None,
     optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
@@ -1043,8 +1180,13 @@ def create_n2v_configuration(
         Size of the patches along the spatial dimensions (e.g. [64, 64]).
     batch_size : int
         Batch size.
-    num_epochs : int
-        Number of epochs.
+    num_epochs : int, default=100
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int, optional
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
     augmentations : list of transforms, default=None
         List of transforms to apply, either both or one of XYFlipModel and
         XYRandomRotate90Model. By default, it applies both XYFlip (on X and Y)
@@ -1063,6 +1205,8 @@ def create_n2v_configuration(
         Axis along which to apply structN2V mask, by default "none".
     struct_n2v_span : int, optional
         Span of the structN2V mask, by default 5.
+    trainer_params : dict, optional
+        Parameters for the trainer, see the relevant documentation.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_params : dict, default=None
@@ -1105,6 +1249,16 @@ def create_n2v_configuration(
     ...     num_epochs=100
     ... )
 
+    You can also limit the number of batches per epoch:
+    >>> config = create_n2v_configuration(
+    ...     experiment_name="n2v_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_steps=100  # limit to 100 batches per epoch
+    ... )
+
     To disable transforms, simply set `augmentations` to an empty list:
     >>> config = create_n2v_configuration(
     ...     experiment_name="n2v_experiment",
@@ -1261,8 +1415,17 @@ def create_n2v_configuration(
     )
 
     # training
+    # Handle trainer parameters with num_epochs and nun_steps
+    final_trainer_params = {} if trainer_params is None else trainer_params.copy()
+
+    # Add num_epochs and nun_steps if provided
+    if num_epochs is not None:
+        final_trainer_params["max_epochs"] = num_epochs
+    if num_steps is not None:
+        final_trainer_params["limit_train_batches"] = num_steps
+
     training_params = _create_training_configuration(
-        num_epochs=num_epochs,
+        trainer_params=final_trainer_params,
         logger=logger,
         checkpoint_params=checkpoint_params,
     )
@@ -1273,3 +1436,664 @@ def create_n2v_configuration(
         data_config=data_params,
         training_config=training_params,
     )
+
+
+def _create_vae_configuration(
+    input_shape: Sequence[int],
+    encoder_conv_strides: tuple[int, ...],
+    decoder_conv_strides: tuple[int, ...],
+    multiscale_count: int,
+    z_dims: tuple[int, ...],
+    output_channels: int,
+    encoder_n_filters: int,
+    decoder_n_filters: int,
+    encoder_dropout: float,
+    decoder_dropout: float,
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ],
+    predict_logvar: Literal[None, "pixelwise"],
+    analytical_kl: bool,
+) -> LVAEModel:
+    """Create a dictionary with the parameters of the vae based algorithm model.
+
+    Parameters
+    ----------
+    input_shape : tuple[int, ...]
+        Shape of the input patch (Z, Y, X) or (Y, X) if the data is 2D.
+    encoder_conv_strides : tuple[int, ...]
+        Strides of the encoder convolutional layers, length also defines 2D or 3D.
+    decoder_conv_strides : tuple[int, ...]
+        Strides of the decoder convolutional layers, length also defines 2D or 3D.
+    multiscale_count : int
+        Number of lateral context layers, specific to MicroSplit.
+    z_dims : tuple[int, ...]
+        Number of hierarchies in the LVAE model.
+    output_channels : int
+        Number of output channels.
+    encoder_n_filters : int
+        Number of filters in the convolutional layers of the encoder.
+    decoder_n_filters : int
+        Number of filters in the convolutional layers of the decoder.
+    encoder_dropout : float
+        Dropout rate for the encoder.
+    decoder_dropout : float
+        Dropout rate for the decoder.
+    nonlinearity : Literal
+        Type of nonlinearity function to use.
+    predict_logvar : Literal # TODO needs review
+        _description_.
+    analytical_kl : bool # TODO needs clarification
+        _description_.
+
+    Returns
+    -------
+    LVAEModel
+        LVAE model with the specified parameters.
+    """
+    return LVAEModel(
+        architecture=SupportedArchitecture.LVAE.value,
+        input_shape=input_shape,
+        encoder_conv_strides=encoder_conv_strides,
+        decoder_conv_strides=decoder_conv_strides,
+        multiscale_count=multiscale_count,
+        z_dims=z_dims,
+        output_channels=output_channels,
+        encoder_n_filters=encoder_n_filters,
+        decoder_n_filters=decoder_n_filters,
+        encoder_dropout=encoder_dropout,
+        decoder_dropout=decoder_dropout,
+        nonlinearity=nonlinearity,
+        predict_logvar=predict_logvar,
+        analytical_kl=analytical_kl,
+    )
+
+
+def _create_vae_based_algorithm(
+    algorithm: Literal["hdn", "microsplit"],
+    loss: LVAELossConfig,
+    input_shape: Sequence[int],
+    encoder_conv_strides: tuple[int, ...],
+    decoder_conv_strides: tuple[int, ...],
+    multiscale_count: int,
+    z_dims: tuple[int, ...],
+    output_channels: int,
+    encoder_n_filters: int,
+    decoder_n_filters: int,
+    encoder_dropout: float,
+    decoder_dropout: float,
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ],
+    predict_logvar: Literal[None, "pixelwise"],
+    analytical_kl: bool,
+    gaussian_likelihood: GaussianLikelihoodConfig | None = None,
+    nm_likelihood: NMLikelihoodConfig | None = None,
+) -> dict:
+    """
+    Create a dictionary with the parameters of the VAE-based algorithm model.
+
+    Parameters
+    ----------
+    algorithm : Literal["hdn"]
+        The algorithm type.
+    loss : Literal["hdn"]
+        The loss function type.
+    input_shape : tuple[int, ...]
+        The shape of the input data.
+    encoder_conv_strides : list[int]
+        The strides of the encoder convolutional layers.
+    decoder_conv_strides : list[int]
+        The strides of the decoder convolutional layers.
+    multiscale_count : int
+        The number of multiscale layers.
+    z_dims : list[int]
+        The dimensions of the latent space.
+    output_channels : int
+        The number of output channels.
+    encoder_n_filters : int
+        The number of filters in the encoder.
+    decoder_n_filters : int
+        The number of filters in the decoder.
+    encoder_dropout : float
+        The dropout rate for the encoder.
+    decoder_dropout : float
+        The dropout rate for the decoder.
+    nonlinearity : Literal
+        The nonlinearity function to use.
+    predict_logvar : Literal[None, "pixelwise"]
+        The type of log variance prediction.
+    analytical_kl : bool
+        Whether to use analytical KL divergence.
+    gaussian_likelihood : Optional[GaussianLikelihoodConfig], optional
+        The Gaussian likelihood model, by default None.
+    nm_likelihood : Optional[NMLikelihoodConfig], optional
+        The noise model likelihood model, by default None.
+
+    Returns
+    -------
+    dict
+        A dictionary with the parameters of the VAE-based algorithm model.
+    """
+    network_model = _create_vae_configuration(
+        input_shape=input_shape,
+        encoder_conv_strides=encoder_conv_strides,
+        decoder_conv_strides=decoder_conv_strides,
+        multiscale_count=multiscale_count,
+        z_dims=z_dims,
+        output_channels=output_channels,
+        encoder_n_filters=encoder_n_filters,
+        decoder_n_filters=decoder_n_filters,
+        encoder_dropout=encoder_dropout,
+        decoder_dropout=decoder_dropout,
+        nonlinearity=nonlinearity,
+        predict_logvar=predict_logvar,
+        analytical_kl=analytical_kl,
+    )
+    assert gaussian_likelihood or nm_likelihood, "Likelihood model must be specified"
+    return {
+        "algorithm": algorithm,
+        "loss": loss,
+        "model": network_model,
+        "gaussian_likelihood": gaussian_likelihood,
+        "noise_model_likelihood": nm_likelihood,
+    }
+
+
+def get_likelihood_config(
+    loss_type: Literal["musplit", "denoisplit", "denoisplit_musplit"],
+    # TODO remove different microsplit loss types, refac
+    predict_logvar: Literal["pixelwise"] | None = None,
+    logvar_lowerbound: float | None = -5.0,
+    nm_paths: list[str] | None = None,
+    data_stats: tuple[float, float] | None = None,
+) -> tuple[
+    GaussianLikelihoodConfig | None,
+    MultiChannelNMConfig | None,
+    NMLikelihoodConfig | None,
+]:
+    """Get the likelihood configuration for split models.
+
+    Returns a tuple containing the following optional entries:
+        - GaussianLikelihoodConfig: Gaussian likelihood configuration for musplit losses
+        - MultiChannelNMConfig: Multi-channel noise model configuration for denoisplit
+        losses
+        - NMLikelihoodConfig: Noise model likelihood configuration for denoisplit losses
+
+    Parameters
+    ----------
+    loss_type : Literal["musplit", "denoisplit", "denoisplit_musplit"]
+        The type of loss function to use.
+    predict_logvar : Literal["pixelwise"] | None, optional
+        Type of log variance prediction, by default None.
+        Required when loss_type is "musplit" or "denoisplit_musplit".
+    logvar_lowerbound : float | None, optional
+        Lower bound for the log variance, by default -5.0.
+        Used when loss_type is "musplit" or "denoisplit_musplit".
+    nm_paths : list[str] | None, optional
+        Paths to the noise model files, by default None.
+        Required when loss_type is "denoisplit" or "denoisplit_musplit".
+    data_stats : tuple[float, float] | None, optional
+        Data statistics (mean, std), by default None.
+        Required when loss_type is "denoisplit" or "denoisplit_musplit".
+
+    Returns
+    -------
+    GaussianLikelihoodConfig or None
+        Configuration for the Gaussian likelihood model.
+    MultiChannelNMConfig or None
+        Configuration for the multi-channel noise model.
+    NMLikelihoodConfig or None
+        Configuration for the noise model likelihood.
+
+    Raises
+    ------
+    ValueError
+        If required parameters are missing for the specified loss_type.
+    """
+    # gaussian likelihood
+    if loss_type in ["musplit", "denoisplit_musplit"]:
+        # if predict_logvar is None:
+        #    raise ValueError(f"predict_logvar is required for loss_type '{loss_type}'")
+        # TODO validators should be in pydantic models
+        gaussian_lik_config = GaussianLikelihoodConfig(
+            predict_logvar=predict_logvar,
+            logvar_lowerbound=logvar_lowerbound,
+        )
+    else:
+        gaussian_lik_config = None
+
+    # noise model likelihood
+    if loss_type in ["denoisplit", "denoisplit_musplit"]:
+        # if nm_paths is None:
+        #     raise ValueError(f"nm_paths is required for loss_type '{loss_type}'")
+        # if data_stats is None:
+        #     raise ValueError(f"data_stats is required for loss_type '{loss_type}'")
+        # TODO validators should be in pydantic models
+        gmm_list = []
+        if nm_paths is not None:
+            for NM_path in nm_paths:
+                gmm_list.append(
+                    GaussianMixtureNMConfig(
+                        model_type="GaussianMixtureNoiseModel",
+                        path=NM_path,
+                    )
+                )
+        noise_model_config = MultiChannelNMConfig(noise_models=gmm_list)
+        nm_lik_config = NMLikelihoodConfig()  # TODO this config isn't needed probably
+    else:
+        noise_model_config = None
+        nm_lik_config = None
+
+    return gaussian_lik_config, noise_model_config, nm_lik_config
+
+
+# TODO wrap parameters into model, loss etc
+# TODO refac likelihood configs to make it 1. Can it be done ?
+def create_hdn_configuration(
+    experiment_name: str,
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: Sequence[int],
+    batch_size: int,
+    num_epochs: int = 100,
+    num_steps: int | None = None,
+    encoder_conv_strides: tuple[int, ...] = (2, 2),
+    decoder_conv_strides: tuple[int, ...] = (2, 2),
+    multiscale_count: int = 1,
+    z_dims: tuple[int, ...] = (128, 128),
+    output_channels: int = 1,
+    encoder_n_filters: int = 32,
+    decoder_n_filters: int = 32,
+    encoder_dropout: float = 0.0,
+    decoder_dropout: float = 0.0,
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ] = "ReLU",
+    analytical_kl: bool = False,
+    predict_logvar: Literal["pixelwise"] | None = None,
+    logvar_lowerbound: Union[float, None] = None,
+    logger: Literal["wandb", "tensorboard", "none"] = "none",
+    trainer_params: dict | None = None,
+    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
+    train_dataloader_params: dict[str, Any] | None = None,
+    val_dataloader_params: dict[str, Any] | None = None,
+) -> Configuration:
+    """
+    Create a configuration for training HDN.
+
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
+    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
+    channels. Likewise, if you set the number of channels, then "C" must be present in
+    `axes`.
+
+    To set the number of output channels, use the `n_channels_out` parameter. If it is
+    not specified, it will be assumed to be equal to `n_channels_in`.
+
+    By default, all channels are trained independently. To train all channels together,
+    set `independent_channels` to False.
+
+    By setting `augmentations` to `None`, the default transformations (flip in X and Y,
+    rotations by 90 degrees in the XY plane) are applied. Rather than the default
+    transforms, a list of transforms can be passed to the `augmentations` parameter. To
+    disable the transforms, simply pass an empty list.
+
+    # TODO revisit the necessity of model_params
+
+    Parameters
+    ----------
+    experiment_name : str
+        Name of the experiment.
+    data_type : Literal["array", "tiff", "custom"]
+        Type of the data.
+    axes : str
+        Axes of the data (e.g. SYX).
+    patch_size : List[int]
+        Size of the patches along the spatial dimensions (e.g. [64, 64]).
+    batch_size : int
+        Batch size.
+    num_epochs : int, default=100
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int, optional
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
+    encoder_conv_strides : tuple[int, ...], optional
+        Strides for the encoder convolutional layers, by default (2, 2).
+    decoder_conv_strides : tuple[int, ...], optional
+        Strides for the decoder convolutional layers, by default (2, 2).
+    multiscale_count : int, optional
+        Number of scales in the multiscale architecture, by default 1.
+    z_dims : tuple[int, ...], optional
+        Dimensions of the latent space, by default (128, 128).
+    output_channels : int, optional
+        Number of output channels, by default 1.
+    encoder_n_filters : int, optional
+        Number of filters in the encoder, by default 32.
+    decoder_n_filters : int, optional
+        Number of filters in the decoder, by default 32.
+    encoder_dropout : float, optional
+        Dropout rate for the encoder, by default 0.0.
+    decoder_dropout : float, optional
+        Dropout rate for the decoder, by default 0.0.
+    nonlinearity : Literal, optional
+        Nonlinearity function to use, by default "ReLU".
+    analytical_kl : bool, optional
+        Whether to use analytical KL divergence, by default False.
+    predict_logvar : Literal[None, "pixelwise"], optional
+        Type of log variance prediction, by default None.
+    logvar_lowerbound : Union[float, None], optional
+        Lower bound for the log variance, by default None.
+    logger : Literal["wandb", "tensorboard", "none"], optional
+        Logger to use for training, by default "none".
+    trainer_params : dict, optional
+        Parameters for the trainer class, see PyTorch Lightning documentation.
+    augmentations : Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]], optional
+        List of augmentations to apply, by default None.
+    train_dataloader_params : Optional[dict[str, Any]], optional
+        Parameters for the training dataloader, by default None.
+    val_dataloader_params : Optional[dict[str, Any]], optional
+        Parameters for the validation dataloader, by default None.
+
+    Returns
+    -------
+    Configuration
+        The configuration object for training HDN.
+
+    Examples
+    --------
+    Minimum example:
+    >>> config = create_hdn_configuration(
+    ...     experiment_name="hdn_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100
+    ... )
+
+    You can also limit the number of batches per epoch:
+    >>> config = create_hdn_configuration(
+    ...     experiment_name="hdn_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_steps=100  # limit to 100 batches per epoch
+    ... )
+    """
+    transform_list = _list_spatial_augmentations(augmentations)
+
+    loss_config = LVAELossConfig(
+        loss_type="hdn", denoisplit_weight=1, musplit_weight=0
+    )  # TODO what are the correct defaults for HDN?
+
+    gaussian_likelihood = GaussianLikelihoodConfig(
+        predict_logvar=predict_logvar, logvar_lowerbound=logvar_lowerbound
+    )
+
+    # algorithm & model
+    algorithm_params = _create_vae_based_algorithm(
+        algorithm="hdn",
+        loss=loss_config,
+        input_shape=patch_size,
+        encoder_conv_strides=encoder_conv_strides,
+        decoder_conv_strides=decoder_conv_strides,
+        multiscale_count=multiscale_count,
+        z_dims=z_dims,
+        output_channels=output_channels,
+        encoder_n_filters=encoder_n_filters,
+        decoder_n_filters=decoder_n_filters,
+        encoder_dropout=encoder_dropout,
+        decoder_dropout=decoder_dropout,
+        nonlinearity=nonlinearity,
+        predict_logvar=predict_logvar,
+        analytical_kl=analytical_kl,
+        gaussian_likelihood=gaussian_likelihood,
+        nm_likelihood=None,
+    )
+
+    # data
+    data_params = _create_data_configuration(
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        batch_size=batch_size,
+        augmentations=transform_list,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
+    )
+
+    # Handle trainer parameters with num_epochs and num_steps
+    final_trainer_params = {} if trainer_params is None else trainer_params.copy()
+
+    # Add num_epochs and num_steps if provided
+    if num_epochs is not None:
+        final_trainer_params["max_epochs"] = num_epochs
+    if num_steps is not None:
+        final_trainer_params["limit_train_batches"] = num_steps
+
+    # training
+    training_params = _create_training_configuration(
+        trainer_params=final_trainer_params,
+        logger=logger,
+    )
+
+    return Configuration(
+        experiment_name=experiment_name,
+        algorithm_config=algorithm_params,
+        data_config=data_params,
+        training_config=training_params,
+    )
+
+
+def create_microsplit_configuration(
+    experiment_name: str,
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: Sequence[int],
+    batch_size: int,
+    num_epochs: int = 100,
+    num_steps: int | None = None,
+    encoder_conv_strides: tuple[int, ...] = (2, 2),
+    decoder_conv_strides: tuple[int, ...] = (2, 2),
+    multiscale_count: int = 3,
+    grid_size: int = 32,  # TODO most likely can be derived from patch size
+    z_dims: tuple[int, ...] = (128, 128),
+    output_channels: int = 1,
+    encoder_n_filters: int = 32,
+    decoder_n_filters: int = 32,
+    encoder_dropout: float = 0.0,
+    decoder_dropout: float = 0.0,
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ] = "ReLU",  # TODO do we need all these?
+    analytical_kl: bool = False,
+    predict_logvar: Literal["pixelwise"] = "pixelwise",
+    logvar_lowerbound: Union[float, None] = None,
+    logger: Literal["wandb", "tensorboard", "none"] = "none",
+    trainer_params: dict | None = None,
+    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]] | None = None,
+    nm_paths: list[str] | None = None,
+    data_stats: tuple[float, float] | None = None,
+    train_dataloader_params: dict[str, Any] | None = None,
+    val_dataloader_params: dict[str, Any] | None = None,
+) -> Configuration:
+    """
+    Create a configuration for training MicroSplit.
+
+    Parameters
+    ----------
+    experiment_name : str
+        Name of the experiment.
+    data_type : Literal["array", "tiff", "custom"]
+        Type of the data.
+    axes : str
+        Axes of the data (e.g. SYX).
+    patch_size : Sequence[int]
+        Size of the patches along the spatial dimensions (e.g. [64, 64]).
+    batch_size : int
+        Batch size.
+    num_epochs : int, default=100
+        Number of epochs to train for. If provided, this will be added to
+        trainer_params.
+    num_steps : int, optional
+        Number of batches in 1 epoch. If provided, this will be added to trainer_params.
+        Translates to `limit_train_batches` in PyTorch Lightning Trainer. See relevant
+        documentation for more details.
+    encoder_conv_strides : tuple[int, ...], optional
+        Strides for the encoder convolutional layers, by default (2, 2).
+    decoder_conv_strides : tuple[int, ...], optional
+        Strides for the decoder convolutional layers, by default (2, 2).
+    multiscale_count : int, optional
+        Number of multiscale levels, by default 1.
+    grid_size : int, optional
+        Size of the grid for the lateral context, by default 32.
+    z_dims : tuple[int, ...], optional
+        List of latent dimensions for each hierarchy level in the LVAE, by default
+        (128, 128).
+    output_channels : int, optional
+        Number of output channels for the model, by default 1.
+    encoder_n_filters : int, optional
+        Number of filters in the encoder, by default 32.
+    decoder_n_filters : int, optional
+        Number of filters in the decoder, by default 32.
+    encoder_dropout : float, optional
+        Dropout rate for the encoder, by default 0.0.
+    decoder_dropout : float, optional
+        Dropout rate for the decoder, by default 0.0.
+    nonlinearity : Literal, optional
+        Nonlinearity to use in the model, by default "ReLU".
+    analytical_kl : bool, optional
+        Whether to use analytical KL divergence, by default False.
+    predict_logvar : Literal["pixelwise"] | None, optional
+        Type of log-variance prediction, by default None.
+    logvar_lowerbound : Union[float, None], optional
+        Lower bound for the log variance, by default None.
+    logger : Literal["wandb", "tensorboard", "none"], optional
+        Logger to use for training, by default "none".
+    trainer_params : dict, optional
+        Parameters for the trainer class, see PyTorch Lightning documentation.
+    augmentations : list[Union[XYFlipModel, XYRandomRotate90Model]] | None, optional
+        List of augmentations to apply, by default None.
+    nm_paths : list[str] | None, optional
+        Paths to the noise model files, by default None.
+    data_stats : tuple[float, float] | None, optional
+        Data statistics (mean, std), by default None.
+    train_dataloader_params : dict[str, Any] | None, optional
+        Parameters for the training dataloader, by default None.
+    val_dataloader_params : dict[str, Any] | None, optional
+        Parameters for the validation dataloader, by default None.
+
+    Returns
+    -------
+    Configuration
+        A configuration object for the microsplit algorithm.
+
+    Examples
+    --------
+    Minimum example:
+    # >>> config = create_microsplit_configuration(
+    # ...     experiment_name="microsplit_experiment",
+    # ...     data_type="array",
+    # ...     axes="YX",
+    # ...     patch_size=[64, 64],
+    # ...     batch_size=32,
+    # ...     num_epochs=100
+
+    # ... )
+
+    # You can also limit the number of batches per epoch:
+    # >>> config = create_microsplit_configuration(
+    # ...     experiment_name="microsplit_experiment",
+    # ...     data_type="array",
+    # ...     axes="YX",
+    # ...     patch_size=[64, 64],
+    # ...     batch_size=32,
+    # ...     num_steps=100  # limit to 100 batches per epoch
+    # ... )
+    """
+    transform_list = _list_spatial_augmentations(augmentations)
+
+    loss_config = LVAELossConfig(
+        loss_type="denoisplit_musplit", denoisplit_weight=0.9, musplit_weight=0.1
+    )  # TODO losses need to be refactored! just for example. Add validator if sum to 1
+
+    # Create likelihood configurations
+    gaussian_likelihood_config, noise_model_config, nm_likelihood_config = (
+        get_likelihood_config(
+            loss_type="denoisplit_musplit",
+            predict_logvar=predict_logvar,
+            logvar_lowerbound=logvar_lowerbound,
+            nm_paths=nm_paths,
+            data_stats=data_stats,
+        )
+    )
+
+    # Create the LVAE model
+    network_model = _create_vae_configuration(
+        input_shape=patch_size,
+        encoder_conv_strides=encoder_conv_strides,
+        decoder_conv_strides=decoder_conv_strides,
+        multiscale_count=multiscale_count,
+        z_dims=z_dims,
+        output_channels=output_channels,
+        encoder_n_filters=encoder_n_filters,
+        decoder_n_filters=decoder_n_filters,
+        encoder_dropout=encoder_dropout,
+        decoder_dropout=decoder_dropout,
+        nonlinearity=nonlinearity,
+        predict_logvar=predict_logvar,
+        analytical_kl=analytical_kl,
+    )
+
+    # Create the MicroSplit algorithm configuration
+    algorithm_params = {
+        "algorithm": "microsplit",
+        "loss": loss_config,
+        "model": network_model,
+        "gaussian_likelihood": gaussian_likelihood_config,
+        "noise_model": noise_model_config,
+        "noise_model_likelihood": nm_likelihood_config,
+    }
+
+    # Convert to MicroSplitAlgorithm instance
+    algorithm_config = MicroSplitAlgorithm(**algorithm_params)
+
+    # data
+    data_params = _create_microsplit_data_configuration(
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        grid_size=grid_size,
+        multiscale_count=multiscale_count,
+        batch_size=batch_size,
+        augmentations=transform_list,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
+    )
+
+    # Handle trainer parameters with num_epochs and num_steps
+    final_trainer_params = {} if trainer_params is None else trainer_params.copy()
+
+    # Add num_epochs and num_steps if provided
+    if num_epochs is not None:
+        final_trainer_params["max_epochs"] = num_epochs
+    if num_steps is not None:
+        final_trainer_params["limit_train_batches"] = num_steps
+
+    # training
+    training_params = _create_training_configuration(
+        trainer_params=final_trainer_params,
+        logger=logger,
+    )
+
+    return Configuration(
+        experiment_name=experiment_name,
+        algorithm_config=algorithm_config,
+        data_config=data_params,
+        training_config=training_params,
+    )