Skip to main content

Implementation Details

Architecture Overview

precip-index/
├── src/
│   ├── indices.py           # SPI/SPEI calculation
│   ├── compute.py           # Core computation algorithms
│   ├── distributions.py     # Multi-distribution fitting (Gamma, Pearson III, etc.)
│   ├── chunked.py           # Memory-efficient chunked processing
│   ├── config.py            # Configuration and constants
│   ├── utils.py             # PET, data validation, memory utilities
│   ├── runtheory.py         # Event analysis (dry & wet)
│   ├── visualization.py     # Plotting functions
│   └── __init__.py          # Package exports
├── input/                   # User data (not in repo)
├── output/                  # Generated results
│   ├── netcdf/
│   ├── csv/
│   └── plots/
└── notebook/               # Tutorials

Core Modules

1. indices.py

Purpose: Calculate SPI and SPEI indices

Key Functions:

  • spi() - Single-scale SPI
  • spi_multi_scale() - Multiple scales at once
  • spei() - Temperature-inclusive index
  • spei_multi_scale() - Multiple SPEI scales
  • spi_global() - Global-scale SPI with automatic chunking
  • spei_global() - Global-scale SPEI with automatic chunking
  • estimate_memory_requirements() - Pre-computation memory estimation
  • save_fitting_params() / load_fitting_params() - Parameter persistence for operational workflows

Technology:

  • NumPy vectorization for gridded operations
  • Numba JIT compilation for fitting loops
  • SciPy for statistical distributions
  • xarray for CF-compliant NetCDF

Performance:

  • 165×244 grid, 12-month scale: ~30 seconds
  • Gamma distribution fitting: Numba-optimized
  • Memory-efficient chunking for large datasets

2. chunked.py

Purpose: Memory-efficient processing for global-scale datasets

Key Classes/Functions:

  • ChunkedProcessor - Main class for chunked computation
    • compute_spi_chunked() - Chunked SPI calculation
    • compute_spei_chunked() - Chunked SPEI calculation
  • estimate_memory() - Estimate memory requirements
  • iter_chunks() - Generate spatial chunk coordinates
  • MemoryEstimate - Named tuple for memory information

Technology:

  • Spatial tiling for memory efficiency
  • Sliced netCDF4 writes: each tile writes only its own slice, rather than reading the whole variable back, modifying it, and writing it out again
  • Tiles are processed serially in the calling process. The n_workers argument is deprecated and ignored: use dask_processor for concurrency
  • Automatic garbage collection
  • Progress callbacks for monitoring

Performance:

  • Global CHIRPS (2160×4320): Processes with ~16GB RAM
  • Chunk size auto-optimization based on available memory
  • ~12x memory reduction vs. in-memory processing

3. dask_processor.py

Purpose: Concurrent tiled processing for global grids, with resume and parameter retention

Key Functions:

  • compute_spi_dask() - Global SPI across a Dask LocalCluster
  • compute_spei_dask() - Global SPEI, same machinery
  • plan_layout() - Size the tile so n_workers tiles fit in memory at once
  • zarr_to_netcdf() - Stream a finished Zarr store out to NetCDF
  • DaskLayout - Dataclass holding tile size, worker count and peak memory

Technology:

  • Zarr region writes, which are lock-free when the regions are disjoint, so workers write concurrently without coordination
  • A run manifest (.manifest.json) keyed on a signature of the run, so an interrupted run resumes and only recomputes missing tiles
  • Per-tile fitting parameters staged to _tileparams/ and assembled at the end
  • Write retry with exponential backoff, which absorbs the transient file locks that antivirus and search indexers produce on Windows

Performance:

  • Peak memory is per worker, not per tile: total is the per-tile figure times n_workers. Always call plan_layout() before launching
  • Global TerraClimate SPEI (4320×8640, 912 months) at tile 640 with 6 workers: 1h 36m to compute all 98 tiles, plus 1h 16m if you export the Zarr store to a single NetCDF. See Table 1 for the phase breakdown

4. compute.py

Purpose: Core computation algorithms for SPI/SPEI

Key Functions:

  • compute_index_parallel() - Main computation with optional memory-efficient mode
  • _rolling_sum_3d() - O(n) cumulative sum algorithm
  • _compute_gamma_params_vectorized() - Vectorized Gamma parameter fitting
  • _compute_params_pearson3_vectorized() - Vectorized Pearson III fitting
  • _transform_to_normal_vectorized() - Memory-efficient transformation
  • compute_index_dask() - Dask-based lazy computation
  • compute_index_dask_to_zarr() - Stream directly to Zarr format

