Includes basic np dlpack function that is tested to work across jax a… - #1067
Conversation
|
@kushalkolar so one issue I've been seeing is that the cuda array interface execution path seems to produce spatially transposed data. See below for a reproducible example. The reason for this is because at the following line, the transpose does different things if the input is a tensor vs. if it is a numpy array. For a tensor, my_tensor.transpose(0, 1) will swap the axes 0 and 1. Otoh, my_numpy_array.transpose(0, 1) will keep the desired shape. Is there any reason we are doing this final transpose step? For an individual ndprocessor, it does not seem necessary. I think there are two options: |
kushalkolar
left a comment
kushalkolar
left a comment
There was a problem hiding this comment.
Where does it say that the from_dlpack requires numpy >= 2.1? Anyways that's fine, v2.1 is 2 years old at this point. pygfx requires min 2.1 anyways too https://github.com/pygfx/pygfx/blob/main/pyproject.toml#L28C1-L28C55
|
|
||
| return cupy.asnumpy(arr) | ||
|
|
||
| data = np.from_dlpack(arr, device='cpu') #This requires numpy >= 2.1 |
There was a problem hiding this comment.
| data = np.from_dlpack(arr, device='cpu') #This requires numpy >= 2.1 | |
| data = np.from_dlpack(arr, device='cpu') |
Annoying that torch behavior is different from numpy. The transpose is necessary by definition, that's how the NDWidget spec is. The specified spatial dims order defines the dim order for the graphic data. Doing the transpose after it's converted to numpy sounds like the way to go. |
|
@kushalkolar Updated this -- so I've delayed the permutation operation until the sliced data is guaranteed to be a numpy array. nd vectors processor, nd image processor, and nd positions processor all subclass ndprocessor, so those are the files where i made this change. One thing to note: with this setup, the spatial functions must be written so they are applied data that has not yet been permuted. for example, if i specify data like this: dims = ['time', 'a', 'b'] The spatial window function received the spatial dimensions in order (a, b) with these changes. I actually think this is ok -- this forces the user to write all their window/spatial function logic using the dimension order specified by dims, but the final order is transposed based on what's in spatial_dims. So we would ofc need to update the specs. I think so long as the user can write their whole pipeline (raw data --> displayed graphic) with a fixed dimension ordering it will be clear. Let me know what you think! |
I think that makes sense. The whole point is that window_funcs and spatial_funcs should be fast and async, and things are naturally async when they remain in torch as long as possible. The Maybe we should rename the arg to |
|
ordered_spatial_dims might be a good one, specifying the two roles of this parameter: (1) it specifies what the spatial dims are and (2) it gives an order to output the spatial dims. thoughts? |
|
@kushalkolar I'll also flag the following potential issue: At this line, a function "squeeze" is used. On a first read, I expected it to fail with torch tensors, since the documented parameters are tensor.squeeze(dim = ...), whereas the code is doing .squeeze(axis =...). Turns out at some point in pytorch the developers allowed axis to make it interoperable with numpy. Anyways it's not clear with other libraries that are compatible with dlpack whether this functionality will break. |
|
where is it documented that torch also accepts |
|
@kushalkolar Example: Looks like under the hood there is some code that parses keywords. |
wow that's hidden deep |
|
@kushalkolar I'm thinking from the standpoint of documentation. We ideally want to be able to tell someone "your dlpack compliant array can work here". Maybe there is not a workaround with the axis squeeze thing though |
|
If it works with torch, cupy and Jax that covers most use cases I'd think? |
|
@kushalkolar Fair enough, we can just make sure to note this in the documentation. Ok so what's left here - just the rename of the spatial_dims param to something else? If you're ok with renaming is ordered_spatial_dims, we'll have to update many of the function signatures in ndwidget, so let's make sure we're both happy with the naming here and I can go ahead and do it. |
|
I added it to the list of renames to do later in bulk on #971 |
|
what's left is the suggestion and updating pyproject.toml and merge latest |
|
@kushalkolar Added the suggested change. Re: pyproject.toml, we are now using np.from_dlpack exclusively to bring gpu arrays to cpu and from there do zero-cost conversion to numpy. So it looks like we don't need to change the pyproject.toml at all. |
For the min version pin? |
…sion across devices
|
@kushalkolar sorry I thought I did that earlier haha. |
| # Axis order of the spatial dimensions to display | ||
| self._spatial_dims_int = tuple( | ||
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | ||
| ) | ||
|
|
There was a problem hiding this comment.
spatial_dims is a mutable property, so this should be a read only property or set in the spatial_dims setter as a private attribute (probably preferable, it doesn't have to be public) and not created once in the constructor
There was a problem hiding this comment.
@kushalkolar good point, updated in latest commit. I also renamed spatial_dims_int to spatial_dims_indices for clarity
| """ | ||
| The ordered sequence of data indices that will be displayed | ||
| """ | ||
| return self._spatial_dims_indices |
There was a problem hiding this comment.
| return self._spatial_dims_indices | |
| return tuple( | |
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | |
| ) |
| @property | ||
| def spatial_dims_indices(self) -> tuple[int, ...]: | ||
| """ | ||
| The ordered sequence of data indices that will be displayed |
There was a problem hiding this comment.
| The ordered sequence of data indices that will be displayed | |
| ordered spatial dim indices that correspond to the named spatial dims |
| ## This is the ordered sequence of data indices that will be displayed | ||
| self._spatial_dims_indices = tuple( | ||
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | ||
| ) | ||
|
|
There was a problem hiding this comment.
| ## This is the ordered sequence of data indices that will be displayed | |
| self._spatial_dims_indices = tuple( | |
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | |
| ) |
| ## This is the ordered sequence of data indices that will be displayed | ||
| self._spatial_dims_indices = tuple( | ||
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | ||
| ) | ||
|
|
||
| @property | ||
| def spatial_dims_indices(self) -> tuple[int, ...]: | ||
| """ | ||
| The ordered sequence of data indices that will be displayed | ||
| """ | ||
| return self._spatial_dims_indices | ||
|
|
There was a problem hiding this comment.
| ## This is the ordered sequence of data indices that will be displayed | |
| self._spatial_dims_indices = tuple( | |
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | |
| ) | |
| @property | |
| def spatial_dims_indices(self) -> tuple[int, ...]: | |
| """ | |
| The ordered sequence of data indices that will be displayed | |
| """ | |
| return self._spatial_dims_indices |
| window_output = await run_in_thread_pool(self._executor, cuda_to_numpy, window_output) | ||
|
|
||
| return window_output | ||
|
|
| ## This is the ordered sequence of data indices that will be displayed | ||
| self._spatial_dims_indices = tuple( | ||
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | ||
| ) | ||
|
|
||
| @property | ||
| def spatial_dims_indices(self) -> tuple[int, ...]: | ||
| """ | ||
| The ordered sequence of data indices that will be displayed | ||
| """ | ||
| return self._spatial_dims_indices | ||
|
|
There was a problem hiding this comment.
| ## This is the ordered sequence of data indices that will be displayed | |
| self._spatial_dims_indices = tuple( | |
| self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims | |
| ) | |
| @property | |
| def spatial_dims_indices(self) -> tuple[int, ...]: | |
| """ | |
| The ordered sequence of data indices that will be displayed | |
| """ | |
| return self._spatial_dims_indices |
| window_output = await run_in_thread_pool(self._executor, cuda_to_numpy, window_output) | ||
|
|
||
| return window_output | ||
|
|
|
|
||
| return window_output | ||
|
|
||
| return window_output.transpose(*self._spatial_dims_indices) |
There was a problem hiding this comment.
| return window_output.transpose(*self._spatial_dims_indices) | |
| return window_output.transpose(*self.spatial_dims_indices) |
it's a public property use it
…in the base class, updates docs, uses public property in the ndprocessor subclasses
|
@kushalkolar thanks -- all makes sense. made the changes, spatial_dims_indices is computed and set in the public property of the base class now. i was hesitant to do it this way at the start because the computation to define spatial dims indices will now be run at every single getitem. but if the number of dimensions is small, the overhead should never be a problem. |
|
It's probably a nanosecond scale operation. |
* start ndprocessors * basic timeseries * add __init__ * heatmap for timeseries works! * NDPositions, basics work, reorganize, increase default scatter size * black * NDPositions working with multi-dim stack of lines, need to test window funcs * scatter collection * progress, need to change to other branch so committing * better * interpolation for heatmap * better place for check * window functions working on n_datapoints dim * p dim window funcs working for single and multiple dims I think * black * index_mappings is working I think, lightly tested on p dim * remove nd_timeseries since nd_positions is sufficient * auto-replace buffers (#974) * remove isolated_buffer * remove isolated_buffer from mixin * basics works for positions data * replaceable buffers for all positions related features * image data buffer can change * resizeable buffers for volume * black * buffer resize condition checked only if new value is an array * gc for buffer managers * uniform colors WIP * switching color modes works! * typo * balck * update tests for color_mode * update examples * backend tests passing * default for all uniforms is True * update examples * forgot * update test * example tests passing * dereferencing test and fixes * simplify texture array tests a bit * image replace buffer tests pass yay * forgot a file * comments, check image graphic * add image reshaping example * add buffer replace imgui thing for manual testing * black * dont call wgpu_obj.destroy(), seems to work and clear VRAM with normal dereferencing * slower changes * update * update example * fixes and tweaks for test * remove unecessary stuff * update * docstrings * fix example * update example * update example * update docs * some basic OOC working * max num of dipslay datapoints * scatter stack, not tested * progress * scatter collection updates * tootip handlers for ndpositions * refactoring, general NDPP_Pandas processor for any dataframe data * separate array logic and graphic logic in `ImageWidget` (#868) * start separating iw plotting and array logic * some more basics down * comment * collapse into just having a window function, no frame_function * progress * placeholder for computing histogram * formatting * remove spaghetti * more progress * basics working :D * black * most of the basics work in iw * fix * progress * progress but still broken * flippin display dims works * camera scale must be positive for MIP rendering * a very difficult to encounter iterator bug! * patch iterator caveats * mostly worksgit status * add ArrayProtocol * rename * fixes * set camera orthogonal to xy plane when going from 3d -> 2d * naming, cleaning * cleanup, correct way to push and pop dims * quality of life improvements * new histogram lut tool * new hlut tool * imagewidget rgb toggle works * more progress * support rgb(a) image volumes * ImageGraphic cleanup * cleanup, docs * fix * updates * new per-data array properties work * black formatting * fixes and other things * typing tweaks * better iterator, fix bugs * fixes * show tooltips in right clck menu * ignore nans and inf for histogram * histogram of zeros * docstrings * fix imgui pixels * iw indices event handlers only get a tuple of the indices * bugfix * fix cmap setter * spatial_func better name * bugfix * hist specify quantile * fix typos (#991) * fix typos * add rendercanvas to intersphinx_mapping * nd-iw backup * correct ImageGraphic w.r.t. ndw * last fixes in ndi * update * basic minimal ndw orchestration working * implement auto-x for timeseries * bugfix update worldobject -> graphic map for image tiles * bugfix linear selector set limits * linear selector for timeseries * return full data if display_window is Noen * arrow key to step indices * imgui separator * fix and ui stuff * both auto x range modes working * progress * moving stuff * much much better organization of things * GlobalIndexVector working with ndpostions and ndimage * examples * progress * do not reset vmin vmax when replacing Image buffer * WIP migrate to xarray * window funcs in NDProcessor class using xarray, WIP, not tested * typo * basic single index slicing working with xarray * window funcs working for NDPositions and NDPP_Pands * display_window window funcs working for NDPositions * imgui stuff * finish migrate NDImage to xarray, basics work * NDImage working mostly, behavior viz is back * better flipping logic * update examples * more progress * update example * histogram working for images * NDProcessor property aliases * more aliasing * fix * update example * fix * ui * cleanup old iw-array, imports, add deprecation warning on old iw * add ndwidget section to deps with xarray * nice repr for NDProcessor * imgui right click menu for ndgraphics * better * controller options separate window * update imgui * fix * fix compute histogram * other features WIP * basics of other features works with ScatterStack for colors, markers, sizes, need to keep testing * require min pygfx v0.16.0 due to gc hash fix necessary for NDWidget * fix PlotArea.y_range * fix to create isolated buffer for colors when buffer replaced * np.empty * cmap_transform_each WIP * progress * fix * replace graphic when data changed, tweak index_mappings * Update installation docs (#1013) * add simplejpeg to notebook deps * Update guide.rst * Update guide.rst * Update README.md * cmap lib handles image colormaps now * multi-windows ndwidget, maintain features like cmap when switching graphics * progress * lighting objects only when a mesh is added * fix * update axes only when camera or view changes * clean heatmap func * stupid print * docstrings, comments * remove unused attr, comments * remove print * add NDGraphic.pause, expose histogram widget * ndg pause in imgui * add helper function to convert heatmap timeseries to postional data shape * index wans't calling handlers * remove * allow image types other than float32 (#1027) * Includes code for doing batched transforms in both directions (#1025) * Includes code for doing batched transforms in both directions * Streamlined parsing logic and updated docstrings for model to world code * Fixes inconsistent error messages * More streamlined type checking * Gets rid of checks and uses np asarray * Some more aesthetic updates to the syntax * Update fastplotlib/graphics/_base.py --------- Co-authored-by: Kushal Kolar <kushalkolar@gmail.com> * async NDProcessor (#1026) * async NDProcessor established, NOT TESTED * ASYNC NDPROC IS WORKING :D :D CELEBRATE * comments * type annot * fix * fix * polish async integration, cuda also integrated * no longer using xarray, allow simpler ArrayProtocol * comments * comments * docs * forgot to set back to False * remove xarray from pyproject.toml * throttling * Adds nd vector graphic (#1034) * Includes nd vector code that works * Faster position assignment, no more for loop * Batched computations for vector set function * Formatting updates * Includes improved annotations and changes ordering of the data slice from the vectors graphic * Some further improvements to the pylinalg code * Fixes remaining formatting and naming issues * Apply suggestions from code review Co-authored-by: Kushal Kolar <kushalkolar@gmail.com> --------- Co-authored-by: Kushal Kolar <kushalkolar@gmail.com> * Adds kwargs so user can customize vector field estimates (#1036) * Adds kwargs so user can customize vector field estimates * Does the kwarg organization the way ndpositions does * fixes graphic constructor * implement yuv and other colorspaces and "bufferless" `TextureArray` (#1033) * implement yuv and 'bufferless' TextureArraY * unbuffered and yuv420 works * warning on tooltip * NDImage always uses unbuffered, support colorspaces in NDIMage * update docstrings * by default disable AA and set pixel_scale=1.0 for performance * unpacked yuv support * independent graphics and texture features for rgb and yuv * docstrings * docstrings * import order * yuv graphic working nicely * add enum to top level namespace * update script to produce add graphics mixin * add yuv example * update ndimage with yuv stuff * fixes * yuv video working well with NDWidget * cleanup * `SelectionVector`, `HighlightSelector`, `VisibilitySelector`, `SelectorCollection` (#1038) * start selection vector stuff, protocol * highlight selector WIP * fixes * progress * lut repeat * highlight and visibility selectors working, selection vector working * PlotArea checks for SelectorProtocol instead, move some stuff around * selector collection, not yet tested * cleanup * bugfix for linear selector limits * add example * better append, remove * cleanup * basically rewrote entire ImageHighlightSelector * black * happy with iamge selectors * toy multi session example * correct mapping stuff * better scalar/vector handling * comments * compute capabilities (#1040) * allow setting texture usage directly * ImageHighlightSelector, VisibilitySelector and ImageVisibilitySelector can handle 'None' as placeholders in the selection * fix example * comments * better axes padding for tick labels, settable axis label (#1049) * better axes padding for tick labels, settable axis label * cleanup * move axes to tools * update * move stuff around * update docs * update * update * use fully fledged async, `NDPositions` improvements (#1050) * use fully fledged async * better throttling * torch.Tensor.tranpose() doesn't like tuples * we need time-based throttling, but it can be gentler * improvements * improvements * fixes * x_range fix * fix NDPandasProcessor * fix * HighlightSelector fix to append None * cleanup better * know ndg current dipslayed indices * improve example * docstrings, cleanup * remove __array_ufunc__ from required attrs * do not require __array__() to be implemented for ndwidget * auto ref ranges * fix imgui * tear out NDTimeseries from NDPositions * imgui refactor (#1066) * start imgui window refactor * update w.r.t. imgui changes * imgui hlut colorbar stuff * better focus logic * imgui popup stuff * update docs * update examples, remove HistogramLUTTool * update examples * imgui menubar example * update docs, tweaks * docs * doc fixes * docs * fix * right click anywhere in colorbar/hlut window * inf line, dashing, thin bool (#1064) * Create partial_camera_linking.py (#1020) * inf line, dashing, thin bool * tests * new screenshots * Selection vector improve (#1061) * Full selection vector implementation * Includes some documentation at top of SelectionVector * Minor typing fix in linear selector selection setter * First working version with selection vector * Reworks the logic for adding selectors, improves some documentation, adds partial instead of lambda functions, improves typing in highlight selector * Fixes casting bug in the integer version of the code * Update fastplotlib/graphics/selectors/_selection_vector.py * Update fastplotlib/graphics/selectors/_selection_vector.py * Removes unused init time parameter and also updates the typing on the default inverse mapping function to always return an integer * Adds a nonneg check in selection vector selection setter * Adds nonneg test for inverse handler --------- Co-authored-by: Kushal Kolar <kushalkolar@gmail.com> * Includes basic np dlpack function that is tested to work across jax a… (#1067) * Includes basic np dlpack function that is tested to work across jax and torch * delays dimension permutation for all code that uses ndprocessor * Eliminates old comment * Updates the numpy dep to be the min version that allows dlpack conversion across devices * Includes code to update spatial dims indices whenever spatial dims is changed * Moves the spatial dims indices computation to the read only property in the base class, updates docs, uses public property in the ndprocessor subclasses * Filters for None values in the inv handler (#1073) * bugfix after last PR * alpha is material.opacity (#1074) * clamp vmin, vmax in colorbar histogram LUT tool (#1076) * Flat graphic features 2nd edition (#1072) * refactor cmap and cmap_transform for positional graphics * color mode stuff * WIP * more WIP * better add_graphic autogen * more progress * mostly done refactoring positions graphics cmap stuff * forgot to update mixin * cmap_transform tweaks * wip, JaggedArray * mostly works, writing tests * remove a print * image collection examples * fix * feature inheritance in PositionsGraphic and cmap_range * cmap handling * cmap tweaks * remove parse_cmap_values, rely on cmap lib * fix example, remove unused line of code * Image just uses cmap lib now too * delete_ndgraphic() * ndtimeseries cmap example * steps for stacks * ndtimeseries sets stack steps * docstrings * docstrings * comments * renames * docstring * fixes * update tests/test_colors_buffer_manager.py * update test_markers_buffer_manager.py * update plot_helpers test * update test_point_rotations_buffer_manager.py * update test_positions_graphics.py * update another test * more tests * update yet more tests * remove unused * infline tests * collections tests * updating examples, wip * more examples updates * more examples updates * more example updates * more example updates * example and docstring * docstring * remove complex line collection slicing example, stick with real usecases * better line stack example * update example * update example * update link * update more examples * examples and fixes * fix * more fixes * fix * updates * Fix * update docs * docs * add image collection to docs conf.py * update docs * Ndw rename consistency (#1078) * rename for consistency * more consistency * more * final renames * fix * cursor fix * missed renames * pandas fix * restore ImageWidget, auto-collapse NDWidget sliders UI if no slider dims and no appended UIs * imagewidget stuff * ImageVolumeGraphic args fix * NDWSubplot dims check only for array-like * collapsed property on ImaguiWindoW * update iw examples * colorbar fix * lingering rename * fix * type * clim_quantile for NDImage, other fixes * public 'max_display_datapoints' property * dw setter * imgui right click for NDGraphics * axes label tweaks * remove unused zarr placeholder * fix docstring * config system (#1080) * basic scaffold done * inheritance * done * config works! * config on graphics * axes config * full config implementation basically works * fix * mixins call Graphic construtors with kwargs nothing is positional * print * config presets * comments, docstrings * docstrings * remove ConfigValue * comments * much better add graphics mixin using descriptors, examples, fix a test * anotehr example * GlobalConfig.to_dict() * docs * reset to default config after each screenshot test * add_<graphics>() stub generator, fix maintain_aspect logic w.r.t. config stuff * change so maintain_aspect can be tested better * docstring * reset config after running each docs gallery examle * better example * docs * remove randon shit * ndw stuff is all top level importable, move types to utils * fix black * fix import * matplotlib migration docs * wording * uncommnet iw for docs conf.py * update example * update * update * temp skip nds_extras in docs * temp comment out vis selector example * fix usage args * just make usage a kwarg to ImageGraphic * update api docs stuff * update example * fix docstrings etc. * type annots * ndw import stuff * docs for nds extras * docs reqs * nds extras * ReferenceIndices.pop_dims * graphics stubs * fix for None data, remove delete_ndgraphic() for now * collections stubs * ndw docs, WIP * ndw guide * types fix * better axes ticks text outline, WIP on ndwidget playback in docs, stubs * ndw animate in docs works * fix display_dims bug * remove tranpose from ArrayProtocol * bugfix, spatial func stuff * ndw tests * black * more balck * add example * fixes * fix * update * other CI fiel * skip ndw and imgui test if not present * garbage collection of graphics in imgui popup * temp imgui ping * fix * update CI * Update screenshots.yml * update screenshots * fix docstrings * update test thresholds, remove line dashing from screenshot tests * update api docs * another example * skip testing an example * higher RMSE threshold on mac because I am so done with this shit * Update ci.yml * Update ci-pygfx-release.yml --------- Co-authored-by: Amol Pasarkar <amolpasarkar@gmail.com> Co-authored-by: Flynn <75346097+FlynnOConnell@users.noreply.github.com> Co-authored-by: Caitlin Lewis <caitlin9165@gmail.com>
2 participants
Fixes #1063
The numpy dlpack fully sidesteps the need to rely on torch, but it requires numpy > 2.1. I think this is ok, since numpy is now at 2.5, and only users who care about high-performance (i.e. computing on GPUs) will use this execution path anyways.