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

careamics/models/lvae/noise_models.py

@@ -1,26 +1,108 @@
-import json
-import os
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import torch
 import torch.nn as nn
 
-from .utils import ModelType
+if TYPE_CHECKING:
+    from careamics.config import GaussianMixtureNMConfig, MultiChannelNMConfig
 
+# TODO this module shouldn't be in lvae folder
 
-class DisentNoiseModel(nn.Module):
 
-    def __init__(self, *nmodels):
-        """
-        Constructor.
+def noise_model_factory(
+    model_config: Optional[MultiChannelNMConfig],
+) -> Optional[MultiChannelNoiseModel]:
+    """Noise model factory.
+
+    Parameters
+    ----------
+    model_config : Optional[MultiChannelNMConfig]
+        Noise model configuration, a `MultiChannelNMConfig` config that defines
+        noise models for the different output channels.
+
+    Returns
+    -------
+    Optional[MultiChannelNoiseModel]
+        A noise model instance.
+
+    Raises
+    ------
+    NotImplementedError
+        If the chosen noise model `model_type` is not implemented.
+        Currently only `GaussianMixtureNoiseModel` is implemented.
+    """
+    if model_config:
+        noise_models = []
+        for nm_config in model_config.noise_models:
+            if nm_config.path:
+                if nm_config.model_type == "GaussianMixtureNoiseModel":
+                    noise_models.append(GaussianMixtureNoiseModel(nm_config))
+                else:
+                    raise NotImplementedError(
+                        f"Model {nm_config.model_type} is not implemented"
+                    )
+
+            else:  # TODO this means signal/obs are provided. Controlled in pydantic model
+                # TODO train a new model. Config should always be provided?
+                if nm_config.model_type == "GaussianMixtureNoiseModel":
+                    trained_nm = train_gm_noise_model(nm_config)
+                    noise_models.append(trained_nm)
+                else:
+                    raise NotImplementedError(
+                        f"Model {nm_config.model_type} is not implemented"
+                    )
+        return MultiChannelNoiseModel(noise_models)
+    return None
+
+
+def train_gm_noise_model(
+    model_config: GaussianMixtureNMConfig,
+) -> GaussianMixtureNoiseModel:
+    """Train a Gaussian mixture noise model.
+
+    Parameters
+    ----------
+    model_config : GaussianMixtureNoiseModel
+        _description_
+
+    Returns
+    -------
+    _description_
+    """
+    # TODO where to put train params?
+    # TODO any training params ? Different channels ?
+    noise_model = GaussianMixtureNoiseModel(model_config)
+    # TODO revisit config unpacking
+    noise_model.train_noise_model(model_config.signal, model_config.observation)
+    return noise_model
+
 
-        This class receives as input a variable number of noise models, each one corresponding to a channel.
+class MultiChannelNoiseModel(nn.Module):
+    def __init__(self, nmodels: list[GaussianMixtureNoiseModel]):
+        """Constructor.
+
+        To handle noise models and the relative likelihood computation for multiple
+        output channels (e.g., muSplit, denoiseSplit).
+
+        This class:
+        - receives as input a variable number of noise models, one for each channel.
+        - computes the likelihood of observations given signals for each channel.
+        - returns the concatenation of these likelihoods.
+
+        Parameters
+        ----------
+        nmodels : list[GaussianMixtureNoiseModel]
+            List of noise models, one for each output channel.
         """
         super().__init__()
-        # self.nmodels = nmodels
         for i, nmodel in enumerate(nmodels):
             if nmodel is not None:
-                self.add_module(f"nmodel_{i}", nmodel)
+                self.add_module(
+                    f"nmodel_{i}", nmodel
+                )  # TODO: wouldn't be easier to use a list?
 
         self._nm_cnt = 0
         for nmodel in nmodels:
@@ -30,181 +112,141 @@ class DisentNoiseModel(nn.Module):
         print(f"[{self.__class__.__name__}] Nmodels count:{self._nm_cnt}")
 
     def likelihood(self, obs: torch.Tensor, signal: torch.Tensor) -> torch.Tensor:
+        """Compute the likelihood of observations given signals for each channel.
 