Technology:

  • Float32 storage (50% memory reduction) with a float64 accumulator inside the rolling sum, see below
  • Cumulative sum for O(n) rolling windows
  • Period-by-period processing to limit peak memory
  • Pearson III fitted by vectorized method of moments rather than a per-cell loop, which is roughly 12x faster and bit-identical to the loop it replaced
NoteWhy the rolling sum accumulates in float64

_rolling_sum_3d() computes windowed sums by differencing a cumulative sum, which is O(n) instead of O(n·scale) but is vulnerable to catastrophic cancellation. For SPEI the SPEI_WATER_BALANCE_OFFSET of 1000 drives the running total toward 8e5, where float32 spacing is 0.0625, so the difference of two large near-equal sums loses precision. The visible symptom is that the same grid cell yields slightly different values depending on how many years of record precede it.

Measured on TerraClimate SPEI-12, comparing a 1950-start against a 1958-start record: float32 accumulation gave mean absolute difference 9.7e-05 and maximum 1.3e-02, while float64 accumulation gave exactly zero. Output stays float32; only the accumulator is widened. The valid-value counter is int32 for the same reason: int16 caps at 32,767, which a daily record longer than about 89 years would overflow.

5. config.py

Purpose: Configuration constants and enumerations

Key Constants:

  • MEMORY_MULTIPLIER / MEMORY_MULTIPLIER_SPEI - Peak memory estimation factors (12x for SPI, 18x for SPEI)
  • MEMORY_SAFETY_FACTOR - Safety margin for chunk sizing (0.7)
  • DEFAULT_CHUNK_LAT/LON - Default chunk dimensions (500)
  • DASK_TILE_MULTIPLIER_SPI / _SPEI - Per-worker peak for the Dask path (9x and 12x), lower because that path is float32 throughout
  • MAX_DASK_TILE - Upper bound on the Dask tile edge (2048)
  • MIN_VALUES_FOR_FIT, MIN_NONZERO_VALUES, MAX_ZERO_PROPORTION - Fitting thresholds, single source of truth for distributions.py
  • Periodicity - Enum for monthly/daily data

6. runtheory.py

Purpose: Climate extreme event analysis using run theory (Yevjevich 1967)

Key Functions:

Event-Based:

  • identify_events() - Extract complete events (dry and wet)
  • summarize_events() - Event statistics

Time-Series:

  • calculate_timeseries() - Month-by-month tracking
  • Real-time monitoring capability

Gridded Statistics:

  • calculate_period_statistics() - Spatial statistics for time period
  • calculate_annual_statistics() - Year-by-year analysis
  • compare_periods() - Multi-period comparison

Implementation:

  • Pure NumPy for single-location (very fast)
  • xarray.apply_ufunc for gridded data
  • Dual magnitude: cumulative + instantaneous

7. utils.py

Purpose: Supporting utilities and memory management

Key Functions:

  • calculate_pet() - Potential evapotranspiration (dispatcher)
    • method='thornthwaite' - Uses only mean temperature (default)
    • method='hargreaves' - Uses Tmin, Tmax, and extraterrestrial radiation
  • eto_thornthwaite() - Thornthwaite PET (low-level)
  • eto_hargreaves() - Hargreaves-Samani PET (low-level)
  • _extraterrestrial_radiation() - Daily Ra calculation (FAO-56)
  • _monthly_mean_extraterrestrial_radiation() - Monthly Ra averages
  • validate_array() - Input checking
  • find_variable() - Resolve a data variable from name patterns, shared by indices, chunked and dask_processor
  • get_optimal_chunk_size() - Calculate optimal chunk dimensions
  • format_bytes() - Human-readable byte formatting
  • get_array_memory_size() - Calculate array memory footprint
  • print_memory_info() - Display current memory usage

8. visualization.py

Purpose: Publication-quality climate extremes visualization

Key Functions:

  • plot_index() - Time series with WMO colors
  • plot_events() - Event timeline with dry/wet differentiation
  • plot_event_characteristics() - Multi-panel analysis
  • plot_event_timeline() - 5-panel evolution
  • plot_spatial_stats() - Spatial mapping
  • generate_location_filename() - Consistent naming

Technology:

  • Matplotlib for all plots
  • Color schemes from scientific literature
  • Automatic folder creation

Data Flow

SPI/SPEI Calculation

Input NetCDF (precip, temp, [tmin, tmax])
    ↓
validate_data()
    ↓
[calculate_pet() if SPEI]
    ├─ pet_method='thornthwaite' → eto_thornthwaite(temp, lat)
    └─ pet_method='hargreaves'  → eto_hargreaves(temp, tmin, tmax, lat)
                                   └─ _extraterrestrial_radiation(lat, doy)
    ↓
