reflectorch 1.0.0__py3-none-any.whl

reflectorch/data_generation/q_generator.py

@@ -0,0 +1,232 @@
+# -*- 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
+
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.utils import uniform_sampler
+from reflectorch.data_generation.priors import BasicParams
+from reflectorch.utils import angle_to_q
+from reflectorch.data_generation.priors.no_constraints import DEFAULT_DEVICE, DEFAULT_DTYPE
+
+__all__ = [
+    "QGenerator",
+    "ConstantQ",
+    "VariableQ",
+    "EquidistantQ",
+    "ConstantAngle",
+]
+
+
+class QGenerator(object):
+    """Base class for momentum transfer (q) generators"""
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        pass
+
+
+class ConstantQ(QGenerator):
+    """Q generator for reflectivity curves with fixed discretization
+
+    Args:
+        q (Union[Tensor, Tuple[float, float, int]], optional): tuple (q_min, q_max, num_q) defining the minimum q value, maximum q value and the number of q points. Defaults to (0., 0.2, 128).
+        device (optional): the Pytorch device. Defaults to DEFAULT_DEVICE.
+        dtype (optional): the Pytorch data type. Defaults to DEFAULT_DTYPE.
+        remove_zero (bool, optional): do not include the upper end of the interval. Defaults to False.
+        fixed_zero (bool, optional): do not include the lower end of the interval. Defaults to False.
+    """
+
+    def __init__(self,
+                 q: Union[Tensor, Tuple[float, float, int]] = (0., 0.2, 128),
+                 device=DEFAULT_DEVICE,
+                 dtype=DEFAULT_DTYPE,
+                 remove_zero: bool = False,
+                 fixed_zero: bool = False,
+                 ):
+        if isinstance(q, (tuple, list)):
+            q = torch.linspace(*q, device=device, dtype=dtype)
+            if remove_zero:
+                if fixed_zero:
+                    q = q[1:]
+                else:
+                    q = q[:-1]
+        self.q = q
+
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        """generate a batch of q values
+
+        Args:
+            batch_size (int): the batch size
+
+        Returns:
+            Tensor: generated batch of q values
+        """
+        return self.q.clone()[None].expand(batch_size, self.q.shape[0])
+    
+
+class VariableQ(QGenerator):
+    """Q generator for reflectivity curves with variable discretization
+
+    Args:
+        q_min_range (list, optional): the range for sampling the minimum q value of the curves, *q_min*. Defaults to [0.01, 0.03].
+        q_max_range (list, optional): the range for sampling the maximum q value of the curves, *q_max*. Defaults to [0.1, 0.5].
+        n_q_range (list, optional): the range for the number of points in the curves (equidistantly sampled between *q_min* and *q_max*, 
+                                    the number of points varies between batches but is constant within a batch). Defaults to [64, 256].
+        device (optional): the Pytorch device. Defaults to DEFAULT_DEVICE.
+        dtype (optional): the Pytorch data type. Defaults to DEFAULT_DTYPE.
+    """
+
+    def __init__(self,
+                 q_min_range: Tuple[float, float] = (0.01, 0.03),
+                 q_max_range: Tuple[float, float] = (0.1, 0.5),
+                 n_q_range: Tuple[int, int] = (64, 256),
+                 device=DEFAULT_DEVICE,
+                 dtype=DEFAULT_DTYPE,
+                 ):
+        self.q_min_range = q_min_range
+        self.q_max_range = q_max_range
+        self.n_q_range = n_q_range
+        self.device = device
+        self.dtype = dtype
+
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        """generate a batch of q values (the number of points varies between batches but is constant within a batch)
+
+        Args:
+            batch_size (int): the batch size
+
+        Returns:
+            Tensor: generated batch of q values
+        """
+        q_min = np.random.uniform(*self.q_min_range, batch_size)
+        q_max = np.random.uniform(*self.q_max_range, batch_size)
+        if self.n_q_range[0] == self.n_q_range[1]:
+            n_q = self.n_q_range[0]
+        else:
+            n_q = np.random.randint(self.n_q_range[0], self.n_q_range[1] + 1)
+        
+        q = torch.from_numpy(np.linspace(q_min, q_max, n_q).T).to(self.device).to(self.dtype)
+        
+        return q
+    
+    def scale_q(self, q):
+        """scales the q values to the range [-1, 1]
+
+        Args:
+            q (Tensor): unscaled q values
+
+        Returns:
+            Tensor: scaled q values
+        """
+        scaled_q_01 = (q - self.q_min_range[0]) / (self.q_max_range[1] - self.q_min_range[0]) 
+
+        return 2.0 * (scaled_q_01 - 0.5)
+
+
+class ConstantAngle(QGenerator):
+    """Q generator for reflectivity curves measured at equidistant angles
+
+    Args:
+        angle_range (Tuple[float, float, int], optional): the range of the incident angles. Defaults to (0., 0.2, 257).
+        wavelength (float, optional): the beam wavelength in units of angstroms. Defaults to 1.
+        device (optional): the Pytorch device. Defaults to DEFAULT_DEVICE.
+        dtype (optional): the Pytorch data type. Defaults to DEFAULT_DTYPE.
+        """
+    def __init__(self,
+                 angle_range: Tuple[float, float, int] = (0., 0.2, 257),
+                 wavelength: float = 1.,
+                 device=DEFAULT_DEVICE,
+                 dtype=DEFAULT_DTYPE,
+                 ):
+        self.q = torch.from_numpy(angle_to_q(np.linspace(*angle_range), wavelength)).to(device).to(dtype)
+
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        """generate a batch of q values
+
+        Args:
+            batch_size (int): the batch size
+
+        Returns:
+            Tensor: generated batch of q values
+        """
+        return self.q.clone()[None].expand(batch_size, self.q.shape[0])
+
+
+class EquidistantQ(QGenerator):
+    def __init__(self,
+                 max_range: Tuple[float, float],
+                 num_values: Union[int, Tuple[int, int]],
+                 device=None,
+                 dtype=torch.float64
+                 ):
+        self.max_range = max_range
+        self._num_values = num_values
+        self.device = device
+        self.dtype = dtype
+
+    @property
+    def num_values(self) -> int:
+        if isinstance(self._num_values, int):
+            return self._num_values
+        return np.random.randint(*self._num_values)
+
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        num_values = self.num_values
+        q_max = uniform_sampler(*self.max_range, batch_size, 1, device=self.device, dtype=self.dtype)
+        norm_qs = torch.linspace(0, 1, num_values + 1, device=self.device, dtype=self.dtype)[1:][None]
+        qs = norm_qs * q_max
+        return qs
+
+
+class TransformerQ(QGenerator):
+    def __init__(self,
+                 q_max: float = 0.2,
+                 num_values: Union[int, Tuple[int, int]] = (30, 512),
+                 min_dq_ratio: float = 5.,
+                 device=None,
+                 dtype=torch.float64,
+                 ):
+        self.min_dq_ratio = min_dq_ratio
+        self.q_max = q_max
+        self._dq_range = q_max / num_values[1], q_max / num_values[0]
+        self._num_values = num_values
+        self.device = device
+        self.dtype = dtype
+
+    def get_batch(self, batch_size: int, context: dict = None) -> Tensor:
+        assert context is not None
+
+        params: BasicParams = context['params']
+        total_thickness = params.thicknesses.sum(-1)
+
+        assert total_thickness.shape[0] == batch_size
+
+        min_dqs = torch.clamp(
+            2 * np.pi / total_thickness / self.min_dq_ratio, self._dq_range[0], self._dq_range[1] * 0.9
+        )
+
+        dqs = torch.rand_like(min_dqs) * (self._dq_range[1] - min_dqs) + min_dqs
+
+        num_q_values = torch.clamp(self.q_max // dqs, *self._num_values).to(torch.int)
+
+        q_values, mask = generate_q_padding_mask(num_q_values, self.q_max)
+
+        context['tgt_key_padding_mask'] = mask
+        context['num_q_values'] = num_q_values
+
+        return q_values
+
+
+def generate_q_padding_mask(num_q_values: Tensor, q_max: float):
+    batch_size = num_q_values.shape[0]
+    dqs = (q_max / num_q_values)[:, None]
+    q_values = torch.arange(1, num_q_values.max().item() + 1)[None].repeat(batch_size, 1) * dqs
+    mask = (q_values > q_max + dqs / 2)
+    q_values[mask] = 0.
+    return q_values, mask

reflectorch/data_generation/reflectivity/__init__.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.reflectivity.abeles import abeles_compiled, abeles
+from reflectorch.data_generation.reflectivity.memory_eff import abeles_memory_eff
+from reflectorch.data_generation.reflectivity.numpy_implementations import (
+    kinematical_approximation_np,
+    abeles_np,
+)
+from reflectorch.data_generation.reflectivity.smearing import abeles_constant_smearing
+from reflectorch.data_generation.reflectivity.kinematical import kinematical_approximation
+
+
+def reflectivity(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+        dq: Tensor = None,
+        gauss_num: int = 51,
+        constant_dq: bool = True,
+        log: bool = False,
+        abeles_func=None,
+):
+    """Function which computes the reflectivity from thin film parameters using 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.
+        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.
+        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.
+
+    Returns:
+        Tensor: the computed reflectivity curves
+    """
+    abeles_func = abeles_func or abeles
+    q = torch.atleast_2d(q)
+
+    if dq is None:
+        reflectivity_curves = abeles_func(q, thickness, roughness, sld)
+    else:
+        reflectivity_curves = abeles_constant_smearing(
+            q, thickness, roughness, sld,
+            dq=dq, gauss_num=gauss_num, constant_dq=constant_dq, abeles_func=abeles_func
+        )
+
+    if log:
+        reflectivity_curves = torch.log10(reflectivity_curves)
+    return reflectivity_curves

reflectorch/data_generation/reflectivity/abeles.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+import math
+
+import torch
+from torch import Tensor
+
+
+def abeles(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+):
+    """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
+
+    Returns:
+        Tensor: simulated reflectivity curves
+    """
+    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]
+    roughness = roughness[:, None] ** 2
+
+    sld = sld * 1e-6 + 1e-30j
+
+    k_z0 = (q / 2).to(c_dtype)
+
+    if k_z0.dim() == 1:
+        k_z0.unsqueeze_(0)
+
+    if k_z0.dim() == 2:
+        k_z0.unsqueeze_(-1)
+
+    k_n = torch.sqrt(k_z0 ** 2 - 4 * math.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 = 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)
+
+    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
+
+    r = (m[..., 1, 0] / m[..., 0, 0]).abs() ** 2
+    r = torch.clamp_max_(r, 1.)
+
+    return r
+
+
+# @torch.jit.script # commented so far due to complex numbers issue
+def abeles_compiled(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+):
+    return abeles(q, thickness, roughness, sld)

