"""Models for observing constraints on an SBDefinition."""
from typing import Annotated, Callable
from astropy import units
from pydantic import AfterValidator, Field, WithJsonSchema, model_validator
from ska_oso_pdm._shared import (
AstropyQuantity,
AstropyUnit,
PdmObject,
StableSet,
TerseStrEnum,
UnitHelpers,
)
from ska_oso_pdm._shared.custom_types import Quantity
__all__ = [
"AngularSeparationConstraint",
"AltitudeConstraint",
"LSTConstraint",
"ObservingConstraints",
]
class AngularSeparationUnits(TerseStrEnum):
ARCMINUTES = "arcmin"
ARCSECONDS = "arcsec"
DEGREES = "deg"
AngularSeparationUnitType = Annotated[
AstropyUnit,
AfterValidator(UnitHelpers.constrain_unit_to(AngularSeparationUnits)),
WithJsonSchema(UnitHelpers.enum_jsonschema(AngularSeparationUnits)),
]
class AngularSeparationQuantity(Quantity):
unit = units.Unit(AngularSeparationUnits.DEGREES)
def constrain_quantity_to_range(
lower: units.Quantity, upper: units.Quantity
) -> Callable[[units.Quantity], units.Quantity]:
if lower > upper:
raise ValueError(f"Lower bound must not exceed upper bound: {lower} > {upper}")
def _closure(value: units.Quantity) -> units.Quantity:
if value < lower or value > upper:
raise ValueError(f"Quantity must be between {lower} and {upper}")
return value
return _closure
ANGULAR_SEPARATION_MIN = units.Quantity(value=0.0, unit=AngularSeparationUnits.DEGREES)
ANGULAR_SEPARATION_MAX = units.Quantity(
value=180.0, unit=AngularSeparationUnits.DEGREES
)
AngularSeparationQuantityType = Annotated[
AstropyQuantity,
AngularSeparationQuantity,
AfterValidator(UnitHelpers.constrain_unit_to(AngularSeparationUnits)),
AfterValidator(
constrain_quantity_to_range(ANGULAR_SEPARATION_MIN, ANGULAR_SEPARATION_MAX)
),
]
class AltitudeUnits(TerseStrEnum):
DEGREES = "deg"
RADIANS = "rad"
AltitudeUnitType = Annotated[
AstropyUnit,
AfterValidator(UnitHelpers.constrain_unit_to(AltitudeUnits)),
WithJsonSchema(UnitHelpers.enum_jsonschema(AltitudeUnits)),
]
class AltitudeQuantity(Quantity):
unit = units.Unit(AltitudeUnits.DEGREES)
ALTITUDE_MIN = units.Quantity(value=0.0, unit=AltitudeUnits.DEGREES)
ALTITUDE_MAX = units.Quantity(value=90.0, unit=AltitudeUnits.DEGREES)
AltitudeQuantityType = Annotated[
AstropyQuantity,
AltitudeQuantity,
AfterValidator(UnitHelpers.constrain_unit_to(AltitudeUnits)),
AfterValidator(constrain_quantity_to_range(ALTITUDE_MIN, ALTITUDE_MAX)),
]
class LSTUnits(TerseStrEnum):
DEGREES = "deg"
RADIANS = "rad"
HOURANGLE = "hourangle"
LSTUnitType = Annotated[
AstropyUnit,
AfterValidator(UnitHelpers.constrain_unit_to(LSTUnits)),
WithJsonSchema(UnitHelpers.enum_jsonschema(LSTUnits)),
]
class LSTQuantity(Quantity):
unit = units.Unit(LSTUnits.HOURANGLE)
LST_MIN = units.Quantity(value=0.0, unit=LSTUnits.DEGREES)
LST_MAX = units.Quantity(value=360.0, unit=LSTUnits.DEGREES)
LSTQuantityType = Annotated[
AstropyQuantity,
LSTQuantity,
AfterValidator(UnitHelpers.constrain_unit_to(LSTUnits)),
AfterValidator(constrain_quantity_to_range(LST_MIN, LST_MAX)),
]
[docs]
class AngularSeparationConstraint(PdmObject):
min: AngularSeparationQuantityType | None = None
[docs]
class AltitudeConstraint(PdmObject):
min: AltitudeQuantityType | None = None
max: AltitudeQuantityType | None = None
@model_validator(mode="after")
def validate_min_less_than_max(self):
if self.min is not None and self.max is not None and self.min >= self.max:
raise ValueError("Altitude min must be less than max")
return self
[docs]
class LSTConstraint(PdmObject):
start: LSTQuantityType
end: LSTQuantityType
[docs]
class ObservingConstraints(PdmObject):
sun_separation: AngularSeparationConstraint = Field(
default_factory=AngularSeparationConstraint
)
moon_separation: AngularSeparationConstraint = Field(
default_factory=AngularSeparationConstraint
)
jupiter_separation: AngularSeparationConstraint = Field(
default_factory=AngularSeparationConstraint
)
altitude: AltitudeConstraint = Field(default_factory=AltitudeConstraint)
lst: LSTConstraint | None = None
minimum_receptors: int | None = None
required_receptors: StableSet[str] = frozenset()