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

reflectorch/inference/inference_model.py

@@ -12,8 +12,10 @@ from IPython.display import display
 from huggingface_hub import hf_hub_download
 
 from reflectorch.data_generation.priors import Params, BasicParams, ExpUniformSubPriorSampler, UniformSubPriorParams
-from reflectorch.data_generation.q_generator import ConstantQ, VariableQ
+from reflectorch.data_generation.priors.parametric_models import NuisanceParamsWrapper
+from reflectorch.data_generation.q_generator import ConstantQ, VariableQ, MaskedVariableQ
 from reflectorch.data_generation.utils import get_density_profiles, get_param_labels
+from reflectorch.inference.preprocess_exp.interpolation import interp_reflectivity
 from reflectorch.paths import CONFIG_DIR, ROOT_DIR, SAVED_MODELS_DIR
 from reflectorch.runs.utils import (
     get_trainer_by_name, train_from_config
@@ -23,10 +25,11 @@ from reflectorch.ml.trainers import PointEstimatorTrainer
 from reflectorch.data_generation.likelihoods import LogLikelihood
 
 from reflectorch.inference.preprocess_exp import StandardPreprocessing
-from reflectorch.inference.scipy_fitter import standard_refl_fit, get_fit_with_growth
+from reflectorch.inference.scipy_fitter import standard_refl_fit, refl_fit, get_fit_with_growth
 from reflectorch.inference.sampler_solution import simple_sampler_solution, get_best_mse_param
 from reflectorch.inference.record_time import print_time
-from reflectorch.utils import to_t
+from reflectorch.inference.plotting import plot_reflectivity, plot_prediction_results, print_prediction_results
+from reflectorch.utils import get_filtering_mask, to_t
 
 class EasyInferenceModel(object):
     """Facilitates the inference process using pretrained models
@@ -103,7 +106,9 @@ class EasyInferenceModel(object):
         self.trainer = get_trainer_by_name(config_name=config_name, config_dir=self.config_dir, model_path=model_path, load_weights=True, inference_device = self.device)
         self.trainer.model.eval()
         
-        print(f'The model corresponds to a `{self.trainer.loader.prior_sampler.param_model.NAME}` parameterization with {self.trainer.loader.prior_sampler.max_num_layers} layers ({self.trainer.loader.prior_sampler.param_dim} predicted parameters)')
+        param_model = self.trainer.loader.prior_sampler.param_model
+        param_model_name = param_model.base_model.NAME if isinstance(param_model, NuisanceParamsWrapper) else param_model.NAME
+        print(f'The model corresponds to a `{param_model_name}` parameterization with {self.trainer.loader.prior_sampler.max_num_layers} layers ({self.trainer.loader.prior_sampler.param_dim} predicted parameters)')
         print("Parameter types and total ranges:")
         for param, range_ in self.trainer.loader.prior_sampler.param_ranges.items():
             print(f"- {param}: {range_}")
@@ -115,23 +120,188 @@ class EasyInferenceModel(object):
             q_min = self.trainer.loader.q_generator.q[0].item()
             q_max = self.trainer.loader.q_generator.q[-1].item()
             n_q = self.trainer.loader.q_generator.q.shape[0]
-            print(f'The model was trained on curves discretized at {n_q} uniform points between between q_min={q_min} and q_max={q_max}')
+            print(f'The model was trained on curves discretized at {n_q} uniform points between q_min={q_min} and q_max={q_max}')
         elif isinstance(self.trainer.loader.q_generator, VariableQ):
             q_min_range = self.trainer.loader.q_generator.q_min_range
             q_max_range = self.trainer.loader.q_generator.q_max_range
             n_q_range = self.trainer.loader.q_generator.n_q_range
-            print(f'The model was trained on curves discretized at a number between {n_q_range[0]} and {n_q_range[1]} of uniform points between between q_min in [{q_min_range[0]}, {q_min_range[1]}] and q_max in [{q_max_range[0]}, {q_max_range[1]}]')
+            if n_q_range[0] == n_q_range[1]:
+                n_q_fixed = n_q_range[0]
+                print(f'The model was trained on curves discretized at exactly {n_q_fixed} uniform points, '
+                    f'between q_min in [{q_min_range[0]}, {q_min_range[1]}] and q_max in [{q_max_range[0]}, {q_max_range[1]}]')
+            else:
+                print(f'The model was trained on curves discretized at a number between {n_q_range[0]} and {n_q_range[1]} '
+                    f'of uniform points between q_min in [{q_min_range[0]}, {q_min_range[1]}] and q_max in [{q_max_range[0]}, {q_max_range[1]}]')
+        
+        if self.trainer.loader.smearing is not None:
+            q_res_min = self.trainer.loader.smearing.sigma_min
+            q_res_max = self.trainer.loader.smearing.sigma_max
+            if self.trainer.loader.smearing.constant_dq == False:
+                print(f"The model was trained with linear resolution smearing (dq/q) in the range [{q_res_min}, {q_res_max}]")
+            elif self.trainer.loader.smearing.constant_dq == True:
+                print(f"The model was trained with constant resolution smearing in the range [{q_res_min}, {q_res_max}]")
+
+        additional_inputs = ["prior bounds"]
+        if self.trainer.train_with_q_input:
+            additional_inputs.append("q values")
+        if self.trainer.condition_on_q_resolutions:
+            additional_inputs.append("the resolution dq/q")
+        if additional_inputs:
+            inputs_str = ", ".join(additional_inputs)
+            print(f"The following quantities are additional inputs to the network: {inputs_str}.")
+
+    def preprocess_and_predict(self, 
+                            reflectivity_curve: Union[np.ndarray, Tensor], 
+                            q_values: Union[np.ndarray, Tensor] = None, 
+                            prior_bounds: Union[np.ndarray, List[Tuple]] = None, 
+                            sigmas: Union[np.ndarray, Tensor] = None,
+                            q_resolution: float = None,
+                            ambient_sld: float = None,
+                            clip_prediction: bool = True, 
+                            polish_prediction: bool = False,
+                            polishing_method: str = 'trf',
+                            polishing_kwargs_reflectivity: dict = None,
+                            use_sigmas_for_polishing: bool = False,
+                            polishing_max_nfev: int = None,
+                            fit_growth: bool = False, 
+                            max_d_change: float = 5.,
+                            calc_pred_curve: bool = True,
+                            calc_pred_sld_profile: bool = False,
+                            calc_polished_sld_profile: bool = False,
+                            sld_profile_padding_left: float = 0.2,
+                            sld_profile_padding_right: float = 1.1,
+                            kwargs_param_labels: dict = {},
+                            
+                            truncate_index_left: int = None,
+                            truncate_index_right: int = None,
+                            enable_error_bars_filtering: bool = True,
+                            filter_threshold=0.3,
+                            filter_remove_singles=True,
+                            filter_remove_consecutives=True,
+                            filter_consecutive=3,
+                            filter_q_start_trunc=0.1,
+                            ):
+        
+        ## Preprocess the data for inference (remove negative intensities, truncation, filer out points with high error bars)
+        (q_values, reflectivity_curve, sigmas, q_resolution, 
+         q_values_original, reflectivity_curve_original, sigmas_original, q_resolution_original) = self._preprocess_input_data(
+            reflectivity_curve=reflectivity_curve,
+            q_values=q_values,
+            sigmas=sigmas,
+            q_resolution=q_resolution,
+            truncate_index_left=truncate_index_left,
+            truncate_index_right=truncate_index_right,
+            enable_error_bars_filtering=enable_error_bars_filtering,
+            filter_threshold=filter_threshold,
+            filter_remove_singles=filter_remove_singles,
+            filter_remove_consecutives=filter_remove_consecutives,
+            filter_consecutive=filter_consecutive,
+            filter_q_start_trunc=filter_q_start_trunc,
+        )
+
+        ### Interpolate the experimental data if needed by the embedding network
+        interp_data = self.interpolate_data_to_model_q(
+            q_exp=q_values,
+            refl_exp=reflectivity_curve,
+            sigmas_exp=sigmas,
+            q_res_exp=q_resolution,
+            as_dict=True
+        )
+
+        q_model = interp_data["q_model"]
+        reflectivity_curve_interp = interp_data["reflectivity"]
+        sigmas_interp = interp_data.get("sigmas")
+        q_resolution_interp = interp_data.get("q_resolution")
+        key_padding_mask = interp_data.get("key_padding_mask")
+        
+        ### Make the prediction
+        prediction_dict = self.predict(
+            reflectivity_curve=reflectivity_curve_interp,
+            q_values=q_model,
+            sigmas=sigmas_interp,
+            q_resolution=q_resolution_interp,
+            key_padding_mask=key_padding_mask,
+            prior_bounds=prior_bounds,
+            ambient_sld=ambient_sld,
+            clip_prediction=clip_prediction,
+            polish_prediction=False, ###do the polishing outside the predict method on the full data
+            supress_sld_amb_back_shift=True, ###do not shift back the slds by the ambient yet
+            calc_pred_curve=calc_pred_curve,
+            calc_pred_sld_profile=calc_pred_sld_profile,
+            sld_profile_padding_left=sld_profile_padding_left,
+            sld_profile_padding_right=sld_profile_padding_right,
+            kwargs_param_labels=kwargs_param_labels,
+        )
+        
+        ### Save interpolated data
+        prediction_dict['q_model'] = q_model
+        prediction_dict['reflectivity_curve_interp'] = reflectivity_curve_interp
+        if q_resolution_interp is not None:
+            prediction_dict['q_resolution_interp'] = q_resolution_interp 
+        if sigmas_interp is not None:
+            prediction_dict['sigmas_interp'] = sigmas_interp
+        if key_padding_mask is not None:
+            prediction_dict['key_padding_mask'] = key_padding_mask
+
+        ### Perform polishing on the original data
+        if polish_prediction:
+            polishing_kwargs = polishing_kwargs_reflectivity or {}
+            polishing_kwargs.setdefault('dq', q_resolution_original)
+            
+            prior_bounds = np.array(prior_bounds)
+            if ambient_sld:
+                sld_indices = self._shift_slds_by_ambient(prior_bounds, ambient_sld)
+
+            polished_dict = self._polish_prediction(
+                q=q_values_original,
+                curve=reflectivity_curve_original,
+                predicted_params=prediction_dict['predicted_params_object'],
+                priors=prior_bounds,
+                ambient_sld_tensor=torch.atleast_2d(torch.as_tensor(ambient_sld)).to(self.device) if ambient_sld is not None else None,
+                calc_polished_sld_profile=calc_polished_sld_profile,
+                sld_x_axis=torch.from_numpy(prediction_dict['predicted_sld_xaxis']),
+                polishing_kwargs_reflectivity = polishing_kwargs,
+                error_bars=sigmas_original if use_sigmas_for_polishing else None,
+                polishing_method=polishing_method,
+                polishing_max_nfev=polishing_max_nfev,
+                fit_growth=fit_growth,
+                max_d_change=max_d_change,
+            )
+
+            prediction_dict.update(polished_dict)
+            if fit_growth and "polished_params_array" in prediction_dict:
+                prediction_dict["param_names"].append("max_d_change")
+        
+        ### Shift back the slds for nonzero ambient
+        if ambient_sld:
+            self._restore_slds_after_ambient_shift(prediction_dict, sld_indices, ambient_sld)   
 
-    def predict(self, reflectivity_curve: Union[np.ndarray, Tensor], 
+        return prediction_dict
+
+
+    def predict(self, 
+                reflectivity_curve: Union[np.ndarray, Tensor], 
                 q_values: Union[np.ndarray, Tensor] = None, 
                 prior_bounds: Union[np.ndarray, List[Tuple]] = None, 
-                clip_prediction: bool = False, 
-                polish_prediction: bool = False, 
+                sigmas: Union[np.ndarray, Tensor] = None,
+                key_padding_mask: Union[np.ndarray, Tensor] = None,
+                q_resolution: float = None,
+                ambient_sld: float = None,
+                clip_prediction: bool = True, 
+                polish_prediction: bool = False,
+                polishing_method: str = 'trf',
+                polishing_kwargs_reflectivity: dict = None,
+                polishing_max_nfev: int = None,
                 fit_growth: bool = False, 
                 max_d_change: float = 5.,
                 use_q_shift: bool = False, 
                 calc_pred_curve: bool = True,
                 calc_pred_sld_profile: bool = False,
+                calc_polished_sld_profile: bool = False,
+                sld_profile_padding_left: float = 0.2,
+                sld_profile_padding_right: float = 1.1,
+                supress_sld_amb_back_shift: bool = False,
+                kwargs_param_labels: dict = {},
                 ):
         """Predict the thin film parameters
 
@@ -153,24 +323,51 @@ class EasyInferenceModel(object):
 
         scaled_curve = self._scale_curve(reflectivity_curve)
         prior_bounds = np.array(prior_bounds)
+
+        if ambient_sld:
+            sld_indices = self._shift_slds_by_ambient(prior_bounds, ambient_sld)
+
         scaled_prior_bounds = self._scale_prior_bounds(prior_bounds)
 
-        if not self.trainer.train_with_q_input:
+        if isinstance(self.trainer.loader.q_generator, ConstantQ):
             q_values = self.trainer.loader.q_generator.q
         else:
             q_values = torch.atleast_2d(to_t(q_values)).to(scaled_curve)
 
+        scaled_q_values = self.trainer.loader.q_generator.scale_q(q_values).to(torch.float32) if self.trainer.train_with_q_input else None
+        
+        if q_resolution is not None:
+            q_resolution_tensor = torch.atleast_2d(torch.as_tensor(q_resolution)).to(scaled_curve)
+            if isinstance(q_resolution, float):
+                unscaled_q_resolutions = q_resolution_tensor
+            else:
+                unscaled_q_resolutions = (q_resolution_tensor / q_values).nanmean(dim=-1, keepdim=True) ##when q_values is padded with 0s, there will be nan at the padded positions
+            scaled_q_resolutions = self.trainer.loader.smearing.scale_resolutions(unscaled_q_resolutions) if self.trainer.condition_on_q_resolutions else None
+            scaled_conditioning_params = scaled_q_resolutions
+            if polishing_kwargs_reflectivity is None:
+                polishing_kwargs_reflectivity = {'dq': q_resolution}
+        else:
+            q_resolution_tensor = None
+            scaled_conditioning_params = None
+
+        if key_padding_mask is not None:
+            key_padding_mask = torch.as_tensor(key_padding_mask, device=self.device)
+            key_padding_mask = key_padding_mask.unsqueeze(0) if key_padding_mask.dim() == 1 else key_padding_mask
+
         if use_q_shift and not self.trainer.train_with_q_input:
             predicted_params = self._qshift_prediction(reflectivity_curve, scaled_prior_bounds, num = 1024, dq_coef = 1.)
-
         else:
             with torch.no_grad():
                 self.trainer.model.eval()
-                if self.trainer.train_with_q_input:
-                    scaled_q = self.trainer.loader.q_generator.scale_q(q_values).float()
-                    scaled_predicted_params = self.trainer.model(scaled_curve, scaled_prior_bounds, scaled_q)
-                else:
-                    scaled_predicted_params = self.trainer.model(scaled_curve, scaled_prior_bounds)
+    
+                scaled_predicted_params = self.trainer.model(
+                    curves=scaled_curve, 
+                    bounds=scaled_prior_bounds, 
+                    q_values=scaled_q_values,
+                    conditioning_params = scaled_conditioning_params,
+                    key_padding_mask = key_padding_mask,
+                    unscaled_q_values = q_values,
+                    )
                 
                 predicted_params = self.trainer.loader.prior_sampler.restore_params(torch.cat([scaled_predicted_params, scaled_prior_bounds], dim=-1))
 
@@ -180,166 +377,80 @@ class EasyInferenceModel(object):
         prediction_dict = {
             "predicted_params_object": predicted_params,
             "predicted_params_array": predicted_params.parameters.squeeze().cpu().numpy(),
-            "param_names" : self.trainer.loader.prior_sampler.param_model.get_param_labels()
+            "param_names" : self.trainer.loader.prior_sampler.param_model.get_param_labels(**kwargs_param_labels)
         }
+        
+        key_padding_mask = None if key_padding_mask is None else key_padding_mask.squeeze().cpu().numpy()
 
         if calc_pred_curve:
-            predicted_curve = predicted_params.reflectivity(q_values).squeeze().cpu().numpy()
-            prediction_dict[ "predicted_curve"] = predicted_curve
+            predicted_curve = predicted_params.reflectivity(q=q_values, dq=q_resolution_tensor).squeeze().cpu().numpy()
+            prediction_dict[ "predicted_curve"] = predicted_curve if key_padding_mask is None else predicted_curve[key_padding_mask]
         
+        ambient_sld_tensor = torch.atleast_2d(torch.as_tensor(ambient_sld, device=self.device)) if ambient_sld is not None else None
         if calc_pred_sld_profile: 
             predicted_sld_xaxis, predicted_sld_profile, _ = get_density_profiles(
-                predicted_params.thicknesses, predicted_params.roughnesses, predicted_params.slds, num=1024, 
+                predicted_params.thicknesses, predicted_params.roughnesses, predicted_params.slds + (ambient_sld_tensor or 0), ambient_sld_tensor, 
+                num=1024, padding_left=sld_profile_padding_left, padding_right=sld_profile_padding_right,
             ) 
             prediction_dict['predicted_sld_profile'] = predicted_sld_profile.squeeze().cpu().numpy()
             prediction_dict['predicted_sld_xaxis'] = predicted_sld_xaxis.squeeze().cpu().numpy()
         else:
             predicted_sld_xaxis = None
+        
+        refl_curve_polish = reflectivity_curve if key_padding_mask is None else reflectivity_curve[key_padding_mask]
+        q_polish = q_values.squeeze().cpu().numpy() if key_padding_mask is None else q_values.squeeze().cpu().numpy()[key_padding_mask]
+        prediction_dict['q_plot_pred'] = q_polish
 
-        if polish_prediction: #only for standard box-model parameterization
-            polished_dict = self._polish_prediction(q = q_values.squeeze().cpu().numpy(), 
-                                                    curve = reflectivity_curve, 
-                                                    predicted_params = predicted_params, 
-                                                    priors = np.array(prior_bounds), 
-                                                    fit_growth = fit_growth,
-                                                    max_d_change = max_d_change, 
-                                                    calc_polished_curve = calc_pred_curve,
-                                                    calc_polished_sld_profile = False,
-                                                    sld_x_axis = predicted_sld_xaxis,
-                                                    )
+        if polish_prediction:
+            if ambient_sld_tensor:
+                ambient_sld_tensor = ambient_sld_tensor.cpu()
+            
+            polished_dict = self._polish_prediction(
+                q = q_polish, 
+                curve = refl_curve_polish, 
+                predicted_params = predicted_params, 
+                priors = np.array(prior_bounds), 
+                error_bars = sigmas,
+                fit_growth = fit_growth,
+                max_d_change = max_d_change, 
+                calc_polished_curve = calc_pred_curve,
+                calc_polished_sld_profile = calc_polished_sld_profile,
+                ambient_sld_tensor=ambient_sld_tensor,
+                sld_x_axis = predicted_sld_xaxis,
+                polishing_method=polishing_method,
+                polishing_max_nfev=polishing_max_nfev,
+                polishing_kwargs_reflectivity=polishing_kwargs_reflectivity,
+            )
             prediction_dict.update(polished_dict)
 
             if fit_growth and "polished_params_array" in prediction_dict:
                 prediction_dict["param_names"].append("max_d_change")
 
-        return prediction_dict
-
-    async def predict_using_widget(self, reflectivity_curve: Union[np.ndarray, torch.Tensor], **kwargs):
-        """Use an interactive Python widget for specifying the prior bounds before the prediction (works only in a Jupyter notebook).
-        The other arguments are the same as for the ``predict`` method.
-        """
-
-        NUM_INTERVALS = self.trainer.loader.prior_sampler.param_dim
-        param_labels = self.trainer.loader.prior_sampler.param_model.get_param_labels()
-        min_bounds = self.trainer.loader.prior_sampler.min_bounds.cpu().numpy().flatten()
-        max_bounds = self.trainer.loader.prior_sampler.max_bounds.cpu().numpy().flatten()
-        max_deltas = self.trainer.loader.prior_sampler.max_delta.cpu().numpy().flatten()
-
-        print(f'Parameter ranges: {self.trainer.loader.prior_sampler.param_ranges}')
-        print(f'Allowed widths of the prior bound intervals (max-min): {self.trainer.loader.prior_sampler.bound_width_ranges}')
-        print(f'Please fill in the values of the minimum and maximum prior bound for each parameter and press the button!')
-
-        def create_interval_widgets(n):
-            intervals = []
-            for i in range(n):
-                interval_label = widgets.Label(value=f'{param_labels[i]}')
-                initial_max = min(max_bounds[i], min_bounds[i] + max_deltas[i])
-                slider = widgets.FloatRangeSlider(
-                    value=[min_bounds[i], initial_max],
-                    min=min_bounds[i],
-                    max=max_bounds[i],
-                    step=0.01,
-                    layout=widgets.Layout(width='400px'),
-                    style={'description_width': '60px'}
-                )
-                
-                def validate_range(change, slider=slider, max_width=max_deltas[i]):
-                    min_val, max_val = change['new']
-                    if max_val - min_val > max_width:
-                        if change['name'] == 'value':
-                            if change['old'][0] != min_val:
-                                max_val = min_val + max_width
-                            else:
-                                min_val = max_val - max_width
-                        slider.value = [min_val, max_val]
-                
-                slider.observe(validate_range, names='value')
-                
-                interval_row = widgets.HBox([interval_label, slider])
-                intervals.append((slider, interval_row))
-            return intervals
-        
-        interval_widgets = create_interval_widgets(NUM_INTERVALS)
-        interval_box = widgets.VBox([widget[1] for widget in interval_widgets])
-        display(interval_box)
-
-        button = widgets.Button(description="Make prediction")
-        display(button)
-
-        prediction_result = None
-
-        def store_values(b, future):
-            print("Debug: Button clicked")
-            values = []
-            for slider, _ in interval_widgets:
-                values.append((slider.value[0], slider.value[1]))
-            array_values = np.array(values)
-            
-            nonlocal prediction_result
-            prediction_result = self.predict(reflectivity_curve=reflectivity_curve, prior_bounds=array_values, **kwargs)
-            print(prediction_result["predicted_params_array"])
-
-            print("Prediction completed. Closing widget.")
-
-            for child in interval_box.children:
-                child.close()
-            button.close()
-
-            future.set_result(prediction_result)
+        if ambient_sld and not supress_sld_amb_back_shift: #Note: the SLD shift will only be reflected in predicted_params_array but not in predicted_params_object; supress_sld_amb_back_shift is required for the 'preprocess_and_predict' method
+            self._restore_slds_after_ambient_shift(prediction_dict, sld_indices, ambient_sld)   
 
-        button.on_click(store_values)
+        return prediction_dict
     
-
-        future = asyncio.Future()
-
-        button.on_click(lambda b: store_values(b, future))
-
-        return await future
-
-
-    def _qshift_prediction(self, curve, scaled_bounds, num: int = 1000, dq_coef: float = 1.) -> BasicParams:
-        assert isinstance(self.trainer.loader.q_generator, ConstantQ), "Prediction with q shifts available only for models with fixed discretization"
-        q = self.trainer.loader.q_generator.q.squeeze().float()
-        dq_max = (q[1] - q[0]) * dq_coef
-        q_shifts = torch.linspace(-dq_max, dq_max, num).to(q)
-
-        curve = to_t(curve).to(scaled_bounds)
-        shifted_curves = _qshift_interp(q.squeeze(), curve, q_shifts)
-
-        assert shifted_curves.shape == (num, q.shape[0])
-
-        scaled_curves = self.trainer.loader.curves_scaler.scale(shifted_curves)
-        scaled_prior_bounds = torch.atleast_2d(scaled_bounds).expand(scaled_curves.shape[0], -1)
-
-        with torch.no_grad():
-            self.trainer.model.eval()
-            scaled_predicted_params = self.trainer.model(scaled_curves, scaled_prior_bounds)
-            restored_params = self.trainer.loader.prior_sampler.restore_params(torch.cat([scaled_predicted_params, scaled_prior_bounds], dim=-1))
-
-            best_param = get_best_mse_param(
-                restored_params,
-                self._get_likelihood(q=self.trainer.loader.q_generator.q, curve=curve),
-            )
-            return best_param
-        
     def _polish_prediction(self,
                            q: np.ndarray,
                            curve: np.ndarray,
                            predicted_params: BasicParams,
                            priors: np.ndarray,
                            sld_x_axis,
+                           ambient_sld_tensor: Tensor = None,
                            fit_growth: bool = False,
                            max_d_change: float = 5.,
                            calc_polished_curve: bool = True,
                            calc_polished_sld_profile: bool = False,
+                           error_bars: np.ndarray = None,
+                           polishing_method: str = 'trf',
+                           polishing_max_nfev: int = None,
+                           polishing_kwargs_reflectivity: dict = None,
                            ) -> dict:
-        params = torch.cat([
-            predicted_params.thicknesses.squeeze(),
-            predicted_params.roughnesses.squeeze(),
-            predicted_params.slds.squeeze()
-        ]).cpu().numpy()
+        params = predicted_params.parameters.squeeze().cpu().numpy()
 
         polished_params_dict = {}
+        polishing_kwargs_reflectivity = polishing_kwargs_reflectivity or {}
 
         try:
             if fit_growth:
@@ -358,18 +469,24 @@ class EasyInferenceModel(object):
                     self.trainer.loader.prior_sampler.param_model
                     )
             else:
