import os import numpy as np import pandas as pd import xarray as xr from . import randint, randn, requires_dask try: import dask import dask.multiprocessing except ImportError: pass os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" class IOSingleNetCDF: """ A few examples that benchmark reading/writing a single netCDF file with xarray """ timeout = 300.0 repeat = 1 number = 5 def make_ds(self): # single Dataset self.ds = xr.Dataset() self.nt = 1000 self.nx = 90 self.ny = 45 self.block_chunks = { "time": self.nt / 4, "lon": self.nx / 3, "lat": self.ny / 3, } self.time_chunks = {"time": int(self.nt / 36)} times = pd.date_range("1970-01-01", periods=self.nt, freq="D") lons = xr.DataArray( np.linspace(0, 360, self.nx), dims=("lon",), attrs={"units": "degrees east", "long_name": "longitude"}, ) lats = xr.DataArray( np.linspace(-90, 90, self.ny), dims=("lat",), attrs={"units": "degrees north", "long_name": "latitude"}, ) self.ds["foo"] = xr.DataArray( randn((self.nt, self.nx, self.ny), frac_nan=0.2), coords={"lon": lons, "lat": lats, "time": times}, dims=("time", "lon", "lat"), name="foo", encoding=None, attrs={"units": "foo units", "description": "a description"}, ) self.ds["bar"] = xr.DataArray( randn((self.nt, self.nx, self.ny), frac_nan=0.2), coords={"lon": lons, "lat": lats, "time": times}, dims=("time", "lon", "lat"), name="bar", encoding=None, attrs={"units": "bar units", "description": "a description"}, ) self.ds["baz"] = xr.DataArray( randn((self.nx, self.ny), frac_nan=0.2).astype(np.float32), coords={"lon": lons, "lat": lats}, dims=("lon", "lat"), name="baz", encoding=None, attrs={"units": "baz units", "description": "a description"}, ) self.ds.attrs = {"history": "created for xarray benchmarking"} self.oinds = { "time": randint(0, self.nt, 120), "lon": randint(0, self.nx, 20), "lat": randint(0, self.ny, 10), } self.vinds = { "time": xr.DataArray(randint(0, self.nt, 120), dims="x"), "lon": xr.DataArray(randint(0, self.nx, 120), dims="x"), "lat": slice(3, 20), } class IOWriteSingleNetCDF3(IOSingleNetCDF): def setup(self): self.format = "NETCDF3_64BIT" self.make_ds() def time_write_dataset_netcdf4(self): self.ds.to_netcdf("test_netcdf4_write.nc", engine="netcdf4", format=self.format) def time_write_dataset_scipy(self): self.ds.to_netcdf("test_scipy_write.nc", engine="scipy", format=self.format) class IOReadSingleNetCDF4(IOSingleNetCDF): def setup(self): self.make_ds() self.filepath = "test_single_file.nc4.nc" self.format = "NETCDF4" self.ds.to_netcdf(self.filepath, format=self.format) def time_load_dataset_netcdf4(self): xr.open_dataset(self.filepath, engine="netcdf4").load() def time_orthogonal_indexing(self): ds = xr.open_dataset(self.filepath, engine="netcdf4") ds = ds.isel(**self.oinds).load() def time_vectorized_indexing(self): ds = xr.open_dataset(self.filepath, engine="netcdf4") ds = ds.isel(**self.vinds).load() class IOReadSingleNetCDF3(IOReadSingleNetCDF4): def setup(self): self.make_ds() self.filepath = "test_single_file.nc3.nc" self.format = "NETCDF3_64BIT" self.ds.to_netcdf(self.filepath, format=self.format) def time_load_dataset_scipy(self): xr.open_dataset(self.filepath, engine="scipy").load() def time_orthogonal_indexing(self): ds = xr.open_dataset(self.filepath, engine="scipy") ds = ds.isel(**self.oinds).load() def time_vectorized_indexing(self): ds = xr.open_dataset(self.filepath, engine="scipy") ds = ds.isel(**self.vinds).load() class IOReadSingleNetCDF4Dask(IOSingleNetCDF): def setup(self): requires_dask() self.make_ds() self.filepath = "test_single_file.nc4.nc" self.format = "NETCDF4" self.ds.to_netcdf(self.filepath, format=self.format) def time_load_dataset_netcdf4_with_block_chunks(self): xr.open_dataset( self.filepath, engine="netcdf4", chunks=self.block_chunks ).load() def time_load_dataset_netcdf4_with_block_chunks_oindexing(self): ds = xr.open_dataset(self.filepath, engine="netcdf4", chunks=self.block_chunks) ds = ds.isel(**self.oinds).load() def time_load_dataset_netcdf4_with_block_chunks_vindexing(self): ds = xr.open_dataset(self.filepath, engine="netcdf4", chunks=self.block_chunks) ds = ds.isel(**self.vinds).load() def time_load_dataset_netcdf4_with_block_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_dataset( self.filepath, engine="netcdf4", chunks=self.block_chunks ).load() def time_load_dataset_netcdf4_with_time_chunks(self): xr.open_dataset(self.filepath, engine="netcdf4", chunks=self.time_chunks).load() def time_load_dataset_netcdf4_with_time_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_dataset( self.filepath, engine="netcdf4", chunks=self.time_chunks ).load() class IOReadSingleNetCDF3Dask(IOReadSingleNetCDF4Dask): def setup(self): requires_dask() self.make_ds() self.filepath = "test_single_file.nc3.nc" self.format = "NETCDF3_64BIT" self.ds.to_netcdf(self.filepath, format=self.format) def time_load_dataset_scipy_with_block_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_dataset( self.filepath, engine="scipy", chunks=self.block_chunks ).load() def time_load_dataset_scipy_with_block_chunks_oindexing(self): ds = xr.open_dataset(self.filepath, engine="scipy", chunks=self.block_chunks) ds = ds.isel(**self.oinds).load() def time_load_dataset_scipy_with_block_chunks_vindexing(self): ds = xr.open_dataset(self.filepath, engine="scipy", chunks=self.block_chunks) ds = ds.isel(**self.vinds).load() def time_load_dataset_scipy_with_time_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_dataset( self.filepath, engine="scipy", chunks=self.time_chunks ).load() class IOMultipleNetCDF: """ A few examples that benchmark reading/writing multiple netCDF files with xarray """ timeout = 300.0 repeat = 1 number = 5 def make_ds(self, nfiles=10): # multiple Dataset self.ds = xr.Dataset() self.nt = 1000 self.nx = 90 self.ny = 45 self.nfiles = nfiles self.block_chunks = { "time": self.nt / 4, "lon": self.nx / 3, "lat": self.ny / 3, } self.time_chunks = {"time": int(self.nt / 36)} self.time_vars = np.split( pd.date_range("1970-01-01", periods=self.nt, freq="D"), self.nfiles ) self.ds_list = [] self.filenames_list = [] for i, times in enumerate(self.time_vars): ds = xr.Dataset() nt = len(times) lons = xr.DataArray( np.linspace(0, 360, self.nx), dims=("lon",), attrs={"units": "degrees east", "long_name": "longitude"}, ) lats = xr.DataArray( np.linspace(-90, 90, self.ny), dims=("lat",), attrs={"units": "degrees north", "long_name": "latitude"}, ) ds["foo"] = xr.DataArray( randn((nt, self.nx, self.ny), frac_nan=0.2), coords={"lon": lons, "lat": lats, "time": times}, dims=("time", "lon", "lat"), name="foo", encoding=None, attrs={"units": "foo units", "description": "a description"}, ) ds["bar"] = xr.DataArray( randn((nt, self.nx, self.ny), frac_nan=0.2), coords={"lon": lons, "lat": lats, "time": times}, dims=("time", "lon", "lat"), name="bar", encoding=None, attrs={"units": "bar units", "description": "a description"}, ) ds["baz"] = xr.DataArray( randn((self.nx, self.ny), frac_nan=0.2).astype(np.float32), coords={"lon": lons, "lat": lats}, dims=("lon", "lat"), name="baz", encoding=None, attrs={"units": "baz units", "description": "a description"}, ) ds.attrs = {"history": "created for xarray benchmarking"} self.ds_list.append(ds) self.filenames_list.append("test_netcdf_%i.nc" % i) class IOWriteMultipleNetCDF3(IOMultipleNetCDF): def setup(self): self.make_ds() self.format = "NETCDF3_64BIT" def time_write_dataset_netcdf4(self): xr.save_mfdataset( self.ds_list, self.filenames_list, engine="netcdf4", format=self.format ) def time_write_dataset_scipy(self): xr.save_mfdataset( self.ds_list, self.filenames_list, engine="scipy", format=self.format ) class IOReadMultipleNetCDF4(IOMultipleNetCDF): def setup(self): requires_dask() self.make_ds() self.format = "NETCDF4" xr.save_mfdataset(self.ds_list, self.filenames_list, format=self.format) def time_load_dataset_netcdf4(self): xr.open_mfdataset(self.filenames_list, engine="netcdf4").load() def time_open_dataset_netcdf4(self): xr.open_mfdataset(self.filenames_list, engine="netcdf4") class IOReadMultipleNetCDF3(IOReadMultipleNetCDF4): def setup(self): requires_dask() self.make_ds() self.format = "NETCDF3_64BIT" xr.save_mfdataset(self.ds_list, self.filenames_list, format=self.format) def time_load_dataset_scipy(self): xr.open_mfdataset(self.filenames_list, engine="scipy").load() def time_open_dataset_scipy(self): xr.open_mfdataset(self.filenames_list, engine="scipy") class IOReadMultipleNetCDF4Dask(IOMultipleNetCDF): def setup(self): requires_dask() self.make_ds() self.format = "NETCDF4" xr.save_mfdataset(self.ds_list, self.filenames_list, format=self.format) def time_load_dataset_netcdf4_with_block_chunks(self): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.block_chunks ).load() def time_load_dataset_netcdf4_with_block_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.block_chunks ).load() def time_load_dataset_netcdf4_with_time_chunks(self): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.time_chunks ).load() def time_load_dataset_netcdf4_with_time_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.time_chunks ).load() def time_open_dataset_netcdf4_with_block_chunks(self): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.block_chunks ) def time_open_dataset_netcdf4_with_block_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.block_chunks ) def time_open_dataset_netcdf4_with_time_chunks(self): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.time_chunks ) def time_open_dataset_netcdf4_with_time_chunks_multiprocessing(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="netcdf4", chunks=self.time_chunks ) class IOReadMultipleNetCDF3Dask(IOReadMultipleNetCDF4Dask): def setup(self): requires_dask() self.make_ds() self.format = "NETCDF3_64BIT" xr.save_mfdataset(self.ds_list, self.filenames_list, format=self.format) def time_load_dataset_scipy_with_block_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="scipy", chunks=self.block_chunks ).load() def time_load_dataset_scipy_with_time_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="scipy", chunks=self.time_chunks ).load() def time_open_dataset_scipy_with_block_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="scipy", chunks=self.block_chunks ) def time_open_dataset_scipy_with_time_chunks(self): with dask.config.set(scheduler="multiprocessing"): xr.open_mfdataset( self.filenames_list, engine="scipy", chunks=self.time_chunks ) def create_delayed_write(): import dask.array as da vals = da.random.random(300, chunks=(1,)) ds = xr.Dataset({"vals": (["a"], vals)}) return ds.to_netcdf("file.nc", engine="netcdf4", compute=False) class IOWriteNetCDFDask: timeout = 60 repeat = 1 number = 5 def setup(self): requires_dask() self.write = create_delayed_write() def time_write(self): self.write.compute() class IOWriteNetCDFDaskDistributed: def setup(self): try: import distributed except ImportError: raise NotImplementedError() self.client = distributed.Client() self.write = create_delayed_write() def cleanup(self): self.client.shutdown() def time_write(self): self.write.compute()