reflectorch 1.5.1__py3-none-any.whl

reflectorch/data_generation/priors/multilayer_models.py

@@ -0,0 +1,311 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+__all__ = [
+    "MULTILAYER_MODELS",
+    "MultilayerModel",
+]
+
+
+class MultilayerModel(object):
+    NAME: str = ''
+    PARAMETER_NAMES: Tuple[str, ...]
+
+    def __init__(self, max_num_layers: int):
+        self.max_num_layers = max_num_layers
+
+    def to_standard_params(self, parametrized_model: Tensor) -> dict:
+        raise NotImplementedError
+
+    def from_standard_params(self, params: dict) -> Tensor:
+        raise NotImplementedError
+
+
+class BasicMultilayerModel1(MultilayerModel):
+    NAME = 'repeating_multilayer_v1'
+
+    PARAMETER_NAMES = (
+        "d_full_rel",
+        "rel_sigmas",
+        "d_block",
+        "s_block_rel",
+        "r_block",
+        "dr",
+        "d3_rel",
+        "s3_rel",
+        "r3",
+        "d_sio2",
+        "s_sio2",
+        "s_si",
+        "r_sio2",
+        "r_si",
+    )
+
+    def to_standard_params(self, parametrized_model: Tensor) -> dict:
+        return multilayer_model1(parametrized_model, self.max_num_layers)
+
+
+class BasicMultilayerModel2(MultilayerModel):
+    NAME = 'repeating_multilayer_v2'
+
+    PARAMETER_NAMES = (
+        "d_full_rel",
+        "rel_sigmas",
+        "dr_sigmoid_rel_pos",
+        "dr_sigmoid_rel_width",
+        "d_block",
+        "s_block_rel",
+        "r_block",
+        "dr",
+        "d3_rel",
+        "s3_rel",
+        "r3",
+        "d_sio2",
+        "s_sio2",
+        "s_si",
+        "r_sio2",
+        "r_si",
+    )
+
+    def to_standard_params(self, parametrized_model: Tensor) -> dict:
+        return multilayer_model2(parametrized_model, self.max_num_layers)
+
+
+class BasicMultilayerModel3(MultilayerModel):
+    NAME = 'repeating_multilayer_v3'
+
+    PARAMETER_NAMES = (
+        "d_full_rel",
+        "rel_sigmas",
+        "dr_sigmoid_rel_pos",
+        "dr_sigmoid_rel_width",
+        "d_block1_rel",
+        "d_block",
+        "s_block_rel",
+        "r_block",
+        "dr",
+        "d3_rel",
+        "s3_rel",
+        "r3",
+        "d_sio2",
+        "s_sio2",
+        "s_si",
+        "r_sio2",
+        "r_si",
+    )
+
+    def to_standard_params(self, parametrized_model: Tensor) -> dict:
+        return multilayer_model3(parametrized_model, self.max_num_layers)
+
+
+MULTILAYER_MODELS = {
+    'repeating_multilayer_v1': BasicMultilayerModel1,
+    'repeating_multilayer_v2': BasicMultilayerModel2,
+    'repeating_multilayer_v3': BasicMultilayerModel3,
+}
+
+
+def multilayer_model1(parametrized_model: Tensor, d_full_rel_max: int = 50) -> dict:
+    n = d_full_rel_max
+
+    (
+        d_full_rel,
+        rel_sigmas,
+        d_block,
+        s_block_rel,
+        r_block,
+        dr,
+        d3_rel,
+        s3_rel,
+        r3,
+        d_sio2,
+        s_sio2,
+        s_si,
+        r_sio2,
+        r_si,
+    ) = parametrized_model.T
+
+    batch_size = parametrized_model.shape[0]
+
+    r_positions = 2 * n - torch.arange(2 * n, dtype=dr.dtype, device=dr.device)[None].repeat(batch_size, 1)
+
+    r_modulations = torch.sigmoid(-(r_positions - 2 * d_full_rel[..., None]) / rel_sigmas[..., None])
+
+    r_block = r_block[:, None].repeat(1, n)
+    dr = dr[:, None].repeat(1, n)
+
+    sld_blocks = torch.stack([r_block, r_block + dr], -1).flatten(1)
+
+    sld_blocks = r_modulations * sld_blocks
+
+    d3 = d3_rel * d_block
+
+    thicknesses = torch.cat(
+        [(d_block / 2)[:, None].repeat(1, n * 2), d3[:, None], d_sio2[:, None]], -1
+    )
+
+    s_block = s_block_rel * d_block
+
+    roughnesses = torch.cat(
+        [s_block[:, None].repeat(1, n * 2), (s3_rel * d3)[:, None], s_sio2[:, None], s_si[:, None]], -1
+    )
+
+    slds = torch.cat(
+        [sld_blocks, r3[:, None], r_sio2[:, None], r_si[:, None]], -1
+    )
+
+    params = dict(
+        thicknesses=thicknesses,
+        roughnesses=roughnesses,
+        slds=slds
+    )
+    return params
+
+
+def multilayer_model2(parametrized_model: Tensor, d_full_rel_max: int = 50) -> dict:
+    n = d_full_rel_max
+
+    (
+        d_full_rel,
+        rel_sigmas,
+        dr_sigmoid_rel_pos,
+        dr_sigmoid_rel_width,
+        d_block,
+        s_block_rel,
+        r_block,
+        dr,
+        d3_rel,
+        s3_rel,
+        r3,
+        d_sio2,
+        s_sio2,
+        s_si,
+        r_sio2,
+        r_si,
+    ) = parametrized_model.T
+
+    batch_size = parametrized_model.shape[0]
+
+    r_positions = 2 * n - torch.arange(2 * n, dtype=dr.dtype, device=dr.device)[None].repeat(batch_size, 1)
+
+    r_modulations = torch.sigmoid(-(r_positions - 2 * d_full_rel[..., None]) / rel_sigmas[..., None])
+
+    r_block = r_block[:, None].repeat(1, n)
+    dr = dr[:, None].repeat(1, n)
+
+    dr_positions = r_positions[:, ::2]
+
+    dr_modulations = torch.sigmoid(
+        -(dr_positions - (2 * d_full_rel * dr_sigmoid_rel_pos)[..., None]) / dr_sigmoid_rel_width[..., None]
+    )
+
+    dr = dr * dr_modulations
+
+    sld_blocks = torch.stack([r_block, r_block + dr], -1).flatten(1)
+
+    sld_blocks = r_modulations * sld_blocks
+
+    d3 = d3_rel * d_block
+
+    thicknesses = torch.cat(
+        [(d_block / 2)[:, None].repeat(1, n * 2), d3[:, None], d_sio2[:, None]], -1
+    )
+
+    s_block = s_block_rel * d_block
+
+    roughnesses = torch.cat(
+        [s_block[:, None].repeat(1, n * 2), (s3_rel * d3)[:, None], s_sio2[:, None], s_si[:, None]], -1
+    )
+
+    slds = torch.cat(
+        [sld_blocks, r3[:, None], r_sio2[:, None], r_si[:, None]], -1
+    )
+
+    params = dict(
+        thicknesses=thicknesses,
+        roughnesses=roughnesses,
+        slds=slds
+    )
+    return params
+
+
+def multilayer_model3(parametrized_model: Tensor, d_full_rel_max: int = 30):
+    n = d_full_rel_max
+
+    (
+        d_full_rel,
+        rel_sigmas,
+        dr_sigmoid_rel_pos,
+        dr_sigmoid_rel_width,
+        d_block1_rel,
+        d_block,
+        s_block_rel,
+        r_block,
+        dr,
+        d3_rel,
+        s3_rel,
+        r3,
+        d_sio2,
+        s_sio2,
+        s_si,
+        r_sio2,
+        r_si,
+    ) = parametrized_model.T
+
+    batch_size = parametrized_model.shape[0]
+
+    r_positions = 2 * n - torch.arange(2 * n, dtype=dr.dtype, device=dr.device)[None].repeat(batch_size, 1)
+
+    r_modulations = torch.sigmoid(
+        -(
+                r_positions - 2 * d_full_rel[..., None]
+        ) / rel_sigmas[..., None]
+    )
+
+    dr_positions = r_positions[:, ::2]
+
+    dr_modulations = dr[..., None] * (1 - torch.sigmoid(
+        -(
+                dr_positions - 2 * d_full_rel[..., None] + 2 * dr_sigmoid_rel_pos[..., None]
+        ) / dr_sigmoid_rel_width[..., None]
+    ))
+
+    r_block = r_block[..., None].repeat(1, n)
+    dr = dr[..., None].repeat(1, n)
+
+    sld_blocks = torch.stack(
+        [
+            r_block + dr_modulations * (1 - d_block1_rel[..., None]),
+            r_block + dr - dr_modulations * d_block1_rel[..., None]
+        ], -1).flatten(1)
+
+    sld_blocks = r_modulations * sld_blocks
+
+    d3 = d3_rel * d_block
+
+    d1, d2 = d_block * d_block1_rel, d_block * (1 - d_block1_rel)
+
+    thickness_blocks = torch.stack([d1[:, None].repeat(1, n), d2[:, None].repeat(1, n)], -1).flatten(1)
+
+    thicknesses = torch.cat(
+        [thickness_blocks, d3[:, None], d_sio2[:, None]], -1
+    )
+
+    s_block = s_block_rel * d_block
+
+    roughnesses = torch.cat(
+        [s_block[:, None].repeat(1, n * 2), (s3_rel * d3)[:, None], s_sio2[:, None], s_si[:, None]], -1
+    )
+
+    slds = torch.cat(
+        [sld_blocks, r3[:, None], r_sio2[:, None], r_si[:, None]], -1
+    )
+
+    params = dict(
+        thicknesses=thicknesses,
+        roughnesses=roughnesses,
+        slds=slds
+    )
+    return params