-                polished_params_arr, curve_polished = standard_refl_fit(
+                polished_params_arr, curve_polished = refl_fit(
                     q = q, 
                     curve = curve, 
                     init_params = params, 
-                    bounds=priors.T)
+                    bounds=priors.T,
+                    prior_sampler=self.trainer.loader.prior_sampler,
+                    error_bars=error_bars,
+                    method=polishing_method,
+                    polishing_max_nfev=polishing_max_nfev,
+                    reflectivity_kwargs=polishing_kwargs_reflectivity,
+                )
                 polished_params = BasicParams(
                     torch.from_numpy(polished_params_arr[None]),
                     torch.from_numpy(priors.T[0][None]),
                     torch.from_numpy(priors.T[1][None]),
                     self.trainer.loader.prior_sampler.max_num_layers,
                     self.trainer.loader.prior_sampler.param_model
-                    )
+                )
         except Exception as err:
             polished_params = predicted_params
             polished_params_arr = get_prediction_array(polished_params)
@@ -379,9 +496,14 @@ class EasyInferenceModel(object):
         if calc_polished_curve:
             polished_params_dict['polished_curve'] = curve_polished
 
+        if ambient_sld_tensor is not None:
+            ambient_sld_tensor = ambient_sld_tensor.to(polished_params.slds.device)
+
+
         if calc_polished_sld_profile:
             _, sld_profile_polished, _ = get_density_profiles(
-                polished_params.thicknesses, polished_params.roughnesses, polished_params.slds, z_axis=sld_x_axis,
+                polished_params.thicknesses, polished_params.roughnesses, polished_params.slds + (ambient_sld_tensor or 0), ambient_sld_tensor, 
+                z_axis=sld_x_axis.to(polished_params.slds.device),
             )
             polished_params_dict['sld_profile_polished'] = sld_profile_polished.squeeze().cpu().numpy()
 
