SCHED-404: First cleanup

requirements.txt

@@ -2,7 +2,7 @@ astropy
 hypothesis
 matplotlib
 more-itertools
-numpy>=1.23.*
+numpy>=1.23.*, <1.24
 openpyxl
 pandas
 pytest

scheduler/core/calculations/targetinfo.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from typing import Dict, Tuple
 
+import numpy as np
 import numpy.typing as npt
 from astropy.coordinates import Angle, SkyCoord
 from astropy.time import TimeDelta
@@ -52,6 +53,9 @@ class TargetInfo:
     rem_visibility_time: TimeDelta
     rem_visibility_frac: float
 
+    def mean_airmass(self, interval: npt.NDArray[int]):
+        return np.mean(self.airmass[interval])
+
 
 # Type aliases for TargetInfo information.
 TargetInfoNightIndexMap = Dict[NightIndex, TargetInfo]

scheduler/core/components/optimizer/__init__.py

@@ -7,6 +7,7 @@
 from scheduler.core.plans import Plans
 
 import numpy.typing as npt
+from lucupy.minimodel import Program
 
 # Convenient type alias for Interval
 Interval = npt.NDArray[int]
@@ -24,6 +25,7 @@ def __init__(self, selection: Selection, algorithm=None):
         # TODO: Assumes that all sites schedule the same amount of nights
         # if num_nights_optimize is None:
         self.period = len(list(self.night_events.values())[0].time_grid)
+        self.selection: Selection = selection
         # else:
         #     self.period = num_nights_optimize
 
@@ -32,3 +34,15 @@ def schedule(self) -> List[Plans]:
         nights = [Plans(self.night_events, night) for night in range(self.period)]
         self.algorithm.schedule(nights)
         return nights
+
+    def _update_score(self, program: Program) -> None:
+        """Update the scores of the incomplete groups in the scheduled program"""
+        program_calculations = self.selection.score_program(program)
+
+        for unique_group_id in program_calculations.top_level_groups:
+            group_data = program_calculations.group_data_map[unique_group_id]
+            group, group_info = group_data
+            schedulable_group = self.selection.schedulable_groups[unique_group_id]
+            # update scores in schedulable_groups if the group is not completely observed
+            if schedulable_group.group.exec_time() >= schedulable_group.group.total_used():
+                schedulable_group.group_info.scores[:] = group_info.scores[:]

scheduler/core/components/optimizer/base.py

@@ -2,9 +2,9 @@
 # For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
 
 from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import Mapping, List, Optional, Tuple
-
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Mapping, List, Optional, Union
 from lucupy.minimodel.program import ProgramID
 
 from scheduler.core.calculations.groupinfo import GroupData
@@ -14,6 +14,20 @@
 from . import Interval
 
 
+@dataclass(frozen=True)
+class MaxGroup:
+    """
+    Store information about the selected group (max score)
+    """
+    group_data: GroupData
+    max_score: float
+    interval: Interval
+    n_min: int
+    n_slots_remaining: int
+    n_std: int
+    exec_sci_nir: timedelta
+
+
 class BaseOptimizer(ABC):
     """
     Base class for all Optimizer components.
@@ -40,5 +54,6 @@ def setup(self, program_info: Mapping[ProgramID, ProgramInfo]):
         ...
 
     @abstractmethod
-    def add(self, group: GroupData, plans: Plans, interval: Optional[Interval] = None):
+    def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]):
         ...
+