Source code for ctao_datamodel.models.dataproducts.simulation

"""Simulation context classes."""

from typing import ClassVar

import astropy.units as u
from astropy.coordinates import Angle, EarthLocation, SkyCoord, angular_separation
from astropy.time import Time

from ..._core import AstroField, ModelBase, Quantity, doc

_NAMESPACE = "CTAO.DataProducts.Context"


__all__ = ["Simulation", "SimulationGridNode"]


class SimulationGridNode(ModelBase):
    """Describes the position within the simulation phase space."""

    _namespace: ClassVar = _NAMESPACE

    zenith: Quantity[u.deg] = AstroField(
        "Fixed zenith angle of the center of the simulated FOV", fits_keyword="SIMZEN"
    )
    azimuth: Quantity[u.deg] = AstroField(
        "Fixed azimuthal angle of the center of the simulated FOV", fits_keyword="SIMAZ"
    )
    hour_angle: Quantity[u.deg] | None = AstroField(
        "Fixed hour angle of the center of the simulated FOV",
        fits_keyword="SIMHA",
    )
    declination: Quantity[u.deg] | None = AstroField(
        "Fixed declination of the center of the simulated FOV",
        default=None,
        fits_keyword="SIMDEC",
    )
    nsb_factor: float = AstroField(
        "NSB rate expressed as a multiplicative factor of the standard dark night"
        " level.",
        default=None,
        fits_keyword="SIMNSB",
    )



    def to_altaz(self, **frame_attrs):
        """Return a SkyCoord from the alt/az coordinates."""
        return SkyCoord(
            alt=Angle("90 deg") - self.zenith,
            az=self.azimuth,
            frame="altaz",
            **frame_attrs,
        )




    def to_hadec(self, **frame_attrs):
        """Return a SkyCoord from the HA/Dec coordinates."""
        return SkyCoord(
            ha=self.hour_angle, dec=self.declination, frame="hadec", **frame_attrs
        )




    def validate_coordinates(self, location: EarthLocation):
        """Check coordinates for consistency using an EarthLocation."""
        altaz = self.to_altaz(location=location, obstime=Time("2020-01-01 00:00:00"))
        expected_hadec = altaz.transform_to("hadec")

        if self.declination is None:
            self.declination = u.Quantity(expected_hadec.dec).to(u.deg)
        if self.hour_angle is None:
            self.hour_angle = u.Quantity(expected_hadec.ha).to(u.deg)

        current_hadec = self.to_hadec()
        separation = angular_separation(
            lon1=current_hadec.spherical.lon,
            lat1=current_hadec.spherical.lat,
            lon2=expected_hadec.spherical.lon,
            lat2=expected_hadec.spherical.lat,
        )

        if separation.to_value(u.deg) > 0.11:
            raise ValueError(
                "In the SimulationGridNode, the given HA/Dec coord is"
                f" {separation.to_value('deg'):.2f} degrees away from given zenith/az"
                " coordinate."
            )



class Simulation(ModelBase):
    """
    Simulation context information.

    Defines additional metadata (outside of the InstanceIdentifier) needed to
    retrieve simulated runs within a simulation production or data products
    derived from them.
    """

    _namespace: ClassVar = _NAMESPACE

    production_name: str = AstroField(
        "Simulation production name", fits_keyword="SIMPROD"
    )

    node: SimulationGridNode | None = AstroField(doc(SimulationGridNode), default=None)