reflectorch 1.5.1__py3-none-any.whl

reflectorch/inference/input_interface.py

@@ -0,0 +1,239 @@
+from reflectorch.data_generation.priors.parametric_models import NuisanceParamsWrapper
+
+class Layer():
+    """Defines a single slab layer with prior bounds for thickness, roughness and SLD.
+
+    The bounds can be given for both real and imaginary parts of the SLD (the latter only if the model supports absorption).
+
+    Args:
+        thickness_bounds (Tuple[float, float]): Minimum and maximum thickness of the layer (in Å).
+        roughness_bounds (Tuple[float, float]): Minimum and maximum interfacial roughness at the top of this layer (in Å).
+        sld_bounds (Tuple[float, float]): Minimum and maximum real SLD of this layer (in 10⁻⁶ Å⁻²).
+        imag_sld_bounds (Tuple[float, float], optional): Minimum and maximum imaginary SLD (in 10⁻⁶ Å⁻²) of this layer. Defaults to None.
+    """
+    def __init__(self, thickness_bounds, roughness_bounds, sld_bounds, imag_sld_bounds=None):
+        self.thickness_bounds = thickness_bounds
+        self.roughness_bounds = roughness_bounds
+        self.sld_bounds = sld_bounds
+        self.imag_sld_bounds = imag_sld_bounds
+
+class Backing():
+    """Defines the backing medium (substrate) for the multilayer structure.
+
+    The backing is assumed to be semi-infinite and has no thickness parameter. 
+    This class ensures compatibility with the layer-based structure definition.
+
+    Args:
+        roughness_bounds (Tuple[float, float]): Minimum and maximum interfacial roughness at the top of the backing medium (in Å).
+        sld_bounds (Tuple[float, float]): Minimum and maximum real SLD of the backing medium (in 10⁻⁶ Å⁻²).
+        imag_sld_bounds (Tuple[float, float], optional): Minimum and maximum imaginary SLD (in 10⁻⁶ Å⁻²) of the backing. Defaults to None.
+    """
+    def __init__(self, roughness_bounds, sld_bounds, imag_sld_bounds=None):
+        self.thickness_bounds = None
+        self.roughness_bounds = roughness_bounds
+        self.sld_bounds = sld_bounds
+        self.imag_sld_bounds = imag_sld_bounds
+
+class Structure():
+    """Defines a multilayer structure and its parameter bounds in a layer-wise manner.
+
+    This class allows the user to define the prior bounds for the full structure (film + backing) in a layer-wise format. It automatically constructs the
+    flattened list of parameter bounds compatible with the inference model’s expected input format.
+
+    Args:
+        layers (List[Union[Layer, Backing]]): Ordered list of layers defining the structure, from the ambient side to the backing. The last element
+            must be a :class:`Backing` instance. Note that the fronting medium (ambient) is not part of this list (since it is not a predicted parameter), 
+            and is treated by default as being 0 (air). For different fronting media one can use the ``ambient_sld`` argument of the prediction method.
+        q_shift_bounds (Tuple[float, float], optional): Bounds for the global ``q_shift`` nuisance parameter. Defaults to None.
+        r_scale_bounds (Tuple[float, float], optional): Bounds for the global reflectivity scale factor ``r_scale``. Defaults to None.
+        log10_background_bounds (Tuple[float, float], optional): Bounds for the background term expressed as log10(background). Defaults to None.
+
+    Attributes:
+        thicknesses_bounds (List[Tuple[float, float]]): Bounds for all thicknesses (excluding backing).
+        roughnesses_bounds (List[Tuple[float, float]]): Bounds for all roughnesses (including backing).
+        slds_bounds (List[Tuple[float, float]]): Bounds for all real SLDs (including backing).
+        imag_slds_bounds (List[Tuple[float, float]]): Bounds for all imaginary SLDs (if provided).
+        prior_bounds (List[Tuple[float, float]]): Flattened list of all parameter bounds in the order expected by the model: thicknesses, 
+            roughnesses, real SLDs, imaginary SLDs (if present), followed by nuisance parameters.
+
+    Example:
+        >>> layer1 = Layer(thickness_bounds=[1, 100], roughness_bounds=[0, 10], sld_bounds=[-2, 2])
+        >>> backing = Backing(roughness_bounds=[0, 15], sld_bounds=[0, 3])
+        >>> structure = Structure(layers=[layer1, backing], r_scale_bounds=[0.9, 1.1])
+        >>> structure.prior_bounds
+    """
+    def __init__(self, layers, q_shift_bounds=None, r_scale_bounds=None, log10_background_bounds=None):
+        self.layers=layers
+        self.q_shift_bounds=q_shift_bounds
+        self.r_scale_bounds = r_scale_bounds
+        self.log10_background_bounds = log10_background_bounds
+        self.thicknesses_bounds = []
+        self.roughnesses_bounds = []
+        self.slds_bounds = []
+        self.imag_slds_bounds = []
+
+        for layer in layers:
+            if layer.thickness_bounds is not None:
+                self.thicknesses_bounds.append(layer.thickness_bounds)
+            self.roughnesses_bounds.append(layer.roughness_bounds)
+            self.slds_bounds.append(layer.sld_bounds)
+            if layer.imag_sld_bounds is not None:
+                self.imag_slds_bounds.append(layer.imag_sld_bounds)
+        
+        self.prior_bounds = self.thicknesses_bounds + self.roughnesses_bounds + self.slds_bounds + self.imag_slds_bounds
+
+        if q_shift_bounds is not None:
+            self.prior_bounds += [q_shift_bounds]
+        if r_scale_bounds is not None:
+            self.prior_bounds += [r_scale_bounds]
+        if log10_background_bounds is not None:
+            self.prior_bounds += [log10_background_bounds]
+    
+    def validate_parameters_and_ranges(self, inference_model):
+        """Validate that all layer bounds and nuisance parameters match the model's configuration.
+
+        This method checks that:
+            * The number of layers matches the model’s expected number.
+            * Each layer’s thickness, roughness, and SLD bounds are within the
+              model’s training ranges.
+            * The SLD bound width does not exceed the maximum training width.
+            * Any nuisance parameters expected by the model (e.g. q_shift, r_scale,
+              log10_background) are provided and within training bounds.
+
+        Args:
+            inference_model (InferenceModel): A loaded instance of :class:`InferenceModel` used to access the model’s metadata.
+
+        Raises:
+            ValueError: If the number of layers, parameter ranges, or nuisance configurations are inconsistent with the model.
+        """
+        if len(self.layers) - 1 != inference_model.trainer.loader.prior_sampler.max_num_layers:
+            raise ValueError(f'Number of layers mismatch: this model expects {inference_model.trainer.loader.prior_sampler.max_num_layers} layers (backing not included) but you provided {len(self.layers) - 1}')
+        
+
+        thickness_ranges = inference_model.trainer.loader.prior_sampler.param_ranges['thicknesses']
+        roughness_ranges = inference_model.trainer.loader.prior_sampler.param_ranges['roughnesses']
+        sld_ranges = inference_model.trainer.loader.prior_sampler.param_ranges['slds']
+        
+        def layer_name(i):
+            if i == inference_model.trainer.loader.prior_sampler.max_num_layers:
+                return 'the backing medium'
+            else:
+                return f'layer {i+1}'
+
+        for i, layer in enumerate(self.layers):
+            if layer.thickness_bounds is not None:
+                if layer.thickness_bounds[0] < thickness_ranges[0] or layer.thickness_bounds[1] > thickness_ranges[1]:
+                    raise ValueError(f"The provided prior bounds for the thickness of layer {i+1} are outside the training range of the network: {thickness_ranges}")
+            if layer.roughness_bounds[0] < roughness_ranges[0] or layer.roughness_bounds[1] > roughness_ranges[1]:
+                raise ValueError(f"The provided prior bounds for the roughness of {layer_name(i)} are outside the training range of the network: {roughness_ranges}")
+            if layer.sld_bounds[0] < sld_ranges[0] or layer.sld_bounds[1] > sld_ranges[1]:
+                raise ValueError(f"The provided prior bounds for the (real) SLD of {layer_name(i)} are outside the training range of the network: {sld_ranges}")
+            
+            max_sld_bounds_width = inference_model.trainer.loader.prior_sampler.bound_width_ranges['slds'][1]
+            if layer.sld_bounds[1] - layer.sld_bounds[0] > max_sld_bounds_width:
+                raise ValueError(f"The provided prior bounds for the (real) SLD of {layer_name(i)} have a width (max - min) exceeding the maximum width used for training: {max_sld_bounds_width}")
+            
+        param_model = inference_model.trainer.loader.prior_sampler.param_model
+        if isinstance(param_model, NuisanceParamsWrapper):
+            nuisance_params_config = inference_model.trainer.loader.prior_sampler.shift_param_config
+                                    
+        if self.q_shift_bounds is not None:
+            if 'q_shift' not in nuisance_params_config:
+                raise ValueError(f'Prior bounds for the q_shift parameter were provided but this parameter is not supported by this model.')
+            q_shift_range = inference_model.trainer.loader.prior_sampler.param_ranges['q_shift']
+            if self.q_shift_bounds[0] < q_shift_range[0] or self.q_shift_bounds[1] > q_shift_range[1]:
+                raise ValueError(f"The provided prior bounds for the q_shift are outside the training range of the network: {q_shift_range}")
+            
+        if self.r_scale_bounds is not None:
+            if 'r_scale' not in nuisance_params_config:
+                raise ValueError(f'Prior bounds for the r_scale parameter were provided but this parameter is not supported by this model.')
+            r_scale_range = inference_model.trainer.loader.prior_sampler.param_ranges['r_scale']
+            if self.r_scale_bounds[0] < r_scale_range[0] or self.r_scale_bounds[1] > r_scale_range[1]:
+                raise ValueError(f"The provided prior bounds for the r_scale are outside the training range of the network: {r_scale_range}")
+            
+        if self.log10_background_bounds is not None:
+            if 'log10_background' not in nuisance_params_config:
+                raise ValueError(f'Prior bounds for the log10_background parameter were provided but this parameter is not supported by this model.')
+            log10_background_range = inference_model.trainer.loader.prior_sampler.param_ranges['log10_background']
+            if self.log10_background_bounds[0] < log10_background_range[0] or self.log10_background_bounds[1] > log10_background_range[1]:
+                raise ValueError(f"The provided prior bounds for the r_scale are outside the training range of the network: {log10_background_range}")
+        
+        if isinstance(param_model, NuisanceParamsWrapper):            
+            if 'q_shift' in nuisance_params_config and self.q_shift_bounds is None:
+                raise ValueError(f'Prior bounds for the q_shift parameter are expected by this model but were not provided.')
+
+            if 'r_scale' in nuisance_params_config and self.r_scale_bounds is None:
+                raise ValueError(f'Prior bounds for the r_scale parameter are expected by this model but were not provided.')
+
+            if 'log10_background' in nuisance_params_config and self.log10_background_bounds is None:
+                raise ValueError(f'Prior bounds for the log10_background parameter are expected by this model but were not provided.')
+        
+        print("All checks passed.")
+
+    def get_huggingface_filtering_query(self):
+        """Constructs a metadata query for selecting compatible pretrained models from Huggingface. Currently it only supports the older (research style) 
+        layout of Huggingface repositories (such as 'valentinsingularity/reflectivity'), but not the newer layout (such as `reflectorch-ILL`).
+
+        Returns:
+            dict: A dictionary describing the structure and parameter bounds, suitable for filtering available model configurations 
+            in a Huggingface repository using :class:`HuggingfaceQueryMatcher`.
+
+        Example:
+            >>> structure = Structure([...])
+            >>> query = structure.get_huggingface_filtering_query()
+            >>> matcher = HuggingfaceQueryMatcher(repo_id='valentinsingularity/reflectivity')
+            >>> configs = matcher.get_matching_configs(query)
+        """
+        query = {'dset.prior_sampler.kwargs.max_num_layers': len(self.layers) - 1}
+        
+        query['dset.prior_sampler.kwargs.param_ranges.thicknesses'] = [min(sl[0] for sl in self.thicknesses_bounds), max(sl[1] for sl in self.thicknesses_bounds)]
+        query['dset.prior_sampler.kwargs.param_ranges.roughnesses'] = [min(sl[0] for sl in self.roughnesses_bounds), max(sl[1] for sl in self.roughnesses_bounds)]
+        query['dset.prior_sampler.kwargs.param_ranges.slds'] = [min(sl[0] for sl in self.slds_bounds), max(sl[1] for sl in self.slds_bounds)]
+
+        if len(self.imag_slds_bounds) > 0:
+            query['dset.prior_sampler.kwargs.model_name'] = 'model_with_absorption'
+            query['dset.prior_sampler.kwargs.param_ranges.islds'] = [min(sl[0] for sl in self.imag_slds_bounds), max(sl[1] for sl in self.imag_slds_bounds)]
+        else:
+            query['dset.prior_sampler.kwargs.model_name'] = 'standard_model'
+
+        if self.q_shift_bounds is not None:
+            query['dset.prior_sampler.kwargs.shift_param_config.q_shift'] = True
+            query['dset.prior_sampler.kwargs.param_ranges.q_shift'] = self.q_shift_bounds
+
+        if self.r_scale_bounds is not None:
+            query['dset.prior_sampler.kwargs.shift_param_config.r_scale'] = True
+            query['dset.prior_sampler.kwargs.param_ranges.r_scale'] = self.r_scale_bounds
+
+        if self.log10_background_bounds is not None:
+            query['dset.prior_sampler.kwargs.shift_param_config.log10_background'] = True
+            query['dset.prior_sampler.kwargs.param_ranges.log10_background'] = self.log10_background_bounds
+
+        return query
+    
+
+if __name__ == '__main__':
+    from reflectorch.inference.query_matcher import HuggingfaceQueryMatcher
+    from reflectorch import EasyInferenceModel
+
+    layer1 = Layer(thickness_bounds=[1, 1000], roughness_bounds=[0, 60], sld_bounds=[-2, 2])
+    layer2 = Layer(thickness_bounds=[1, 50], roughness_bounds=[0, 10], sld_bounds=[1, 4])
+    backing = Backing(roughness_bounds=[0, 15], sld_bounds=[0, 3])
+
+    structure = Structure(
+        layers=[layer1, layer2, backing], 
+        r_scale_bounds=[0.9, 1.1],
+        log10_background_bounds=[-8, -5],
+    )
+
+    print(structure.prior_bounds)
+
+    query_matcher = HuggingfaceQueryMatcher(repo_id='valentinsingularity/reflectivity')
+    filtering_query = structure.get_huggingface_filtering_query()
+    print(filtering_query)
+
+    matching_configs = query_matcher.get_matching_configs(filtering_query)
+    print(f'Matching configs: {matching_configs}')
+
+
+    inference_model = EasyInferenceModel(config_name=matching_configs[0])
+    structure.validate_parameters_and_ranges(inference_model)

