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 "
f"an object of type {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