reflectorch 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

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,104 @@ 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)
+        ambient_sld = ambient_sld.expand(bs, 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')