conversions.py

from astropy.time import Time
import astropy.units
from datetime import datetime, timedelta, time as dtime
from astropy.coordinates.earth import EarthLocation
from astropy.coordinates import Angle, get_body
import astropy.time
from functools import lru_cache
from lofar.sas.tmss.tmss.tmssapp.models.calculations import StationTimeline
from lofar.sas.tmss.tmss.tmssapp.models.specification import CommonSchemaTemplate
from django.db.utils import IntegrityError

from django.core.exceptions import ObjectDoesNotExist, MultipleObjectsReturned

import logging
logger = logging.getLogger(__name__)


def create_astroplan_observer_for_station(station: str) -> 'Observer':
    '''
    returns an astroplan observer for object for a given station, located in the LBA center of the given station
    :param station: a station name, e.g. "CS002"
    :return: astroplan.observer.Observer object
    '''
    from astroplan.observer import Observer # import here in the function to prevent this module from doing an astroplan setup-and-check-which-takes-to-long-at-startup
    from lofar.lta.sip import station_coordinates
    coords = station_coordinates.parse_station_coordinates()["%s_LBA" % station.upper()]
    location = EarthLocation.from_geocentric(x=coords['x'], y=coords['y'], z=coords['z'],  unit=astropy.units.m)
    observer = Observer(location, name="LOFAR", timezone="UTC")
    return observer


# default angle to the horizon at which the sunset/sunrise starts and ends, as per LOFAR definition.
SUN_SET_RISE_ANGLE_TO_HORIZON = Angle(10, unit=astropy.units.deg)
# default n_grid_points; higher is more precise but very costly; astropy defaults to 150, errors now can be in the minutes, increase if this is not good enough
# TODO: To be considered, now we store the sunset/sunrise data in advanced, we can increase the number of points!!
SUN_SET_RISE_PRECISION = 30


