What is the most efficient/quickest way to assign new values to a subset of a dask DataArray?

[GrowlingM1ke]

I have video data that I have loaded into a dask DataArray. Some of the frames in the video are erroneous and are either missing or contain weird artifacts. For analysis purposes I want to get rid of these problematic frames by interpolating or copying over values from frames adjacent to them.

Whatever I attempt to do seems to be very slow or results in using too much memory, which I think is due to the fact that the entire video is loaded into memory before the fix is executed. Below is a sample of my code:

def _fill_right(self, click):
    if self._last > self._first:
        self.update_status.value = "Started Filling"
        fill = xr.apply_ufunc(
            _fill,
            self.varr_copy.chunk(dict(frame=-1)),
            input_core_dims=[["frame", "height", "width"]],
            output_core_dims=[["frame", "height", "width"]],
            dask="parallelized",
            kwargs=dict(direction="right", first=self._first, last=self._last),
            output_dtypes=[self.varr_copy.dtype],
        ).compute()
        # Convert to dataset and update self.ds
        self.varr_copy = fill
        self.ds = fill.to_dataset().compute()
        self.update_ds()

def _fill_interpolate(self, click):
    if self._last > self._first:
        self.update_status.value = "Started Filling"
        fill = xr.apply_ufunc(
            _fill,
            self.varr_copy.chunk(dict(frame=-1)),
            input_core_dims=[["frame", "height", "width"]],
            output_core_dims=[["frame", "height", "width"]],
            dask="parallelized",
            kwargs=dict(direction="interpolate", first=self._first, last=self._last),
            output_dtypes=[self.varr_copy.dtype],
        ).compute()
        # Convert to dataset and update self.ds
        self.varr_copy = fill
        self.ds = fill.to_dataset().compute()
        self.update_ds()

Below is the fill function:

def _fill(x: np.ndarray, first: int, last: int, direction: str = 'left') -> np.ndarray:
    x = x.copy()
    if x.any():
        if direction == "left":
            x[first:last] = x[first]
        elif direction == "right":
            x[first+1:last] = x[last]
        elif direction == "interpolate":
            frames = last-first-1
            first_frame = x[first].astype("float")
            last_frame = x[last].astype("float")
            coefficients = (np.arange(1, frames+1)) / (frames + 2)
            # Due to memory issues, we need to do this per frame
            for i, coef in enumerate(coefficients):
                x[first+i+1] = np.rint(first_frame - (first_frame - last_frame) * coef).astype(x.dtype)
    return x

The process of writing this code has been mostly trial and error so I fully apologize if the above is incoherent. Is there any approach I could take that would allow me to apply my fix to a subset of the dataArray? Could I load only a specific chunk of the dataArray and apply the fix to it?

Thank you in advance!

[andersy005]

@GrowlingM1ke,

• you're calling .compute() multiple times (fill.compute() and fill.to_dataset().compute()). Each .compute() call triggers a full evaluation of the computation graph up to that point. i'd try to structure your code so that you call .compute() only once at the end of your processing pipeline
• additionally, your chunking strategy (self.varr_copy.chunk(dict(frame=-1))) is potentially an issue. chunking the entire frame dimension (frame=-1) can lead to large chunks if your video has many frames. Consider chunking into smaller sizes to balance memory usage and parallel processing efficiency

def _fill_interpolate(self, click):
    if self._last > self._first:
        self.update_status.value = "Started Filling"

        # adjust chunk size as needed
        # use `auto` chunking to let dask choose appropriate chunk sizes for height/width dimensions
        chunk_size = {'frame': 10, 'height': 'auto', 'width': 'auto'}
        self.varr_copy = self.varr_copy.chunk(chunk_size)

        fill = xr.apply_ufunc(
            _fill,
            self.varr_copy,
            input_core_dims=[["frame", "height", "width"]],
            output_core_dims=[["frame", "height", "width"]],
            dask="parallelized",
            kwargs=dict(direction="interpolate", first=self._first, last=self._last),
            output_dtypes=[self.varr_copy.dtype],
        )

        # delay compute call
        self.varr_copy = fill
        self.ds = fill.to_dataset()

.... 

and then when you after calling self._fill_interpolate(....), you can trigger computation by calling self.ds.compute().

considering the above points, i recommend trying the revised _fill_interpolate() function and see if there's an improvement

[GrowlingM1ke]

Hey there,

First of all, I would like to say thank you for your response!

I've spent the past few days working out the kinks of DataArray and reading up what I can do. To address your points:

• You're absolutely correct about me using too many compute() calls so I went ahead and got rid of a good chunk of them.
• This was definitely part of the problem but I had the following issue; if I only needed to interpolate a small section of the video then how would I communicate to other chunks that don't require interpolation not to do anything?

My solution was as follows, create a 1-dimensional array size of frames that would contain values of 1 or 0 to indicate whether a given section required interpolation/filling or not, then each chunk would quickly check indices.any() if there is any work to be done on that section. Fair warning to anyone else who wants to try this solution, I am only doing this since I am unfortunately tied to on an older version of xarray and dask (2021.2.0 to be specific) before numpy-esque item assignment was a thing.

One thing I'll quickly mention is that for whatever reason when I ran with vectorize=True I would get an item assignment is not allowed error.

I'm still running it one issue but it is not related to this part. Tomorrow I'll finish up my code and if everything is working I'll post a solution.