reflectorch 1.0.0__py3-none-any.whl

reflectorch/data_generation/reflectivity/smearing.py

@@ -0,0 +1,123 @@
+# -*- coding: utf-8 -*-
+from math import pi, sqrt, log
+
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.reflectivity.abeles import abeles
+from torch.nn.functional import conv1d, pad
+
+
+def abeles_constant_smearing(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+        dq: Tensor = None,
+        gauss_num: int = 51,
+        constant_dq: bool = True,
+        abeles_func=None,
+):
+    abeles_func = abeles_func or abeles
+    q_lin = _get_q_axes(q, dq, gauss_num, constant_dq=constant_dq)
+    kernels = _get_t_gauss_kernels(dq, gauss_num)
+
+    curves = abeles_func(q_lin, thickness, roughness, sld)
+
+    padding = (kernels.shape[-1] - 1) // 2
+    smeared_curves = conv1d(
+        pad(curves[None], (padding, padding), 'reflect'), kernels[:, None], groups=kernels.shape[0],
+    )[0]
+
+    if q.shape[0] != smeared_curves.shape[0]:
+        q = q.expand(smeared_curves.shape[0], *q.shape[1:])
+
+    smeared_curves = _batch_linear_interp1d(q_lin, smeared_curves, q)
+
+    return smeared_curves
+
+
+_FWHM = 2 * sqrt(2 * log(2.0))
+_2PI_SQRT = 1. / sqrt(2 * pi)
+
+
+def _batch_linspace(start: Tensor, end: Tensor, num: int):
+    return torch.linspace(0, 1, int(num), device=end.device, dtype=end.dtype)[None] * (end - start) + start
+
+
+def _torch_gauss(x, s):
+    return _2PI_SQRT / s * torch.exp(-0.5 * x ** 2 / s / s)
+
+
+def _get_t_gauss_kernels(resolutions: Tensor, gaussnum: int = 51):
+    gauss_x = _batch_linspace(-1.7 * resolutions, 1.7 * resolutions, gaussnum)
+    gauss_y = _torch_gauss(gauss_x, resolutions / _FWHM) * (gauss_x[:, 1] - gauss_x[:, 0])[:, None]
+    return gauss_y
+
+
+def _get_q_axes(q: Tensor, resolutions: Tensor, gaussnum: int = 51, constant_dq: bool = True):
+    if constant_dq:
+        return _get_q_axes_for_constant_dq(q, resolutions, gaussnum)
+    else:
+        return _get_q_axes_for_linear_dq(q, resolutions, gaussnum)
+
+
+def _get_q_axes_for_linear_dq(q: Tensor, resolutions: Tensor, gaussnum: int = 51):
+    gaussgpoint = (gaussnum - 1) / 2
+
+    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))
+
+    interpnums = torch.abs(
+        (torch.abs(end - start)) / (1.7 * resolutions / _FWHM / gaussgpoint)
+    ).round().to(int)
+
+    q_lin = 10 ** _batch_linspace_with_padding(start, end, interpnums)
+
+    return q_lin
+
+
+def _get_q_axes_for_constant_dq(q: Tensor, resolutions: Tensor, gaussnum: int = 51) -> Tensor:
+    gaussgpoint = (gaussnum - 1) / 2
+
+    start = q.min(1).values[:, None] - resolutions * 1.7
+    end = q.max(1).values[:, None] + resolutions * 1.7
+
+    interpnums = torch.abs(
+        (torch.abs(end - start)) / (1.7 * resolutions / gaussgpoint)
+    ).round().to(int)
+
+    q_lin = _batch_linspace_with_padding(start, end, interpnums)
+    q_lin = torch.clamp_min_(q_lin, 1e-6)
+
+    return q_lin
+
+
+def _batch_linspace_with_padding(start: Tensor, end: Tensor, nums: Tensor) -> Tensor:
+    max_num = nums.max().int().item()
+
+    deltas = 1 / (nums - 1)
+
+    x = torch.clamp_min_(_batch_linspace(deltas * (nums - max_num), torch.ones_like(deltas), max_num), 0)
+
+    x = x * (end - start) + start
+
+    return x
+
+
+def _batch_linear_interp1d(x: Tensor, y: Tensor, x_new: Tensor) -> Tensor:
+    eps = torch.finfo(y.dtype).eps
+
+    ind = torch.searchsorted(x.contiguous(), x_new.contiguous())
+
+    ind = torch.clamp_(ind - 1, 0, x.shape[-1] - 2)
+    slopes = (y[..., 1:] - y[..., :-1]) / (eps + (x[..., 1:] - x[..., :-1]))
+    ind_y = ind + torch.arange(slopes.shape[0], device=slopes.device)[:, None] * y.shape[1]
+    ind_slopes = ind + torch.arange(slopes.shape[0], device=slopes.device)[:, None] * slopes.shape[1]
+
+    y_new = y.flatten()[ind_y] + slopes.flatten()[ind_slopes] * (x_new - x.flatten()[ind_y])
+
+    return y_new

