[WIP] spherical volume rendering experiment (do not merge)

.gitignore

@@ -1,3 +1,5 @@
+# temp files
+*.swp
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

examples/amr_spherical_volume_rendering.py

@@ -0,0 +1,50 @@
+import sys
+
+import numpy as np
+import yt
+
+import yt_idv
+
+# yt reminder: phi is the polar angle (0 to 2pi)
+# theta is the angle from north (0 to pi)
+
+
+# coord ordering here will be r, phi, theta
+
+bbox_options = {
+    "partial": np.array([[0.5, 1.0], [0.0, np.pi / 4], [np.pi / 4, np.pi / 2]]),
+    "whole": np.array([[0.1, 1.0], [0.0, 2 * np.pi], [0, np.pi]]),
+    "north_hemi": np.array([[0.1, 1.0], [0.0, 2 * np.pi], [0, 0.5 * np.pi]]),
+    "south_hemi": np.array([[0.1, 1.0], [0.0, 2 * np.pi], [0.5 * np.pi, np.pi]]),
+    "ew_hemi": np.array([[0.1, 1.0], [0.0, np.pi], [0.0, np.pi]]),
+}
+
+
+sz = (256, 256, 256)
+fake_data = {"density": np.random.random(sz)}
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        bbox_type = sys.argv[1]
+    else:
+        bbox_type = "partial"
+
+    bbox = bbox_options[bbox_type]
+
+    ds = yt.load_uniform_grid(
+        fake_data,
+        sz,
+        bbox=bbox,
+        nprocs=4096,
+        geometry=("spherical", ("r", "phi", "theta")),
+        length_unit="m",
+    )
+
+    rc = yt_idv.render_context(height=800, width=800, gui=True)
+    dd = ds.all_data()
+    dd.max_level = 1
+    # sg = rc.add_scene(ds, ("index", "r"), no_ghost=True)
+    # sg = rc.add_scene(ds, ("index", "theta"), no_ghost=True)
+    sg = rc.add_scene(ds, ("index", "phi"), no_ghost=True)
+    sg.camera.focus = [0.0, 0.0, 0.0]
+    rc.run()

examples/spherical_subset_volume_rendering.py

@@ -0,0 +1,33 @@
+import numpy as np
+import yt
+
+import yt_idv
+
+# Spherical Test (to line 20)
+
+NDIM = 32
+
+bbox = np.array(
+    [[0.0, 0.5], [np.pi / 8, 2 * np.pi / 8], [2 * np.pi / 8, 3 * np.pi / 8]]
+)
+
+fake_data = {"density": np.random.random((NDIM, NDIM, NDIM))}
+ds = yt.load_uniform_grid(
+    fake_data,
+    [NDIM, NDIM, NDIM],
+    bbox=bbox,
+    geometry="spherical",
+)
+
+rc = yt_idv.render_context(height=800, width=800, gui=True)
+dd = ds.all_data()
+dd.max_level = 1
+sg = rc.add_scene(ds, ("index", "r"), no_ghost=True)
+sg.camera.focus = [0.0, 0.0, 0.0]
+rc.run()
+
+# Cartesian Test (to line 25)
+# ds = yt.load_sample("IsolatedGalaxy")
+# rc = yt_idv.render_context(height=800, width=800, gui=True)
+# sg = rc.add_scene(ds, "density", no_ghost=True)
+# rc.run()

tests/test_spherical_vol_rendering.py

@@ -0,0 +1,44 @@
+import os
+
+import numpy as np
+import pytest
+import yt
+
+import yt_idv
+
+
+@pytest.fixture()
+def osmesa_fake_spherical():
+    """Return an OSMesa context that has a "fake" AMR dataset added, with "radius"
+    as the field.
+    """
+
+    sz = (10, 10, 10)
+    fake_data = {"density": np.random.random(sz)}
+
+    bbox = np.array([[0.1, 1.0], [0.0, 2 * np.pi], [0, np.pi]])
+
+    ds = yt.load_uniform_grid(
+        fake_data,
+        sz,
+        bbox=bbox,
+        nprocs=1,
+        geometry=("spherical", ("r", "phi", "theta")),
+        length_unit="m",
+    )
+    dd = ds.all_data()
+    rc = yt_idv.render_context("osmesa", width=1024, height=1024)
+    rc.add_scene(dd, "density", no_ghost=True)
+    yield rc
+    rc.osmesa.OSMesaDestroyContext(rc.context)
+
+
+def test_spherical(osmesa_fake_spherical, tmp_path):
+    # just checking that it runs here
+    image = osmesa_fake_spherical.run()
+
+    outdir = tmp_path / "snapshots"
+    outdir.mkdir()
+    fn = str(outdir / "testout.png")
+    _ = yt.write_bitmap(image, fn)
+    assert os.path.exists(fn)

yt_idv/scene_components/base_component.py

@@ -62,7 +62,7 @@ class SceneComponent(traitlets.HasTraits):
     # These attributes are
     cmap_min = traitlets.CFloat(None, allow_none=True)
     cmap_max = traitlets.CFloat(None, allow_none=True)
