reflectorch 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

reflectorch/data_generation/scale_curves.py

@@ -1,112 +1,112 @@
-from pathlib import Path
-
-import torch
-from torch import Tensor
-
-from reflectorch.data_generation.priors import PriorSampler
-from reflectorch.paths import SAVED_MODELS_DIR
-
-
-class CurvesScaler(object):
-    """Base class for curve scalers"""
-    def scale(self, curves: Tensor):
-        raise NotImplementedError
-
-    def restore(self, curves: Tensor):
-        raise NotImplementedError
-
-
-class LogAffineCurvesScaler(CurvesScaler):
-    """    Curve scaler which scales the reflectivity curves according to the logarithmic affine transformation:
-    :math:`\log_{10}(R + eps) \cdot weight + bias`.
-
-    Args:
-        weight (float): multiplication factor in the transformation
-        bias (float): addition term in the transformation
-        eps (float): sets the minimum intensity value of the reflectivity curves which is considered
-    """
-    def __init__(self, weight: float = 0.1, bias: float = 0.5, eps: float = 1e-10):
-        self.weight = weight
-        self.bias = bias
-        self.eps = eps
-
-    def scale(self, curves: Tensor):
-        """scales the reflectivity curves to a ML-friendly range
-
-        Args:
-            curves (Tensor): original reflectivity curves
-
-        Returns:
-            Tensor: reflectivity curves scaled to a ML-friendly range
-        """
-        return torch.log10(curves + self.eps) * self.weight + self.bias
-
-    def restore(self, curves: Tensor):
-        """restores the physical reflectivity curves
-
-        Args:
-            curves (Tensor): scaled reflectivity curves
-
-        Returns:
-            Tensor: reflectivity curves restored to the physical range
-        """
-        return 10 ** ((curves - self.bias) / self.weight) - self.eps
-
-
-class MeanNormalizationCurvesScaler(CurvesScaler):
-    """Curve scaler which scales the reflectivity curves by the precomputed mean of a batch of curves
-
-    Args:
-        path (str, optional): path to the precomputed mean of the curves, only used if ``curves_mean`` is None. Defaults to None.
-        curves_mean (Tensor, optional): the precomputed mean of the curves. Defaults to None.
-        device (torch.device, optional): the Pytorch device. Defaults to 'cuda'.
-    """
-
-    def __init__(self, path: str = None, curves_mean: Tensor = None, device: torch.device = 'cuda'):
-        if curves_mean is None:
-            curves_mean = torch.load(self.get_path(path))
-        self.curves_mean = curves_mean.to(device)
-
-    def scale(self, curves: Tensor):
-        """scales the reflectivity curves to a ML-friendly range
-
-        Args:
-            curves (Tensor): original reflectivity curves
-
-        Returns:
-            Tensor: reflectivity curves scaled to a ML-friendly range
-        """
-        self.curves_mean = self.curves_mean.to(curves)
-        return curves / self.curves_mean - 1
-
-    def restore(self, curves: Tensor):
-        """restores the physical reflectivity curves
-
-        Args:
-            curves (Tensor): scaled reflectivity curves
-
-        Returns:
-            Tensor: reflectivity curves restored to the physical range
-        """
-        self.curves_mean = self.curves_mean.to(curves)
-        return (curves + 1) * self.curves_mean
-
-    @staticmethod
-    def save(prior_sampler: PriorSampler, q: Tensor, path: str, num: int = 16384):
-        """computes the mean of a batch of reflectivity curves and saves it
-
-        Args:
-            prior_sampler (PriorSampler): the prior sampler
-            q (Tensor): the q values
-            path (str): the path for saving the mean of the curves
-            num (int, optional): the number of curves used to compute the mean. Defaults to 16384.
-        """
-        params = prior_sampler.sample(num)
-        curves_mean = params.reflectivity(q, log=False).mean(0).cpu()
-        torch.save(curves_mean, MeanNormalizationCurvesScaler.get_path(path))
-
-    @staticmethod
-    def get_path(path: str) -> Path:
-        if not path.endswith('.pt'):
-            path = path + '.pt'
-        return SAVED_MODELS_DIR / path
+from pathlib import Path
+
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.priors import PriorSampler
+from reflectorch.paths import SAVED_MODELS_DIR
+
+
+class CurvesScaler(object):
+    """Base class for curve scalers"""
+    def scale(self, curves: Tensor):
+        raise NotImplementedError
+
+    def restore(self, curves: Tensor):
+        raise NotImplementedError
+
+
+class LogAffineCurvesScaler(CurvesScaler):
+    """    Curve scaler which scales the reflectivity curves according to the logarithmic affine transformation:
+    :math:`\log_{10}(R + eps) \cdot weight + bias`.
+
+    Args:
+        weight (float): multiplication factor in the transformation
+        bias (float): addition term in the transformation
+        eps (float): sets the minimum intensity value of the reflectivity curves which is considered
+    """
+    def __init__(self, weight: float = 0.1, bias: float = 0.5, eps: float = 1e-10):
+        self.weight = weight
+        self.bias = bias
+        self.eps = eps
+
+    def scale(self, curves: Tensor):
+        """scales the reflectivity curves to a ML-friendly range
+
+        Args:
+            curves (Tensor): original reflectivity curves
+
+        Returns:
+            Tensor: reflectivity curves scaled to a ML-friendly range
+        """
+        return torch.log10(curves + self.eps) * self.weight + self.bias
+
+    def restore(self, curves: Tensor):
+        """restores the physical reflectivity curves
+
+        Args:
+            curves (Tensor): scaled reflectivity curves
+
+        Returns:
+            Tensor: reflectivity curves restored to the physical range
+        """
+        return 10 ** ((curves - self.bias) / self.weight) - self.eps
+
+
+class MeanNormalizationCurvesScaler(CurvesScaler):
+    """Curve scaler which scales the reflectivity curves by the precomputed mean of a batch of curves
+
+    Args:
+        path (str, optional): path to the precomputed mean of the curves, only used if ``curves_mean`` is None. Defaults to None.
+        curves_mean (Tensor, optional): the precomputed mean of the curves. Defaults to None.
+        device (torch.device, optional): the Pytorch device. Defaults to 'cuda'.
+    """
+
+    def __init__(self, path: str = None, curves_mean: Tensor = None, device: torch.device = 'cuda'):
+        if curves_mean is None:
+            curves_mean = torch.load(self.get_path(path))
+        self.curves_mean = curves_mean.to(device)
+
+    def scale(self, curves: Tensor):
+        """scales the reflectivity curves to a ML-friendly range
+
+        Args:
+            curves (Tensor): original reflectivity curves
+
+        Returns:
+            Tensor: reflectivity curves scaled to a ML-friendly range
+        """
+        self.curves_mean = self.curves_mean.to(curves)
+        return curves / self.curves_mean - 1
+
+    def restore(self, curves: Tensor):
+        """restores the physical reflectivity curves
+
+        Args:
+            curves (Tensor): scaled reflectivity curves
+
+        Returns:
+            Tensor: reflectivity curves restored to the physical range
+        """
+        self.curves_mean = self.curves_mean.to(curves)
+        return (curves + 1) * self.curves_mean
+
+    @staticmethod
+    def save(prior_sampler: PriorSampler, q: Tensor, path: str, num: int = 16384):
+        """computes the mean of a batch of reflectivity curves and saves it
+
+        Args:
+            prior_sampler (PriorSampler): the prior sampler
+            q (Tensor): the q values
+            path (str): the path for saving the mean of the curves
+            num (int, optional): the number of curves used to compute the mean. Defaults to 16384.
+        """
+        params = prior_sampler.sample(num)
+        curves_mean = params.reflectivity(q, log=False).mean(0).cpu()
+        torch.save(curves_mean, MeanNormalizationCurvesScaler.get_path(path))
+
+    @staticmethod
+    def get_path(path: str) -> Path:
+        if not path.endswith('.pt'):
+            path = path + '.pt'
+        return SAVED_MODELS_DIR / path

