How to plot a huge image out of memory?

Datashader
1

It loads the entire dataset before doing anything (and crashes), so I was wondering if it could just skip every 10 at farthest zoom (e.g. is there a decimate function for datashader outside of hv)?

import numpy as np
import dask.array as da
import xarray as xr
from datashader import transfer_functions as tf

def sample(range_=(0.0,2.4)):
    xs = da.arange(700000, dtype=np.float32)
    ys = da.arange(280000, dtype=np.float32)
    return xr.DataArray(
        da.random.random((len(ys), len(xs)), chunks=(10000, 10000)).astype(np.float32),
        coords=[('y', ys), ('x', xs)])

ds = sample(range_=(0,5))
tf.shade(ds)

Then it fails:

distributed.nanny - WARNING - Restarting worker
distributed.nanny - WARNING - Worker exceeded 95% memory budget. Restarting
distributed.nanny - WARNING - Worker exceeded 95% memory budget. Restarting
distributed.nanny - WARNING - Restarting worker
distributed.nanny - WARNING - Restarting worker

KilledWorker: ("('random_sample-astype-', 13, 11)", <WorkerState 'tcp://ip', name: 0, memory: 0, processing: 601>)
3

Are you saying that even with distributed, operating on a cluster, the total memory available on the cluster isn’t sufficient for this image? Or that it’s not really distributed, and that the image is larger than this one machine’s memory? In either case, I’m not sure how to ensure that your DataArray is backed by Dask, but it seems like it should be feasible (if slow) to work out of core using such a Dask array as long as you don’t persist it. Probably needs some arguments to DataArray?

1 Like
4

It’s like 720 GBs total; opens fine, runs for a while fine, then randomly crashes. The minimal example should reproduce the issue.

I am also using:

client = Client(n_workers=4, threads_per_worker=2, memory_limit="8GiB")
5

Similar:

6

Expanded a bit on your example from Larger than memory plotting: how to set initial slice to plot - #8 by ahuang11 and am very happy with the result :slight_smile:

Here my zarr file contains two arrays that weigh ~50GB each. Depending on the zoom-level (controlled through the max_size), only every ith pixel is sampled for datashading.

One thing I haven’t managed to figure out yet is how I could make the cmap and cbar_limit options interactive.

import numpy as np
import holoviews as hv
import holoviews.streams
import xarray as xr
from holoviews.operation.datashader import rasterize
import param
import panel as pn
import numpy as np

hv.extension("bokeh")


class BigDataViz(param.Parameterized):

    dm_view = param.ClassSelector(class_=hv.DynamicMap, precedence=0)
    dm_shaded_view = param.ClassSelector(class_=hv.DynamicMap, precedence=0)

    ds: xr.DataArray = param.ClassSelector(class_=xr.Dataset, precedence=0)

    array: str = param.Selector()

    cmap: str = param.Selector(
        default="Viridis",
        objects=["Viridis", "gray", "PiYG", "flag", "Set1"],
        precedence=0,
    )
    cbar_limit: tuple = param.Tuple(default=None, length=2, precedence=0)

    max_size: int = param.Integer(default=int(7e7))

    i: int = param.Integer(default=1)
    npoints: int = param.Integer(precedence=0)

    def __init__(self, **params):

        params["ds"] = params["ds"].sortby(list(params["ds"].coords), ascending=True)

        super().__init__(**params)

        spatial_vars = [x for x in ds.data_vars if len(ds[x].dims) == 2]
        self.param["array"].objects = spatial_vars
        self.param.update(array=spatial_vars[0])

        self.init_range = self.determine_middle()

        self.range_xy = holoviews.streams.RangeXY(
            x_range=self.init_range[0], y_range=self.init_range[1]
        )

        self.dm_view = hv.DynamicMap(
            self.load_data,
            streams=[
                self.range_xy,
                self.param["array"],
                self.param["max_size"],
            ],
        )

        self.dm_shaded_view = rasterize(self.dm_view).opts(
            cmap=self.cmap,
            clim=self.cbar_limit,
            colorbar=True,
            tools=["hover"],
            responsive=True,
            framewise=True,
        )

    def determine_middle(self):
        blocksize = np.sqrt(self.max_size).astype(int)
        bb = {dim: int((ds[dim].size - blocksize) / 2) for dim in ds[self.array].dims}
        x = self.ds[self.array].isel(
            {dim: slice(idx, idx + blocksize) for dim, idx in bb.items()}
        )
        return [
            (float(x[dim].min().values), float(x[dim].max().values))
            for dim in ds[self.array].dims
        ]

    def load_data(self, x_range, y_range, array, max_size):
        i = 0
        data: xr.DataArray = self.ds[array]
        dims = self.ds[array].dims
        while data.size > max_size:
            i += 1
            kwargs = {
                dim: slice(*range_, i) for dim, range_ in zip(dims, [x_range, y_range])
            }
            data = self.ds[array].sel(**kwargs)
        self.param.update(i=i, npoints=data.size)
        image = hv.Image(data.compute(), kdims=list(dims))
        return image

    def serve(self):
        pn.serve(
            pn.Row(
                pn.pane.HoloViews(
                        self.dm_shaded_view, 
                        sizing_mode="stretch_both"
                    ),
                pn.Param(
                    self.param, 
                    display_threshold=2),
                sizing_mode="stretch_both",
                styles={"background-color": "yellow"}
            ),
        )


if __name__ == "__main__":

    file_path = "my_data.zarr3"

    ds = xr.open_zarr(file_path).transpose(
        # NOTE first dimension of variable will map to X-axis, second to Y-axis.
        "x", "y",
    ).drop_vars("spatial_ref")

    bdv = BigDataViz(ds=ds, cbar_limit=(-2, 1))

    bdv.serve()

My dataset looks like this:

<xarray.Dataset> Size: 91GB
Dimensions:                     (x: 120960, y: 47040)
Coordinates:
  * x                           (x) float64 968kB -180.0 -180.0 ... 180.0 180.0
  * y                           (y) float64 376kB 80.0 80.0 ... -59.99 -60.0
Data variables:
    KGE  (x, y) float64 46GB dask.array<chunksize=(512, 512), meta=np.ndarray>
    NSE  (x, y) float64 46GB dask.array<chunksize=(512, 512), meta=np.ndarray>
Attributes:
    units:      -
    name:       r
    long_name:  Pearson Correlation Coefficient (WAPOR - CLMS, etlook)
1 Like