From 216216ae5c318d9a4f7c267fdce30fff3e89e309 Mon Sep 17 00:00:00 2001 From: kushalkolar Date: Thu, 9 Mar 2023 01:38:23 -0500 Subject: [PATCH] make feature buffer public, make isolated buffer optional for Image and Heatmap --- fastplotlib/graphics/features/_base.py | 10 ++--- fastplotlib/graphics/features/_colors.py | 8 ++-- fastplotlib/graphics/features/_data.py | 31 ++++++++----- fastplotlib/graphics/image.py | 57 +++++++++++++++++++++--- 4 files changed, 80 insertions(+), 26 deletions(-) diff --git a/fastplotlib/graphics/features/_base.py b/fastplotlib/graphics/features/_base.py index 7177b7bae..80029180e 100644 --- a/fastplotlib/graphics/features/_base.py +++ b/fastplotlib/graphics/features/_base.py @@ -4,7 +4,7 @@ from typing import * import numpy as np -from pygfx import Buffer +from pygfx import Buffer, Texture supported_dtypes = [ @@ -226,7 +226,7 @@ def _update_range(self, key): @property @abstractmethod - def _buffer(self) -> Buffer: + def buffer(self) -> Union[Buffer, Texture]: pass @property @@ -238,21 +238,21 @@ def _update_range_indices(self, key): key = cleanup_slice(key, self._upper_bound) if isinstance(key, int): - self._buffer.update_range(key, size=1) + self.buffer.update_range(key, size=1) return # else if it's a slice obj if isinstance(key, slice): if key.step == 1: # we cleaned up the slice obj so step of None becomes 1 # update range according to size using the offset - self._buffer.update_range(offset=key.start, size=key.stop - key.start) + self.buffer.update_range(offset=key.start, size=key.stop - key.start) else: step = key.step # convert slice to indices ixs = range(key.start, key.stop, step) for ix in ixs: - self._buffer.update_range(ix, size=1) + self.buffer.update_range(ix, size=1) else: raise TypeError("must pass int or slice to update range") diff --git a/fastplotlib/graphics/features/_colors.py b/fastplotlib/graphics/features/_colors.py index 7833e0b2c..a5147b95e 100644 --- a/fastplotlib/graphics/features/_colors.py +++ b/fastplotlib/graphics/features/_colors.py @@ -7,11 +7,11 @@ class ColorFeature(GraphicFeatureIndexable): @property - def _buffer(self): + def buffer(self): return self._parent.world_object.geometry.colors def __getitem__(self, item): - return self._buffer.data[item] + return self.buffer.data[item] def __init__(self, parent, colors, n_colors: int, alpha: float = 1.0, collection_index: int = None): """ @@ -113,7 +113,7 @@ def __setitem__(self, key, value): raise ValueError("fancy indexing for colors must be 2-dimension, i.e. [n_datapoints, RGBA]") # set the user passed data directly - self._buffer.data[key] = value + self.buffer.data[key] = value # update range # first slice obj is going to be the indexing so use key[0] @@ -162,7 +162,7 @@ def __setitem__(self, key, value): else: raise ValueError("numpy array passed to color must be of shape (4,) or (n_colors_modify, 4)") - self._buffer.data[key] = new_colors + self.buffer.data[key] = new_colors self._update_range(key) self._feature_changed(key, new_colors) diff --git a/fastplotlib/graphics/features/_data.py b/fastplotlib/graphics/features/_data.py index 95e549247..8e9599fa6 100644 --- a/fastplotlib/graphics/features/_data.py +++ b/fastplotlib/graphics/features/_data.py @@ -1,7 +1,7 @@ from typing import * import numpy as np -from pygfx import Buffer, Texture +from pygfx import Buffer, Texture, TextureView from ._base import GraphicFeatureIndexable, cleanup_slice, FeatureEvent, to_gpu_supported_dtype @@ -16,11 +16,11 @@ def __init__(self, parent, data: Any, collection_index: int = None): super(PointsDataFeature, self).__init__(parent, data, collection_index=collection_index) @property - def _buffer(self) -> Buffer: + def buffer(self) -> Buffer: return self._parent.world_object.geometry.positions def __getitem__(self, item): - return self._buffer.data[item] + return self.buffer.data[item] def _fix_data(self, data, parent): graphic_type = parent.__class__.__name__ @@ -54,7 +54,7 @@ def __setitem__(self, key, value): # otherwise assume that they have the right shape # numpy will throw errors if it can't broadcast - self._buffer.data[key] = value + self.buffer.data[key] = value self._update_range(key) # avoid creating dicts constantly if there are no events to handle if len(self._event_handlers) > 0: @@ -97,21 +97,25 @@ def __init__(self, parent, data: Any): "``[x_dim, y_dim]`` or ``[x_dim, y_dim, rgb]``" ) - data = to_gpu_supported_dtype(data) super(ImageDataFeature, self).__init__(parent, data) @property - def _buffer(self) -> Texture: + def buffer(self) -> Texture: + """Texture buffer for the image data""" return self._parent.world_object.geometry.grid.texture + def update_gpu(self): + """Update the GPU with the buffer""" + self._update_range(None) + def __getitem__(self, item): - return self._buffer.data[item] + return self.buffer.data[item] def __setitem__(self, key, value): # make sure float32 value = to_gpu_supported_dtype(value) - self._buffer.data[key] = value + self.buffer.data[key] = value self._update_range(key) # avoid creating dicts constantly if there are no events to handle @@ -119,7 +123,7 @@ def __setitem__(self, key, value): self._feature_changed(key, value) def _update_range(self, key): - self._buffer.update_range((0, 0, 0), size=self._buffer.size) + self.buffer.update_range((0, 0, 0), size=self.buffer.size) def _feature_changed(self, key, new_data): if key is not None: @@ -144,9 +148,14 @@ def _feature_changed(self, key, new_data): class HeatmapDataFeature(ImageDataFeature): @property - def _buffer(self) -> List[Texture]: + def buffer(self) -> List[Texture]: + """list of Texture buffer for the image data""" return [img.geometry.grid.texture for img in self._parent.world_object.children] + def update_gpu(self): + """Update the GPU with the buffer""" + self._update_range(None) + def __getitem__(self, item): return self._data[item] @@ -162,7 +171,7 @@ def __setitem__(self, key, value): self._feature_changed(key, value) def _update_range(self, key): - for buffer in self._buffer: + for buffer in self.buffer: buffer.update_range((0, 0, 0), size=buffer.size) def _feature_changed(self, key, new_data): diff --git a/fastplotlib/graphics/image.py b/fastplotlib/graphics/image.py index 854f757f2..83cae3de8 100644 --- a/fastplotlib/graphics/image.py +++ b/fastplotlib/graphics/image.py @@ -2,11 +2,13 @@ from math import ceil from itertools import product +import numpy as np import pygfx from pygfx.utils import unpack_bitfield from ._base import Graphic, Interaction, PreviouslyModifiedData from .features import ImageCmapFeature, ImageDataFeature, HeatmapDataFeature, HeatmapCmapFeature +from .features._base import to_gpu_supported_dtype from ..utils import quick_min_max @@ -23,6 +25,7 @@ def __init__( vmax: int = None, cmap: str = 'plasma', filter: str = "nearest", + isolated_buffer: bool = True, *args, **kwargs ): @@ -43,6 +46,10 @@ def __init__( colormap to use to display the image data, ignored if data is RGB filter: str, optional, default "nearest" interpolation filter, one of "nearest" or "linear" + isolated_buffer: bool, default True + If True, initialize a buffer with the same shape as the input data and then + set the data, useful if the data arrays are ready-only such as memmaps. + If False, the input array is itself used as the buffer. args: additional arguments passed to Graphic kwargs: @@ -65,20 +72,29 @@ def __init__( super().__init__(*args, **kwargs) - self.data = ImageDataFeature(self, data) + data = to_gpu_supported_dtype(data) + + # TODO: we need to organize and do this better + if isolated_buffer: + # initialize a buffer with the same shape as the input data + # we do not directly use the input data array as the buffer + # because if the input array is a read-only type, such as + # numpy memmaps, we would not be able to change the image data + buffer_init = np.zeros(shape=data.shape, dtype=data.dtype) + else: + buffer_init = data if (vmin is None) or (vmax is None): vmin, vmax = quick_min_max(data) - texture_view = pygfx.Texture(self.data(), dim=2).get_view(filter=filter) + texture_view = pygfx.Texture(buffer_init, dim=2).get_view(filter=filter) geometry = pygfx.Geometry(grid=texture_view) # if data is RGB - if self.data().ndim == 3: + if data.ndim == 3: self.cmap = None material = pygfx.ImageBasicMaterial(clim=(vmin, vmax)) - # if data is just 2D without color information, use colormap LUT else: self.cmap = ImageCmapFeature(self, cmap) @@ -89,6 +105,13 @@ def __init__( material ) + self.data = ImageDataFeature(self, data) + # TODO: we need to organize and do this better + if isolated_buffer: + # if the buffer was initialized with zeros + # set it with the actual data + self.data = data + @property def vmin(self) -> float: """Minimum contrast limit.""" @@ -176,6 +199,7 @@ def __init__( cmap: str = 'plasma', filter: str = "nearest", chunk_size: int = 8192, + isolated_buffer: bool = True, *args, **kwargs ): @@ -198,6 +222,10 @@ def __init__( interpolation filter, one of "nearest" or "linear" chunk_size: int, default 8192, max 8192 chunk size for each tile used to make up the heatmap texture + isolated_buffer: bool, default True + If True, initialize a buffer with the same shape as the input data and then + set the data, useful if the data arrays are ready-only such as memmaps. + If False, the input array is itself used as the buffer. args: additional arguments passed to Graphic kwargs: @@ -223,7 +251,17 @@ def __init__( if chunk_size > 8192: raise ValueError("Maximum chunk size is 8192") - self.data = HeatmapDataFeature(self, data) + data = to_gpu_supported_dtype(data) + + # TODO: we need to organize and do this better + if isolated_buffer: + # initialize a buffer with the same shape as the input data + # we do not directly use the input data array as the buffer + # because if the input array is a read-only type, such as + # numpy memmaps, we would not be able to change the image data + buffer_init = np.zeros(shape=data.shape, dtype=data.dtype) + else: + buffer_init = data row_chunks = range(ceil(data.shape[0] / chunk_size)) col_chunks = range(ceil(data.shape[1] / chunk_size)) @@ -249,7 +287,7 @@ def __init__( # x and y positions of the Tile in world space coordinates y_pos, x_pos = row_start, col_start - tex_view = pygfx.Texture(data[row_start:row_stop, col_start:col_stop], dim=2).get_view(filter=filter) + tex_view = pygfx.Texture(buffer_init[row_start:row_stop, col_start:col_stop], dim=2).get_view(filter=filter) geometry = pygfx.Geometry(grid=tex_view) # material = pygfx.ImageBasicMaterial(clim=(0, 1), map=self.cmap()) @@ -264,6 +302,13 @@ def __init__( self.world_object.add(img) + self.data = HeatmapDataFeature(self, buffer_init) + # TODO: we need to organize and do this better + if isolated_buffer: + # if the buffer was initialized with zeros + # set it with the actual data + self.data = data + @property def vmin(self) -> float: """Minimum contrast limit."""