reflectorch/data_generation/priors/multilayer_structures.py

@@ -0,0 +1,104 @@
+from typing import Tuple, Dict
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.priors.base import PriorSampler
+from reflectorch.data_generation.priors.params import Params
+from reflectorch.data_generation.priors.no_constraints import (
+    DEFAULT_DEVICE,
+    DEFAULT_DTYPE,
+)
+
+from reflectorch.data_generation.priors.multilayer_models import MULTILAYER_MODELS, MultilayerModel
+from reflectorch.utils import to_t
+
+
+class MultilayerStructureParams(Params):
+    pass
+
+
+class SimpleMultilayerSampler(PriorSampler):
+    PARAM_CLS = MultilayerStructureParams
+
+    def __init__(self,
+                 params: Dict[str, Tuple[float, float]],
+                 model_name: str,
+                 device: torch.device = DEFAULT_DEVICE,
+                 dtype: torch.dtype = DEFAULT_DTYPE,
+                 max_num_layers: int = 50,
+                 ):
+        self.multilayer_model: MultilayerModel = MULTILAYER_MODELS[model_name](max_num_layers)
+        self.device = device
+        self.dtype = dtype
+        self.num_layers = max_num_layers
+        ordered_bounds = [params[k] for k in self.multilayer_model.PARAMETER_NAMES]
+        self._np_bounds = np.array(ordered_bounds).T
+        self.min_bounds, self.max_bounds = torch.tensor(ordered_bounds, device=device, dtype=dtype).T[:, None]
+        self._param_dim = len(params)
+
+    @property
+    def max_num_layers(self) -> int:
+        return self.num_layers
+
+    @property
+    def param_dim(self) -> int:
+        return self._param_dim
+
+    def sample(self, batch_size: int) -> MultilayerStructureParams:
+        return self.optimized_sample(batch_size)[0]
+
+    def optimized_sample(self, batch_size: int) -> Tuple[MultilayerStructureParams, Tensor]:
+        scaled_params = torch.rand(
+            batch_size,
+            self.min_bounds.shape[-1],
+            device=self.min_bounds.device,
+            dtype=self.min_bounds.dtype,
+        )
+
+        targets = self.restore_params(scaled_params)
+
+        return targets, scaled_params
+
+    def get_np_bounds(self):
+        return np.array(self._np_bounds)
+
+    def restore_np_params(self, params: np.ndarray):
+        p = self.multilayer_model.to_standard_params(
+            torch.atleast_2d(to_t(params))
+        )
+
+        return {
+            'thickness': p['thicknesses'].squeeze().cpu().numpy(),
+            'roughness': p['roughnesses'].squeeze().cpu().numpy(),
+            'sld': p['slds'].squeeze().cpu().numpy()
+        }
+
+    def restore_params2parametrized(self, scaled_params: Tensor) -> Tensor:
+        return scaled_params * (self.max_bounds - self.min_bounds) + self.min_bounds
+
+    def restore_params(self, scaled_params: Tensor) -> MultilayerStructureParams:
+        return self.to_standard_params(self.restore_params2parametrized(scaled_params))
+
+    def to_standard_params(self, params: Tensor) -> MultilayerStructureParams:
+        return MultilayerStructureParams(**self.multilayer_model.to_standard_params(params))
+
+    def scale_params(self, params: Params) -> Tensor:
+        raise NotImplementedError
+
+    def log_prob(self, params: Params) -> Tensor:
+        raise NotImplementedError
+
+    def get_indices_within_domain(self, params: Params) -> Tensor:
+        raise NotImplementedError
+
+    def get_indices_within_bounds(self, params: Params) -> Tensor:
+        raise NotImplementedError
+
+    def filter_params(self, params: Params) -> Params:
+        indices = self.get_indices_within_domain(params)
+        return params[indices]
+
+    def clamp_params(self, params: Params) -> Params:
+        raise NotImplementedError