reflectorch/data_generation/scale_curves.py

@@ -0,0 +1,118 @@
+# -*- 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
+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

@@ -0,0 +1,67 @@
+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 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.
+        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 = (1e-4, 5e-3),
+                 constant_dq: bool = True,
+                 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 get_curves(self, q_values: Tensor, params: BasicParams):
+        dq, indices = self.generate_resolutions(params.batch_size, device=params.device, dtype=params.dtype)
+
+        if dq is None:
+            return params.reflectivity(q_values, log=False)
+
+        curves = torch.empty(params.batch_size, q_values.shape[-1], device=params.device, dtype=params.dtype)
+
+        if (~indices).sum().item():
+            if q_values.dim() == 2 and q_values.shape[0] > 1:
+                q = q_values[~indices]
+            else:
+                q = q_values
+
+            curves[~indices] = params[~indices].reflectivity(q, log=False)
+
+        if indices.sum().item():
+            if q_values.dim() == 2 and q_values.shape[0] > 1:
+                q = q_values[indices]
+            else:
+                q = q_values
+
+            curves[indices] = params[indices].reflectivity(
+                q, dq=dq, constant_dq=self.constant_dq, log=False, gauss_num=self.gauss_num
+            )
+
+        return curves

reflectorch/data_generation/utils.py