spi() or spei()
    ├─ Temporal aggregation (rolling sum)
    ├─ Distribution fitting (Numba-optimized)
    └─ Probability transformation
    ↓
Output NetCDF (SPI/SPEI)

Climate Extreme Event Analysis

SPI/SPEI time series
    ↓
identify_events()
    ├─ Threshold crossing detection
    ├─ Duration calculation
    ├─ Magnitude (cumulative + instantaneous)
    ├─ Intensity (magnitude/duration)
    └─ Peak identification
    ↓
Events DataFrame or
Period Statistics Dataset

Optimization Strategies

1. Numba JIT Compilation

Where: Distribution fitting loops in indices.py

@numba.jit(nopython=True)
def fit_gamma_numba(precip_2d):
    # Fast parameter estimation
    # 10-20x faster than pure NumPy

Benefit: Critical for operational speed

2. Vectorization

Where: All array operations

# Vectorized temporal aggregation
precip_sum = precip.rolling(time=scale).sum()

# Vectorized threshold operations
is_event = (spi < threshold).astype(int)

Benefit: Leverages NumPy/xarray efficiency

3. Chunking (Dask-ready)

Where: Large dataset processing

# Open with chunks for lazy loading
precip = xr.open_dataset('data.nc', chunks={'time': 100})

# Operations are lazy until .compute()
result = spi(precip, scale=12)  # No computation yet
result.to_netcdf('output.nc')   # Computes and saves

Benefit: Handle datasets larger than RAM

4. Parameter Caching

Where: Operational/production systems

# Fit once on calibration period
params = spi(precip_calibration, scale=12,
             calibration_start_year=1991,
             calibration_end_year=2020)

# Save parameters
params.to_netcdf('spi_12_params.nc')

# Reuse for operational updates
spi_new = spi(precip_new, scale=12,
              fitting_params=params)

Benefit: Near-instant operational updates

Memory Management

Memory Requirements by Data Size

Grid Size Input Data Peak Memory (Old) Peak Memory (New) Recommended Approach
Single Location <1 MB <10 MB <10 MB In-memory
Small (100×100) ~50 MB ~500 MB ~100 MB In-memory
Medium (500×500) ~500 MB ~5 GB ~1 GB In-memory or chunked
Large (1000×1000) ~2 GB ~20 GB ~4 GB Chunked
Global (2160×4320) ~36 GB ~400 GB ~16 GB Chunked (required)

Memory-Efficient Processing (New in 2026.1)

The new chunked processing module reduces peak memory by ~12x through:

  1. Float32 Internal Processing
    • Uses float32 instead of float64 internally
    • 50% memory reduction for intermediate arrays
    • Output precision maintained
  2. Cumulative Sum Algorithm
    • O(n) rolling sum instead of O(n×scale)
    • Single pass through time dimension
    • No intermediate rolling arrays
  3. Period-by-Period Transformation
    • Process one calendar period at a time
    • Garbage collection between periods
    • Limits peak memory during standardization
  4. Spatial Chunking
    • Process data in spatial tiles
    • Stream results to disk
    • Memory bounded by chunk size, not total grid

Chunk Size Guidelines

Available RAM Recommended Chunk Size Notes
16 GB 200 × 200 Minimum recommended
32 GB 300 × 300 Good for workstations
64 GB 400 × 400 Standard servers
128 GB 600 × 600 High-memory systems
256 GB+ 800 × 800 Large servers

Usage Examples

In-memory (small data):

from indices import spi
spi_12 = spi(precip, scale=12)

Chunked processing (large data):

from indices import spi_global
result = spi_global(
    'global_precip.nc',
    'spi_12_global.nc',
    scale=12,
    chunk_size=500
)

Memory estimation before processing:

from indices import estimate_memory_requirements
mem = estimate_memory_requirements('global_precip.nc')
print(mem)  # Shows input size, peak memory, recommendations

Concurrent global processing with resume and saved parameters:

from dask_processor import compute_spei_dask, plan_layout

layout = plan_layout(n_time=912, n_lat=4320, n_lon=8640,
                     n_workers=6, index='spei')
print(layout)   # check peak total against available RAM before launching

compute_spei_dask(
    'terraclimate_ppt.nc',
    'terraclimate_pet.nc',
    'spei_12.zarr',
    scale=12,
    distribution='pearson3',
    tile=layout.tile,
    n_workers=layout.n_workers,
    params_path='spei_12_params.nc',
    resume=True
)

Output Standards

NetCDF Files

Compliance: CF Convention 1.8

