cleaning up duplication, unused functions, etc.

examples/amr_spherical_volume_rendering.py

@@ -1,4 +1,4 @@
-import sys
+import argparse
 
 import numpy as np
 import yt
@@ -7,45 +7,109 @@
 
 # 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 / 3], [np.pi / 4, np.pi / 2]]),
-    "whole": np.array([[0., 1.0], [0.0, 2 * np.pi], [0, np.pi]]),
+    "whole": np.array([[0.0, 1.0], [0.0, 2 * np.pi], [0, np.pi]]),
+    "shell": np.array([[0.7, 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]]),
+    "north_shell": np.array([[0.8, 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]
+
+    parser = argparse.ArgumentParser(
+        prog="amr_spherical_volume_rendering",
+        description="Loads an example spherical dataset in yt_idv",
+    )
+
+    msg = f"The geometry subset to generate: one of {list(bbox_options.keys())}"
+    parser.add_argument("-d", "--domain", default="partial", help=msg)
+    msg = (
+        "The field to plot. Provide a comma-separated string with field_type,field "
+        "e.g., to plot the field tuple ('index', 'phi'): \n "
+        "    $ python amr_spherical_volume_rendering.py -f index,x "
+        "\nIf a single string is provided, a field type of gas is assumed."
+    )
+    parser.add_argument("-f", "--field", default="index,phi", help=msg)
+
+    args = parser.parse_args()
+
+    field = str(args.field).split(",")
+    if len(field) == 1:
+        field = ("gas", str(field).strip())
+    elif len(field) == 2:
+        field = (field[0].strip(), field[1].strip())
     else:
-        bbox_type = "partial"
+        raise RuntimeError(
+            "Unexpected field formatting. Provide a single string"
+            " to provide just a field (will assume field type "
+            " of 'gas', or a comma separated string to provide a "
+            "field type and a field"
+        )
 
-    bbox = bbox_options[bbox_type]
+    if args.domain not in bbox_options:
+        raise RuntimeError(
+            f"domain must be one of {list(bbox_options.keys())}, found {args.domain}"
+        )
+    bbox = bbox_options[args.domain]
 
     ds = yt.load_uniform_grid(
         fake_data,
         sz,
         bbox=bbox,
-        nprocs=64,
-        geometry=("spherical", ("r", "phi", "theta")),
-        length_unit="m",
+        nprocs=128,
+        geometry="spherical",
+        axis_order=("r", "phi", "theta"),
+        length_unit=1,
+    )
+
+    phi_c = ds.quan(ds.domain_center[ds.coordinates.axis_id["phi"]].d, "")
+    theta_c = ds.quan(ds.domain_center[ds.coordinates.axis_id["theta"]].d, "")
+    rmax = ds.domain_right_edge[ds.coordinates.axis_id["r"]]
+    phi_f = ds.quan(15.0 * np.pi / 180.0, "")
+    theta_f = ds.quan(15.0 * np.pi / 180.0, "")
+    min_val = ds.quan(0.1, "")
+
+    def _tube(field, data):
+        phi = data["index", "phi"]
+        theta = data["index", "theta"]
+        tube = (1 - min_val) * np.exp(-(((theta - theta_c) / theta_f) ** 2))
+        tube = tube * np.exp(-(((phi - phi_c) / phi_f) ** 2))
+        return tube + min_val
+
+    ds.add_field(
+        name=("stream", "tube"),
+        function=_tube,
+        sampling_type="local",
     )
 
+    def _r_rev(field, data):
+        r = data["index", "r"]
+        return rmax - r
+
+    ds.add_field(
+        name=("stream", "r_rev"),
+        function=_r_rev,
+        sampling_type="local",
+    )
+
+    if field not in ds.field_list + ds.derived_field_list:
+        spaces = " " * 8
+        fld_list_str = f"\n{spaces}".join(str(fld) for fld in ds.field_list)
+        drv_fld_list_str = f"\n{spaces}".join(str(fld) for fld in ds.derived_field_list)
+        raise RuntimeError(
+            f"field {field} not in field_list or derived_field_list:\n"
+            f"\n    ds.field_list:\n{spaces}{fld_list_str}"
+            f"\n    ds.derived_field_list:\n{spaces}{drv_fld_list_str}"
+        )
+
     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 = rc.add_scene(ds, ("stream", "density"), no_ghost=True)
-    sg.camera.focus = [0.0, 0.0, 0.0]
+    sg = rc.add_scene(ds, field, no_ghost=True)
+    rc.scene.components[0].sample_factor = 20.0
+
     rc.run()

yt_idv/coordinate_utilities.pyx

@@ -301,42 +301,3 @@ def cartesian_bboxes(MixedCoordBBox bbox_handler,
             dx[i] = dxyz_i[0]
             dy[i] = dxyz_i[1]
             dz[i] = dxyz_i[2]
-
-
-@cython.cdivision(True)
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def one_cartesian_bbox(MixedCoordBBox bbox_handler,
-                      np.float64_t pos0,
-                      np.float64_t pos1,
-                      np.float64_t pos2,
-                      np.float64_t dpos0,
-                      np.float64_t dpos1,
-                      np.float64_t dpos2,
-                                      ):
-    # calculates the cartesian bounding boxes around non-cartesian
-    # volume elements
-    #
-    # bbox_handler : a MixedCoordBBox child instance
-    # pos0, pos1, pos2: native coordinates of element centers
-    # dpos0, dpos1, dpos2: element widths in native coordinates
-    # x, y, z: cartesian centers of bounding boxes, modified in place
-    # dx, dy, dz : full-widths of the cartesian bounding boxes, modified in place
-
-    cdef int i
-    cdef np.float64_t xyz_i[3]
-    cdef np.float64_t dxyz_i[3]
-
-    with nogil:
-
-        bbox_handler.get_cartesian_bbox(pos0,
-                                        pos1,
-                                        pos2,
-                                        dpos0,
-                                        dpos1,
-                                        dpos2,
-                                        xyz_i,
-                                        dxyz_i)
-
-
-    return xyz_i, dxyz_i

yt_idv/scene_components/blocks.py

@@ -140,9 +140,9 @@ def draw(self, scene, program):
 
     def _set_uniforms(self, scene, shader_program):
         shader_program._set_uniform(
-            "is_spherical", self.data.data_source.ds.geometry == "spherical"
+            "is_spherical", self.data._yt_geom_str == "spherical"
         )
-        if self.data.data_source.ds.geometry == "spherical":
+        if self.data._yt_geom_str == "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"])

yt_idv/scene_data/block_collection.py

@@ -76,21 +76,21 @@ def add_data(self, field, no_ghost=False):
         if hasattr(self.max_val, "in_units"):
             self.max_val = self.max_val.d
 
-        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")
         dx = np.array(dx, dtype="f4")
-        le = np.array(le, dtype="f4")
-        re = np.array(re, dtype="f4")
+        le = np.array(le)
+        re = np.array(re)
         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
+            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())
+
         self._set_geometry_attributes(le, re)
         self.vertex_array.attributes.append(
             VertexAttribute(name="model_vertex", data=vert)
@@ -133,8 +133,8 @@ def _set_geometry_attributes(self, le, re):
             self.vertex_array.attributes.append(
                 VertexAttribute(name="phi_plane_re", data=phi_plane_re)
             )
-            # cartesian bbox: very ugly, just testing...
-            print("calculating cartesian bounding boxes.")
+            # cartesian bbox calcualtions
+            # TODO: clean this up by rewriting the cython a bit...
             widths = re - le
             centers = (le + re) / 2.0
             bbox_handler = SphericalMixedCoordBBox()
@@ -156,7 +156,7 @@ def _set_geometry_attributes(self, le, re):
             le_cart = np.column_stack([x - dx / 2.0, y - dy / 2.0, z - dz / 2.0])
             re_cart = np.column_stack([x + dx / 2.0, y + dy / 2.0, z + dz / 2.0])
 
-            # normalize to viewport in (0, 1), preserving ratio of the bounding box
+            # cartesian le, re, width of whole domain
             domain_le = []
             domain_re = []
             domain_wid = []
@@ -165,13 +165,18 @@ def _set_geometry_attributes(self, le, re):
                 domain_re.append(re_cart[:, idim].max())
                 domain_wid.append(domain_re[idim] - domain_le[idim])
 
+            # normalize to viewport in (0, 1), preserving ratio of the bounding box
             max_wid = np.max(domain_wid)
             for idim in range(3):
                 le_cart[:, idim] = (le_cart[:, idim] - domain_le[idim]) / max_wid
                 re_cart[:, idim] = (re_cart[:, idim] - domain_le[idim]) / max_wid
 
             le_cart = np.asarray(le_cart)
             re_cart = np.asarray(re_cart)
+
+            # these will get passed down as uniforms to go from screen coords of
+            # 0,1 to cartesian coords of domain_le to domain_re from which full
+            # spherical coords can be calculated.
             self.cart_bbox_max_width = max_wid
             self.cart_bbox_le = np.array(domain_le).astype("f4")
 
@@ -182,6 +187,9 @@ def _set_geometry_attributes(self, le, re):
                 VertexAttribute(name="re_cart", data=re_cart.astype("f4"))
             )
 
