Skip to content

better buffer handling #150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions fastplotlib/graphics/features/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from typing import *

import numpy as np
from pygfx import Buffer
from pygfx import Buffer, Texture


supported_dtypes = [
Expand Down Expand Up @@ -226,7 +226,7 @@ def _update_range(self, key):

@property
@abstractmethod
def _buffer(self) -> Buffer:
def buffer(self) -> Union[Buffer, Texture]:
pass

@property
Expand All @@ -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")

8 changes: 4 additions & 4 deletions fastplotlib/graphics/features/_colors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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)
Expand Down
31 changes: 20 additions & 11 deletions fastplotlib/graphics/features/_data.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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__
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -97,29 +97,33 @@ 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
if len(self._event_handlers) > 0:
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:
Expand All @@ -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]

Expand All @@ -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):
Expand Down
57 changes: 51 additions & 6 deletions fastplotlib/graphics/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -23,6 +25,7 @@ def __init__(
vmax: int = None,
cmap: str = 'plasma',
filter: str = "nearest",
isolated_buffer: bool = True,
*args,
**kwargs
):
Expand All @@ -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:
Expand All @@ -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
Copy link
Member

@clewis7 clewis7 Mar 10, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so, the point of having the option for using an isolated buffer is to be able to change data in the event that the type of the data is read only? will this occur when using improv?

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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

before, Texture and Geometry were all under the hood correct?


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)
Expand All @@ -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."""
Expand Down Expand Up @@ -176,6 +199,7 @@ def __init__(
cmap: str = 'plasma',
filter: str = "nearest",
chunk_size: int = 8192,
isolated_buffer: bool = True,
*args,
**kwargs
):
Expand All @@ -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:
Expand All @@ -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))
Expand All @@ -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())

Expand All @@ -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."""
Expand Down