careamics 0.0.3__py3-none-any.whl → 0.0.4.1__py3-none-any.whl

careamics/config/configuration_model.py

@@ -124,7 +124,6 @@ class Configuration(BaseModel):
     >>> config_dict = {
     ...         "experiment_name": "N2V_experiment",
     ...         "algorithm_config": {
-    ...             "algorithm_type": "fcn",
     ...             "algorithm": "n2v",
     ...             "loss": "n2v",
     ...             "model": {
@@ -158,7 +157,7 @@ class Configuration(BaseModel):
 
     # Sub-configurations
     algorithm_config: Union[FCNAlgorithmConfig, VAEAlgorithmConfig] = Field(
-        discriminator="algorithm_type"
+        discriminator="algorithm"
     )
     """Algorithm configuration, holding all parameters required to configure the
     model."""

careamics/config/data_model.py

@@ -10,7 +10,6 @@ from numpy.typing import NDArray
 from pydantic import (
     BaseModel,
     ConfigDict,
-    Discriminator,
     Field,
     PlainSerializer,
     field_validator,
@@ -19,9 +18,7 @@ from pydantic import (
 from typing_extensions import Annotated, Self
 
 from .support import SupportedTransform
-from .transformations.n2v_manipulate_model import N2VManipulateModel
-from .transformations.xy_flip_model import XYFlipModel
-from .transformations.xy_random_rotate90_model import XYRandomRotate90Model
+from .transformations import TRANSFORMS_UNION, N2VManipulateModel
 from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
 
@@ -48,17 +45,6 @@ Float = Annotated[float, PlainSerializer(np_float_to_scientific_str, return_type
 """Annotated float type, used to serialize floats to strings."""
 
 
-TRANSFORMS_UNION = Annotated[
-    Union[
-        XYFlipModel,
-        XYRandomRotate90Model,
-        N2VManipulateModel,
-    ],
-    Discriminator("name"),  # used to tell the different transform models apart
-]
-"""Available transforms in CAREamics."""
-
-
 class DataConfig(BaseModel):
     """
     Data configuration.

careamics/config/fcn_algorithm_model.py

@@ -24,11 +24,11 @@ class FCNAlgorithmConfig(BaseModel):
 
     Attributes
     ----------
-    algorithm : Literal["n2v", "custom"]
+    algorithm : {"n2v", "care", "n2n", "custom"}
         Algorithm to use.
-    loss : Literal["n2v", "mae", "mse"]
+    loss : {"n2v", "mae", "mse"}
         Loss function to use.
-    model : Union[UNetModel, LVAEModel, CustomModel]
+    model : UNetModel or CustomModel
         Model architecture to use.
     optimizer : OptimizerModel, optional
         Optimizer to use.
@@ -48,7 +48,6 @@ class FCNAlgorithmConfig(BaseModel):
     >>> from careamics.config import FCNAlgorithmConfig
     >>> config_dict = {
     ...     "algorithm": "n2v",
-    ...     "algorithm_type": "fcn",
     ...     "loss": "n2v",
     ...     "model": {
     ...         "architecture": "UNet",
@@ -65,11 +64,6 @@ class FCNAlgorithmConfig(BaseModel):
     )
 
     # Mandatory fields
-    # defined in SupportedAlgorithm
-    algorithm_type: Literal["fcn"]
-    """Algorithm type must be `fcn` (fully convolutional network) to differentiate this
-    configuration from LVAE."""
-
     algorithm: Literal["n2v", "care", "n2n", "custom"]
     """Name of the algorithm, as defined in SupportedAlgorithm. Use `custom` for custom
     model architecture."""
@@ -145,3 +139,14 @@ class FCNAlgorithmConfig(BaseModel):
             Pretty string.
         """
         return pformat(self.model_dump())
+
+    @classmethod
+    def get_compatible_algorithms(cls) -> list[str]:
+        """Get the list of compatible algorithms.
+
+        Returns
+        -------
+        list of str
+            List of compatible algorithms.
+        """
+        return ["n2v", "care", "n2n"]

careamics/config/likelihood_model.py

@@ -2,16 +2,30 @@
 
 from typing import Literal, Optional, Union
 
+import numpy as np
 import torch
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator
+from typing_extensions import Annotated
 
 from careamics.models.lvae.noise_models import (
     GaussianMixtureNoiseModel,
     MultiChannelNoiseModel,
 )
+from careamics.utils.serializers import _array_to_json, _to_torch
 
 NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
 
+# TODO: this is a temporary solution to serialize and deserialize tensor fields
+# in pydantic models. Specifically, the aim is to enable saving and loading configs
+# with such tensors to/from JSON files during, resp., training and evaluation.
+Tensor = Annotated[
+    Union[np.ndarray, torch.Tensor],
+    PlainSerializer(_array_to_json, return_type=str),
+    PlainValidator(_to_torch),
+]
+"""Annotated tensor type, used to serialize arrays or tensors to JSON strings
+and deserialize them back to tensors."""
+
 
 class GaussianLikelihoodConfig(BaseModel):
     """Gaussian likelihood configuration."""
@@ -31,13 +45,16 @@ class NMLikelihoodConfig(BaseModel):
 
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
-    data_mean: Union[torch.Tensor] = torch.zeros(1)
+    # TODO remove and use as parameters to the likelihood functions?
+    data_mean: Tensor = torch.zeros(1)
     """The mean of the data, used to unnormalize data for noise model evaluation.
     Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
 
-    data_std: Union[torch.Tensor] = torch.ones(1)
+    # TODO remove and use as parameters to the likelihood functions?
+    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))."""
 
-    noise_model: Union[NoiseModel, None] = None
+    # 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/nm_model.py

@@ -4,8 +4,30 @@ from pathlib import Path
 from typing import Literal, Optional, Union
 
 import numpy as np
-from pydantic import BaseModel, ConfigDict, Field, model_validator
-from typing_extensions import Self
+import torch
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PlainSerializer,
+    PlainValidator,
+    model_validator,
+)
+from typing_extensions import Annotated, Self
+
+from careamics.utils.serializers import _array_to_json, _to_numpy
+
+# TODO: this is a temporary solution to serialize and deserialize array fields
+# in pydantic models. Specifically, the aim is to enable saving and loading configs
+# with such arrays to/from JSON files during, resp., training and evaluation.
+Array = Annotated[
+    Union[np.ndarray, torch.Tensor],
+    PlainSerializer(_array_to_json, return_type=str),
+    PlainValidator(_to_numpy),
+]
+"""Annotated array type, used to serialize arrays or tensors to JSON strings
+and deserialize them back to arrays."""
+
 
 # TODO: add histogram-based noise model
 
@@ -26,13 +48,17 @@ class GaussianMixtureNMConfig(BaseModel):
     """Path to the directory where the trained noise model (*.npz) is saved in the
     `train` method."""
 
-    signal: Optional[Union[str, Path, np.ndarray]] = None
+    # TODO remove and use as parameters to the NM functions?
+    signal: Optional[Union[str, Path, np.ndarray]] = Field(default=None, exclude=True)
     """Path to the file containing signal or respective numpy array."""
 
-    observation: Optional[Union[str, Path, np.ndarray]] = None
+    # TODO remove and use as parameters to the NM functions?
+    observation: Optional[Union[str, Path, np.ndarray]] = Field(
+        default=None, exclude=True
+    )
     """Path to the file containing observation or respective numpy array."""
 
-    weight: Optional[np.ndarray] = None
+    weight: Optional[Array] = None
     """A [3*n_gaussian, n_coeff] sized array containing the values of the weights
     describing the GMM noise model, with each row corresponding to one
     parameter of each gaussian, namely [mean, standard deviation and weight].

careamics/config/optimizer_models.py

@@ -44,7 +44,9 @@ class OptimizerModel(BaseModel):
     )
 
     # Mandatory field
-    name: Literal["Adam", "SGD"] = Field(default="Adam", validate_default=True)
+    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

careamics/config/support/supported_optimizers.py

@@ -19,7 +19,7 @@ class SupportedOptimizer(str, BaseEnum):
     # Adagrad = "Adagrad"
     ADAM = "Adam"
     # AdamW = "AdamW"
-    # Adamax = "Adamax"
+    ADAMAX = "Adamax"
     # LBFGS = "LBFGS"
     # NAdam = "NAdam"
     # RAdam = "RAdam"

careamics/config/support/supported_transforms.py

@@ -9,3 +9,4 @@ class SupportedTransform(str, BaseEnum):
     XY_FLIP = "XYFlip"
     XY_RANDOM_ROTATE90 = "XYRandomRotate90"
     N2V_MANIPULATE = "N2VManipulate"
+    NORMALIZE = "Normalize"

careamics/config/training_model.py

@@ -3,13 +3,9 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-)
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from .callback_model import CheckpointModel, EarlyStoppingModel
 
@@ -37,6 +33,20 @@ class TrainingConfig(BaseModel):
     num_epochs: int = Field(default=20, ge=1)
     """Number of epochs, greater than 0."""
 
+    precision: Literal["64", "32", "16-mixed", "bf16-mixed"] = Field(default="32")
+    """Numerical precision"""
+    max_steps: int = Field(default=-1, ge=-1)
+    """Maximum number of steps to train for. -1 means no limit."""
+    check_val_every_n_epoch: int = Field(default=1, ge=1)
+    """Validation step frequency."""
+    enable_progress_bar: bool = Field(default=True)
+    """Whether to enable the progress bar."""
+    accumulate_grad_batches: int = Field(default=1, ge=1)
+    """Number of batches to accumulate gradients over before stepping the optimizer."""
+    gradient_clip_val: Optional[Union[int, float]] = None
+    """The value to which to clip the gradient"""
+    gradient_clip_algorithm: Literal["value", "norm"] = "norm"
+    """The algorithm to use for gradient clipping (see lightning `Trainer`)."""
     logger: Optional[Literal["wandb", "tensorboard"]] = None
     """Logger to use during training. If None, no logger will be used. Available
     loggers are defined in SupportedLogger."""
@@ -70,3 +80,22 @@ class TrainingConfig(BaseModel):
             Whether the logger is defined or not.
         """
         return self.logger is not None
+
+    @field_validator("max_steps")
+    @classmethod
+    def validate_max_steps(cls, max_steps: int) -> int:
+        """Validate the max_steps parameter.
+
+        Parameters
+        ----------
+        max_steps : int
+            Maximum number of steps to train for. -1 means no limit.
+
+        Returns
+        -------
+        int
+            Validated max_steps.
+        """
+        if max_steps == 0:
+            raise ValueError("max_steps must be greater than 0. Use -1 for no limit.")
+        return max_steps

careamics/config/transformations/__init__.py

@@ -5,11 +5,14 @@ __all__ = [
     "XYFlipModel",
     "NormalizeModel",
     "XYRandomRotate90Model",
-    "XorYFlipModel",
+    "TransformModel",
+    "TRANSFORMS_UNION",
 ]
 
 
 from .n2v_manipulate_model import N2VManipulateModel
 from .normalize_model import NormalizeModel
+from .transform_model import TransformModel
+from .transform_union import TRANSFORMS_UNION
 from .xy_flip_model import XYFlipModel
 from .xy_random_rotate90_model import XYRandomRotate90Model

careamics/config/transformations/transform_union.py

@@ -0,0 +1,20 @@
+"""Type used to represent all transformations users can create."""
+
+from typing import Union
+
+from pydantic import Discriminator
+from typing_extensions import Annotated
+
+from .n2v_manipulate_model import N2VManipulateModel
+from .xy_flip_model import XYFlipModel
+from .xy_random_rotate90_model import XYRandomRotate90Model
+
+TRANSFORMS_UNION = Annotated[
+    Union[
+        XYFlipModel,
+        XYRandomRotate90Model,
+        N2VManipulateModel,
+    ],
+    Discriminator("name"),  # used to tell the different transform models apart
+]
+"""Available transforms in CAREamics."""

careamics/config/vae_algorithm_model.py

@@ -19,40 +19,7 @@ from .optimizer_models import LrSchedulerModel, OptimizerModel
 class VAEAlgorithmConfig(BaseModel):
     """Algorithm configuration.
 
-    This Pydantic model validates the parameters governing the components of the
-    training algorithm: which algorithm, loss function, model architecture, optimizer,
-    and learning rate scheduler to use.
-
-    Currently, we only support N2V, CARE, N2N and custom models. The `n2v` algorithm is
-    only compatible with `n2v` loss and `UNet` architecture. The `custom` algorithm
-    allows you to register your own architecture and select it using its name as
-    `name` in the custom pydantic model.
-
-    Attributes
-    ----------
-    algorithm : algorithm: Literal["musplit", "denoisplit", "custom"]
-        Algorithm to use.
-    loss : Literal["musplit", "denoisplit", "denoisplit_musplit"]
-        Loss function to use.
-    model : Union[LVAEModel, CustomModel]
-        Model architecture to use.
-    noise_model: Optional[MultiChannelNmModel]
-        Noise model to use.
-    noise_model_likelihood_model: Optional[NMLikelihoodModel]
-        Noise model likelihood model to use.
-    gaussian_likelihood_model: Optional[GaussianLikelihoodModel]
-        Gaussian likelihood model to use.
-    optimizer : OptimizerModel, optional
-        Optimizer to use.
-    lr_scheduler : LrSchedulerModel, optional
-        Learning rate scheduler to use.
-
-    Raises
-    ------
-    ValueError
-        Algorithm parameter type validation errors.
-    ValueError
-        If the algorithm, loss and model are not compatible.
+    # TODO
 
     Examples
     --------
@@ -70,8 +37,7 @@ class VAEAlgorithmConfig(BaseModel):
     # defined in SupportedAlgorithm
     # TODO: Use supported Enum classes for typing?
     #   - values can still be passed as strings and they will be cast to Enum
-    algorithm_type: Literal["vae"]
-    algorithm: Literal["musplit", "denoisplit", "custom"]
+    algorithm: Literal["musplit", "denoisplit"]
     loss: Literal["musplit", "denoisplit", "denoisplit_musplit"]
     model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
 

careamics/dataset/tiling/lvae_tiled_patching.py

@@ -8,6 +8,7 @@ import numpy as np
 from numpy.typing import NDArray
 
 from careamics.config.tile_information import TileInformation
+from careamics.lvae_training.dataset.utils.index_manager import GridIndexManager
 
 
 def extract_tiles(
@@ -66,10 +67,12 @@ def extract_tiles(
         # itertools.product is equivalent of nested loops
 
         stitch_size = tile_size - overlaps
-        for tile_grid_coords in itertools.product(*[range(n) for n in tile_grid_shape]):
+        for tile_grid_indices in itertools.product(
+            *[range(n) for n in tile_grid_shape]
+        ):
 
             # calculate crop coordinates
-            crop_coords_start = np.array(tile_grid_coords) * stitch_size
+            crop_coords_start = np.array(tile_grid_indices) * stitch_size
             crop_slices: tuple[Union[builtins.ellipsis, slice], ...] = (
                 ...,
                 *[
@@ -80,7 +83,7 @@ def extract_tiles(
             tile = sample[crop_slices]
 
             tile_info = compute_tile_info(
-                np.array(tile_grid_coords),
+                np.array(tile_grid_indices),
                 np.array(data_shape),
                 np.array(tile_size),
                 np.array(overlaps),
@@ -93,19 +96,98 @@ def extract_tiles(
             yield tile, tile_info
 
 
+def compute_tile_info_legacy(
+    grid_index_manager: GridIndexManager, index: int
+) -> TileInformation:
+    """
+    Compute the tile information for a tile at a given dataset index.
+
+    Parameters
+    ----------
+    grid_index_manager : GridIndexManager
+        The grid index manager that keeps track of tile locations.
+    index : int
+        The dataset index.
+
+    Returns
+    -------
+    TileInformation
+        Information that describes how to crop and stitch a tile to create a full image.
+
+    Raises
+    ------
+    ValueError
+        If `grid_index_manager.data_shape` does not have 4 or 5 dimensions.
+    """
+    data_shape = np.array(grid_index_manager.data_shape)
+    if len(data_shape) == 5:
+        n_spatial_dims = 3
+    elif len(data_shape) == 4:
+        n_spatial_dims = 2
+    else:
+        raise ValueError("Data shape must have 4 or 5 dimensions, equating to SC(Z)YX.")
+
+    stitch_coords_start = np.array(
+        grid_index_manager.get_location_from_dataset_idx(index)
+    )
+    stitch_coords_end = stitch_coords_start + np.array(grid_index_manager.grid_shape)
+
+    tile_coords_start = stitch_coords_start - grid_index_manager.patch_offset()
+
+    # --- replace out of bounds indices
+    out_of_lower_bound = stitch_coords_start < 0
+    out_of_upper_bound = stitch_coords_end > data_shape
+    stitch_coords_start[out_of_lower_bound] = 0
+    stitch_coords_end[out_of_upper_bound] = data_shape[out_of_upper_bound]
+
+    # TODO: TilingMode not in current version
+    # if grid_index_manager.tiling_mode == TilingMode.ShiftBoundary:
+    #     for dim in range(len(stitch_coords_start)):
+    #         if tile_coords_start[dim] == 0:
+    #             stitch_coords_start[dim] = 0
+    #         if tile_coords_end[dim] == grid_index_manager.data_shape[dim]:
+    #             tile_coords_end [dim]= grid_index_manager.data_shape[dim]
+
+    # --- calculate overlap crop coords
+    overlap_crop_coords_start = stitch_coords_start - tile_coords_start
+    overlap_crop_coords_end = overlap_crop_coords_start + (
+        stitch_coords_end - stitch_coords_start
+    )
+
+    last_tile = index == grid_index_manager.total_grid_count() - 1
+
+    # --- combine start and end
+    stitch_coords = tuple(
+        (start, end) for start, end in zip(stitch_coords_start, stitch_coords_end)
+    )
+    overlap_crop_coords = tuple(
+        (start, end)
+        for start, end in zip(overlap_crop_coords_start, overlap_crop_coords_end)
+    )
+
+    tile_info = TileInformation(
+        array_shape=data_shape[1:],  # remove S dim
+        last_tile=last_tile,
+        overlap_crop_coords=overlap_crop_coords[-n_spatial_dims:],
+        stitch_coords=stitch_coords[-n_spatial_dims:],
+        sample_id=0,
+    )
+    return tile_info
+
+
 def compute_tile_info(
-    tile_grid_coords: NDArray[np.int_],
+    tile_grid_indices: NDArray[np.int_],
     data_shape: NDArray[np.int_],
     tile_size: NDArray[np.int_],
     overlaps: NDArray[np.int_],
     sample_id: int = 0,
 ) -> TileInformation:
     """
-    Compute the tile information for a tile with the coordinates `tile_grid_coords`.
+    Compute the tile information for a tile with the coordinates `tile_grid_indices`.
 
     Parameters
     ----------
-    tile_grid_coords : 1D np.array of int
+    tile_grid_indices : 1D np.array of int
         The coordinates of the tile within the tile grid, ((Z), Y, X), i.e. for 2D
         tiling the coordinates for the second tile in the first row of tiles would be
         (0, 1).
@@ -127,7 +209,7 @@ def compute_tile_info(
 
     # The extent of the tile which will make up part of the stitched image.
     stitch_size = tile_size - overlaps
-    stitch_coords_start = tile_grid_coords * stitch_size
+    stitch_coords_start = tile_grid_indices * stitch_size
     stitch_coords_end = stitch_coords_start + stitch_size
 
     tile_coords_start = stitch_coords_start - overlaps // 2
@@ -155,7 +237,7 @@ def compute_tile_info(
 
     # --- Check if last tile
     tile_grid_shape = np.array(compute_tile_grid_shape(data_shape, tile_size, overlaps))
-    last_tile = (tile_grid_coords == (tile_grid_shape - 1)).all()
+    last_tile = (tile_grid_indices == (tile_grid_shape - 1)).all()
 
     tile_info = TileInformation(
         array_shape=data_shape,

careamics/lightning/lightning_module.py

@@ -1,6 +1,6 @@
 """CAREamics Lightning module."""
 
-from typing import Any, Callable, Literal, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -271,6 +271,12 @@ class VAEModule(L.LightningModule):
         self.noise_model: NoiseModel = noise_model_factory(
             self.algorithm_config.noise_model
         )
+        # TODO: here we can add some code to check whether the noise model is not None
+        # and `self.algorithm_config.noise_model_likelihood_model.noise_model` is,
+        # instead, None. In that case we could assign the noise model to the latter.
+        # This is particular useful when loading an algorithm config from file.
+        # Indeed, in that case the noise model in the nm likelihood is likely
+        # not available since excluded from serializaion.
         self.noise_model_likelihood: NoiseModelLikelihood = likelihood_factory(
             self.algorithm_config.noise_model_likelihood_model
         )
@@ -550,7 +556,6 @@ class VAEModule(L.LightningModule):
 
 # TODO: make this LVAE compatible (?)
 def create_careamics_module(
-    algorithm_type: Literal["fcn"],
     algorithm: Union[SupportedAlgorithm, str],
     loss: Union[SupportedLoss, str],
     architecture: Union[SupportedArchitecture, str],
@@ -567,8 +572,6 @@ def create_careamics_module(
 
     Parameters
     ----------
-    algorithm_type : Literal["fcn"]
-        Algorithm type to use for training.
     algorithm : SupportedAlgorithm or str
         Algorithm to use for training (see SupportedAlgorithm).
     loss : SupportedLoss or str
@@ -604,7 +607,6 @@ def create_careamics_module(
     if model_parameters is None:
         model_parameters = {}
     algorithm_configuration: dict[str, Any] = {
-        "algorithm_type": algorithm_type,
         "algorithm": algorithm,
         "loss": loss,
         "optimizer": {
@@ -623,10 +625,10 @@ def create_careamics_module(
     algorithm_configuration["model"] = model_configuration
 
     # call the parent init using an AlgorithmModel instance
-    if algorithm_configuration["algorithm_type"] == "fcn":
+    algorithm_str = algorithm_configuration["algorithm"]
+    if algorithm_str in FCNAlgorithmConfig.get_compatible_algorithms():
         return FCNModule(FCNAlgorithmConfig(**algorithm_configuration))
     else:
         raise NotImplementedError(
-            f"Model {algorithm_configuration['model']['architecture']} is not"
-            f"implemented or unknown."
+            f"Model {algorithm_str} is not implemented or unknown."
         )

careamics/lightning/train_data_module.py

@@ -9,8 +9,8 @@ from numpy.typing import NDArray
 from torch.utils.data import DataLoader
 
 from careamics.config import DataConfig
-from careamics.config.data_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedData
+from careamics.config.transformations import TransformModel
 from careamics.dataset.dataset_utils import (
     get_files_size,
     list_files,
@@ -472,7 +472,7 @@ def create_train_datamodule(
     axes: str,
     batch_size: int,
     val_data: Optional[Union[str, Path, NDArray]] = None,
-    transforms: Optional[list[TRANSFORMS_UNION]] = None,
+    transforms: Optional[list[TransformModel]] = None,
     train_target_data: Optional[Union[str, Path, NDArray]] = None,
     val_target_data: Optional[Union[str, Path, NDArray]] = None,
     read_source_func: Optional[Callable] = None,

careamics/losses/loss_factory.py

@@ -56,9 +56,9 @@ class LVAELossParameters:
     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`)."""
-    musplit_weight: float = 0.0
-    """Weight for the muSplit loss (used in the muSplit-deonoiSplit loss)."""
-    denoisplit_weight: float = 1.0
+    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_type: Literal["kl", "kl_restricted", "kl_spatial", "kl_channelwise"] = "kl"
     """Type of KL divergence used as KL loss."""

careamics/losses/lvae/losses.py

@@ -137,8 +137,8 @@ def reconstruction_loss_musplit_denoisplit(
     recons_loss : torch.Tensor
         The reconstruction loss. Shape is (1, ).
     """
-    # TODO: is this safe to check for predict_logvar value?
-    # otherwise use `gaussian_likelihood.predict_logvar` (or both)
+    # TODO: refactor this function to make it closer to `get_reconstruction_loss`
+    # (or viceversa)
     if predictions.shape[1] == 2 * targets.shape[1]:
         # predictions contain both mean and log-variance
         out_mean, _ = predictions.chunk(2, dim=1)

careamics/lvae_training/dataset/__init__.py

@@ -0,0 +1,15 @@
+from .multich_dataset import MultiChDloader
+from .lc_dataset import LCMultiChDloader
+from .multifile_dataset import MultiFileDset
+from .config import DatasetConfig
+from .types import DataType, DataSplitType, TilingMode
+
+__all__ = [
+    "DatasetConfig",
+    "MultiChDloader",
+    "LCMultiChDloader",
+    "MultiFileDset",
+    "DataType",
+    "DataSplitType",
+    "TilingMode",
+]