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

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')