reflectorch 1.0.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

reflectorch/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from reflectorch.data_generation import *
 from reflectorch.ml import *
 from reflectorch.models import *

reflectorch/data_generation/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from reflectorch.data_generation.dataset import BasicDataset, BATCH_DATA_TYPE
 from reflectorch.data_generation.priors import (
     Params,

reflectorch/data_generation/dataset.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Dict, Union
 import warnings
 

reflectorch/data_generation/likelihoods.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Union, Tuple
 
 import torch

reflectorch/data_generation/noise.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Union, Tuple
 from math import log10
 

reflectorch/data_generation/priors/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from reflectorch.data_generation.priors.params import Params
 from reflectorch.data_generation.priors.base import PriorSampler
 from reflectorch.data_generation.priors.no_constraints import BasicPriorSampler

reflectorch/data_generation/priors/base.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from torch import Tensor
 
 from reflectorch.data_generation.priors.params import Params

reflectorch/data_generation/priors/exp_subprior_sampler.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Tuple, Union, List
 
 import torch

reflectorch/data_generation/priors/independent_priors.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Union, Tuple
 from math import log
 

reflectorch/data_generation/priors/multilayer_structures.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Tuple, Dict
 
 import numpy as np

reflectorch/data_generation/priors/no_constraints.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 import logging
 from functools import lru_cache
 from typing import Tuple

reflectorch/data_generation/priors/parametric_subpriors.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Tuple, Dict, Type, List
 
 import torch

reflectorch/data_generation/priors/params.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import List, Tuple
 
 import torch

reflectorch/data_generation/priors/subprior_sampler.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from functools import lru_cache
 from typing import Tuple
 

reflectorch/data_generation/priors/utils.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Tuple
 
 import torch

reflectorch/data_generation/process_data.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Any
 
 __all__ = [

reflectorch/data_generation/q_generator.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Tuple, Union
 
 import numpy as np

reflectorch/data_generation/reflectivity/__init__.py

@@ -24,21 +24,25 @@ def reflectivity(
         log: bool = False,
         abeles_func=None,
 ):
-    """Function which computes the reflectivity from thin film parameters using the Abeles matrix formalism
+    """Function which computes the reflectivity curves from thin film parameters. 
+    By default it uses the fast implementation of the Abeles matrix formalism.
 
     Args:
-        q (Tensor): the momentum transfer (q) values
-        thickness (Tensor): the layer thicknesses
-        roughness (Tensor): the interlayer roughnesses
-        sld (Tensor): the SLDs of the layers
-        dq (Tensor, optional): the resolution for curve smearing. Defaults to None.
+        q (Tensor): tensor of momentum transfer (q) values with shape [batch_size, n_points] or [n_points]
+        thickness (Tensor): tensor containing the layer thicknesses (ordered from top to bottom) with shape [batch_size, n_layers]
+        roughness (Tensor): tensor containing the interlayer roughnesses (ordered from top to bottom) with shape [batch_size, n_layers + 1]
+        sld (Tensor): tensors containing the layer SLDs (real or complex; ordered from top to bottom) with shape [batch_size, n_layers + 1]. 
+                    It includes the substrate but excludes the ambient medium which is assumed to have an SLD of 0.
+        dq (Tensor, optional): tensor of resolutions used for curve smearing with shape [batch_size, 1].
+                            Either dq if ``constant_dq`` is ``True`` or dq/q if ``constant_dq`` is ``False``. Defaults to None.
         gauss_num (int, optional): the number of gaussians for curve smearing. Defaults to 51.
-        constant_dq (bool, optional): whether the smearing is constant. Defaults to True.
+        constant_dq (bool, optional): if ``True`` the smearing is constant (constant dq at each point in the curve) 
+                                    otherwise the smearing is linear (constant dq/q at each point in the curve). Defaults to True.
         log (bool, optional): if True the base 10 logarithm of the reflectivity curves is returned. Defaults to False.
-        abeles_func (Callable, optional): a function implementing the simulation of the reflectivity curves, if different than the default implementation. Defaults to None.
+        abeles_func (Callable, optional): a function implementing the simulation of the reflectivity curves, if different than the default Abeles matrix implementation. Defaults to None.
 
     Returns:
-        Tensor: the computed reflectivity curves
+        Tensor: tensor containing the simulated reflectivity curves with shape [batch_size, n_points]
     """
     abeles_func = abeles_func or abeles
     q = torch.atleast_2d(q)

reflectorch/data_generation/reflectivity/abeles.py

@@ -14,13 +14,14 @@ def abeles(
     """Simulates reflectivity curves for SLD profiles with box model parameterization using the Abeles matrix method
 
     Args:
-        q (Tensor): q values
-        thickness (Tensor): layer thicknesses
-        roughness (Tensor): interlayer roughnesses
-        sld (Tensor): layer SLDs
+        q (Tensor): tensor of momentum transfer (q) values with shape [batch_size, n_points] or [n_points]
+        thickness (Tensor): tensor containing the layer thicknesses (ordered from top to bottom) with shape [batch_size, n_layers]
+        roughness (Tensor): tensor containing the interlayer roughnesses (ordered from top to bottom) with shape [batch_size, n_layers + 1]
+        sld (Tensor): tensors containing the layer SLDs (real or complex; ordered from top to bottom) with shape [batch_size, n_layers + 1]. 
+                    It includes the substrate but excludes the ambient medium which is assumed to have an SLD of 0.
 
     Returns:
-        Tensor: simulated reflectivity curves
+        Tensor: tensor containing the simulated reflectivity curves with shape [batch_size, n_points]
     """
     c_dtype = torch.complex128 if q.dtype is torch.float64 else torch.complex64
 

reflectorch/data_generation/reflectivity/kinematical.py

@@ -13,6 +13,20 @@ def kinematical_approximation(
         apply_fresnel: bool = True,
         log: bool = False,
 ):
+    """Simulates reflectivity curves for SLD profiles with box model parameterization using the kinematical approximation
+
+    Args:
+        q (Tensor): tensor of momentum transfer (q) values with shape [batch_size, n_points] or [n_points]
+        thickness (Tensor): tensor containing the layer thicknesses (ordered from top to bottom) with shape [batch_size, n_layers]
+        roughness (Tensor): tensor containing the interlayer roughnesses (ordered from top to bottom) with shape [batch_size, n_layers + 1]
+        sld (Tensor): tensors containing the layer SLDs (real or complex; ordered from top to bottom) with shape [batch_size, n_layers + 1]. 
+                    It includes the substrate but excludes the ambient medium which is assumed to have an SLD of 0.
+        apply_fresnel (bool, optional): whether to use the Fresnel coefficient in the computation. Defaults to ``True``.
+        log (bool, optional): if True the base 10 logarithm of the reflectivity curves is returned. Defaults to ``False``.
+
+    Returns:
+        Tensor: tensor containing the simulated reflectivity curves with shape [batch_size, n_points]
+    """
     c_dtype = torch.complex128 if q.dtype is torch.float64 else torch.complex64
 
     batch_size, num_layers = thickness.shape

reflectorch/data_generation/reflectivity/memory_eff.py

@@ -12,6 +12,19 @@ def abeles_memory_eff(
         roughness: Tensor,
         sld: Tensor,
 ):
+    """Simulates reflectivity curves for SLD profiles with box model parameterization using a memory-efficient implementation the Abeles matrix method. 
+    It is computationally slower compared to the implementation in the 'abeles' function.
+
+    Args:
+        q (Tensor): tensor of momentum transfer (q) values with shape [batch_size, n_points] or [n_points]
+        thickness (Tensor): tensor containing the layer thicknesses (ordered from top to bottom) with shape [batch_size, n_layers]
+        roughness (Tensor): tensor containing the interlayer roughnesses (ordered from top to bottom) with shape [batch_size, n_layers + 1]
+        sld (Tensor): tensors containing the layer SLDs (real or complex; ordered from top to bottom) with shape [batch_size, n_layers + 1]. 
+                    It includes the substrate but excludes the ambient medium which is assumed to have an SLD of 0.
+
+    Returns:
+        Tensor: tensor containing the simulated reflectivity curves with shape [batch_size, n_points]
+    """
     c_dtype = torch.complex128 if q.dtype is torch.float64 else torch.complex64
 
     batch_size, num_layers = thickness.shape

reflectorch/data_generation/reflectivity/smearing.py

@@ -68,8 +68,8 @@ def _get_q_axes_for_linear_dq(q: Tensor, resolutions: Tensor, gaussnum: int = 51
     lowq = torch.clamp_min_(q.min(1).values, 1e-6)
     highq = q.max(1).values
 
-    start = torch.log10(lowq) - 6 * resolutions / _FWHM
-    end = torch.log10(highq * (1 + 6 * resolutions / _FWHM))
+    start = torch.log10(lowq)[:, None] - 6 * resolutions / _FWHM
+    end = torch.log10(highq[:, None] * (1 + 6 * resolutions / _FWHM))
 
     interpnums = torch.abs(
         (torch.abs(end - start)) / (1.7 * resolutions / _FWHM / gaussgpoint)

reflectorch/data_generation/scale_curves.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from pathlib import Path
 
 import torch

reflectorch/data_generation/smearing.py

@@ -9,8 +9,9 @@ class Smearing(object):
     The intensity at a q point will be the average of the intensities of neighbouring q points, weighted by a gaussian profile.
 
     Args:
-        sigma_range (tuple, optional): the range for sampling the standard deviation of the gaussians. Defaults to (1e-4, 5e-3).
-        constant_dq (bool, optional): whether the smearing is constant for each q point. Defaults to True.
+        sigma_range (tuple, optional): the range for sampling the resolutions. Defaults to (1e-4, 5e-3).
+        constant_dq (bool, optional): if ``True`` the smearing is constant (the resolution is given by the constant dq at each point in the curve) 
+                                    otherwise the smearing is linear (the resolution is given by the constant dq/q at each point in the curve). Defaults to True.
         gauss_num (int, optional): the number of interpolating gaussian profiles. Defaults to 31.
         share_smeared (float, optional): the share of curves in the batch for which the resolution smearing is applied. Defaults to 0.2.
     """

reflectorch/data_generation/utils.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import List, Union
 from math import sqrt, pi, log10
 

reflectorch/extensions/__init__.py

@@ -1,6 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-

reflectorch/extensions/jupyter/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from .callbacks import JPlotLoss
 
 

reflectorch/extensions/jupyter/callbacks.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from IPython.display import clear_output
 
 from ...ml import TrainerCallback, Trainer

reflectorch/extensions/matplotlib/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from reflectorch.extensions.matplotlib.losses import plot_losses
 
 __all__ = [

reflectorch/extensions/matplotlib/losses.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 import matplotlib.pyplot as plt
 
 

reflectorch/inference/__init__.py

@@ -1,4 +1,5 @@
 from reflectorch.inference.inference_model import InferenceModel, EasyInferenceModel
+from reflectorch.inference.query_matcher import HuggingfaceQueryMatcher
 from reflectorch.inference.multilayer_inference_model import MultilayerInferenceModel
 from reflectorch.inference.preprocess_exp import (
     StandardPreprocessing,
@@ -13,6 +14,7 @@ __all__ = [
     "InferenceModel",
     "EasyInferenceModel",
     "MultilayerInferenceModel",
+    "HuggingfaceQueryMatcher",
     "StandardPreprocessing",
     "standard_preprocessing",
     "ReflGradientFit",

reflectorch/inference/inference_model.py

@@ -35,15 +35,17 @@ class EasyInferenceModel(object):
         config_name (str, optional): the name of the configuration file used to initialize the model (either with or without the '.yaml' extension). Defaults to None.
         model_name (str, optional): the name of the file containing the weights of the model (either with or without the '.pt' extension), only required if different than: `'model_' + config_name + '.pt'`. Defaults to None 
         root_dir (str, optional): path to root directory containing the 'configs' and 'saved_models' subdirectories, if different from the package root directory (ROOT_DIR). Defaults to None.
+        weights_format (str, optional): format (extension) of the weights file, either 'pt' or 'safetensors'. Defaults to 'safetensors'.
         repo_id (str, optional): the id of the Huggingface repository from which the configuration files and model weights should be downloaded automatically if not found locally (in the 'configs' and 'saved_models' subdirectories of the root directory). Defaults to 'valentinsingularity/reflectivity'.
         trainer (PointEstimatorTrainer, optional): if provided, this trainer instance is used directly instead of being initialized from the configuration file. Defaults to None.
         device (str, optional): the Pytorch device ('cuda' or 'cpu'). Defaults to 'cuda'.
     """
-    def __init__(self, config_name: str = None, model_name: str = None, root_dir:str = None, repo_id: str = 'valentinsingularity/reflectivity', 
-                 trainer: PointEstimatorTrainer = None, device='cuda'):
+    def __init__(self, config_name: str = None, model_name: str = None, root_dir:str = None, weights_format: str = 'safetensors',
+                 repo_id: str = 'valentinsingularity/reflectivity', trainer: PointEstimatorTrainer = None, device='cuda'):
         self.config_name = config_name
         self.model_name = model_name
         self.root_dir = root_dir
+        self.weights_format = weights_format
         self.repo_id = repo_id
         self.trainer = trainer
         self.device = device
@@ -58,7 +60,7 @@ class EasyInferenceModel(object):
 
         Args:
             config_name (str): the name of the configuration file used to initialize the model (either with or without the '.yaml' extension).
-            model_name (str): the name of the file containing the weights of the model (either with or without the '.pt' extension), only required if different than: `'model_' + config_name + '.pt'`.
+            model_name (str): the name of the file containing the weights of the model (either with or without the '.pt' or '.safetensors'  extension), only required if different than: `'model_' + config_name + extension`.
             root_dir (str): path to root directory containing the 'configs' and 'saved_models' subdirectories, if different from the package root directory (ROOT_DIR).
         """
         if self.config_name == config_name and self.trainer is not None:
@@ -72,9 +74,10 @@ class EasyInferenceModel(object):
             self.config_name = config_name
         
         self.config_dir = Path(root_dir) / 'configs' if root_dir else CONFIG_DIR
-        self.model_name = model_name or 'model_' + config_name_no_extension + '.pt'
-        if not self.model_name.endswith('.pt'):
-            self.model_name += '.pt'
+        weights_extension = '.' + self.weights_format
+        self.model_name = model_name or 'model_' + config_name_no_extension + weights_extension
+        if not self.model_name.endswith(weights_extension):
+            self.model_name += weights_extension
         self.model_dir = Path(root_dir) / 'saved_models' if root_dir else SAVED_MODELS_DIR
 
         config_path = Path(self.config_dir) / self.config_name

reflectorch/inference/query_matcher.py

@@ -0,0 +1,82 @@
+import os
+import tempfile
+import yaml
+from huggingface_hub import hf_hub_download, list_repo_files
+
+class HuggingfaceQueryMatcher:
+    """Downloads the available configurations files to a temporary directory and provides functionality for filtering those configuration files matching user specified queries.
+    
+    Args:
+        repo_id (str): The Hugging Face repository ID.
+        config_dir (str): Directory within the repo where YAML files are stored.
+    """
+    def __init__(self, repo_id='valentinsingularity/reflectivity', config_dir='configs'):
+        self.repo_id = repo_id
+        self.config_dir = config_dir
+        self.cache = {
+            'parsed_configs': None,
+            'temp_dir': None
+        }
+        self._renew_cache()
+
+    def _renew_cache(self):
+        temp_dir = tempfile.mkdtemp()
+        print(f"Temporary directory created at: {temp_dir}")
+
+        repo_files = list_repo_files(self.repo_id, repo_type='model')
+        config_files = [file for file in repo_files if file.startswith(self.config_dir) and file.endswith('.yaml')]
+
+        downloaded_files = []
+        for file in config_files:
+            file_path = hf_hub_download(repo_id=self.repo_id, filename=file, local_dir=temp_dir, repo_type='model')
+            downloaded_files.append(file_path)
+        
+        parsed_configs = {}
+        for file_path in downloaded_files:
+            with open(file_path, 'r') as file:
+                config_data = yaml.safe_load(file)
+                file_name = os.path.basename(file_path)
+                parsed_configs[file_name] = config_data
+        
+        self.cache['parsed_configs'] = parsed_configs
+        self.cache['temp_dir'] = temp_dir
+
+    def get_matching_configs(self, query):
+        """retrieves configuration files that match the user specified query.
+        
+        Args:
+            query (dict): Dictionary of key-value pairs to filter configurations, e.g. ``query = {'dset.prior_sampler.kwargs.max_num_layers': 3, 'dset.prior_sampler.kwargs.param_ranges.slds': [0., 100.]}``.
+                          For keys containing the ``param_ranges`` subkey a configuration is selected if the value of the query (i.e. desired parameter range) 
+                          is a subrange of the parameter range in the configuration, in all other cases the values must match exactly.
+        
+        Returns:
+            list: List of file names that match the query.
+        """
+        
+        filtered_configs = []
+        
+        for file_name, config_data in self.cache['parsed_configs'].items():
+            if self.matches_query(config_data, query):
+                filtered_configs.append(file_name)
+        
+        return filtered_configs
+    
+    def matches_query(self, config_data, query):
+        for q_key, q_value in query.items():
+            keys = q_key.split('.')
+            value = self.deep_get(config_data, keys)
+            if 'param_ranges' in keys:
+                if q_value[0] < value[0] or q_value[1] > value[1]:
+                    return False
+            else: 
+                if value != q_value:
+                    return False
+                
+        return True
+
+    def deep_get(self, d, keys):
+        for key in keys:
+            if isinstance(d, dict):
+                d = d.get(key, None)
+
+        return d

reflectorch/ml/__init__.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from reflectorch.ml.basic_trainer import *
 from reflectorch.ml.callbacks import *
 from reflectorch.ml.trainers import *
@@ -32,6 +26,5 @@ __all__ = [
     'MultilayerDataLoader',
     'RealTimeSimTrainer',
     'DenoisingAETrainer',
-    'VAETrainer',
     'PointEstimatorTrainer',
 ]

reflectorch/ml/basic_trainer.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from typing import Optional, Tuple, Iterable, Any, Union, Type
 from collections import defaultdict
 

reflectorch/ml/callbacks.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 import torch
 
 import numpy as np

reflectorch/ml/loggers.py

@@ -1,10 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
-
 __all__ = [
     'Logger',
     'Loggers',

reflectorch/ml/schedulers.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 from torch.optim import lr_scheduler
 
 import numpy as np

reflectorch/ml/trainers.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -16,7 +10,6 @@ from reflectorch.ml.dataloaders import ReflectivityDataLoader
 __all__ = [
     'RealTimeSimTrainer',
     'DenoisingAETrainer',
-    'VAETrainer',
     'PointEstimatorTrainer',
 ]
 
@@ -97,30 +90,4 @@ class DenoisingAETrainer(RealTimeSimTrainer):
         restored_curves = self.model(scaled_noisy_curves)
         loss = self.criterion(scaled_curves, restored_curves)
         return {'loss': loss}
-    
-
-class VAETrainer(DenoisingAETrainer):
-    """Trainer which can be used for training a denoising autoencoder model. Overrides _get_batch and get_loss_dict methods """
-    def init(self):
-        self.loader.calc_denoised_curves = True
-        self.freebits = 0.05
-    
-    def calc_kl(self, z_mu, z_logvar):
-        return 0.5*(z_mu**2 + torch.exp(z_logvar) - 1 - z_logvar)
-    
-    def gaussian_log_prob(self, z, mu, logvar):
-        return -0.5*(np.log(2*np.pi) + logvar + (z-mu)**2/torch.exp(logvar))
-
-    def get_loss_dict(self, batch_data):
-        """returns the reconstruction loss of the autoencoder"""
-        scaled_noisy_curves, scaled_curves = batch_data
-        _, (z_mu, z_logvar, restored_curves_mu, restored_curves_logvar) = self.model(scaled_noisy_curves)
-
-        l_rec = -torch.mean(self.gaussian_log_prob(scaled_curves, restored_curves_mu, restored_curves_logvar), dim=-1)  
-        l_kl = torch.mean(F.relu(self.calc_kl(z_mu, z_logvar) - self.freebits*np.log(2)) + self.freebits*np.log(2), dim=-1)
-        loss = torch.mean(l_rec + l_kl)/np.log(2)
-
-        l_rec = torch.mean(l_rec)
-        l_kl = torch.mean(l_kl)
-
-        return {'loss': loss}
+    

reflectorch/ml/utils.py

@@ -1,9 +1,2 @@
-# -*- coding: utf-8 -*-
-#
-#
-# This source code is licensed under the GPL license found in the
-# LICENSE file in the root directory of this source tree.
-
-
 def is_divisor(num: int, div: int):
     return num and not num % div