reflectorch 1.1.0__py3-none-any.whl → 1.2.1__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): tensor containing the layer SLDs (real or complex; ordered from top to bottom) with shape 
+            [batch_size, n_layers + 1] (excluding ambient SLD which is assumed to be 0)  or [batch_size, n_layers + 2] (including ambient SLD; only for the default ``abeles_func='abeles'``) 
+        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 ('abeles'). 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

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import math
+from functools import reduce
 
 import torch
 from torch import Tensor
@@ -14,23 +15,38 @@ 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): tensor containing the layer SLDs (real or complex; ordered from top to bottom). The tensor shape should be one of the following:
+            - [batch_size, n_layers + 1]: in this case, the ambient SLD is not included but assumed to be 0
+            - [batch_size, n_layers + 2]: this shape includes the ambient SLD as the first element in the tensor
 
     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
 
     batch_size, num_layers = thickness.shape
 
-    sld = torch.cat([torch.zeros(batch_size, 1).to(sld), sld], -1)[:, None]
-    thickness = torch.cat([torch.zeros(batch_size, 1).to(thickness), thickness], -1)[:, None]
+    if sld.shape[-1] == num_layers + 1:
+        # add zero ambient sld
+        sld = torch.cat([torch.zeros(batch_size, 1).to(sld), sld], -1)
+    if sld.shape[-1] != num_layers + 2:
+        raise ValueError(
+            "Number of SLD values does not equal to num_layers + 2 (substrate + ambient)."
+        )
+
+    sld = sld[:, None]
+
+    # add zero thickness for ambient layer:
+    thickness = torch.cat([torch.zeros(batch_size, 1).to(thickness), thickness], -1)[
+        :, None
+    ]
+
     roughness = roughness[:, None] ** 2
 
-    sld = sld * 1e-6 + 1e-30j
+    sld = (sld - sld[..., :1]) * 1e-6 + 1e-36j
 
     k_z0 = (q / 2).to(c_dtype)
 
@@ -40,7 +56,7 @@ def abeles(
     if k_z0.dim() == 2:
         k_z0.unsqueeze_(-1)
 
-    k_n = torch.sqrt(k_z0 ** 2 - 4 * math.pi * sld)
+    k_n = torch.sqrt(k_z0**2 - 4 * math.pi * sld)
 
     # k_n.shape - (batch, q, layers)
 
@@ -51,22 +67,22 @@ def abeles(
     exp_beta = torch.exp(beta)
     exp_m_beta = torch.exp(-beta)
 
-    rn = (k_n - k_np1) / (k_n + k_np1) * torch.exp(- 2 * k_n * k_np1 * roughness)
+    rn = (k_n - k_np1) / (k_n + k_np1) * torch.exp(-2 * k_n * k_np1 * roughness)
 
-    c_matrices = torch.stack([
-        torch.stack([exp_beta, rn * exp_m_beta], -1),
-        torch.stack([rn * exp_beta, exp_m_beta], -1),
-    ], -1)
+    c_matrices = torch.stack(
+        [
+            torch.stack([exp_beta, rn * exp_m_beta], -1),
+            torch.stack([rn * exp_beta, exp_m_beta], -1),
+        ],
+        -1,
+    )
 
     c_matrices = [c.squeeze(-3) for c in c_matrices.split(1, -3)]
 
-    m, c_matrices = c_matrices[0], c_matrices[1:]
-
-    for c in c_matrices:
-        m = m @ c
+    m = reduce(torch.matmul, c_matrices)
 
     r = (m[..., 1, 0] / m[..., 0, 0]).abs() ** 2
-    r = torch.clamp_max_(r, 1.)
+    r = torch.clamp_max_(r, 1.0)
 
     return r
 

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): tensor 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): tensor 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/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
@@ -351,7 +354,8 @@ class EasyInferenceModel(object):
                     torch.from_numpy(polished_params_arr[:-1][None]),
                     torch.from_numpy(priors.T[0][None]),
                     torch.from_numpy(priors.T[1][None]),
-                    #self.trainer.loader.prior_sampler.param_model
+                    self.trainer.loader.prior_sampler.max_num_layers,
+                    self.trainer.loader.prior_sampler.param_model
                     )
             else:
                 polished_params_arr, curve_polished = standard_refl_fit(
@@ -363,7 +367,8 @@ class EasyInferenceModel(object):
                     torch.from_numpy(polished_params_arr[None]),
                     torch.from_numpy(priors.T[0][None]),
                     torch.from_numpy(priors.T[1][None]),
-                    #self.trainer.loader.prior_sampler.param_model
+                    self.trainer.loader.prior_sampler.max_num_layers,
+                    self.trainer.loader.prior_sampler.param_model
                     )
         except Exception as err:
             polished_params = predicted_params

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

reflectorch/models/__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.models.encoders import *
 from reflectorch.models.networks import *
 
