Coverage for o2/models/state.py: 85%

1import datetime 

2import xml.etree.ElementTree as ET 

3from dataclasses import dataclass, replace 

4from typing import TYPE_CHECKING, Any 

5 

6import pytz 

7from prosimos.simulation_setup import SimDiffSetup 

8 

9from o2.models.evaluation import Evaluation 

10from o2.models.settings import Settings 

11from o2.simulation_runner import SimulationRunner 

12from o2.util.logger import warn 

13from o2.util.sim_diff_setup_fileless import SimDiffSetupFileless 

14 

15if TYPE_CHECKING: 

16 from o2.models.timetable import TimetableType 

17 

18 

19@dataclass(frozen=True) 

20class State: 

21 """A state in the optimization process. 

22 

23 It's mainly a container class for the timetable. 

24 

25 The State should contain all information to run a simulation -- that's why it 

26 contains the BPMN definition as well. 

27 """ 

28 

29 bpmn_definition: str 

30 

31 timetable: "TimetableType" 

32 # TODO: Move to setting class 

33 for_testing: bool = False 

34 

35 def replace_timetable(self, /, **changes: Any) -> "State": # noqa: ANN401 

36 """Replace the timetable with the given changes.""" 

37 return replace(self, timetable=replace(self.timetable, **changes)) 

38 

39 def evaluate(self) -> Evaluation: 

40 """Evaluate the current state.""" 

41 if not self.is_valid(): 

42 warn("Trying to evaluate an invalid state.") 

43 return Evaluation.empty() 

44 try: 

45 if Settings.USE_MEDIAN_SIMULATION_FOR_EVALUATION: 

46 result = SimulationRunner.run_simulation_median(self) 

47 else: 

48 result = SimulationRunner.run_simulation(self) 

49 except Exception as e: 

50 if Settings.RAISE_SIMULATION_ERRORS: 

51 raise e 

52 return Evaluation.empty() 

53 return Evaluation.from_run_simulation_result( 

54 self.timetable.get_hourly_rates(), 

55 self.timetable.get_fixed_cost_fns(), 

56 self.timetable.batching_rules_exist, 

57 result, 

58 ) 

59 

60 def to_sim_diff_setup(self) -> SimDiffSetup: 

61 """Convert the state to a SimDiffSetup.""" 

62 setup = SimDiffSetupFileless( 

63 "test", 

64 self.bpmn_definition, 

65 self.timetable, 

66 False, 

67 self.timetable.total_cases, 

68 ) 

69 if self.for_testing: 

70 # For testing we start on 03.01.2000, a Monday 

71 starting_at_datetime = pytz.utc.localize(datetime.datetime(2000, 1, 3)) 

72 else: 

73 starting_at_datetime = pytz.utc.localize(datetime.datetime.now()) 

74 

75 setup.set_starting_datetime(starting_at_datetime) 

76 

77 return setup 

78 

79 def get_name_of_task(self, task_id: str) -> str: 

80 """Get the name of a task.""" 

81 bpmn_tree = ET.fromstring(self.bpmn_definition) 

82 node = bpmn_tree.find(f".//*[@id='{task_id}']") 

83 assert node is not None 

84 return node.attrib["name"] 

85 

86 def get_task_names(self) -> dict[str, str]: 

87 """Get a mapping of task IDs to task names.""" 

88 bpmn_tree = ET.fromstring(self.bpmn_definition) 

89 task_nodes = bpmn_tree.findall(".//*[@id]") 

90 return {node.attrib["id"]: node.attrib.get("name", node.attrib["id"]) for node in task_nodes} 

91 

92 def is_valid(self) -> bool: 

93 """Check if the state is valid.""" 

94 return self.timetable.is_valid() 

95 

96 

97class TabuState(State): 

98 """A state that is tabu.""" 

99 

100 def __init__(self) -> None: 

101 """Initialize the tabu state.""" 

102 

103 def is_valid(self) -> bool: 

104 """Check if the state is valid.""" 

105 return False