reflectorch 1.5.1__py3-none-any.whl

reflectorch/data_generation/q_generator.py

@@ -0,0 +1,280 @@
+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",
+    "MaskedVariableQ",
+]
+
+
+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_min = q.min().item()
+        self.q_max = q.max().item()
+        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])
+    
+    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) / (self.q_max - self.q_min)
+        return 2.0 * (scaled_q_01 - 0.5)
+    
+
+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),
+                 mode: str = 'equidistant',
+                 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.mode = mode
+        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 = torch.rand(batch_size, device=self.device, dtype=self.dtype) * (self.q_min_range[1] - self.q_min_range[0]) + self.q_min_range[0]
+        q_max = torch.rand(batch_size, device=self.device, dtype=self.dtype) * (self.q_max_range[1] - self.q_max_range[0]) + self.q_max_range[0]
+
+        n_q = torch.randint(self.n_q_range[0], self.n_q_range[1] + 1, (1,), device=self.device).item()
+
+        if self.mode == 'equidistant':
+            q = torch.linspace(0, 1, n_q, device=self.device, dtype=self.dtype)
+        elif self.mode == 'random':
+            q = torch.rand(n_q, device=self.device, dtype=self.dtype).sort().values
+        elif self.mode == 'logspace':
+            q = torch.logspace(
+                start=torch.log10(torch.tensor(1e-4, dtype=self.dtype, device=self.device)),
+                end=torch.log10(torch.tensor(1.0, dtype=self.dtype, device=self.device)),
+                steps=n_q, dtype=self.dtype, device=self.device)
+
+        q = q_min[:, None] + q * (q_max - q_min)[:, None]
+        
+        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 MaskedVariableQ:
+    def __init__(self,
+                 q_min_range=(0.01, 0.03),
+                 q_max_range=(0.1, 0.5),
+                 n_q_range=(64, 256),
+                 mode='equidistant',
+                 shuffle_mask=False,
+                 total_thickness_constraint=True,
+                 min_points_per_fringe=4,
+                 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
+        self.mode = mode
+        self.shuffle_mask = shuffle_mask
+        self.total_thickness_constraint = total_thickness_constraint
+        self.min_points_per_fringe = min_points_per_fringe
+    
+    def get_batch(self, batch_size, context):
+        assert context is not None
+
+        q_min = torch.rand(batch_size, device=self.device, dtype=self.dtype) * (self.q_min_range[1] - self.q_min_range[0]) + self.q_min_range[0]
+        q_max = torch.rand(batch_size, device=self.device, dtype=self.dtype) * (self.q_max_range[1] - self.q_max_range[0]) + self.q_max_range[0]
+
+        max_n_q = self.n_q_range[1]
+
+        if self.mode == 'equidistant':
+            positions = torch.linspace(0, 1, max_n_q, device=self.device, dtype=self.dtype).expand(batch_size, max_n_q)
+        elif self.mode == 'random':
+            positions = torch.rand(batch_size, max_n_q, device=self.device, dtype=self.dtype)
+            positions, _ = positions.sort(dim=-1)
+        elif self.mode == 'mixed':
+            positions = torch.empty(batch_size, max_n_q, device=self.device, dtype=self.dtype)
+
+            half = batch_size // 2 # half batch gets equidistant
+            eq_pos = torch.linspace(0, 1, max_n_q, device=self.device, dtype=self.dtype).expand(half, max_n_q)
+            positions[:half] = eq_pos
+
+            rand_pos = torch.rand(batch_size - half, max_n_q, device=self.device, dtype=self.dtype) # other half gets sorted random
+            rand_pos, _ = rand_pos.sort(dim=-1)
+            positions[half:] = rand_pos
+        else:
+            raise ValueError(f"Unknown spacing mode: {self.mode}")
+
+        q = q_min[:, None] + positions * (q_max - q_min)[:, None]
+
+        n_qs = torch.randint(self.n_q_range[0], self.n_q_range[1] + 1, (batch_size,), device=self.device)
+
+        if 'params' in context and self.total_thickness_constraint: ### N_points > 1 + (Q_spread * total_thickness * min_np_per_kiessing_fringe) / (2*pi)
+            d_total = context['params'].thicknesses.sum(-1)
+            limit = 1 + ((q_max - q_min) * d_total * self.min_points_per_fringe) / (2*np.pi)
+            limit = limit.ceil().int()
+            n_qs = torch.maximum(n_qs, limit)
+            n_qs = torch.clamp(n_qs, max=self.n_q_range[1])
+
+        indices = torch.arange(max_n_q, device=self.device).expand(batch_size, max_n_q)
+        valid_mask = indices < n_qs[:, None] # right side padding
+        
+        if self.shuffle_mask: # shuffle valid positions (inter-spread padding)
+            perm = torch.argsort(torch.rand(batch_size, max_n_q, device=self.device), dim=-1)
+            valid_mask = torch.gather(valid_mask, dim=1, index=perm)
+
+        context['key_padding_mask'] = valid_mask
+        context['n_points'] = valid_mask.sum(dim=-1)
+
+        return q
+    
+    def scale_q(self, q):
+        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)

