diff --git a/MANIFEST.in b/MANIFEST.in index 1687be6d..59519138 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -4,6 +4,5 @@ include pysteps/pystepsrc_schema.json include pysteps/io/mch_lut_8bit_Metranet_AZC_V104.txt include pysteps/io/mch_lut_8bit_Metranet_v103.txt recursive-include pysteps *.pyx +recursive-include pysteps/tests/data *.json *.npz *.rst include pyproject.toml - - \ No newline at end of file diff --git a/doc/source/index.rst b/doc/source/index.rst index daf8678a..c31b6555 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -34,6 +34,9 @@ __ https://github.com/pySTEPS/pysteps Gallery <../auto_examples/index> My first nowcast (Colab Notebook) API Reference + Parametric rainfall estimators + Calibrated stochastic rainfall fields + Parametric nowcasting and simulation Example data Configuration file (pystepsrc) Machine learning applications diff --git a/doc/source/pysteps_reference/noise.rst b/doc/source/pysteps_reference/noise.rst index 62ba24f9..b67aadb5 100644 --- a/doc/source/pysteps_reference/noise.rst +++ b/doc/source/pysteps_reference/noise.rst @@ -2,10 +2,11 @@ pysteps.noise ============= -Implementation of deterministic and ensemble nowcasting methods. +Methods for generating spatially correlated noise and velocity perturbations. -.. automodule:: pysteps.noise.interface .. automodule:: pysteps.noise.fftgenerators +.. automodule:: pysteps.noise.interface .. automodule:: pysteps.noise.motion +.. automodule:: pysteps.noise.parametric .. automodule:: pysteps.noise.utils diff --git a/doc/source/pysteps_reference/nowcasts.rst b/doc/source/pysteps_reference/nowcasts.rst index 35d95cbb..cb4287cf 100644 --- a/doc/source/pysteps_reference/nowcasts.rst +++ b/doc/source/pysteps_reference/nowcasts.rst @@ -8,8 +8,9 @@ Implementation of deterministic and ensemble nowcasting methods. .. automodule:: pysteps.nowcasts.interface .. automodule:: pysteps.nowcasts.anvil .. automodule:: pysteps.nowcasts.extrapolation -.. automodule:: pysteps.nowcasts.linda .. automodule:: pysteps.nowcasts.lagrangian_probability +.. automodule:: pysteps.nowcasts.linda +.. automodule:: pysteps.nowcasts.parametric .. automodule:: pysteps.nowcasts.sprog .. automodule:: pysteps.nowcasts.sseps .. automodule:: pysteps.nowcasts.steps diff --git a/doc/source/pysteps_reference/timeseries.rst b/doc/source/pysteps_reference/timeseries.rst index 74e9087f..efb9848d 100644 --- a/doc/source/pysteps_reference/timeseries.rst +++ b/doc/source/pysteps_reference/timeseries.rst @@ -7,4 +7,5 @@ Methods and models for time series analysis. .. automodule:: pysteps.timeseries.autoregression .. automodule:: pysteps.timeseries.correlation +.. automodule:: pysteps.timeseries.parametric diff --git a/doc/source/pysteps_reference/utils.rst b/doc/source/pysteps_reference/utils.rst index 6682dc66..78447824 100644 --- a/doc/source/pysteps_reference/utils.rst +++ b/doc/source/pysteps_reference/utils.rst @@ -16,5 +16,6 @@ Implementation of miscellaneous utility functions. .. automodule:: pysteps.utils.pca .. automodule:: pysteps.utils.reprojection .. automodule:: pysteps.utils.spectral +.. automodule:: pysteps.utils.statistics .. automodule:: pysteps.utils.tapering .. automodule:: pysteps.utils.transformation diff --git a/doc/source/user_guide/calibrated_generator.rst b/doc/source/user_guide/calibrated_generator.rst new file mode 100644 index 00000000..6ecddc00 --- /dev/null +++ b/doc/source/user_guide/calibrated_generator.rst @@ -0,0 +1,121 @@ +Calibrated stochastic rainfall fields +===================================== + +The calibrated generator creates random rainfall patterns from a rain fraction, +a rain-intensity distribution, and spectral slopes describing spatial structure. +It returns a NumPy array in transformed rainfall units. In parametric nowcasting, +these fields provide new spatial variability at each forecast step. + +Here, the marginal distribution means the distribution of values across pixels, +without regard to their positions. Spectral slopes describe how variability is +distributed across spatial scales. See :ref:`sphx_glr_auto_examples_plot_calibrated_noise_generator.py` +for a radar-data example using rainfall statistics, parameter selection, +calibration diagnostics, rainfall distributions, spectra, and normalized +cascade innovations. It uses the same Swiss composite as the original +stochastic-noise tutorial on a complete subdomain. + +Generating a field +------------------ + +Supply signed spectral slopes, rain fraction, and caller-selected Normal location +and scale in transformed rainfall units. The example values illustrate the +interface; the parametric engine estimates and selects these parameters from +the observations:: + + import numpy as np + from pysteps.noise.parametric import generate_calibrated_field + + parameters = { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.08, + "mean_db": 6.0, + "stdev_db": 3.5, + } + field, diagnostics = generate_calibrated_field( + (256, 256), parameters, + kmperpixel=1.0, scale_break=15.0, threshold=0.0, + seed=0, return_diagnostics=True, + ) + assert field.shape == (256, 256) and np.isfinite(field).all() + +``shape`` is in ``(y, x)`` order; pixel spacing and break wavelength are in km. +The threshold is in the field's transformed units, so 1 mm/h rainfall occurrence +corresponds to a 0 dBR threshold. Location and scale are typically latent Normal +parameters, but this helper consumes the active moments supplied by the caller +and performs no latent-versus-ordinary moment selection or regional slope repair. +The input mapping is not modified. + +The default result is the field alone. Optional diagnostics report estimated +spectral slopes, distance, number of attempts, selected attempt, tolerance status, +target and calibrated parameters, settings, and all candidate estimates and +distances. A finite returned field does not imply that the requested spectral +tolerance was met. + +Generation and spectral correction +---------------------------------- + +Independent ``U(0,1)`` samples are FFT-filtered using half the signed power slopes +as amplitude exponents. The resulting correlated field is mapped to the target +rain-intensity distribution. This mapping also changes its spectrum. To +compensate, both filter slopes receive an additive correction, default -0.5, +before generation; the Normal location and scale are used as supplied. + +The histogram marginal mapper uses 500 bins and clips probabilities to +``[0.001, 0.999]``. It selects a wet-area cutoff using a bin's upper edge and maps +the selected sample to Normal quantiles using bin-centre interpolation. The +returned dry anomalies remain continuous. Mapping can change the realised wet +fraction and spectrum, and the clipped tails constrain extremes; occurrence and +marginal matching are approximate. + +Spectra are assessed on a copy with values at or below the transformed rain +threshold replaced by the dry value. The distance is +``sqrt(sum((estimated_beta-target_beta)**2/beta_variances))`` with default +variances ``(0.067738, 0.076405)``. These are empirical weights. Fractions at or +below 0.1 permit at most five attempts; higher fractions use one. The first +candidate with distance at or below 1 is accepted. Otherwise the closest finite +candidate is selected, retaining the first exact tie. If all spectral assessments +are unavailable, the last field is returned with an explicit failure status. + +The correction, weights, acceptance threshold, retry settings, bin count, and +quantile clip can all be overridden. The correction and error weights are +empirical and may need calibration for other rainfall data. + +Random streams and dry cases +---------------------------- + +``seed`` creates a local NumPy random stream without changing global random state. +Alternatively, pass a caller-owned ``rng``; it advances in place and must be +independent for each ensemble member. Supply ``rng`` or ``seed``. Retry counts +affect stream consumption, so reproducibility requires matching parameters, +settings, and numerical implementation as well as the initial seed. + +An explicit zero rain fraction produces a constant dry field with zero attempts +and no spectral-success claim. Positive-fraction mappings with too few distinct +wet pixels raise an error. Missing/nonfinite input and invalid parameters also +raise errors; they are not classified as dry rainfall. + +Cascade innovations and marginal adjustment +------------------------------------------- + +The raw field retains its rainfall marginal. The parametric engine generates it +first, then obtains normalized cascade innovations through decomposition:: + + from pysteps.cascade.bandpass_filters import filter_gaussian + from pysteps.cascade.decomposition import decomposition_fft + + bandpass = filter_gaussian(field.shape, 6) + innovations = decomposition_fft( + field, bandpass, compute_stats=True, normalize=True + ) + assert np.isfinite(innovations["cascade_levels"]).all() + +The raw field is not guaranteed to have global mean zero and variance one. +Existing native STEPS noise-method registrations retain their perturbation +interface; the calibrated parametric engine consumes this helper directly. + +The shared :func:`pysteps.postprocessing.probmatching.match_rainfall_distribution` +also supports the post-AR marginal adjustment. Its default suppresses no positive +rain fraction. The forecast passes ``min_rain_fraction=0.025`` to suppress +very small target fractions after an AR update. Stochastic generation itself +does not apply that cutoff. diff --git a/doc/source/user_guide/parametric_estimators.rst b/doc/source/user_guide/parametric_estimators.rst new file mode 100644 index 00000000..192fee04 --- /dev/null +++ b/doc/source/user_guide/parametric_estimators.rst @@ -0,0 +1,115 @@ +Parametric rainfall estimators +============================== + +The parametric estimators provide NumPy-based spatial statistics and temporal +model utilities. They support the three-observation forecast interface described in :doc:`parametric_nowcast` +and can also be used independently. + +Spatial estimation +------------------ + +Use a finite rainfall field and retain the threshold and dry-value convention +when transforming it. Occurrence and marginal fitting use explicitly separate +cutoffs: one identifies rainy pixels, and the other selects the intensity +samples used for fitting. This example uses illustrative settings:: + + import numpy as np + import xarray as xr + from pysteps.utils.transformation import dB_transform + from pysteps.utils.statistics import compute_field_statistics + from pysteps.utils.spectral import rapsd_logarithmic, fit_power_law_spectrum + + rng = np.random.default_rng(42) + rain = rng.lognormal(1.0, 0.8, size=(256, 256)) + rain[rng.uniform(size=rain.shape) > 0.3] = 0.0 + dataset = xr.Dataset( + {"rain": (("y", "x"), rain, {"units": "mm/h"})}, + attrs={"precip_var": "rain"}, + ) + transformed_dataset = dB_transform(dataset, threshold=1.0, zerovalue=-0.1) + transformed = transformed_dataset["rain"].values + metadata = transformed_dataset["rain"].attrs + statistics = compute_field_statistics( + transformed, metadata["threshold"], lower_truncation=1.25 + ) + frequency, power = rapsd_logarithmic(transformed) + pixel_size_km = 1.0 + scale_break_km = 15.0 + spectrum_fit = fit_power_law_spectrum( + frequency, power, scale_break=scale_break_km / pixel_size_km + ) + +``ordinary_mean`` and ``ordinary_stdev`` describe pixels at or above the +occurrence threshold. ``latent_mean`` and ``latent_stdev`` estimate the parent +Normal from samples at or above the fitting cutoff. They are different estimates; +the fitting cutoff does not change the wet fraction. ``nonzero_fit`` reports +sample count and fitting status, including finite bounded solutions whose +parameters should not be interpreted as an unrestricted interior fit. + +The radial estimator subtracts the field mean, divides by its standard +deviation, and averages unscaled FFT power in logarithmic frequency bins. It differs from +:func:`pysteps.utils.spectral.rapsd`, which uses different binning and +normalization. Its returned arrays contain linear frequencies in cycles per +pixel and mean linear power per Fourier mode, not annulus-integrated energy. +The fit uses signed slopes, an optional prescribed break, and a continuous +piecewise-linear model in 10*log10 units. The two initial bins and frequencies +at or above 10**(-0.4) cycles per pixel are excluded by default. The fixed-break +model is linear in its coefficients and is solved by unweighted least squares. + +Insufficient variance produces empty spectrum arrays. Insufficient or +unidentified fits produce explicit statuses and unavailable coefficients, not +climatological substitutions. Field statistics and spectra reject missing input +pixels; callers must apply a deliberate missing-data policy before using these +complete-grid estimators. Real inputs are converted to float64 for estimation +without being modified. Neither function clips fitted parameters. + +Temporal estimation and modelling +--------------------------------- + +:func:`pysteps.timeseries.parametric.lagrangian_autocorrelation` accepts prepared +cascade histories with shape ``(time, level, y, x)`` and a fixed velocity array +with shape ``(2, y, x)`` in pixels per timestep. It uses successive unit-step +advection. Its default ``boundary="zero"`` fills values outside the grid with +zero after each step and correlates the whole grid. Zero is the mean anomaly +of a normalized cascade. The ``boundary="mask"`` alternative excludes pixels +that are outside the valid support of any aligned history. These two boundary +choices produce different correlation estimates. + +:func:`pysteps.timeseries.parametric.ar2_correlation_time` converts observed lag +pairs into the stationary AR(2) one-sided discrete ACF sum including lag zero. +Output time units equal the supplied timestep units. Invalid and nonstationary +pairs remain unavailable. This quantity is not a forecast-skill horizon. + +Reference-time estimation rescales the selected level times to one reference +wavelength and takes their geometric mean. By default it uses the first three +levels; every contributing time must be positive and finite. Use ``levels=(0,)`` +to estimate from the largest level alone, as the forecast does by default. +Supply the scaling exponent separately. Use nominal wavelengths from the +cascade filter:: + + from pysteps.cascade.bandpass_filters import filter_gaussian + from pysteps.timeseries.parametric import ( + estimate_reference_correlation_time, model_cascade_correlations + ) + + bandpass = filter_gaussian((256, 256), 6, d=1 / pixel_size_km) + wavelengths_km = 1 / bandpass["central_freqs"] + observed_times_minutes = np.array([180.0, 70.0, 40.0, 20.0, 10.0, 5.0]) + exponent = 0.8 # illustrative temporal scaling exponent + reference_minutes = estimate_reference_correlation_time( + observed_times_minutes, wavelengths_km, exponent + ) + correlations, diagnostics = model_cascade_correlations( + reference_minutes, wavelengths_km, exponent, + alpha=2.72, timestep=5.0, return_diagnostics=True + ) + +The first band's nominal wavelength is a label for a broad low-pass response, +not an energy-weighted effective wavelength. The temporal model applies the +nonperiodic lag-2 adjustment when inverting the target timescale. Requested times +at or below one timestep use white-noise lag pairs, achieving one timestep. +Diagnostics retain requested and achieved times and identify sub-timestep limits. + +The estimators leave regional parameter selection and bounds to the caller. +See :doc:`parametric_nowcast` for their use within a complete forecast and +:doc:`calibrated_generator` for generating fields from the estimated parameters. diff --git a/doc/source/user_guide/parametric_nowcast.rst b/doc/source/user_guide/parametric_nowcast.rst new file mode 100644 index 00000000..735493b4 --- /dev/null +++ b/doc/source/user_guide/parametric_nowcast.rst @@ -0,0 +1,384 @@ +Parametric nowcasting and simulation +==================================== + +``pysteps.nowcasts.parametric.forecast`` takes three consecutive rain-rate +observations and returns an ensemble in mm/h. Transformation, motion estimation, +cascade decomposition, statistical estimation, and forecast evolution happen +inside the call. Each ensemble member represents one possible evolution of the +rainfall field. For simulations with prescribed rainfall statistics and motion, +use :func:`pysteps.nowcasts.parametric.simulate` instead. Both interfaces share +the same numerical evolution. + +The method describes rainfall using the fraction of rainy pixels, the +rain-intensity distribution, spectral slopes describing spatial structure, and +correlation times describing persistence at different spatial scales. It holds +these statistics and motion fixed while individual rainfall patterns evolve. +Rain intensity is modelled in dBR, defined as ``10*log10(R)`` for rain rate R in +mm/h. A cascade separates this transformed field into bands of spatial scales; +each band evolves with a second-order autoregressive, or AR(2), model. + +Nowcasting from observations +---------------------------- + +For observations stored in a v2 ``dataset`` with three time slices, ordered +oldest to newest:: + + from pysteps import nowcasts + + forecast = nowcasts.get_method("parametric") + ensemble, diagnostics = forecast( + dataset, + timesteps=12, + n_ens_members=20, + seed=42, + return_diagnostics=True, + ) + rain = ensemble[ensemble.attrs["precip_var"]] + +The precipitation variable must have dimensions ``(time, y, x)``, in any order, +and its name must be stored in ``dataset.attrs["precip_var"]``. Its ``units`` +attribute must be ``mm/h`` and ``transform`` must be absent or ``None``. +Observations must be floating point, finite, and nonnegative. Missing pixels +are rejected, including when the latest observation is dry. + +The spatial coordinates must be regularly spaced with square pixels, ``units`` +of ``m`` or ``km``, and signed ``stepsize`` attributes matching the coordinates. +Descending y coordinates are supported without reversing the data. The three +observation timestamps must be increasing, regularly spaced datetime64 values. +``time.attrs["stepsize"]`` must give their interval in seconds. Pixel size and +cadence are derived from these coordinates. Optional ``kmperpixel`` (km) and +``timestep`` (minutes) overrides must agree with the dataset. In particular, +450-second observations retain their physical interval of 7.5 minutes. + +The result is an xarray dataset whose precipitation has dimensions +``(ens_number, time, y, x)``, with ensemble members numbered from one. It +preserves spatial coordinates and projection metadata and adds a scalar +``forecast_reference_time`` equal to the newest observation time. Forecast +valid times start one observation interval later. Outputs are thresholded +linear rain rates in mm/h, with the occurrence threshold and zero dry value +recorded in their attributes. Untransformed output omits the ``transform`` +attribute, allowing native xarray NetCDF serialization. The numerical engine +loads arrays eagerly. + +Fixed motion can be supplied in ``velocity_x`` and ``velocity_y`` dataset +variables, each with dimensions ``(y, x)`` in any order. Their values are x/y +pixel displacements per observation interval. Velocity units can be omitted, +as in the upstream motion methods, or recorded as ``pixels/timestep``. +Both components must be present together and finite. Otherwise dense +Lucas-Kanade motion is estimated from the latest two transformed observations. +Statistical parameters and motion remain fixed throughout each forecast. + +For the complete Swiss radar nowcasting workflow, including clipping, +aggregation, plotting, and a lossless NetCDF round trip, see +:ref:`sphx_glr_auto_examples_plot_parametric_nowcast_swiss.py`. The common +dataset model is described in :py:mod:`pysteps.io.importers`. + +The Swiss example also writes ``parametric_nowcast_swiss_member0.gif`` into +``rcparams.outputs["path_outputs"]`` (the current directory by default). It +plays the three QPE observations followed by the six forecasts of member index +zero, which has ``ens_number=1``. All frames share a map and colour scale, and +titles identify the data source and valid time. QPE retains measured rain +rates, while forecasts apply the configured occurrence threshold. The frames +show the actual five-minute fields, with no temporal interpolation. + +NetCDF export +-------------- + +The returned dataset can be exported directly with xarray:: + + import xarray as xr + + precip_var = ensemble.attrs["precip_var"] + ensemble.to_netcdf( + "forecast.nc", + engine="netcdf4", + encoding={precip_var: {"zlib": True, "complevel": 1}}, + ) + with xr.open_dataset("forecast.nc", engine="netcdf4") as reopened: + xr.testing.assert_identical(ensemble, reopened.load()) + +This uses lossless compression and retains float64 rainfall, member identifiers, +valid times, forecast reference time, and available geographic metadata. +Projection metadata are preserved as supplied; this does not add missing CF +projection mappings. The older ``pysteps.io.exporters`` interface still expects +arrays and a separate metadata dictionary. + +Callbacks continue to receive one NumPy array per lead. Native NetCDF export +here writes the assembled dataset; incremental writing through the older +exporter remains a separate interface. + +Settings and regional calibration +--------------------------------- + +The defaults provide a starting configuration. Rainfall thresholds, spectral +corrections, temporal scaling, and parameter bounds may need adjustment for the +region, spatial resolution, and observation interval. The low-rain-fraction +replacement values below were derived from Auckland rainfall data; they are +regional estimates rather than universal constants. + +Top-level settings: + +.. list-table:: + :header-rows: 1 + + * - Setting + - Default + - Meaning + * - ``n_cascade_levels`` + - 6 + - Normalized Gaussian cascades + * - ``rain_threshold`` + - 1.0 + - Occurrence threshold, mm/h + * - ``scale_break`` + - 20.0 + - Fixed spectral break, km + * - ``zerovalue`` + - threshold dBR - 0.1 + - Dry transformed value + * - ``taper_sigma`` + - 5.0 + - Boundary-mask Gaussian width, pixels + * - ``min_rain_fraction`` + - 0.025 + - Return dry rainfall for target fractions strictly below this value + +Estimation settings are keyword arguments to ``forecast()``: + +=========================== ======================== ============================== +Key Default Meaning +=========================== ======================== ============================== +``lower_truncation`` 1.25 Latent Normal fit cutoff, dBR +``temporal_exponent`` 0.81 Timescale versus wavelength +``alpha`` 2.72 Pre-adjustment lag-pair power +``reference_levels`` ``(0,)`` Levels used for reference time +``reference_time`` None Optional time prior, minutes +``moment_distance_limit`` 0.5 Latent/ordinary selection +``beta_1_bounds`` (-3.9, just below -1.25) Large-scale slope bounds +``beta_2_min`` -4.0 Small-scale lower bound +``minimum_beta_gap`` 0.1 beta_2 <= beta_1 - gap +``correlation_time_bounds`` (60.0, 720.0) Reference-time bounds, minutes +``low_rain_fraction`` 0.05 Regional fallback below this +``fallback_beta_1`` -1.45 Low-coverage slope replacement +``fallback_beta_2`` -3.253958088759979 Regional median replacement +``fallback_reference_time`` 330.61752959848894 Regional median, minutes +=========================== ======================== ============================== + +Ordinary wet-pixel moments and fitted parent-Normal moments are returned +separately in diagnostics. The latent pair is selected only if finite, with +positive scale, and within the configured normalized Euclidean distance of +the ordinary pair. Otherwise the ordinary pair is used. Rain fraction and +marginal moments are not clipped to a climatology. Below the low-coverage +threshold, both slopes and reference time are replaced together before joint +bounds are applied. The replacement small-scale slope and reference time are +medians from the Auckland archive; all three replacement values are configurable. + +The forecast estimates reference time from the largest cascade level by default. +The standalone estimator averages scale-adjusted times from three levels by +default. Set ``reference_levels=(0, 1, 2)`` to use that approach in a forecast, +and choose a temporal scaling exponent appropriate to the data. Every selected +level must have a valid estimate; failed levels are not dropped automatically. +Timescales use the one-sided discrete AR(2) ACF sum including lag zero; they are +persistence summaries rather than forecast-skill horizons. Nominal wavelengths +come from the actual Gaussian filter. Scale-specific times at or below one +timestep resolve to white noise, with the limitation recorded in diagnostics. + +``noise_kwargs`` accepts ``spectral_bias=-0.5``, ``low_war_threshold=0.1``, +``max_attempts=5``, ``distance_threshold=1.0``, +``beta_variances=(0.067738, 0.076405)``, ``nbins=500``, and ``eps=0.001``. +The last two also control the post-AR marginal mapping. See +:doc:`calibrated_generator` for calibration and finite-grid limitations. + +``motion_kwargs`` overrides dense Lucas-Kanade settings. Mapping values such +as ``lk_kwargs`` or ``fd_kwargs`` update the corresponding default submapping. +The defaults use window (31,31), three pyramid levels, minimum eigenvalue +3e-5, Shi-Tomasi corners (400, quality 0.025, separation 5), IDW interpolation +(8 neighbors, power 1.25), outlier settings 3.5 and 20, and declustering scale 8. +It is estimated from two images, once, with no velocity perturbations. + +Dry and unavailable estimates +----------------------------- + +If the newest field has no rain at the occurrence threshold, the ensemble is +zero, including when earlier observations were wet. Motion/noise generation is +unnecessary and the diagnostic status is ``dry_latest_observation``. Subthreshold +rain counts as dry for this method. Earlier constant fields give zero normalized +anomalies rather than a division by zero. + +Positive rainfall requires at least two distinct wet-pixel values to estimate +its marginal. Failure raises a clear error. If spectral or temporal estimates +are unavailable outside the configured low-coverage fallback, the engine raises +instead of inventing estimates or reporting dry output. For example, partly dry +histories can lack a persistence estimate; users with an investigated temporal +prior can supply ``reference_time=200.0``. The raw failed +estimate remains visible, and bounds still apply to the supplied prior. + +The post-AR fraction cutoff is separate: a valid positive target below 0.025 +produces dry output. Set ``min_rain_fraction=0.0`` to disable that suppression. +Missing data is always an error. + +Configured temporal fallback +---------------------------- + +A supplied temporal scaling model can replace unavailable persistence estimates:: + + ensemble, diagnostics = forecast( + dataset, + timesteps=12, + temporal_fallback={"lzero": 100.0, "tzero": 200.0}, + seed=42, + return_diagnostics=True, + ) + +These example values are illustrative. ``lzero`` is a positive reference +wavelength in km, and ``tzero`` is its positive correlation time in minutes. +If any observed AR(2) correlation time or the estimated reference time is invalid, +the method constructs correlation times for all levels from +``tzero * (wavelength / lzero)**temporal_exponent``. Requested times at or below +one observation interval resolve to white noise, as in the ordinary temporal +model. Diagnostics under ``temporal_selection`` record the selected source, +reason, and configuration. + +Supplying ``temporal_fallback`` disables the combined low-rain-fraction +replacement of both slopes and reference time, even when temporal fallback is +not activated. Spatial estimates must therefore be available. Missing input, +invalid motion, and unavailable marginal estimates remain errors. An activated +fallback uses its configured time without clipping to +``correlation_time_bounds``; spectral bounds still apply. This option cannot be +combined with an explicit ``reference_time`` override. + +Simulation with prescribed parameters +------------------------------------- + +:func:`pysteps.nowcasts.parametric.simulate` uses supplied rainfall parameters +and motion to generate stochastic time series. The initial fields provide the +cascade states; they are not used to estimate the generation parameters. +This is useful for controlled experiments with specified rainfall statistics. +See :ref:`sphx_glr_auto_examples_plot_parametric_simulation.py` for a complete +example that constructs its initial history and a uniform advection field. + +Supply the same three-field dataset used by ``forecast``, including both +``velocity_x`` and ``velocity_y`` variables. Their values are displacements in +pixels per observation interval. Motion is required even for a dry target; +use zero-valued components for no advection. The latest two fields initialize +the AR(2) state. The oldest field is accepted and validated for consistency +with the forecast input interface. + +For an initial ``dataset`` prepared in this way:: + + from pysteps.nowcasts.parametric import simulate + + parameters = { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.2, + "mean_db": 6.0, + "stdev_db": 3.0, + } + ensemble, diagnostics = simulate( + dataset, + timesteps=12, + parameters=parameters, + reference_time=400.0, # minutes + reference_wavelength=256.0, # km + temporal_exponent=0.81, + alpha=2.72, + n_ens_members=1, + seed=42, + return_diagnostics=True, + ) + +The five rainfall parameters are finite real scalars. ``rain_fraction`` must +lie in [0, 1] and ``stdev_db`` must be positive. The two spectral slopes are +signed; ``mean_db`` and ``stdev_db`` are Normal location and scale in dBR. +These are generation targets, distinct from the initial fields' sample moments. +The mapping must contain exactly these five keys. Values are held fixed without +moment selection, regional replacements, or clipping to estimation bounds. +Spectral calibration still uses the configurable controls described in +:doc:`calibrated_generator`. + +Reference time and reference wavelength are required, positive inputs. In this +example, correlation time is 400 minutes at 256 km, and each cascade level uses + +.. math:: + + T_j = 400 \left(\frac{\lambda_j}{256}\right)^{0.81}, + +where :math:`\lambda_j` is its nominal wavelength in km. The defaults are +``temporal_exponent=0.81`` and ``alpha=2.72``. The helper +:func:`pysteps.timeseries.parametric.model_cascade_correlations` finds lag-1 and +lag-2 correlations for these times. The relation +``rho2 = rho1**alpha`` is adjusted where needed for a nonperiodic AR(2) model; +the correlations are then converted to AR coefficients and innovation weights. +These quantities are reported separately in diagnostics. + +Correlation time is the one-sided discrete AR(2) autocorrelation sum including +lag zero. It is not an e-folding time. Requested times at or below the observation +interval use white noise, with the achieved times recorded. Distribution +matching and boundary feedback can change the correlations measured from the +resulting rainfall fields. + +A zero target rain fraction produces dry output even from wet initial fields. +A positive target can generate rain from dry initial fields. Unlike ``forecast``, +``simulate`` defaults to ``min_rain_fraction=0.0`` so positive target fractions +are not suppressed. Missing inputs and failed positive-fraction distribution +mappings remain errors. + +Shifted copies of one stochastic field can provide the initial history, as in +the gallery example. They describe pure translation, so the initial history +is not stationary under the prescribed temporal model. Interpretation of +simulated statistics should account for this initial transient. + +The returned dataset, member numbering, random streams, callbacks, and output +retention follow the forecast conventions below. For compatibility, the output +``forecast_reference_time`` coordinate identifies the newest initial field. +With diagnostics requested, ``parameter_source`` is ``supplied`` and +``model_correlations`` and ``ar_coefficients`` identify the model driving the +simulation. No observation-based fitting diagnostics are produced. + +State evolution and boundaries +------------------------------ + +Each member starts from the newest and previous initial cascade states. +One- and two-step semi-Lagrangian backward trajectories are cached once. The +two-step trajectory retains two unit integration steps. Histories are sampled +with linear interpolation and zero normalized anomalies outside the domain. +Normalized stochastic innovations are then added everywhere, including inflow. +There is no second output-advection pass or periodic wrap of rainfall. + +After normalization across all cascade levels and spatial pixels together, +recomposition uses the innovation cascade's moments. Histogram-based marginal +mapping is applied to the recomposed field. The engine returns its inverse dBR +transform, then multiplies the transformed field by the smoothed one-step +domain mask and decomposes it again to form the next state. This nonlinear +feedback differs from native STEPS, which keeps its AR state separate from +output probability matching. Raw calibrated fields retain a rainfall marginal; +their cascade decomposition supplies normalized innovations. + +All calculations and results are float64, and generated values must be finite. +Input observations and settings are not modified. + +Randomness, diagnostics, and incremental output +----------------------------------------------- + +A local NumPy SeedSequence spawns independent member streams. With the same +seed and settings, the leading members agree when ensemble size increases. +Global NumPy random state and supplied SeedSequence objects are unchanged. +When seed is omitted, diagnostics record the entropy needed to repeat the run. +Numerical implementations, calibration retry decisions, and method settings +are part of reproducibility. + +``return_diagnostics=True`` returns ``(ensemble, diagnostics)`` with effective +settings, seed information, motion source/field, raw marginal and spectral +estimates, observed correlations/times, selected bounded parameters, safeguards, +AR coefficients, and calibration diagnostics indexed ``[member][lead]``. + +A callback receives one complete ``(member, y, x)`` array in mm/h per lead. +Its copy can be modified without affecting state or retained output. Set +``return_output=False`` to return None while streaming; this retains cascade +histories and current fields rather than the full forecast horizon. Diagnostics +can still be requested. Callback exceptions propagate to the caller. + +The method supports integer lead counts at the observation cadence, +fixed motion and parameters, complete grids, and AR(2). Irregular output times, +motion perturbations, and lead-dependent parameter blending remain extensions. diff --git a/examples/plot_calibrated_noise_generator.py b/examples/plot_calibrated_noise_generator.py new file mode 100644 index 00000000..8bce1bbc --- /dev/null +++ b/examples/plot_calibrated_noise_generator.py @@ -0,0 +1,282 @@ +""" +Calibrated stochastic rainfall fields +===================================== + +This example creates random rainfall patterns with statistics similar to a +Swiss radar observation. It estimates the rainy area, rain-intensity +distribution, and spatial structure, then uses them to generate three fields. +The figures compare the generated fields with the observation. + +It uses the same radar image as +:ref:`sphx_glr_auto_examples_plot_noise_generators.py`. The final section shows +how to separate a generated field into spatial scales for use in a nowcast. + +Download the `pysteps example data `_ +and configure the ``mch`` data source in ``pystepsrc``. MeteoSwiss provides the +radar observation used here. This example generates independent spatial fields; +it does not perform temporal nowcasting. +""" + +from pathlib import Path +from pprint import pprint + +import matplotlib.pyplot as plt +import numpy as np + +from pysteps import io, rcparams +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.cascade.decomposition import decomposition_fft +from pysteps.noise.parametric import ( + generate_calibrated_field, + generate_stochastic_field, +) +from pysteps.utils import conversion, dimension, transformation +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.visualization import basemaps, plot_precip_field +from pysteps.xarray_helpers import geodata_from_dataset + +rain_threshold = 1.0 # mm/h +scale_break = 20.0 # km, prescribed rather than estimated +lower_truncation = 1.25 # dBR, separate from the occurrence threshold +seed = 42 +num_realizations = 3 + +############################################################################### +# Read the Swiss radar composite +# ------------------------------ +# +# Use the 11 July 2016, 21:00 UTC image from ``plot_noise_generators.py``. +# Select a complete 300 km subdomain in Swiss LV03 coordinates, then aggregate +# to 2 km. Missing radar coverage is not interpreted as dry rainfall. + +filename = ( + Path(rcparams.data_sources["mch"]["root_path"]) + / "20160711" + / "AQC161932100V_00005.801.gif" +) +dataset = io.import_mch_gif(str(filename), product="AQC", unit="mm", accutime=5.0) + +# Convert to mm/h, select the domain, and average to a 2 km grid. +dataset = conversion.to_rainrate(dataset) +extent = (450000, 750000, 0, 300000) +dataset = dimension.clip_domain(dataset, extent) +dataset = dimension.aggregate_fields_space(dataset, 2000) +precip_var = dataset.attrs["precip_var"] +rain = dataset[precip_var].values +geodata = geodata_from_dataset(dataset) +if not np.isfinite(rain).all() or np.any(rain < 0): + raise ValueError("The example requires complete nonnegative rain-rate data") +kmperpixel = abs(float(dataset.x.attrs["stepsize"])) / 1000.0 + +# Work in dBR = 10*log10(rain rate in mm/h). Extract arrays for the estimators. +transformed_dataset = transformation.dB_transform( + dataset, + threshold=rain_threshold, + zerovalue=10 * np.log10(rain_threshold) - 0.1, +) +observed = transformed_dataset[precip_var].values +threshold = transformed_dataset[precip_var].attrs["threshold"] +zerovalue = transformed_dataset[precip_var].attrs["zerovalue"] + +############################################################################### +# Estimate rainfall statistics +# ---------------------------- +# +# Calculate the fraction of rainy pixels and their mean and standard deviation. +# Also fit a Normal distribution that accounts for the lower fitting cutoff. +# Use its fitted parameters if they are close to the sample estimates; +# otherwise, use the sample estimates. This is the selection used by the nowcast. +# +# Fit spectral slopes on either side of the prescribed 20 km break to describe +# spatial structure. Convert the break to pixels for this estimator. + +statistics = compute_field_statistics(observed, threshold, lower_truncation) +frequency, power = rapsd_logarithmic(observed) +spectrum = fit_power_law_spectrum(frequency, power, scale_break / kmperpixel) +if spectrum["status"] != "ok": + raise ValueError(f"Unavailable spectral estimate: {spectrum['status']}") +ordinary = np.array([statistics["ordinary_mean"], statistics["ordinary_stdev"]]) +latent = np.array([statistics["latent_mean"], statistics["latent_stdev"]]) +if not np.isfinite(ordinary).all() or ordinary[1] <= 0: + raise ValueError("Raining-pixel moments are unavailable") +# Measure the difference between parameter pairs in units of the sample +# standard deviation. The selection limit is configurable in the nowcast. +valid_latent = np.isfinite(latent).all() and latent[1] > 0 +moment_distance = ( + np.linalg.norm(latent - ordinary) / ordinary[1] if valid_latent else np.inf +) +use_latent = valid_latent and moment_distance <= 0.5 +location, scale = latent if use_latent else ordinary +parameters = dict( + beta_1=spectrum["beta_1"], + beta_2=spectrum["beta_2"], + rain_fraction=statistics["rain_fraction"], + mean_db=float(location), + stdev_db=float(scale), +) +print("Ordinary wet-pixel mean/stdev (dBR):", ordinary) +print("Fitted parent-Normal mean/stdev (dBR):", latent) +print("Moment source:", "latent" if use_latent else "ordinary") +print("Generation targets:") +pprint(parameters) + +############################################################################### +# Generate calibrated rainfall fields +# ----------------------------------- +# +# Generate three fields in dBR using separate seeds for repeatable results. +# Matching rain intensities changes spatial structure, so the generator adjusts +# the input spectral slopes and checks the resulting spectrum. Diagnostics +# report whether each field is close enough to its target slopes. +# +# Also generate a field without the correction to compare the spectra later. +# It uses the first calibrated run's initial seed, although retries may make +# that run select a different random field. + +fields, diagnostics = [], [] +for member in range(num_realizations): + field, info = generate_calibrated_field( + rain.shape, + parameters, + kmperpixel, + scale_break, + threshold, + seed=seed + member, + zerovalue=zerovalue, + return_diagnostics=True, + ) + fields.append(field) + diagnostics.append(info) + print( + f"Realization {member + 1}: {info['status']}, " + f"distance={info['distance']:.3f}, attempts={info['attempts']}, " + f"selected={info['selected_attempt']}" + ) +uncorrected = generate_stochastic_field( + rain.shape, + parameters, + kmperpixel, + scale_break, + threshold, + seed=seed, + zerovalue=zerovalue, +) + +# Convert to mm/h for plotting. Keep the original dBR arrays for decomposition. +rain_fields = [] +for field in fields: + generated = transformed_dataset.copy(deep=True) + generated[precip_var].data = field.copy() + generated = transformation.dB_transform(generated, inverse=True) + rain_fields.append(generated[precip_var].values) +observed_rain = transformation.dB_transform(transformed_dataset, inverse=True)[ + precip_var +].values + +# Plot the observation and generated rainfall fields. +fig = plt.figure(figsize=(10, 8)) +for index, (field, title) in enumerate( + zip( + [observed_rain, *rain_fields], + ["Observation", "Realization 1", "Realization 2", "Realization 3"], + ) +): + ax = basemaps.plot_geography( + dataset.attrs["projection"], extent, scale="110m", subplot=(2, 2, index + 1) + ) + plot_precip_field(field, ax=ax, geodata=geodata, axis="off", title=title) +plt.tight_layout() +plt.show() + +############################################################################### +# Compare rainfall occurrence and intensity distributions +# ------------------------------------------------------- +# +# For each rain-rate threshold, plot the fraction of all pixels at or above it. +# At 1 mm/h, this is the rainy fraction. Distribution matching is approximate, +# so generated fractions and intensities can differ from their targets. + +colors = ["tab:blue", "tab:orange", "tab:green"] +plotted_fields = [observed_rain, *rain_fields] +levels = np.geomspace(rain_threshold, max(f.max() for f in plotted_fields), 100) +fig, ax = plt.subplots() +for index, (field, label) in enumerate( + zip(plotted_fields, ["Observed", "Calibrated 1", "Calibrated 2", "Calibrated 3"]) +): + ordered = np.sort(field.ravel()) + probability = ( + ordered.size - np.searchsorted(ordered, levels, side="left") + ) / ordered.size + ax.semilogx( + levels, + probability, + color="black" if index == 0 else colors[index - 1], + label=label, + ) +ax.set_xlabel("Rain-rate threshold (mm/h)") +ax.set_ylabel("Fraction of pixels at or above threshold") +ax.set_title("Rainfall occurrence and intensity") +ax.legend() +ax.grid(True, alpha=0.25) +plt.show() + +############################################################################### +# Compare the realized radial spectra +# ----------------------------------- +# +# Compare how variability is distributed across spatial scales. As in the +# calibration check, replace values at or below the rain threshold on a copy. +# The estimator subtracts the mean and divides by the standard deviation before +# averaging FFT power over directions; it does not describe rainfall orientation. + +fig, ax = plt.subplots() +ax.loglog(kmperpixel / frequency, power, color="black", label="Observed") +for index, (field, label) in enumerate( + zip( + [uncorrected, *fields], + ["Without correction", "Calibrated 1", "Calibrated 2", "Calibrated 3"], + ) +): + assessment = field.copy() + assessment[assessment <= threshold] = zerovalue + frequencies, powers = rapsd_logarithmic(assessment) + ax.loglog( + kmperpixel / frequencies, + powers, + color="gray" if index == 0 else colors[index - 1], + linestyle="--" if index == 0 else "-", + label=label, + ) +ax.axvline(scale_break, color="black", linestyle=":", label="Prescribed break") +ax.invert_xaxis() +ax.set_xlabel("Wavelength (km)") +ax.set_ylabel("Radial mean FFT power of standardized dBR") +ax.set_title("Spectral assessment after marginal mapping") +ax.legend() +ax.grid(True, alpha=0.25) +plt.show() + +############################################################################### +# Separate the generated field into spatial scales +# ------------------------------------------------ +# +# A cascade splits the field into bands of spatial scales. Each band is then +# normalized to mean zero and standard deviation one. The nowcast uses these +# normalized bands as random innovations in its autoregressive update. +# Decompose the dBR field directly, retaining its values below the rain threshold. + +bandpass = filter_gaussian(fields[0].shape, 6) +cascade = decomposition_fft(fields[0], bandpass, compute_stats=True, normalize=True) +innovations = cascade["cascade_levels"] +print("Innovation means:", innovations.mean(axis=(1, 2))) +print("Innovation standard deviations:", innovations.std(axis=(1, 2))) +fig, axes = plt.subplots(1, 2, figsize=(9, 4), constrained_layout=True) +for ax, level in zip(axes, [0, 5]): + image = ax.imshow(innovations[level], cmap="RdBu_r", vmin=-3, vmax=3) + ax.set_title(f"Normalized cascade level {level + 1}") + ax.set_axis_off() +fig.colorbar(image, ax=axes, label="Standardized anomaly", shrink=0.8) +plt.show() + +# sphinx_gallery_thumbnail_number = 1 diff --git a/examples/plot_parametric_nowcast_swiss.py b/examples/plot_parametric_nowcast_swiss.py new file mode 100644 index 00000000..77ab6892 --- /dev/null +++ b/examples/plot_parametric_nowcast_swiss.py @@ -0,0 +1,293 @@ +""" +Parametric nowcast with Swiss radar data +======================================== + +This tutorial uses the same Swiss radar sequence as the STEPS nowcast example. +Three rain-rate observations are passed directly to the parametric engine, +which performs transformation, motion estimation, parameter estimation, and +ensemble generation internally. + +Download the `pysteps example data `_ +and configure the ``mch`` data source in ``pystepsrc`` before running this example. +The radar observations are provided by MeteoSwiss. +""" + +from datetime import datetime +from pprint import pprint +from pathlib import Path +from tempfile import TemporaryDirectory + +import matplotlib.pyplot as plt +from matplotlib.animation import FuncAnimation, PillowWriter +import numpy as np +import xarray as xr + +from pysteps import io, nowcasts, rcparams +from pysteps.utils import conversion, dimension +from pysteps.visualization import basemaps, plot_precip_field +from pysteps.xarray_helpers import geodata_from_dataset + +# Set ensemble size, forecast length, and a seed for repeatable results. +n_ens_members = 20 +n_leadtimes = 6 +seed = 24 +rain_threshold = 1.0 # mm/h +date = datetime(2017, 1, 31, 12, 0) + +############################################################################### +# Read three observations in rain-rate units +# ------------------------------------------ +# +# Load the observations at 11:50, 11:55, and 12:00 UTC using the usual pySTEPS +# data-source configuration. The imported accumulation is converted to mm/h. + +source = rcparams.data_sources["mch"] +timestep = source["timestep"] # minutes between observations +filenames = io.find_by_date( + date, + source["root_path"], + source["path_fmt"], + source["fn_pattern"], + source["fn_ext"], + timestep, + num_prev_files=2, +) +importer = io.get_method(source["importer"], "importer") +observations = io.read_timeseries(filenames, importer, **source["importer_kwargs"]) +observations = conversion.to_rainrate(observations) +precip_var = observations.attrs["precip_var"] + +############################################################################### +# Prepare a complete spatial domain +# --------------------------------- +# +# The full Swiss composite has missing coverage outside the radar footprint. +# The parametric method requires complete finite grids. Select a 300 km +# square subdomain with complete coverage at all three times, in Swiss LV03 +# coordinates (metres). Clip at the native resolution, then aggregate to 2 km +# to limit computation and memory. Missing pixels are not replaced with zero. + +extent = (450000, 750000, 0, 300000) # left, right, bottom, top in metres +observations = dimension.clip_domain(observations, extent) +observations = dimension.aggregate_fields_space(observations, 2000) +rain = observations[precip_var] +if ( + observations.sizes["time"] != 3 + or not np.isfinite(rain.values).all() + or (rain.values < 0).any() +): + raise ValueError("The tutorial requires three complete nonnegative rain fields") +geodata = geodata_from_dataset(observations) + + +# Use the same geographic metadata for observations and forecasts. A coarse +# basemap limits the auxiliary Natural Earth map data needed by the example. +def make_map(subplot=(1, 1, 1)): + """Create the geographic axes before plotting the precipitation raster.""" + return basemaps.plot_geography( + observations.attrs["projection"], extent, scale="110m", subplot=subplot + ) + + +# Inspect the newest observation. +plt.figure() +plot_precip_field( + rain.isel(time=-1).values, + ax=make_map(), + geodata=geodata, + axis="off", + title="Swiss radar observation, 31 January 2017, 12:00 UTC", +) +plt.show() + +############################################################################### +# Generate the parametric ensemble +# -------------------------------- +# +# Pass linear rain rates directly: no external dBR transform or motion estimate +# is needed. The returned dataset contains rain rates in mm/h with dimensions +# (ens_number, time, y, x). Its first valid time is 12:05 and its last is 12:30. +# Grid cell size and time step are derived from the observation coordinates. +# Motion and rainfall statistics are estimated once and held fixed. Each member's +# rainfall pattern evolves separately using its own random perturbations. +# Setting the seed makes these results reproducible. +# +# The defaults include a 1 mm/h rain threshold and a 20 km spectral break. +# This example demonstrates the workflow; regional calibration and forecast +# skill require a separate assessment. + +forecast = nowcasts.get_method("parametric") +ensemble, diagnostics = forecast( + observations, + timesteps=n_leadtimes, + n_ens_members=n_ens_members, + rain_threshold=rain_threshold, + seed=seed, + return_diagnostics=True, +) +forecast_rain = ensemble[ensemble.attrs["precip_var"]] +print("Ensemble dimensions:", dict(forecast_rain.sizes)) +print("Generation parameters:") +pprint(diagnostics["parameters"]) +print("Marginal moment source:", diagnostics["safeguards"]["moment_source"]) +print("Persistence definition:", diagnostics["temporal_model"]["definition"]) + +############################################################################### +# Plot the ensemble mean +# ---------------------- +# +# Averaging realizations smooths rainfall features whose position or intensity +# differs between members. Individual members retain more spatial variability. + +lead_minutes = float( + (ensemble.time.values[-1] - ensemble.forecast_reference_time.values) + / np.timedelta64(1, "m") +) +ensemble_mean = forecast_rain.isel(time=-1).mean("ens_number", skipna=False) +plt.figure() +plot_precip_field( + ensemble_mean.values, + ax=make_map(), + geodata=geodata, + axis="off", + title=f"Ensemble mean parametric nowcast (+{lead_minutes:g} min)", +) +plt.show() + +############################################################################### +# Plot four ensemble members +# -------------------------- + +fig = plt.figure(figsize=(10, 8)) +for member in range(4): + ax = make_map((2, 2, member + 1)) + ax = plot_precip_field( + forecast_rain.isel(ens_number=member, time=-1).values, + ax=ax, + geodata=geodata, + colorbar=False, + axis="off", + ) + ax.set_title(f"Member {member + 1} (+{lead_minutes:g} min)") +plt.tight_layout() +plt.show() + +############################################################################### +# Derive exceedance probabilities +# ------------------------------- +# +# Estimate the probability of rain at or above 1 mm/h as the fraction of members +# meeting that threshold. Use the same threshold as the forecast. + +probability_threshold = 1.0 # mm/h +probabilities = (forecast_rain.isel(time=-1) >= probability_threshold).mean( + "ens_number" +) +plt.figure() +plot_precip_field( + probabilities.values, + ax=make_map(), + geodata=geodata, + axis="off", + ptype="prob", + units="mm/h", + probthr=probability_threshold, + title=f"Probability of rain >= {probability_threshold:g} mm/h (+{lead_minutes:g} min)", +) +plt.show() + +############################################################################### +# Animate observations followed by the first ensemble member +# ---------------------------------------------------------- +# +# Show the three radar observations followed by six forecasts on the same map +# and colour scale. Array index zero selects ensemble member number one. +# Display the valid times and pause at the last observation and forecast. + +animation_member_index = 0 +animation_fields = np.concatenate( + [rain.values, forecast_rain.isel(ens_number=animation_member_index).values] +) +animation_times = np.concatenate([observations.time.values, ensemble.time.values]) +animation_observations = observations.sizes["time"] +animation_member_number = int(ensemble.ens_number.values[animation_member_index]) +animation_figure = plt.figure(figsize=(7, 7)) +animation_axis = make_map() +plot_precip_field(animation_fields[0], ax=animation_axis, geodata=geodata, axis="off") +animation_raster = animation_axis.images[-1] +animation_figure.text( + 0.5, + 0.025, + f"Forecast occurrence threshold: {rain_threshold:g} mm/h. QPE retains measured rates.", + ha="center", + fontsize=9, +) + + +def update_animation(frame): + """Update rainfall and its timestamp while retaining geography and colours.""" + animation_raster.set_data(animation_fields[frame]) + valid_time = np.datetime_as_string(animation_times[frame], unit="m").replace( + "T", " " + ) + if frame < animation_observations: + label = "QPE observation" + else: + lead = float( + (animation_times[frame] - ensemble.forecast_reference_time.values) + / np.timedelta64(1, "m") + ) + label = ( + f"Forecast member index {animation_member_index} " + f"(ens_number={animation_member_number}), +{lead:g} min" + ) + animation_axis.set_title(f"{label}\n{valid_time} UTC") + return (animation_raster, animation_axis.title) + + +animation_frames = ( + list(range(animation_observations)) + + [animation_observations - 1] + + list(range(animation_observations, len(animation_times))) + + [len(animation_times) - 1] * 2 +) +animation = FuncAnimation( + animation_figure, + update_animation, + frames=animation_frames, + interval=800, + repeat=True, + cache_frame_data=False, +) +animation_path = ( + Path(rcparams.outputs["path_outputs"]).expanduser() + / "parametric_nowcast_swiss_member0.gif" +) +animation_path.parent.mkdir(parents=True, exist_ok=True) +animation.save(animation_path, writer=PillowWriter(fps=1.25), dpi=110) +plt.close(animation_figure) +print("QPE and first-member animation:", animation_path.resolve()) + +############################################################################### +# Export the complete forecast dataset +# ------------------------------------- +# +# Save the dataset with lossless compression and check that reopening preserves +# its values, coordinates, and attributes. Choose a persistent output path to +# keep the file. Projection metadata are retained as supplied; export does not +# add missing CF projection mappings. + +with TemporaryDirectory(prefix="pysteps-parametric-swiss-") as output_directory: + output_path = Path(output_directory) / "parametric_nowcast_swiss.nc" + ensemble.to_netcdf( + output_path, + engine="netcdf4", + encoding={precip_var: {"zlib": True, "complevel": 1}}, + ) + with xr.open_dataset(output_path, engine="netcdf4") as reopened: + xr.testing.assert_identical(ensemble, reopened.load()) + print( + "NetCDF round trip preserved all forecast values, coordinates, and attributes." + ) + +# sphinx_gallery_thumbnail_number = 3 diff --git a/examples/plot_parametric_simulation.py b/examples/plot_parametric_simulation.py new file mode 100644 index 00000000..7e573e97 --- /dev/null +++ b/examples/plot_parametric_simulation.py @@ -0,0 +1,424 @@ +""" +Simulate a time series of rainfall fields +========================================= + +Generate rainfall patterns that move and change with time using prescribed +rainfall statistics. This example sets the rain fraction, intensity distribution, +spatial structure, and temporal persistence explicitly. It runs without radar +observations. + +The correlation time is 400 minutes at a wavelength of 256 km. Smaller rainfall +features have shorter correlation times and change more quickly. A uniform +motion field moves the rain one pixel east every five minutes. +""" + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.cascade.decomposition import decomposition_fft +from pysteps.noise.parametric import generate_calibrated_field +from pysteps.nowcasts.parametric import simulate +from pysteps.timeseries.parametric import ( + ar2_correlation_time, + estimate_reference_correlation_time, + lagrangian_autocorrelation, + model_cascade_correlations, +) +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.utils.transformation import dB_transform + +############################################################################### +# Choose rainfall and temporal parameters +# --------------------------------------- +# +# Describe the rainy fraction and intensity distribution in dBR, where +# dBR = 10*log10(rain rate in mm/h). The two spectral slopes describe spatial +# variability on either side of the 20 km break. These parameters stay fixed; +# the positions and intensities of individual rainfall features evolve. + +shape = (128, 128) +kmperpixel = 1.0 +timestep = 5.0 # minutes +n_cascade_levels = 6 +rain_threshold = 1.0 # mm/h +scale_break = 20.0 # km +parameters = dict( + beta_1=-2.4, + beta_2=-3.1, + rain_fraction=0.2, + mean_db=6.0, + stdev_db=3.0, +) +reference_time = 400.0 # minutes +reference_wavelength = 256.0 # km +temporal_exponent = 0.81 +alpha = 2.72 + +############################################################################### +# Calculate lag correlations at each spatial scale +# ------------------------------------------------ +# +# Separate the field into six bands of spatial scales, called cascade levels. +# Assign a correlation time to each level using +# T = 400 * (wavelength / 256)**0.81, with wavelengths in km and times in minutes. +# The temporal helper finds lag-1 and lag-2 correlations for each time using +# rho2 = rho1**alpha, adjusted where necessary for a nonperiodic AR(2) process. +# Correlation time here is a discrete autocorrelation sum including lag zero. + +bandpass = filter_gaussian(shape, n_cascade_levels) +wavelengths = kmperpixel / bandpass["central_freqs"] +lag_correlations, temporal_diagnostics = model_cascade_correlations( + reference_time, + wavelengths, + temporal_exponent, + alpha, + timestep, + reference_wavelength=reference_wavelength, + return_diagnostics=True, +) +print("Prescribed temporal model (wavelength in km, time in minutes):") +print(f"{'Wavelength':>12} {'Time':>12} {'Lag-1':>12} {'Lag-2':>12}") +for row in np.column_stack( + [wavelengths, temporal_diagnostics["requested_times"], lag_correlations] +): + print(" ".join(f"{value:12.4f}" for value in row)) + +############################################################################### +# Construct an initial history and uniform motion +# ----------------------------------------------- +# +# Generate one random rainfall field, then create two earlier fields by shifting +# it west. np.roll wraps the array edges to construct this artificial history. +# The simulation itself uses open boundaries, with random variability added at +# inflow. Shifted copies describe pure translation; the initial history is not +# a stationary realization of the prescribed temporal model. + +transformed = generate_calibrated_field( + shape, + parameters, + kmperpixel, + scale_break, + threshold=10 * np.log10(rain_threshold), + seed=9, +) +rain = 10 ** (transformed / 10) +rain[rain < rain_threshold] = 0 +initial_fields = np.stack([np.roll(rain, -2, axis=1), np.roll(rain, -1, axis=1), rain]) + +# Attach physical coordinates and the supplied velocity. Rainfall is in mm/h; +# velocity is in pixels per interval, and the time stepsize attribute is seconds. +dataset = xr.Dataset( + { + "rain": (("time", "y", "x"), initial_fields, {"units": "mm/h"}), + "velocity_x": (("y", "x"), np.ones(shape), {"units": "pixels/timestep"}), + "velocity_y": (("y", "x"), np.zeros(shape), {"units": "pixels/timestep"}), + }, + coords={ + "time": ( + "time", + np.array( + ["2026-01-01T00:00", "2026-01-01T00:05", "2026-01-01T00:10"], + dtype="datetime64[s]", + ), + {"stepsize": timestep * 60}, + ), + "x": ( + "x", + (np.arange(shape[1]) + 0.5) * kmperpixel, + {"units": "km", "stepsize": kmperpixel}, + ), + "y": ( + "y", + (np.arange(shape[0]) + 0.5) * kmperpixel, + {"units": "km", "stepsize": kmperpixel}, + ), + }, + attrs={"precip_var": "rain"}, +) + +############################################################################### +# Simulate one hour of rainfall +# ----------------------------- +# +# simulate() computes the same lag correlations shown above, then converts them +# to AR coefficients and innovation weights. At each step, it advects the two +# previous cascade states and combines them with a new random field. Rainfall +# distribution matching and boundary tapering feed into the next state. +# The supplied parameters and velocity are used throughout, without estimation. + +ensemble, diagnostics = simulate( + dataset, + timesteps=12, + parameters=parameters, + reference_time=reference_time, + reference_wavelength=reference_wavelength, + temporal_exponent=temporal_exponent, + alpha=alpha, + n_cascade_levels=n_cascade_levels, + rain_threshold=rain_threshold, + scale_break=scale_break, + n_ens_members=1, + seed=42, + return_diagnostics=True, +) +simulated_rain = ensemble[ensemble.attrs["precip_var"]].isel(ens_number=0) + +############################################################################### +# Recover parameters from the simulated rainfall +# ---------------------------------------------- +# +# Estimate statistics from all twelve generated fields of the first realization. +# Exclude the initial field and its shifted copies. Use the returned rain rates, +# transformed back to dBR, rather than the internal cascade states or generation +# diagnostics. Dry output pixels receive the same transformed dry value as the +# initial data. +# +# Estimate the fraction of rainy pixels and fit the parent Normal distribution +# above 1.25 dBR. Its fitted location and scale can be compared with mean_db and +# stdev_db. Ordinary wet-pixel moments describe the truncated sample and are +# different quantities. Fit both spectral slopes with the prescribed break. + +threshold_db = 10 * np.log10(rain_threshold) +lower_truncation = 1.25 # dBR, fitting cutoff distinct from the rain threshold +transformed_ensemble = dB_transform( + ensemble, threshold=rain_threshold, zerovalue=threshold_db - 0.1 +) +validation_fields = ( + transformed_ensemble[ensemble.attrs["precip_var"]].isel(ens_number=0).values +) +field_estimates = [] +marginal_statuses = [] +spectral_statuses = [] +cascades = [] +for field in validation_fields: + statistics = compute_field_statistics(field, threshold_db, lower_truncation) + frequency, power = rapsd_logarithmic(field) + spectrum = fit_power_law_spectrum( + frequency, power, scale_break=scale_break / kmperpixel + ) + field_estimates.append( + [ + statistics["rain_fraction"], + statistics["latent_mean"], + statistics["latent_stdev"], + spectrum["beta_1"], + spectrum["beta_2"], + ] + ) + marginal_statuses.append(statistics["nonzero_fit"]["status"]) + spectral_statuses.append(spectrum["status"]) + cascades.append( + decomposition_fft(field, bandpass, compute_stats=True, normalize=True)[ + "cascade_levels" + ] + ) +field_estimates = np.asarray(field_estimates) +cascades = np.stack(cascades) + +# Show variability between fields alongside the target. These fields are +# temporally correlated: their standard deviation is not a confidence interval. +# Retain unavailable estimates as NaN and report finite counts; do not replace +# them with the targets or silently discard failed fits. +labels = ( + "Rain fraction", + "Normal location (dBR)", + "Normal scale (dBR)", + "Large-scale slope", + "Small-scale slope", +) +targets = [ + parameters[name] + for name in ("rain_fraction", "mean_db", "stdev_db", "beta_1", "beta_2") +] +print(f"\nParameter recovery from {len(validation_fields)} fields of one realization:") +print( + f"{'Parameter':<23} {'Target':>12} {'Mean':>12} {'SD':>12} {'Finite / total':>15}" +) +for label, target, values in zip(labels, targets, field_estimates.T): + count = np.count_nonzero(np.isfinite(values)) + print( + f"{label:<23} {target:12.4f} {np.mean(values):12.4f} " + f"{np.std(values):12.4f} {str(count) + '/' + str(len(values)):>15}" + ) +for name, statuses in [ + ("Marginal", marginal_statuses), + ("Spectral", spectral_statuses), +]: + for status in sorted(set(statuses)): + print(f"{name} fit status {status}: {statuses.count(status)}/{len(statuses)}") + +############################################################################### +# Assess noisy spectral estimates with Monte Carlo sampling +# --------------------------------------------------------- +# +# A spectral slope fitted to one field can be noisy, especially at large scales +# represented by few Fourier modes. Use a Monte Carlo comparison to assess the +# spectral targets: generate independent initial fields and simulate independent +# one-hour realizations with the same parameters. +# +# Average the fitted slopes over the twelve fields within each run, then report +# the mean, standard deviation, and 5th/95th percentiles across independent runs. +# This treats each run as one replicate; its correlated frames are not separate +# independent samples. The percentile range describes the spread of run means, +# not a confidence interval for the overall mean. Calibration and nonlinear +# evolution can still leave systematic differences from the target slopes. + +n_monte_carlo = 32 +monte_carlo_slopes = [] +for run_seed in np.random.SeedSequence(2026).spawn(n_monte_carlo): + initial_seed, evolution_seed = run_seed.spawn(2) + mc_transformed = generate_calibrated_field( + shape, parameters, kmperpixel, scale_break, threshold_db, seed=initial_seed + ) + mc_rain = 10 ** (mc_transformed / 10) + mc_rain[mc_rain < rain_threshold] = 0 + mc_history = dataset.copy(deep=True) + mc_history[dataset.attrs["precip_var"]].data = np.stack( + [np.roll(mc_rain, -2, axis=1), np.roll(mc_rain, -1, axis=1), mc_rain] + ) + mc_ensemble = simulate( + mc_history, + timesteps=len(validation_fields), + parameters=parameters, + reference_time=reference_time, + reference_wavelength=reference_wavelength, + temporal_exponent=temporal_exponent, + alpha=alpha, + n_cascade_levels=n_cascade_levels, + rain_threshold=rain_threshold, + scale_break=scale_break, + n_ens_members=1, + seed=evolution_seed, + ) + mc_fields = ( + dB_transform( + mc_ensemble, threshold=rain_threshold, zerovalue=threshold_db - 0.1 + )[ensemble.attrs["precip_var"]] + .isel(ens_number=0) + .values + ) + mc_slopes = [] + for mc_field in mc_fields: + frequency, power = rapsd_logarithmic(mc_field) + spectrum = fit_power_law_spectrum( + frequency, power, scale_break=scale_break / kmperpixel + ) + mc_slopes.append([spectrum["beta_1"], spectrum["beta_2"]]) + # A failed field fit leaves its run mean unavailable; do not silently drop it. + monte_carlo_slopes.append(np.mean(mc_slopes, axis=0)) +monte_carlo_slopes = np.asarray(monte_carlo_slopes) +print(f"\nMonte Carlo spectral recovery: {n_monte_carlo} independent runs") +print( + f"{'Parameter':<23} {'Target':>12} {'MC mean':>12} {'MC SD':>12} " + f"{'P05':>12} {'P95':>12} {'Finite runs':>12}" +) +for label, name, values in zip( + ["Large-scale slope", "Small-scale slope"], + ["beta_1", "beta_2"], + monte_carlo_slopes.T, +): + lower, upper = np.percentile(values, [5, 95]) + count = np.count_nonzero(np.isfinite(values)) + print( + f"{label:<23} {parameters[name]:12.4f} {np.mean(values):12.4f} " + f"{np.std(values):12.4f} {lower:12.4f} {upper:12.4f} " + f"{str(count) + '/' + str(n_monte_carlo):>12}" + ) + +############################################################################### +# Compare prescribed and estimated temporal correlations +# ------------------------------------------------------ +# +# Align each overlapping three-field history with the known velocity and +# estimate its spatial lag-1 and lag-2 correlations at every cascade level. +# Use the same zero-anomaly boundary convention as the forecast estimator. +# Average the lag pairs, then convert that average to an implied AR(2) time. +# These are estimates from the displayed realization, separate from the Monte +# Carlo spectral check. One hour is short relative to the 400-minute reference +# time, and shifted initial fields introduce a transient. Distribution matching +# and boundary feedback can also change the measured correlations. + +velocity = np.stack([dataset.velocity_x.values, dataset.velocity_y.values]) +measured_correlations = np.stack( + [ + lagrangian_autocorrelation(cascades[end - 2 : end + 1], velocity) + for end in range(2, len(cascades)) + ] +) +average_correlations = np.mean(measured_correlations, axis=0) +estimated_times = ar2_correlation_time( + average_correlations[:, 0], average_correlations[:, 1], timestep +) +print( + f"\nTemporal recovery from {len(measured_correlations)} overlapping histories " + "(km and minutes):" +) +print( + f"{'Wavelength':>12} {'Model T':>12} {'Implied T':>12} " + f"{'Model lag-1':>12} {'Est. lag-1':>12} " + f"{'Model lag-2':>12} {'Est. lag-2':>12}" +) +for row in np.column_stack( + [ + wavelengths, + temporal_diagnostics["achieved_times"], + estimated_times, + lag_correlations[:, 0], + average_correlations[:, 0], + lag_correlations[:, 1], + average_correlations[:, 1], + ] +): + print(" ".join(f"{value:12.4f}" for value in row)) + +# Rescale the first three level times to 256 km with the prescribed exponent. +# All three must be valid; the estimator returns NaN rather than dropping one. +# The spectral break, temporal exponent, and alpha are held fixed in this check, +# not independently recovered from the simulation. +estimated_reference_time = estimate_reference_correlation_time( + estimated_times, + wavelengths, + temporal_exponent, + levels=(0, 1, 2), + reference_wavelength=reference_wavelength, +) +print( + f"\n{'Reference scale (km)':<23} {'Target time (min)':>18} {'Implied time (min)':>19}" +) +print( + f"{reference_wavelength:<23.4f} {reference_time:18.4f} {estimated_reference_time:19.4f}" +) + +# The implied times are diagnostic extrapolations from two lags, rather than +# direct measurements of long correlation times. Longer runs and a separate +# Monte Carlo temporal assessment are needed to assess long-run persistence. + +############################################################################### +# Show the evolution of one realization +# ------------------------------------- +# +# Plot the initial field and five times from the same realization. Use a common +# colour scale to see the translation and changes in rainfall structure. + +steps = (0, 1, 3, 6, 9, 12) +fig, axes = plt.subplots( + 2, 3, figsize=(11, 7), sharex=True, sharey=True, constrained_layout=True +) +for ax, step in zip(axes.flat, steps): + field = rain if step == 0 else simulated_rain.isel(time=step - 1).values + image = ax.imshow( + field, + origin="lower", + vmin=0, + vmax=20, + cmap="Blues", + extent=(0, shape[1] * kmperpixel, 0, shape[0] * kmperpixel), + ) + ax.set_title("Initial field" if step == 0 else f"+{step * timestep:g} min") + ax.set_xlabel("x (km)") + ax.set_ylabel("y (km)") +fig.colorbar(image, ax=axes, label="Rain rate (mm/h)", shrink=0.8) +plt.show() diff --git a/pysteps/noise/__init__.py b/pysteps/noise/__init__.py index 7be00590..0e11951d 100644 --- a/pysteps/noise/__init__.py +++ b/pysteps/noise/__init__.py @@ -5,4 +5,4 @@ """ from .interface import get_method -from . import utils, motion, fftgenerators +from . import utils, motion, fftgenerators, parametric diff --git a/pysteps/noise/parametric.py b/pysteps/noise/parametric.py new file mode 100644 index 00000000..2d142903 --- /dev/null +++ b/pysteps/noise/parametric.py @@ -0,0 +1,440 @@ +""" +pysteps.noise.parametric +======================== + +Statistically parameterized rainfall fields with inverse spectral calibration. + +.. autosummary:: + :toctree: ../generated/ + + power_law_filter + generate_stochastic_field + generate_calibrated_field +""" + +from collections.abc import Mapping + +import numpy as np + +from pysteps.postprocessing.probmatching import match_rainfall_distribution +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic + +GENERATOR_ID = "inverse_calibrated_v2" +SPECTRAL_BIAS = -0.5 +LOW_WAR_THRESHOLD = 0.1 +MAX_CALIBRATION_ATTEMPTS = 5 +SPECTRAL_DISTANCE_THRESHOLD = 1.0 +BETA_VARIANCES = (0.067738, 0.076405) + + +def _finite(name, value): + """Validate and return a finite real scalar.""" + if not np.isscalar(value) or np.iscomplexobj(value) or not np.isfinite(value): + raise ValueError(f"{name} must be a finite real scalar") + return float(value) + + +def _shape(shape): + """Validate and return two positive grid dimensions in (y, x) order.""" + try: + shape = tuple(shape) + except TypeError: + raise ValueError( + "shape must contain two positive integer dimensions (y, x)" + ) from None + if len(shape) != 2 or any( + not isinstance(n, (int, np.integer)) or n < 1 for n in shape + ): + raise ValueError("shape must contain two positive integer dimensions (y, x)") + return tuple(int(n) for n in shape) + + +def _parameters(parameters): + """Copy and validate the five rainfall-generation parameters.""" + if not isinstance(parameters, Mapping): + raise ValueError("parameters must be a mapping") + names = ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + if any(name not in parameters for name in names): + raise ValueError(f"parameters must contain {', '.join(names)}") + values = {name: _finite(name, parameters[name]) for name in names} + if not 0 <= values["rain_fraction"] <= 1: + raise ValueError("rain_fraction must lie in [0, 1]") + if values["stdev_db"] <= 0: + raise ValueError("stdev_db must be positive") + return values + + +def _random_generator(rng, seed): + """Return the supplied random stream or create a local seeded stream.""" + if rng is not None: + if seed is not None: + raise ValueError("supply rng or seed, not both") + if not isinstance(rng, (np.random.Generator, np.random.RandomState)): + raise ValueError("rng must be a NumPy Generator or RandomState") + return rng + return np.random.default_rng(seed) + + +def power_law_filter(shape, beta_1, beta_2=None, kmperpixel=1.0, scale_break=None): + """Construct a full-FFT amplitude filter from signed power-spectrum slopes. + + Parameters + ---------- + shape: tuple of int + Output shape (y, x), each dimension positive. + beta_1, beta_2: float + Signed power slopes, power proportional to frequency**beta. The second + slope is optional; provide beta_2 and scale_break together. + kmperpixel: float, optional + Positive isotropic pixel spacing in km (default 1). + scale_break: float, optional + Positive break wavelength in km. beta_1 applies below its reciprocal + frequency, beta_2 at and above it, with continuity at the break. + + Returns + ------- + filter: ndarray + Finite full-FFT amplitude filter of shape (y, x), using half the power + exponents. DC amplitude equals one. The reference frequency is the + smallest nonzero frequency along the longer grid dimension. + """ + shape = _shape(shape) + beta_1 = _finite("beta_1", beta_1) + spacing = _finite("kmperpixel", kmperpixel) + if spacing <= 0: + raise ValueError("kmperpixel must be positive") + if (beta_2 is None) != (scale_break is None): + raise ValueError("provide beta_2 and scale_break together") + # Power is squared amplitude, so the filter uses half the power exponent. + fy = np.fft.fftfreq(shape[0], d=spacing) + fx = np.fft.fftfreq(shape[1], d=spacing) + radius = np.sqrt(fy[:, None] ** 2 + fx[None, :] ** 2) + reference = 1 / (max(shape) * spacing) + result = np.ones(shape) + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + if beta_2 is None: + selected = radius > 0 + result[selected] = (radius[selected] / reference) ** (beta_1 / 2) + else: + beta_2 = _finite("beta_2", beta_2) + scale_break = _finite("scale_break", scale_break) + if scale_break <= 0: + raise ValueError("scale_break must be positive") + breakpoint = 1 / scale_break + low = radius < breakpoint + result[low] = (radius[low] / reference) ** (beta_1 / 2) + result[~low] = (breakpoint / reference) ** (beta_1 / 2) * ( + radius[~low] / breakpoint + ) ** (beta_2 / 2) + result[radius == 0] = 1.0 + if not np.isfinite(result).all(): + raise ValueError("parameters produce a nonfinite spectral filter") + return result + + +def _field_from_filter( + spectrum_filter, parameters, threshold, rng, zerovalue, nbins, eps +): + """Generate a correlated field and match its rainfall distribution. + + Apply the full-FFT amplitude filter to uniform white noise, then map the + result to the supplied rain fraction and Normal intensity parameters. + Return a (y, x) array in the same transformed units as threshold; values + below the threshold remain continuous for later cascade decomposition. + """ + # Filtering introduces spatial correlation before intensity matching. + white_noise = rng.uniform(low=0, high=1, size=spectrum_filter.shape) + filtered = np.fft.ifft2( + np.fft.fft2(white_noise, spectrum_filter.shape) * spectrum_filter + ).real + return match_rainfall_distribution( + filtered, + parameters["rain_fraction"], + parameters["mean_db"], + parameters["stdev_db"], + threshold, + zerovalue=zerovalue, + nbins=nbins, + eps=eps, + ) + + +def generate_stochastic_field( + shape, + parameters, + kmperpixel, + scale_break, + threshold, + *, + rng=None, + seed=None, + zerovalue=None, + nbins=500, + eps=0.001, +): + """Generate a random rainfall field without spectral bias correction. + + Parameters + ---------- + shape: tuple of int + Output array shape (y, x). + parameters: mapping + Finite beta_1, beta_2, rain_fraction, mean_db, and positive stdev_db. + The last two are caller-selected parent-Normal generation moments. + The mapping is not modified; no regional parameter selection or clipping occurs. + kmperpixel, scale_break: float + Positive pixel spacing and spectral break wavelength, both in km. + threshold: float + Finite rain/no-rain threshold in transformed units (dBR when used for + parametric nowcasting). The returned field has the same transformed units. + rng: numpy.random.Generator or RandomState, optional + Caller-owned random stream; advanced in place. + seed: int or SeedSequence, optional + Seed for a local default_rng when rng is absent. Supply only one of them. + zerovalue: float, optional + Finite dry value below threshold, default threshold-0.1. + nbins, eps: int, float, optional + Marginal-mapping bins and quantile clip (defaults 500 and 0.001). + + Returns + ------- + field: ndarray + Float64 field of shape (y, x). Independent U(0,1) inputs are spectrally + filtered, then mapped to the requested occurrence and Normal marginal. + Dry anomalies remain continuous. This field is not globally standardized + or converted to linear rainfall. Occurrence and marginal matching are + approximate. An explicit zero fraction returns a constant dry field. + """ + parameters = _parameters(parameters) + threshold = _finite("threshold", threshold) + spectrum_filter = power_law_filter( + shape, parameters["beta_1"], parameters["beta_2"], kmperpixel, scale_break + ) + rng = _random_generator(rng, seed) + if parameters["rain_fraction"] == 0: + # Validate mapping settings without consuming a random draw. + return match_rainfall_distribution( + np.zeros(spectrum_filter.shape), + 0, + parameters["mean_db"], + parameters["stdev_db"], + threshold, + zerovalue=zerovalue, + nbins=nbins, + eps=eps, + ) + return _field_from_filter( + spectrum_filter, parameters, threshold, rng, zerovalue, nbins, eps + ) + + +def _estimate_beta(field, threshold, zerovalue, scale_break_pixels): + """Estimate two spectral slopes after replacing subthreshold values. + + Use a copy for assessment so the generated field retains its continuous + values below threshold. The break wavelength is in pixels. Return two + NaNs if the spectrum cannot be estimated or fitted. + """ + truncated = field.copy() + truncated[truncated <= threshold] = zerovalue + try: + frequency, power = rapsd_logarithmic(truncated) + result = fit_power_law_spectrum(frequency, power, scale_break_pixels) + except (RuntimeError, TypeError, ValueError, np.linalg.LinAlgError): + return np.nan, np.nan + return float(result["beta_1"]), float(result["beta_2"]) + + +def generate_calibrated_field( + shape, + parameters, + kmperpixel, + scale_break, + threshold, + *, + rng=None, + seed=None, + zerovalue=None, + spectral_bias=SPECTRAL_BIAS, + low_war_threshold=LOW_WAR_THRESHOLD, + max_attempts=MAX_CALIBRATION_ATTEMPTS, + distance_threshold=SPECTRAL_DISTANCE_THRESHOLD, + beta_variances=BETA_VARIANCES, + nbins=500, + eps=0.001, + return_diagnostics=False, +): + """Generate a random rainfall field with spectral bias correction. + + Parameters + ---------- + shape: tuple of int + Output shape (y, x), with at least four pixels along each axis. + parameters: mapping + Target beta_1, beta_2, rain_fraction, mean_db, and stdev_db, as in + :func:`generate_stochastic_field`. The mapping is not modified. + kmperpixel, scale_break: float + Positive pixel spacing and spectral break wavelength, both in km. + threshold: float + Rain/no-rain threshold in transformed units, also used for assessment. + rng: numpy.random.Generator or RandomState, optional + Caller-owned random stream, advanced in place. Supply rng or seed. + seed: int or SeedSequence, optional + Seed for a local random stream when rng is absent. + zerovalue: float, optional + Dry value below threshold, default threshold-0.1. + spectral_bias: float, optional + Additive correction to both input power slopes (default -0.5). The Normal + location and scale are used as supplied. + low_war_threshold: float, optional + Fractions at or below this value permit retries (default 0.1). + max_attempts: int, optional + Maximum attempts for low fractions (default 5). Higher fractions use one. + distance_threshold: float, optional + Accept when sqrt(sum((estimated_beta-target_beta)**2/beta_variances)) is + at or below this nonnegative tolerance (default 1). + beta_variances: array_like, optional + Two positive finite calibration error variances. Defaults are + (0.067738, 0.076405); these weights are empirical, not universal errors. + nbins: int, optional + Number of distribution-matching histogram bins (default 500). + eps: float, optional + Normal-quantile probability clip (default 0.001). + return_diagnostics: bool, optional + Also return selected spectral estimates, distance, attempted/selected + candidate indices, tolerance status, targets, effective calibration + settings, and per-candidate estimates and distances. + + Returns + ------- + field: ndarray + Selected float64 rainfall-marginal field, not standardized. Slopes are + assessed on a copy with values <= threshold replaced by zerovalue. + diagnostics: dict, optional + First in-tolerance candidate is selected. Otherwise use the closest + finite-distance candidate, retaining the first tie. If all spectra are + unavailable, use the last candidate and report unavailable_spectrum. + + Notes + ----- + Distribution matching can change a filtered field's spectrum. The additive + slope correction compensates for this effect empirically; its default and + the error variances may need calibration for other rainfall data. + Calibration need not meet its tolerance; diagnostics report the outcome. + A confirmed zero target returns dry output with zero attempts and no + spectral-success claim. Invalid parameters or an unavailable + positive-fraction marginal mapping raise errors. + This function does not perform moment selection or slope clipping. + """ + parameters = _parameters(parameters) + shape = _shape(shape) + if min(shape) < 4: + raise ValueError("calibrated shape must have at least four pixels per axis") + threshold = _finite("threshold", threshold) + spectral_bias = _finite("spectral_bias", spectral_bias) + low_war_threshold = _finite("low_war_threshold", low_war_threshold) + distance_threshold = _finite("distance_threshold", distance_threshold) + if not 0 <= low_war_threshold <= 1 or distance_threshold < 0: + raise ValueError( + "low_war_threshold must lie in [0,1] and distance_threshold must be nonnegative" + ) + if not isinstance(max_attempts, (int, np.integer)) or max_attempts < 1: + raise ValueError("max_attempts must be a positive integer") + if np.iscomplexobj(beta_variances): + raise ValueError("beta_variances must contain real values") + variances = np.asarray(beta_variances, dtype=float) + if ( + variances.shape != (2,) + or not np.isfinite(variances).all() + or np.any(variances <= 0) + ): + raise ValueError("beta_variances must contain two positive finite values") + zerovalue = ( + threshold - 0.1 if zerovalue is None else _finite("zerovalue", zerovalue) + ) + # Adjust input slopes before mapping, then assess against the original targets. + calibrated = dict( + parameters, + beta_1=parameters["beta_1"] + spectral_bias, + beta_2=parameters["beta_2"] + spectral_bias, + ) + spectrum_filter = power_law_filter( + shape, calibrated["beta_1"], calibrated["beta_2"], kmperpixel, scale_break + ) + rng = _random_generator(rng, seed) + settings = dict( + kmperpixel=float(kmperpixel), + scale_break=float(scale_break), + threshold=threshold, + zerovalue=zerovalue, + spectral_bias=spectral_bias, + low_war_threshold=low_war_threshold, + max_attempts=int(max_attempts), + distance_threshold=distance_threshold, + beta_variances=variances.copy(), + nbins=nbins, + eps=eps, + ) + estimates, distances = [], [] + + def output(candidate, status, tolerance_achieved): + field, beta_1, beta_2, distance, selected = candidate + diagnostics = dict( + generator=GENERATOR_ID, + status=status, + estimated_beta_1=beta_1, + estimated_beta_2=beta_2, + distance=distance, + attempts=len(distances), + selected_attempt=selected, + tolerance_achieved=tolerance_achieved, + target_parameters=parameters.copy(), + calibrated_parameters=calibrated.copy(), + settings=settings, + candidate_estimates=np.asarray(estimates).reshape(-1, 2), + candidate_distances=np.asarray(distances), + ) + return (field, diagnostics) if return_diagnostics else field + + if parameters["rain_fraction"] == 0: + field = match_rainfall_distribution( + np.zeros(shape), + 0, + parameters["mean_db"], + parameters["stdev_db"], + threshold, + zerovalue=zerovalue, + nbins=nbins, + eps=eps, + ) + return output((field, np.nan, np.nan, np.nan, 0), "dry", False) + # Retry sparse fields, whose realised spectra vary more between draws. + best = last = None + limit = max_attempts if parameters["rain_fraction"] <= low_war_threshold else 1 + for attempt in range(1, limit + 1): + field = _field_from_filter( + spectrum_filter, calibrated, threshold, rng, zerovalue, nbins, eps + ) + beta_1, beta_2 = _estimate_beta( + field, threshold, zerovalue, scale_break / kmperpixel + ) + error = np.array([beta_1 - parameters["beta_1"], beta_2 - parameters["beta_2"]]) + distance = ( + float(np.sqrt(np.sum(error**2 / variances))) + if np.isfinite(error).all() + else np.inf + ) + estimates.append([beta_1, beta_2]) + distances.append(distance) + last = (field, beta_1, beta_2, distance, attempt) + if distance <= distance_threshold: + return output(last, "tolerance_achieved", True) + if np.isfinite(distance) and (best is None or distance < best[3]): + best = last + # If tolerance is missed, keep the closest measurable candidate. If none + # can be assessed, return the last field with an unavailable-spectrum status. + return output( + best if best is not None else last, + "best_available" if best is not None else "unavailable_spectrum", + False, + ) diff --git a/pysteps/nowcasts/_cascade_advection.py b/pysteps/nowcasts/_cascade_advection.py new file mode 100644 index 00000000..f54920a0 --- /dev/null +++ b/pysteps/nowcasts/_cascade_advection.py @@ -0,0 +1,73 @@ +"""Fixed-flow trajectory caching for Eulerian cascade evolution. + +Zero outside the domain represents a zero normalized anomaly. Innovations +are added after advection, including at inflow pixels. +""" + +import numpy as np +from scipy.ndimage import map_coordinates + +from pysteps.extrapolation.semilagrangian import extrapolate + + +def prepare_advection_coordinates(velocity, ar_order=2): + """Prepare backward trajectories for sampling previous cascade states. + + Parameters + ---------- + velocity: array_like + Fixed x and y displacements with shape (2, y, x), in pixels per + observation interval. All values must be finite. + ar_order: {1, 2}, optional + Number of previous states to advect (default 2). + + Returns + ------- + coordinates: dict + Read-only arrays of shape (2, y, x), keyed by integer lag. Components + are in (y, x) order for scipy.ndimage.map_coordinates. The two-step + trajectory uses two unit integration steps with the same velocity. + """ + velocity = np.asarray(velocity, dtype=float) + if ( + velocity.ndim != 3 + or velocity.shape[0] != 2 + or min(velocity.shape[1:]) < 2 + or not np.isfinite(velocity).all() + ): + raise ValueError("velocity must be finite with shape (2, y, x)") + if ar_order not in (1, 2): + raise ValueError("ar_order must be 1 or 2") + yy, xx = np.indices(velocity.shape[1:]) + xy = np.stack([xx, yy]) + coordinates = {} + for lag in range(1, ar_order + 1): + _, displacement = extrapolate(None, velocity, lag, return_displacement=True) + # Extrapolation returns (x, y) displacements; SciPy samples in (y, x). + coordinates[lag] = (xy + displacement)[::-1].copy() + coordinates[lag].setflags(write=False) + return coordinates + + +def advect_field(field, coordinates): + """Advect a normalized cascade field using backward sampling coordinates. + + Parameters + ---------- + field: ndarray + Finite two-dimensional normalized cascade field in (y, x) order. + coordinates: ndarray + Sampling positions with shape (2, y, x), in (y, x) component order. + + Returns + ------- + advected: ndarray + Linearly interpolated field. Samples outside the grid are zero, + representing the cascade level's mean. Random innovations are added + separately during forecasting, including at these inflow pixels. + """ + if not np.isfinite(field).all(): + raise ValueError("field contains non-finite values") + return map_coordinates( + field, coordinates, order=1, mode="constant", cval=0.0, prefilter=False + ) diff --git a/pysteps/nowcasts/interface.py b/pysteps/nowcasts/interface.py index 3f87dbf6..0ba67745 100644 --- a/pysteps/nowcasts/interface.py +++ b/pysteps/nowcasts/interface.py @@ -5,24 +5,23 @@ Interface for the nowcasts module. It returns a callable function for computing nowcasts. -The methods in the nowcasts module implement the following interface: +The v2 nowcasting interface takes an xarray dataset: - ``forecast(precip, velocity, timesteps, **keywords)`` + ``forecast(dataset, timesteps, **keywords)`` -where precip is a (m,n) array with input precipitation field to be advected and -velocity is a (2,m,n) array containing the x- and y-components of -the m x n advection field. -timesteps can be an integer or a list. An integer specifies the number of time -steps to forecast, where the output time step is taken from the inputs. -Irregular time steps can be given in a list. -The interface accepts optional keyword arguments specific to the given method. +The dataset contains precipitation, spatial coordinates, observation times, +projection metadata and, when supplied, velocity_x and velocity_y variables. +The global ``precip_var`` attribute identifies the precipitation variable. +The output is an xarray.Dataset with dimensions (time, y, x) for deterministic methods +or (ens_number, time, y, x) for ensemble methods. Method-specific forecast +functions document their input history and optional settings. -The output depends on the type of the method. -For deterministic methods, the output is a three-dimensional array of shape -(num_timesteps,m,n) containing a time series of nowcast precipitation fields. -For stochastic methods that produce an ensemble, the output is a -four-dimensional array of shape (num_ensemble_members,num_timesteps,m,n). -The time step of the output is taken from the inputs. +The ``parametric`` method requires exactly three complete, untransformed +rain-rate observations in mm/h on a regular grid with square pixels. Pixel +size and time step are derived from coordinates; the time coordinate's +``stepsize`` attribute is in seconds. +It estimates motion if velocity variables are absent and returns rain rates in +mm/h. It accepts integer lead counts at the observation interval. .. autosummary:: :toctree: ../generated/ @@ -37,6 +36,7 @@ extrapolation, lagrangian_probability, linda, + parametric, sprog, sseps, steps, @@ -51,6 +51,7 @@ _nowcast_methods["lagrangian"] = extrapolation.forecast _nowcast_methods["lagrangian_probability"] = lagrangian_probability.forecast _nowcast_methods["linda"] = linda.forecast +_nowcast_methods["parametric"] = parametric.forecast _nowcast_methods["probability"] = lagrangian_probability.forecast _nowcast_methods["sprog"] = sprog.forecast _nowcast_methods["sseps"] = sseps.forecast @@ -81,6 +82,9 @@ def get_method(name): +-----------------+-------------------------------------------------------+ | linda | the LINDA method developed in :cite:`PCN2021` | +-----------------+-------------------------------------------------------+ + | parametric | three rain-rate observations to a parametric ensemble | + | | with internal parameter and motion estimation | + +-----------------+-------------------------------------------------------+ | lagrangian\_ | this approach computes local lagrangian probability | | probability | forecasts of threshold exceedences | +-----------------+-------------------------------------------------------+ diff --git a/pysteps/nowcasts/parametric.py b/pysteps/nowcasts/parametric.py new file mode 100644 index 00000000..4cffc355 --- /dev/null +++ b/pysteps/nowcasts/parametric.py @@ -0,0 +1,1300 @@ +""" +pysteps.nowcasts.parametric +=========================== + +Generate rainfall time series from estimated or prescribed parameters. + +.. autosummary:: + :toctree: ../generated/ + + forecast + simulate +""" + +from collections.abc import Mapping +from copy import deepcopy + +import numpy as np +import xarray as xr +from scipy.ndimage import gaussian_filter + +from pysteps import motion +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.cascade.decomposition import decomposition_fft, recompose_fft +from pysteps.noise.parametric import generate_calibrated_field, GENERATOR_ID +from pysteps.postprocessing.probmatching import match_rainfall_distribution +from pysteps.timeseries.autoregression import estimate_ar_params_yw +from pysteps.timeseries.parametric import ( + ar2_correlation_time, + estimate_reference_correlation_time, + lagrangian_autocorrelation, + model_cascade_correlations, +) +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.utils.transformation import dB_transform +from pysteps.xarray_helpers import convert_output_to_xarray_dataset + +from ._cascade_advection import advect_field, prepare_advection_coordinates + +_MOTION_DEFAULTS = dict( + lk_kwargs=dict(winsize=(31, 31), nr_levels=3, min_eig_thr=3e-5), + fd_method="shitomasi", + fd_kwargs=dict(max_corners=400, quality_level=0.025, min_distance=5), + interp_method="idwinterp2d", + interp_kwargs=dict(k=8, power=1.25), + nr_std_outlier=3.5, + k_outlier=20, + decl_scale=8, +) +_NOISE_DEFAULTS = dict( + spectral_bias=-0.5, + low_war_threshold=0.1, + max_attempts=5, + distance_threshold=1.0, + beta_variances=(0.067738, 0.076405), + nbins=500, + eps=0.001, +) + + +def _finite(name, value): + """Validate and return a finite real scalar.""" + if not np.isscalar(value) or np.iscomplexobj(value): + raise ValueError(f"{name} must be a finite real scalar") + try: + valid = np.isfinite(value) + except TypeError: + valid = False + if not valid: + raise ValueError(f"{name} must be a finite real scalar") + return float(value) + + +def _positive(name, value): + """Validate and return a positive real scalar.""" + value = _finite(name, value) + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def _integer(name, value, minimum=1): + """Validate an integer at or above the specified minimum.""" + if ( + isinstance(value, (bool, np.bool_)) + or not isinstance(value, (int, np.integer)) + or value < minimum + ): + raise ValueError(f"{name} must be an integer >= {minimum}") + return int(value) + + +def _options(name, supplied, defaults): + """Copy defaults and apply overrides, rejecting unknown setting names.""" + if supplied is None: + return deepcopy(defaults) + if not isinstance(supplied, Mapping): + raise ValueError(f"{name} must be a mapping") + unknown = set(supplied) - set(defaults) + if unknown: + raise ValueError(f"Unknown {name}: {sorted(unknown)}") + result = deepcopy(defaults) + result.update(deepcopy(dict(supplied))) + return result + + +def _estimation_settings(settings, threshold, n_levels): + """Validate estimation settings against the threshold and cascade count.""" + settings = deepcopy(settings) + for key in settings: + if key not in ( + "reference_levels", + "reference_time", + "beta_1_bounds", + "correlation_time_bounds", + ): + settings[key] = _finite(key, settings[key]) + for key in ("alpha", "minimum_beta_gap", "fallback_reference_time"): + _positive(key, settings[key]) + if settings["moment_distance_limit"] < 0: + raise ValueError("moment_distance_limit must be nonnegative") + if not 0 <= settings["low_rain_fraction"] <= 1: + raise ValueError("low_rain_fraction must lie in [0, 1]") + if settings["lower_truncation"] < threshold: + raise ValueError("lower_truncation must be at or above the dBR threshold") + for key in ("beta_1_bounds", "correlation_time_bounds"): + bounds = np.asarray(settings[key], float) + if ( + bounds.shape != (2,) + or not np.isfinite(bounds).all() + or bounds[0] > bounds[1] + ): + raise ValueError(f"{key} must be two finite increasing bounds") + settings[key] = tuple(float(v) for v in bounds) + if settings["correlation_time_bounds"][0] <= 0: + raise ValueError("correlation_time_bounds must be positive") + if ( + settings["beta_2_min"] + > settings["beta_1_bounds"][0] - settings["minimum_beta_gap"] + ): + raise ValueError("beta_2_min and beta_1_bounds cannot satisfy minimum_beta_gap") + levels = np.asarray(settings["reference_levels"]) + if ( + levels.ndim != 1 + or not levels.size + or not np.issubdtype(levels.dtype, np.integer) + or np.any(levels < 0) + or np.any(levels >= n_levels) + or np.unique(levels).size != levels.size + ): + raise ValueError("reference_levels must contain distinct valid level indices") + settings["reference_levels"] = tuple(int(v) for v in levels) + if settings["reference_time"] is not None: + settings["reference_time"] = _positive( + "reference_time", settings["reference_time"] + ) + return settings + + +def _noise_settings(supplied): + """Validate spectral-calibration and distribution-matching settings.""" + settings = _options("noise_kwargs", supplied, _NOISE_DEFAULTS) + for key in ("spectral_bias", "low_war_threshold", "distance_threshold", "eps"): + settings[key] = _finite(key, settings[key]) + if ( + not 0 <= settings["low_war_threshold"] <= 1 + or settings["distance_threshold"] < 0 + ): + raise ValueError( + "noise fractions must lie in [0, 1] and distance must be nonnegative" + ) + if not 0 < settings["eps"] < 0.5: + raise ValueError("noise eps must lie in (0, 0.5)") + settings["max_attempts"] = _integer("max_attempts", settings["max_attempts"]) + settings["nbins"] = _integer("nbins", settings["nbins"], 2) + variances = np.asarray(settings["beta_variances"], float) + if ( + variances.shape != (2,) + or not np.isfinite(variances).all() + or np.any(variances <= 0) + ): + raise ValueError("beta_variances must contain two positive finite values") + settings["beta_variances"] = tuple(float(v) for v in variances) + return settings + + +def _decompose(field, bp_filter): + """Normalize each cascade level, using zero anomalies for constant levels.""" + result = decomposition_fft(field, bp_filter, compute_stats=True, normalize=False) + levels = result["cascade_levels"] + if np.ptp(field) == 0: + levels.fill(0) + result["stds"] = [0.0] * len(levels) + result["normalized"] = True + return result + for k, (mean, stdev) in enumerate(zip(result["means"], result["stds"])): + levels[k] = (levels[k] - mean) / stdev if stdev > 0 else 0 + result["normalized"] = True + return result + + +def _generation_parameters(statistics, spectrum, reference_time, settings): + """Select intensity parameters and constrain slopes and persistence. + + Prefer fitted Normal parameters when close to the wet-pixel sample moments. + For low rain fractions, replace both spectral slopes and reference time + together. Apply the slope and time bounds afterwards, retaining the raw + estimates and each selection in the returned audit dictionary. + """ + mu, sigma = statistics["ordinary_mean"], statistics["ordinary_stdev"] + if not np.isfinite(mu) or not np.isfinite(sigma) or sigma <= 0: + raise ValueError( + "Cannot estimate rainfall marginal: newest field needs at least two " + "distinct raining values; constant/single-pixel rain is not dry" + ) + latent = np.array([statistics["latent_mean"], statistics["latent_stdev"]]) + valid = np.isfinite(latent).all() and latent[1] > 0 + distance = ( + float(np.hypot(latent[0] - mu, latent[1] - sigma) / sigma) if valid else None + ) + use_latent = valid and distance <= settings["moment_distance_limit"] + selected = dict( + rain_fraction=statistics["rain_fraction"], + mean_db=float(latent[0] if use_latent else mu), + stdev_db=float(latent[1] if use_latent else sigma), + beta_1=spectrum["beta_1"], + beta_2=spectrum["beta_2"], + reference_time=reference_time, + ) + original = selected.copy() + low_war = selected["rain_fraction"] < settings["low_rain_fraction"] + if low_war: + selected.update( + beta_1=settings["fallback_beta_1"], + beta_2=settings["fallback_beta_2"], + reference_time=settings["fallback_reference_time"], + ) + if not all( + np.isfinite(selected[k]) for k in ("beta_1", "beta_2", "reference_time") + ): + raise ValueError( + "Unavailable spectral/persistence estimates outside the low-rainfall " + "fallback; inspect observations, scale_break, or reference_time" + ) + if selected["reference_time"] <= 0: + raise ValueError("Estimated reference_time must be positive") + before_bounds = selected.copy() + selected["beta_1"] = float(np.clip(selected["beta_1"], *settings["beta_1_bounds"])) + selected["beta_2"] = float( + np.clip( + selected["beta_2"], + settings["beta_2_min"], + selected["beta_1"] - settings["minimum_beta_gap"], + ) + ) + selected["reference_time"] = float( + np.clip(selected["reference_time"], *settings["correlation_time_bounds"]) + ) + audit = dict( + moment_source="latent" if use_latent else "ordinary", + moment_distance=distance, + moment_reason=( + "within_limit" + if use_latent + else ("distance" if valid else "invalid_latent") + ), + low_rain_fraction_fallback=low_war, + before_safeguards=original, + before_bounds=before_bounds, + bounded_parameters=[k for k in selected if selected[k] != before_bounds[k]], + ) + return selected, audit + + +def _advance( + histories, + parameters, + phi, + bp_filter, + coordinates, + rng, + threshold, + zerovalue, + kmperpixel, + scale_break, + noise_settings, + min_rain_fraction, +): + """Advance one member by one observation interval. + + Parameters + ---------- + histories: sequence of dict + Two normalized cascade decompositions, newest first. Each contains + cascade_levels with shape (level, y, x). Histories are not modified. + parameters: mapping + Selected rain fraction, Normal intensity parameters, and spectral slopes. + phi: ndarray + Array of shape (level, 3): lag-1 coefficient, lag-2 coefficient, and + innovation coefficient returned by estimate_ar_params_yw. + bp_filter: dict + Gaussian cascade filter used for both histories and random fields. + coordinates: dict + Backward sampling coordinates for lags 1 and 2, in (y, x) order. + rng: numpy.random.Generator + Member-specific random stream, advanced in place. + threshold, zerovalue: float + Rain occurrence threshold and dry value in dBR. + kmperpixel, scale_break: float + Pixel spacing and spectral break wavelength in km. + noise_settings: dict + Calibration and histogram-mapping options. + min_rain_fraction: float + Target fraction below which distribution matching returns dry output. + + Returns + ------- + field: ndarray + Forecast field in dBR with shape (y, x), before boundary tapering. + diagnostics: dict + Calibration results for the random field used in this update. + """ + # Generate rainfall variability, then normalize its cascade levels to obtain + # the innovations used in the autoregressive update. + noise, diagnostics = generate_calibrated_field( + bp_filter["shape"], + parameters, + kmperpixel, + scale_break, + threshold, + rng=rng, + zerovalue=zerovalue, + return_diagnostics=True, + **noise_settings, + ) + innovation = _decompose(noise, bp_filter) + levels = np.empty_like(innovation["cascade_levels"]) + # Bring both previous states to the forecast time before the AR(2) update. + # Zero anomalies outside the grid allow innovations to populate inflow. + for k in range(len(levels)): + levels[k] = ( + phi[k, 0] * advect_field(histories[0]["cascade_levels"][k], coordinates[1]) + + phi[k, 1] + * advect_field(histories[1]["cascade_levels"][k], coordinates[2]) + + phi[k, 2] * innovation["cascade_levels"][k] + ) + # Normalize over all levels and spatial pixels together. This couples the + # levels and differs from normalizing each updated level separately. + stdev = np.std(levels) + if stdev > 0.1: + levels = (levels - np.mean(levels)) / stdev + updated = dict( + domain="spatial", + normalized=True, + compact_output=False, + cascade_levels=levels, + means=innovation["means"], + stds=innovation["stds"], + ) + # Recompose with the innovation moments and restore the target rainfall + # distribution. The caller feeds this adjusted field into the next update. + field = match_rainfall_distribution( + recompose_fft(updated), + parameters["rain_fraction"], + parameters["mean_db"], + parameters["stdev_db"], + threshold, + zerovalue=zerovalue, + min_rain_fraction=min_rain_fraction, + nbins=noise_settings["nbins"], + eps=noise_settings["eps"], + ) + return field, diagnostics + + +def _velocity_from_dataset(dataset): + """Read fixed x/y displacements as an array with shape (2, y, x). + + Return None if both velocity variables are absent. Otherwise require both + components, finite values, and dimensions (y, x) in either order. Values + are displacements in pixels per observation interval. Extract arrays before + stacking to avoid automatic xarray coordinate alignment. + """ + names = ("velocity_x", "velocity_y") + present = [name in dataset.data_vars for name in names] + if not any(present): + return None + if not all(present): + raise ValueError("Supply both velocity_x and velocity_y, or neither") + components = [] + for name in names: + field = dataset[name] + if set(field.dims) != {"y", "x"} or field.ndim != 2: + raise ValueError(f"{name} must have dimensions (y, x)") + if field.attrs.get("units") not in ( + None, + "pixels/timestep", + "pixels / timestep", + ): + raise ValueError(f"{name} must be in pixels per observation timestep") + values = field.transpose("y", "x").values + if not np.issubdtype(values.dtype, np.number) or np.iscomplexobj(values): + raise ValueError(f"{name} must contain finite real values") + if not np.isfinite(values).all(): + raise ValueError(f"{name} must contain finite real values") + components.append(values) + return np.array(components, dtype=float, copy=True) + + +def _prepare_dataset(dataset, kmperpixel, timestep): + """Validate the grid and copy observations into (time, y, x) order. + + Require three finite rain-rate fields, regularly spaced observation times, + and a regular grid with square pixels. Coordinate stepsize attributes must + match the actual spacing. Return the copied dataset, optional velocity, + pixel size in km, and observation interval in minutes. Explicit pixel-size + or interval overrides must agree with the coordinates. + """ + if not isinstance(dataset, xr.Dataset): + raise ValueError("dataset must be an xarray.Dataset") + precip_var = dataset.attrs.get("precip_var") + if not isinstance(precip_var, str) or precip_var not in dataset.data_vars: + raise ValueError( + "dataset.attrs['precip_var'] must name a precipitation variable" + ) + precip = dataset[precip_var] + if precip.ndim != 3 or set(precip.dims) != {"time", "y", "x"}: + raise ValueError("Precipitation must have dimensions (time, y, x)") + precip = precip.transpose("time", "y", "x") + if ( + precip.shape[0] != 3 + or min(precip.shape[1:]) < 2 + or max(precip.shape[1:]) < 3 + or not np.issubdtype(precip.dtype, np.floating) + ): + raise ValueError("Precipitation must be floating point with shape (3, y, x)") + if precip.attrs.get("units") != "mm/h" or precip.attrs.get("transform") is not None: + raise ValueError("Precipitation must be untransformed rain rates in mm/h") + rain = precip.values + if not np.isfinite(rain).all() or np.any(rain < 0): + raise ValueError( + "Precipitation must contain only finite nonnegative rain rates" + ) + + spacings = [] + for axis in ("x", "y"): + if axis not in dataset.coords or dataset[axis].dims != (axis,): + raise ValueError(f"A one-dimensional {axis} coordinate is required") + coord = dataset[axis] + if not np.issubdtype(coord.dtype, np.number) or np.iscomplexobj(coord): + raise ValueError(f"{axis} coordinates must be finite real numbers") + values = np.asarray(coord.values, dtype=float) + differences = np.diff(values) + step = differences[0] + if ( + not np.isfinite(values).all() + or step == 0 + or not np.allclose(differences, step, rtol=1e-7, atol=abs(step) * 1e-10) + ): + raise ValueError( + f"{axis} coordinates must be regularly spaced and monotonic" + ) + declared = _finite(f"{axis}.stepsize", coord.attrs.get("stepsize")) + if not np.isclose(declared, step, rtol=1e-7, atol=abs(step) * 1e-10): + raise ValueError(f"{axis}.stepsize disagrees with coordinate spacing") + units = coord.attrs.get("units") + if units not in ("m", "km"): + raise ValueError(f"{axis} coordinate units must be 'm' or 'km'") + spacings.append(abs(step) / (1000.0 if units == "m" else 1.0)) + if not np.isclose(*spacings, rtol=1e-7, atol=0): + raise ValueError("The parametric forecast requires square pixels") + + if "time" not in dataset.coords or dataset.time.dims != ("time",): + raise ValueError("A one-dimensional datetime time coordinate is required") + times = dataset.time.values + if not np.issubdtype(times.dtype, np.datetime64) or np.isnat(times).any(): + raise ValueError("time must contain valid datetime64 observation times") + intervals = np.diff(times) / np.timedelta64(1, "s") + seconds = intervals[0] + if seconds <= 0 or not np.allclose(intervals, seconds, rtol=1e-9, atol=1e-9): + raise ValueError("Observation times must be increasing and regularly spaced") + declared = _positive("time.stepsize", dataset.time.attrs.get("stepsize")) + if not np.isclose(declared, seconds, rtol=1e-9, atol=1e-9): + raise ValueError("time.stepsize must match observation spacing in seconds") + derived_timestep = float(seconds / 60.0) + for name, supplied, derived in ( + ("kmperpixel", kmperpixel, spacings[0]), + ("timestep", timestep, derived_timestep), + ): + if supplied is not None and not np.isclose( + _positive(name, supplied), derived, rtol=1e-7, atol=0 + ): + raise ValueError(f"{name} disagrees with dataset coordinates") + + velocity = _velocity_from_dataset(dataset) + prepared = dataset.copy(deep=True) + prepared[precip_var] = precip.copy(deep=True).astype(float) + return prepared, velocity, float(spacings[0]), derived_timestep + + +def _random_streams(seed, n_ens_members): + """Create independent member streams and record their seed identities.""" + # Give each member a separate stream; increasing ensemble size leaves the + # existing member streams unchanged. Copy a supplied SeedSequence first. + if isinstance(seed, np.random.SeedSequence): + root_seed = np.random.SeedSequence( + seed.entropy, spawn_key=seed.spawn_key, pool_size=seed.pool_size + ) + else: + if seed is not None: + _integer("seed", seed, 0) + root_seed = np.random.SeedSequence(seed) + member_seeds = root_seed.spawn(n_ens_members) + streams = [np.random.default_rng(s) for s in member_seeds] + information = dict( + entropy=root_seed.entropy, + spawn_key=root_seed.spawn_key, + pool_size=root_seed.pool_size, + member_spawn_keys=[s.spawn_key for s in member_seeds], + ) + return streams, information + + +def _evolution_settings( + timesteps, + n_ens_members, + n_cascade_levels, + rain_threshold, + scale_break, + zerovalue, + taper_sigma, + min_rain_fraction, + callback, + return_output, + return_diagnostics, +): + """Validate shared evolution controls and return settings with explicit units.""" + timesteps = _integer("timesteps", timesteps) + n_ens_members = _integer("n_ens_members", n_ens_members) + n_cascade_levels = _integer("n_cascade_levels", n_cascade_levels, 3) + rain_threshold = _positive("rain_threshold", rain_threshold) + scale_break = _positive("scale_break", scale_break) + threshold = float(10 * np.log10(rain_threshold)) + zerovalue = ( + threshold - 0.1 if zerovalue is None else _finite("zerovalue", zerovalue) + ) + if zerovalue >= threshold: + raise ValueError("zerovalue must be below the dBR threshold") + taper_sigma = _finite("taper_sigma", taper_sigma) + min_rain_fraction = _finite("min_rain_fraction", min_rain_fraction) + if taper_sigma < 0 or not 0 <= min_rain_fraction <= 1: + raise ValueError("taper_sigma must be nonnegative; min_rain_fraction in [0, 1]") + if callback is not None and not callable(callback): + raise ValueError("callback must be callable") + for name, flag in [ + ("return_output", return_output), + ("return_diagnostics", return_diagnostics), + ]: + if not isinstance(flag, (bool, np.bool_)): + raise ValueError(f"{name} must be boolean") + return dict( + timesteps=timesteps, + n_ens_members=n_ens_members, + n_cascade_levels=n_cascade_levels, + rain_threshold=rain_threshold, + scale_break_km=scale_break, + threshold_dbr=threshold, + zerovalue_dbr=zerovalue, + taper_sigma_pixels=taper_sigma, + min_rain_fraction=min_rain_fraction, + ) + + +def _run_ensemble( + dataset, + observed, + velocity, + parameters, + phi, + bp_filter, + streams, + settings, + diagnostics, + callback, + return_output, + return_diagnostics, +): + """Evolve initialized members and collect or stream rain-rate fields. + + Both forecast and simulate use this loop. The observed sequence contains + normalized cascade decompositions, oldest first. Settings contain validated + physical units and calibration controls. A missing parameter set or zero target rain + fraction returns dry fields without drawing noise or constructing trajectories. + """ + shape = dataset[dataset.attrs["precip_var"]].shape[1:] + n_ens_members = len(streams) + timesteps = settings["timesteps"] + threshold = settings["threshold_dbr"] + zerovalue = settings["zerovalue_dbr"] + rain_threshold = settings["rain_threshold"] + kmperpixel = settings["kmperpixel"] + scale_break = settings["scale_break_km"] + noise_settings = settings["noise"] + min_rain_fraction = settings["min_rain_fraction"] + dry = parameters is None or parameters["rain_fraction"] == 0 + output = np.empty((n_ens_members, timesteps, *shape)) if return_output else None + if not dry: + coordinates = prepare_advection_coordinates(velocity) + taper = gaussian_filter( + advect_field(np.ones(shape), coordinates[1]), settings["taper_sigma_pixels"] + ) + # Histories share read-only observations until each member replaces its + # own list entries. The numerical kernels never mutate these observations. + histories = [[observed[-1], observed[-2]] for _ in range(n_ens_members)] + members = range(n_ens_members) if not dry else () + for lead in range(timesteps): + fields = ( + np.zeros((n_ens_members, *shape)) + if dry + else np.empty((n_ens_members, *shape)) + ) + for member in members: + field, generation = _advance( + histories[member], + parameters, + phi, + bp_filter, + coordinates, + streams[member], + threshold, + zerovalue, + kmperpixel, + scale_break, + noise_settings, + min_rain_fraction, + ) + with np.errstate(over="ignore"): + physical = 10 ** (field / 10) + physical[physical < rain_threshold] = 0 + if not np.isfinite(physical).all(): + raise ValueError("Generated rainfall contains nonfinite values") + # Return physical rainfall before tapering. Only the feedback state + # is tapered and decomposed for the next forecast step. + fields[member] = physical + histories[member] = [ + _decompose(field * taper, bp_filter), + histories[member][0], + ] + if return_diagnostics: + diagnostics["generation"][member].append(generation) + if output is not None: + output[:, lead] = fields + if callback is not None: + callback(fields.copy()) + if output is not None: + output = convert_output_to_xarray_dataset(dataset, timesteps, output) + return (output, diagnostics) if return_diagnostics else output + + +def forecast( + dataset, + timesteps=1, + *, + n_ens_members=24, + kmperpixel=None, + timestep=None, + n_cascade_levels=6, + rain_threshold=1.0, + scale_break=20.0, + zerovalue=None, + taper_sigma=5.0, + min_rain_fraction=0.025, + seed=None, + motion_kwargs=None, + lower_truncation=1.25, + temporal_exponent=0.81, + alpha=2.72, + reference_levels=(0,), + reference_time=None, + moment_distance_limit=0.5, + beta_1_bounds=(-3.9, float(np.nextafter(-1.25, -np.inf))), + beta_2_min=-4.0, + minimum_beta_gap=0.1, + correlation_time_bounds=(60.0, 720.0), + low_rain_fraction=0.05, + fallback_beta_1=-1.45, + fallback_beta_2=-3.253958088759979, + fallback_reference_time=330.61752959848894, + temporal_fallback=None, + noise_kwargs=None, + callback=None, + return_output=True, + return_diagnostics=False, +): + """Generate a parametric nowcast ensemble from three rain-rate observations. + + Parameters + ---------- + dataset: xarray.Dataset + Three complete, finite, nonnegative rain-rate observations in mm/h, + ordered oldest first. The precipitation variable is named by the global + ``precip_var`` attribute and has dimensions (time, y, x), in any order. + Time must be datetime64 and regularly spaced; its ``stepsize`` is in + seconds. Spatial coordinates must be regular with square pixels, units + 'm' or 'km', and signed ``stepsize`` attributes. Optional velocity_x and + velocity_y variables have dimensions (y, x) and contain displacements + in pixels per observation timestep. If absent, estimate Lucas-Kanade + motion from the latest two dBR fields. + timesteps: int, optional + Positive number of leads at the observation interval. Lead zero omitted. + n_ens_members, n_cascade_levels: int, optional + Ensemble size (24) and Gaussian cascade levels (6, minimum 3). + kmperpixel, timestep: float, optional + Pixel size in km and interval in minutes, derived from coordinates by + default. Explicit values must agree with the dataset. Fractional + minutes are preserved. + rain_threshold, scale_break: float, optional + Rain occurrence threshold in mm/h (1) and spectral break in km (20). + zerovalue: float, optional + Dry dBR value below 10*log10(rain_threshold), default threshold_dbr-0.1. + taper_sigma: float, optional + Gaussian boundary-mask smoothing in pixels (5). Applied to transformed + forecasts before feedback decomposition, after producing the output. + min_rain_fraction: float, optional + Post-AR dry-output cutoff (0.025), strictly applied below this fraction. + Positive low fractions are still distinguished from missing estimates. + seed: int or numpy.random.SeedSequence, optional + Seed for independent member streams. Global RNG and supplied SeedSequence + are not mutated. Prefix members agree when ensemble size is increased. + motion_kwargs: dict, optional + Overrides of the default dense Lucas-Kanade configuration. Ignored + when velocity is supplied; diagnostics record that motion source. + lower_truncation: float, optional + Lower cutoff for fitting the parent Normal distribution, in dBR (1.25). + Must be at or above the transformed rain occurrence threshold. + temporal_exponent: float, optional + Exponent relating correlation time to cascade wavelength (0.81). + alpha: float, optional + Power controlling the model lag-1 and lag-2 correlations (2.72). + reference_levels: tuple of int, optional + Zero-based cascade levels used to estimate reference time, default (0,). + Every selected level must yield a valid estimate. + reference_time: float, optional + Reference correlation time in minutes. By default, estimate it from the + observations. A supplied value replaces that estimate before applying + low-coverage replacements and correlation-time bounds. + moment_distance_limit: float, optional + Maximum Euclidean distance between fitted and sample intensity moments, + normalized by the sample standard deviation (0.5). Use the fitted + moments only when valid and within this limit. + beta_1_bounds: tuple of float, optional + Lower and upper bounds for the large-scale spectral slope, default + (-3.9, just below -1.25). + beta_2_min: float, optional + Lower bound for the small-scale spectral slope (-4.0). + minimum_beta_gap: float, optional + Minimum separation enforcing beta_2 <= beta_1 - gap (0.1). + correlation_time_bounds: tuple of float, optional + Lower and upper reference-time bounds in minutes, default (60, 720). + low_rain_fraction: float, optional + Rain fraction below which both slopes and reference time are replaced + by the fallback values (0.05). Set to zero to disable replacement. + fallback_beta_1, fallback_beta_2: float, optional + Low-coverage spectral slope replacements, default -1.45 and + -3.253958088759979. The small-scale value is an Auckland rainfall + archive median; see :doc:`/user_guide/parametric_nowcast`. + fallback_reference_time: float, optional + Low-coverage reference-time replacement in minutes + (330.61752959848894), the Auckland rainfall archive median. + temporal_fallback: dict, optional + Configured temporal-only fallback with positive ``lzero`` (km) and + ``tzero`` (minutes). If any observed AR(2) pair or the estimated + reference time is invalid, use tzero*(wavelength/lzero)**temporal_exponent + to construct all model correlations. Complete input grids and valid + marginal/spatial estimates remain required. This mode disables the + combined low-rain-fraction slope and reference-time replacement. Cannot + be combined with an explicit reference_time override. + noise_kwargs: dict, optional + Calibration controls for generate_calibrated_field; nbins and eps also + control post-AR marginal mapping. Randomness is controlled by seed. + callback: callable, optional + Called once per lead with a fresh (member, y, x) rain-rate array. + return_output, return_diagnostics: bool, optional + Return full ensemble (default True), and optional diagnostics (False). + With return_output=False, callback receives output without retaining + the full forecast. Errors from a callback propagate. + + Returns + ------- + ensemble: xarray.Dataset or None + Float64 mm/h precipitation with dimensions (ens_number, time, y, x), + retaining spatial coordinates and projection metadata. Valid times + start one observation interval after forecast_reference_time. + None if return_output=False. + diagnostics: dict, optional + When requested, return (ensemble, diagnostics): raw estimates, selected + and bounded generation parameters, settings, motion, seed information, + and per-member/per-lead calibration results. + + See also + -------- + simulate + + Notes + ----- + Fixed statistical parameters and motion are used throughout the forecast. + The two previous cascade states are advected to the forecast time before + adding normalized random innovations. Distribution matching, boundary taper, + and re-decomposition feed back into the next state. In STEPS, output + probability matching is separate from the autoregressive state. Zero + normalized anomalies at inflow are supplemented by innovations. A newest + field with no rain at the occurrence threshold gives a dry ensemble even + if older fields were wet. Earlier dry fields use + zero anomalies. Unavailable positive-rain moments raise rather than becoming + dry forecasts. Dataset data are loaded eagerly; the numerical engine uses + NumPy arrays. + """ + # Resolve grid geometry and cadence before working with NumPy arrays. + dataset, velocity, kmperpixel, timestep = _prepare_dataset( + dataset, kmperpixel, timestep + ) + precip_var = dataset.attrs["precip_var"] + rain = dataset[precip_var].values + shape = rain.shape[1:] + settings = _evolution_settings( + timesteps, + n_ens_members, + n_cascade_levels, + rain_threshold, + scale_break, + zerovalue, + taper_sigma, + min_rain_fraction, + callback, + return_output, + return_diagnostics, + ) + timesteps = settings["timesteps"] + n_ens_members = settings["n_ens_members"] + n_cascade_levels = settings["n_cascade_levels"] + rain_threshold = settings["rain_threshold"] + threshold = settings["threshold_dbr"] + zerovalue = settings["zerovalue_dbr"] + estimation = _estimation_settings( + dict( + lower_truncation=lower_truncation, + temporal_exponent=temporal_exponent, + alpha=alpha, + reference_levels=reference_levels, + reference_time=reference_time, + moment_distance_limit=moment_distance_limit, + beta_1_bounds=beta_1_bounds, + beta_2_min=beta_2_min, + minimum_beta_gap=minimum_beta_gap, + correlation_time_bounds=correlation_time_bounds, + low_rain_fraction=low_rain_fraction, + fallback_beta_1=fallback_beta_1, + fallback_beta_2=fallback_beta_2, + fallback_reference_time=fallback_reference_time, + ), + threshold, + n_cascade_levels, + ) + if temporal_fallback is not None: + if not isinstance(temporal_fallback, Mapping) or set(temporal_fallback) != { + "lzero", + "tzero", + }: + raise ValueError( + "temporal_fallback requires exactly lzero (km) and tzero (minutes)" + ) + temporal_fallback = {k: _positive(k, v) for k, v in temporal_fallback.items()} + if estimation["reference_time"] is not None: + raise ValueError( + "temporal_fallback cannot be combined with reference_time override" + ) + # Temporal fallback must never mask an unavailable spatial estimate. + estimation["low_rain_fraction"] = 0.0 + noise_settings = _noise_settings(noise_kwargs) + if motion_kwargs is not None and not isinstance(motion_kwargs, Mapping): + raise ValueError("motion_kwargs must be a mapping") + motion_settings = deepcopy(_MOTION_DEFAULTS) + for key, value in dict(motion_kwargs or {}).items(): + if isinstance(value, Mapping) and isinstance(motion_settings.get(key), Mapping): + motion_settings[key].update(deepcopy(dict(value))) + else: + motion_settings[key] = deepcopy(value) + if velocity is not None: + if np.iscomplexobj(velocity) or np.ma.isMaskedArray(velocity): + raise ValueError("velocity must be a finite real array") + velocity = np.array(velocity, dtype=float, copy=True) + if velocity.shape != (2, *shape) or not np.isfinite(velocity).all(): + raise ValueError("velocity must be finite with shape (2, y, x)") + streams, seed_information = _random_streams(seed, n_ens_members) + settings.update( + timestep_minutes=timestep, + kmperpixel=kmperpixel, + estimation=estimation, + temporal_fallback=temporal_fallback, + noise=noise_settings, + motion=motion_settings, + ) + diagnostics = dict( + method="parametric", + generator=GENERATOR_ID, + settings=deepcopy(settings), + motion_source="supplied" if velocity is not None else "lucaskanade", + seed=seed_information, + generation=[[] for _ in range(n_ens_members)] if return_diagnostics else None, + ) + transformed_dataset = dB_transform( + dataset, threshold=rain_threshold, zerovalue=zerovalue + ) + transformed = transformed_dataset[precip_var].values + # Untransformed rates omit the transform attribute so native NetCDF export + # does not encounter a non-serializable None value. + dataset[precip_var].attrs.pop("transform", None) + dataset[precip_var].attrs.update( + units="mm/h", threshold=rain_threshold, zerovalue=0.0 + ) + # Estimate occurrence and intensity from the newest observation. A dry + # newest field has a defined result without estimating motion or spectra. + statistics = compute_field_statistics( + transformed[-1], threshold, estimation["lower_truncation"] + ) + diagnostics["statistics"] = statistics + if statistics["rain_fraction"] == 0: + diagnostics.update( + status="dry_latest_observation", velocity=velocity, parameters=None + ) + if velocity is None: + diagnostics["motion_source"] = "not_required_for_dry_input" + return _run_ensemble( + dataset, + None, + velocity, + None, + None, + None, + streams, + settings, + diagnostics, + callback, + return_output, + return_diagnostics, + ) + if velocity is None: + motion_dataset = motion.get_method("LK")( + transformed_dataset.isel(time=slice(-2, None)), **motion_settings + ) + try: + velocity = ( + _velocity_from_dataset(motion_dataset) + if isinstance(motion_dataset, xr.Dataset) + else None + ) + except ValueError as error: + raise ValueError("Motion estimator returned invalid velocity") from error + if velocity is None or velocity.shape != (2, *shape): + raise ValueError("Motion estimator returned invalid velocity") + diagnostics["velocity"] = velocity.copy() + # Estimate scale-dependent persistence from the three aligned observations. + bp_filter = filter_gaussian(shape, n_cascade_levels) + observed = [_decompose(field, bp_filter) for field in transformed] + wavelengths = kmperpixel / bp_filter["central_freqs"] + correlations = lagrangian_autocorrelation( + np.stack([d["cascade_levels"] for d in observed]), velocity + ) + times = ar2_correlation_time(correlations[:, 0], correlations[:, 1], timestep) + estimated_time = estimate_reference_correlation_time( + times, + wavelengths, + estimation["temporal_exponent"], + levels=estimation["reference_levels"], + ) + reference_time = ( + estimated_time + if estimation["reference_time"] is None + else estimation["reference_time"] + ) + temporal_selection = dict(source="observed", reason=None) + if estimation["reference_time"] is not None: + temporal_selection["source"] = "override" + use_fallback = temporal_fallback is not None and ( + not np.isfinite(times).all() + or np.any(times <= 0) + or not np.isfinite(estimated_time) + or estimated_time <= 0 + ) + if use_fallback: + reference_time = ( + temporal_fallback["tzero"] + * (wavelengths[0] / temporal_fallback["lzero"]) + ** estimation["temporal_exponent"] + ) + temporal_selection.update( + source="configured_fallback", + reason="invalid_observed_ar2_or_reference_time", + configuration=deepcopy(temporal_fallback), + ) + temporal_selection["selected_reference_time_minutes"] = reference_time + # Fit spatial structure and select the parameters held fixed for all leads. + frequency, power = rapsd_logarithmic(transformed[-1]) + spectrum = fit_power_law_spectrum(frequency, power, scale_break / kmperpixel) + parameters, safeguards = _generation_parameters( + statistics, spectrum, reference_time, estimation + ) + if use_fallback: + # The explicit configured scale law is not clipped to regional bounds. + parameters["reference_time"] = reference_time + safeguards["bounded_parameters"] = [ + k for k in safeguards["bounded_parameters"] if k != "reference_time" + ] + model_lags, temporal_diagnostics = model_cascade_correlations( + parameters["reference_time"], + wavelengths, + estimation["temporal_exponent"], + estimation["alpha"], + timestep, + return_diagnostics=True, + ) + # Convert model lag correlations to AR coefficients and innovation weights. + phi = np.array([estimate_ar_params_yw(pair) for pair in model_lags]) + diagnostics.update( + status="ok", + spectrum=spectrum, + observed_correlations=correlations, + observed_correlation_times_minutes=times, + estimated_reference_time_minutes=estimated_time, + wavelengths_km=wavelengths, + parameters=parameters.copy(), + safeguards=safeguards, + temporal_model=temporal_diagnostics, + temporal_selection=temporal_selection, + model_correlations=model_lags, + ar_coefficients=phi, + constant_observed_levels=[ + np.flatnonzero(np.asarray(d["stds"]) == 0) for d in observed + ], + ) + return _run_ensemble( + dataset, + observed, + velocity, + parameters, + phi, + bp_filter, + streams, + settings, + diagnostics, + callback, + return_output, + return_diagnostics, + ) + + +def simulate( + dataset, + timesteps=1, + *, + parameters, + reference_time, + reference_wavelength, + temporal_exponent=0.81, + alpha=2.72, + n_ens_members=24, + kmperpixel=None, + timestep=None, + n_cascade_levels=6, + rain_threshold=1.0, + scale_break=20.0, + zerovalue=None, + taper_sigma=5.0, + min_rain_fraction=0.0, + seed=None, + noise_kwargs=None, + callback=None, + return_output=True, + return_diagnostics=False, +): + """Simulate rainfall time series using prescribed statistics and motion. + + Parameters + ---------- + dataset: xarray.Dataset + Three complete, finite, nonnegative rain-rate fields in mm/h, ordered + oldest first, on a regular grid with square pixels and at least four + pixels along each spatial dimension. Use the dataset + conventions described in :func:`forecast`: a global precip_var attribute, + precipitation dimensions (time, y, x) in any order, spatial units 'm' or + 'km', and coordinate stepsize attributes matching the grid and time + interval. Time stepsize is in seconds. Supply both velocity_x and + velocity_y with dimensions (y, x), in either order, and displacements + in pixels per interval. Motion is not estimated. The latest two fields + initialize the AR(2) state; the oldest is retained for consistency with + the forecast input interface. + timesteps: int, optional + Positive number of intervals to simulate (default 1). The first output + is one observation interval after the newest initial field. + parameters: mapping + Exactly five finite real values: beta_1 and beta_2 are signed spectral + slopes; rain_fraction lies in [0, 1]; mean_db and positive stdev_db are + the Normal location and scale in dBR. These describe the generation + targets, not necessarily the initial fields' sample statistics. Values + are copied and held fixed, without estimation, clipping, or regional + replacements. + reference_time: float + Positive correlation time in minutes at reference_wavelength. This is + the one-sided discrete AR(2) autocorrelation sum including lag zero, + not an e-folding time or a forecast-skill horizon. + reference_wavelength: float + Positive reference wavelength in km. It need not equal a cascade level's + nominal wavelength or the domain size. + temporal_exponent: float, optional + Finite exponent in T = reference_time * + (wavelength/reference_wavelength)**temporal_exponent (default 0.81). + alpha: float, optional + Positive exponent in rho2 = rho1**alpha before the nonperiodic AR(2) + adjustment (default 2.72). Lag correlations are converted to AR + coefficients and innovation weights separately. + n_ens_members: int, optional + Number of independent realizations (default 24). + kmperpixel, timestep: float, optional + Pixel spacing in km and interval in minutes. Derived from coordinates + by default; explicit values must agree. Fractional minutes are retained. + n_cascade_levels: int, optional + Number of Gaussian cascade levels (default 6, minimum 3). + rain_threshold: float, optional + Rain occurrence threshold in mm/h (default 1). + scale_break: float, optional + Prescribed spectral break wavelength in km (default 20). + zerovalue: float, optional + Dry dBR value below the transformed rain threshold. The default is + 10*log10(rain_threshold)-0.1. + taper_sigma: float, optional + Gaussian boundary-mask smoothing width in pixels (default 5). Tapering + affects the feedback state after the output field has been produced. + min_rain_fraction: float, optional + Return dry output for target fractions strictly below this value. + Default 0 preserves positive target fractions; forecast defaults to + 0.025. Failed positive-fraction mappings raise an error. + seed: int or numpy.random.SeedSequence, optional + Seed for separate member streams. Global random state and a supplied + SeedSequence are unchanged. Increasing the ensemble size preserves the + leading members for identical settings. + noise_kwargs: dict, optional + Overrides for spectral calibration and distribution matching, as in + :func:`forecast`. Calibration may miss its tolerance; diagnostics + report the selected candidate and its spectral error. + callback: callable, optional + Called once per interval with a fresh (member, y, x) array in mm/h. + Modifying it does not affect the state or retained output. Exceptions + propagate to the caller. + return_output: bool, optional + Retain and return the full time series (default True). Set False to + stream through callback without retaining the full output. + return_diagnostics: bool, optional + Also return diagnostics (default False). + + Returns + ------- + ensemble: xarray.Dataset or None + Float64 rain rates with dimensions (ens_number, time, y, x), preserving + spatial coordinates and projection metadata. Members are numbered from + one. The forecast_reference_time coordinate identifies the newest + initial field, following the forecast output convention. None if + return_output=False. + diagnostics: dict, optional + When requested, return (ensemble, diagnostics). Record the prescribed + parameters, effective settings, seed identities, supplied velocity, + nominal wavelengths, requested and achieved correlation times, model + lag correlations, AR coefficients, and calibration results indexed + [member][interval]. + + See also + -------- + forecast, pysteps.timeseries.parametric.model_cascade_correlations + + Notes + ----- + Simulation and forecasting share the same advection, stochastic AR(2) + updates, distribution matching, and boundary feedback. Parameters and motion + stay fixed while each member's rainfall pattern evolves. Matching of rain + fraction, intensity distribution, and spectra is approximate on finite + grids. Distribution matching and boundary feedback can change correlations + measured from the resulting rainfall. Requested level times at or below one + interval resolve to white noise, with achieved times recorded in diagnostics. + + A zero target rain fraction produces dry output even from wet initial + fields. Conversely, a positive target can generate rain from dry initial + fields. Missing initial values remain errors. Shifted copies of one field + can initialize a simulation, but they represent pure translation; they do + not establish a stationary history for the prescribed temporal model. + """ + dataset, velocity, kmperpixel, timestep = _prepare_dataset( + dataset, kmperpixel, timestep + ) + if velocity is None: + raise ValueError("simulate requires both velocity_x and velocity_y") + precip_var = dataset.attrs["precip_var"] + shape = dataset[precip_var].shape[1:] + if min(shape) < 4: + raise ValueError("Simulation requires at least four pixels per axis") + settings = _evolution_settings( + timesteps, + n_ens_members, + n_cascade_levels, + rain_threshold, + scale_break, + zerovalue, + taper_sigma, + min_rain_fraction, + callback, + return_output, + return_diagnostics, + ) + if not isinstance(parameters, Mapping): + raise ValueError("parameters must be a mapping") + names = ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + if set(parameters) != set(names): + raise ValueError(f"parameters must contain exactly {', '.join(names)}") + parameters = {name: _finite(name, parameters[name]) for name in names} + if not 0 <= parameters["rain_fraction"] <= 1: + raise ValueError("rain_fraction must lie in [0, 1]") + _positive("stdev_db", parameters["stdev_db"]) + reference_time = _positive("reference_time", reference_time) + reference_wavelength = _positive("reference_wavelength", reference_wavelength) + temporal_exponent = _finite("temporal_exponent", temporal_exponent) + alpha = _positive("alpha", alpha) + noise_settings = _noise_settings(noise_kwargs) + streams, seed_information = _random_streams(seed, settings["n_ens_members"]) + settings.update( + timestep_minutes=timestep, + kmperpixel=kmperpixel, + reference_time_minutes=reference_time, + reference_wavelength_km=reference_wavelength, + temporal_exponent=temporal_exponent, + alpha=alpha, + noise=noise_settings, + ) + + # Construct temporal correlations from the prescribed scale law, using the + # actual filter wavelengths. No correlations are estimated from the history. + bp_filter = filter_gaussian(shape, settings["n_cascade_levels"]) + wavelengths = kmperpixel / bp_filter["central_freqs"] + model_lags, temporal_diagnostics = model_cascade_correlations( + reference_time, + wavelengths, + temporal_exponent, + alpha, + timestep, + reference_wavelength=reference_wavelength, + return_diagnostics=True, + ) + phi = np.array([estimate_ar_params_yw(pair) for pair in model_lags]) + + # Initial rainfall supplies cascade states only. It does not alter the + # supplied intensity, occurrence, spectral, or persistence parameters. + transformed_dataset = dB_transform( + dataset, + threshold=settings["rain_threshold"], + zerovalue=settings["zerovalue_dbr"], + ) + observed = [ + _decompose(field, bp_filter) for field in transformed_dataset[precip_var].values + ] + dataset[precip_var].attrs.pop("transform", None) + dataset[precip_var].attrs.update( + units="mm/h", threshold=settings["rain_threshold"], zerovalue=0.0 + ) + diagnostics = dict( + method="parametric_simulation", + generator=GENERATOR_ID, + status="dry_target" if parameters["rain_fraction"] == 0 else "ok", + settings=deepcopy(settings), + parameter_source="supplied", + parameters=parameters.copy(), + motion_source="supplied", + velocity=velocity.copy(), + seed=seed_information, + wavelengths_km=wavelengths, + temporal_model=temporal_diagnostics, + model_correlations=model_lags, + ar_coefficients=phi, + generation=[[] for _ in streams] if return_diagnostics else None, + ) + return _run_ensemble( + dataset, + observed, + velocity, + parameters, + phi, + bp_filter, + streams, + settings, + diagnostics, + callback, + return_output, + return_diagnostics, + ) diff --git a/pysteps/postprocessing/probmatching.py b/pysteps/postprocessing/probmatching.py index ba493c51..457b2a45 100644 --- a/pysteps/postprocessing/probmatching.py +++ b/pysteps/postprocessing/probmatching.py @@ -14,11 +14,134 @@ pmm_compute shift_scale resample_distributions + match_rainfall_distribution """ import numpy as np from scipy import interpolate as sip from scipy import optimize as sop +from scipy import stats as sstats + + +def match_rainfall_distribution( + field, + rain_fraction, + mean, + stdev, + threshold, + *, + zerovalue=None, + min_rain_fraction=0.0, + nbins=500, + eps=0.001, +): + """Map a finite field to a prescribed occurrence and Normal rainfall marginal. + + Parameters + ---------- + field: array_like + Nonempty, finite real 2D field, converted to float64 without modification. + rain_fraction: float + Requested wet fraction in [0, 1]. A histogram cutoff selects the wet + pixels before marginal mapping; the final wet fraction is approximate. + mean, stdev: float + Finite parent-Normal location and positive scale in transformed units. + These are caller-selected generation moments, not necessarily ordinary + retained wet-pixel moments. No moment selection is performed here. + threshold: float + Rain/no-rain threshold in the same transformed units. + zerovalue: float, optional + Finite dry value below threshold, default threshold-0.1. Used for an + explicit dry result; positive-fraction mapping retains continuous dry + values below threshold for subsequent cascade decomposition. + min_rain_fraction: float, optional + Explicit dry-output policy for fractions below this value. Default 0 + suppresses no positive fraction. The parametric forecast uses 0.025. + nbins: int, optional + Number of histogram bins (default 500). + eps: float, optional + Normal-quantile probability clip in (0, 0.5), default 0.001. + + Returns + ------- + mapped: ndarray + Field of the same shape and transformed units. Zero requested fraction + or a fraction below min_rain_fraction returns the configured dry value. + + Notes + ----- + Find the last bin + whose cumulative probability is <= 1-rain_fraction, shift its upper edge + to threshold, and map the strictly-above-threshold sample's binned CDF to + clipped Normal quantiles using bin-centre linear interpolation. Wet mapped + values can cross back below threshold, and spectra can change. Neither + occurrence nor marginal moments are exact guarantees. Unavailable or + degenerate positive-fraction mappings raise ValueError, not a dry fallback. + """ + if np.iscomplexobj(field): + raise ValueError("field must contain real values") + field = np.asarray(field, dtype=float) + if field.ndim != 2 or field.size == 0 or not np.isfinite(field).all(): + raise ValueError("field must be a nonempty finite two-dimensional array") + for name, value in [ + ("rain_fraction", rain_fraction), + ("mean", mean), + ("stdev", stdev), + ("threshold", threshold), + ("min_rain_fraction", min_rain_fraction), + ("eps", eps), + ]: + if not np.isscalar(value) or np.iscomplexobj(value) or not np.isfinite(value): + raise ValueError(f"{name} must be finite") + if not 0 <= rain_fraction <= 1 or not 0 <= min_rain_fraction <= 1: + raise ValueError("rain fractions must lie in [0, 1]") + if stdev <= 0: + raise ValueError("stdev must be positive") + if not isinstance(nbins, (int, np.integer)) or nbins < 2: + raise ValueError("nbins must be an integer of at least two") + if not 0 < eps < 0.5: + raise ValueError("eps must lie in (0, 0.5)") + zerovalue = threshold - 0.1 if zerovalue is None else zerovalue + if ( + not np.isscalar(zerovalue) + or np.iscomplexobj(zerovalue) + or not np.isfinite(zerovalue) + or zerovalue >= threshold + ): + raise ValueError("zerovalue must be finite and below threshold") + if rain_fraction == 0 or rain_fraction < min_rain_fraction: + return np.full_like(field, zerovalue) + + # Shift a histogram cutoff to the occurrence threshold to select wet pixels. + # Values below it remain continuous rather than being set to zerovalue. + distribution = sstats.cumfreq(field, numbins=nbins) + edges = distribution.lowerlimit + np.arange(nbins + 1) * distribution.binsize + cdf = distribution.cumcount / distribution.cumcount[-1] + below = np.flatnonzero(cdf <= 1 - rain_fraction) + rain_bin = int(below[-1]) if below.size else 0 + mapped = field + (threshold - edges[rain_bin + 1]) + wet = mapped > threshold + sample = mapped[wet] + if sample.size < 2 or np.ptp(sample) == 0: + raise ValueError("mapping produced too few distinct raining pixels") + # Map the wet sample through its conditional CDF to Normal quantiles. + # Clipping avoids infinite quantiles; mapped values may cross the threshold. + distribution = sstats.cumfreq(sample, numbins=nbins) + edges = distribution.lowerlimit + np.arange(nbins + 1) * distribution.binsize + centres = 0.5 * (edges[:-1] + edges[1:]) + cdf = distribution.cumcount / distribution.cumcount[-1] + quantiles = sstats.norm.ppf(np.clip(cdf, eps, 1 - eps), loc=mean, scale=stdev) + transform = sip.interp1d( + centres, + quantiles, + kind="linear", + bounds_error=False, + fill_value=(quantiles[0], quantiles[-1]), + ) + mapped[wet] = transform(sample) + if not np.isfinite(mapped).all(): + raise ValueError("marginal mapping produced nonfinite values") + return mapped def compute_empirical_cdf(bin_edges, hist): diff --git a/pysteps/tests/data/README.rst b/pysteps/tests/data/README.rst new file mode 100644 index 00000000..a562df0f --- /dev/null +++ b/pysteps/tests/data/README.rst @@ -0,0 +1,114 @@ +Parametric regression data +========================== + +The parametric tests combine independent numerical checks with frozen results +from the original WRNZ implementation. Normal test runs read the committed JSON +and NPZ files; they need neither the external repositories nor a database. +The simulation recovery reference described below records a separate pySTEPS +end-to-end regression case. + +The source-comparison JSON ``provenance`` entries identify the original revisions, +file SHA-256 hashes, and numerical-library versions. Case entries record parameters, +settings, and seeds. The NPZ files contain complete expected fields and histories. +Retain these identities when updating fixtures and explain intended numerical +changes. For comparisons with the original implementation, expected values must +continue to come from that implementation rather than being replaced with +pySTEPS results. + +Fixture generation +------------------ + +Run these commands from the pySTEPS repository root in an environment containing +the original project's dependencies, with this checkout on ``PYTHONPATH``:: + + python pysteps/tests/data/generate_parametric_reference.py \ + /path/to/pysteps_param /path/to/psteps_param_paper + python pysteps/tests/data/generate_calibrated_reference.py \ + /path/to/pysteps_param + python pysteps/tests/data/generate_nowcast_reference.py \ + /path/to/pysteps_param + +These commands overwrite the corresponding ``*_reference.json`` and +``*_reference.npz`` files beside the scripts. They extract original numerical +definitions without importing the operational database components. Original +function names and repository paths are intentional here: they establish where +the expected results came from. + +The committed references identify ``pysteps_param`` revision +``94bd4c1571637a8f246ac391470e1a820d6376a6`` and, where used, +``psteps_param_paper`` revision ``335f7660775480d865d9771820a81923c461832a``. +The generator and forecast metadata also identify settings and kernels used in +the comparisons. Compatibility with later source revisions is not assumed. + +Comparisons with saved rainfall cases +------------------------------------- + +The comparison scripts require the external paper repository's saved inputs. +They read those inputs and write a JSON report to the requested output path:: + + python pysteps/tests/data/compare_parametric_paper_fields.py \ + /path/to/pysteps_param /path/to/psteps_param_paper \ + --output /tmp/parametric_estimators_validation.json + python pysteps/tests/data/compare_calibrated_paper_fields.py \ + /path/to/pysteps_param /path/to/psteps_param_paper \ + --output /tmp/calibrated_generator_validation.json + python pysteps/tests/data/compare_parametric_nowcast_paper.py \ + /path/to/pysteps_param /path/to/psteps_param_paper \ + --leads 3 --output /tmp/parametric_nowcast_validation.json + +The estimator and generator scripts use ``figs/individual_field_fits/*.npz``. +The forecast script uses the saved case histories under ``figs/extreme_parameter_experiment``; +its module docstring explains the chronological selection of observations. +Reports include source identity, input hashes, settings, and numerical +differences. An existing report at the output path is overwritten. + +Compare identical-parameter evolution separately from the complete chain of +estimation and generation. Equivalent spectral regression solvers can produce +small coefficient differences that histogram interpolation and forecast +feedback amplify. Test comments and assertions explain the comparison +tolerances. These checks establish numerical agreement for the specified cases; +they do not measure forecast skill or establish calibration in another region. + +The current pySTEPS API and independent numerical tests define the maintained +behaviour. A deliberate method change may differ from the WRNZ references, but +must explain the difference and its validation. + +Simulation parameter recovery +----------------------------- + +``parametric_simulation_reference.json`` records the deterministic example in +``examples/plot_parametric_simulation.py``: a 128 by 128 grid at 1 km, a generated +initial field with seed 9, two translated copies, uniform motion of one pixel +per interval, and one member with seed 42 over twelve five-minute steps. +The reference time is 400 minutes at 256 km. The JSON records all inputs, the +source revision and example hash, numerical-library versions, and expected +recovered statistics. This reference was generated from the pySTEPS example; +it is not an additional comparison with WRNZ. + +Run the complete case with:: + + python -m pytest -q pysteps/tests/test_simulation_parametric.py::test_simulation_example_end_to_end_parameter_recovery + +The test generates the history, calls the public ``simulate`` function, checks +the output coordinates and units, and estimates occurrence, fitted Normal +parameters, spectral slopes, and Lagrangian correlations from the returned +rainfall. All twelve fields and all ten overlapping three-field histories are +included. It does not run plotting or the example's 32-run Monte Carlo check. + +The translated initialization creates a temporal transient. Expected implied +correlation times therefore describe this fixed setup, including its bias; +the recovered reference time is approximately 1001.54 minutes. The test does +not require recovery of the prescribed 400 minutes and does not define a +universal acceptable bias or replace the paper's scientific evaluation. + +Regression tolerances allow two pixels in occurrence, small numerical changes +in fitted moments and slopes, and 1e-6 absolute error in mean lag correlations. +Implied times have a relative tolerance of 0.1%, since the conversion is sensitive +near unit correlation. These are fixed-seed numerical tolerances, not Monte Carlo +acceptance bounds. The test gives the exact tolerances for each quantity. + +A deliberate change to the generation or estimation method should compare +its results with this baseline before updating the expected statistics. +Reproduce the generation and recovery steps documented in the example with +the stored case settings; review changes, then update the JSON and its +provenance explicitly. Normal tests only read the reference file. diff --git a/pysteps/tests/data/calibrated_generator_reference.json b/pysteps/tests/data/calibrated_generator_reference.json new file mode 100644 index 00000000..34cf2dc2 --- /dev/null +++ b/pysteps/tests/data/calibrated_generator_reference.json @@ -0,0 +1,100 @@ +{ + "provenance": { + "source_revision": "94bd4c1571637a8f246ac391470e1a820d6376a6", + "source_sha256": { + "rain_stats_compute.py": "2fd6f0b622e88e143d8128f80cb5d0920f0111c1021c6c448222c90bff92c889", + "stochastic_generator.py": "94c491e4fd90da0c733ce65b48227b6b73dc31a8fc6a23a44e211ad36b3f9343" + }, + "numpy_version": "2.5.1", + "scipy_version": "1.18.0" + }, + "cases": [ + { + "shape": [ + 128, + 96 + ], + "parameters": { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.08, + "mean_db": 6.0, + "stdev_db": 3.5 + }, + "seed": 4, + "kmperpixel": 1.0, + "scale_break": 20.0, + "threshold": 0.0, + "settings": {}, + "diagnostics": { + "estimated_beta_1": -2.699411978542842, + "estimated_beta_2": -2.968097079810876, + "distance": 1.2454548636920237, + "attempts": 5, + "selected_attempt": 3, + "tolerance_achieved": false + } + }, + { + "shape": [ + 128, + 96 + ], + "parameters": { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.3, + "mean_db": 6.0, + "stdev_db": 3.5 + }, + "seed": 17, + "kmperpixel": 1.0, + "scale_break": 20.0, + "threshold": 0.0, + "settings": {}, + "diagnostics": { + "estimated_beta_1": -2.7242621897368915, + "estimated_beta_2": -3.2148751587431503, + "distance": 1.3133773650788025, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + } + }, + { + "shape": [ + 95, + 127 + ], + "parameters": { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.05, + "mean_db": 6.0, + "stdev_db": 3.5 + }, + "seed": 11, + "kmperpixel": 1.0, + "scale_break": 20.0, + "threshold": 0.0, + "settings": { + "distance_threshold": 0.01 + }, + "diagnostics": { + "estimated_beta_1": -2.3629217611415716, + "estimated_beta_2": -2.8660387998408443, + "distance": 0.8583197619445808, + "attempts": 5, + "selected_attempt": 3, + "tolerance_achieved": false + } + } + ], + "post_ar_parameters": { + "rain_fraction": 0.35, + "mean": 6.5, + "stdev": 1.2, + "threshold": 0.0, + "zerovalue": -0.1 + } +} diff --git a/pysteps/tests/data/calibrated_generator_reference.npz b/pysteps/tests/data/calibrated_generator_reference.npz new file mode 100644 index 00000000..90432857 Binary files /dev/null and b/pysteps/tests/data/calibrated_generator_reference.npz differ diff --git a/pysteps/tests/data/compare_calibrated_paper_fields.py b/pysteps/tests/data/compare_calibrated_paper_fields.py new file mode 100644 index 00000000..0ddd2614 --- /dev/null +++ b/pysteps/tests/data/compare_calibrated_paper_fields.py @@ -0,0 +1,173 @@ +"""Check original and ported generator outputs on saved paper rainfall targets. + +Requires the source project's dependencies, including xarray. Inputs and outputs +are local files; no database is accessed. Raw estimates are passed to both +generators without the regional parameter selection used by the forecast. +""" + +import argparse +import hashlib +import json +from pathlib import Path +import subprocess + +import numpy as np +import xarray as xr +from scipy.optimize import minimize_scalar +from scipy.special import log_ndtr + +from pysteps.noise.parametric import generate_calibrated_field +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.utils.transformation import dB_transform +from generate_calibrated_reference import source_generator +from generate_parametric_reference import _load_functions + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + parser.add_argument("paper_repo", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + source, provenance = source_generator(args.parameter_repo) + source.update(minimize_scalar=minimize_scalar, log_ndtr=log_ndtr) + _load_functions( + args.parameter_repo / "src/pysteps_param/rain_stats_compute.py", + {"fit_lower_truncated_normal", "compute_field_stats"}, + source, + ) + report = dict( + provenance=provenance, + paper_revision=subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=args.paper_repo, text=True + ).strip(), + seeds=[0, 20260908], + scale_break_km=15.0, + pixel_size_km=1.0, + threshold_mm_per_hour=1.0, + lower_truncation_dbr=1.25, + target_policy="Raw estimator outputs; regional generation QA not applied", + coupled_maximum_absolute_error_tolerance_dbr=1e-4, + cases=[], + ) + paths = sorted((args.paper_repo / "figs/individual_field_fits").glob("*.npz")) + if not paths: + raise ValueError("No saved paper fields found") + for path in paths: + with np.load(path) as stored: + rain = stored["rain"].astype(float) + dataset = xr.Dataset( + {"rain": (("y", "x"), rain, {"units": "mm/h"})}, + attrs={"precip_var": "rain"}, + ) + transformed_dataset = dB_transform(dataset, threshold=1.0, zerovalue=-0.1) + transformed = transformed_dataset["rain"].values + metadata = transformed_dataset["rain"].attrs + _, spectrum = source["power_spectrum_1D"](transformed, 15) + marginal = source["compute_field_stats"]( + transformed, metadata, fit_truncated=True, lower_truncation=1.25 + ) + parameters = dict( + beta_1=spectrum["beta_1"], + beta_2=spectrum["beta_2"], + rain_fraction=marginal["nonzero_fraction"], + mean_db=marginal["latent_mean"], + stdev_db=marginal["latent_stdev"], + ) + frequency, power = rapsd_logarithmic(transformed) + port_spectrum = fit_power_law_spectrum(frequency, power, 15) + port_marginal = compute_field_statistics( + transformed, metadata["threshold"], 1.25 + ) + port_parameters = dict( + beta_1=port_spectrum["beta_1"], + beta_2=port_spectrum["beta_2"], + rain_fraction=port_marginal["rain_fraction"], + mean_db=port_marginal["latent_mean"], + stdev_db=port_marginal["latent_stdev"], + ) + case = dict( + name=path.stem, + input_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + parameters=parameters, + comparisons=[], + ) + for seed in report["seeds"]: + original = source["gen_calibrated_stoch_field"]( + source["StepsParameters"](**parameters), + rain.shape[1], + rain.shape[0], + 1, + 15, + metadata["threshold"], + rng=np.random.default_rng(seed), + zerovalue=metadata["zerovalue"], + ) + fixed, fixed_diagnostics = generate_calibrated_field( + rain.shape, + parameters, + 1, + 15, + metadata["threshold"], + seed=seed, + zerovalue=metadata["zerovalue"], + return_diagnostics=True, + ) + coupled, coupled_diagnostics = generate_calibrated_field( + rain.shape, + port_parameters, + 1, + 15, + metadata["threshold"], + seed=seed, + zerovalue=metadata["zerovalue"], + return_diagnostics=True, + ) + np.testing.assert_allclose(fixed, original.field, atol=1e-10, rtol=1e-12) + # Equivalent spectral regression solvers produce slightly different + # slopes. Histogram interpolation can amplify those differences. + # Check a 1e-4 dBR absolute limit (about 2.3e-5 relative rain rate) + # separately from the tight identical-target generator comparison. + np.testing.assert_allclose(coupled, original.field, atol=1e-4, rtol=0) + for diagnostics in (fixed_diagnostics, coupled_diagnostics): + for name in ("attempts", "selected_attempt", "tolerance_achieved"): + assert diagnostics[name] == getattr(original, name) + for name in ("estimated_beta_1", "estimated_beta_2", "distance"): + np.testing.assert_allclose( + diagnostics[name], getattr(original, name), atol=3e-5, rtol=1e-7 + ) + case["comparisons"].append( + dict( + seed=seed, + attempts=original.attempts, + selected_attempt=original.selected_attempt, + tolerance_achieved=original.tolerance_achieved, + distance=original.distance, + identical_targets_maximum_field_difference=float( + np.max(abs(fixed - original.field)) + ), + coupled_estimators_maximum_field_difference=float( + np.max(abs(coupled - original.field)) + ), + coupled_distance_difference=float( + abs(coupled_diagnostics["distance"] - original.distance) + ), + coupled_threshold_classification_changes=int( + np.count_nonzero( + (coupled > metadata["threshold"]) + != (original.field > metadata["threshold"]) + ) + ), + ) + ) + report["cases"].append(case) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + print( + f"Compared {len(paths)} paper cases at {len(report['seeds'])} seeds; report: {args.output}" + ) + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/compare_parametric_nowcast_paper.py b/pysteps/tests/data/compare_parametric_nowcast_paper.py new file mode 100644 index 00000000..a0b990ea --- /dev/null +++ b/pysteps/tests/data/compare_parametric_nowcast_paper.py @@ -0,0 +1,283 @@ +"""Compare complete forecasts on frozen three-observation paper histories. + +The archives store [current, previous, future1, ...]; select [1, 0, 2] as +chronological observations and estimate motion anew for the latest pair. +Nothing is read from a database or written to the research repository. +""" + +import argparse +import hashlib +import json +from pathlib import Path +import subprocess + +import numpy as np +import xarray as xr +from scipy.ndimage import gaussian_filter + +from pysteps import motion +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.nowcasts import parametric +from pysteps.nowcasts._cascade_advection import ( + advect_field, + prepare_advection_coordinates, +) +from pysteps.timeseries.autoregression import estimate_ar_params_yw +from pysteps.timeseries.parametric import lagrangian_autocorrelation +from pysteps.utils.spectral import rapsd_logarithmic, fit_power_law_spectrum + +from generate_nowcast_reference import source_engine, config, prepare, evolve + + +def fixed_parameter_forecast(rain, velocity, cfg, parameters, seed, n_members, n_leads): + """Isolate evolution using the source's selected parameters and lag pairs.""" + bp = filter_gaussian(rain.shape[1:], 6) + transformed = 10 * np.log10(np.maximum(rain, 1)) + transformed[rain < 1] = -0.1 + observed = [parametric._decompose(f, bp) for f in transformed] + coordinates = prepare_advection_coordinates(velocity) + taper = gaussian_filter(advect_field(np.ones(rain.shape[1:]), coordinates[1]), 5) + target = { + k: getattr(parameters, k) + for k in ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + } + phi = np.array( + [ + estimate_ar_params_yw(pair) + for pair in zip(parameters.lag_1, parameters.lag_2) + ] + ) + output = np.empty((n_members, n_leads, *rain.shape[1:])) + for member, child in enumerate(np.random.SeedSequence(seed).spawn(n_members)): + rng = np.random.default_rng(child) + histories = [observed[-1], observed[-2]] + for lead in range(n_leads): + field, _ = parametric._advance( + histories, + target, + phi, + bp, + coordinates, + rng, + 0, + -0.1, + 1, + 20, + parametric._noise_settings(None), + 0.025, + ) + physical = 10 ** (field / 10) + physical[physical < 1] = 0 + output[member, lead] = physical + histories = [parametric._decompose(field * taper, bp), histories[0]] + return output + + +def metrics(forecast, observations, velocity): + bp = filter_gaussian(forecast.shape[-2:], 6) + result = [] + for member in forecast: + history = [] + for field in observations[-2:]: + db = 10 * np.log10(np.maximum(field, 1)) + db[field < 1] = -0.1 + history.append(parametric._decompose(db, bp)["cascade_levels"]) + for field in member: + wet = field >= 1 + db = 10 * np.log10(np.maximum(field, 1)) + db[~wet] = -0.1 + frequency, power = rapsd_logarithmic(db) + slopes = fit_power_law_spectrum(frequency, power, 20) + history.append(parametric._decompose(db, bp)["cascade_levels"]) + correlations = lagrangian_autocorrelation(np.stack(history[-3:]), velocity) + result.append( + [ + float(wet.mean()), + float(field.mean()), + float(db[wet].mean()) if wet.any() else np.nan, + float(db[wet].std()) if wet.any() else np.nan, + slopes["beta_1"], + slopes["beta_2"], + *correlations.ravel(), + ] + ) + return np.asarray(result) + + +def finite_max(values): + values = np.asarray(values) + return ( + float(np.max(values[np.isfinite(values)])) + if np.isfinite(values).any() + else None + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + parser.add_argument("paper_repo", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--leads", type=int, default=12) + args = parser.parse_args() + source, provenance = source_engine(args.parameter_repo) + import scipy + + report = dict( + provenance=provenance, + paper_revision=subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=args.paper_repo, text=True + ).strip(), + numpy_version=np.__version__, + scipy_version=scipy.__version__, + leads=args.leads, + members=2, + seeds=[0, 20261015], + observation_policy="Archive indices [1,0,2], oldest first; newest is original cycle +5 minutes", + motion_policy="Reference Lucas-Kanade settings, re-estimated from latest pair", + config=dict(config((256, 256)), n_forecasts=args.leads), + comparisons=[], + ) + for name in ("control_1", "rain_fraction_low_2", "beta_1_low_2"): + folder = args.paper_repo / "figs/extreme_parameter_experiment" / name + path = folder / "inputs.npz" + with np.load(path) as saved: + rain = saved["rain"][[1, 0, 2]].astype(float) + metadata = json.loads((folder / "input_metadata.json").read_text()) + observation_times = np.array( + [metadata["rain"][i]["valid_time"] for i in (1, 0, 2)], + dtype="datetime64[ns]", + ) + # The archives contain raster values, not georeferencing. A local 1 km + # Cartesian grid retains their array orientation for numerical comparison. + dataset = xr.Dataset( + {"precip_intensity": (("time", "y", "x"), rain, {"units": "mm/h"})}, + coords={ + "time": ("time", observation_times, {"stepsize": 300}), + "x": ( + "x", + (np.arange(rain.shape[2]) + 0.5) * 1000, + {"units": "m", "stepsize": 1000}, + ), + "y": ( + "y", + (np.arange(rain.shape[1]) + 0.5) * 1000, + {"units": "m", "stepsize": 1000}, + ), + }, + attrs={"precip_var": "precip_intensity"}, + ) + transformed = source["DBTransformer"](threshold=1).transform(rain) + motion_input = dataset.isel(time=slice(-2, None)).copy(deep=True) + motion_input["precip_intensity"].data = transformed[-2:] + motion_input["precip_intensity"].attrs["transform"] = "dB" + motion_output = motion.get_method("LK")( + motion_input, **parametric._MOTION_DEFAULTS + ) + velocity = np.stack( + [motion_output["velocity_x"].values, motion_output["velocity_y"].values] + ) + cfg = config(rain.shape[1:]) + cfg["n_forecasts"] = args.leads + parameters, _, _, _ = prepare(source, rain, velocity, cfg) + for seed in report["seeds"]: + original, calibration = evolve( + source, + rain, + velocity, + cfg, + parameters, + seed, + n_members=2, + n_leads=args.leads, + ) + fixed = fixed_parameter_forecast( + rain, velocity, cfg, parameters, seed, 2, args.leads + ) + # Exercise the public call including its internal motion estimation. + coupled_dataset, diagnostics = parametric.forecast( + dataset, + timesteps=args.leads, + n_ens_members=2, + kmperpixel=1, + timestep=5, + seed=seed, + return_diagnostics=True, + ) + coupled = coupled_dataset[coupled_dataset.attrs["precip_var"]].values + np.testing.assert_array_equal(velocity, diagnostics["velocity"]) + for key in ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db"): + np.testing.assert_allclose( + diagnostics["parameters"][key], + getattr(parameters, key), + atol=1e-6, + rtol=0, + ) + np.testing.assert_allclose( + diagnostics["parameters"]["reference_time"], + parameters.corl_zero, + atol=1e-7, + rtol=0, + ) + np.testing.assert_allclose( + diagnostics["model_correlations"], + np.array([parameters.lag_1, parameters.lag_2]).T, + atol=1e-12, + rtol=0, + ) + np.testing.assert_allclose(fixed, original, atol=1e-8, rtol=1e-10) + difference = abs(coupled - original) + rms = float(np.sqrt(np.mean((coupled - original) ** 2))) + original_metrics = metrics(original, rain, velocity) + coupled_metrics = metrics(coupled, rain, velocity) + deltas = abs(coupled_metrics - original_metrics) + assert np.array_equal( + np.isfinite(original_metrics), np.isfinite(coupled_metrics) + ) + assert np.isfinite(coupled).all() and np.all(coupled >= 0) + selection_changes = sum( + diagnostics["generation"][m][lead]["selected_attempt"] + != calibration[m][lead]["selected_attempt"] + for m in range(2) + for lead in range(args.leads) + ) + report["comparisons"].append( + dict( + case=name, + seed=seed, + input_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + observation_times=[ + metadata["rain"][i]["valid_time"] for i in (1, 0, 2) + ], + parameters=diagnostics["parameters"], + source_parameters=parameters.active_parameters(), + moment_source=diagnostics["safeguards"]["moment_source"], + low_coverage_fallback=diagnostics["safeguards"][ + "low_rain_fraction_fallback" + ], + identical_parameters_maximum_error_mm_per_hour=float( + abs(fixed - original).max() + ), + coupled_maximum_error_mm_per_hour=float(difference.max()), + coupled_rmse_mm_per_hour=rms, + rain_classification_changes=int( + np.count_nonzero((coupled >= 1) != (original >= 1)) + ), + total_output_pixels=int(coupled.size), + candidate_selection_changes=selection_changes, + maximum_wet_fraction_difference=finite_max(deltas[:, 0]), + maximum_mean_rain_difference=finite_max(deltas[:, 1]), + maximum_wet_dbr_mean_difference=finite_max(deltas[:, 2]), + maximum_wet_dbr_stdev_difference=finite_max(deltas[:, 3]), + maximum_spectral_slope_difference=finite_max(deltas[:, 4:6]), + maximum_lagrangian_correlation_difference=finite_max(deltas[:, 6:]), + ) + ) + print( + f"{name}, seed {seed}: fixed max {abs(fixed-original).max():.3g}, coupled max {difference.max():.3g}, RMSE {rms:.3g}" + ) + args.output.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/compare_parametric_paper_fields.py b/pysteps/tests/data/compare_parametric_paper_fields.py new file mode 100644 index 00000000..a62c7966 --- /dev/null +++ b/pysteps/tests/data/compare_parametric_paper_fields.py @@ -0,0 +1,115 @@ +"""Compare source and ported estimators on saved paper rainfall fields. + +Requires the source project's environment, including xarray, and the saved +figs/individual_field_fits/*.npz inputs. Does not connect to a database. +""" + +import argparse +import hashlib +import json +from pathlib import Path +import subprocess + +import numpy as np +from scipy.optimize import curve_fit, minimize_scalar +from scipy.special import log_ndtr + +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.utils.transformation import dB_transform +from generate_parametric_reference import _load_functions + + +def main(): + import xarray as xr + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + parser.add_argument("paper_repo", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + namespace = dict( + np=np, + xr=xr, + curve_fit=curve_fit, + minimize_scalar=minimize_scalar, + log_ndtr=log_ndtr, + ) + source = args.parameter_repo / "src/pysteps_param/rain_stats_compute.py" + source_hash = _load_functions( + source, + {"power_spectrum_1D", "compute_field_stats", "fit_lower_truncated_normal"}, + namespace, + ) + report = dict( + source_sha256=source_hash, + source_revision=subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=args.parameter_repo, text=True + ).strip(), + threshold_mm_per_hour=1.0, + lower_truncation_dbr=1.25, + scale_break_km=15.0, + pixel_size_km=1.0, + cases=[], + ) + paths = sorted((args.paper_repo / "figs/individual_field_fits").glob("*.npz")) + if not paths: + raise ValueError("No saved paper fields found") + for path in paths: + with np.load(path) as saved: + rain = saved["rain"] + dataset = xr.Dataset( + {"precip_intensity": (("y", "x"), rain.astype(float), {"units": "mm/h"})}, + attrs={"precip_var": "precip_intensity"}, + ) + transformed_dataset = dB_transform(dataset, threshold=1.0, zerovalue=-0.1) + transformed = transformed_dataset["precip_intensity"].values + metadata = transformed_dataset["precip_intensity"].attrs + source_spectrum, source_model = namespace["power_spectrum_1D"](transformed, 15) + source_statistics = namespace["compute_field_stats"]( + transformed, metadata, fit_truncated=True, lower_truncation=1.25 + ) + frequency, power = rapsd_logarithmic(transformed) + model = fit_power_law_spectrum(frequency, power, 15) + statistics = compute_field_statistics(transformed, metadata["threshold"], 1.25) + np.testing.assert_allclose( + 10 * np.log10(frequency), source_spectrum.psd_bins.values, atol=1e-12 + ) + np.testing.assert_allclose( + 10 * np.log10(power), source_spectrum.psd.values, atol=1e-10 + ) + differences = { + name: abs(model[name] - source_model[name]) + for name in ("beta_1", "beta_2", "c1", "c2") + } + assert max(differences.values()) < 3e-5 + for name in ("ordinary_mean", "ordinary_stdev", "latent_mean", "latent_stdev"): + np.testing.assert_allclose( + statistics[name], source_statistics[name], rtol=1e-12 + ) + np.testing.assert_allclose( + statistics["rain_fraction"], + source_statistics["nonzero_fraction"], + rtol=1e-12, + ) + report["cases"].append( + dict( + name=path.stem, + input_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + shape=list(rain.shape), + maximum_spectrum_error_db=float( + np.max(abs(10 * np.log10(power) - source_spectrum.psd.values)) + ), + spectral_parameter_absolute_differences=differences, + rain_fraction=statistics["rain_fraction"], + marginal_fit_status=statistics["nonzero_fit"]["status"], + spectrum_fit_status=model["status"], + ) + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + print(f"Compared {len(paths)} saved fields; report: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/generate_calibrated_reference.py b/pysteps/tests/data/generate_calibrated_reference.py new file mode 100644 index 00000000..d047085f --- /dev/null +++ b/pysteps/tests/data/generate_calibrated_reference.py @@ -0,0 +1,159 @@ +"""Freeze original generator outputs on synthetic targets without database access. + +Run with the source project's dependencies (including xarray): +python generate_calibrated_reference.py /path/to/pysteps_param + +Normal tests read only the JSON and NPZ fixtures, not the source repository. +""" + +import argparse +import ast +from dataclasses import dataclass, make_dataclass, replace +import hashlib +import json +from pathlib import Path +import subprocess +from typing import Optional + +import numpy as np +from scipy import interpolate, stats +from scipy.optimize import curve_fit + +from generate_parametric_reference import _load_functions + + +def source_generator(parameter_repo): + """Load the actual source functions, result class, and constants. + + Supply a five-scalar dataclass in place of the operational StepsParameters. + The numerical functions access only these five attributes; replace() retains + its original semantics. No QA selection or database identity is introduced. + """ + import xarray as xr + + source = parameter_repo / "src/pysteps_param" + parameters = make_dataclass( + "ScalarParameters", + [ + (name, float) + for name in ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + ], + ) + namespace = dict( + __name__=__name__, + np=np, + xr=xr, + stats=stats, + interpolate=interpolate, + curve_fit=curve_fit, + dataclass=dataclass, + replace=replace, + Optional=Optional, + StepsParameters=parameters, + ) + stats_hash = _load_functions( + source / "rain_stats_compute.py", {"power_spectrum_1D"}, namespace + ) + path = source / "stochastic_generator.py" + original = path.read_text() + nodes = [ + node + for node in ast.parse(original).body + if isinstance(node, ast.FunctionDef) + or isinstance(node, ast.ClassDef) + and node.name == "CalibratedStochasticField" + or isinstance(node, ast.Assign) + and all( + isinstance(target, ast.Name) and target.id.isupper() + for target in node.targets + ) + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(path), "exec"), namespace) + provenance = dict( + source_revision=subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=parameter_repo, text=True + ).strip(), + source_sha256={ + "rain_stats_compute.py": stats_hash, + "stochastic_generator.py": hashlib.sha256(original.encode()).hexdigest(), + }, + ) + return namespace, provenance + + +def main(): + import scipy + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + args = parser.parse_args() + source, provenance = source_generator(args.parameter_repo) + provenance.update(numpy_version=np.__version__, scipy_version=scipy.__version__) + metadata = dict(provenance=provenance, cases=[]) + arrays = {} + cases = [ + ((128, 96), 0.08, 4, {}), + ((128, 96), 0.3, 17, {}), + ((95, 127), 0.05, 11, {"distance_threshold": 0.01}), + ] + for index, (shape, fraction, seed, settings) in enumerate(cases): + parameters = dict( + beta_1=-2.4, beta_2=-3.1, rain_fraction=fraction, mean_db=6.0, stdev_db=3.5 + ) + target = source["StepsParameters"](**parameters) + result = source["gen_calibrated_stoch_field"]( + target, + shape[1], + shape[0], + 1.0, + 20.0, + 0.0, + rng=np.random.default_rng(seed), + **settings, + ) + arrays[f"calibrated_{index}"] = result.field + arrays[f"uncalibrated_{index}"] = source["gen_stoch_field"]( + target, shape[1], shape[0], 1.0, 20.0, 0.0, rng=np.random.default_rng(seed) + ) + metadata["cases"].append( + dict( + shape=shape, + parameters=parameters, + seed=seed, + kmperpixel=1.0, + scale_break=20.0, + threshold=0.0, + settings=settings, + diagnostics={ + name: getattr(result, name) + for name in ( + "estimated_beta_1", + "estimated_beta_2", + "distance", + "attempts", + "selected_attempt", + "tolerance_achieved", + ) + }, + ) + ) + target = source["StepsParameters"]( + beta_1=-2.0, beta_2=-3.0, rain_fraction=0.35, mean_db=6.5, stdev_db=1.2 + ) + arrays["post_ar_input"] = np.random.default_rng(19).normal(size=(48, 64)) + arrays["post_ar_output"] = source["normalize_db_field"]( + arrays["post_ar_input"], target, 0.0, -0.1 + ) + metadata["post_ar_parameters"] = dict( + rain_fraction=0.35, mean=6.5, stdev=1.2, threshold=0.0, zerovalue=-0.1 + ) + prefix = Path(__file__).with_name("calibrated_generator_reference") + np.savez_compressed(prefix.with_suffix(".npz"), **arrays) + prefix.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, allow_nan=False) + "\n" + ) + print(prefix) + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/generate_nowcast_reference.py b/pysteps/tests/data/generate_nowcast_reference.py new file mode 100644 index 00000000..3473b7cd --- /dev/null +++ b/pysteps/tests/data/generate_nowcast_reference.py @@ -0,0 +1,315 @@ +"""Freeze the original preparation, QA, and multi-lead forecast numerics. + +Run in the source project's environment: +python generate_nowcast_reference.py /path/to/pysteps_param + +No database is accessed. AST extraction loads original numerical definitions; +only a relative import inside QA is removed to resolve the extracted helpers. +Reference independent streams are supplied explicitly to the original generator. +""" + +import argparse +import ast +from copy import deepcopy +from dataclasses import dataclass, field +import hashlib +import inspect +import json +from pathlib import Path +from typing import Any, ClassVar, Mapping, TypedDict, cast +import warnings +import sys + +import numpy as np +from scipy.ndimage import gaussian_filter, map_coordinates +from scipy.optimize import minimize_scalar +from scipy.special import log_ndtr + +from pysteps import extrapolation +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.cascade.decomposition import decomposition_fft +from pysteps.timeseries.autoregression import ( + estimate_ar_params_yw, + adjust_lag2_corrcoef2, +) + +from generate_calibrated_reference import source_generator +from generate_parametric_reference import _load_functions + + +def source_engine(parameter_repo): + namespace, provenance = source_generator(parameter_repo) + source = parameter_repo / "src/pysteps_param" + namespace.update( + __file__=str(source / "steps_params.py"), + deepcopy=deepcopy, + copy=__import__("copy"), + np=np, + Path=Path, + json=json, + dataclass=dataclass, + field=field, + ClassVar=ClassVar, + Mapping=Mapping, + warnings=warnings, + Any=Any, + TypedDict=TypedDict, + cast=cast, + minimize_scalar=minimize_scalar, + log_ndtr=log_ndtr, + extrapolation=extrapolation, + map_coordinates=map_coordinates, + estimate_ar_params_yw=estimate_ar_params_yw, + adjust_lag2_corrcoef2=adjust_lag2_corrcoef2, + filter_gaussian=filter_gaussian, + decomposition_fft=decomposition_fft, + ) + from pysteps.cascade.decomposition import recompose_fft + + namespace["recompose_fft"] = recompose_fft + for filename, names in [ + ( + "rain_stats_compute.py", + { + "compute_field_stats", + "fit_lower_truncated_normal", + "correlation_length", + "lagr_auto_cor", + }, + ), + ("cascade_utils.py", {"calculate_wavelengths"}), + ( + "shared_utils.py", + { + "update_field", + "prepare_advection_coordinates", + "_adjust_model_lag2", + "fit_auto_cors", + "calc_model_auto_cors", + "model_correlation_diagnostics", + }, + ), + ]: + provenance["source_sha256"][filename] = _load_functions( + source / filename, names, namespace + ) + for filename, wanted in [ + ("model_config.py", None), + ("transformer.py", {"BaseTransformer", "DBTransformer"}), + ("steps_params.py", {"StepsParameters"}), + ]: + path = source / filename + tree = ast.parse(path.read_text()) + if wanted is not None: + tree.body = [ + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name in wanted + ] + # QA's relative import resolves to the same original functions above. + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + node.body = [ + item + for item in node.body + if not isinstance(item, ast.ImportFrom) + ] + exec(compile(tree, str(path), "exec"), namespace) + provenance["source_sha256"][filename] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + namespace["get_transformer"] = lambda name, threshold: namespace["DBTransformer"]( + threshold=threshold + ) + namespace["CORRELATION_TIME_DEFINITION"] = "ar_acf_sum_including_zero_v2" + provenance["source_sha256"]["cli/nowcast.py"] = hashlib.sha256( + (source / "cli/nowcast.py").read_bytes() + ).hexdigest() + provenance["source_sha256"]["config/low_war_fallback.json"] = hashlib.sha256( + (source / "config/low_war_fallback.json").read_bytes() + ).hexdigest() + import scipy + + provenance.update( + python_version=sys.version.split()[0], + numpy_version=np.__version__, + scipy_version=scipy.__version__, + native_pysteps_sha256={ + name: hashlib.sha256( + Path(inspect.getsourcefile(function)).read_bytes() + ).hexdigest() + for name, function in ( + ("cascade/decomposition.py", decomposition_fft), + ("cascade/bandpass_filters.py", filter_gaussian), + ("timeseries/autoregression.py", estimate_ar_params_yw), + ( + "extrapolation/semilagrangian.py", + extrapolation.get_method("semilagrangian"), + ), + ) + }, + ) + return namespace, provenance + + +def config(shape, timestep=5.0, kmperpixel=1.0, scale_break=20.0): + return dict( + n_rows=shape[0], + n_cols=shape[1], + n_cascade_levels=6, + n_forecasts=3, + n_ensembles=2, + p_size=kmperpixel * 1000, + ht=0.81, + alpha=2.72, + timestep=timestep * 60, + ar_order=2, + rain_threshold=1.0, + scale_break_km=scale_break, + transform="dbr", + ) + + +def prepare(source, rain, velocity, cfg): + transformer = source["DBTransformer"](threshold=cfg["rain_threshold"]) + transformed = transformer.transform(rain) + bp_filter = filter_gaussian(rain.shape[1:], cfg["n_cascade_levels"]) + cascades = [ + decomposition_fft(f, bp_filter, compute_stats=True, normalize=True) + for f in transformed + ] + marginal = source["compute_field_stats"]( + transformed[-1], transformer.metadata, fit_truncated=True, lower_truncation=1.25 + ) + _, spectrum = source["power_spectrum_1D"]( + transformed[-1], cfg["scale_break_km"] / (cfg["p_size"] / 1000) + ) + correlations = source["lagr_auto_cor"]( + np.stack([c["cascade_levels"] for c in cascades]), velocity + ) + reference_time = source["correlation_length"]( + *correlations[0], cfg["timestep"] / 60 + ) + parameters = source["StepsParameters"]( + metadata={}, + ordinary_mean_db=marginal["ordinary_mean"], + ordinary_stdev_db=marginal["ordinary_stdev"], + latent_mean_db=marginal["latent_mean"], + latent_stdev_db=marginal["latent_stdev"], + rain_fraction=marginal["nonzero_fraction"], + beta_1=spectrum["beta_1"], + beta_2=spectrum["beta_2"], + corl_zero=reference_time, + ).apply_qa(cfg) + return parameters, cascades, bp_filter, transformer + + +def evolve(source, rain, velocity, cfg, parameters, seed, n_members=2, n_leads=3): + bp_filter = filter_gaussian(rain.shape[1:], cfg["n_cascade_levels"]) + transformer = source["DBTransformer"](threshold=cfg["rain_threshold"]) + init = [ + decomposition_fft( + transformer.transform(f), bp_filter, compute_stats=True, normalize=True + ) + for f in rain[-1:-3:-1] + ] + coordinates = source["prepare_advection_coordinates"](velocity) + taper = gaussian_filter( + extrapolation.get_method("semilagrangian")( + np.ones(rain.shape[1:]), velocity, 1, outval=0 + )[0], + 5.0, + ) + seeds = np.random.SeedSequence(seed).spawn(n_members) + original_generator = source["gen_calibrated_stoch_field"] + output = np.empty((n_members, n_leads, *rain.shape[1:])) + calibration = [[] for _ in range(n_members)] + try: + for member in range(n_members): + histories = deepcopy(init) + rng = np.random.default_rng(seeds[member]) + + def seeded_generator(*args, **kwargs): + result = original_generator(*args, **kwargs, rng=rng) + calibration[member].append( + {k: v for k, v in vars(result).items() if k != "field"} + ) + return result + + source["gen_calibrated_stoch_field"] = seeded_generator + for lead in range(n_leads): + field = source["update_field"]( + histories, + velocity, + parameters, + bp_filter, + cfg, + advection_coordinates=coordinates, + ) + output[member, lead] = transformer.inverse_transform(field) + histories = [ + decomposition_fft( + field * taper, bp_filter, compute_stats=True, normalize=True + ), + histories[0], + ] + finally: + source["gen_calibrated_stoch_field"] = original_generator + return output, calibration + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + args = parser.parse_args() + source, provenance = source_engine(args.parameter_repo) + arrays, cases = {}, [] + for index, fraction in enumerate([0.2, 0.04]): + shape = (128, 96) + cfg = config(shape, timestep=7.5 if index else 5.0) + target = source["StepsParameters"]( + metadata={}, + beta_1=-2.4, + beta_2=-3.1, + rain_fraction=fraction, + mean_db=6, + stdev_db=3, + ) + dbr = source["gen_stoch_field"]( + target, shape[1], shape[0], 1, 20, 0, rng=np.random.default_rng(9 + index) + ) + rain = 10 ** (dbr / 10) + rain[rain < 1] = 0 + history = np.stack([np.roll(rain, -2, axis=1), np.roll(rain, -1, axis=1), rain]) + velocity = np.zeros((2, *shape)) + velocity[0] = 1 + parameters, _, _, _ = prepare(source, history, velocity, cfg) + seed = 20261015 + index + output, calibration = evolve(source, history, velocity, cfg, parameters, seed) + arrays.update( + { + f"history_{index}": history, + f"velocity_{index}": velocity, + f"forecast_{index}": output, + } + ) + cases.append( + dict( + config=cfg, + seed=seed, + parameters=parameters.active_parameters(), + moment_source=parameters.qa["moment_source"], + calibration=calibration, + ) + ) + folder = Path(__file__).parent + np.savez_compressed(folder / "parametric_nowcast_reference.npz", **arrays) + (folder / "parametric_nowcast_reference.json").write_text( + json.dumps(dict(provenance=provenance, cases=cases), indent=2, allow_nan=False) + + "\n" + ) + print(f"Wrote {len(cases)} source forecast cases") + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/generate_parametric_reference.py b/pysteps/tests/data/generate_parametric_reference.py new file mode 100644 index 00000000..b1440919 --- /dev/null +++ b/pysteps/tests/data/generate_parametric_reference.py @@ -0,0 +1,164 @@ +"""Regenerate numerical references from the two original source repositories. + +Run in the source project's environment (including xarray): +python generate_parametric_reference.py /path/to/pysteps_param /path/to/psteps_param_paper + +Tests read only the generated JSON and do not need either external repository. +""" + +import argparse +import ast +import hashlib +import json +from pathlib import Path +import subprocess + +import numpy as np +from scipy.optimize import brentq, curve_fit, minimize_scalar +from scipy.special import log_ndtr + +from pysteps import extrapolation +from pysteps.cascade.bandpass_filters import filter_gaussian + + +def _load_functions(path, names, namespace): + # Load the original numerical definitions without executing package-level + # imports of database or operational modules. + source = path.read_text() + definitions = [ + node + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) and node.name in names + ] + if {node.name for node in definitions} != set(names): + raise ValueError(f"Missing source functions in {path}") + exec( + compile(ast.Module(body=definitions, type_ignores=[]), str(path), "exec"), + namespace, + ) + return hashlib.sha256(source.encode()).hexdigest() + + +def main(): + import xarray as xr + import scipy + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("parameter_repo", type=Path) + parser.add_argument("paper_repo", type=Path) + args = parser.parse_args() + source = args.parameter_repo / "src/pysteps_param" + namespace = dict( + np=np, + xr=xr, + curve_fit=curve_fit, + minimize_scalar=minimize_scalar, + log_ndtr=log_ndtr, + brentq=brentq, + extrapolation=extrapolation, + ) + hashes = {} + hashes["pysteps_param/src/pysteps_param/rain_stats_compute.py"] = _load_functions( + source / "rain_stats_compute.py", + { + "power_spectrum_1D", + "fit_lower_truncated_normal", + "compute_field_stats", + "correlation_length", + "lagr_auto_cor", + }, + namespace, + ) + hashes["pysteps_param/src/pysteps_param/shared_utils.py"] = _load_functions( + source / "shared_utils.py", {"fit_auto_cors", "_adjust_model_lag2"}, namespace + ) + hashes["psteps_param_paper/workflows/reference_timescale_experiment.py"] = ( + _load_functions( + args.paper_repo / "workflows/reference_timescale_experiment.py", + {"reference_log"}, + namespace, + ) + ) + result = dict( + provenance=dict( + source_revisions={ + name: subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, text=True + ).strip() + for name, root in [ + ("pysteps_param", args.parameter_repo), + ("psteps_param_paper", args.paper_repo), + ] + }, + source_sha256=hashes, + numpy_version=np.__version__, + scipy_version=scipy.__version__, + description="Numerical references evaluated from original function definitions on deterministic synthetic inputs; no database access.", + ), + spectra=[], + ar2_inverse=[], + ) + for shape in [(128, 128), (95, 127)]: + y, x = np.indices(shape) + field = np.random.default_rng(23).normal(size=shape) + 4 * np.cos( + 2 * np.pi * (5 * y / shape[0] + 3 * x / shape[1]) + ) + models = [] + for scale_break in (None, 20): + spectrum, model = namespace["power_spectrum_1D"](field, scale_break) + models.append(model) + result["spectra"].append( + dict( + shape=shape, + frequency_db=spectrum.psd_bins.values.tolist(), + power_db=spectrum.psd.values.tolist(), + models=models, + ) + ) + rng = np.random.default_rng(37) + field = rng.normal(4, 3, size=(128, 96)) + field[rng.uniform(size=field.shape) > 0.3] = -0.1 + stats = namespace["compute_field_stats"]( + field, + {"threshold": 0, "zerovalue": -0.1, "transform": "dB"}, + fit_truncated=True, + lower_truncation=1.25, + ) + result["marginal"] = { + name: stats[name] + for name in [ + "ordinary_mean", + "ordinary_stdev", + "latent_mean", + "latent_stdev", + "mean", + "stdev", + ] + } + result["marginal"]["rain_fraction"] = stats["nonzero_fraction"] + history = np.random.default_rng(45).normal(size=(3, 2, 24, 32)) + velocity = np.full((2, 24, 32), 0.25) + velocity[0] = 0.6 + result["lags"] = namespace["lagr_auto_cor"](history, velocity).tolist() + for alpha in (0.8, 2.72, 4): + for time in (5, 10, 60, 420): + r1, r2, _ = namespace["fit_auto_cors"](time, alpha, 5) + result["ar2_inverse"].append(dict(time=time, alpha=alpha, pair=[r1, r2])) + wavelengths = 1 / filter_gaussian((256, 256), 6, d=1)["central_freqs"] + times = np.array([[60, 42, 25, 10, 8, 6], [120, 70, 31, 20, 10, 5]]) + exponent = 0.7718612877449645 + result["paper_reference"] = dict( + wavelengths=wavelengths.tolist(), + times=times.tolist(), + exponent=exponent, + reference_times=np.exp( + namespace["reference_log"](times, wavelengths, exponent) + ).tolist(), + ) + path = Path(__file__).with_name("parametric_estimators_reference.json") + path.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print(path) + + +if __name__ == "__main__": + main() diff --git a/pysteps/tests/data/parametric_estimators_reference.json b/pysteps/tests/data/parametric_estimators_reference.json new file mode 100644 index 00000000..be59727d --- /dev/null +++ b/pysteps/tests/data/parametric_estimators_reference.json @@ -0,0 +1,420 @@ +{ + "provenance": { + "source_revisions": { + "pysteps_param": "94bd4c1571637a8f246ac391470e1a820d6376a6", + "psteps_param_paper": "335f7660775480d865d9771820a81923c461832a" + }, + "source_sha256": { + "pysteps_param/src/pysteps_param/rain_stats_compute.py": "2fd6f0b622e88e143d8128f80cb5d0920f0111c1021c6c448222c90bff92c889", + "pysteps_param/src/pysteps_param/shared_utils.py": "4e7346cda389b99971246e4cf060adeec9c51cc9444704446a071ccc44a5288f", + "psteps_param_paper/workflows/reference_timescale_experiment.py": "f93769bb6b4f94701b682c2947bfec661d0fa089fafa0ad2132d4744e7d6cee5" + }, + "numpy_version": "2.5.1", + "scipy_version": "1.18.0", + "description": "Numerical references evaluated from original function definitions on deterministic synthetic inputs; no database access." + }, + "spectra": [ + { + "shape": [ + 128, + 128 + ], + "frequency_db": [ + -17.667251038327528, + -16.62037767221858, + -16.358659330691346, + -16.096940989164107, + -15.573504306109635, + -15.050067623055163, + -14.788349281527926, + -14.52663094000069, + -14.003194256946216, + -13.74147591541898, + -13.479757573891744, + -13.218039232364507, + -12.956320890837269, + -12.694602549310032, + -12.432884207782795, + -12.17116586625556, + -11.909447524728321, + -11.647729183201086, + -11.386010841673848, + -11.12429250014661, + -10.862574158619376, + -10.600855817092139, + -10.3391374755649, + -10.077419134037664, + -9.815700792510428, + -9.553982450983193, + -9.292264109455955, + -9.03054576792872, + -8.768827426401483, + -8.507109084874246, + -8.24539074334701, + -7.983672401819772, + -7.721954060292536, + -7.460235718765299, + -7.198517377238063, + -6.936799035710824, + -6.675080694183588, + -6.413362352656352, + -6.151644011129115, + -5.889925669601878, + -5.628207328074641, + -5.366488986547405, + -5.104770645020168, + -4.843052303492932, + -4.581333961965695, + -4.319615620438459, + -4.057897278911222, + -3.796178937383985, + -3.5344605958567494, + -3.2727422543295113, + -3.011023912802275, + -2.7493055712750385, + -2.4875872297478026, + -2.225868888220565, + -1.9641505466933276, + -1.702432205166092 + ], + "power_db": [ + 32.21404823693171, + 32.89115806726739, + 31.577373883302656, + 28.82632685065227, + 32.53204791407693, + 18.227053814532425, + 32.34731566908866, + 32.59450536851558, + 34.32142466901202, + 32.22744056651604, + 72.98771267445036, + 30.565395736421802, + 34.95378940097872, + 29.793454900506084, + 30.17849205522203, + 30.86793520242505, + 33.25565618937198, + 32.28764048250997, + 32.474320741622556, + 34.88645980840497, + 32.45018941268981, + 33.86022169320155, + 34.03311950786575, + 31.828145357820013, + 32.2523420785126, + 31.985915536416858, + 33.33305235105088, + 33.17467990441208, + 32.26714354192855, + 31.98576568056385, + 33.20549716177654, + 33.1596442083763, + 32.81630883012787, + 32.45988258537234, + 32.84419641634696, + 32.56861411042712, + 32.71016110087893, + 32.67319800011376, + 32.67739642508208, + 32.97680110663359, + 32.82724636105732, + 32.345353412324805, + 32.5310842412989, + 32.897376674033836, + 32.27736692836629, + 32.67337497580174, + 32.38010211913294, + 32.533425911263386, + 32.38907888191223, + 32.34062182235669, + 32.384596188693415, + 32.47910959397415, + 32.85629115161199, + 32.712122674743746, + 32.84431812332657, + 33.18497928300023 + ], + "models": [ + { + "beta_1": -0.05268062227348369, + "beta_2": null, + "c1": 32.54838859499382, + "c2": null, + "scale_break": null + }, + { + "beta_1": 1.750219739903674, + "beta_2": -0.40906661474982897, + "c1": 57.923321834892, + "c2": 29.830358668570593, + "scale_break": 20.0 + } + ] + }, + { + "shape": [ + 95, + 127 + ], + "frequency_db": [ + -16.553811597148872, + -16.225674374749133, + -15.89753715234939, + -14.913125485150168, + -14.584988262750427, + -14.256851040350684, + -13.928713817950946, + -13.600576595551207, + -13.272439373151464, + -12.944302150751723, + -12.616164928351983, + -12.28802770595224, + -11.959890483552499, + -11.63175326115276, + -11.303616038753017, + -10.975478816353279, + -10.647341593953536, + -10.319204371553797, + -9.991067149154055, + -9.662929926754313, + -9.334792704354573, + -9.00665548195483, + -8.67851825955509, + -8.350381037155351, + -8.02224381475561, + -7.694106592355868, + -7.365969369956128, + -7.0378321475563865, + -6.709694925156645, + -6.381557702756905, + -6.053420480357165, + -5.725283257957425, + -5.397146035557682, + -5.069008813157943, + -4.740871590758201, + -4.41273436835846, + -4.084597145958718, + -3.756459923558978, + -3.428322701159238, + -3.1001854787594976, + -2.7720482563597564, + -2.4439110339600143, + -2.115773811560275, + -1.7876365891605324 + ], + "power_db": [ + 29.596752108332932, + 28.599293553596823, + 33.77561208788202, + 31.812027504318486, + 29.098075789303095, + 30.40169398250597, + 29.327640206538494, + 27.363420838411884, + 31.733827101431466, + 31.984511553230636, + 32.370948559104335, + 69.0752136877862, + 31.664778170206, + 30.87052803355375, + 30.2926202517147, + 32.23407752743059, + 31.10302473421841, + 31.76113827303327, + 31.237543669781328, + 32.21429105170373, + 32.443496881228896, + 32.101719729101546, + 31.952552126310287, + 31.450292614378963, + 31.70056341231345, + 30.785298020765595, + 31.38208548693949, + 30.76054457607694, + 30.415209154698513, + 31.404183893420626, + 31.685923903583767, + 31.664054109928223, + 31.64058740674395, + 31.220371947476643, + 31.384272922719155, + 30.951877101954558, + 31.203472798749754, + 31.306650889532584, + 31.07704516767248, + 31.3949229114742, + 31.41658267109576, + 31.19917051893485, + 31.11708521309568, + 30.934879607893023 + ], + "models": [ + { + "beta_1": -0.21547126372428127, + "beta_2": null, + "c1": 30.27008964716364, + "c2": null, + "scale_break": null + }, + { + "beta_1": 1.8738479155680257, + "beta_2": -0.5102178985934316, + "c1": 59.072933167953764, + "c2": 28.0555218093425, + "scale_break": 20.0 + } + ] + } + ], + "ar2_inverse": [ + { + "time": 5, + "alpha": 0.8, + "pair": [ + 0.0, + 0.0 + ] + }, + { + "time": 10, + "alpha": 0.8, + "pair": [ + 0.25029703934242475, + 0.3301904970651707 + ] + }, + { + "time": 60, + "alpha": 0.8, + "pair": [ + 0.7516432139329509, + 0.7958100010789851 + ] + }, + { + "time": 420, + "alpha": 0.8, + "pair": [ + 0.9547653195683252, + 0.9636455237978131 + ] + }, + { + "time": 5, + "alpha": 2.72, + "pair": [ + 0.0, + 0.0 + ] + }, + { + "time": 10, + "alpha": 2.72, + "pair": [ + 0.5700065264871803, + 0.25880799846894853 + ] + }, + { + "time": 60, + "alpha": 2.72, + "pair": [ + 0.9580235217567927, + 0.889904060320336 + ] + }, + { + "time": 420, + "alpha": 2.72, + "pair": [ + 0.9943443348521103, + 0.9846913121295453 + ] + }, + { + "time": 5, + "alpha": 4, + "pair": [ + 0.0, + 0.0 + ] + }, + { + "time": 10, + "alpha": 4, + "pair": [ + 0.5700065264871803, + 0.25880799846894853 + ] + }, + { + "time": 60, + "alpha": 4, + "pair": [ + 0.9852101493475673, + 0.9498681838104814 + ] + }, + { + "time": 420, + "alpha": 4, + "pair": [ + 0.9997132719080527, + 0.9988800690605281 + ] + } + ], + "marginal": { + "ordinary_mean": 4.5544651679144215, + "ordinary_stdev": 2.544133913407154, + "latent_mean": 4.02706361980481, + "latent_stdev": 2.9781258381020006, + "mean": 1.139150098359972, + "stdev": 2.5022682644201866, + "rain_fraction": 0.2740885416666667 + }, + "lags": [ + [ + 0.007973669863881415, + 0.01433718006760386 + ], + [ + -0.005838299935575144, + 0.005083562758421983 + ] + ], + "paper_reference": { + "wavelengths": [ + 256.0, + 70.28517926862416, + 31.308488057340607, + 13.946345938598434, + 6.212390859719535, + 2.0 + ], + "times": [ + [ + 60, + 42, + 25, + 10, + 8, + 6 + ], + [ + 120, + 70, + 31, + 20, + 10, + 5 + ] + ], + "exponent": 0.7718612877449645, + "reference_times": [ + 95.28171104914436, + 152.91251252988084 + ] + } +} diff --git a/pysteps/tests/data/parametric_nowcast_reference.json b/pysteps/tests/data/parametric_nowcast_reference.json new file mode 100644 index 00000000..ef566ae0 --- /dev/null +++ b/pysteps/tests/data/parametric_nowcast_reference.json @@ -0,0 +1,221 @@ +{ + "provenance": { + "source_revision": "94bd4c1571637a8f246ac391470e1a820d6376a6", + "source_sha256": { + "rain_stats_compute.py": "2fd6f0b622e88e143d8128f80cb5d0920f0111c1021c6c448222c90bff92c889", + "stochastic_generator.py": "94c491e4fd90da0c733ce65b48227b6b73dc31a8fc6a23a44e211ad36b3f9343", + "cascade_utils.py": "77c362a23aa46b214f1d7f613fe0e747a5f6e7f861c17f85094c5b91db13d403", + "shared_utils.py": "4e7346cda389b99971246e4cf060adeec9c51cc9444704446a071ccc44a5288f", + "model_config.py": "3504c8492f8158cd84b155753cb205e0a7810cd04c37d35f253c9d25fc374128", + "transformer.py": "8693d646602538cd75332cd4367170b4387c9844102b011b7090dae033482703", + "steps_params.py": "265c31239ac8647fd599648e0f00dbdbc8c963e4a7bdaadfa89edb177c4c6c28", + "cli/nowcast.py": "d892c72f8daa608669fce0e185397a9664644b53663eeb1d5cc3708ceb7bec26", + "config/low_war_fallback.json": "bf0b3e0c9c12d4ce0f1cf64d3fe997a8e3364f285666071c24835c3141c45e87" + }, + "python_version": "3.12.13", + "numpy_version": "2.5.1", + "scipy_version": "1.18.0", + "native_pysteps_sha256": { + "cascade/decomposition.py": "2963be8c5399afa8c066ab6df93d199f8362f476a93d24add96a508aecb2b84e", + "cascade/bandpass_filters.py": "0527beb2aa67a620425479fa09fa18dbf8ef311109a1de72e9df6d436ee7403b", + "timeseries/autoregression.py": "08205f5de8d6dbe6fea0d3021c929d50990a6379098359a1aab6c4315808d851", + "extrapolation/semilagrangian.py": "5f870efe3b9152fefd417ecc61ffa188d8e9a306515f676b66890d9f1443e9c3" + } + }, + "cases": [ + { + "config": { + "n_rows": 128, + "n_cols": 96, + "n_cascade_levels": 6, + "n_forecasts": 3, + "n_ensembles": 2, + "p_size": 1000.0, + "ht": 0.81, + "alpha": 2.72, + "timestep": 300.0, + "ar_order": 2, + "rain_threshold": 1.0, + "scale_break_km": 20.0, + "transform": "dbr" + }, + "seed": 20261015, + "parameters": { + "mean_db": 6.0028472842915725, + "stdev_db": 3.0012414749875695, + "rain_fraction": 0.19710286458333334, + "beta_1": -2.304747110331049, + "beta_2": -2.5631814182884716, + "corl_zero": 456.36535657886674, + "lag_1": [ + 0.9947989543385891, + 0.9871752674307953, + 0.9771449287175584, + 0.9587433096366739, + 0.9237438162041633, + 0.7935185195767728 + ], + "lag_2": [ + 0.9859163541951547, + 0.9655002791814661, + 0.9390493841270735, + 0.89172384927532, + 0.8059353402996958, + 0.5395489647808457 + ] + }, + "moment_source": "latent", + "calibration": [ + [ + { + "estimated_beta_1": -2.97392072940944, + "estimated_beta_2": -2.483368847914942, + "distance": 2.58728403725768, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -1.817123223192054, + "estimated_beta_2": -2.6302184889764453, + "distance": 1.8891966298004956, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -2.9519417030723605, + "estimated_beta_2": -2.4530007388125172, + "distance": 2.518418223677037, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + } + ], + [ + { + "estimated_beta_1": -1.7058699586645432, + "estimated_beta_2": -2.9044931766578634, + "distance": 2.6113997297501057, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -2.8656737433833874, + "estimated_beta_2": -2.5431414850604748, + "distance": 2.1564304711388225, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -2.8926687128107043, + "estimated_beta_2": -2.625601365428789, + "distance": 2.270191677685425, + "attempts": 1, + "selected_attempt": 1, + "tolerance_achieved": false + } + ] + ] + }, + { + "config": { + "n_rows": 128, + "n_cols": 96, + "n_cascade_levels": 6, + "n_forecasts": 3, + "n_ensembles": 2, + "p_size": 1000.0, + "ht": 0.81, + "alpha": 2.72, + "timestep": 450.0, + "ar_order": 2, + "rain_threshold": 1.0, + "scale_break_km": 20.0, + "transform": "dbr" + }, + "seed": 20261016, + "parameters": { + "mean_db": 6.0135757197788715, + "stdev_db": 3.0093428695820377, + "rain_fraction": 0.03955078125, + "beta_1": -1.45, + "beta_2": -3.253958088759979, + "corl_zero": 330.61752959848894, + "lag_1": [ + 0.9891292643911687, + 0.9728202376077736, + 0.9506635983379176, + 0.9078567874780119, + 0.816699134701302, + 0.37407983990030463 + ], + "lag_2": [ + 0.9707073074247313, + 0.9277878117174473, + 0.8714311340224543, + 0.7687890306712478, + 0.5776929978968777, + 0.10753778774958642 + ] + }, + "moment_source": "latent", + "calibration": [ + [ + { + "estimated_beta_1": -1.523354904476627, + "estimated_beta_2": -2.8895547119152427, + "distance": 1.3481136783625316, + "attempts": 5, + "selected_attempt": 3, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -1.5548026590261312, + "estimated_beta_2": -2.975659648574481, + "distance": 1.0843549412999052, + "attempts": 5, + "selected_attempt": 3, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -0.6578986279771498, + "estimated_beta_2": -2.7596856788529567, + "distance": 3.5298757819809303, + "attempts": 5, + "selected_attempt": 2, + "tolerance_achieved": false + } + ], + [ + { + "estimated_beta_1": -1.2562055284061402, + "estimated_beta_2": -2.827728216458951, + "distance": 1.7123619898017537, + "attempts": 5, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -2.054098896544064, + "estimated_beta_2": -2.944043379979809, + "distance": 2.577699519586813, + "attempts": 5, + "selected_attempt": 1, + "tolerance_achieved": false + }, + { + "estimated_beta_1": -1.1774689342257276, + "estimated_beta_2": -2.769690712172061, + "distance": 2.0410397132918026, + "attempts": 5, + "selected_attempt": 1, + "tolerance_achieved": false + } + ] + ] + } + ] +} diff --git a/pysteps/tests/data/parametric_nowcast_reference.npz b/pysteps/tests/data/parametric_nowcast_reference.npz new file mode 100644 index 00000000..e2390ffa Binary files /dev/null and b/pysteps/tests/data/parametric_nowcast_reference.npz differ diff --git a/pysteps/tests/data/parametric_simulation_reference.json b/pysteps/tests/data/parametric_simulation_reference.json new file mode 100644 index 00000000..a2e69c7b --- /dev/null +++ b/pysteps/tests/data/parametric_simulation_reference.json @@ -0,0 +1,205 @@ +{ + "provenance": { + "description": "Fixed-seed pySTEPS simulation and parameter-recovery regression baseline. The translated initial history retains a persistence transient; recovered times are not required to equal the prescribed times.", + "pysteps_revision": "02c52c6976b2b2013c9a710af73d731a7fd2aff0", + "example_sha256": "95bdc5835cf810b09a93e9b03531ce1820dded3c0d548e66f91c88de999a2931", + "python_version": "3.13.18", + "numpy_version": "2.5.4", + "scipy_version": "1.18.1" + }, + "case": { + "shape": [ + 128, + 128 + ], + "kmperpixel": 1.0, + "timestep_minutes": 5.0, + "n_cascade_levels": 6, + "rain_threshold": 1.0, + "scale_break_km": 20.0, + "parameters": { + "beta_1": -2.4, + "beta_2": -3.1, + "rain_fraction": 0.2, + "mean_db": 6.0, + "stdev_db": 3.0 + }, + "reference_time_minutes": 400.0, + "reference_wavelength_km": 256.0, + "temporal_exponent": 0.81, + "alpha": 2.72, + "initial_seed": 9, + "evolution_seed": 42, + "n_timesteps": 12, + "n_ens_members": 1, + "velocity_pixels_per_timestep": [ + 1.0, + 0.0 + ], + "lower_truncation_dbr": 1.25 + }, + "expected": { + "field_parameter_names": [ + "rain_fraction", + "latent_mean", + "latent_stdev", + "beta_1", + "beta_2" + ], + "field_parameters": [ + [ + 0.19671630859375, + 5.986249096139553, + 3.012794880676619, + -2.2557237980868403, + -3.1141331601444824 + ], + [ + 0.19573974609375, + 6.016000339472105, + 2.9937188072117156, + -2.3026626012007245, + -3.2753025735088093 + ], + [ + 0.19732666015625, + 6.002121903317441, + 3.0059439970753967, + -2.197762094766152, + -3.1665853591479802 + ], + [ + 0.196533203125, + 6.01078589534421, + 3.000267695191747, + -1.9024561361129966, + -3.3591817187947495 + ], + [ + 0.19647216796875, + 6.007605236995318, + 3.001684868552406, + -1.381710667204899, + -3.128221868555097 + ], + [ + 0.19671630859375, + 5.995297095394097, + 3.009845705291915, + -1.4967524364344298, + -3.0656365574002926 + ], + [ + 0.19677734375, + 6.01321965613408, + 2.9974893524016766, + -1.8162638082199871, + -3.189509935219112 + ], + [ + 0.197021484375, + 6.006500846058137, + 3.0011801577347486, + -2.209983980362383, + -3.27064696478013 + ], + [ + 0.196533203125, + 6.022383951062275, + 2.9906279041013786, + -2.0206483070286563, + -3.03297340660577 + ], + [ + 0.197021484375, + 6.019270015751998, + 2.9913060498988484, + -2.1704243610864493, + -3.174610558183172 + ], + [ + 0.196044921875, + 5.992851718774125, + 3.0109195870655547, + -2.2807586082454363, + -3.1638150624423744 + ], + [ + 0.19610595703125, + 6.007345349393529, + 3.0015274954276343, + -1.9180022593139838, + -3.1622929516961777 + ] + ], + "wavelengths_km": [ + 128.0, + 42.666666666666664, + 21.333333333333332, + 10.666666666666666, + 5.333333333333333, + 2.0 + ], + "model_correlations": [ + [ + 0.9895046127483391, + 0.9717095674103806 + ], + [ + 0.9737829005847584, + 0.9302871714238284 + ], + [ + 0.9524697863763386, + 0.8759418701868156 + ], + [ + 0.9114361898321935, + 0.7770616058765629 + ], + [ + 0.8260762701358226, + 0.5946957013548607 + ], + [ + 0.40098044282778983, + 0.12403262116727216 + ] + ], + "mean_observed_correlations": [ + [ + 0.9956493182859601, + 0.989635285906074 + ], + [ + 0.9855593938093655, + 0.9649218639462054 + ], + [ + 0.9720016927658678, + 0.9301018111857937 + ], + [ + 0.9429181836364829, + 0.8560155629373574 + ], + [ + 0.8550468554132097, + 0.6631306333051005 + ], + [ + 0.5520161033975197, + 0.27420614707810975 + ] + ], + "implied_correlation_times_minutes": [ + 776.992666271944, + 220.70619946046185, + 104.59312856023948, + 48.49459426084739, + 21.58225649541986, + 10.432811132441214 + ], + "implied_reference_time_minutes": 1001.5425576214981 + } +} diff --git a/pysteps/tests/test_noise_parametric.py b/pysteps/tests/test_noise_parametric.py new file mode 100644 index 00000000..e4fc5e5c --- /dev/null +++ b/pysteps/tests/test_noise_parametric.py @@ -0,0 +1,354 @@ +"""Calibrated generator source parity, distribution mapping, and selection tests.""" + +import copy +import json +from pathlib import Path + +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal + +from pysteps.noise import parametric +from pysteps.postprocessing.probmatching import match_rainfall_distribution + + +def _parameters(rain_fraction=0.08): + return dict( + beta_1=-2.4, beta_2=-3.1, rain_fraction=rain_fraction, mean_db=6.0, stdev_db=3.5 + ) + + +@pytest.fixture(scope="module") +def reference(): + root = Path(__file__).parent / "data/calibrated_generator_reference" + metadata = json.loads(root.with_suffix(".json").read_text()) + with np.load(root.with_suffix(".npz")) as stored: + arrays = {name: stored[name] for name in stored.files} + return metadata, arrays + + +def test_full_field_parity_with_original_generator(reference): + metadata, arrays = reference + for index, case in enumerate(metadata["cases"]): + field, diagnostics = parametric.generate_calibrated_field( + case["shape"], + case["parameters"], + case["kmperpixel"], + case["scale_break"], + case["threshold"], + seed=case["seed"], + return_diagnostics=True, + **case["settings"], + ) + assert_allclose(field, arrays[f"calibrated_{index}"], atol=1e-10, rtol=1e-12) + for name in ("attempts", "selected_attempt", "tolerance_achieved"): + assert diagnostics[name] == case["diagnostics"][name] + for name in ("estimated_beta_1", "estimated_beta_2", "distance"): + assert diagnostics[name] == pytest.approx( + case["diagnostics"][name], abs=3e-5 + ) + raw = parametric.generate_stochastic_field( + case["shape"], + case["parameters"], + case["kmperpixel"], + case["scale_break"], + case["threshold"], + seed=case["seed"], + ) + assert_allclose(raw, arrays[f"uncalibrated_{index}"], atol=1e-10, rtol=1e-12) + + +def test_shared_mapping_matches_post_ar_source(reference): + metadata, arrays = reference + original = arrays["post_ar_input"].copy() + result = match_rainfall_distribution( + original, **metadata["post_ar_parameters"], min_rain_fraction=0.025 + ) + assert_allclose(result, arrays["post_ar_output"], atol=1e-12, rtol=1e-12) + assert_array_equal(original, arrays["post_ar_input"]) + + +@pytest.mark.parametrize("shape", [(64, 48), (48, 64), (63, 37)]) +def test_filter_uses_half_power_exponents_and_physical_units(shape): + spacing = 0.5 + beta_1, beta_2, wavelength = -2.0, -3.5, 10.0 + filt = parametric.power_law_filter(shape, beta_1, beta_2, spacing, wavelength) + fy = np.fft.fftfreq(shape[0], spacing) + fx = np.fft.fftfreq(shape[1], spacing) + radius = np.hypot(fy[:, None], fx[None, :]) + reference = 1 / (max(shape) * spacing) + for y, x in [(0, 1), (1, 0), (shape[0] // 2, shape[1] // 2)]: + f = radius[y, x] + expected_power = ( + (f / reference) ** beta_1 + if f < 1 / wavelength + else ((1 / wavelength / reference) ** beta_1 * (f * wavelength) ** beta_2) + ) + assert filt[y, x] ** 2 == pytest.approx(expected_power, rel=1e-13) + assert filt[0, 0] == 1 and np.isfinite(filt).all() + # Scaling pixel spacing and break together leaves dimensionless responses. + assert_allclose( + filt, + parametric.power_law_filter(shape, beta_1, beta_2, spacing * 2, wavelength * 2), + ) + assert_allclose( + filt.T, + parametric.power_law_filter(shape[::-1], beta_1, beta_2, spacing, wavelength), + ) + + +@pytest.mark.parametrize("threshold", [-2, 0, 2]) +def test_mapping_upper_bin_edge_and_threshold_translation(threshold): + data = np.arange(100, dtype=float).reshape(10, 10) + result = match_rainfall_distribution(data, 0.3, 20, 1, threshold, nbins=100) + assert_array_equal(result > threshold, data >= 70) + base = match_rainfall_distribution(data, 0.3, 20, 1, 0, nbins=100) + assert_allclose(result[data < 70], base[data < 70] + threshold) + assert_allclose(result[data >= 70], base[data >= 70]) + assert_array_equal(data, np.arange(100).reshape(10, 10)) + + +def test_marginal_quantile_bounds_and_approximate_statistics(): + parameters = dict(_parameters(0.3), mean_db=6, stdev_db=1) + field = parametric.generate_stochastic_field( + (128, 128), parameters, 1, 20, 0, seed=123 + ) + wet = field[field > 0] + assert abs(wet.size / field.size - 0.3) < 0.02 + assert abs(wet.mean() - 6) < 0.1 and abs(wet.std() - 1) < 0.1 + from scipy.stats import norm + + assert wet.min() >= norm.ppf(0.001, loc=6, scale=1) + assert wet.max() <= norm.ppf(0.999, loc=6, scale=1) + assert field.std() != pytest.approx(1) # Raw marginal field is not standardized. + + +def test_first_acceptable_candidate_and_weighted_distance(monkeypatch): + estimates = iter([(-1.5, -2.2), (-2.3, -3), (-2.4, -3.1)]) + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: next(estimates)) + _, diagnostics = parametric.generate_calibrated_field( + (32, 32), _parameters(0.1), 1, 20, 0, seed=4, return_diagnostics=True + ) + assert diagnostics["attempts"] == diagnostics["selected_attempt"] == 2 + assert diagnostics["tolerance_achieved"] + assert diagnostics["distance"] == pytest.approx( + np.sqrt(0.1**2 / 0.067738 + 0.1**2 / 0.076405) + ) + + +def test_distance_tolerance_is_inclusive(monkeypatch): + parameters = dict(_parameters(0.1), beta_1=-2, beta_2=-3) + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: (-1, -3)) + _, diagnostics = parametric.generate_calibrated_field( + (32, 32), + parameters, + 1, + 20, + 0, + seed=4, + beta_variances=(1, 1), + distance_threshold=1, + return_diagnostics=True, + ) + assert diagnostics["distance"] == 1 and diagnostics["tolerance_achieved"] + assert diagnostics["attempts"] == 1 + + +def test_best_candidate_after_attempt_limit(monkeypatch): + estimates = iter([(-1, -2), (-2, -2.8), (-2.2, -2.9), (-1.8, -2.5), (-1.2, -2.1)]) + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: next(estimates)) + _, diagnostics = parametric.generate_calibrated_field( + (32, 32), + _parameters(0.05), + 1, + 20, + 0, + seed=4, + distance_threshold=0.01, + return_diagnostics=True, + ) + assert ( + not diagnostics["tolerance_achieved"] + and diagnostics["status"] == "best_available" + ) + assert diagnostics["attempts"] == 5 and diagnostics["selected_attempt"] == 3 + assert diagnostics["candidate_estimates"].shape == (5, 2) + + +def test_best_candidate_retains_first_tie(monkeypatch): + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: (-2, -3)) + _, diagnostics = parametric.generate_calibrated_field( + (32, 32), + _parameters(), + 1, + 20, + 0, + seed=4, + distance_threshold=0, + return_diagnostics=True, + ) + assert diagnostics["attempts"] == 5 and diagnostics["selected_attempt"] == 1 + + +def test_all_unavailable_spectra_select_last_candidate(monkeypatch): + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: (np.nan, np.nan)) + field, diagnostics = parametric.generate_calibrated_field( + (32, 32), _parameters(), 1, 20, 0, seed=4, return_diagnostics=True + ) + assert np.isfinite(field).all() and np.isinf(diagnostics["distance"]) + assert diagnostics["attempts"] == diagnostics["selected_attempt"] == 5 + assert ( + diagnostics["status"] == "unavailable_spectrum" + and not diagnostics["tolerance_achieved"] + ) + + +def test_numerical_spectrum_failure_is_reported(monkeypatch): + def unavailable(*_): + raise ValueError("unavailable numerical spectrum") + + monkeypatch.setattr(parametric, "rapsd_logarithmic", unavailable) + field, diagnostics = parametric.generate_calibrated_field( + (32, 32), _parameters(0.3), 1, 20, 0, seed=4, return_diagnostics=True + ) + assert np.isfinite(field).all() and diagnostics["status"] == "unavailable_spectrum" + + +def test_higher_fraction_uses_one_candidate(monkeypatch): + monkeypatch.setattr(parametric, "_estimate_beta", lambda *_: (-1, -2)) + _, diagnostics = parametric.generate_calibrated_field( + (32, 32), + _parameters(0.1001), + 1, + 20, + 0, + seed=4, + distance_threshold=0.01, + return_diagnostics=True, + ) + assert diagnostics["attempts"] == 1 + + +def test_parameters_seed_stream_and_calibration_controls(): + parameters = _parameters(0.3) + original = copy.deepcopy(parameters) + global_state = np.random.get_state() + first, diagnostics = parametric.generate_calibrated_field( + (64, 96), parameters, 1, 20, 0, seed=0, return_diagnostics=True + ) + assert_array_equal( + first, + parametric.generate_calibrated_field((64, 96), parameters, 1, 20, 0, seed=0), + ) + after = np.random.get_state() + assert global_state[0] == after[0] and global_state[2:] == after[2:] + assert_array_equal(global_state[1], after[1]) + assert parameters == original + assert diagnostics["calibrated_parameters"]["mean_db"] == parameters["mean_db"] + assert diagnostics["calibrated_parameters"]["stdev_db"] == parameters["stdev_db"] + assert diagnostics["calibrated_parameters"]["beta_1"] == parameters["beta_1"] - 0.5 + rng = np.random.default_rng(7) + a = parametric.generate_stochastic_field((64, 96), parameters, 1, 20, 0, rng=rng) + b = parametric.generate_stochastic_field((64, 96), parameters, 1, 20, 0, rng=rng) + assert not np.array_equal(a, b) + custom, custom_diagnostics = parametric.generate_calibrated_field( + (64, 96), + parameters, + 1, + 20, + 0, + seed=0, + spectral_bias=-0.2, + return_diagnostics=True, + ) + assert ( + custom_diagnostics["calibrated_parameters"]["beta_1"] + == parameters["beta_1"] - 0.2 + ) + assert not np.array_equal(first, custom) + + +def test_dry_target_does_not_consume_random_draw(): + rng = np.random.default_rng(7) + before = copy.deepcopy(rng.bit_generator.state) + field, diagnostics = parametric.generate_calibrated_field( + (32, 32), + _parameters(0), + 1, + 20, + 0, + rng=rng, + zerovalue=-1, + return_diagnostics=True, + ) + assert_array_equal(field, -1) + assert diagnostics["status"] == "dry" and diagnostics["attempts"] == 0 + assert not diagnostics["tolerance_achieved"] and np.isnan(diagnostics["distance"]) + assert rng.bit_generator.state == before + + +def test_post_ar_dry_policy_is_explicit_and_invalid_data_is_not_dry(): + field = np.random.default_rng(2).normal(size=(64, 64)) + suppressed = match_rainfall_distribution( + field, 0.01, 6, 2, 0, min_rain_fraction=0.025 + ) + assert_array_equal(suppressed, -0.1) + generated = match_rainfall_distribution(field, 0.01, 6, 2, 0) + assert np.any(generated > 0) + with pytest.raises(ValueError): + match_rainfall_distribution(np.full((8, 8), np.nan), 0, 6, 2, 0) + with pytest.raises(ValueError): + match_rainfall_distribution(np.ones((8, 8)), 0.1, 6, 2, 0) + + +@pytest.mark.parametrize( + "changes", + [ + {"max_attempts": 0}, + {"max_attempts": 1.5}, + {"distance_threshold": -1}, + {"distance_threshold": np.inf}, + {"low_war_threshold": 2}, + {"beta_variances": [0, 1]}, + {"beta_variances": [1]}, + {"beta_variances": np.array([1 + 1j, 1])}, + {"nbins": 1}, + {"eps": 0.5}, + {"zerovalue": 0}, + {"seed": 0, "rng": np.random.default_rng(0)}, + ], +) +def test_invalid_settings_raise(changes): + with pytest.raises(ValueError): + parametric.generate_calibrated_field( + (32, 32), _parameters(), 1, 20, 0, **changes + ) + + +@pytest.mark.parametrize( + "changes", + [ + {"rain_fraction": -1}, + {"rain_fraction": 1.1}, + {"mean_db": np.nan}, + {"stdev_db": 0}, + {"beta_1": np.inf}, + {"beta_2": None}, + ], +) +def test_invalid_parameters_raise(changes): + with pytest.raises(ValueError): + parametric.generate_calibrated_field( + (32, 32), dict(_parameters(), **changes), 1, 20, 0 + ) + + +def test_degenerate_geometry_and_too_sparse_mapping_are_explicit(): + with pytest.raises(ValueError): + parametric.power_law_filter((16, 16), -2, -3, 1, None) + with pytest.raises(ValueError): + parametric.generate_calibrated_field((2, 8), _parameters(), 1, 20, 0) + with pytest.raises(ValueError): + parametric.generate_stochastic_field( + (16, 16), _parameters(1e-12), 1, 20, 0, seed=0 + ) diff --git a/pysteps/tests/test_nowcasts_parametric.py b/pysteps/tests/test_nowcasts_parametric.py new file mode 100644 index 00000000..3dae377d --- /dev/null +++ b/pysteps/tests/test_nowcasts_parametric.py @@ -0,0 +1,801 @@ +"""Forecast contracts and source numerical parity, without operational services.""" + +from copy import deepcopy +import json +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr +from scipy.ndimage import gaussian_filter + +from pysteps import nowcasts +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.extrapolation.semilagrangian import extrapolate +from pysteps.nowcasts import parametric +from pysteps.nowcasts._cascade_advection import ( + advect_field, + prepare_advection_coordinates, +) +from pysteps.timeseries.autoregression import estimate_ar_params_yw + + +def _dataset(rain, velocity=None, timestep=5.0, kmperpixel=1.0): + """Attach physical coordinates without changing the saved reference arrays.""" + dataset = xr.Dataset( + { + "precip_intensity": ( + ("time", "y", "x"), + rain.copy(), + { + "units": "mm/h", + "transform": None, + "threshold": 0.1, + "zerovalue": 0.0, + "accutime": timestep, + }, + ), + }, + coords={ + "time": ( + "time", + np.datetime64("2026-01-01T00:00:00", "ns") + + np.arange(rain.shape[0]) + * np.timedelta64(round(timestep * 60e9), "ns"), + {"stepsize": timestep * 60}, + ), + "x": ( + "x", + (np.arange(rain.shape[2]) + 0.5) * kmperpixel * 1000, + {"units": "m", "stepsize": kmperpixel * 1000}, + ), + "y": ( + "y", + (np.arange(rain.shape[1]) + 0.5) * kmperpixel * 1000, + {"units": "m", "stepsize": kmperpixel * 1000}, + ), + }, + attrs={ + "precip_var": "precip_intensity", + "projection": "+proj=stere +lat_0=90 +lon_0=0 +lat_ts=60 +datum=WGS84 +units=m", + "institution": "reference fixture", + }, + ) + if velocity is not None: + dataset["velocity_x"] = (("y", "x"), velocity[0].copy()) + dataset["velocity_y"] = (("y", "x"), velocity[1].copy()) + return dataset + + +def _forecast_array(rain, velocity=None, timesteps=1, **kwargs): + """Run original numerical assertions through the public dataset interface.""" + if not isinstance(rain, np.ndarray): + return parametric.forecast(rain, timesteps, **kwargs) + dataset = _dataset( + rain, + velocity, + timestep=kwargs.get("timestep", 5.0), + kmperpixel=kwargs.get("kmperpixel", 1.0), + ) + result = parametric.forecast(dataset, timesteps, **kwargs) + if isinstance(result, tuple): + output, diagnostics = result + return ( + output["precip_intensity"].values if output is not None else None, + diagnostics, + ) + return result["precip_intensity"].values if result is not None else None + + +@pytest.fixture(scope="module") +def reference(): + folder = Path(__file__).parent / "data" + metadata = json.loads((folder / "parametric_nowcast_reference.json").read_text()) + with np.load(folder / "parametric_nowcast_reference.npz") as arrays: + yield metadata, {k: arrays[k].copy() for k in arrays.files} + + +def _call(reference, index=0, **kwargs): + metadata, arrays = reference + case = metadata["cases"][index] + defaults = dict( + n_ens_members=2, + kmperpixel=1, + timestep=case["config"]["timestep"] / 60, + seed=case["seed"], + ) + defaults.update(kwargs) + return _forecast_array( + arrays[f"history_{index}"], arrays[f"velocity_{index}"], 3, **defaults + ) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_forecast_source_preparation_and_coupled_output(reference, index): + metadata, arrays = reference + case = metadata["cases"][index] + result, diagnostics = _call(reference, index, return_diagnostics=True) + assert result.shape == (2, 3, 128, 96) + assert np.isfinite(result).all() and np.all(result >= 0) + assert np.all((result == 0) | (result >= 1)) + for name in ("beta_1", "beta_2", "mean_db", "stdev_db", "rain_fraction"): + assert diagnostics["parameters"][name] == pytest.approx( + case["parameters"][name], abs=1e-6 + ) + assert diagnostics["parameters"]["reference_time"] == pytest.approx( + case["parameters"]["corl_zero"], abs=1e-7 + ) + np.testing.assert_allclose( + diagnostics["model_correlations"][:, 0], case["parameters"]["lag_1"], atol=1e-12 + ) + np.testing.assert_allclose( + diagnostics["model_correlations"][:, 1], case["parameters"]["lag_2"], atol=1e-12 + ) + assert diagnostics["safeguards"]["moment_source"] == case["moment_source"] + assert diagnostics["safeguards"]["low_rain_fraction_fallback"] == bool(index) + # Equivalent spectral solvers differ in the last digits. Histogram mapping + # and its feedback can amplify that perturbation over multiple leads. + expected = arrays[f"forecast_{index}"] + assert np.max(abs(result - expected)) < 0.05 + assert np.sqrt(np.mean((result - expected) ** 2)) < 0.001 + assert np.count_nonzero((result > 0) != (expected > 0)) <= 2 + for member in range(2): + for lead in range(3): + actual = diagnostics["generation"][member][lead] + original = case["calibration"][member][lead] + for name in ("attempts", "selected_attempt", "tolerance_achieved"): + assert actual[name] == original[name] + + +@pytest.mark.parametrize("index", [0, 1]) +def test_source_evolution_with_identical_parameters(reference, index): + metadata, arrays = reference + case = metadata["cases"][index] + rain, velocity = arrays[f"history_{index}"], arrays[f"velocity_{index}"] + transformed = rain.copy() + wet = transformed >= 1 + transformed[wet] = 10 * np.log10(transformed[wet]) + transformed[~wet] = -0.1 + bp = filter_gaussian(rain.shape[1:], 6) + observed = [parametric._decompose(f, bp) for f in transformed] + coordinates = prepare_advection_coordinates(velocity) + taper = gaussian_filter(advect_field(np.ones(rain.shape[1:]), coordinates[1]), 5) + original = case["parameters"] + target = { + k: original[k] + for k in ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + } + phi = np.array( + [ + estimate_ar_params_yw(pair) + for pair in zip(original["lag_1"], original["lag_2"]) + ] + ) + actual = np.empty_like(arrays[f"forecast_{index}"]) + for member, child in enumerate(np.random.SeedSequence(case["seed"]).spawn(2)): + rng = np.random.default_rng(child) + histories = [observed[-1], observed[-2]] + for lead in range(3): + field, _ = parametric._advance( + histories, + target, + phi, + bp, + coordinates, + rng, + 0, + -0.1, + 1, + 20, + parametric._noise_settings(None), + 0.025, + ) + physical = 10 ** (field / 10) + physical[physical < 1] = 0 + actual[member, lead] = physical + histories = [parametric._decompose(field * taper, bp), histories[0]] + np.testing.assert_allclose( + actual, arrays[f"forecast_{index}"], atol=1e-8, rtol=1e-10 + ) + + +@pytest.mark.parametrize("kind", ["zero", "translation", "varying"]) +def test_cached_advection_exactly_matches_native(kind): + rng = np.random.default_rng(45) + shape = (47, 64) + field = rng.normal(size=shape) + velocity = np.zeros((2, *shape)) + if kind == "translation": + velocity[0], velocity[1] = 1.3, -0.4 + elif kind == "varying": + yy, xx = np.indices(shape) + velocity[0] = 0.1 + xx / 30 + velocity[1] = -0.2 + yy / 40 + coordinates = prepare_advection_coordinates(velocity) + for lag in (1, 2): + native = extrapolate(field, velocity, lag, outval=0)[lag - 1] + np.testing.assert_array_equal(advect_field(field, coordinates[lag]), native) + assert not coordinates[lag].flags.writeable + + +def test_streams_member_isolation_and_input_purity(reference): + metadata, arrays = reference + original = {k: v.copy() for k, v in arrays.items()} + seed = np.random.SeedSequence(72) + np.random.seed(23) + global_state = deepcopy(np.random.get_state()) + result = _call(reference, seed=seed) + again = _call(reference, seed=seed) + np.testing.assert_array_equal(result, again) + assert seed.n_children_spawned == 0 + current_state = np.random.get_state() + assert global_state[0] == current_state[0] + np.testing.assert_array_equal(global_state[1], current_state[1]) + assert global_state[2:] == current_state[2:] + prefix = _call(reference, n_ens_members=1, seed=seed) + np.testing.assert_array_equal(prefix[0], result[0]) + larger = _call(reference, n_ens_members=3, seed=seed) + np.testing.assert_array_equal(larger[:2], result) + assert not np.array_equal(result[0], result[1]) + assert not np.array_equal(result, _call(reference, seed=73)) + for k in arrays: + np.testing.assert_array_equal(arrays[k], original[k]) + + +def test_callback_streaming_matches_output_and_cannot_mutate_state(reference): + expected = _call(reference) + leads = [] + + def callback(fields): + leads.append(fields.copy()) + fields.fill(0) + + result = _call(reference, callback=callback) + np.testing.assert_array_equal(result, expected) + np.testing.assert_array_equal(np.stack(leads, axis=1), expected) + leads.clear() + result, diagnostics = _call( + reference, callback=callback, return_output=False, return_diagnostics=True + ) + assert result is None + np.testing.assert_array_equal(np.stack(leads, axis=1), expected) + assert len(diagnostics["generation"][0]) == 3 + + +def test_callback_failure_propagates(reference): + def callback(fields): + raise RuntimeError("writer failed") + + with pytest.raises(RuntimeError, match="writer failed"): + _call(reference, callback=callback) + + +def test_internal_motion_uses_latest_pair_and_reference_configuration( + reference, monkeypatch +): + _, arrays = reference + history = arrays["history_0"] + calls = [] + + def motion(images, **kwargs): + calls.append((images.copy(deep=True), kwargs)) + images["velocity_x"] = (("y", "x"), arrays["velocity_0"][0].copy()) + images["velocity_y"] = (("y", "x"), arrays["velocity_0"][1].copy()) + return images + + monkeypatch.setattr(parametric.motion, "get_method", lambda name: motion) + overrides = dict(decl_scale=12) + output, diagnostics = _forecast_array( + history, + timesteps=2, + n_ens_members=1, + kmperpixel=1, + timestep=5, + seed=4, + motion_kwargs=overrides, + return_diagnostics=True, + ) + assert output.shape == (1, 2, *history.shape[1:]) + assert len(calls) == 1 and calls[0][0].sizes["time"] == 2 + expected = history[-2:].copy() + wet = expected >= 1 + expected[wet] = 10 * np.log10(expected[wet]) + expected[~wet] = -0.1 + np.testing.assert_array_equal(calls[0][0]["precip_intensity"].values, expected) + assert calls[0][0]["precip_intensity"].attrs["transform"] == "dB" + assert calls[0][1]["decl_scale"] == 12 + assert overrides == {"decl_scale": 12} + assert diagnostics["motion_source"] == "lucaskanade" + + +def test_bad_estimated_motion_rejected(reference, monkeypatch): + _, arrays = reference + monkeypatch.setattr( + parametric.motion, + "get_method", + lambda name: lambda *args, **kw: np.full((2, 128, 96), np.nan), + ) + with pytest.raises(ValueError, match="invalid velocity"): + _forecast_array(arrays["history_0"], kmperpixel=1, timestep=5) + + +@pytest.mark.parametrize("prior_wet", [False, True]) +def test_dry_latest_history_skips_estimation_and_motion(prior_wet, monkeypatch): + rain = np.zeros((3, 64, 64)) + if prior_wet: + rain[0] = 10 + + def unexpected(*args, **kwargs): + raise AssertionError("Dry history should not require motion or noise") + + monkeypatch.setattr(parametric.motion, "get_method", unexpected) + monkeypatch.setattr(parametric, "generate_calibrated_field", unexpected) + leads = [] + output, diagnostics = _forecast_array( + rain, + timesteps=2, + kmperpixel=1, + timestep=7.5, + n_ens_members=3, + callback=lambda f: leads.append(f.copy()), + return_diagnostics=True, + ) + assert output.shape == (3, 2, 64, 64) and not output.any() + assert len(leads) == 2 and not np.stack(leads).any() + assert diagnostics["status"] == "dry_latest_observation" + assert diagnostics["parameters"] is None + + +def test_partly_dry_history_requires_explicit_temporal_prior(reference): + _, arrays = reference + history = arrays["history_0"].copy() + history[:2] = 0 + kwargs = dict(kmperpixel=1, timestep=5, n_ens_members=1, seed=4) + with pytest.raises(ValueError, match="Unavailable spectral/persistence"): + _forecast_array(history, arrays["velocity_0"], 2, **kwargs) + result, diagnostics = _forecast_array( + history, + arrays["velocity_0"], + 2, + reference_time=200, + return_diagnostics=True, + **kwargs, + ) + assert np.isfinite(result).all() + assert diagnostics["parameters"]["reference_time"] == 200 + assert np.isnan(diagnostics["estimated_reference_time_minutes"]) + assert all( + len(levels) == 6 for levels in diagnostics["constant_observed_levels"][:2] + ) + + +def test_low_fraction_post_ar_suppression_is_explicit(reference): + _, arrays = reference + history = arrays["history_1"].copy() + wet_indices = np.flatnonzero(history[-1] >= 1) + history[-1].flat[wet_indices[200:]] = 0 + result, diagnostics = _forecast_array( + history, + arrays["velocity_1"], + 2, + kmperpixel=1, + timestep=7.5, + n_ens_members=1, + seed=42, + return_diagnostics=True, + ) + assert not result.any() + assert 0 < diagnostics["statistics"]["rain_fraction"] < 0.025 + assert ( + diagnostics["status"] == "ok" + and diagnostics["safeguards"]["low_rain_fraction_fallback"] + ) + + +@pytest.mark.parametrize("kind", ["constant", "single_pixel"]) +def test_unestimable_wet_input_is_not_dry(kind): + rain = np.zeros((3, 64, 64)) + if kind == "constant": + rain[:] = 5 + else: + rain[-1, 32, 32] = 5 + with pytest.raises(ValueError, match="at least two distinct"): + _forecast_array( + rain, np.zeros((2, 64, 64)), kmperpixel=1, timestep=5, n_ens_members=1 + ) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"timesteps": [1, 2]}, + {"timesteps": 0}, + {"n_ens_members": True}, + {"n_ens_members": 0}, + {"timestep": 0}, + {"kmperpixel": 0}, + {"rain_threshold": 0}, + {"scale_break": np.nan}, + {"zerovalue": 0}, + {"taper_sigma": -1}, + {"min_rain_fraction": 1.1}, + {"seed": -1}, + {"callback": 4}, + {"return_output": "yes"}, + {"n_cascade_levels": 2}, + {"alpha": 0}, + {"lower_truncation": -1}, + {"reference_levels": (6,)}, + {"reference_levels": (0, 0)}, + {"reference_time": 0}, + {"noise_kwargs": {"seed": 1}}, + {"noise_kwargs": {"eps": 0}}, + {"noise_kwargs": {"beta_variances": (0, 1)}}, + {"motion_kwargs": 3}, + ], +) +def test_invalid_options_even_when_dry(kwargs): + options = dict(kmperpixel=1, timestep=5, n_ens_members=1) + options.update(kwargs) + with pytest.raises(ValueError): + _forecast_array(np.zeros((3, 64, 64)), **options) + + +@pytest.mark.parametrize( + "kind", ["shape", "integer", "nan", "negative", "mask", "complex", "list"] +) +def test_invalid_observations(kind): + rain = np.zeros((3, 64, 64)) + if kind == "shape": + rain = rain[:2] + elif kind == "integer": + rain = rain.astype(int) + elif kind == "nan": + rain[0, 0, 0] = np.nan + elif kind == "negative": + rain[0, 0, 0] = -1 + elif kind == "mask": + rain = np.ma.array(rain, mask=False) + rain.mask[0, 0, 0] = True + elif kind == "complex": + rain = rain.astype(complex) + elif kind == "list": + rain = rain.tolist() + with pytest.raises(ValueError): + _forecast_array(rain, kmperpixel=1, timestep=5) + + +def test_method_registration(): + assert nowcasts.get_method("parametric") is parametric.forecast + + +def test_explicit_estimation_settings_are_retained(): + options = dict( + lower_truncation=1.5, + temporal_exponent=0.9, + alpha=2.5, + reference_levels=(0, 1, 2), + reference_time=250, + moment_distance_limit=0.6, + beta_1_bounds=(-3.8, -1.3), + beta_2_min=-4.2, + minimum_beta_gap=0.2, + correlation_time_bounds=(50, 800), + low_rain_fraction=0.04, + fallback_beta_1=-1.5, + fallback_beta_2=-3.2, + fallback_reference_time=320, + ) + original = deepcopy(options) + _, diagnostics = _forecast_array( + np.zeros((3, 8, 8)), + n_ens_members=1, + return_diagnostics=True, + **options, + ) + assert diagnostics["settings"]["estimation"] == options + assert options == original + + +def test_reference_time_conflicts_with_temporal_fallback(reference): + with pytest.raises(ValueError, match="cannot be combined with reference_time"): + _call( + reference, + reference_time=200, + temporal_fallback={"lzero": 256, "tzero": 420}, + ) + + +def test_safeguards_moment_selection_and_joint_bounds(): + statistics = dict( + ordinary_mean=6, + ordinary_stdev=3, + latent_mean=20, + latent_stdev=3, + rain_fraction=0.3, + ) + spectrum = dict(beta_1=-1, beta_2=-1) + _, diagnostics = _forecast_array( + np.zeros((3, 8, 8)), n_ens_members=1, return_diagnostics=True + ) + settings = diagnostics["settings"]["estimation"] + parameters, audit = parametric._generation_parameters( + statistics, spectrum, 20, settings + ) + assert audit["moment_source"] == "ordinary" and audit["moment_reason"] == "distance" + assert parameters["reference_time"] == 60 + assert parameters["beta_1"] < -1.25 + assert parameters["beta_2"] <= parameters["beta_1"] - 0.1 + assert spectrum == dict(beta_1=-1, beta_2=-1) + statistics["latent_mean"] = np.nan + _, audit = parametric._generation_parameters(statistics, spectrum, 200, settings) + assert audit["moment_reason"] == "invalid_latent" + + +@pytest.mark.parametrize("invalid_pair", [[np.nan, np.nan], [1.0, 1.0], [0.2, -0.99]]) +def test_configured_temporal_fallback(reference, monkeypatch, invalid_pair): + def correlations(cascades, velocity): + return np.tile(invalid_pair, (cascades.shape[1], 1)) + + monkeypatch.setattr(parametric, "lagrangian_autocorrelation", correlations) + result, diag = _call( + reference, + temporal_fallback={"lzero": 256, "tzero": 420}, + timestep=7.5, + return_diagnostics=True, + ) + expected_time = 420 * (diag["wavelengths_km"][0] / 256) ** 0.81 + assert diag["temporal_selection"]["source"] == "configured_fallback" + assert diag["parameters"]["reference_time"] == pytest.approx(expected_time) + np.testing.assert_allclose( + diag["temporal_model"]["requested_times"], + 420 * (diag["wavelengths_km"] / 256) ** 0.81, + ) + assert not diag["safeguards"]["low_rain_fraction_fallback"] + assert np.isfinite(result).all() + + +def test_configured_fallback_preserves_valid_observation(reference): + result, diag = _call( + reference, + temporal_fallback={"lzero": 256, "tzero": 420}, + return_diagnostics=True, + ) + assert diag["temporal_selection"]["source"] == "observed" + original = _call(reference) + np.testing.assert_array_equal(result, original) + + +def test_temporal_fallback_does_not_replace_spatial_estimates(reference, monkeypatch): + monkeypatch.setattr( + parametric, + "fit_power_law_spectrum", + lambda *args: {"beta_1": np.nan, "beta_2": np.nan}, + ) + with pytest.raises(ValueError, match="Unavailable spectral"): + _call(reference, temporal_fallback={"lzero": 256, "tzero": 420}) + + +@pytest.mark.parametrize( + "fallback", [{}, {"lzero": 0, "tzero": 420}, {"lzero": 256, "tzero": float("nan")}] +) +def test_invalid_temporal_fallback(reference, fallback): + with pytest.raises(ValueError): + _call(reference, temporal_fallback=fallback) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_dataset_coordinates_metadata_and_input_purity(reference, monkeypatch, index): + metadata, arrays = reference + seconds = metadata["cases"][index]["config"]["timestep"] + dataset = _dataset( + arrays[f"history_{index}"], arrays[f"velocity_{index}"], timestep=seconds / 60 + ) + original = dataset.copy(deep=True) + + def unexpected(*args, **kwargs): + raise AssertionError("Supplied velocity must bypass motion estimation") + + monkeypatch.setattr(parametric.motion, "get_method", unexpected) + output, diagnostics = parametric.forecast( + dataset, + 3, + n_ens_members=2, + seed=metadata["cases"][index]["seed"], + return_diagnostics=True, + ) + xr.testing.assert_identical(dataset, original) + field = output[output.attrs["precip_var"]] + assert field.dims == ("ens_number", "time", "y", "x") + assert field.dtype == np.float64 + expected_attrs = {**dataset.precip_intensity.attrs, "threshold": 1.0} + expected_attrs.pop("transform") + assert field.attrs == expected_attrs + assert output.attrs == dataset.attrs + assert diagnostics["motion_source"] == "supplied" + assert diagnostics["settings"]["timestep_minutes"] == seconds / 60 + assert diagnostics["settings"]["kmperpixel"] == 1.0 + np.testing.assert_array_equal(output.ens_number, [1, 2]) + np.testing.assert_array_equal( + output.time.values, + dataset.time.values[-1] + np.arange(1, 4) * np.timedelta64(int(seconds), "s"), + ) + assert output.forecast_reference_time.values == dataset.time.values[-1] + xr.testing.assert_identical(output.x.variable, dataset.x.variable) + xr.testing.assert_identical(output.y.variable, dataset.y.variable) + assert output.time.attrs["stepsize"] == seconds + assert "velocity_x" not in output and "velocity_y" not in output + # Returned coordinates must not expose mutable aliases into the input. + output.x.attrs["units"] = "changed on output only" + xr.testing.assert_identical(dataset, original) + + +def test_dataset_dimension_order_and_coordinate_units(reference): + _, arrays = reference + dataset = _dataset(arrays["history_0"], arrays["velocity_0"]) + expected = parametric.forecast(dataset, 1, n_ens_members=1, seed=4) + for axis in ("x", "y"): + dataset[axis] = dataset[axis] / 1000 + dataset[axis].attrs.update(units="km", stepsize=1.0) + dataset = dataset.transpose("x", "time", "y") + actual = parametric.forecast(dataset, 1, n_ens_members=1, seed=4) + np.testing.assert_array_equal(actual.precip_intensity, expected.precip_intensity) + xr.testing.assert_identical(actual.x.variable, dataset.x.variable) + xr.testing.assert_identical(actual.y.variable, dataset.y.variable) + + +def test_descending_y_and_dry_forecast_metadata(): + dataset = _dataset(np.zeros((3, 8, 9)), timestep=7.5) + dataset = dataset.isel(y=slice(None, None, -1)) + dataset.y.attrs["stepsize"] = -1000.0 + output = parametric.forecast(dataset, 2, n_ens_members=1) + xr.testing.assert_identical(output.y.variable, dataset.y.variable) + assert not output.precip_intensity.values.any() + assert "transform" not in output.precip_intensity.attrs + np.testing.assert_array_equal( + np.diff(output.time.values) / np.timedelta64(1, "s"), [450] + ) + + +@pytest.mark.parametrize("overrides", [{"kmperpixel": 2}, {"timestep": 5}]) +def test_explicit_geometry_must_agree_with_dataset(overrides): + dataset = _dataset(np.zeros((3, 8, 9)), timestep=7.5) + with pytest.raises(ValueError, match="disagrees with dataset coordinates"): + parametric.forecast(dataset, **overrides) + + +@pytest.mark.parametrize( + "kind, message", + [ + ("array", "xarray.Dataset"), + ("precip_var", "precip_var"), + ("units", "untransformed rain rates"), + ("transform", "untransformed rain rates"), + ("integer", "floating point"), + ("missing", "finite nonnegative"), + ("negative", "finite nonnegative"), + ("coordinate_units", "coordinate units"), + ("irregular_x", "regularly spaced"), + ("square", "square pixels"), + ("spatial_stepsize", "stepsize disagrees"), + ("missing_coordinate", "coordinate is required"), + ("time_stepsize", "spacing in seconds"), + ("irregular_time", "increasing and regularly spaced"), + ("backwards_time", "increasing and regularly spaced"), + ("nat", "valid datetime64"), + ("numeric_time", "valid datetime64"), + ("partial_velocity", "both velocity_x and velocity_y"), + ("missing_velocity", "finite real"), + ("velocity_units", "pixels per observation timestep"), + ("velocity_time_dimension", r"dimensions \(y, x\)"), + ], +) +def test_dataset_rejects_invalid_physical_inputs_before_dry_shortcut(kind, message): + dataset = _dataset(np.zeros((3, 8, 9)), np.zeros((2, 8, 9))) + if kind == "array": + dataset = dataset.precip_intensity.values + elif kind == "precip_var": + dataset.attrs["precip_var"] = "absent" + elif kind in ("units", "transform"): + dataset.precip_intensity.attrs[kind] = "dB" if kind == "transform" else "mm" + elif kind == "integer": + dataset["precip_intensity"] = dataset.precip_intensity.astype(int) + elif kind in ("missing", "negative"): + dataset.precip_intensity.values[0, 0, 0] = np.nan if kind == "missing" else -1 + elif kind == "coordinate_units": + dataset.x.attrs["units"] = "degrees_east" + elif kind == "irregular_x": + values = dataset.x.values.copy() + values[2] += 100 + dataset = dataset.assign_coords(x=("x", values, dataset.x.attrs)) + elif kind == "square": + dataset["y"] = dataset.y * 2 + dataset.y.attrs.update(units="m", stepsize=2000) + elif kind == "spatial_stepsize": + dataset.y.attrs["stepsize"] = -1000 + elif kind == "missing_coordinate": + dataset = dataset.drop_vars("x") + elif kind == "time_stepsize": + dataset.time.attrs["stepsize"] = 5 # Minutes must not be read as seconds. + elif kind == "irregular_time": + values = dataset.time.values.copy() + values[1] += np.timedelta64(1, "s") + dataset = dataset.assign_coords(time=("time", values, dataset.time.attrs)) + elif kind == "backwards_time": + dataset = dataset.isel(time=slice(None, None, -1)) + elif kind == "nat": + values = dataset.time.values.copy() + values[0] = np.datetime64("NaT", "ns") + dataset = dataset.assign_coords(time=("time", values, dataset.time.attrs)) + elif kind == "numeric_time": + dataset = dataset.assign_coords(time=[0, 300, 600]) + elif kind == "partial_velocity": + dataset = dataset.drop_vars("velocity_y") + elif kind == "missing_velocity": + dataset.velocity_x.values[0, 0] = np.nan + elif kind == "velocity_units": + dataset.velocity_x.attrs["units"] = "m/s" + elif kind == "velocity_time_dimension": + dataset["velocity_x"] = (("time", "y", "x"), np.zeros((3, 8, 9))) + with pytest.raises(ValueError, match=message): + parametric.forecast(dataset, n_ens_members=1) + + +@pytest.mark.parametrize("case", [0, 1, "dry"]) +def test_forecast_native_netcdf_roundtrip(reference, tmp_path, case): + """Exercise actual importer metadata and lossless CF decoding at both cadences.""" + from datetime import datetime + + from pysteps.xarray_helpers import convert_input_to_xarray_dataset + + metadata, arrays = reference + index = 0 if case == "dry" else case + seconds = 450 if case == "dry" else metadata["cases"][index]["config"]["timestep"] + rain = arrays[f"history_{index}"].copy() + if case == "dry": + rain.fill(0) + height, width = rain.shape[-2:] + dataset = convert_input_to_xarray_dataset( + rain, + None, + { + "unit": "mm/h", + "threshold": 0.1, + "zerovalue": 0.0, + "accutime": seconds / 60, + "projection": _dataset(rain).attrs["projection"] + " +x_0=0 +y_0=0", + "institution": "reference fixture", + "x1": 0, + "x2": width * 1000, + "y1": 0, + "y2": height * 1000, + "xpixelsize": 1000, + "ypixelsize": 1000, + "cartesian_unit": "m", + "yorigin": "upper", + }, + startdate=datetime(2026, 1, 1), + timestep=seconds, + ) + dataset["velocity_x"] = (("y", "x"), arrays[f"velocity_{index}"][0]) + dataset["velocity_y"] = (("y", "x"), arrays[f"velocity_{index}"][1]) + original = dataset.copy(deep=True) + output = parametric.forecast(dataset, 3, n_ens_members=2, seed=4) + precip_var = output.attrs["precip_var"] + output_path = tmp_path / "forecast.nc" + output.to_netcdf( + output_path, + engine="netcdf4", + encoding={precip_var: {"zlib": True, "complevel": 1}}, + ) + with xr.open_dataset(output_path, engine="netcdf4") as reopened: + loaded = reopened.load() + xr.testing.assert_identical(output, loaded) + xr.testing.assert_identical(dataset, original) + np.testing.assert_array_equal(output[precip_var], loaded[precip_var]) + assert loaded[precip_var].dtype == np.float64 + assert loaded.time.attrs["stepsize"] == seconds + assert "transform" not in loaded[precip_var].attrs + assert "lat" in loaded.coords and "lon" in loaded.coords + assert loaded[precip_var].attrs["grid_mapping"] in loaded.coords + assert "velocity_x" not in loaded and "velocity_y" not in loaded diff --git a/pysteps/tests/test_parametric_estimators.py b/pysteps/tests/test_parametric_estimators.py new file mode 100644 index 00000000..30c3205b --- /dev/null +++ b/pysteps/tests/test_parametric_estimators.py @@ -0,0 +1,327 @@ +"""Independent numerical checks and frozen source-implementation comparisons.""" + +import json +from pathlib import Path + +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal +from scipy.stats import truncnorm + +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.timeseries.parametric import ( + ar2_correlation_time, + estimate_reference_correlation_time, + fit_ar2_correlations, + lagrangian_autocorrelation, + model_cascade_correlations, +) +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import ( + compute_field_statistics, + fit_lower_truncated_normal, +) + + +def _directional_field(shape): + y, x = np.indices(shape) + return np.random.default_rng(23).normal(size=shape) + 4 * np.cos( + 2 * np.pi * (5 * y / shape[0] + 3 * x / shape[1]) + ) + + +def _marginal_field(): + rng = np.random.default_rng(37) + field = rng.normal(4, 3, size=(128, 96)) + field[rng.uniform(size=field.shape) > 0.3] = -0.1 + return field + + +def _lag_history(): + rng = np.random.default_rng(45) + history = rng.normal(size=(3, 2, 24, 32)) + velocity = np.full((2, 24, 32), 0.25) + velocity[0] = 0.6 + return history, velocity + + +@pytest.fixture(scope="module") +def reference(): + return json.loads( + ( + Path(__file__).parent / "data/parametric_estimators_reference.json" + ).read_text() + ) + + +def test_source_marginal_parity(reference): + field = _marginal_field() + original = field.copy() + estimated = compute_field_statistics(field, 0, lower_truncation=1.25) + for name, expected in reference["marginal"].items(): + # The likelihood optimizer's latent moments vary slightly across + # numerical libraries and platforms; direct moments remain tighter. + if name in ("latent_mean", "latent_stdev"): + rel_tol, abs_tol = 1e-7, 1e-8 + else: + rel_tol, abs_tol = 1e-12, 1e-12 + assert estimated[name] == pytest.approx(expected, rel=rel_tol, abs=abs_tol) + assert_array_equal(field, original) + + +def test_source_spectrum_parity(reference): + for case in reference["spectra"]: + frequency, power = rapsd_logarithmic(_directional_field(tuple(case["shape"]))) + assert_allclose(10 * np.log10(frequency), case["frequency_db"], atol=1e-12) + assert_allclose(10 * np.log10(power), case["power_db"], atol=1e-10) + for expected in case["models"]: + estimated = fit_power_law_spectrum( + frequency, power, expected["scale_break"] + ) + assert estimated["status"] == "ok" + for name in ("beta_1", "beta_2", "c1", "c2"): + if expected[name] is None: + assert estimated[name] is None + else: + assert estimated[name] == pytest.approx(expected[name], abs=3e-5) + + +def test_source_temporal_and_paper_reference_parity(reference): + history, velocity = _lag_history() + assert_allclose( + lagrangian_autocorrelation(history, velocity), reference["lags"], atol=1e-14 + ) + for case in reference["ar2_inverse"]: + pair = fit_ar2_correlations(case["time"], case["alpha"], 5) + assert_allclose(pair, case["pair"], atol=1e-12) + paper = reference["paper_reference"] + assert_allclose( + estimate_reference_correlation_time( + paper["times"], paper["wavelengths"], paper["exponent"] + ), + paper["reference_times"], + rtol=1e-14, + ) + + +@pytest.mark.parametrize("shape", [(128, 128), (127, 127), (96, 128), (95, 127)]) +def test_radial_mean_matches_independent_full_fft(shape): + field = _directional_field(shape) + original = field.copy() + normalized = (field - field.mean()) / field.std() + full_power = abs(np.fft.fft2(normalized)) ** 2 + y, x = np.meshgrid( + np.fft.fftfreq(shape[0]), np.fft.fftfreq(shape[1]), indexing="ij" + ) + radius = np.hypot(y, x) + n = min(shape) // 2 + limit = np.hypot( + abs(np.fft.fftfreq(shape[0])[: shape[0] // 2]).max(), + abs(np.fft.fftfreq(shape[1])[: shape[1] // 2]).max(), + ) + # Preserve the historical endpoint rounding: exactly-boundary Fourier modes + # can move bins if geomspace is substituted for the source logspace edges. + edges = np.logspace(np.log10(1 / n), np.log10(limit), n) + expected_frequency, expected_power = [], [] + for lo, hi in zip(edges[:-1], edges[1:]): + mask = (radius >= lo) & (radius < hi) + if mask.any(): + expected_frequency.append((lo + hi) / 2) + expected_power.append(full_power[mask].mean()) + frequency, power = rapsd_logarithmic(field) + assert_allclose(frequency, expected_frequency, atol=1e-14) + assert_allclose(power, expected_power, rtol=1e-12) + for changed in (field.T, field[:, ::-1], np.rot90(field)): + actual_frequency, actual_power = rapsd_logarithmic(changed) + assert_allclose(actual_frequency, frequency, atol=1e-14) + assert_allclose(actual_power, power, rtol=1e-12) + assert_array_equal(field, original) + + +def test_exact_continuous_power_law_recovery(): + frequency = np.geomspace(0.005, 0.5, 120) + breakpoint = 1 / 20 + power = np.where( + frequency <= breakpoint, frequency**-2, breakpoint**1.5 * frequency**-3.5 + ) + result = fit_power_law_spectrum(frequency, power, scale_break=20) + assert result["status"] == "ok" + assert_allclose([result["beta_1"], result["beta_2"]], [-2, -3.5], atol=1e-7) + assert result["rmse_db"] < 1e-7 + assert result["c1"] + result["beta_1"] * 10 * np.log10(breakpoint) == pytest.approx( + result["c2"] + result["beta_2"] * 10 * np.log10(breakpoint) + ) + + +def test_unavailable_spectral_fits_are_explicit(): + frequency, power = rapsd_logarithmic(np.full((32, 32), -0.1)) + assert frequency.size == power.size == 0 + assert ( + fit_power_law_spectrum(frequency, power, 20)["status"] == "insufficient_samples" + ) + frequency = np.geomspace(0.02, 0.2, 30) + assert ( + fit_power_law_spectrum(frequency, np.ones(30), 1000)["status"] + == "unidentified_model" + ) + assert ( + fit_power_law_spectrum(frequency, np.zeros(30))["status"] + == "insufficient_samples" + ) + + +def test_latent_normal_recovery_and_occurrence_cutoff(): + mean, scale, cutoff = 2.0, 3.0, 1.25 + probabilities = (np.arange(20000) + 0.5) / 20000 + sample = truncnorm.ppf( + probabilities, (cutoff - mean) / scale, np.inf, loc=mean, scale=scale + ) + fitted_mean, fitted_scale, diagnostics = fit_lower_truncated_normal( + sample, cutoff, True + ) + assert_allclose([fitted_mean, fitted_scale], [mean, scale], atol=0.003) + assert diagnostics["status"] == "ok" + field = np.array([[-0.1, 0, 0.5, 2], [1, 3, 4, 5]]) + a = compute_field_statistics(field, 0, lower_truncation=1.25) + b = compute_field_statistics(field, 0, lower_truncation=2) + assert a["rain_fraction"] == b["rain_fraction"] == 7 / 8 + assert a["ordinary_mean"] == b["ordinary_mean"] + assert a["nonzero_fit"]["n_samples"] == 4 + + +def test_dry_constant_and_bounded_marginal_fits(): + dry = compute_field_statistics(np.full((4, 4), -0.1), 0, lower_truncation=1.25) + assert dry["rain_fraction"] == 0 + assert np.isnan(dry["ordinary_mean"]) and np.isnan(dry["latent_mean"]) + assert dry["nonzero_fit"]["status"] == "insufficient_samples" + assert fit_lower_truncated_normal([2, 2], 1, True)[2]["status"] == "constant_sample" + _, _, bounded = fit_lower_truncated_normal([0.01, 0.01, 0.01, 100], 0, True) + assert bounded["status"] == "bounded_solution" and bounded["boundary_limited"] + + +def _summed_acf(rho1, rho2, timestep): + phi1, phi2 = np.linalg.solve([[1, rho1], [rho1, 1]], [rho1, rho2]) + previous, current, total = 1.0, rho1, 1.0 + rho1 + for _ in range(100000): + previous, current = current, phi1 * current + phi2 * previous + total += current + if abs(previous) + abs(current) < 1e-13: + return timestep * total + raise AssertionError("ACF recurrence did not converge") + + +@pytest.mark.parametrize( + "phi1,phi2", [(1.2, -0.3), (0.4, 0.3), (-0.6, -0.3), (0.9, 0), (0, 0)] +) +def test_forward_acf_sum(phi1, phi2): + rho1 = phi1 / (1 - phi2) + rho2 = phi1 * rho1 + phi2 + assert ar2_correlation_time(rho1, rho2, 5) == pytest.approx( + _summed_acf(rho1, rho2, 5), rel=1e-10 + ) + assert ar2_correlation_time(rho1, rho2, 300) == pytest.approx( + 60 * ar2_correlation_time(rho1, rho2, 5) + ) + + +@pytest.mark.parametrize("alpha", [0.8, 1.25, 2, 2.72, 4]) +@pytest.mark.parametrize("target", [5, 5.00001, 10, 60, 420, 720]) +def test_inverse_matches_independent_acf_sum(alpha, target): + pair, diagnostics = fit_ar2_correlations(target, alpha, 5, True) + assert _summed_acf(*pair, 5) == pytest.approx(target, rel=1e-7) + assert diagnostics["achieved_time"] == pytest.approx(target, rel=1e-7) + + +def test_reference_selection_and_nominal_scale_model(): + wavelength = 1 / filter_gaussian((256, 256), 6, d=1)["central_freqs"] + exponent = 0.7718612877449645 + observed = 120 * (wavelength / wavelength[0]) ** exponent + assert estimate_reference_correlation_time( + observed, wavelength, exponent + ) == pytest.approx(120) + altered = observed.copy() + altered[-1] = np.nan + assert estimate_reference_correlation_time( + altered, wavelength, exponent + ) == pytest.approx(120) + altered[1] = np.nan + assert np.isnan(estimate_reference_correlation_time(altered, wavelength, exponent)) + assert estimate_reference_correlation_time( + altered, wavelength, exponent, levels=(0,) + ) == pytest.approx(120) + pairs, diagnostics = model_cascade_correlations( + 60, wavelength, exponent, 2.72, 5, return_diagnostics=True + ) + assert_allclose( + diagnostics["achieved_times"], + np.maximum(diagnostics["requested_times"], 5), + rtol=1e-7, + ) + assert_array_equal(pairs[diagnostics["limited_below_timestep"]], 0) + # Changing the reference convention must not change the predicted level times. + rescaled = estimate_reference_correlation_time( + observed, wavelength, exponent, reference_wavelength=128 + ) + alternate = model_cascade_correlations( + rescaled, wavelength, exponent, 2.72, 5, reference_wavelength=128 + ) + original = model_cascade_correlations(120, wavelength, exponent, 2.72, 5) + assert_allclose(alternate, original, atol=1e-12) + flat = model_cascade_correlations(120, wavelength, 0, 2.72, 5) + assert_allclose(ar2_correlation_time(flat[:, 0], flat[:, 1], 5), 120) + + +def test_lagrangian_boundary_conventions_and_input_preservation(): + base = np.random.default_rng(7).normal(size=(32, 40)) + history = np.stack([base, np.roll(base, 1, axis=1), np.roll(base, 2, axis=1)])[ + :, None + ] + velocity = np.zeros((2, 32, 40)) + velocity[0] = 1 + original = history.copy() + assert_allclose( + lagrangian_autocorrelation(history, velocity, boundary="mask"), + [[1, 1]], + atol=1e-14, + ) + zero = lagrangian_autocorrelation(history, velocity) + assert (zero < 1).all() + assert_array_equal(history, original) + assert np.isnan(lagrangian_autocorrelation(np.ones_like(history), velocity)).all() + assert_allclose( + lagrangian_autocorrelation(history, np.zeros_like(velocity)), + [ + [ + np.corrcoef(history[-1, 0].ravel(), history[-2, 0].ravel())[0, 1], + np.corrcoef(history[-1, 0].ravel(), history[-3, 0].ravel())[0, 1], + ] + ], + ) + + +@pytest.mark.parametrize( + "operation", + [ + lambda: compute_field_statistics([[np.nan]], 0), + lambda: compute_field_statistics([[1j]], 0), + lambda: compute_field_statistics([[1]], 1, lower_truncation=0), + lambda: fit_lower_truncated_normal([1, 2], np.nan), + lambda: rapsd_logarithmic(np.zeros((2, 2))), + lambda: rapsd_logarithmic(np.full((8, 8), np.nan)), + lambda: rapsd_logarithmic(np.full((8, 8), 1j)), + lambda: fit_power_law_spectrum([0.2, 0.1], [1, 2]), + lambda: fit_ar2_correlations(4, 2.72, 5), + lambda: fit_ar2_correlations(1e20, 2.72, 5), + lambda: estimate_reference_correlation_time([1, 2], [10, 5], 0.8), + lambda: model_cascade_correlations(60, [10, 5], 0.8, 0, 5), + lambda: lagrangian_autocorrelation(np.zeros((2, 1, 8, 8)), np.zeros((2, 8, 8))), + ], +) +def test_invalid_inputs_raise(operation): + with pytest.raises(ValueError): + operation() + + +def test_invalid_ar_pairs_are_not_repaired(): + assert np.isnan(ar2_correlation_time([np.nan, 1, 0.5], [0, 0.5, 1.5], 5)).all() + assert np.isnan(ar2_correlation_time(0.5, 0.2, 0)) diff --git a/pysteps/tests/test_simulation_parametric.py b/pysteps/tests/test_simulation_parametric.py new file mode 100644 index 00000000..b4e293bb --- /dev/null +++ b/pysteps/tests/test_simulation_parametric.py @@ -0,0 +1,502 @@ +"""Check prescribed rainfall simulation and its shared forecast evolution.""" + +from copy import deepcopy +import json +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr + +from pysteps.cascade.bandpass_filters import filter_gaussian +from pysteps.cascade.decomposition import decomposition_fft +from pysteps.noise.parametric import generate_calibrated_field +from pysteps.nowcasts import parametric +from pysteps.timeseries.parametric import ( + ar2_correlation_time, + estimate_reference_correlation_time, + lagrangian_autocorrelation, +) +from pysteps.utils.spectral import fit_power_law_spectrum, rapsd_logarithmic +from pysteps.utils.statistics import compute_field_statistics +from pysteps.utils.transformation import dB_transform + + +@pytest.fixture +def initial(): + """Build a translating history with known motion and a fractional cadence.""" + rng = np.random.default_rng(17) + rain = rng.lognormal(1.0, 0.8, (32, 40)) + rain[rng.random(rain.shape) > 0.3] = 0 + return xr.Dataset( + { + "rain": ( + ("time", "y", "x"), + np.stack([np.roll(rain, -2, 1), np.roll(rain, -1, 1), rain]), + {"units": "mm/h", "transform": None}, + ), + "velocity_x": ( + ("y", "x"), + np.ones(rain.shape), + {"units": "pixels/timestep"}, + ), + "velocity_y": (("y", "x"), np.zeros(rain.shape)), + }, + coords={ + "time": ( + "time", + np.datetime64("2026-01-01T00:00:00") + + np.arange(3) * np.timedelta64(450, "s"), + {"stepsize": 450}, + ), + "x": ("x", (np.arange(40) + 0.5) * 1000, {"units": "m", "stepsize": 1000}), + "y": ( + "y", + (np.arange(32)[::-1] + 0.5) * 1000, + {"units": "m", "stepsize": -1000}, + ), + }, + attrs={"precip_var": "rain", "projection": "+proj=longlat +datum=WGS84"}, + ) + + +@pytest.fixture +def targets(): + """Return generation targets distinct from the history's sample statistics.""" + return dict(beta_1=-2.4, beta_2=-3.1, rain_fraction=0.2, mean_db=6.0, stdev_db=3.0) + + +def _simulate(initial, targets, **kwargs): + """Supply example settings while allowing individual contract checks.""" + settings = dict( + parameters=targets, + reference_time=400.0, + reference_wavelength=256.0, + timesteps=3, + n_ens_members=2, + seed=42, + return_diagnostics=True, + ) + settings.update(kwargs) + return parametric.simulate(initial, **settings) + + +def test_prescribed_scale_law_and_output_coordinates(initial, targets): + output, diagnostics = _simulate(initial.transpose("x", "time", "y"), targets) + expected = 400 * (diagnostics["wavelengths_km"] / 256) ** 0.81 + np.testing.assert_allclose( + diagnostics["temporal_model"]["requested_times"], expected + ) + assert diagnostics["settings"]["timestep_minutes"] == 7.5 + assert diagnostics["parameters"] == targets + assert diagnostics["parameter_source"] == "supplied" + assert output.rain.dims == ("ens_number", "time", "y", "x") + assert output.rain.shape == (2, 3, 32, 40) + assert np.isfinite(output.rain).all() and (output.rain >= 0).all() + assert np.all((output.rain == 0) | (output.rain >= 1)) + np.testing.assert_array_equal(output.x, initial.x) + np.testing.assert_array_equal(output.y, initial.y) + np.testing.assert_array_equal(output.ens_number, [1, 2]) + np.testing.assert_array_equal( + output.time, + initial.time.values[-1] + np.arange(1, 4) * np.timedelta64(450, "s"), + ) + assert output.forecast_reference_time == initial.time.values[-1] + assert output.attrs["projection"] == initial.attrs["projection"] + assert "transform" not in output.rain.attrs + + # Sum the AR recurrence independently to check the achieved discrete times. + phi = diagnostics["ar_coefficients"] + previous = np.ones(len(phi)) + current = diagnostics["model_correlations"][:, 0].copy() + total = previous + current + for _ in range(10000): + previous, current = current, phi[:, 0] * current + phi[:, 1] * previous + total += current + np.testing.assert_allclose(total * 7.5, np.maximum(expected, 7.5), atol=1e-7) + + +def test_simulation_does_not_estimate_or_replace_parameters( + initial, targets, monkeypatch +): + def forbidden(*args, **kwargs): + raise AssertionError("Simulation called an observation-based estimator") + + for name in [ + "compute_field_statistics", + "fit_power_law_spectrum", + "rapsd_logarithmic", + "lagrangian_autocorrelation", + "_generation_parameters", + ]: + monkeypatch.setattr(parametric, name, forbidden) + monkeypatch.setattr(parametric.motion, "get_method", forbidden) + initial.rain.values.fill(0) + # These valid supplied slopes violate the forecast's regional selection bounds. + targets.update(beta_1=-1.0, beta_2=-0.7, rain_fraction=0.04) + output, diagnostics = _simulate(initial, targets, reference_time=1200) + assert diagnostics["parameters"] == targets + assert diagnostics["settings"]["reference_time_minutes"] == 1200 + assert diagnostics["settings"]["min_rain_fraction"] == 0 + assert (output.rain > 0).any() + + +def test_simulation_matches_forecast_with_identical_parameters(): + folder = Path(__file__).parent / "data" + case = json.loads((folder / "parametric_nowcast_reference.json").read_text())[ + "cases" + ][0] + with np.load(folder / "parametric_nowcast_reference.npz") as data: + rain, velocity = data["history_0"], data["velocity_0"] + history = xr.Dataset( + { + "rain": (("time", "y", "x"), rain, {"units": "mm/h"}), + "velocity_x": (("y", "x"), velocity[0]), + "velocity_y": (("y", "x"), velocity[1]), + }, + coords={ + "time": ( + "time", + np.datetime64("2026-01-01") + np.arange(3) * np.timedelta64(5, "m"), + {"stepsize": 300}, + ), + "x": ("x", np.arange(rain.shape[2]), {"units": "km", "stepsize": 1}), + "y": ("y", np.arange(rain.shape[1]), {"units": "km", "stepsize": 1}), + }, + attrs={"precip_var": "rain"}, + ) + forecast, diagnostics = parametric.forecast( + history, + timesteps=3, + n_ens_members=2, + seed=case["seed"], + return_diagnostics=True, + ) + targets = { + name: diagnostics["parameters"][name] + for name in ("beta_1", "beta_2", "rain_fraction", "mean_db", "stdev_db") + } + simulation, supplied = _simulate( + history, + targets, + seed=case["seed"], + min_rain_fraction=0.025, + reference_time=diagnostics["parameters"]["reference_time"], + reference_wavelength=diagnostics["wavelengths_km"][0], + ) + xr.testing.assert_identical(simulation, forecast) + np.testing.assert_array_equal( + supplied["ar_coefficients"], diagnostics["ar_coefficients"] + ) + + +def test_streaming_purity_and_member_streams(initial, targets): + saved = initial.copy(deep=True) + supplied_targets = targets.copy() + noise_settings = dict(max_attempts=2) + seed = np.random.SeedSequence(42) + seed.spawn(2) + seed_state = deepcopy(seed.state) + global_state = np.random.get_state() + output, diagnostics = _simulate( + initial, targets, seed=seed, noise_kwargs=noise_settings + ) + streamed = [] + + def callback(field): + streamed.append(field.copy()) + field.fill(-1) + + result, streamed_diagnostics = _simulate( + initial, + targets, + seed=seed, + noise_kwargs=noise_settings, + callback=callback, + return_output=False, + ) + assert result is None + np.testing.assert_array_equal(np.stack(streamed, axis=1), output.rain) + assert streamed_diagnostics["seed"] == diagnostics["seed"] + larger, _ = _simulate( + initial, targets, seed=seed, n_ens_members=3, noise_kwargs=noise_settings + ) + np.testing.assert_array_equal(larger.rain[:2], output.rain) + xr.testing.assert_identical(initial, saved) + assert targets == supplied_targets + assert noise_settings == dict(max_attempts=2) + assert seed.state == seed_state + np.testing.assert_equal(np.random.get_state(), global_state) + + +def test_dry_target_validates_input_and_skips_generation(initial, targets, monkeypatch): + targets["rain_fraction"] = 0 + + def forbidden(*args, **kwargs): + raise AssertionError("A dry target consumed stochastic generation") + + monkeypatch.setattr(parametric, "generate_calibrated_field", forbidden) + output, diagnostics = _simulate(initial, targets) + assert (output.rain == 0).all() + assert diagnostics["status"] == "dry_target" + assert diagnostics["generation"] == [[], []] + initial.rain.values[0, 0, 0] = np.nan + with pytest.raises(ValueError, match="finite"): + _simulate(initial, targets) + + +@pytest.mark.parametrize( + "name,value", + [ + ("reference_time", 0), + ("reference_time", np.inf), + ("reference_wavelength", -1), + ("alpha", 0), + ("temporal_exponent", np.nan), + ("timesteps", 0), + ("n_ens_members", 0), + ("callback", 1), + ("return_output", 1), + ], +) +def test_invalid_simulation_settings(initial, targets, name, value): + with pytest.raises(ValueError): + _simulate(initial, targets, **{name: value}) + + +@pytest.mark.parametrize( + "change", + [ + dict(stdev_db=0), + dict(rain_fraction=-0.1), + dict(rain_fraction=1.1), + dict(mean_db=np.nan), + dict(beta_1=1j), + dict(unknown=1), + ], +) +def test_invalid_simulation_targets(initial, targets, change): + targets.update(change) + with pytest.raises(ValueError): + _simulate(initial, targets) + + +def test_simulation_requires_motion_and_complete_parameter_mapping(initial, targets): + with pytest.raises(ValueError, match="requires both"): + _simulate(initial.drop_vars(["velocity_x", "velocity_y"]), targets) + targets.pop("stdev_db") + with pytest.raises(ValueError, match="exactly"): + _simulate(initial, targets) + + +def test_subinterval_times_resolve_to_white_noise(initial, targets): + _, diagnostics = _simulate(initial, targets, reference_time=1.0) + assert diagnostics["temporal_model"]["limited_below_timestep"].all() + np.testing.assert_array_equal(diagnostics["model_correlations"], 0) + np.testing.assert_array_equal(diagnostics["temporal_model"]["achieved_times"], 7.5) + + +def test_simulation_example_end_to_end_parameter_recovery(): + """Regress the fixed-seed gallery workflow through output parameter recovery.""" + reference = json.loads( + ( + Path(__file__).parent / "data/parametric_simulation_reference.json" + ).read_text() + ) + case, expected = reference["case"], reference["expected"] + shape = tuple(case["shape"]) + threshold_db = 10 * np.log10(case["rain_threshold"]) + interval = np.timedelta64(round(case["timestep_minutes"] * 60), "s") + + # Start from a generated field and its translated copies, as in the example. + # The test uses public functions without importing plotting code or running + # the example's Monte Carlo experiment. + transformed = generate_calibrated_field( + shape, + case["parameters"], + case["kmperpixel"], + case["scale_break_km"], + threshold_db, + seed=case["initial_seed"], + ) + rain = 10 ** (transformed / 10) + rain[rain < case["rain_threshold"]] = 0 + velocity = np.array( + [ + np.full(shape, component) + for component in case["velocity_pixels_per_timestep"] + ] + ) + history = xr.Dataset( + { + "rain": ( + ("time", "y", "x"), + np.stack([np.roll(rain, -2, axis=1), np.roll(rain, -1, axis=1), rain]), + {"units": "mm/h"}, + ), + "velocity_x": (("y", "x"), velocity[0], {"units": "pixels/timestep"}), + "velocity_y": (("y", "x"), velocity[1], {"units": "pixels/timestep"}), + }, + coords={ + "time": ( + "time", + np.datetime64("2026-01-01T00:00:00") + np.arange(3) * interval, + {"stepsize": case["timestep_minutes"] * 60}, + ), + "x": ( + "x", + (np.arange(shape[1]) + 0.5) * case["kmperpixel"], + {"units": "km", "stepsize": case["kmperpixel"]}, + ), + "y": ( + "y", + (np.arange(shape[0]) + 0.5) * case["kmperpixel"], + {"units": "km", "stepsize": case["kmperpixel"]}, + ), + }, + attrs={"precip_var": "rain"}, + ) + output, diagnostics = parametric.simulate( + history, + timesteps=case["n_timesteps"], + parameters=case["parameters"], + reference_time=case["reference_time_minutes"], + reference_wavelength=case["reference_wavelength_km"], + temporal_exponent=case["temporal_exponent"], + alpha=case["alpha"], + n_cascade_levels=case["n_cascade_levels"], + rain_threshold=case["rain_threshold"], + scale_break=case["scale_break_km"], + n_ens_members=case["n_ens_members"], + seed=case["evolution_seed"], + return_diagnostics=True, + ) + assert output.rain.dims == ("ens_number", "time", "y", "x") + assert output.rain.shape == (case["n_ens_members"], case["n_timesteps"], *shape) + assert output.rain.dtype == np.float64 + assert np.isfinite(output.rain).all() and (output.rain >= 0).all() + assert np.all((output.rain == 0) | (output.rain >= case["rain_threshold"])) + assert output.rain.attrs["units"] == "mm/h" + assert output.rain.attrs["threshold"] == case["rain_threshold"] + assert output.rain.attrs["zerovalue"] == 0 + assert "transform" not in output.rain.attrs + xr.testing.assert_identical(output.x.variable, history.x.variable) + xr.testing.assert_identical(output.y.variable, history.y.variable) + np.testing.assert_array_equal(output.ens_number, [1]) + np.testing.assert_array_equal( + output.time, + history.time.values[-1] + np.arange(1, case["n_timesteps"] + 1) * interval, + ) + assert output.time.attrs["stepsize"] == case["timestep_minutes"] * 60 + assert output.forecast_reference_time == history.time.values[-1] + assert diagnostics["parameters"] == case["parameters"] + + bandpass = filter_gaussian(shape, case["n_cascade_levels"]) + wavelengths = case["kmperpixel"] / bandpass["central_freqs"] + np.testing.assert_allclose(wavelengths, expected["wavelengths_km"], rtol=1e-12) + np.testing.assert_allclose( + diagnostics["temporal_model"]["requested_times"], + case["reference_time_minutes"] + * (wavelengths / case["reference_wavelength_km"]) ** case["temporal_exponent"], + rtol=1e-12, + ) + np.testing.assert_allclose( + diagnostics["model_correlations"], + expected["model_correlations"], + atol=1e-10, + rtol=1e-10, + ) + + # Recover statistics from every returned field, rather than accepting the + # supplied parameters in diagnostics as evidence of the generated behaviour. + fields = ( + dB_transform( + output, + threshold=case["rain_threshold"], + zerovalue=threshold_db - 0.1, + ) + .rain.isel(ens_number=0) + .values + ) + recovered = [] + cascades = [] + for field in fields: + statistics = compute_field_statistics( + field, threshold_db, case["lower_truncation_dbr"] + ) + frequency, power = rapsd_logarithmic(field) + spectrum = fit_power_law_spectrum( + frequency, + power, + scale_break=case["scale_break_km"] / case["kmperpixel"], + ) + assert statistics["nonzero_fit"]["status"] == "ok" + assert spectrum["status"] == "ok" + recovered.append( + [ + statistics["rain_fraction"], + statistics["latent_mean"], + statistics["latent_stdev"], + spectrum["beta_1"], + spectrum["beta_2"], + ] + ) + cascades.append( + decomposition_fft( + field, + bandpass, + compute_stats=True, + normalize=True, + )["cascade_levels"] + ) + recovered = np.asarray(recovered) + assert np.isfinite(recovered).all() + # Fixed seeds remove Monte Carlo sampling variability from this comparison. + # Allow two threshold-crossing pixels and small floating-point differences + # in fitted moments/slopes; these are numerical regression tolerances. + tolerances = [2 / np.prod(shape), 1e-4, 1e-4, 1e-4, 1e-4] + for column, (name, tolerance) in enumerate( + zip(expected["field_parameter_names"], tolerances) + ): + np.testing.assert_allclose( + recovered[:, column], + np.asarray(expected["field_parameters"])[:, column], + rtol=1e-5, + atol=tolerance, + err_msg=name, + ) + + # Include only simulated fields: the artificial initialization is excluded + # from the ten overlapping three-field histories, but its transient persists. + cascades = np.stack(cascades) + observed = np.mean( + [ + lagrangian_autocorrelation(cascades[end - 2 : end + 1], velocity) + for end in range(2, len(cascades)) + ], + axis=0, + ) + assert np.isfinite(observed).all() + np.testing.assert_allclose( + observed, expected["mean_observed_correlations"], atol=1e-6, rtol=0 + ) + times = ar2_correlation_time( + observed[:, 0], observed[:, 1], case["timestep_minutes"] + ) + reference_time = estimate_reference_correlation_time( + times, + wavelengths, + case["temporal_exponent"], + levels=(0, 1, 2), + reference_wavelength=case["reference_wavelength_km"], + ) + # The reference describes this initialization and one-hour run, not stationary + # recovery of the prescribed 400 minutes. Do not make a bias factor a model + # requirement. High correlations amplify roundoff in the implied times, so + # check the lag pairs tightly and allow 0.1% on their nonlinear conversion. + np.testing.assert_allclose( + times, expected["implied_correlation_times_minutes"], rtol=1e-3, atol=1e-6 + ) + assert reference_time == pytest.approx( + expected["implied_reference_time_minutes"], rel=1e-3 + ) diff --git a/pysteps/timeseries/parametric.py b/pysteps/timeseries/parametric.py new file mode 100644 index 00000000..21cfa3f3 --- /dev/null +++ b/pysteps/timeseries/parametric.py @@ -0,0 +1,373 @@ +""" +pysteps.timeseries.parametric +============================= + +Correlation estimators and temporal scaling for parametric nowcasting. + +.. autosummary:: + :toctree: ../generated/ + + lagrangian_autocorrelation + ar2_correlation_time + fit_ar2_correlations + estimate_reference_correlation_time + model_cascade_correlations +""" + +import numpy as np +from scipy.optimize import brentq + +from pysteps.extrapolation.semilagrangian import extrapolate + +CORRELATION_TIME_DEFINITION = "ar_acf_sum_including_zero_v2" + + +def _finite(name, value): + """Require a finite scalar.""" + if not np.isscalar(value) or not np.isfinite(value): + raise ValueError(f"{name} must be finite") + + +def _positive(name, value): + """Require a positive finite scalar.""" + _finite(name, value) + if value <= 0: + raise ValueError(f"{name} must be positive and finite") + + +def _wavelengths(values): + """Return a nonempty one-dimensional array of positive wavelengths.""" + values = np.asarray(values, dtype=float) + if ( + values.ndim != 1 + or values.size == 0 + or not np.isfinite(values).all() + or np.any(values <= 0) + ): + raise ValueError("wavelengths must be a nonempty positive finite 1D array") + return values + + +def ar2_correlation_time(lag1, lag2, timestep, tol=1e-5): + """Calculate the one-sided stationary AR(2) discrete ACF sum including zero. + + Parameters + ---------- + lag1, lag2: float or array_like + Observed lag-1 and lag-2 correlations. Arrays are broadcast together. + timestep: float + Sampling interval. The returned times have the same unit. + tol: float, optional + Guard for singular or near-nonstationary pairs (default 1e-5). + + Returns + ------- + time: float or ndarray + timestep * sum(rho[k], k=0..infinity). Invalid, nonstationary, and + numerically singular pairs return NaN, without AR(1) substitution. + + Notes + ----- + The sum is timestep*(1+rho1-phi1)/(1-phi1-phi2), with Yule-Walker AR + coefficients. White noise has time equal to one timestep. This differs + from a trapezoidal ACF integral by timestep/2 and is not a skill horizon. + """ + _positive("tol", tol) + if tol >= 1: + raise ValueError("tol must be less than one") + lag1, lag2 = np.broadcast_arrays(np.asarray(lag1, float), np.asarray(lag2, float)) + result = np.full(lag1.shape, np.nan) + if not np.isscalar(timestep) or not np.isfinite(timestep) or timestep <= 0: + return float(result) if result.ndim == 0 else result + with np.errstate(divide="ignore", invalid="ignore"): + denominator = 1 - lag1**2 + phi1 = lag1 * (1 - lag2) / denominator + phi2 = (lag2 - lag1**2) / denominator + total_denominator = 1 - phi1 - phi2 + # Require stationary AR coefficients and denominators separated from + # zero; near-unit roots make the infinite correlation sum unstable. + valid = ( + np.isfinite(lag1) + & np.isfinite(lag2) + & (np.abs(lag1) < 1 - tol) + & (np.abs(lag2) < 1 - tol) + & (np.abs(denominator) >= tol) + & (np.abs(phi2) < 1 - tol) + & (phi1 + phi2 < 1 - tol) + & (phi2 - phi1 < 1 - tol) + & (total_denominator > tol) + ) + np.copyto(result, timestep * (1 + lag1 - phi1) / total_denominator, where=valid) + return float(result) if result.ndim == 0 else result + + +def _adjust_nonperiodic_lag2(lag1, lag2): + """Raise the lag-2 correlation to the nonperiodic AR(2) lower bound.""" + # Algebraically equivalent to adjust_lag2_corrcoef2, but avoids division + # by lag1**2 and cancellation near zero. + root = np.sqrt(1 - lag1**2) + lower = lag1**2 * (1 + 2 * root) / (1 + root) ** 2 + return float(max(lag2, 2 * lag1 * lag2 - 1, lower)) + + +def fit_ar2_correlations(correlation_time, alpha, timestep, return_diagnostics=False): + """Invert a discrete timescale in the nonnegative, nonperiodic AR(2) family. + + Parameters + ---------- + correlation_time, timestep: float + Positive target time and sampling interval in the same unit. A target + below timestep is unrepresentable and raises ValueError. + alpha: float + Positive exponent in rho2 = rho1**alpha before nonperiodic adjustment. + return_diagnostics: bool, optional + Also return requested and achieved times and their absolute difference. + + Returns + ------- + correlations: ndarray + The adjusted [rho1, rho2] used by the inversion. Exactly one timestep + returns [0, 0]. The exponent relation can change at the adjustment bound. + diagnostics: dict, optional + Definition, requested_time, achieved_time, and absolute_error. + """ + _positive("correlation_time", correlation_time) + _positive("alpha", alpha) + _positive("timestep", timestep) + if correlation_time < timestep: + raise ValueError("correlation_time is below the discrete minimum timestep") + + def pair(rho1): + return _adjust_nonperiodic_lag2(rho1, rho1**alpha) + + def residual(rho1): + return ar2_correlation_time(rho1, pair(rho1), timestep) - correlation_time + + rho1 = rho2 = 0.0 + if correlation_time > timestep: + # Find a finite sign-changing bracket below the stationarity limit. + # Nonfinite residuals lower the upper search bound. + lower, upper, candidate = 0.0, 1.0, 0.5 + bracket = None + for _ in range(80): + value = residual(candidate) + if np.isfinite(value): + if value >= 0: + bracket = (0.0, candidate) + break + lower = candidate + else: + upper = candidate + candidate = (lower + upper) / 2 + if candidate == lower or candidate == upper: + break + if bracket is None: + raise ValueError( + "correlation_time exceeds the numerically stable AR(2) range" + ) + rho1 = float(brentq(residual, *bracket, xtol=1e-14, rtol=1e-14)) + rho2 = pair(rho1) + achieved = ar2_correlation_time(rho1, rho2, timestep) + if not np.isfinite(achieved) or not np.isclose( + achieved, correlation_time, rtol=1e-7, atol=1e-8 + ): + raise ValueError("AR(2) inverse failed to achieve the requested timescale") + correlations = np.array([rho1, rho2]) + diagnostics = dict( + definition=CORRELATION_TIME_DEFINITION, + requested_time=float(correlation_time), + achieved_time=achieved, + absolute_error=float(abs(achieved - correlation_time)), + ) + return (correlations, diagnostics) if return_diagnostics else correlations + + +def estimate_reference_correlation_time( + times, wavelengths, exponent, levels=(0, 1, 2), reference_wavelength=None +): + """Estimate a reference time by geometric averaging selected scale-adjusted times. + + Parameters + ---------- + times: array_like + Observed correlation times with shape (..., level), in one time unit. + wavelengths: array_like + Positive nominal wavelengths, one per level, from the actual cascade + filter (reciprocals of its central frequencies in consistent units). + exponent: float + Fixed finite temporal scaling exponent, not estimated here. + levels: sequence of int, optional + Zero-based contributing levels; defaults to (0, 1, 2). Use (0,) to + estimate the reference time from the largest cascade level alone. + reference_wavelength: float, optional + Reference length, default wavelengths[0]. + + Returns + ------- + time: float or ndarray + exp(mean(log(T_j)-exponent*log(wavelength_j/reference_wavelength))). + Every selected time must be positive and finite; otherwise NaN is + returned for that sample. No level dropping or implicit fallback occurs. + """ + wavelengths = _wavelengths(wavelengths) + times = np.asarray(times, dtype=float) + if times.ndim < 1 or times.shape[-1] != wavelengths.size: + raise ValueError("times' last dimension must match wavelengths") + _finite("exponent", exponent) + levels = np.asarray(levels) + if ( + levels.ndim != 1 + or levels.size == 0 + or not np.issubdtype(levels.dtype, np.integer) + or np.any(levels < 0) + or np.any(levels >= wavelengths.size) + or np.unique(levels).size != levels.size + ): + raise ValueError("levels must contain distinct valid integer indices") + reference = wavelengths[0] if reference_wavelength is None else reference_wavelength + _positive("reference_wavelength", reference) + selected = times[..., levels] + valid = (np.isfinite(selected) & (selected > 0)).all(axis=-1) + result = np.full(times.shape[:-1], np.nan) + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + values = np.exp( + np.mean( + np.log(selected) - exponent * np.log(wavelengths[levels] / reference), + axis=-1, + ) + ) + np.copyto(result, values, where=valid & np.isfinite(values)) + return float(result) if result.ndim == 0 else result + + +def model_cascade_correlations( + reference_time, + wavelengths, + exponent, + alpha, + timestep, + reference_wavelength=None, + return_diagnostics=False, +): + """Construct AR(2) lag pairs from a reference time and nominal wavelengths. + + Parameters + ---------- + reference_time, timestep: float + Positive reference time and sampling interval, in the same unit. + wavelengths: array_like + Positive nominal wavelengths, from the actual cascade filter. + exponent, alpha: float + Finite scale-time exponent and positive pre-adjustment lag-pair exponent. + reference_wavelength: float, optional + Reference length, default wavelengths[0]. + return_diagnostics: bool, optional + Also return requested and achieved times, and sub-timestep limit flags. + + Returns + ------- + correlations: ndarray + Shape (level, 2), containing adjusted lag-1/lag-2 pairs. Requested times + at or below timestep use white noise [0, 0], achieving timestep. + diagnostics: dict, optional + The chosen ACF definition and per-level requested/achieved times. Length + and time units are caller supplied; no implicit conversion is performed. + """ + wavelengths = _wavelengths(wavelengths) + for name, value in [ + ("reference_time", reference_time), + ("alpha", alpha), + ("timestep", timestep), + ]: + _positive(name, value) + _finite("exponent", exponent) + reference = wavelengths[0] if reference_wavelength is None else reference_wavelength + _positive("reference_wavelength", reference) + requested = reference_time * (wavelengths / reference) ** exponent + if not np.isfinite(requested).all() or np.any(requested <= 0): + raise ValueError("scale-adjusted times must be positive and finite") + correlations = np.array( + [ + fit_ar2_correlations(max(time, timestep), alpha, timestep) + for time in requested + ] + ) + diagnostics = dict( + definition=CORRELATION_TIME_DEFINITION, + timestep=float(timestep), + requested_times=requested, + achieved_times=ar2_correlation_time( + correlations[:, 0], correlations[:, 1], timestep + ), + limited_below_timestep=requested < timestep, + ) + return (correlations, diagnostics) if return_diagnostics else correlations + + +def lagrangian_autocorrelation( + cascades, velocity, ar_order=2, min_stdev=0.1, boundary="zero" +): + """Estimate cascade correlations by successive unit-step backward advection. + + Parameters + ---------- + cascades: array_like + Finite history with shape (time, level, y, x), oldest first, containing + at least ar_order+1 times. Only the latest required history is used. + velocity: array_like + Finite fixed flow with shape (2, y, x), in pixels per timestep. + ar_order: {1, 2}, optional + Number of observed lags to estimate (default 2). + min_stdev: float, optional + Both aligned fields must exceed this population standard deviation. + boundary: {'zero', 'mask'}, optional + 'zero' fills out-of-domain pixels with zero after each unit step and + correlates the whole grid. For normalized cascades, zero represents + the level mean. 'mask' preserves NaNs through advection and correlates + only pixels that are finite across all aligned lags. + + Returns + ------- + correlations: ndarray + Shape (level, ar_order), in ascending lag order. Insufficient support or + variance produces NaN. There is no clipping or stationarity repair. + """ + cascades, velocity = np.asarray(cascades, float), np.asarray(velocity, float) + if not isinstance(ar_order, (int, np.integer)) or ar_order not in (1, 2): + raise ValueError("ar_order must be 1 or 2") + if ( + cascades.ndim != 4 + or cascades.shape[0] < ar_order + 1 + or any(size == 0 for size in cascades.shape) + ): + raise ValueError( + "cascades must have shape (time, level, y, x) with sufficient history" + ) + if velocity.shape != (2, *cascades.shape[-2:]): + raise ValueError("velocity must have shape (2, y, x) matching cascades") + if not np.isfinite(cascades).all() or not np.isfinite(velocity).all(): + raise ValueError("cascades and velocity must contain only finite values") + if boundary not in ("zero", "mask"): + raise ValueError("boundary must be 'zero' or 'mask'") + if not np.isfinite(min_stdev) or min_stdev < 0: + raise ValueError("min_stdev must be finite and nonnegative") + result = np.full((cascades.shape[1], ar_order), np.nan) + for level in range(cascades.shape[1]): + aligned = [cascades[-1, level]] + # Apply unit steps successively; a two-step field is interpolated twice. + for lag in range(1, ar_order + 1): + field = cascades[-(lag + 1), level] + for _ in range(lag): + field = extrapolate(field, velocity, 1, allow_nonfinite_values=True)[0] + if boundary == "zero": + field = np.where(np.isfinite(field), field, 0) + aligned.append(field) + support = np.isfinite(np.stack(aligned)).all(axis=0) + if np.count_nonzero(support) < 2: + continue + latest = aligned[0][support] + for lag in range(1, ar_order + 1): + previous = aligned[lag][support] + if np.std(latest) > min_stdev and np.std(previous) > min_stdev: + result[level, lag - 1] = np.corrcoef(latest, previous)[0, 1] + return result diff --git a/pysteps/utils/__init__.py b/pysteps/utils/__init__.py index a1a37c39..b5543a6f 100644 --- a/pysteps/utils/__init__.py +++ b/pysteps/utils/__init__.py @@ -11,5 +11,6 @@ from .pca import * from .reprojection import * from .spectral import * +from .statistics import * from .tapering import * from .transformation import * diff --git a/pysteps/utils/spectral.py b/pysteps/utils/spectral.py index a304b4ab..7803e969 100644 --- a/pysteps/utils/spectral.py +++ b/pysteps/utils/spectral.py @@ -11,6 +11,8 @@ corrcoef mean rapsd + rapsd_logarithmic + fit_power_law_spectrum remove_rain_norain_discontinuity std """ @@ -19,6 +21,185 @@ from . import arrays +def rapsd_logarithmic(field, min_stdev=0.1): + """Compute a standardized field's radial spectrum in logarithmic bins. + + Parameters + ---------- + field: array_like + Finite, real two-dimensional field in (y, x) order. Dimensions must + each contain at least four pixels. The input is not modified. + min_stdev: float, optional + Minimum population standard deviation required for estimation. + + Returns + ------- + frequency, power: ndarray + Bin-centre frequencies in cycles per pixel and mean power per Fourier + mode, both in linear units. Empty arrays indicate insufficient variance + or bin support. + + Notes + ----- + Use min(shape)//2 logarithmic edges, starting at 1/(min(shape)//2) + cycles per pixel. The upper edge is the diagonal radius formed from the + largest absolute frequencies in the first half of each FFT frequency array. + This differs from :func:`rapsd` and omits the lowest frequencies. Empty bins + are omitted. Power is the squared unscaled FFT of a zero-mean, unit-variance + field, not annulus-integrated energy. Real-FFT multiplicities reproduce the + full Fourier-plane mean. + """ + if np.iscomplexobj(field): + raise ValueError("field must contain real values") + field = np.asarray(field, dtype=float) + if field.ndim != 2 or min(field.shape) < 4: + raise ValueError( + "field must be two-dimensional with at least 4 pixels per axis" + ) + if not np.isfinite(field).all(): + raise ValueError("field must contain only finite values") + if not np.isfinite(min_stdev) or min_stdev < 0: + raise ValueError("min_stdev must be finite and nonnegative") + stdev = np.std(field) + if stdev == 0 or stdev < min_stdev: + return np.empty(0), np.empty(0) + normalized = (field - np.mean(field)) / stdev + power = np.abs(np.fft.rfft2(normalized)) ** 2 + rows, cols = field.shape + fx, fy = np.fft.rfftfreq(cols), np.fft.fftfreq(rows) + radius = np.hypot(fy[:, None], fx[None, :]) + # The real FFT omits conjugate modes. Count them twice, except in the + # zero-frequency and (for even widths) Nyquist columns. + weights = np.full(power.shape, 2.0) + weights[:, 0] = 1.0 + if cols % 2 == 0: + weights[:, -1] = 1.0 + n_bins = min(rows, cols) // 2 + maximum = np.hypot( + np.max(np.abs(np.fft.fftfreq(rows)[: rows // 2])), + np.max(np.abs(np.fft.fftfreq(cols)[: cols // 2])), + ) + # Keep logarithmic endpoints explicit: roundoff can move a Fourier mode + # between bins when it falls exactly on an edge. + edges = np.logspace(np.log10(1 / n_bins), np.log10(maximum), n_bins) + frequency, radial_power = [], [] + for lower, upper in zip(edges[:-1], edges[1:]): + selected = (radius >= lower) & (radius < upper) + if np.any(selected): + frequency.append(0.5 * (lower + upper)) + radial_power.append(np.average(power[selected], weights=weights[selected])) + return np.asarray(frequency), np.asarray(radial_power) + + +def fit_power_law_spectrum( + frequency, power, scale_break=None, skip_bins=2, max_frequency=10 ** (-0.4) +): + """Fit signed spectral slopes, optionally at a prescribed wavelength break. + + Parameters + ---------- + frequency, power: array_like + Equally sized one-dimensional arrays in linear units, with frequencies + positive and strictly increasing. Power must be finite and nonnegative. + scale_break: float, optional + Fixed break wavelength in pixels when frequencies are cycles per pixel. + For physical frequencies, use the corresponding length unit. None fits + one slope. The fit is continuous at frequency 1/scale_break. + skip_bins: int, optional + Number of initial bins excluded from regression (default 2). + max_frequency: float, optional + Exclusive upper fit limit, default 10**(-0.4) cycles per pixel. + Set this in matching units if supplying physical frequencies. + + Returns + ------- + result: dict + ``beta_1``, ``beta_2``, ``c1``, ``c2``, and ``scale_break`` describe + 10*log10(power) = beta*10*log10(frequency) + c. Slopes are signed. + ``status``, ``n_samples``, and ``rmse_db`` describe the fit. Unavailable + coefficients are NaN; beta_2 and c2 are None for a single-slope model. + + Notes + ----- + Zero-power bins cannot enter log-space regression. Insufficient samples, + an unidentified break, or failed optimization are reported without inventing + parameter estimates. The model is linear in its coefficients and is solved + by unweighted least squares. No climatological clipping or calibration is + applied. + """ + frequency, power = np.asarray(frequency, float), np.asarray(power, float) + if frequency.ndim != 1 or power.shape != frequency.shape: + raise ValueError("frequency and power must be equally sized 1D arrays") + if ( + not np.isfinite(frequency).all() + or np.any(frequency <= 0) + or np.any(np.diff(frequency) <= 0) + or not np.isfinite(power).all() + or np.any(power < 0) + ): + raise ValueError( + "frequencies must increase and be positive; power must be nonnegative" + ) + if not isinstance(skip_bins, (int, np.integer)) or skip_bins < 0: + raise ValueError("skip_bins must be a nonnegative integer") + if not np.isfinite(max_frequency) or max_frequency <= 0: + raise ValueError("max_frequency must be positive and finite") + if scale_break is not None and (not np.isfinite(scale_break) or scale_break <= 0): + raise ValueError("scale_break must be positive and finite") + selected = (np.arange(frequency.size) >= skip_bins) & (frequency < max_frequency) + selected &= power > 0 + x, y = 10 * np.log10(frequency[selected]), 10 * np.log10(power[selected]) + result = dict( + beta_1=np.nan, + beta_2=np.nan if scale_break is not None else None, + c1=np.nan, + c2=np.nan if scale_break is not None else None, + scale_break=scale_break, + status="insufficient_samples", + n_samples=int(x.size), + rmse_db=np.nan, + ) + if scale_break is None: + + def model(x, beta, intercept): + return beta * x + intercept + + design = np.column_stack([x, np.ones(x.size)]) + else: + breakpoint = 10 * np.log10(1 / scale_break) + + def model(x, beta_1, beta_2, c1): + c2 = (beta_1 - beta_2) * breakpoint + c1 + return np.where(x <= breakpoint, beta_1 * x + c1, beta_2 * x + c2) + + # This basis fits both slopes and a shared intercept while enforcing + # continuity at the fixed breakpoint. + design = np.column_stack( + [np.minimum(x, breakpoint), np.maximum(x - breakpoint, 0), np.ones(x.size)] + ) + if x.size < design.shape[1]: + return result + if np.linalg.matrix_rank(design) < design.shape[1]: + result["status"] = "unidentified_model" + return result + try: + fitted, _, _, _ = np.linalg.lstsq(design, y, rcond=None) + except np.linalg.LinAlgError as error: + result.update(status="linear_solver_failed", message=str(error)) + return result + if not np.isfinite(fitted).all(): + result["status"] = "linear_solver_failed" + return result + result.update(beta_1=float(fitted[0]), c1=float(fitted[-1]), status="ok") + if scale_break is not None: + result.update( + beta_2=float(fitted[1]), + c2=float((fitted[0] - fitted[1]) * breakpoint + fitted[2]), + ) + result["rmse_db"] = float(np.sqrt(np.mean((model(x, *fitted) - y) ** 2))) + return result + + def corrcoef(X, Y, shape, use_full_fft=False): """ Compute the correlation coefficient between two-dimensional arrays in diff --git a/pysteps/utils/statistics.py b/pysteps/utils/statistics.py new file mode 100644 index 00000000..8f511b86 --- /dev/null +++ b/pysteps/utils/statistics.py @@ -0,0 +1,180 @@ +""" +pysteps.utils.statistics +======================== + +Marginal parameter estimators for transformed precipitation. + +.. autosummary:: + :toctree: ../generated/ + + fit_lower_truncated_normal + compute_field_statistics +""" + +import numpy as np +from scipy.optimize import minimize_scalar +from scipy.special import log_ndtr + +__all__ = ["fit_lower_truncated_normal", "compute_field_statistics"] + + +def fit_lower_truncated_normal(data, lower_truncation, return_diagnostics=False): + """Fit a parent Normal distribution to samples at or above a lower cutoff. + + Parameters + ---------- + data: array_like + Observations in transformed units. Nonfinite and below-cutoff values + are excluded. The input is not modified. + lower_truncation: float + Finite lower fitting cutoff in the same units as data. + return_diagnostics: bool, optional + Also return a dictionary describing sample count and fitting status. + + Returns + ------- + location, scale: float + Parent-Normal location and scale, distinct from retained-sample moments. + Both are NaN for insufficient or constant samples or failed optimization. + diagnostics: dict, optional + Sample count, status, standardized cutoff, and boundary indicators. + + Notes + ----- + Profile the scale analytically for a = (cutoff-location)/scale and minimize + the conditional negative log likelihood over -12 <= a <= 12. Use log_ndtr + to evaluate the retained Normal tail probability without underflow. + A finite bounded solution is not proof of an unrestricted finite MLE: + ``boundary_limited`` identifies proximity to the bound or shifted-sample + coefficient of variation >= 1. The bounded numerical solution is returned + with this diagnostic so callers can decide whether to use it. Pixels have + equal marginal weights; spatial dependence and parameter uncertainty are + not estimated here. + """ + if not np.isfinite(lower_truncation): + raise ValueError("lower_truncation must be finite") + if np.iscomplexobj(data): + raise ValueError("data must contain real values") + sample = np.asarray(data, dtype=float) + sample = sample[np.isfinite(sample) & (sample >= lower_truncation)] + diagnostics = dict( + n_samples=int(sample.size), + status="insufficient_samples", + boundary_limited=False, + ) + + def output(location=np.nan, scale=np.nan): + result = (float(location), float(scale)) + return (*result, diagnostics) if return_diagnostics else result + + if sample.size < 2: + return output() + if np.ptp(sample) <= 0: + diagnostics["status"] = "constant_sample" + return output() + # Express retained observations relative to the cutoff. For a fixed + # standardized cutoff a, the likelihood gives scale analytically. + shifted = sample - lower_truncation + mean_shifted = float(np.mean(shifted)) + second_moment = float(np.mean(shifted * shifted)) + + def scale_for_cutoff(a): + """Return the positive profile-likelihood scale for a fixed cutoff a.""" + return 0.5 * ( + a * mean_shifted + np.sqrt(a**2 * mean_shifted**2 + 4 * second_moment) + ) + + def objective(a): + """Evaluate conditional negative log likelihood per retained sample.""" + scale = scale_for_cutoff(a) + if not np.isfinite(scale) or scale <= 0: + return np.inf + standardized_second_moment = ( + second_moment / scale**2 + 2 * a * mean_shifted / scale + a**2 + ) + return float(np.log(scale) + 0.5 * standardized_second_moment + log_ndtr(-a)) + + # A bounded scalar search avoids fitting location and scale independently. + result = minimize_scalar( + objective, + bounds=(-12.0, 12.0), + method="bounded", + options={"xatol": 1e-8, "maxiter": 200}, + ) + a = float(result.x) + scale = float(scale_for_cutoff(a)) + location = float(lower_truncation - a * scale) + if not result.success or not np.isfinite([location, scale]).all() or scale <= 0: + diagnostics["status"] = "optimization_failed" + return output() + cv = float(np.std(shifted) / mean_shifted) + boundary = abs(a) >= 12 - 1e-5 or cv >= 1 + diagnostics.update( + status="bounded_solution" if boundary else "ok", + boundary_limited=bool(boundary), + standardized_cutoff=a, + shifted_sample_cv=cv, + ) + return output(location, scale) + + +def compute_field_statistics(data, threshold, lower_truncation=None): + """Estimate rainfall occurrence and ordinary and latent marginal parameters. + + Parameters + ---------- + data: array_like + Nonempty, finite two-dimensional field. Units may be transformed rainfall + (dBR for parametric nowcasting). The input is not modified. + threshold: float + Finite rain-occurrence threshold in the same units as data. Values at + or above the threshold are wet. The denominator is the whole valid grid. + lower_truncation: float, optional + Lower cutoff for fitting a latent Normal, at or above threshold. None + skips fitting. This setting does not change rainfall occurrence. + + Returns + ------- + result: dict + Rain fraction, whole-field mean and standard deviation, ``ordinary_mean`` + and ``ordinary_stdev`` of wet pixels (population moments), ``latent_mean`` + and ``latent_stdev``, pixel counts, and ``nonzero_fit`` diagnostics. + Unavailable moments are NaN. Estimates are not clipped or replaced. + """ + if np.iscomplexobj(data): + raise ValueError("data must contain real values") + data = np.asarray(data, dtype=float) + if data.ndim != 2 or data.size == 0: + raise ValueError("data must be a nonempty two-dimensional field") + if not np.isfinite(data).all(): + raise ValueError("data must contain only finite values") + if not np.isfinite(threshold): + raise ValueError("threshold must be finite") + if lower_truncation is not None and ( + not np.isfinite(lower_truncation) or lower_truncation < threshold + ): + raise ValueError("lower_truncation must be finite and at or above threshold") + wet = data[data >= threshold] + mean, stdev = ( + (float(np.mean(wet)), float(np.std(wet))) if wet.size else (np.nan, np.nan) + ) + result = dict( + rain_fraction=float(wet.size / data.size), + mean=float(np.mean(data)), + stdev=float(np.std(data)), + ordinary_mean=mean, + ordinary_stdev=stdev, + latent_mean=np.nan, + latent_stdev=np.nan, + n_pixels=int(data.size), + n_wet=int(wet.size), + threshold=float(threshold), + nonzero_fit={"status": "not_requested", "lower_truncation": None}, + ) + if lower_truncation is not None: + location, scale, diagnostics = fit_lower_truncated_normal( + wet, lower_truncation, return_diagnostics=True + ) + diagnostics["lower_truncation"] = float(lower_truncation) + result.update(latent_mean=location, latent_stdev=scale, nonzero_fit=diagnostics) + return result