Source code for ska_ost_senscalc.common.model

import dataclasses
from dataclasses import dataclass, field
from enum import Enum, IntEnum
from typing import List, Optional, TypedDict, Union

import astropy.units as u
import numpy as np
from astropy.coordinates import Latitude
from astropy.coordinates.sky_coordinate import SkyCoord
from astropy.units import Quantity

from ska_ost_senscalc.common.spectropolarimetry import SpectropolarimetryResults
from ska_ost_senscalc.low.model import CalculatorInput
from ska_ost_senscalc.mid.calculator import DEFAULT_ALPHA, DEFAULT_EL, DEFAULT_PWV
from ska_ost_senscalc.subarray import LOWArrayConfiguration, MIDArrayConfiguration
from ska_ost_senscalc.utilities import Telescope




@dataclass
class ConfusionNoiseResponse:  # TODO: merge it with ConfusionNoise model later
    """
    ConfusionNoiseResponse is a typed dictionary constrained to match the
    schema of a confusion noise JSON object, as contained in the parent JSON
    result of a weighting endpoint query.
    """

    value: float | int
    limit_type: str





@dataclass
class BeamSizeResponse:
    """
    BeamSizeResponse is a typed dictionary constrained to match the schema of
    the JSON object outlining the synthesized beam size, as contained in the
    parent JSON result of a weighting endpoint query.
    """

    beam_maj_scaled: float
    beam_min_scaled: float
    beam_pa: float





@dataclass
class SingleWeightingResponse:
    """
    SingleWeightingResponse is a typed dictionary constrained to match the
    schema of a single weighting calculation, as performed for the main
    continuum or zoom weighting calculation and for each subband frequency.
    """

    weighting_factor: float
    sbs_conv_factor: float
    confusion_noise: ConfusionNoiseResponse
    beam_size: BeamSizeResponse





@dataclass
class SubbandWeightingResponse(SingleWeightingResponse):
    subband_freq_centre: u.Quantity





@dataclass
class ContinuumWeightingResponse(SingleWeightingResponse):
    """
    A typed dictionary constrained to match the schema of
    the weighting endpoint query for continuum.
    """

    subbands: list[SubbandWeightingResponse] = field(default_factory=list)





@dataclass
class SingleZoomWeightingResponse(SingleWeightingResponse):
    """
    A typed dictionary constrained to match the schema of
    the weighting endpoint query for each line.
    """

    freq_centre: u.Quantity





class SubbandResponse(TypedDict):
    subband_freq_centre: Quantity
    sensitivity: Quantity





@dataclass
class WeightingContinuumSpectralResults:
    """
    WeightingContinuumSpectralResults is a data class containing a continuum weighting calculation
    and a spectral weighting calculation.
    """

    continuum_weighting: ContinuumWeightingResponse
    spectral_weighting: SingleZoomWeightingResponse





@dataclass
class SensitivityTransformationsSubbandsInput:
    """
    Data class representing the input for the subbands transformations for continuum sensitivities
    (supplied integration time) and integration times (supplied sensitivities)
    """

    calculate_subbands: list[SubbandResponse]
    weighting_subbands: list[SubbandWeightingResponse] = field(default_factory=list)

    # used for checking sensitivity limit for supplied sensitivity subbands
    continuum_thermal_sensitivity: Quantity = None





class ContinuumSensitivityResults(TypedDict):
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    results of a single continuum sensitivity calculation.
    """

    continuum_sensitivity: Quantity
    continuum_subband_sensitivities: Optional[list[SubbandResponse]]
    spectral_sensitivity: Quantity

    spectropolarimetry_results: SpectropolarimetryResults

    warnings: list[str]





@dataclass
class SingleZoomSensitivityResponse:
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    response body of a single zoom sensitivity calculation.
    """

    freq_centre: Quantity
    spectral_sensitivity: Quantity
    spectropolarimetry_results: SpectropolarimetryResults
    warnings: list[str] = field(default_factory=list)