+            # does not seem that diagonal is used anywhere, but recalculating to
+            # be safe...
+            self.diagonal = np.sqrt(((re_cart - le_cart) ** 2).sum())
         else:
             raise NotImplementedError(
                 f"{self.name} does not implement {self._yt_geom_str} geometries."

yt_idv/scene_graph.py

@@ -260,13 +260,7 @@ def from_ds(ds, field, no_ghost=True):
         else:
             ds, data_source = ds.ds, ds
 
-        if str(ds.geometry) == "cartesian":
-            center = ds.domain_center
-            pos = center + 1.5 * ds.domain_width.in_units("unitary")
-            near_plane = 3.0 * ds.index.get_smallest_dx().min().in_units("unitary").d
-            near_plane = max(near_plane, 1e-5)
-        else:
-            center, pos, near_plane = _get_camera_for_geometry(data_source, ds)
+        center, pos, near_plane = _get_camera_for_geometry(data_source, ds)
 
         c = TrackballCamera(position=pos, focus=center, near_plane=near_plane)
         c.update_orientation(0, 0, 0, 0)
@@ -285,14 +279,17 @@ def _get_camera_for_geometry(data_source, ds):
         # spherical coordinates
         center = np.array([0.5, 0.5, 0.5])
         wid = np.array([1.0, 1.0, 1.0])
-
         pos = center + 1.5 * wid
-
         dx_aprox = wid[0] / np.max(ds.domain_dimensions)
         near_plane = 3.0 * dx_aprox
         near_plane = max(near_plane, 1e-5)
-        return center, pos, near_plane
+    elif str(ds.geometry) == "cartesian":
+        center = ds.domain_center
+        pos = center + 1.5 * ds.domain_width.in_units("unitary")
+        near_plane = 3.0 * ds.index.get_smallest_dx().min().in_units("unitary").d
+        near_plane = max(near_plane, 1e-5)
     else:
         raise NotImplementedError(
             "Only cartesian and spherical geometries are supported at present."
         )
+    return center, pos, near_plane