explore, document, and type annotate supported array-like objects

kushalkolar commented 7 months ago

We should look into the various array-like types that can work with pygfx.Buffer and pygfx.Texture, document them, and then figure out the best way to put type annotations to indicate the various array-like types that are supported. Some types to check to get started:

torch.Tensor
jax arrays @BalzaniEdoardo
xarray?
dask?
h5py datasets (lazy loading), I think this works
zarr arrays, I think this works too
various pynapple types, @gviejo
tf.Tensor - this does work, but does anyone still use tensorflow :laughing:

Should be enough for now.

@apasarkar relates to your question on type annotations

BalzaniEdoardo commented 6 months ago

"""
Line Drawing
============

Drawing a line with a shape that makes it interesting for demonstrating/testing
various aspects of line rendering. Use the middle-mouse button to set the
position of the last point. Use '1' and '2' to toggle between normal and dashed
mode.
"""

# sphinx_gallery_pygfx_docs = 'screenshot'
# sphinx_gallery_pygfx_test = 'run'

import numpy as np
from wgpu.gui.auto import WgpuCanvas, run
import pygfx as gfx
import pylinalg as la
import zarr, h5py, jax
from pathlib import Path
import shutil

def create_scene(positions):
    scene = gfx.Scene()
    line = gfx.Line(
        gfx.Geometry(positions=positions),
        gfx.LineMaterial(thickness=22.0, color=(0.8, 0.7, 0.0), opacity=0.5),
    )
    scene.add(line)
    return line, scene

canvas = WgpuCanvas(size=(1000, 800))
renderer = gfx.WgpuRenderer(canvas)
renderer_svg = gfx.SvgRenderer(640, 480, "~/line.svg")

renderer.blend_mode = "weighted"

positions = [[200 + np.sin(i) * i * 6, 200 + np.cos(i) * i * 6, 0] for i in range(20)]
positions += [[np.nan, np.nan, np.nan]]
positions += [[400 - np.sin(i) * i * 6, 200 + np.cos(i) * i * 6, 0] for i in range(20)]
positions += [[np.nan, np.nan, np.nan]]
positions += [
    [100, 450, 0],
    [102, 450, 0],
    [104, 450, 0],
    [106, 450, 0],
    [200, 450, 0],
    [200, 445, 0],
    [400, 440, 0],
    [300, 400, 0],
    [300, 390, 0],
    [400, 370, 0],
    [350, 350, 0],
]

# Spiral away in z (to make the depth buffer less boring)
for i in range(len(positions)):
    positions[i][2] = i

camera = gfx.OrthographicCamera(600, 500)
camera.local.position = (300, 250, 0)

controller = gfx.PanZoomController(camera, register_events=renderer)

alpha = 0
d_alpha = 0.05

def animate(line, scene):
    line.material.dash_offset += 0.1
    renderer.render(scene, camera)
    canvas.request_draw()

def zarr_array_test(positions, path="", dtype=np.float32):
    # Convert list to a NumPy array
    np_array = np.asarray(positions, dtype=dtype)

    # Create a Zarr array from the NumPy array
    zarr_array = zarr.array(np_array, chunks=(2,), dtype=np_array.dtype)
    fh_path = Path(path) / 'store.zarr'
    try:
        with zarr.open(fh_path, mode='w') as root:
            root['positions'] = zarr_array
            # print(root["positions"])
            # print(type(root["positions"]))
            line, scene = create_scene(root["positions"])
            renderer_svg.render(scene, camera)
            canvas.request_draw(lambda: animate(line, scene))
            run()
    finally:
        if fh_path.exists():
            shutil.rmtree(fh_path)

def hdf5_array_test(positions, path="", dtype=np.float32):
    # Convert list to a NumPy array
    np_array = np.asarray(positions, dtype=dtype)

    # Prepare file path for the HDF5 store
    fh_path = Path(path) / 'store.h5'

    try:
        # Open the HDF5 file for writing
        with h5py.File(fh_path, 'w') as f:
            # Create dataset from the NumPy array
            dset = f.create_dataset('positions', data=np_array)
            # print(dset)
            # print(type(dset))

            # Suppose functions to create a scene and render it (not defined in this snippet)
            line, scene = create_scene(dset) 
            renderer_svg.render(scene, camera)
            canvas.request_draw(lambda: animate(line, scene))
            run()
    finally:
        # Clean up: remove the HDF5 file
        if fh_path.exists():
            fh_path.unlink()  # Remove the file directly

