SCHED-406: Changing List[T] to Dict[NightIndex, T] and massive cleanup.

scheduler/core/calculations/groupinfo.py

@@ -2,17 +2,15 @@
 # For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
 
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict
 
-from lucupy.decorators import immutable
-from lucupy.minimodel import Conditions, Group, UniqueGroupID
+from lucupy.minimodel import Conditions, Group, NightIndex, UniqueGroupID
 import numpy.typing as npt
 
 from .scores import Scores
 
 
-@immutable
-@dataclass(frozen=True)
+@dataclass
 class GroupInfo:
     """
     Information regarding Groups that can only be calculated in the Selector.
@@ -36,10 +34,10 @@ class GroupInfo:
     minimum_conditions: Conditions
     is_splittable: bool
     standards: float
-    night_filtering: npt.NDArray[bool]
-    conditions_score: List[npt.NDArray[float]]
-    wind_score: List[npt.NDArray[float]]
-    schedulable_slot_indices: List[npt.NDArray[int]]
+    night_filtering: Dict[NightIndex, bool]
+    conditions_score: Dict[NightIndex, npt.NDArray[float]]
+    wind_score: Dict[NightIndex, npt.NDArray[float]]
+    schedulable_slot_indices: Dict[NightIndex, npt.NDArray[int]]
     scores: Scores
 
 

scheduler/core/calculations/scores.py

@@ -1,13 +1,14 @@
 # Copyright (c) 2016-2022 Association of Universities for Research in Astronomy, Inc. (AURA)
 # For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
 
-from typing import List
+from typing import Dict
 
+from lucupy.minimodel import NightIndex
 import numpy.typing as npt
 
 # Scores for the timeslots in a specific night.
 NightTimeSlotScores = npt.NDArray[float]
 
 # Scores across all nights per timeslot.
 # Indexed by night index, and then timeslot index.
-Scores = List[NightTimeSlotScores]
+Scores = Dict[NightIndex, NightTimeSlotScores]

scheduler/core/components/collector/__init__.py

@@ -488,13 +488,11 @@ def load_programs(self, program_provider_class: Type[ProgramProvider], data: Ite
 
     def night_configurations(self,
                              site: Site,
-                             night_indices: NightIndices) -> List[NightConfiguration]:
+                             night_indices: NightIndices) -> Dict[NightIndices, NightConfiguration]:
         """
         Return the list of NightConfiguration for the site and nights under configuration.
         """
-        return [
-            Collector._resource_service.get_night_configuration(
-                site,
-                self.get_night_events(site).time_grid[night_idx].datetime.date() - Collector._DAY
-            )
-            for night_idx in night_indices]
+        return {night_idx: Collector._resource_service.get_night_configuration(
+            site,
+            self.get_night_events(site).time_grid[night_idx].datetime.date() - Collector._DAY
+        ) for night_idx in night_indices}

scheduler/core/components/optimizer/dummy.py

@@ -61,7 +61,7 @@ def add(self, group: GroupData, plans: Plans, interval: Optional[Interval] = Non
                 obs_len = plan.time2slots(plan.time_slot_length, observation.exec_time())
                 if plan.time_left() >= obs_len and observation not in plan:
                     start, start_time_slot = DummyOptimizer._first_free_time(plan)
-                    visit_score = sum(group.group_info.scores[plans.night][start_time_slot:start_time_slot+obs_len])
+                    visit_score = sum(group.group_info.scores[plans.night_idx][start_time_slot:start_time_slot + obs_len])
                     plan.add(observation, start, start_time_slot, obs_len, visit_score)
                     return True
                 else:

scheduler/core/components/optimizer/timeline.py

@@ -6,7 +6,7 @@
 from typing import ClassVar, List, Mapping, Optional, Sequence, Tuple
 
 import numpy as np
-from lucupy.minimodel import Observation, ObservationID, Site
+from lucupy.minimodel import NightIndex, Observation, ObservationID, Site
 
 from scheduler.core.calculations.nightevents import NightEvents
 from . import Interval
@@ -141,10 +141,10 @@ class Timelines:
     A collection of Timeline from all sites for a specific night
     """
 
