Calculating Standardized Precipitation Index (SPI)

One of the key advantages of the XEE package is the ability to use Python packages to work with data from Google Earth Engine. This package allows you to easily extract data from the vast data catalog of GEE and use widely used implementations available as Python libraries.

Overview of the Task

This notebook shows how to extract a precipitation time-series images from Earth Engine and then use the Climate Indicies Python package to calculate the Calculating Standardized Precipitation (SPI).

Input Layers:

• ne_10m_admin_0_countries_ind.zip: A shapefile of country boundaries

Output Layers:

• 2024-01.tif, 2024-02.tif …. , 2024-12.tif : GeoTIFF files of monthly SPI for the selected country.

Data Credit:

• Made with Natural Earth. Free vector and raster map data @ naturalearthdata.com.

• Muñoz Sabater, J., (2019): ERA5-Land monthly averaged data from 1981 to present. Copernicus Climate Change Service (C3S) Climate Data Store (CDS). (2026-06-02), doi:10.24381/cds.68d2bb30

Running the Notebook: The preferred way to run this notebook is on Google Colab.

Setup and Data Download

The following blocks of code will install the required packages and download the datasets to your Colab environment.

%%capture
if 'google.colab' in str(get_ipython()):
    !pip install 'climate-indices>=3.0.0' xee rioxarray dask['distributed']

import datetime
import ee
import geopandas as gpd
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import rioxarray as rxr
import xarray as xr
from climate_indices import spi
from climate_indices.exceptions import GoodnessOfFitWarning, MissingDataWarning
from climate_indices.indices import Distribution
from xee import helpers
import warnings

# quiet the INFO log messages from climate_indices
logging.getLogger('climate_indices').setLevel(logging.WARNING)

data_folder = 'data'
output_folder = 'output'

if not os.path.exists(data_folder):
    os.mkdir(data_folder)
if not os.path.exists(output_folder):
    os.mkdir(output_folder)

def download(url):
    filename = os.path.join(data_folder, os.path.basename(url))
    if not os.path.exists(filename):
        from urllib.request import urlretrieve
        local, _ = urlretrieve(url, filename)
        print('Downloaded ' + local)

data_url = 'https://naciscdn.org/naturalearth/10m/cultural/'


shapefile = 'ne_10m_admin_0_countries_ind.zip'

download(data_url + shapefile)

Initialize EE and Dask Cluster

Initialize EE with the High-Volume EndPoint which is recommended to be used with XEE for workflows that do not use a lot of server side processing and are primarily for extracting data from stored collections. Replace the value of the cloud_project variable with your own project id that is linked with GEE.

cloud_project = 'spatialthoughts' # replace with your project id

try:
    ee.Initialize(
        project=cloud_project,
        opt_url='https://earthengine-highvolume.googleapis.com')
except:
    ee.Authenticate()
    ee.Initialize(
        project=cloud_project,
        opt_url='https://earthengine-highvolume.googleapis.com')

Setup a local Dask cluster.

from dask.distributed import Client

client = Client()  # set up local cluster on the machine
client

If you are running this notebook in Colab, you will need to create and use a proxy URL to see the dashboard running on the local server.

if 'google.colab' in str(get_ipython()):
    from google.colab import output
    port_to_expose = 8787  # This is the default port for Dask dashboard
    print(output.eval_js(f'google.colab.kernel.proxyPort({port_to_expose})'))

Each of our Dask workers need Earth Engine authentication. Initialize Dask workers using ee.Initialize().

from dask.distributed import WorkerPlugin

class EEPlugin(WorkerPlugin):
    def __init__(self):
        pass
    def setup(self, worker):
        self.worker = worker
        try:
            ee.Initialize(
                project=cloud_project,
                opt_url='https://earthengine-highvolume.googleapis.com')
        except:
            ee.Authenticate()
            ee.Initialize(
                project=cloud_project,
                opt_url='https://earthengine-highvolume.googleapis.com')

