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