@@ -390,26 +512,540 @@ class EasyInferenceModel(object):
     def _scale_curve(self, curve: Union[np.ndarray, Tensor]):
         if not isinstance(curve, Tensor):
             curve = torch.from_numpy(curve).float()
-        curve = torch.atleast_2d(curve).to(self.device)
+        curve = curve.unsqueeze(0).to(self.device)
         scaled_curve = self.trainer.loader.curves_scaler.scale(curve)
         return scaled_curve
     
     def _scale_prior_bounds(self, prior_bounds: List[Tuple]):
-        prior_bounds = torch.tensor(prior_bounds)
-        prior_bounds = prior_bounds.to(self.device).T
-        min_bounds, max_bounds = prior_bounds[:, None]
-
-        scaled_bounds = torch.cat([
-            self.trainer.loader.prior_sampler.scale_bounds(min_bounds), 
-            self.trainer.loader.prior_sampler.scale_bounds(max_bounds)
-        ], -1)
-
-        return scaled_bounds.float()
+        try:
+            prior_bounds = torch.tensor(prior_bounds)
+            prior_bounds = prior_bounds.to(self.device).T
+            min_bounds, max_bounds = prior_bounds[:, None]
+
+            scaled_bounds = torch.cat([
+                self.trainer.loader.prior_sampler.scale_bounds(min_bounds), 
+                self.trainer.loader.prior_sampler.scale_bounds(max_bounds)
+            ], -1)
+
+            return scaled_bounds.float()
+
+        except RuntimeError as e:
+            expected_param_dim = self.trainer.loader.prior_sampler.param_dim
+            actual_param_dim = prior_bounds.shape[1] if prior_bounds.ndim == 2 else len(prior_bounds)
+
+            msg = (
+                f"\n **Parameter dimension mismatch during inference!**\n"
+                f"- Model expects **{expected_param_dim}** parameters.\n"
+                f"- You provided **{actual_param_dim}** prior bounds.\n\n"
+                f"💡This often occurs when:\n"
+                f"- The model was trained with additional nuisance parameters like `r_scale`, `q_shift`, or `log10_background`,\n"
+                f"  but they were not included in the `prior_bounds` passed to `.predict()`.\n"
+                f"- The number of layers or parameterization type differs from the one used during training.\n\n"
+                f" Check the configuration or the summary of expected parameters."
+            )
+            raise ValueError(msg) from e    
     
