adding tests, example

Author: chrishavlin

examples/block_collection_normalization.py

@@ -0,0 +1,55 @@
+"""
+Sometimes you want more control over data normalization and color-mapping. This
+example demonstrates how to do that with BlockCollection data.
+"""
+
+import numpy as np
+import yt
+
+import yt_idv
+from yt_idv.cameras.trackball_camera import TrackballCamera
+from yt_idv.scene_components.blocks import BlockRendering
+from yt_idv.scene_data.block_collection import BlockCollection
+from yt_idv.scene_graph import SceneGraph
+
+# create some data
+shp = (32, 32, 32)
+data = {"constant_field": np.full(shape=shp, fill_value=5.0)}
+ds = yt.load_uniform_grid(data, shp, length_unit=1)
+
+
+rc = yt_idv.render_context(height=800, width=800, gui=True)
+c = TrackballCamera.from_dataset(ds)
+rc.scene = SceneGraph(camera=c)
+
+# the following BlockCollection initialization does the following:
+#  1. sets the min, max values to be used for data normalization when
+#     loading textures
+#  2. specifies that those min, max vals should not be re-computed
+#     (compute_min_max=False)
+#  3. specifies that the data should always be normalized -- usually
+#     normalization is skipped for constant fields.
+block_coll = BlockCollection(
+    data_source=ds.all_data(),
+    min_val=0.0,
+    max_val=10.0,
+    compute_min_max=False,
+    always_normalize=True,
+)
+block_coll.add_data(("stream", "constant_field"), no_ghost=True)
+rc.scene.data_objects.append(block_coll)
+
+# set fixed cmap ranges. Data are in normalized space. A constant data value of 5.
+# (the fill value used above) with [min val, max val] of [0., 10.] normalizes to 0.5.
+# So setting the fixed cmap min and max to 0 and 1 will result in a constant color
+# at the middle of whatever colormap is selected.
+block_rendering = BlockRendering(
+    data=block_coll,
+    fixed_cmap_min=0.0,
+    fixed_cmap_max=1.0,
+    cmap_log=False,
+)
+
+rc.scene.components.append(block_rendering)
+
+rc.run()

yt_idv/scene_components/base_component.py

@@ -77,6 +77,8 @@ class SceneComponent(traitlets.HasTraits):
     _final_pass_invalid = True
 
     # These attributes are just for colormap application
+    fixed_cmap_min = traitlets.CFloat(None, allow_none=True)
+    fixed_cmap_max = traitlets.CFloat(None, allow_none=True)
     cmap_min = traitlets.CFloat(None, allow_none=True)
     cmap_max = traitlets.CFloat(None, allow_none=True)
     cmap_log = traitlets.Bool(True)
@@ -584,16 +586,27 @@ def _construct_isolayer_table(self, imgui) -> bool:
         return changed
 
     def _reset_cmap_bounds(self, print_new_bounds=True):
+
         data = self.fb.data
         if self.use_db:
             data[:, :, :3] = self.fb.depth_data[:, :, None]
         data = data[data[:, :, 3] > 0][:, 0]
         if data.size > 0:
             self.cmap_min = data.min()
             self.cmap_max = data.max()
+
         if data.size == 0:
             self.cmap_min = 0.0
             self.cmap_max = 1.0
-        elif print_new_bounds:
+
+        # over-ride with the fixed values if they're set
+        if self.fixed_cmap_max is not None:
+            self.cmap_max = self.fixed_cmap_max
+
+        if self.fixed_cmap_min is not None:
+            self.cmap_min = self.fixed_cmap_min
+
+        if print_new_bounds:
             print(f"Computed new cmap values {self.cmap_min} - {self.cmap_max}")
+
         self._cmap_bounds_invalid = False

yt_idv/scene_data/block_collection.py

@@ -18,7 +18,8 @@ class BlockCollection(SceneData):
     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)
+    compute_min_max = traitlets.Bool(True)
+    always_normalize = traitlets.Bool(False)
 
     @traitlets.default("vertex_array")
     def _default_vertex_array(self):
@@ -78,7 +79,7 @@ def add_data(self, field, no_ghost=False):
             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)
 