+        Parameters
+        ----------
+        obs : torch.Tensor
+            Noisy observations, i.e., the target(s). Specifically, the input noisy
+            image for HDN, or the noisy unmixed images used for supervision
+            for denoiSplit. Shape: (B, C, [Z], Y, X), where C is the number of
+            unmixed channels.
+        signal : torch.Tensor
+            Underlying signals, i.e., the (clean) output of the model. Specifically, the
+            denoised image for HDN, or the unmixed images for denoiSplit.
+            Shape: (B, C, [Z], Y, X), where C is the number of unmixed channels.
+        """
+        # Case 1: obs and signal have a single channel (e.g., denoising)
         if obs.shape[1] == 1:
             assert signal.shape[1] == 1
-            assert self.n2model is None
             return self.nmodel_0.likelihood(obs, signal)
 
-        assert obs.shape[1] == self._nm_cnt, f"{obs.shape[1]} != {self._nm_cnt}"
-
+        # Case 2: obs and signal have multiple channels (e.g., denoiSplit)
+        assert obs.shape[1] == self._nm_cnt, (
+            "The number of channels in `obs` must match the number of noise models."
+            f" Got instead: obs={obs.shape[1]},  nm={self._nm_cnt}"
+        )
         ll_list = []
         for ch_idx in range(obs.shape[1]):
             nmodel = getattr(self, f"nmodel_{ch_idx}")
             ll_list.append(
                 nmodel.likelihood(
                     obs[:, ch_idx : ch_idx + 1], signal[:, ch_idx : ch_idx + 1]
-                )
+                )  # slicing to keep the channel dimension
             )
-
         return torch.cat(ll_list, dim=1)
 
 
-def last2path(fpath: str):
-    return os.path.join(*fpath.split("/")[-2:])
-
-
-def get_nm_config(noise_model_fpath: str):
-    config_fpath = os.path.join(os.path.dirname(noise_model_fpath), "config.json")
-    with open(config_fpath) as f:
-        noise_model_config = json.load(f)
-    return noise_model_config
-
-
+# TODO: is this needed?
 def fastShuffle(series, num):
+    """_summary_.
+
+    Parameters
+    ----------
+    series : _type_
+        _description_
+    num : _type_
+        _description_
+
+    Returns
+    -------
+    _type_
+        _description_
+    """
     length = series.shape[0]
-    for i in range(num):
+    for _ in range(num):
         series = series[np.random.permutation(length), :]
     return series
 
 
-def get_noise_model(
-    enable_noise_model: bool,
-    model_type: ModelType,
-    noise_model_type: str,
-    noise_model_ch1_fpath: str,
-    noise_model_ch2_fpath: str,
-    noise_model_learnable: bool = False,
-    denoise_channel: str = "input",
-):
-    if enable_noise_model:
-        nmodels = []
-        # HDN -> one single output -> one single noise model
-        if model_type == ModelType.Denoiser:
-            if noise_model_type == "hist":
-                raise NotImplementedError(
-                    '"hist" noise model is not supported for now.'
-                )
-            elif noise_model_type == "gmm":
-                if denoise_channel == "Ch1":
-                    nmodel_fpath = noise_model_ch1_fpath
-                    print(f"Noise model Ch1: {nmodel_fpath}")
-                    nmodel1 = GaussianMixtureNoiseModel(params=np.load(nmodel_fpath))
-                    nmodel2 = None
-                    nmodels = [nmodel1, nmodel2]
-                elif denoise_channel == "Ch2":
-                    nmodel_fpath = noise_model_ch2_fpath
-                    print(f"Noise model Ch2: {nmodel_fpath}")
-                    nmodel1 = GaussianMixtureNoiseModel(params=np.load(nmodel_fpath))
-                    nmodel2 = None
-                    nmodels = [nmodel1, nmodel2]
-                elif denoise_channel == "input":
-                    nmodel_fpath = noise_model_ch1_fpath
-                    print(f"Noise model input: {nmodel_fpath}")
-                    nmodel1 = GaussianMixtureNoiseModel(params=np.load(nmodel_fpath))
-                    nmodel2 = None
-                    nmodels = [nmodel1, nmodel2]
-                else:
-                    raise ValueError(f"Invalid denoise_channel: {denoise_channel}")
-        # muSplit -> two outputs -> two noise models
-        elif noise_model_type == "gmm":
-            print(f"Noise model Ch1: {noise_model_ch1_fpath}")
-            print(f"Noise model Ch2: {noise_model_ch2_fpath}")
-
-            nmodel1 = GaussianMixtureNoiseModel(params=np.load(noise_model_ch1_fpath))
-            nmodel2 = GaussianMixtureNoiseModel(params=np.load(noise_model_ch2_fpath))
-
-            nmodels = [nmodel1, nmodel2]
-
-            # if 'noise_model_ch3_fpath' in config.model:
-            #     print(f'Noise model Ch3: {config.model.noise_model_ch3_fpath}')
-            #     nmodel3 = GaussianMixtureNoiseModel(params=np.load(config.model.noise_model_ch3_fpath))
-            #     nmodels = [nmodel1, nmodel2, nmodel3]
-            # else:
-            #     nmodels = [nmodel1, nmodel2]
-        else:
-            raise ValueError(f"Invalid noise_model_type: {noise_model_type}")
-
-        if noise_model_learnable:
-            for nmodel in nmodels:
-                if nmodel is not None:
-                    nmodel.make_learnable()
-
-        return DisentNoiseModel(*nmodels)
-    return None
-
-
 class GaussianMixtureNoiseModel(nn.Module):
-    """
-    The GaussianMixtureNoiseModel class describes a noise model which is parameterized as a mixture of gaussians.
-    If you would like to initialize a new object from scratch, then set `params`= None and specify the other parameters as keyword arguments.
-    If you are instead loading a model, use only `params`.
+    """Define a noise model parameterized as a mixture of gaussians.
+
+    If `config.path` is not provided a new object is initialized from scratch.
+    Otherwise, a model is loaded from `config.path`.
 
     Parameters
     ----------
-    **kwargs: keyworded, variable-length argument dictionary.
-    Arguments include:
-        min_signal : float
-            Minimum signal intensity expected in the image.
-        max_signal : float
-            Maximum signal intensity expected in the image.
-        path: string
-            Path to the directory where the trained noise model (*.npz) is saved in the `train` method.
-        weight : array
-            A [3*n_gaussian, n_coeff] sized array containing the values of the weights describing the noise model.
-            Each gaussian contributes three parameters (mean, standard deviation and weight), hence the number of rows in `weight` are 3*n_gaussian.
-            If `weight=None`, the weight array is initialized using the `min_signal` and `max_signal` parameters.
-        n_gaussian: int
-            Number of gaussians.
-        n_coeff: int
-            Number of coefficients to describe the functional relationship between gaussian parameters and the signal.
-            2 implies a linear relationship, 3 implies a quadratic relationship and so on.
-        device: device
-            GPU device.
-        min_sigma: int
-            All values of sigma (`standard deviation`) below min_sigma are clamped to become equal to min_sigma.
-        params: dictionary
-            Use `params` if one wishes to load a model with trained weights.
-            While initializing a new object of the class `GaussianMixtureNoiseModel` from scratch, set this to `None`.
+    config : GaussianMixtureNMConfig
+        A `pydantic` model that defines the configuration of the GMM noise model.
+
+    Attributes
+    ----------
+    min_signal : float
+        Minimum signal intensity expected in the image.
+    max_signal : float
+        Maximum signal intensity expected in the image.
+    path: Union[str, Path]
+        Path to the directory where the trained noise model (*.npz) is saved in the `train` method.
+    weight : torch.nn.Parameter
+        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].
+        Specifically, rows are organized as follows:
+        - first n_gaussian rows correspond to the means
+        - next n_gaussian rows correspond to the weights
+        - last n_gaussian rows correspond to the standard deviations
+        If `weight=None`, the weight array is initialized using the `min_signal`
+        and `max_signal` parameters.
+    n_gaussian: int
+        Number of gaussians in the mixture.
+    n_coeff: int
+        Number of coefficients to describe the functional relationship between gaussian
+        parameters and the signal. 2 implies a linear relationship, 3 implies a quadratic
+        relationship and so on.
+    device: device
+        GPU device.
+    min_sigma: float
+        All values of `standard deviation` below this are clamped to this value.
     """
 
