remove sink_only from get_forecaster signature

src/pymgrid/forecast/forecaster.py

@@ -5,7 +5,7 @@
 from pymgrid.utils.space import ModuleSpace
 
 
-def get_forecaster(forecaster, observation_space, forecast_shape, sink_only=False, time_series=None, increase_uncertainty=False):
+def get_forecaster(forecaster, observation_space, forecast_shape, time_series=None, increase_uncertainty=False):
     """
     Get the forecasting function for the time series module.
 
@@ -32,13 +32,13 @@ def get_forecaster(forecaster, observation_space, forecast_shape, sink_only=Fals
 
         * If ``None``, no forecast.
 
+    forecast_shape : int or tuple of int
+        Expected shape of forecasts. If an integer, will return forecasts of shape (shape, 1).
+
     observation_space : :class:`ModuleSpace <pymgrid.utils.space.ModuleSpace>`
         Observation space; used to determine values to pad missing forecasts when we are forecasting past the
         end of the time series.
 
-    sink_only : bool
-        Whether the module is a sink and is not a source.
-
     time_series: ndarray[float] or None, default None.
         The underlying time series, used to validate UserDefinedForecaster.
         Only used if callable(forecaster).
@@ -55,27 +55,22 @@ def get_forecaster(forecaster, observation_space, forecast_shape, sink_only=Fals
     """
 
     if forecaster is None:
-        return NoForecaster(observation_space, forecast_shape, sink_only)
+        return NoForecaster(observation_space, forecast_shape)
     elif isinstance(forecaster, (UserDefinedForecaster, OracleForecaster, GaussianNoiseForecaster)):
         return forecaster
     elif callable(forecaster):
-        return UserDefinedForecaster(forecaster, observation_space, forecast_shape, sink_only, time_series)
+        return UserDefinedForecaster(forecaster, observation_space, forecast_shape, time_series)
     elif forecaster == "oracle":
-        return OracleForecaster(observation_space, forecast_shape, sink_only)
+        return OracleForecaster(observation_space, forecast_shape)
     elif is_number(forecaster):
-        return GaussianNoiseForecaster(
-            forecaster,
-            observation_space,
-            forecast_shape,
-            sink_only,
-            increase_uncertainty=increase_uncertainty
-        )
+        return GaussianNoiseForecaster(forecaster, observation_space, forecast_shape,
+                                       increase_uncertainty=increase_uncertainty)
     else:
         raise ValueError(f"Unable to parse forecaster of type {type(forecaster)}")
 
 
 class Forecaster:
-    def __init__(self, observation_space, forecast_shape, sink_only):
+    def __init__(self, observation_space, forecast_shape):
         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
@@ -176,7 +171,7 @@ def __repr__(self):
 
 
 class UserDefinedForecaster(Forecaster):
-    def __init__(self, forecaster_function, observation_space, forecast_shape, sink_only, time_series):
+    def __init__(self, forecaster_function, observation_space, forecast_shape, time_series):
         self.is_vectorized_forecaster, self.cast_to_arr = \
             _validate_callable_forecaster(forecaster_function, time_series)
 
@@ -185,7 +180,7 @@ def __init__(self, forecaster_function, observation_space, forecast_shape, sink_
 
         self._forecaster = forecaster_function
 
-        super().__init__(observation_space, forecast_shape, sink_only)
+        super().__init__(observation_space, forecast_shape)
 
     def _cast_to_arr(self, forecast, val_c_n):
         if self.cast_to_arr:
@@ -203,8 +198,8 @@ def _forecast(self, val_c, val_c_n, n):
 
 
 class GaussianNoiseForecaster(Forecaster):
-    def __init__(self, noise_std, observation_space, forecast_shape, sink_only, increase_uncertainty=False):
-        super().__init__(observation_space, forecast_shape, sink_only)
+    def __init__(self, noise_std, observation_space, forecast_shape, increase_uncertainty=False):
+        super().__init__(observation_space, forecast_shape)
 
         self.input_noise_std = noise_std
         self.increase_uncertainty = increase_uncertainty

src/pymgrid/modules/base/timeseries/base_timeseries_module.py

@@ -38,7 +38,6 @@ def __init__(self,
         self._forecaster = get_forecaster(forecaster,
                                           self._get_observation_spaces(),
                                           forecast_shape=(self.forecast_horizon, len(self.state_components)),
-                                          sink_only=self.is_sink and not self.is_source,
                                           time_series=self.time_series,
                                           increase_uncertainty=forecaster_increase_uncertainty)
 
@@ -228,7 +227,6 @@ def set_forecaster(self,
         self._forecaster = get_forecaster(forecaster,
                                           self._observation_space,
                                           (self.forecast_horizon, len(self.state_components)),
-                                          self.is_sink and not self.is_source,
                                           self.time_series,
                                           increase_uncertainty=forecaster_increase_uncertainty)