-    cmap_log = traitlets.Bool(True)
+    cmap_log = traitlets.Bool(False)
     scale = traitlets.CFloat(1.0)
 
     @traitlets.observe("display_bounds")

yt_idv/scene_components/blocks.py

@@ -139,6 +139,15 @@ def draw(self, scene, program):
                         GL.glDrawArrays(GL.GL_POINTS, tex_ind * each, each)
 
     def _set_uniforms(self, scene, shader_program):
+        shader_program._set_uniform(
+            "is_spherical", self.data.data_source.ds.geometry == "spherical"
+        )
+        if self.data.data_source.ds.geometry == "spherical":
+            axis_id = self.data.data_source.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"])
+
         shader_program._set_uniform("box_width", self.box_width)
         shader_program._set_uniform("sample_factor", self.sample_factor)
         shader_program._set_uniform("ds_tex", np.array([0, 0, 0, 0, 0, 0]))

yt_idv/scene_data/_geometry_utils.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+
+def phi_normal_planes(edge_coordinates, axis_id, cast_type: str = None):
+    # for spherical geometries, calculates the cartesian normals and constants
+    # defining the planes normal to a fixed-phi value. The phi normal plane for
+    # a given spherical coordinate (r, theta, phi) will contain the given
+    # coordinate and the z-axis.
+    #
+    # edge_coordinates: 3D array of shape (N, 3) containing the spherical
+    #                   coordinates for which we want the phi-normal.
+    # axis_id: dictionary that maps the spherical coordinate axis names to the
+    #          index number.
+
+    phi = edge_coordinates[:, axis_id["phi"]]
+    theta = edge_coordinates[:, axis_id["theta"]]
+    r = edge_coordinates[:, axis_id["r"]]
+
+    # get the cartesian values of the coordinates
+    z = r * np.cos(theta)
+    r_xy = r * np.sin(theta)
+    x = r_xy * np.cos(phi)
+    y = r_xy * np.sin(phi)
+    xyz = np.column_stack([x, y, z])
+
+    # construct the planes
+    z_hat = np.array([0, 0, 1])
+    # cross product is vectorized, result is shape (N, 3):
+    normal_vec = np.cross(xyz, z_hat)
+    # dot product is not vectorized, do it manually via an elemntwise multiplication
+    # then summation. result will have shape (N,)
+    d = (normal_vec * xyz).sum(axis=1)  # manual dot product
+
+    normals_d = np.column_stack([normal_vec, d])
+    if cast_type is not None:
+        normals_d = normals_d.astype(cast_type)
+    return normals_d

yt_idv/scene_data/block_collection.py

@@ -5,6 +5,7 @@
 from yt.data_objects.data_containers import YTDataContainer
 
 from yt_idv.opengl_support import Texture3D, VertexArray, VertexAttribute
+from yt_idv.scene_data._geometry_utils import phi_normal_planes
 from yt_idv.scene_data.base_data import SceneData
 
 
@@ -41,6 +42,7 @@ def add_data(self, field, no_ghost=False):
         # 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 = [], [], [], []
+
         self.min_val = +np.inf
         self.max_val = -np.inf
         if self.scale:
@@ -77,27 +79,60 @@ def add_data(self, field, no_ghost=False):
         LE = np.array([b.LeftEdge for i, b in self.blocks.values()]).min(axis=0)
         RE = np.array([b.RightEdge for i, b in self.blocks.values()]).max(axis=0)
         self.diagonal = np.sqrt(((RE - LE) ** 2).sum())
+
         # Now we set up our buffer
         vert = np.array(vert, dtype="f4")
-        units = self.data_source.ds.units
-        ratio = (units.code_length / units.unitary).base_value
-        dx = np.array(dx, dtype="f4") * ratio
-        le = np.array(le, dtype="f4") * ratio
-        re = np.array(re, dtype="f4") * ratio
+        dx = np.array(dx, dtype="f4")
+        le = np.array(le, dtype="f4")
+        re = np.array(re, dtype="f4")
+        if self._yt_geom_str == "cartesian":
+            units = self.data_source.ds.units
+            ratio = (units.code_length / units.unitary).base_value
+            dx = dx * ratio
+            le = le * ratio
+            re = re * ratio
+        self._set_geometry_attributes(le, re)
         self.vertex_array.attributes.append(
             VertexAttribute(name="model_vertex", data=vert)
         )
         self.vertex_array.attributes.append(VertexAttribute(name="in_dx", data=dx))
         self.vertex_array.attributes.append(
-            VertexAttribute(name="in_left_edge", data=le)
+            VertexAttribute(name="in_left_edge", data=le.astype("f4"))
         )
         self.vertex_array.attributes.append(
-            VertexAttribute(name="in_right_edge", data=re)
+            VertexAttribute(name="in_right_edge", data=re.astype("f4"))
         )
 
         # Now we set up our textures
         self._load_textures()
 
