Merge pull request #159 from chrishavlin/sph_vol_render_working

Volume Rendering datasets in spherical coordinates

Author: chrishavlin

.gitignore

@@ -1,3 +1,5 @@
+# temp files
+*.swp
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -8,6 +10,8 @@ __pycache__/
 
 # Cythonized files
 yt_idv/utilities.c
+yt_idv/coordinate_utilities.c
+
 
 # imgui files
 imgui.ini

README.md

@@ -29,6 +29,7 @@ loaded in yt.  It is written to provide both scripting and interactive access.
     * Block and grid outlines
 * Support for sub-selections of data via the yt data selection interface
 * Integration with the [ipywidgets](https://ipywidgets.readthedocs.org/) ``Image`` widget.
+* Direct volume rendering of block-AMR data in spherical coordinates.
 
 ## Examples
 

docs/coordinate_systems.rst

@@ -0,0 +1,99 @@
+.. highlight:: shell
+
+================================
+Non-Cartesian Coordinate Systems
+================================
+
+Initial support for volume rendering of data defined in 3D non-cartesian coordinate systems
+was added in yt_idv 0.5.0 for block-AMR data in spherical coordinates. While not all
+rendering methods and annotations are supported for spherical coordinates, yt_idv can
+directly calculate maximum intensity projections, integrative projections and projection with custom
+transfer functions without any pre-interpolation or re-gridding.
+
+.. image:: resources/sph-vr-example.png
+    :align: center
+
+The approach for handling data defined in non-cartesian coordinates is to pre-calculate cartesian
+bounding boxes of the data blocks. The rendering pipeline then uses the cartesian bounding boxes
+to calculate ray entry/exit points during ray tracing. Between the entry/exit points, the
+cartesian coordinates of the ray position are converted to the native coordinates of the data,
+which is used to sample the texture maps (which are stored in native coordinates). The initial
+implementaiton of the algorithm uses the ray entry/exit points of the cartesian bounding boxes,
+and so not all points along the ray are gauranteed to lie within the bounds of the data -- though
+these data points will be discarded during ray tracing, it does mean that a larger number of samples
+along the ray may be required to ensure the underlying non-cartesian element is properly sampled compared
+to the standard cartesian data ray tracing (this is controlled by the ``sample_factor`` attribute of
+the ``BlockRendering`` component).
+
+At present, supported non-cartesian coordinate systems include Spherical Coordinates with (r, theta, phi), where r is radius, theta is co-latitude (between 0, pi)
+and phi is azimuth (between 0, 2pi), following yt conventions.
+
+----------------------------
+Notes on further development
+----------------------------
+
+Further contributions are welcome for adding support for the remaining 3d non-cartesian coordinate systems
+that yt supports that are not yet supported here (3d cylindrical, 3d geographic) as well as for adding
+support for non-cartesian coordinate systems in additional yt_idv components.
+
+To add support for additional non-cartesian coordinate systems requires two steps:
+
+#. Add methods for calcualting cartesian bounding boxes
+#. Add support to the rendering pipeline
+
+****************************************************
+Add methods for calcualting cartesian bounding boxes
+****************************************************
+
+The cartesian bounding box methods are defined in ``yt_idv/coordinate_utilities.pyx``. When adding support
+for a new coordiante system, implement a new ``MixedCoordBBox`` child class, following the
+``SphericalMixedCoordBBox`` example. The main method is ``get_cartesian_bbox``, which returns the
+cartesian bounding box for a single non-cartesian element. Once implemented, ``yt_idv.cartesian_bboxes``
+can be used to calculate bounding boxes for arrays of elements.
+
+Next, ``yt_idv.scene_data.block_collection.BlockCollection`` should be updated (mostly in
+``add_data`` and ``_set_geometry_attributes``) to handle the new coordinate system following the
+example for spherical coordinates.
+
+These two steps will handle the CPU-side of the calculations, next you need to update
+the rendering pipeline.
+
+*************************************
+Add support to the rendering pipeline
+*************************************
+
+A number of changes are required related to the rendering pipeline for the ``block_rendering``
+shader program. In order to avoid shader code duplication for different coordinate systems and
+also avoid overly branching code within the shader itself, pre-processor directives are used
+to provide switches between shader behavior for different coordinate systems.
+
+The ``NONCARTESIAN_GEOM`` directive is for functionality that should cover shared behavior
+between all non-cartesian coordinate systems. For example, ``grid_position.vert.glsl``,
+``grid_expand.geom.glsl`` and ``ray_tracing.frag.glsl`` all use ``NONCARTESIAN_GEOM`` for
+passing along the additional vertex attributes related to the cartesian bounding boxes.
+
+Functionality specific to a coordinate system should be defined with a new pre-processor
+directive flag. For spherical coordinates, this is ``SPHERICAL_GEOM``, and is used in
+``ray_tracing.frag.glsl`` for defining and using the functions for transforming from
+cartesian to spherical coordinates. To add a new one for use in the shaders, add a new
+entry to the ``yt_idv.scene_components.base_component._geom_directives`` dictionary and
+then wrap any functionality specific to your new coordinate system within pre-processor checks,
+for example::
+
+    #ifdef SPHERICAL_GEOM
+    vec3 cart_to_sphere_vec3(vec3 v) {
+        // transform a single point in cartesian coords to spherical
+        vec3 vout = vec3(0.,0.,0.);
+
+        // ---- code trimmed for clarity ---- //
+
+        return vout;
+
+    }
+    #endif
+
+When the shader program compiles, the above function will only be defined for rendering
+data in spherical coordinates.
+
+At a minimum, you will need to add a function to handle the coordinate conversion from the
+cartesian position of the ray to the native coordinates of your data.

docs/index.rst

@@ -9,6 +9,7 @@ Welcome to yt_idv's documentation!
    installation
    usage
    scene
+   coordinate_systems
    examples
    modules
    contributing

docs/resources/sph-vr-example.png

@@ -0,0 +1,130 @@
+import argparse
+
+import numpy as np
+import yt
+
+import yt_idv
+
+# yt reminder: phi is the azimuthal angle (0 to 2pi)
+# theta is the co-latitude, the angle from north (0 to pi)
+# coord ordering here will be r, phi, theta
+
+ax_order = ("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.0, 1.0], [0.0, 2 * np.pi], [0, np.pi]]),
+    "shell": np.array([[0.5, 1.0], [0.0, 2 * np.pi], [0, np.pi]]),
+    "north_hemi": np.array([[0.0, 1.0], [0.0, 2 * np.pi], [0, 0.5 * np.pi]]),
+    "north_shell": np.array([[0.5, 1.0], [0.0, 2 * np.pi], [0, 0.5 * np.pi]]),
+    "south_hemi": np.array([[0.0, 1.0], [0.0, 2 * np.pi], [0.5 * np.pi, np.pi]]),
+    "ew_hemi": np.array([[0.0, 1.0], [0.0, np.pi], [0.0, np.pi]]),
+}
+
+max_levs = {
+    "partial": 4,
+    "whole": 5,
+    "shell": 4,
+    "north_hemi": 5,
+    "north_shell": 2,
+    "south_hemi": 5,
+    "ew_hemi": 5,
+}
+
+
+def _build_ds(bbox_key):
+    bbox = bbox_options[bbox_key]
+
+    bbox_wid = bbox[:, 1] - bbox[:, 0]
+    bbox_c = np.mean(bbox, axis=1)
+    sz_0 = np.array((64, 64, 64))
+    max_lev = max_levs[bbox_key]
+
+    # divide grid by 2 every time
+    dd0 = bbox_wid / sz_0
+    sz_i = sz_0.copy()
+    grids = []
+    for lev in range(max_lev):
+
+        box_wid_factor = 2.0 * int(lev > 0) + int(lev == 0) * 1.0
+        bbox_wid = bbox_wid / box_wid_factor
+        le_i = bbox_c - bbox_wid / 2.0
+        re_i = bbox_c + bbox_wid / 2.0
+
+        # find closest start/end index in lev 0 grid
+        start_i = np.round(le_i / dd0).astype(int)
+        end_i = np.round(re_i / dd0).astype(int)
+        sz_0 = end_i - start_i
+
+        # recompute for rounding errors
+        le_i = start_i * dd0
+        re_i = le_i + sz_0 * dd0
+
+        sz_i = sz_0 * 2**lev
+
+        levp1 = np.full(sz_i, lev + 1.0)
+        grid = {
+            "left_edge": le_i,
+            "right_edge": re_i,
+            "dimensions": sz_i,
+            "level": lev,
+            ("stream", "density"): np.random.random(sz_i) * (lev + 1),
+            ("stream", "lev_p1"): levp1,
+        }
+        grids.append(grid)
+
+    ds = yt.load_amr_grids(
+        grids,
+        sz_0,
+        bbox=bbox,
+        length_unit=1,
+        geometry="spherical",
+        axis_order=ax_order,
+    )
+
+    return ds
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        prog="spherical_amr_rendering_with_refinement",
+        description="Loads an example spherical dataset with grid refinement 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="stream,density", help=msg)
+
+    parser.add_argument(
+        "--listfields", action=argparse.BooleanOptionalAction, default=True
+    )
+
+    args = parser.parse_args()
+
+    if args.domain not in bbox_options:
+        raise RuntimeError(
+            f"domain must be one of {list(bbox_options.keys())}, found {args.domain}"
+        )
+
+    ds = _build_ds(args.domain)
+
+    if args.listfields:
+        print("available fields:")
+        print(ds.field_list)
+
+    fld = tuple(str(args.field).split(","))
+
+    rc = yt_idv.render_context(height=800, width=800, gui=True)
+    sg = rc.add_scene(ds, fld, no_ghost=True)
+    rc.scene.components[0].cmap_log = False
+    rc.scene.components[0].sample_factor = 5.0
+
+    rc.run()

examples/spherical_amr_rendering_with_refinement.py

@@ -0,0 +1,127 @@
+import argparse
+
+import numpy as np
+import yt
+
+import yt_idv
+
+# yt reminder: phi is the azimuthal angle (0 to 2pi)
+# theta is the co-latitude, 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.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]]),
+    "quadrant_shell": np.array([[0.6, 1.0], [0.0, np.pi / 2], [0.0, np.pi / 2]]),
+}
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        prog="spherical_amr_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)
+    parser.add_argument(
+        "-np", "--nprocs", default=64, help="number of grids to decompose domain to"
+    )
+    parser.add_argument(
+        "-sz", "--size", default=256, help="dimensions, will be (size, size size)"
+    )
+
+    args = parser.parse_args()
+
+    sz = (int(args.size),) * 3
+    fake_data = {"density": np.random.random(sz)}
+
+    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:
+        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"
+        )
+
+    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]
+
+    nprocs = int(args.nprocs)
+
+    ds = yt.load_uniform_grid(
+        fake_data,
+        sz,
+        bbox=bbox,
+        nprocs=nprocs,
+        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)
+    sg = rc.add_scene(ds, field, no_ghost=True)
+    rc.scene.components[0].sample_factor = 5.0
+    rc.scene.components[0].cmap_log = False
+    rc.scene.components[0]._reset_cmap_bounds()
+
+    rc.run()