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

reflectorch/data_generation/reflectivity/__init__.py

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

@@ -1,97 +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)
+# -*- 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

@@ -1,71 +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
-
+# -*- 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)