def jax_array_test(positions, dtype=np.float32):
    # Convert list to a NumPy array
    np_array = jax.numpy.asarray(positions, dtype=dtype)

    # Suppose functions to create a scene and render it (not defined in this snippet)
    line, scene = create_scene(np_array) 
    renderer_svg.render(scene, camera)
    canvas.request_draw(lambda: animate(line, scene))
    run()

if __name__ == "__main__":
    tests = [zarr_array_test, hdf5_array_test, jax_array_test]
    print("\n")
    for func in tests:
        try:
            func(positions)
            print(f"success {func}\n")
        except Exception as e:
            print(f"fail {func} with exception")
            print(e)
            print("\n")

I put together this test; This fails for all array type

Zarr array

Traceback (most recent call last):
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 142, in <module>
    func(positions)
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 91, in zarr_array_test
    line, scene = create_scene(root["positions"])
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 25, in create_scene
    gfx.Geometry(positions=positions),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/geometries/_base.py", line 54, in __init__
    resource = Buffer(val)
               ^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/resources/_buffer.py", line 54, in __init__
    self._mem = mem = memoryview(data)
                      ^^^^^^^^^^^^^^^^
TypeError: memoryview: a bytes-like object is required, not 'Array'

h5py Dataset


Traceback (most recent call last):
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 142, in <module>
    func(positions)
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 116, in hdf5_array_test
    line, scene = create_scene(dset)  # Access data with slicing
                  ^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 25, in create_scene
    gfx.Geometry(positions=positions),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/geometries/_base.py", line 54, in __init__
    resource = Buffer(val)
               ^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/resources/_buffer.py", line 54, in __init__
    self._mem = mem = memoryview(data)
                      ^^^^^^^^^^^^^^^^
TypeError: memoryview: a bytes-like object is required, not 'Dataset'

JAX array

