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.frontends.stream.api import load_uniform_grid prng = RandomState(0x4d3d3d3) if not iterable(ndims): ndims = [ndims, ndims, ndims] else: assert (len(ndims) == 3) if not iterable(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": data['number_of_particles'] = v.size 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( (particles, 3)), unit) else: data['io', field] = (prng.random_sample(size=particles), unit) else: for f in ('particle_position_%s' % ax for ax in 'xyz'): data['io', f] = (prng.random_sample(size=particles), 'code_length') for f in ('particle_velocity_%s' % 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') data['number_of_particles'] = particles ug = load_uniform_grid(data, ndims, length_unit=length_unit, nprocs=nprocs, unit_system=unit_system, bbox=bbox) return ug
def _sanitize_camera_property_units(value, scene): if iterable(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], string_types)): return scene.arr([scene.arr(value[0], value[1]).in_units('unitary')]*3) if len(value) == 3: if all([iterable(v) for v in value]): if all([isinstance(v[0], numeric_type) and isinstance(v[1], string_types) for v in value]): return scene.arr( [scene.arr(v[0], v[1]) for v in value]) else: raise RuntimeError( "Cannot set camera width to invalid value '%s'" % (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( "Cannot set camera width to invalid value '%s'" % (value, ))
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 iterable(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 iterable(width): width = validate_iterable_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_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('center keyword "%s" not recognized' % center) elif isinstance(center, YTArray): return self.ds.arr(center), self.convert_to_cartesian(center) elif iterable(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('center keyword "%s" not recognized' % center) center = self.ds.arr(center, "code_length") elif iterable(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('center keyword "%s" not recognized' % center) # This has to return both a center and a display_center display_center = self.convert_to_cartesian(center) return center, display_center
def sanitize_center(self, center, axis): if isinstance(center, string_types): 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": center = (self.ds.domain_left_edge + self.ds.domain_right_edge) / 2 else: raise RuntimeError('center keyword \"%s\" not recognized' % center) elif isinstance(center, YTArray): return self.ds.arr(center), self.convert_to_cartesian(center) elif iterable(center): if isinstance(center[0], string_types) and isinstance(center[1], string_types): 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("center keyword \"%s\" not recognized" % center) center = self.ds.arr(center, 'code_length') elif iterable(center[0]) and isinstance(center[1], string_types): center = self.ds.arr(center[0], center[1]) else: center = self.ds.arr(center, 'code_length') else: raise RuntimeError("center keyword \"%s\" not recognized" % center) # 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 iterable(width): validate_width_tuple(width) width = self.ds.quan(width[0], width[1]) if height is None: height = width elif iterable(height): validate_width_tuple(height) height = self.ds.quan(height[0], height[1]) if not iterable(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 construct_image(ds, axis, data_source, center, width=None, image_res=None): 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 image_res is None: ddims = ds.domain_dimensions*ds.refine_by**ds.index.max_level if iterable(axis): nx = ddims.max() ny = ddims.max() else: nx, ny = [ddims[idx] for idx in axis_wcs[axis]] else: if iterable(image_res): nx, ny = image_res else: nx, ny = image_res, image_res dx, dy = width[0]/nx, width[1]/ny crpix = [0.5*(nx+1), 0.5*(ny+1)] if hasattr(ds, "wcs") and not iterable(axis): # This is a FITS dataset, so we use it to construct the WCS cunit = [str(ds.wcs.wcs.cunit[idx]) for idx in axis_wcs[axis]] ctype = [ds.wcs.wcs.ctype[idx] for idx in axis_wcs[axis]] cdelt = [ds.wcs.wcs.cdelt[idx] for idx in axis_wcs[axis]] ctr_pix = center.in_units("code_length")[:ds.dimensionality].v crval = ds.wcs.wcs_pix2world(ctr_pix.reshape(1, ds.dimensionality))[0] crval = [crval[idx] for idx in axis_wcs[axis]] else: 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) cunit = [unit]*2 ctype = ["LINEAR"]*2 cdelt = [dx.in_units(unit)]*2 if iterable(axis): crval = center.in_units(unit) else: crval = [center[idx].in_units(unit) for idx in axis_wcs[axis]] if hasattr(data_source, 'to_frb'): if iterable(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 = 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
def construct_image(ds, axis, data_source, center, width=None, image_res=None): 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 image_res is None: ddims = ds.domain_dimensions*ds.refine_by**ds.index.max_level if iterable(axis): nx = ddims.max() ny = ddims.max() else: nx, ny = [ddims[idx] for idx in axis_wcs[axis]] else: if iterable(image_res): nx, ny = image_res else: nx, ny = image_res, image_res dx = width[0]/nx crpix = [0.5*(nx+1), 0.5*(ny+1)] if hasattr(ds, "wcs") and not iterable(axis): # This is a FITS dataset, so we use it to construct the WCS cunit = [str(ds.wcs.wcs.cunit[idx]) for idx in axis_wcs[axis]] ctype = [ds.wcs.wcs.ctype[idx] for idx in axis_wcs[axis]] cdelt = [ds.wcs.wcs.cdelt[idx] for idx in axis_wcs[axis]] ctr_pix = center.in_units("code_length")[:ds.dimensionality].v crval = ds.wcs.wcs_pix2world(ctr_pix.reshape(1, ds.dimensionality))[0] crval = [crval[idx] for idx in axis_wcs[axis]] else: 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) cunit = [unit]*2 ctype = ["LINEAR"]*2 cdelt = [dx.in_units(unit)]*2 if iterable(axis): crval = center.in_units(unit) else: crval = [center[idx].in_units(unit) for idx in axis_wcs[axis]] if hasattr(data_source, 'to_frb'): if iterable(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
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 iterable(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 iterable(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 iterable(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 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.frontends.stream.api import load_particles prng = RandomState(0x4d3d3d3) if not iterable(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): 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 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 iterable(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 _yt_array_hdf5_attr(fh, attr, val): r"""Save a YTArray or YTQuantity as an hdf5 attribute. Save an hdf5 attribute. If it has units, save an additional attribute with the units. Parameters ---------- fh : an open hdf5 file, group, or dataset The hdf5 file, group, or dataset to which the attribute will be written. attr : str The name of the attribute to be saved. val : anything The value to be saved. """ if val is None: val = "None" if hasattr(val, "units"): fh.attrs["%s_units" % attr] = str(val.units) # The following is a crappy workaround for getting # Unicode strings into HDF5 attributes in Python 3 if iterable(val): val = np.array(val) if val.dtype.kind == 'U': val = val.astype('|S') try: fh.attrs[str(attr)] = val # This is raised if no HDF5 equivalent exists. # In that case, save its string representation. except TypeError: fh.attrs[str(attr)] = str(val)
def fset(self, value): if iterable(value): if len(value) != 2: raise RuntimeError else: value = (value, value) self._resolution = value
def __init__(self, fsize, axrect, figure, axes): """Initialize PlotMPL class""" import matplotlib.figure self._plot_valid = True if figure is None: if not iterable(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 _reconstruct_clump(parent, field, mi, ma, valid, children, data, clump_info, function=None): obj = object.__new__(Clump) if iterable(parent): try: parent = parent[1] except KeyError: parent = parent if children is None: children = [] obj.parent, obj.field, obj.min_val, obj.max_val, \ obj.valid, obj.children, obj.clump_info, obj.function = \ parent, field, mi, ma, valid, children, clump_info, function # Now we override, because the parent/child relationship seems a bit # unreliable in the unpickling for child in children: child.parent = obj obj.data = data[1] # Strip out the PF obj.quantities = obj.data.quantities if obj.parent is None: return (data[0], obj) return obj
def __init__(self, x_data, y_data, data, x_scale, y_scale, z_scale, cmap, zlim, figure_size, fontsize, figure, axes, cax): self._initfinished = False self._draw_colorbar = True self._draw_axes = True self._figure_size = figure_size # Compute layout fontscale = float(fontsize) / 18.0 if fontscale < 1.0: fontscale = np.sqrt(fontscale) if iterable(figure_size): self._cb_size = 0.0375 * figure_size[0] else: self._cb_size = 0.0375 * figure_size self._ax_text_size = [1.1 * fontscale, 0.9 * fontscale] self._top_buff_size = 0.30 * fontscale self._aspect = 1.0 size, axrect, caxrect = self._get_best_layout() super(PhasePlotMPL, self).__init__(size, axrect, caxrect, zlim, figure, axes, cax) self._init_image(x_data, y_data, data, x_scale, y_scale, z_scale, zlim, cmap) self._initfinished = True
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 = ensure_list(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 iterable(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 _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 iterable(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 _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 iterable(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 __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 iterable(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 compare(self, old_result): if hasattr(self.result, 'tostring'): self.compare_array_delta(self.result, old_result, 1e-7) return elif iterable(self.result): a1 = na.array(self.result) a2 = na.array(old_result) self.compare_array_delta(a1, a2, 1e-7) else: if self.result != old_result: raise FieldHashesDontMatch
def _get_best_layout(self): if self._draw_colorbar: cb_size = self._cb_size cb_text_size = self._ax_text_size[1] + 0.45 else: cb_size = 0.0 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 = 0.0 y_axis_size = 0.0 if self._draw_axes or self._draw_colorbar: top_buff_size = self._top_buff_size else: top_buff_size = 0.0 # Ensure the figure size along the long axis is always equal to _figure_size if iterable(self._figure_size): x_fig_size = self._figure_size[0] y_fig_size = self._figure_size[1] else: if self._aspect >= 1.0: x_fig_size = self._figure_size y_fig_size = self._figure_size / self._aspect if self._aspect < 1.0: x_fig_size = self._figure_size * self._aspect y_fig_size = self._figure_size 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 = ( x_frac_widths[0], y_frac_widths[0], x_frac_widths[1], y_frac_widths[1], ) 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
def parse_value(value, default_units, ds=None): if ds is None: quan = YTQuantity else: quan = ds.quan if isinstance(value, YTQuantity): return quan(value).in_units(default_units) elif iterable(value): return quan(value[0], value[1]).in_units(default_units) else: return quan(value, default_units)
def sanitize_depth(self, depth): if iterable(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 load_object(self, name): """ Load and return and object from the data_file using the Pickle protocol, under the name *name* on the node /Objects. """ obj = self.get_data("/Objects", name) if obj is None: return obj = cPickle.loads(obj.value) if iterable(obj) and len(obj) == 2: obj = obj[1] # Just the object, not the ds if hasattr(obj, '_fix_pickle'): obj._fix_pickle() return obj
def from_sizes(cls, sizes): sizes = ensure_list(sizes) pool = cls() rank = pool.comm.rank for i,size in enumerate(sizes): if iterable(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 from_sizes(cls, sizes): sizes = ensure_list(sizes) pool = cls() rank = pool.comm.rank for i, size in enumerate(sizes): if iterable(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 iterable(width) and len(width) > 1 and isinstance( width[1], string_types): width = data_source.ds.quan(width[0], input_units=width[1]) # Now convert back to code length for subsequent manipulation width = width.in_units("code_length") # .value if not iterable(width): width = data_source.ds.arr([width, width, width], input_units='code_length') # left/right, top/bottom, front/back if not isinstance(width, YTArray): width = data_source.ds.arr(width, input_units="code_length") if not isinstance(focus, YTArray): focus = data_source.ds.arr(focus, input_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(Camera, self).__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 iterable(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 sphere_selector(hc, ds2, radius_field, factor=1, min_radius=None): r""" Select halos within a sphere around the target halo. Parameters ---------- hc : halo container object Halo container associated with the target halo. ds2 : halo catalog-type dataset The dataset of the ancestor halos. radius_field : str Name of the field to be used to get the halo radius. factor : float, optional Multiplicative factor of the halo radius in which potential halos will be gathered. Default: 1. min_radius : YTQuantity or tuple of (value, unit) An absolute minimum radius for the sphere. Returns ------- my_ids : list of ints List of ids of potential halos. """ if min_radius is not None: if isinstance(min_radius, YTQuantity): pass elif iterable(min_radius) and len(min_radius) == 2: min_radius = ds2.quan(min_radius[0], min_radius[1]) else: raise RuntimeError( "min_radius should be YTQuantity or (value, unit) tuple.") # Never mix code units from multiple datasets!!! center = ds2.arr(hc.position.to("code_length").d, "code_length") radius = factor * hc[radius_field] radius = ds2.quan(radius.to("code_length").d[0], "code_length") if min_radius is not None: radius = max(radius, min_radius) try: sp = ds2.sphere(center, radius) my_ids = sp[(hc.ptype, "particle_identifier")] my_ids = my_ids[np.argsort(sp[(hc.ptype, "particle_radius")])] return my_ids.d.astype(np.int64) except YTSphereTooSmall: return []
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 iterable(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 __init__(self, data_source, figure_size, fontsize): self.data_source = data_source if iterable(figure_size): self.figure_size = float(figure_size[0]), float(figure_size[1]) else: self.figure_size = float(figure_size) self.plots = PlotDictionary(data_source) self._callbacks = [] self._field_transform = {} self._colormaps = defaultdict(lambda: 'algae') 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._colorbar_label = PlotDictionary( self.data_source, lambda: None)
def __init__(self, ds, normal, fields, center='c', width=(1.0, 'unitary'), weight_field=None, image_res=512, depth_res=256, north_vector=None, depth=(1.0,"unitary"), no_ghost=False, method='integrate'): fields = ensure_list(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 iterable(image_res): image_res = (image_res, image_res) res = (image_res[0], image_res[1], depth_res) for field in fields: buf[field] = off_axis_projection(ds, center, normal, wd, res, field, no_ghost=no_ghost, north_vector=north_vector, method=method, weight=weight_field).swapaxes(0, 1) center = ds.arr([0.0] * 2, 'code_length') w, not_an_frb = construct_image(ds, normal, buf, center, width=width, image_res=image_res) super(FITSOffAxisProjection, self).__init__(buf, fields=fields, wcs=w)
def __init__(self, data_source, figure_size, fontsize): self.data_source = data_source if iterable(figure_size): self.figure_size = float(figure_size[0]), float(figure_size[1]) else: self.figure_size = float(figure_size) self.plots = PlotDictionary(data_source) self._callbacks = [] self._field_transform = {} self._colormaps = defaultdict(lambda: 'algae') 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._cbar_minorticks = {} self._colorbar_label = PlotDictionary( self.data_source, lambda: None)
def _deserialize_from_h5(g, ds): result = {} for item in g: if item == "chunks": continue if "units" in g[item].attrs: if iterable(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], h5.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, data, fields=None, units=None, width=None, wcs=None): r""" Initialize a FITSImageData object. FITSImageData contains a collection of FITS ImageHDU instances and WCS information, along with units for each of the images. FITSImageData instances can be constructed from ImageArrays, NumPy arrays, dicts of such arrays, FixedResolutionBuffers, and YTCoveringGrids. The latter two are the most powerful because WCS information can be constructed automatically from their coordinates. Parameters ---------- data : FixedResolutionBuffer or a YTCoveringGrid. Or, an ImageArray, an numpy.ndarray, or dict of such arrays The data to be made into a FITS image or images. fields : single string or list of strings, optional The field names for the data. If *fields* is none and *data* has keys, it will use these for the fields. If *data* is just a single array one field name must be specified. units : string The units of the WCS coordinates. Defaults to "cm". width : float or YTQuantity The width of the image. Either a single value or iterable of values. If a float, assumed to be in *units*. Only used if this information is not already provided by *data*. wcs : `astropy.wcs.WCS` instance, optional Supply an AstroPy WCS instance. Will override automatic WCS creation from FixedResolutionBuffers and YTCoveringGrids. Examples -------- >>> # This example uses a FRB. >>> ds = load("sloshing_nomag2_hdf5_plt_cnt_0150") >>> prj = ds.proj(2, "kT", weight_field="density") >>> frb = prj.to_frb((0.5, "Mpc"), 800) >>> # This example just uses the FRB and puts the coords in kpc. >>> f_kpc = FITSImageData(frb, fields="kT", units="kpc") >>> # This example specifies a specific WCS. >>> from astropy.wcs import WCS >>> w = WCS(naxis=self.dimensionality) >>> w.wcs.crval = [30., 45.] # RA, Dec in degrees >>> w.wcs.cunit = ["deg"]*2 >>> nx, ny = 800, 800 >>> w.wcs.crpix = [0.5*(nx+1), 0.5*(ny+1)] >>> w.wcs.ctype = ["RA---TAN","DEC--TAN"] >>> scale = 1./3600. # One arcsec per pixel >>> w.wcs.cdelt = [-scale, scale] >>> f_deg = FITSImageData(frb, fields="kT", wcs=w) >>> f_deg.writeto("temp.fits") """ if units is None: units = "cm" if width is None: width = 1.0 exclude_fields = ['x','y','z','px','py','pz', 'pdx','pdy','pdz','weight_field'] super(FITSImageData, self).__init__() if isinstance(fields, string_types): fields = [fields] if hasattr(data, 'keys'): img_data = data if fields is None: fields = list(img_data.keys()) elif isinstance(data, np.ndarray): if fields is None: mylog.warning("No field name given for this array. Calling it 'image_data'.") fn = 'image_data' fields = [fn] else: fn = fields[0] img_data = {fn: data} self.fields = [] for fd in fields: if isinstance(fd, tuple): self.fields.append(fd[1]) else: self.fields.append(fd) first = True self.field_units = {} for key in fields: if key not in exclude_fields: if hasattr(img_data[key], "units"): self.field_units[key] = img_data[key].units else: self.field_units[key] = "dimensionless" mylog.info("Making a FITS image of field %s" % key) if first: hdu = pyfits.PrimaryHDU(np.array(img_data[key])) first = False else: hdu = pyfits.ImageHDU(np.array(img_data[key])) hdu.name = key hdu.header["btype"] = key if hasattr(img_data[key], "units"): hdu.header["bunit"] = re.sub('()', '', str(img_data[key].units)) self.append(hdu) self.shape = self[0].shape self.dimensionality = len(self.shape) if wcs is None: w = pywcs.WCS(header=self[0].header, naxis=self.dimensionality) if isinstance(img_data, FixedResolutionBuffer): # FRBs are a special case where we have coordinate # information, so we take advantage of this and # construct the WCS object dx = (img_data.bounds[1]-img_data.bounds[0]).in_units(units).v/self.shape[0] dy = (img_data.bounds[3]-img_data.bounds[2]).in_units(units).v/self.shape[1] xctr = 0.5*(img_data.bounds[1]+img_data.bounds[0]).in_units(units).v yctr = 0.5*(img_data.bounds[3]+img_data.bounds[2]).in_units(units).v center = [xctr, yctr] cdelt = [dx,dy] elif isinstance(img_data, YTCoveringGridBase): cdelt = img_data.dds.in_units(units).v center = 0.5*(img_data.left_edge+img_data.right_edge).in_units(units).v else: # If img_data is just an array, we assume the center is the origin # and use the image width to determine the cell widths if not iterable(width): width = [width]*self.dimensionality if isinstance(width[0], YTQuantity): cdelt = [wh.in_units(units).v/n for wh, n in zip(width, self.shape)] else: cdelt = [float(wh)/n for wh, n in zip(width, self.shape)] center = [0.0]*self.dimensionality w.wcs.crpix = 0.5*(np.array(self.shape)+1) w.wcs.crval = center w.wcs.cdelt = cdelt w.wcs.ctype = ["linear"]*self.dimensionality w.wcs.cunit = [units]*self.dimensionality self.set_wcs(w) else: self.set_wcs(wcs)
def from_data_source(cls, data_source, redshift, area, exp_time, source_model, point_sources=False, parameters=None, center=None, dist=None, cosmology=None, velocity_fields=None): r""" Initialize a :class:`~pyxsim.photon_list.PhotonList` from a yt data source. The redshift, collecting area, exposure time, and cosmology are stored in the *parameters* dictionary which is passed to the *source_model* function. Parameters ---------- data_source : :class:`~yt.data_objects.data_containers.YTSelectionContainer` The data source from which the photons will be generated. redshift : float The cosmological redshift for the photons. area : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity` The collecting area to determine the number of photons. If units are not specified, it is assumed to be in cm^2. exp_time : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity` The exposure time to determine the number of photons. If units are not specified, it is assumed to be in seconds. source_model : :class:`~pyxsim.source_models.SourceModel` A source model used to generate the photons. point_sources : boolean, optional If True, the photons will be assumed to be generated from the exact positions of the cells or particles and not smeared around within a volume. Default: False parameters : dict, optional A dictionary of parameters to be passed for the source model to use, if necessary. center : string or array_like, optional The origin of the photon spatial coordinates. Accepts "c", "max", or a coordinate. If not specified, pyxsim attempts to use the "center" field parameter of the data_source. dist : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity` The angular diameter distance, used for nearby sources. This may be optionally supplied instead of it being determined from the *redshift* and given *cosmology*. If units are not specified, it is assumed to be in kpc. To use this, the redshift must be set to zero. cosmology : :class:`~yt.utilities.cosmology.Cosmology`, optional Cosmological information. If not supplied, we try to get the cosmology from the dataset. Otherwise, LCDM with the default yt parameters is assumed. velocity_fields : list of fields The yt fields to use for the velocity. If not specified, the following will be assumed: ['velocity_x', 'velocity_y', 'velocity_z'] for grid datasets ['particle_velocity_x', 'particle_velocity_y', 'particle_velocity_z'] for particle datasets Examples -------- >>> thermal_model = ThermalSourceModel(apec_model, Zmet=0.3) >>> redshift = 0.05 >>> area = 6000.0 # assumed here in cm**2 >>> time = 2.0e5 # assumed here in seconds >>> sp = ds.sphere("c", (500., "kpc")) >>> my_photons = PhotonList.from_data_source(sp, redshift, area, ... time, thermal_model) """ ds = data_source.ds if parameters is None: parameters = {} if cosmology is None: if hasattr(ds, 'cosmology'): cosmo = ds.cosmology else: cosmo = Cosmology() else: cosmo = cosmology if dist is None: if redshift <= 0.0: msg = "If redshift <= 0.0, you must specify a distance to the " \ "source using the 'dist' argument!" mylog.error(msg) raise ValueError(msg) D_A = cosmo.angular_diameter_distance(0.0, redshift).in_units("Mpc") else: D_A = parse_value(dist, "kpc") if redshift > 0.0: mylog.warning("Redshift must be zero for nearby sources. " "Resetting redshift to 0.0.") redshift = 0.0 if isinstance(center, string_types): if center == "center" or center == "c": parameters["center"] = ds.domain_center elif center == "max" or center == "m": parameters["center"] = ds.find_max("density")[-1] elif iterable(center): if isinstance(center, YTArray): parameters["center"] = center.in_units("code_length") elif isinstance(center, tuple): if center[0] == "min": parameters["center"] = ds.find_min(center[1])[-1] elif center[0] == "max": parameters["center"] = ds.find_max(center[1])[-1] else: raise RuntimeError else: parameters["center"] = ds.arr(center, "code_length") elif center is None: if hasattr(data_source, "left_edge"): parameters["center"] = 0.5*(data_source.left_edge+data_source.right_edge) else: parameters["center"] = data_source.get_field_parameter("center") parameters["fid_exp_time"] = parse_value(exp_time, "s") parameters["fid_area"] = parse_value(area, "cm**2") parameters["fid_redshift"] = redshift parameters["fid_d_a"] = D_A parameters["hubble"] = cosmo.hubble_constant parameters["omega_matter"] = cosmo.omega_matter parameters["omega_lambda"] = cosmo.omega_lambda if redshift > 0.0: mylog.info("Cosmology: h = %g, omega_matter = %g, omega_lambda = %g" % (cosmo.hubble_constant, cosmo.omega_matter, cosmo.omega_lambda)) else: mylog.info("Observing local source at distance %s." % D_A) D_A = parameters["fid_d_a"].in_cgs() dist_fac = 1.0/(4.*np.pi*D_A.value*D_A.value*(1.+redshift)**2) spectral_norm = parameters["fid_area"].v*parameters["fid_exp_time"].v*dist_fac source_model.setup_model(data_source, redshift, spectral_norm) p_fields, v_fields, w_field = determine_fields(ds, source_model.source_type, point_sources) if velocity_fields is not None: v_fields = velocity_fields if p_fields[0] == ("index", "x"): parameters["data_type"] = "cells" else: parameters["data_type"] = "particles" citer = data_source.chunks([], "io") photons = defaultdict(list) for chunk in parallel_objects(citer): chunk_data = source_model(chunk) if chunk_data is not None: ncells, number_of_photons, idxs, energies = chunk_data photons["num_photons"].append(number_of_photons) photons["energy"].append(energies) photons["pos"].append(np.array([chunk[p_fields[0]].d[idxs], chunk[p_fields[1]].d[idxs], chunk[p_fields[2]].d[idxs]])) photons["vel"].append(np.array([chunk[v_fields[0]].d[idxs], chunk[v_fields[1]].d[idxs], chunk[v_fields[2]].d[idxs]])) if w_field is None: photons["dx"].append(np.zeros(ncells)) else: photons["dx"].append(chunk[w_field].d[idxs]) source_model.cleanup_model() photon_units = {"pos": ds.field_info[p_fields[0]].units, "vel": ds.field_info[v_fields[0]].units, "energy": "keV"} if w_field is None: photon_units["dx"] = "kpc" else: photon_units["dx"] = ds.field_info[w_field].units concatenate_photons(ds, photons, photon_units) c = parameters["center"].to("kpc") if sum(ds.periodicity) > 0: # Fix photon coordinates for regions crossing a periodic boundary dw = ds.domain_width.to("kpc") le, re = find_object_bounds(data_source) for i in range(3): if ds.periodicity[i] and photons["pos"].shape[0] > 0: tfl = photons["pos"][:,i] < le[i] tfr = photons["pos"][:,i] > re[i] photons["pos"][tfl,i] += dw[i] photons["pos"][tfr,i] -= dw[i] # Re-center all coordinates if photons["pos"].shape[0] > 0: photons["pos"] -= c mylog.info("Finished generating photons.") mylog.info("Number of photons generated: %d" % int(np.sum(photons["num_photons"]))) mylog.info("Number of cells with photons: %d" % photons["dx"].size) return cls(photons, parameters, cosmo)
def off_axis(self, L, center="c", width=(1, "unitary"), nx=800, source=None): r""" Make an off-axis projection of the SZ signal. Parameters ---------- L : array_like The normal vector of the projection. center : A sequence of floats, a string, or a tuple. The coordinate of the center of the image. If set to 'c', 'center' or left blank, the plot is centered on the middle of the domain. If set to 'max' or 'm', the center will be located at the maximum of the ('gas', 'density') field. Centering on the max or min of a specific field is supported by providing a tuple such as ("min","temperature") or ("max","dark_matter_density"). Units can be specified by passing in *center* as a tuple containing a coordinate and string unit name or by passing in a YTArray. If a list or unitless array is supplied, code units are assumed. width : float, tuple, or YTQuantity. The width of the projection. A float will assume the width is in code units. A (value, unit) tuple or YTQuantity allows for the units of the width to be specified. nx : integer, optional The dimensions on a side of the projection image. source : yt.data_objects.data_containers.YTSelectionContainer, optional If specified, this will be the data source used for selecting regions to project. Currently unsupported in yt 2.x. Examples -------- >>> L = np.array([0.5, 1.0, 0.75]) >>> szprj.off_axis(L, center="c", width=(2.0, "Mpc")) """ if iterable(width): w = self.ds.quan(width[0], width[1]).in_units("code_length").value elif isinstance(width, YTQuantity): w = width.in_units("code_length").value else: w = width ctr, dctr = self.ds.coordinates.sanitize_center(center, L) if source is not None: mylog.error("Source argument is not currently supported for off-axis S-Z projections.") raise NotImplementedError beta_par = generate_beta_par(L) self.ds.add_field(("gas","beta_par"), function=beta_par, units="g/cm**3") setup_sunyaev_zeldovich_fields(self.ds) dens = off_axis_projection(self.ds, ctr, L, w, nx, "density") Te = off_axis_projection(self.ds, ctr, L, w, nx, "t_sz")/dens bpar = off_axis_projection(self.ds, ctr, L, w, nx, "beta_par")/dens omega1 = off_axis_projection(self.ds, ctr, L, w, nx, "t_squared")/dens omega1 = omega1/(Te*Te) - 1. if self.high_order: bperp2 = off_axis_projection(self.ds, ctr, L, w, nx, "beta_perp_squared")/dens sigma1 = off_axis_projection(self.ds, ctr, L, w, nx, "t_beta_par")/dens sigma1 = sigma1/Te - bpar kappa1 = off_axis_projection(self.ds, ctr, L, w, nx, "beta_par_squared")/dens kappa1 -= bpar else: bperp2 = np.zeros((nx,nx)) sigma1 = np.zeros((nx,nx)) kappa1 = np.zeros((nx,nx)) tau = sigma_thompson*dens*self.mueinv/mh self.bounds = np.array([-0.5*w, 0.5*w, -0.5*w, 0.5*w]) self.dx = w/nx self.dy = w/nx self.nx = nx self._compute_intensity(np.array(tau), np.array(Te), np.array(bpar), np.array(omega1), np.array(sigma1), np.array(kappa1), np.array(bperp2)) self.ds.field_info.pop(("gas","beta_par"))
def __init__(self, ds, normal, field, velocity_bounds, center="c", width=(1.0,"unitary"), dims=100, thermal_broad=False, atomic_weight=56., depth=(1.0,"unitary"), depth_res=256, method="integrate", weight_field=None, no_shifting=False, north_vector=None, no_ghost=True): r""" Initialize a PPVCube object. Parameters ---------- ds : dataset The dataset. normal : array_like or string The normal vector along with to make the projections. If an array, it will be normalized. If a string, it will be assumed to be along one of the principal axes of the domain ("x", "y", or "z"). field : string The field to project. velocity_bounds : tuple A 4-tuple of (vmin, vmax, nbins, units) for the velocity bounds to integrate over. center : A sequence of floats, a string, or a tuple. The coordinate of the center of the image. If set to 'c', 'center' or left blank, the plot is centered on the middle of the domain. If set to 'max' or 'm', the center will be located at the maximum of the ('gas', 'density') field. Centering on the max or min of a specific field is supported by providing a tuple such as ("min","temperature") or ("max","dark_matter_density"). Units can be specified by passing in *center* as a tuple containing a coordinate and string unit name or by passing in a YTArray. If a list or unitless array is supplied, code units are assumed. width : float, tuple, or YTQuantity. The width of the projection. A float will assume the width is in code units. A (value, unit) tuple or YTQuantity allows for the units of the width to be specified. Implies width = height, e.g. the aspect ratio of the PPVCube's spatial dimensions is 1. dims : integer, optional The spatial resolution of the cube. Implies nx = ny, e.g. the aspect ratio of the PPVCube's spatial dimensions is 1. thermal_broad : boolean, optional Whether or not to broaden the line using the gas temperature. Default: False. atomic_weight : float, optional Set this value to the atomic weight of the particle that is emitting the line if *thermal_broad* is True. Defaults to 56 (Fe). depth : A tuple or a float, optional A tuple containing the depth to project through and the string key of the unit: (width, 'unit'). If set to a float, code units are assumed. Only for off-axis cubes. depth_res : integer, optional The resolution of integration along the line of sight for off-axis cubes. Default: 256 method : string, optional Set the projection method to be used. "integrate" : line of sight integration over the line element. "sum" : straight summation over the line of sight. weight_field : string, optional The name of the weighting field. Set to None for no weight. no_shifting : boolean, optional If set, no shifting due to velocity will occur but only thermal broadening. Should not be set when *thermal_broad* is False, otherwise nothing happens! north_vector : a sequence of floats A vector defining the 'up' direction. This option sets the orientation of the plane of projection. If not set, an arbitrary grid-aligned north_vector is chosen. Ignored in the case of on-axis cubes. no_ghost: bool, optional Optimization option for off-axis cases. 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 Examples -------- >>> i = 60*np.pi/180. >>> L = [0.0,np.sin(i),np.cos(i)] >>> cube = PPVCube(ds, L, "density", (-5.,4.,100,"km/s"), width=(10.,"kpc")) """ self.ds = ds self.field = field self.width = width self.particle_mass = atomic_weight*mh self.thermal_broad = thermal_broad self.no_shifting = no_shifting if not isinstance(normal, string_types): width = ds.coordinates.sanitize_width(normal, width, depth) width = tuple(el.in_units('code_length').v for el in width) if no_shifting and not thermal_broad: raise RuntimeError("no_shifting cannot be True when thermal_broad is False!") self.center = ds.coordinates.sanitize_center(center, normal)[0] self.nx = dims self.ny = dims self.nv = velocity_bounds[2] if method not in ["integrate","sum"]: raise RuntimeError("Only the 'integrate' and 'sum' projection +" "methods are supported in PPVCube.") dd = ds.all_data() fd = dd._determine_fields(field)[0] self.field_units = ds._get_field_info(fd).units self.vbins = ds.arr(np.linspace(velocity_bounds[0], velocity_bounds[1], velocity_bounds[2]+1), velocity_bounds[3]) self._vbins = self.vbins.copy() self.vmid = 0.5*(self.vbins[1:]+self.vbins[:-1]) self.vmid_cgs = self.vmid.in_cgs().v self.dv = self.vbins[1]-self.vbins[0] self.dv_cgs = self.dv.in_cgs().v self.current_v = 0.0 _vlos = create_vlos(normal, self.no_shifting) self.ds.add_field(("gas","v_los"), function=_vlos, units="cm/s") _intensity = self._create_intensity() self.ds.add_field(("gas","intensity"), function=_intensity, units=self.field_units) if method == "integrate" and weight_field is None: self.proj_units = str(ds.quan(1.0, self.field_units+"*cm").units) elif method == "sum": self.proj_units = self.field_units storage = {} pbar = get_pbar("Generating cube.", self.nv) for sto, i in parallel_objects(range(self.nv), storage=storage): self.current_v = self.vmid_cgs[i] if isinstance(normal, string_types): prj = ds.proj("intensity", ds.coordinates.axis_id[normal], method=method, weight_field=weight_field) buf = prj.to_frb(width, self.nx, center=self.center)["intensity"] else: buf = off_axis_projection(ds, self.center, normal, width, (self.nx, self.ny, depth_res), "intensity", north_vector=north_vector, no_ghost=no_ghost, method=method, weight=weight_field).swapaxes(0,1) sto.result_id = i sto.result = buf pbar.update(i) pbar.finish() self.data = ds.arr(np.zeros((self.nx,self.ny,self.nv)), self.proj_units) if is_root(): for i, buf in sorted(storage.items()): self.data[:,:,i] = buf.transpose() self.axis_type = "velocity" # Now fix the width if iterable(self.width): self.width = ds.quan(self.width[0], self.width[1]) elif not isinstance(self.width, YTQuantity): self.width = ds.quan(self.width, "code_length") self.ds.field_info.pop(("gas","intensity")) self.ds.field_info.pop(("gas","v_los"))
def _get_best_layout(self): # Ensure the figure size along the long axis is always equal to _figure_size if iterable(self._figure_size): x_fig_size = self._figure_size[0] y_fig_size = self._figure_size[1] else: x_fig_size = self._figure_size y_fig_size = self._figure_size/self._aspect if hasattr(self, '_unit_aspect'): y_fig_size = y_fig_size * self._unit_aspect 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