In [1]:

import numpy as np
import xarray as xr

import rasterio.features
import stackstac
import pystac_client
import planetary_computer

import xrspatial.multispectral as ms

In [3]:

# from dask_gateway import GatewayCluster

# # just for speed. This runs fine on a single machine.
# cluster = GatewayCluster()
# cluster.scale(24)
# client = cluster.get_client()
# cluster

Problem setup¶

Our task is to create a cloud-free composite of some Sentinel-2 imagery by taking a median over time. The "easy" way of doing things is to use stackstac to create the DataArray and call .median(dim="time").

In [92]:

area_of_interest = {
    "type": "Polygon",
    "coordinates": [
        [
            [-122.27508544921875, 47.54687159892238],
            [-121.96128845214844, 47.54687159892238],
            [-121.96128845214844, 47.745787772920934],
            [-122.27508544921875, 47.745787772920934],
            [-122.27508544921875, 47.54687159892238],
        ]
    ],
}
bbox = rasterio.features.bounds(area_of_interest)

stac = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

search = stac.search(
    bbox=bbox,
#     datetime="2018-01-01/2020-12-31",
    datetime="2020-06-01/2020-12-31",
    collections=["sentinel-2-l2a"],
    limit=500,  # fetch items in batches of 500
    query={"eo:cloud_cover": {"lt": 25}},
)

items = list(search.get_items())
print(len(items))

signed_items = [planetary_computer.sign(item).to_dict() for item in items]

bad_data = (
    stackstac.stack(
        signed_items,
        assets=["B04", "B03", "B02"],  # red, green, blue
        chunksize=2048,
    )
    .where(lambda x: x > 0, other=np.nan)  # sentinel-2 uses 0 as nodata
#     .assign_coords(band=lambda x: x.common_name.rename("band"))  # use common names
)
bad_data

Out[92]:

<xarray.DataArray 'stackstac-6add7ccc542226180faed750e3208a58' (time: 26, band: 3, y: 10980, x: 10980)>
dask.array<where, shape=(26, 3, 10980, 10980), dtype=float64, chunksize=(1, 1, 2048, 2048), chunktype=numpy.ndarray>
Coordinates:
  * time                                     (time) datetime64[ns] 2020-06-26...
    id                                       (time) <U54 'S2A_MSIL2A_20200626...
  * band                                     (band) <U3 'B04' 'B03' 'B02'
  * x                                        (x) float64 5e+05 ... 6.098e+05
  * y                                        (y) float64 5.3e+06 ... 5.19e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    s2:generation_time                       (time) <U24 '2020-08-24T04:44:17...
    s2:not_vegetated_percentage              (time) float64 6.29 6.48 ... 1.764
    s2:cloud_shadow_percentage               (time) float64 0.095 ... 8.791
    s2:snow_ice_percentage                   (time) float64 1.324 ... 1.404
    sat:orbit_state                          <U10 'descending'
    s2:thin_cirrus_percentage                (time) float64 5.034 ... 0.1651
    s2:nodata_pixel_percentage               (time) float64 3.652 ... 3.947
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:unclassified_percentage               (time) float64 1.278 0.992 ... 8.58
    s2:water_percentage                      (time) float64 10.82 ... 9.798
    s2:medium_proba_clouds_percentage        (time) float64 0.366 ... 6.466
    s2:product_uri                           (time) <U65 'S2A_MSIL2A_20200626...
    platform                                 (time) <U11 'Sentinel-2A' ... 'S...
    eo:cloud_cover                           (time) float64 6.042 ... 22.94
    s2:high_proba_clouds_percentage          (time) float64 0.6417 ... 16.31
    s2:granule_id                            (time) <U62 'S2A_OPER_MSI_L2A_TL...
    s2:dark_features_percentage              (time) float64 0.5517 ... 14.56
    sat:relative_orbit                       (time) int64 13 56 56 ... 56 56 13
    s2:mean_solar_zenith                     (time) float64 26.92 ... 71.85
    s2:datatake_type                         <U8 'INS-NOBS'
    s2:vegetation_percentage                 (time) float64 73.6 68.68 ... 32.16
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    s2:reflectance_conversion_factor         (time) float64 0.968 ... 1.034
    s2:datastrip_id                          (time) <U64 'S2A_OPER_MSI_L2A_DS...
    s2:mean_solar_azimuth                    (time) float64 148.0 ... 165.6
    s2:datatake_id                           (time) <U34 'GS2A_20200626T18592...
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    (band) <U20 'Band 4 - Red - 10m'...
    gsd                                      int64 10
    common_name                              (band) <U5 'red' 'green' 'blue'
    center_wavelength                        (band) float64 0.665 0.56 0.49
    full_width_half_max                      (band) float64 0.038 0.045 0.098
    epsg                                     int64 32610
