Coverage for middle_layer/allocate/domain_layer/services/prioritize_osd_service.py: 69.39%

1# Copyright 2024 Associated Universities, Inc. 

2# 

3# This file is part of Telescope Time Allocation Tools (TTAT). 

4# 

5# TTAT is free software: you can redistribute it and/or modify 

6# it under the terms of the GNU General Public License as published by 

7# the Free Software Foundation, either version 3 of the License, or 

8# any later version. 

9# 

10# TTAT is distributed in the hope that it will be useful, 

11# but WITHOUT ANY WARRANTY; without even the implied warranty of 

12# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

13# GNU General Public License for more details. 

14# 

15# You should have received a copy of the GNU General Public License 

16# along with TTAT. If not, see <https://www.gnu.org/licenses/>. 

17import itertools 

18import logging 

19from collections import defaultdict 

20from enum import Enum 

21from random import choice 

22from typing import Callable, Iterator 

23 

24from allocate.domain_layer.entities.allocation_version import AllocationVersion 

25from allocate.domain_layer.entities.available_time_model_version import AvailableTimeModelVersion 

26from allocate.domain_layer.entities.observation_specification_disposition import ObservationSpecificationDisposition 

27 

28AVAILABLE_PRIORITY_NAMES = ["N", "A", "B", "C", "D"] 

29 

30 

31Prioritizer = Callable[[Iterator[ObservationSpecificationDisposition]], None] 

32""" 

33Prioritizers are functions that work on lists of OSDs 

34""" 

35 

36 

37def gbt_algorithm(osds: Iterator[ObservationSpecificationDisposition]): 

38 logging.info("Running the GBT prioritization algorithm") 

39 for osd in osds: 

40 osd.scheduling_priority_name = "N" 

41 

42 

43def vla_algorithm(osds: Iterator[ObservationSpecificationDisposition]): 

44 logging.info("Running the VLA prioritization algorithm") 

45 for osd in osds: 

46 osd.scheduling_priority_name = "N" 

47 

48 

49def vlba_algorithm(osds: Iterator[ObservationSpecificationDisposition]): 

50 logging.info("Running the VLBA prioritization algorithm") 

51 for osd in osds: 

52 osd.scheduling_priority_name = "N" 

53 

54 

55def unknown_facility_algorithm(osds: Iterator[ObservationSpecificationDisposition]): 

56 logging.info("Running the prioritization algorithm for an unknown facility") 

57 for osd in osds: 

58 osd.scheduling_priority_name = "N" 

59 

60 

61FACILITY_ALGORITHMS: defaultdict[str, Prioritizer] = defaultdict( 

62 lambda: unknown_facility_algorithm, GBT=gbt_algorithm, VLA=vla_algorithm, VLBA=vlba_algorithm 

63) 

64 

65 

66class PrioritizerMode(Enum): 

67 ALGORITHM = 0 # Use defined algorithm 

68 ALL_N = 1 # All OSDs assigned N priority 

69 RANDOM = 2 # Random priority assignment 

70 

71 def prioritize(self, facility: str, osds: Iterator[ObservationSpecificationDisposition]): 

72 match self: 

73 case self.ALGORITHM: 

74 FACILITY_ALGORITHMS[facility](osds) 

75 

76 case self.ALL_N: 

77 for osd in osds: 

78 osd.scheduling_priority_name = "N" 

79 

80 case self.RANDOM: 

81 for osd in osds: 

82 osd.scheduling_priority_name = choice(AVAILABLE_PRIORITY_NAMES) 

83 

84 

85def prioritize_osds( 

86 av: AllocationVersion, atmv: AvailableTimeModelVersion, mode: PrioritizerMode = PrioritizerMode.ALGORITHM 

87) -> tuple[AllocationVersion, AvailableTimeModelVersion]: 

88 """Prioritize observation specification dispositions 

89 :param av: AllocationVersion containing OSDs to prioritize 

90 :param atmv: AvailableTimeModelVersion for prioritization 

91 :param mode: Prioritization mode to use 

92 :return: Allocation Version with prioritized OSDs set to read-only 

93 """ 

94 

95 # get all the OSDs that are not locked with priority NP, grouped by facility 

96 osds_by_facility = itertools.groupby( 

97 ( 

98 osd 

99 for ad in av.allocation_dispositions 

100 for osd in ad.observation_specification_dispositions 

101 if osd.scheduling_priority_name == "NP" and not osd.scheduling_priority_locked 

102 ), 

103 lambda osd: osd.facility.facility_name, 

104 ) 

105 

106 # run the appropriate prioritizer for the mode and facility 

107 for facility, osds in osds_by_facility: 

108 mode.prioritize(facility, osds) 

109 

110 av.is_read_only = atmv.is_read_only = True 

111 return av, atmv