Coverage for agentlib_flexquant/modules/shadow_mpc.py: 62%

1from typing import Dict, Union 

2from agentlib.core.datamodels import AgentVariable 

3from agentlib_mpc.modules import mpc_full, minlp_mpc 

4from agentlib_flexquant.utils.data_handling import strip_multi_index, fill_nans, MEAN 

5from agentlib_flexquant.data_structures.globals import full_trajectory_prefix, full_trajectory_suffix 

6 

7 

8class FlexibilityShadowMPC(mpc_full.MPC): 

9 

10 config: mpc_full.MPCConfig 

11 

12 def __init__(self, *args, **kwargs): 

13 # create instance variable 

14 self._full_controls: Dict[str, Union[AgentVariable, None]] = {} 

15 super().__init__(*args, **kwargs) 

16 

17 def register_callbacks(self): 

18 for control_var in self.config.controls: 

19 self.agent.data_broker.register_callback( 

20 name=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}", 

21 alias=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}", 

22 callback=self.calc_flex_callback, 

23 ) 

24 for input_var in self.config.inputs: 

25 if input_var.name.replace(full_trajectory_prefix, "", 1).replace( 

26 full_trajectory_suffix, "" 

27 ) in [control_var.name for control_var in self.config.controls]: 

28 self._full_controls[input_var.name] = input_var 

29 

30 super().register_callbacks() 

31 

32 def calc_flex_callback(self, inp: AgentVariable, name: str): 

33 """Set the control trajectories before calculating the flexibility offer. 

34 

35 self.model should account for flexibility in its cost function. 

36 

37 """ 

38 # during provision dont calculate flex 

39 if self.get("in_provision").value: 

40 return 

41 

42 # do not trigger callback on self set variables 

43 if self.agent.config.id == inp.source.agent_id: 

44 return 

45 

46 vals = strip_multi_index(inp.value) 

47 if vals.isna().any(): 

48 vals = fill_nans(series=vals, method=MEAN) 

49 

50 # the MPC Predictions starts at t=env.now not t=0 

51 vals.index += self.env.time 

52 self._full_controls[name].value = vals 

53 self.set(name, vals) 

54 # make sure all controls are set 

55 if all(x.value is not None for x in self._full_controls.values()): 

56 self.do_step() 

57 for name in self._full_controls.keys(): 

58 self._full_controls[name].value = None 

59 

60 def process(self): 

61 # the shadow mpc should only be run after the results of the baseline are sent 

62 yield self.env.event() 

63 

64 

65class FlexibilityShadowMINLPMPC(minlp_mpc.MINLPMPC): 

66 

67 config: minlp_mpc.MINLPMPCConfig 

68 

69 def __init__(self, *args, **kwargs): 

70 # create instance variable 

71 self._full_controls: Dict[str, Union[AgentVariable, None]] = {} 

72 super().__init__(*args, **kwargs) 

73 

74 def register_callbacks(self): 

75 for control_var in self.config.controls + self.config.binary_controls: 

76 self.agent.data_broker.register_callback( 

77 name=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}", 

78 alias=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}", 

79 callback=self.calc_flex_callback, 

80 ) 

81 for input_var in self.config.inputs: 

82 if input_var.name.replace(full_trajectory_prefix, "", 1).replace( 

83 full_trajectory_suffix, "" 

84 ) in [control_var.name for control_var in self.config.controls + self.config.binary_controls]: 

85 self._full_controls[input_var.name] = input_var 

86 

87 super().register_callbacks() 

88 

89 def calc_flex_callback(self, inp: AgentVariable, name: str): 

90 """Set the control trajectories before calculating the flexibility offer. 

91 

92 self.model should account for flexibility in its cost function 

93 

94 """ 

95 # during provision dont calculate flex 

96 if self.get("in_provision").value: 

97 return 

98 

99 # do not trigger callback on self set variables 

100 if self.agent.config.id == inp.source.agent_id: 

101 return 

102 

103 vals = strip_multi_index(inp.value) 

104 if vals.isna().any(): 

105 vals = fill_nans(vals, method=MEAN) 

106 

107 # the MPC Predictions starts at t=env.now not t=0 

108 vals.index += self.env.time 

109 self._full_controls[name].value = vals 

110 self.set(name, vals) 

111 # make sure all controls are set 

112 if all(x.value is not None for x in self._full_controls.values()): 

113 self.do_step() 

114 for name in self._full_controls.keys(): 

115 self._full_controls[name].value = None 

116 

117 def process(self): 

118 # the shadow mpc should only be run after the results of the baseline are sent 

119 yield self.env.event()