-        if self._compute_min_max:
+        if self.compute_min_max:
             if hasattr(min_val, "in_units"):
                 min_val = min_val.d
             if hasattr(max_val, "in_units"):
@@ -220,7 +221,7 @@ def _load_textures(self):
             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:
+            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.

yt_idv/tests/conftest.py

@@ -5,6 +5,8 @@
 import yt
 from pytest_html import extras as html_extras
 
+import yt_idv
+
 
 def pytest_configure(config):
     # this will get run before all tests, before collection and
@@ -22,3 +24,12 @@ def _snap_image(rc):
         extras.append(html_extras.html("<br clear='all'/>"))
 
     return _snap_image
+
+
+@pytest.fixture()
+def osmesa_empty_rc():
+    """yield an OSMesa empy context then destroy"""
+
+    rc = yt_idv.render_context("osmesa", width=1024, height=1024)
+    yield rc
+    rc.osmesa.OSMesaDestroyContext(rc.context)

yt_idv/tests/test_block_normalization.py

@@ -0,0 +1,68 @@
+import numpy as np
+import pytest
+import yt
+
+from yt_idv.cameras.trackball_camera import TrackballCamera
+from yt_idv.scene_components.blocks import BlockRendering
+from yt_idv.scene_data.block_collection import BlockCollection
+from yt_idv.scene_graph import SceneGraph
+
+
+@pytest.fixture()
+def ds_yt_ugrid():
+    # create some data
+    shp = (8, 8, 8)
+    data = {"constant_field": np.full(shape=shp, fill_value=5.0)}
+    ds = yt.load_uniform_grid(data, shp, length_unit=1)
+    return ds
+
+
+def test_block_collection_normalization(osmesa_empty_rc, ds_yt_ugrid):
+
+    block_coll: BlockCollection = BlockCollection(
+        data_source=ds_yt_ugrid.all_data(),
+        min_val=0.0,
+        max_val=10.0,
+        compute_min_max=False,
+        always_normalize=True,
+    )
+    block_coll.add_data(("stream", "constant_field"), no_ghost=True)
+
+    assert np.allclose(block_coll.texture_objects[0].data, 0.5)
+
+
+def test_block_rendering_cmap_norms(osmesa_empty_rc, ds_yt_ugrid, image_store):
+
+    block_coll: BlockCollection = BlockCollection(
+        data_source=ds_yt_ugrid.all_data(),
+        min_val=0.0,
+        max_val=10.0,
+        compute_min_max=False,
+        always_normalize=True,
+    )
+    block_coll.add_data(("stream", "constant_field"), no_ghost=True)
+
+    block_rendering: BlockRendering = BlockRendering(
+        data=block_coll,
+        fixed_cmap_min=0.2,
+        fixed_cmap_max=0.8,
+        cmap_log=False,
+    )
+
+    block_rendering._reset_cmap_bounds()
+
+    assert block_rendering.cmap_min == 0.2
+    assert block_rendering.cmap_max == 0.8
+
+    block_rendering.fixed_cmap_min = 0.0
+    block_rendering.fixed_cmap_max = 1.0
+    block_rendering._reset_cmap_bounds()
+    assert block_rendering.cmap_min == 0.0
+    assert block_rendering.cmap_max == 1.0
+
+    c = TrackballCamera.from_dataset(ds_yt_ugrid)
+    osmesa_empty_rc.scene = SceneGraph(camera=c)
+    osmesa_empty_rc.scene.data_objects.append(block_coll)
+    osmesa_empty_rc.scene.components.append(block_rendering)
+
+    image_store(osmesa_empty_rc)

yt_idv/tests/test_spherical_vol_rendering.py

@@ -2,18 +2,6 @@
 import pytest
 import yt
 
-import yt_idv
-
-
-@pytest.fixture()
-def osmesa_empty_rc():
-    """yield an OSMesa empy context then destroy"""
-
-    rc = yt_idv.render_context("osmesa", width=1024, height=1024)
-    yield rc
-    rc.osmesa.OSMesaDestroyContext(rc.context)
-
-
 bbox_options = {
     "partial": {
         "bbox": np.array([[0.5, 1.0], [0.0, np.pi / 3], [np.pi / 4, np.pi / 2]]),