ee_plugin = EEPlugin()
client.register_plugin(ee_plugin)

Procedure

Load the Natural Earth Admin0 boundaries and extract the polygon for India.

shapefile_path = os.path.join(data_folder, shapefile)
gdf = gpd.read_file(shapefile_path, encoding='utf-8')
country = gdf[gdf['SOV_A3'] == 'IND']
geometry = country.geometry.union_all()

Now we have an ImageCollection that we want to get it as a XArray Dataset. We define the region of interest and extract the ImageCollection using the ‘ee’ engine.

era5 = ee.ImageCollection('ECMWF/ERA5_LAND/MONTHLY_AGGR')
filtered = era5 \
  .filter(ee.Filter.date('1995-01-01', '2025-01-01')) \
  .select('total_precipitation_sum')

We now read the filtered collecting using XEE. XEE requires explicit grid parameters. We extract these using the helper function extract_grid_params.

grid_params = helpers.extract_grid_params(filtered)
grid_params

For SPI processing, we need the whole time-series in a single chunk, so we set the time: -1.

ds = xr.open_dataset(
    filtered,
    engine='ee',
    **grid_params,
    chunks={'time': -1, 'y': 128, 'x': 128}
)
ds

Clip the pixels outside the geometry.

clipped_ds = ds \
  .rio.clip(country.geometry.values)
clipped_ds

Check the size of the resulting dataset.

print(f'DataSet size: {clipped_ds.nbytes/1e6:.2f} MB.')

Select the total_precipitation_sum variable. Many XArray functions require all the dimensions to be sorted in ascending order. Make sure y and x are sorted.

We also re-chunk the data into smaller spatial chunks for the SPI calculation. climate_indices processes an entire chunk in one vectorized call and needs roughly 50 KB of memory per pixel, so the recommended budget is about 1,600 pixels per chunk. Smaller chunks also give the Dask workers more tasks to run in parallel. We keep the larger 128 x 128 chunks when reading the data, as they require fewer requests to Earth Engine.

clipped_da = clipped_ds.total_precipitation_sum
clipped_da = clipped_da \
  .sortby(['y', 'x']) \
  .chunk({'time': -1, 'y': 32, 'x': 32})

Configure and run the SPI calculation.

Starting with v3.0.0, climate_indices accepts XArray DataArrays directly and is vectorized over space. Instead of fitting the distribution one pixel at a time in a Python loop (xr.apply_ufunc() with vectorize=True), it fits a whole spatial chunk in a single NumPy call. With a Dask-backed input, it runs once per chunk in parallel on our Dask workers. The data_start_year and periodicity are inferred from the time coordinate.

scale = 3
distribution = Distribution.gamma
calibration_year_initial = 1995
calibration_year_final = 2024

da_spi = spi(
    values=clipped_da,
    scale=scale,
    distribution=distribution,
    calibration_year_initial=calibration_year_initial,
    calibration_year_final=calibration_year_final,
)
da_spi

%%time
da_spi = da_spi.compute()

Visualize the Results

We use a faceted plot to visualize all months of a selected year.

# select a year to visualize
selected = da_spi.sel(time='2024')

legend_levels = [-3,-2,-1,0,1,2,3]
selected.plot(cmap='RdBu', col='time', col_wrap=4, levels=legend_levels)
plt.show()

Save Output as GeoTIFFs

We can use rioxarray to clip and save the resulting arrays as GeoTIFF files.

# select a year to save
selected = da_spi.sel(time='2024')

for time in selected.time.values:
    date_string = np.datetime_as_string(time, unit='M')
    image = selected.sel(time=time) \
      .rio.write_crs('EPSG:4326') \
      .rio.clip(country.geometry.values)
    output_file = f'{date_string}.tif'
    output_path = os.path.join(output_folder, output_file)
    image.rio.to_raster(output_path, driver='COG')
    print(f'Saved {output_path}')

---

If you want to give feedback or share your experience with this tutorial, please comment below. (requires GitHub account)