reflectorch/inference/loading_data.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+import numpy as np
+
+def load_mft_data(filepath):
+    """
+    Load q, reflectivity, reflectivity error, and q-resolution from an .mft file.
+    
+    Parameters:
+        filepath (str or Path): Path to the .mft file
+
+    Returns:
+        q, refl, refl_err, q_res : np.ndarray
+    """
+    filepath = Path(filepath)
+    
+    with filepath.open('r', encoding='utf-8') as f:
+        lines = f.readlines()
+
+    start_idx = next(
+        i for i, line in enumerate(lines)
+        if line.strip().startswith('q') and 'q_res' in line
+    ) + 1
+
+    data = []
+    for line in lines[start_idx:]:
+        parts = line.strip().split()
+        if len(parts) == 4:
+            try:
+                data.append([float(p.replace('E', 'e')) for p in parts])
+            except ValueError:
+                continue
+
+    if not data:
+        raise ValueError(f"No valid data found in {filepath}")
+
+    data_array = np.array(data)
+    return data_array.T
+
+
+def load_ort_data(filepath):
+    raise NotImplementedError("Loading ORT data is not implemented yet")
+
+def load_csv_data(filepath):
+    raise NotImplementedError("Loading CSV data is not implemented yet")
+
+
+def load_data(filepath):
+    if filepath.endswith('.mft'):
+        return load_mft_data(filepath)
+    elif filepath.endswith('.ort'):
+        return load_ort_data(filepath)
+    elif filepath.endswith('.csv'):
+        return load_csv_data(filepath)
+    else:
+        raise ValueError(f"Unsupported file format: {filepath}")