+    def _shift_slds_by_ambient(self, prior_bounds: np.ndarray, ambient_sld: float):
+        n_layers = self.trainer.loader.prior_sampler.max_num_layers
+        sld_indices = slice(2*n_layers+1, 3*n_layers+2)
+        prior_bounds[sld_indices, ...] -= ambient_sld
+
+        training_min_bounds = self.trainer.loader.prior_sampler.min_bounds.squeeze().cpu().numpy()
+        training_max_bounds = self.trainer.loader.prior_sampler.max_bounds.squeeze().cpu().numpy()
+        lower_bound_check = (prior_bounds[sld_indices, 0] >= training_min_bounds[sld_indices]).all()
+        upper_bound_check = (prior_bounds[sld_indices, 1] <= training_max_bounds[sld_indices]).all()
+        assert lower_bound_check and upper_bound_check, "Shifting the layer SLDs by the ambient SLD exceeded the training ranges."
+
+        return sld_indices
+    
+    def _restore_slds_after_ambient_shift(self, prediction_dict, sld_indices, ambient_sld):
+        prediction_dict["predicted_params_array"][sld_indices] += ambient_sld
+        if "polished_params_array" in prediction_dict:
+            prediction_dict["polished_params_array"][sld_indices] += ambient_sld
+
     def _get_likelihood(self, q, curve, rel_err: float = 0.1, abs_err: float = 1e-12):
         return LogLikelihood(
             q, curve, self.trainer.loader.prior_sampler, curve * rel_err + abs_err
         )
