diff --git a/.gitignore b/.gitignore index 5532662..718e79c 100644 --- a/.gitignore +++ b/.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 diff --git a/README.md b/README.md index b46d397..ce8ac84 100644 --- a/README.md +++ b/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 diff --git a/docs/coordinate_systems.rst b/docs/coordinate_systems.rst new file mode 100644 index 0000000..42bd53b --- /dev/null +++ b/docs/coordinate_systems.rst @@ -0,0 +1,55 @@ +.. 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 (with some caveats). + +PUT A NICE SCREEN SHOT HERE. + +--------------------------------------------------- +Volume Rendering of Data in a Spherical Coordinates +--------------------------------------------------- + +Overview + +#. pre-calculation of cartesian bounding boxes of the amr blocks, accelerated with cython. +#. load block data as textures in normalized native coordinates (just like cartesian) +#. pass down cartesian boudning boxes and spherical bounding boxes of blocks as vertex attributes +#. standard slab test with the cartesian bounding boxes for ray-element intersections +#. step along ray, calculate spherical coordinates (in the shader!), evalulate texture (discarding points outside the actual speherical volume element). + + +Mention: Pre-processor switches in order to re-use carteisan shaders efficiently. + +The main limitation in this approach is in stepping along the ray: intersection with the + cartesian bounding box of a spherical element does not gaurentee intersection with the enclosed +spherical elemnt. So a fairly large number of sample points along the ray must be used, increasing +the computational cost to achieve the same fidelity as rendering data that is natively in +cartesian coordinates. + +---------------------------- +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. + +**************************************** +Supporting additional coordinate systems +**************************************** + +Describe coordinate_utilities.pyx (and .pxd), how to add a new coordinate system. + +*************************************** +Supporting additional rendering methods +*************************************** + +Adding additional coordiante system conversions to shaders (with pre-processor +directives). diff --git a/docs/index.rst b/docs/index.rst index 2518a3c..b65d603 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -9,6 +9,7 @@ Welcome to yt_idv's documentation! installation usage scene + coordinate_systems examples modules contributing diff --git a/examples/amr_spherical_volume_rendering.py b/examples/amr_spherical_volume_rendering.py new file mode 100644 index 0000000..90ea76b --- /dev/null +++ b/examples/amr_spherical_volume_rendering.py @@ -0,0 +1,127 @@ +import argparse + +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 / 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="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) + 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 = 20.0 + rc.scene.components[0].cmap_log = False + rc.scene.components[0]._reset_cmap_bounds() + + rc.run() diff --git a/yt_idv/coordinate_utilities.pxd b/yt_idv/coordinate_utilities.pxd new file mode 100644 index 0000000..2de53e0 --- /dev/null +++ b/yt_idv/coordinate_utilities.pxd @@ -0,0 +1,57 @@ +cimport numpy as np + + +cdef inline np.float64_t fmax(np.float64_t f0, np.float64_t f1) noexcept nogil: + if f0 > f1: return f0 + return f1 + +cdef inline np.float64_t fmin(np.float64_t f0, np.float64_t f1) noexcept nogil: + if f0 < f1: return f0 + return f1 + +cdef (np.float64_t, np.float64_t, np.float64_t) _spherical_to_cartesian(np.float64_t r, + np.float64_t theta, + np.float64_t phi) noexcept nogil + + +cdef (np.float64_t, np.float64_t, np.float64_t) _cartesian_to_spherical(np.float64_t x, + np.float64_t y, + np.float64_t z) noexcept nogil + +cdef class MixedCoordBBox: + cdef int get_cartesian_bbox(self, + np.float64_t pos0, + np.float64_t pos1, + np.float64_t pos2, + np.float64_t dpos0, + np.float64_t dpos1, + np.float64_t dpos2, + np.float64_t xyz_i[3], + np.float64_t dxyz_i[3] + ) noexcept nogil + + +cdef class SphericalMixedCoordBBox(MixedCoordBBox): + cdef int get_cartesian_bbox( + self, + np.float64_t pos0, + np.float64_t pos1, + np.float64_t pos2, + np.float64_t dpos0, + np.float64_t dpos1, + np.float64_t dpos2, + np.float64_t[3] xyz_i, + np.float64_t[3] dxyz_i + ) noexcept nogil + + cdef void _get_cartesian_bbox( + self, + np.float64_t pos0, + np.float64_t pos1, + np.float64_t pos2, + np.float64_t dpos0, + np.float64_t dpos1, + np.float64_t dpos2, + np.float64_t[3] xyz_i, + np.float64_t[3] dxyz_i + ) noexcept nogil diff --git a/yt_idv/coordinate_utilities.pyx b/yt_idv/coordinate_utilities.pyx new file mode 100644 index 0000000..0f7df45 --- /dev/null +++ b/yt_idv/coordinate_utilities.pyx @@ -0,0 +1,526 @@ +cimport cython + +import numpy as np + +cimport numpy as np +from libc.math cimport acos, atan2, cos, sin, sqrt +from numpy.math cimport INFINITY as NPY_INF, PI as NPY_PI + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +cdef (np.float64_t, np.float64_t, np.float64_t) _spherical_to_cartesian(np.float64_t r, + np.float64_t theta, + np.float64_t phi) noexcept nogil: + # transform a single point in spherical coords to cartesian + # r : radius + # theta: colatitude + # phi: azimuthal (longitudinal) angle + cdef np.float64_t x, y, xy, z + + if r == 0.0: + return 0.0, 0.0, 0.0 + + xy = r * sin(theta) + x = xy * cos(phi) + y = xy * sin(phi) + z = r * cos(theta) + return x, y, z + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +cdef (np.float64_t, np.float64_t, np.float64_t) _cartesian_to_spherical(np.float64_t x, + np.float64_t y, + np.float64_t z) noexcept nogil: + # transform a single point in cartesian coords to spherical, returns + # r : radius + # theta: colatitude + # phi: azimuthal angle in range (0, 2pi) + cdef np.float64_t r, theta, phi + r = sqrt(x*x + y*y + z*z) + theta = acos(z / r) + phi = atan2(y, x) + # atan2 returns -pi to pi, adjust to (0, 2pi) + if phi < 0: + phi = phi + 2 * NPY_PI + return r, theta, phi + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +def cartesian_to_spherical(np.ndarray x, + np.ndarray y, + np.ndarray z): + # transform an array of points in cartesian coords to spherical, returns + # r : radius + # theta: colatitude + # phi: azimuthal angle in range (0, 2pi) + cdef np.ndarray[np.float64_t, ndim=1] r1d, th1d, phi1d + cdef np.ndarray[np.float64_t, ndim=1] x1d, y1d, z1d + cdef int i, n_x, ndim + cdef np.int64_t[:] shp + + ndim = x.ndim + + shp = np.zeros((ndim,), dtype=np.int64) + for i in range(ndim): + shp[i] = x.shape[i] + + x1d = x.reshape(-1) + y1d = y.reshape(-1) + z1d = z.reshape(-1) + + n_x = x1d.size + r1d = np.zeros((n_x,), dtype=np.float64) + th1d = np.zeros((n_x,), dtype=np.float64) + phi1d = np.zeros((n_x,), dtype=np.float64) + + + with nogil: + for i in range(n_x): + r1d[i], th1d[i], phi1d[i] = _cartesian_to_spherical(x1d[i], y1d[i], z1d[i]) + + r = r1d.reshape(shp) + theta = th1d.reshape(shp) + phi = phi1d.reshape(shp) + return r, theta, phi + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +def spherical_to_cartesian(np.ndarray r, + np.ndarray theta, + np.ndarray phi): + + # transform an array of points in spherical coords to cartesian + cdef np.ndarray[np.float64_t, ndim=1] r1d, th1d, phi1d + cdef np.ndarray[np.float64_t, ndim=1] x1d, y1d, z1d + + cdef np.int64_t[:] shp + cdef int i, n_r, ndim + + ndim = r.ndim + + shp = np.zeros((ndim,), dtype=np.int64) + for i in range(ndim): + shp[i] = r.shape[i] + + r1d = r.reshape(-1) + th1d = theta.reshape(-1) + phi1d = phi.reshape(-1) + + n_r = r1d.size + x1d = np.zeros((n_r,), dtype=np.float64) + y1d = np.zeros((n_r,), dtype=np.float64) + z1d = np.zeros((n_r,), dtype=np.float64) + + + with nogil: + for i in range(n_r): + x1d[i], y1d[i], z1d[i] = _spherical_to_cartesian(r1d[i], th1d[i], phi1d[i]) + + x = x1d.reshape(shp) + y = y1d.reshape(shp) + z = z1d.reshape(shp) + return x, y, z + + +cdef void _reduce_2_bboxes(np.float64_t[3] xyz_0, + np.float64_t[3] dxyz_0, + np.float64_t[3] xyz_1, + np.float64_t[3] dxyz_1, + np.float64_t[3] xyz, + np.float64_t[3] dxyz) noexcept nogil: + + # find the effective bounding box given 2 others + cdef np.float64_t le_i, re_i + cdef int idim + cdef int ndim = 3 + + for idim in range(ndim): + le_i = fmin(xyz_0[idim] - dxyz_0[idim]/2., + xyz_1[idim] - dxyz_1[idim]/2.) + re_i = fmax(xyz_0[idim] + dxyz_0[idim]/2., + xyz_1[idim] + dxyz_1[idim]/2.) + xyz[idim] = (le_i + re_i)/2. + dxyz[idim] = re_i - le_i + + +cdef class MixedCoordBBox: + # abstract class for calculating cartesian bounding boxes + # from non-cartesian grid elements. + cdef int get_cartesian_bbox(self, + np.float64_t pos0, + np.float64_t pos1, + np.float64_t pos2, + np.float64_t dpos0, + np.float64_t dpos1, + np.float64_t dpos2, + np.float64_t xyz_i[3], + np.float64_t dxyz_i[3] + ) noexcept nogil: + pass + + +cdef class SphericalMixedCoordBBox(MixedCoordBBox): + # Cartesian bounding boxes of spherical grid elements + cdef int get_cartesian_bbox(self, + np.float64_t pos0, # r + np.float64_t pos1, # theta + np.float64_t pos2, # phi + np.float64_t dpos0, # r + np.float64_t dpos1, # theta + np.float64_t dpos2, # phi + np.float64_t[3] xyz_i, + np.float64_t[3] dxyz_i, + ) noexcept nogil: + """ + Calculate the cartesian bounding box for a single spherical volume element. + + Parameters + ---------- + pos0: float + radius value at element center + pos1: float + co-latitude angle value at element center, in range (0, pi) (theta in yt) + pos2: float + azimuthal angle value at element center, in range (0, 2pi) (phi in yt) + dpos0: float + radial width of element + dpos1: float + co-latitude angular width of element + pos2: float + azimuthal anglular width of element + xyz_i: + 3-element array to store the center of the resulting bounding box + dxyz_i: + 3-element array to store the width of the resulting bounding box + + Returns + ------- + int + error code: 0 for success, 1 if something went wrong... + + Note: this function is recursive! If either angular width is above a cutoff (of + 0.2 radians or about 11 degrees), the element will be recursively divided. + """ + + cdef np.float64_t max_angle = 0.2 # about 11 degrees, if smaller, will subsample + cdef np.float64_t dpos2_2, dpos1_2 # half-widths + cdef np.float64_t pos2_2, pos1_2 # centers of half elements + cdef np.float64_t[3] xyz_0, dxyz_0 # additional 1st sample point if needed + cdef np.float64_t[3] xyz_1, dxyz_1 # additional 2nd sample point if needed + + if dpos1 < max_angle and dpos2 < max_angle: + self._get_cartesian_bbox( + pos0, pos1,pos2, + dpos0, dpos1, dpos2, + xyz_i, dxyz_i + ) + elif dpos1 >= max_angle: + # split into two, recursively get bbox, then get min/max of + # the bboxes + + # first sample + dpos1_2 = dpos1 / 2.0 + pos1_2 = pos1 - dpos1_2 / 2. + self.get_cartesian_bbox( + pos0, pos1_2, pos2, dpos0, dpos1_2, dpos2, xyz_0, dxyz_0 + ) + + # second sample + pos1_2 = pos1 + dpos1_2 / 2. + self.get_cartesian_bbox( + pos0, pos1_2, pos2, dpos0, dpos1_2, dpos2, xyz_1, dxyz_1 + ) + + _reduce_2_bboxes(xyz_0, dxyz_0, xyz_1, dxyz_1, xyz_i, dxyz_i) + + elif dpos2 >= max_angle: + # first sample + dpos2_2 = dpos2 / 2.0 + pos2_2 = pos2 - dpos2_2 / 2. + self.get_cartesian_bbox( + pos0, pos1, pos2_2, dpos0, dpos1, dpos2_2, xyz_0, dxyz_0 + ) + + # second sample + pos2_2 = pos2 + dpos2_2 / 2. + self.get_cartesian_bbox( + pos0, pos1, pos2_2, dpos0, dpos1, dpos2_2, xyz_1, dxyz_1 + ) + + _reduce_2_bboxes(xyz_0, dxyz_0, xyz_1, dxyz_1, xyz_i, dxyz_i) + else: + # this should be unreachable. Do not need to check for the case where + # both dpos2 and dpos1 > max angle because of the recursive call! + return 1 + return 0 + + + cdef void _get_cartesian_bbox(self, + np.float64_t pos0, + np.float64_t pos1, + np.float64_t pos2, + np.float64_t dpos0, + np.float64_t dpos1, + np.float64_t dpos2, + np.float64_t[3] xyz_i, + np.float64_t[3] dxyz_i + ) noexcept nogil: + + cdef np.float64_t r_i, theta_i, phi_i, dr_i, dtheta_i, dphi_i + cdef np.float64_t h_dphi, h_dtheta, h_dr, r_r + cdef np.float64_t xi, yi, zi, r_lr, theta_lr, phi_lr, phi_lr2, theta_lr2 + cdef np.float64_t xli, yli, zli, xri, yri, zri, r_xy, r_xy2 + cdef int isign_r, isign_ph, isign_th + cdef np.float64_t sign_r, sign_th, sign_ph + + cdef np.float64_t NPY_PI_2 = NPY_PI / 2.0 + cdef np.float64_t NPY_PI_3_2 = 3. * NPY_PI / 2.0 + cdef np.float64_t NPY_2xPI = 2. * NPY_PI + + r_i = pos0 + theta_i = pos1 + phi_i = pos2 + dr_i = dpos0 + dtheta_i = dpos1 + dphi_i = dpos2 + + # initialize the left/right values + xli = NPY_INF + yli = NPY_INF + zli = NPY_INF + xri = -1.0 * NPY_INF + yri = -1.0 * NPY_INF + zri = -1.0 * NPY_INF + + # find the min/max bounds over the 8 corners of the + # spherical volume element. + h_dphi = dphi_i / 2.0 + h_dtheta = dtheta_i / 2.0 + h_dr = dr_i / 2.0 + for isign_r in range(2): + for isign_ph in range(2): + for isign_th in range(2): + sign_r = 1.0 - 2.0 * isign_r + sign_th = 1.0 - 2.0 * isign_th + sign_ph = 1.0 - 2.0 * isign_ph + r_lr = r_i + sign_r * h_dr + theta_lr = theta_i + sign_th * h_dtheta + phi_lr = phi_i + sign_ph * h_dphi + + xi, yi, zi = _spherical_to_cartesian(r_lr, theta_lr, phi_lr) + + xli = fmin(xli, xi) + yli = fmin(yli, yi) + zli = fmin(zli, zi) + xri = fmax(xri, xi) + yri = fmax(yri, yi) + zri = fmax(zri, zi) + + # need to correct for special cases: + # if polar angle, phi, spans pi/2, pi or 3pi/2 then just + # taking the min/max of the corners will miss the cusp of the + # element. When this condition is met, the x/y min/max will + # equal +/- the projection of the max r in the xy plane -- in this case, + # the theta angle that gives the max projection of r in + # the x-y plane will change depending on the whether theta < or > pi/2, + # so the following calculates for the min/max theta value of the element + # and takes the max. + # ALSO note, that the following does check for when an edge aligns with the + # phi=0/2pi, it does not handle an element spanning the periodic boundary. + # Oh and this may break down for large elements that span whole + # quadrants... + phi_lr = phi_i - h_dphi + phi_lr2 = phi_i + h_dphi + theta_lr = theta_i - h_dtheta + theta_lr2 = theta_i + h_dtheta + r_r = r_i + h_dr + if theta_lr < NPY_PI_2 and theta_lr2 > NPY_PI_2: + r_xy = r_r + else: + r_xy = r_r * sin(theta_lr) + r_xy2 = r_r * sin(theta_lr2) + r_xy = fmax(r_xy, r_xy2) + + if phi_lr == 0.0 or phi_lr2 == NPY_2xPI: + # need to re-check this, for when theta spans equator + xri = r_xy + elif phi_lr < NPY_PI_2 and phi_lr2 > NPY_PI_2: + yri = r_xy + elif phi_lr < NPY_PI and phi_lr2 > NPY_PI: + xli = -r_xy + elif phi_lr < NPY_PI_3_2 and phi_lr2 > NPY_PI_3_2: + yli = -r_xy + + xyz_i[0] = (xri+xli)/2. + xyz_i[1] = (yri+yli)/2. + xyz_i[2] = (zri+zli)/2. + dxyz_i[0] = xri-xli + dxyz_i[1] = yri-yli + dxyz_i[2] = zri-zli + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +def cartesian_bboxes(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, + ): + """ + Calculate the cartesian bounding boxes around non-cartesian volume elements + + Parameters + ---------- + bbox_handler + a MixedCoordBBox child instance + pos0, pos1, pos2 + native coordinates of element centers + dpos0, dpos1, dpos2 + element widths in native coordinates + + Returns + ------- + x_y_z + a 3-element tuples containing the cartesian bounding box centers + dx_y_z + a 3-element tuples containing the cartesian bounding box widths + + """ + + cdef int i, n_pos, i_result + cdef np.float64_t[3] xyz_i + cdef np.float64_t[3] dxyz_i + cdef int failure = 0 + cdef np.float64_t[:] x, y, z # bbox centers + cdef np.float64_t[:] dx, dy, dz # bbox full widths + + n_pos = pos0.size + x = np.empty_like(pos0) + y = np.empty_like(pos0) + z = np.empty_like(pos0) + dx = np.empty_like(pos0) + dy = np.empty_like(pos0) + dz = np.empty_like(pos0) + + with nogil: + for i in range(n_pos): + + i_result = bbox_handler.get_cartesian_bbox(pos0[i], + pos1[i], + pos2[i], + dpos0[i], + dpos1[i], + dpos2[i], + xyz_i, + dxyz_i) + failure += i_result + + x[i] = xyz_i[0] + y[i] = xyz_i[1] + z[i] = xyz_i[2] + dx[i] = dxyz_i[0] + dy[i] = dxyz_i[1] + dz[i] = dxyz_i[2] + + if failure > 0: + raise RuntimeError("Unexpected error in get_cartesian_bbox.") + + x_y_z = (np.asarray(x), np.asarray(y), np.asarray(z)) + dx_y_z = (np.asarray(dx), np.asarray(dy), np.asarray(dz)) + return x_y_z, dx_y_z + + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +def cartesian_bboxes_edges(MixedCoordBBox bbox_handler, + np.float64_t[:] le0, + np.float64_t[:] le1, + np.float64_t[:] le2, + np.float64_t[:] re0, + np.float64_t[:] re1, + np.float64_t[:] re2, + ): + """ + Calculate the cartesian bounding boxes around non-cartesian volume elements + + Same as cartesian_bboxes, but supplying and returning element left/right edges. + + Parameters + ---------- + bbox_handler + a MixedCoordBBox child instance + le0, le1, le2 + native coordinates of element left edges + re0, re1, re2 + native coordinates of element right edges + + Returns + ------- + xyz_le + a 3-element tuples containing the cartesian bounding box left edges + xyz_re + a 3-element tuples containing the cartesian bounding box right edges + + """ + cdef int i, n_pos, i_result + cdef np.float64_t[3] xyz_i + cdef np.float64_t[3] dxyz_i + cdef np.float64_t pos0, pos1, pos2, dpos0, dpos1, dpos2 + cdef int failure = 0 + cdef np.float64_t[:] x_le, y_le, z_le # bbox left edges + cdef np.float64_t[:] x_re, y_re, z_re # bbox right edges + + n_pos = le0.size + x_le = np.empty_like(le0) + y_le = np.empty_like(le0) + z_le = np.empty_like(le0) + x_re = np.empty_like(le0) + y_re = np.empty_like(le0) + z_re = np.empty_like(le0) + + with nogil: + for i in range(n_pos): + pos0 = (le0[i] + re0[i]) / 2. + pos1 = (le1[i] + re1[i]) / 2. + pos2 = (le2[i] + re2[i]) / 2. + dpos0 = re0[i] - le0[i] + dpos1 = re1[i] - le1[i] + dpos2 = re2[i] - le2[i] + i_result = bbox_handler.get_cartesian_bbox(pos0, + pos1, + pos2, + dpos0, + dpos1, + dpos2, + xyz_i, + dxyz_i) + failure += i_result + + x_le[i] = xyz_i[0] - dxyz_i[0] / 2. + x_re[i] = xyz_i[0] + dxyz_i[0] / 2. + y_le[i] = xyz_i[1] - dxyz_i[1] / 2. + y_re[i] = xyz_i[1] + dxyz_i[1] / 2. + z_le[i] = xyz_i[2] - dxyz_i[2] / 2. + z_re[i] = xyz_i[2] + dxyz_i[2] / 2. + + + if failure > 0: + raise RuntimeError("Unexpected error in get_cartesian_bbox.") + + xyz_le = (np.asarray(x_le), np.asarray(y_le), np.asarray(z_le)) + xyz_re = (np.asarray(x_re), np.asarray(y_re), np.asarray(z_re)) + + return xyz_le, xyz_re diff --git a/yt_idv/scene_components/base_component.py b/yt_idv/scene_components/base_component.py index 1c60f03..e70c499 100644 --- a/yt_idv/scene_components/base_component.py +++ b/yt_idv/scene_components/base_component.py @@ -64,6 +64,7 @@ class SceneComponent(traitlets.HasTraits): _program1_invalid = True _program2_invalid = True _cmap_bounds_invalid = True + _data_geometry = traitlets.Unicode(default_value="cartesian") display_name = traitlets.Unicode(allow_none=True) @@ -81,6 +82,35 @@ class SceneComponent(traitlets.HasTraits): # This attribute determines whether or not this component is "active" active = traitlets.Bool(True) + @traitlets.observe("_data_geometry") + def _update_geometry_pp_directives(self, change): + if change["new"] == "spherical" and change["old"] == "cartesian": + # this should only happen once on initial data load, so only handling the + # change from the default value. + current_shader = component_shaders[self.name][self.render_method] + with self.hold_trait_notifications(): + for shd in ("vertex", "geometry", "fragment"): + # update the preprocessor directive state for this shader + self._program1_pp_defs.add_definition(shd, ("SPHERICAL_GEOM", "")) + pp_defs = self._program1_pp_defs[shd] + + # update shader attributes of self + new_shader_tuple = (current_shader[f"first_{shd}"], pp_defs) + setattr(self, f"{shd}_shader", new_shader_tuple) + + # all of the above shader attribute changes will trigger further + # traitlet changes that cause _recompile_shader to run + + self._data_geometry = change["new"] + + @traitlets.observe("data") + def _update_geometry(self, change): + if ( + hasattr(change["new"], "_yt_geom_str") + and change["new"]._yt_geom_str == "spherical" + ): + self._data_geometry = change["new"]._yt_geom_str + @traitlets.observe("display_bounds") def _change_display_bounds(self, change): # We need to update the framebuffer if the width or height has changed diff --git a/yt_idv/scene_components/blocks.py b/yt_idv/scene_components/blocks.py index ab9f503..52aa6c7 100644 --- a/yt_idv/scene_components/blocks.py +++ b/yt_idv/scene_components/blocks.py @@ -8,7 +8,7 @@ from yt_idv.opengl_support import TransferFunctionTexture from yt_idv.scene_components.base_component import SceneComponent from yt_idv.scene_data.block_collection import BlockCollection -from yt_idv.shader_objects import component_shaders +from yt_idv.shader_objects import component_shaders, get_shader_combos class BlockRendering(SceneComponent): @@ -34,34 +34,40 @@ class BlockRendering(SceneComponent): def render_gui(self, imgui, renderer, scene): changed = super().render_gui(imgui, renderer, scene) + + max_sf = 50.0 if self._yt_geom_str == "spherical" else 20.0 _, sample_factor = imgui.slider_float( - "Sample Factor", self.sample_factor, 1.0, 20.0 + "Sample Factor", self.sample_factor, 1.0, max_sf ) if _: self.sample_factor = sample_factor # Now, shaders - shader_combos = list(sorted(component_shaders[self.name])) + valid_shaders = get_shader_combos( + self.name, coord_system=self.data._yt_geom_str + ) descriptions = [ - component_shaders[self.name][_]["description"] for _ in shader_combos + component_shaders[self.name][_]["description"] for _ in valid_shaders ] - selected = shader_combos.index(self.render_method) + selected = valid_shaders.index(self.render_method) _, shader_ind = imgui.listbox("Shader", selected, descriptions) if _: - self.render_method = shader_combos[shader_ind] + self.render_method = valid_shaders[shader_ind] changed = changed or _ - if imgui.button("Add Block Outline"): - from ..scene_annotations.block_outline import BlockOutline - - block_outline = BlockOutline(data=self.data) - scene.annotations.append(block_outline) - if imgui.button("Add Grid Outline"): - from ..scene_annotations.grid_outlines import GridOutlines - from ..scene_data.grid_positions import GridPositions - - grids = self.data.data_source.ds.index.grids.tolist() - gp = GridPositions(grid_list=grids) - scene.data_objects.append(gp) - scene.components.append(GridOutlines(data=gp)) + if self.data._yt_geom_str == "cartesian": + # the following only work for cartesian data at present + if imgui.button("Add Block Outline"): + from ..scene_annotations.block_outline import BlockOutline + + block_outline = BlockOutline(data=self.data) + scene.annotations.append(block_outline) + if imgui.button("Add Grid Outline"): + from ..scene_annotations.grid_outlines import GridOutlines + from ..scene_data.grid_positions import GridPositions + + grids = self.data.data_source.ds.index.grids.tolist() + gp = GridPositions(grid_list=grids) + scene.data_objects.append(gp) + scene.components.append(GridOutlines(data=gp)) if self.render_method == "transfer_function": # Now for the transfer function stuff imgui.image_button( @@ -139,6 +145,16 @@ def draw(self, scene, program): GL.glDrawArrays(GL.GL_POINTS, tex_ind * each, each) def _set_uniforms(self, scene, shader_program): + 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"]) + shader_program._set_uniform("id_phi", axis_id["phi"]) + shader_program._set_uniform( + "cart_bbox_max_width", self.data.cart_bbox_max_width + ) + shader_program._set_uniform("cart_bbox_le", self.data.cart_bbox_le) + 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])) @@ -149,3 +165,7 @@ def _set_uniforms(self, scene, shader_program): shader_program._set_uniform("tf_log", float(self.tf_log)) shader_program._set_uniform("slice_normal", np.array(self.slice_normal)) shader_program._set_uniform("slice_position", np.array(self.slice_position)) + + @property + def _yt_geom_str(self): + return self.data._yt_geom_str diff --git a/yt_idv/scene_components/octree_blocks.py b/yt_idv/scene_components/octree_blocks.py index 8b98d69..9457ff8 100644 --- a/yt_idv/scene_components/octree_blocks.py +++ b/yt_idv/scene_components/octree_blocks.py @@ -7,7 +7,7 @@ from yt_idv.opengl_support import TransferFunctionTexture from yt_idv.scene_components.base_component import SceneComponent from yt_idv.scene_data.octree_block_collection import OctreeBlockCollection -from yt_idv.shader_objects import component_shaders +from yt_idv.shader_objects import component_shaders, get_shader_combos class OctreeBlockRendering(SceneComponent): @@ -37,7 +37,7 @@ def render_gui(self, imgui, renderer, scene): if _: self.sample_factor = sample_factor # Now, shaders - shader_combos = list(sorted(component_shaders[self.name])) + shader_combos = get_shader_combos(self.name) descriptions = [ component_shaders[self.name][_]["description"] for _ in shader_combos ] diff --git a/yt_idv/scene_data/block_collection.py b/yt_idv/scene_data/block_collection.py index e0b9f57..8a5dd5b 100644 --- a/yt_idv/scene_data/block_collection.py +++ b/yt_idv/scene_data/block_collection.py @@ -17,6 +17,7 @@ class BlockCollection(SceneData): scale = traitlets.Bool(False) blocks_by_grid = traitlets.Instance(defaultdict, (list,)) grids_by_block = traitlets.Dict(default_value=()) + _yt_geom_str = traitlets.Unicode("cartesian") @traitlets.default("vertex_array") def _default_vertex_array(self): @@ -38,12 +39,19 @@ def add_data(self, field, no_ghost=False): Should we speed things up by skipping ghost zone generation? """ self.data_source.tiles.set_fields([field], [False], no_ghost=no_ghost) + + self._yt_geom_str = str(self.data_source.ds.geometry) + # 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 + # 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: + if self.scale and self._yt_geom_str == "cartesian": left_min = np.ones(3, "f8") * np.inf right_max = np.ones(3, "f8") * -np.inf for block in self.data_source.tiles.traverse(): @@ -74,35 +82,100 @@ 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()) - - # Note: the block LeftEdge and RightEdge arrays are plain np arrays in - # units of code_length, so need to convert to unitary units (in range 0,1) - # after the fact: - 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 - # Now we set up our buffer vert = np.array(vert, dtype="f4") + dx = np.array(dx, dtype="f4") + le = np.array(le) + re = np.array(re) + if self._yt_geom_str == "cartesian": + # Note: the block LeftEdge and RightEdge arrays are plain np arrays in + # units of code_length, so need to convert to unitary units (in range 0,1) + # after the fact: + 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) ) 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() + 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": + from yt_idv.coordinate_utilities import ( + SphericalMixedCoordBBox, + cartesian_bboxes_edges, + ) + + axis_id = self.data_source.ds.coordinates.axis_id + # cartesian bbox calculations + bbox_handler = SphericalMixedCoordBBox() + le_cart, re_cart = cartesian_bboxes_edges( + bbox_handler, + le[:, axis_id["r"]], + le[:, axis_id["theta"]], + le[:, axis_id["phi"]], + re[:, axis_id["r"]], + re[:, axis_id["theta"]], + re[:, axis_id["phi"]], + ) + le_cart = np.column_stack(le_cart) + re_cart = np.column_stack(re_cart) + + # cartesian le, re, width of whole domain + domain_le = le_cart.min(axis=0) + domain_re = re_cart.max(axis=0) + domain_wid = domain_re - domain_le + + # 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 = domain_le.astype("f4") + + self.vertex_array.attributes.append( + VertexAttribute(name="le_cart", data=le_cart.astype("f4")) + ) + self.vertex_array.attributes.append( + 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." + ) + def viewpoint_iter(self, camera): for block in self.data_source.tiles.traverse(viewpoint=camera.position): vbo_i, _ = self.blocks[id(block)] diff --git a/yt_idv/scene_graph.py b/yt_idv/scene_graph.py index 656b51f..c7b1f18 100644 --- a/yt_idv/scene_graph.py +++ b/yt_idv/scene_graph.py @@ -1,5 +1,6 @@ import contextlib +import numpy as np import traitlets from OpenGL import GL from yt.data_objects.static_output import Dataset @@ -258,10 +259,8 @@ def from_ds(ds, field, no_ghost=True): data_source = ds.all_data() else: ds, data_source = ds.ds, ds - 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) + + 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) @@ -270,3 +269,27 @@ def from_ds(ds, field, no_ghost=True): if field is not None: scene.add_volume(data_source, field, no_ghost=no_ghost) return scene + + +def _get_camera_for_geometry(data_source, ds): + + if str(ds.geometry) == "spherical": + # always return the in-screen coordinate focus here. The shader + # will handle transforming to the expected full cartesian and + # 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) + 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 diff --git a/yt_idv/shader_objects.py b/yt_idv/shader_objects.py index c05fa2b..5af0f5d 100644 --- a/yt_idv/shader_objects.py +++ b/yt_idv/shader_objects.py @@ -504,3 +504,19 @@ def reset(self, shader_type: Optional[str] = None): default_shader_combos.update( {_: component_shaders[_].pop("default_value") for _ in component_shaders} ) + + +def get_shader_combos(component_name, coord_system="cartesian"): + shader_combos = list(sorted(component_shaders[component_name])) + if coord_system == "cartesian": + return shader_combos + + valid_shader_combos = [] + for shader_name in shader_combos: + shader = component_shaders[component_name][shader_name] + if ( + "coordinate_systems" in shader + and coord_system in shader["coordinate_systems"] + ): + valid_shader_combos.append(shader_name) + return valid_shader_combos diff --git a/yt_idv/shaders/grid_expand.geom.glsl b/yt_idv/shaders/grid_expand.geom.glsl index edf0955..2f47266 100644 --- a/yt_idv/shaders/grid_expand.geom.glsl +++ b/yt_idv/shaders/grid_expand.geom.glsl @@ -5,9 +5,17 @@ flat in vec3 vdx[]; flat in vec3 vleft_edge[]; flat in vec3 vright_edge[]; +#ifdef SPHERICAL_GEOM +flat in vec3 vleft_edge_cart[]; +flat in vec3 vright_edge_cart[]; +flat out vec3 left_edge_cart; +flat out vec3 right_edge_cart; +#endif + 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; @@ -36,22 +44,39 @@ uniform vec3 arrangement[8] = vec3[]( uniform int aindex[14] = int[](6, 7, 4, 5, 1, 7, 3, 6, 2, 4, 0, 1, 2, 3); + void main() { vec4 center = gl_in[0].gl_Position; - - vec3 width = vright_edge[0] - vleft_edge[0]; - + vec3 width, le; vec4 newPos; + #ifdef SPHERICAL_GEOM + width = vright_edge_cart[0] - vleft_edge_cart[0]; + le = vleft_edge_cart[0]; + #else + width = vright_edge[0] - vleft_edge[0]; + le = vleft_edge[0]; + #endif + for (int i = 0; i < 14; i++) { - newPos = vec4(vleft_edge[0] + width * arrangement[aindex[i]], 1.0); + // walks through each vertex of the triangle strip, emit it. need to + // emit gl_Position in cartesian, pass along other attributes in native + // coordinates + + newPos = vec4(le + width * arrangement[aindex[i]], 1.0); gl_Position = projection * modelview * newPos; 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]; + + #ifdef SPHERICAL_GEOM + left_edge_cart = vleft_edge_cart[0]; + right_edge_cart = vright_edge_cart[0]; + #endif + dx = vdx[0]; v_model = newPos; texture_offset = ivec3(0); diff --git a/yt_idv/shaders/grid_position.vert.glsl b/yt_idv/shaders/grid_position.vert.glsl index 7aa95c7..8c93cd6 100644 --- a/yt_idv/shaders/grid_position.vert.glsl +++ b/yt_idv/shaders/grid_position.vert.glsl @@ -1,7 +1,13 @@ -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; + + + flat out vec4 vv_model; flat out mat4 vinverse_proj; flat out mat4 vinverse_mvm; @@ -10,16 +16,35 @@ flat out vec3 vdx; flat out vec3 vleft_edge; flat out vec3 vright_edge; +#ifdef SPHERICAL_GEOM +// pre-computed cartesian le, re +in vec3 le_cart; +in vec3 re_cart; + +flat out vec3 vleft_edge_cart; +flat out vec3 vright_edge_cart; +#endif + + 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; + + // native coordinates vdx = vec3(in_dx); vleft_edge = vec3(in_left_edge); vright_edge = vec3(in_right_edge); + #ifdef SPHERICAL_GEOM + // cartesian bounding boxes + vleft_edge_cart = vec3(le_cart); + vright_edge_cart = vec3(re_cart); + #endif } diff --git a/yt_idv/shaders/header.inc.glsl b/yt_idv/shaders/header.inc.glsl index 5ae7f43..5a3e146 100644 --- a/yt_idv/shaders/header.inc.glsl +++ b/yt_idv/shaders/header.inc.glsl @@ -1 +1,4 @@ #version 330 core + +const float INFINITY = 1. / 0.; +const float PI = 3.1415926535897932384626433832795; diff --git a/yt_idv/shaders/known_uniforms.inc.glsl b/yt_idv/shaders/known_uniforms.inc.glsl index 0700651..bb04ae3 100644 --- a/yt_idv/shaders/known_uniforms.inc.glsl +++ b/yt_idv/shaders/known_uniforms.inc.glsl @@ -60,6 +60,13 @@ uniform float iso_layers[32]; uniform float iso_layer_tol[32]; uniform float iso_alphas[32]; +// spherical coordinates +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 +uniform vec3 cart_bbox_le; +uniform float cart_bbox_max_width; + // draw outline control uniform float draw_boundary; uniform vec4 boundary_color; diff --git a/yt_idv/shaders/ray_tracing.frag.glsl b/yt_idv/shaders/ray_tracing.frag.glsl index db3c15e..574d6cd 100644 --- a/yt_idv/shaders/ray_tracing.frag.glsl +++ b/yt_idv/shaders/ray_tracing.frag.glsl @@ -6,6 +6,12 @@ flat in mat4 inverse_proj; flat in mat4 inverse_mvm; flat in mat4 inverse_pmvm; flat in ivec3 texture_offset; + +#ifdef SPHERICAL_GEOM +flat in vec3 left_edge_cart; +flat in vec3 right_edge_cart; +#endif + out vec4 output_color; bool within_bb(vec3 pos) @@ -15,6 +21,32 @@ bool within_bb(vec3 pos) return all(left) && all(right); } +#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.); + + // 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, cart_bbox_* variables are also uniforms + float x = v[0] * cart_bbox_max_width + cart_bbox_le[0]; + float y = v[1] * cart_bbox_max_width + cart_bbox_le[1]; + float z = v[2] * cart_bbox_max_width + cart_bbox_le[2]; + vout[id_r] = x * x + y * y + z * z; + vout[id_r] = sqrt(vout[id_r]); + vout[id_theta] = acos(z / vout[id_r]); + float phi = atan(y, x); + // atan2 returns -pi to pi, adjust to (0, 2pi) + if (phi < 0 ){ + phi = phi + 2.0 * PI; + } + vout[id_phi] = phi; + + return vout; + +} +#endif + vec3 get_offset_texture_position(sampler3D tex, vec3 tex_curr_pos) { ivec3 texsize = textureSize(tex, 0); // lod (mipmap level) always 0? @@ -26,18 +58,20 @@ bool sample_texture(vec3 tex_curr_pos, inout vec4 curr_color, float tdelta, vec4 cleanup_phase(in vec4 curr_color, in vec3 dir, in float t0, in float t1); // This main() function will call a function called sample_texture at every -// step along the ray. It must be of the form +// step along the ray. sample_texture must be of the form // void (vec3 tex_curr_pos, inout vec4 curr_color, float tdelta, float t, // vec3 direction); - void main() { + // Obtain screen coordinates // https://www.opengl.org/wiki/Compute_eye_space_from_window_space#From_gl_FragCoord vec3 ray_position = v_model.xyz; + vec3 ray_position_native; + + output_color = vec4(0.); // Five samples - vec3 step_size = dx/sample_factor; vec3 dir = -normalize(camera_pos.xyz - ray_position); dir = max(abs(dir), 0.0001) * sign(dir); vec4 curr_color = vec4(0.0); @@ -46,8 +80,18 @@ void main() // This will help solve the left/right edge issues. vec3 idir = 1.0/dir; - vec3 tl = (left_edge - camera_pos)*idir; - vec3 tr = (right_edge - camera_pos)*idir; + vec3 tl, tr, dx_cart; + #ifdef SPHERICAL_GEOM + tl = (left_edge_cart - camera_pos)*idir; + tr = (right_edge_cart - camera_pos)*idir; + dx_cart = right_edge_cart - left_edge_cart; + #else + tl = (left_edge - camera_pos)*idir; + tr = (right_edge - camera_pos)*idir; + dx_cart = dx; + #endif + vec3 step_size = dx_cart/ sample_factor; + vec3 tmin, tmax; bvec3 temp_x, temp_y; // These 't' prefixes actually mean 'parameter', as we use in grid_traversal.pyx. @@ -76,7 +120,7 @@ void main() float tdelta = min(temp_t.x, temp_t.y); float t = t0; - vec3 range = (right_edge + dx/2.0) - (left_edge - dx/2.0); + vec3 range = (right_edge + dx/2.0) - (left_edge - dx/2.0); // texture range in native coords vec3 nzones = range / dx; vec3 ndx = 1.0/nzones; @@ -84,6 +128,7 @@ void main() bool sampled; bool ever_sampled = false; + bool within_el = true; vec4 v_clip_coord; float f_ndc_depth; @@ -92,18 +137,27 @@ void main() ray_position = p0; while(t <= t1) { - tex_curr_pos = (ray_position - left_edge) / range; // Scale from 0 .. 1 - // But, we actually need it to be 0 + normalized dx/2 to 1 - normalized dx/2 - tex_curr_pos = (tex_curr_pos * (1.0 - ndx)) + ndx/2.0; - sampled = sample_texture(tex_curr_pos, curr_color, tdelta, t, dir); + // texture position + #ifdef SPHERICAL_GEOM + ray_position_native = cart_to_sphere_vec3(ray_position); + within_el = within_bb(ray_position_native); + #else + ray_position_native = ray_position; + #endif + + if (within_el) { + tex_curr_pos = (ray_position_native - left_edge) / range; // Scale from 0 .. 1 + // But, we actually need it to be 0 + normalized dx/2 to 1 - normalized dx/2 + tex_curr_pos = (tex_curr_pos * (1.0 - ndx)) + ndx/2.0; + sampled = sample_texture(tex_curr_pos, curr_color, tdelta, t, dir); + } if (sampled) { ever_sampled = true; v_clip_coord = projection * modelview * vec4(ray_position, 1.0); f_ndc_depth = v_clip_coord.z / v_clip_coord.w; depth = min(depth, (1.0 - 0.0) * 0.5 * f_ndc_depth + (1.0 + 0.0) * 0.5); - } t += tdelta; diff --git a/yt_idv/shaders/shaderlist.yaml b/yt_idv/shaders/shaderlist.yaml index 4b91943..2b871ed 100644 --- a/yt_idv/shaders/shaderlist.yaml +++ b/yt_idv/shaders/shaderlist.yaml @@ -208,6 +208,7 @@ component_shaders: first_fragment: max_intensity second_vertex: passthrough second_fragment: apply_colormap + coordinate_systems: [cartesian, spherical] projection: description: Projective integration first_vertex: grid_position @@ -215,6 +216,7 @@ component_shaders: first_fragment: projection second_vertex: passthrough second_fragment: apply_colormap + coordinate_systems: [cartesian, spherical] transfer_function: description: Color transfer function first_vertex: grid_position @@ -222,6 +224,7 @@ component_shaders: first_fragment: transfer_function second_vertex: passthrough second_fragment: passthrough + coordinate_systems: [cartesian, spherical] isocontours: description: Isocontours first_vertex: grid_position @@ -236,6 +239,14 @@ component_shaders: first_fragment: slice_sample second_vertex: passthrough second_fragment: apply_colormap + constant: + description: Constant Color + first_vertex: grid_position + first_geometry: grid_expand + first_fragment: constant + second_vertex: passthrough + second_fragment: apply_colormap + coordinate_systems: [cartesian, spherical] octree_block_rendering: default_value: max_intensity max_intensity: diff --git a/yt_idv/tests/conftest.py b/yt_idv/tests/conftest.py index 8c9a39d..b7aa091 100644 --- a/yt_idv/tests/conftest.py +++ b/yt_idv/tests/conftest.py @@ -1,7 +1,24 @@ +import base64 import os +import pytest +import yt +from pytest_html import extras as html_extras + def pytest_configure(config): # this will get run before all tests, before collection and # any opengl imports that happen within test files. os.environ["PYOPENGL_PLATFORM"] = "osmesa" + + +@pytest.fixture() +def image_store(request, extras, tmpdir): + def _snap_image(rc): + image = rc.run() + img = yt.write_bitmap(image, None) + content = base64.b64encode(img).decode("ascii") + extras.append(html_extras.png(content)) + extras.append(html_extras.html("
")) + + return _snap_image diff --git a/yt_idv/tests/test_component_shaders.py b/yt_idv/tests/test_component_shaders.py new file mode 100644 index 0000000..11bd13b --- /dev/null +++ b/yt_idv/tests/test_component_shaders.py @@ -0,0 +1,23 @@ +import pytest + +from yt_idv.shader_objects import component_shaders, get_shader_combos + + +@pytest.mark.parametrize("render_name", component_shaders.keys()) +def test_get_shader_combos_all_components(render_name): + # default call to get_shader_combos should return all shaders + expected = component_shaders[render_name] + actual = get_shader_combos(render_name) + assert set(expected) == set(actual) + + +def test_get_shader_combos_coord(): + # only block_rendering supports non-cartesian for now, check that the supported + # shaders are returned as expected. + shaders = get_shader_combos("block_rendering", coord_system="spherical") + not_supported = ["isocontours", "slice"] + assert all(s not in shaders for s in not_supported) + supported = ["max_intensity", "projection", "transfer_function", "constant"] + for s in supported: + assert s in shaders + assert all(s in shaders for s in supported) diff --git a/yt_idv/tests/test_spherical_bboxes.py b/yt_idv/tests/test_spherical_bboxes.py new file mode 100644 index 0000000..81734fe --- /dev/null +++ b/yt_idv/tests/test_spherical_bboxes.py @@ -0,0 +1,235 @@ +import numpy as np +import pytest + +from yt_idv.coordinate_utilities import ( + SphericalMixedCoordBBox, + cartesian_bboxes, + cartesian_bboxes_edges, + cartesian_to_spherical, + spherical_to_cartesian, +) + + +def test_cartesian_bboxes_for_spherical(): + + # this test checks the special cases where + # a spherical volume element crosses an axis + # or when an element edge is lined up with an axis + + # check element that includes theta=0 as an edge + r = np.array([0.95]) + dr = np.array([0.1]) + theta = np.array([0.05]) + dtheta = np.array([0.1]) + phi = np.array([0.05]) + dphi = np.array([0.05]) + + bbox_handler = SphericalMixedCoordBBox() + + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + assert xyz[2] + dxyz[2] / 2 == 1.0 + assert np.allclose(xyz[0] - dxyz[0] / 2, 0.0) + assert np.allclose(xyz[1] - dxyz[1] / 2, 0.0) + + # now theta = np.pi + theta = np.array([np.pi - dtheta[0] / 2]) + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + assert xyz[2] - dxyz[2] / 2 == -1.0 + assert np.allclose(xyz[0] - dxyz[0] / 2, 0.0) + assert np.allclose(xyz[1] - dxyz[1] / 2, 0.0) + + # element at equator, overlapping the +y axis + theta = np.array([np.pi / 2]) + phi = np.array([np.pi / 2]) + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + + assert xyz[1] + dxyz[1] / 2 == 1.0 + assert np.allclose(xyz[0], 0.0) + assert np.allclose(xyz[2], 0.0) + + # element at equator, overlapping the -x axis + phi = np.array([np.pi]) + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + + assert xyz[0] - dxyz[0] / 2 == -1.0 + assert np.allclose(xyz[1], 0.0) + assert np.allclose(xyz[2], 0.0) + + # element at equator, overlapping the -y axis + phi = np.array([3 * np.pi / 2]) + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + + assert xyz[1] - dxyz[1] / 2 == -1.0 + assert np.allclose(xyz[0], 0.0) + assert np.allclose(xyz[2], 0.0) + + # element at equator, overlapping +x axis + phi = dphi / 2.0 + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + assert xyz[0] + dxyz[0] / 2 == 1.0 + + # element with edge on +x axis in -theta direction + theta = np.pi / 2 - dtheta / 2 + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + assert xyz[0] + dxyz[0] / 2 == 1.0 + + # element with edge on +x axis in +theta direction + theta = np.pi / 2 + dtheta / 2 + xyz, dxyz = cartesian_bboxes(bbox_handler, r, theta, phi, dr, dtheta, dphi) + assert xyz[0] + dxyz[0] / 2 == 1.0 + + # finally, check that things work OK with a wide range of + # angles + + r_edges = np.linspace(0.4, 1.0, 10, dtype="float64") + theta_edges = np.linspace(0, np.pi, 10, dtype="float64") + phi_edges = np.linspace(0.0, 2 * np.pi, 10, dtype="float64") + + r = (r_edges[0:-1] + r_edges[1:]) / 2.0 + theta = (theta_edges[0:-1] + theta_edges[1:]) / 2.0 + phi = (phi_edges[0:-1] + phi_edges[1:]) / 2.0 + + dr = r_edges[1:] - r_edges[:-1] + dtheta = theta_edges[1:] - theta_edges[:-1] + dphi = phi_edges[1:] - phi_edges[:-1] + + r_th_ph = np.meshgrid(r, theta, phi) + d_r_th_ph = np.meshgrid(dr, dtheta, dphi) + r_th_ph = [r_th_ph[i].ravel() for i in range(3)] + d_r_th_ph = [d_r_th_ph[i].ravel() for i in range(3)] + + x_y_z, d_x_y_z = cartesian_bboxes( + bbox_handler, + r_th_ph[0], + r_th_ph[1], + r_th_ph[2], + d_r_th_ph[0], + d_r_th_ph[1], + d_r_th_ph[2], + ) + + assert np.all(np.isfinite(x_y_z)) + assert np.all(np.isfinite(d_x_y_z)) + + # and check the extents again for completeness... + for i in range(3): + max_val = np.max(x_y_z[i] + d_x_y_z[i] / 2.0) + min_val = np.min(x_y_z[i] - d_x_y_z[i] / 2.0) + assert max_val == 1.0 + assert min_val == -1.0 + + +def test_spherical_cartesian_roundtrip(): + xyz = [np.linspace(0, 1, 10) for _ in range(3)] + xyz = np.meshgrid(*xyz) + xyz = [xyzi.ravel() for xyzi in xyz] + x, y, z = xyz + + r, theta, phi = cartesian_to_spherical(x, y, z) + x_out, y_out, z_out = spherical_to_cartesian(r, theta, phi) + + assert np.allclose(x_out, x) + assert np.allclose(y_out, y) + assert np.allclose(z_out, z) + assert np.max(r) == np.sqrt(3.0) + + +@pytest.mark.parametrize("n_angles", (2, 4, 8, 16)) +def test_large_elements(n_angles): + + bbox_handler = SphericalMixedCoordBBox() + + r_edges = np.linspace(0.4, 1.0, 10, dtype="float64") + theta_edges = np.linspace(0, np.pi, n_angles, dtype="float64") + phi_edges = np.linspace(0.0, 2 * np.pi, n_angles, dtype="float64") + + r = (r_edges[0:-1] + r_edges[1:]) / 2.0 + theta = (theta_edges[0:-1] + theta_edges[1:]) / 2.0 + phi = (phi_edges[0:-1] + phi_edges[1:]) / 2.0 + + dr = r_edges[1:] - r_edges[:-1] + dtheta = theta_edges[1:] - theta_edges[:-1] + dphi = phi_edges[1:] - phi_edges[:-1] + + r_th_ph = np.meshgrid(r, theta, phi) + d_r_th_ph = np.meshgrid(dr, dtheta, dphi) + r_th_ph = [r_th_ph[i].ravel() for i in range(3)] + d_r_th_ph = [d_r_th_ph[i].ravel() for i in range(3)] + + x_y_z, d_x_y_z = cartesian_bboxes( + bbox_handler, + r_th_ph[0], + r_th_ph[1], + r_th_ph[2], + d_r_th_ph[0], + d_r_th_ph[1], + d_r_th_ph[2], + ) + + x_y_z = np.column_stack(x_y_z) + d_x_y_z = np.column_stack(d_x_y_z) + + assert np.all(np.isfinite(x_y_z)) + assert np.all(np.isfinite(d_x_y_z)) + + le = np.min(x_y_z - d_x_y_z / 2.0, axis=0) + re = np.max(x_y_z + d_x_y_z / 2.0, axis=0) + + assert np.all(le == -1.0) + assert np.all(re == 1.0) + + +def test_spherical_boxes_edges(): + # check that you get the same result from supplying centers + # vs edges. + + bbox_handler = SphericalMixedCoordBBox() + + r_edges = np.linspace(0.0, 10.0, 20, dtype="float64") + theta_edges = np.linspace(0, np.pi, 20, dtype="float64") + phi_edges = np.linspace(0.0, 2 * np.pi, 20, dtype="float64") + + r = (r_edges[0:-1] + r_edges[1:]) / 2.0 + theta = (theta_edges[0:-1] + theta_edges[1:]) / 2.0 + phi = (phi_edges[0:-1] + phi_edges[1:]) / 2.0 + + dr = r_edges[1:] - r_edges[:-1] + dtheta = theta_edges[1:] - theta_edges[:-1] + dphi = phi_edges[1:] - phi_edges[:-1] + + r_th_ph = np.meshgrid(r, theta, phi) + d_r_th_ph = np.meshgrid(dr, dtheta, dphi) + r_th_ph = [r_th_ph[i].ravel() for i in range(3)] + d_r_th_ph = [d_r_th_ph[i].ravel() for i in range(3)] + + x_y_z, d_x_y_z = cartesian_bboxes( + bbox_handler, + r_th_ph[0], + r_th_ph[1], + r_th_ph[2], + d_r_th_ph[0], + d_r_th_ph[1], + d_r_th_ph[2], + ) + x_y_z = np.column_stack(x_y_z) + d_x_y_z = np.column_stack(d_x_y_z) + + le = [r_th_ph[i] - d_r_th_ph[i] / 2.0 for i in range(3)] + re = [r_th_ph[i] + d_r_th_ph[i] / 2.0 for i in range(3)] + xyz_le, xyz_re = cartesian_bboxes_edges( + bbox_handler, + le[0], + le[1], + le[2], + re[0], + re[1], + re[2], + ) + xyz_le = np.column_stack(xyz_le) + xyz_re = np.column_stack(xyz_re) + + centers = (xyz_le + xyz_re) / 2.0 + widths = xyz_re - xyz_le + + assert np.allclose(centers, x_y_z) + assert np.allclose(widths, d_x_y_z) diff --git a/yt_idv/tests/test_spherical_vol_rendering.py b/yt_idv/tests/test_spherical_vol_rendering.py new file mode 100644 index 0000000..46fd3d8 --- /dev/null +++ b/yt_idv/tests/test_spherical_vol_rendering.py @@ -0,0 +1,92 @@ +import numpy as np +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]]), + "field": ("index", "r"), + "camera_position": [-0.5, -1, 2.5], + }, + "whole": { + "bbox": np.array([[0.0, 1.0], [0.0, 2 * np.pi], [0, np.pi]]), + "field": ("index", "phi"), + "camera_position": [0.5, 0.5, 2], + }, + "quadrant_shell": { + "bbox": np.array([[0.6, 1.0], [0.0, np.pi / 2], [0.0, np.pi / 2]]), + "field": ("index", "theta"), + }, +} + + +@pytest.mark.parametrize("bbox_option", bbox_options.keys()) +def test_spherical_bounds(osmesa_empty_rc, image_store, bbox_option): + + sz = (32, 32, 32) + fake_data = {"density": np.random.random(sz)} + + bbox = bbox_options[bbox_option]["bbox"] + + ds = yt.load_uniform_grid( + fake_data, + sz, + bbox=bbox, + nprocs=8, + geometry=("spherical", ("r", "phi", "theta")), + length_unit="m", + ) + dd = ds.all_data() + + field = bbox_options[bbox_option]["field"] + osmesa_empty_rc.add_scene(dd, field, no_ghost=True) + osmesa_empty_rc.scene.components[0].sample_factor = 20.0 + osmesa_empty_rc.scene.components[0].cmap_log = False + cpos = bbox_options[bbox_option].get("camera_position", None) + if cpos: + osmesa_empty_rc.scene.camera.set_position(cpos) + + image_store(osmesa_empty_rc) + + +@pytest.mark.parametrize("nprocs", [1, 2, 4, 16]) +def test_spherical_nprocs(osmesa_empty_rc, image_store, nprocs): + + sz = (32, 32, 32) + fake_data = {"density": np.random.random(sz)} + + bbox_option = "whole" + bbox = bbox_options[bbox_option]["bbox"] + + ds = yt.load_uniform_grid( + fake_data, + sz, + bbox=bbox, + nprocs=nprocs, + geometry=("spherical", ("r", "phi", "theta")), + length_unit="m", + ) + dd = ds.all_data() + + field = bbox_options[bbox_option]["field"] + osmesa_empty_rc.add_scene(dd, field, no_ghost=True) + osmesa_empty_rc.scene.components[0].sample_factor = 20.0 + osmesa_empty_rc.scene.components[0].cmap_log = False + osmesa_empty_rc.scene.components[0]._reset_cmap_bounds() + cpos = bbox_options[bbox_option].get("camera_position", None) + if cpos: + osmesa_empty_rc.scene.camera.set_position(cpos) + + image_store(osmesa_empty_rc) diff --git a/yt_idv/tests/test_yt_idv.py b/yt_idv/tests/test_yt_idv.py index 07f68e5..b38dbdd 100644 --- a/yt_idv/tests/test_yt_idv.py +++ b/yt_idv/tests/test_yt_idv.py @@ -1,12 +1,9 @@ """Tests for `yt_idv` package.""" -import base64 - import numpy as np import pytest import yt import yt.testing -from pytest_html import extras as html_extras import yt_idv from yt_idv import shader_objects @@ -38,18 +35,6 @@ def osmesa_empty(): rc.osmesa.OSMesaDestroyContext(rc.context) -@pytest.fixture() -def image_store(request, extras, tmpdir): - def _snap_image(rc): - image = rc.run() - img = yt.write_bitmap(image, None) - content = base64.b64encode(img).decode("ascii") - extras.append(html_extras.png(content)) - extras.append(html_extras.html("
")) - - return _snap_image - - def test_snapshots(osmesa_fake_amr, image_store): """Check that we can make some snapshots.""" osmesa_fake_amr.scene.components[0].render_method = "max_intensity"