reflectorch/data_generation/priors/no_constraints.py

@@ -0,0 +1,206 @@
+import logging
+from functools import lru_cache
+from typing import Tuple
+from math import sqrt
+
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.utils import (
+    get_slds_from_d_rhos,
+    uniform_sampler,
+)
+
+from reflectorch.data_generation.priors.utils import (
+    get_allowed_roughness_indices,
+    generate_roughnesses,
+    params_within_bounds,
+)
+from reflectorch.data_generation.priors.scaler_mixin import ScalerMixin
+from reflectorch.data_generation.priors.base import PriorSampler
+from reflectorch.data_generation.priors.params import Params
+
+__all__ = [
+    "BasicPriorSampler",
+    "DEFAULT_ROUGHNESS_RANGE",
+    "DEFAULT_THICKNESS_RANGE",
+    "DEFAULT_SLD_RANGE",
+    "DEFAULT_NUM_LAYERS",
+    "DEFAULT_DEVICE",
+    "DEFAULT_DTYPE",
+    "DEFAULT_SCALED_RANGE",
+    "DEFAULT_USE_DRHO",
+]
+
+DEFAULT_THICKNESS_RANGE: Tuple[float, float] = (1., 500.)
+DEFAULT_ROUGHNESS_RANGE: Tuple[float, float] = (0., 50.)
+DEFAULT_SLD_RANGE: Tuple[float, float] = (-10., 30.)
+DEFAULT_NUM_LAYERS: int = 5
+DEFAULT_USE_DRHO: bool = False
+DEFAULT_DEVICE: torch.device = torch.device('cuda')
+DEFAULT_DTYPE: torch.dtype = torch.float64
+DEFAULT_SCALED_RANGE: Tuple[float, float] = (-sqrt(3.), sqrt(3.))
+
+
+class BasicPriorSampler(PriorSampler, ScalerMixin):
+    """Prior samplers for thicknesses, roughnesses and slds"""
+    def __init__(self,
+                 thickness_range: Tuple[float, float] = DEFAULT_THICKNESS_RANGE,
+                 roughness_range: Tuple[float, float] = DEFAULT_ROUGHNESS_RANGE,
+                 sld_range: Tuple[float, float] = DEFAULT_SLD_RANGE,
+                 num_layers: int = DEFAULT_NUM_LAYERS,
+                 use_drho: bool = DEFAULT_USE_DRHO,
+                 device: torch.device = DEFAULT_DEVICE,
+                 dtype: torch.dtype = DEFAULT_DTYPE,
+                 scaled_range: Tuple[float, float] = DEFAULT_SCALED_RANGE,
+                 restrict_roughnesses: bool = True,
+                 ):
+        self.logger = logging.getLogger(__name__)
+        self.thickness_range = thickness_range
+        self.roughness_range = roughness_range
+        self.sld_range = sld_range
+        self.num_layers = num_layers
+        self.device = device
+        self.dtype = dtype
+        self.scaled_range = scaled_range
+        self.use_drho = use_drho
+        self.restrict_roughnesses = restrict_roughnesses
+
+    @property
+    def max_num_layers(self) -> int:
+        return self.num_layers
+
+    @lru_cache()
+    def min_vector(self, layers_num, drho: bool = False):
+        if drho:
+            sld_min = self.sld_range[0] - self.sld_range[1]
+        else:
+            sld_min = self.sld_range[0]
+
+        return torch.tensor(
+            [self.thickness_range[0]] * layers_num +
+            [self.roughness_range[0]] * (layers_num + 1) +
+            [sld_min] * (layers_num + 1),
+            device=self.device,
+            dtype=self.dtype
+        )
+
+    @lru_cache()
+    def max_vector(self, layers_num, drho: bool = False):
+        if drho:
+            sld_max = self.sld_range[1] - self.sld_range[0]
+        else:
+            sld_max = self.sld_range[1]
+        return torch.tensor(
+            [self.thickness_range[1]] * layers_num +
+            [self.roughness_range[1]] * (layers_num + 1) +
+            [sld_max] * (layers_num + 1),
+            device=self.device,
+            dtype=self.dtype
+        )
+
+    @lru_cache()
+    def delta_vector(self, layers_num, drho: bool = False):
+        return self._get_delta_vector(self.min_vector(layers_num, drho), self.max_vector(layers_num, drho))
+
+    def restore_params(self, scaled_params: Tensor) -> Params:
+        layers_num = self.PARAM_CLS.size2layers_num(scaled_params.shape[-1])
+
+        params_t = self._restore(
+            scaled_params,
+            self.min_vector(layers_num, drho=self.use_drho).to(scaled_params),
+            self.max_vector(layers_num, drho=self.use_drho).to(scaled_params),
+        )
+
+        params = self.PARAM_CLS.from_tensor(params_t)
+
+        if self.use_drho:
+            params.slds = get_slds_from_d_rhos(params.slds)
+
+        return params
+
+    def scale_params(self, params: Params) -> Tensor:
+        layers_num = params.max_layer_num
+
+        return self._scale(
+            params.as_tensor(use_drho=self.use_drho),
+            self.min_vector(layers_num, drho=self.use_drho).to(params.thicknesses),
+            self.max_vector(layers_num, drho=self.use_drho).to(params.thicknesses),
+        )
+
+    def get_indices_within_bounds(self, params: Params) -> Tensor:
+        layer_num = params.max_layer_num
+
+        return params_within_bounds(
+            params.as_tensor(),
+            self.min_vector(layer_num),
+            self.max_vector(layer_num),
+        )
+
+    def get_indices_within_domain(self, params: Params) -> Tensor:
+        if self.restrict_roughnesses:
+            indices = (
+                    self.get_indices_within_bounds(params) &
+                    self.get_allowed_roughness_indices(params)
+            )
+        else:
+            indices = self.get_indices_within_bounds(params)
+        return indices
+
+    def clamp_params(self, params: Params) -> Params:
+        layer_num = params.max_layer_num
+        params = params.as_tensor()
+        params = torch.clamp(
+            params,
+            self.min_vector(layer_num),
+            self.max_vector(layer_num),
+        )
+        params = Params.from_tensor(params)
+        return params
+
+    @staticmethod
+    def get_allowed_roughness_indices(params: Params) -> Tensor:
+        return get_allowed_roughness_indices(params.thicknesses, params.roughnesses)
+
+    def log_prob(self, params: Params) -> Tensor:
+        # so far we ignore non-uniform distribution of roughnesses and slds.
+        log_prob = torch.zeros(params.batch_size, device=params.device, dtype=params.dtype)
+        indices = self.get_indices_within_bounds(params)
+        log_prob[~indices] = float('-inf')
+        return log_prob
+
+    def sample(self, batch_size: int) -> Params:
+        slds = self.generate_slds(batch_size)
+        thicknesses = self.generate_thicknesses(batch_size)
+        roughnesses = self.generate_roughnesses(thicknesses)
+
+        params = Params(thicknesses, roughnesses, slds)
+
+        return params
+
+    def generate_slds(self, batch_size: int):
+        return uniform_sampler(
+            *self.sld_range, batch_size,
+            self.num_layers + 1,
+            device=self.device,
+            dtype=self.dtype
+        )
+
+    def generate_thicknesses(self, batch_size: int):
+        return uniform_sampler(
+            *self.thickness_range, batch_size,
+            self.num_layers,
+            device=self.device,
+            dtype=self.dtype
+        )
+
+    def generate_roughnesses(self, thicknesses: Tensor) -> Tensor:
+        if self.restrict_roughnesses:
+            return generate_roughnesses(thicknesses, self.roughness_range)
+        else:
+            return uniform_sampler(
+                *self.roughness_range, thicknesses.shape[0],
+                self.num_layers + 1,
+                device=self.device,
+                dtype=self.dtype
+            )