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

reflectorch/data_generation/__init__.py

@@ -41,6 +41,7 @@ from reflectorch.data_generation.noise import (
     ScalingNoise,
     BackgroundNoise,
     BasicExpIntensityNoise,
+    GaussianExpIntensityNoise,
     BasicQNoiseGenerator,
 )
 from reflectorch.data_generation.scale_curves import (
@@ -111,6 +112,7 @@ __all__ = [
     "LogLikelihood",
     "PoissonLogLikelihood",
     "BasicExpIntensityNoise",
+    "GaussianExpIntensityNoise",
     "BasicQNoiseGenerator",
     "ConstantAngle",
     "SubpriorParametricSampler",

reflectorch/data_generation/dataset.py

@@ -25,6 +25,7 @@ class BasicDataset(object):
         curves_scaler (CurvesScaler, optional): the reflectivity curve scaler. Defaults to an instance of LogAffineCurvesScaler, 
                                                 which scales the curves to the range [-1, 1], the minimum considered intensity being 1e-10.
         calc_denoised_curves (bool, optional): whether to add the curves without noise to the dictionary. Defaults to False.
+        calc_nonsmeared_curves (bool, optional): whether to add the curves without smearing to the dictionary (only relevant when smearing is applied). Defaults to False.
         smearing (Smearing, optional): curve smearing generator. Defaults to None.
     """
     def __init__(self,
@@ -34,6 +35,7 @@ class BasicDataset(object):
                  q_noise: QNoiseGenerator = None,
                  curves_scaler: CurvesScaler = None,
                  calc_denoised_curves: bool = False,
+                 calc_nonsmeared_curves: bool = False,
                  smearing: Smearing = None,
                  ):
         self.q_generator = q_generator
@@ -43,6 +45,7 @@ class BasicDataset(object):
         self.prior_sampler = prior_sampler
         self.smearing = smearing
         self.calc_denoised_curves = calc_denoised_curves
+        self.calc_nonsmeared_curves = calc_nonsmeared_curves
 
     def update_batch_data(self, batch_data: BATCH_DATA_TYPE) -> None:
         """implement in a subclass to edit batch_data dict inplace"""
@@ -74,7 +77,15 @@ class BasicDataset(object):
 
         batch_data['q_values'] = q_values
 
-        curves = self._calc_curves(q_values, params)
+        refl_kwargs = {}
+
+        curves, q_resolutions, nonsmeared_curves = self._calc_curves(q_values, params, refl_kwargs)
+
+        if torch.is_tensor(q_resolutions):
+            batch_data['q_resolutions'] = q_resolutions
+
+        if torch.is_tensor(nonsmeared_curves):
+            batch_data['nonsmeared_curves'] = nonsmeared_curves
 
         if self.calc_denoised_curves:
             batch_data['curves'] = curves
@@ -88,10 +99,13 @@ class BasicDataset(object):
         batch_data['scaled_noisy_curves'] = scaled_noisy_curves
 
         is_finite = torch.all(torch.isfinite(scaled_noisy_curves), -1)
+
         if not torch.all(is_finite).item():
             infinite_indices = ~is_finite
-            warnings.warn(f'Batch with {infinite_indices.sum().item()} curves with infinities skipped.')
-            return self.get_batch(batch_size = batch_size)
+            to_recalculate = infinite_indices.sum().item()
+            warnings.warn(f'Infinite number appeared in the curve simulation! Recalculate {to_recalculate} curves.')
+            recalculated_batch_data = self.get_batch(to_recalculate)
+            _insert_batch_data(batch_data, recalculated_batch_data, infinite_indices)
 
         is_finite = torch.all(torch.isfinite(batch_data['scaled_noisy_curves']), -1)
         assert torch.all(is_finite).item()
@@ -100,13 +114,19 @@ class BasicDataset(object):
 
         return batch_data
 
-    def _calc_curves(self, q_values: Tensor, params: BasicParams):
+    def _calc_curves(self, q_values: Tensor, params: BasicParams, refl_kwargs):
+        nonsmeared_curves = None
+
         if self.smearing:
-            curves = self.smearing.get_curves(q_values, params)
+            if self.calc_nonsmeared_curves:
+                nonsmeared_curves = params.reflectivity(q_values, **refl_kwargs)
+            curves, q_resolutions = self.smearing.get_curves(q_values, params, refl_kwargs)
         else:
-            curves = params.reflectivity(q_values)
+            curves = params.reflectivity(q_values, **refl_kwargs)
+            q_resolutions = None
+
         curves = curves.to(q_values)
-        return curves
+        return curves, q_resolutions, nonsmeared_curves
 
 
 def _insert_batch_data(tgt_batch_data, add_batch_data, indices):

reflectorch/data_generation/noise.py

@@ -1,11 +1,11 @@
-from typing import Union, Tuple
+from typing import List, Union, Tuple
 from math import log10
 
 import torch
 from torch import Tensor
 
 from reflectorch.data_generation.process_data import ProcessData
-from reflectorch.data_generation.utils import uniform_sampler
+from reflectorch.data_generation.utils import logdist_sampler, uniform_sampler
 
 __all__ = [
     "QNoiseGenerator",
@@ -18,6 +18,7 @@ __all__ = [
     "ShiftNoise",
     "BackgroundNoise",
     "BasicExpIntensityNoise",
+    "GaussianExpIntensityNoise",
     "BasicQNoiseGenerator",
 ]
 
@@ -102,18 +103,22 @@ class BasicQNoiseGenerator(QNoiseGenerator):
                                     Defaults to (0, 1e-3).
     """
     def __init__(self,
+                 apply_systematic_shifts: bool = True,
                  shift_std: float = 1e-3,
+                 apply_gaussian_noise: bool = False,
                  noise_std: Union[float, Tuple[float, float]] = (0, 1e-3),
                  add_to_context: bool = False,
                  ):
-        self.q_shift = QSystematicShiftGenerator(shift_std, add_to_context=add_to_context)
-        self.q_noise = QNormalNoiseGenerator(noise_std, add_to_context=add_to_context)
+        self.q_shift = QSystematicShiftGenerator(shift_std, add_to_context=add_to_context) if apply_systematic_shifts else None
+        self.q_noise = QNormalNoiseGenerator(noise_std, add_to_context=add_to_context) if apply_gaussian_noise else None
 
     def apply(self, qs: Tensor, context: dict = None):
-        """applies random noise to the q values"""
+        """applies noise to the q values"""
         qs = torch.atleast_2d(qs)
-        qs = self.q_shift.apply(qs, context)
-        qs = self.q_noise.apply(qs, context)
+        if self.q_shift:
+            qs = self.q_shift.apply(qs, context)
+        if self.q_noise:
+            qs = self.q_noise.apply(qs, context)
         return qs
 
 
@@ -154,6 +159,51 @@ class MultiplicativeLogNormalNoiseGenerator(IntensityNoiseGenerator):
 
         return noise * curves
 
+class GaussianNoiseGenerator(IntensityNoiseGenerator):
+    """Noise generator which applies noise as R_n = R + eps, with eps~N(0, sigmas) and sigmas = relative_errors * R
+
+    Args:
+        relative_errors (Union[float, Tuple[float, float], List[Tuple[float, float]]])
+        consistent_relative_errors (bool): If True the relative_error is the same for all point of a curve, otherwise it is sampled uniformly.
+    """
+
+    def __init__(self, relative_errors: Union[float, Tuple[float, float], List[float], List[Tuple[float, float]]], 
+                 consistent_rel_err: bool = False,
+                 add_to_context: bool = False):
+        self.relative_errors = relative_errors
+        self.consistent_rel_err = consistent_rel_err
+        self.add_to_context = add_to_context
+
+    def apply(self, curves: Tensor, context: dict = None):
+        """Applies Gaussian noise to the curves."""
+        relative_errors = self.relative_errors
+        num_channels = curves.shape[1] if curves.dim() == 3 else 1
+
+        if isinstance(relative_errors, float):
+            relative_errors = torch.ones_like(curves) * relative_errors
+
+        elif isinstance(relative_errors, (list, tuple)) and isinstance(relative_errors[0], float):
+            if self.consistent_rel_err:
+                relative_errors = uniform_sampler(*relative_errors, curves.shape[0], num_channels, device=curves.device, dtype=curves.dtype)
+                if num_channels > 1:
+                    relative_errors = relative_errors.unsqueeze(-1)
+            else:
+                relative_errors = uniform_sampler(*relative_errors, *curves.shape, device=curves.device, dtype=curves.dtype)
+
+        else:
+            if self.consistent_rel_err:
+                relative_errors = torch.stack([uniform_sampler(*item, curves.shape[0], 1, device=curves.device, dtype=curves.dtype) for item in relative_errors], dim=1)
+            else:
+                relative_errors = torch.stack([uniform_sampler(*item, curves.shape[0], curves.shape[-1], device=curves.device, dtype=curves.dtype) for item in relative_errors], dim=1)
+
+        sigmas = relative_errors * curves
+        noise = torch.normal(mean=0., std=sigmas).clamp_min_(0.0)
+
+        if self.add_to_context and context is not None:
+            context['relative_errors'] = relative_errors
+            context['sigmas'] = sigmas
+            
+        return curves + noise
 
 class PoissonNoiseGenerator(IntensityNoiseGenerator):
     """Noise generator which applies Poisson noise to the reflectivity curves
@@ -273,6 +323,7 @@ class BackgroundNoise(IntensityNoiseGenerator):
 
     Args:
         background_range (tuple, optional): The range from which the background value is sampled. Defaults to (1.0e-10, 1.0e-8).
+        add_to_context (bool, optional): If True, adds generated noise parameters to the context dictionary. Defaults to False.
     """
     def __init__(self,
                  background_range: tuple = (1.0e-10, 1.0e-8),
@@ -283,7 +334,7 @@ class BackgroundNoise(IntensityNoiseGenerator):
 
     def apply(self, curves: Tensor, context: dict = None) -> Tensor:
         """applies background noise to the curves"""
-        backgrounds = uniform_sampler(
+        backgrounds = logdist_sampler(
             *self.background_range, curves.shape[0], 1,
             device=curves.device, dtype=curves.dtype
         )
@@ -294,6 +345,61 @@ class BackgroundNoise(IntensityNoiseGenerator):
 
         return curves
 
+class GaussianExpIntensityNoise(IntensityNoiseGenerator):
+    """
+    A composite noise generator that applies Gaussian, shift and background noise to reflectivity curves.
+
+    This class combines three types of noise:
+    1. Gaussian noise: Applies Gaussian noise (to account for count-based Poisson noise as well as other sources of error)
+    2. Shift noise: Applies a multiplicative shift to the curves, equivalent to a vertical shift in logarithmic space.
+    3. Background noise: Adds a constant background value to the curves.
+
+    Args:
+        relative_errors (Union[float, Tuple[float, float], List[Tuple[float, float]]]): The range of relative errors for Gaussian noise. Defaults to (0.001, 0.15).
+        consistent_rel_err (bool, optional): If True, uses a consistent relative error for Gaussian noise across all points in a curve. Defaults to False.
+        shift_range (tuple, optional): The range of shift factors for shift noise. Defaults to (-0.1, 0.2e-2).
+        background_range (tuple, optional): The range from which the background value is sampled. Defaults to (1.0e-10, 1.0e-8).
+        apply_shift (bool, optional): If True, applies shift noise to the curves. Defaults to False.
+        apply_background (bool, optional): If True, applies background noise to the curves. Defaults to False.
+        same_background_across_channels(bool, optional): If True, the same background is applied to all channels of a multi-channel curve. Defaults to False.
+        add_to_context (bool, optional): If True, adds generated noise parameters to the context dictionary. Defaults to False.
+    """
+    def __init__(self,
+                 relative_errors: Tuple[float, float] = (0.001, 0.15),
+                 consistent_rel_err: bool = False,
+                 apply_shift: bool = False,
+                 shift_range: tuple = (-0.1, 0.2e-2),
+                 apply_background: bool = False,
+                 background_range: tuple = (1.0e-10, 1.0e-8),
+                 same_background_across_channels: bool = False,
+                 add_to_context: bool = False,
+                 ):
+        
+        self.gaussian_noise = GaussianNoiseGenerator(
+            relative_errors=relative_errors,
+            consistent_rel_err=consistent_rel_err,
+            add_to_context=add_to_context,
+        )
+
+        self.shift_noise = ShiftNoise(
+            shift_range=shift_range, add_to_context=add_to_context
+        ) if apply_shift else None
+
+        self.background_noise = BackgroundNoise(
+            background_range=background_range, add_to_context=add_to_context
+        ) if apply_background else None
+
+    def apply(self, curves: Tensor, context: dict = None):
+        """applies the specified types of noise to the input curves"""
+        if self.shift_noise:
+            curves = self.shift_noise(curves, context)
+        
+        if self.background_noise:
+            curves = self.background_noise.apply(curves, context)
+
+        curves = self.gaussian_noise(curves, context)
+        
+        return curves
 
 class BasicExpIntensityNoise(IntensityNoiseGenerator):
     """
@@ -362,4 +468,4 @@ class BasicExpIntensityNoise(IntensityNoiseGenerator):
         if self.background_noise:
             curves = self.background_noise.apply(curves, context)
 
-        return curves
+        return curves

reflectorch/data_generation/priors/parametric_models.py

@@ -10,7 +10,6 @@ from reflectorch.data_generation.reflectivity import (
 )
 from reflectorch.data_generation.utils import (
     get_param_labels,
-    get_param_labels_absorption_model,
 )
 from reflectorch.data_generation.priors.sampler_strategies import (
     SamplerStrategy,
@@ -44,7 +43,7 @@ class ParametricModel(object):
     @property
     def param_dim(self) -> int:
         """get the number of parameters
-        
+
         Returns:
             int:
         """
@@ -106,7 +105,7 @@ class ParametricModel(object):
 
         return min_bounds, max_bounds, min_deltas, max_deltas
 
-    def get_param_labels(self) -> List[str]:
+    def get_param_labels(self, **kwargs) -> List[str]:
         """get the list with the name of the parameters
 
         Returns:
@@ -158,9 +157,10 @@ class StandardModel(ParametricModel):
     def _init_sampler_strategy(self,
                                constrained_roughness: bool = True,
                                max_thickness_share: float = 0.5,
+                               nuisance_params_dim: int = 0,
                                **kwargs):
         if constrained_roughness:
-            num_params = self.param_dim
+            num_params = self.param_dim + nuisance_params_dim
             thickness_mask = torch.zeros(num_params, dtype=torch.bool)
             roughness_mask = torch.zeros(num_params, dtype=torch.bool)
             thickness_mask[:self.max_num_layers] = True
@@ -204,8 +204,8 @@ class StandardModel(ParametricModel):
 
         return min_bounds, max_bounds, min_deltas, max_deltas
 
-    def get_param_labels(self) -> List[str]:
-        return get_param_labels(self.max_num_layers)
+    def get_param_labels(self, **kwargs) -> List[str]:
+        return get_param_labels(self.max_num_layers, **kwargs)
 
     @staticmethod
     def _params2dict(parametrized_model: Tensor):
@@ -250,9 +250,10 @@ class ModelWithAbsorption(StandardModel):
                                constrained_isld: bool = True,
                                max_thickness_share: float = 0.5,
                                max_sld_share: float = 0.2,
+                               nuisance_params_dim: int = 0,
                                **kwargs):
         if constrained_roughness:
-            num_params = self.param_dim
+            num_params = self.param_dim + nuisance_params_dim
             thickness_mask = torch.zeros(num_params, dtype=torch.bool)
             roughness_mask = torch.zeros(num_params, dtype=torch.bool)
             thickness_mask[:self.max_num_layers] = True
@@ -262,10 +263,11 @@ class ModelWithAbsorption(StandardModel):
                 sld_mask = torch.zeros(num_params, dtype=torch.bool)
                 isld_mask = torch.zeros(num_params, dtype=torch.bool)
                 sld_mask[2 * self.max_num_layers + 1:3 * self.max_num_layers + 2] = True
-                isld_mask[3 * self.max_num_layers + 2:] = True
+                isld_mask[3 * self.max_num_layers + 2:4 * self.max_num_layers + 3] = True
                 return ConstrainedRoughnessAndImgSldSamplerStrategy(
                     thickness_mask, roughness_mask, sld_mask, isld_mask,
-                    max_thickness_share=max_thickness_share, max_sld_share=max_sld_share
+                    max_thickness_share=max_thickness_share, max_sld_share=max_sld_share,
+                    **kwargs
                 )
             else:
                 return ConstrainedRoughnessSamplerStrategy(
@@ -305,9 +307,9 @@ class ModelWithAbsorption(StandardModel):
 
         return min_bounds, max_bounds, min_deltas, max_deltas
 
-    def get_param_labels(self) -> List[str]:
-        return get_param_labels_absorption_model(self.max_num_layers)
-
+    def get_param_labels(self, **kwargs) -> List[str]:
+        return get_param_labels(self.max_num_layers, parameterization_type='absorption', **kwargs)
+    
     @staticmethod
     def _params2dict(parametrized_model: Tensor):
         num_params = parametrized_model.shape[-1]
@@ -355,8 +357,9 @@ class ModelWithShifts(StandardModel):
 
         return params
 
-    def get_param_labels(self) -> List[str]:
-        return get_param_labels(self.max_num_layers) + [r"$\Delta q$ (Å$^{{-1}}$)", r"$\Delta I$"]
+    def get_param_labels(self, **kwargs) -> List[str]:
+        return get_param_labels(self.max_num_layers, **kwargs) + [r"$\Delta q$ (Å$^{{-1}}$)", r"$\Delta I$"]
+
 
     @staticmethod
     def _params2dict(parametrized_model: Tensor):
@@ -384,7 +387,7 @@ class ModelWithShifts(StandardModel):
 
 def reflectivity_with_shifts(q, thickness, roughness, sld, q_shift, norm_shift, **kwargs):
     q = torch.atleast_2d(q) + q_shift
-    return reflectivity(q, thickness, roughness, sld, **kwargs) * norm_shift
+    return reflectivity(q, thickness, roughness, sld, **kwargs) * norm_shift   
 
 class NoFresnelModel(StandardModel):
     NAME = 'no_fresnel_model'
@@ -765,3 +768,75 @@ def multilayer_model3(parametrized_model: Tensor, d_full_rel_max: int = 30):
         sld=slds
     )
     return params
+
+
+class NuisanceParamsWrapper(ParametricModel):
+    """
+    Wraps a base model (e.g. StandardModel) to add nuisance parameters, allowing independent enabling/disabling.
+
+    Args:
+        base_model (ParametricModel): The base parametric model.
+        nuisance_params_config (Dict[str, bool]): Dictionary where keys are parameter names
+                                                  and values are `True` (enable) or `False` (disable).
+    """
+
+    def __init__(self, base_model: ParametricModel, nuisance_params_config: Dict[str, bool] = None, **kwargs):
+        self.base_model = base_model
+        self.nuisance_params_config = nuisance_params_config or {}
+
+        self.enabled_nuisance_params = [name for name, is_enabled in self.nuisance_params_config.items() if is_enabled]
+
+        self.PARAMETER_NAMES = self.base_model.PARAMETER_NAMES + tuple(self.enabled_nuisance_params)
+        self._param_dim = self.base_model.param_dim + len(self.enabled_nuisance_params)
+        
+        super().__init__(base_model.max_num_layers, **kwargs)
+
+    def _init_sampler_strategy(self, **kwargs):
+        return self.base_model._init_sampler_strategy(nuisance_params_dim=len(self.enabled_nuisance_params), **kwargs)
+
+    @property
+    def param_dim(self) -> int:
+        return self._param_dim
+
+    def to_standard_params(self, parametrized_model: Tensor) -> dict:
+        """Extracts base model parameters only."""
+        base_dim = self.base_model.param_dim
+        base_part = parametrized_model[..., :base_dim]
+        return self.base_model.to_standard_params(base_part)
+
+    def reflectivity(self, q, parametrized_model: Tensor, **kwargs) -> Tensor:
+        """Computes reflectivity with optional nuisance parameter shifts."""
+        base_dim = self.base_model.param_dim
+        base_params = parametrized_model[..., :base_dim]
+        nuisance_part = parametrized_model[..., base_dim:]
+
+        nuisance_dict = {param: nuisance_part[..., i].unsqueeze(-1) for i, param in enumerate(self.enabled_nuisance_params)}
+        if "log10_background" in nuisance_dict:
+            nuisance_dict["background"] = 10 ** nuisance_dict.pop("log10_background")
+
+        return self.base_model.reflectivity(q, base_params, **nuisance_dict, **kwargs)
+
+    def init_bounds(self, param_ranges: Dict[str, Tuple[float, float]],
+                    bound_width_ranges: Dict[str, Tuple[float, float]], device=None, dtype=None):
+        """Initialize bounds for enabled nuisance parameters."""
+        min_bounds_base, max_bounds_base, min_deltas_base, max_deltas_base = self.base_model.init_bounds(
+            param_ranges, bound_width_ranges, device, dtype)
+
+        ordered_bounds_nuisance = [param_ranges[k] for k in self.enabled_nuisance_params]
+        delta_bounds_nuisance = [bound_width_ranges[k] for k in self.enabled_nuisance_params]
+
+        if ordered_bounds_nuisance:
+            min_bounds_nuisance, max_bounds_nuisance = torch.tensor(ordered_bounds_nuisance, device=device, dtype=dtype).T[:, None]
+            min_deltas_nuisance, max_deltas_nuisance = torch.tensor(delta_bounds_nuisance, device=device, dtype=dtype).T[:, None]
+
+            min_bounds = torch.cat([min_bounds_base, min_bounds_nuisance], dim=-1)
+            max_bounds = torch.cat([max_bounds_base, max_bounds_nuisance], dim=-1)
+            min_deltas = torch.cat([min_deltas_base, min_deltas_nuisance], dim=-1)
+            max_deltas = torch.cat([max_deltas_base, max_deltas_nuisance], dim=-1)
+        else:
+            min_bounds, max_bounds, min_deltas, max_deltas = min_bounds_base, max_bounds_base, min_deltas_base, max_deltas_base
+
+        return min_bounds, max_bounds, min_deltas, max_deltas
+    
+    def get_param_labels(self, **kwargs) -> List[str]:
+        return self.base_model.get_param_labels(**kwargs) + self.enabled_nuisance_params

reflectorch/data_generation/priors/parametric_subpriors.py

@@ -12,6 +12,7 @@ from reflectorch.data_generation.priors.no_constraints import (
 
 from reflectorch.data_generation.priors.parametric_models import (
     MULTILAYER_MODELS,
+    NuisanceParamsWrapper,
     ParametricModel,
 )
 from reflectorch.data_generation.priors.scaler_mixin import ScalerMixin
@@ -54,9 +55,9 @@ class BasicParams(AbstractParams):
         self.min_bounds = min_bounds
         self.max_bounds = max_bounds
 
-    def get_param_labels(self) -> List[str]:
+    def get_param_labels(self, **kwargs) -> List[str]:
         """gets the parameter labels"""
-        return self.param_model.get_param_labels()
+        return self.param_model.get_param_labels(**kwargs)
 
     def reflectivity(self, q: Tensor, log: bool = False, **kwargs):
         r"""computes the reflectivity curves directly from the parameters
@@ -97,6 +98,18 @@ class BasicParams(AbstractParams):
         """gets the slds"""
         params = self.param_model.to_standard_params(self.parameters)
         return params['sld']
+        
+    @property
+    def real_slds(self):
+        """gets the real part of the slds"""
+        params = self.param_model.to_standard_params(self.parameters)
+        return params['sld'].real
+    
+    @property
+    def imag_slds(self):
+        """gets the imaginary part of the slds (only for complex dtypes)"""
+        params = self.param_model.to_standard_params(self.parameters)
+        return params['sld'].imag
 
     @staticmethod
     def rearrange_context_from_params(
@@ -201,11 +214,19 @@ class SubpriorParametricSampler(PriorSampler, ScalerMixin):
             scaled_range (Tuple[float, float], optional): the range for scaling the parameters. Defaults to (-1., 1.)
         """
         self.scaled_range = scaled_range
-        self.param_model: ParametricModel = MULTILAYER_MODELS[model_name](
-            max_num_layers,
-            logdist=logdist,
-            **kwargs
-        )
+
+        self.shift_param_config = kwargs.pop('shift_param_config', {})
+
+        base_model: ParametricModel = MULTILAYER_MODELS[model_name](max_num_layers, logdist=logdist, **kwargs)
+        if any(self.shift_param_config.values()):
+            self.param_model = NuisanceParamsWrapper(
+                base_model=base_model,
+                nuisance_params_config=self.shift_param_config,
+                **kwargs,
+            )
+        else:
+            self.param_model = base_model
+
         self.device = device
         self.dtype = dtype
         self.num_layers = max_num_layers

reflectorch/data_generation/priors/sampler_strategies.py

@@ -73,11 +73,13 @@ class ConstrainedRoughnessSamplerStrategy(BasicSamplerStrategy):
                  roughness_mask: Tensor,
                  logdist: bool = False,
                  max_thickness_share: float = 0.5,
+                 max_total_thickness: float = None,
                  ):
         super().__init__(logdist=logdist)
         self.thickness_mask = thickness_mask
         self.roughness_mask = roughness_mask
         self.max_thickness_share = max_thickness_share
+        self.max_total_thickness = max_total_thickness
 
     def sample(self, batch_size: int,
                total_min_bounds: Tensor,
@@ -106,7 +108,8 @@ class ConstrainedRoughnessSamplerStrategy(BasicSamplerStrategy):
             thickness_mask=self.thickness_mask.to(device),
             roughness_mask=self.roughness_mask.to(device),
             widths_sampler_func=self.widths_sampler_func,
-            coef=self.max_thickness_share,
+            coef_roughness=self.max_thickness_share,
+            max_total_thickness=self.max_total_thickness,
         )
 
 class ConstrainedRoughnessAndImgSldSamplerStrategy(BasicSamplerStrategy):
@@ -129,6 +132,7 @@ class ConstrainedRoughnessAndImgSldSamplerStrategy(BasicSamplerStrategy):
                  logdist: bool = False,
                  max_thickness_share: float = 0.5,
                  max_sld_share: float = 0.2,
+                 max_total_thickness: float = None,
                  ):
         super().__init__(logdist=logdist)
         self.thickness_mask = thickness_mask
@@ -137,6 +141,7 @@ class ConstrainedRoughnessAndImgSldSamplerStrategy(BasicSamplerStrategy):
         self.isld_mask = isld_mask
         self.max_thickness_share = max_thickness_share
         self.max_sld_share = max_sld_share
+        self.max_total_thickness = max_total_thickness
 
     def sample(self, batch_size: int,
                total_min_bounds: Tensor,
@@ -169,6 +174,7 @@ class ConstrainedRoughnessAndImgSldSamplerStrategy(BasicSamplerStrategy):
             widths_sampler_func=self.widths_sampler_func,
             coef_roughness=self.max_thickness_share,
             coef_isld=self.max_sld_share,
+            max_total_thickness=self.max_total_thickness,
         )
 
 def basic_sampler(
@@ -214,15 +220,44 @@ def constrained_roughness_sampler(
         thickness_mask: Tensor,
         roughness_mask: Tensor,
         widths_sampler_func,
-        coef: float = 0.5,
+        coef_roughness: float = 0.5,
+        max_total_thickness: float = None,
 ):
     params, min_bounds, max_bounds = basic_sampler(
         batch_size, total_min_bounds, total_max_bounds, total_min_delta, total_max_delta,
         widths_sampler_func=widths_sampler_func,
     )
 
+    if max_total_thickness is not None:
+        total_thickness = max_bounds[:, thickness_mask].sum(-1)
+        indices = total_thickness > max_total_thickness
+
+        if indices.any():
+            eps = 0.01
+            rand_scale = torch.rand_like(total_thickness) * eps + 1 - eps
+            scale_coef = max_total_thickness / total_thickness * rand_scale
+            scale_coef[~indices] = 1.0
+            min_bounds[:, thickness_mask] *= scale_coef[:, None]
+            max_bounds[:, thickness_mask] *= scale_coef[:, None]
+            params[:, thickness_mask] *= scale_coef[:, None]
+
+            min_bounds[:, thickness_mask] = torch.clamp_min(
+                min_bounds[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
+            max_bounds[:, thickness_mask] = torch.clamp_min(
+                max_bounds[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
+            params[:, thickness_mask] = torch.clamp_min(
+                params[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
     max_roughness = torch.minimum(
-        get_max_allowed_roughness(thicknesses=params[..., thickness_mask], coef=coef),
+        get_max_allowed_roughness(thicknesses=params[..., thickness_mask], coef=coef_roughness),
         total_max_bounds[..., roughness_mask]
     )
     min_roughness = total_min_bounds[..., roughness_mask]
@@ -256,12 +291,41 @@ def constrained_roughness_and_isld_sampler(
         widths_sampler_func,
         coef_roughness: float = 0.5,
         coef_isld: float = 0.2,
+        max_total_thickness: float = None,
 ):
     params, min_bounds, max_bounds = basic_sampler(
         batch_size, total_min_bounds, total_max_bounds, total_min_delta, total_max_delta,
         widths_sampler_func=widths_sampler_func,
     )
 
+    if max_total_thickness is not None:
+        total_thickness = max_bounds[:, thickness_mask].sum(-1)
+        indices = total_thickness > max_total_thickness
+
+        if indices.any():
+            eps = 0.01
+            rand_scale = torch.rand_like(total_thickness) * eps + 1 - eps
+            scale_coef = max_total_thickness / total_thickness * rand_scale
+            scale_coef[~indices] = 1.0
+            min_bounds[:, thickness_mask] *= scale_coef[:, None]
+            max_bounds[:, thickness_mask] *= scale_coef[:, None]
+            params[:, thickness_mask] *= scale_coef[:, None]
+
+            min_bounds[:, thickness_mask] = torch.clamp_min(
+                min_bounds[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
+            max_bounds[:, thickness_mask] = torch.clamp_min(
+                max_bounds[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
+            params[:, thickness_mask] = torch.clamp_min(
+                params[:, thickness_mask],
+                total_min_bounds[:, thickness_mask],
+            )
+
     max_roughness = torch.minimum(
         get_max_allowed_roughness(thicknesses=params[..., thickness_mask], coef=coef_roughness),
         total_max_bounds[..., roughness_mask]