@@ -0,0 +1,154 @@
+# -*- 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
+
+import torch
+from torch import Tensor
+
+__all__ = [
+    "get_reversed_params",
+    "get_density_profiles",
+    "uniform_sampler",
+    "logdist_sampler",
+    "triangular_sampler",
+    "get_param_labels",
+    "get_d_rhos",
+    "get_slds_from_d_rhos",
+]
+
+
+def uniform_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
+    if isinstance(low, Tensor):
+        device, dtype = low.device, low.dtype
+    return torch.rand(*shape, device=device, dtype=dtype) * (high - low) + low
+
+
+def logdist_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
+    if isinstance(low, Tensor):
+        device, dtype = low.device, low.dtype
+        low, high = map(torch.log10, (low, high))
+    else:
+        low, high = map(log10, (low, high))
+    return 10 ** (torch.rand(*shape, device=device, dtype=dtype) * (high - low) + low)
+
+
+def triangular_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
+    if isinstance(low, Tensor):
+        device, dtype = low.device, low.dtype
+
+    x = torch.rand(*shape, device=device, dtype=dtype)
+
+    return (high - low) * (1 - torch.sqrt(x)) + low
+
+
+def get_reversed_params(thicknesses: Tensor, roughnesses: Tensor, slds: Tensor):
+    reversed_slds = torch.cumsum(
+        torch.flip(
+            torch.diff(
+                torch.cat([torch.zeros(slds.shape[0], 1).to(slds), slds], dim=-1),
+                dim=-1
+            ), (-1,)
+        ),
+        dim=-1
+    )
+    reversed_thicknesses = torch.flip(thicknesses, [-1])
+    reversed_roughnesses = torch.flip(roughnesses, [-1])
+    reversed_params = torch.cat([reversed_thicknesses, reversed_roughnesses, reversed_slds], -1)
+
+    return reversed_params
+
+
+def get_density_profiles(
+        thicknesses: Tensor,
+        roughnesses: Tensor,
+        slds: Tensor,
+        z_axis: Tensor = None,
+        num: int = 1000
+):
+    """Generates SLD profiles (and their derivative) based on batches of thicknesses, roughnesses and layer SLDs. 
+    
+    The axis has its zero at the top (ambient medium) interface and is positive inside the film.
+
+    Args:
+        thicknesses (Tensor): the layer thicknesses (top to bottom)
+        roughnesses (Tensor): the interlayer roughnesses (top to bottom)
+        slds (Tensor): the layer SLDs (top to bottom)
+        z_axis (Tensor, optional): a custom depth (z) axis. Defaults to None.
+        num (int, optional): number of discretization points for the profile. Defaults to 1000.
+
+    Returns:
+        tuple: the z axis, the computed density profile rho(z) and the derivative of the density profile drho/dz(z)
+    """
+    assert torch.all(roughnesses >= 0), 'Negative roughness happened'
+    assert torch.all(thicknesses >= 0), 'Negative thickness happened'
+
+    sample_num = thicknesses.shape[0]
+
+    d_rhos = get_d_rhos(slds)
+
+    zs = torch.cumsum(torch.cat([torch.zeros(sample_num, 1).to(thicknesses), thicknesses], dim=-1), dim=-1)
+
+    if z_axis is None:
+        z_axis = torch.linspace(- zs.max() * 0.1, zs.max() * 1.1, num, device=thicknesses.device)[None]
+    elif len(z_axis.shape) == 1:
+        z_axis = z_axis[None]
+
+    sigmas = roughnesses * sqrt(2)
+
+    profile = get_erf(z_axis[:, None], zs[..., None], sigmas[..., None], d_rhos[..., None]).sum(1)
+
+    d_profile = get_gauss(z_axis[:, None], zs[..., None], sigmas[..., None], d_rhos[..., None]).sum(1)
+
+    z_axis = z_axis[0]
+
+    return z_axis, profile, d_profile
+
+
+def get_d_rhos(slds: Tensor) -> Tensor:
+    d_rhos = torch.cat([slds[:, 0][:, None], torch.diff(slds, dim=-1)], -1)
+    return d_rhos
+
+
+def get_slds_from_d_rhos(d_rhos: Tensor) -> Tensor:
+    slds = torch.cumsum(d_rhos, dim=-1)
+    return slds
+
+
+def get_erf(z, z0, sigma, amp):
+    return (torch.erf((z - z0) / sigma) + 1) * amp / 2
+
+
+def get_gauss(z, z0, sigma, amp):
+    return amp / (sigma * sqrt(2 * pi)) * torch.exp(- (z - z0) ** 2 / 2 / sigma ** 2)
+
+
+def get_param_labels(
+        num_layers: int, *,
+        thickness_name: str = 'Thickness',
+        roughness_name: str = 'Roughness',
+        sld_name: str = 'SLD',
+        substrate_name: str = 'sub',
+) -> List[str]:
+    thickness_labels = [f'{thickness_name} L{num_layers - i}' for i in range(num_layers)]
+    roughness_labels = [f'{roughness_name} L{num_layers - i}' for i in range(num_layers)] + [f'{roughness_name} {substrate_name}']
+    sld_labels = [f'{sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{sld_name} {substrate_name}']
+    return thickness_labels + roughness_labels + sld_labels
+
+def get_param_labels_absorption_model(
+        num_layers: int, *,
+        thickness_name: str = 'Thickness',
+        roughness_name: str = 'Roughness',
+        real_sld_name: str = 'SLD real',
+        imag_sld_name: str = 'SLD imag',
+        substrate_name: str = 'sub',
+) -> List[str]:
+    thickness_labels = [f'{thickness_name} L{num_layers - i}' for i in range(num_layers)]
+    roughness_labels = [f'{roughness_name} L{num_layers - i}' for i in range(num_layers)] + [f'{roughness_name} {substrate_name}']
+    real_sld_labels = [f'{real_sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{real_sld_name} {substrate_name}']
+    imag_sld_labels = [f'{imag_sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{imag_sld_name} {substrate_name}']
+    return thickness_labels + roughness_labels + real_sld_labels + imag_sld_labels