-    def __init__(self, **kwargs):
+    # TODO training a NM relies on getting a clean data(N2V e.g,)
+    def __init__(self, config: GaussianMixtureNMConfig):
         super().__init__()
         self._learnable = False
 
-        if kwargs.get("params") is None:
-            weight = kwargs.get("weight")
-            n_gaussian = kwargs.get("n_gaussian")
-            n_coeff = kwargs.get("n_coeff")
-            min_signal = kwargs.get("min_signal")
-            max_signal = kwargs.get("max_signal")
+        if config.path is None:
+            # TODO this is (probably) to train a nm. We leave it for later refactoring
+            weight = config.weight
+            n_gaussian = config.n_gaussian
+            n_coeff = config.n_coeff
+            min_signal = config.min_signal
+            max_signal = config.max_signal
             # self.device = kwargs.get('device')
-            self.path = kwargs.get("path")
-            self.min_sigma = kwargs.get("min_sigma")
+            # TODO min_sigma cant be None ?
+            self.min_sigma = config.min_sigma
             if weight is None:
                 weight = np.random.randn(n_gaussian * 3, n_coeff)
                 weight[n_gaussian : 2 * n_gaussian, 1] = np.log(max_signal - min_signal)
-                weight = torch.from_numpy(
-                    weight.astype(np.float32)
-                ).float()  # .to(self.device)
-                weight = nn.Parameter(weight, requires_grad=True)
+                weight = torch.from_numpy(weight.astype(np.float32)).float().cuda()
+                weight.requires_grad = True
 
             self.n_gaussian = weight.shape[0] // 3
             self.n_coeff = weight.shape[1]
             self.weight = weight
