gammasimtools 0.24.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

simtools/ray_tracing/psf_parameter_optimisation.py

@@ -11,6 +11,7 @@ to quantify the difference between measured and simulated PSF curves.
 
 import logging
 from collections import OrderedDict
+from dataclasses import dataclass
 
 import astropy.units as u
 import numpy as np
@@ -20,6 +21,7 @@ from scipy import stats
 from simtools.data_model import model_data_writer as writer
 from simtools.ray_tracing.ray_tracing import RayTracing
 from simtools.utils import general as gen
+from simtools.utils import names
 from simtools.visualization import plot_psf
 from simtools.visualization.plot_psf import DEFAULT_FRACTION, get_psf_diameter_label
 
@@ -32,6 +34,766 @@ CUMULATIVE_PSF = "Cumulative PSF"
 KS_STATISTIC_NAME = "KS statistic"
 
 
+@dataclass
+class GradientStepResult:
+    """
+    Result from a gradient descent step attempt.
+
+    Attributes
+    ----------
+    params : dict or None
+        Updated parameter values, None if step failed.
+    psf_diameter : float or None
+        PSF containment diameter in cm, None if step failed.
+    metric : float or None
+        Optimization metric value (RMSD or KS statistic), None if step failed.
+    p_value : float or None
+        P-value from KS test (None if using RMSD or if step failed).
+    simulated_data : numpy.ndarray or None
+        Simulated PSF data, None if step failed.
+    step_accepted : bool
+        True if the step improved the metric, False otherwise.
+    learning_rate : float
+        Final learning rate after adjustments.
+    """
+
+    params: dict | None
+    psf_diameter: float | None
+    metric: float | None
+    p_value: float | None
+    simulated_data: np.ndarray | None
+    step_accepted: bool
+    learning_rate: float
+
+
+class PSFParameterOptimizer:
+    """
+    Gradient descent optimizer for PSF parameters.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object containing parameter configurations.
+    site_model : SiteModel
+        Site model object with environmental conditions.
+    args_dict : dict
+        Dictionary containing simulation configuration arguments.
+    data_to_plot : dict
+        Dictionary containing measured PSF data under "measured" key.
+    radius : numpy.ndarray
+        Radius values in cm for PSF evaluation.
+    output_dir : Path
+        Directory for saving optimization results and plots.
+
+    Attributes
+    ----------
+    simulation_cache : dict
+        Cache for simulation results to avoid redundant ray tracing simulations.
+        Key: frozenset of (param_name, tuple(values)) items
+        Value: (psf_diameter, metric, p_value, simulated_data)
+    """
+
+    # Learning rate adjustment constants
+    LR_REDUCTION_FACTOR = 0.7
+    LR_INCREASE_FACTOR = 2.0
+    LR_MINIMUM_THRESHOLD = 1e-6
+    LR_MAXIMUM_THRESHOLD = 0.01
+    LR_RESET_VALUE = 0.0001
+
+    def __init__(self, tel_model, site_model, args_dict, data_to_plot, radius, output_dir):
+        """Initialize the PSF parameter optimizer."""
+        self.tel_model = tel_model
+        self.site_model = site_model
+        self.args_dict = args_dict
+        self.data_to_plot = data_to_plot
+        self.radius = radius
+        self.output_dir = output_dir
+        self.use_ks_statistic = args_dict.get("ks_statistic", False)
+        self.fraction = args_dict.get("fraction", DEFAULT_FRACTION)
+        self.simulation_cache = {}
+        self.cache_hits = 0
+        self.cache_misses = 0
+
+    def _params_to_cache_key(self, params):
+        """Convert parameters dict to a hashable cache key."""
+        items = []
+        for key in sorted(params.keys()):
+            value = params[key]
+            if isinstance(value, list):
+                items.append((key, tuple(value)))
+            else:
+                items.append((key, value))
+        return frozenset(items)
+
+    def _reduce_learning_rate(self, current_lr):
+        """
+        Reduce learning rate with minimum threshold and reset.
+
+        Parameters
+        ----------
+        current_lr : float
+            Current learning rate.
+
+        Returns
+        -------
+        float
+            Reduced learning rate, reset to LR_RESET_VALUE if below threshold.
+        """
+        new_lr = current_lr * self.LR_REDUCTION_FACTOR
+        if new_lr < self.LR_MINIMUM_THRESHOLD:
+            return self.LR_RESET_VALUE
+        return new_lr
+
+    def _increase_learning_rate(self, current_lr):
+        """
+        Increase learning rate.
+
+        Parameters
+        ----------
+        current_lr : float
+            Current learning rate.
+
+        Returns
+        -------
+        float
+            Increased learning rate, capped at LR_MAXIMUM_THRESHOLD.
+        """
+        new_lr = current_lr * self.LR_INCREASE_FACTOR
+        return min(new_lr, self.LR_MAXIMUM_THRESHOLD)
+
+    def get_initial_parameters(self):
+        """
+        Get current PSF parameter values from the telescope model.
+
+        Returns
+        -------
+        dict
+            Dictionary of current parameter values.
+        """
+        return get_previous_values(self.tel_model)
+
+    def run_simulation(
+        self, pars, pdf_pages=None, is_best=False, use_cache=True, use_ks_statistic=None
+    ):
+        """
+        Run PSF simulation for given parameters with optional caching.
+
+        Parameters
+        ----------
+        pars : dict
+            Dictionary of parameter values to test.
+        pdf_pages : PdfPages, optional
+            PDF pages object for saving plots.
+        is_best : bool, optional
+            Flag indicating if this is the best parameter set.
+        use_cache : bool, optional
+            If True, use cached results if available.
+        use_ks_statistic : bool, optional
+            If provided, override self.use_ks_statistic for this simulation.
+
+        Returns
+        -------
+        tuple
+            (psf_diameter, metric, p_value, simulated_data)
+        """
+        # Determine which statistic to use
+        ks_stat = use_ks_statistic if use_ks_statistic is not None else self.use_ks_statistic
+
+        if use_cache and pdf_pages is None and not is_best:
+            cache_key = self._params_to_cache_key(pars)
+            if cache_key in self.simulation_cache:
+                self.cache_hits += 1
+                return self.simulation_cache[cache_key]
+            self.cache_misses += 1
+
+        result = run_psf_simulation(
+            self.tel_model,
+            self.site_model,
+            self.args_dict,
+            pars,
+            self.data_to_plot,
+            self.radius,
+            pdf_pages,
+            is_best,
+            ks_stat,
+        )
+
+        # Cache the result if caching is enabled and not plotting
+        if use_cache and pdf_pages is None and not is_best:
+            cache_key = self._params_to_cache_key(pars)
+            self.simulation_cache[cache_key] = result
+
+        return result
+
+    def calculate_gradient(self, current_params, current_metric, epsilon=0.0005):
+        """
+        Calculate numerical gradients for all optimization parameters.
+
+        Parameters
+        ----------
+        current_params : dict
+            Dictionary of current parameter values.
+        current_metric : float
+            Current RMSD or KS statistic value.
+        epsilon : float, optional
+            Perturbation value for finite difference calculation.
+
+        Returns
+        -------
+        dict or None
+            Dictionary mapping parameter names to their gradient values.
+            Returns None if gradient calculation fails for any parameter.
+        """
+        gradients = {}
+
+        for param_name, param_values in current_params.items():
+            param_gradient = self._calculate_param_gradient(
+                current_params,
+                current_metric,
+                param_name,
+                param_values,
+                epsilon,
+            )
+            if param_gradient is None:
+                return None
+            gradients[param_name] = param_gradient
+
+        return gradients
+
+    def apply_gradient_step(self, current_params, gradients, learning_rate):
+        """
+        Apply gradient descent step to update parameters while preserving constraints.
+
+        This function applies the standard gradient descent update and preserves
+        zenith angle components (index 1) for mirror alignment parameters.
+
+        Note: Use _are_all_parameters_within_allowed_range() to validate the result before
+        accepting the step.
+
+        Parameters
+        ----------
+        current_params : dict
+            Dictionary of current parameter values.
+        gradients : dict
+            Dictionary of gradient values for each parameter.
+        learning_rate : float
+            Step size for the gradient descent update.
+
+        Returns
+        -------
+        dict
+            Dictionary of updated parameter values after applying the gradient step.
+        """
+        new_params = {}
+        for param_name, param_values in current_params.items():
+            param_gradients = gradients[param_name]
+
+            if isinstance(param_values, list):
+                updated_values = []
+                for i, (value, gradient) in enumerate(zip(param_values, param_gradients)):
+                    # Apply gradient descent update
+                    new_value = value - learning_rate * gradient
+
+                    # Enforce constraint: preserve zenith angle (index 1) for mirror alignment
+                    if (
+                        param_name
+                        in ["mirror_align_random_horizontal", "mirror_align_random_vertical"]
+                        and i == 1
+                    ):
+                        new_value = value  # Keep original zenith angle value
+
+                    updated_values.append(new_value)
+
+                new_params[param_name] = updated_values
+            else:
+                new_value = param_values - learning_rate * param_gradients
+                new_params[param_name] = new_value
+
+        return new_params
+
+    def _calculate_param_gradient(
+        self, current_params, current_metric, param_name, param_values, epsilon
+    ):
+        """
+        Calculate numerical gradient for a single parameter using finite differences.
+
+        Parameters
+        ----------
+        current_params : dict
+            Dictionary of current parameter values for all optimization parameters.
+        current_metric : float
+            Current RMSD or KS statistic value.
+        param_name : str
+            Name of the parameter for which to calculate the gradient.
+        param_values : float or list
+            Current value(s) of the parameter.
+        epsilon : float
+            Small perturbation value for finite difference calculation.
+
+        Returns
+        -------
+        float or list or None
+            Gradient value(s) for the parameter.
+            Returns None if simulation fails for any component.
+        """
+        param_gradients = []
+        values_list = param_values if isinstance(param_values, list) else [param_values]
+
+        for i, value in enumerate(values_list):
+            perturbed_params = _create_perturbed_params(
+                current_params, param_name, param_values, i, value, epsilon
+            )
+
+            # Calculate gradient for this parameter component using cached simulations
+            try:
+                _, perturbed_metric, _, _ = self.run_simulation(perturbed_params, use_cache=True)
+                gradient = (perturbed_metric - current_metric) / epsilon
+            except (ValueError, RuntimeError) as e:
+                logger.warning(f"Simulation failed for {param_name}[{i}] gradient calculation: {e}")
+                return None
+
+            param_gradients.append(gradient)
+
+        return param_gradients[0] if not isinstance(param_values, list) else param_gradients
+
+    def perform_gradient_step_with_retries(
+        self, current_params, current_metric, learning_rate, max_retries=3
+    ):
+        """
+        Attempt gradient descent step with adaptive learning rate reduction.
+
+        Parameters
+        ----------
+        current_params : dict
+            Dictionary of current parameter values.
+        current_metric : float
+            Current optimization metric value.
+        learning_rate : float
+            Initial learning rate for the gradient descent step.
+        max_retries : int, optional
+            Maximum number of attempts with learning rate reduction.
+
+        Returns
+        -------
+        GradientStepResult
+            Result object containing updated parameters and step status.
+        """
+        current_lr = learning_rate
+
+        for attempt in range(max_retries):
+            try:
+                gradients = self.calculate_gradient(current_params, current_metric)
+
+                if gradients is None:
+                    logger.warning(
+                        f"Gradient calculation failed on attempt {attempt + 1}, skipping step"
+                    )
+                    return GradientStepResult(
+                        params=None,
+                        psf_diameter=None,
+                        metric=None,
+                        p_value=None,
+                        simulated_data=None,
+                        step_accepted=False,
+                        learning_rate=current_lr,
+                    )
+
+                new_params = self.apply_gradient_step(current_params, gradients, current_lr)
+
+                # Validate that all parameters are within allowed ranges
+                if not _are_all_parameters_within_allowed_range(new_params):
+                    logger.info(
+                        f"Step rejected: parameters would go out of bounds with learning rate "
+                        f"{current_lr:.6f}, reducing to {current_lr * self.LR_REDUCTION_FACTOR:.6f}"
+                    )
+                    current_lr = self._reduce_learning_rate(current_lr)
+                    continue
+
+                new_psf_diameter, new_metric, new_p_value, new_simulated_data = self.run_simulation(
+                    new_params, use_cache=True
+                )
+
+                if new_metric < current_metric:
+                    return GradientStepResult(
+                        params=new_params,
+                        psf_diameter=new_psf_diameter,
+                        metric=new_metric,
+                        p_value=new_p_value,
+                        simulated_data=new_simulated_data,
+                        step_accepted=True,
+                        learning_rate=current_lr,
+                    )
+
+                logger.info(
+                    f"Step rejected (RMSD {current_metric:.6f} -> {new_metric:.6f}), "
+                    f"reducing learning rate to {current_lr * self.LR_REDUCTION_FACTOR:.6f}"
+                )
+                current_lr = self._reduce_learning_rate(current_lr)
+
+            except (ValueError, RuntimeError, KeyError) as e:
+                logger.warning(f"Simulation failed on attempt {attempt + 1}: {e}")
+                continue
+
+        return GradientStepResult(
+            params=None,
+            psf_diameter=None,
+            metric=None,
+            p_value=None,
+            simulated_data=None,
+            step_accepted=False,
+            learning_rate=current_lr,
+        )
+
+    def _create_step_plot(
+        self,
+        pdf_pages,
+        current_params,
+        new_psf_diameter,
+        new_metric,
+        new_p_value,
+        new_simulated_data,
+    ):
+        """Create plot for an accepted gradient step."""
+        if (
+            pdf_pages is None
+            or not self.args_dict.get("plot_all", False)
+            or new_simulated_data is None
+        ):
+            return
+
+        self.data_to_plot["simulated"] = new_simulated_data
+        plot_psf.create_psf_parameter_plot(
+            self.data_to_plot,
+            current_params,
+            new_psf_diameter,
+            new_metric,
+            False,
+            pdf_pages,
+            fraction=self.fraction,
+            p_value=new_p_value,
+            use_ks_statistic=self.use_ks_statistic,
+        )
+        del self.data_to_plot["simulated"]
+
+    def _create_final_plot(self, pdf_pages, best_params, best_psf_diameter):
+        """Create final plot for best parameters."""
+        if pdf_pages is None or best_params is None:
+            return
+
+        logger.info("Creating final plot for best parameters with both RMSD and KS statistic...")
+        _, best_ks_stat, best_p_value, best_simulated_data = self.run_simulation(
+            best_params,
+            pdf_pages=None,
+            is_best=False,
+            use_cache=False,
+            use_ks_statistic=True,
+        )
+        best_rmsd = calculate_rmsd(
+            self.data_to_plot["measured"][CUMULATIVE_PSF], best_simulated_data[CUMULATIVE_PSF]
+        )
+
+        self.data_to_plot["simulated"] = best_simulated_data
+        plot_psf.create_psf_parameter_plot(
+            self.data_to_plot,
+            best_params,
+            best_psf_diameter,
+            best_rmsd,
+            True,
+            pdf_pages,
+            fraction=self.fraction,
+            p_value=best_p_value,
+            use_ks_statistic=False,
+            second_metric=best_ks_stat,
+        )
+        del self.data_to_plot["simulated"]
+        pdf_pages.close()
+        logger.info("Cumulative PSF plots saved")
+
+    def run_gradient_descent(self, rmsd_threshold, learning_rate, max_iterations=200):
+        """
+        Run gradient descent optimization to minimize PSF fitting metric.
+
+        Parameters
+        ----------
+        rmsd_threshold : float
+            Convergence threshold for RMSD improvement.
+        learning_rate : float
+            Initial learning rate for gradient descent steps.
+        max_iterations : int, optional
+            Maximum number of optimization iterations.
+
+        Returns
+        -------
+        tuple
+            (best_params, best_psf_diameter, results)
+        """
+        if self.data_to_plot is None or self.radius is None:
+            logger.error("No PSF measurement data provided. Cannot run optimization.")
+            return None, None, []
+
+        current_params = self.get_initial_parameters()
+        pdf_pages = plot_psf.setup_pdf_plotting(
+            self.args_dict, self.output_dir, self.tel_model.name
+        )
+        results = []
+
+        # Evaluate initial parameters
+        current_psf_diameter, current_metric, current_p_value, simulated_data = self.run_simulation(
+            current_params,
+            pdf_pages=pdf_pages if self.args_dict.get("plot_all", False) else None,
+            is_best=False,
+        )
+
+        results.append(
+            (
+                current_params.copy(),
+                current_metric,
+                current_p_value,
+                current_psf_diameter,
+                simulated_data,
+            )
+        )
+        best_metric, best_params, best_psf_diameter = (
+            current_metric,
+            current_params.copy(),
+            current_psf_diameter,
+        )
+
+        logger.info(
+            f"Initial RMSD: {current_metric:.6f}, PSF diameter: {current_psf_diameter:.6f} cm"
+        )
+
+        iteration = 0
+        current_lr = learning_rate
+
+        while iteration < max_iterations:
+            if current_metric <= rmsd_threshold:
+                logger.info(
+                    f"Optimization converged: RMSD {current_metric:.6f} <= "
+                    f"threshold {rmsd_threshold:.6f}"
+                )
+                break
+
+            iteration += 1
+            logger.info(f"Gradient descent iteration {iteration}")
+
+            step_result = self.perform_gradient_step_with_retries(
+                current_params,
+                current_metric,
+                current_lr,
+            )
+
+            if not step_result.step_accepted or step_result.params is None:
+                current_lr = self._increase_learning_rate(current_lr)
+                logger.info(f"No step accepted, increasing learning rate to {current_lr:.6f}")
+                continue
+
+            # Step was accepted - update state
+            current_params = step_result.params
+            current_metric = step_result.metric
+            current_psf_diameter = step_result.psf_diameter
+            current_lr = step_result.learning_rate
+
+            results.append(
+                (
+                    current_params.copy(),
+                    current_metric,
+                    None,
+                    current_psf_diameter,
+                    step_result.simulated_data,
+                )
+            )
+
+            if current_metric < best_metric:
+                best_metric, best_params, best_psf_diameter = (
+                    current_metric,
+                    current_params.copy(),
+                    current_psf_diameter,
+                )
+
+            self._create_step_plot(
+                pdf_pages,
+                current_params,
+                step_result.psf_diameter,
+                step_result.metric,
+                step_result.p_value,
+                step_result.simulated_data,
+            )
+            logger.info(f"  Accepted step: improved to {step_result.metric:.6f}")
+
+        self._create_final_plot(pdf_pages, best_params, best_psf_diameter)
+        return best_params, best_psf_diameter, results
+
+    def analyze_monte_carlo_error(self, n_simulations=500):
+        """
+        Analyze Monte Carlo uncertainty in PSF optimization metrics.
+
+        Parameters
+        ----------
+        n_simulations : int, optional
+            Number of Monte Carlo simulations to run.
+
+        Returns
+        -------
+        tuple
+            Monte Carlo analysis results.
+        """
+        if self.data_to_plot is None or self.radius is None:
+            logger.error("No PSF measurement data provided. Cannot analyze Monte Carlo error.")
+            return None, None, [], None, None, [], None, None, []
+
+        initial_params = self.get_initial_parameters()
+        for param_name, param_values in initial_params.items():
+            logger.info(f"  {param_name}: {param_values}")
+
+        metric_values, p_values, psf_diameter_values = [], [], []
+
+        for i in range(n_simulations):
+            try:
+                psf_diameter, metric, p_value, _ = self.run_simulation(
+                    initial_params,
+                    use_cache=False,
+                )
+                metric_values.append(metric)
+                psf_diameter_values.append(psf_diameter)
+                p_values.append(p_value)
+            except (ValueError, RuntimeError) as e:
+                logger.warning(f"WARNING: Simulation {i + 1} failed: {e}")
+
+        if not metric_values:
+            logger.error("All Monte Carlo simulations failed.")
+            return None, None, [], None, None, [], None, None, []
+
+        mean_metric, std_metric = np.mean(metric_values), np.std(metric_values, ddof=1)
+        mean_psf_diameter, std_psf_diameter = (
+            np.mean(psf_diameter_values),
+            np.std(psf_diameter_values, ddof=1),
+        )
+
+        if self.use_ks_statistic:
+            valid_p_values = [p for p in p_values if p is not None]
+            mean_p_value = np.mean(valid_p_values) if valid_p_values else None
+            std_p_value = np.std(valid_p_values, ddof=1) if valid_p_values else None
+        else:
+            mean_p_value = std_p_value = None
+
+        return (
+            mean_metric,
+            std_metric,
+            metric_values,
+            mean_p_value,
+            std_p_value,
+            p_values,
+            mean_psf_diameter,
+            std_psf_diameter,
+            psf_diameter_values,
+        )
+
+
+def _is_parameter_within_allowed_range(param_name, param_index, value):
+    """
+    Check if a parameter value is within the allowed range defined in its schema.
+
+    Parameters
+    ----------
+    param_name : str
+        Name of the parameter to check.
+    param_index : int
+        Index within the parameter array (for multi-component parameters).
+    value : float
+        The parameter value to check.
+
+    Returns
+    -------
+    bool
+        True if the parameter is within allowed range, False otherwise.
+        Returns True if no range constraints are defined for the parameter.
+    """
+    try:
+        param_schema = names.model_parameters().get(param_name)
+        if param_schema is None:
+            return True
+
+        data = param_schema.get("data")
+        if not isinstance(data, list) or len(data) <= param_index:
+            return True
+
+        param_data = data[param_index]
+        allowed_range = param_data.get("allowed_range")
+        if not allowed_range:
+            return True
+
+        min_val = allowed_range.get("min")
+        max_val = allowed_range.get("max")
+
+        if min_val is not None and value < min_val:
+            return False
+        return not (max_val is not None and value > max_val)
+
+    except (KeyError, IndexError) as e:
+        logger.warning(f"Error reading schema for {param_name}[{param_index}]: {e}")
+        return True
+
+
+def _are_all_parameters_within_allowed_range(params):
+    """
+    Check if all parameters in the parameter dictionary are within allowed ranges.
+
+    Parameters
+    ----------
+    params : dict
+        Dictionary of parameter values to validate.
+
+    Returns
+    -------
+    bool
+        True if all parameters are within allowed ranges, False otherwise.
+    """
+    for name, values in params.items():
+        values = values if isinstance(values, list) else [values]
+        for i, v in enumerate(values):
+            if not _is_parameter_within_allowed_range(name, i, v):
+                logger.debug(f"{name}[{i}]={v:.6f} out of range")
+                return False
+    return True
+
+
+def _create_perturbed_params(current_params, param_name, param_values, param_index, value, epsilon):
+    """
+    Create parameter dictionary with one parameter perturbed by epsilon.
+
+    Parameters
+    ----------
+    current_params : dict
+        Dictionary of current parameter values.
+    param_name : str
+        Name of the parameter to perturb.
+    param_values : float or list
+        Current parameter values.
+    param_index : int
+        Index of the parameter to perturb (for list parameters).
+    value : float
+        Current value of the specific parameter component.
+    epsilon : float
+        Perturbation amount.
+
+    Returns
+    -------
+    dict
+        New parameter dictionary with the specified parameter perturbed.
+    """
+    perturbed_params = {
+        k: v.copy() if isinstance(v, list) else v for k, v in current_params.items()
+    }
+
+    if isinstance(param_values, list):
+        perturbed_params[param_name][param_index] = value + epsilon
+    else:
+        perturbed_params[param_name] = value + epsilon
+
+    return perturbed_params
+
+
 def _create_log_header_and_format_value(title, tel_model, additional_info=None, value=None):
     """Create log header and format parameter values."""
     if value is not None:  # Format value mode
