Spherical volume rendering

.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()

yt_idv/scene_components/base_component.py

@@ -57,7 +57,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", 0)

yt_idv/scene_data/block_collection.py

@@ -41,6 +41,39 @@ 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 = [], [], [], []
+
+        # spherical-only: SHOULD ONLY DO THIS IF DATASET IS SPHERICAL
+        # normal vectors and plane constants for min/max phi face
+        phi_plane_le = []
+        phi_plane_re = []
+
+        axis_id = self.data_source.ds.coordinates.axis_id
+
+        def calculate_phi_normal_plane(le_or_re):
+            # calculates the parameters for a plane parallel to the z-axis and
+            # passing through the x-y values defined by the polar angle phi
+            # the result plane defines the boundary on +/-phi faces.
+            #
+            # eq of plane:
+            # X dot N = d where N is a normal vector, d is a constant,
+            #             X is a position vector
+            # this function returns N and d.
+            phi = le_or_re[axis_id["phi"]]
+            theta = le_or_re[axis_id["theta"]]
+            r = le_or_re[axis_id["r"]]
+
+            # prob can/should use a yt func here to ensure consistency....
+            z = r * np.cos(theta)
+            r_xy = r * np.sin(theta)
+            x = r_xy * np.cos(phi)
+            y = r_xy * np.sin(phi)
+            pt = np.array([x, y, z])
+
+            z_hat = np.array([0, 0, 1])
+            normal_vec = np.cross(pt, z_hat)
+            d = np.dot(normal_vec, pt)
+            return normal_vec, d
+
         self.min_val = +np.inf
         self.max_val = -np.inf
         if self.scale:
@@ -64,6 +97,13 @@ def add_data(self, field, no_ghost=False):
             dx.append(dds.tolist())
             le.append(block.LeftEdge.tolist())
             re.append(block.RightEdge.tolist())
+
+            normal, const = calculate_phi_normal_plane(le[-1])
+            phi_plane_le.append(np.array([*normal, const]))
+
+            normal, const = calculate_phi_normal_plane(re[-1])
+            phi_plane_re.append(np.array([*normal, const]))
+
         for (g, node, (sl, _dims, _gi)) in self.data_source.tiles.slice_traverse():
             block = node.data
             self.blocks_by_grid[g.id - g._id_offset].append((id(block), i))
@@ -94,6 +134,15 @@ def add_data(self, field, no_ghost=False):
             VertexAttribute(name="in_right_edge", data=re)
         )
 
+        phi_plane_re = np.array(phi_plane_re, dtype="f4")
+        phi_plane_le = np.array(phi_plane_le, dtype="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)
+        )
+
         # Now we set up our textures
         self._load_textures()
 

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();
     }
-
 }

yt_idv/shaders/grid_position.vert.glsl

@@ -1,25 +1,59 @@
-in vec4 model_vertex; // The location of the vertex in model space
+// note: all in/out variables below are always in native coordinates (e.g.,
+// spherical or cartesian) except when noted.
+
+in vec4 model_vertex;
 in vec3 in_dx;
 in vec3 in_left_edge;
 in vec3 in_right_edge;
+
+in vec4 phi_plane_le;  // first 3 elements are the normal vector, final is constant
+in vec4 phi_plane_re;
+
 flat out vec4 vv_model;
-flat out mat4 vinverse_proj;
-flat out mat4 vinverse_mvm;
-flat out mat4 vinverse_pmvm;
+flat out mat4 vinverse_proj;  // always cartesian
+flat out mat4 vinverse_mvm;  // always cartesian
+flat out mat4 vinverse_pmvm; // always cartesian
 flat out vec3 vdx;
 flat out vec3 vleft_edge;
 flat out vec3 vright_edge;
 
+flat out vec4 vphi_plane_le;
+flat out vec4 vphi_plane_re;
+
+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;
+    }
+}
+
+vec4 transform_vec4(vec4 v) {
+    vec3 v3 = transform_vec3(vec3(v));
+    return vec4(v3[0], v3[1], v3[2], v[3]);
+}
+
 void main()
 {
+    // camera uniforms:
+    // projection, modelview
     vv_model = model_vertex;
     vinverse_proj = inverse(projection);
     // inverse model-view-matrix
     vinverse_mvm = inverse(modelview);
     vinverse_pmvm = inverse(projection * modelview);
-    gl_Position = projection * modelview * model_vertex;
+    gl_Position = projection * modelview * transform_vec4(model_vertex);
     vdx = vec3(in_dx);
     vleft_edge = vec3(in_left_edge);
     vright_edge = vec3(in_right_edge);
 
+    vphi_plane_le = vec4(phi_plane_le);
+    vphi_plane_re = vec4(phi_plane_re);
 }

yt_idv/shaders/known_uniforms.inc.glsl

@@ -64,3 +64,9 @@ uniform float iso_layers[32];
 uniform bool iso_log;
 uniform float iso_min;
 uniform float iso_max;
+
+// spherical coordinates
+uniform bool is_spherical;
+uniform int id_theta;  // azimuthal angle (0 to pi) index in the yt dataset
+uniform int id_r;  // radial index in the yt dataset
+uniform int id_phi;  // polar angle (0 to 2pi) indexi n the yt dataset