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

reflectorch/data_generation/reflectivity/memory_eff.py

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

reflectorch/data_generation/reflectivity/numpy_implementations.py

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