From 36334aa9adb9b4954f94f3eb90bf47edbb9ae015 Mon Sep 17 00:00:00 2001 From: Sebastian Raaphorst Date: Wed, 26 Jul 2023 16:28:27 -1000 Subject: [PATCH] SCHED-406: Changing List[T] to Dict[NightIndex, T] and massive cleanup. --- scheduler/core/calculations/groupinfo.py | 16 +-- scheduler/core/calculations/scores.py | 5 +- .../core/components/collector/__init__.py | 12 +- scheduler/core/components/optimizer/dummy.py | 2 +- .../core/components/optimizer/greedymax.py | 127 ++++++++++-------- .../core/components/optimizer/timeline.py | 6 +- scheduler/core/components/ranker/base.py | 16 +-- scheduler/core/components/ranker/default.py | 30 +++-- .../core/components/selector/__init__.py | 73 +++++----- scheduler/core/output/__init__.py | 2 +- scheduler/core/plans/__init__.py | 19 ++- scheduler/core/service/service.py | 2 +- scheduler/graphql_mid/types.py | 2 +- scheduler/scripts/run_greedymax.py | 9 +- 14 files changed, 170 insertions(+), 151 deletions(-) diff --git a/scheduler/core/calculations/groupinfo.py b/scheduler/core/calculations/groupinfo.py index 03650fb1..d3785fe6 100644 --- a/scheduler/core/calculations/groupinfo.py +++ b/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 diff --git a/scheduler/core/calculations/scores.py b/scheduler/core/calculations/scores.py index e23d4b15..79f34d3a 100644 --- a/scheduler/core/calculations/scores.py +++ b/scheduler/core/calculations/scores.py @@ -1,8 +1,9 @@ # 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. @@ -10,4 +11,4 @@ # Scores across all nights per timeslot. # Indexed by night index, and then timeslot index. -Scores = List[NightTimeSlotScores] +Scores = Dict[NightIndex, NightTimeSlotScores] diff --git a/scheduler/core/components/collector/__init__.py b/scheduler/core/components/collector/__init__.py index 823fa772..c4fff2f5 100644 --- a/scheduler/core/components/collector/__init__.py +++ b/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} diff --git a/scheduler/core/components/optimizer/dummy.py b/scheduler/core/components/optimizer/dummy.py index 6fae9326..823c99c0 100644 --- a/scheduler/core/components/optimizer/dummy.py +++ b/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: diff --git a/scheduler/core/components/optimizer/greedymax.py b/scheduler/core/components/optimizer/greedymax.py index d52686fe..4d78dad3 100644 --- a/scheduler/core/components/optimizer/greedymax.py +++ b/scheduler/core/components/optimizer/greedymax.py @@ -5,16 +5,17 @@ from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Dict, FrozenSet, List, Optional, Tuple, Union +from typing import Dict, FrozenSet, List, Optional, Tuple from scheduler.core.calculations.selection import Selection from scheduler.core.calculations import GroupData from scheduler.core.plans import Plan, Plans from scheduler.core.components.optimizer.timeline import Timelines +from scheduler.services import logger_factory from .base import BaseOptimizer, MaxGroup from . import Interval -from lucupy.minimodel import Program, Group, Observation, Sequence +from lucupy.minimodel import Group, NightIndex, Observation, Program, Sequence from lucupy.minimodel import ObservationID, Site, UniqueGroupID, QAState, ObservationClass, ObservationStatus from lucupy.minimodel.resource import Resource import numpy as np @@ -22,6 +23,9 @@ import matplotlib.pyplot as plt +logger = logger_factory.create_logger(__name__) + + @dataclass(frozen=True) class ObsPlanData: """ @@ -43,7 +47,7 @@ def __init__(self, min_visit_len: timedelta = timedelta(minutes=30), show_plots: self.group_data_list: List[GroupData] = [] self.group_ids: List[UniqueGroupID] = [] self.obs_group_ids: List[UniqueGroupID] = [] - self.timelines: List[Timelines] = [] + self.timelines: Dict[NightIndex, Timelines] = {} self.sites: FrozenSet[Site] = frozenset() self.obs_in_plan: Dict = {} self.min_visit_len = min_visit_len @@ -59,7 +63,8 @@ def setup(self, selection: Selection) -> GreedyMaxOptimizer: self.group_data_list = list(selection.schedulable_groups.values()) # self._process_group_data(self.group_data_list) self.obs_group_ids = list(selection.obs_group_ids) # noqa - self.timelines = [Timelines(selection.night_events, night_idx) for night_idx in selection.night_indices] + self.timelines = {night_idx: Timelines(selection.night_events, night_idx) + for night_idx in selection.night_indices} self.sites = selection.sites self.time_slot_length = selection.time_slot_length for site in self.sites: @@ -261,7 +266,7 @@ def _find_max_group(self, plans: Plans) -> Optional[MaxGroup]: only_first_interval = False # Get the unscheduled, available intervals (time slots) - open_intervals = {site: self.timelines[plans.night][site].get_available_intervals(only_first_interval) + open_intervals = {site: self.timelines[plans.night_idx][site].get_available_intervals(only_first_interval) for site in self.sites} max_scores = [] @@ -275,13 +280,13 @@ def _find_max_group(self, plans: Plans) -> Optional[MaxGroup]: # Make a list of scores in the remaining groups for group_data in self.group_data_list: site = group_data.group.observations()[0].site - if not self.timelines[plans.night][site].is_full: + if not self.timelines[plans.night_idx][site].is_full: for interval_idx, interval in enumerate(open_intervals[site]): # print(f'Interval: {iint}') # scores = group_data.group_info.scores[plans.night] # Get the maximum score over the interval. - smax = np.max(group_data.group_info.scores[plans.night][interval]) + smax = np.max(group_data.group_info.scores[plans.night_idx][interval]) if smax > 0.0: # Check if the interval is long enough to be useful (longer than min visit length). # Remaining time for the group. @@ -292,14 +297,14 @@ def _find_max_group(self, plans: Plans) -> Optional[MaxGroup]: # Evaluate sub-intervals (e.g. timing windows, gaps in the score). # Find time slot locations where the score > 0. # interval is a numpy array that indexes into the scores for the night to return a sub-array. - check_interval = group_data.group_info.scores[plans.night][interval] + check_interval = group_data.group_info.scores[plans.night_idx][interval] group_intervals = self.non_zero_intervals(check_interval) max_score_on_interval = 0.0 max_interval = None for group_interval in group_intervals: grp_interval_length = group_interval[1] - group_interval[0] - max_score = np.max(group_data.group_info.scores[plans.night] + max_score = np.max(group_data.group_info.scores[plans.night_idx] [interval[group_interval[0]:group_interval[1]]]) # Find the max_score in the group intervals with non-zero scores @@ -378,14 +383,10 @@ def _find_max_group(self, plans: Plans) -> Optional[MaxGroup]: return max_group_info @staticmethod - def _integrate_score(night_idx: int, + def _integrate_score(night_idx: NightIndex, max_group_info: MaxGroup) -> Interval: - """Use the score array to find the best location in the timeline - - group_data: Group data of group with maximum score - interval: the timeline interval, where to place group_data - n_time: length of the group in time steps - night: night counter + """ + Use the score array to find the best location in the timeline """ start = max_group_info.interval[0] end = max_group_info.interval[-1] @@ -430,7 +431,7 @@ def _integrate_score(night_idx: int, return best_interval - def _find_group_position(self, night_idx: int, max_group_info: MaxGroup) -> Interval: + def _find_group_position(self, night_idx: NightIndex, max_group_info: MaxGroup) -> Interval: """Find the best location in the timeline""" best_interval = max_group_info.interval @@ -489,18 +490,23 @@ def nir_slots(self, science_obs, n_slots_filled, len_interval) -> Tuple[int, int return slot_start_nir, slot_end_nir, obs_id_nir - def mean_airmass(self, obsid, interval, night=0): + def mean_airmass(self, obs_id: ObservationID, interval: Interval, night_idx: NightIndex): """ Calculate the mean airmass of an observation over the given interval """ - programid = obsid.program_id() + programid = obs_id.program_id() # print(obsid.id, programid.id) - airmass = self.selection.program_info[programid].target_info[obsid][night].airmass[interval] + airmass = self.selection.program_info[programid].target_info[obs_id][night_idx].airmass[interval] return np.mean(airmass) - def place_standards(self, night, interval, science_obs, partner_obs, n_std, verbose: bool = True) \ - -> Tuple[List, List]: + def place_standards(self, + night_idx: NightIndex, + interval: Interval, + science_obs, + partner_obs, + n_std, + verbose: bool = True) -> Tuple[List, List]: """ Pick the standards that best match the NIR science observations by airmass """ @@ -527,7 +533,7 @@ def place_standards(self, night, interval, science_obs, partner_obs, n_std, verb slot_start = n_slots_acq slot_end = n_slots_cal - 1 - xmean_cal = self.mean_airmass(partcal_obs.id, interval[slot_start:slot_end + 1], night=night) + xmean_cal = self.mean_airmass(partcal_obs.id, interval[slot_start:slot_end + 1], night_idx=night_idx) if verbose: print(f'Standard {partcal_obs.id.id}') @@ -544,7 +550,7 @@ def place_standards(self, night, interval, science_obs, partner_obs, n_std, verb slot_start_nir = slot_end + idx_start_nir slot_end_nir = slot_end + idx_end_nir - xmean_nir = self.mean_airmass(obs_id_nir, interval[slot_start_nir:slot_end_nir + 1], night=night) + xmean_nir = self.mean_airmass(obs_id_nir, interval[slot_start_nir:slot_end_nir + 1], night_idx=night_idx) xdiff_before = np.abs(xmean_nir - xmean_cal) if verbose: @@ -563,7 +569,7 @@ def place_standards(self, night, interval, science_obs, partner_obs, n_std, verb slot_start = len_int - 1 - n_slots_cal + n_slots_acq slot_end = slot_start + n_slots_cal - n_slots_acq - 1 - xmean_cal = self.mean_airmass(partcal_obs.id, interval[slot_start:slot_end + 1], night=night) + xmean_cal = self.mean_airmass(partcal_obs.id, interval[slot_start:slot_end + 1], night_idx=night_idx) if verbose: print(f'\n\t Try std after') @@ -577,7 +583,7 @@ def place_standards(self, night, interval, science_obs, partner_obs, n_std, verb slot_start_nir = idx_start_nir slot_end_nir = idx_end_nir - xmean_nir = self.mean_airmass(obs_id_nir, interval[slot_start_nir:slot_end_nir + 1], night=night) + xmean_nir = self.mean_airmass(obs_id_nir, interval[slot_start_nir:slot_end_nir + 1], night_idx=night_idx) xdiff_after = np.abs(xmean_nir - xmean_cal) if verbose: @@ -649,17 +655,13 @@ def _charge_time(observation: Observation, atom_start: int = 0, atom_end: int = observation.sequence[n_atom].observed = True observation.sequence[n_atom].qa_state = QAState.PASS - def plot_airmass(self, obsid, interval=None, night=0) -> None: + def plot_airmass(self, obs_id: ObservationID, interval=None, night_idx: NightIndex = 0) -> None: """ Plot airmass vs time slot - :param obsid: - :param interval: - :param night: - :return: None """ - programid = obsid.program_id() + programid = obs_id.program_id() - airmass = self.selection.program_info[programid].target_info[obsid][night].airmass + airmass = self.selection.program_info[programid].target_info[obs_id][night_idx].airmass x = np.array([i for i in range(len(airmass))], dtype=int) p = plt.plot(airmass) @@ -667,7 +669,7 @@ def plot_airmass(self, obsid, interval=None, night=0) -> None: if interval is not None: plt.plot(x[interval], airmass[interval], color=colour, linewidth=4) plt.ylim(2.5, 0.95) - plt.title(obsid.id) + plt.title(obs_id.id) plt.xlabel('Time Slot') plt.ylabel('Airmass') plt.show() @@ -736,7 +738,7 @@ def plot_timelines(self, night: int = 0) -> None: plt.show() - def _update_score(self, program: Program, night=0) -> None: + def _update_score(self, program: Program, night_idx: NightIndex) -> None: """Update the scores of the incomplete groups in the scheduled program""" print("Running score_program") @@ -748,19 +750,20 @@ def _update_score(self, program: Program, night=0) -> None: group, group_info = group_data schedulable_group = self.selection.schedulable_groups[unique_group_id] print(f"{unique_group_id} {schedulable_group.group.exec_time()} {schedulable_group.group.total_used()}") - print(f"\tOld max score: {np.max(schedulable_group.group_info.scores[night]):7.2f} new max score[0]: " - f"{np.max(group_info.scores[night]):7.2f}") + print(f"\tOld max score: {np.max(schedulable_group.group_info.scores[night_idx]):7.2f} new max score[0]: " + f"{np.max(group_info.scores[night_idx]):7.2f}") # 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[:] - print(f"\tUpdated max score: {np.max(schedulable_group.group_info.scores[night]):7.2f}") + schedulable_group.group_info.scores = group_info.scores + # schedulable_group.group_info.scores[:] = group_info.scores[:] + print(f"\tUpdated max score: {np.max(schedulable_group.group_info.scores[night_idx]):7.2f}") def _run(self, plans: Plans) -> None: # Fill plans for all sites on one night - while not self.timelines[plans.night].all_done() and len(self.group_data_list) > 0: + while not self.timelines[plans.night_idx].all_done() and len(self.group_data_list) > 0: - print(f"\nNight {plans.night + 1}") + print(f"\nNight {plans.night_idx + 1}") # Find the group with the max score in an open interval max_group_info = self._find_max_group(plans) @@ -768,7 +771,7 @@ def _run(self, plans: Plans) -> None: # If something found, add it to the timeline and plan if max_group_info is not None: # max_score, max_group, max_interval = max_group_info - added = self.add(plans.night, max_group_info) + added = self.add(plans.night_idx, max_group_info) if added: print(f"Group {max_group_info.group_data.group.unique_id.id} with " f"max score {max_group_info.max_score} added.") @@ -779,20 +782,27 @@ def _run(self, plans: Plans) -> None: # Nothing remaining can be scheduled # for plan in plans: # plan.is_full = True - for timeline in self.timelines[plans.night]: + # TODO NOTE: Does this really mean the timeline is full? + for timeline in self.timelines[plans.night_idx]: + logger.warning(f'Setting timelines corresponding to {plans.night_idx} to full (no max_group_info).') timeline.is_full = True if self.show_plots: - self.plot_timelines(plans.night) + self.plot_timelines(plans.night_idx) # Write observations from the timelines to the output plan self.output_plans(plans) - def _add_visit(self, night, obs, max_group_info, best_interval, n_slots_filled) -> int: + def _add_visit(self, + night_idx: NightIndex, + obs: Observation, + max_group_info: GroupData | MaxGroup, + best_interval, + n_slots_filled) -> int: """Add and observation to the timeline and do pseudo-time accounting""" site = max_group_info.group_data.group.observations()[0].site - timeline = self.timelines[night][site] + timeline = self.timelines[night_idx][site] program = self.selection.program_info[max_group_info.group_data.group.program_id].program iobs = self.obs_group_ids.index(obs.to_unique_group_id) @@ -813,7 +823,7 @@ def _add_visit(self, night, obs, max_group_info, best_interval, n_slots_filled) start_time_slot, start = timeline.add(iobs, visit_length, best_interval) # Get visit score and store information for the output plans - visit_score = sum(max_group_info.group_data.group_info.scores[night][start_time_slot:start_time_slot+visit_length]) + visit_score = sum(max_group_info.group_data.group_info.scores[night_idx][start_time_slot:start_time_slot + visit_length]) self.obs_in_plan[site][start_time_slot] = ObsPlanData( obs=obs, obs_start=start, @@ -829,7 +839,7 @@ def _add_visit(self, night, obs, max_group_info, best_interval, n_slots_filled) return n_slots_filled - def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: + def add(self, night_idx: NightIndex, max_group_info: GroupData | MaxGroup) -> bool: """ Add a group to a Plan - find the best location within the interval (maximize the score) and select standards """ @@ -841,7 +851,7 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: # to place the group in the timeline site = max_group_info.group_data.group.observations()[0].site - timeline = self.timelines[night][site] + timeline = self.timelines[night_idx][site] program = self.selection.program_info[max_group_info.group_data.group.program_id].program # visit = [] # list of observations in visit result = False @@ -850,12 +860,12 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: if not timeline.is_full: # Find the best location in timeline for the group - best_interval = self._find_group_position(night, max_group_info) + best_interval = self._find_group_position(night_idx, max_group_info) if self.show_plots: - self._plot_interval(max_group_info.group_data.group_info.scores[night], max_group_info.interval, + self._plot_interval(max_group_info.group_data.group_info.scores[night_idx], max_group_info.interval, best_interval, - label=f'Night {night + 1}: {max_group_info.group_data.group.unique_id.id}') + label=f'Night {night_idx + 1}: {max_group_info.group_data.group.unique_id.id}') # When/If we eventually support splitting NIR observations, then we need to calculate the # NIR science time in best_interval and the number of basecal (e.g. telluric standards) needed. @@ -871,7 +881,7 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: if max_group_info.n_std > 0: if max_group_info.exec_sci_nir > timedelta(0): - standards, place_before = self.place_standards(night, best_interval, prog_obs, part_obs, + standards, place_before = self.place_standards(night_idx, best_interval, prog_obs, part_obs, max_group_info.n_std) for ii, std in enumerate(standards): n_slots_cal += Plan.time2slots(self.time_slot_length, std.exec_time()) @@ -887,14 +897,14 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: if before_std is not None: obs = before_std print(f"Adding before_std: {obs.to_unique_group_id} {obs.id.id}") - n_slots_filled = self._add_visit(night, obs, max_group_info, best_interval, n_slots_filled) + n_slots_filled = self._add_visit(night_idx, obs, max_group_info, best_interval, n_slots_filled) # split at atoms for obs in prog_obs: # Reserve space for the cals, otherwise the science observes will fill the interval n_slots_filled = n_slots_cal print(f"Adding science: {obs.to_unique_group_id} {obs.id.id}") - n_slots_filled = self._add_visit(night, obs, max_group_info, best_interval, n_slots_filled) + n_slots_filled = self._add_visit(night_idx, obs, max_group_info, best_interval, n_slots_filled) if after_std is not None: # "put back" time for the final standard n_slots_filled -= Plan.time2slots(self.time_slot_length, standards[-1].exec_time()) @@ -902,14 +912,15 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: if after_std is not None: obs = after_std print(f"Adding after_std: {obs.to_unique_group_id} {obs.id.id}") - n_slots_filled = self._add_visit(night, obs, max_group_info, best_interval, n_slots_filled) + n_slots_filled = self._add_visit(night_idx, obs, max_group_info, best_interval, n_slots_filled) # TODO: Shift to remove any gaps in the plan? # Re-score program (pseudo time accounting) - self._update_score(program, night=night) + self._update_score(program, night_idx=night_idx) if timeline.slots_unscheduled() <= 0: + logger.warning(f'Timeline for {night_idx} is full: no slots remain unscheduled.') timeline.is_full = True result = True @@ -919,7 +930,7 @@ def add(self, night: int, max_group_info: Union[GroupData, MaxGroup]) -> bool: def output_plans(self, plans: Plans) -> None: """Write visit information from timelines to output plans, ensures chronological order""" - for timeline in self.timelines[plans.night]: + for timeline in self.timelines[plans.night_idx]: obs_order = timeline.get_observation_order() for idx, start_time_slot, end_time_slot in obs_order: if idx > -1: diff --git a/scheduler/core/components/optimizer/timeline.py b/scheduler/core/components/optimizer/timeline.py index 49941bc0..a15cfeb1 100644 --- a/scheduler/core/components/optimizer/timeline.py +++ b/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], diff --git a/scheduler/core/components/ranker/base.py b/scheduler/core/components/ranker/base.py index ed01d9d9..b33e33bd 100644 --- a/scheduler/core/components/ranker/base.py +++ b/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: """ diff --git a/scheduler/core/components/ranker/default.py b/scheduler/core/components/ranker/default.py index 01d3ca54..c2d18559 100644 --- a/scheduler/core/components/ranker/default.py +++ b/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 diff --git a/scheduler/core/components/selector/__init__.py b/scheduler/core/components/selector/__init__.py index 0410587a..b8b201b7 100644 --- a/scheduler/core/components/selector/__init__.py +++ b/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,27 +418,30 @@ 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([ @@ -442,7 +449,7 @@ def _calculate_and_group(self, 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) diff --git a/scheduler/core/output/__init__.py b/scheduler/core/output/__init__.py index 952956c0..ec07304f 100644 --- a/scheduler/core/output/__init__.py +++ b/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: diff --git a/scheduler/core/plans/__init__.py b/scheduler/core/plans/__init__.py index 90076d07..12eeac97 100644 --- a/scheduler/core/plans/__init__.py +++ b/scheduler/core/plans/__init__.py @@ -6,7 +6,7 @@ from math import ceil from typing import List, Mapping, Optional, Tuple -from lucupy.minimodel import Observation, ObservationID, Site, Conditions, Band, Resource +from lucupy.minimodel import Band, Conditions, NightIndex, Observation, ObservationID, Resource, Site import numpy as np import numpy.typing as npt @@ -108,7 +108,7 @@ def _place_standards(self, science_obs: List[Observation], partner_obs: List[Observation], target_info: TargetInfoNightIndexMap, - night: int, + night_idx: NightIndex, n_std) -> Tuple[List, List]: """Pick the standards that best match the NIR science observations by airmass @@ -117,8 +117,6 @@ def _place_standards(self, standards = [] placement = [] - # print(f'Running place_standards') - xdiff_min = xdiff_before_min = xdiff_after_min = 99999. std_before = None std_after = None @@ -135,14 +133,14 @@ def _place_standards(self, slot_start = n_slots_acq slot_end = n_slots_cal - 1 - xmean_cal = target_info[night][partcal_obs.id].mean_airmass(interval[slot_start:slot_end + 1]) + xmean_cal = target_info[night_idx][partcal_obs.id].mean_airmass(interval[slot_start:slot_end + 1]) # Mean NIR science airmass idx_start_nir, idx_end_nir, obs_id_nir = self.nir_slots(science_obs, n_slots_cal, len(interval)) slot_start_nir = slot_end + idx_start_nir slot_end_nir = slot_end + idx_end_nir - xmean_nir = target_info[night][obs_id_nir].mean_airmass(interval[slot_start_nir:slot_end_nir + 1]) + xmean_nir = target_info[night_idx][obs_id_nir].mean_airmass(interval[slot_start_nir:slot_end_nir + 1]) xdiff_before = np.abs(xmean_nir - xmean_cal) # Try std last @@ -151,13 +149,14 @@ def _place_standards(self, slot_start = len_int - 1 - n_slots_cal + n_slots_acq slot_end = slot_start + n_slots_cal - n_slots_acq - 1 - xmean_cal = target_info[night][partcal_obs.id][night].mean_airmass(interval[slot_start:slot_end + 1]) + xmean_cal = (target_info[night_idx][partcal_obs.id][night_idx] + .mean_airmass(interval[slot_start:slot_end + 1])) # Mean NIR science airmass slot_start_nir = idx_start_nir slot_end_nir = idx_end_nir - xmean_nir = target_info[night][obs_id_nir].mean_airmass(interval[slot_start_nir:slot_end_nir + 1]) + xmean_nir = target_info[night_idx][obs_id_nir].mean_airmass(interval[slot_start_nir:slot_end_nir + 1]) xdiff_after = np.abs(xmean_nir - xmean_cal) if n_std == 1: @@ -222,9 +221,9 @@ class Plans: A collection of Plan 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.plans = {} - self.night = night_idx + self.night_idx = night_idx for site, ne in night_events.items(): if ne is not None: self.plans[site] = Plan(ne.local_times[night_idx][0], diff --git a/scheduler/core/service/service.py b/scheduler/core/service/service.py index 6f3b6e0f..cbea56c6 100644 --- a/scheduler/core/service/service.py +++ b/scheduler/core/service/service.py @@ -129,7 +129,7 @@ def calculate_plans_stats(all_plans: List[Plans], # Calculate altitude data ti = collector.get_target_info(visit.obs_id) end_time_slot = visit.start_time_slot + visit.time_slots - values = ti[plans.night].alt[visit.start_time_slot: end_time_slot] + values = ti[plans.night_idx].alt[visit.start_time_slot: end_time_slot] alt_degs = [val.dms[0] + (val.dms[1]/60) + (val.dms[2]/3600) for val in values] plan.alt_degs.append(alt_degs) diff --git a/scheduler/graphql_mid/types.py b/scheduler/graphql_mid/types.py index f094349e..a63eb465 100644 --- a/scheduler/graphql_mid/types.py +++ b/scheduler/graphql_mid/types.py @@ -100,7 +100,7 @@ class SPlans: @staticmethod def from_computed_plans(plans: Plans, sites: FrozenSet[Site]) -> 'SPlans': return SPlans( - night_idx=plans.night, + night_idx=plans.night_idx, plans_per_site=[SPlan.from_computed_plan(plans[site]) for site in sites]) def for_site(self, site: Site) -> 'SPlans': diff --git a/scheduler/scripts/run_greedymax.py b/scheduler/scripts/run_greedymax.py index 07826330..6d7d379c 100644 --- a/scheduler/scripts/run_greedymax.py +++ b/scheduler/scripts/run_greedymax.py @@ -49,7 +49,10 @@ # Execute the Selector. # Not sure the best way to display the output. selector = SchedulerBuilder.build_selector(collector, num_nights_to_schedule=3) - selection = selector.select() + + # TODO: Loop here on num_nights_to_schedule with select, schedule, and time accounting. + selection = selector.select(night_indices=np.array([0])) + # selection = selector.select() # Notes for data access: # The Selector returns all the data that an Optimizer needs in order to generate plans. @@ -169,8 +172,8 @@ print('') # Timeline tests - for tl in optimizer_blueprint.algorithm.timelines: - print(f'Night {tl.night + 1}') + for tl in optimizer_blueprint.algorithm.timelines.values(): + print(f'Night {tl.night_idx + 1}') # for site, ne in gm_optimizer.night_events.items(): for site in optimizer.night_events.keys(): print(f'\t {site}')