Required Attributes:

attrs = {
    'standard_name': 'standardized_precipitation_index',
    'long_name': 'SPI-12 (Gamma distribution)',
    'units': '1',
    'scale': 12,
    'distribution': 'gamma',
    'calibration_period': '1991-2020',
    'threshold': -1.0
}

Coordinates:

  • time: Datetime64 with monthly frequency
  • lat: WGS84 latitude
  • lon: WGS84 longitude

CSV Files

Format: Standard pandas DataFrame export

Columns:

  • Event-based: event_id, start_date, end_date, duration, magnitude, etc.
  • Time-series: time, index_value, is_event, duration, magnitude_cumulative, etc.

Plots

Format: PNG (300 dpi, recommended)

Naming: Location/period-based

  • plot_*_lat*.##_lon*.##.png
  • plot_spatial_*_YYYY-YYYY.png

Testing Strategy

Unit Tests (Future)

  • Individual function validation
  • Edge case handling
  • Parameter validation

Integration Tests

  • test_drought_characteristics.py
  • End-to-end workflow
  • Real data validation

Validation

  • Compare with published results
  • Cross-check with original climate-indices
  • Visual inspection of outputs

Dependencies

Required

numpy>=1.20.0       # Array operations
scipy>=1.7.0        # Statistical distributions
xarray>=0.19.0      # NetCDF handling
netCDF4>=1.5.7      # NetCDF I/O
numba>=0.54.0       # JIT compilation
matplotlib>=3.4.0   # Visualization
pandas>=1.1.0       # DataFrame operations

Required for Global Processing

dask[complete]>=2021.10.0  # Parallel/distributed computing
zarr>=2.10.0               # Efficient chunked array storage
psutil>=5.8.0              # Memory detection and monitoring

Optional

cartopy>=0.20.0     # Map projections for visualization
jupyter>=1.0.0      # Notebook support

Error Handling

Data Validation

  • Check for NaN values
  • Verify time dimension
  • Ensure positive precipitation
  • Validate coordinate systems

Distribution Fitting

  • Fallback to default parameters if fitting fails
  • Warning messages for problematic locations
  • NaN preservation in outputs

File Operations

  • Auto-create output directories
  • Handle existing file overwrites
  • Validate paths

Performance Benchmarks

Standard Processing (Laptop: 4 cores, 16 GB RAM)

Operation Grid Size Time Memory
SPI-12 calculation 165×244×408 30 sec 500 MB
Drought events (single) 408 months <0.1 sec <1 MB
Period statistics 165×244, 5 years 30 sec 200 MB
Annual statistics 165×244, 10 years 5 min 300 MB
Plot generation Single location 2 sec 50 MB

Global-Scale Processing (Workstation: 128 GB RAM)

Operation Dataset Grid Size Time Steps Chunks Total Time
SPI-12 (Gamma), serial chunked CHIRPS v3 (0.05°) 2400 × 7200 (17.3M cells) 539 months 12 ~2h 47m
SPEI-12 (Pearson III), serial chunked TerraClimate (0.04°) 4320 × 8640 (37.3M cells) 816 months 32 ~26h
SPEI-12 (Pearson III), Dask TerraClimate v1.1 (0.04°) 4320 × 8640 (37.3M cells) 912 months 98 tiles, 6 workers 1h 36m
SPEI-2 (Pearson III), Dask TerraClimate v1.1 (0.04°) 4320 × 8640 (37.3M cells) 912 months 98 tiles, 6 workers 1h 40m

The Dask row is the same computation on dask_processor at tile 640 with 6 concurrent workers, rather than the serial chunked path. It also benefits from the vectorized Pearson III fitting described below.

Dask Run, Phase by Phase

Read off the console timestamps of two complete runs, identical except for the accumulation scale. TerraClimate v1.1, Pearson III, 912 × 4320 × 8640, tile 640, 6 workers, ~350 GB RAM available and 100.2 GB peak planned. Both completed all 98 tiles with no failures:

Table 1: Measured phase timings for two global Dask SPEI runs at different accumulation scales.
Phase SPEI-12 SPEI-2
Cluster startup through first completed tile 0h 05m 35s 0h 06m 02s
Tiles 1 to 98 1h 28m 34s 1h 32m 18s
Teardown and parameter assembly (6.67 GB) 0h 00m 38s 0h 00m 36s
Write parameter file 0h 01m 24s 0h 01m 23s
Dask compute total 1h 36m 14s 1h 40m 21s
Zarr to NetCDF export (136 GB array) 1h 16m 29s 1h 17m 46s
End to end 2h 52m 43s 2h 58m 07s