def timestamps_and_stations_to_sun_rise_and_set(timestamps: tuple, stations: tuple, angle_to_horizon: Angle=SUN_SET_RISE_ANGLE_TO_HORIZON,
                                                create_when_not_found=False) -> dict:
    """
    Retrieve for given stations and given timestamps the sunrise/sunset/day/night data as dictionary
    If station/timestamp is already calculated it will be retrieved from database otherwise it will be calculated
    and added to the database for possible future retrieval (optional parameter must be true).
    Storing the pre-calculated data into a database makes retrieval faster.

    The day/sunrise/sunset is always on the date of the timestamp.
    The night is usually the one _starting_ on the date of the time stamp, unless the given timestamp falls before
    sunrise, in which case it is the night _ending_ on the timestamp date.

    :param timestamps: tuple of datetimes, e.g. datetime(2020, 1, 1)
    :param stations: tuple of station names, e.g. ("CS002")
    :param angle_to_horizon: the angle between horizon and given coordinates for which rise and set times are returned
    :param: create_when_not_found: Add data to database if not found in database and so calculated for first time
    :return A dict that maps station names to a nested dict that contains lists of start and end times for sunrise,
            sunset, day and night, on each requested date.
        E.g.
        {"CS002":
            {   "sunrise": [{"start": datetime(2020, 1, 1, 6, 0, 0)), "end": datetime(2020, 1, 1, 6, 30, 0)},
                            {"start": datetime(2020, 1, 2, 6, 0, 0)), "end": datetime(2020, 1, 2, 6, 30, 0)}],
                "sunset": [{"start": datetime(2020, 1, 1, 18, 0, 0)), "end": datetime(2020, 1, 1, 18, 30, 0)},
                           {"start": datetime(2020, 1, 2, 18, 0, 0)), "end": datetime(2020, 1, 2, 18, 30, 0)}],
                "day": [{"start": datetime(2020, 1, 1, 6, 30, 0)), "end": datetime(2020, 1, 1, 18, 00, 0)},
                        {"start": datetime(2020, 1, 2, 6, 30, 0)), "end": datetime(2020, 1, 2, 18, 00, 0)}],
                "night": [{"start": datetime(2020, 1, 1, 18, 30, 0)), "end": datetime(2020, 1, 2, 6, 0, 0)},
                          {"start": datetime(2020, 1, 2, 18,3 0, 0)), "end": datetime(2020, 1, 3, 6, 0, 0)}],
            }
        }
    """
    return_dict = {}
    for station in stations:
        observer = create_astroplan_observer_for_station(station)
        for timestamp in timestamps:
            # We can also check if ALL stations/timestamps are in DB once. Do it now in a loop for each
            # station/timestamp, because we might missing something
            station_timestamp_found = False
            try:
                obj = StationTimeline.objects.get(station_name=station, timestamp=datetime.date(timestamp))
                station_timestamp_found = True
            except ObjectDoesNotExist:
                station_timestamp_found = False

            if station_timestamp_found:
                logger.debug("StationTimeline data found in DB for station=%s, timestamp=%s" % (station,timestamp))
                sunrise_dict = {"start": obj.sunrise_start, "end": obj.sunrise_end}
                sunset_dict = {"start": obj.sunset_start, "end": obj.sunset_end}
            else:
                # Not found in database so calculate it
                try:
                    sunrise_dict, sunset_dict = calculate_and_get_sunrise_and_sunset_of_observer_day(observer, timestamp, angle_to_horizon)
                except Exception as exp:
                    logger.warning("Can not calculate sunrise/sunset for station=%s, timestamp=%s" % (station,timestamp))
                    # raise exp
                    # Don't let it crash for now
                    # The stations SE607 and LV614 station has problems calculation on 2021-07-01....
                    # The SE607 also on 2021-06-04 ??
                    break
                # Add to database
                if create_when_not_found:
                    try:
                        station_timeline = StationTimeline.objects.create(
                                                    station_name=station,
                                                    timestamp=timestamp,
                                                    sunrise_start=sunrise_dict['start'],
                                                    sunrise_end=sunrise_dict['end'],
                                                    sunset_start=sunset_dict['start'],
                                                    sunset_end=sunset_dict['end'])
                        logger.debug("StationTimeline %s calculated and created for station=%s, timestamp=%s" %
                                    (station_timeline, station, timestamp))
                    except IntegrityError as e:
                        if 'unique_station_time_line' in str(e):
                            logger.info("StationTimeline with station=%s and timestamp=%s already exists, "
                                        "so not added to database",  station, timestamp)
                        else:
                            raise

            # Derive day/night from sunset/sunrise
            day_dict = {"start": sunrise_dict["end"], "end": sunset_dict["start"]}

            if timestamp >= sunrise_dict["start"]:
                # Determine next sunrise start
                try:
                    obj_next = StationTimeline.objects.get(station_name=station,
                                                           timestamp=datetime.date(timestamp + timedelta(days=1)))
                    sunrise_next_start = obj_next.sunrise_start
                except:
                    sunrise_next_start = observer.sun_rise_time(time=Time(sunrise_dict["end"]), horizon=-angle_to_horizon,
                                                                which='next',
                                                                n_grid_points=SUN_SET_RISE_PRECISION).to_datetime()
                night_dict = {"start": sunset_dict["end"], "end": sunrise_next_start}
            else:
                # Determine previous sunset end
                try:
                    obj_prev = StationTimeline.objects.get(station_name=station,
                                                           timestamp=datetime.date(timestamp - timedelta(days=1)))
                    sunset_previous_end = obj_prev.sunrise_start
                except:
                    sunset_previous_end = observer.sun_set_time(time=Time(sunrise_dict["start"]), horizon=-angle_to_horizon,
                                                                which='previous',
                                                                n_grid_points=SUN_SET_RISE_PRECISION).to_datetime()
                night_dict = {"start": sunset_previous_end, "end": sunrise_dict["start"]}

            # Create overall result
            return_dict.setdefault(station, {})
            return_dict[station].setdefault("sunrise", []).append(sunrise_dict)
            return_dict[station].setdefault("sunset", []).append(sunset_dict)
            return_dict[station].setdefault("day", []).append(day_dict)
            return_dict[station].setdefault("night", []).append(night_dict)

    return return_dict


@lru_cache(maxsize=256, typed=False)
def calculate_and_get_sunrise_and_sunset_of_observer_day(observer, timestamp: datetime, angle_to_horizon: Angle) -> dict:
    """
    Compute sunrise, sunset of the given observer object (station) at the given timestamp.
    :param observer: observer object
    :param timestamp: Datetime of a day (datetime(2020, 1, 1)
    :param the angle between horizon and given coordinates for which rise and set times are returned
    :return: dictionaries (with 'start' and 'end' defined) of sunrise, sunset
    """
    sunrise_start = observer.sun_rise_time(time=Time(datetime.combine(timestamp.date(), dtime(12, 0, 0))),
                                           horizon=-angle_to_horizon, which='previous',
                                           n_grid_points=SUN_SET_RISE_PRECISION)
    sunrise_end = observer.sun_rise_time(time=Time(sunrise_start), horizon=angle_to_horizon, which='next',
                                         n_grid_points=SUN_SET_RISE_PRECISION)
    sunset_start = observer.sun_set_time(time=sunrise_end, horizon=angle_to_horizon, which='next',
                                         n_grid_points=SUN_SET_RISE_PRECISION)
    sunset_end = observer.sun_set_time(time=sunset_start, horizon=-angle_to_horizon, which='next',
                                       n_grid_points=SUN_SET_RISE_PRECISION)

    sunrise_dict = {"start": sunrise_start.to_datetime(), "end": sunrise_end.to_datetime()}
    sunset_dict = {"start": sunset_start.to_datetime(), "end": sunset_end.to_datetime()}

    return sunrise_dict, sunset_dict