reflectorch/data_generation/reflectivity/kinematical.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from torch import Tensor
+
+
+def kinematical_approximation(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+        *,
+        apply_fresnel: bool = True,
+        log: bool = False,
+):
+    c_dtype = torch.complex128 if q.dtype is torch.float64 else torch.complex64
+
+    batch_size, num_layers = thickness.shape
+
+    q = q.to(c_dtype)
+
+    if q.dim() == 1:
+        q.unsqueeze_(0)
+
+    if q.dim() == 2:
+        q.unsqueeze_(-1)
+
+    sld = sld * 1e-6 + 1e-30j
+
+    drho = torch.cat([sld[..., 0][..., None], sld[..., 1:] - sld[..., :-1]], -1)[:, None]
+    thickness = torch.cumsum(torch.cat([torch.zeros(batch_size, 1).to(thickness), thickness], -1), -1)[:, None]
+    roughness = roughness[:, None]
+
+    r = (drho * torch.exp(- (roughness * q) ** 2 / 2 + 1j * (q * thickness))).sum(-1).abs().float() ** 2
+
+    if apply_fresnel:
+
+        substrate_sld = sld[:, -1:]
+
+        rf = _get_resnel_reflectivity(q, substrate_sld[:, None])
+
+        r = torch.clamp_max_(r * rf / substrate_sld.real ** 2, 1.)
+
+    if log:
+        r = torch.log10(r)
+
+    return r
+
+
+def _get_resnel_reflectivity(q, substrate_slds):
+    _RE_CONST = 0.28174103675406496
+
+    q_c = torch.sqrt(substrate_slds + 0j) / _RE_CONST * 2
+    q_prime = torch.sqrt(q ** 2 - q_c ** 2 + 0j)
+    r_f = ((q - q_prime) / (q + q_prime)).abs().float() ** 2
+
+    return r_f.squeeze(-1)
+