reflectorch/inference/multilayer_fitter.py

@@ -0,0 +1,171 @@
+from tqdm import trange
+
+from time import perf_counter
+
+import torch
+from torch import nn, Tensor
+
+from reflectorch.data_generation.reflectivity import kinematical_approximation
+from reflectorch.data_generation.priors.multilayer_structures import SimpleMultilayerSampler
+
+
+class MultilayerFit(object):
+    def __init__(self,
+                 q: Tensor,
+                 exp_curve: Tensor,
+                 params: Tensor,
+                 convert_func,
+                 scale_curve_func,
+                 min_bounds: Tensor,
+                 max_bounds: Tensor,
+                 optim_cls=None,
+                 lr: float = 5e-2,
+                 ):
+
+        self.q = q
+        self.scale_curve_func = scale_curve_func
+        self.scaled_exp_curve = self.scale_curve_func(torch.atleast_2d(exp_curve))
+        self.params_to_fit = nn.Parameter(params.clone())
+        self.convert_func = convert_func
+
+        self.min_bounds = min_bounds
+        self.max_bounds = max_bounds
+
+        optim_cls = optim_cls or torch.optim.Adam
+        self.optim = optim_cls([self.params_to_fit], lr)
+
+        self.best_loss = float('inf')
+        self.best_params, self.best_curve = None, None
+
+        self.get_best_solution()
+        self.losses = []
+
+    @classmethod
+    def from_prior_sampler(
+            cls,
+            q: Tensor,
+            exp_curve: Tensor,
+            prior_sampler: SimpleMultilayerSampler,
+            batch_size: int = 2 ** 13,
+            **kwargs
+    ):
+        _, scaled_params = prior_sampler.optimized_sample(batch_size)
+        params = prior_sampler.restore_params2parametrized(scaled_params)
+
+        return cls(
+            q.float(), exp_curve.float(), params.float(),
+            convert_func=get_convert_func(prior_sampler.multilayer_model),
+            scale_curve_func=_save_log, min_bounds=prior_sampler.min_bounds,
+            max_bounds=prior_sampler.max_bounds,
+            **kwargs
+        )
+
+    @classmethod
+    def from_prediction(
+            cls,
+            param: Tensor,
+            prior_sampler,
+            q: Tensor,
+            exp_curve: Tensor,
+            batch_size: int = 2 ** 13,
+            rel_bounds: float = 0.3,
+            **kwargs
+    ):
+
+        num_params = param.shape[-1]
+
+        deltas = prior_sampler.restore_params2parametrized(
+            (torch.rand(batch_size, num_params, device=param.device) * 2 - 1) * rel_bounds
+        ) - prior_sampler.min_bounds
+
+        deltas[0] = 0.
+
+        params = torch.atleast_2d(param) + deltas
+
+        return cls(
+            q.float(), exp_curve.float(), params.float(),
+            convert_func=get_convert_func(prior_sampler.multilayer_model),
+            scale_curve_func=_save_log, min_bounds=prior_sampler.min_bounds,
+            max_bounds=prior_sampler.max_bounds,
+            **kwargs
+        )
+
+    def get_clipped_params(self):
+        return torch.clamp(self.params_to_fit.detach().clone(), self.min_bounds, self.max_bounds)
+
+    def calc_loss(self, reduce=True, clip: bool = False):
+        if clip:
+            params = self.get_clipped_params()
+        else:
+            params = self.params_to_fit
+        curves = self.get_curves(params=params)
+        losses = ((self.scale_curve_func(curves) - self.scaled_exp_curve) ** 2)
+
+        if reduce:
+            return losses.sum()
+        else:
+            return losses.sum(-1)
+
+    def get_curves(self, params: Tensor = None):
+        if params is None:
+            params = self.params_to_fit
+        curves = kinematical_approximation(self.q, **self.convert_func(params))
+        return torch.atleast_2d(curves)
+
+    def run(self, num_iterations: int = 500, disable_tqdm: bool = False):
+        pbar = trange(num_iterations, disable=disable_tqdm)
+
+        for _ in pbar:
+            self.optim.zero_grad()
+            loss = self.calc_loss()
+            loss.backward()
+            self.optim.step()
+            self.losses.append(loss.item())
+            pbar.set_description(f'Loss = {loss.item():.2e}')
+
+    def clear(self):
+        self.losses.clear()
+
+    def run_fixed_time(self, time_limit: float = 2.):
+
+        start = perf_counter()
+
+        while True:
+            self.optim.zero_grad()
+            loss = self.calc_loss()
+            loss.backward()
+            self.optim.step()
+            self.losses.append(loss.item())
+            time_spent = perf_counter() - start
+
+            if time_spent > time_limit:
+                break
+
+    @torch.no_grad()
+    def get_best_solution(self, clip: bool = True):
+        losses = self.calc_loss(clip=clip, reduce=False)
+        idx = torch.argmin(losses)
+        best_loss = losses[idx].item()
+
+        if best_loss < self.best_loss:
+            best_curve = self.get_curves()[idx]
+            best_params = self.params_to_fit.detach()[idx].clone()
+            self.best_params, self.best_curve, self.best_loss = best_params, best_curve, best_loss
+
+        return self.best_params
+
+
+def get_convert_func(multilayer_model):
+    def func(params):
+        params = multilayer_model.to_standard_params(params)
+        return {
+            'thickness': params['thicknesses'],  # very useful transformation indeed ...
+            'roughness': params['roughnesses'],
+            'sld': params['slds'],
+        }
+
+    return func
+
+
+def _save_log(curves, eps: float = 1e-10):
+    return torch.log10(curves + eps)

