gammasimtools 0.26.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

simtools/ray_tracing/mirror_panel_psf.py

@@ -1,293 +1,413 @@
-"""Mirror panel PSF calculation."""
+"""Mirror panel PSF calculation with per-mirror PSF diameter optimization."""
 
+import json
 import logging
+import os
+from dataclasses import asdict, dataclass
+from pathlib import Path
 
-import astropy.units as u
 import numpy as np
-from astropy.table import QTable, Table
+from astropy.table import Table
 
-import simtools.data_model.model_data_writer as writer
 import simtools.utils.general as gen
-from simtools.data_model.metadata_collector import MetadataCollector
+from simtools.data_model import model_data_writer
+from simtools.job_execution.process_pool import process_pool_map_ordered
 from simtools.model.model_utils import initialize_simulation_models
 from simtools.ray_tracing.ray_tracing import RayTracing
+from simtools.utils import names
+from simtools.visualization import plot_psf
+
+
+@dataclass
+class MirrorOptimizationResult:
+    """Dataclass to store the result of a single mirror RNDA optimization."""
+
+    mirror: int  # Zero-based mirror index
+    measured_psf_mm: float  # Measured PSF diameter in mm
+    optimized_rnda: list[float]  # Optimized RNDA values [sigma1, fraction2, sigma2]
+    simulated_psf_mm: float  # Simulated PSF diameter after optimization
+    percentage_diff: float  # Absolute percentage difference
+    containment_fraction: float  # PSF containment fraction used for the diameter
+
+
+def _optimize_single_mirror_worker(args):
+    """Worker wrapper for optimizing a single mirror in parallel."""
+    instance, mirror_idx, measured = args
+    return instance.optimize_single_mirror(mirror_idx, measured)
 
 
 class MirrorPanelPSF:
     """
-    Mirror panel PSF and random reflection angle calculation.
+    Mirror panel PSF and random reflection angle (RNDA) calculation.
 
-    This class is used to derive the random reflection angle for the mirror panels in the telescope.
-
-    Known limitations: single Gaussian PSF model, no support for multiple PSF components (as allowed
-    in the model parameters).
+    This class derives the RNDA for mirror panels by optimizing
+    per-mirror PSF diameters (at a given containment fraction)
+    using a percentage difference metric.
+    Optimization uses a gradient descent with finite-difference gradients.
 
     Parameters
     ----------
-    label: str
+    label : str
         Application label.
-    args_dict: dict
-        Dictionary with input arguments.
+    args_dict : dict
+        Dictionary with input arguments, e.g. site, telescope, data path, model version.
     """
 
+    # Hard guardrails
+    GRAD_CLIP = 1e4
+    MAX_LOG_STEP = 0.25
+    MAX_FRAC_STEP = 0.1
+    MAX_ITER = 100
+
     def __init__(self, label, args_dict):
-        """Initialize the MirrorPanelPSF class."""
         self._logger = logging.getLogger(__name__)
-        self._logger.debug("Initializing MirrorPanelPSF")
-
+        self.label = label
         self.args_dict = args_dict
-        self.telescope_model, self.site_model = self._define_telescope_model(label)
 
-        if self.args_dict["test"]:
-            self.args_dict["number_of_mirrors_to_test"] = 2
+        self.fraction = args_dict.get("fraction", 0.8)
+
+        self.telescope_model, self.site_model, _ = initialize_simulation_models(
+            label=label,
+            site=args_dict["site"],
+            telescope_name=args_dict["telescope"],
+            model_version=args_dict["model_version"],
+        )
 
-        if self.args_dict["psf_measurement"]:
-            self._get_psf_containment()
-        if not self.args_dict["psf_measurement_containment_mean"]:
-            raise ValueError("Missing PSF measurement")
+        self.measured_data = self._load_measured_data()
+        self.rnda_start = list(
+            self.telescope_model.get_parameter_value("mirror_reflection_random_angle")
+        )
 
-        self.mean_d80 = None
-        self.sig_d80 = None
-        self.rnda_start = self._get_starting_value()
+        self.per_mirror_results = []
         self.rnda_opt = None
-        self.results_rnda = []
-        self.results_mean = []
-        self.results_sig = []
+        self.final_percentage_diff = None
 