Attributes:
    spec:        RasterSpec(epsg=32610, bounds=(499980.0, 5190240.0, 609780.0...
    crs:         epsg:32610
    transform:   | 10.00, 0.00, 499980.00|\n| 0.00,-10.00, 5300040.00|\n| 0.0...
    resolution:  10.0

xarray.DataArray

'stackstac-6add7ccc542226180faed750e3208a58'

time: 26
band: 3
y: 10980
x: 10980

dask.array<chunksize=(1, 1, 2048, 2048), meta=np.ndarray>

	Array	Chunk
Bytes	70.06 GiB	32.00 MiB
Shape	(26, 3, 10980, 10980)	(1, 1, 2048, 2048)
Count	8616 Tasks	2808 Chunks
Type	float64	numpy.ndarray

Coordinates: (46)

Attributes: (4)
spec :
RasterSpec(epsg=32610, bounds=(499980.0, 5190240.0, 609780.0, 5300040.0), resolutions_xy=(10.0, 10.0))
crs :
epsg:32610
transform :
| 10.00, 0.00, 499980.00| | 0.00,-10.00, 5300040.00| | 0.00, 0.00, 1.00|
resolution :
10.0

Unfortunately, this median computation requires an expensive shuffle. Currently we have chunk sizes of (time=1, band=1, y=2048, x=2048). To compute the median, all of the values in time need to be in the same chunk. Dask will automatically rechunk to something like (time=36, band=1, y=256, x=256). This means we won't blow up memory on any single task, but it requires shuffling a bunch of data, which can be dangerous to your cluster's health.

In [93]:

bad_result = bad_data.median(dim="time")
bad_result

Out[93]:

<xarray.DataArray 'stackstac-6add7ccc542226180faed750e3208a58' (band: 3, y: 10980, x: 10980)>
dask.array<nanmedian, shape=(3, 10980, 10980), dtype=float64, chunksize=(1, 512, 512), chunktype=numpy.ndarray>
Coordinates:
  * band                                     (band) <U3 'B04' 'B03' 'B02'
  * x                                        (x) float64 5e+05 ... 6.098e+05
  * y                                        (y) float64 5.3e+06 ... 5.19e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    sat:orbit_state                          <U10 'descending'
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:datatake_type                         <U8 'INS-NOBS'
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    (band) <U20 'Band 4 - Red - 10m'...
    gsd                                      int64 10
    common_name                              (band) <U5 'red' 'green' 'blue'
    center_wavelength                        (band) float64 0.665 0.56 0.49
    full_width_half_max                      (band) float64 0.038 0.045 0.098
    epsg                                     int64 32610

xarray.DataArray

'stackstac-6add7ccc542226180faed750e3208a58'

band: 3
y: 10980
x: 10980

dask.array<chunksize=(1, 512, 512), meta=np.ndarray>

	Array	Chunk
Bytes	2.69 GiB	2.00 MiB
Shape	(3, 10980, 10980)	(1, 512, 512)
Count	22440 Tasks	1452 Chunks
Type	float64	numpy.ndarray

Coordinates: (22)
- band
  (band)
  <U3
  'B04' 'B03' 'B02'
```
array(['B04', 'B03', 'B02'], dtype='<U3')
```
- x
  (x)
  float64
  5e+05 5e+05 ... 6.098e+05 6.098e+05
```
array([499980., 499990., 500000., ..., 609750., 609760., 609770.])
```
- y
  (y)
  float64
  5.3e+06 5.3e+06 ... 5.19e+06
```
array([5300040., 5300030., 5300020., ..., 5190270., 5190260., 5190250.])
```
- s2:mgrs_tile
  ()
  <U5
  '10TET'
```
array('10TET', dtype='<U5')
```
- instruments
  ()
  <U3
  'msi'
```
array('msi', dtype='<U3')
```
- proj:epsg
  ()
  int64
  32610
```
array(32610)
```
- sat:orbit_state
  ()
  <U10
  'descending'
```
array('descending', dtype='<U10')
```
- s2:degraded_msi_data_percentage
  ()
  float64
  0.0
```
array(0.)
```
- s2:processing_baseline
  ()
  <U5
  '02.12'
```
array('02.12', dtype='<U5')
```
- s2:product_type
  ()
  <U7
  'S2MSI2A'
```
array('S2MSI2A', dtype='<U7')
```
- s2:datatake_type
  ()
  <U8
  'INS-NOBS'
```
array('INS-NOBS', dtype='<U8')
```
- constellation
  ()
  <U10
  'Sentinel 2'
```
array('Sentinel 2', dtype='<U10')
```
- s2:saturated_defective_pixel_percentage
  ()
  float64
  0.0
```
array(0.)
```
- proj:transform
  ()
  object
  {0.0, 5300040.0, 10.0, 499980.0,...
```
array({0.0, 5300040.0, 10.0, 499980.0, -10.0}, dtype=object)
```
- proj:bbox
  ()
  object
  {5190240.0, 609780.0, 499980.0, ...
```
array({5190240.0, 609780.0, 499980.0, 5300040.0}, dtype=object)
```
- proj:shape
  ()
  object
  {10980}
```
array({10980}, dtype=object)
```
- title
  (band)
  <U20
  'Band 4 - Red - 10m' ... 'Band 2...
```
array(['Band 4 - Red - 10m', 'Band 3 - Green - 10m',
       'Band 2 - Blue - 10m'], dtype='<U20')
```
- gsd
  ()
  int64
  10
```
array(10)
```
- common_name
  (band)
  <U5
  'red' 'green' 'blue'
```
array(['red', 'green', 'blue'], dtype='<U5')
```
- center_wavelength
  (band)
  float64
  0.665 0.56 0.49
```
array([0.665, 0.56 , 0.49 ])
```
- full_width_half_max
  (band)
  float64
  0.038 0.045 0.098
```
array([0.038, 0.045, 0.098])
```
- epsg
  ()
  int64
  32610
```
array(32610)
```
Attributes: (0)

As a secondary issue, this task graph is also relatively complex. The bad_result DataArray has ~25,000 tasks, which reflects the intermediate rechunk operations. Can we do better?

Rechunking "pushdown" / fusion¶

What if we loaded the data how we want it, rather than rechunking after the fact? GDAL / COGs support windowed reads, it's technically feasible to "push" the rechunking operation into the I/O operation.

In the "old" workflow, each task roughly looks like

Load data from a single file
Rechunk that data it into smaller pieces:
- Some small pieces were sent to other workers; Some small pieces were kept
- Some small pieces were kept locally, to build a contiguous "timeseries" for the small window by requesting data from other workers and concatenating it.

In this proposed workflow, each data reading task will instead look like:

Load data from many files; Reading a small window from each file in the timeseries and concatenating the results.

And that's it. No rechunking, no communication and (hopefully) no memory issues.

Unforunately, the user-level code for this is much more complicated than a simple stackstac.median! I'm hopeful that Dask could become aware of this type of optimization in the future, to push this kind of rechunk down to the I/O layer if the I/O layer supports it (COG does; I think there's proposed work in Zarr to support something like this). Then we might be able to have our cake (store a single copy of the data) and eat it too (fast-ish access for both spatial and timeseries analysis).

The next section implements this stragegy. If you're just intersted in the result, I recommend skipping to comparison below.

First we need the slices. We know we want chunks of size something like (len(items), 1, ..., ...). We can let Dask figure that out for us.

In [94]:

import dask.array

chunks = dask.array.empty(bad_data.shape, chunks=(-1, 1, "auto", "auto"))[0, 0].chunks
slices = dask.array.core.slices_from_chunks(chunks)
slices[:2]

Out[94]:

[(slice(0, 732, None), slice(0, 732, None)),
 (slice(0, 732, None), slice(732, 1464, None))]

We have a few ways to handle reading the actual data. We could use stackstac, rioxarray, rasterio directly, etc. We'll use stackstac since it will handle the time dimension and other metadata for us. It supports cropping with a bounds argument, so we need to find the bounding box (in native CRS) for each window.

In [95]:

bounds = []
for x_slice, y_slice in slices:
    x = bad_data.x[x_slice]
    y = bad_data.y[y_slice]
    # TODO: our meta is off by 1; maybe backwards?
    x_range = x.min().item(), x.max().item() + 0.1
    y_range = y.min().item(), y.max().item() + 0.1
    bounds.append((x_range[0], y_range[0], x_range[1] + 1, y_range[1] + 1))

Now we need a DataArray per window. This is a bit awkward, since the data will be coming from a lazy / delayed Dask Array, but the metadata needs to be computed eagerly. This could be improved, but seems to work.

In [132]:

@dask.delayed
def read_window(stac_items, asset, bounds):
    return stackstac.stack(stac_items.to_dict()["features"], assets=[asset], bounds=bounds).compute(scheduler="single-threaded").where(lambda x: x > 0).data

def read_metadata(stac_items, asset, bounds):
    return stackstac.stack(stac_items.to_dict()["features"], assets=[asset], bounds=bounds).where(lambda x: x > 0)

In [133]:

import pystac
%time ic = pystac.ItemCollection(signed_items, clone_items=True)

CPU times: user 23.2 ms, sys: 7.57 ms, total: 30.8 ms
Wall time: 28.6 ms

Next, let's get the List of DataArrays. We'll have one DataArray per "timeseries" / window.

This section needs to be optimzied / parallelized. We're nominally just operating on metadata here, but it's still taking too long for interactive analysis.

In [161]:

metas[0].time

In [134]:

%%time
import dask.array as da
import xarray as xr

# might want to parallize this
timeseries = []
assets = ["B04", "B03", "B02"]

for bound in bounds:
    metas = [read_metadata(ic, asset, bound) for asset in assets]
    data = [read_window(ic, asset, bound) for asset in assets]
    xarrays = []
    for m, d in zip(metas, data):
        xarrays.append(xr.DataArray(
            da.from_delayed(d, shape=m.shape, dtype=m.dtype),
            coords=m.coords,
            dims=m.dims,
            attrs=m.attrs,
        ))
    ts = xr.concat(xarrays, dim="band")
    timeseries.append(ts)

CPU times: user 23.6 s, sys: 1.17 s, total: 24.8 s
Wall time: 23.2 s

Now that we have a List[DataArray], we need to assemble them into a single large DataArray. To do that, we'll make a series of "strips" where each strip shares the same x coordinates. We'll then concatenate those strips together to get our final DataArray.

In [218]:

stac_items = ic
asset = "B03"
m = stackstac.stack(stac_items.to_dict()["features"], assets=[asset], bounds=bounds[0]).where(lambda x: x > 0)
m

Out[218]:

<xarray.DataArray 'stackstac-e3d8fc4409a42a643b4985fd3694f302' (time: 26, band: 1, y: 732, x: 732)>
dask.array<where, shape=(26, 1, 732, 732), dtype=float64, chunksize=(1, 1, 732, 732), chunktype=numpy.ndarray>
Coordinates:
  * time                                     (time) datetime64[ns] 2020-06-26...
    id                                       (time) <U54 'S2A_MSIL2A_20200626...
  * band                                     (band) <U3 'B03'
  * x                                        (x) float64 5e+05 ... 5.073e+05
  * y                                        (y) float64 5.3e+06 ... 5.293e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    s2:generation_time                       (time) <U24 '2020-08-24T04:44:17...
    s2:not_vegetated_percentage              (time) float64 6.29 6.48 ... 1.764
    s2:cloud_shadow_percentage               (time) float64 0.095 ... 8.791
    s2:snow_ice_percentage                   (time) float64 1.324 ... 1.404
    sat:orbit_state                          <U10 'descending'
    s2:thin_cirrus_percentage                (time) float64 5.034 ... 0.1651
    s2:nodata_pixel_percentage               (time) float64 3.652 ... 3.947
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:unclassified_percentage               (time) float64 1.278 0.992 ... 8.58
    s2:water_percentage                      (time) float64 10.82 ... 9.798
    s2:medium_proba_clouds_percentage        (time) float64 0.366 ... 6.466
    s2:product_uri                           (time) <U65 'S2A_MSIL2A_20200626...
    platform                                 (time) <U11 'Sentinel-2A' ... 'S...
    eo:cloud_cover                           (time) float64 6.042 ... 22.94
    s2:high_proba_clouds_percentage          (time) float64 0.6417 ... 16.31
    s2:granule_id                            (time) <U62 'S2A_OPER_MSI_L2A_TL...
    s2:dark_features_percentage              (time) float64 0.5517 ... 14.56
    sat:relative_orbit                       (time) int64 13 56 56 ... 56 56 13
    s2:mean_solar_zenith                     (time) float64 26.92 ... 71.85
    s2:datatake_type                         <U8 'INS-NOBS'
    s2:vegetation_percentage                 (time) float64 73.6 68.68 ... 32.16
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    s2:reflectance_conversion_factor         (time) float64 0.968 ... 1.034
    s2:datastrip_id                          (time) <U64 'S2A_OPER_MSI_L2A_DS...
    s2:mean_solar_azimuth                    (time) float64 148.0 ... 165.6
    s2:datatake_id                           (time) <U34 'GS2A_20200626T18592...
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    <U20 'Band 3 - Green - 10m'
    gsd                                      int64 10
    common_name                              <U5 'green'
    center_wavelength                        float64 0.56
    full_width_half_max                      float64 0.045
    epsg                                     int64 32610
Attributes:
    spec:        RasterSpec(epsg=32610, bounds=(499980.0, 5292730.0, 507300.0...
    crs:         epsg:32610
    transform:   | 10.00, 0.00, 499980.00|\n| 0.00,-10.00, 5300050.00|\n| 0.0...
    resolution:  10.0

xarray.DataArray

'stackstac-e3d8fc4409a42a643b4985fd3694f302'

time: 26
band: 1
y: 732
x: 732

dask.array<chunksize=(1, 1, 732, 732), meta=np.ndarray>

	Array	Chunk
Bytes	106.29 MiB	4.09 MiB
Shape	(26, 1, 732, 732)	(1, 1, 732, 732)
Count	131 Tasks	26 Chunks
Type	float64	numpy.ndarray

Coordinates: (46)

Attributes: (4)
spec :
RasterSpec(epsg=32610, bounds=(499980.0, 5292730.0, 507300.0, 5300050.0), resolutions_xy=(10.0, 10.0))
crs :
epsg:32610
transform :
| 10.00, 0.00, 499980.00| | 0.00,-10.00, 5300050.00| | 0.00, 0.00, 1.00|
resolution :
10.0

In [220]:

bad_data[:, 0, :732, :732]

Out[220]:

<xarray.DataArray 'stackstac-6add7ccc542226180faed750e3208a58' (time: 26, y: 732, x: 732)>
dask.array<getitem, shape=(26, 732, 732), dtype=float64, chunksize=(1, 732, 732), chunktype=numpy.ndarray>
Coordinates:
  * time                                     (time) datetime64[ns] 2020-06-26...
    id                                       (time) <U54 'S2A_MSIL2A_20200626...
    band                                     <U3 'B04'
  * x                                        (x) float64 5e+05 ... 5.073e+05
  * y                                        (y) float64 5.3e+06 ... 5.293e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    s2:generation_time                       (time) <U24 '2020-08-24T04:44:17...
    s2:not_vegetated_percentage              (time) float64 6.29 6.48 ... 1.764
    s2:cloud_shadow_percentage               (time) float64 0.095 ... 8.791
    s2:snow_ice_percentage                   (time) float64 1.324 ... 1.404
    sat:orbit_state                          <U10 'descending'
    s2:thin_cirrus_percentage                (time) float64 5.034 ... 0.1651
    s2:nodata_pixel_percentage               (time) float64 3.652 ... 3.947
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:unclassified_percentage               (time) float64 1.278 0.992 ... 8.58
    s2:water_percentage                      (time) float64 10.82 ... 9.798
    s2:medium_proba_clouds_percentage        (time) float64 0.366 ... 6.466
    s2:product_uri                           (time) <U65 'S2A_MSIL2A_20200626...
    platform                                 (time) <U11 'Sentinel-2A' ... 'S...
    eo:cloud_cover                           (time) float64 6.042 ... 22.94
    s2:high_proba_clouds_percentage          (time) float64 0.6417 ... 16.31
    s2:granule_id                            (time) <U62 'S2A_OPER_MSI_L2A_TL...
    s2:dark_features_percentage              (time) float64 0.5517 ... 14.56
    sat:relative_orbit                       (time) int64 13 56 56 ... 56 56 13
    s2:mean_solar_zenith                     (time) float64 26.92 ... 71.85
    s2:datatake_type                         <U8 'INS-NOBS'
    s2:vegetation_percentage                 (time) float64 73.6 68.68 ... 32.16
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    s2:reflectance_conversion_factor         (time) float64 0.968 ... 1.034
    s2:datastrip_id                          (time) <U64 'S2A_OPER_MSI_L2A_DS...
    s2:mean_solar_azimuth                    (time) float64 148.0 ... 165.6
    s2:datatake_id                           (time) <U34 'GS2A_20200626T18592...
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    <U20 'Band 4 - Red - 10m'
    gsd                                      int64 10
    common_name                              <U5 'red'
    center_wavelength                        float64 0.665
    full_width_half_max                      float64 0.038
    epsg                                     int64 32610
Attributes:
    spec:        RasterSpec(epsg=32610, bounds=(499980.0, 5190240.0, 609780.0...
    crs:         epsg:32610
    transform:   | 10.00, 0.00, 499980.00|\n| 0.00,-10.00, 5300040.00|\n| 0.0...
    resolution:  10.0

xarray.DataArray

'stackstac-6add7ccc542226180faed750e3208a58'

time: 26
y: 732
x: 732

dask.array<chunksize=(1, 732, 732), meta=np.ndarray>

	Array	Chunk
Bytes	106.29 MiB	4.09 MiB
Shape	(26, 732, 732)	(1, 732, 732)
Count	8642 Tasks	26 Chunks
Type	float64	numpy.ndarray

Coordinates: (46)

Attributes: (4)
spec :
RasterSpec(epsg=32610, bounds=(499980.0, 5190240.0, 609780.0, 5300040.0), resolutions_xy=(10.0, 10.0))
crs :
epsg:32610
transform :
| 10.00, 0.00, 499980.00| | 0.00,-10.00, 5300040.00| | 0.00, 0.00, 1.00|
resolution :
10.0

In [135]:

import itertools

key = lambda arr: arr.x[0].item()

strips = [
    xr.concat(list(v), dim="y")
    for _, v in itertools.groupby(timeseries, key=key)
]
good_data = xr.concat(strips, dim="x")
good_data

Out[135]:

<xarray.DataArray 'from-value-5bb16d51b8ec4f402aed7f070a7fbdc8' (time: 26, band: 3, y: 10980, x: 10980)>
dask.array<concatenate, shape=(26, 3, 10980, 10980), dtype=float64, chunksize=(26, 1, 732, 732), chunktype=numpy.ndarray>
Coordinates:
  * time                                     (time) datetime64[ns] 2020-06-26...
    id                                       (time) <U54 'S2A_MSIL2A_20200626...
  * band                                     (band) <U3 'B04' 'B03' 'B02'
  * x                                        (x) float64 5e+05 ... 6.098e+05
  * y                                        (y) float64 5.3e+06 ... 5.19e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    s2:generation_time                       (time) <U24 '2020-08-24T04:44:17...
    s2:not_vegetated_percentage              (time) float64 6.29 6.48 ... 1.764
    s2:cloud_shadow_percentage               (time) float64 0.095 ... 8.791
    s2:snow_ice_percentage                   (time) float64 1.324 ... 1.404
    sat:orbit_state                          <U10 'descending'
    s2:thin_cirrus_percentage                (time) float64 5.034 ... 0.1651
    s2:nodata_pixel_percentage               (time) float64 3.652 ... 3.947
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:unclassified_percentage               (time) float64 1.278 0.992 ... 8.58
    s2:water_percentage                      (time) float64 10.82 ... 9.798
    s2:medium_proba_clouds_percentage        (time) float64 0.366 ... 6.466
    s2:product_uri                           (time) <U65 'S2A_MSIL2A_20200626...
    platform                                 (time) <U11 'Sentinel-2A' ... 'S...
    eo:cloud_cover                           (time) float64 6.042 ... 22.94
    s2:high_proba_clouds_percentage          (time) float64 0.6417 ... 16.31
    s2:granule_id                            (time) <U62 'S2A_OPER_MSI_L2A_TL...
    s2:dark_features_percentage              (time) float64 0.5517 ... 14.56
    sat:relative_orbit                       (time) int64 13 56 56 ... 56 56 13
    s2:mean_solar_zenith                     (time) float64 26.92 ... 71.85
    s2:datatake_type                         <U8 'INS-NOBS'
    s2:vegetation_percentage                 (time) float64 73.6 68.68 ... 32.16
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    s2:reflectance_conversion_factor         (time) float64 0.968 ... 1.034
    s2:datastrip_id                          (time) <U64 'S2A_OPER_MSI_L2A_DS...
    s2:mean_solar_azimuth                    (time) float64 148.0 ... 165.6
    s2:datatake_id                           (time) <U34 'GS2A_20200626T18592...
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    (band) <U20 'Band 4 - Red - 10m'...
    gsd                                      int64 10
    common_name                              (band) <U5 'red' 'green' 'blue'
    center_wavelength                        (band) float64 0.665 0.56 0.49
    full_width_half_max                      (band) float64 0.038 0.045 0.098
    epsg                                     int64 32610
Attributes:
    spec:        RasterSpec(epsg=32610, bounds=(499980.0, 5292730.0, 507300.0...
    crs:         epsg:32610
    transform:   | 10.00, 0.00, 499980.00|\n| 0.00,-10.00, 5300050.00|\n| 0.0...
    resolution:  10.0

xarray.DataArray

'from-value-5bb16d51b8ec4f402aed7f070a7fbdc8'

time: 26
band: 3
y: 10980
x: 10980

dask.array<chunksize=(26, 1, 732, 732), meta=np.ndarray>

	Array	Chunk
Bytes	70.06 GiB	106.29 MiB
Shape	(26, 3, 10980, 10980)	(26, 1, 732, 732)
Count	3375 Tasks	675 Chunks
Type	float64	numpy.ndarray

Coordinates: (46)

Attributes: (4)
spec :
RasterSpec(epsg=32610, bounds=(499980.0, 5292730.0, 507300.0, 5300050.0), resolutions_xy=(10.0, 10.0))
crs :
epsg:32610
transform :
| 10.00, 0.00, 499980.00| | 0.00,-10.00, 5300050.00| | 0.00, 0.00, 1.00|
resolution :
10.0

In [207]:

x = read_window(ic, "B02", bounds[0]).compute()

In [136]:

a = bad_data.data.rechunk((-1, 1, 732, 732)).blocks[0, 0, 0, 0]

In [137]:

b = good_data.data.blocks[0, 0, 0, 0]

In [139]:

%time aa = a.compute()
%time bb = b.compute()

CPU times: user 545 ms, sys: 164 ms, total: 710 ms
Wall time: 1.94 s
CPU times: user 331 ms, sys: 139 ms, total: 471 ms
Wall time: 2.95 s

And now we can compute the median:

In [41]:

good_result = good_data.median(dim="time")
good_result

Out[41]:

<xarray.DataArray 'from-value-546cce3c085aed0303a22d26a0a6d390' (band: 3, y: 10980, x: 10980)>
dask.array<nanmedian, shape=(3, 10980, 10980), dtype=float64, chunksize=(1, 732, 732), chunktype=numpy.ndarray>
Coordinates:
  * band                                     (band) <U3 'B04' 'B03' 'B02'
  * x                                        (x) float64 5e+05 ... 6.098e+05
  * y                                        (y) float64 5.3e+06 ... 5.19e+06
    s2:mgrs_tile                             <U5 '10TET'
    instruments                              <U3 'msi'
    proj:epsg                                int64 32610
    sat:orbit_state                          <U10 'descending'
    s2:degraded_msi_data_percentage          float64 0.0
    s2:processing_baseline                   <U5 '02.12'
    s2:product_type                          <U7 'S2MSI2A'
    s2:datatake_type                         <U8 'INS-NOBS'
    constellation                            <U10 'Sentinel 2'
    s2:saturated_defective_pixel_percentage  float64 0.0
    proj:transform                           object {0.0, 5300040.0, 10.0, 49...
    proj:bbox                                object {5190240.0, 609780.0, 499...
    proj:shape                               object {10980}
    title                                    (band) <U20 'Band 4 - Red - 10m'...
    gsd                                      int64 10
    common_name                              (band) <U5 'red' 'green' 'blue'
    center_wavelength                        (band) float64 0.665 0.56 0.49
    full_width_half_max                      (band) float64 0.038 0.045 0.098
    epsg                                     int64 32610

xarray.DataArray

'from-value-546cce3c085aed0303a22d26a0a6d390'

band: 3
y: 10980
x: 10980

dask.array<chunksize=(1, 732, 732), meta=np.ndarray>

	Array	Chunk
Bytes	2.69 GiB	4.09 MiB
Shape	(3, 10980, 10980)	(1, 732, 732)
Count	4050 Tasks	675 Chunks
Type	float64	numpy.ndarray

Coordinates: (22)
- band
  (band)
  <U3
  'B04' 'B03' 'B02'
```
array(['B04', 'B03', 'B02'], dtype='<U3')
```
- x
  (x)
  float64
  5e+05 5e+05 ... 6.098e+05 6.098e+05
```
array([499980., 499990., 500000., ..., 609750., 609760., 609770.])
```
- y
  (y)
  float64
  5.3e+06 5.3e+06 ... 5.19e+06
```
array([5300050., 5300040., 5300030., ..., 5190280., 5190270., 5190260.])
```
- s2:mgrs_tile
  ()
  <U5
  '10TET'
```
array('10TET', dtype='<U5')
```
- instruments
  ()
  <U3
  'msi'
```
array('msi', dtype='<U3')
```
- proj:epsg
  ()
  int64
  32610
```
array(32610)
```
- sat:orbit_state
  ()
  <U10
  'descending'
```
array('descending', dtype='<U10')
```
- s2:degraded_msi_data_percentage
  ()
  float64
  0.0
```
array(0.)
```
- s2:processing_baseline
  ()
  <U5
  '02.12'
```
array('02.12', dtype='<U5')
```
- s2:product_type
  ()
  <U7
  'S2MSI2A'
```
array('S2MSI2A', dtype='<U7')
```
- s2:datatake_type
  ()
  <U8
  'INS-NOBS'
```
array('INS-NOBS', dtype='<U8')
```
- constellation
  ()
  <U10
  'Sentinel 2'
```
array('Sentinel 2', dtype='<U10')
```
- s2:saturated_defective_pixel_percentage
  ()
  float64
  0.0
```
array(0.)
```
- proj:transform
  ()
  object
  {0.0, 5300040.0, 10.0, 499980.0,...
```
array({0.0, 5300040.0, 10.0, 499980.0, -10.0}, dtype=object)
```
- proj:bbox
  ()
  object
  {5190240.0, 609780.0, 499980.0, ...
```
array({5190240.0, 609780.0, 499980.0, 5300040.0}, dtype=object)
```
- proj:shape
  ()
  object
  {10980}
```
array({10980}, dtype=object)
```
- title
  (band)
  <U20
  'Band 4 - Red - 10m' ... 'Band 2...
```
array(['Band 4 - Red - 10m', 'Band 3 - Green - 10m',
       'Band 2 - Blue - 10m'], dtype='<U20')
```
- gsd
  ()
  int64
  10
```
array(10)
```
- common_name
  (band)
  <U5
  'red' 'green' 'blue'
```
array(['red', 'green', 'blue'], dtype='<U5')
```
- center_wavelength
  (band)
  float64
  0.665 0.56 0.49
```
array([0.665, 0.56 , 0.49 ])
```
- full_width_half_max
  (band)
  float64
  0.038 0.045 0.098
```
array([0.038, 0.045, 0.098])
```
- epsg
  ()
  int64
  32610
```
array(32610)
```
Attributes: (0)

Comparison¶

A few things to point out:

The number of tasks is much smaller: about 5,000 compared to 25,000 from before.
The computation represented by this task graph is much simpler. Initially, each output chunk of the median would require many input chunks, as data is gathered from all many arrays that were read individually.

In [13]:

client.wait_for_workers(24)

In [14]:

bad_result.data.blocks[0, 0, 0].visualize(optimize_graph=True)

Out[14]:

By contrast, each output chunk of the good_result median is extremely simple:

In [15]:

good_result.data.blocks[0, 0, 0].visualize(optimize_graph=True)

Out[15]:

Simpler tasks mean less (or no) communication, which tends to be the bane of Dask clusters everywhere.

That said, it's not a panacea. This kind of "read small bits from many files" is slower than reading a large byte stream and then using Dask to move data around. I haven't profiled it yet, but it's like some combination of 1.) per file overhead (high latency filesystem, parsing metada, etc.) 2.) HTTPs vs. TCP. So if your computation does work

In [16]:

from distributed import performance_report

In [ ]:

with performance_report("good-median.html"):
    a = good_result.compute()

In [ ]:

with performance_report("bad-median.html"):
    b = bad_result.compute()

In [ ]:

cluster.close()

In [ ]:

assert a.equals(b)

...