@@ -12,7 +6,6 @@ __all__ = [
     "ConvDecoder",
     "ConvAutoencoder",
     "ConvVAE",
-    "TransformerEncoder",
     "FnoEncoder",
     "SpectralConv1d",
     "ConvResidualNet1D",

reflectorch/models/encoders/__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.models.encoders.conv_encoder import (
     ConvEncoder,
     ConvDecoder,
@@ -11,12 +5,10 @@ from reflectorch.models.encoders.conv_encoder import (
     ConvVAE,
 )
 from reflectorch.models.encoders.fno import FnoEncoder, SpectralConv1d
-from reflectorch.models.encoders.transformers import TransformerEncoder
 from reflectorch.models.encoders.conv_res_net import ConvResidualNet1D
 
 
 __all__ = [
-    "TransformerEncoder",
     "ConvEncoder",
     "ConvDecoder",
     "ConvAutoencoder",

reflectorch/models/encoders/conv_encoder.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 pathlib import Path
 

reflectorch/models/encoders/conv_res_net.py

@@ -1,8 +1,4 @@
-# -*- 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 nn
 from torch.nn import functional as F

reflectorch/models/networks/__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.models.networks.mlp_networks import (
     NetworkWithPriorsConvEmb,
     NetworkWithPriorsFnoEmb,

reflectorch/paths.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
 from pathlib import Path
 

reflectorch/runs/__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.runs.train import (
     run_train,
     run_train_on_cluster,
@@ -17,6 +11,7 @@ from reflectorch.runs.utils import (
     get_callbacks_from_config,
     get_trainer_by_name,
     get_callbacks_by_name,
+    convert_pt_to_safetensors,
 )
 
 from reflectorch.runs.config import load_config
@@ -31,5 +26,6 @@ __all__ = [
     'get_callbacks_from_config',
     'get_trainer_by_name',
     'get_callbacks_by_name',
+    'convert_pt_to_safetensors',
     'load_config',
 ]

reflectorch/runs/config.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 yaml
 
 from pathlib import Path

reflectorch/runs/slurm_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, Union
 from pathlib import Path
 import subprocess

reflectorch/runs/train.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 click
 
 from reflectorch.runs.slurm_utils import save_sbatch_and_run

reflectorch/runs/utils.py

@@ -2,6 +2,8 @@ from pathlib import Path
 from typing import Tuple
 
 import torch
+import safetensors.torch
+import os
 
 from reflectorch import *
 from reflectorch.runs.config import load_config
@@ -13,6 +15,7 @@ __all__ = [
     "get_callbacks_from_config",
     "get_trainer_by_name",
     "get_callbacks_by_name",
+    "convert_pt_to_safetensors",
 ]
 
 
@@ -214,15 +217,25 @@ def get_trainer_by_name(config_name, config_dir=None, model_path=None, load_weig
         model_name = f'model_{config_name}.pt'
         model_path = SAVED_MODELS_DIR / model_name
 
-    try:
-        state_dict = torch.load(model_path, map_location=inference_device)
-    except Exception as err:
-        raise RuntimeError(f'Could not load model from {model_path}') from err
-
-    if 'model' in state_dict:
-        trainer.model.load_state_dict(state_dict['model'])
+    if str(model_path).endswith('.pt'):
+        try:
+            state_dict = torch.load(model_path, map_location=inference_device)
+        except Exception as err:
+            raise RuntimeError(f'Could not load model from {model_path}') from err
+
+        if 'model' in state_dict:
+            trainer.model.load_state_dict(state_dict['model'])
+        else:
+            trainer.model.load_state_dict(state_dict)
+
+    elif str(model_path).endswith('.safetensors'):
+        try:
+            load_state_dict_safetensors(model=trainer.model, filename=model_path, device=inference_device)
+        except Exception as err:
+            raise RuntimeError(f'Could not load model from {model_path}') from err
+        
     else:
-        trainer.model.load_state_dict(state_dict)
+        raise RuntimeError('Weigths file with unknown extension')
 
     return trainer
 
@@ -297,3 +310,59 @@ def init_dset(config: dict):
     )
 
     return dset
+
+def split_complex_tensors(state_dict):
+    new_state_dict = {}
+    for key, tensor in state_dict.items():
+        if tensor.is_complex():
+            new_state_dict[f"{key}_real"] = tensor.real.clone()
+            new_state_dict[f"{key}_imag"] = tensor.imag.clone()
+        else:
+            new_state_dict[key] = tensor
+    return new_state_dict
+
+def recombine_complex_tensors(state_dict):
+    new_state_dict = {}
+    keys = list(state_dict.keys())
+    visited = set()
+
+    for key in keys:
+        if key.endswith('_real') or key.endswith('_imag'):
+            base_key = key[:-5]
+            new_state_dict[base_key] = torch.complex(state_dict[base_key + '_real'], state_dict[base_key + '_imag'])
+            visited.add(base_key + '_real')
+            visited.add(base_key + '_imag')
+        elif key not in visited:
+            new_state_dict[key] = state_dict[key]
+
+    return new_state_dict
+
+def convert_pt_to_safetensors(input_dir):
+    """Creates '.safetensors' files for all the model state dictionaries inside '.pt' files in the specified directory.
+
+    Args:
+        input_dir (str): directory containing model weights 
+    """
+    if not os.path.isdir(input_dir):
+        raise ValueError(f"Input directory {input_dir} does not exist")
+
+    for file_name in os.listdir(input_dir):
+        if file_name.endswith('.pt'):
+            pt_file_path = os.path.join(input_dir, file_name)
+            safetensors_file_path = os.path.join(input_dir, file_name[:-3] + '.safetensors')
+
+            if os.path.exists(safetensors_file_path):
+                print(f"Skipping {pt_file_path}, corresponding .safetensors file already exists.")
+                continue
+
+            print(f"Converting {pt_file_path} to .safetensors format.")
+            data_pt = torch.load(pt_file_path)
+            model_state_dict = data_pt["model"]
+            model_state_dict = split_complex_tensors(model_state_dict) #handle tensors with complex dtype which are not natively supported by safetensors
+
+            safetensors.torch.save_file(tensors=model_state_dict, filename=safetensors_file_path)
+
+def load_state_dict_safetensors(model, filename, device):
+    state_dict = safetensors.torch.load_file(filename=filename, device=device)
+    state_dict = recombine_complex_tensors(state_dict)
+    model.load_state_dict(state_dict)

reflectorch/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.
-
 import numpy as np
 from numpy import ndarray
 

{reflectorch-1.1.0.dist-info → reflectorch-1.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: reflectorch
-Version: 1.1.0
+Version: 1.2.1
 Summary: A Pytorch-based package for the analysis of reflectometry data
 Author-email: Vladimir Starostin <vladimir.starostin@uni-tuebingen.de>, Valentin Munteanu <valentin.munteanu@uni-tuebingen.de>
 Maintainer-email: Valentin Munteanu <valentin.munteanu@uni-tuebingen.de>, Vladimir Starostin <vladimir.starostin@uni-tuebingen.de>, Alexander Hinderhofer <alexander.hinderhofer@uni-tuebingen.de>
@@ -27,6 +27,7 @@ Requires-Dist: click
 Requires-Dist: matplotlib
 Requires-Dist: ipywidgets
 Requires-Dist: huggingface-hub
+Requires-Dist: safetensors
 Provides-Extra: build
 Requires-Dist: build ; extra == 'build'
 Requires-Dist: twine ; extra == 'build'

reflectorch-1.2.1.dist-info/RECORD

@@ -0,0 +1,83 @@
+reflectorch/__init__.py,sha256=N98TX0LJtHBT8Q8MbUedjVlHSQJXii0EMmwuhSKHvOQ,736
+reflectorch/paths.py,sha256=Z_VRVknkqRn03ULShc8YCp0czqqMJ9w2CRuUi8e2OV0,814
+reflectorch/test_config.py,sha256=1T7pMJ-WYLEu-4WtYMQxcJrqXvgdpvJ1yi2qINd0kNA,99
+reflectorch/train.py,sha256=-c8ac1fpjrCiEwnAaXg_rcBNl1stO1V5p5afx_78xHs,87
+reflectorch/train_on_cluster.py,sha256=aG3_g5_rzL8iL1tvtdY9ueJTo1f2Pn8lGJgudrrRknU,109
+reflectorch/utils.py,sha256=ehG5dU02_WIFPIVGPsSjHl3Ji10MQ8-jsCwRkwUO7D4,1993
+reflectorch/data_generation/__init__.py,sha256=TcF7kf16GCsQpp6cIXj-s4vb_gjrtv7tgXwLSeI8Xy8,3353
+reflectorch/data_generation/dataset.py,sha256=wEHMIzA1XQGQkzkN8WQRNqG53lyzxIVc80nMJKs0mgI,7273
+reflectorch/data_generation/likelihoods.py,sha256=gnqmsEfsZnCC6WuPcIe9rFl4SeiLv_LyNlTVM-5YTE8,2862
+reflectorch/data_generation/noise.py,sha256=c4ytClId3t3T5FO8r0NqAVz-x_zYWAJ2VkpxcvssrVY,16159
+reflectorch/data_generation/process_data.py,sha256=iBvWB4_X5TTIZTacsq9baBZPyGCRyI1G-J3l1BVMuD4,1180
+reflectorch/data_generation/q_generator.py,sha256=Etlp_6Tj2it2nK8PQSWgQZBLh8_459P985GrZs9U1lE,8529
+reflectorch/data_generation/scale_curves.py,sha256=hNZoGkA9FMdLUarp0fwj5nCDYvpy9NbSiO85_rCaNr8,4084
+reflectorch/data_generation/smearing.py,sha256=MfHQ2sa-wPu5stO835m8i1-3NnJYvxS3qpzNs1jXWAE,3109
+reflectorch/data_generation/utils.py,sha256=3Di6I5Ihy2rm8NMeCEvIVqH0IXFEX2bcxkHSc6y9WgU,5607
+reflectorch/data_generation/priors/__init__.py,sha256=ZUaQUgNR44MQGYcPVJaSzzoE270710RqdKAPcaBNHWo,1941
+reflectorch/data_generation/priors/base.py,sha256=JNa2A3F4IWaEwV7gyTQjVdCRevUI43UUpTOm1U3XA8k,1705
+reflectorch/data_generation/priors/exp_subprior_sampler.py,sha256=hjHey_32HGgZ4ojOAhmmx8jcQPCyvcPRgxMlfthSdOo,11551
+reflectorch/data_generation/priors/independent_priors.py,sha256=ZdFCW4Ea6cK9f0Pnk2F-F6JLOgZSMsc77ZWrbQl1gkE,7241
+reflectorch/data_generation/priors/multilayer_models.py,sha256=lAf-HJPbIDnbD1ecDTvx03TfA4jLN5tTLbwaBCiYgWM,7763
+reflectorch/data_generation/priors/multilayer_structures.py,sha256=b1LTkzMK_R2hgbAmif96zRRzZRV2V-dD7eSTaNalMU8,3693
+reflectorch/data_generation/priors/no_constraints.py,sha256=qioTThJ17NixCaMIIf4dusPtvcK_uI2jOoysyXKnkZ4,7292
+reflectorch/data_generation/priors/parametric_models.py,sha256=vaAGcm9Ky-coidNliE4R1YvoqCcbNwEwJVCwwmkKI4E,25139
+reflectorch/data_generation/priors/parametric_subpriors.py,sha256=mN4MELBOQurNponknB_1n46fiNncXkI5PnnPL3hJee4,14352
+reflectorch/data_generation/priors/params.py,sha256=JpH-LRxTztc0QP-3QwjLdu0MsVc1rSixtcscYGs_2Ew,8238
+reflectorch/data_generation/priors/sampler_strategies.py,sha256=U-v5dXpFLJq939aQKvNl7n1Lih-y97K8WJCXBFRQiA0,13059
+reflectorch/data_generation/priors/scaler_mixin.py,sha256=gI64v2KOZugSJWaLKASR34fn6qVFl-aoeVA1BR5yXNg,2648
+reflectorch/data_generation/priors/subprior_sampler.py,sha256=TE8DOQhxVr69VmWSjwHyElpgVjOhZKBIPChFWNVYzRc,14769
+reflectorch/data_generation/priors/utils.py,sha256=bmIZYHq95gG68kETfNH6ygR39oitUEJ0eCO_Fb8maH0,3836
+reflectorch/data_generation/reflectivity/__init__.py,sha256=MxVtNqO52iORdual2TtmSRN41TC9VmVSFH5ov-wuAJc,3202
+reflectorch/data_generation/reflectivity/abeles.py,sha256=Pcov-m0HImw1DOR57t0OsBOLZs3MA6fk7RePf_VCEZU,3097
+reflectorch/data_generation/reflectivity/kinematical.py,sha256=78E128l3DQif_vz2Ae6hjRg-69kRz7cbN0Ak5dOc5PM,2625
+reflectorch/data_generation/reflectivity/memory_eff.py,sha256=pbsXkcGwDorey6qVx2SlBK6IlzNJY1FeeF84cEJB_hY,4023
+reflectorch/data_generation/reflectivity/numpy_implementations.py,sha256=QBzn4yVnOdlkHeeR-ZFPS115GnLdO9lMTGO2d3YhG9I,3177
+reflectorch/data_generation/reflectivity/smearing.py,sha256=pc95Lig9NIWtHDCKnKbLp7G_kmL7_3YB6ZpDphgC2D8,4001
+reflectorch/extensions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+reflectorch/extensions/jupyter/__init__.py,sha256=bs6RQQ8LRcKAJu2I9K8BZEMQSll3kLdirRvxLbgsY9w,72
+reflectorch/extensions/jupyter/callbacks.py,sha256=K4Z8Mlm_Pa80YDsNOdMZLIJZyVrg09kBVzVdFHKgQFk,1074
+reflectorch/extensions/matplotlib/__init__.py,sha256=ARmc228xbZUj4fWNq5fVxZ8PN444fjDNxNP4ytVCyfo,104
+reflectorch/extensions/matplotlib/losses.py,sha256=H2vrjNqzmQ0BP4377uN4uhlH88aS5OAfA7pwRM2WHqA,686
+reflectorch/inference/__init__.py,sha256=i0KNn83XN33mLrV7bpHdLd0SXxuGCKfbQcIoa247Uts,834
+reflectorch/inference/inference_model.py,sha256=uxzdL2S01CN6lN737B4x5L_NIFnaCkHrNUQKwh4xNW0,36739
+reflectorch/inference/multilayer_fitter.py,sha256=0CxDpLOEp1terR4N39yFlxhvA8qAbHf_01NbmvYadck,5510
+reflectorch/inference/multilayer_inference_model.py,sha256=hH_-dJGdMOox8GHXdM_nODXDlNgh_v449xW5FmklRdo,7575
+reflectorch/inference/query_matcher.py,sha256=Dk49dW0XreeCjufzYBTKchfTdVbG6759ryV6I-wQL60,3387
+reflectorch/inference/record_time.py,sha256=3er-aoR8Sd_Kc4qNwUmRqkEz4FYhVxdi1ARnBohybzM,1140
+reflectorch/inference/sampler_solution.py,sha256=DeJM3EXEb6S5EiASj3mmdNI-Y07Cr5UzzA5oq-vEB-Q,2288
+reflectorch/inference/scipy_fitter.py,sha256=339M33OdmfgOpifJGLYk4KVcnnNJrY6_aH7Lz6Vtt24,5404
+reflectorch/inference/torch_fitter.py,sha256=j1NzkzLCmQ4H6LfIi82LsSBmIdunnWzm3kbGx-hqvDs,3391
+reflectorch/inference/preprocess_exp/__init__.py,sha256=bR6H-xgBD96z4P9P1T2ECnWvalrimdMTfTNArIWPLy0,383
+reflectorch/inference/preprocess_exp/attenuation.py,sha256=UKDfUjCKKMgAuEs5Ttyo0KEQmpvHZI52UgVflh7T81A,1518
+reflectorch/inference/preprocess_exp/cut_with_q_ratio.py,sha256=CbtwIw7iJNkoVxqTHKzONBgGFwOsCUyFoTIQ8XMLTfY,1149
+reflectorch/inference/preprocess_exp/footprint.py,sha256=xc409M5X-QW0Ce_6dEZdj8NkOY1vd0LaGpPQFxiOOR0,2625
+reflectorch/inference/preprocess_exp/interpolation.py,sha256=o2v-mlfRYzeaaikeQVeY7EA7j-K42dMfT12oN3mk51k,694
+reflectorch/inference/preprocess_exp/normalize.py,sha256=09v7nZdtw6SnW_67xFrPnqzOA5AtFBGarjK4Pfn4VoE,695
+reflectorch/inference/preprocess_exp/preprocess.py,sha256=pyyq8fSvcm1bWAShzGHYnKOc55Rofh4FIw1AC7Smq-U,5111
+reflectorch/ml/__init__.py,sha256=TJERkE_itNOH3GwtC-9Xv0rZ70qEkMukFLD1qXsMdOQ,730
+reflectorch/ml/basic_trainer.py,sha256=2z_3Iy_9dEgOha8da69RjiMQ_89C2SoJM9omh0WO-ek,9647
+reflectorch/ml/callbacks.py,sha256=blZNFubtRQkcx5sQrTNkARiPj75T8VdCn6CYJNN7hKg,2743
+reflectorch/ml/dataloaders.py,sha256=IvKmsH5gX_b-00KRFeL-x3keEfBcvYkQFWGWd8Caj-I,1073
+reflectorch/ml/loggers.py,sha256=oYZskGMbbWfW3sOXOC-4F-DIL-cLhYrRFeggrmPtPGM,743
+reflectorch/ml/schedulers.py,sha256=AO-dS1bZgn7p-gGJJhocmL6Vc8XLMwynZxvko4BCEsw,10048
+reflectorch/ml/trainers.py,sha256=DXO0_ue3iVQKi8vDttEQZbNLcn3xwkRBIFrw0XTomn4,3643
+reflectorch/ml/utils.py,sha256=_OKkc6o5od7vnfBNZKWsGApJxA62TrAlLok92s0nG4k,71
+reflectorch/models/__init__.py,sha256=vsGXi1BoLVyki_OF_Ezv9GXCpuJ22liGwolVZ0rxyyM,335
+reflectorch/models/activations.py,sha256=LDiIxCnLFb8r_TRBZSt6vdOZmexCWAGa5DfE_SotUL8,1431
+reflectorch/models/encoders/__init__.py,sha256=X9cHeWjJGVXNFwii9ZJasNbqrV4kt9QyJYPdKfxgp04,443
+reflectorch/models/encoders/conv_encoder.py,sha256=Xa9Yo2lDctFWlnz8vIIdzwakcE_cgBR_J92UL59XhbA,7653
+reflectorch/models/encoders/conv_res_net.py,sha256=VP1CCKdNqsva2R4l-IV6d63mQk6bI4Aresfd6b6_YKU,3343
+reflectorch/models/encoders/fno.py,sha256=s_S7hnpLE7iGfyvnQ-QvTh0rKO5KFiy5tUYau4sJbvI,4693
+reflectorch/models/networks/__init__.py,sha256=zHUvrlb0KVOpRrwVjwjR1g8sVWX6VH80FLGqu4jiux8,291
+reflectorch/models/networks/mlp_networks.py,sha256=C7py6qCBVaYYt0FMEf8gbT4lndArKpUYYgTN1001-T8,11614
+reflectorch/models/networks/residual_net.py,sha256=msDJaDw7qL9ebEW1Avw6Qw0lgni68AMgF4kXiJKzeaQ,4637
+reflectorch/runs/__init__.py,sha256=xjuZGjqZuEovlpe9Jj5d8Nn5ii-5jvAYdeHT5oYaYVI,723
+reflectorch/runs/config.py,sha256=8YtUOXr_DvNvgpu59CNAr3KrQijx1AGDE95gYrsuwsM,804
+reflectorch/runs/slurm_utils.py,sha256=Zyj4_K5YpiWNJhgpFLWYHsSaaI-mgVEWsN15Gd_BhI0,2600
+reflectorch/runs/train.py,sha256=e-Jj0fwYlUB2NLDxCzy0cLsSrJbNqd5pN5T1L7-Eiig,2560
+reflectorch/runs/utils.py,sha256=bp3Nwd6pfX5yFlTVX28zDtSTo2gg2JZh9HCpZpvnJWk,12637
+reflectorch-1.2.1.dist-info/LICENSE.txt,sha256=2kX9kLKiIRiQRqUXwk3J-Ba3fqmztNu8ORskLBlAuKM,1098
+reflectorch-1.2.1.dist-info/METADATA,sha256=ANjKriMGr9vTl5pvhgqpsCL40aZxh1ryKrfy66Eh_wU,7805
+reflectorch-1.2.1.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+reflectorch-1.2.1.dist-info/top_level.txt,sha256=2EyIWrt4SeZ3hNadLXvEVpPFhyoZ4An7YflP4y_E3Fc,12
+reflectorch-1.2.1.dist-info/RECORD,,

{reflectorch-1.1.0.dist-info → reflectorch-1.2.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (71.0.4)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

reflectorch/models/encoders/transformers.py

@@ -1,56 +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.
-
-
-import torch
-from torch import nn
-
-
-class TransformerEncoder(nn.Module):
-    def __init__(
-            self,
-            dim: int = 64,
-            nhead: int = 8,
-            num_encoder_layers: int = 4,
-            num_decoder_layers: int = 2,
-            dim_feedforward: int = 512,
-            dropout: float = 0.01,
-            activation: str = 'gelu',
-            in_dim: int = 2,
-            out_dim: int = None,
-    ):
-
-        super().__init__()
-
-        self.in_projector = nn.Linear(in_dim, dim)
-
-        self.dim = dim
-
-        self.transformer = nn.Transformer(
-            dim, nhead=nhead,
-            num_encoder_layers=num_encoder_layers,
-            num_decoder_layers=num_decoder_layers,
-            dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            activation=activation
-        )
-
-        if out_dim:
-            self.out_projector = nn.Linear(dim, out_dim)
-        else:
-            self.out_projector = None
-
-    def forward(self, src, tgt, src_key_padding_mask=None, tgt_key_padding_mask=None, **kwargs):
-        src = self.in_projector(src.transpose(1, 2)).transpose(0, 1)
-
-        res = self.transformer(
-            src, tgt, src_key_padding_mask=src_key_padding_mask, tgt_key_padding_mask=tgt_key_padding_mask, **kwargs
-        )
-
-        if self.out_projector:
-            res = self.out_projector(res).squeeze(-1)
-
-        return res.squeeze(0)

reflectorch-1.1.0.dist-info/RECORD

@@ -1,84 +0,0 @@
-reflectorch/__init__.py,sha256=Xs-ZhOJ1D_yluLD98aj6Pr2ssy-tDFwtCnIIukDFOW4,895
-reflectorch/paths.py,sha256=8xAtonoxLxfwnZF4rriUPidPaaz9pNvOEibFXiXu9yc,973
-reflectorch/test_config.py,sha256=1T7pMJ-WYLEu-4WtYMQxcJrqXvgdpvJ1yi2qINd0kNA,99
-reflectorch/train.py,sha256=-c8ac1fpjrCiEwnAaXg_rcBNl1stO1V5p5afx_78xHs,87
-reflectorch/train_on_cluster.py,sha256=aG3_g5_rzL8iL1tvtdY9ueJTo1f2Pn8lGJgudrrRknU,109
-reflectorch/utils.py,sha256=TIRilK7px-AAw3MWNmhBPC8dzq7RI5YbkNbEcqDvrNs,2152
-reflectorch/data_generation/__init__.py,sha256=7lbWlxj4tMuhgTzz_lFgrRkYp2AVJm262ZLoJx_Riyw,3512
-reflectorch/data_generation/dataset.py,sha256=O8R_RaRBXFYfvsI3co67fDeclEG1pMO50RXZTIPUf_4,7432
-reflectorch/data_generation/likelihoods.py,sha256=mXplS5nwoH4sAeHqp3ciEf--GMA3ulPsnxasix0HMn0,3021
-reflectorch/data_generation/noise.py,sha256=cZDTYPlteIclEhnRcV5DNKG64cvz-Q4VHiWGGdZkeDk,16318
-reflectorch/data_generation/process_data.py,sha256=0kpVWqkpDkCDEybIu2uMYdr-ytfT0sVWKfY-BbXVo9c,1339
-reflectorch/data_generation/q_generator.py,sha256=8_TLXamHF-5Lsr1g3WgAd6gcNuQgXPZ9hSJNazY247A,8688
-reflectorch/data_generation/scale_curves.py,sha256=vMiyq9Y3S5nFaf5RA9g_Id96lrcr4D_O0s_Amz_3fWM,4243
-reflectorch/data_generation/smearing.py,sha256=UqcWBATSyvds-P_Soq0hUfBipe3VJstBE88m_v0z-rc,2929
-reflectorch/data_generation/utils.py,sha256=aBICDneVaZhlEqCyMJAzuDXbCe0kLUyJ_9VHxUTBrao,5766
-reflectorch/data_generation/priors/__init__.py,sha256=5tN-0V8GduS8N1riZ6sDUXp_Wr7WM6840tlka4SsFqU,2100
-reflectorch/data_generation/priors/base.py,sha256=y7e6AWxbMekw5MHtW_h3_VhmzcGS4T6hnhqso4p2MDA,1864
-reflectorch/data_generation/priors/exp_subprior_sampler.py,sha256=WPjc7aGLr_qJiiNHA1yfKPd1EHJKcmOIixkbhhy7QKg,11710
-reflectorch/data_generation/priors/independent_priors.py,sha256=T4_uvX74iHScPGa5u_-fvdwurUu97AuRssfDL3tYxKY,7400
-reflectorch/data_generation/priors/multilayer_models.py,sha256=lAf-HJPbIDnbD1ecDTvx03TfA4jLN5tTLbwaBCiYgWM,7763
-reflectorch/data_generation/priors/multilayer_structures.py,sha256=nGVoMstkn--0kdKlT5o3VsGZ0dHUdMoeSXeI9t9V61Q,3852
-reflectorch/data_generation/priors/no_constraints.py,sha256=jLzsKXyQFHW2dtKViCpJpEbNjk2njjWpOaUxD8Hc3wE,7451
-reflectorch/data_generation/priors/parametric_models.py,sha256=vaAGcm9Ky-coidNliE4R1YvoqCcbNwEwJVCwwmkKI4E,25139
-reflectorch/data_generation/priors/parametric_subpriors.py,sha256=hSuSlZO1KPVAftgYzdsF6CtSDGV709v_DFjcPNKzc0g,14511
-reflectorch/data_generation/priors/params.py,sha256=fND4ZNlplNertLHvIimfW0KKc2uWtPTUpwygCRskIu4,8397
-reflectorch/data_generation/priors/sampler_strategies.py,sha256=U-v5dXpFLJq939aQKvNl7n1Lih-y97K8WJCXBFRQiA0,13059
-reflectorch/data_generation/priors/scaler_mixin.py,sha256=gI64v2KOZugSJWaLKASR34fn6qVFl-aoeVA1BR5yXNg,2648
-reflectorch/data_generation/priors/subprior_sampler.py,sha256=6kaLIvJ_dSNfPCAbOzW1vyYKU3zQ6Qdnuu8nMRFkEoQ,14928
-reflectorch/data_generation/priors/utils.py,sha256=-iCLSkb_MrDIvmfEZKxpWMJDwj9PtZGP8LwbbM9JYos,3995
-reflectorch/data_generation/reflectivity/__init__.py,sha256=whto-vWKjWt3ZziqajgJ5M2RLZWXavSZpJLvyYcuom4,2280
-reflectorch/data_generation/reflectivity/abeles.py,sha256=4G42f1XPcu6hvYbvV3HCfivHx7oRYq_a6AzHyGK_UeM,2091
-reflectorch/data_generation/reflectivity/kinematical.py,sha256=tNmh1aYHA9-eFxnY-hJZBbctkpC89tht4o9rxbDBMdU,1475
-reflectorch/data_generation/reflectivity/memory_eff.py,sha256=OzPZJoFupRdlz8Qchndz0A5aYhO33yU_yYMdnVxgsNg,2997
-reflectorch/data_generation/reflectivity/numpy_implementations.py,sha256=QBzn4yVnOdlkHeeR-ZFPS115GnLdO9lMTGO2d3YhG9I,3177
-reflectorch/data_generation/reflectivity/smearing.py,sha256=XVtFfQFT-Ouyl8wa_IKc0p8ZZg8PCNtZIqe_RFy6E6E,3983
-reflectorch/extensions/__init__.py,sha256=XeuLafCqNwwmfWTcJgbuzCzpiypBG7ZatbIZrT9TvBA,159
-reflectorch/extensions/jupyter/__init__.py,sha256=inEXUpeVWeAhkW5nkW_dASBzsAlv4htvj7GIS7svIGk,231
-reflectorch/extensions/jupyter/callbacks.py,sha256=piDR4ax6JFSOPyqfkk-nxrhyWYdMrxgC8ocoaJbbbu8,1233
-reflectorch/extensions/matplotlib/__init__.py,sha256=8II5pU8015VrMjFI8szCKBP1zjz0dFAzBn7smNQzGuA,263
-reflectorch/extensions/matplotlib/losses.py,sha256=TqcyrFrls1N6RXotFyXDF64Xz6nJGg7n5XMSXFdeRtQ,845
-reflectorch/inference/__init__.py,sha256=i0KNn83XN33mLrV7bpHdLd0SXxuGCKfbQcIoa247Uts,834
-reflectorch/inference/inference_model.py,sha256=QvnQDRRcZByHDJh84T5W8O3X_aJLZI6AmlskE6BaBlU,36265
-reflectorch/inference/multilayer_fitter.py,sha256=0CxDpLOEp1terR4N39yFlxhvA8qAbHf_01NbmvYadck,5510
-reflectorch/inference/multilayer_inference_model.py,sha256=hH_-dJGdMOox8GHXdM_nODXDlNgh_v449xW5FmklRdo,7575
-reflectorch/inference/query_matcher.py,sha256=Dk49dW0XreeCjufzYBTKchfTdVbG6759ryV6I-wQL60,3387
-reflectorch/inference/record_time.py,sha256=3er-aoR8Sd_Kc4qNwUmRqkEz4FYhVxdi1ARnBohybzM,1140
-reflectorch/inference/sampler_solution.py,sha256=DeJM3EXEb6S5EiASj3mmdNI-Y07Cr5UzzA5oq-vEB-Q,2288
-reflectorch/inference/scipy_fitter.py,sha256=339M33OdmfgOpifJGLYk4KVcnnNJrY6_aH7Lz6Vtt24,5404
-reflectorch/inference/torch_fitter.py,sha256=j1NzkzLCmQ4H6LfIi82LsSBmIdunnWzm3kbGx-hqvDs,3391
-reflectorch/inference/preprocess_exp/__init__.py,sha256=bR6H-xgBD96z4P9P1T2ECnWvalrimdMTfTNArIWPLy0,383
-reflectorch/inference/preprocess_exp/attenuation.py,sha256=UKDfUjCKKMgAuEs5Ttyo0KEQmpvHZI52UgVflh7T81A,1518
-reflectorch/inference/preprocess_exp/cut_with_q_ratio.py,sha256=CbtwIw7iJNkoVxqTHKzONBgGFwOsCUyFoTIQ8XMLTfY,1149
-reflectorch/inference/preprocess_exp/footprint.py,sha256=xc409M5X-QW0Ce_6dEZdj8NkOY1vd0LaGpPQFxiOOR0,2625
-reflectorch/inference/preprocess_exp/interpolation.py,sha256=o2v-mlfRYzeaaikeQVeY7EA7j-K42dMfT12oN3mk51k,694
-reflectorch/inference/preprocess_exp/normalize.py,sha256=09v7nZdtw6SnW_67xFrPnqzOA5AtFBGarjK4Pfn4VoE,695
-reflectorch/inference/preprocess_exp/preprocess.py,sha256=pyyq8fSvcm1bWAShzGHYnKOc55Rofh4FIw1AC7Smq-U,5111
-reflectorch/ml/__init__.py,sha256=wdItiY13KD6PlCrHnHVcdpQOgTB5iUSj_qn4BZFM_uU,908
-reflectorch/ml/basic_trainer.py,sha256=Kr-oVAlmZjkL9MuJDxHAKA_1tTqUvX-3Q2BETWWlsmE,9806
-reflectorch/ml/callbacks.py,sha256=YxA_VUlfsE9Uh9MotPe2tXq6rbCyoG52LfI3e_YQy3w,2902
-reflectorch/ml/dataloaders.py,sha256=IvKmsH5gX_b-00KRFeL-x3keEfBcvYkQFWGWd8Caj-I,1073
-reflectorch/ml/loggers.py,sha256=XC7KwqHDTSr_2iWyBatOQO6EuFtK1bvwUVBcoA-D7fg,904
-reflectorch/ml/schedulers.py,sha256=xIloPpmCSnB35YniyzcDZoXHJFMT_rz0CWh2xiXnDak,10207
-reflectorch/ml/trainers.py,sha256=36R_oU33UHoebd7F1eNVlQ1GdhJXeGMgWsg-RrId2Mg,5014
-reflectorch/ml/utils.py,sha256=VfgWVjnXTrvw8eIMhFJXEaf7gkmp3rTUHrZvy42b_2k,232
-reflectorch/models/__init__.py,sha256=4k6JTr4XOhxtchCIlkcYNW51CmdIPsVAGfAwuhhTgYI,521
-reflectorch/models/activations.py,sha256=LDiIxCnLFb8r_TRBZSt6vdOZmexCWAGa5DfE_SotUL8,1431
-reflectorch/models/encoders/__init__.py,sha256=9PT31292CtfXlm1jucd7-2h69M_2vQNYQeaFX0lM2EM,702
-reflectorch/models/encoders/conv_encoder.py,sha256=Ns5df_baTh-7lu-xRaO_jnnar1apsXGKNDfbaFIHv0U,7812
-reflectorch/models/encoders/conv_res_net.py,sha256=_TYbF9GMThOtYuGmiyzIkClbq8wwRA251IFzUlxMwdU,3497
-reflectorch/models/encoders/fno.py,sha256=s_S7hnpLE7iGfyvnQ-QvTh0rKO5KFiy5tUYau4sJbvI,4693
-reflectorch/models/encoders/transformers.py,sha256=hfgGr2HiTj7DvaQnm_5RU_osPxVZn-L0r5OGqF8ZJZ4,1610
-reflectorch/models/networks/__init__.py,sha256=_NBjIl4QNLAuzBb2IaOIGG37iWwGzVQwuQhbcP9lxpI,450
-reflectorch/models/networks/mlp_networks.py,sha256=C7py6qCBVaYYt0FMEf8gbT4lndArKpUYYgTN1001-T8,11614
-reflectorch/models/networks/residual_net.py,sha256=msDJaDw7qL9ebEW1Avw6Qw0lgni68AMgF4kXiJKzeaQ,4637
-reflectorch/runs/__init__.py,sha256=2BcdMJul5yd726p8w4iqlKhygAAxiu1zu0MKDe96bWk,816
-reflectorch/runs/config.py,sha256=6aEub3NV0jmoREdegV7S3Nz-5o1xPZnmPpNgYfMpdys,963
-reflectorch/runs/slurm_utils.py,sha256=T5vsWrcduq_N9mS9XAXjAbx7PHcYiiiwjdS0iiXh_TI,2759
-reflectorch/runs/train.py,sha256=NaHMUYApjOCeajyS5UMQkeCVyxVtroohXK5ceHNLOkM,2719
-reflectorch/runs/utils.py,sha256=j_gJYrw4fIZvKJWXPdt1mOR0d_Ht6pg0rDjE2iOTLc8,9737
-reflectorch-1.1.0.dist-info/LICENSE.txt,sha256=2kX9kLKiIRiQRqUXwk3J-Ba3fqmztNu8ORskLBlAuKM,1098
-reflectorch-1.1.0.dist-info/METADATA,sha256=Jj8WKCgTrNn8_TT7GZS6gbLX2BxLauF3J4FJbo-18ZM,7777
-reflectorch-1.1.0.dist-info/WHEEL,sha256=rWxmBtp7hEUqVLOnTaDOPpR-cZpCDkzhhcBce-Zyd5k,91
-reflectorch-1.1.0.dist-info/top_level.txt,sha256=2EyIWrt4SeZ3hNadLXvEVpPFhyoZ4An7YflP4y_E3Fc,12
-reflectorch-1.1.0.dist-info/RECORD,,