09 Dask demo

UW Geospatial Data Analysis
CEE467/CEWA567
David Shean, Eric Gagliano, Quinn Brencher

from pathlib import Path
import xarray as xr
import os
import time

from dask.distributed import Client
client = Client()  
client

Client

Client-ee37f65f-fd5b-11ef-aaa0-00155de26d4e

Connection method: Cluster object	Cluster type: distributed.LocalCluster
Dashboard: http://127.0.0.1:8787/status

Launch dashboard in JupyterLab

Cluster Info

LocalCluster

2c6fd39e

Dashboard: http://127.0.0.1:8787/status	Workers: 4
Total threads: 16	Total memory: 15.48 GiB
Status: running	Using processes: True

Scheduler Info

Scheduler

Scheduler-616b447b-e31e-4b1e-a8f0-2004966e025f

Comm: tcp://127.0.0.1:36977	Workers: 4
Dashboard: http://127.0.0.1:8787/status	Total threads: 16
Started: Just now	Total memory: 15.48 GiB

Workers

Worker: 0

Comm: tcp://127.0.0.1:43959	Total threads: 4
Dashboard: http://127.0.0.1:38163/status	Memory: 3.87 GiB
Nanny: tcp://127.0.0.1:36147
Local directory: /tmp/dask-scratch-space/worker-8005rux3

Worker: 1

Comm: tcp://127.0.0.1:34647	Total threads: 4
Dashboard: http://127.0.0.1:39721/status	Memory: 3.87 GiB
Nanny: tcp://127.0.0.1:39347
Local directory: /tmp/dask-scratch-space/worker-sav5tp3c

Worker: 2

Comm: tcp://127.0.0.1:41485	Total threads: 4
Dashboard: http://127.0.0.1:36713/status	Memory: 3.87 GiB
Nanny: tcp://127.0.0.1:35701
Local directory: /tmp/dask-scratch-space/worker-03ksmeqo

Worker: 3

Comm: tcp://127.0.0.1:40267	Total threads: 4
Dashboard: http://127.0.0.1:34875/status	Memory: 3.87 GiB
Nanny: tcp://127.0.0.1:44511
Local directory: /tmp/dask-scratch-space/worker-46pneafw

era5_data_dir = f'{Path.home()}/gda_demo_data/era5_data'
anom_fn = os.path.join(era5_data_dir, '1month_anomaly_Global_ea_2t.nc') 

def coarsen_and_polyfit(ds, coarsen_factor=2):
    ds = ds.assign_coords(longitude=(((ds.longitude + 180) % 360) - 180)).sortby('longitude')
    ds_coarsened = ds.coarsen(latitude=coarsen_factor, longitude=coarsen_factor, boundary='trim').mean()
    ds_coarsened_polyfit = ds_coarsened.polyfit(dim='time',deg=1)
    return ds_coarsened_polyfit

anom_nochunks_ds = xr.open_dataset(anom_fn, chunks=None)
anom_nochunks_ds

<xarray.Dataset> Size: 2GB
Dimensions:    (time: 517, latitude: 721, longitude: 1440)
Coordinates:
  * time       (time) datetime64[ns] 4kB 1979-01-01 1979-02-01 ... 2022-01-01
  * latitude   (latitude) float64 6kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0
  * longitude  (longitude) float64 12kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8
Data variables:
    t2m        (time, latitude, longitude) float32 2GB ...
Attributes:
    GRIB_edition:            1
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2022-02-28T07:59 GRIB to CDM+CF via cfgrib-0.9.1...

xarray.Dataset

• Dimensions:
  • time: 517
  • latitude: 721
  • longitude: 1440
• Coordinates: (3)
  • time
    (time)
    datetime64[ns]
    1979-01-01 ... 2022-01-01

    long_name :

    initial time of forecast

    standard_name :

    forecast_reference_time

    array(['1979-01-01T00:00:00.000000000', '1979-02-01T00:00:00.000000000',
           '1979-03-01T00:00:00.000000000', ..., '2021-11-01T00:00:00.000000000',
           '2021-12-01T00:00:00.000000000', '2022-01-01T00:00:00.000000000'],
          dtype='datetime64[ns]')

  • latitude
    (latitude)
    float64
    90.0 89.75 89.5 ... -89.75 -90.0

    units :

    degrees_north

    standard_name :

    latitude

    long_name :

    latitude

    stored_direction :

    decreasing

    array([ 90.  ,  89.75,  89.5 , ..., -89.5 , -89.75, -90.  ])

  • longitude
    (longitude)
    float64
    0.0 0.25 0.5 ... 359.2 359.5 359.8

    units :

    degrees_east

    standard_name :

    longitude

    long_name :

    longitude

    array([0.0000e+00, 2.5000e-01, 5.0000e-01, ..., 3.5925e+02, 3.5950e+02,
           3.5975e+02])

