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

careamics/models/lvae/likelihoods.py

@@ -5,7 +5,7 @@ Script containing modules for defining different likelihood functions (as nn.Mod
 from __future__ import annotations
 
 import math
-from typing import Literal, Union, TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import numpy as np
 import torch
@@ -26,7 +26,8 @@ if TYPE_CHECKING:
 
 
 def likelihood_factory(
-    config: Union[GaussianLikelihoodConfig, NMLikelihoodConfig, None]
+    config: Optional[Union[GaussianLikelihoodConfig, NMLikelihoodConfig]],
+    noise_model: Optional[NoiseModel] = None,
 ):
     """
     Factory function for creating likelihood modules.
@@ -35,6 +36,8 @@ def likelihood_factory(
     ----------
     config: Union[GaussianLikelihoodConfig, NMLikelihoodConfig]
         The configuration object for the likelihood module.
+    noise_model: Optional[NoiseModel]
+        The noise model instance used to define the `NoiseModelLikelihood`.
 
     Returns
     -------
@@ -53,7 +56,7 @@ def likelihood_factory(
         return NoiseModelLikelihood(
             data_mean=config.data_mean,
             data_std=config.data_std,
-            noiseModel=config.noise_model,
+            noise_model=noise_model,
         )
     else:
         raise ValueError(f"Invalid likelihood model type: {config.model_type}")
@@ -99,8 +102,8 @@ class LikelihoodModule(nn.Module):
         self, input_: torch.Tensor, x: Union[torch.Tensor, None]
     ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         """
-        Parameters:
-        -----------
+        Parameters
+        ----------
         input_: torch.Tensor
             The output of the top-down pass (e.g., reconstructed image in HDN,
             or the unmixed images in 'Split' models).
@@ -181,7 +184,6 @@ class GaussianLikelihood(LikelihoodModule):
             log-variance. If the attribute `predict_logvar` is `None` then the second
             element will be `None`.
         """
-
         # if LadderVAE.predict_logvar is None, dim 1 of `x`` has no. of target channels
         if self.predict_logvar is None:
             return x, None
@@ -290,7 +292,7 @@ class NoiseModelLikelihood(LikelihoodModule):
         self,
         data_mean: Union[np.ndarray, torch.Tensor],
         data_std: Union[np.ndarray, torch.Tensor],
-        noiseModel: NoiseModel,
+        noise_model: NoiseModel,
     ):
         """Constructor.
 
@@ -307,7 +309,7 @@ class NoiseModelLikelihood(LikelihoodModule):
         super().__init__()
         self.data_mean = torch.Tensor(data_mean)
         self.data_std = torch.Tensor(data_std)
-        self.noiseModel = noiseModel
+        self.noiseModel = noise_model
 
     def _set_params_to_same_device_as(
         self, correct_device_tensor: torch.Tensor