gammasimtools 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

simtools/ray_tracing/psf_parameter_optimisation.py

@@ -0,0 +1,792 @@
+"""
+PSF parameter optimisation and fitting routines for mirror alignment and reflection parameters.
+
+This module provides functions for loading PSF data, generating random parameter sets,
+running PSF simulations, calculating RMSD, and finding the best-fit parameters for a given
+telescope model.
+
+PSF (Point Spread Function) describes how a point source of light is spread out by the
+optical system, and RMSD (Root Mean Squared Deviation) is used as the optimization metric
+to quantify the difference between measured and simulated PSF curves.
+"""
+
+import logging
+from collections import OrderedDict
+
+import astropy.units as u
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.backends.backend_pdf import PdfPages
+
+from simtools.data_model import model_data_writer as writer
+from simtools.model import model_utils
+from simtools.ray_tracing.ray_tracing import RayTracing
+from simtools.utils import general as gen
+from simtools.visualization import visualize
+
+logger = logging.getLogger(__name__)
+
+# Constants
+RADIUS_CM = "Radius [cm]"
+CUMULATIVE_PSF = "Cumulative PSF"
+
+MRRA_RANGE_DEFAULT = 0.004  # Mirror reflection random angle range
+MRF_RANGE_DEFAULT = 0.1  # Mirror reflection fraction range
+MRRA2_RANGE_DEFAULT = 0.03  # Second mirror reflection random angle range
+MAR_RANGE_DEFAULT = 0.005  # Mirror alignment random range
+MAX_OFFSET_DEFAULT = 4.5  # Maximum off-axis angle in degrees
+OFFSET_STEPS_DEFAULT = 0.1  # Step size for off-axis angle sampling
+
+
+def load_psf_data(data_file):
+    """
+    Load data from a text file containing cumulative PSF measurements.
+
+    Parameters
+    ----------
+    data_file : str
+        Name of the data file with the measured cumulative PSF.
+        Expected format:
+        Column 0: radial distance in mm
+        Column 2: cumulative PSF values
+
+    Returns
+    -------
+    numpy.ndarray
+        Loaded and processed data with radius in cm and normalized cumulative PSF.
+    """
+    d_type = {"names": (RADIUS_CM, CUMULATIVE_PSF), "formats": ("f8", "f8")}
+    data = np.loadtxt(data_file, dtype=d_type, usecols=(0, 2))
+    data[RADIUS_CM] *= 0.1  # Convert from mm to cm
+    data[CUMULATIVE_PSF] /= np.max(np.abs(data[CUMULATIVE_PSF]))  # Normalize to max = 1.0
+    return data
+
+
+def calculate_rmsd(data, sim):
+    """Calculate Root Mean Squared Deviation to be used as metric to find the best parameters."""
+    return np.sqrt(np.mean((data - sim) ** 2))
+
+
+def add_parameters(
+    all_parameters,
+    mirror_reflection,
+    mirror_align,
+    mirror_reflection_fraction=0.15,
+    mirror_reflection_2=0.035,
+):
+    """
+    Transform and add parameters to the all_parameters list.
+
+    Parameters
+    ----------
+    mirror_reflection : float
+        The random angle of mirror reflection.
+    mirror_align : float
+        The random angle for mirror alignment (both horizontal and vertical).
+    mirror_reflection_fraction : float, optional
+        The fraction of the mirror reflection. Default is 0.15.
+    mirror_reflection_2 : float, optional
+        A secondary random angle for mirror reflection. Default is 0.035.
+
+    Returns
+    -------
+    None
+        Updates the all_parameters list in place.
+    """
+    pars = {
+        "mirror_reflection_random_angle": [
+            mirror_reflection,
+            mirror_reflection_fraction,
+            mirror_reflection_2,
+        ],
+        "mirror_align_random_horizontal": [mirror_align, 28.0, 0.0, 0.0],
+        "mirror_align_random_vertical": [mirror_align, 28.0, 0.0, 0.0],
+    }
+    all_parameters.append(pars)
+
+
+def get_previous_values(tel_model):
+    """
+    Retrieve previous parameter values from the telescope model.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object.
+
+    Returns
+    -------
+    tuple
+        Tuple containing the previous values of mirror_reflection_random_angle (first entry),
+        mirror_reflection_fraction, second entry), mirror_reflection_random_angle (third entry),
+        and mirror_align_random_horizontal/vertical.
+    """
+    split_par = tel_model.get_parameter_value("mirror_reflection_random_angle")
+    mrra_0, mfr_0, mrra2_0 = split_par[0], split_par[1], split_par[2]
+    mar_0 = tel_model.get_parameter_value("mirror_align_random_horizontal")[0]
+    logger.debug(
+        "Previous parameter values:\n"
+        f"MRRA = {mrra_0!s}\n"
+        f"MRF = {mfr_0!s}\n"
+        f"MRRA2 = {mrra2_0!s}\n"
+        f"MAR = {mar_0!s}\n"
+    )
+    return mrra_0, mfr_0, mrra2_0, mar_0
+
+
+def generate_random_parameters(
+    all_parameters, n_runs, args_dict, mrra_0, mfr_0, mrra2_0, mar_0, tel_model
+):
+    """
+    Generate random parameters for tuning.
+
+    The parameter ranges around the previous values are configurable via module constants.
+
+    Parameters
+    ----------
+    all_parameters : list
+        List to store all parameter sets.
+    n_runs : int
+        Number of random parameter combinations to test.
+    args_dict : dict
+        Dictionary containing parsed command-line arguments.
+    mrra_0 : float
+        Initial value of mirror_reflection_random_angle.
+    mfr_0 : float
+        Initial value of mirror_reflection_fraction.
+    mrra2_0 : float
+        Initial value of the second mirror_reflection_random_angle.
+    mar_0 : float
+        Initial value of mirror_align_random_horizontal/vertical.
+    tel_model : TelescopeModel
+        Telescope model object to check if it's a dual mirror telescope.
+    """
+    if args_dict["fixed"]:
+        logger.debug("fixed=True - First entry of mirror_reflection_random_angle is kept fixed.")
+
+    is_dual_mirror = model_utils.is_two_mirror_telescope(tel_model.name)
+    if is_dual_mirror:
+        mar_fixed_value = 0.0
+    else:
+        mar_fixed_value = None
+
+    for _ in range(n_runs):
+        mrra_range = MRRA_RANGE_DEFAULT if not args_dict["fixed"] else 0
+        mrf_range = MRF_RANGE_DEFAULT
+        mrra2_range = MRRA2_RANGE_DEFAULT
+        mar_range = MAR_RANGE_DEFAULT
+        rng = np.random.default_rng(seed=args_dict.get("random_seed"))
+        mrra = rng.uniform(max(mrra_0 - mrra_range, 0), mrra_0 + mrra_range)
+        mrf = rng.uniform(max(mfr_0 - mrf_range, 0), mfr_0 + mrf_range)
+        mrra2 = rng.uniform(max(mrra2_0 - mrra2_range, 0), mrra2_0 + mrra2_range)
+
+        # Set mar to 0 for dual mirror telescopes, otherwise use random value
+        if mar_fixed_value is not None:
+            mar = mar_fixed_value
+        else:
+            mar = rng.uniform(max(mar_0 - mar_range, 0), mar_0 + mar_range)
+
+        add_parameters(all_parameters, mrra, mar, mrf, mrra2)
+
+
+def _run_ray_tracing_simulation(tel_model, site_model, args_dict, pars):
+    """
+    Run a ray tracing simulation with the given telescope parameters.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object.
+    site_model : SiteModel
+        Site model object.
+    args_dict : dict
+        Dictionary containing parsed command-line arguments.
+    pars : dict
+        Parameter set dictionary.
+
+    Returns
+    -------
+    tuple
+        (d80, simulated_data) - D80 value and simulated data from ray tracing.
+    """
+    if pars is not None:
+        tel_model.change_multiple_parameters(**pars)
+    else:
+        raise ValueError("No best parameters found")
+
+    ray = RayTracing(
+        telescope_model=tel_model,
+        site_model=site_model,
+        simtel_path=args_dict["simtel_path"],
+        zenith_angle=args_dict["zenith"] * u.deg,
+        source_distance=args_dict["src_distance"] * u.km,
+        off_axis_angle=[0.0] * u.deg,
+    )
+    ray.simulate(test=args_dict.get("test", False), force=True)
+    ray.analyze(force=True, use_rx=False)
+    im = ray.images()[0]
+    d80 = im.get_psf()
+
+    return d80, im
+
+
+def _create_psf_simulation_plot(data_to_plot, pars, d80, rmsd, is_best, pdf_pages):
+    """
+    Create a plot for PSF simulation results.
+
+    Parameters
+    ----------
+    data_to_plot : dict
+        Data dictionary for plotting.
+    pars : dict
+        Parameter set dictionary.
+    d80 : float
+        D80 value.
+    rmsd : float
+        RMSD value.
+    is_best : bool
+        Whether this is the best parameter set.
+    pdf_pages : PdfPages
+        PDF pages object for saving plots.
+    """
+    fig = visualize.plot_1d(
+        data_to_plot,
+        plot_difference=True,
+        no_markers=True,
+    )
+    ax = fig.get_axes()[0]
+    ax.set_ylim(0, 1.05)
+    ax.set_ylabel(CUMULATIVE_PSF)
+
+    title_prefix = "* " if is_best else ""
+    ax.set_title(
+        f"{title_prefix}refl_rnd = "
+        f"{pars['mirror_reflection_random_angle'][0]:.5f}, "
+        f"{pars['mirror_reflection_random_angle'][1]:.5f}, "
+        f"{pars['mirror_reflection_random_angle'][2]:.5f}\n"
+        f"align_rnd = {pars['mirror_align_random_vertical'][0]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][1]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][2]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][3]:.5f}"
+    )
+
+    d80_color = "red" if is_best else "black"
+    d80_weight = "bold" if is_best else "normal"
+    d80_text = f"D80 = {d80:.5f} cm"
+
+    ax.text(
+        0.5,
+        0.3,
+        f"{d80_text}\nRMSD = {rmsd:.4f}",
+        verticalalignment="center",
+        horizontalalignment="left",
+        transform=ax.transAxes,
+        color=d80_color,
+        weight=d80_weight,
+        bbox={"boxstyle": "round,pad=0.3", "facecolor": "yellow", "alpha": 0.7}
+        if is_best
+        else None,
+    )
+
+    if is_best:
+        fig.text(
+            0.02,
+            0.02,
+            "* Best parameter set (lowest RMSD)",
+            fontsize=8,
+            style="italic",
+            color="red",
+        )
+
+    pdf_pages.savefig(fig, bbox_inches="tight")
+    plt.clf()
+
+
+def run_psf_simulation(
+    tel_model,
+    site_model,
+    args_dict,
+    pars,
+    data_to_plot,
+    radius,
+    pdf_pages=None,
+    is_best=False,
+    return_simulated_data=False,
+):
+    """
+    Run the simulation for one set of parameters and return D80, RMSD.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object.
+    site_model : SiteModel
+        Site model object.
+    args_dict : dict
+        Dictionary containing parsed command-line arguments.
+    pars : dict
+        Parameter set dictionary.
+    data_to_plot : dict
+        Data dictionary for plotting.
+    radius : array-like
+        Radius data.
+    pdf_pages : PdfPages, optional
+        PDF pages object for plotting. If None, no plotting is done.
+    is_best : bool, optional
+        Whether this is the best parameter set for highlighting in plots.
+    return_simulated_data : bool, optional
+        If True, returns simulated data as third element in return tuple.
+
+    Returns
+    -------
+    tuple
+        (d80, rmsd) if return_simulated_data=False
+        (d80, rmsd, simulated_data) if return_simulated_data=True
+    """
+    d80, im = _run_ray_tracing_simulation(tel_model, site_model, args_dict, pars)
+
+    if radius is None:
+        raise ValueError("Radius data is not available.")
+
+    simulated_data = im.get_cumulative_data(radius * u.cm)
+    rmsd = calculate_rmsd(data_to_plot["measured"][CUMULATIVE_PSF], simulated_data[CUMULATIVE_PSF])
+
+    # Handle plotting if requested
+    if pdf_pages is not None and args_dict.get("plot_all", False):
+        data_to_plot["simulated"] = simulated_data
+        _create_psf_simulation_plot(data_to_plot, pars, d80, rmsd, is_best, pdf_pages)
+        del data_to_plot["simulated"]
+
+    return (d80, rmsd, simulated_data) if return_simulated_data else (d80, rmsd)
+
+
+def load_and_process_data(args_dict):
+    """
+    Load and process data if specified in the command-line arguments.
+
+    Returns
+    -------
+    - data_to_plot: OrderedDict containing loaded and processed data.
+    - radius: Radius data from loaded data (if available).
+    """
+    data_to_plot = OrderedDict()
+    radius = None
+    if args_dict["data"] is not None:
+        data_file = gen.find_file(args_dict["data"], args_dict["model_path"])
+        data_to_plot["measured"] = load_psf_data(data_file)
+        radius = data_to_plot["measured"][RADIUS_CM]
+    return data_to_plot, radius
+
+
+def _create_plot_for_parameters(pars, rmsd, d80, simulated_data, data_to_plot, is_best, pdf_pages):
+    """
+    Create a single plot for a parameter set.
+
+    Parameters
+    ----------
+    pars : dict
+        Parameter set dictionary
+    rmsd : float
+        RMSD value for this parameter set
+    d80 : float
+        D80 value for this parameter set
+    simulated_data : array
+        Simulated data for plotting
+    data_to_plot : dict
+        Data dictionary for plotting
+    is_best : bool
+        Whether this is the best parameter set
+    pdf_pages : PdfPages
+        PDF pages object to save the plot
+    """
+    original_simulated = data_to_plot.get("simulated")
+    data_to_plot["simulated"] = simulated_data
+
+    fig = visualize.plot_1d(
+        data_to_plot,
+        plot_difference=True,
+        no_markers=True,
+    )
+    ax = fig.get_axes()[0]
+    ax.set_ylim(0, 1.05)
+    ax.set_ylabel(CUMULATIVE_PSF)
+
+    title_prefix = "* " if is_best else ""
+
+    ax.set_title(
+        f"{title_prefix}reflection = "
+        f"{pars['mirror_reflection_random_angle'][0]:.5f}, "
+        f"{pars['mirror_reflection_random_angle'][1]:.5f}, "
+        f"{pars['mirror_reflection_random_angle'][2]:.5f}\n"
+        f"align_vertical = {pars['mirror_align_random_vertical'][0]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][1]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][2]:.5f}, "
+        f"{pars['mirror_align_random_vertical'][3]:.5f}\n"
+        f"align_horizontal = {pars['mirror_align_random_horizontal'][0]:.5f}, "
+        f"{pars['mirror_align_random_horizontal'][1]:.5f}, "
+        f"{pars['mirror_align_random_horizontal'][2]:.5f}, "
+        f"{pars['mirror_align_random_horizontal'][3]:.5f}"
+    )
+
+    d80_color = "red" if is_best else "black"
+    d80_weight = "bold" if is_best else "normal"
+
+    ax.text(
+        0.5,
+        0.3,
+        f"D80 = {d80:.5f} cm\nRMSD = {rmsd:.4f}",
+        verticalalignment="center",
+        horizontalalignment="left",
+        transform=ax.transAxes,
+        color=d80_color,
+        weight=d80_weight,
+        bbox={"boxstyle": "round,pad=0.3", "facecolor": "yellow", "alpha": 0.7}
+        if is_best
+        else None,
+    )
+
+    if is_best:
+        fig.text(
+            0.02,
+            0.02,
+            "* Best parameter set (lowest RMSD)",
+            fontsize=8,
+            style="italic",
+            color="red",
+        )
+
+    pdf_pages.savefig(fig, bbox_inches="tight")
+    plt.clf()
+
+    if original_simulated is not None:
+        data_to_plot["simulated"] = original_simulated
+
+
+def _create_all_plots(results, best_pars, data_to_plot, pdf_pages):
+    """
+    Create plots for all parameter sets if requested.
+
+    Parameters
+    ----------
+    results : list
+        List of (pars, rmsd, d80, simulated_data) tuples
+    best_pars : dict
+        Best parameter set for highlighting
+    data_to_plot : dict
+        Data dictionary for plotting
+    pdf_pages : PdfPages
+        PDF pages object to save plots
+    """
+    logger.info("Creating plots for all parameter sets...")
+
+    for i, (pars, rmsd, d80, simulated_data) in enumerate(results):
+        is_best = pars is best_pars
+        logger.info(f"Creating plot {i + 1}/{len(results)}{' (BEST)' if is_best else ''}")
+
+        _create_plot_for_parameters(
+            pars, rmsd, d80, simulated_data, data_to_plot, is_best, pdf_pages
+        )
+
+
+def find_best_parameters(
+    all_parameters, tel_model, site_model, args_dict, data_to_plot, radius, pdf_pages=None
+):
+    """
+    Find the best parameters by running simulations for all parameter sets.
+
+    Loop over all parameter sets, run the simulation, compute RMSD,
+    and return the best parameters and their RMSD.
+    """
+    best_rmsd = float("inf")
+    best_pars = None
+    best_d80 = None
+    results = []  # Store (pars, rmsd, d80, simulated_data)
+
+    logger.info(f"Running {len(all_parameters)} simulations...")
+
+    for i, pars in enumerate(all_parameters):
+        try:
+            logger.info(f"Running simulation {i + 1}/{len(all_parameters)}")
+            d80, rmsd, simulated_data = run_psf_simulation(
+                tel_model,
+                site_model,
+                args_dict,
+                pars,
+                data_to_plot,
+                radius,
+                return_simulated_data=True,
+                pdf_pages=None,
+            )
+        except (ValueError, RuntimeError) as e:
+            logger.warning(f"Simulation failed for parameters {pars}: {e}")
+            continue
+
+        results.append((pars, rmsd, d80, simulated_data))
+        if rmsd < best_rmsd:
+            best_rmsd = rmsd
+            best_pars = pars
+            best_d80 = d80
+
+    logger.info(f"Best RMSD found: {best_rmsd:.5f}")
+
+    # Create all plots if requested
+    if pdf_pages is not None and args_dict.get("plot_all", False) and results:
+        _create_all_plots(results, best_pars, data_to_plot, pdf_pages)
+
+    return best_pars, best_d80, results
+
+
+def create_d80_vs_offaxis_plot(tel_model, site_model, args_dict, best_pars, output_dir):
+    """
+    Create D80 vs off-axis angle plot using the best parameters.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object.
+    site_model : SiteModel
+        Site model object.
+    args_dict : dict
+        Dictionary containing parsed command-line arguments.
+    best_pars : dict
+        Best parameter set.
+    output_dir : Path
+        Output directory for saving plots.
+    """
+    logger.info("Creating D80 vs off-axis angle plot with best parameters...")
+
+    # Apply best parameters to telescope model
+    tel_model.change_multiple_parameters(**best_pars)
+
+    # Create off-axis angle array
+    max_offset = args_dict.get("max_offset", MAX_OFFSET_DEFAULT)
+    offset_steps = args_dict.get("offset_steps", OFFSET_STEPS_DEFAULT)
+    off_axis_angles = np.linspace(
+        0,
+        max_offset,
+        int(max_offset / offset_steps) + 1,
+    )
+
+    ray = RayTracing(
+        telescope_model=tel_model,
+        site_model=site_model,
+        simtel_path=args_dict["simtel_path"],
+        zenith_angle=args_dict["zenith"] * u.deg,
+        source_distance=args_dict["src_distance"] * u.km,
+        off_axis_angle=off_axis_angles * u.deg,
+    )
+
+    logger.info(f"Running ray tracing for {len(off_axis_angles)} off-axis angles...")
+    ray.simulate(test=args_dict.get("test", False), force=True)
+    ray.analyze(force=True)
+
+    for key in ["d80_cm", "d80_deg"]:
+        plt.figure(figsize=(10, 6), tight_layout=True)
+
+        ray.plot(key, marker="o", linestyle="-", color="blue", linewidth=2, markersize=6)
+
+        plt.title(
+            f"PSF D80 vs Off-axis Angle - {tel_model.name}\n"
+            f"Best Parameters: \n"
+            f"reflection=[{best_pars['mirror_reflection_random_angle'][0]:.4f},"
+            f"{best_pars['mirror_reflection_random_angle'][1]:.4f},"
+            f"{best_pars['mirror_reflection_random_angle'][2]:.4f}],\n"
+            f"align_horizontal={best_pars['mirror_align_random_horizontal'][0]:.4f}\n"
+            f"align_vertical={best_pars['mirror_align_random_vertical'][0]:.4f}\n"
+        )
+        plt.xlabel("Off-axis Angle (degrees)")
+        plt.ylabel("D80 (cm)" if key == "d80_cm" else "D80 (degrees)")
+        plt.ylim(bottom=0)
+        plt.xticks(rotation=45)
+        plt.xlim(0, max_offset)
+        plt.grid(True, alpha=0.3)
+
+        plot_file_name = f"tune_psf_{tel_model.name}_best_params_{key}.pdf"
+        plot_file = output_dir.joinpath(plot_file_name)
+        visualize.save_figure(plt, plot_file, log_title=f"D80 vs off-axis ({key})")
+
+    plt.close("all")
+
+
+def write_tested_parameters_to_file(results, best_pars, best_d80, output_dir, tel_model):
+    """
+    Write all tested parameters and their metrics to a text file.
+
+    Parameters
+    ----------
+    results : list
+        List of (pars, rmsd, d80, simulated_data) tuples
+    best_pars : dict
+        Best parameter set
+    best_d80 : float
+        Best D80 value
+    output_dir : Path
+        Output directory path
+    tel_model : TelescopeModel
+        Telescope model object for filename generation
+    """
+    param_file = output_dir.joinpath(f"psf_optimization_{tel_model.name}.log")
+    with open(param_file, "w", encoding="utf-8") as f:
+        f.write("# PSF Parameter Optimization Log\n")
+        f.write(f"# Telescope: {tel_model.name}\n")
+        f.write(f"# Total parameter sets tested: {len(results)}\n")
+        f.write("#" + "=" * 60 + "\n\n")
+
+        f.write("PARAMETER TESTING RESULTS:\n")
+        for i, (pars, rmsd, d80, _) in enumerate(results):
+            is_best = pars is best_pars
+            status = "BEST" if is_best else "TESTED"
+            f.write(f"[{status}] Set {i + 1:03d}: RMSD={rmsd:.5f}, D80={d80:.5f} cm\n")
+            for par, value in pars.items():
+                f.write(f"    {par}: {value}\n")
+            f.write("\n")
+
+        f.write("OPTIMIZATION SUMMARY:\n")
+        f.write(f"Best RMSD: {min(result[1] for result in results):.5f}\n")
+        f.write(f"Best D80: {best_d80:.5f} cm\n")
+        f.write("\nOPTIMIZED PARAMETERS:\n")
+        for par, value in best_pars.items():
+            f.write(f"{par}: {value}\n")
+    return param_file
+
+
+def _add_units_to_psf_parameters(best_pars):
+    """
+    Add proper astropy units to PSF parameters based on their schemas.
+
+    Parameters
+    ----------
+    best_pars : dict
+        Dictionary with PSF parameter names as keys and values as lists
+
+    Returns
+    -------
+    dict
+        Dictionary with same keys but values converted to astropy quantities with units
+    """
+    psf_pars_with_units = {}
+
+    for param_name, param_values in best_pars.items():
+        if param_name == "mirror_reflection_random_angle":
+            psf_pars_with_units[param_name] = [
+                param_values[0] * u.deg,
+                param_values[1] * u.dimensionless_unscaled,
+                param_values[2] * u.deg,
+            ]
+        elif param_name in ["mirror_align_random_horizontal", "mirror_align_random_vertical"]:
+            psf_pars_with_units[param_name] = [
+                param_values[0] * u.deg,
+                param_values[1] * u.deg,
+                param_values[2] * u.dimensionless_unscaled,
+                param_values[3] * u.dimensionless_unscaled,
+            ]
+        else:
+            psf_pars_with_units[param_name] = param_values
+
+    return psf_pars_with_units
+
+
+def export_psf_parameters(best_pars, tel_model, parameter_version, output_dir):
+    """
+    Export PSF parameters as simulation model parameter files.
+
+    Parameters
+    ----------
+    best_pars : dict
+        Best parameter set
+    tel_model : TelescopeModel
+        Telescope model object
+    parameter_version : str
+        Parameter version string
+    output_dir : Path
+        Output directory path
+    """
+    try:
+        logger.info("Exporting best PSF parameters as simulation model parameter files")
+        psf_pars_with_units = _add_units_to_psf_parameters(best_pars)
+        parameter_output_path = output_dir / tel_model.name
+        for parameter_name, parameter_value in psf_pars_with_units.items():
+            writer.ModelDataWriter.dump_model_parameter(
+                parameter_name=parameter_name,
+                value=parameter_value,
+                instrument=tel_model.name,
+                parameter_version=parameter_version,
+                output_file=f"{parameter_name}-{parameter_version}.json",
+                output_path=parameter_output_path,
+                use_plain_output_path=True,
+            )
+        logger.info(f"simulation model parameter files exported to {output_dir}")
+    except ImportError as e:
+        logger.warning(f"Could not export simulation parameters: {e}")
+    except (ValueError, KeyError, OSError) as e:
+        logger.error(f"Error exporting simulation parameters: {e}")
+
+
+def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
+    """
+    Run the complete PSF parameter optimization workflow.
+
+    This function consolidates the main optimization logic to make the application lighter.
+
+    Parameters
+    ----------
+    tel_model : TelescopeModel
+        Telescope model object
+    site_model : SiteModel
+        Site model object
+    args_dict : dict
+        Dictionary containing parsed command-line arguments
+    output_dir : Path
+        Output directory path
+
+    Returns
+    -------
+    None
+        All results are saved to files and printed to console
+    """
+    # Generate parameter sets
+    all_parameters = []
+    mrra_0, mfr_0, mrra2_0, mar_0 = get_previous_values(tel_model)
+
+    n_runs = args_dict.get("n_runs")
+    generate_random_parameters(
+        all_parameters, n_runs, args_dict, mrra_0, mfr_0, mrra2_0, mar_0, tel_model
+    )
+
+    data_to_plot, radius = load_and_process_data(args_dict)
+
+    # Preparing figure name and PDF pages for plotting
+    plot_file_name = "_".join(("tune_psf", tel_model.name + ".pdf"))
+    plot_file = output_dir.joinpath(plot_file_name)
+    pdf_pages = PdfPages(plot_file)
+
+    # Find best parameters
+    best_pars, best_d80, results = find_best_parameters(
+        all_parameters, tel_model, site_model, args_dict, data_to_plot, radius, pdf_pages
+    )
+
+    plt.close()
+    pdf_pages.close()
+
+    # Write all tested parameters and their metrics to a file
+    param_file = write_tested_parameters_to_file(
+        results, best_pars, best_d80, output_dir, tel_model
+    )
+    print(f"\nParameter results written to {param_file}")
+
+    # Automatically create D80 vs off-axis angle plot for best parameters
+    create_d80_vs_offaxis_plot(tel_model, site_model, args_dict, best_pars, output_dir)
+    print("D80 vs off-axis angle plots created successfully")
+
+    print("\nBest parameters:")
+    for par, value in best_pars.items():
+        print(f"{par} = {value}")
+
+    # Export best parameters as simulation model parameter files (if flag is provided)
+    if args_dict.get("write_psf_parameters", False):
+        export_psf_parameters(
+            best_pars,
+            tel_model,
+            args_dict.get("parameter_version", "0.0.0"),
+            output_dir.parent,
+        )