@@ -137,7 +899,7 @@ def run_psf_simulation(
         Dictionary of parameter values to test in the simulation.
     data_to_plot : dict
         Dictionary containing measured PSF data under "measured" key.
-    radius : array-like
+    radius : numpy.ndarray
         Radius values in cm for PSF evaluation.
     pdf_pages : PdfPages, optional
         PDF pages object for saving plots (default: None).
@@ -340,7 +1102,7 @@ def export_psf_parameters(best_pars, telescope, parameter_version, output_dir):
     """
     try:
         psf_pars_with_units = _add_units_to_psf_parameters(best_pars)
-        parameter_output_path = output_dir.parent / telescope
+        parameter_output_path = output_dir.joinpath(telescope)
         for parameter_name, parameter_value in psf_pars_with_units.items():
             writer.ModelDataWriter.dump_model_parameter(
                 parameter_name=parameter_name,
@@ -356,534 +1118,6 @@ def export_psf_parameters(best_pars, telescope, parameter_version, output_dir):
         logger.error(f"Error exporting simulation parameters: {e}")
 
 
-def _calculate_param_gradient(
-    tel_model,
-    site_model,
-    args_dict,
-    current_params,
-    data_to_plot,
-    radius,
-    current_rmsd,
-    param_name,
-    param_values,
-    epsilon,
-    use_ks_statistic,
-):
-    """
-    Calculate numerical gradient for a single parameter using finite differences.
-
-    The gradient is calculated using forward finite differences:
-    gradient = (f(x + epsilon) - f(x)) / epsilon
-
-    Parameters
-    ----------
-    tel_model : TelescopeModel
-        The telescope model object containing the current parameter configuration.
-    site_model : SiteModel
-        The site model object with environmental conditions.
-    args_dict : dict
-        Dictionary containing simulation arguments and configuration options.
-    current_params : dict
-        Dictionary of current parameter values for all optimization parameters.
-    data_to_plot : dict
-        Dictionary containing measured PSF data with "measured" key.
-    radius : array-like
-        Radius values in cm for PSF evaluation.
-    current_rmsd : float
-        Current RMSD at the current parameter configuration.
-    param_name : str
-        Name of the parameter for which to calculate the gradient.
-    param_values : float or list
-        Current value(s) of the parameter. Can be a single value or list of values.
-    epsilon : float
-        Small perturbation value for finite difference calculation.
-    use_ks_statistic : bool
-        If True, calculate gradient with respect to KS statistic; if False, use RMSD.
-
-    Returns
-    -------
-    float or list
-        Gradient value(s) for the parameter. Returns a single float if param_values
-        is a single value, or a list of gradients if param_values is a list.
-
-    If a simulation fails during gradient calculation, a gradient of 0.0 is assigned
-    for that component to ensure the optimization can continue.
-    """
-    param_gradients = []
-    values_list = param_values if isinstance(param_values, list) else [param_values]
-
-    for i, value in enumerate(values_list):
-        perturbed_params = {
-            k: v.copy() if isinstance(v, list) else v for k, v in current_params.items()
-        }
-
-        if isinstance(param_values, list):
-            perturbed_params[param_name][i] = value + epsilon
-        else:
-            perturbed_params[param_name] = value + epsilon
-
-        try:
-            _, perturbed_rmsd, _, _ = run_psf_simulation(
-                tel_model,
-                site_model,
-                args_dict,
-                perturbed_params,
-                data_to_plot,
-                radius,
-                pdf_pages=None,
-                is_best=False,
-                use_ks_statistic=use_ks_statistic,
-            )
-            param_gradients.append((perturbed_rmsd - current_rmsd) / epsilon)
-        except (ValueError, RuntimeError):
-            param_gradients.append(0.0)
-
-    return param_gradients[0] if not isinstance(param_values, list) else param_gradients
-
-
-def calculate_gradient(
-    tel_model,
-    site_model,
-    args_dict,
-    current_params,
-    data_to_plot,
-    radius,
-    current_rmsd,
-    epsilon=0.0005,
-    use_ks_statistic=False,
-):
-    """
-    Calculate numerical gradients for all optimization parameters.
-
-    Parameters
-    ----------
-    tel_model : TelescopeModel
-        Telescope model object for simulations.
-    site_model : SiteModel
-        Site model object with environmental conditions.
-    args_dict : dict
-        Dictionary containing simulation configuration arguments.
-    current_params : dict
-        Dictionary of current parameter values for all optimization parameters.
-    data_to_plot : dict
-        Dictionary containing measured PSF data.
-    radius : array-like
-        Radius values in cm for PSF evaluation.
-    current_rmsd : float
-        Current RMSD or KS statistic value.
-    epsilon : float, optional
-        Perturbation value for finite difference calculation (default: 0.0005).
-    use_ks_statistic : bool, optional
-        If True, calculate gradients for KS statistic; if False, use RMSD (default: False).
-
-    Returns
-    -------
-    dict
-        Dictionary mapping parameter names to their gradient values.
-        For parameters with multiple components, gradients are returned as lists.
-    """
-    gradients = {}
-
-    for param_name, param_values in current_params.items():
-        gradients[param_name] = _calculate_param_gradient(
-            tel_model,
-            site_model,
-            args_dict,
-            current_params,
-            data_to_plot,
-            radius,
-            current_rmsd,
-            param_name,
-            param_values,
-            epsilon,
-            use_ks_statistic,
-        )
-
-    return gradients
-
-
-def apply_gradient_step(current_params, gradients, learning_rate):
-    """
-    Apply gradient descent step to update parameters.
-
-    Parameters
-    ----------
-    current_params : dict
-        Dictionary of current parameter values.
-    gradients : dict
-        Dictionary of gradient values for each parameter.
-    learning_rate : float
-        Step size for the gradient descent update.
-
-    Returns
-    -------
-    dict
-        Dictionary of updated parameter values after applying the gradient step.
-    """
-    new_params = {}
-    for param_name, param_values in current_params.items():
-        param_gradients = gradients[param_name]
-
-        if isinstance(param_values, list):
-            new_params[param_name] = [
-                value - learning_rate * gradient
-                for value, gradient in zip(param_values, param_gradients)
-            ]
-        else:
-            new_params[param_name] = param_values - learning_rate * param_gradients
-
-    return new_params
-
-
-def _perform_gradient_step_with_retries(
-    tel_model,
-    site_model,
-    args_dict,
-    current_params,
-    current_metric,
-    data_to_plot,
-    radius,
-    learning_rate,
-    max_retries=3,
-):
-    """
-    Attempt gradient descent step with adaptive learning rate reduction on rejection.
-
-    The learning rate reduction strategy follows these rules:
-    - If step is rejected: learning_rate *= 0.7
-    - If attempt number < number of max retries then try again
-    - If learning_rate drops below 1e-5: reset to 0.001
-    - If all retries fail: returns None values with step_accepted=False
-
-    This adaptive approach helps navigate local minima and ensures robust convergence
-    by automatically adjusting the step size based on optimization progress.
-
-    Parameters
-    ----------
-    tel_model : TelescopeModel
-        Telescope model object containing the current parameter configuration.
-    site_model : SiteModel
-        Site model object with environmental conditions for ray tracing simulations.
-    args_dict : dict
-        Dictionary containing simulation configuration arguments and settings.
-    current_params : dict
-        Dictionary of current parameter values for all optimization parameters.
-    current_metric : float
-        Current optimization metric value (RMSD or KS statistic) to improve upon.
-    data_to_plot : dict
-        Dictionary containing measured PSF data under "measured" key for comparison.
-    radius : array-like
-        Radius values in cm for PSF evaluation and comparison.
-    learning_rate : float
-        Initial learning rate for the gradient descent step.
-    max_retries : int, optional
-        Maximum number of attempts with learning rate reduction (default: 3).
-
-    Returns
-    -------
-    tuple of (dict, float, float, float or None, array, bool, float)
-        - new_params: Updated parameter dictionary if step accepted, None if rejected
-        - new_psf_diameter: PSF containment diameter in cm for new parameters, None if step rejected
-        - new_metric: New optimization metric value, None if step rejected
-        - new_p_value: p-value from KS test if applicable, None otherwise
-        - new_simulated_data: Simulated PSF data array, None if step rejected
-        - step_accepted: Boolean indicating if any step was accepted
-        - final_learning_rate: Learning rate after potential reductions
-
-    """
-    current_lr = learning_rate
-
-    for attempt in range(max_retries):
-        try:
-            gradients = calculate_gradient(
-                tel_model,
-                site_model,
-                args_dict,
-                current_params,
-                data_to_plot,
-                radius,
-                current_metric,
-                use_ks_statistic=False,
-            )
-            new_params = apply_gradient_step(current_params, gradients, current_lr)
-
-            new_psf_diameter, new_metric, new_p_value, new_simulated_data = run_psf_simulation(
-                tel_model,
-                site_model,
-                args_dict,
-                new_params,
-                data_to_plot,
-                radius,
-                pdf_pages=None,
-                is_best=False,
-                use_ks_statistic=False,
-            )
-
-            if new_metric < current_metric:
-                return (
-                    new_params,
-                    new_psf_diameter,
-                    new_metric,
-                    new_p_value,
-                    new_simulated_data,
-                    True,
-                    current_lr,
-                )
-
-            logger.info(
-                f"Step rejected (RMSD {current_metric:.6f} -> {new_metric:.6f}), "
-                f"reducing learning rate {current_lr:.6f} -> {current_lr * 0.7:.6f}"
-            )
-            current_lr *= 0.7
-
-            if current_lr < 1e-5:
-                current_lr = 0.001
-
-        except (ValueError, RuntimeError, KeyError) as e:
-            logger.warning(f"Simulation failed on attempt {attempt + 1}: {e}")
-            continue
-
-    return None, None, None, None, None, False, current_lr
-
-
-def _create_step_plot(
-    pdf_pages,
-    args_dict,
-    data_to_plot,
-    current_params,
-    new_psf_diameter,
-    new_metric,
-    new_p_value,
-    new_simulated_data,
-):
-    """Create plot for an accepted gradient step."""
-    if pdf_pages is None or not args_dict.get("plot_all", False) or new_simulated_data is None:
-        return
-
-    data_to_plot["simulated"] = new_simulated_data
-    plot_psf.create_psf_parameter_plot(
-        data_to_plot,
-        current_params,
-        new_psf_diameter,
-        new_metric,
-        False,
-        pdf_pages,
-        fraction=args_dict.get("fraction", DEFAULT_FRACTION),
-        p_value=new_p_value,
-        use_ks_statistic=False,
-    )
-    del data_to_plot["simulated"]
-
-
-def _create_final_plot(
-    pdf_pages,
-    tel_model,
-    site_model,
-    args_dict,
-    best_params,
-    data_to_plot,
-    radius,
-    best_psf_diameter,
-):
-    """Create final plot for best parameters."""
-    if pdf_pages is None or best_params is None:
-        return
-
-    logger.info("Creating final plot for best parameters with both RMSD and KS statistic...")
-    _, best_ks_stat, best_p_value, best_simulated_data = run_psf_simulation(
-        tel_model,
-        site_model,
-        args_dict,
-        best_params,
-        data_to_plot,
-        radius,
-        pdf_pages=None,
-        is_best=False,
-        use_ks_statistic=True,
-    )
-    best_rmsd = calculate_rmsd(
-        data_to_plot["measured"][CUMULATIVE_PSF], best_simulated_data[CUMULATIVE_PSF]
-    )
-
-    data_to_plot["simulated"] = best_simulated_data
-    plot_psf.create_psf_parameter_plot(
-        data_to_plot,
-        best_params,
-        best_psf_diameter,
-        best_rmsd,
-        True,
-        pdf_pages,
-        fraction=args_dict.get("fraction", DEFAULT_FRACTION),
-        p_value=best_p_value,
-        use_ks_statistic=False,
-        second_metric=best_ks_stat,
-    )
-    del data_to_plot["simulated"]
-    pdf_pages.close()
-    logger.info("Cumulative PSF plots saved")
-
-
-def run_gradient_descent_optimization(
-    tel_model,
-    site_model,
-    args_dict,
-    data_to_plot,
-    radius,
-    rmsd_threshold,
-    learning_rate,
-    output_dir,
-):
-    """
-    Run gradient descent optimization to minimize PSF fitting metric.
-
-    Parameters
-    ----------
-    tel_model : TelescopeModel
-        Telescope model object to be optimized.
-    site_model : SiteModel
-        Site model object with environmental conditions.
-    args_dict : dict
-        Dictionary containing simulation configuration arguments.
-    data_to_plot : dict
-        Dictionary containing measured PSF data under "measured" key.
-    radius : array-like
-        Radius values in cm for PSF evaluation.
-    rmsd_threshold : float
-        Convergence threshold for RMSD improvement.
-    learning_rate : float
-        Initial learning rate for gradient descent steps.
-    output_dir : Path
-        Directory for saving optimization plots and results.
-
-    Returns
-    -------
-    tuple of (dict, float, list)
-        - best_params: Dictionary of optimized parameter values
-        - best_psf_diameter: PSF containment diameter in cm for the best parameters
-        - results: List of (params, metric, p_value, psf_diameter, simulated_data)
-          for each iteration
-
-    Returns None values if optimization fails or no measurement data is provided.
-    """
-    if data_to_plot is None or radius is None:
-        logger.error("No PSF measurement data provided. Cannot run optimization.")
-        return None, None, []
-
-    current_params = get_previous_values(tel_model)
-    pdf_pages = plot_psf.setup_pdf_plotting(args_dict, output_dir, tel_model.name)
-    results = []
-
-    # Evaluate initial parameters
-    current_psf_diameter, current_metric, current_p_value, simulated_data = run_psf_simulation(
-        tel_model,
-        site_model,
-        args_dict,
-        current_params,
-        data_to_plot,
-        radius,
-        pdf_pages=pdf_pages if args_dict.get("plot_all", False) else None,
-        is_best=False,
-        use_ks_statistic=False,
-    )
-
-    results.append(
-        (
-            current_params.copy(),
-            current_metric,
-            current_p_value,
-            current_psf_diameter,
-            simulated_data,
-        )
-    )
-    best_metric, best_params, best_psf_diameter = (
-        current_metric,
-        current_params.copy(),
-        current_psf_diameter,
-    )
-
-    logger.info(f"Initial RMSD: {current_metric:.6f}, PSF diameter: {current_psf_diameter:.6f} cm")
-
-    iteration = 0
-    max_total_iterations = 100
-
-    while iteration < max_total_iterations:
-        if current_metric <= rmsd_threshold:
-            logger.info(
-                f"Optimization converged: RMSD {current_metric:.6f} <= "
-                f"threshold {rmsd_threshold:.6f}"
-            )
-            break
-
-        iteration += 1
-        logger.info(f"Gradient descent iteration {iteration}")
-
-        step_result = _perform_gradient_step_with_retries(
-            tel_model,
-            site_model,
-            args_dict,
-            current_params,
-            current_metric,
-            data_to_plot,
-            radius,
-            learning_rate,
-        )
-        (
-            new_params,
-            new_psf_diameter,
-            new_metric,
-            new_p_value,
-            new_simulated_data,
-            step_accepted,
-            learning_rate,
-        ) = step_result
-
-        if not step_accepted or new_params is None:
-            learning_rate *= 2.0
-            logger.info(f"No step accepted, increasing learning rate to {learning_rate:.6f}")
-            continue
-
-        # Step was accepted - update state
-        current_params, current_metric, current_psf_diameter = (
-            new_params,
-            new_metric,
-            new_psf_diameter,
-        )
-        results.append(
-            (current_params.copy(), current_metric, None, current_psf_diameter, new_simulated_data)
-        )
-
-        if current_metric < best_metric:
-            best_metric, best_params, best_psf_diameter = (
-                current_metric,
-                current_params.copy(),
-                current_psf_diameter,
-            )
-
-        _create_step_plot(
-            pdf_pages,
-            args_dict,
-            data_to_plot,
-            current_params,
-            new_psf_diameter,
-            new_metric,
-            new_p_value,
-            new_simulated_data,
-        )
-        logger.info(f"  Accepted step: improved to {new_metric:.6f}")
-
-    _create_final_plot(
-        pdf_pages,
-        tel_model,
-        site_model,
-        args_dict,
-        best_params,
-        data_to_plot,
-        radius,
-        best_psf_diameter,
-    )
-    return best_params, best_psf_diameter, results
-
-
 def _write_log_interpretation(f, use_ks_statistic):
     """Write interpretation section for the log file."""
     if use_ks_statistic:
@@ -1027,103 +1261,6 @@ def write_gradient_descent_log(
     return param_file
 
 
-def analyze_monte_carlo_error(
-    tel_model, site_model, args_dict, data_to_plot, radius, n_simulations=500
-):
-    """
-    Analyze Monte Carlo uncertainty in PSF optimization metrics.
-
-    Runs multiple simulations with the same parameters to quantify the
-    statistical uncertainty in the optimization metric due to Monte Carlo
-    noise in the ray tracing simulations. Returns None values if no
-    measurement data is provided or all simulations fail.
-
-    Parameters
-    ----------
-    tel_model : TelescopeModel
-        Telescope model object with current parameter configuration.
-    site_model : SiteModel
-        Site model object with environmental conditions.
-    args_dict : dict
-        Dictionary containing simulation configuration arguments.
-    data_to_plot : dict
-        Dictionary containing measured PSF data under "measured" key.
-    radius : array-like
-        Radius values in cm for PSF evaluation.
-    n_simulations : int, optional
-        Number of Monte Carlo simulations to run (default: 500).
-
-    Returns
-    -------
-    tuple of (float, float, list, float, float, list, float, float, list)
-        - mean_metric: Mean RMSD or KS statistic value
-        - std_metric: Standard deviation of metric values
-        - metric_values: List of all metric values from simulations
-        - mean_p_value: Mean p-value (None if using RMSD)
-        - std_p_value: Standard deviation of p-values (None if using RMSD)
-        - p_values: List of all p-values from simulations
-        - mean_psf_diameter: Mean PSF containment diameter in cm
-        - std_psf_diameter: Standard deviation of PSF diameter values
-        - psf_diameter_values: List of all PSF diameter values from simulations
-    """
-    if data_to_plot is None or radius is None:
-        logger.error("No PSF measurement data provided. Cannot analyze Monte Carlo error.")
-        return None, None, []
-
-    initial_params = get_previous_values(tel_model)
-    for param_name, param_values in initial_params.items():
-        logger.info(f"  {param_name}: {param_values}")
-
-    use_ks_statistic = args_dict.get("ks_statistic", False)
-    metric_values, p_values, psf_diameter_values = [], [], []
-
-    for i in range(n_simulations):
-        try:
-            psf_diameter, metric, p_value, _ = run_psf_simulation(
-                tel_model,
-                site_model,
-                args_dict,
-                initial_params,
-                data_to_plot,
-                radius,
-                use_ks_statistic=use_ks_statistic,
-            )
-            metric_values.append(metric)
-            psf_diameter_values.append(psf_diameter)
-            p_values.append(p_value)
-        except (ValueError, RuntimeError) as e:
-            logger.warning(f"WARNING: Simulation {i + 1} failed: {e}")
-
-    if not metric_values:
-        logger.error("All Monte Carlo simulations failed.")
-        return None, None, [], None, None, []
-
-    mean_metric, std_metric = np.mean(metric_values), np.std(metric_values, ddof=1)
-    mean_psf_diameter, std_psf_diameter = (
-        np.mean(psf_diameter_values),
-        np.std(psf_diameter_values, ddof=1),
-    )
-
-    if use_ks_statistic:
-        valid_p_values = [p for p in p_values if p is not None]
-        mean_p_value = np.mean(valid_p_values) if valid_p_values else None
-        std_p_value = np.std(valid_p_values, ddof=1) if valid_p_values else None
-    else:
-        mean_p_value = std_p_value = None
-
-    return (
-        mean_metric,
-        std_metric,
-        metric_values,
-        mean_p_value,
-        std_p_value,
-        p_values,
-        mean_psf_diameter,
-        std_psf_diameter,
-        psf_diameter_values,
-    )
-
-
 def write_monte_carlo_analysis(
     mc_results, output_dir, tel_model, use_ks_statistic=False, fraction=DEFAULT_FRACTION
 ):
@@ -1222,12 +1359,7 @@ def write_monte_carlo_analysis(
             zip(metric_values, p_values, psf_diameter_values)
         ):
             if use_ks_statistic and p_value is not None:
-                if p_value > 0.05:
-                    significance = "GOOD"
-                elif p_value > 0.01:
-                    significance = "FAIR"
-                else:
-                    significance = "POOR"
+                significance = plot_psf.get_significance_label(p_value)
                 f.write(
                     f"Simulation {i + 1:2d}: {metric_name}={metric_val:.6f}, "
                     f"p_value={p_value:.6f} ({significance}), {psf_label}={psf_diameter:.6f} cm\n"
@@ -1241,63 +1373,71 @@ def write_monte_carlo_analysis(
     return mc_file
 
 
-def _handle_monte_carlo_analysis(
-    tel_model, site_model, args_dict, data_to_plot, radius, output_dir, use_ks_statistic
-):
-    """Handle Monte Carlo analysis if requested."""
-    if not args_dict.get("monte_carlo_analysis", False):
-        return False
-
-    mc_results = analyze_monte_carlo_error(tel_model, site_model, args_dict, data_to_plot, radius)
-    if mc_results[0] is not None:
-        mc_file = write_monte_carlo_analysis(
-            mc_results,
-            output_dir,
-            tel_model,
-            use_ks_statistic,
-            args_dict.get("fraction", DEFAULT_FRACTION),
-        )
-        logger.info(f"Monte Carlo analysis results written to {mc_file}")
-        mc_plot_file = output_dir.joinpath(f"monte_carlo_uncertainty_{tel_model.name}.pdf")
-        plot_psf.create_monte_carlo_uncertainty_plot(
-            mc_results, mc_plot_file, args_dict.get("fraction", DEFAULT_FRACTION), use_ks_statistic
-        )
-    return True
+def cleanup_intermediate_files(output_dir):
+    """
+    Remove intermediate log and list files from the output directory.
+
+    Parameters
+    ----------
+    output_dir : Path
+        Directory containing output files to clean up.
+    """
+    patterns = ["*.log", "*.lis*"]
+    files_removed = 0
+
+    for pattern in patterns:
+        for file_path in output_dir.glob(pattern):
+            file_path.unlink()
+            files_removed += 1
+            logger.debug(f"Removed: {file_path.name}")
+
+    if files_removed > 0:
+        logger.info(f"Cleanup: removed {files_removed} intermediate files")
 
 
 def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
-    """Run the complete PSF parameter optimization workflow using gradient descent."""
+    """
+    Run the complete PSF parameter optimization workflow using gradient descent.
+
+    This function creates a PSFParameterOptimizer instance and orchestrates
+    the optimization process.
+    """
     data_to_plot, radius = load_and_process_data(args_dict)
-    use_ks_statistic = args_dict.get("ks_statistic", False)
 
-    if _handle_monte_carlo_analysis(
-        tel_model, site_model, args_dict, data_to_plot, radius, output_dir, use_ks_statistic
-    ):
+    # Create optimizer instance to encapsulate state and methods
+    optimizer = PSFParameterOptimizer(
+        tel_model, site_model, args_dict, data_to_plot, radius, output_dir
+    )
+
+    # Handle Monte Carlo analysis if requested
+    if args_dict.get("monte_carlo_analysis", False):
+        mc_results = optimizer.analyze_monte_carlo_error()
+        if mc_results[0] is not None:
+            mc_file = write_monte_carlo_analysis(
+                mc_results,
+                output_dir,
+                tel_model,
+                optimizer.use_ks_statistic,
+                optimizer.fraction,
+            )
+            logger.info(f"Monte Carlo analysis results written to {mc_file}")
+            mc_plot_file = output_dir.joinpath(f"monte_carlo_uncertainty_{tel_model.name}.pdf")
+            plot_psf.create_monte_carlo_uncertainty_plot(
+                mc_results, mc_plot_file, optimizer.fraction, optimizer.use_ks_statistic
+            )
         return
 
     # Run gradient descent optimization
     threshold = args_dict.get("rmsd_threshold")
     learning_rate = args_dict.get("learning_rate")
 
-    best_pars, best_psf_diameter, gd_results = run_gradient_descent_optimization(
-        tel_model,
-        site_model,
-        args_dict,
-        data_to_plot,
-        radius,
-        rmsd_threshold=threshold,
-        learning_rate=learning_rate,
-        output_dir=output_dir,
+    best_pars, best_psf_diameter, gd_results = optimizer.run_gradient_descent(
+        threshold, learning_rate
     )
 
     # Check if optimization was successful
     if not gd_results or best_pars is None:
-        logger.error("Gradient descent optimization failed. No valid results found.")
-        if radius is None:
-            logger.error(
-                "Possible cause: No PSF measurement data provided. "
-                "Use --data argument to provide PSF data."
-            )
+        logger.error("Gradient descent optimization failed to produce results.")
         return
 
     plot_psf.create_optimization_plots(args_dict, gd_results, tel_model, data_to_plot, output_dir)
@@ -1309,8 +1449,8 @@ def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
         gd_results,
         threshold,
         convergence_plot_file,
-        args_dict.get("fraction", DEFAULT_FRACTION),
-        use_ks_statistic,
+        optimizer.fraction,
+        optimizer.use_ks_statistic,
     )
 
     param_file = write_gradient_descent_log(
@@ -1319,8 +1459,8 @@ def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
         best_psf_diameter,
         output_dir,
         tel_model,
-        use_ks_statistic,
-        args_dict.get("fraction", DEFAULT_FRACTION),
+        optimizer.use_ks_statistic,
+        optimizer.fraction,
     )
     logger.info(f"\nGradient descent progression written to {param_file}")
 
@@ -1331,3 +1471,6 @@ def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
         export_psf_parameters(
             best_pars, args_dict.get("telescope"), args_dict.get("parameter_version"), output_dir
         )
+
+    if args_dict.get("cleanup", False):
+        cleanup_intermediate_files(output_dir)