reflectorch/extensions/__init__.py

@@ -0,0 +1,6 @@
+# -*- 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

@@ -0,0 +1,12 @@
+# -*- 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
+
+
+__all__ = [
+    'JPlotLoss',
+]

reflectorch/extensions/jupyter/callbacks.py

@@ -0,0 +1,40 @@
+# -*- 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
+
+from ..matplotlib import plot_losses
+
+
+class JPlotLoss(TrainerCallback):
+    """Callback for plotting the loss in a Jupyter notebook
+    """
+    def __init__(self, frequency: int, log: bool = True, clear: bool = True, **kwargs):
+        """
+
+        Args:
+            frequency (int): plotting frequency
+            log (bool, optional): if True, the plot is on a logarithmic scale. Defaults to True.
+            clear (bool, optional):
+        """
+        self.frequency = frequency
+        self.log = log
+        self.kwargs = kwargs
+        self.clear = clear
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        if not batch_num % self.frequency:
+            if self.clear:
+                clear_output(wait=True)
+
+            plot_losses(
+                trainer.losses,
+                log=self.log,
+                best_epoch=trainer.callback_params.get('saved_iteration', None),
+                **self.kwargs
+            )

reflectorch/extensions/matplotlib/__init__.py

@@ -0,0 +1,11 @@
+# -*- 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__ = [
+    "plot_losses",
+]

reflectorch/extensions/matplotlib/losses.py

@@ -0,0 +1,38 @@
+# -*- 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
+
+
+def plot_losses(
+        losses: dict,
+        log: bool = False,
+        show: bool = True,
+        title: str = 'Losses',
+        x_label: str = 'Iterations',
+        best_epoch: float = None,
+        **kwargs
+):
+    func = plt.semilogy if log else plt.plot
+
+    if len(losses) <= 2:
+        losses = {'loss': losses['total_loss']}
+
+    for k, data in losses.items():
+        func(data, label=k, **kwargs)
+
+    if best_epoch is not None:
+        plt.axvline(best_epoch, ls='--', color='red')
+
+    plt.xlabel(x_label)
+
+    if len(losses) > 2:
+        plt.legend()
+
+    plt.title(title)
+
+    if show:
+        plt.show()

reflectorch/inference/__init__.py

@@ -0,0 +1,22 @@
+from reflectorch.inference.inference_model import InferenceModel, EasyInferenceModel
+from reflectorch.inference.multilayer_inference_model import MultilayerInferenceModel
+from reflectorch.inference.preprocess_exp import (
+    StandardPreprocessing,
+    standard_preprocessing,
+    interp_reflectivity,
+    apply_attenuation_correction,
+    apply_footprint_correction,
+)
+from reflectorch.inference.torch_fitter import ReflGradientFit
+
+__all__ = [
+    "InferenceModel",
+    "EasyInferenceModel",
+    "MultilayerInferenceModel",
+    "StandardPreprocessing",
+    "standard_preprocessing",
+    "ReflGradientFit",
+    "interp_reflectivity",
+    "apply_attenuation_correction",
+    "apply_footprint_correction",
+]