def sanitize_width(self, axis, width, depth):
if width is None:
# initialize the index if it is not already initialized
self.ds.index
# Default to code units
if not is_sequence(axis):
xax = self.x_axis[axis]
yax = self.y_axis[axis]
w = self.ds.domain_width[np.array([xax, yax])]
else:
# axis is actually the normal vector
# for an off-axis data object.
mi = np.argmin(self.ds.domain_width)
w = self.ds.domain_width[np.array((mi, mi))]
width = (w[0], w[1])
elif is_sequence(width):
width = validate_sequence_width(width, self.ds)
elif isinstance(width, YTQuantity):
width = (width, width)
elif isinstance(width, Number):
width = (
self.ds.quan(width, "code_length"),
self.ds.quan(width, "code_length"),
)
else:
raise YTInvalidWidthError(width)
if depth is not None:
depth = self.sanitize_depth(depth)
return width + depth
return width
def _sanitize_camera_property_units(value, scene):
if is_sequence(value):
if len(value) == 1:
return _sanitize_camera_property_units(value[0], scene)
elif isinstance(value, YTArray) and len(value) == 3:
return scene.arr(value).in_units("unitary")
elif (len(value) == 2 and isinstance(value[0], numeric_type)
and isinstance(value[1], str)):
return scene.arr(
[scene.arr(value[0], value[1]).in_units("unitary")] * 3)
if len(value) == 3:
if all([is_sequence(v) for v in value]):
if all([
isinstance(v[0], numeric_type)
and isinstance(v[1], str) for v in value
]):
return scene.arr([scene.arr(v[0], v[1]) for v in value])
else:
raise RuntimeError(
f"Cannot set camera width to invalid value '{value}'")
return scene.arr(value, "unitary")
else:
if isinstance(value, (YTQuantity, YTArray)):
return scene.arr([value.d] * 3, value.units).in_units("unitary")
elif isinstance(value, numeric_type):
return scene.arr([value] * 3, "unitary")
raise RuntimeError(f"Cannot set camera width to invalid value '{value}'")
def sanitize_center(self, center, axis):
if isinstance(center, str):
if center.lower() == "m" or center.lower() == "max":
v, center = self.ds.find_max(("gas", "density"))
center = self.ds.arr(center, "code_length")
elif center.lower() == "c" or center.lower() == "center":
# domain_left_edge and domain_right_edge might not be
# initialized until we create the index, so create it
self.ds.index
center = (self.ds.domain_left_edge +
self.ds.domain_right_edge) / 2
else:
raise RuntimeError(f'center keyword "{center}" not recognized')
elif isinstance(center, YTArray):
return self.ds.arr(center), self.convert_to_cartesian(center)
elif is_sequence(center):
if isinstance(center[0], str) and isinstance(center[1], str):
if center[0].lower() == "min":
v, center = self.ds.find_min(center[1])
elif center[0].lower() == "max":
v, center = self.ds.find_max(center[1])
else:
raise RuntimeError(
f'center keyword "{center}" not recognized')
center = self.ds.arr(center, "code_length")
elif is_sequence(center[0]) and isinstance(center[1], str):
center = self.ds.arr(center[0], center[1])
else:
center = self.ds.arr(center, "code_length")
else:
raise RuntimeError(f'center keyword "{center}" not recognized')
# This has to return both a center and a display_center
display_center = self.convert_to_cartesian(center)
return center, display_center
def to_frb(self, width, resolution, height=None, periodic=False):
r"""This function returns a FixedResolutionBuffer generated from this
object.
An ObliqueFixedResolutionBuffer is an object that accepts a
variable-resolution 2D object and transforms it into an NxM bitmap that
can be plotted, examined or processed. This is a convenience function
to return an FRB directly from an existing 2D data object. Unlike the
corresponding to_frb function for other YTSelectionContainer2D objects,
this does not accept a 'center' parameter as it is assumed to be
centered at the center of the cutting plane.
Parameters
----------
width : width specifier
This can either be a floating point value, in the native domain
units of the simulation, or a tuple of the (value, unit) style.
This will be the width of the FRB.
height : height specifier, optional
This will be the height of the FRB, by default it is equal to width.
resolution : int or tuple of ints
The number of pixels on a side of the final FRB.
periodic : boolean
This can be true or false, and governs whether the pixelization
will span the domain boundaries.
Returns
-------
frb : :class:`~yt.visualization.fixed_resolution.ObliqueFixedResolutionBuffer`
A fixed resolution buffer, which can be queried for fields.
Examples
--------
>>> v, c = ds.find_max("density")
>>> sp = ds.sphere(c, (100.0, 'au'))
>>> L = sp.quantities.angular_momentum_vector()
>>> cutting = ds.cutting(L, c)
>>> frb = cutting.to_frb( (1.0, 'pc'), 1024)
>>> write_image(np.log10(frb["Density"]), 'density_1pc.png')
"""
if is_sequence(width):
validate_width_tuple(width)
width = self.ds.quan(width[0], width[1])
if height is None:
height = width
elif is_sequence(height):
validate_width_tuple(height)
height = self.ds.quan(height[0], height[1])
if not is_sequence(resolution):
resolution = (resolution, resolution)
from yt.visualization.fixed_resolution import FixedResolutionBuffer
bounds = (-width / 2.0, width / 2.0, -height / 2.0, height / 2.0)
frb = FixedResolutionBuffer(self,
bounds,
resolution,
periodic=periodic)
return frb
def fake_random_ds(
ndims,
peak_value=1.0,
fields=("density", "velocity_x", "velocity_y", "velocity_z"),
units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
particle_fields=None,
particle_field_units=None,
negative=False,
nprocs=1,
particles=0,
length_unit=1.0,
unit_system="cgs",
bbox=None,
):
from yt.loaders import load_uniform_grid
prng = RandomState(0x4D3D3D3)
if not is_sequence(ndims):
ndims = [ndims, ndims, ndims]
else:
assert len(ndims) == 3
if not is_sequence(negative):
negative = [negative for f in fields]
assert len(fields) == len(negative)
offsets = []
for n in negative:
if n:
offsets.append(0.5)
else:
offsets.append(0.0)
data = {}
for field, offset, u in zip(fields, offsets, units):
v = (prng.random_sample(ndims) - offset) * peak_value
if field[0] == "all":
v = v.ravel()
data[field] = (v, u)
if particles:
if particle_fields is not None:
for field, unit in zip(particle_fields, particle_field_units):
if field in ("particle_position", "particle_velocity"):
data["io", field] = (prng.random_sample((int(particles), 3)), unit)
else:
data["io", field] = (prng.random_sample(size=int(particles)), unit)
else:
for f in (f"particle_position_{ax}" for ax in "xyz"):
data["io", f] = (prng.random_sample(size=particles), "code_length")
for f in (f"particle_velocity_{ax}" for ax in "xyz"):
data["io", f] = (prng.random_sample(size=particles) - 0.5, "cm/s")
data["io", "particle_mass"] = (prng.random_sample(particles), "g")
ug = load_uniform_grid(
data,
ndims,
length_unit=length_unit,
nprocs=nprocs,
unit_system=unit_system,
bbox=bbox,
)
return ug
def construct_image(ds, axis, data_source, center, image_res, width,
length_unit):
if width is None:
width = ds.domain_width[axis_wcs[axis]]
unit = ds.get_smallest_appropriate_unit(width[0])
mylog.info("Making an image of the entire domain, "
"so setting the center to the domain center.")
else:
width = ds.coordinates.sanitize_width(axis, width, None)
unit = str(width[0].units)
if is_sequence(image_res):
nx, ny = image_res
else:
nx, ny = image_res, image_res
dx = width[0] / nx
dy = width[1] / ny
crpix = [0.5 * (nx + 1), 0.5 * (ny + 1)]
if unit == "unitary":
unit = ds.get_smallest_appropriate_unit(ds.domain_width.max())
elif unit == "code_length":
unit = ds.get_smallest_appropriate_unit(ds.quan(1.0, "code_length"))
unit = sanitize_fits_unit(unit)
if length_unit is None:
length_unit = unit
if any(char.isdigit() for char in length_unit) and "*" in length_unit:
length_unit = length_unit.split("*")[-1]
cunit = [length_unit] * 2
ctype = ["LINEAR"] * 2
cdelt = [dx.in_units(length_unit), dy.in_units(length_unit)]
if is_sequence(axis):
crval = center.in_units(length_unit)
else:
crval = [center[idx].in_units(length_unit) for idx in axis_wcs[axis]]
if hasattr(data_source, "to_frb"):
if is_sequence(axis):
frb = data_source.to_frb(width[0], (nx, ny), height=width[1])
else:
frb = data_source.to_frb(width[0], (nx, ny),
center=center,
height=width[1])
else:
frb = None
w = _astropy.pywcs.WCS(naxis=2)
w.wcs.crpix = crpix
w.wcs.cdelt = cdelt
w.wcs.crval = crval
w.wcs.cunit = cunit
w.wcs.ctype = ctype
return w, frb, length_unit
def set_fields(self, fields, log_fields, no_ghost, force=False):
new_fields = self.data_source._determine_fields(fields)
regenerate_data = (
self.fields is None
or len(self.fields) != len(new_fields)
or self.fields != new_fields
or force
)
if not is_sequence(log_fields):
log_fields = [log_fields]
new_log_fields = list(log_fields)
self.tree.trunk.set_dirty(regenerate_data)
self.fields = new_fields
if self.log_fields is not None and not regenerate_data:
flip_log = list(map(operator.ne, self.log_fields, new_log_fields))
else:
flip_log = [False] * len(new_log_fields)
self.log_fields = new_log_fields
self.no_ghost = no_ghost
del self.bricks, self.brick_dimensions
self.brick_dimensions = []
bricks = []
for b in self.traverse():
list(map(_apply_log, b.my_data, flip_log, self.log_fields))
bricks.append(b)
self.bricks = np.array(bricks)
self.brick_dimensions = np.array(self.brick_dimensions)
self._initialized = True
def add_field(self,
name,
function,
sampling_type,
*,
force_override=False,
**kwargs):
sampling_type = self._sanitize_sampling_type(sampling_type)
if isinstance(name, str) or not is_sequence(name):
if sampling_type == "particle":
ftype = "all"
else:
ftype = "gas"
name = (ftype, name)
# Handle the case where the field has already been added.
if not force_override and name in self:
mylog.warning(
"Field %s already exists. To override use `force_override=True`.",
name,
)
return super().add_field(name,
function,
sampling_type,
force_override=force_override,
**kwargs)
def __init__(
self,
outputs,
parallel=True,
setup_function=None,
mixed_dataset_types=False,
**kwargs,
):
# This is needed to properly set _pre_outputs for Simulation subclasses.
self._mixed_dataset_types = mixed_dataset_types
if is_sequence(outputs) and not isinstance(outputs, str):
self._pre_outputs = outputs[:]
self.tasks = AnalysisTaskProxy(self)
self.params = TimeSeriesParametersContainer(self)
if setup_function is None:
def _null(x):
return None
setup_function = _null
self._setup_function = setup_function
for type_name in data_object_registry:
setattr(self, type_name,
functools.partial(DatasetSeriesObject, self, type_name))
self.parallel = parallel
self.kwargs = kwargs
def add_field(self,
name,
function,
sampling_type,
*,
force_override=False,
**kwargs):
if isinstance(name, str) or not is_sequence(name):
# the base method only accepts proper tuple field keys
# and is only used internally, while this method is exposed to users
# and is documented as usable with single strings as name
if sampling_type == "particle":
ftype = "all"
else:
ftype = "gas"
name = (ftype, name)
# Handle the case where the field has already been added.
if not force_override and name in self:
mylog.warning(
"Field %s already exists. To override use `force_override=True`.",
name,
)
return super().add_field(name,
function,
sampling_type,
force_override=force_override,
**kwargs)
def __init__(self, fsize, axrect, figure, axes):
"""Initialize PlotMPL class"""
import matplotlib.figure
self._plot_valid = True
if figure is None:
if not is_sequence(fsize):
fsize = (fsize, fsize)
self.figure = matplotlib.figure.Figure(figsize=fsize, frameon=True)
else:
figure.set_size_inches(fsize)
self.figure = figure
if axes is None:
self._create_axes(axrect)
else:
axes.cla()
axes.set_position(axrect)
self.axes = axes
canvas_classes = self._set_canvas()
self.canvas = canvas_classes[0](self.figure)
if len(canvas_classes) > 1:
self.manager = canvas_classes[1](self.canvas, 1)
for which in ["major", "minor"]:
for axis in "xy":
self.axes.tick_params(
which=which, axis=axis, direction="in", top=True, right=True
)
def __init__(self, fsize, axrect, figure, axes):
"""Initialize PlotMPL class"""
import matplotlib.figure
self._plot_valid = True
if figure is None:
if not is_sequence(fsize):
fsize = (fsize, fsize)
self.figure = matplotlib.figure.Figure(figsize=fsize, frameon=True)
else:
figure.set_size_inches(fsize)
self.figure = figure
if axes is None:
self._create_axes(axrect)
else:
axes.cla()
axes.set_position(axrect)
self.axes = axes
self.interactivity = get_interactivity()
figure_canvas, figure_manager = self._get_canvas_classes()
self.canvas = figure_canvas(self.figure)
if figure_manager is not None:
self.manager = figure_manager(self.canvas, 1)
self.axes.tick_params(which="both",
axis="both",
direction="in",
top=True,
right=True)
def resolution(self, value):
if is_sequence(value):
if len(value) != 2:
raise RuntimeError
else:
value = (value, value)
self._resolution = value
def _get_plot_instance(self, field):
fontscale = self._font_properties._size / 14.0
top_buff_size = 0.35 * fontscale
x_axis_size = 1.35 * fontscale
y_axis_size = 0.7 * fontscale
right_buff_size = 0.2 * fontscale
if is_sequence(self.figure_size):
figure_size = self.figure_size
else:
figure_size = (self.figure_size, self.figure_size)
xbins = np.array([x_axis_size, figure_size[0], right_buff_size])
ybins = np.array([y_axis_size, figure_size[1], top_buff_size])
size = [xbins.sum(), ybins.sum()]
x_frac_widths = xbins / size[0]
y_frac_widths = ybins / size[1]
axrect = (
x_frac_widths[0],
y_frac_widths[0],
x_frac_widths[1],
y_frac_widths[1],
)
try:
plot = self.plots[field]
except KeyError:
plot = PlotMPL(self.figure_size, axrect, None, None)
self.plots[field] = plot
return plot
def _prepare_grid(self):
"""Copies all the appropriate attributes from the index."""
# This is definitely the slowest part of generating the index
# Now we give it pointers to all of its attributes
# Note that to keep in line with Enzo, we have broken PEP-8
h = self.index # cache it
my_ind = self.id - self._id_offset
self.ActiveDimensions = h.grid_dimensions[my_ind]
self.LeftEdge = h.grid_left_edge[my_ind]
self.RightEdge = h.grid_right_edge[my_ind]
# This can be expensive so we allow people to disable this behavior
# via a config option
if RECONSTRUCT_INDEX:
if is_sequence(self.Parent) and len(self.Parent) > 0:
p = self.Parent[0]
else:
p = self.Parent
if p is not None and p != []:
# clamp grid edges to an integer multiple of the parent cell
# width
clamp_edges(self.LeftEdge, p.LeftEdge, p.dds)
clamp_edges(self.RightEdge, p.RightEdge, p.dds)
h.grid_levels[my_ind, 0] = self.Level
# This might be needed for streaming formats
# self.Time = h.gridTimes[my_ind,0]
self.NumberOfParticles = h.grid_particle_count[my_ind, 0]
def sanitize_depth(self, depth):
if is_sequence(depth):
validate_width_tuple(depth)
depth = (self.ds.quan(depth[0], fix_unitary(depth[1])), )
elif isinstance(depth, Number):
depth = (self.ds.quan(depth,
"code_length",
registry=self.ds.unit_registry), )
elif isinstance(depth, YTQuantity):
depth = (depth, )
else:
raise YTInvalidWidthError(depth)
return depth
def from_sizes(cls, sizes):
pool = cls()
rank = pool.comm.rank
for i, size in enumerate(always_iterable(sizes)):
if is_sequence(size):
size, name = size
else:
name = "workgroup_%02i" % i
pool.add_workgroup(size, name=name)
for wg in pool.workgroups:
if rank in wg.ranks:
workgroup = wg
return pool, workgroup
def set_defaults_from_data_source(self, data_source):
"""Resets the camera attributes to their default values"""
position = data_source.ds.domain_right_edge
width = 1.5 * data_source.ds.domain_width.max()
(xmi, xma), (ymi, yma), (zmi, zma) = data_source.quantities["Extrema"](
["x", "y", "z"])
width = np.sqrt((xma - xmi)**2 + (yma - ymi)**2 + (zma - zmi)**2)
focus = data_source.get_field_parameter("center")
if is_sequence(width) and len(width) > 1 and isinstance(width[1], str):
width = data_source.ds.quan(width[0], units=width[1])
# Now convert back to code length for subsequent manipulation
width = width.in_units("code_length") # .value
if not is_sequence(width):
width = data_source.ds.arr([width, width, width],
units="code_length")
# left/right, top/bottom, front/back
if not isinstance(width, YTArray):
width = data_source.ds.arr(width, units="code_length")
if not isinstance(focus, YTArray):
focus = data_source.ds.arr(focus, units="code_length")
# We can't use the property setters yet, since they rely on attributes
# that will not be set up until the base class initializer is called.
# See Issue #1131.
self._width = width
self._focus = focus
self._position = position
self._domain_center = data_source.ds.domain_center
self._domain_width = data_source.ds.domain_width
super().__init__(self.focus - self.position,
self.north_vector,
steady_north=False)
self._moved = True
def _los_field(field, data):
if data.has_field_parameter(f"bulk_{basename}"):
fns = [(fc[0], f"relative_{fc[1]}") for fc in field_comps]
else:
fns = field_comps
ax = data.get_field_parameter("axis")
if is_sequence(ax):
# Make sure this is a unit vector
ax /= np.sqrt(np.dot(ax, ax))
ret = data[fns[0]] * ax[0] + data[fns[1]] * ax[1] + data[fns[2]] * ax[2]
elif ax in [0, 1, 2]:
ret = data[fns[ax]]
else:
raise NeedsParameter(["axis"])
return ret
def _validate_point(point, ds, start=False):
if not is_sequence(point):
raise RuntimeError("Input point must be array-like")
if not isinstance(point, YTArray):
point = ds.arr(point, "code_length", dtype=np.float64)
if len(point.shape) != 1:
raise RuntimeError("Input point must be a 1D array")
if point.shape[0] < ds.dimensionality:
raise RuntimeError("Input point must have an element for each dimension")
# need to pad to 3D elements to avoid issues later
if point.shape[0] < 3:
if start:
val = 0
else:
val = 1
point = np.append(point.d, [val] * (3 - ds.dimensionality)) * point.uq
return point
def __init__(self, data_source, figure_size, fontsize):
from matplotlib.font_manager import FontProperties
self.data_source = data_source
self.ds = data_source.ds
self.ts = self._initialize_dataset(self.ds)
if is_sequence(figure_size):
self.figure_size = float(figure_size[0]), float(figure_size[1])
else:
self.figure_size = float(figure_size)
font_path = matplotlib.get_data_path() + "/fonts/ttf/STIXGeneral.ttf"
self._font_properties = FontProperties(size=fontsize, fname=font_path)
self._font_color = None
self._xlabel = None
self._ylabel = None
self._minorticks = {}
self._field_transform = {}
def fake_particle_ds(
fields=None,
units=None,
negative=None,
npart=16 ** 3,
length_unit=1.0,
data=None,
):
from yt.loaders import load_particles
prng = RandomState(0x4D3D3D3)
if negative is not None and not is_sequence(negative):
negative = [negative for f in fields]
fields, units, negative = _check_field_unit_args_helper(
{
"fields": fields,
"units": units,
"negative": negative,
},
{
"fields": _fake_particle_ds_default_fields,
"units": _fake_particle_ds_default_units,
"negative": _fake_particle_ds_default_negative,
},
)
offsets = []
for n in negative:
if n:
offsets.append(0.5)
else:
offsets.append(0.0)
data = data if data else {}
for field, offset, u in zip(fields, offsets, units):
if field in data:
v = data[field]
continue
if "position" in field:
v = prng.normal(loc=0.5, scale=0.25, size=npart)
np.clip(v, 0.0, 1.0, v)
v = prng.random_sample(npart) - offset
data[field] = (v, u)
bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]])
ds = load_particles(data, 1.0, bbox=bbox)
return ds
def __init__(
self,
ds,
normal,
fields,
center="c",
width=(1.0, "unitary"),
weight_field=None,
image_res=512,
data_source=None,
north_vector=None,
depth=(1.0, "unitary"),
method="integrate",
length_unit=None,
):
fields = list(iter_fields(fields))
center, dcenter = ds.coordinates.sanitize_center(center, 4)
buf = {}
width = ds.coordinates.sanitize_width(normal, width, depth)
wd = tuple(el.in_units("code_length").v for el in width)
if not is_sequence(image_res):
image_res = (image_res, image_res)
res = (image_res[0], image_res[1])
if data_source is None:
source = ds
else:
source = data_source
for field in fields:
buf[field] = off_axis_projection(
source,
center,
normal,
wd,
res,
field,
north_vector=north_vector,
method=method,
weight=weight_field,
).swapaxes(0, 1)
center = ds.arr([0.0] * 2, "code_length")
w, not_an_frb, lunit = construct_image(
ds, normal, buf, center, image_res, width, length_unit
)
super(FITSOffAxisProjection, self).__init__(
buf, fields=fields, wcs=w, length_unit=lunit, ds=ds
)
def _deserialize_from_h5(g, ds):
result = {}
for item in g:
if item == "chunks":
continue
if "units" in g[item].attrs:
if is_sequence(g[item]):
result[item] = ds.arr(g[item][:], g[item].attrs["units"])
else:
result[item] = ds.quan(g[item][()], g[item].attrs["units"])
elif isinstance(g[item], h5py.Group):
result[item] = _deserialize_from_h5(g[item], ds)
elif g[item] == "None":
result[item] = None
else:
try:
result[item] = g[item][:] # try array
except ValueError:
result[item] = g[item][()] # fallback to scalar
return result
def __init__(self, positions, colors=None, color_stride=1, radii=None):
assert positions.ndim == 2 and positions.shape[1] == 3
if colors is not None:
assert colors.ndim == 2 and colors.shape[1] == 4
assert colors.shape[0] == positions.shape[0]
if not is_sequence(radii):
if radii is not None: # broadcast the value
radii = radii * np.ones(positions.shape[0], dtype="int64")
else: # default radii to 0 pixels (i.e. point is 1 pixel wide)
radii = np.zeros(positions.shape[0], dtype="int64")
else:
assert radii.ndim == 1
assert radii.shape[0] == positions.shape[0]
self.positions = positions
# If colors aren't individually set, make black with full opacity
if colors is None:
colors = np.ones((len(positions), 4))
self.colors = colors
self.color_stride = color_stride
self.radii = radii
def fake_particle_ds(
fields=(
"particle_position_x",
"particle_position_y",
"particle_position_z",
"particle_mass",
"particle_velocity_x",
"particle_velocity_y",
"particle_velocity_z",
),
units=("cm", "cm", "cm", "g", "cm/s", "cm/s", "cm/s"),
negative=(False, False, False, False, True, True, True),
npart=16 ** 3,
length_unit=1.0,
data=None,
):
from yt.loaders import load_particles
prng = RandomState(0x4D3D3D3)
if not is_sequence(negative):
negative = [negative for f in fields]
assert len(fields) == len(negative)
offsets = []
for n in negative:
if n:
offsets.append(0.5)
else:
offsets.append(0.0)
data = data if data else {}
for field, offset, u in zip(fields, offsets, units):
if field in data:
v = data[field]
continue
if "position" in field:
v = prng.normal(loc=0.5, scale=0.25, size=npart)
np.clip(v, 0.0, 1.0, v)
v = prng.random_sample(npart) - offset
data[field] = (v, u)
bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]])
ds = load_particles(data, 1.0, bbox=bbox)
return ds
def __init__(self, data_source, figure_size, fontsize):
self.data_source = data_source
self.ds = data_source.ds
self.ts = self._initialize_dataset(self.ds)
if is_sequence(figure_size):
self.figure_size = float(figure_size[0]), float(figure_size[1])
else:
self.figure_size = float(figure_size)
if sys.version_info >= (3, 9):
font_dict = DEFAULT_FONT_PROPERTIES | {"size": fontsize}
else:
font_dict = {**DEFAULT_FONT_PROPERTIES, "size": fontsize}
self._font_properties = FontProperties(**font_dict)
self._font_color = None
self._xlabel = None
self._ylabel = None
self._minorticks = {}
self._field_transform = {}
self.setup_defaults()
def _initialize_dataset(self, ts):
if not isinstance(ts, DatasetSeries):
if not is_sequence(ts):
ts = [ts]
ts = DatasetSeries(ts)
return ts
def off_axis_projection(
data_source,
center,
normal_vector,
width,
resolution,
item,
weight=None,
volume=None,
no_ghost=False,
interpolated=False,
north_vector=None,
num_threads=1,
method="integrate",
):
r"""Project through a dataset, off-axis, and return the image plane.
This function will accept the necessary items to integrate through a volume
at an arbitrary angle and return the integrated field of view to the user.
Note that if a weight is supplied, it will multiply the pre-interpolated
values together, then create cell-centered values, then interpolate within
the cell to conduct the integration.
Parameters
----------
data_source : ~yt.data_objects.static_output.Dataset
or ~yt.data_objects.data_containers.YTSelectionDataContainer
This is the dataset or data object to volume render.
center : array_like
The current 'center' of the view port -- the focal point for the
camera.
normal_vector : array_like
The vector between the camera position and the center.
width : float or list of floats
The current width of the image. If a single float, the volume is
cubical, but if not, it is left/right, top/bottom, front/back
resolution : int or list of ints
The number of pixels in each direction.
item: string
The field to project through the volume
weight : optional, default None
If supplied, the field will be pre-multiplied by this, then divided by
the integrated value of this field. This returns an average rather
than a sum.
volume : `yt.extensions.volume_rendering.AMRKDTree`, optional
The volume to ray cast through. Can be specified for finer-grained
control, but otherwise will be automatically generated.
no_ghost: bool, optional
Optimization option. If True, homogenized bricks will
extrapolate out from grid instead of interpolating from
ghost zones that have to first be calculated. This can
lead to large speed improvements, but at a loss of
accuracy/smoothness in resulting image. The effects are
less notable when the transfer function is smooth and
broad. Default: True
interpolated : optional, default False
If True, the data is first interpolated to vertex-centered data,
then tri-linearly interpolated along the ray. Not suggested for
quantitative studies.
north_vector : optional, array_like, default None
A vector that, if specified, restricts the orientation such that the
north vector dotted into the image plane points "up". Useful for rotations
num_threads: integer, optional, default 1
Use this many OpenMP threads during projection.
method : string
The method of projection. Valid methods are:
"integrate" with no weight_field specified : integrate the requested
field along the line of sight.
"integrate" with a weight_field specified : weight the requested
field by the weighting field and integrate along the line of sight.
"sum" : This method is the same as integrate, except that it does not
multiply by a path length when performing the integration, and is
just a straight summation of the field along the given axis. WARNING:
This should only be used for uniform resolution grid datasets, as other
datasets may result in unphysical images.
or camera movements.
Returns
-------
image : array
An (N,N) array of the final integrated values, in float64 form.
Examples
--------
>>> image = off_axis_projection(ds, [0.5, 0.5, 0.5], [0.2,0.3,0.4],
... 0.2, N, "temperature", "density")
>>> write_image(np.log10(image), "offaxis.png")
"""
if method not in ("integrate", "sum"):
raise NotImplementedError(
"Only 'integrate' or 'sum' methods are valid for off-axis-projections"
)
if interpolated:
raise NotImplementedError(
"Only interpolated=False methods are currently implemented "
"for off-axis-projections")
data_source = data_source_or_all(data_source)
item = data_source._determine_fields([item])[0]
# Assure vectors are numpy arrays as expected by cython code
normal_vector = np.array(normal_vector, dtype="float64")
if north_vector is not None:
north_vector = np.array(north_vector, dtype="float64")
# Add the normal as a field parameter to the data source
# so line of sight fields can use it
data_source.set_field_parameter("axis", normal_vector)
# Sanitize units
if not hasattr(center, "units"):
center = data_source.ds.arr(center, "code_length")
if not hasattr(width, "units"):
width = data_source.ds.arr(width, "code_length")
if hasattr(data_source.ds, "_sph_ptypes"):
if method != "integrate":
raise NotImplementedError("SPH Only allows 'integrate' method")
sph_ptypes = data_source.ds._sph_ptypes
fi = data_source.ds.field_info[item]
raise_error = False
ptype = sph_ptypes[0]
ppos = [f"particle_position_{ax}" for ax in "xyz"]
# Assure that the field we're trying to off-axis project
# has a field type as the SPH particle type or if the field is an
# alias to an SPH field or is a 'gas' field
if item[0] in data_source.ds.known_filters:
if item[0] not in sph_ptypes:
raise_error = True
else:
ptype = item[0]
ppos = ["x", "y", "z"]
elif fi.alias_field:
if fi.alias_name[0] not in sph_ptypes:
raise_error = True
elif item[0] != "gas":
ptype = item[0]
else:
if fi.name[0] not in sph_ptypes and fi.name[0] != "gas":
raise_error = True
if raise_error:
raise RuntimeError(
"Can only perform off-axis projections for SPH fields, "
"Received '%s'" % (item, ))
normal = np.array(normal_vector)
normal = normal / np.linalg.norm(normal)
# If north_vector is None, we set the default here.
# This is chosen so that if normal_vector is one of the
# cartesian coordinate axes, the projection will match
# the corresponding on-axis projection.
if north_vector is None:
vecs = np.identity(3)
t = np.cross(vecs, normal).sum(axis=1)
ax = t.argmax()
east_vector = np.cross(vecs[ax, :], normal).ravel()
north = np.cross(normal, east_vector).ravel()
else:
north = np.array(north_vector)
north = north / np.linalg.norm(north)
east_vector = np.cross(north, normal).ravel()
# if weight is None:
buf = np.zeros((resolution[0], resolution[1]), dtype="float64")
x_min = center[0] - width[0] / 2
x_max = center[0] + width[0] / 2
y_min = center[1] - width[1] / 2
y_max = center[1] + width[1] / 2
z_min = center[2] - width[2] / 2
z_max = center[2] + width[2] / 2
finfo = data_source.ds.field_info[item]
ounits = finfo.output_units
bounds = [x_min, x_max, y_min, y_max, z_min, z_max]
if weight is None:
for chunk in data_source.chunks([], "io"):
off_axis_projection_SPH(
chunk[ptype, ppos[0]].to("code_length").d,
chunk[ptype, ppos[1]].to("code_length").d,
chunk[ptype, ppos[2]].to("code_length").d,
chunk[ptype, "mass"].to("code_mass").d,
chunk[ptype, "density"].to("code_density").d,
chunk[ptype, "smoothing_length"].to("code_length").d,
bounds,
center.to("code_length").d,
width.to("code_length").d,
chunk[item].in_units(ounits),
buf,
normal_vector,
north,
)
# Assure that the path length unit is in the default length units
# for the dataset by scaling the units of the smoothing length
path_length_unit = data_source.ds._get_field_info(
(ptype, "smoothing_length")).units
path_length_unit = Unit(path_length_unit,
registry=data_source.ds.unit_registry)
default_path_length_unit = data_source.ds.unit_system["length"]
buf *= data_source.ds.quan(
1, path_length_unit).in_units(default_path_length_unit)
item_unit = data_source.ds._get_field_info(item).units
item_unit = Unit(item_unit, registry=data_source.ds.unit_registry)
funits = item_unit * default_path_length_unit
else:
# if there is a weight field, take two projections:
# one of field*weight, the other of just weight, and divide them
weight_buff = np.zeros((resolution[0], resolution[1]),
dtype="float64")
wounits = data_source.ds.field_info[weight].output_units
for chunk in data_source.chunks([], "io"):
off_axis_projection_SPH(
chunk[ptype, ppos[0]].to("code_length").d,
chunk[ptype, ppos[1]].to("code_length").d,
chunk[ptype, ppos[2]].to("code_length").d,
chunk[ptype, "mass"].to("code_mass").d,
chunk[ptype, "density"].to("code_density").d,
chunk[ptype, "smoothing_length"].to("code_length").d,
bounds,
center.to("code_length").d,
width.to("code_length").d,
chunk[item].in_units(ounits),
buf,
normal_vector,
north,
weight_field=chunk[weight].in_units(wounits),
)
for chunk in data_source.chunks([], "io"):
off_axis_projection_SPH(
chunk[ptype, ppos[0]].to("code_length").d,
chunk[ptype, ppos[1]].to("code_length").d,
chunk[ptype, ppos[2]].to("code_length").d,
chunk[ptype, "mass"].to("code_mass").d,
chunk[ptype, "density"].to("code_density").d,
chunk[ptype, "smoothing_length"].to("code_length").d,
bounds,
center.to("code_length").d,
width.to("code_length").d,
chunk[weight].to(wounits),
weight_buff,
normal_vector,
north,
)
normalization_2d_utility(buf, weight_buff)
item_unit = data_source.ds._get_field_info(item).units
item_unit = Unit(item_unit, registry=data_source.ds.unit_registry)
funits = item_unit
myinfo = {
"field": item,
"east_vector": east_vector,
"north_vector": north_vector,
"normal_vector": normal_vector,
"width": width,
"units": funits,
"type": "SPH smoothed projection",
}
return ImageArray(buf,
funits,
registry=data_source.ds.unit_registry,
info=myinfo)
sc = Scene()
data_source.ds.index
if item is None:
field = data_source.ds.field_list[0]
mylog.info("Setting default field to %s", field.__repr__())
funits = data_source.ds._get_field_info(item).units
vol = KDTreeVolumeSource(data_source, item)
vol.num_threads = num_threads
if weight is None:
vol.set_field(item)
else:
# This is a temporary field, which we will remove at the end.
weightfield = ("index", "temp_weightfield")
def _make_wf(f, w):
def temp_weightfield(a, b):
tr = b[f].astype("float64") * b[w]
return tr.d
return temp_weightfield
data_source.ds.field_info.add_field(weightfield,
sampling_type="cell",
function=_make_wf(item, weight))
# Now we have to tell the dataset to add it and to calculate
# its dependencies..
deps, _ = data_source.ds.field_info.check_derived_fields([weightfield])
data_source.ds.field_dependencies.update(deps)
vol.set_field(weightfield)
vol.set_weight_field(weight)
ptf = ProjectionTransferFunction()
vol.set_transfer_function(ptf)
camera = sc.add_camera(data_source)
camera.set_width(width)
if not is_sequence(resolution):
resolution = [resolution] * 2
camera.resolution = resolution
if not is_sequence(width):
width = data_source.ds.arr([width] * 3)
normal = np.array(normal_vector)
normal = normal / np.linalg.norm(normal)
camera.position = center - width[2] * normal
camera.focus = center
# If north_vector is None, we set the default here.
# This is chosen so that if normal_vector is one of the
# cartesian coordinate axes, the projection will match
# the corresponding on-axis projection.
if north_vector is None:
vecs = np.identity(3)
t = np.cross(vecs, normal).sum(axis=1)
ax = t.argmax()
east_vector = np.cross(vecs[ax, :], normal).ravel()
north = np.cross(normal, east_vector).ravel()
else:
north = np.array(north_vector)
north = north / np.linalg.norm(north)
camera.switch_orientation(normal, north)
sc.add_source(vol)
vol.set_sampler(camera, interpolated=False)
assert vol.sampler is not None
fields = [vol.field]
if vol.weight_field is not None:
fields.append(vol.weight_field)
mylog.debug("Casting rays")
for (grid, mask) in data_source.blocks:
data = []
for f in fields:
# strip units before multiplying by mask for speed
grid_data = grid[f]
units = grid_data.units
data.append(
data_source.ds.arr(grid_data.d * mask, units, dtype="float64"))
pg = PartitionedGrid(
grid.id,
data,
mask.astype("uint8"),
grid.LeftEdge,
grid.RightEdge,
grid.ActiveDimensions.astype("int64"),
)
grid.clear_data()
vol.sampler(pg, num_threads=num_threads)
image = vol.finalize_image(camera, vol.sampler.aimage)
image = ImageArray(image,
funits,
registry=data_source.ds.unit_registry,
info=image.info)
if weight is not None:
data_source.ds.field_info.pop(("index", "temp_weightfield"))
if method == "integrate":
if weight is None:
dl = width[2].in_units(data_source.ds.unit_system["length"])
image *= dl
else:
mask = image[:, :, 1] == 0
image[:, :, 0] /= image[:, :, 1]
image[mask] = 0
return image[:, :, 0]
def _get_best_layout(self):
# Ensure the figure size along the long axis is always equal to _figure_size
unit_aspect = getattr(self, "_unit_aspect", 1)
if is_sequence(self._figure_size):
x_fig_size, y_fig_size = self._figure_size
y_fig_size *= unit_aspect
else:
x_fig_size = y_fig_size = self._figure_size
scaling = self._aspect / unit_aspect
if scaling < 1:
x_fig_size *= scaling
else:
y_fig_size /= scaling
if self._draw_colorbar:
cb_size = self._cb_size
cb_text_size = self._ax_text_size[1] + 0.45
else:
cb_size = x_fig_size * 0.04
cb_text_size = 0.0
if self._draw_axes:
x_axis_size = self._ax_text_size[0]
y_axis_size = self._ax_text_size[1]
else:
x_axis_size = x_fig_size * 0.04
y_axis_size = y_fig_size * 0.04
top_buff_size = self._top_buff_size
if not self._draw_axes and not self._draw_colorbar:
x_axis_size = 0.0
y_axis_size = 0.0
cb_size = 0.0
cb_text_size = 0.0
top_buff_size = 0.0
xbins = np.array([x_axis_size, x_fig_size, cb_size, cb_text_size])
ybins = np.array([y_axis_size, y_fig_size, top_buff_size])
size = [xbins.sum(), ybins.sum()]
x_frac_widths = xbins / size[0]
y_frac_widths = ybins / size[1]
# axrect is the rectangle defining the area of the
# axis object of the plot. Its range goes from 0 to 1 in
# x and y directions. The first two values are the x,y
# start values of the axis object (lower left corner), and the
# second two values are the size of the axis object. To get
# the upper right corner, add the first x,y to the second x,y.
axrect = (
x_frac_widths[0],
y_frac_widths[0],
x_frac_widths[1],
y_frac_widths[1],
)
# caxrect is the rectangle defining the area of the colorbar
# axis object of the plot. It is defined just as the axrect
# tuple is.
caxrect = (
x_frac_widths[0] + x_frac_widths[1],
y_frac_widths[0],
x_frac_widths[2],
y_frac_widths[1],
)
return size, axrect, caxrect
