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

careamics/dataset/dataset_utils/running_stats.py

@@ -34,36 +34,35 @@ def update_iterative_stats(
     Parameters
     ----------
     count : NDArray
-        Number of elements in the array.
+        Number of elements in the array. Shape: (C,).
     mean : NDArray
-        Mean of the array.
+        Mean of the array. Shape: (C,).
     m2 : NDArray
-        Variance of the array.
+        Variance of the array. Shape: (C,).
     new_values : NDArray
-        New values to add to the mean and variance.
+        New values to add to the mean and variance. Shape: (C, 1, 1, Z, Y, X).
 
     Returns
     -------
     tuple[NDArray, NDArray, NDArray]
         Updated count, mean, and variance.
     """
-    count += np.array([np.prod(channel.shape) for channel in new_values])
-    # newvalues - oldMean
-    delta = [
-        np.subtract(v.flatten(), [m] * len(v.flatten()))
-        for v, m in zip(new_values, mean)
-    ]
+    num_channels = len(new_values)
 
-    mean += np.array([np.sum(d / c) for d, c in zip(delta, count)])
-    # newvalues - newMeant
-    delta2 = [
-        np.subtract(v.flatten(), [m] * len(v.flatten()))
-        for v, m in zip(new_values, mean)
-    ]
+    # --- update channel-wise counts ---
+    count += np.ones_like(count) * np.prod(new_values.shape[1:])
 
-    m2 += np.array([np.sum(d * d2) for d, d2 in zip(delta, delta2)])
+    # --- update channel-wise mean ---
+    # compute (new_values - old_mean) -> shape: (C, Z*Y*X)
+    delta = new_values.reshape(num_channels, -1) - mean.reshape(num_channels, 1)
+    mean += np.sum(delta / count.reshape(num_channels, 1), axis=1)
 
-    return (count, mean, m2)
+    # --- update channel-wise SoS ---
+    # compute (new_values - new_mean) -> shape: (C, Z*Y*X)
+    delta2 = new_values.reshape(num_channels, -1) - mean.reshape(num_channels, 1)
+    m2 += np.sum(delta * delta2, axis=1)
+
+    return count, mean, m2
 
 
 def finalize_iterative_stats(
@@ -74,18 +73,18 @@ def finalize_iterative_stats(
     Parameters
     ----------
     count : NDArray
-        Number of elements in the array.
+        Number of elements in the array. Shape: (C,).
     mean : NDArray
-        Mean of the array.
+        Mean of the array. Shape: (C,).
     m2 : NDArray
-        Variance of the array.
+        Variance of the array. Shape: (C,).
 
     Returns
     -------
     tuple[NDArray, NDArray]
-        Final mean and standard deviation.
+        Final channel-wise mean and standard deviation.
     """
-    std = np.array([np.sqrt(m / c) for m, c in zip(m2, count)])
+    std = np.sqrt(m2 / count)
     if any(c < 2 for c in count):
         return np.full(mean.shape, np.nan), np.full(std.shape, np.nan)
     else:

careamics/lightning/lightning_module.py

@@ -14,8 +14,8 @@ from careamics.config.support import (
     SupportedOptimizer,
     SupportedScheduler,
 )
+from careamics.config.tile_information import TileInformation
 from careamics.losses import loss_factory
-from careamics.losses.loss_factory import LVAELossParameters
 from careamics.models.lvae.likelihoods import (
     GaussianLikelihood,
     NoiseModelLikelihood,
@@ -164,7 +164,17 @@ class FCNModule(L.LightningModule):
         Any
             Model output.
         """
-        if self._trainer.datamodule.tiled:
+        # TODO refactor when redoing datasets
+        # hacky way to determine if it is PredictDataModule, otherwise there is a
+        # circular import to solve with isinstance
+        from_prediction = hasattr(self._trainer.datamodule, "tiled")
+        is_tiled = (
+            len(batch) > 1
+            and isinstance(batch[1], list)
+            and isinstance(batch[1][0], TileInformation)
+        )
+
+        if is_tiled:
             x, *aux = batch
         else:
             x = batch
@@ -172,7 +182,10 @@ class FCNModule(L.LightningModule):
 
         # apply test-time augmentation if available
         # TODO: probably wont work with batch size > 1
-        if self._trainer.datamodule.prediction_config.tta_transforms:
+        if (
+            from_prediction
+            and self._trainer.datamodule.prediction_config.tta_transforms
+        ):
             tta = ImageRestorationTTA()
             augmented_batch = tta.forward(x)  # list of augmented tensors
             augmented_output = []
@@ -184,9 +197,18 @@ class FCNModule(L.LightningModule):
             output = self.model(x)
 
         # Denormalize the output
+        # TODO incompatible API between predict and train datasets
         denorm = Denormalize(
-            image_means=self._trainer.datamodule.predict_dataset.image_means,
-            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
+            image_means=(
+                self._trainer.datamodule.predict_dataset.image_means
+                if from_prediction
+                else self._trainer.datamodule.train_dataset.image_stats.means
+            ),
+            image_stds=(
+                self._trainer.datamodule.predict_dataset.image_stds
+                if from_prediction
+                else self._trainer.datamodule.train_dataset.image_stats.stds
+            ),
         )
         denormalized_output = denorm(patch=output.cpu().numpy())
 
@@ -266,29 +288,27 @@ class VAEModule(L.LightningModule):
         # TODO: log algorithm config
         # self.save_hyperparameters(self.algorithm_config.model_dump())
 
-        # create model and loss function
+        # create model
         self.model: nn.Module = model_factory(self.algorithm_config.model)
-        self.noise_model: NoiseModel = noise_model_factory(
+
+        # create loss function
+        self.noise_model: Optional[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
+
+        self.noise_model_likelihood: Optional[NoiseModelLikelihood] = (
+            likelihood_factory(
+                config=self.algorithm_config.noise_model_likelihood,
+                noise_model=self.noise_model,
+            )
         )
-        self.gaussian_likelihood: GaussianLikelihood = likelihood_factory(
-            self.algorithm_config.gaussian_likelihood_model
+
+        self.gaussian_likelihood: Optional[GaussianLikelihood] = likelihood_factory(
+            self.algorithm_config.gaussian_likelihood
         )
-        self.loss_parameters = LVAELossParameters(
-            noise_model_likelihood=self.noise_model_likelihood,
-            gaussian_likelihood=self.gaussian_likelihood,
-            # TODO: musplit/denoisplit weights ?
-        )  # type: ignore
-        self.loss_func = loss_factory(self.algorithm_config.loss)
+
+        self.loss_parameters = self.algorithm_config.loss
+        self.loss_func = loss_factory(self.algorithm_config.loss.loss_type)
 
         # save optimizer and lr_scheduler names and parameters
         self.optimizer_name = self.algorithm_config.optimizer.name
@@ -344,11 +364,16 @@ class VAEModule(L.LightningModule):
         out = self.model(x)
 
         # Update loss parameters
-        # TODO rethink loss parameters
-        self.loss_parameters.current_epoch = self.current_epoch
+        self.loss_parameters.kl_params.current_epoch = self.current_epoch
 
         # Compute loss
-        loss = self.loss_func(out, target, self.loss_parameters)  # TODO ugly ?
+        loss = self.loss_func(
+            model_outputs=out,
+            targets=target,
+            config=self.loss_parameters,
+            gaussian_likelihood=self.gaussian_likelihood,
+            noise_model_likelihood=self.noise_model_likelihood,
+        )
 
         # Logging
         # TODO: implement a separate logging method?
@@ -376,7 +401,13 @@ class VAEModule(L.LightningModule):
         out = self.model(x)
 
         # Compute loss
-        loss = self.loss_func(out, target, self.loss_parameters)
+        loss = self.loss_func(
+            model_outputs=out,
+            targets=target,
+            config=self.loss_parameters,
+            gaussian_likelihood=self.gaussian_likelihood,
+            noise_model_likelihood=self.noise_model_likelihood,
+        )
 
         # Logging
         # Rename val_loss dict

careamics/lightning/train_data_module.py

@@ -2,11 +2,12 @@
 
 from pathlib import Path
 from typing import Any, Callable, Literal, Optional, Union
+from warnings import warn
 
 import numpy as np
 import pytorch_lightning as L
 from numpy.typing import NDArray
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, IterableDataset
 
 from careamics.config import DataConfig
 from careamics.config.support import SupportedData
@@ -446,6 +447,19 @@ class TrainDataModule(L.LightningDataModule):
         Any
             Training dataloader.
         """
+        # check because iterable dataset cannot be shuffled
+        if not isinstance(self.train_dataset, IterableDataset):
+            if ("shuffle" in self.dataloader_params) and (
+                not self.dataloader_params["shuffle"]
+            ):
+                warn(
+                    "Dataloader parameters include `shuffle=False`, this will be "
+                    "passed to the training dataloader and may result in bad results.",
+                    stacklevel=1,
+                )
+            else:
+                self.dataloader_params["shuffle"] = True
+
         return DataLoader(
             self.train_dataset, batch_size=self.batch_size, **self.dataloader_params
         )

careamics/losses/loss_factory.py

@@ -7,7 +7,7 @@ This module contains a factory function for creating loss functions.
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+from typing import Callable, Union
 
 from torch import Tensor as tensor
 
@@ -15,18 +15,6 @@ from ..config.support import SupportedLoss
 from .fcn.losses import mae_loss, mse_loss, n2v_loss
 from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss
 
-if TYPE_CHECKING:
-    from careamics.models.lvae.likelihoods import (
-        GaussianLikelihood,
-        NoiseModelLikelihood,
-    )
-    from careamics.models.lvae.noise_models import (
-        GaussianMixtureNoiseModel,
-        MultiChannelNoiseModel,
-    )
-
-    NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
-
 
 @dataclass
 class FCNLossParameters:
@@ -40,78 +28,6 @@ class FCNLossParameters:
     loss_weight: float
 
 
-@dataclass  # TODO why not pydantic?
-class LVAELossParameters:
-    """Dataclass for LVAE loss."""
-
-    # TODO: refactor in more modular blocks (otherwise it gets messy very easily)
-    # e.g., - weights, - kl_params, ...
-
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None
-    """Noise model likelihood instance."""
-    gaussian_likelihood: Optional[GaussianLikelihood] = None
-    """Gaussian likelihood instance."""
-    current_epoch: int = 0
-    """Current epoch in the training loop."""
-    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.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."""
-    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`)."""
-    kl_annealing: bool = False
-    """Whether to apply KL loss annealing."""
-    kl_start: int = -1
-    """Epoch at which KL loss annealing starts."""
-    kl_annealtime: int = 10
-    """Number of epochs for which KL loss annealing is applied."""
-    non_stochastic: bool = False
-    """Whether to sample latents and compute KL."""
-
-
-# TODO: really needed?
-# like it is now, it is difficult to use, we need a way to specify the
-# loss parameters in a more user-friendly way.
-def loss_parameters_factory(
-    type: SupportedLoss,
-) -> Union[FCNLossParameters, LVAELossParameters]:
-    """Return loss parameters.
-
-    Parameters
-    ----------
-    type : SupportedLoss
-        Requested loss.
-
-    Returns
-    -------
-    Union[FCNLossParameters, LVAELossParameters]
-        Loss parameters.
-
-    Raises
-    ------
-    NotImplementedError
-        If the loss is unknown.
-    """
-    if type in [SupportedLoss.N2V, SupportedLoss.MSE, SupportedLoss.MAE]:
-        return FCNLossParameters
-
-    elif type in [
-        SupportedLoss.MUSPLIT,
-        SupportedLoss.DENOISPLIT,
-        SupportedLoss.DENOISPLIT_MUSPLIT,
-    ]:
-        return LVAELossParameters  # it returns the class, not an instance
-
-    else:
-        raise NotImplementedError(f"Loss {type} is not yet supported.")
-
-
 def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
     """Return loss function.