reflectorch 1.5.1__py3-none-any.whl

reflectorch/data_generation/reflectivity/numpy_implementations.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+
+def abeles_np(
+        q: np.ndarray,
+        thickness: np.ndarray,
+        roughness: np.ndarray,
+        sld: np.ndarray,
+):
+    c_dtype = np.complex128 if q.dtype is np.float64 else np.complex64
+
+    if q.ndim == thickness.ndim == roughness.ndim == sld.ndim == 1:
+        zero_batch = True
+    else:
+        zero_batch = False
+
+    thickness = np.atleast_2d(thickness)
+    roughness = np.atleast_2d(roughness)
+    sld = np.atleast_2d(sld)
+
+    batch_size, num_layers = thickness.shape
+
+    sld = np.concatenate([np.zeros((batch_size, 1)).astype(sld.dtype), sld], -1)[:, None]
+    thickness = np.concatenate([np.zeros((batch_size, 1)).astype(thickness.dtype), thickness], -1)[:, None]
+    roughness = roughness[:, None] ** 2
+
+    sld = sld * 1e-6 + 1e-30j
+
+    k_z0 = (q / 2).astype(c_dtype)
+
+    if len(k_z0.shape) == 1:
+        k_z0 = k_z0[None]
+
+    if len(k_z0.shape) == 2:
+        k_z0 = k_z0[..., None]
+
+    k_n = np.sqrt(k_z0 ** 2 - 4 * np.pi * sld)
+
+    # k_n.shape - (batch, q, layers)
+
+    k_n, k_np1 = k_n[..., :-1], k_n[..., 1:]
+
+    beta = 1j * thickness * k_n
+
+    exp_beta = np.exp(beta)
+    exp_m_beta = np.exp(-beta)
+
+    rn = (k_n - k_np1) / (k_n + k_np1) * np.exp(- 2 * k_n * k_np1 * roughness)
+
+    c_matrices = np.stack([
+        np.stack([exp_beta, rn * exp_m_beta], -1),
+        np.stack([rn * exp_beta, exp_m_beta], -1),
+    ], -1)
+
+    c_matrices = np.moveaxis(c_matrices, -3, 0)
+
+    m, c_matrices = c_matrices[0], c_matrices[1:]
+
+    for c in c_matrices:
+        m = m @ c
+
+    r = np.abs(m[..., 1, 0] / m[..., 0, 0]) ** 2
+    r = np.clip(r, None, 1.)
+
+    if zero_batch:
+        r = r[0]
+
+    return r
+
+
+def kinematical_approximation_np(
+        q: np.ndarray,
+        thickness: np.ndarray,
+        roughness: np.ndarray,
+        sld: np.ndarray,
+):
+    if q.ndim == thickness.ndim == roughness.ndim == sld.ndim == 1:
+        zero_batch = True
+    else:
+        zero_batch = False
+
+    thickness = np.atleast_2d(thickness)
+    roughness = np.atleast_2d(roughness)
+    sld = np.atleast_2d(sld) * 1e-6 + 1e-30j
+    substrate_sld = sld[:, -1:]
+
+    batch_size, num_layers = thickness.shape
+
+    if q.ndim == 1:
+        q = q[None]
+
+    if q.ndim == 2:
+        q = q[..., None]
+
+    drho = np.concatenate([sld[..., 0][..., None], sld[..., 1:] - sld[..., :-1]], -1)[:, None]
+    thickness = np.cumsum(np.concatenate([np.zeros((batch_size, 1)), thickness], -1), -1)[:, None]
+    roughness = roughness[:, None]
+
+    r = np.abs((drho * np.exp(- (roughness * q) ** 2 / 2 + 1j * (q * thickness))).sum(-1)).astype(float) ** 2
+
+    rf = _get_resnel_reflectivity_np(q, substrate_sld[:, None])
+
+    r = np.clip(r * rf / np.real(substrate_sld) ** 2, None, 1.)
+
+    if zero_batch:
+        r = r[0]
+
+    return r
+
+
+def _get_resnel_reflectivity_np(q, substrate_slds):
+    _RE_CONST = 0.28174103675406496
+
+    q_c = np.sqrt(substrate_slds + 0j) / _RE_CONST * 2
+    q_prime = np.sqrt(q ** 2 - q_c ** 2 + 0j)
+    r_f = np.abs((q - q_prime) / (q + q_prime)).astype(float) ** 2
+
+    return r_f[..., 0]

reflectorch/data_generation/reflectivity/smearing.py