+    
+    def get_param_labels(self, **kwargs):
+        return self.trainer.loader.prior_sampler.param_model.get_param_labels(**kwargs)
+    
+    @staticmethod
+    def _preprocess_input_data(
+                           reflectivity_curve,
+                           q_values,
+                           sigmas=None,
+                           q_resolution=None,
+                           truncate_index_left=None,
+                           truncate_index_right=None,
+                           enable_error_bars_filtering=True,
+                           filter_threshold=0.3,
+                           filter_remove_singles=True,
+                           filter_remove_consecutives=True,
+                           filter_consecutive=3,
+                           filter_q_start_trunc=0.1):
+    
+        # Save originals for polishing
+        reflectivity_curve_original = reflectivity_curve.copy()
+        q_values_original = q_values.copy() if q_values is not None else None
+        q_resolution_original = q_resolution.copy() if isinstance(q_resolution, np.ndarray) else q_resolution
+        sigmas_original = sigmas.copy() if sigmas is not None else None
+
+        # Remove points with non-positive intensities
+        nonnegative_mask = reflectivity_curve > 0.0
+        reflectivity_curve = reflectivity_curve[nonnegative_mask]
+        q_values = q_values[nonnegative_mask]
+        if sigmas is not None:
+            sigmas = sigmas[nonnegative_mask]
+        if isinstance(q_resolution, np.ndarray):
+            q_resolution = q_resolution[nonnegative_mask]
+
+        # Truncate arrays
+        if truncate_index_left is not None or truncate_index_right is not None:
+            slice_obj = slice(truncate_index_left, truncate_index_right)
+            reflectivity_curve = reflectivity_curve[slice_obj]
+            q_values = q_values[slice_obj]
+            if sigmas is not None:
+                sigmas = sigmas[slice_obj]
+            if isinstance(q_resolution, np.ndarray):
+                q_resolution = q_resolution[slice_obj]
+
+        # Filter high-error points
+        if enable_error_bars_filtering and sigmas is not None:
+            valid_mask = get_filtering_mask(
+                q_values,
+                reflectivity_curve,
+                sigmas,
+                threshold=filter_threshold,
+                consecutive=filter_consecutive,
+                remove_singles=filter_remove_singles,
+                remove_consecutives=filter_remove_consecutives,
+                q_start_trunc=filter_q_start_trunc
+            )
+            reflectivity_curve = reflectivity_curve[valid_mask]
+            q_values = q_values[valid_mask]
+            sigmas = sigmas[valid_mask]
+            if isinstance(q_resolution, np.ndarray):
+                q_resolution = q_resolution[valid_mask]
+
+        return (q_values, reflectivity_curve, sigmas, q_resolution,
+                q_values_original, reflectivity_curve_original,
+                sigmas_original, q_resolution_original)
+    
+    def interpolate_data_to_model_q(
+        self,
+        q_exp,
+        refl_exp,
+        sigmas_exp=None,
+        q_res_exp=None,
+        as_dict=False
+    ):
+        q_generator = self.trainer.loader.q_generator
+
+        def _pad(arr, pad_to, value=0.0):
+            if arr is None:
+                return None
+            return np.pad(arr, (0, pad_to - len(arr)), constant_values=value)
+
+        def _interp_or_keep(q_model, q_exp, arr):
+            """Interpolate arrays, keep floats or None unchanged."""
+            if arr is None:
+                return None
+            return np.interp(q_model, q_exp, arr) if isinstance(arr, np.ndarray) else arr
+
+        def _pad_or_keep(arr, max_n):
+            """Pad arrays, keep floats or None unchanged."""
+            if arr is None:
+                return None
+            return _pad(arr, max_n, 0.0) if isinstance(arr, np.ndarray) else arr
+
+        def _prepare_return(q, refl, sigmas=None, q_res=None, mask=None, as_dict=False):
+            if as_dict:
+                result = {"q_model": q, "reflectivity": refl}
+                if sigmas is not None: result["sigmas"] = sigmas
+                if q_res is not None: result["q_resolution"] = q_res
+                if mask is not None: result["key_padding_mask"] = mask
+                return result
+            result = [q, refl]
+            if sigmas is not None: result.append(sigmas)
+            if q_res is not None: result.append(q_res)
+            if mask is not None: result.append(mask)
+            return tuple(result)
+
+        # ConstantQ
+        if isinstance(q_generator, ConstantQ):
+            q_model = q_generator.q.cpu().numpy()
+            refl_out = interp_reflectivity(q_model, q_exp, refl_exp)
+            sigmas_out = _interp_or_keep(q_model, q_exp, sigmas_exp)
+            q_res_out = _interp_or_keep(q_model, q_exp, q_res_exp)
+            return _prepare_return(q_model, refl_out, sigmas_out, q_res_out, None, as_dict)
+
+        # VariableQ
+        elif isinstance(q_generator, VariableQ):
+            if q_generator.n_q_range[0] == q_generator.n_q_range[1]:
+                n_q_model = q_generator.n_q_range[0]
+                q_min = max(q_exp.min(), q_generator.q_min_range[0])
+                q_max = min(q_exp.max(), q_generator.q_max_range[1])
+                if self.trainer.loader.q_generator.mode == 'logspace':
+                    q_model = torch.logspace(start=torch.log10(torch.tensor(q_min, device=self.device)),
+                                             end=torch.log10(torch.tensor(q_max, device=self.device)),
+                                             steps=n_q_model, device=self.device).to('cpu')
+                    logspace = True
+                else:
+                    q_model = np.linspace(q_min, q_max, n_q_model)
+                    logspace = False
+            else:
+                return _prepare_return(q_exp, refl_exp, sigmas_exp, q_res_exp, None, as_dict)
+
+            refl_out = interp_reflectivity(q_model, q_exp, refl_exp, logspace=logspace)
+            sigmas_out = _interp_or_keep(q_model, q_exp, sigmas_exp)
+            q_res_out = _interp_or_keep(q_model, q_exp, q_res_exp)
+            return _prepare_return(q_model, refl_out, sigmas_out, q_res_out, None, as_dict)
+
+        # MaskedVariableQ
+        elif isinstance(q_generator, MaskedVariableQ):
+            min_n, max_n = q_generator.n_q_range
+            n_exp = len(q_exp)
+
+            if min_n <= n_exp <= max_n:
+                # Pad only
+                q_model = _pad(q_exp, max_n, 0.0)
+                refl_out = _pad(refl_exp, max_n, 0.0)
+                sigmas_out = _pad_or_keep(sigmas_exp, max_n)
+                q_res_out = _pad_or_keep(q_res_exp, max_n)
+                key_padding_mask = np.zeros(max_n, dtype=bool)
+                key_padding_mask[:n_exp] = True
+
+            else:
+                # Interpolate + pad
+                n_interp = min(max(n_exp, min_n), max_n)
+                q_min = max(q_exp.min(), q_generator.q_min_range[0])
+                q_max = min(q_exp.max(), q_generator.q_max_range[1])
+                q_interp = np.linspace(q_min, q_max, n_interp)
+
+                refl_interp = interp_reflectivity(q_interp, q_exp, refl_exp)
+                sigmas_interp = _interp_or_keep(q_interp, q_exp, sigmas_exp)
+                q_res_interp = _interp_or_keep(q_interp, q_exp, q_res_exp)
+
+                q_model = _pad(q_interp, max_n, 0.0)
+                refl_out = _pad(refl_interp, max_n, 0.0)
+                sigmas_out = _pad_or_keep(sigmas_interp, max_n)
+                q_res_out = _pad_or_keep(q_res_interp, max_n)
+                key_padding_mask = np.zeros(max_n, dtype=bool)
+                key_padding_mask[:n_interp] = True
+
+            return _prepare_return(q_model, refl_out, sigmas_out, q_res_out, key_padding_mask, as_dict)
+
+        else:
+            raise TypeError(f"Unsupported QGenerator type: {type(q_generator)}")
+        
+    def _qshift_prediction(self, curve, scaled_bounds, num: int = 1000, dq_coef: float = 1.) -> BasicParams:
+        assert isinstance(self.trainer.loader.q_generator, ConstantQ), "Prediction with q shifts available only for models with fixed discretization"
+        q = self.trainer.loader.q_generator.q.squeeze().float()
+        dq_max = (q[1] - q[0]) * dq_coef
+        q_shifts = torch.linspace(-dq_max, dq_max, num).to(q)
+
+        curve = to_t(curve).to(scaled_bounds)
+        shifted_curves = _qshift_interp(q.squeeze(), curve, q_shifts)
+
+        assert shifted_curves.shape == (num, q.shape[0])
+
+        scaled_curves = self.trainer.loader.curves_scaler.scale(shifted_curves)
+        scaled_prior_bounds = torch.atleast_2d(scaled_bounds).expand(scaled_curves.shape[0], -1)
+
+        with torch.no_grad():
+            self.trainer.model.eval()
+            scaled_predicted_params = self.trainer.model(scaled_curves, scaled_prior_bounds)
+            restored_params = self.trainer.loader.prior_sampler.restore_params(torch.cat([scaled_predicted_params, scaled_prior_bounds], dim=-1))
+
+            best_param = get_best_mse_param(
+                restored_params,
+                self._get_likelihood(q=self.trainer.loader.q_generator.q, curve=curve),
+            )
+            return best_param
+
+    def predict_using_widget(self, reflectivity_curve, **kwargs):
+        """
+        """
+
+        NUM_INTERVALS = self.trainer.loader.prior_sampler.param_dim
+        param_labels = self.trainer.loader.prior_sampler.param_model.get_param_labels()
+        min_bounds = self.trainer.loader.prior_sampler.min_bounds.cpu().numpy().flatten()
+        max_bounds = self.trainer.loader.prior_sampler.max_bounds.cpu().numpy().flatten()
+        max_deltas = self.trainer.loader.prior_sampler.max_delta.cpu().numpy().flatten()
+
+        print(f'Adjust the sliders for each parameter and press "Predict". Repeat as desired. Press "Close Widget" to finish.')
+
+        interval_widgets = []
+        for i in range(NUM_INTERVALS):
+            label = widgets.Label(value=f'{param_labels[i]}')
+            initial_max = min(max_bounds[i], min_bounds[i] + max_deltas[i])
+            slider = widgets.FloatRangeSlider(
+                value=[min_bounds[i], initial_max],
+                min=min_bounds[i],
+                max=max_bounds[i],
+                step=0.01,
+                layout=widgets.Layout(width='400px'),
+                style={'description_width': '60px'}
+            )
+
+            def validate_range(change, slider=slider, max_width=max_deltas[i]):
+                min_val, max_val = change['new']
+                if max_val - min_val > max_width:
+                    old_min_val, old_max_val = change['old']
+                    if abs(old_min_val - min_val) > abs(old_max_val - max_val):
+                        max_val = min_val + max_width
+                    else:
+                        min_val = max_val - max_width
+                    slider.value = [min_val, max_val]
+
+            slider.observe(validate_range, names='value')
+            interval_widgets.append((slider, widgets.HBox([label, slider])))
+
+        sliders_box = widgets.VBox([iw[1] for iw in interval_widgets])
+
+        output = widgets.Output()
+        predict_button = widgets.Button(description="Predict")
+        close_button = widgets.Button(description="Close Widget")
+
+        container = widgets.VBox([sliders_box, widgets.HBox([predict_button, close_button]), output])
+        display(container)
+
+        @output.capture(clear_output=True)
+        def on_predict_click(_):
+            if 'prior_bounds' in kwargs:
+                array_values = kwargs.pop('prior_bounds')
+                for i, (s, _) in enumerate(interval_widgets):
+                    s.value = tuple(array_values[i])
+            else:
+                values = [(s.value[0], s.value[1]) for s, _ in interval_widgets]
+                array_values = np.array(values)
+
+            prediction_result = self.predict(reflectivity_curve=reflectivity_curve,
+                                            prior_bounds=array_values,
+                                            **kwargs)
+            param_names = self.trainer.loader.prior_sampler.param_model.get_param_labels()
+            print_prediction_results(prediction_result)
+
+            plot_prediction_results(
+                prediction_result,
+                q_exp=kwargs['q_values'],
+                curve_exp=reflectivity_curve,
+            )
+            self.widget_prediction_result = prediction_result
+
+        def on_close_click(_):
+            container.close()
+            print("Widget closed.")
+
+        predict_button.on_click(on_predict_click)
+        close_button.on_click(on_close_click)
+
+    def preprocess_and_predict_using_widget(self,
+                                            reflectivity_curve,
+                                            q_values=None,
+                                            sigmas=None,
+                                            q_resolution=None,
+                                            prior_bounds=None,
+                                            ambient_sld=None,
+                                            ):
+        """
+        Interactive widget around `preprocess_and_predict`
+        Results are stored in `self.widget_prediction_result`
+        """
+
+        if q_values is None:
+            raise ValueError("q_values must be provided for this widget.")
+
+        N = len(reflectivity_curve)
+
+        # ---------- Priors sliders ----------
+        param_labels = self.trainer.loader.prior_sampler.param_model.get_param_labels()
+        min_bounds = self.trainer.loader.prior_sampler.min_bounds.cpu().numpy().flatten()
+        max_bounds = self.trainer.loader.prior_sampler.max_bounds.cpu().numpy().flatten()
+        max_deltas = self.trainer.loader.prior_sampler.max_delta.cpu().numpy().flatten()
+        NUM = len(param_labels)
+
+        sliders, rows = [], []
+        init_pb = np.array(prior_bounds) if prior_bounds is not None else None
+        for i in range(NUM):
+            init_min = float(init_pb[i, 0]) if init_pb is not None else float(min_bounds[i])
+            init_max = float(init_pb[i, 1]) if init_pb is not None else float(min(min_bounds[i] + max_deltas[i], max_bounds[i]))
+            lab = widgets.Label(value=param_labels[i])
+            s = widgets.FloatRangeSlider(
+                value=[init_min, init_max],
+                min=float(min_bounds[i]),
+                max=float(max_bounds[i]),
+                step=0.01,
+                layout=widgets.Layout(width='420px'),
+                readout_format='.3f'
+            )
+            # Constrain slider widths
+            def _mk_validator(slider, max_width=float(max_deltas[i])):
+                def _validate(change):
+                    a, b = change['new']
+                    if b - a > max_width:
+                        oa, ob = change['old']
+                        if abs(oa - a) > abs(ob - b):
+                            b = a + max_width
+                        else:
+                            a = b - max_width
+                        slider.value = (a, b)
+                return _validate
+            s.observe(_mk_validator(s), names='value')
+            sliders.append(s)
+            rows.append(widgets.HBox([lab, s]))
+        priors_box = widgets.VBox([widgets.HTML("<b>Priors</b>")] + rows)
+
+        # ---------- Preprocess & Predict controls ----------
+        # Preprocessing
+        trunc_L = widgets.IntSlider(description='truncate left',  min=0, max=max(0, N-1), step=1, value=0)
+        trunc_R = widgets.IntSlider(description='truncate right', min=1, max=N,           step=1, value=N)
+        enable_filt = widgets.Checkbox(description='filter error bars', value=True)
+        thr = widgets.FloatSlider(description='filtering threshold', min=0.0, max=1.0, step=0.01, value=0.3)
+        rem_single = widgets.Checkbox(description='remove singles', value=True)
+        rem_cons = widgets.Checkbox(description='remove consecutives', value=True)
+        consec = widgets.IntSlider(description='num. consecutive', min=1, max=10, step=1, value=3)
+        qstart = widgets.FloatSlider(description='q_start_trunc', min=0.0, max=1.0, step=0.01, value=0.1)
+
+        # Polishing
+        polish = widgets.Checkbox(description='polish prediction', value=True)
+        use_sigmas_polish = widgets.Checkbox(description='use sigmas during polishing', value=True)
+
+        # Plotting
+        pred_show_yerr = widgets.Checkbox(description='show error bars', value=True)
+        pred_show_xerr = widgets.Checkbox(description='show q-resolution', value=False)
+        pred_logx = widgets.Checkbox(description='log x-axis', value=False)
+        plot_sld = widgets.Checkbox(description='plot SLD profile', value=True)
+        sld_pad_left = widgets.FloatText(description='SLD pad left',  value=0.2, step=0.1)
+        sld_pad_right = widgets.FloatText(description='SLD pad right', value=1.1, step=0.1)
+
+        # Color pickers
+        exp_color_picker = widgets.ColorPicker(description='exp color',       value='#0000FF')  # blue
+        exp_errcolor_picker = widgets.ColorPicker(description='errbar color',    value='#800080')  # purple
+        pred_color_picker = widgets.ColorPicker(description='pred color',      value='#FF0000')  # red
+        pol_color_picker = widgets.ColorPicker(description='polished color',  value='#FFA500')  # orange
+        sld_pred_color_picker = widgets.ColorPicker(description='SLD pred color',  value='#FF0000')  # red
+        sld_pol_color_picker = widgets.ColorPicker(description='SLD pol color',   value='#FFA500')  # orange
+
+        # Compute toggles
+        calc_curve = widgets.Checkbox(description='calc curve', value=True)
+        calc_sld_pred = widgets.Checkbox(description='calc predicted SLD', value=True)
+        calc_sld_pol = widgets.Checkbox(description='calc polished SLD', value=True)
+
+        btn_predict = widgets.Button(description='Predict')
+        btn_close = widgets.Button(description='Close')
+
+        controls_box = widgets.VBox([
+            widgets.HTML("<b>Preprocess & Predict</b>"),
+            widgets.HTML("<i>Preprocessing</i>"),
+            widgets.HBox([trunc_L, trunc_R]),
+            widgets.HBox([enable_filt, rem_single, rem_cons]),
+            widgets.HBox([thr, consec, qstart]),
+            widgets.HTML("<i>Polishing</i>"),
+            widgets.HBox([polish, use_sigmas_polish]),
+            widgets.HTML("<i>Plotting</i>"),
+            widgets.HBox([pred_show_yerr, pred_show_xerr, pred_logx, plot_sld]),
+            widgets.HBox([sld_pad_left, sld_pad_right]),
+            widgets.HTML("<i>Colors</i>"),
+            widgets.HBox([exp_color_picker, exp_errcolor_picker]),
+            widgets.HBox([pred_color_picker, pol_color_picker]),
+            widgets.HBox([sld_pred_color_picker, sld_pol_color_picker]),
+            widgets.HTML("<i>Compute</i>"),
+            widgets.HBox([calc_curve, calc_sld_pred, calc_sld_pol]),
+            widgets.HBox([btn_predict, btn_close]),
+        ])
+
+        out_predict = widgets.Output()
+        container = widgets.VBox([priors_box, controls_box, out_predict])
+        display(container)
+
+        def _sync_trunc(_):
+            if trunc_L.value >= trunc_R.value:
+                trunc_L.value = max(0, trunc_R.value - 1)
+        trunc_L.observe(_sync_trunc, names='value')
+        trunc_R.observe(_sync_trunc, names='value')
+
+        def _current_priors():
+            return np.array([s.value for s in sliders], dtype=np.float32)  # (param_dim, 2)
+
+        @out_predict.capture(clear_output=True)
+        def _on_predict(_):
+            out_predict.clear_output(wait=True)
+
+            res = self.preprocess_and_predict(
+                reflectivity_curve=reflectivity_curve,
+                q_values=q_values,
+                prior_bounds=_current_priors(),
+                sigmas=sigmas,
+                q_resolution=q_resolution,
+                ambient_sld=ambient_sld,
+                clip_prediction=True,
+                polish_prediction=polish.value,
+                use_sigmas_for_polishing=use_sigmas_polish.value,
+                calc_pred_curve=calc_curve.value,
+                calc_pred_sld_profile=(calc_sld_pred.value or plot_sld.value),
+                calc_polished_sld_profile=(calc_sld_pol.value or plot_sld.value),
+                sld_profile_padding_left=float(sld_pad_left.value),
+                sld_profile_padding_right=float(sld_pad_right.value),
+
+                truncate_index_left=trunc_L.value,
+                truncate_index_right=trunc_R.value,
+                enable_error_bars_filtering=enable_filt.value,
+                filter_threshold=thr.value,
+                filter_remove_singles=rem_single.value,
+                filter_remove_consecutives=rem_cons.value,
+                filter_consecutive=consec.value,
+                filter_q_start_trunc=qstart.value,
+            )
+
+            # Full experimental data as scatter
+            q_exp_plot = q_values
+            r_exp_plot = reflectivity_curve
+            yerr_plot = (sigmas if pred_show_yerr.value else None)
+            xerr_plot = (q_resolution if pred_show_xerr.value else None)
+
+            # Predicted curve only on the model region
+            q_pred = res.get('q_plot_pred', None)
+            r_pred = res.get('predicted_curve', None)
+
+            # Polished curve on the full experimental grid
+            q_pol = q_values if ('polished_curve' in res) else None
+            r_pol = res.get('polished_curve', None)
+
+            # SLD profiles
+            z_sld = res.get('predicted_sld_xaxis', None)
+            sld_pred = res.get('predicted_sld_profile', None)
+            sld_pol = res.get('sld_profile_polished', None)
+
+            print_prediction_results(res)
+
+            plot_reflectivity(
+                q_exp=q_exp_plot, r_exp=r_exp_plot,
+                yerr=yerr_plot, xerr=xerr_plot,
+                exp_style=('errorbar' if pred_show_yerr.value or pred_show_xerr.value else 'scatter'),
+                exp_color=exp_color_picker.value,
+                exp_errcolor=exp_errcolor_picker.value,
+                q_pred=q_pred, r_pred=r_pred, pred_color=pred_color_picker.value,
+                q_pol=q_pol, r_pol=r_pol, pol_color=pol_color_picker.value,
+                z_sld=z_sld, sld_pred=sld_pred, sld_pol=sld_pol,
+                sld_pred_color=sld_pred_color_picker.value,
+                sld_pol_color=sld_pol_color_picker.value,
+                plot_sld_profile=plot_sld.value,
+                logx=pred_logx.value, logy=True,
+                figsize=(12,6),
+                legend=True
+            )
+
+            self.widget_prediction_result = res
+
+        def _on_close(_):
+            container.close()
+            print("Widget closed.")
+
+        btn_predict.on_click(_on_predict)
+        btn_close.on_click(_on_close)
 
 class InferenceModel(object):
     def __init__(self, name: str = None, trainer: PointEstimatorTrainer = None, preprocessing_parameters: dict = None,