reflectorch/inference/multilayer_inference_model.py

@@ -0,0 +1,193 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor
+import numpy as np
+
+from reflectorch.inference.inference_model import (
+    InferenceModel,
+)
+from reflectorch.data_generation.reflectivity import kinematical_approximation_np, abeles_np
+
+from reflectorch.data_generation.priors import (
+    MultilayerStructureParams,
+    SimpleMultilayerSampler,
+)
+from reflectorch.inference.record_time import print_time
+from reflectorch.inference.scipy_fitter import standard_refl_fit
+from reflectorch.inference.multilayer_fitter import MultilayerFit
+
+
+class MultilayerInferenceModel(InferenceModel):
+    def predict(self,
+                intensity: np.ndarray,
+                scattering_angle: np.ndarray,
+                attenuation: np.ndarray,
+                priors: np.ndarray = None,
+                preprocessing_parameters: dict = None,
+                polish: bool = True,
+                use_raw_q: bool = False,
+                **kwargs
+                ) -> dict:
+
+        with print_time("everything"):
+            with print_time("preprocess"):
+                preprocessed_dict = self.preprocess(
+                    intensity, scattering_angle, attenuation, **(preprocessing_parameters or {})
+                )
+
+            preprocessed_curve = preprocessed_dict["curve_interp"]
+
+            raw_curve, raw_q = preprocessed_dict["curve"], preprocessed_dict["q_values"]
+
+            with print_time("predict_from_preprocessed_curve"):
+                preprocessed_dict.update(self.predict_from_preprocessed_curve(
+                    preprocessed_curve, priors,
+                    raw_curve=(raw_curve if use_raw_q else None),
+                    raw_q=raw_q,
+                    polish=polish,
+                    use_raw_q=use_raw_q,
+                    **kwargs
+                ))
+
+            return preprocessed_dict
+
+    def predict_from_preprocessed_curve(self,
+                                        curve: np.ndarray,
+                                        priors: np.ndarray = None, *,  # ignore the priors so far
+                                        polish: bool = True,
+                                        raw_curve: np.ndarray = None,
+                                        raw_q: np.ndarray = None,
+                                        clip_prediction: bool = True,
+                                        use_raw_q: bool = False,
+                                        use_sampler: bool = False,
+                                        fitted_time_limit: float = 3.,
+                                        sampler_rel_bounds: float = 0.3,
+                                        polish_with_abeles: bool = False,
+                                        **kwargs
+                                        ) -> dict:
+
+        scaled_curve = self._scale_curve(curve)
+
+        predicted_params, parametrized = self._simple_prediction(scaled_curve)
+
+        if use_sampler:
+            parametrized: Tensor = self._sampler_solution(
+                curve, parametrized,
+                time_limit=fitted_time_limit,
+                rel_bounds=sampler_rel_bounds,
+            )
+
+        init_raw_q = raw_q
+
+        if raw_curve is None:
+            raw_curve = curve
+            raw_q = self.q.squeeze().cpu().numpy()
+            raw_q_t = self.q
+        else:
+            raw_q_t = torch.from_numpy(raw_q).to(self.q)
+
+        # if q_ratio != 1.:
+        #     predicted_params.scale_with_q(q_ratio)
+        #     raw_q = raw_q * q_ratio
+        #     raw_q_t = raw_q_t * q_ratio
+
+        prediction_dict = {
+            "params": parametrized.squeeze().cpu().numpy(),
+            "param_names": list(self._prior_sampler.multilayer_model.PARAMETER_NAMES),
+            "curve_predicted": predicted_params.reflectivity(raw_q_t).squeeze().cpu().numpy()
+        }
+
+        # sld_x_axis, sld_profile, _ = get_density_profiles(
+        #     predicted_params.thicknesses, predicted_params.roughnesses, predicted_params.slds, num=1024,
+        # )
+        #
+        # prediction_dict['sld_profile'] = sld_profile.squeeze().cpu().numpy()
+        # prediction_dict['sld_x_axis'] = sld_x_axis.squeeze().cpu().numpy()
+
+        if polish:
+            prediction_dict.update(self._polish_prediction(
+                raw_q, raw_curve, parametrized, q_values=init_raw_q, use_kinematical=True
+            ))
+        if polish_with_abeles:
+            prediction_dict.update(self._polish_prediction(
+                raw_q, raw_curve, parametrized, q_values=init_raw_q, use_kinematical=False
+            ))
+
+        return prediction_dict
+
+    def _simple_prediction(self, scaled_curve) -> Tuple[MultilayerStructureParams, Tensor]:
+        with torch.no_grad():
+            self.trainer.model.eval()
+            scaled_params = self.trainer.model(scaled_curve)
+
+        predicted_params, parametrized = self._restore_predicted_params(scaled_params)
+        return predicted_params, parametrized
+
+    def _restore_predicted_params(self, scaled_params: Tensor) -> Tuple[MultilayerStructureParams, Tensor]:
+        parametrized = self._prior_sampler.restore_params2parametrized(scaled_params)
+        predicted_params: MultilayerStructureParams = self._prior_sampler.restore_params(scaled_params)
+        return predicted_params, parametrized
+
+    @print_time
+    def _sampler_solution(
+            self,
+            curve: Tensor or np.ndarray,
+            predicted_params: Tensor,
+            batch_size: int = 2 ** 13,
+            time_limit: float = 3.,
+            rel_bounds: float = 0.3,
+    ) -> Tensor:
+
+        fit_obj = MultilayerFit.from_prediction(
+            predicted_params, self._prior_sampler, self.q, torch.as_tensor(curve).to(self.q),
+            batch_size=batch_size, rel_bounds=rel_bounds,
+        )
+
+        fit_obj.run_fixed_time(time_limit)
+
+        best_params = fit_obj.get_best_solution()
+
+        return best_params
+
+    @property
+    def _prior_sampler(self) -> SimpleMultilayerSampler:
+        return self.trainer.loader.prior_sampler
+
+    @print_time
+    def _polish_prediction(self,
+                           q: np.ndarray,
+                           curve: np.ndarray,
+                           predicted_params: Tensor,
+                           q_values: np.ndarray,
+                           use_kinematical: bool = True,
+                           ) -> dict:
+
+        params = predicted_params.squeeze().cpu().numpy()
+        polished_params_dict = {}
+
+        if use_kinematical:
+            refl_generator = kinematical_approximation_np
+        else:
+            refl_generator = abeles_np
+
+        try:
+            polished_params_arr, curve_polished = standard_refl_fit(
+                q, curve, params, restore_params_func=self._prior_sampler.restore_np_params,
+                refl_generator=refl_generator,
+                bounds=self._prior_sampler.get_np_bounds(),
+            )
+            params = self._prior_sampler.restore_np_params(polished_params_arr)
+            if q_values is None:
+                q_values = q
+            curve_polished = abeles_np(q_values, **params)
+
+        except Exception as err:
+            self.log.exception(err)
+            polished_params_arr = params
+            curve_polished = np.zeros_like(q)
+
+        polished_params_dict['params_polished'] = polished_params_arr
+        polished_params_dict['curve_polished'] = curve_polished
+
+        return polished_params_dict