Merge pull request #452 from lsst/tickets/SP-2082

ehneilsen · web-flow · commit cc8dc54279f2 · 2025-03-28T15:52:18.000-05:00
tickets/SP-2082: Support getting obsloctac visits from the pre-night sim archive
diff --git a/rubin_sim/maf/utils/opsim_utils.py b/rubin_sim/maf/utils/opsim_utils.py
@@ -6,6 +6,8 @@
 
 import os
 import sqlite3
+import urllib
+from contextlib import closing
 
 import numpy as np
 import pandas as pd
@@ -15,8 +17,8 @@
 
 def get_sim_data(
     db_con,
-    sqlconstraint,
-    dbcols,
+    sqlconstraint=None,
+    dbcols=None,
     stackers=None,
     table_name=None,
     full_sql_query=None,
@@ -77,26 +79,35 @@ def get_sim_data(
         # that's probably fine, keep people from getting fancy with old sims
         table_name = "observations"
 
-    if isinstance(db_con, str):
-        con = sqlite3.connect(db_con)
-    else:
-        con = db_con
-
     if full_sql_query is None:
-        col_str = ""
-        for colname in dbcols:
-            col_str += colname + ", "
-        col_str = col_str[0:-2] + " "
+        col_str = "*" if dbcols is None else ", ".join(dbcols)
 
         query = "SELECT %s FROM %s" % (col_str, table_name)
         if len(sqlconstraint) > 0:
             query += " WHERE %s" % (sqlconstraint)
         query += ";"
-        sim_data = pd.read_sql(query, con).to_records(index=False)
-
     else:
         query = full_sql_query
-        sim_data = pd.read_sql(query, con).to_records(index=False)
+
+    if isinstance(db_con, sqlite3.Connection):
+        sim_data = pd.read_sql(query, db_con).to_records(index=False)
+    elif isinstance(db_con, str) and os.path.isfile(db_con):
+        with closing(sqlite3.connect(db_con)) as con:
+            sim_data = pd.read_sql(query, con).to_records(index=False)
+    elif (not isinstance(db_con, str)) or urllib.parse.urlparse(db_con).scheme != "":
+        try:
+            from lsst.resources import ResourcePath
+
+            with ResourcePath(db_con).as_local() as local_db_path:
+                with closing(sqlite3.connect(local_db_path.ospath)) as con:
+                    sim_data = pd.read_sql(query, con).to_records(index=False)
+        except ModuleNotFoundError:
+            raise RuntimeError(
+                f"Cannot read visits from {db_con}."
+                "Maybe it does not exist, or maybe you need to install lsst.resources."
+            )
+    else:
+        raise RuntimeError("Cannot find {db_con}.")
 
     if len(sim_data) == 0:
         raise UserWarning("No data found matching sqlconstraint %s" % (sqlconstraint))
diff --git a/rubin_sim/sim_archive/sim_archive.py b/rubin_sim/sim_archive/sim_archive.py
@@ -33,6 +33,9 @@
 from astropy.time import Time
 from rubin_scheduler.scheduler import sim_runner
 from rubin_scheduler.scheduler.utils import SchemaConverter
+from rubin_scheduler.site_models.almanac import Almanac
+
+from rubin_sim.maf.utils.opsim_utils import get_sim_data
 
 LOGGER = logging.getLogger(__name__)
 
@@ -931,3 +934,189 @@ def compile_sim_archive_metadata_cli(*args):
         compilation_resource = ResourcePath(compilation_uri)
 
     compilation_resource = compile_sim_metadata(archive_uri, compilation_resource, append=append)
+
+
+def find_latest_prenight_sim_for_nights(
+    first_day_obs: str | None = None,
+    last_day_obs: str | None = None,
+    tags: tuple[str] = ("ideal", "nominal"),
+    max_simulation_age: int = 2,
+    archive_uri: str = "s3://rubin:rubin-scheduler-prenight/opsim/",
+    compilation_uri: str = "s3://rubin:rubin-scheduler-prenight/opsim/compiled_metadata_cache.h5",
+) -> pd.DataFrame:
+    """Find the most recent prenight simulation that covers a night.
+
+    Parameters
+    ----------
+    first_day_obs : `str` or  `None`
+        The date of the evening for the first night for which to get
+        a simulation. If `None`, then the current date will be used.
+    last_day_obs : `str` or  `None`
+        The date of the evening for the last night for which to get
+        a simulation. If `None`, then the current date will be used.
+    tags : `tuple[str]`
+        A tuple of tags to filter simulations by.
+        Defaults to ``('ideal', 'nominal')``.
+    max_simulation_age : `int`
+        The maximum age of simulations to consider, in days.
+        Simulations older than ``max_simulation_age`` will not be considered.
+        Defaults to 2.
+    archive_uri : `str`
+        The URI of the archive from which to fetch the simulation.
+        Defaults to ``s3://rubin:rubin-scheduler-prenight/opsim/``.
+    compilation_uri : `str`
+        The URI of the compiled metadata HDF5 file for efficient querying.
+        Defaults to
+        ``s3://rubin:rubin-scheduler-prenight/opsim/compiled_metadata_cache.h5``.
+
+    Returns
+    -------
+    sim_metadata : `dict`
+        A dictionary with metadata for the simulation.
+    """
+
+    if first_day_obs is None:
+        first_day_obs = Time(Time.now().mjd - 0.5, format="mjd").iso[:10]
+    if last_day_obs is None:
+        last_day_obs = first_day_obs
+
+    sim_metadata = read_archived_sim_metadata(
+        archive_uri, num_nights=max_simulation_age, compilation_resource=compilation_uri
+    )
+
+    best_sim = None
+    for uri, sim in sim_metadata.items():
+        sim["uri"] = uri
+        sim["exec_date"] = uri.split("/")[-3]
+        sim["date_index"] = int(uri.split("/")[-2])
+
+        if sim["simulated_dates"]["first"] > first_day_obs:
+            continue
+        if sim["simulated_dates"]["last"] < last_day_obs:
+            continue
+        if not set(tags).issubset(sim["tags"]):
+            continue
+        if best_sim is not None:
+            if sim["exec_date"] < best_sim["exec_date"]:
+                continue
+            if sim["date_index"] < best_sim["date_index"]:
+                continue
+        best_sim = sim
+
+    if best_sim is not None:
+        best_sim["opsim_rp"] = (
+            ResourcePath(archive_uri)
+            .join(best_sim["exec_date"], forceDirectory=True)
+            .join(f"{best_sim['date_index']}", forceDirectory=True)
+            .join(best_sim["files"]["observations"]["name"])
+        )
+
+    return best_sim
+
+
+def fetch_latest_prenight_sim_for_nights(
+    first_day_obs: str | None = None,
+    last_day_obs: str | None = None,
+    tags: tuple[str] = ("ideal", "nominal"),
+    max_simulation_age: int = 2,
+    archive_uri: str = "s3://rubin:rubin-scheduler-prenight/opsim/",
+    compilation_uri: str = "s3://rubin:rubin-scheduler-prenight/opsim/compiled_metadata_cache.h5",
+    **kwargs,
+) -> pd.DataFrame:
+    """Fetches visit parameters from the latest archived pre-night simulation
+    with requested tags for a specified day of observing.
+
+    Parameters
+    ----------
+    first_day_obs : `str` or  `None`
+        The date of the evening for the first night for which to get
+        a simulation. If `None`, then the current date will be used.
+    last_day_obs : `str` or  `None`
+        The date of the evening for the last night for which to get
+        a simulation. If `None`, then the current date will be used.
+    tags : `tuple[str]`
+        A tuple of tags to filter simulations by.
+        Defaults to ``('ideal', 'nominal')``.
+    max_simulation_age : `int`
+        The maximum age of simulations to consider, in days. Simulations older
+        than ``max_simulation_age`` will not be considered. Defaults to 2.
+    archive_uri : `str`
+        The URI of the archive from which to fetch the simulation.
+        Defaults to ``s3://rubin:rubin-scheduler-prenight/opsim/``.
+    compilation_uri : `str`
+        The URI of the compiled metadata HDF5 file for efficient querying.
+        Defaults to ``s3://rubin:rubin-scheduler-prenight/opsim/compiled_metadata_cache.h5``.
+    **kwargs
+        Additional keyword arguments passed to
+        `rubin_sim.maf.get_sim_data`.
+
+    Returns
+    -------
+    visits : `pd.DataFrame`
+        A pandas DataFrame containing visit parameters.
+    """
+
+    sim_metadata = find_latest_prenight_sim_for_nights(
+        first_day_obs, last_day_obs, tags, max_simulation_age, archive_uri, compilation_uri
+    )
+    visits = get_sim_data(sim_metadata["opsim_rp"], **kwargs)
+
+    return visits
+
+
+def fetch_obsloctap_visits(day_obs: str | None = None, nights: int = 2) -> pd.DataFrame:
+    """Return visits from latest nominal prenight briefing simulation.
+
+    Parameters
+    ----------
+    day_obs : `str`
+        The day_obs of the night, in YYYY-MM-DD format (e.g. 2025-03-26).
+        Default None will use the date of the next sunset.
+    nights : `int`
+        The number of nights of observations to return.
+        Defaults to 2.
+
+    Returns
+    -------
+    visits : `pd.DataFrame`
+        The visits from the prenight simulation.
+    """
+    dbcols = [
+        "observationStartMJD",
+        "fieldRA",
+        "fieldDec",
+        "rotSkyPos",
+        "band",
+        "visitExposureTime",
+        "night",
+        "target_name",
+    ]
+
+    # Start with the first night that starts after the reference time,
+    # which is the current time by default.
+    # So, if the reference time is during a night, it starts with the
+    # following night.
+    night_bounds = pd.DataFrame(Almanac().sunsets)
+    reference_time = Time.now() if day_obs is None else Time(day_obs, format="iso", scale="utc")
+    first_night = night_bounds.query(f"sunset > {reference_time.mjd}").night.min()
+    last_night = first_night + nights - 1
+
+    night_bounds.set_index("night", inplace=True)
+    start_mjd = night_bounds.loc[first_night, "sunset"]
+    end_mjd = night_bounds.loc[last_night, "sunrise"]
+
+    sqlconstraint = (f"observationStartMJD BETWEEN {start_mjd} AND {end_mjd}",)
+    max_simulation_age = int(np.ceil(Time.now().mjd - reference_time.mjd)) + 1
+
+    first_day_obs = Time(start_mjd - 0.5, format="mjd").iso[:10]
+    last_day_obs = Time(end_mjd - 0.5, format="mjd").iso[:10]
+    visits = fetch_latest_prenight_sim_for_nights(
+        first_day_obs,
+        last_day_obs,
+        tags=("ideal", "nominal"),
+        max_simulation_age=max_simulation_age,
+        sqlconstraint=sqlconstraint,
+        dbcols=dbcols,
+    )
+
+    return visits
diff --git a/tests/sim_archive/test_sim_archive.py b/tests/sim_archive/test_sim_archive.py
@@ -6,6 +6,9 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
+import numpy as np
+import pandas as pd
+from astropy.time import Time
 from rubin_scheduler.scheduler import sim_runner
 from rubin_scheduler.scheduler.example import example_scheduler
 from rubin_scheduler.scheduler.model_observatory import ModelObservatory
@@ -18,6 +21,9 @@
     from rubin_sim.sim_archive.sim_archive import (
         check_opsim_archive_resource,
         compile_sim_metadata,
+        fetch_latest_prenight_sim_for_nights,
+        fetch_obsloctap_visits,
+        find_latest_prenight_sim_for_nights,
         make_sim_archive_cli,
         make_sim_archive_dir,
         read_archived_sim_metadata,
@@ -134,6 +140,28 @@ def test_cli(self):
                 test_archive_uri,
             )
 
+    @unittest.skipIf(not HAVE_LSST_RESOURCES, "No lsst.resources")
+    def test_find_latest_prenight_sim_for_night(self):
+        day_obs = "2025-03-25"
+        max_simulation_age = int(np.ceil(Time.now().mjd - Time(day_obs).mjd)) + 1
+        sim_metadata = find_latest_prenight_sim_for_nights(day_obs, max_simulation_age=max_simulation_age)
+        assert sim_metadata["simulated_dates"]["first"] <= day_obs <= sim_metadata["simulated_dates"]["last"]
+
+    @unittest.skipIf(not HAVE_LSST_RESOURCES, "No lsst.resources")
+    def test_fetch_latest_prenight_sim_for_night(self):
+        day_obs = "2025-03-25"
+        max_simulation_age = int(np.ceil(Time.now().mjd - Time(day_obs).mjd)) + 1
+        visits = fetch_latest_prenight_sim_for_nights(day_obs, max_simulation_age=max_simulation_age)
+        assert len(visits) > 0
+
+    @unittest.skipIf(not HAVE_LSST_RESOURCES, "No lsst.resources")
+    def test_fetch_obsloctap_visits(self):
+        day_obs = "2025-03-25"
+        num_nights = 2
+        visits = pd.DataFrame(fetch_obsloctap_visits(day_obs, nights=num_nights))
+        assert np.floor(visits["observationStartMJD"].min() - 0.5) == Time(day_obs).mjd
+        assert np.floor(visits["observationStartMJD"].max() - 0.5) == Time(day_obs).mjd + num_nights - 1
+
 
 if __name__ == "__main__":
     unittest.main()