reflectorch/data_generation/reflectivity/memory_eff.py

@@ -0,0 +1,92 @@
+# -*- coding: utf-8 -*-
+
+from math import pi
+
+import torch
+from torch import Tensor
+
+
+def abeles_memory_eff(
+        q: Tensor,
+        thickness: Tensor,
+        roughness: Tensor,
+        sld: Tensor,
+):
+    c_dtype = torch.complex128 if q.dtype is torch.float64 else torch.complex64
+
+    batch_size, num_layers = thickness.shape
+
+    sld = sld * 1e-6 + 1e-30j
+
+    num_interfaces = num_layers + 1
+
+    k_z0 = (q / 2).to(c_dtype)
+
+    if len(k_z0.shape) == 1:
+        k_z0.unsqueeze_(0)
+
+    thickness_prev_layer = 1.  # ambient
+
+    for interface_num in range(num_interfaces):
+
+        prev_layer_idx = interface_num - 1
+        next_layer_idx = interface_num
+
+        if interface_num == 0:
+            k_z_previous_layer = _get_relative_k_z(k_z0, torch.zeros(batch_size, 1).to(sld))
+        else:
+            thickness_prev_layer = thickness[:, prev_layer_idx].unsqueeze(1)
+            k_z_previous_layer = _get_relative_k_z(k_z0, sld[:, prev_layer_idx].unsqueeze(1))
+
+        k_z_next_layer = _get_relative_k_z(k_z0, sld[:, next_layer_idx].unsqueeze(1))  # (batch_num, q_num)
+
+        reflection_matrix = _make_reflection_matrix(
+            k_z_previous_layer, k_z_next_layer, roughness[:, interface_num].unsqueeze(1)
+        )
+
+        if interface_num == 0:
+            total_reflectivity_matrix = reflection_matrix
+        else:
+            translation_matrix = _make_translation_matrix(k_z_previous_layer, thickness_prev_layer)
+
+            total_reflectivity_matrix = torch.einsum(
+                'bnmr, bmlr, bljr -> bnjr', total_reflectivity_matrix, translation_matrix, reflection_matrix
+            )
+
+    r = total_reflectivity_matrix[:, 0, 1] / total_reflectivity_matrix[:, 1, 1]
+
+    reflectivity = torch.clamp_max_(torch.abs(r) ** 2, 1.).flatten(1)
+
+    return reflectivity
+
+
+def _get_relative_k_z(k_z0, scattering_length_density):
+    return torch.sqrt(k_z0 ** 2 - 4 * pi * scattering_length_density)
+
+
+def _make_reflection_matrix(k_z_previous_layer, k_z_next_layer, interface_roughness):
+    p = _safe_div((k_z_previous_layer + k_z_next_layer), (2 * k_z_previous_layer)) * \
+        torch.exp(-(k_z_previous_layer - k_z_next_layer) ** 2 * 0.5 * interface_roughness ** 2)
+
+    m = _safe_div((k_z_previous_layer - k_z_next_layer), (2 * k_z_previous_layer)) * \
+        torch.exp(-(k_z_previous_layer + k_z_next_layer) ** 2 * 0.5 * interface_roughness ** 2)
+
+    return _stack_mtx(p, m, m, p)
+
+
+def _stack_mtx(a11, a12, a21, a22):
+    return torch.stack([
+        torch.stack([a11, a12], dim=1),
+        torch.stack([a21, a22], dim=1),
+    ], dim=1)
+
+
+def _make_translation_matrix(k_z, thickness):
+    return _stack_mtx(
+        torch.exp(-1j * k_z * thickness), torch.zeros_like(k_z),
+        torch.zeros_like(k_z), torch.exp(1j * k_z * thickness)
+    )
+
+
+def _safe_div(numerator, denominator):
+    return torch.where(denominator == 0, numerator, torch.divide(numerator, denominator))

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]