-            self.min_signal = torch.Tensor([min_signal])  # .to(self.device)
-            self.max_signal = torch.Tensor([max_signal])  # .to(self.device)
-            self.tol = torch.Tensor([1e-10])  # .to(self.device)
+            self.min_signal = torch.Tensor([min_signal])
+            self.max_signal = torch.Tensor([max_signal])
+            self.tol = torch.Tensor([1e-10])
         else:
-            params = kwargs.get("params")
+            params = np.load(config.path)
             # self.device = kwargs.get('device')
 
-            self.min_signal = torch.Tensor(params["min_signal"])  # .to(self.device)
-            self.max_signal = torch.Tensor(params["max_signal"])  # .to(self.device)
+            self.min_signal = torch.Tensor(params["min_signal"])
+            self.max_signal = torch.Tensor(params["max_signal"])
 
             self.weight = torch.nn.Parameter(
                 torch.Tensor(params["trained_weight"]), requires_grad=False
-            )  # .to(self.device)
+            )
             self.min_sigma = params["min_sigma"].item()
             self.n_gaussian = self.weight.shape[0] // 3
             self.n_coeff = self.weight.shape[1]
@@ -216,19 +258,17 @@ class GaussianMixtureNoiseModel(nn.Module):
 
     def make_learnable(self):
         print(f"[{self.__class__.__name__}] Making noise model learnable")
-
         self._learnable = True
         self.weight.requires_grad = True
 
-        #
-
     def to_device(self, cuda_tensor):
+        # TODO wtf is this ?
         # move everything to GPU
         if self.min_signal.device != cuda_tensor.device:
-            self.max_signal = self.max_signal.to(cuda_tensor.device)
-            self.min_signal = self.min_signal.to(cuda_tensor.device)
-            self.tol = self.tol.to(cuda_tensor.device)
-            self.weight = self.weight.to(cuda_tensor.device)
+            self.max_signal = self.max_signal.cuda()
+            self.min_signal = self.min_signal.cuda()
+            self.tol = self.tol.cuda()
+            # self.weight = self.weight.cuda()
             if self._learnable:
                 self.weight.requires_grad = True
 
@@ -254,21 +294,24 @@ class GaussianMixtureNoiseModel(nn.Module):
             )
         return value
 
-    def normalDens(self, x, m_=0.0, std_=None):
-        """Evaluates the normal probability density at `x` given the mean `m` and standard deviation `std`.
+    def normalDens(
+        self, x: torch.Tensor, m_: torch.Tensor = 0.0, std_: torch.Tensor = None
+    ) -> torch.Tensor:
+        """Evaluates the normal probability density at `x` given the mean `m` and
+        standard deviation `std`.
 
         Parameters
         ----------
-        x: torch.cuda.FloatTensor
-            Observations
-        m_: torch.cuda.FloatTensor
-            Mean
-        std_: torch.cuda.FloatTensor
-            Standard-deviation
+        x: torch.Tensor
+            Observations (i.e., noisy image).
+        m_: torch.Tensor
+            Pixel-wise mean.
+        std_: torch.Tensor
+            Pixel-wise standard deviation.
 
         Returns
         -------