Traceback (most recent call last):
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 142, in <module>
    func(positions)
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 131, in jax_array_test
    line, scene = create_scene(np_array)  # Access data with slicing
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/Code/fastplotlib/_script/line_pygfx_test.py", line 25, in create_scene
    gfx.Geometry(positions=positions),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/geometries/_base.py", line 54, in __init__
    resource = Buffer(val)
               ^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/resources/_buffer.py", line 55, in __init__
    subformat = get_item_format_from_memoryview(mem)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ebalzani/.pyenv/versions/fastplotlib/lib/python3.11/site-packages/pygfx/resources/_base.py", line 52, in get_item_format_from_memoryview
    raise TypeError(
TypeError: Cannot convert '=f' to wgpu format. You should provide data with a different dtype.

kushalkolar commented 6 months ago

From Edoardo, seems like np.asarray(<jax array of any dtype>).astype(np.float32) works for all cases (if original data is int or float jax) and there is no copy operation.

BalzaniEdoardo commented 6 months ago

yes, this example works

"""
Line Drawing
============

Drawing a line with a shape that makes it interesting for demonstrating/testing
various aspects of line rendering. Use the middle-mouse button to set the
position of the last point. Use '1' and '2' to toggle between normal and dashed
mode.
"""

# sphinx_gallery_pygfx_docs = 'screenshot'
# sphinx_gallery_pygfx_test = 'run'

import numpy as np
from wgpu.gui.auto import WgpuCanvas, run
import pygfx as gfx
import pylinalg as la
import zarr, h5py, jax
from pathlib import Path
import shutil

def create_scene(positions):
    scene = gfx.Scene()
    line = gfx.Line(
        gfx.Geometry(positions=positions),
        gfx.LineMaterial(thickness=22.0, color=(0.8, 0.7, 0.0), opacity=0.5),
    )
    scene.add(line)
    return line, scene

canvas = WgpuCanvas(size=(1000, 800))
renderer = gfx.WgpuRenderer(canvas)
renderer_svg = gfx.SvgRenderer(640, 480, "~/line.svg")

renderer.blend_mode = "weighted"

positions = [[200 + np.sin(i) * i * 6, 200 + np.cos(i) * i * 6, 0] for i in range(20)]
positions += [[np.nan, np.nan, np.nan]]
positions += [[400 - np.sin(i) * i * 6, 200 + np.cos(i) * i * 6, 0] for i in range(20)]
positions += [[np.nan, np.nan, np.nan]]
positions += [
    [100, 450, 0],
    [102, 450, 0],
    [104, 450, 0],
    [106, 450, 0],
    [200, 450, 0],
    [200, 445, 0],
    [400, 440, 0],
    [300, 400, 0],
    [300, 390, 0],
    [400, 370, 0],
    [350, 350, 0],
]

# Spiral away in z (to make the depth buffer less boring)
for i in range(len(positions)):
    positions[i][2] = i

camera = gfx.OrthographicCamera(600, 500)
camera.local.position = (300, 250, 0)

controller = gfx.PanZoomController(camera, register_events=renderer)

alpha = 0
d_alpha = 0.05

def animate(line, scene):
    line.material.dash_offset += 0.1
    renderer.render(scene, camera)
    canvas.request_draw()

def zarr_array_test(positions, path="", dtype=np.float32):
    # Convert list to a NumPy array
    np_array = np.asarray(positions, dtype=dtype)

    # Create a Zarr array from the NumPy array
    zarr_array = zarr.array(np_array, chunks=(2,), dtype=np_array.dtype)
    fh_path = Path(path) / 'store.zarr'
    try:
        with zarr.open(fh_path, mode='w') as root:
            root['positions'] = zarr_array
            # print(root["positions"])
            # print(type(root["positions"]))
            line, scene = create_scene(np.asarray(root["positions"]).astype(np.float32))  # call asarray
            renderer_svg.render(scene, camera)
            canvas.request_draw(lambda: animate(line, scene))
            run()
    finally:
        if fh_path.exists():
            shutil.rmtree(fh_path)

def hdf5_array_test(positions, path="", dtype=np.float32):
    # Convert list to a NumPy array
    np_array = np.asarray(positions, dtype=dtype)

    # Prepare file path for the HDF5 store
    fh_path = Path(path) / 'store.h5'

    try:
        # Open the HDF5 file for writing
        with h5py.File(fh_path, 'w') as f:
            # Create dataset from the NumPy array
            dset = f.create_dataset('positions', data=np_array)
            # print(dset)
            # print(type(dset))

            # Suppose functions to create a scene and render it (not defined in this snippet)
            line, scene = create_scene(np.asarray(dset).astype(np.float32))  # Access data with slicing
            renderer_svg.render(scene, camera)
            canvas.request_draw(lambda: animate(line, scene))
            run()
    finally:
        # Clean up: remove the HDF5 file
        if fh_path.exists():
            fh_path.unlink()  # Remove the file directly

def jax_array_test(positions, dtype=np.float32):
    # Convert list to a NumPy array
    positions = np.round(positions)
    np_array = np.asarray(jax.numpy.asarray(positions, dtype=int)).astype(np.float32)

    # Suppose functions to create a scene and render it (not defined in this snippet)
    line, scene = create_scene(np_array)  # Access data with slicing
    renderer_svg.render(scene, camera)
    canvas.request_draw(lambda: animate(line, scene))
    run()

if __name__ == "__main__":
    tests = [hdf5_array_test, zarr_array_test, jax_array_test]
    print("\n")
    for func in tests:
        try:
            func(positions)
            print(f"success {func}\n")
        except ValueError as e:
            print(f"fail {func} with exception")
            print(e)
            print("\n")

BalzaniEdoardo commented 6 months ago

@kushalkolar note that jax array's are immutable and arr = numpy.asarray(jax.numpy.array([...])) is also immutable

kushalkolar commented 6 months ago

So with type annotations, we can just use numpy arrays everywhere in fpl, and use np.asarray(<memview-of-other-array>) when the user passes them in fpl. I think we can do this in a small PR after #511

kushalkolar commented 3 months ago

checked some thing with Amol, seems like using np.asarray(torch_tensor) works without copying :D

kushalkolar commented 3 months ago

@almarklein you were wondering about this

fastplotlib / fastplotlib