• Data variables: (1)
  • t2m
    (time, latitude, longitude)
    float32
    ...

    GRIB_paramId :

    167

    GRIB_dataType :

    an

    GRIB_numberOfPoints :

    1038240

    GRIB_typeOfLevel :

    surface

    GRIB_stepUnits :

    1

    GRIB_stepType :

    avgua

    GRIB_gridType :

    regular_ll

    GRIB_NV :

    0

    GRIB_Nx :

    1440

    GRIB_Ny :

    721

    GRIB_cfName :

    unknown

    GRIB_cfVarName :

    t2m

    GRIB_gridDefinitionDescription :

    Latitude/Longitude Grid

    GRIB_iDirectionIncrementInDegrees :

    0.25

    GRIB_iScansNegatively :

    0

    GRIB_jDirectionIncrementInDegrees :

    0.25

    GRIB_jPointsAreConsecutive :

    0

    GRIB_jScansPositively :

    0

    GRIB_latitudeOfFirstGridPointInDegrees :

    90.0

    GRIB_latitudeOfLastGridPointInDegrees :

    -90.0

    GRIB_longitudeOfFirstGridPointInDegrees :

    0.0

    GRIB_longitudeOfLastGridPointInDegrees :

    359.75

    GRIB_missingValue :

    9999

    GRIB_name :

    2 metre temperature

    GRIB_shortName :

    2t

    GRIB_totalNumber :

    0

    GRIB_units :

    K

    long_name :

    2 metre temperature

    units :

    K

    standard_name :

    unknown

    [536770080 values with dtype=float32]

• Indexes: (3)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['1979-01-01', '1979-02-01', '1979-03-01', '1979-04-01',
                   '1979-05-01', '1979-06-01', '1979-07-01', '1979-08-01',
                   '1979-09-01', '1979-10-01',
                   ...
                   '2021-04-01', '2021-05-01', '2021-06-01', '2021-07-01',
                   '2021-08-01', '2021-09-01', '2021-10-01', '2021-11-01',
                   '2021-12-01', '2022-01-01'],
                  dtype='datetime64[ns]', name='time', length=517, freq=None))

  • latitude
    PandasIndex

    PandasIndex(Index([  90.0,  89.75,   89.5,  89.25,   89.0,  88.75,   88.5,  88.25,   88.0,
            87.75,
           ...
           -87.75,  -88.0, -88.25,  -88.5, -88.75,  -89.0, -89.25,  -89.5, -89.75,
            -90.0],
          dtype='float64', name='latitude', length=721))

  • longitude
    PandasIndex

    PandasIndex(Index([   0.0,   0.25,    0.5,   0.75,    1.0,   1.25,    1.5,   1.75,    2.0,
             2.25,
           ...
            357.5, 357.75,  358.0, 358.25,  358.5, 358.75,  359.0, 359.25,  359.5,
           359.75],
          dtype='float64', name='longitude', length=1440))

• Attributes: (7)

  GRIB_edition :

  1

  GRIB_centre :

  ecmf

  GRIB_centreDescription :

  European Centre for Medium-Range Weather Forecasts

  GRIB_subCentre :

  0

  Conventions :

  CF-1.7

  institution :

  European Centre for Medium-Range Weather Forecasts

  history :

  2022-02-28T07:59 GRIB to CDM+CF via cfgrib-0.9.10.0/ecCodes-2.24.2 with {"source": "ecv-for-climate-change_t2m_anomaly/1month_anomaly_Global_ea_2t_197901_1981-2010_v02.grib", "filter_by_keys": {}, "encode_cf": ["parameter", "time", "geography", "vertical"]}

start = time.time()
coarsen_and_polyfit(anom_nochunks_ds, coarsen_factor=6)
end = time.time()
print(f"Time taken without chunks: {end-start:.2f} seconds")

Time taken without chunks: 11.90 seconds

anom_nochunks_ds = None

anom_autochunks_ds = xr.open_dataset(anom_fn, chunks="auto")
anom_autochunks_ds

<xarray.Dataset> Size: 2GB
Dimensions:    (time: 517, latitude: 721, longitude: 1440)
Coordinates:
  * time       (time) datetime64[ns] 4kB 1979-01-01 1979-02-01 ... 2022-01-01
  * latitude   (latitude) float64 6kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0
  * longitude  (longitude) float64 12kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8
Data variables:
    t2m        (time, latitude, longitude) float32 2GB dask.array<chunksize=(205, 286, 571), meta=np.ndarray>
Attributes:
    GRIB_edition:            1
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2022-02-28T07:59 GRIB to CDM+CF via cfgrib-0.9.1...

xarray.Dataset

• Dimensions:
  • time: 517
  • latitude: 721
  • longitude: 1440