-        tmp: torch.cuda.FloatTensor
+        tmp: torch.Tensor
             Normal probability density of `x` given `m_` and `std_`
         """
         tmp = -((x - m_) ** 2)
@@ -277,72 +320,73 @@ class GaussianMixtureNoiseModel(nn.Module):
         tmp = tmp / torch.sqrt((2.0 * np.pi) * std_ * std_)
         return tmp
 
-    def likelihood(self, observations, signals):
-        """Evaluates the likelihood of observations given the signals and the corresponding gaussian parameters.
+    def likelihood(
+        self, observations: torch.Tensor, signals: torch.Tensor
+    ) -> torch.Tensor:
+        """Evaluate the likelihood of observations given the signals and the
+        corresponding gaussian parameters.
 
         Parameters
         ----------
         observations : torch.cuda.FloatTensor
-            Noisy observations
+            Noisy observations.
         signals : torch.cuda.FloatTensor
-            Underlying signals
+            Underlying signals.
 
         Returns
         -------
         value :p + self.tol
             Likelihood of observations given the signals and the GMM noise model
         """
-        self.to_device(signals)
+        self.to_device(signals)  # move al needed stuff to the same device as `signals``
         gaussianParameters = self.getGaussianParameters(signals)
         p = 0
         for gaussian in range(self.n_gaussian):
             p += (
                 self.normalDens(
-                    observations,
-                    gaussianParameters[gaussian],
-                    gaussianParameters[self.n_gaussian + gaussian],
+                    x=observations,
+                    m_=gaussianParameters[gaussian],
+                    std_=gaussianParameters[self.n_gaussian + gaussian],
                 )
                 * gaussianParameters[2 * self.n_gaussian + gaussian]
             )
         return p + self.tol
 
-    def getGaussianParameters(self, signals):
-        """Returns the noise model for given signals
+    def getGaussianParameters(self, signals: torch.Tensor) -> list[torch.Tensor]:
+        """Returns the noise model for given signals.
 
         Parameters
         ----------
-        signals : torch.cuda.FloatTensor
+        signals : torch.Tensor
             Underlying signals
 
         Returns
         -------
-        noiseModel: list of torch.cuda.FloatTensor
-            Contains a list of `mu`, `sigma` and `alpha` for the `signals`
+        gmmParams: list[torch.Tensor]
+            A list containing tensors representing `mu`, `sigma` and `alpha`
+            parameters for the `n_gaussian` gaussians in the mixture.
 
         """
-        noiseModel = []
+        gmmParams = []
         mu = []
         sigma = []
         alpha = []
         kernels = self.weight.shape[0] // 3
         for num in range(kernels):
+            # For each Gaussian in the mixture, evaluate mean, std and weight
             mu.append(self.polynomialRegressor(self.weight[num, :], signals))
-            # expval = torch.exp(torch.clamp(self.weight[kernels + num, :], max=MAX_VAR_W))
+
             expval = torch.exp(self.weight[kernels + num, :])
-            # self.maxval = max(self.maxval, expval.max().item())
+            # TODO: why taking the exp? it is not in PPN2V paper...
             sigmaTemp = self.polynomialRegressor(expval, signals)
             sigmaTemp = torch.clamp(sigmaTemp, min=self.min_sigma)
             sigma.append(torch.sqrt(sigmaTemp))
 
-            # expval = torch.exp(
-            #     torch.clamp(
-            #         self.polynomialRegressor(self.weight[2 * kernels + num, :], signals) + self.tol, MAX_ALPHA_W))
             expval = torch.exp(
                 self.polynomialRegressor(self.weight[2 * kernels + num, :], signals)
                 + self.tol
             )
-            # self.maxval = max(self.maxval, expval.max().item())
-            alpha.append(expval)
+            alpha.append(expval)  # NOTE: these are the numerators of weights
 
         sum_alpha = 0
         for al in range(kernels):
@@ -357,24 +401,24 @@ class GaussianMixtureNoiseModel(nn.Module):
         for ker in range(kernels):
             sum_means = alpha[ker] * mu[ker] + sum_means
 
-        mu_shifted = []
         # subtracting the alpha weighted average of the means from the means
         # ensures that the GMM has the inclination to have the mean=signals.
-        # its like a residual conection. I don't understand why we need to learn the mean?
+        # TODO: I don't understand why we need to learn the mean?
         for ker in range(kernels):
             mu[ker] = mu[ker] - sum_means + signals
 
         for i in range(kernels):
-            noiseModel.append(mu[i])
+            gmmParams.append(mu[i])
         for j in range(kernels):
-            noiseModel.append(sigma[j])
+            gmmParams.append(sigma[j])
         for k in range(kernels):
-            noiseModel.append(alpha[k])
+            gmmParams.append(alpha[k])
 
-        return noiseModel
+        return gmmParams
 
+    # TODO: this is to train the noise model
     def getSignalObservationPairs(self, signal, observation, lowerClip, upperClip):
-        """Returns the Signal-Observation pixel intensities as a two-column array
+        """Returns the Signal-Observation pixel intensities as a two-column array.
 
         Parameters
         ----------
@@ -389,7 +433,7 @@ class GaussianMixtureNoiseModel(nn.Module):
 
         Returns
         -------
-        noiseModel: list of torch floats
+        gmmParams: list of torch floats
             Contains a list of `mu`, `sigma` and `alpha` for the `signals`
         """
         lb = np.percentile(signal, lowerClip)
