Source code for reflectorch.data_generation.utils

# -*- coding: utf-8 -*-
#
#
# This source code is licensed under the GPL license found in the
# LICENSE file in the root directory of this source tree.

from typing import List, Union
from math import sqrt, pi, log10

import torch
from torch import Tensor

__all__ = [
    "get_reversed_params",
    "get_density_profiles",
    "uniform_sampler",
    "logdist_sampler",
    "triangular_sampler",
    "get_param_labels",
    "get_d_rhos",
    "get_slds_from_d_rhos",
]


def uniform_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
    if isinstance(low, Tensor):
        device, dtype = low.device, low.dtype
    return torch.rand(*shape, device=device, dtype=dtype) * (high - low) + low


def logdist_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
    if isinstance(low, Tensor):
        device, dtype = low.device, low.dtype
        low, high = map(torch.log10, (low, high))
    else:
        low, high = map(log10, (low, high))
    return 10 ** (torch.rand(*shape, device=device, dtype=dtype) * (high - low) + low)


def triangular_sampler(low: Union[float, Tensor], high: Union[float, Tensor], *shape, device=None, dtype=None):
    if isinstance(low, Tensor):
        device, dtype = low.device, low.dtype

    x = torch.rand(*shape, device=device, dtype=dtype)

    return (high - low) * (1 - torch.sqrt(x)) + low


def get_reversed_params(thicknesses: Tensor, roughnesses: Tensor, slds: Tensor):
    reversed_slds = torch.cumsum(
        torch.flip(
            torch.diff(
                torch.cat([torch.zeros(slds.shape[0], 1).to(slds), slds], dim=-1),
                dim=-1
            ), (-1,)
        ),
        dim=-1
    )
    reversed_thicknesses = torch.flip(thicknesses, [-1])
    reversed_roughnesses = torch.flip(roughnesses, [-1])
    reversed_params = torch.cat([reversed_thicknesses, reversed_roughnesses, reversed_slds], -1)

    return reversed_params


[docs] def get_density_profiles( thicknesses: Tensor, roughnesses: Tensor, slds: Tensor, z_axis: Tensor = None, num: int = 1000 ): """Generates SLD profiles (and their derivative) based on batches of thicknesses, roughnesses and layer SLDs. The axis has its zero at the top (ambient medium) interface and is positive inside the film. Args: thicknesses (Tensor): the layer thicknesses (top to bottom) roughnesses (Tensor): the interlayer roughnesses (top to bottom) slds (Tensor): the layer SLDs (top to bottom) z_axis (Tensor, optional): a custom depth (z) axis. Defaults to None. num (int, optional): number of discretization points for the profile. Defaults to 1000. Returns: tuple: the z axis, the computed density profile rho(z) and the derivative of the density profile drho/dz(z) """ assert torch.all(roughnesses >= 0), 'Negative roughness happened' assert torch.all(thicknesses >= 0), 'Negative thickness happened' sample_num = thicknesses.shape[0] d_rhos = get_d_rhos(slds) zs = torch.cumsum(torch.cat([torch.zeros(sample_num, 1).to(thicknesses), thicknesses], dim=-1), dim=-1) if z_axis is None: z_axis = torch.linspace(- zs.max() * 0.1, zs.max() * 1.1, num, device=thicknesses.device)[None] elif len(z_axis.shape) == 1: z_axis = z_axis[None] sigmas = roughnesses * sqrt(2) profile = get_erf(z_axis[:, None], zs[..., None], sigmas[..., None], d_rhos[..., None]).sum(1) d_profile = get_gauss(z_axis[:, None], zs[..., None], sigmas[..., None], d_rhos[..., None]).sum(1) z_axis = z_axis[0] return z_axis, profile, d_profile
def get_d_rhos(slds: Tensor) -> Tensor: d_rhos = torch.cat([slds[:, 0][:, None], torch.diff(slds, dim=-1)], -1) return d_rhos def get_slds_from_d_rhos(d_rhos: Tensor) -> Tensor: slds = torch.cumsum(d_rhos, dim=-1) return slds def get_erf(z, z0, sigma, amp): return (torch.erf((z - z0) / sigma) + 1) * amp / 2 def get_gauss(z, z0, sigma, amp): return amp / (sigma * sqrt(2 * pi)) * torch.exp(- (z - z0) ** 2 / 2 / sigma ** 2) 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', ) -> 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