Merge pull request #101 from lsst/tickets/PREOPS-5360

tickets/PREOPS-5360: clean up the nightsum and prenight Times Square notebooks by updating schedview to support them better

schedview/collect/opsim.py

@@ -1,6 +1,8 @@
 import sqlite3
+from warnings import warn
 
 import pandas as pd
+import rubin_scheduler
 import yaml
 from astropy.time import Time
 from lsst.resources import ResourcePath
@@ -90,7 +92,17 @@ def read_opsim(
                 ]
 
             try:
-                visits = pd.DataFrame(maf.get_sim_data(sim_connection, constraint, dbcols, **kwargs))
+                try:
+                    visits = pd.DataFrame(maf.get_sim_data(sim_connection, constraint, dbcols, **kwargs))
+                except UserWarning:
+                    warn("No visits match constraints.")
+                    visits = (
+                        SchemaConverter()
+                        .obs2opsim(rubin_scheduler.scheduler.utils.empty_observation())
+                        .drop(index=0)
+                    )
+                    if "observationId" not in visits.columns and "ID" in visits.columns:
+                        visits.rename(columns={"ID": "observationId"}, inplace=True)
             except NameError as e:
                 if e.name == "maf" and e.args == ("name 'maf' is not defined",):
                     if len(kwargs) > 0:

schedview/compute/__init__.py

@@ -1,6 +1,7 @@
 __all__ = [
     "convert_evening_date_to_night_of_survey",
     "night_events",
+    "compute_sun_moon_positions",
     "LsstCameraFootprintPerimeter",
     "replay_visits",
     "compute_basis_function_reward_at_time",
@@ -15,7 +16,7 @@
     "visits",
 ]
 