+    @property
+    def _yt_geom_str(self):
+        # 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
+        return str(self.data_source.ds.geometry)
+
+    def _set_geometry_attributes(self, le, re):
+        # set any vertex_array attributes that depend on the yt geometry type
+
+        if self._yt_geom_str == "cartesian":
+            return
+        elif self._yt_geom_str == "spherical":
+            axis_id = self.data_source.ds.coordinates.axis_id
+            phi_plane_le = phi_normal_planes(le, axis_id, cast_type="f4")
+            phi_plane_re = phi_normal_planes(re, axis_id, cast_type="f4")
+            self.vertex_array.attributes.append(
+                VertexAttribute(name="phi_plane_le", data=phi_plane_le)
+            )
+            self.vertex_array.attributes.append(
+                VertexAttribute(name="phi_plane_re", data=phi_plane_re)
+            )
+        else:
+            raise NotImplementedError(
+                f"{self.name} does not implement {self._yt_geom_str} geometries."
+            )
+
     def viewpoint_iter(self, camera):
         for block in self.data_source.tiles.traverse(viewpoint=camera.position):
             vbo_i, _ = self.blocks[id(block)]

yt_idv/shaders/grid_expand.geom.glsl

@@ -1,25 +1,32 @@
 layout ( points ) in;
 layout ( triangle_strip, max_vertices = 14 ) out;
 
+// note: all in/out variables below are always in native coordinates (e.g.,
+// spherical or cartesian) except when noted.
 flat in vec3 vdx[];
 flat in vec3 vleft_edge[];
 flat in vec3 vright_edge[];
+flat in vec4 vphi_plane_le[];
+flat in vec4 vphi_plane_re[];
 
 flat out vec3 dx;
 flat out vec3 left_edge;
 flat out vec3 right_edge;
-flat out mat4 inverse_proj;
-flat out mat4 inverse_mvm;
-flat out mat4 inverse_pmvm;
+flat out mat4 inverse_proj; // always cartesian
+flat out mat4 inverse_mvm; // always cartesian
+flat out mat4 inverse_pmvm; // always cartesian
 out vec4 v_model;
 
-flat in mat4 vinverse_proj[];
-flat in mat4 vinverse_mvm[];
-flat in mat4 vinverse_pmvm[];
+flat in mat4 vinverse_proj[];  // always cartesian
+flat in mat4 vinverse_mvm[]; // always cartesian
+flat in mat4 vinverse_pmvm[]; // always cartesian
 flat in vec4 vv_model[];
 
 flat out ivec3 texture_offset;
 
+flat out vec4 phi_plane_le;
+flat out vec4 phi_plane_re;
+
 // https://stackoverflow.com/questions/28375338/cube-using-single-gl-triangle-strip
 // suggests that the triangle strip we want for the cube is
 
@@ -36,26 +43,46 @@ uniform vec3 arrangement[8] = vec3[](
 
 uniform int aindex[14] = int[](6, 7, 4, 5, 1, 7, 3, 6, 2, 4, 0, 1, 2, 3);
 
+vec3 transform_vec3(vec3 v) {
+    if (is_spherical) {
+        // in yt, phi is the polar angle from (0, 2pi), theta is the azimuthal
+        // angle (0, pi). the id_ values below are uniforms that depend on the
+        // yt dataset coordinate ordering
+        return vec3(
+            v[id_r] * sin(v[id_theta]) * cos(v[id_phi]),
+            v[id_r] * sin(v[id_theta]) * sin(v[id_phi]),
+            v[id_r] * cos(v[id_theta])
+        );
+    } else {
+        return v;
+    }
+}
+
 void main() {
 
-    vec4 center = gl_in[0].gl_Position;
+    vec4 center = gl_in[0].gl_Position;  // always cartesian
 
     vec3 width = vright_edge[0] - vleft_edge[0];
 
     vec4 newPos;
+    vec3 current_pos;
 
     for (int i = 0; i < 14; i++) {
-        newPos = vec4(vleft_edge[0] + width * arrangement[aindex[i]], 1.0);
-        gl_Position = projection * modelview * newPos;
+        // walks through each vertex of the triangle strip, emit it. need to
+        // emit gl_Position in cartesian, but pass native coords out in v_model
+        current_pos = vec3(vleft_edge[0] + width * arrangement[aindex[i]]);
+        newPos = vec4(transform_vec3(current_pos), 1.0); // cartesian
+        gl_Position = projection * modelview * newPos;  // cartesian
         left_edge = vleft_edge[0];
         right_edge = vright_edge[0];
         inverse_pmvm = vinverse_pmvm[0];
         inverse_proj = vinverse_proj[0];
         inverse_mvm = vinverse_mvm[0];
+        phi_plane_le = vphi_plane_le[0];
+        phi_plane_re = vphi_plane_re[0];
         dx = vdx[0];
-        v_model = newPos;
+        v_model = vec4(current_pos, 1.0);
         texture_offset = ivec3(0);
         EmitVertex();
     }
-
 }