-    def __init__(self, night_events: Mapping[Site, NightEvents], night_idx: int):
+    def __init__(self, night_events: Mapping[Site, NightEvents], night_idx: NightIndex):
 
         self.timelines = {}
-        self.night = night_idx
+        self.night_idx = night_idx
         for site, ne in night_events.items():
             if ne is not None:
                 self.timelines[site] = Timeline(ne.local_times[night_idx][0],

scheduler/core/components/ranker/base.py

@@ -2,11 +2,11 @@
 # For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
 
 from abc import abstractmethod
-from typing import Dict, FrozenSet, List
+from typing import Dict, FrozenSet
 
 import numpy as np
 import numpy.typing as npt
-from lucupy.minimodel import ALL_SITES, AndGroup, OrGroup, Group, NightIndices, Observation, Program, Site
+from lucupy.minimodel import ALL_SITES, AndGroup, OrGroup, Group, NightIndex, NightIndices, Observation, Program, Site
 
 from scheduler.core.calculations import Scores, GroupDataMap
 from scheduler.core.components.collector import Collector
@@ -27,19 +27,19 @@ def __init__(self,
         # 1. An empty observation score array.
         # 2. An empty group scoring array, used to collect the group scores.
         # This allows us to avoid having to store a reference to the Collector in the Ranker.
-        self._empty_obs_scores: Dict[Site, List[npt.NDArray[float]]] = {}
-        self._empty_group_scores: Dict[Site, List[npt.NDArray[float]]] = {}
+        self._empty_obs_scores: Dict[Site, Dict[NightIndex, npt.NDArray[float]]] = {}
+        self._empty_group_scores: Dict[Site, Dict[NightIndex, npt.NDArray[float]]] = {}
         for site in self.sites:
             night_events = collector.get_night_events(site)
 
             # Create a full zero score that fits the sites, nights, and time slots for observations.
-            self._empty_obs_scores[site] = [np.zeros(len(night_events.times[night_idx]), dtype=float)
-                                            for night_idx in self.night_indices]
+            self._empty_obs_scores[site] = {night_idx: np.zeros(len(night_events.times[night_idx]), dtype=float)
+                                            for night_idx in self.night_indices}
 
             # Create a full zero score that fits the sites, nights, and time slots for group calculations.
             # As this must collect the subgroups, the dimensions are different from observation scores.
-            self._empty_group_scores[site] = [np.zeros((0, len(night_events.times[night_idx])), dtype=float)
-                                              for night_idx in self.night_indices]
+            self._empty_group_scores[site] = {night_idx: np.zeros((0, len(night_events.times[night_idx])), dtype=float)
+                                              for night_idx in self.night_indices}
 
     def score_group(self, group: Group, group_data_map: GroupDataMap) -> Scores:
         """

scheduler/core/components/ranker/default.py

@@ -204,33 +204,35 @@ def score_observation(self, program: Program, obs: Observation) -> Scores:
         print(f'   cplt: {cplt:.2f}  metric: {metric[0]:.2f}')
 
         # Declination for the base target per night.
-        dec = [target_info[night_idx].coord.dec for night_idx in self.night_indices]
+        dec = {night_idx: target_info[night_idx].coord.dec for night_idx in self.night_indices}
 
         # Hour angle / airmass
-        ha = [target_info[night_idx].hourangle for night_idx in self.night_indices]
+        ha = {night_idx: target_info[night_idx].hourangle for night_idx in self.night_indices}
 
         # Get the latitude associated with the site.
         site_latitude = obs.site.location.lat
         if site_latitude < 0. * u.deg:
-            dec_diff = [np.abs(site_latitude - np.max(dec[night_idx])) for night_idx in self.night_indices]
+            dec_diff = {night_idx: np.abs(site_latitude - np.max(dec[night_idx])) for night_idx in self.night_indices}
         else:
-            dec_diff = [np.abs(np.min(dec[night_idx]) - site_latitude) for night_idx in self.night_indices]
+            dec_diff = {night_idx: np.abs(np.min(dec[night_idx]) - site_latitude) for night_idx in self.night_indices}
 
-        c = np.array([self.params.dec_diff_less_40 if angle < 40. * u.deg
-                      else self.params.dec_diff for angle in dec_diff])
+        c = {night_idx: self.params.dec_diff_less_40 if angle < 40. * u.deg else self.params.dec_diff
+             for night_idx, angle in dec_diff.items()}
+        # c = np.array([self.params.dec_diff_less_40 if angle < 40. * u.deg
+        #               else self.params.dec_diff for angle in dec_diff])
 
-        wha = [c[night_idx][0] + c[night_idx][1] * ha[night_idx] / u.hourangle
+        wha = {night_idx: c[night_idx][0] + c[night_idx][1] * ha[night_idx] / u.hourangle
                + (c[night_idx][2] / u.hourangle ** 2) * ha[night_idx] ** 2
-               for night_idx in self.night_indices]
-        kk = [np.where(wha[night_idx] <= 0.)[0] for night_idx in self.night_indices]
+               for night_idx in self.night_indices}
+        kk = {night_idx: np.where(wha[night_idx] <= 0.)[0] for night_idx in self.night_indices}
         for night_idx in self.night_indices:
             wha[night_idx][kk[night_idx]] = 0.
-        print(f'   max wha: {np.max(wha[0]):.2f}  visfrac: {target_info[0].rem_visibility_frac:.5f}')
+        # print(f'   max wha: {np.max(wha[0]):.2f}  visfrac: {target_info[0].rem_visibility_frac:.5f}')
 
-        p = [(metric[0] ** self.params.met_power) *
+        p = {night_idx: (metric[0] ** self.params.met_power) *
              (target_info[night_idx].rem_visibility_frac ** self.params.vis_power) *
              (wha[night_idx] ** self.params.wha_power)
-             for night_idx in self.night_indices]
+             for night_idx in self.night_indices}
 
         # Assign scores in p to all indices where visibility constraints are met.
         # They will otherwise be 0 as originally defined.
@@ -270,8 +272,8 @@ def _score_and_group(self, group: AndGroup, group_data_map: GroupDataMap) -> Sco
 
         # Combine the scores as per the score_combiner and return.
         # apply_along_axis results in a (1, #timeslots in night) array, so we have to take index 0.
-        return [np.apply_along_axis(self.params.score_combiner, 0, scores[night_idx])[0]
-                for night_idx in self.night_indices]
+        return {night_idx: np.apply_along_axis(self.params.score_combiner, 0, scores[night_idx])[0]
+                for night_idx in self.night_indices}
 
     def _score_or_group(self, group: OrGroup, group_data_map: GroupDataMap) -> Scores:
         raise NotImplementedError

scheduler/core/components/selector/__init__.py

@@ -88,9 +88,10 @@ def select(self,
         if ranker is None:
             ranker = DefaultRanker(self.collector, night_indices, sites)
 
-        # The night_indices in the Selector and Ranker must be the same.
-        if not np.array_equal(night_indices, ranker.night_indices):
-            raise ValueError(f'The Ranker must have the same night indices as the Selector select method.')
+        # TODO: ERASE?
+        # # The night_indices in the Selector and Ranker must be the same.
+        # if not np.array_equal(night_indices, ranker.night_indices):
+        #     raise ValueError(f'The Ranker must have the same night indices as the Selector select method.')
 
         # Create the structure to hold the mapping fom program ID to its group info.
         program_info_map: Dict[ProgramID, ProgramInfo] = {}
@@ -187,8 +188,9 @@ def score_program(self,
                                                           night_configurations,
                                                           ranker)
 
+        # We want to check if there are any time slots where a group can be scheduled: otherwise, we omit it.
         group_data_map = {gp_id: gp_data for gp_id, gp_data in unfiltered_group_data_map.items()
-                          if any(len(indices) > 0 for indices in gp_data.group_info.schedulable_slot_indices)}
+                          if any(len(indices) > 0 for indices in gp_data.group_info.schedulable_slot_indices.values())}
 
         # In an observation group, the only child is an Observation:
         # hence, references here to group.children are simply the Observation.
@@ -286,26 +288,26 @@ def _calculate_observation_group(self,
 
         # Calculate a numpy array of bool indexed by night to determine when the group can be added to the plan
         # based on the night configuration filtering.
-        night_filtering = np.array([night_configurations[obs.site][night_idx].filter.group_filter(group)
-                                    for night_idx in night_indices])
+        night_filtering = {night_idx: night_configurations[obs.site][night_idx].filter.group_filter(group)
+                           for night_idx in night_indices}
 
         if obs.obs_class in [ObservationClass.SCIENCE, ObservationClass.PROGCAL]:
             # If we are science or progcal, then the neg HA value for a night is if the first HA for the night
             # is negative.
-            neg_ha = np.array([target_info[night_idx].hourangle[0].value < 0 for night_idx in ranker.night_indices])
+            neg_ha = {night_idx: target_info[night_idx].hourangle[0].value < 0 for night_idx in night_indices}
         else:
-            neg_ha = np.array([False] * len(ranker.night_indices))
+            neg_ha = {night_idx: False for night_idx in night_indices}
         too_type = obs.too_type
 
         # Calculate when the conditions are met and an adjustment array if the conditions are better than needed.
         # TODO: Maybe we only need to concern ourselves with the night indices where the resources are in place.
-        conditions_score = []
-        wind_score = []
+        conditions_score = {}
+        wind_score = {}
 
         # We need the night_events for the night for timing information.
         night_events = self.collector.get_night_events(obs.site)
 
-        for night_idx in ranker.night_indices:
+        for night_idx in night_indices:
             # Get the conditions for the night.
             start_time = night_events.times[night_idx][0]
             end_time = night_events.times[night_idx][-1]
@@ -325,33 +327,35 @@ def _calculate_observation_group(self,
             # If we can obtain the conditions variant, calculate the conditions and wind mapping.
             # Otherwise, use arrays of all zeros to indicate that we cannot calculate this information.
             if actual_conditions is not None:
-                conditions_score.append(Selector._match_conditions(mrc, actual_conditions, neg_ha[night_idx], too_type))
-                wind_score.append(Selector._wind_conditions(actual_conditions, target_info[night_idx].az))
+                conditions_score[night_idx] = Selector._match_conditions(mrc,
+                                                                         actual_conditions,
+                                                                         neg_ha[night_idx],
+                                                                         too_type)
+                wind_score[night_idx] = Selector._wind_conditions(actual_conditions, target_info[night_idx].az)
             else:
                 zero = np.zeros(len(night_events.times[night_idx]))
-                conditions_score.append(zero.copy())
-                wind_score.append(zero.copy())
+                conditions_score[night_idx] = zero.copy()
+                wind_score[night_idx] = zero.copy()
 
         # Calculate the schedulable slot indices.
         # These are the indices where the observation has:
         # 1. Visibility
         # 2. Resources available
         # 3. Conditions that are met
-        schedulable_slot_indices = []
-        for night_idx in ranker.night_indices:
+        schedulable_slot_indices = {}
+        for night_idx in night_indices:
             vis_idx = target_info[night_idx].visibility_slot_idx
             if night_filtering[night_idx]:
-                schedulable_slot_indices.append(np.where(conditions_score[night_idx][vis_idx] > 0)[0])
+                schedulable_slot_indices[night_idx] = np.where(conditions_score[night_idx][vis_idx] > 0)[0]
             else:
-                schedulable_slot_indices.append(np.array([]))
+                schedulable_slot_indices[night_idx] = np.array([])
 
         obs_scores = ranker.score_observation(program, obs)
 
-        # Calculate the scores for the observation across all nights across all timeslots.
-        # To avoid the issue of ragged arrays (which are illegal in NumPy 1.24), we must do this night-by-night in
-        # order to end up with a List[npt.NDArray[float]] instead of an npt.NDArray[npt.NDArray[float]].
-        scores = [np.multiply(np.multiply(conditions_score[night_idx], obs_scores[night_idx]), wind_score[night_idx])
-                  for night_idx in night_indices]
+        # Calculate the scores for the observation across all night indices across all timeslots.
+        scores = {night_idx: np.multiply(
+                np.multiply(conditions_score[night_idx], obs_scores[night_idx]),
+                wind_score[night_idx]) for night_idx in night_indices}
 
         # These scores might differ from the observation score in the ranker since they have been adjusted for
         # conditions and wind.
@@ -414,35 +418,38 @@ def _calculate_and_group(self,
         # standards = np.sum([group_info_map[sg.unique_id].standards for sg in group.children])
         standards = 0.
 
-        # The filtering for this group is the product of filtering for the subgroups.
-        sg_night_filtering = [group_data_map[sg.unique_id].group_info.night_filtering for sg in group.children]
-        night_filtering = np.multiply.reduce(sg_night_filtering).astype(bool)
+        # The group is filtered in for a night_idx if all its subgroups are filtered in for that night_idx.
+        night_filtering = {night_idx: all(
+                                  group_data_map[sg.unique_id].group_info.night_filtering[night_idx]
+                                  for sg in group.children
+        ) for night_idx in night_indices}
 
         # The conditions score is the product of the conditions scores for each subgroup across each night.
-        conditions_score = []
+        conditions_score = {}
         for night_idx in night_indices:
             conditions_scores_for_night = [group_data_map[sg.unique_id].group_info.conditions_score[night_idx]
                                            for sg in group.children]
-            conditions_score.append(np.multiply.reduce(conditions_scores_for_night))
+            conditions_score[night_idx] = np.multiply.reduce(conditions_scores_for_night)
 
         # The wind score is the product of the wind scores for each subgroup across each night.
-        wind_score = []
+        wind_score = {}
         for night_idx in night_indices:
             wind_scores_for_night = [group_data_map[sg.unique_id].group_info.wind_score[night_idx]
                                      for sg in group.children]
-            wind_score.append(np.multiply.reduce(wind_scores_for_night))
+            wind_score[night_idx] = np.multiply.reduce(wind_scores_for_night)
 
         # The schedulable slot indices are the unions of the schedulable slot indices for each subgroup
         # across each night.
-        schedulable_slot_indices = [
+        schedulable_slot_indices = {
+            night_idx:
             # For each night, take the concatenation of the schedulable time slots for all children of the group
             # and make it unique, which also puts it in sorted order.
             np.unique(np.concatenate([
                 group_data_map[sg.unique_id].group_info.schedulable_slot_indices[night_idx]
                 for sg in group.children
             ]))
             for night_idx in night_indices
-        ]
+        }
 
         # Calculate the scores for the group across all nights across all timeslots.
         scores = ranker.score_group(group, group_data_map)

scheduler/core/output/__init__.py

@@ -123,7 +123,7 @@ def print_plans(all_plans: List[Plans]) -> None:
     """
 
     for plans in all_plans:
-        print(f'\n\n+++++ NIGHT {plans.night + 1} +++++')
+        print(f'\n\n+++++ NIGHT {plans.night_idx + 1} +++++')
         for plan in plans:
             print(f'Plan for site: {plan.site.name}')
             for visit in plan.visits: