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

reflectorch/data_generation/priors/multilayer_models.py

@@ -1,311 +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
+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