Skip to content

Commit

Permalink
move date stackers into their own submodule
Browse files Browse the repository at this point in the history
  • Loading branch information
ehneilsen committed Mar 29, 2024
1 parent aa7be99 commit c202434
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 133 deletions.
1 change: 1 addition & 0 deletions rubin_sim/maf/stackers/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from .base_stacker import *
from .coord_stackers import *
from .date_stackers import *
from .general_stackers import *
from .get_col_info import *
from .label_stackers import *
Expand Down
140 changes: 140 additions & 0 deletions rubin_sim/maf/stackers/date_stackers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
__all__ = (
"ObservationStartDatetime64Stacker",
"DayObsStacker",
"DayObsMJDStacker",
"DayObsISOStacker",
)

import numpy as np
from astropy.time import Time

from .base_stacker import BaseStacker


class ObservationStartDatetime64Stacker(BaseStacker):
"""Add the observation start time as a numpy.datetime64."""

cols_added = ["observationStartDatetime64"]

def __init__(
self,
mjd_col="observationStartMJD",
):
self.mjd_col = mjd_col
self.units = [None]
self.cols_added_dtypes = ["datetime64[ns]"]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data["observationStartDatetime64"] = Time(sim_data[self.mjd_col], format="mjd").datetime64

return sim_data


def _compute_day_obs_mjd(mjd):
day_obs_mjd = np.floor(mjd - 0.5).astype("int")
return day_obs_mjd


def _compute_day_obs_astropy_time(mjd):
day_obs_time = Time(_compute_day_obs_mjd(mjd), format="mjd")
return day_obs_time


def _compute_day_obs_iso8601(mjd):
iso_times = _compute_day_obs_astropy_time(mjd).iso

# Work both for mjd as a scalar and a numpy array
if isinstance(iso_times, str):
day_obs_iso = iso_times[:10]
else:
day_obs_iso = np.array([d[:10] for d in iso_times])

return day_obs_iso


def _compute_day_obs_int(mjd):
day_obs_int = np.array([d.replace("-", "") for d in _compute_day_obs_iso8601(mjd)])

return day_obs_int


class DayObsStacker(BaseStacker):
"""Add dayObs as, as defined by SITCOMTN-32.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD.
"""

cols_added = ["dayObs"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = ["days"]
self.cols_added_dtypes = [int]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_int(sim_data[self.mjd_col])
return sim_data


class DayObsMJDStacker(BaseStacker):
"""Add dayObs defined by SITCOMTN-32, as an MJD.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD.
"""

cols_added = ["day_obs_mjd"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = ["days"]
self.cols_added_dtypes = [int]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_mjd(sim_data[self.mjd_col])
return sim_data


class DayObsISOStacker(BaseStacker):
"""Add dayObs as defined by SITCOMTN-32, in ISO 8601 format.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD."""

cols_added = ["day_obs_iso8601"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = [None]
self.cols_added_dtypes = [(str, 10)]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_iso8601(sim_data[self.mjd_col])
return sim_data
134 changes: 1 addition & 133 deletions rubin_sim/maf/stackers/general_stackers.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,6 @@
"DcrStacker",
"FiveSigmaStacker",
"SaturationStacker",
"ObservationStartDatetime64Stacker",
"DayObsStacker",
"DayObsMJDStacker",
"DayObsISOStacker",
"OverheadStacker",
)

Expand All @@ -25,6 +21,7 @@
from rubin_sim.maf.utils import load_inst_zeropoints

from .base_stacker import BaseStacker
from .date_stackers import DayObsMJDStacker


class SaturationStacker(BaseStacker):
Expand Down Expand Up @@ -510,135 +507,6 @@ def _run(self, sim_data, cols_present=False):
return sim_data


class ObservationStartDatetime64Stacker(BaseStacker):
"""Add the observation start time as a numpy.datetime64."""

cols_added = ["observationStartDatetime64"]

def __init__(
self,
mjd_col="observationStartMJD",
):
self.mjd_col = mjd_col
self.units = [None]
self.cols_added_dtypes = ["datetime64[ns]"]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data["observationStartDatetime64"] = Time(sim_data[self.mjd_col], format="mjd").datetime64

return sim_data


def _compute_day_obs_mjd(mjd):
day_obs_mjd = np.floor(mjd - 0.5).astype("int")
return day_obs_mjd


def _compute_day_obs_astropy_time(mjd):
day_obs_time = Time(_compute_day_obs_mjd(mjd), format="mjd")
return day_obs_time


def _compute_day_obs_iso8601(mjd):
iso_times = _compute_day_obs_astropy_time(mjd).iso

# Work both for mjd as a scalar and a numpy array
if isinstance(iso_times, str):
day_obs_iso = iso_times[:10]
else:
day_obs_iso = np.array([d[:10] for d in iso_times])

return day_obs_iso


def _compute_day_obs_int(mjd):
day_obs_int = np.array([d.replace("-", "") for d in _compute_day_obs_iso8601(mjd)])

return day_obs_int


class DayObsStacker(BaseStacker):
"""Add dayObs as, as defined by SITCOMTN-32.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD.
"""

cols_added = ["dayObs"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = ["days"]
self.cols_added_dtypes = [int]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_int(sim_data[self.mjd_col])
return sim_data


class DayObsMJDStacker(BaseStacker):
"""Add dayObs defined by SITCOMTN-32, as an MJD.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD.
"""

cols_added = ["day_obs_mjd"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = ["days"]
self.cols_added_dtypes = [int]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_mjd(sim_data[self.mjd_col])
return sim_data


class DayObsISOStacker(BaseStacker):
"""Add dayObs as defined by SITCOMTN-32, in ISO 8601 format.
Parameters
----------
mjd_col : `str`
The column with the observatin start MJD."""

cols_added = ["day_obs_iso8601"]

def __init__(self, mjd_col="observationStartMJD"):
self.mjd_col = mjd_col
self.units = [None]
self.cols_added_dtypes = [(str, 10)]

def _run(self, sim_data, cols_present=False):
if cols_present:
# Column already present in data; assume it is correct and does not
# need recalculating.
return sim_data

sim_data[self.cols_added[0]] = _compute_day_obs_iso8601(sim_data[self.mjd_col])
return sim_data


class OverheadStacker(BaseStacker):
"""Add time between visits in seconds.
Expand Down

0 comments on commit c202434

Please sign in to comment.