diff --git a/examples/grid_collections_enzo_64.py b/examples/grid_collections_enzo_64.py new file mode 100644 index 0000000..a0e1117 --- /dev/null +++ b/examples/grid_collections_enzo_64.py @@ -0,0 +1,27 @@ +import yt + +import yt_idv +from yt_idv.cameras.trackball_camera import TrackballCamera +from yt_idv.scene_components.blocks import GridCollectionRendering +from yt_idv.scene_data.block_collection import GridCollection +from yt_idv.scene_graph import SceneGraph + +ds = yt.load_sample("Enzo_64") + +# define a couple of arbitrary grids. the selected field will be sampled on +# each grid then loaded in as 3D textures (without any refinement). +le = ds.domain_center - ds.domain_width / 4 +re = ds.domain_center + ds.domain_width / 4 +ag1 = ds.arbitrary_grid(le, re, [64, 64, 64]) +ag2 = ds.arbitrary_grid(re, ds.domain_right_edge, [32, 32, 32]) + +rc = yt_idv.render_context(height=800, width=800, gui=True) + +c = TrackballCamera.from_dataset(ds) +rc.scene = SceneGraph(camera=c) +grid_coll = GridCollection(data_source=[ag1, ag2]) +grid_coll.add_data(("gas", "density"), no_ghost=True) +rc.scene.data_objects.append(grid_coll) +rc.scene.components.append(GridCollectionRendering(data=grid_coll)) +c.set_position([-1.0727880001068115, 1.6017001867294312, 2.250051736831665]) +rc.run() diff --git a/examples/grid_collections_fake_data.py b/examples/grid_collections_fake_data.py new file mode 100644 index 0000000..ffed8a7 --- /dev/null +++ b/examples/grid_collections_fake_data.py @@ -0,0 +1,117 @@ +import argparse + +import numpy as np +import yt + +import yt_idv +from yt_idv.cameras.trackball_camera import TrackballCamera +from yt_idv.scene_components.blocks import GridCollectionRendering +from yt_idv.scene_data.block_collection import GridCollection +from yt_idv.scene_graph import SceneGraph + +_bboxes_by_geom = { + "cartesian": np.array([[-1, 1], [-1, 1], [-1, 1]]), + "spherical": np.array([[0, 1], [0, np.pi], [0, 2 * np.pi]]), +} + +_camera_pos = { + "cartesian": [-0.13453812897205353, -1.2374168634414673, 2.3244969844818115], + "spherical": [2.618459939956665, 3.529810905456543, -2.2845287322998047], +} + +_fields = { + "cartesian": (("stream", "density"), False), + "spherical": (("index", "theta"), False), +} + + +def _get_ags(ds, geom): + ags = [] + if geom == "spherical": + le = ds.domain_center.copy() + le[0] = le[0] + ds.domain_width[0] / 4 + ags.append(ds.arbitrary_grid(le, ds.domain_right_edge, [64, 64, 64])) + + hwid = ds.domain_width / 2 + le = ds.domain_left_edge + hwid * np.array([0.0, 0.0, 1.0]) + re = le + hwid + ags.append(ds.arbitrary_grid(le, re, [32, 32, 32])) + else: + ags.append( + ds.arbitrary_grid(ds.domain_center, ds.domain_right_edge, [64, 64, 64]) + ) + ags.append( + ds.arbitrary_grid(ds.domain_left_edge, ds.domain_center, [32, 32, 32]) + ) + return ags + + +def _get_cgs(ds, geom): + cgs = [] + if geom == "spherical": + le = ds.domain_center.copy() + le[0] = le[0] + ds.domain_width[0] / 4 + cgs.append(ds.covering_grid(0, le, 8)) + + hwid = ds.domain_width / 2 + le = ds.domain_left_edge + hwid * np.array([0.0, 0.0, 1.0]) + cgs.append(ds.covering_grid(0, le, [16, 16, 16])) + else: + cgs.append(ds.covering_grid(0, ds.domain_left_edge, [16, 16, 16])) + cgs.append(ds.covering_grid(0, ds.domain_center, [16, 16, 16])) + return cgs + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser( + prog="grid_collections", + description="Methods of loading lists of covering or arbitrary grids", + ) + + msg = "The geometry to use: cartesian (default) or spherical" + parser.add_argument("-g", "--geometry", default="cartesian", help=msg) + msg = "arbitrary (default) or covering" + parser.add_argument("--grid_type", default="arbitrary", help=msg) + + args = parser.parse_args() + geometry = args.geometry + + sz = (32, 32, 32) + fake_data = {"density": np.random.random(sz)} + ds = yt.load_uniform_grid( + fake_data, + sz, + bbox=_bboxes_by_geom[geometry], + geometry=geometry, + length_unit="m", + ) + + if args.grid_type == "arbitrary": + grids = _get_ags(ds, geometry) + else: + grids = _get_cgs(ds, geometry) + + rc = yt_idv.render_context(height=800, width=800, gui=True) + + c = TrackballCamera.from_dataset(ds) + rc.scene = SceneGraph(camera=c) + grid_coll = GridCollection(data_source=grids) + grid_coll.add_data(_fields[geometry][0], no_ghost=True) + rc.scene.data_objects.append(grid_coll) + rc.scene.components.append(GridCollectionRendering(data=grid_coll)) + + if _fields[geometry][1] is False: + rc.scene.components[0].cmap_log = False + + cpos = _camera_pos.get(geometry, None) + if cpos: + rc.scene.camera.set_position(cpos) + + rc.scene.camera.focus = ( + 0, + 0, + 0, + ) + + rc.run() diff --git a/yt_idv/scene_components/blocks.py b/yt_idv/scene_components/blocks.py index 3e51f70..4978329 100644 --- a/yt_idv/scene_components/blocks.py +++ b/yt_idv/scene_components/blocks.py @@ -7,7 +7,7 @@ from yt_idv.gui_support import add_popup_help from yt_idv.opengl_support import TransferFunctionTexture from yt_idv.scene_components.base_component import SceneComponent -from yt_idv.scene_data.block_collection import BlockCollection +from yt_idv.scene_data.block_collection import BlockCollection, GridCollection from yt_idv.shader_objects import component_shaders, get_shader_combos @@ -30,6 +30,7 @@ class BlockRendering(SceneComponent): slice_position = traitlets.Tuple((0.5, 0.5, 0.5)).tag(trait=traitlets.CFloat()) slice_normal = traitlets.Tuple((1.0, 0.0, 0.0)).tag(trait=traitlets.CFloat()) + _allow_block_outlines = True priority = 10 def render_gui(self, imgui, renderer, scene): @@ -55,20 +56,22 @@ def render_gui(self, imgui, renderer, scene): if _: self.render_method = valid_shaders[shader_ind] changed = changed or _ - if imgui.button("Add Block Outline"): - if self.data._yt_geom_str == "cartesian": - from ..scene_annotations.block_outline import BlockOutline - block_outline = BlockOutline(data=self.data) - scene.annotations.append(block_outline) - elif self.data._yt_geom_str == "spherical": - from ..scene_data.block_collection import _block_collection_outlines + if self._allow_block_outlines: + if imgui.button("Add Block Outline"): + if self.data._yt_geom_str == "cartesian": + from ..scene_annotations.block_outline import BlockOutline - cc, cc_render = _block_collection_outlines( - self.data, outline_type="blocks" - ) - scene.data_objects.append(cc) - scene.components.append(cc_render) + block_outline = BlockOutline(data=self.data) + scene.annotations.append(block_outline) + elif self.data._yt_geom_str == "spherical": + from ..scene_data.block_collection import _block_collection_outlines + + cc, cc_render = _block_collection_outlines( + self.data, outline_type="blocks" + ) + scene.data_objects.append(cc) + scene.components.append(cc_render) if imgui.button("Add Grid Outline"): if self.data._yt_geom_str == "cartesian": @@ -165,7 +168,7 @@ def draw(self, scene, program): def _set_uniforms(self, scene, shader_program): if self.data._yt_geom_str == "spherical": - axis_id = self.data.data_source.ds.coordinates.axis_id + axis_id = self.data._get_ds().coordinates.axis_id shader_program._set_uniform("id_theta", axis_id["theta"]) shader_program._set_uniform("id_r", axis_id["r"]) shader_program._set_uniform("id_phi", axis_id["phi"]) @@ -184,3 +187,9 @@ def _set_uniforms(self, scene, shader_program): @property def _yt_geom_str(self): return self.data._yt_geom_str + + +class GridCollectionRendering(BlockRendering): + name = "block_rendering" + data = traitlets.Instance(GridCollection) + _allow_block_outlines = False diff --git a/yt_idv/scene_data/block_collection.py b/yt_idv/scene_data/block_collection.py index c708b46..e537634 100644 --- a/yt_idv/scene_data/block_collection.py +++ b/yt_idv/scene_data/block_collection.py @@ -1,4 +1,5 @@ from collections import defaultdict +from typing import Union import numpy as np import traitlets @@ -6,78 +7,84 @@ from yt_idv.opengl_support import Texture3D, VertexArray, VertexAttribute from yt_idv.scene_data.base_data import SceneData +from yt_idv.utilities._vector_operations import sort_points_along_ray -class BlockCollection(SceneData): - name = "block_collection" - data_source = traitlets.Instance(YTDataContainer) +class AbstractDataCollection(SceneData): + name = "abstract_data_collection" texture_objects = traitlets.Dict(value_trait=traitlets.Instance(Texture3D)) bitmap_objects = traitlets.Dict(value_trait=traitlets.Instance(Texture3D)) - blocks = traitlets.Dict(default_value=()) - scale = traitlets.Bool(False) - blocks_by_grid = traitlets.Instance(defaultdict, (list,)) - grids_by_block = traitlets.Dict(default_value=()) - _yt_geom_str = traitlets.Unicode("cartesian") compute_min_max = traitlets.Bool(True) always_normalize = traitlets.Bool(False) + blocks = traitlets.Dict( + default_value=() + ) # Note: "blocks" here are generic volumes. + _yt_geom_str = traitlets.Unicode("cartesian") + scale = traitlets.Bool(False) @traitlets.default("vertex_array") def _default_vertex_array(self): return VertexArray(name="block_info", each=1) + def _initialize_data_source(self, field, no_ghost=False): + raise NotImplementedError("need to implement") + + def _volume_iterator(self): + # a generator that yields blocks + raise NotImplementedError("need to implement") + + def _get_volume_data(self, volume): + raise NotImplementedError("need to implement") + + def _get_ds(self): + raise NotImplementedError("need to implement") + + def _post_process_blocks(self): + # a hook that gets called after all blocks are iterated over + pass + + def _finalize_add_data(self): + # a hook that is called as the final step of add_data + # useful for clearing cached data or other clean up + # operations + pass + def add_data(self, field, no_ghost=False): - r"""Adds a source of data for the block collection. - - Given a `data_source` and a `field` to populate from, adds the data - to the block collection so that is able to be rendered. - - Parameters - ---------- - data_source : YTRegion - A YTRegion object to use as a data source. - field : string - A field to populate from. - no_ghost : bool (False) - Should we speed things up by skipping ghost zone generation? - """ - self.data_source.tiles.set_fields([field], [False], no_ghost=no_ghost) - self._yt_geom_str = str(self.data_source.ds.geometry) - # note: casting to string for compatibility with new and old geometry - # attributes (now an enum member in latest yt), - # see https://github.com/yt-project/yt/pull/4244 + self._initialize_data_source(field, no_ghost=no_ghost) - # Every time we change our data source, we wipe all existing ones. - # We now set up our vertices into our current data source. vert, dx, le, re = [], [], [], [] - min_val = +np.inf max_val = -np.inf + if self.scale and self._yt_geom_str == "cartesian": left_min = np.ones(3, "f8") * np.inf right_max = np.ones(3, "f8") * -np.inf - for block in self.data_source.tiles.traverse(): + for block in self._volume_iterator(): np.minimum(left_min, block.LeftEdge, left_min) np.maximum(right_max, block.LeftEdge, right_max) scale = right_max.max() - left_min.min() - for block in self.data_source.tiles.traverse(): + for block in self._volume_iterator(): block.LeftEdge -= left_min block.LeftEdge /= scale block.RightEdge -= left_min block.RightEdge /= scale - for i, block in enumerate(self.data_source.tiles.traverse()): - min_val = min(min_val, np.nanmin(np.abs(block.my_data[0])).min()) - max_val = max(max_val, np.nanmax(np.abs(block.my_data[0])).max()) + + for i, block in enumerate(self._volume_iterator()): + min_val = min( + min_val, np.nanmin(np.abs(self._get_volume_data(block))).min() + ) + max_val = max( + max_val, np.nanmax(np.abs(self._get_volume_data(block))).max() + ) self.blocks[id(block)] = (i, block) vert.append([1.0, 1.0, 1.0, 1.0]) dds = (block.RightEdge - block.LeftEdge) / block.source_mask.shape dx.append(dds.tolist()) le.append(block.LeftEdge.tolist()) re.append(block.RightEdge.tolist()) - for g, node, (sl, _, gi) in self.data_source.tiles.slice_traverse(): - block = node.data - self.blocks_by_grid[g.id - g._id_offset].append((id(block), gi)) - self.grids_by_block[id(node.data)] = (g.id - g._id_offset, sl) + + self._post_process_blocks() if self.compute_min_max: if hasattr(min_val, "in_units"): @@ -92,11 +99,12 @@ def add_data(self, field, no_ghost=False): dx = np.array(dx, dtype="f4") le = np.array(le) re = np.array(re) + ds = self._get_ds() if self._yt_geom_str == "cartesian": # Note: the block LeftEdge and RightEdge arrays are plain np arrays in # units of code_length, so need to convert to unitary units (in range 0,1) # after the fact: - units = self.data_source.ds.units + units = ds.units ratio = (units.code_length / units.unitary).base_value dx = dx * ratio le = le * ratio @@ -105,8 +113,8 @@ def add_data(self, field, no_ghost=False): RE = np.array([b.RightEdge for i, b in self.blocks.values()]).max(axis=0) self.diagonal = np.sqrt(((RE - LE) ** 2).sum()) elif self._yt_geom_str == "spherical": - rad_index = self.data_source.ds.coordinates.axis_id["r"] - max_r = self.data_source.ds.domain_right_edge[rad_index] + rad_index = ds.coordinates.axis_id["r"] + max_r = ds.domain_right_edge[rad_index] le[:, rad_index] = le[:, rad_index] / max_r re[:, rad_index] = re[:, rad_index] / max_r dx[:, rad_index] = dx[:, rad_index] / max_r @@ -125,6 +133,27 @@ def add_data(self, field, no_ghost=False): # Now we set up our textures self._load_textures() + self._finalize_add_data() + + def _load_textures(self): + for block_id in sorted(self.blocks): + vbo_i, block = self.blocks[block_id] + n_data = self._get_volume_data(block) + n_data = np.abs(n_data).copy(order="F").astype("float32").d + # Avoid setting to NaNs + if self.max_val != self.min_val or self.always_normalize: + n_data = self._normalize_by_min_max(n_data) + # blocks filled with identically 0 values will be + # skipped by the shader, so offset by a tiny value. + # see https://github.com/yt-project/yt_idv/issues/171 + n_data[n_data == 0.0] += np.finfo(np.float32).eps + + data_tex = Texture3D(data=n_data) + bitmap_tex = Texture3D( + data=block.source_mask * 255, min_filter="nearest", mag_filter="nearest" + ) + self.texture_objects[vbo_i] = data_tex + self.bitmap_objects[vbo_i] = bitmap_tex def _set_geometry_attributes(self, le, re, dx): # set any vertex_array attributes that depend on the yt geometry type @@ -133,6 +162,7 @@ def _set_geometry_attributes(self, le, re, dx): # should already be normalized in the range of (0, 1) if self._yt_geom_str == "cartesian": + self.cartesian_center = (le + re) / 2 return elif self._yt_geom_str == "spherical": from yt_idv.utilities.coordinate_utilities import ( @@ -140,7 +170,8 @@ def _set_geometry_attributes(self, le, re, dx): cartesian_bboxes_edges, ) - axis_id = self.data_source.ds.coordinates.axis_id + ds = self._get_ds() + axis_id = ds.coordinates.axis_id # first, we need an approximation of the grid spacing # in cartesian coordinates. this is used by the @@ -187,6 +218,7 @@ def _set_geometry_attributes(self, le, re, dx): self.cart_bbox_le = domain_le self.cart_bbox_center = (domain_re + domain_le) / 2.0 self.cart_min_dx = np.min(np.linalg.norm(dx_cart)) + self.cartesian_center = (le_cart + re_cart) / 2 self.vertex_array.attributes.append( VertexAttribute(name="le_cart", data=le_cart.astype("f4")) @@ -206,6 +238,32 @@ def _set_geometry_attributes(self, le, re, dx): f"{self.name} does not implement {self._yt_geom_str} geometries." ) + def viewpoint_iter(self, camera): + raise NotImplementedError("oops, need this one for sure.") + + def filter_callback(self, callback): + raise NotImplementedError("oops, need this one for sure.") + + +class BlockCollection(AbstractDataCollection): + name = "block_collection" + data_source = traitlets.Instance(YTDataContainer) + blocks_by_grid = traitlets.Instance(defaultdict, (list,)) + grids_by_block = traitlets.Dict(default_value=()) + + def _initialize_data_source(self, field, no_ghost=False): + self.data_source.tiles.set_fields([field], [False], no_ghost=no_ghost) + self._yt_geom_str = str(self.data_source.ds.geometry) + + def _volume_iterator(self): + yield from self.data_source.tiles.traverse() + + def _get_volume_data(self, block): + return block.my_data[0] + + def _get_ds(self): + return self.data_source.ds + def viewpoint_iter(self, camera): for block in self.data_source.tiles.traverse(viewpoint=camera.position): vbo_i, _ = self.blocks[id(block)] @@ -225,24 +283,11 @@ def filter_callback(self, callback): vbo_i, _ = self.blocks[b_id] self.bitmap_objects[vbo_i].data = new_bitmap[sl] - def _load_textures(self): - for block_id in sorted(self.blocks): - vbo_i, block = self.blocks[block_id] - n_data = np.abs(block.my_data[0]).copy(order="F").astype("float32").d - # Avoid setting to NaNs - if self.max_val != self.min_val or self.always_normalize: - n_data = self._normalize_by_min_max(n_data) - # blocks filled with identically 0 values will be - # skipped by the shader, so offset by a tiny value. - # see https://github.com/yt-project/yt_idv/issues/171 - n_data[n_data == 0.0] += np.finfo(np.float32).eps - - data_tex = Texture3D(data=n_data) - bitmap_tex = Texture3D( - data=block.source_mask * 255, min_filter="nearest", mag_filter="nearest" - ) - self.texture_objects[vbo_i] = data_tex - self.bitmap_objects[vbo_i] = bitmap_tex + def _post_process_blocks(self): + for g, node, (sl, _, gi) in self.data_source.tiles.slice_traverse(): + block = node.data + self.blocks_by_grid[g.id - g._id_offset].append((id(block), gi)) + self.grids_by_block[id(node.data)] = (g.id - g._id_offset, sl) _grid_id_list = None @@ -259,8 +304,65 @@ def intersected_grids(self): return [self.data_source.ds.index.grids[gid] for gid in self.grid_id_list] +class _DummyBlock: + def __init__(self, grid) -> None: + self.grid = grid + self.LeftEdge = grid.LeftEdge + self.RightEdge = grid.RightEdge + + @property + def source_mask(self): + return np.ones(self.grid.ActiveDimensions, dtype="uint8") + + +class GridCollection(AbstractDataCollection): + name = "block_collection" + data_source = traitlets.List(trait=traitlets.Instance(YTDataContainer)) + + _field = None + _no_ghost = None + + def _initialize_data_source(self, field, no_ghost=False): + self._field = field + self._no_ghost = no_ghost + self._yt_geom_str = str(self.data_source[0].ds.geometry) + + def _volume_iterator(self): + for grid in self.data_source: + yield _DummyBlock(grid) + + def _get_volume_data(self, block): + return block.grid[self._field] + + def _get_ds(self): + return self.data_source[0].ds + + def _sorted_block_indices(self, camera): + centers = self.cartesian_center + ray_origin = camera.position + ray_direction = camera.focus - camera.position + ray_direction = ray_direction / np.linalg.norm(ray_direction) + sorted_indices = sort_points_along_ray( + centers, ray_origin, ray_direction, back_to_front=True + ) + return sorted_indices + + def viewpoint_iter(self, camera): + for vbo_i in self._sorted_block_indices(camera): + yield (vbo_i, self.texture_objects[vbo_i], self.bitmap_objects[vbo_i]) + + def _finalize_add_data(self): + # data already loaded into textures, no need to hold onto the field data + for grid in self.data_source: + grid.clear_data() + + @property + def intersected_grids(self): + return self.data_source # the data_source is already a list of grids + + def _block_collection_outlines( - block_collection: BlockCollection, + block_collection: Union[BlockCollection, GridCollection], display_name: str = "block outlines", segments_per_edge: int = 20, outline_type: str = "blocks", @@ -284,25 +386,21 @@ def _block_collection_outlines( data_collection = CurveCollection() - if outline_type == "blocks": + if outline_type == "blocks" and isinstance(block_collection, BlockCollection): block_iterator = block_collection.data_source.tiles.traverse() else: - # note this can be simplified after - # https://github.com/yt-project/yt_idv/pull/179 - gids = [gid for gid, _ in block_collection.grids_by_block.values()] - gids = np.unique(gids) - ds = block_collection.data_source.ds - block_iterator = [ds.index.grids[gid] for gid in gids] + block_iterator = block_collection.intersected_grids if block_collection._yt_geom_str == "spherical": from yt_idv.utilities.coordinate_utilities import spherical_to_cartesian # should move this down to cython to speed it up - axis_id = block_collection.data_source.ds.coordinates.axis_id + ds = block_collection._get_ds() + axis_id = ds.coordinates.axis_id n_verts = segments_per_edge + 1 rad_index = axis_id["r"] - max_r = block_collection.data_source.ds.domain_right_edge[rad_index] + max_r = ds.domain_right_edge[rad_index] for block in block_iterator: le_i = block.LeftEdge diff --git a/yt_idv/tests/test_grid_collection.py b/yt_idv/tests/test_grid_collection.py new file mode 100644 index 0000000..11fac1d --- /dev/null +++ b/yt_idv/tests/test_grid_collection.py @@ -0,0 +1,122 @@ +import numpy as np +import pytest +import yt + +from yt_idv.cameras.trackball_camera import TrackballCamera +from yt_idv.scene_annotations.grid_outlines import GridOutlines +from yt_idv.scene_components.blocks import GridCollectionRendering +from yt_idv.scene_data.block_collection import ( + GridCollection, + _block_collection_outlines, +) +from yt_idv.scene_data.grid_positions import GridPositions +from yt_idv.scene_graph import SceneGraph + +_bboxes_by_geom = { + "cartesian": np.array([[-1, 1], [-1, 1], [-1, 1]]), + "spherical": np.array([[0, 1], [0, np.pi], [0, 2 * np.pi]]), +} + +_camera_pos = { + "cartesian": [-0.13453812897205353, -1.2374168634414673, 2.3244969844818115], + "spherical": [2.618459939956665, 3.529810905456543, -2.2845287322998047], +} + +_fields = { + "cartesian": (("stream", "density"), False), + "spherical": (("index", "theta"), False), +} + + +def _get_yt_ds(geometry: str): + sz = (32, 32, 32) + fake_data = {"density": np.random.random(sz)} + ds = yt.load_uniform_grid( + fake_data, + sz, + bbox=_bboxes_by_geom[geometry], + nprocs=8, + geometry=geometry, + length_unit="m", + ) + return ds + + +def _get_ags(ds, geom): + ags = [] + if geom == "spherical": + le = ds.domain_center.copy() + le[0] = le[0] + ds.domain_width[0] / 4 + ags.append(ds.arbitrary_grid(le, ds.domain_right_edge, [64, 64, 64])) + + hwid = ds.domain_width / 2 + le = ds.domain_left_edge + hwid * np.array([0.0, 0.0, 1.0]) + re = le + hwid + ags.append(ds.arbitrary_grid(le, re, [32, 32, 32])) + else: + ags.append( + ds.arbitrary_grid(ds.domain_center, ds.domain_right_edge, [64, 64, 64]) + ) + ags.append( + ds.arbitrary_grid(ds.domain_left_edge, ds.domain_center, [32, 32, 32]) + ) + return ags + + +def _get_cgs(ds, geom): + cgs = [] + if geom == "spherical": + le = ds.domain_center.copy() + le[0] = le[0] + ds.domain_width[0] / 4 + cgs.append(ds.covering_grid(0, le, 8)) + + hwid = ds.domain_width / 2 + le = ds.domain_left_edge + hwid * np.array([0.0, 0.0, 1.0]) + cgs.append(ds.covering_grid(0, le, [16, 16, 16])) + else: + cgs.append(ds.covering_grid(0, ds.domain_left_edge, [16, 16, 16])) + cgs.append(ds.covering_grid(0, ds.domain_center, [16, 16, 16])) + return cgs + + +@pytest.mark.parametrize("geometry", ["cartesian", "spherical"]) +@pytest.mark.parametrize("grid_type", ["arbitrary", "covering"]) +def test_grid_list(osmesa_empty_rc, image_store, geometry, grid_type): + + ds = _get_yt_ds(geometry) + + if grid_type == "arbitrary": + ags = _get_ags(ds, geometry) + else: + ags = _get_cgs(ds, geometry) + + c = TrackballCamera.from_dataset(ds) + + osmesa_empty_rc.scene = SceneGraph(camera=c) + + grid_coll = GridCollection(data_source=ags) + grid_coll.add_data(_fields[geometry][0], no_ghost=True) + + osmesa_empty_rc.scene.data_objects.append(grid_coll) + osmesa_empty_rc.scene.components.append(GridCollectionRendering(data=grid_coll)) + + if _fields[geometry][1] is False: + osmesa_empty_rc.scene.components[0].cmap_log = False + + if geometry == "spherical": + outlines, outlines_render = _block_collection_outlines(grid_coll) + outlines_render.curve_rgba = (1.0, 0.0, 0.0, 1.0) + else: + outlines = GridPositions(grid_list=grid_coll.intersected_grids) + outlines_render = GridOutlines(data=outlines, box_color=(1, 0.0, 0.0)) + + osmesa_empty_rc.scene.data_objects.append(outlines) + osmesa_empty_rc.scene.components.append(outlines_render) + + cpos = _camera_pos.get(geometry, None) + if cpos: + osmesa_empty_rc.scene.camera.set_position(cpos) + osmesa_empty_rc.scene.camera.focus = ds.domain_center.d + osmesa_empty_rc.scene.render() + + image_store(osmesa_empty_rc) diff --git a/yt_idv/tests/test_vector_operations.py b/yt_idv/tests/test_vector_operations.py new file mode 100644 index 0000000..50e8bab --- /dev/null +++ b/yt_idv/tests/test_vector_operations.py @@ -0,0 +1,86 @@ +import numpy as np +import pytest + +from yt_idv.utilities._vector_operations import dist_along_ray, sort_points_along_ray + + +@pytest.mark.parametrize("idim", range(3)) +def test_projection_onto_dim(idim): + + # project along x from 0, should get the x input + pts = (np.random.random((20, 3)) - 0.5) * 2 + + ray0 = np.array([0.0, 0.0, 1.0]) + ray_dir = np.array( + [ + 0.0, + 0.0, + 0.0, + ] + ) + ray_dir[idim] = 1.0 + + tvals = dist_along_ray(pts, ray0, ray_dir) + expected = pts[:, idim] - ray0[idim] + assert np.allclose(tvals, expected) + + indx = sort_points_along_ray(pts, ray0, ray_dir, back_to_front=False) + sorted_t = tvals[indx] + assert np.all(sorted_t[1:] >= sorted_t[:-1]) + + indx = sort_points_along_ray(pts, ray0, ray_dir, back_to_front=True) + sorted_t = tvals[indx] + assert np.all(sorted_t[:-1] >= sorted_t[1:]) + + +def test_general_projection(): + pts = (np.random.random((20, 3)) - 0.5) * 2 + + ray0 = np.array([0.0, 0.0, 1.0]) + ray_dir = np.array( + [ + 0.1, + 0.5, + 5.0, + ] + ) + ray_dir = ray_dir / np.linalg.norm(ray_dir) + tvals = dist_along_ray(pts, ray0, ray_dir) + assert np.all(np.isreal(tvals)) + + indx = sort_points_along_ray(pts, ray0, ray_dir, back_to_front=False) + sorted_t = tvals[indx] + assert np.all(sorted_t[1:] >= sorted_t[:-1]) + + indx = sort_points_along_ray(pts, ray0, ray_dir, back_to_front=True) + sorted_t = tvals[indx] + assert np.all(sorted_t[:-1] >= sorted_t[1:]) + + +def test_single_point(): + pts = np.array([5.0, 10.0, 15.0]) + ray0 = np.array([0.0, 0.0, 1.0]) + ray_dir = np.array( + [ + 0.0, + 0.0, + 1.0, + ] + ) + tval = dist_along_ray(pts, ray0, ray_dir) + assert tval == pts[2] - ray0[2] + + +def test_dist_along_ray_edge_errors(): + + pts = np.random.random((10, 3, 2)) + ray0 = np.array([0.0, 0.0, 1.0]) + ray_dir = np.array( + [ + 1.0, + 0.0, + 0.0, + ] + ) + with pytest.raises(ValueError, match="pts must be at most a 2D array"): + dist_along_ray(pts, ray0, ray_dir) diff --git a/yt_idv/utilities/_vector_operations.py b/yt_idv/utilities/_vector_operations.py new file mode 100644 index 0000000..9a65497 --- /dev/null +++ b/yt_idv/utilities/_vector_operations.py @@ -0,0 +1,89 @@ +from typing import List + +import numpy as np +import numpy.typing as npt + + +def dist_along_ray( + pts: npt.NDArray[np.float64], + ray_orgin: npt.NDArray[np.float64], + ray_unit_dir: npt.NDArray[np.float64], +) -> npt.NDArray[np.float64]: + """Project an array of 3D points onto an arbitrary ray + + Parameters + ---------- + pts : npt.NDArray[np.float64] + The points to project, shape (3, N) where N is the number of + points. + ray_orgin : npt.NDArray[np.float64] + The 3D ray origin + ray_unit_dir : npt.NDArray[np.float64] + The ray direction vector, already a unit normal vector + + Returns + ------- + npt.NDArray[np.float64] + the scalar distance of each point from the ray origin, of + shape (N,) where N is the number of points + """ + if pts.ndim > 2: + raise ValueError( + f"pts must be at most a 2D array of shape (N, 3), found {pts.shape}" + ) + + if pts.ndim == 2: + assert pts.shape[1] == 3 + npts = pts.shape[0] + else: + assert pts.size == 3 + npts = 1 + + vec = pts - ray_orgin + assert vec.shape == pts.shape + + if pts.ndim == 1: + t_vals = np.dot(vec, ray_unit_dir) + else: + t_vals = np.inner(vec, ray_unit_dir) + + if pts.ndim > 1: + assert t_vals.shape == (npts,) + + return t_vals + + +def sort_points_along_ray( + pts: npt.NDArray[np.float64], + ray_orgin: npt.NDArray[np.float64], + ray_unit_dir: npt.NDArray[np.float64], + back_to_front: bool = True, +) -> List[int]: + """Project and sort an array of 3D locations along a given ray + + Parameters + ---------- + pts : npt.NDArray[np.float64] + The points to project, shape (3, N) where N is the number of + points. + ray_orgin : npt.NDArray[np.float64] + The 3D ray origin + ray_unit_dir : npt.NDArray[np.float64] + The ray direction vector, already a unit normal vector + back_to_front : bool + If True (default), then points are sorted in reverse order + + + Returns + ------- + List[int] + the indices of the pts array sorted by the distance along + the provided ray. + """ + + # project centers onto a ray + t_vals = dist_along_ray(pts, ray_orgin, ray_unit_dir) + sorted_indices = np.argsort(t_vals).tolist() + if back_to_front: + sorted_indices.reverse() + return sorted_indices