# todo: Depending on usage patterns, we should consider refactoring this a little so that we cache on a function with a single timestamp as input. Requests with similar (but not identical) timestamps or bodies currently make no use of cached results for the subset computed in previous requests.
@lru_cache(maxsize=256, typed=False)  # does not like lists, so use tuples to allow caching
def coordinates_and_timestamps_to_separation_from_bodies(angle1: float, angle2: float, direction_type: str, timestamps: tuple, bodies: tuple) -> dict:
    """
    compute angular distances of the given sky coordinates from the given solar system bodies at the given timestamps (seen from LOFAR core)
    :param angle1: first angle of celectial coordinates, e.g. RA
    :param angle2: second angle of celectial coordinates, e.g. Dec
    :param direction_type: direction_type of celectial coordinates, e.g. 'J2000'
    :param timestamps: tuple of datetimes, e.g. (datetime(2020, 1, 1, 15, 0, 0), datetime(2020, 1, 1, 16, 0, 0))
    :param bodies: tuple of solar system bodies, e.g. ('sun', 'moon', 'jupiter')
    :return A dict that maps each body to a dict that maps the given timestamp to a separation angle from the given coordinate.
        E.g.
        {
           "sun": {datetime(2020, 1, 1, 6, 0, 0): Angle("0.7rad"), datetime(2020, 1, 1, 7, 0, 0): Angle("0.7rad")},
           "moon": {datetime(2020, 1, 1, 6, 0, 0): Angle("0.4rad"), datetime(2020, 1, 1, 7, 0, 0): Angle("0.4rad")},
           "jupiter": {datetime(2020, 1, 1, 6, 0, 0): Angle("2.7rad"), datetime(2020, 1, 1, 7, 0, 0): Angle("2.7rad")}
        }
    """
    if direction_type == "J2000":
        coord = astropy.coordinates.SkyCoord(ra=angle1, dec=angle2, unit=astropy.units.rad)
    else:
        raise ValueError("Do not know how to convert direction_type=%s to SkyCoord" % direction_type)
    return_dict = {}
    for body in bodies:
        location = create_astroplan_observer_for_station("CS002").location
        for timestamp in timestamps:
            # get body coords at timestamp
            body_coord = get_body(body=body, time=astropy.time.Time(timestamp), location=location)
            angle = coord.separation(body_coord)
            return_dict.setdefault(body, {})[timestamp] = angle
    return return_dict


# default angle above horizon, above which the target it reporte as 'up'
TARGET_SET_RISE_ANGLE_TO_HORIZON = Angle(0, unit=astropy.units.deg)  # if default should be non-zero, should we include it explicitly in response?
# default n_grid_points; higher is more precise but very costly; astropy defaults to 150, note that errors can be in the minutes with a lower values
TARGET_SET_RISE_PRECISION = 150

@lru_cache(maxsize=256, typed=False)  # does not like lists, so use tuples to allow caching
def coordinates_timestamps_and_stations_to_target_rise_and_set(angle1: float, angle2: float, direction_type: str, timestamps: tuple, stations: tuple, angle_to_horizon: Angle=TARGET_SET_RISE_ANGLE_TO_HORIZON) -> dict:
    """
    Compute rise and set times of the given coordinates above the provided horizon, for each given station and timestamp.
    The set time is always the one following the provided timestamp.
    This implies that if the target is up at a given timestamp, the surrounding rise and set times are returned.
    Otherwise both rise and set times follow the timestamp.
    :param angle1: first angle of celectial coordinates, e.g. RA
    :param angle2: second angle of celectial coordinates, e.g. Dec
    :param direction_type: direction_type of celectial coordinates, e.g. 'J2000'
    :param timestamps: tuple of datetimes, e.g. (datetime(2020, 1, 1), datetime(2020, 1, 2))
    :param stations: tuple of station names, e.g. ("CS002",)
    :param angle_to_horizon: the angle between horizon and given coordinates for which rise and set times are returned
    :return A dict that maps station names to a list of dicts with rise and set times for each requested date.
            If rise and set are None, the target is always above or below horizon, and the respective boolean is True.
        E.g.
        {"CS002": [{"rise": datetime(2020, 1, 1, 4, 0, 0), "set": datetime(2020, 1, 1, 11, 0, 0), "always_above_horizon": False, "always_below_horizon": False},
                   {"rise": datetime(2020, 1, 2, 4, 0, 0), "set": datetime(2020, 1, 2, 11, 0, 0), "always_above_horizon": False, "always_below_horizon": False}]
        }
    """
    if direction_type == "J2000":
        coord = astropy.coordinates.SkyCoord(ra=angle1, dec=angle2, unit=astropy.units.rad)
    else:
        raise ValueError("Do not know how to convert direction_type=%s to SkyCoord" % direction_type)
    return_dict = {}
    for station in stations:
        for timestamp in timestamps:
            # todo: this can probably be made faster by moving the following logic to an own function with single station/timestamp as input and putting the lru_cache on there.
            observer = create_astroplan_observer_for_station(station)
            try:
                target_set = observer.target_set_time(target=coord, time=Time(timestamp), horizon=angle_to_horizon, which='next', n_grid_points=TARGET_SET_RISE_PRECISION)
                target_rise = observer.target_rise_time(target=coord, time=Time(target_set), horizon=angle_to_horizon, which='previous', n_grid_points=TARGET_SET_RISE_PRECISION)
                return_dict.setdefault(station, []).append(
                    {"rise": target_rise.to_datetime(),
                     "set": target_set.to_datetime(),
                     "always_above_horizon": False,
                     "always_below_horizon": False})
            except TypeError as e:
                if "numpy.float64" in str(e):
                    # Note: when the target is always above or below horizon, astroplan excepts with the not very
                    #  meaningful error: 'numpy.float64' object does not support item assignment
                    # Determine whether the target is always above or below horizon so that we can return some useful
                    # additional info, e.g. for scheduling purposes.
                    is_up = observer.target_is_up(target=coord, time=Time(timestamp), horizon=angle_to_horizon)
                    return_dict.setdefault(station, []).append(
                        {"rise": None,
                         "set": None,
                         "always_above_horizon": is_up,
                         "always_below_horizon": not is_up})
                else:
                    raise


    return return_dict



