Downsample zarr

Demonstrates downsampling zarr using xarray and dask with data from SMAP
Author

Aimee Barciauskas

Published

January 3, 2023

Run this notebook

You can launch this notebook in VEDA JupyterHub by clicking the link below.

Launch in VEDA JupyterHub (requires access)

Learn more

Inside the Hub

This notebook was written on the VEDA JupyterHub and as such is designed to be run on a jupyterhub which is associated with an AWS IAM role which has been granted permissions to the VEDA data store via its bucket policy. The instance used provided 16GB of RAM.

See (VEDA Analytics JupyterHub Access)[https://nasa-impact.github.io/veda-docs/veda-jh-access.html] for information about how to gain access.

Outside the Hub

The data is in a protected bucket. Please request access by emailng aimee@developmentseed.org or alexandra@developmentseed.org and providing your affiliation, interest in or expected use of the dataset and an AWS IAM role or user Amazon Resource Name (ARN). The team will help you configure the cognito client.

You should then run:

%run -i 'cognito_login.py'

Approach

This notebook demonstrates 2 strategies for to subselect data from a Zarr dataset in order to visualize using the memory of a notebook.

  1. Downsample the temporal resolution of the data using xarray.DataArray.resample
  2. Coarsening the spatial aspect of the data using xarray.DataArray.coarsen

A strategy for visualizing any large amount of data is Datashader which bins data into a fixed 2-D array. The call to rasterize ensures the use of the datashader library to bin the data.

About the data

The SMAP mission is an orbiting observatory that measures the amount of water in the surface soil everywhere on Earth.

Load libraries

!pip install -q hvplot==0.8.4
import xarray as xr
import hvplot.xarray
import geoviews as gv
import datashader as dsh
from holoviews.operation.datashader import rasterize

gv.output(size=300)

Optional: Create and Scale a Dask Cluster

We create a separate Dask cluster to speed up reprojecting the data (and other potential computations which could be required and are parallelizable).

Note if you skip this cell you will still be using Dask, you’ll just be using the machine where you are running this notebook.

from dask_gateway import GatewayCluster, Gateway

gateway = Gateway()
clusters = gateway.list_clusters()

# connect to an existing cluster - this is useful when the kernel shutdown in the middle of an interactive session
if clusters:
    cluster = gateway.connect(clusters[0].name)
else:
    cluster = GatewayCluster(shutdown_on_close=True)

cluster.scale(16)
client = cluster.get_client()
client

Client

Client-2c42768f-d18f-11ee-89f9-02b3484b28f4

Connection method: Cluster object Cluster type: dask_gateway.GatewayCluster
Dashboard: /services/dask-gateway/clusters/prod.ce053103b0e44d4d9c26bd8f50584072/status

Cluster Info

GatewayCluster

Open the dataset from S3

ds = xr.open_zarr("s3:///veda-data-store-staging/EIS/zarr/SPL3SMP.zarr")
ds
<xarray.Dataset>
Dimensions:                        (northing_m: 406, easting_m: 964,
                                    datetime: 1679)
Coordinates:
  * datetime                       (datetime) datetime64[ns] 2018-01-01 ... 2...
  * easting_m                      (easting_m) float64 -1.735e+07 ... 1.735e+07
  * northing_m                     (northing_m) float64 7.297e+06 ... -7.297e+06
Data variables: (12/26)
    albedo                         (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    albedo_pm                      (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    bulk_density                   (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    bulk_density_pm                (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    clay_fraction                  (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    clay_fraction_pm               (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    ...                             ...
    static_water_body_fraction     (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    static_water_body_fraction_pm  (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    surface_flag                   (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    surface_flag_pm                (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    surface_temperature            (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    surface_temperature_pm         (northing_m, easting_m, datetime) float32 dask.array<chunksize=(100, 100, 100), meta=np.ndarray>
    • datetime
      PandasIndex
      PandasIndex(DatetimeIndex(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-04',
               '2018-01-05', '2018-01-06', '2018-01-07', '2018-01-08',
               '2018-01-09', '2018-01-10',
               ...
               '2022-08-31', '2022-09-01', '2022-09-02', '2022-09-03',
               '2022-09-04', '2022-09-05', '2022-09-06', '2022-09-07',
               '2022-09-08', '2022-09-09'],
              dtype='datetime64[ns]', name='datetime', length=1679, freq=None))
    • easting_m
      PandasIndex
      PandasIndex(Index([      -17349514.34,       -17313482.12,        -17277449.9,
             -17241417.68,       -17205385.46,       -17169353.24,
             -17133321.02,        -17097288.8,       -17061256.58,
             -17025224.36,
       ...
       17025223.540000003,        17061255.76,        17097287.98,
               17133320.2, 17169352.419999998, 17205384.640000004,
       17241416.860000003, 17277449.080000002,         17313481.3,
              17349513.52],
      dtype='float64', name='easting_m', length=964))
    • northing_m
      PandasIndex
      PandasIndex(Index([        7296524.72,          7260492.5,  7224460.279999999,
               7188428.06,         7152395.84,         7116363.62,
        7080331.399999999,         7044299.18,         7008266.96,
               6972234.74,
       ...
       -6972234.400000001,        -7008266.62, -7044298.840000001,
       -7080331.060000001,        -7116363.28, -7152395.500000001,
       -7188427.720000002,        -7224459.94, -7260492.160000001,
              -7296524.38],
      dtype='float64', name='northing_m', length=406))

    • Select the variable of interest (soil moisture for this example).

    soil_moisture = ds.soil_moisture
soil_moisture
    <xarray.DataArray 'soil_moisture' (northing_m: 406, easting_m: 964,
                                   datetime: 1679)>
dask.array<open_dataset-soil_moisture, shape=(406, 964, 1679), dtype=float32, chunksize=(100, 100, 100), chunktype=numpy.ndarray>
Coordinates:
  * datetime    (datetime) datetime64[ns] 2018-01-01 2018-01-02 ... 2022-09-09
  * easting_m   (easting_m) float64 -1.735e+07 -1.731e+07 ... 1.735e+07
  * northing_m  (northing_m) float64 7.297e+06 7.26e+06 ... -7.26e+06 -7.297e+06
Attributes:
    long_name:  Representative DCA soil moisture measurement for the Earth ba...
    units:      cm**3/cm**3
    valid_max:  0.5
    valid_min:  0.019999999552965164

    Strategy 1: Downsample the temporal resolution of the data

    To plot one day from every month, resample the data to 1 observation a month.

    somo_one_month = soil_moisture.resample(datetime="1ME").nearest()

    Quick plot

    We can generate a quick plot using hvplot and datashader.

    # workaround to avoid warnings that are triggered within Dask.
import warnings

warnings.filterwarnings(
    "ignore", message="All-NaN slice encountered", category=RuntimeWarning
)
    somo_one_month.hvplot(
    x="easting_m",
    y="northing_m",
    groupby="datetime",
    crs="EPSG:6933",
    coastline=True,
    rasterize=True,
    aggregator="mean",
    frame_height=150,
    widget_location="bottom",
)

    Reproject before plotting

    Reproject the data for map visualization.

    somo_one_month = somo_one_month.transpose("datetime", "northing_m", "easting_m")
somo_one_month = somo_one_month.rio.set_spatial_dims(
    x_dim="easting_m", y_dim="northing_m"
)
somo_one_month = somo_one_month.rio.write_crs("EPSG:6933")
somo_reprojected = somo_one_month.rio.reproject("EPSG:4326")
somo_reprojected
    <xarray.DataArray 'soil_moisture' (datetime: 57, y: 1046, x: 2214)>
array([[[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
...
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]]], dtype=float32)
Coordinates:
  * x            (x) float64 -179.9 -179.8 -179.6 -179.4 ... 179.6 179.8 179.9
  * y            (y) float64 84.96 84.8 84.64 84.48 ... -84.64 -84.8 -84.96
  * datetime     (datetime) datetime64[ns] 2018-01-31 2018-02-28 ... 2022-09-30
    spatial_ref  int64 0
Attributes:
    long_name:  Representative DCA soil moisture measurement for the Earth ba...
    units:      cm**3/cm**3
    valid_max:  0.5
    valid_min:  0.019999999552965164

    Note that this is now a fully materialized data array - when we reproject we trigger an implicit compute.

    Create a geoviews dataset and visualize the data on a map.

    kdims = ["datetime", "x", "y"]
vdims = ["soil_moisture"]
xr_dataset = gv.Dataset(somo_reprojected, kdims=kdims, vdims=vdims)
images = xr_dataset.to(gv.Image, ["x", "y"])

rasterize(
    images, precompute=True, aggregator=dsh.mean("soil_moisture")
) * gv.feature.coastline
    /srv/conda/envs/notebook/lib/python3.11/site-packages/holoviews/core/data/xarray.py:340: UserWarning: The `squeeze` kwarg to GroupBy is being removed.Pass .groupby(..., squeeze=False) to disable squeezing, which is the new default, and to silence this warning.
  for k, v in dataset.data.groupby(index_dims[0].name):

    Strategy 2: Coarsen spatial resolution of the data

    Below, we coarsen the spatial resolution of the data by a factor of 4 in the x and 2 in the y. These values were chosen because they can be used with the exact boundary argument as the dimensions size is a multiple of these values.

    You can also coarsen by datetime, using the same strategy as below but replacing easting_m and northing_m with datetime. If {datetime: n} is the value give to the dim argument, this would create a mean of the soil moisture average for n days.

    Once the data has been coarsned, again it is reprojected for map visualization and then visualized using Geoviews.

    coarsened = soil_moisture.coarsen(dim={"easting_m": 4, "northing_m": 2}).mean()

coarsened = coarsened.transpose("datetime", "northing_m", "easting_m")
coarsened = coarsened.rio.set_spatial_dims(x_dim="easting_m", y_dim="northing_m")
coarsened = coarsened.rio.write_crs("epsg:6933")
coarsened_reprojected = coarsened.rio.reproject("EPSG:4326")
coarsened_reprojected
    <xarray.DataArray 'soil_moisture' (datetime: 1679, y: 315, x: 667)>
array([[[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
...
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]],

       [[nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        ...,
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan],
        [nan, nan, nan, ..., nan, nan, nan]]], dtype=float32)
Coordinates:
  * x            (x) float64 -179.7 -179.2 -178.7 -178.1 ... 178.6 179.2 179.7
  * y            (y) float64 84.77 84.23 83.7 83.16 ... -83.64 -84.18 -84.72
  * datetime     (datetime) datetime64[ns] 2018-01-01 2018-01-02 ... 2022-09-09
    spatial_ref  int64 0
Attributes:
    long_name:   Representative DCA soil moisture measurement for the Earth b...
    units:       cm**3/cm**3
    valid_max:   0.5
    valid_min:   0.019999999552965164
    _FillValue:  3.402823466e+38
    kdims = ["datetime", "x", "y"]
vdims = ["soil_moisture"]
xr_dataset = gv.Dataset(coarsened_reprojected, kdims=kdims, vdims=vdims)
images = xr_dataset.to(gv.Image, ["x", "y"])
    /srv/conda/envs/notebook/lib/python3.11/site-packages/holoviews/core/data/xarray.py:340: UserWarning: The `squeeze` kwarg to GroupBy is being removed.Pass .groupby(..., squeeze=False) to disable squeezing, which is the new default, and to silence this warning.
  for k, v in dataset.data.groupby(index_dims[0].name):
    rasterize(
    images, precompute=True, aggregator=dsh.mean("soil_moisture")
) * gv.feature.coastline

    Cleanup

    When using a remote Dask cluster it is recommented to explicitly close the cluster.

    client.shutdown()