@@ -0,0 +1,138 @@
+# -*- 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 = 31,
+        constant_dq: bool = False,
+        abeles_func=None,
+        **abeles_kwargs
+):
+    abeles_func = abeles_func or abeles
+
+    if dq.dtype != thickness.dtype:
+        q = q.to(thickness)
+    
+    if dq.dtype != thickness.dtype:
+        dq = dq.to(thickness)
+
+    if q.shape[0] == 1:
+        q = q.repeat(thickness.shape[0], 1)
+
+    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, **abeles_kwargs)
+
+    padding = (kernels.shape[-1] - 1) // 2
+    padded_curves = pad(curves, (padding, padding), 'reflect')
+
+    smeared_curves = conv1d(
+        padded_curves, kernels[:, None], groups=kernels.shape[0],
+    )
+
+    if q.shape[0] != smeared_curves.shape[0]:
+        repeat_factor = smeared_curves.shape[0] // q.shape[0]
+        q = q.repeat(repeat_factor, 1)
+        q_lin = q_lin.repeat(repeat_factor, 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 = False):
+    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)[:, 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)
+    ).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/reflectivity/smearing_pointwise.py

@@ -0,0 +1,110 @@
+import torch
+import scipy
+import numpy as np
+from functools import lru_cache
+from typing import Tuple
+
+from reflectorch.data_generation.reflectivity.abeles import abeles
+
+#Pytorch version based on the JAX implementation of pointwise smearing in the refnx package.
+
+@lru_cache(maxsize=128)
+def gauss_legendre(n: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Calculate Gaussian quadrature abscissae and weights.
+
+    Args:
+        n (int): Gaussian quadrature order.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The abscissae and weights for Gauss-Legendre integration.
+    """
+    return scipy.special.p_roots(n)
+
+def gauss(x: torch.Tensor) -> torch.Tensor:
+    """
+    Calculate the Gaussian function.
+
+    Args:
+        x (torch.Tensor): Input tensor.
+
+    Returns:
+        torch.Tensor: Output tensor after applying the Gaussian function.
+    """
+    return torch.exp(-0.5 * x * x)
+
+def abeles_pointwise_smearing(
+    q: torch.Tensor,
+    dq: torch.Tensor,
+    thickness: torch.Tensor,
+    roughness: torch.Tensor,
+    sld: torch.Tensor,
+    gauss_num: int = 17,
+    abeles_func=None,
+    **abeles_kwargs,
+) -> torch.Tensor:
+    """
+    Compute reflectivity with variable smearing using Gaussian quadrature.
+
+    Args:
+        q (torch.Tensor): The momentum transfer (q) values.
+        dq (torch.Tensor): The resolution for curve smearing.
+        thickness (torch.Tensor): The layer thicknesses.
+        roughness (torch.Tensor): The interlayer roughnesses.
+        sld (torch.Tensor): The SLDs of the layers.
+        sld_magnetic (torch.Tensor, optional): The magnetic SLDs of the layers.
+        magnetization_angle (torch.Tensor, optional): The magnetization angles.
+        polarizer_eff (torch.Tensor, optional): The polarizer efficiency.
+        analyzer_eff (torch.Tensor, optional): The analyzer efficiency.
+        abeles_func (Callable, optional): A function implementing the simulation of the reflectivity curves.
+        gauss_num (int, optional): Gaussian quadrature order. Defaults to 17.
+
+    Returns:
+        torch.Tensor: The computed reflectivity curves.
+    """
+    abeles_func = abeles_func or abeles
+
+    if q.shape[0] == 1:
+        q = q.repeat(thickness.shape[0], 1)
+
+    _FWHM = 2 * np.sqrt(2 * np.log(2.0))
+    _INTLIMIT = 3.5
+
+    bs = q.shape[0]
+    nq = q.shape[-1]
+    device = q.device
+
+    quad_order = gauss_num
+    abscissa, weights = gauss_legendre(quad_order)
+    abscissa = torch.tensor(abscissa)[None, :, None].to(device)
+    weights = torch.tensor(weights)[None, :, None].to(device)
+    prefactor = 1.0 / np.sqrt(2 * np.pi)
+
+    gaussvals = prefactor * gauss(abscissa * _INTLIMIT)
+
+    va = q[:, None, :] - _INTLIMIT * dq[:, None, :] / _FWHM
+    vb = q[:, None, :] + _INTLIMIT * dq[:, None, :] / _FWHM
+
+    qvals_for_res_0 = (abscissa * (vb - va) + vb + va) / 2
+    qvals_for_res = qvals_for_res_0.reshape(bs, -1)
+
+    refl_curves = abeles_func(
+        q=qvals_for_res,
+        thickness=thickness,
+        roughness=roughness,
+        sld=sld,
+        **abeles_kwargs
+    )
+
+    # Handle multiple channels
+    if refl_curves.dim() == 3:
+        n_channels = refl_curves.shape[1]
+        refl_curves = refl_curves.reshape(bs, n_channels, quad_order, nq)
+        refl_curves = refl_curves * gaussvals.unsqueeze(1) * weights.unsqueeze(1)
+        refl_curves = torch.sum(refl_curves, dim=2) * _INTLIMIT
+    else:
+        refl_curves = refl_curves.reshape(bs, quad_order, nq)
+        refl_curves = refl_curves * gaussvals * weights
+        refl_curves = torch.sum(refl_curves, dim=1) * _INTLIMIT
+
+    return refl_curves

reflectorch/data_generation/scale_curves.py

@@ -0,0 +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

reflectorch/data_generation/smearing.py

@@ -0,0 +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
+
+        return curves, q_resolutions