-from .astro import convert_evening_date_to_night_of_survey, night_events
+from .astro import compute_sun_moon_positions, convert_evening_date_to_night_of_survey, night_events
 from .camera import LsstCameraFootprintPerimeter
 from .scheduler import (
     compute_basis_function_reward_at_time,

schedview/compute/astro.py

@@ -129,3 +129,40 @@ def compute_central_night(visits, site=None, timezone="Chile/Continental"):
     central_night = Time(central_mjd + mjd_shift, format="mjd", scale="utc").datetime.date()
 
     return central_night
+
+
+def compute_sun_moon_positions(observatory: ModelObservatory) -> pd.DataFrame:
+    """Create a DataFrame of sun and moon positions with one row per body,
+    one column per coordinate, at the time set for the observatory.
+
+    Parameters
+    ----------
+    observatory : `ModelObservatory`
+        The model observatory.
+
+    Returns
+    -------
+    body_positions : `pandas.DataFrame`
+        The table of body positions.
+    """
+    body_positions_wide = pd.DataFrame(observatory.almanac.get_sun_moon_positions(observatory.mjd))
+
+    body_positions_wide.index.name = "r"
+    body_positions_wide.reset_index(inplace=True)
+
+    angle_columns = ["RA", "dec", "alt", "az"]
+    all_columns = angle_columns + ["phase"]
+    body_positions = (
+        pd.wide_to_long(
+            body_positions_wide,
+            stubnames=("sun", "moon"),
+            suffix=r".*",
+            sep="_",
+            i="r",
+            j="coordinate",
+        )
+        .droplevel("r")
+        .T[all_columns]
+    )
+    body_positions[angle_columns] = np.degrees(body_positions[angle_columns])
+    return body_positions

schedview/compute/maf.py

@@ -1,4 +1,3 @@
-import sqlite3
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
@@ -10,12 +9,12 @@
 
 
 def _visits_to_opsim(visits, opsim):
-    # Only write columns in visits that are in opsim databases,
-    # thereby avoiding added columns that might not be types
-    # that can be written to sqlite databases.
-    opsim_columns = list(SchemaConverter().convert_dict.keys())
-    with sqlite3.connect(opsim) as connection:
-        visits.reset_index()[opsim_columns].to_sql("observations", connection, index=False)
+    # Take advantage of the schema migration code in SchemaConverter to make
+    # sure we have the correct column names
+
+    schema_converter = SchemaConverter()
+    obs = schema_converter.opsimdf2obs(visits)
+    schema_converter.obs2opsim(obs, filename=opsim)
 
 
 def compute_metric(visits, metric_bundle):

schedview/plot/rewards.py

@@ -1,6 +1,12 @@
 import warnings
+from collections.abc import Callable
 
 import bokeh
+import bokeh.io
+import bokeh.layouts
+import bokeh.models
+import bokeh.plotting
+import bokeh.transform
 import colorcet
 import holoviews as hv
 import numpy as np
@@ -532,3 +538,61 @@ def reward_timeline_for_surveys(rewards_df, day_obs_mjd, show=True, **figure_kwa
         plot = None
 
     return plot
+
+
+def nested_tier_reward_timeline_plot(
+    rewards_df: pd.DataFrame,
+    plot_func: Callable,
+    day_obs_mjd: int,
+    max_basis_functions_per_tab: int = 200,
+    show: bool = True,
+) -> bokeh.models.UIElement:
+    """Plot rewards timelines tabs by tier.
+
+    Parameters
+    ----------
+    rewards_df : `pandas.DataFrame`
+        The table of rewards data.
+    plot_func : `Callable`
+        Plot function to be called in each tab.
+    day_obs_mjd : `int`
+        The MJD of the day_obs of the night to plot
+    max_basis_functions_per_tab : `int`
+        Maximum basis functions to plot in each tab, by default 200
+    show : `bool`
+        Actually show the plot? Defaults to `True`.
+
+    Returns
+    -------
+    plot : `bokeh.models.UIElement`
+        The bokeh plot.
+    """
+
+    tier_plot = {}
+    tier_indexes = np.sort(rewards_df.reset_index().list_index.unique())
+    sorted_rewards_df = rewards_df.sort_index()
+    for tier_index in tier_indexes:
+        num_basis_functions = len(
+            sorted_rewards_df.loc[tier_index, ("basis_function", "survey_label")].drop_duplicates()
+        )
+        if num_basis_functions > max_basis_functions_per_tab:
+            tier_plot[tier_index] = bokeh.models.Div(text="Too many basis functions to plot.")
+        else:
+            try:
+                tier_plot[tier_index] = plot_func(rewards_df, tier_index, day_obs_mjd, show=False)
+            except Exception as e:
+                print(f"Not showing tier {tier_index} due to an exception: {str(e)}")
+                tier_plot[tier_index] = None
+
+    plot = bokeh.models.Tabs(
+        tabs=[
+            bokeh.models.TabPanel(child=tier_plot[t], title=f"Tier {t}")
+            for t in tier_indexes
+            if tier_plot[t] is not None
+        ]
+    )
+
+    if show:
+        bokeh.io.show(plot)
+
+    return plot

schedview/plot/survey.py

@@ -1,18 +1,24 @@
+import datetime
 import re
 from functools import partial
 
 import astropy
 import bokeh
+import bokeh.io
 import colorcet
 import healpy as hp
 import numpy as np
 import pandas as pd
 import rubin_scheduler.scheduler.features
 import rubin_scheduler.scheduler.surveys  # noqa: F401
 from astropy.time import Time
+from bokeh.models.ui.ui_element import UIElement
+from matplotlib.figure import Figure
 from uranography.api import HorizonMap, Planisphere, make_zscale_linear_cmap
 
+import schedview.compute.astro
 from schedview.compute.camera import LsstCameraFootprintPerimeter
+from schedview.compute.maf import compute_hpix_metric_in_bands
 
 
 def map_survey_healpix(
@@ -411,3 +417,57 @@ def create_hpix_visit_map_grid(hpix_maps, visits, conditions, **kwargs):
 
     map_grid = bokeh.layouts.gridplot(map_lists)
     return map_grid
+
+
+def create_metric_visit_map_grid(
+    metric, metric_visits, visits, observatory, nside=32, use_matplotlib=False, **kwargs
+) -> Figure | UIElement | None:
+    """Create a grid of maps of metric values with visits overplotted.
+
+    Parameters
+    ----------
+    metric : `numpy.array`
+        An array of healpix values
+    metric_visits : `pd.DataFrame`
+        The visits to use to compute the metric
+    visits : `pd.DataFrame`
+        The table of visits to plot, with columns matching the opsim database
+        definitions.
+    observatory : `ModelObservatory`
+        The model observotary to use.
+    nside : `int`
+        The nside with which to compute the metric.
+    use_matplotlib: `bool`
+        Use matplotlib instead of bokeh? Defaults to False.
+
+    Returns
+    -------
+    plot : `bokeh.models.plots.Plot`
+        The plot with the map
+    """
+
+    if len(metric_visits):
+        metric_hpix = compute_hpix_metric_in_bands(metric_visits, metric, nside=nside)
+    else:
+        metric_hpix = np.zeros(hp.nside2npix(nside))
+
+    if len(visits):
+        if use_matplotlib:
+            from schedview.plot import survey_skyproj
+
+            day_obs_mjd = np.floor(observatory.mjd - 0.5).astype("int")
+            day_obs_dt = Time(day_obs_mjd, format="mjd").datetime
+            day_obs_date = datetime.date(day_obs_dt.year, day_obs_dt.month, day_obs_dt.day)
+            night_events = schedview.compute.astro.night_events(day_obs_date)
+            fig = survey_skyproj.create_hpix_visit_map_grid(
+                visits, metric_hpix, observatory, night_events, **kwargs
+            )
+            return fig
+        else:
+            map_grid = create_hpix_visit_map_grid(
+                metric_hpix, visits, observatory.return_conditions(), **kwargs
+            )
+            bokeh.io.show(map_grid)
+            return map_grid
+    else:
+        print("No visits")