reflectorch/data_generation/reflectivity/__init__.py

@@ -0,0 +1,102 @@
+# -*- 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.smearing_pointwise import abeles_pointwise_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 = False,
+        log: bool = False,
+        q_shift: Tensor = 0.0,
+        r_scale: Tensor = 1.0,
+        background: Tensor = 0.0,
+        solvent_vf = None,
+        solvent_mode = 'fronting',
+        abeles_func = None,
+        **abeles_kwargs
+):
+    """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): 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): 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 False.
+        log (bool, optional): if True the base 10 logarithm of the reflectivity curves is returned. Defaults to False.
+        q_shift (float or Tensor, optional): misalignment in q.
+        r_scale (float or Tensor, optional): normalization factor (scales reflectivity).
+        background (float or Tensor, optional): background intensity.
+        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.
+        abeles_kwargs: Additional arguments specific to the chosen `abeles_func`.
+    Returns:
+        Tensor: the computed reflectivity curves
+    """
+    abeles_func = abeles_func or abeles
+    q = torch.atleast_2d(q) + q_shift
+    q = torch.clamp(q, min=0.0)
+    
+    if solvent_vf is not None:
+        num_layers = thickness.shape[-1]
+        if solvent_mode == 'fronting':
+            assert sld.shape[-1] == num_layers + 2
+            assert solvent_vf.shape[-1] == num_layers
+            solvent_sld = sld[..., [0]]
+            idx = slice(1, num_layers)
+            sld[..., idx] = solvent_vf * solvent_sld + (1.0 - solvent_vf) * sld[..., idx]
+        elif solvent_mode == 'backing':
+            solvent_sld = sld[..., [-1]]
+            idx = slice(1, num_layers) if sld.shape[-1] == num_layers + 2 else slice(0, num_layers)
+            sld[..., idx] = solvent_vf * solvent_sld + (1.0 - solvent_vf) * sld[..., idx]
+        else:
+            raise NotImplementedError
+
+    if dq is None:
+        reflectivity_curves = abeles_func(q, thickness, roughness, sld, **abeles_kwargs)
+    else:
+        if dq.shape[-1] > 1:
+            reflectivity_curves = abeles_pointwise_smearing(
+                    q=q, dq=dq, thickness=thickness, roughness=roughness, sld=sld, 
+                    abeles_func=abeles_func, gauss_num=gauss_num,
+                    **abeles_kwargs,
+            )
+        else:
+            reflectivity_curves = abeles_constant_smearing(
+                q, thickness, roughness, sld,
+                dq=dq, gauss_num=gauss_num, constant_dq=constant_dq, abeles_func=abeles_func,
+                **abeles_kwargs,
+            )
+
+    if isinstance(r_scale, Tensor):
+        r_scale = r_scale.view(-1, *[1] * (reflectivity_curves.dim() - 1))
+    if isinstance(background, Tensor):
+        background = background.view(-1, *[1] * (reflectivity_curves.dim() - 1))
+
+    reflectivity_curves = reflectivity_curves * r_scale + background
+
+    if log:
+        reflectivity_curves = torch.log10(reflectivity_curves)
+
+    return reflectivity_curves

reflectorch/data_generation/reflectivity/abeles.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+import math
+from functools import reduce
+
+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): 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: 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
+
+    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 - sld[..., :1]) * 1e-6 + 1e-36j
+
+    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 = reduce(torch.matmul, c_matrices)
+
+    r = (m[..., 1, 0] / m[..., 0, 0]).abs() ** 2
+    r = torch.clamp_max_(r, 1.0)
+
+    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,71 @@
+# -*- 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,
+):
+    """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
+
+    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_fresnel_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_fresnel_reflectivity(q, substrate_slds):
+    _RE_CONST = 0.28174103675406496 # 2/sqrt(16*pi)
+
+    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,105 @@
+# -*- 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,
+):
+    """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
+
+    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))