def local_sidereal_time_for_utc_and_station(timestamp: datetime = None,
                                            station: str = 'CS002',
                                            field: str = 'LBA',
                                            kind: str = "apparent"):
    """
    calculate local sidereal time for given utc time and station
    :param timestamp: timestamp as datetime object
    :param station: station name
    :param field: antennafield, 'LBA' or 'HBA'
    :param kind: 'mean' or 'apparent'
    :return:
    """
    from lofar.lta.sip import station_coordinates

    if timestamp is None:
        timestamp = datetime.utcnow()
    station_coords = station_coordinates.parse_station_coordinates()
    field_coords = station_coords["%s_%s" % (station, field)]
    location = EarthLocation.from_geocentric(x=field_coords['x'], y=field_coords['y'], z=field_coords['z'], unit=astropy.units.m)
    return local_sidereal_time_for_utc_and_longitude(timestamp=timestamp, longitude=location.lon.to_string(decimal=True), kind=kind)


def local_sidereal_time_for_utc_and_longitude(timestamp: datetime = None,
                                              longitude: float = 6.8693028,
                                              kind: str = "apparent"):
    """
    :param timestamp: timestamp as datetime object
    :param longitude: decimal longitude of observer location (defaults to CS002 LBA center)
    :param kind: 'mean' or 'apparent'
    :return:
    """
    if timestamp is None:
        timestamp = datetime.utcnow()
    t = Time(timestamp, format='datetime', scale='utc')
    return t.sidereal_time(kind=kind, longitude=longitude)


def antennafields_for_antennaset_and_station(antennaset:str, station:str) -> list:
    """
    convert an antennaset to a list of antennafields
    :param antennaset: A string identifier for an antennaset, like 'HBA_DUAL'
    :param station: A string identifier for a station, like 'CS001'
    :return: a list of antennafields that the station uses for the given antennaset ['HBA0', 'HBA1']
    """
    if antennaset.startswith('LBA'):
        fields = ['LBA']
    elif antennaset.startswith('HBA') and not station.startswith('CS'):
        fields = ['HBA']
    elif antennaset.startswith('HBA_DUAL'):
        fields = ['HBA0', 'HBA1']
    elif antennaset.startswith('HBA_ZERO'):
        fields = ['HBA0']
    elif antennaset.startswith('HBA_ONE'):
        fields = ['HBA1']
    else:
        raise ValueError('Cannot determine antennafields for station=%s antennaset=%s' % (station, antennaset))

    return fields


def get_all_stations():
    """
    returns all possible stations.
    Retrieve station names from station template by getting the Dutch and International stations,
    then you should have it all.
    """
    lst_stations = []
    for station_group in ["Dutch", "International"]:
        try:
            station_schema_template = CommonSchemaTemplate.objects.get(name="stations", version=1)
            groups = station_schema_template.schema['definitions']['station_group']['anyOf']
            selected_group = next(g for g in groups if g['title'].lower() == station_group.lower())
            lst_stations.extend(selected_group['properties']['stations']['enum'][0])
        except Exception:
            logger.warning("No stations schema found, sorry can not determine station list, return empty list")
    return lst_stations