Source code for yt.visualization.volume_rendering.utils

import numpy as np

from yt.data_objects.selection_objects.data_selection_objects import (
    YTSelectionContainer3D,
)
from yt.data_objects.static_output import Dataset
from yt.utilities.lib import bounding_volume_hierarchy
from yt.utilities.lib.image_samplers import (
    InterpolatedProjectionSampler,
    ProjectionSampler,
    VolumeRenderSampler,
)
from yt.utilities.on_demand_imports import NotAModule

try:
    from yt.utilities.lib.embree_mesh import mesh_traversal  # type: ignore
# Catch ValueError in case size of objects in Cython change
except (ImportError, ValueError):
    mesh_traversal = NotAModule("pyembree")


[docs] def data_source_or_all(data_source): if isinstance(data_source, Dataset): data_source = data_source.all_data() if not isinstance(data_source, (YTSelectionContainer3D, type(None))): raise RuntimeError( "The data_source is not a valid 3D data container.\n" "Expected an object of type YTSelectionContainer3D but received " "an object of type %s." % type(data_source) ) return data_source
[docs] def new_mesh_sampler(camera, render_source, engine): params = ensure_code_unit_params(camera._get_sampler_params(render_source)) args = ( np.atleast_3d(params["vp_pos"]), np.atleast_3d(params["vp_dir"]), params["center"], params["bounds"], np.atleast_3d(params["image"]).astype("float64"), params["x_vec"], params["y_vec"], params["width"], render_source.volume_method, ) kwargs = {"lens_type": params["lens_type"]} if engine == "embree": sampler = mesh_traversal.EmbreeMeshSampler(*args, **kwargs) elif engine == "yt": sampler = bounding_volume_hierarchy.BVHMeshSampler(*args, **kwargs) return sampler
[docs] def new_volume_render_sampler(camera, render_source): params = ensure_code_unit_params(camera._get_sampler_params(render_source)) params.update(transfer_function=render_source.transfer_function) params.update(transfer_function=render_source.transfer_function) params.update(num_samples=render_source.num_samples) args = ( np.atleast_3d(params["vp_pos"]), np.atleast_3d(params["vp_dir"]), params["center"], params["bounds"], params["image"], params["x_vec"], params["y_vec"], params["width"], render_source.volume_method, params["transfer_function"], params["num_samples"], ) kwargs = { "lens_type": params["lens_type"], } if "camera_data" in params: kwargs["camera_data"] = params["camera_data"] if render_source.zbuffer is not None: kwargs["zbuffer"] = render_source.zbuffer.z args[4][:] = np.reshape( render_source.zbuffer.rgba[:], (camera.resolution[0], camera.resolution[1], 4), ) else: kwargs["zbuffer"] = np.ones(params["image"].shape[:2], "float64") sampler = VolumeRenderSampler(*args, **kwargs) return sampler
[docs] def new_interpolated_projection_sampler(camera, render_source): params = ensure_code_unit_params(camera._get_sampler_params(render_source)) params.update(transfer_function=render_source.transfer_function) params.update(num_samples=render_source.num_samples) args = ( np.atleast_3d(params["vp_pos"]), np.atleast_3d(params["vp_dir"]), params["center"], params["bounds"], params["image"], params["x_vec"], params["y_vec"], params["width"], render_source.volume_method, params["num_samples"], ) kwargs = {"lens_type": params["lens_type"]} if render_source.zbuffer is not None: kwargs["zbuffer"] = render_source.zbuffer.z else: kwargs["zbuffer"] = np.ones(params["image"].shape[:2], "float64") sampler = InterpolatedProjectionSampler(*args, **kwargs) return sampler
[docs] def new_projection_sampler(camera, render_source): params = ensure_code_unit_params(camera._get_sampler_params(render_source)) params.update(transfer_function=render_source.transfer_function) params.update(num_samples=render_source.num_samples) args = ( np.atleast_3d(params["vp_pos"]), np.atleast_3d(params["vp_dir"]), params["center"], params["bounds"], params["image"], params["x_vec"], params["y_vec"], params["width"], render_source.volume_method, params["num_samples"], ) kwargs = { "lens_type": params["lens_type"], } if render_source.zbuffer is not None: kwargs["zbuffer"] = render_source.zbuffer.z else: kwargs["zbuffer"] = np.ones(params["image"].shape[:2], "float64") sampler = ProjectionSampler(*args, **kwargs) return sampler
[docs] def get_corners(le, re): return np.array( [ [le[0], le[1], le[2]], [re[0], le[1], le[2]], [re[0], re[1], le[2]], [le[0], re[1], le[2]], [le[0], le[1], re[2]], [re[0], le[1], re[2]], [re[0], re[1], re[2]], [le[0], re[1], re[2]], ], dtype="float64", )
[docs] def ensure_code_unit_params(params): for param_name in ["center", "vp_pos", "vp_dir", "width"]: param = params[param_name] if hasattr(param, "in_units"): params[param_name] = param.in_units("code_length") bounds = params["bounds"] if hasattr(bounds[0], "units"): params["bounds"] = tuple(b.in_units("code_length").d for b in bounds) return params