-    def _define_telescope_model(self, label):
+    def _load_measured_data(self):
         """
-        Define telescope model.
-
-        This includes updating the configuration with mirror list and/or random focal length given
-        as input.
+        Load measured PSF diameter from ECSV file.
 
         Returns
         -------
-        tel : TelescopeModel
-            The telescope model.
-        site_model : SiteModel
-            The site model.
+        astropy.table.Column
+            Column containing PSF diameter values in mm.
         """
-        tel_model, site_model, _ = initialize_simulation_models(
-            label=label,
-            site=self.args_dict["site"],
-            telescope_name=self.args_dict["telescope"],
-            model_version=self.args_dict["model_version"],
-        )
-        if self.args_dict["mirror_list"] is not None:
-            mirror_list_file = gen.find_file(
-                name=self.args_dict["mirror_list"], loc=self.args_dict["model_path"]
-            )
-            tel_model.overwrite_model_parameter("mirror_list", self.args_dict["mirror_list"])
-            tel_model.overwrite_model_file("mirror_list", mirror_list_file)
-        if self.args_dict["random_focal_length"] is not None:
-            tel_model.overwrite_model_parameter(
-                "random_focal_length", str(self.args_dict["random_focal_length"])
-            )
+        data_file = gen.find_file(self.args_dict["data"], self.args_dict.get("model_path", "."))
+        table = Table.read(data_file)
+        if "psf_opt" in table.colnames:
+            return table["psf_opt"]
+        if "d80" in table.colnames:
+            fraction = float(self.args_dict.get("fraction", 0.8))
+            if not np.isclose(fraction, 0.8, rtol=1e-09, atol=1e-09):
+                self._logger.warning(
+                    "Input table provides 'd80' column, but --fraction=%.3f was requested. "
+                    "Make sure the measured column matches the selected containment fraction.",
+                    fraction,
+                )
+            return table["d80"]
+        raise ValueError("Data file must contain either 'psf_opt' or 'd80' column")
 
-        return tel_model, site_model
+    def _simulate_single_mirror_psf(self, mirror_idx, rnda_values):
+        """
+        Simulate a single mirror and return its PSF diameter.
 
-    def _get_psf_containment(self):
-        """Read measured single-mirror point-spread function from file and return mean and sigma."""
-        # If this is a test, read just the first few lines since we only simulate those mirrors
-        data_end = (
-            self.args_dict["number_of_mirrors_to_test"] + 1 if self.args_dict["test"] else None
-        )
-        _psf_list = Table.read(
-            self.args_dict["psf_measurement"], format="ascii.ecsv", data_end=data_end
+        Parameters
+        ----------
+        mirror_idx : int
+            Zero-based index of the mirror.
+        rnda_values : list of float
+            RNDA values [sigma1, fraction2, sigma2].
+
+        Returns
+        -------
+        float
+            Simulated PSF diameter in mm.
+        """
+        fraction = self.fraction
+        self.telescope_model.overwrite_model_parameter(
+            "mirror_reflection_random_angle", rnda_values
         )
-        try:
-            self.args_dict["psf_measurement_containment_mean"] = np.nanmean(
-                np.array(_psf_list["psf_opt"].to("cm").value)
-            )
-            self.args_dict["psf_measurement_containment_sigma"] = np.nanstd(
-                np.array(_psf_list["psf_opt"].to("cm").value)
-            )
-        except KeyError as exc:
-            raise KeyError(
-                f"Missing column psf measurement (psf_opt) in {self.args_dict['psf_measurement']}"
-            ) from exc
 
