Forecasting updates

src/pymgrid/forecast/forecaster.py

@@ -2,8 +2,10 @@
 from pandas.api.types import is_number, is_numeric_dtype
 from abc import abstractmethod
 
+from pymgrid.utils.space import ModuleSpace
 
-def get_forecaster(forecaster, observation_space, sink_only=False, time_series=None, increase_uncertainty=False):
+
+def get_forecaster(forecaster, observation_space, forecast_shape, sink_only=False, time_series=None, increase_uncertainty=False):
     """
     Get the forecasting function for the time series module.
 
@@ -53,17 +55,18 @@ def get_forecaster(forecaster, observation_space, sink_only=False, time_series=N
     """
 
     if forecaster is None:
-        return NoForecaster(observation_space, sink_only)
+        return NoForecaster(observation_space, forecast_shape, sink_only)
     elif isinstance(forecaster, (UserDefinedForecaster, OracleForecaster, GaussianNoiseForecaster)):
         return forecaster
     elif callable(forecaster):
-        return UserDefinedForecaster(forecaster, observation_space, sink_only, time_series)
+        return UserDefinedForecaster(forecaster, observation_space, forecast_shape, sink_only, time_series)
     elif forecaster == "oracle":
-        return OracleForecaster(observation_space, sink_only)
+        return OracleForecaster(observation_space, forecast_shape, sink_only)
     elif is_number(forecaster):
         return GaussianNoiseForecaster(
             forecaster,
             observation_space,
+            forecast_shape,
             sink_only,
             increase_uncertainty=increase_uncertainty
         )
