Coverage for agentlib_flexquant/modules/shadow_mpc.py: 62%
63 statements
« prev ^ index » next coverage.py v7.4.4, created at 2025-08-01 15:10 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2025-08-01 15:10 +0000
1from agentlib_mpc.modules import mpc_full, minlp_mpc
2from agentlib_flexquant.utils.data_handling import strip_multi_index, fill_nans, MEAN, INTERPOLATE
3from agentlib_flexquant.data_structures.globals import (
4 full_trajectory_prefix,
5 full_trajectory_suffix,
6)
7from typing import Dict, Union
8from agentlib.core.datamodels import AgentVariable
11class FlexibilityShadowMPC(mpc_full.MPC):
13 config: mpc_full.MPCConfig
15 def __init__(self, *args, **kwargs):
16 # create instance variable
17 self._full_controls: Dict[str, Union[AgentVariable, None]] = {}
18 super().__init__(*args, **kwargs)
20 def register_callbacks(self):
21 for control_var in self.config.controls:
22 self.agent.data_broker.register_callback(
23 name=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}",
24 alias=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}",
25 callback=self.calc_flex_callback,
26 )
27 for input_var in self.config.inputs:
28 if input_var.name.replace(full_trajectory_prefix, "", 1).replace(
29 full_trajectory_suffix, ""
30 ) in [control_var.name for control_var in self.config.controls]:
31 self._full_controls[input_var.name] = input_var
33 super().register_callbacks()
35 def calc_flex_callback(self, inp, name):
36 """set the control trajectories before calculating the flexibility offer.
37 self.model should account for flexibility in its cost function
39 """
40 # during provision dont calculate flex
41 if self.get("in_provision").value:
42 return
44 # do not trigger callback on self set variables
45 if self.agent.config.id == inp.source.agent_id:
46 return
48 vals = strip_multi_index(inp.value)
49 if vals.isna().any():
50 vals = fill_nans(series=vals, method=MEAN)
52 # the MPC Predictions starts at t=env.now not t=0
53 vals.index += self.env.time
54 self._full_controls[name].value = vals
55 self.set(name, vals)
56 # make sure all controls are set
57 if all(x.value is not None for x in self._full_controls.values()):
58 self.do_step()
59 for name in self._full_controls.keys():
60 self._full_controls[name].value = None
62 def process(self):
63 # the shadow mpc should only be run after the results of the baseline are sent
64 yield self.env.event()
67class FlexibilityShadowMINLPMPC(minlp_mpc.MINLPMPC):
69 config: minlp_mpc.MINLPMPCConfig
71 def __init__(self, *args, **kwargs):
72 # create instance variable
73 self._full_controls: Dict[str, Union[AgentVariable, None]] = {}
74 super().__init__(*args, **kwargs)
76 def register_callbacks(self):
77 for control_var in self.config.controls + self.config.binary_controls:
78 self.agent.data_broker.register_callback(
79 name=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}",
80 alias=f"{full_trajectory_prefix}{control_var.name}{full_trajectory_suffix}",
81 callback=self.calc_flex_callback,
82 )
83 for input_var in self.config.inputs:
84 if input_var.name.replace(full_trajectory_prefix, "", 1).replace(
85 full_trajectory_suffix, ""
86 ) in [control_var.name for control_var in self.config.controls + self.config.binary_controls]:
87 self._full_controls[input_var.name] = input_var
89 super().register_callbacks()
91 def calc_flex_callback(self, inp, name):
92 """set the control trajectories before calculating the flexibility offer.
93 self.model should account for flexibility in its cost function
95 """
96 # during provision dont calculate flex
97 if self.get("in_provision").value:
98 return
100 # do not trigger callback on self set variables
101 if self.agent.config.id == inp.source.agent_id:
102 return
104 vals = strip_multi_index(inp.value)
105 if vals.isna().any():
106 vals = fill_nans(vals, method=MEAN)
108 # the MPC Predictions starts at t=env.now not t=0
109 vals.index += self.env.time
110 self._full_controls[name].value = vals
111 self.set(name, vals)
112 # make sure all controls are set
113 if all(x.value is not None for x in self._full_controls.values()):
114 self.do_step()
115 for name in self._full_controls.keys():
116 self._full_controls[name].value = None
118 def process(self):
119 # the shadow mpc should only be run after the results of the baseline are sent
120 yield self.env.event()