reflectorch 1.2.1__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.
@@ -37,24 +44,59 @@ def reflectivity(
                             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.
+        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/kinematical.py

@@ -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/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

reflectorch/data_generation/utils.py

@@ -57,7 +57,7 @@ def get_reversed_params(thicknesses: Tensor, roughnesses: Tensor, slds: Tensor):
     return reversed_params
 
 
-def get_density_profiles(
+def get_density_profiles_sld(
         thicknesses: Tensor,
         roughnesses: Tensor,
         slds: Tensor,
@@ -120,29 +120,103 @@ def get_erf(z, z0, sigma, amp):
 def get_gauss(z, z0, sigma, amp):
     return amp / (sigma * sqrt(2 * pi)) * torch.exp(- (z - z0) ** 2 / 2 / sigma ** 2)
 
+def get_density_profiles(
+    thicknesses: torch.Tensor,
+    roughnesses: torch.Tensor,
+    slds: torch.Tensor,
+    ambient_sld: torch.Tensor = None,
+    z_axis: torch.Tensor = None,
+    num: int = 1000,
+    padding_left: float = 0.2,
+    padding_right: float = 1.1,
+):
+    """
+    Args:
+        thicknesses (Tensor): finite layer thicknesses.
+        roughnesses (Tensor): interface roughnesses for all transitions (ambient→layer1 ... layerN→substrate).
+        slds (Tensor): SLDs for the finite layers + substrate.
+        ambient_sld (Tensor, optional): SLD for the top ambient. Defaults to 0.0 if None.
+        z_axis (Tensor, optional): a custom depth axis. If None, a linear axis is generated.
+        num (int): number of points in the generated z-axis (if z_axis is None).
+        padding_left (float): factor to extend the negative (above the surface) portion of z-axis.
+        padding_right (float): factor to extend the positive (into the sample) portion of z-axis.
+
+    Returns:
+        (z_axis, profile, d_profile)
+          z_axis:    1D Tensor of shape (num, ) with the depth coordinates.
+          profile:   2D Tensor of shape (batch_size, num) giving the SLD at each depth.
+          d_profile: 2D Tensor of shape (batch_size, num) giving d(SLD)/dz at each depth.
+    """
+
+    bs, n = thicknesses.shape
+    assert roughnesses.shape == (bs, n + 1),  (
+        f"Roughnesses must be (batch_size, num_layers+1). Found {roughnesses.shape} instead."
+    )
+    assert slds.shape == (bs, n + 1), (
+        f"SLDs must be (batch_size, num_layers+1). Found {slds.shape} instead."
+    )
+    assert torch.all(thicknesses >= 0), "Negative thickness encountered."
+    assert torch.all(roughnesses >= 0), "Negative roughness encountered."
+
+    if ambient_sld is None:
+        ambient_sld = torch.zeros((bs, 1), device=thicknesses.device)
+    else:
+        if ambient_sld.ndim == 1:
+            ambient_sld = ambient_sld.unsqueeze(-1)
+
+    slds_all = torch.cat([ambient_sld, slds], dim=-1)  # new dimension: n+2
+    d_rhos = torch.diff(slds_all, dim=-1)  # (bs, n+1)
+
+    interfaces = torch.cat([
+        torch.zeros((bs, 1), device=thicknesses.device),  # z=0 for ambient→layer1
+        thicknesses
+    ], dim=-1).cumsum(dim=-1)  # now shape => (bs, n+1)
+
+    total_thickness = interfaces[..., -1].max()
+    if z_axis is None:
+        z_axis = torch.linspace(
+            -padding_left * total_thickness,
+            padding_right * total_thickness,
+            num,
+            device=thicknesses.device
+        )  # shape => (num,)
+    if z_axis.ndim == 1:
+        z_axis = z_axis.unsqueeze(0)  # shape => (1, num)
+
+    z_b = z_axis.repeat(bs, 1).unsqueeze(1)           # (bs, 1, num)
+    interfaces_b = interfaces.unsqueeze(-1)           # (bs, n+1, 1)
+    sigmas_b = (roughnesses * sqrt(2)).unsqueeze(-1)  # (bs, n+1, 1)
+    d_rhos_b = d_rhos.unsqueeze(-1)                   # (bs, n+1, 1)
+
+    profile = get_erf(z_b, interfaces_b, sigmas_b, d_rhos_b).sum(dim=1)    # (bs, num)
+    if ambient_sld is not None:
+        profile = profile + ambient_sld
+
+    d_profile = get_gauss(z_b, interfaces_b, sigmas_b, d_rhos_b).sum(dim=1) # (bs, num)
+
+    return z_axis.squeeze(0), profile, d_profile
 
 def get_param_labels(
         num_layers: int, *,
         thickness_name: str = 'Thickness',
         roughness_name: str = 'Roughness',
         sld_name: str = 'SLD',
-        substrate_name: str = 'sub',
-) -> List[str]:
-    thickness_labels = [f'{thickness_name} L{num_layers - i}' for i in range(num_layers)]
-    roughness_labels = [f'{roughness_name} L{num_layers - i}' for i in range(num_layers)] + [f'{roughness_name} {substrate_name}']
-    sld_labels = [f'{sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{sld_name} {substrate_name}']
-    return thickness_labels + roughness_labels + sld_labels
-
-def get_param_labels_absorption_model(
-        num_layers: int, *,
-        thickness_name: str = 'Thickness',
-        roughness_name: str = 'Roughness',
-        real_sld_name: str = 'SLD real',
         imag_sld_name: str = 'SLD imag',
         substrate_name: str = 'sub',
+        parameterization_type: str = 'standard',
+        number_top_to_bottom: bool = False,
 ) -> List[str]:
-    thickness_labels = [f'{thickness_name} L{num_layers - i}' for i in range(num_layers)]
-    roughness_labels = [f'{roughness_name} L{num_layers - i}' for i in range(num_layers)] + [f'{roughness_name} {substrate_name}']
-    real_sld_labels = [f'{real_sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{real_sld_name} {substrate_name}']
-    imag_sld_labels = [f'{imag_sld_name} L{num_layers - i}' for i in range(num_layers)] + [f'{imag_sld_name} {substrate_name}']
-    return thickness_labels + roughness_labels + real_sld_labels + imag_sld_labels
+    def pos(i):
+        return i + 1 if number_top_to_bottom else num_layers - i
+        
+    thickness_labels = [f'{thickness_name} L{pos(i)}' for i in range(num_layers)]
+    roughness_labels = [f'{roughness_name} L{pos(i)}' for i in range(num_layers)] + [f'{roughness_name} {substrate_name}']
+    sld_labels = [f'{sld_name} L{pos(i)}' for i in range(num_layers)] + [f'{sld_name} {substrate_name}']
+
+    all_labels = thickness_labels + roughness_labels + sld_labels
+
+    if parameterization_type == 'absorption':
+        imag_sld_labels = [f'{imag_sld_name} L{pos(i)}' for i in range(num_layers)] + [f'{imag_sld_name} {substrate_name}']
+        all_labels = all_labels + imag_sld_labels
+        
+    return all_labels

reflectorch/extensions/refnx/refnx_conversion.py

@@ -0,0 +1,77 @@
+import numpy as np
+from functools import reduce
+from operator import or_
+
+from reflectorch.inference.inference_model import EasyInferenceModel
+from reflectorch import BasicParams
+
+import refnx
+from refnx.dataset import ReflectDataset, Data1D
+from refnx.analysis import Transform, CurveFitter, Objective, Model, Parameter
+from refnx.reflect import SLD, Slab, ReflectModel
+
+def covert_reflectorch_prediction_to_refnx_structure(inference_model: EasyInferenceModel, pred_params_object: BasicParams,  prior_bounds: np.array):
+    assert inference_model.trainer.loader.prior_sampler.param_model.__class__.__name__ == 'StandardModel'
+
+    n_layers = inference_model.trainer.loader.prior_sampler.max_num_layers
+    init_thicknesses = pred_params_object.thicknesses.squeeze().tolist()
+    init_roughnesses = pred_params_object.roughnesses.squeeze().tolist()
+    init_slds = pred_params_object.slds.squeeze().tolist()
+
+    sld_objects = []
+
+    for sld in init_slds:
+        sld_objects.append(SLD(value=sld))
+
+    layer_objects = [SLD(0)()]
+    for i in range(n_layers):
+        layer_objects.append(sld_objects[i](init_thicknesses[i], init_roughnesses[i]))
+
+    layer_objects.append(sld_objects[-1](0, init_roughnesses[-1]))
+
+    thickness_bounds = prior_bounds[:n_layers]
+    roughness_bounds = prior_bounds[n_layers:2*n_layers+1]
+    sld_bounds = prior_bounds[2*n_layers+1:]
+
+    for i, layer in enumerate(layer_objects):
+        if i == 0:
+            print("Ambient (air)")
+            print(80 * '-')
+        elif i < n_layers+1:
+            layer.thick.setp(bounds=thickness_bounds[i-1], vary=True)
+            layer.rough.setp(bounds=roughness_bounds[i-1], vary=True)
+            layer.sld.real.setp(bounds=sld_bounds[i-1], vary=True)
+
+            print(f'Layer {i}')
+            print(f'Thickness: value {layer.thick.value}, vary {layer.thick.vary}, bounds {layer.thick.bounds}')
+            print(f'Roughness: value {layer.rough.value}, vary {layer.rough.vary}, bounds {layer.rough.bounds}')
+            print(f'SLD: value {layer.sld.real.value}, vary {layer.sld.real.vary}, bounds {layer.sld.real.bounds}')
+            print(80 * '-')
+        else: #substrate
+            layer.rough.setp(bounds=roughness_bounds[i-1], vary=True)
+            layer.sld.real.setp(bounds=sld_bounds[i-1], vary=True)
+
+            print(f'Substrate')
+            print(f'Thickness: value {layer.thick.value}, vary {layer.thick.vary}, bounds {layer.thick.bounds}')
+            print(f'Roughness: value {layer.rough.value}, vary {layer.rough.vary}, bounds {layer.rough.bounds}')
+            print(f'SLD: value {layer.sld.real.value}, vary {layer.sld.real.vary}, bounds {layer.sld.real.bounds}')
+    
+    refnx_structure = reduce(or_, layer_objects)
+
+    return refnx_structure
+
+
+###Example usage:
+# refnx_structure = covert_reflectorch_prediction_to_refnx_structure(inference_model, pred_params_object, prior_bounds)
+
+# refnx_reflect_model = ReflectModel(refnx_structure, bkg=1e-10, dq=0.0)
+# refnx_reflect_model.scale.setp(bounds=(0.8, 1.2), vary=True)
+# refnx_reflect_model.q_offset.setp(bounds=(-0.01, 0.01), vary=True)
+# refnx_reflect_model.bkg.setp(bounds=(1e-10, 1e-8), vary=True)
+
+
+# data = Data1D(data=(q_model, exp_curve_interp))
+
+# refnx_objective = Objective(refnx_reflect_model, data, transform=Transform("logY"))
+# fitter = CurveFitter(refnx_objective)
+# fitter.fit('least_squares')