Both columns were read from the dask_processor console output of the run itself. The SPEI-12 run is 2026-09-20 09:23 to 12:16; the SPEI-2 run is 2026-09-19 23:52 to 2026-09-20 02:50.

Tile completions arrived every 30 seconds at the median for SPEI-12 (mean 55s, range 2s to 265s) with 6 workers in flight. The spread is real: tiles over ocean finish almost immediately because every cell is NaN, while dense land tiles carry the full fitting cost.

ImportantThe export costs as much as the computation

Converting the finished Zarr store to a single NetCDF took 1h 16m, 44% of the end-to-end time, to write one 136 GB array. That step is single-threaded and I/O-bound, so no amount of extra workers helps it.

If your downstream tooling can read Zarr, skip zarr_to_netcdf() entirely and halve your wall clock. Convert only when you need a single portable file.

NoteScale does not drive the cost, as the two columns show

SPEI-12 accumulates a window six times longer than SPEI-2 and finished 4 minutes faster, a difference inside run-to-run noise. This is the expected result rather than a surprise: _rolling_sum_3d() differences a cumulative sum, so it is O(n) in the number of time steps whatever the window length, and the distribution fitting is per calendar period, not per scale.

Budget a global run by grid size and record length. The scale you pick is free.

Why Does TerraClimate SPEI Take ~9.5× Longer?

The difference comes from compounding factors across every dimension:

Factor CHIRPS SPI-12 TerraClimate SPEI-12 Impact
Grid cells 17.3M 37.3M 2.2× more cells
Time steps 539 months (1981–2025) 816 months (1958–2026) 1.5× longer series
Dataset size ~69 GB ~226 GB 3.2× more I/O
Distribution Gamma (Method of Moments) Pearson III (Method of Moments) Slower fitting
Extra computation — PET calculation Additional overhead
Spatial chunks 12 32 2.7× more I/O cycles

Key takeaway: the cell count, the series length and the I/O volume all grow together. Combined, 2.2× more grid cells and 1.5× longer series already put the work near 3.3×, and the I/O overhead of 32 chunks plus the PET computation account for the rest of the ~9.5× ratio.

The fitting itself is no longer the bottleneck it was when these numbers were taken. Both distributions now use closed-form method of moments, and Pearson III is fitted with array reductions rather than one SciPy call per cell per calendar period, which replaces roughly 24.9 million scalar calls per 1440×1440 tile with 12 array operations. Timings above predate that change: see the Dask row for a current figure.

TipRegional Processing Is Fast

These benchmarks are for global-scale datasets. Country or regional subsets (e.g., Indonesia, East Africa) typically complete in seconds to minutes on a standard laptop with 16 GB RAM.

Chunk Size vs. Performance Trade-offs

Chunk Size Memory per Chunk I/O Operations Speed
200×200 ~2 GB Many (slow I/O) Slower
500×500 ~8 GB Moderate Balanced
800×800 ~20 GB Few (fast I/O) Faster

Recommendation: Use the largest chunk size your system can handle for best performance.

Scalability

Horizontal Scaling

  • dask_processor runs tiles concurrently on a local cluster, with resume
  • Grid cells are independent, so tiles need no coordination
  • Batch processing capability

Vertical Scaling

  • Chunked I/O for huge datasets
  • Zarr region writes for concurrent output
  • Memory-mapped arrays for extreme sizes

Future Optimizations

  1. GPU Acceleration - CuPy for distribution fitting
  2. Parallel Processing - Multiprocessing for annual statistics
  3. Caching - Redis for parameter storage
  4. Streaming - Process data as it arrives

Code Style

Follows:

  • PEP 8 conventions
  • NumPy docstring format
  • Type hints (future)

Standards:

  • Functions: snake_case
  • Classes: PascalCase (minimal usage)
  • Constants: UPPER_CASE
  • Private: _leading_underscore

Version Control

Current: 2026.1 (January 2026)

Key Changes:

  • 2026.1: Complete package with SPI/SPEI, run theory, and dual magnitude

See: CHANGELOG for details

References

  • McKee, T.B., et al. (1993). SPI methodology
  • Vicente-Serrano, S.M., et al. (2010). SPEI methodology
  • Yevjevich, V. (1967). Run theory
  • Adams, J. (2017). Original climate-indices package

Contributing

To extend functionality:

  1. Add function to appropriate module
  2. Export in __init__.py
  3. Document in user guide
  4. Add test case
  5. Update CHANGELOG

Maintain:

  • CF Convention compliance
  • NumPy vectorization
  • Clear documentation
  • Consistent API

See Also

Back to top