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

reflectorch/data_generation/q_generator.py

@@ -50,6 +50,8 @@ class ConstantQ(QGenerator):
                     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:
@@ -63,14 +65,26 @@ class ConstantQ(QGenerator):
         """
         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*, 
+        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.
@@ -80,12 +94,14 @@ class VariableQ(QGenerator):
                  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
 
@@ -98,14 +114,18 @@ class VariableQ(QGenerator):
         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)
+
+        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
+
+        q = q_min[:, None] + q * (q_max - q_min)[:, None]
         
         return q
     

reflectorch/data_generation/reflectivity/__init__.py

@@ -10,6 +10,7 @@ from reflectorch.data_generation.reflectivity.numpy_implementations import (
     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
 
 
@@ -20,9 +21,15 @@ def reflectivity(
         sld: Tensor,
         dq: Tensor = None,
         gauss_num: int = 51,
-        constant_dq: bool = True,
+        constant_dq: bool = False,
         log: bool = False,
-        abeles_func=None,
+        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.
@@ -31,30 +38,65 @@ def reflectivity(
         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): tensors 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.
+        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 True.
+                                    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.
-        abeles_func (Callable, optional): a function implementing the simulation of the reflectivity curves, if different than the default Abeles matrix implementation. Defaults to None.
-
+        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: tensor containing the simulated reflectivity curves with shape [batch_size, n_points]
+        Tensor: the computed reflectivity curves
     """
     abeles_func = abeles_func or abeles
-    q = torch.atleast_2d(q)
+    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)
+        reflectivity_curves = abeles_func(q, thickness, roughness, sld, **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
-        )
+        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,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import math
+from functools import reduce
 
 import torch
 from torch import Tensor
@@ -17,8 +18,9 @@ def abeles(
         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): tensors 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.
+        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]
@@ -27,11 +29,24 @@ def abeles(
 
     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]
+    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 * 1e-6 + 1e-30j
+    sld = (sld - sld[..., :1]) * 1e-6 + 1e-36j
 
     k_z0 = (q / 2).to(c_dtype)
 
@@ -41,7 +56,7 @@ def abeles(
     if k_z0.dim() == 2:
         k_z0.unsqueeze_(-1)
 
-    k_n = torch.sqrt(k_z0 ** 2 - 4 * math.pi * sld)
+    k_n = torch.sqrt(k_z0**2 - 4 * math.pi * sld)
 
     # k_n.shape - (batch, q, layers)
 
@@ -52,22 +67,22 @@ def abeles(
     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)
+    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 = 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
+    m = reduce(torch.matmul, c_matrices)
 
     r = (m[..., 1, 0] / m[..., 0, 0]).abs() ** 2
-    r = torch.clamp_max_(r, 1.)
+    r = torch.clamp_max_(r, 1.0)
 
     return r
 

reflectorch/data_generation/reflectivity/kinematical.py

@@ -19,7 +19,7 @@ def kinematical_approximation(
         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): tensors containing the layer SLDs (real or complex; 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``.
@@ -51,7 +51,7 @@ def kinematical_approximation(
 
         substrate_sld = sld[:, -1:]
 
-        rf = _get_resnel_reflectivity(q, substrate_sld[:, None])
+        rf = _get_fresnel_reflectivity(q, substrate_sld[:, None])
 
         r = torch.clamp_max_(r * rf / substrate_sld.real ** 2, 1.)
 
@@ -61,12 +61,11 @@ def kinematical_approximation(
     return r
 
 
-def _get_resnel_reflectivity(q, substrate_slds):
-    _RE_CONST = 0.28174103675406496
+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)
-
+    return r_f.squeeze(-1)

reflectorch/data_generation/reflectivity/memory_eff.py

@@ -19,7 +19,7 @@ def abeles_memory_eff(
         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): tensors containing the layer SLDs (real or complex; 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:

reflectorch/data_generation/reflectivity/smearing.py

@@ -14,26 +14,41 @@ def abeles_constant_smearing(
         roughness: Tensor,
         sld: Tensor,
         dq: Tensor = None,
-        gauss_num: int = 51,
-        constant_dq: bool = True,
+        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)
+    
+    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(
-        pad(curves[None], (padding, padding), 'reflect'), kernels[:, None], groups=kernels.shape[0],
-    )[0]
+        padded_curves, kernels[:, None], groups=kernels.shape[0],
+    )
 
     if q.shape[0] != smeared_curves.shape[0]:
-        q = q.expand(smeared_curves.shape[0], *q.shape[1:])
-
+        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
 
 
@@ -55,7 +70,7 @@ def _get_t_gauss_kernels(resolutions: Tensor, gaussnum: int = 51):
     return gauss_y
 
 
-def _get_q_axes(q: Tensor, resolutions: Tensor, gaussnum: int = 51, constant_dq: bool = True):
+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:

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/smearing.py

@@ -9,15 +9,15 @@ class Smearing(object):
     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 (1e-4, 5e-3).
+        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 = (1e-4, 5e-3),
-                 constant_dq: bool = True,
+                 sigma_range: tuple = (0.01, 0.1),
+                 constant_dq: bool = False,
                  gauss_num: int = 31,
                  share_smeared: float = 0.2,
                  ):
@@ -38,31 +38,62 @@ class Smearing(object):
         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 {}
 
-    def get_curves(self, q_values: Tensor, params: BasicParams):
         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, log=False)
-
-        curves = torch.empty(params.batch_size, q_values.shape[-1], device=params.device, dtype=params.dtype)
+            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
 
-            curves[~indices] = params[~indices].reflectivity(q, log=False)
+            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
 
-            curves[indices] = params[indices].reflectivity(
-                q, dq=dq, constant_dq=self.constant_dq, log=False, gauss_num=self.gauss_num
+            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
+        return curves, q_resolutions