Add ensemble functionality

CHANGELOG.md

@@ -9,13 +9,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added support for running an ensemble. Default output is ensemble statistics, but --full_ensemble can be specified [!17](https://github.com/dmidk/sunflow/pull/17), @KristianHMoller
 - Added support for using a different, smaller output domain than the satellite input domain [!10](https://github.com/dmidk/sunflow/pull/10), @KristianHMoller
 - Added a check for the number of ensemble members, as the code currently supports only one [!13](https://github.com/dmidk/sunflow/pull/13), @KristianHMoller
 - Subsetting to bounding box is now also done in the `s3` and `files` code paths [!11](https://github.com/dmidk/sunflow/pull/11), @JoachimKoenigslieb
 - Subsetting to bounding box now correctly handles both ascending and descending lat/lon coordinates [!11](https://github.com/dmidk/sunflow/pull/11), @JoachimKoenigslieb
 
 ### Changed
 
+- Change number of ensemble members from being read from environment variable to being a CLI argument. [!17](https://github.com/dmidk/sunflow/pull/17), @KristianHMoller
+- Ensemble output statistics are now configurable via the `ENSEMBLE_STATISTICS` environment variable, and selected statistics are written as separate NetCDF variables. [!17](https://github.com/dmidk/sunflow/pull/17), @KristianHMoller
+- Update SolarSTEPS dependency to use the version with arctan2 [!19](https://github.com/dmidk/sunflow/pull/19), @KristianHMoller
 - Modified subset_to_bbox function to slice based on edges rather than centers [!10](https://github.com/dmidk/sunflow/pull/10), @KristianHMoller
 - Pass down number of ensembles from configurations to `ProbabilisticAdvection`. This gives a roughly 3x speedup in the no ensembles case [!12](https://github.com/dmidk/sunflow/pull/12), @JoachimKoenigslieb
 

README.md

@@ -101,13 +101,13 @@ podman run -it --rm --entrypoint="" sunflow bash
 |----------|---------|-------------|
 | `NOWCAST_DIRECTORY` | `.` | Directory for forecast output files |
 | `SATELLITE_DATA_DIRECTORY` | `.` | Directory for input satellite data archive |
-| `ENS_MEMBERS` | `1` | Number of ensemble members |
 | `PAST_STEPS` | `4` | Number of past time steps for motion field |
 | `FUTURE_STEPS` | `24` | Number of forecast time steps |
 | `INPUT_DATA_AVAILABILITY_DELAY_MINUTES` | `24` | Data availability delay (minutes) |
 | `INPUT_DATA_FREQUENCY_MINUTES` | `15` | Data frequency (minutes) |
 | `MAX_WAITING_TIME_MINUTES` | `27` | Maximum wait time for data (minutes) |
 | `MAX_CLEARSKY_FALLBACK_DAYS` | `3` | Days back to search for fallback clear-sky data |
+| `ENSEMBLE_STATISTICS` | `median,mean,p10,p25,p75,p90` | Comma-separated list of statistics for ensemble output. Allowed: `median`, `mean`, `p10`, `p25`, `p75`, `p90` (aliases `10th_percentile`, `25th_percentile`, `75th_percentile`, `90th_percentile` are accepted). |
 
 #### Data Source Configuration
 
@@ -142,6 +142,12 @@ podman run -it --rm --entrypoint="" sunflow bash
 - `--start_time` - Start of a time range in ISO8601 format (use with `--end_time`)
 - `--end_time` - End of a time range in ISO8601 format, inclusive (use with `--start_time`)
 - `--run_mode` - Specify run mode: `download` (fetch from API), `files` (local files), or `s3` (object storage)
+- `--ensemble_members` - Number of ensemble members (Default 1)
+- `--full_ensemble` - Specify that the full ensemble is the desired output rather than ensemble statistics
+
+For ensemble runs (`--ensemble_members > 1`), default output is the configured
+ensemble statistics from `ENSEMBLE_STATISTICS`. Use `--full_ensemble` to output
+all members instead.
 
 ## Data Sources
 

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "pvlib>=0.10.0",
     "pysteps>=1.7.0",
     "opencv-python-headless>=4.5.0",
-    "SolarSTEPS @ git+https://github.com/dmidk/SolarSTEPS@18012d7b11c56895e14abc39c3800f68a9a44ecb",
+    "SolarSTEPS @ git+https://github.com/dmidk/SolarSTEPS@3def2b36974bc2fb92f2ff14fc40dc8da2a9c547",
     "loguru>=0.7.3",
     "isodate>=0.7.2",
     "fsspec>=2023.1.0",

sunflow/config.py

@@ -3,6 +3,27 @@
 from dataclasses import dataclass
 from typing import Self
 
+DEFAULT_ENSEMBLE_STATISTICS = "median,mean,p10,p25,p75,p90"
+_ALLOWED_STATISTICS = {
+    "median",
+    "mean",
+    "p10",
+    "p25",
+    "p75",
+    "p90",
+    "10th_percentile",
+    "25th_percentile",
+    "75th_percentile",
+    "90th_percentile",
+}
+_STATISTIC_ALIASES = {
+    "10th_percentile": "p10",
+    "25th_percentile": "p25",
+    "75th_percentile": "p75",
+    "90th_percentile": "p90",
+}
+
+
 # Predefined domain options
 # Format: lon_min,lat_min,lon_max,lat_max
 DOMAIN_OPTIONS: dict[str, str | None] = {
@@ -50,33 +71,53 @@ class NowcastConfig:
 
     nowcast_directory: str
     ens_members: int
+    alpha: float
+    beta: float
     past_steps: int
     future_steps: int
     input_data_availability_delay_minutes: int
     input_data_frequency_minutes: int
     max_waiting_time_minutes: int
     satellite_data_directory: str
     max_clearsky_fallback_days: int
+    ensemble_statistics: list[str]
 
     @classmethod
-    def from_env(cls) -> Self:
+    def from_env(cls, ensemble_members: int = 1) -> Self:
         """Load nowcast configuration from environment variables with defaults.
 
         Reads the following environment variables:
 
         - NOWCAST_DIRECTORY (default: .)
-        - ENS_MEMBERS (default: 1)
+        - ALPHA (default: 0.0 for ENS_MEMBERS=1, 9.23 for ENS_MEMBERS>1)
+        - BETA (default: 0.0 for ENS_MEMBERS=1, 0.15 for ENS_MEMBERS>1)
         - PAST_STEPS (default: 4)
         - FUTURE_STEPS (default: 24)
         - INPUT_DATA_AVAILABILITY_DELAY_MINUTES (default: 24)
         - INPUT_DATA_FREQUENCY_MINUTES (default: 15)
         - MAX_WAITING_TIME_MINUTES (default: 27)
         - SATELLITE_DATA_DIRECTORY (default: .)
         - MAX_CLEARSKY_FALLBACK_DAYS (default: 3)
+        - ENSEMBLE_STATISTICS (default: median,mean,p10,p25,p75,p90)
         """
+
+        ens_members = ensemble_members
+        # Reference for default noise values:
+        # A. Carpentieri, D. Folini, D. Nerini, S. Pulkkinen, M. Wild, A. Meyer,
+        # "Intraday probabilistic forecasts of surface solar radiation with cloud
+        # scale-dependent autoregressive advection,"
+        # Applied Energy, Volume 351, 2023
+        default_alpha = 0.0 if ens_members == 1 else 9.23
+        default_beta = 0.0 if ens_members == 1 else 0.15
+        statistics = _parse_ensemble_statistics(
+            os.getenv("ENSEMBLE_STATISTICS", DEFAULT_ENSEMBLE_STATISTICS)
+        )
+
         return cls(
             nowcast_directory=os.getenv("NOWCAST_DIRECTORY", "."),
-            ens_members=int(os.getenv("ENS_MEMBERS", "1")),
+            ens_members=ens_members,
+            alpha=float(os.getenv("ALPHA", str(default_alpha))),
+            beta=float(os.getenv("BETA", str(default_beta))),
             past_steps=int(os.getenv("PAST_STEPS", "4")),
             future_steps=int(os.getenv("FUTURE_STEPS", "24")),
             input_data_availability_delay_minutes=int(
@@ -88,4 +129,25 @@ def from_env(cls) -> Self:
             max_waiting_time_minutes=int(os.getenv("MAX_WAITING_TIME_MINUTES", "27")),
             satellite_data_directory=os.getenv("SATELLITE_DATA_DIRECTORY", "."),
             max_clearsky_fallback_days=int(os.getenv("MAX_CLEARSKY_FALLBACK_DAYS", "3")),
+            ensemble_statistics=statistics,
+        )
+
+
+def _parse_ensemble_statistics(raw_statistics: str) -> list[str]:
+    """Parse and validate requested ensemble statistics from environment."""
+    statistics = [
+        token.strip().lower() for token in raw_statistics.split(",") if token.strip()
+    ]
+    if not statistics:
+        raise ValueError("ENSEMBLE_STATISTICS must contain at least one statistic")
+
+    invalid = [stat for stat in statistics if stat not in _ALLOWED_STATISTICS]
+    if invalid:
+        raise ValueError(
+            "Invalid ENSEMBLE_STATISTICS value(s): "
+            f"{', '.join(invalid)}. Allowed values are: "
+            "median, mean, p10, p25, p75, p90, "
+            "10th_percentile, 25th_percentile, 75th_percentile, 90th_percentile"
         )
+
+    return [_STATISTIC_ALIASES.get(stat, stat) for stat in statistics]

sunflow/data_io.py

@@ -459,60 +459,89 @@ def fetch_clearsky_with_fallback(
 
 
 def save_forecast(
-    forecast: np.ndarray,
+    forecast: np.ndarray | dict[str, np.ndarray],
     time_step: datetime,
     n_steps: int,
     latitudes: np.ndarray,
     longitudes: np.ndarray,
     dataset_name: str,
     nowcast_config: NowcastConfig,
     model_version: str,
+    output_mode: str,
     run_mode: str = "files",
     s3_config: S3Config | None = None,
 ) -> str:
-    """Save forecast array to a CF-compliant NetCDF4 file.
+    """Save forecast data to a CF-compliant NetCDF4 file.
 
     Writes the probabilistic advection forecast to either a local file or S3,
     depending on `run_mode`. The time coordinate is stored as CF-convention
     numeric values (float64, minutes since the forecast reference time).
 
     Args:
-        forecast: Forecast array, shape [time, lat, lon] or
-            [ensemble, time, lat, lon].
+        forecast: Forecast array with shape [ensemble, time, lat, lon],
+            or mapping of statistic name to arrays with the same shape.
         time_step: Forecast reference time (start of the forecast window).
         n_steps: Number of forecast time steps to write.
         latitudes: 1-D array of latitude values (degrees).
         longitudes: 1-D array of longitude values (degrees).
         dataset_name: Name of the source dataset (options: KNMI, DWD).
-        nowcast_config: NowcastConfig object supplying output directory,
-            ensemble size, and input data frequency.
+        nowcast_config: NowcastConfig object supplying output directory
+            and input data frequency.
         model_version: Model version string written as a global attribute.
+        output_mode: Output aggregation mode label written to global
+            NetCDF attrs (expected: 'deterministic', 'ensemble_statistics',
+            or 'full_ensemble').
         run_mode: One of 'files' (local) or 's3'. Defaults to 'files'.
         s3_config: S3Config object; required when run_mode is 's3'.
 
     Returns:
         Filename (basename only) of the written NetCDF file.
     """
     input_data_frequency_minutes = nowcast_config.input_data_frequency_minutes
-    ens_members = nowcast_config.ens_members
     filename = f"SolarNowcast_{time_step.strftime('%Y%m%d%H%M')}.nc"
 
-    # Add ensemble dimension if needed (forecast should be [ensemble, time, lat, lon])
-    if forecast.ndim == 3:
-        forecast = forecast[np.newaxis, :, :, :]  # Now [1, time, lat, lon]
+    data_vars: dict[str, tuple[list[str], np.ndarray, dict[str, str]]]
+    statistics_attr = ""
 
-    # Build time coordinate (CF-convention: minutes since forecast reference time)
+    if isinstance(forecast, dict):
+        if not forecast:
+            raise ValueError("forecast statistics mapping cannot be empty")
 
-    time_step_naive = time_step.replace(tzinfo=None)
-    _time_units = f"minutes since {time_step_naive.strftime('%Y-%m-%d %H:%M:%S')}"
-    time_datetimes = [
-        time_step_naive + timedelta(minutes=input_data_frequency_minutes * i)
-        for i in range(0, n_steps)
-    ]
+        data_vars = {}
+        first_shape: tuple[int, ...] | None = None
+        for statistic, values in forecast.items():
+            if values.ndim != 4:
+                raise ValueError(
+                    "Each statistic array must have shape " "(ensemble, time, lat, lon)."
+                )
 
-    ds = xr.Dataset(
-        {
-            "probabilistic_advection": (
+            if first_shape is None:
+                first_shape = values.shape
+            elif values.shape != first_shape:
+                raise ValueError("All statistic arrays must share the same shape")
+
+            variable_name = f"GHI_probabilistic_advection_{statistic}"
+            data_vars[variable_name] = (
+                ["ensemble", "time", "lat", "lon"],
+                values,
+                {
+                    "description": (
+                        f"Probabilistic advection solar forecast ({statistic})"
+                    ),
+                    "long_name": "Surface downwelling solar radiation",
+                    "units": "W m-2",
+                },
+            )
+
+        ens_members = first_shape[0] if first_shape is not None else 0
+        statistics_attr = ",".join(forecast.keys())
+    else:
+        if forecast.ndim != 4:
+            raise ValueError("forecast must have shape (ensemble, time, lat, lon).")
+
+        ens_members = forecast.shape[0]
+        data_vars = {
+            "GHI_probabilistic_advection": (
                 ["ensemble", "time", "lat", "lon"],
                 forecast,
                 {
@@ -521,7 +550,19 @@ def save_forecast(
                     "units": "W m-2",
                 },
             ),
-        },
+        }
+
+    # Build time coordinate (CF-convention: minutes since forecast reference time)
+
+    time_step_naive = time_step.replace(tzinfo=None)
+    _time_units = f"minutes since {time_step_naive.strftime('%Y-%m-%d %H:%M:%S')}"
+    time_datetimes = [
+        time_step_naive + timedelta(minutes=input_data_frequency_minutes * i)
+        for i in range(0, n_steps)
+    ]
+
+    ds = xr.Dataset(
+        data_vars,
         coords={
             "time": (
                 ["time"],
@@ -549,11 +590,13 @@ def save_forecast(
                 f"Simple Probabilistic Advection solar forecast "
                 f"using {dataset_name} data"
             ),
+            "output_mode": output_mode,
             "history": (
                 f"Created "
                 f"{datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC"
             ),
             "model_version": model_version,
+            **({"statistics": statistics_attr} if statistics_attr else {}),
         },
     )