reflectorch/data_generation/smearing.py

@@ -1,99 +1,99 @@
-import torch
-from torch import Tensor
-
-from reflectorch.data_generation.priors.parametric_subpriors import BasicParams
-
-
-class Smearing(object):
-    """Class which applies resolution smearing to the reflectivity curves. 
-    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 resolutions. Defaults to (0.01, 0.1).
-        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.
-    """
-    def __init__(self,
-                 sigma_range: tuple = (0.01, 0.1),
-                 constant_dq: bool = False,
-                 gauss_num: int = 31,
-                 share_smeared: float = 0.2,
-                 ):
-        self.sigma_min, self.sigma_max = sigma_range
-        self.sigma_delta = self.sigma_max - self.sigma_min
-        self.constant_dq = constant_dq
-        self.gauss_num = gauss_num
-        self.share_smeared = share_smeared
-
-    def __repr__(self):
-        return f'Smearing(({self.sigma_min}, {self.sigma_max})'
-
-    def generate_resolutions(self, batch_size: int, device=None, dtype=None):
-        num_smeared = int(batch_size * self.share_smeared)
-        if not num_smeared:
-            return None, None
-        dq = torch.rand(num_smeared, 1, device=device, dtype=dtype) * self.sigma_delta + self.sigma_min
-        indices = torch.zeros(batch_size, device=device, dtype=torch.bool)
-        indices[torch.randperm(batch_size, device=device)[:num_smeared]] = True
-        return dq, indices
-    
-    def scale_resolutions(self, resolutions: Tensor) -> Tensor:
-        """Scales the q-resolution values to [-1,1] range using the internal sigma range"""
-        sigma_min = 0.0 if self.share_smeared != 1.0 else self.sigma_min
-        return 2 * (resolutions - sigma_min) / (self.sigma_max - sigma_min) - 1
-
-    def get_curves(self, q_values: Tensor, params: BasicParams, refl_kwargs:dict = None):
-        refl_kwargs = refl_kwargs or {}
-
-        dq, indices = self.generate_resolutions(params.batch_size, device=params.device, dtype=params.dtype)
-        q_resolutions = torch.zeros(q_values.shape[0], 1, dtype=q_values.dtype, device=q_values.device)
-
-        if dq is None:
-            return params.reflectivity(q_values, **refl_kwargs), q_resolutions
-        
-        refl_kwargs_not_smeared = {}
-        refl_kwargs_smeared = {}
-        for key, value in refl_kwargs.items():
-            if isinstance(value, torch.Tensor) and value.shape[0] == params.batch_size:
-                refl_kwargs_not_smeared[key] = value[~indices]
-                refl_kwargs_smeared[key] = value[indices]
-            else:
-                refl_kwargs_not_smeared[key] = value
-                refl_kwargs_smeared[key] = value
-
-         # Compute unsmeared reflectivity
-        if (~indices).sum().item():
-            if q_values.dim() == 2 and q_values.shape[0] > 1:
-                q = q_values[~indices]
-            else:
-                q = q_values
-
-            reflectivity_not_smeared = params[~indices].reflectivity(q, **refl_kwargs_not_smeared)
-        else:
-            reflectivity_not_smeared = None
-
-        # Compute smeared reflectivity
-        if indices.sum().item():
-            if q_values.dim() == 2 and q_values.shape[0] > 1:
-                q = q_values[indices]
-            else:
-                q = q_values
-
-            reflectivity_smeared = params[indices].reflectivity(
-                q, dq=dq, constant_dq=self.constant_dq, gauss_num=self.gauss_num, **refl_kwargs_smeared
-            )
-        else:
-            reflectivity_smeared = None
-
-        curves = torch.empty(params.batch_size, q_values.shape[-1], device=params.device, dtype=params.dtype)
-        
-        if (~indices).sum().item():
-            curves[~indices] = reflectivity_not_smeared
-        
-        curves[indices] = reflectivity_smeared
-
-        q_resolutions[indices] = dq
-
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.priors.parametric_subpriors import BasicParams
+
+
+class Smearing(object):
+    """Class which applies resolution smearing to the reflectivity curves. 
+    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 resolutions. Defaults to (0.01, 0.1).
+        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.
+    """
+    def __init__(self,
+                 sigma_range: tuple = (0.01, 0.1),
+                 constant_dq: bool = False,
+                 gauss_num: int = 31,
+                 share_smeared: float = 0.2,
+                 ):
+        self.sigma_min, self.sigma_max = sigma_range
+        self.sigma_delta = self.sigma_max - self.sigma_min
+        self.constant_dq = constant_dq
+        self.gauss_num = gauss_num
+        self.share_smeared = share_smeared
+
+    def __repr__(self):
+        return f'Smearing(({self.sigma_min}, {self.sigma_max})'
+
+    def generate_resolutions(self, batch_size: int, device=None, dtype=None):
+        num_smeared = int(batch_size * self.share_smeared)
+        if not num_smeared:
+            return None, None
+        dq = torch.rand(num_smeared, 1, device=device, dtype=dtype) * self.sigma_delta + self.sigma_min
+        indices = torch.zeros(batch_size, device=device, dtype=torch.bool)
+        indices[torch.randperm(batch_size, device=device)[:num_smeared]] = True
+        return dq, indices
+    
+    def scale_resolutions(self, resolutions: Tensor) -> Tensor:
+        """Scales the q-resolution values to [-1,1] range using the internal sigma range"""
+        sigma_min = 0.0 if self.share_smeared != 1.0 else self.sigma_min
+        return 2 * (resolutions - sigma_min) / (self.sigma_max - sigma_min) - 1
+
+    def get_curves(self, q_values: Tensor, params: BasicParams, refl_kwargs:dict = None):
+        refl_kwargs = refl_kwargs or {}
+
+        dq, indices = self.generate_resolutions(params.batch_size, device=params.device, dtype=params.dtype)
+        q_resolutions = torch.zeros(q_values.shape[0], 1, dtype=q_values.dtype, device=q_values.device)
+
+        if dq is None:
+            return params.reflectivity(q_values, **refl_kwargs), q_resolutions
+        
+        refl_kwargs_not_smeared = {}
+        refl_kwargs_smeared = {}
+        for key, value in refl_kwargs.items():
+            if isinstance(value, torch.Tensor) and value.shape[0] == params.batch_size:
+                refl_kwargs_not_smeared[key] = value[~indices]
+                refl_kwargs_smeared[key] = value[indices]
+            else:
+                refl_kwargs_not_smeared[key] = value
+                refl_kwargs_smeared[key] = value
+
+         # Compute unsmeared reflectivity
+        if (~indices).sum().item():
+            if q_values.dim() == 2 and q_values.shape[0] > 1:
+                q = q_values[~indices]
+            else:
+                q = q_values
+
+            reflectivity_not_smeared = params[~indices].reflectivity(q, **refl_kwargs_not_smeared)
+        else:
+            reflectivity_not_smeared = None
+
+        # Compute smeared reflectivity
+        if indices.sum().item():
+            if q_values.dim() == 2 and q_values.shape[0] > 1:
+                q = q_values[indices]
+            else:
+                q = q_values
+
+            reflectivity_smeared = params[indices].reflectivity(
+                q, dq=dq, constant_dq=self.constant_dq, gauss_num=self.gauss_num, **refl_kwargs_smeared
+            )
+        else:
+            reflectivity_smeared = None
+
+        curves = torch.empty(params.batch_size, q_values.shape[-1], device=params.device, dtype=params.dtype)
+        
+        if (~indices).sum().item():
+            curves[~indices] = reflectivity_not_smeared
+        
+        curves[indices] = reflectivity_smeared
+
+        q_resolutions[indices] = dq
+
         return curves, q_resolutions