diff --git a/CHANGELOG.md b/CHANGELOG.md index 335010f..90b7918 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ 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 @@ -16,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### 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 diff --git a/README.md b/README.md index 18b8ffb..f37441f 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 57f5e34..fe97637 100644 --- a/pyproject.toml +++ b/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", diff --git a/sunflow/config.py b/sunflow/config.py index 58e67f9..b8faa5c 100644 --- a/sunflow/config.py +++ b/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,6 +71,8 @@ class NowcastConfig: nowcast_directory: str ens_members: int + alpha: float + beta: float past_steps: int future_steps: int input_data_availability_delay_minutes: int @@ -57,15 +80,17 @@ class NowcastConfig: 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) @@ -73,10 +98,26 @@ def from_env(cls) -> Self: - 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] diff --git a/sunflow/data_io.py b/sunflow/data_io.py index 3fcc9b3..ec744ee 100644 --- a/sunflow/data_io.py +++ b/sunflow/data_io.py @@ -459,7 +459,7 @@ 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, @@ -467,26 +467,30 @@ def save_forecast( 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'. @@ -494,25 +498,50 @@ def save_forecast( 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 {}), }, ) diff --git a/sunflow/forecast.py b/sunflow/forecast.py index 3abe325..f3f8bc4 100644 --- a/sunflow/forecast.py +++ b/sunflow/forecast.py @@ -5,7 +5,7 @@ import xarray as xr from Models.ProbabilisticAdvection import ProbabilisticAdvection -from .geospatial import get_coordinates +from .geospatial import crop_forecast_to_domain, get_coordinates def preprocess_data( @@ -61,16 +61,19 @@ def preprocess_data( ) -def simple_advection_forecast( - ratio_data: np.ndarray, motion_field: np.ndarray, n_steps: int, ens_members: int +def probabilistic_advection_forecast( + ratio_data: np.ndarray, + motion_field: np.ndarray, + n_steps: int, + ens_members: int, + alpha: float, + beta: float, ) -> np.ndarray: - """Run a deterministic advection forecast on solar irradiance ratios. + """Run a probabilistic advection forecast on solar irradiance ratios. - Uses ProbabilisticAdvection with noise parameters alpha=0 and beta=0, - which disables Gaussian noise on the motion field norm and - von Mises noise on the direction, yielding a purely deterministic - advection result. The ensemble dimension added by the model is removed - before returning. + Uses ProbabilisticAdvection with configurable noise parameters: + alpha controls Gaussian noise on motion field norm and beta controls + von Mises noise on motion field direction. Args: ratio_data: Input array of shape (time, lat, lon) containing @@ -78,25 +81,24 @@ def simple_advection_forecast( motion_field: Optical flow field of shape (2, lat, lon) as produced by dense_lucaskanade. n_steps: Number of forecast timesteps to produce. + ens_members: Number of ensemble members. + alpha: Gaussian noise strength on motion field norm. + beta: von Mises noise strength on motion field angle. Returns: Forecast array of shape (n_steps, lat, lon). """ - # Initialize ProbabilisticAdvection with NO noise (alpha=0, beta=0) + # Initialize ProbabilisticAdvection with configured noise settings. pa = ProbabilisticAdvection( - alpha=0.0, # No Gaussian noise on motion field norm - beta=0.0, # No von Mises noise on motion field angle + alpha=alpha, + beta=beta, return_motion_field=False, ens_members=ens_members, ) # Run probabilistic advection using the correct method name forecast = pa.maps_forecast(n_steps, ratio_data, motion_field) - # Remove ensemble dimension if present (squeeze to get shape: [time, lat, lon]) - if forecast.ndim == 4: # [ensemble, time, lat, lon] - forecast = forecast[0] # Take first (and only) ensemble member - return forecast @@ -115,7 +117,7 @@ def multiply_clearsky( Args: ratio_forecast: Forecast array of shape (n_steps, lat, lon) - containing SDS/SDS_CS ratios. + or (ensemble, n_steps, lat, lon) containing SDS/SDS_CS ratios. clearsky_data: xarray Dataset with a 'time' dimension containing the clearsky variable for each forecast step. previous_day_time_steps: List of datetimes (one per forecast step) @@ -124,26 +126,153 @@ def multiply_clearsky( NetCDF variable name in the datasets. Returns: - Solar irradiance forecast array of shape (n_steps, lat, lon) - in W m⁻². + Solar irradiance forecast array with the same shape as + ratio_forecast, in W m⁻². Raises: RuntimeError: If clearsky data is missing for any forecast timestep. """ - solar_forecast = np.zeros_like(ratio_forecast) + clearsky_steps: list[np.ndarray] = [] - for i, time_step in enumerate(previous_day_time_steps): + for time_step in previous_day_time_steps: try: sds_cs = clearsky_data.sel(time=time_step.replace(tzinfo=None))[ nc_variable_names["sds_cs"] ].values - - # Multiply ratio by clearsky - solar_forecast[i] = ratio_forecast[i] * sds_cs + clearsky_steps.append(sds_cs) except KeyError: raise RuntimeError( f"No clearsky data for {time_step.strftime('%Y-%m-%dT%H:%M:%SZ')}, " "cannot compute solar forecast for this step." ) + clearsky_stack = np.stack(clearsky_steps, axis=0) + + if ratio_forecast.ndim == 3: + if ratio_forecast.shape[0] != clearsky_stack.shape[0]: + raise ValueError( + "ratio_forecast time dimension does not match clearsky timesteps " + f"({ratio_forecast.shape[0]} != {clearsky_stack.shape[0]})." + ) + solar_forecast = ratio_forecast * clearsky_stack + return solar_forecast[ + np.newaxis, :, :, : + ] # Add ensemble dimension for consistency + + if ratio_forecast.ndim == 4: + if ratio_forecast.shape[1] != clearsky_stack.shape[0]: + raise ValueError( + "ratio_forecast time dimension does not match clearsky timesteps " + f"({ratio_forecast.shape[1]} != {clearsky_stack.shape[0]})." + ) + return ratio_forecast * clearsky_stack[np.newaxis, :, :, :] + + raise ValueError( + "ratio_forecast must have shape (time, lat, lon) or " + "(ensemble, time, lat, lon)." + ) + + +def prepend_t0( + clearsky_data: xr.Dataset, + ratio_data: np.ndarray, + solar_forecast: np.ndarray, + config: dict, + clearsky_t0_time: datetime, +) -> np.ndarray: + """Prepend analysis timestep (t=0) to an ensemble solar forecast. + + Computes the t=0 solar field as the latest observed ratio + (ratio_data[-1]) multiplied by clearsky irradiance at clearsky_t0_time, + then prepends that field to all ensemble members in solar_forecast. + + Args: + clearsky_data: Dataset containing clearsky irradiance values on + a time axis. + ratio_data: Ratio history array with shape (time, lat, lon). + solar_forecast: Forecast array with shape + (ensemble, forecast_time, lat, lon). + config: Runtime configuration dict containing + config["nc_variable_names"]["sds_cs"]. + clearsky_t0_time: Timestamp for the analysis clearsky field, + typically one day before the nowcast time. + + Returns: + Array of shape (ensemble, forecast_time + 1, lat, lon) + with the analysis field inserted at index 0 along the time axis. + """ + # Prepend timestep 0: current observation (ratio_data[-1]) × clearsky at t=0 + sds_cs_t0 = clearsky_data.sel(time=clearsky_t0_time.replace(tzinfo=None))[ + config["nc_variable_names"]["sds_cs"] + ].values + solar_t0 = ratio_data[-1] * sds_cs_t0 + + # Broadcast the same t=0 clearsky-based analysis field to all ensembles. + solar_t0_ens = np.broadcast_to( + solar_t0, + (solar_forecast.shape[0],) + solar_t0.shape, + ) + solar_forecast = np.concatenate( + [solar_t0_ens[:, np.newaxis, :, :], solar_forecast], + axis=1, + ) return solar_forecast + + +def compute_ensemble_statistics( + forecast: np.ndarray, + statistics: list[str], + latitudes: np.ndarray, + longitudes: np.ndarray, + domain_nowcast: str, +) -> dict[str, np.ndarray]: + """Compute requested statistics over ensemble members (axis 0).""" + computed: dict[str, np.ndarray] = {} + for statistic in statistics: + match statistic: + case "median": + computed["median"], latitudes, longitudes = crop_forecast_to_domain( + np.median(forecast, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case "mean": + computed["mean"], latitudes, longitudes = crop_forecast_to_domain( + np.mean(forecast, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case "p10": + computed["p10"], latitudes, longitudes = crop_forecast_to_domain( + np.percentile(forecast, 10, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case "p25": + computed["p25"], latitudes, longitudes = crop_forecast_to_domain( + np.percentile(forecast, 25, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case "p75": + computed["p75"], latitudes, longitudes = crop_forecast_to_domain( + np.percentile(forecast, 75, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case "p90": + computed["p90"], latitudes, longitudes = crop_forecast_to_domain( + np.percentile(forecast, 90, axis=0, keepdims=True), + latitudes, + longitudes, + domain_nowcast, + ) + case _: # Defensive check; config parsing validates these values. + raise ValueError(f"Unsupported ensemble statistic: {statistic}") + + return computed, latitudes, longitudes diff --git a/sunflow/geospatial.py b/sunflow/geospatial.py index 1fc1560..2b91f9a 100644 --- a/sunflow/geospatial.py +++ b/sunflow/geospatial.py @@ -234,10 +234,10 @@ def crop_forecast_to_domain( longitudes: np.ndarray, domain_bbox: str, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Crop a [time, lat, lon] forecast and its coordinates to a domain bbox. + """Crop a forecast and its coordinates to a domain bbox. Args: - forecast: Forecast array with shape [time, lat, lon]. + forecast: Forecast array with shape [ensemble, time, lat, lon]. latitudes: 1-D latitude array. longitudes: 1-D longitude array. domain_bbox: Requested bbox string lon_min,lat_min,lon_max,lat_max. @@ -246,12 +246,13 @@ def crop_forecast_to_domain( Tuple (cropped_forecast, cropped_latitudes, cropped_longitudes). Raises: - RuntimeError: If forecast dimensionality is not [time, lat, lon] or the + RuntimeError: If forecast dimensionality is not [ensemble, time, lat, lon] or the requested domain has no overlap with the provided coordinates. """ - if forecast.ndim != 3: + if forecast.ndim != 4: raise RuntimeError( - f"Expected forecast shape [time, lat, lon], got {forecast.shape}." + "Expected forecast shape [ensemble, time, lat, lon], " + f"got {forecast.shape}." ) lon_min, lat_min, lon_max, lat_max = parse_bbox(domain_bbox) @@ -278,7 +279,7 @@ def crop_forecast_to_domain( f"Requested domain_nowcast={domain_bbox} does not overlap forecast grid." ) - cropped_forecast = forecast[:, lat_idx, :][:, :, lon_idx] + cropped_forecast = forecast[:, :, lat_idx, :][:, :, :, lon_idx] return cropped_forecast, latitudes[lat_idx], longitudes[lon_idx] diff --git a/sunflow/main.py b/sunflow/main.py index 4671cdd..d5539d6 100644 --- a/sunflow/main.py +++ b/sunflow/main.py @@ -7,7 +7,6 @@ from enum import Enum import isodate -import numpy as np import yaml from loguru import logger from pysteps.motion.lucaskanade import dense_lucaskanade @@ -22,7 +21,13 @@ save_forecast, ) from .downloaders import download_past_data -from .forecast import multiply_clearsky, preprocess_data, simple_advection_forecast +from .forecast import ( + compute_ensemble_statistics, + multiply_clearsky, + prepend_t0, + preprocess_data, + probabilistic_advection_forecast, +) from .geospatial import ( check_solar_elevation, crop_forecast_to_domain, @@ -38,7 +43,6 @@ validate_config, validate_custom_domain, validate_data_shape, - validate_nowcast_config, validate_run_mode, verify_environment_variables, ) @@ -156,6 +160,20 @@ def parse_datetime_with_timezone(datetime_str: str) -> datetime: help="End of time span in ISO8601 format (inclusive). Use with --start_time.", default=None, ) + parser.add_argument( + "--ensemble_members", + type=int, + default=1, + help="Number of ensemble members (default: 1)", + ) + parser.add_argument( + "--full_ensemble", + action="store_true", + help=( + "Save full ensemble output. By default, ensemble statistics are " + "saved for ensemble runs." + ), + ) args = parser.parse_args() @@ -199,6 +217,7 @@ def run_nowcast( domain_satellite_name: str, nowcast_config: NowcastConfig, s3_config: S3Config, + full_ensemble: bool = False, custom_time: bool = True, ) -> RunResult: """Run a single nowcast for the given (already-rounded) time step. @@ -213,6 +232,8 @@ def run_nowcast( domain_satellite_name: Domain identifier used for input filenames. nowcast_config: NowcastConfig object. s3_config: S3Config object. + full_ensemble: If True, save all ensemble members. If False, + save configured ensemble statistics over ensemble members. custom_time: If True, skip the retry wait loop on missing data. Returns: @@ -309,12 +330,14 @@ def run_nowcast( # Compute motion field motion_field = dense_lucaskanade(ratio_data) - # Simple forecast (ratio forecast) - ratio_forecast = simple_advection_forecast( + # Probabilistic advection forecast (ratio forecast) + ratio_forecast = probabilistic_advection_forecast( ratio_data, motion_field, nowcast_config.future_steps, ens_members=nowcast_config.ens_members, + alpha=nowcast_config.alpha, + beta=nowcast_config.beta, ) # Generate previous day time steps for clearsky lookup @@ -373,23 +396,50 @@ def run_nowcast( config["nc_variable_names"], ) - # Prepend timestep 0: current observation (ratio_data[-1]) × clearsky at t=0 - sds_cs_t0 = clearsky_data.sel(time=clearsky_t0_time.replace(tzinfo=None))[ - config["nc_variable_names"]["sds_cs"] - ].values - solar_t0 = ratio_data[-1] * sds_cs_t0 - solar_forecast = np.concatenate([solar_t0[np.newaxis, :, :], solar_forecast], axis=0) - - solar_forecast, latitudes, longitudes = crop_forecast_to_domain( - solar_forecast, - latitudes, - longitudes, - domain_nowcast, + solar_forecast = prepend_t0( + clearsky_data, ratio_data, solar_forecast, config, clearsky_t0_time ) + if full_ensemble: + output_forecast, latitudes, longitudes = crop_forecast_to_domain( + solar_forecast, + latitudes, + longitudes, + domain_nowcast, + ) + output_mode = "full_ensemble" + logger.info("Saving full ensemble forecast") + else: + if solar_forecast.shape[0] == 1: + output_forecast, latitudes, longitudes = crop_forecast_to_domain( + solar_forecast, + latitudes, + longitudes, + domain_nowcast, + ) + output_mode = "deterministic" + logger.info( + "Saving deterministic forecast " + "(single ensemble member, kept as singleton ensemble dimension)" + ) + else: + output_forecast, latitudes, longitudes = compute_ensemble_statistics( + solar_forecast, + nowcast_config.ensemble_statistics, + latitudes, + longitudes, + domain_nowcast, + ) + output_mode = "ensemble_statistics" + logger.info( + "Saving ensemble statistics across members: " + f"{', '.join(nowcast_config.ensemble_statistics)} " + "(each with singleton ensemble dimension)" + ) + # Save forecast (now contains actual solar irradiance, not ratios) filename = save_forecast( - solar_forecast, + output_forecast, time_step, nowcast_config.future_steps + 1, # +1 for the t=0 analysis step latitudes, @@ -397,6 +447,7 @@ def run_nowcast( dataset_name, nowcast_config, model_version, + output_mode, run_mode, s3_config, ) @@ -442,9 +493,9 @@ def cli() -> None: ) # Load configuration - nowcast_config = NowcastConfig.from_env() - s3_config = S3Config.from_env() args = parse_arguments() + nowcast_config = NowcastConfig.from_env(ensemble_members=args.ensemble_members) + s3_config = S3Config.from_env() run_mode = args.run_mode dataset_name = args.dataset @@ -490,10 +541,33 @@ def cli() -> None: logger.info(f"Using {dataset_name} dataset") logger.info(f"Using satellite domain {domain_satellite_name}: {domain_satellite}") logger.info(f"Using nowcast domain {domain_nowcast_name}: {domain_nowcast}") + logger.info(f"Number of ensemble members: {nowcast_config.ens_members}") + if args.full_ensemble: + logger.info("Output mode: full ensemble") + elif nowcast_config.ens_members == 1: + logger.info("Output mode: deterministic") + else: + logger.info( + "Output mode: ensemble statistics " + f"({', '.join(nowcast_config.ensemble_statistics)})" + ) + logger.info( + f"Using probabilistic advection noise parameters alpha={nowcast_config.alpha}, " + f"beta={nowcast_config.beta}" + ) + + if nowcast_config.ens_members == 1 and ( + nowcast_config.alpha != 0.0 or nowcast_config.beta != 0.0 + ): + logger.warning( + "Running with a single ensemble member, but non-zero probabilistic advection" + "noise parameters alpha and/or beta. This is generally not recommended as it" + "simply adds noise to the nowcast without providing any ensemble spread. " + "Consider setting alpha=0.0 and beta=0.0 for a single-member run." + ) validate_run_mode(run_mode, dataset_name) validate_config(config, dataset_name) - validate_nowcast_config(nowcast_config) verify_environment_variables(run_mode, dataset_name) # Determine the time steps to run @@ -545,6 +619,7 @@ def cli() -> None: domain_satellite_name, nowcast_config, s3_config, + full_ensemble=args.full_ensemble, custom_time=custom_time, ) results.append(result) diff --git a/sunflow/validation.py b/sunflow/validation.py index 3f94c33..7e7ad62 100644 --- a/sunflow/validation.py +++ b/sunflow/validation.py @@ -10,7 +10,6 @@ import xarray as xr from loguru import logger -from .config import NowcastConfig from .geospatial import parse_bbox @@ -69,27 +68,6 @@ def validate_config(config: dict[str, Any], dataset_name: str) -> None: sys.exit(1) -def validate_nowcast_config(nowcast_config: NowcastConfig) -> None: - """Validate that the options selected for the nowcast config are valid. - - Checks the nowcast config created from imported environment variables. - Exits immediately for invalid choices. - - Args: - nowcast_config: Instance of the NowcastConfig class - loaded from environment variables in config.py. - - Raises: - SystemExit: If any invalid choice is detected. - """ - if nowcast_config.ens_members != 1: - logger.error( - f"Invalid nowcast configuration: Currently, only ens_members=1 is supported. " - f"Current value: {nowcast_config.ens_members}. Exiting.\n" - ) - sys.exit(1) - - def validate_run_mode(run_mode: str, dataset_name: str) -> None: """Validate that the run mode is compatible with the dataset.