class PSSSensitivityResponse(TypedDict, total=False):
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    response body of a single PSS sensitivity calculation.
    """

    folded_pulse_sensitivity: Quantity
    single_pulse_sensitivity: Quantity
    warning: Optional[str]





class Weighting(Enum):
    """
    Enumeration for different weighting
    """

    NATURAL = "natural"
    ROBUST = "robust"
    UNIFORM = "uniform"





class WeightingSpectralMode(Enum):
    """
    Enumeration spectral modes supported by the calculator, which are used in the look up table.
    """

    # TODO: Should it be renamed to Zoom?
    LINE = "line"
    CONTINUUM = "continuum"





@dataclass
class ContinuumCalculatorAndWeightingInput:
    """
    This dataclass represents the internal model of the Calculator and Weighting for continuum and
    zoom modes.
    """

    _: dataclasses.KW_ONLY
    freq_centre: Quantity
    bandwidth_mhz: float
    num_stations: int
    pointing_centre: str
    integration_time_h: float
    elevation_limit: float
    spectral_mode: WeightingSpectralMode
    telescope: Telescope
    spectral_averaging_factor: float = 1
    n_subbands: int
    weighting_mode: Weighting
    subarray_configuration: LOWArrayConfiguration | MIDArrayConfiguration | None
    taper: Quantity = 0.0 * u.arcsec
    robustness: int
    subband_freq_centres: Optional[list[Quantity]] = None





@dataclass
class ZoomCalculatorAndWeightingInput:
    """
    This dataclass represents the internal model of the Calculator and Weighting for zoom modes.
    """

    _: dataclasses.KW_ONLY
    dec: Latitude
    num_stations: int
    integration_time_h: int
    elevation_limit: float
    spectral_averaging_factor: float = 1
    spectral_resolutions_hz: list[float]
    total_bandwidths_khz: list[float]
    freq_centres: list[float]
    pointing_centre: str
    subarray_configuration: LOWArrayConfiguration | MIDArrayConfiguration | None
    weighting_mode: Weighting
    robustness: int = 0
    taper: Quantity = 0.0 * u.arcsec
    telescope: Telescope





@dataclass
class ContinuumSensitivitiesTransformationsInput:
    """
    This dataclass represents the input for the transformations for continuum sensitivities
    (supplied integration time)
    """

    continuum_sensitivity: Quantity
    spectral_sensitivity: Quantity
    continuum_weighting_factor: np.float64
    spectral_weighting_factor: np.float64
    continuum_sbs_conv_factor: np.float64
    spectral_sbs_conv_factor: np.float64
    continuum_conf_noise: np.float64
    spectral_conf_noise: np.float64
    continuum_beam_min: np.float64
    continuum_beam_maj: np.float64
    spectral_beam_min: np.float64
    spectral_beam_maj: np.float64





@dataclass
class SensitivitiesTransformationsWithoutWeightingInput:
    """
    This dataclass represents the input for the transformations for continuum sensitivities
    (supplied integration time) with a custom subarray config (custom sub-array).
    No weighting is calculated for custom sub-array.
    """

    spectral_sensitivity: Quantity
    continuum_sensitivity: Quantity = None
    weighting_factor: int = (
        1  # there is no weighting factor, so we always use 1 for the calculations
    )





@dataclass
class ZoomSensitivitiesTransformationsInput:
    """
    This dataclass represents the input for the transformations for zoom sensitivities for low
    """

    spectral_sensitivity: Quantity
    weighting_factor: np.float64
    min_beam_size: np.float64
    maj_beam_size: np.float64
    sbs_conv_factor: np.float64
    confusion_noise: np.float64





@dataclass
class ContinuumWeightingRequestParams:
    """
    Represents the parameters sent to the continuum weighting calculation, after they have been deserialised into enums and astropy objects.
    """

    spectral_mode: WeightingSpectralMode
    telescope: Telescope
    weighting_mode: Weighting
    subarray_configuration: LOWArrayConfiguration | MIDArrayConfiguration
    dec: Latitude
    freq_centre: Quantity
    taper: Quantity = 0.0 * u.arcsec
    robustness: int = 0
    subband_freq_centres: Optional[list[Quantity]] = None





class SpectralAveragingFactor(IntEnum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4
    SIX = 6
    EIGHT = 8
    TWELVE = 12
    TWENTY_FOUR = 24





class Robustness(IntEnum):
    MINUS_TWO = -2
    MINUS_ONE = -1
    ZERO = 0
    ONE = 1
    TWO = 2





class PulsarMode(Enum):
    """
    Enumeration for the different pulsar modes
    """

    SINGLE_PULSE = "single_pulse"
    FOLDED_PULSE = "folded_pulse"





@dataclass
class ZoomRequest:
    """Mandatory and optional parameters for zoom/calculate request"""

    rx_band: str
    freq_centres_hz: list[int]
    pointing_centre: str | SkyCoord
    spectral_resolutions_hz: list[float]
    total_bandwidths_hz: list[float]

    subarray_configuration: str | MIDArrayConfiguration = None
    n_ska: int = None
    n_meer: int = None
    pwv: float = DEFAULT_PWV
    el: float = DEFAULT_EL.value
    spectral_averaging_factor: SpectralAveragingFactor = 1
    supplied_sensitivities: list[float] = None
    sensitivity_unit: u.Unit = u.Jy / u.beam
    integration_time_s: int = None
    eta_system: float = None
    eta_pointing: float = None
    eta_coherence: float = None
    eta_digitisation: float = None
    eta_correlation: float = None
    eta_bandpass: float = None
    t_sys_ska: float = None
    t_rx_ska: float = None
    t_spl_ska: float = None
    t_sys_meer: float = None
    t_rx_meer: float = None
    t_spl_meer: float = None
    t_sky_ska: float = None
    t_sky_meer: float = None
    t_gal_ska: float = None
    t_gal_meer: float = None
    alpha: float = DEFAULT_ALPHA
    eta_meer: float = None
    eta_ska: float = None
    weighting_mode: Weighting = Weighting.UNIFORM
    robustness: Robustness = Robustness.ZERO
    taper: float | Quantity = 0.0





@dataclass
class ZoomRequestPrepared(ZoomRequest):
    telescope: str = "n/a"  # will be changed during the validation
    target: SkyCoord = None  # will be changed during the calculations
    freq_centre_hz: float = None  # will be changed during the calculations
    freq_centres: list[Quantity] = field(
        default_factory=list
    )  # will be changed during the calculations
    bandwidth_hz: float = None  # will be changed during the calculations





@dataclass
class PssRequestMid:
    """Mandatory and optional parameters for Mid pss/calculate request"""

    rx_band: str
    freq_centre_hz: int
    pointing_centre: str | SkyCoord

    pulsar_mode: str
    bandwidth_hz: float = 300e6
    spectral_resolution_hz: float = 107.5e3
    pulse_period: float = 33.0  # Assume crab as default (33 ms)
    intrinsic_pulse_width: float = 0.004  # Assume crab as default (4 us)

    dm: float = 0.0
    weighting_mode: Weighting = None
    subarray_configuration: str | MIDArrayConfiguration = None
    pwv: float = DEFAULT_PWV
    el: float = DEFAULT_EL.value
    integration_time_s: int = None
    eta_system: float = None
    eta_pointing: float = None
    eta_coherence: float = None
    eta_digitisation: float = None
    eta_correlation: float = None
    eta_bandpass: float = None
    t_sys_ska: float = None
    t_rx_ska: float = None
    t_spl_ska: float = None
    t_sys_meer: float = None
    t_rx_meer: float = None
    t_spl_meer: float = None
    t_sky_ska: float = None
    t_sky_meer: float = None
    t_gal_ska: float = None
    t_gal_meer: float = None
    alpha: float = DEFAULT_ALPHA
    eta_meer: float = None
    eta_ska: float = None





@dataclass
class ContinuumCalculateRequestParams:
    """
    Represents the parameters sent to the continuum calculation, after they have been deserialised into enums and astropy objects.
    """

    # keeping both since we haven't completely combine weighting/calculate
    freq_centre_hz: float
    rx_band: str
    bandwidth_hz: float
    pointing_centre: str | SkyCoord
    weighting_mode: Weighting

    subarray_configuration: LOWArrayConfiguration | MIDArrayConfiguration = None

    robustness: Robustness = Robustness.ZERO
    n_subbands: int = 1
    subband_freq_centres: list[Quantity] = field(default_factory=list)
    spectral_averaging_factor: float = 1
    integration_time_s: float = None
    supplied_sensitivity: float = None
    sensitivity_unit: u.Unit = u.Jy / u.beam
    n_ska: int = None
    n_meer: int = None
    pwv: float = DEFAULT_PWV
    el: float = DEFAULT_EL.value
    eta_system: float = None
    eta_pointing: float = None
    eta_coherence: float = None
    eta_digitisation: float = None
    eta_correlation: float = None
    eta_bandpass: float = None
    t_sys_ska: float = None
    t_rx_ska: float = None
    t_spl_ska: float = None
    t_sys_meer: float = None
    t_rx_meer: float = None
    t_spl_meer: float = None
    t_sky_ska: float = None
    t_sky_meer: float = None
    t_gal_ska: float = None
    t_gal_meer: float = None
    alpha: float = DEFAULT_ALPHA
    eta_meer: float = None
    eta_ska: float = None
    weighting_mode: Weighting = Weighting.UNIFORM
    robustness: Robustness = Robustness.ZERO
    taper: float | Quantity = 0.0

    freq_centre: Quantity = None
    telescope: Telescope = None





@dataclass
class ContinuumRequest(ContinuumCalculateRequestParams):
    subband_freq_centres_hz: list[float] = field(default_factory=list)
    subband_supplied_sensitivities: list[float] = field(default_factory=list)
    subband_supplied_sensitivities_unit: u.Unit = u.Jy / u.beam
    target: SkyCoord = None  # will be changed during the calculations





class Limit(Enum):
    """
    Enumeration for different types of limit
    """

    UPPER = "upper limit"
    LOWER = "lower limit"
    VALUE = "value"





@dataclasses.dataclass
class BeamSize:
    beam_maj: u.Quantity
    beam_min: u.Quantity
    beam_pa: u.Quantity





@dataclasses.dataclass
class ConfusionNoise:
    value: u.Quantity
    limit: Limit



@dataclasses.dataclass(kw_only=True)
class _WeightingInput:
    dec: Latitude
    weighting_mode: Weighting
    subarray_configuration: Union[MIDArrayConfiguration, LOWArrayConfiguration]
    telescope: Telescope
    spectral_mode: WeightingSpectralMode
    robustness: int = 0
    taper: u.Quantity




@dataclasses.dataclass(kw_only=True)
class WeightingInput(_WeightingInput):
    freq_centre: u.Quantity





@dataclasses.dataclass(kw_only=True)
class WeightingMultiInput(_WeightingInput):
    freq_centres: list[u.Quantity]





@dataclasses.dataclass
class WeightingResult:
    weighting_factor: float
    surface_brightness_conversion_factor: u.Quantity
    beam_size: BeamSize
    confusion_noise: ConfusionNoise





@dataclasses.dataclass
class WeightingMultiResultElement(WeightingResult):
    freq_centre: u.Quantity





class PulsarSamplingTime(Enum):
    """
    Enumeration for the different PSS and PST sampling times
    """

    LOW_PSS = 69.12e-6 * u.second  # 69.12 us
    LOW_PST = 207.36e-6 * u.second  # 207.36 us
    MID_PSS = 65e-6 * u.second  # 65e-6
    MID_PST = 18.6e-6 * u.second  # 18.6 us





@dataclasses.dataclass
class CalculatorInputPSS(CalculatorInput):
    """
    This dataclass represents the internal model of the Calculator for the PSS mode.

    The following units are implicitly assumed (first three parameters are inherited
    from CalculatorInput):
    - freq_centre, bandwidth: [MHz]
    - duration: [h]
    - chan_width: [Hz]
    - dm: [pc/cm^3]
    - pulse_period, intrinsic_pulse_width: [ms]
    """

    _: dataclasses.KW_ONLY
    chan_width: float
    dm: float
    pulse_period: float
    intrinsic_pulse_width: float





@dataclass
class BeamSizeResult:
    """
    The beam size result returned
    as part of the transformed results of a sensitivity calculation
    """

    beam_maj: Quantity
    beam_min: Quantity





@dataclass
class SubbandsItemsTransformedResults:
    """
    Data class representing the transformed results items for the subbands of a single continuum sensitivity calculation.
    """

    min_value: Quantity
    max_value: Quantity





@dataclass
class SubbandsBeamSizeTransformedResults:
    """
    Data class representing the transformed results items for the subbands of a single continuum sensitivity calculation.
    """

    min_value: BeamSizeResult
    max_value: BeamSizeResult





@dataclass
class SubbandsContinuumSensitivityTransformedResults:
    """
    Data class representing the transformed results for the subbands of a single continuum sensitivity calculation.
    """

    weighted_sensitivity_per_subband: SubbandsItemsTransformedResults
    confusion_noise_per_subband: SubbandsItemsTransformedResults = None
    total_sensitivity_per_subband: SubbandsItemsTransformedResults = None
    synthesized_beam_size_per_subband: SubbandsBeamSizeTransformedResults = None
    surface_brightness_sensitivity_per_subband: SubbandsItemsTransformedResults = None
    warnings: list[str] = field(default_factory=list)





@dataclass
class ContinuumSensitivityTransformedResults:
    """
    Data class representing the transformed results of a single continuum sensitivity calculation.
    """

    weighted_continuum_sensitivity: Quantity = None
    continuum_confusion_noise: Quantity = None
    total_continuum_sensitivity: Quantity = None
    continuum_synthesized_beam_size: BeamSizeResult = None
    continuum_surface_brightness_sensitivity: Quantity = None
    weighted_spectral_sensitivity: Quantity = None
    spectral_confusion_noise: Quantity = None
    total_spectral_sensitivity: Quantity = None
    spectral_synthesized_beam_size: BeamSizeResult = None
    spectral_surface_brightness_sensitivity: Quantity = None
    warnings: list[str] = field(default_factory=list)





@dataclass
class IntegrationTimeTransformationsWithoutWeightingInput:
    """
    This dataclass represents the input for the transformations for continuum integration
    time (supplied sensitivity) with a custom subarray config (custom sub-array).
    """

    spectral_integration_time: Quantity
    continuum_integration_time: Quantity = None





@dataclass
class ContinuumIntegrationTimeTransformedResults:
    """
    Data class representing the transformed results of a single continuum integration time calculation.
    """

    continuum_confusion_noise: Quantity = None
    continuum_synthesized_beam_size: BeamSizeResult = None
    continuum_integration_time: Quantity = None

    spectral_confusion_noise: Quantity = None
    spectral_synthesized_beam_size: BeamSizeResult = None
    spectral_integration_time: Quantity = None

    warnings: list[str] = field(default_factory=list)





@dataclass
class ZoomSensitivityTransformedResults:
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    transformed results of a single zoom sensitivity calculation.
    """

    weighted_spectral_sensitivity: Quantity = None
    spectral_confusion_noise: Quantity = None
    total_spectral_sensitivity: Quantity = None
    spectral_synthesized_beam_size: BeamSizeResult = None
    spectral_surface_brightness_sensitivity: Quantity = None
    warnings: List[Union[str, dict]] = field(default_factory=list)