@@ -72,12 +75,32 @@ def get_forecaster(forecaster, observation_space, sink_only=False, time_series=N
 
 
 class Forecaster:
-    def __init__(self, observation_space, sink_only):
+    def __init__(self, observation_space, forecast_shape, sink_only):
         self._observation_space = observation_space
+        self._forecast_shaped_space = self._get_forecast_shaped_space(forecast_shape)
         self._fill_arr = (self._observation_space.unnormalized.high + self._observation_space.unnormalized.low) / 2
         if sink_only:
             self._fill_arr *= -1
 
+    def _get_forecast_shaped_space(self, shape):
+        if len(shape) == 1:
+            shape = (*shape, 1)
+        elif len(shape) > 2:
+            raise ValueError(f'shape must be one- or two-dimensional.')
+
+        n_in_forecast = shape[0]*shape[1]
+
+        if n_in_forecast:
+            unnormalized_low = self._observation_space.unnormalized.low[-n_in_forecast:]
+            unnormalized_high = self._observation_space.unnormalized.high[-n_in_forecast:]
+        else:
+            unnormalized_low = np.array([])
+            unnormalized_high = np.array([])
+
+        return ModuleSpace(unnormalized_low=unnormalized_low.reshape(shape),
+                           unnormalized_high=unnormalized_high.reshape(shape),
+                           shape=shape)
+
     @abstractmethod
     def _forecast(self, val_c, val_c_n, n):
         pass
@@ -88,6 +111,11 @@ def _pad(self, forecast, n):
         else:
             pad_amount = n - forecast.shape[0]
             pad = self._fill_arr.reshape((-1, forecast.shape[1]))[-pad_amount:]
+
+            if pad.shape[0] < pad_amount:
+                raise RuntimeError("Attempting to pad a forecast to a value larger than the module's observation space "
+                                   "implies.")
+
             return np.concatenate((forecast, pad))
 
     def full_pad(self, shape, forecast_horizon):
@@ -96,6 +124,20 @@ def full_pad(self, shape, forecast_horizon):
         empty_forecast = np.array([]).reshape((0, shape[1]))
         return self._pad(empty_forecast, forecast_horizon)
 
+    @property
+    def observation_space(self):
+        return self._observation_space
+
+    @observation_space.setter
+    def observation_space(self, value):
+        self._observation_space = value
+        self._fill_arr = (self._observation_space.unnormalized.high + self._observation_space.unnormalized.low) / 2
+        new_shape = (
+            int((value.shape[0] - self._forecast_shaped_space.shape[1]) / self._forecast_shaped_space.shape[1]),
+            self._forecast_shaped_space.shape[1]
+        )
+        self._forecast_shaped_space = self._get_forecast_shaped_space(new_shape)
+
     def __eq__(self, other):
         if type(self) != type(other):
             return NotImplemented
@@ -106,6 +148,10 @@ def __eq__(self, other):
     def __call__(self, val_c, val_c_n, n):
         if len(val_c_n.shape) == 1:
             val_c_n = val_c_n.reshape((-1, 1))
+
+        if val_c_n.shape[0] > self._forecast_shaped_space.shape[0]:
+            raise RuntimeError(f'val_c_n shape {val_c_n.shape} is too large for space {self._forecast_shaped_space.shape}')
+
         forecast = self._forecast(val_c, val_c_n, n)
 
         if forecast is None:
@@ -120,7 +166,7 @@ def __repr__(self):
 
 
 class UserDefinedForecaster(Forecaster):
-    def __init__(self, forecaster_function, observation_space, sink_only, time_series):
+    def __init__(self, forecaster_function, observation_space, forecast_shape, sink_only, time_series):
         self.is_vectorized_forecaster, self.cast_to_arr = \
             _validate_callable_forecaster(forecaster_function, time_series)
 
@@ -129,7 +175,7 @@ def __init__(self, forecaster_function, observation_space, sink_only, time_serie
 
         self._forecaster = forecaster_function
 
-        super().__init__(observation_space, sink_only)
+        super().__init__(observation_space, forecast_shape, sink_only)
 
     def _cast_to_arr(self, forecast, val_c_n):
         if self.cast_to_arr:
@@ -147,37 +193,47 @@ def _forecast(self, val_c, val_c_n, n):
 
 
 class GaussianNoiseForecaster(Forecaster):
-    def __init__(self, noise_std, observation_space, sink_only, increase_uncertainty=False):
+    def __init__(self, noise_std, observation_space, forecast_shape, sink_only, increase_uncertainty=False):
+        super().__init__(observation_space, forecast_shape, sink_only)
+
         self.input_noise_std = noise_std
         self.increase_uncertainty = increase_uncertainty
-        self._noise_size = None
-        self._noise_std = None
 
-        super().__init__(observation_space, sink_only)
+        self._noise_size = self._forecast_shaped_space.shape
+        self._noise_std = self._get_noise_std()
 
     def _get_noise_std(self):
         if self.increase_uncertainty:
-            return self.input_noise_std * (1 + np.log(1 + np.arange(self._noise_size)))
+            return self.input_noise_std * np.outer(
+                1 + np.log(1 + np.arange(self._noise_size[0])),
+                np.ones(self._noise_size[-1])
+            )
         else:
             return self.input_noise_std
 
     def _get_noise(self, size):
-        if self._noise_size is None:
-            self._noise_size = size
-        if size != self._noise_size:
-            raise ValueError(f"size {size} incompatible with previous size {self._noise_size}")
-        return np.random.normal(scale=self.noise_std, size=size)
+        return np.random.normal(scale=self._noise_std, size=size)
+
+    def _forecast(self, val_c, val_c_n, n):
+        forecast = val_c_n + self._get_noise(val_c_n.shape).reshape(val_c_n.shape)
+
+        lb = self._forecast_shaped_space.unnormalized.low[-val_c_n.shape[0]:]
+        ub = self._forecast_shaped_space.unnormalized.high[-val_c_n.shape[0]:]
+        lt_lb = forecast < lb
+        gt_ub = forecast > ub
+
+        forecast[lt_lb] = lb[lt_lb]
+        forecast[gt_ub] = ub[gt_ub]
+
+        return forecast
 
     @property
     def noise_std(self):
-        if self._noise_std is None:
-            self._noise_std = self._get_noise_std()
         return self._noise_std
 
-    def _forecast(self, val_c, val_c_n, n):
-        forecast = val_c_n + self._get_noise(len(val_c_n)).reshape(val_c_n.shape)
-        forecast[(forecast * val_c_n) < 0] = 0
-        return forecast
+    @noise_std.setter
+    def noise_std(self, value):
+        pass
 
     def __repr__(self):
         return f'GaussianNoiseForecaster(noise_std={self.input_noise_std}, ' \
@@ -188,6 +244,7 @@ class NoForecaster(Forecaster):
     def _forecast(self, val_c, val_c_n, n):
         return None
 
+
 def _validate_callable_forecaster(forecaster, time_series):
     val_c = time_series[0]
     n = np.random.randint(2, len(time_series))

src/pymgrid/microgrid/microgrid.py

@@ -403,6 +403,48 @@ def get_log(self, as_frame=True, drop_singleton_key=False):
 
         return df.to_dict()
 
+    def set_forecaster(self,
+                       forecaster,
+                       forecast_horizon=DEFAULT_HORIZON,
+                       forecaster_increase_uncertainty=False):
+        """
+        TODO write this docstring
+        Then get this into your MPC/RBC/RL config somehow
+
+        Parameters
+        ----------
+        forecaster
+        forecast_horizon
+        forecaster_increase_uncertainty
+
+        Returns
+        -------
+
+        """
+        if isinstance(forecaster, dict):
+            for module_name, _forecaster in forecaster.items():
+                if module_name not in self._modules.names():
+                    raise NameError(f'Unrecognized module {module_name}.')
+
+                try:
+                    self._modules[module_name].set_forecaster(
+                        _forecaster,
+                        forcast_horizon=forecast_horizon,
+                        forecaster_increase_uncertainty=forecaster_increase_uncertainty
+                    )
+                except AttributeError:
+                    pass
+        else:
+            for module in self._modules.iterlist():
+                try:
+                    module.set_forecaster(
+                        forecaster,
+                        forecast_horizon=forecast_horizon,
+                        forecaster_increase_uncertainty=forecaster_increase_uncertainty
+                    )
+                except AttributeError:
+                    pass
+
     def get_forecast_horizon(self):
         """
         Get the forecast horizon of timeseries modules contained in the microgrid.

src/pymgrid/modules/base/base_module.py

@@ -70,7 +70,7 @@ def reset(self):
             Normalized observation after resetting.
 
         """
-        self._current_step = 0
+        self._update_step(reset=True)
         self._logger.flush()
         return self.to_normalized(self.state, obs=True)
 
@@ -143,7 +143,7 @@ def step(self, action, normalized=True):
         state_dict = self.state_dict
         reward, done, info = self._unnormalized_step(denormalized_action)
         self._log(state_dict, reward=reward, **info)
-        self._current_step += 1
+        self._update_step()
 
         obs = self.to_normalized(self.state, obs=True)
 
@@ -280,6 +280,11 @@ def _log(self, state_dict_pre_step, provided_energy=None, absorbed_energy=None,
         _info.update(state_dict_pre_step)
         self._logger.log(**_info)
 
+    def _update_step(self, reset=False):
+        if reset:
+            self._current_step = 0
+        self._current_step += 1
+
     @abstractmethod
     def update(self, external_energy_change, as_source=False, as_sink=False):
         """