From d43f0662a2a76516e5a5479604c9e951f8a72ac3 Mon Sep 17 00:00:00 2001
From: Peter Yoachim <yoachim@uw.edu>
Date: Fri, 1 Sep 2023 16:34:39 -0700
Subject: [PATCH] fix m5 diff basis function

---
 .../basis_functions/basis_functions.py        | 38 ++++++++++++++++---
 rubin_sim/scheduler/features/features.py      |  2 +-
 .../model_observatory/model_observatory.py    | 15 ++++++--
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/rubin_sim/scheduler/basis_functions/basis_functions.py b/rubin_sim/scheduler/basis_functions/basis_functions.py
index 67f7b52da..4d871ccd0 100644
--- a/rubin_sim/scheduler/basis_functions/basis_functions.py
+++ b/rubin_sim/scheduler/basis_functions/basis_functions.py
@@ -49,8 +49,9 @@
 
 from rubin_sim.scheduler import features, utils
 from rubin_sim.scheduler.utils import IntRounded
+from rubin_sim.site_models import SeeingModel
 from rubin_sim.skybrightness_pre import dark_sky
-from rubin_sim.utils import _hpid2_ra_dec
+from rubin_sim.utils import _hpid2_ra_dec, m5_flat_sed
 
 
 class BaseBasisFunction:
@@ -836,15 +837,42 @@ class M5DiffBasisFunction(BaseBasisFunction):
     """Basis function based on the 5-sigma depth.
     Look up the best depth a healpixel achieves, and compute
     the limiting depth difference given current conditions
+
+    Parameters
+    ----------
+    fiducial_FWHMEff : `float`
+        The zenith seeing to assume for "good" conditions
     """
 
-    def __init__(self, filtername="r", nside=None):
+    def __init__(self, filtername="r", fiducial_FWHMEff=0.7, nside=None):
         super(M5DiffBasisFunction, self).__init__(nside=nside, filtername=filtername)
-        # Need to look up the deepest m5 values for all the healpixels
-
-        self.dark_map = dark_sky(nside)[filtername]
+        # The dark sky surface brightness values
+        self.dark_sky = dark_sky(nside)[filtername]
+        self.dark_map = None
+        self.fiducial_FWHMEff = fiducial_FWHMEff
+        self.filtername = filtername
 
     def _calc_value(self, conditions, indx=None):
+        if self.dark_map is None:
+            # compute the maximum altitude each HEALpix reaches
+            sindec = np.sin(conditions.dec)
+            sinlat = np.sin(conditions.site.latitude_rad)
+            sinalt = sindec * sinlat + np.cos(conditions.dec) * np.cos(conditions.site.latitude_rad)
+            sinalt = np.clip(sinalt, -1, 1)
+            alt_max = np.arcsin(sinalt)
+            airmass_min = 1.0 / np.cos(np.pi / 2.0 - alt_max)
+            sm = SeeingModel(filter_list=[self.filtername])
+            fwhm_eff = sm(self.fiducial_FWHMEff, airmass_min)["fwhmEff"][0]
+            self.dark_map = m5_flat_sed(
+                self.filtername,
+                musky=self.dark_sky,
+                fwhm_eff=fwhm_eff,
+                exp_time=30.0,
+                airmass=airmass_min,
+                nexp=1,
+                tau_cloud=0,
+            )
+
         # No way to get the sign on this right the first time.
         result = conditions.m5_depth[self.filtername] - self.dark_map
         return result
diff --git a/rubin_sim/scheduler/features/features.py b/rubin_sim/scheduler/features/features.py
index 4b4d86e82..989b7cc79 100644
--- a/rubin_sim/scheduler/features/features.py
+++ b/rubin_sim/scheduler/features/features.py
@@ -127,7 +127,7 @@ def add_observations_array(self, observations_array, observations_hpid):
         if self.filtername is None:
             self.feature += np.size(observations_array)
         else:
-            in_filt = np.where(observations_array["filter"] == self.filter)[0]
+            in_filt = np.where(observations_array["filter"] == self.filtername)[0]
             self.feature += np.size(in_filt)
 
     def add_observation(self, observation, indx=None):
diff --git a/rubin_sim/scheduler/model_observatory/model_observatory.py b/rubin_sim/scheduler/model_observatory/model_observatory.py
index a6a6e1b03..3d96b029a 100644
--- a/rubin_sim/scheduler/model_observatory/model_observatory.py
+++ b/rubin_sim/scheduler/model_observatory/model_observatory.py
@@ -1,5 +1,6 @@
 __all__ = ("ModelObservatory", "NoClouds", "NominalSeeing")
 
+
 import healpy as hp
 import numpy as np
 from astropy.coordinates import EarthLocation
@@ -215,11 +216,12 @@ def __init__(
         sun_moon_info = self.almanac.get_sun_moon_positions(self.mjd)
         season_offset = create_season_offset(self.nside, sun_moon_info["sun_RA"])
         self.sun_ra_start = sun_moon_info["sun_RA"] + 0
+        self.season_offset = season_offset
         # Conditions object to update and return on request
         self.conditions = Conditions(
             nside=self.nside,
-            mjd_start=mjd_start,
-            season_offset=season_offset,
+            mjd_start=self.mjd_start,
+            season_offset=self.season_offset,
             sun_ra_start=self.sun_ra_start,
         )
 
@@ -246,8 +248,13 @@ def return_conditions(self):
         -------
         rubin_sim.scheduler.features.conditions object
         """
-
-        self.conditions.mjd = self.mjd
+        self.conditions = Conditions(
+            nside=self.nside,
+            mjd_start=self.mjd_start,
+            season_offset=self.season_offset,
+            sun_ra_start=self.sun_ra_start,
+            mjd=self.mjd,
+        )
 
         self.conditions.night = int(self.night)
         # Current time as astropy time