@dataclass
class ZoomSensitivityResponse:
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    results of a zoom sensitivity calculation.
    """

    calculate: list[SingleZoomSensitivityResponse]
    weighting: list[SingleZoomWeightingResponse]
    transformed_result: list[ZoomSensitivityTransformedResults]





class ContinuumSensitivityCalculateResponse(ContinuumSensitivityResults):
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    response of a single continuum sensitivity calculation.
    Extends on the base class ContinuumSensitivityResults and append
    the transformed results and continuum and spectral weighting results.
    """

    continuum_weighting: ContinuumWeightingResponse
    spectral_weighting: ContinuumWeightingResponse
    transformed_result: ContinuumSensitivityTransformedResults





class ContinuumSensitivityResponse(TypedDict):
    """
    Typed dictionary constrained to match the OpenAPI schema for the
    results of a single continuum sensitivity calculation.
    """

    calculate: ContinuumSensitivityResults
    weighting: WeightingContinuumSpectralResults
    transformed_result: ContinuumSensitivityTransformedResults





@dataclass
class SubbandsContinuumIntegrationTimeTransformedResults:
    """
    Data class representing the transformed results for the subbands of a single continuum integration time calculation.
    """

    confusion_noise_per_subband: SubbandsItemsTransformedResults = None
    synthesized_beam_size_per_subband: SubbandsBeamSizeTransformedResults = None
    integration_time_per_subband: SubbandsItemsTransformedResults = None
    warnings: list[Union[str, dict]] = field(default_factory=list)





class EnumConversion:
    """
    Utility class to convert OpenAPI enumeration members to Python Enum
    members.

    OpenAPI enums generally map to Python enums but their name will usually
    be formatted differently, as the Python convention is for enumeration
    member names to be all upper case. This class decouples the OpenAPI naming
    convention from that all-caps requirement.
    """

    @staticmethod
    def to_weighting(val: str, msgs: list[str]) -> Weighting:
        return EnumConversion._convert(Weighting, val, msgs)

    @staticmethod
    def to_array_configuration(val: str, msgs) -> LOWArrayConfiguration:
        return EnumConversion._convert(LOWArrayConfiguration, val, msgs)

    @staticmethod
    def _convert(cls, val: str, msgs: list[str]):
        try:
            return cls[val.upper()]
        except (ValueError, KeyError):
            msg = f"{val} could not be mapped to a {cls.__name__} enum member"
            msgs.append(msg)





class MidBand(Enum):
    """
    Enumeration of the receiver bands available for MID
    """

    BAND_1 = "Band 1"
    BAND_2 = "Band 2"
    BAND_3 = "Band 3"
    BAND_4 = "Band 4"
    BAND_5A = "Band 5a"
    BAND_5B = "Band 5b"