• Coordinates: (3)
  • time
    (time)
    datetime64[ns]
    1979-01-01 ... 2022-01-01

    long_name :

    initial time of forecast

    standard_name :

    forecast_reference_time

    array(['1979-01-01T00:00:00.000000000', '1979-02-01T00:00:00.000000000',
           '1979-03-01T00:00:00.000000000', ..., '2021-11-01T00:00:00.000000000',
           '2021-12-01T00:00:00.000000000', '2022-01-01T00:00:00.000000000'],
          dtype='datetime64[ns]')

  • latitude
    (latitude)
    float64
    90.0 89.75 89.5 ... -89.75 -90.0

    units :

    degrees_north

    standard_name :

    latitude

    long_name :

    latitude

    stored_direction :

    decreasing

    array([ 90.  ,  89.75,  89.5 , ..., -89.5 , -89.75, -90.  ])

  • longitude
    (longitude)
    float64
    0.0 0.25 0.5 ... 359.2 359.5 359.8

    units :

    degrees_east

    standard_name :

    longitude

    long_name :

    longitude

    array([0.0000e+00, 2.5000e-01, 5.0000e-01, ..., 3.5925e+02, 3.5950e+02,
           3.5975e+02])

• Data variables: (1)
  • t2m
    (time, latitude, longitude)
    float32
    dask.array<chunksize=(205, 286, 571), meta=np.ndarray>

    GRIB_paramId :

    167

    GRIB_dataType :

    an

    GRIB_numberOfPoints :

    1038240

    GRIB_typeOfLevel :

    surface

    GRIB_stepUnits :

    1

    GRIB_stepType :

    avgua

    GRIB_gridType :

    regular_ll

    GRIB_NV :

    0

    GRIB_Nx :

    1440

    GRIB_Ny :

    721

    GRIB_cfName :

    unknown

    GRIB_cfVarName :

    t2m

    GRIB_gridDefinitionDescription :

    Latitude/Longitude Grid

    GRIB_iDirectionIncrementInDegrees :

    0.25

    GRIB_iScansNegatively :

    0

    GRIB_jDirectionIncrementInDegrees :

    0.25

    GRIB_jPointsAreConsecutive :

    0

    GRIB_jScansPositively :

    0

    GRIB_latitudeOfFirstGridPointInDegrees :

    90.0

    GRIB_latitudeOfLastGridPointInDegrees :

    -90.0

    GRIB_longitudeOfFirstGridPointInDegrees :

    0.0

    GRIB_longitudeOfLastGridPointInDegrees :

    359.75

    GRIB_missingValue :

    9999

    GRIB_name :

    2 metre temperature

    GRIB_shortName :

    2t

    GRIB_totalNumber :

    0

    GRIB_units :

    K

    long_name :

    2 metre temperature

    units :

    K

    standard_name :

    unknown

    Array Chunk Bytes 2.00 GiB 127.71 MiB Shape (517, 721, 1440) (205, 286, 571) Dask graph 27 chunks in 2 graph layers Data type float32 numpy.ndarray

• Indexes: (3)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['1979-01-01', '1979-02-01', '1979-03-01', '1979-04-01',
                   '1979-05-01', '1979-06-01', '1979-07-01', '1979-08-01',
                   '1979-09-01', '1979-10-01',
                   ...
                   '2021-04-01', '2021-05-01', '2021-06-01', '2021-07-01',
                   '2021-08-01', '2021-09-01', '2021-10-01', '2021-11-01',
                   '2021-12-01', '2022-01-01'],
                  dtype='datetime64[ns]', name='time', length=517, freq=None))

  • latitude
    PandasIndex

    PandasIndex(Index([  90.0,  89.75,   89.5,  89.25,   89.0,  88.75,   88.5,  88.25,   88.0,
            87.75,
           ...
           -87.75,  -88.0, -88.25,  -88.5, -88.75,  -89.0, -89.25,  -89.5, -89.75,
            -90.0],
          dtype='float64', name='latitude', length=721))

  • longitude
    PandasIndex

    PandasIndex(Index([   0.0,   0.25,    0.5,   0.75,    1.0,   1.25,    1.5,   1.75,    2.0,
             2.25,
           ...
            357.5, 357.75,  358.0, 358.25,  358.5, 358.75,  359.0, 359.25,  359.5,
           359.75],
          dtype='float64', name='longitude', length=1440))

• Attributes: (7)

  GRIB_edition :

  1

  GRIB_centre :

  ecmf

  GRIB_centreDescription :

  European Centre for Medium-Range Weather Forecasts

  GRIB_subCentre :

  0

  Conventions :

  CF-1.7

  institution :

  European Centre for Medium-Range Weather Forecasts

  history :

  2022-02-28T07:59 GRIB to CDM+CF via cfgrib-0.9.10.0/ecCodes-2.24.2 with {"source": "ecv-for-climate-change_t2m_anomaly/1month_anomaly_Global_ea_2t_197901_1981-2010_v02.grib", "filter_by_keys": {}, "encode_cf": ["parameter", "time", "geography", "vertical"]}

start = time.time()
coarsen_and_polyfit(anom_autochunks_ds, coarsen_factor=6).compute()
end = time.time()
print(f"Time taken with auto chunks: {end-start:.2f} seconds")

Time taken with auto chunks: 6.94 seconds