-        self._logger.info(
-            f"Determined PSF containment to {self.args_dict['psf_measurement_containment_mean']:.4}"
-            f" +- {self.args_dict['psf_measurement_containment_sigma']:.4} cm"
+        rt_label = f"{self.label}_m{mirror_idx}"
+        ray = RayTracing(
+            telescope_model=self.telescope_model,
+            site_model=self.site_model,
+            label=rt_label,
+            single_mirror_mode=True,
+            mirror_numbers=[mirror_idx],
         )
+        ray.simulate(test=self.args_dict.get("test", False), force=True)
+        ray.analyze(force=True, containment_fraction=fraction)
+
+        return float(ray.get_psf_mm())
+
+    @staticmethod
+    def _signed_pct_diff(measured, simulated):
+        """Compute signed percentage difference."""
+        if measured <= 0:
+            raise ValueError("Measured PSF diameter must be positive")
+        return 100.0 * (simulated - measured) / measured
 
-    def derive_random_reflection_angle(self, save_figures=False):
+    def _evaluate(self, mirror_idx, measured, rnda):
         """
-        Minimize the difference between measured and simulated PSF for reflection angle.
+        Evaluate the objective function for a mirror.
 
-        Main loop of the optimization process. The method iterates over different values of the
-        random reflection angle until the difference in the mean value of the D80 containment
-        is minimal.
+        Parameters
+        ----------
+        mirror_idx : int
+            Mirror index (0-based).
+        measured : float
+            Measured PSF diameter (mm).
+        rnda : list of float
+            Current RNDA values.
+
+        Returns
+        -------
+        tuple of (float, float, float)
+            Simulated PSF diameter, signed percentage difference, squared percentage difference.
         """
-        if self.args_dict["no_tuning"]:
-            self.rnda_opt = self.rnda_start
-        else:
-            self._optimize_reflection_angle()
+        sim = self._simulate_single_mirror_psf(mirror_idx, rnda)
+        pct = self._signed_pct_diff(measured, sim)
+        return sim, pct, pct * pct
 
-        self.mean_d80, self.sig_d80 = self.run_simulations_and_analysis(
-            self.rnda_opt, save_figures=save_figures
-        )
+    @staticmethod
+    def _rnda_bounds():
+        """
+        Get allowed bounds for RNDA parameters.
 
-    def _optimize_reflection_angle(self, step_size=0.1, max_iteration=100):
+        Returns
+        -------
+        list of tuple of float
+            [(min1, max1), (min2, max2), (min3, max3)]
         """
-        Optimize the random reflection angle to minimize the difference in D80 containment.
+        schema = names.model_parameters()["mirror_reflection_random_angle"]["data"]
+        return [
+            (max(schema[0]["allowed_range"]["min"], 1e-12), schema[0]["allowed_range"]["max"]),
+            (schema[1]["allowed_range"]["min"], schema[1]["allowed_range"]["max"]),
+            (max(schema[2]["allowed_range"]["min"], 1e-12), schema[2]["allowed_range"]["max"]),
+        ]
+
+    def optimize_single_mirror(self, mirror_idx, measured_psf):
+        """
+        Optimize RNDA for a single mirror using gradient descent.
 
         Parameters
         ----------
-        step_size: float
-            Initial step size for optimization.
-        max_iteration: int
-            Maximum number of iterations.
-
-        Raises
-        ------
-        ValueError
-            If the optimization reaches the maximum number of iterations without converging.
+        mirror_idx : int
+            Zero-based mirror index.
+        measured_psf : float
+            Measured PSF diameter in mm.
 
+        Returns
+        -------
+        MirrorOptimizationResult
+            Optimization results for the mirror.
         """
-        relative_tolerance_d80 = self.args_dict["rtol_psf_containment"]
+        threshold_pct = 100 * float(self.args_dict.get("threshold", 0.05))
+        learning_rate = float(self.args_dict.get("learning_rate", 1e-3))
+        bounds = self._rnda_bounds()
+        fraction = self.fraction
+
+        rnda = list(self.rnda_start)
+        sim, pct, obj = self._evaluate(mirror_idx, measured_psf, rnda)
+        best = {"rnda": list(rnda), "sim": sim, "pct": abs(pct), "obj": obj}
+
         self._logger.info(
-            "Optimizing random reflection angle "
-            f"(relative tolerance = {relative_tolerance_d80}, "
-            f"step size = {step_size}, max iteration = {max_iteration})"
+            "Mirror %d | initial PSF %.3f mm (f=%.2f) | pct %.2f%%",
+            mirror_idx,
+            sim,
+            fraction,
+            pct,
         )
 
-        def collect_results(rnda, mean, sig):
-            self.results_rnda.append(rnda)
-            self.results_mean.append(mean)
-            self.results_sig.append(sig)
+        for iteration in range(self.MAX_ITER):
+            old_rnda = list(rnda)
+            for param_idx, (param_value, param_bounds) in enumerate(zip(old_rnda, bounds)):
+                rnda[param_idx] = self._update_single_rnda_parameter(
+                    mirror_idx=mirror_idx,
+                    measured_psf=measured_psf,
+                    rnda=rnda,
+                    param_idx=param_idx,
+                    param_value=param_value,
+                    param_bounds=param_bounds,
+                    learning_rate=learning_rate,
+                )
 
-        reference_d80 = self.args_dict["psf_measurement_containment_mean"]
-        rnda = self.rnda_start
-        prev_error_d80 = float("inf")
-        iteration = 0
+            sim, pct, obj = self._evaluate(mirror_idx, measured_psf, rnda)
+            pct = abs(pct)
 
-        while True:
-            mean_d80, sig_d80 = self.run_simulations_and_analysis(rnda)
-            error_d80 = abs(1 - mean_d80 / reference_d80)
-            collect_results(rnda, mean_d80, sig_d80)
+            old_rnda_str = "[" + ", ".join(f"{v:.4g}" for v in old_rnda) + "]"
+            rnda_str = "[" + ", ".join(f"{v:.4g}" for v in rnda) + "]"
 
-            if error_d80 < relative_tolerance_d80:
-                break
+            self._logger.info(
+                "Iter %d | rnda %s -> %s | pct %.2f -> %.2f | lr %.2g",
+                iteration + 1,
+                old_rnda_str,
+                rnda_str,
+                best["pct"],
+                pct,
+                learning_rate,
+            )
 
-            if mean_d80 < reference_d80:
-                rnda += step_size * self.rnda_start
+            if obj < best["obj"]:
+                best.update(rnda=list(rnda), sim=sim, pct=pct, obj=obj)
+                learning_rate *= 1.1
             else:
-                rnda -= step_size * self.rnda_start
-
-            if error_d80 >= prev_error_d80:
-                step_size = step_size / 2
-            prev_error_d80 = error_d80
-            iteration += 1
-            if iteration > max_iteration:
-                raise ValueError(
-                    f"Maximum iterations ({max_iteration}) reached without convergence."
-                )
+                rnda = list(old_rnda)
+                learning_rate *= 0.5
 
-        self.rnda_opt = rnda
-
-    def _get_starting_value(self):
-        """Get optimization starting value from command line or previous model."""
-        if self.args_dict["rnda"] != 0:
-            rnda_start = self.args_dict["rnda"]
-        else:
-            rnda_start = self.telescope_model.get_parameter_value("mirror_reflection_random_angle")[
-                0
-            ]
+            if best["pct"] <= threshold_pct or learning_rate < 1e-12:
+                break
 
-        self._logger.info(f"Start value for mirror_reflection_random_angle: {rnda_start} deg")
-        return rnda_start
+        return MirrorOptimizationResult(
+            mirror=mirror_idx,
+            measured_psf_mm=float(measured_psf),
+            optimized_rnda=best["rnda"],
+            simulated_psf_mm=best["sim"],
+            percentage_diff=best["pct"],
+            containment_fraction=fraction,
+        )
 
-    def run_simulations_and_analysis(self, rnda, save_figures=False):
+    def _update_single_rnda_parameter(
+        self,
+        mirror_idx,
+        measured_psf,
+        rnda,
+        param_idx,
+        param_value,
+        param_bounds,
+        learning_rate,
+    ):
         """
-        Run ray tracing simulations and analysis for one given value of rnda.
+        Update a single RNDA parameter using finite-difference gradient descent.
 
         Parameters
         ----------
-        rnda: float
-            Random reflection angle in degrees.
-        save_figures: bool
-            Save figures.
+        mirror_idx : int
+        measured_psf : float
+        rnda : list of float
+        param_idx : int
+        param_value : float
+        param_bounds : tuple of float
+        learning_rate : float
 
         Returns
         -------
-        mean_d80: float
-            Mean value of D80 in cm.
-        sig_d80: float
-            Standard deviation of D80 in cm.
+        float
+            Updated RNDA parameter value.
         """
-        self.telescope_model.overwrite_model_parameter("mirror_reflection_random_angle", rnda)
-        ray = RayTracing(
-            telescope_model=self.telescope_model,
-            site_model=self.site_model,
-            single_mirror_mode=True,
-            mirror_numbers=(
-                list(range(1, self.args_dict["number_of_mirrors_to_test"] + 1))
-                if self.args_dict["test"]
-                else "all"
-            ),
-            use_random_focal_length=self.args_dict["use_random_focal_length"],
-            random_focal_length_seed=self.args_dict.get("random_focal_length_seed"),
-        )
-        ray.simulate(test=self.args_dict["test"], force=True)  # force has to be True, always
-        ray.analyze(force=True)
-        if save_figures:
-            ray.plot("d80_cm", save=True, d80=self.args_dict["psf_measurement_containment_mean"])
-
-        return (
-            ray.get_mean("d80_cm").to(u.cm).value,
-            ray.get_std_dev("d80_cm").to(u.cm).value,
-        )
+        param_min, param_max = param_bounds
+        is_log_param = param_idx != 1
+        epsilon = max(1e-6, 0.05 * param_value) if is_log_param else 0.05
 
-    def print_results(self):
-        """Print results to stdout."""
-        containment_fraction_percent = int(self.args_dict["containment_fraction"] * 100)
+        rnda[param_idx] = min(param_max, param_value + epsilon)
+        _, _, f_plus = self._evaluate(mirror_idx, measured_psf, rnda)
 
-        print(f"\nMeasured D{containment_fraction_percent}:")
-        if self.args_dict["psf_measurement_containment_sigma"] is not None:
-            print(
-                f"Mean = {self.args_dict['psf_measurement_containment_mean']:.3f} cm, "
-                f"StdDev = {self.args_dict['psf_measurement_containment_sigma']:.3f} cm"
-            )
-        else:
-            print(f"Mean = {self.args_dict['psf_measurement_containment_mean']:.3f} cm")
-        print(f"\nSimulated D{containment_fraction_percent}:")
-        print(f"Mean = {self.mean_d80:.3f} cm, StdDev = {self.sig_d80:.3f} cm")
-        print("\nmirror_random_reflection_angle")
-        print(f"Previous value = {self.rnda_start:.6f}")
-        print(f"New value = {self.rnda_opt:.6f}\n")
+        rnda[param_idx] = max(param_min, param_value - epsilon)
+        _, _, f_minus = self._evaluate(mirror_idx, measured_psf, rnda)
 
-    def write_optimization_data(self):
+        rnda[param_idx] = param_value
+        gradient = (f_plus - f_minus) / (2 * epsilon)
+
+        # Clip the gradient to avoid excessively large steps.
+        if is_log_param:
+            gradient = np.clip(gradient * param_value, -self.GRAD_CLIP, self.GRAD_CLIP)
+            step = np.clip(-learning_rate * gradient, -self.MAX_LOG_STEP, self.MAX_LOG_STEP)
+            # clip the parameters so they are within allowed range
+            return np.clip(param_value * np.exp(step), param_min, param_max)
+
+        gradient = np.clip(gradient, -self.GRAD_CLIP, self.GRAD_CLIP)
+        step = np.clip(-learning_rate * gradient, -self.MAX_FRAC_STEP, self.MAX_FRAC_STEP)
+        return np.clip(param_value + step, param_min, param_max)
+
+    def optimize_with_gradient_descent(self):
         """
-        Write optimization results to an astropy table (ecsv file).
+        Optimize all mirrors in parallel using process pool.
 
-        Used mostly for debugging of the optimization process.
-        The first entry of the table is the best fit result.
+        Sets
+        ----
+        self.per_mirror_results : list of dict
+        self.rnda_opt : list of float
+        self.final_percentage_diff : float
         """
-        containment_fraction_percent = int(self.args_dict["containment_fraction"] * 100)
-
-        result_table = QTable(
-            [
-                [True] + [False] * len(self.results_rnda),
-                [self.rnda_opt, *self.results_rnda] * u.deg,
-                ([0.0] * (len(self.results_rnda) + 1)),
-                ([0.0] * (len(self.results_rnda) + 1)) * u.deg,
-                [self.mean_d80, *self.results_mean] * u.cm,
-                [self.sig_d80, *self.results_sig] * u.cm,
-            ],
-            names=(
-                "best_fit",
-                "mirror_reflection_random_angle_sigma1",
-                "mirror_reflection_random_angle_fraction2",
-                "mirror_reflection_random_angle_sigma2",
-                f"containment_radius_D{containment_fraction_percent}",
-                f"containment_radius_sigma_D{containment_fraction_percent}",
-            ),
+        n_mirrors = len(self.measured_data)
+        if self.args_dict.get("test"):
+            n_mirrors = min(n_mirrors, self.args_dict.get("number_of_mirrors_to_test"))
+
+        n_workers = int(self.args_dict.get("n_workers") or os.cpu_count())
+        parent = MirrorPanelPSF(self.label, dict(self.args_dict))
+        worker_args = [(parent, i, parent.measured_data[i]) for i in range(n_mirrors)]
+
+        self.per_mirror_results = process_pool_map_ordered(
+            _optimize_single_mirror_worker,
+            worker_args,
+            max_workers=n_workers,
+            mp_start_method="fork",
         )
-        writer.ModelDataWriter.dump(
-            args_dict=self.args_dict,
-            metadata=MetadataCollector(args_dict=self.args_dict),
-            product_data=result_table,
+
+        self.rnda_opt = np.mean(
+            [r.optimized_rnda for r in self.per_mirror_results], axis=0
+        ).tolist()
+
+        self.final_percentage_diff = float(
+            np.mean([r.percentage_diff for r in self.per_mirror_results])
         )
+
+    def write_optimization_data(self):
+        """Write optimization results and optionally export as a model parameter."""
+        output_dir = Path(self.args_dict.get("output_path", "."))
+        telescope = self.args_dict.get("telescope")
+        parameter_version = self.args_dict.get("parameter_version")
+        parameter_name = "mirror_reflection_random_angle"
+        parameter_output_path = output_dir / str(telescope)
+        parameter_output_path.mkdir(parents=True, exist_ok=True)
+        output_file = parameter_output_path / "per_mirror_rnda.json"
+
+        per_mirror_results_out = []
+        for result in self.per_mirror_results:
+            result_dict = asdict(result)
+            for k, v in result_dict.items():
+                if k == "optimized_rnda":
+                    result_dict[k] = [float(f"{x:.4f}") for x in v]
+                elif isinstance(v, (int, float)):
+                    result_dict[k] = float(f"{v:.4f}")
+            per_mirror_results_out.append(result_dict)
+
+        output_file.write_text(
+            json.dumps(
+                {
+                    "telescope": telescope,
+                    "model_version": self.args_dict.get("model_version"),
+                    "per_mirror_results": per_mirror_results_out,
+                },
+                indent=2,
+            ),
+            encoding="utf-8",
+        )
+        self._logger.info("Results written to %s", output_file)
+
+        # Export averaged RNDA
+        if telescope and parameter_version and self.rnda_opt is not None:
+            try:
+                model_data_writer.ModelDataWriter.dump_model_parameter(
+                    parameter_name=parameter_name,
+                    value=[float(f"{v:.4f}") for v in self.rnda_opt],
+                    instrument=str(telescope),
+                    parameter_version=str(parameter_version),
+                    output_file=f"{parameter_name}-{parameter_version}.json",
+                    output_path=parameter_output_path,
+                    unit=["deg", "dimensionless", "deg"],
+                )
+                self._logger.info(
+                    "Exported model parameter %s (%s) to %s",
+                    parameter_name,
+                    parameter_version,
+                    parameter_output_path,
+                )
+            except (OSError, ValueError, TypeError) as e:
+                self._logger.warning("Failed to export model parameter %s: %s", parameter_name, e)
+
+    def write_psf_histogram(self):
+        """Plot histogram of measured vs simulated psf values."""
+        measured = [r.measured_psf_mm for r in self.per_mirror_results]
+        simulated = [r.simulated_psf_mm for r in self.per_mirror_results]
+        return plot_psf.plot_psf_histogram(measured, simulated, self.args_dict)