@@ -407,3 +451,91 @@ class GaussianMixtureNoiseModel(nn.Module):
             (sig_obs_pairs[:, 0] > lb) & (sig_obs_pairs[:, 0] < ub)
         ]
         return fastShuffle(sig_obs_pairs, 2)
+
+    # TODO: what's the use of this method?
+    def forward(self, x, y):
+        """Temporary dummy forward method."""
+        return x, y
+
+    # TODO taken from pn2v. Ashesh needs to clarify this
+    def train_noise_model(
+        self,
+        signal,
+        observation,
+        learning_rate=1e-1,
+        batchSize=250000,
+        n_epochs=2000,
+        name="GMMNoiseModel.npz",
+        lowerClip=0,
+        upperClip=100,
+    ):
+        """Training to learn the noise model from signal - observation pairs.
+
+        Parameters
+        ----------
+        signal: numpy array
+            Clean Signal Data
+        observation: numpy array
+            Noisy Observation Data
+        learning_rate: float
+            Learning rate. Default = 1e-1.
+        batchSize: int
+            Nini-batch size. Default = 250000.
+        n_epochs: int
+            Number of epochs. Default = 2000.
+        name: string
+
+            Model name. Default is `GMMNoiseModel`. This model after being trained is saved at the location `path`.
+
+        lowerClip : int
+            Lower percentile for clipping. Default is 0.
+        upperClip : int
+            Upper percentile for clipping. Default is 100.
+
+
+        """
+        sig_obs_pairs = self.getSignalObservationPairs(
+            signal, observation, lowerClip, upperClip
+        )
+        counter = 0
+        optimizer = torch.optim.Adam([self.weight], lr=learning_rate)
+        for t in range(n_epochs):
+
+            jointLoss = 0
+            if (counter + 1) * batchSize >= sig_obs_pairs.shape[0]:
+                counter = 0
+                sig_obs_pairs = fastShuffle(sig_obs_pairs, 1)
+
+            batch_vectors = sig_obs_pairs[
+                counter * batchSize : (counter + 1) * batchSize, :
+            ]
+            observations = batch_vectors[:, 1].astype(np.float32)
+            signals = batch_vectors[:, 0].astype(np.float32)
+            # TODO do we absolutely need to move to GPU?
+            observations = (
+                torch.from_numpy(observations.astype(np.float32)).float().cuda()
+            )
+            signals = torch.from_numpy(signals).float().cuda()
+            p = self.likelihood(observations, signals)
+            loss = torch.mean(-torch.log(p))
+            jointLoss = jointLoss + loss
+
+            if t % 100 == 0:
+                print(t, jointLoss.item())
+
+            if t % (int(n_epochs * 0.5)) == 0:
+                trained_weight = self.weight.cpu().detach().numpy()
+                min_signal = self.min_signal.cpu().detach().numpy()
+                max_signal = self.max_signal.cpu().detach().numpy()
+                # TODO do we need to save?
+                # np.savez(self.path+name, trained_weight=trained_weight, min_signal = min_signal, max_signal = max_signal, min_sigma = self.min_sigma)
+
+            optimizer.zero_grad()
+            jointLoss.backward()
+            optimizer.step()
+            counter += 1
+
+        print("===================\n")
+        # print("The trained parameters (" + name + ") is saved at location: "+ self.path)
+        # TODO return istead of save ?
+        return trained_weight, min_signal, max_signal, self.min_sigma