class ParticleTrajectories(object):
r"""A collection of particle trajectories in time over a series of
datasets.
The ParticleTrajectories object contains a collection of
particle trajectories for a specified set of particle indices.
Parameters
----------
outputs : `yt.data_objects.time_series.DatasetSeries` or list of strings
DatasetSeries object, or a time-sorted list of filenames to
construct a new DatasetSeries object.
indices : array_like
An integer array of particle indices whose trajectories we
want to track. If they are not sorted they will be sorted.
fields : list of strings, optional
A set of fields that is retrieved when the trajectory
collection is instantiated.
Default : None (will default to the fields 'particle_position_x',
'particle_position_y', 'particle_position_z')
suppress_logging : boolean
Suppress yt's logging when iterating over the simulation time
series.
Default : False
Examples
________
>>> from yt.mods import *
>>> my_fns = glob.glob("orbit_hdf5_chk_00[0-9][0-9]")
>>> my_fns.sort()
>>> fields = ["particle_position_x", "particle_position_y",
>>> "particle_position_z", "particle_velocity_x",
>>> "particle_velocity_y", "particle_velocity_z"]
>>> ds = load(my_fns[0])
>>> init_sphere = ds.sphere(ds.domain_center, (.5, "unitary"))
>>> indices = init_sphere["particle_index"].astype("int")
>>> trajs = ParticleTrajectories(my_fns, indices, fields=fields)
>>> for t in trajs :
>>> print t["particle_velocity_x"].max(), t["particle_velocity_x"].min()
"""
def __init__(self, outputs, indices, fields=None, suppress_logging=False):
indices.sort() # Just in case the caller wasn't careful
self.field_data = YTFieldData()
if isinstance(outputs, DatasetSeries):
self.data_series = outputs
else:
self.data_series = DatasetSeries(outputs)
self.masks = []
self.sorts = []
self.array_indices = []
self.indices = indices
self.num_indices = len(indices)
self.num_steps = len(outputs)
self.times = []
self.suppress_logging = suppress_logging
# Default fields
if fields is None: fields = []
fields.append("particle_position_x")
fields.append("particle_position_y")
fields.append("particle_position_z")
fields = list(OrderedDict.fromkeys(fields))
if self.suppress_logging:
old_level = int(ytcfg.get("yt", "loglevel"))
mylog.setLevel(40)
my_storage = {}
pbar = get_pbar("Constructing trajectory information",
len(self.data_series))
for i, (sto,
ds) in enumerate(self.data_series.piter(storage=my_storage)):
dd = ds.all_data()
idx_field = dd._determine_fields("particle_index")[0]
newtags = dd[idx_field].ndarray_view().astype("int64")
mask = np.in1d(newtags, indices, assume_unique=True)
sorts = np.argsort(newtags[mask])
self.array_indices.append(
np.where(np.in1d(indices, newtags, assume_unique=True))[0])
self.masks.append(mask)
self.sorts.append(sorts)
sto.result_id = ds.parameter_filename
sto.result = ds.current_time
pbar.update(i)
pbar.finish()
if self.suppress_logging:
mylog.setLevel(old_level)
times = []
for fn, time in sorted(my_storage.items()):
times.append(time)
self.times = self.data_series[0].arr([time for time in times],
times[0].units)
self.particle_fields = []
# Instantiate fields the caller requested
for field in fields:
self._get_data(field)
def has_key(self, key):
return (key in self.field_data)
def keys(self):
return self.field_data.keys()
def __getitem__(self, key):
"""
Get the field associated with key.
"""
if key == "particle_time":
return self.times
if key not in self.field_data:
self._get_data(key)
return self.field_data[key]
def __setitem__(self, key, val):
"""
Sets a field to be some other value.
"""
self.field_data[key] = val
def __delitem__(self, key):
"""
Delete the field from the trajectory
"""
del self.field_data[key]
def __iter__(self):
"""
This iterates over the trajectories for
the different particles, returning dicts
of fields for each trajectory
"""
for idx in range(self.num_indices):
traj = {}
traj["particle_index"] = self.indices[idx]
traj["particle_time"] = self.times
for field in self.field_data.keys():
traj[field] = self[field][idx, :]
yield traj
def __len__(self):
"""
The number of individual trajectories
"""
return self.num_indices
def add_fields(self, fields):
"""
Add a list of fields to an existing trajectory
Parameters
----------
fields : list of strings
A list of fields to be added to the current trajectory
collection.
Examples
________
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.add_fields(["particle_mass", "particle_gpot"])
"""
for field in fields:
if field not in self.field_data:
self._get_data(field)
def _get_data(self, field):
"""
Get a field to include in the trajectory collection.
The trajectory collection itself is a dict of 2D numpy arrays,
with shape (num_indices, num_steps)
"""
if field not in self.field_data:
if self.suppress_logging:
old_level = int(ytcfg.get("yt", "loglevel"))
mylog.setLevel(40)
ds_first = self.data_series[0]
dd_first = ds_first.all_data()
fd = dd_first._determine_fields(field)[0]
if field not in self.particle_fields:
if self.data_series[0].field_info[fd].particle_type:
self.particle_fields.append(field)
particles = np.empty((self.num_indices, self.num_steps))
particles[:] = np.nan
step = int(0)
pbar = get_pbar("Generating field %s in trajectories." % (field),
self.num_steps)
my_storage = {}
for i, (sto, ds) in enumerate(
self.data_series.piter(storage=my_storage)):
mask = self.masks[i]
sort = self.sorts[i]
if field in self.particle_fields:
# This is easy... just get the particle fields
dd = ds.all_data()
pfield = dd[fd].ndarray_view()[mask][sort]
else:
# This is hard... must loop over grids
pfield = np.zeros((self.num_indices))
x = self["particle_position_x"][:, step].ndarray_view()
y = self["particle_position_y"][:, step].ndarray_view()
z = self["particle_position_z"][:, step].ndarray_view()
# This will fail for non-grid index objects
particle_grids, particle_grid_inds = ds.index._find_points(
x, y, z)
for grid in particle_grids:
cube = grid.retrieve_ghost_zones(1, [fd])
CICSample_3(
x, y, z, pfield, self.num_indices, cube[fd],
np.array(grid.LeftEdge).astype(np.float64),
np.array(grid.ActiveDimensions).astype(np.int32),
grid.dds[0])
sto.result_id = ds.parameter_filename
sto.result = (self.array_indices[i], pfield)
pbar.update(step)
step += 1
pbar.finish()
for i, (fn, (indices,
pfield)) in enumerate(sorted(my_storage.items())):
particles[indices, i] = pfield
self.field_data[field] = array_like_field(dd_first, particles, fd)
if self.suppress_logging:
mylog.setLevel(old_level)
return self.field_data[field]
def trajectory_from_index(self, index):
"""
Retrieve a single trajectory corresponding to a specific particle
index
Parameters
----------
index : int
This defines which particle trajectory from the
ParticleTrajectories object will be returned.
Returns
-------
A dictionary corresponding to the particle's trajectory and the
fields along that trajectory
Examples
--------
>>> from yt.mods import *
>>> import matplotlib.pylab as pl
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> traj = trajs.trajectory_from_index(indices[0])
>>> pl.plot(traj["particle_time"], traj["particle_position_x"], "-x")
>>> pl.savefig("orbit")
"""
mask = np.in1d(self.indices, (index, ), assume_unique=True)
if not np.any(mask):
print("The particle index %d is not in the list!" % (index))
raise IndexError
fields = [field for field in sorted(self.field_data.keys())]
traj = {}
traj["particle_time"] = self.times
traj["particle_index"] = index
for field in fields:
traj[field] = self[field][mask, :][0]
return traj
@parallel_root_only
def write_out(self, filename_base):
"""
Write out particle trajectories to tab-separated ASCII files (one
for each trajectory) with the field names in the file header. Each
file is named with a basename and the index number.
Parameters
----------
filename_base : string
The prefix for the outputted ASCII files.
Examples
--------
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.write_out("orbit_trajectory")
"""
fields = [field for field in sorted(self.field_data.keys())]
num_fields = len(fields)
first_str = "# particle_time\t" + "\t".join(fields) + "\n"
template_str = "%g\t" * num_fields + "%g\n"
for ix in range(self.num_indices):
outlines = [first_str]
for it in range(self.num_steps):
outlines.append(
template_str %
tuple([self.times[it]] +
[self[field][ix, it] for field in fields]))
fid = open(filename_base + "_%d.dat" % self.indices[ix], "w")
fid.writelines(outlines)
fid.close()
del fid
@parallel_root_only
def write_out_h5(self, filename):
"""
Write out all the particle trajectories to a single HDF5 file
that contains the indices, the times, and the 2D array for each
field individually
Parameters
----------
filename : string
The output filename for the HDF5 file
Examples
--------
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.write_out_h5("orbit_trajectories")
"""
fid = h5py.File(filename, "w")
fields = [field for field in sorted(self.field_data.keys())]
fid.create_dataset("particle_indices",
dtype=np.int32,
data=self.indices)
fid.create_dataset("particle_time", data=self.times)
for field in fields:
fid.create_dataset("%s" % field, data=self[field])
fid.close()
class ParticleTrajectories(object):
r"""A collection of particle trajectories in time over a series of
datasets.
The ParticleTrajectories object contains a collection of
particle trajectories for a specified set of particle indices.
Parameters
----------
outputs : `yt.data_objects.time_series.DatasetSeries` or list of strings
DatasetSeries object, or a time-sorted list of filenames to
construct a new DatasetSeries object.
indices : array_like
An integer array of particle indices whose trajectories we
want to track. If they are not sorted they will be sorted.
fields : list of strings, optional
A set of fields that is retrieved when the trajectory
collection is instantiated.
Default : None (will default to the fields 'particle_position_x',
'particle_position_y', 'particle_position_z')
suppress_logging : boolean
Suppress yt's logging when iterating over the simulation time
series.
Default : False
Examples
________
>>> from yt.mods import *
>>> my_fns = glob.glob("orbit_hdf5_chk_00[0-9][0-9]")
>>> my_fns.sort()
>>> fields = ["particle_position_x", "particle_position_y",
>>> "particle_position_z", "particle_velocity_x",
>>> "particle_velocity_y", "particle_velocity_z"]
>>> ds = load(my_fns[0])
>>> init_sphere = ds.sphere(ds.domain_center, (.5, "unitary"))
>>> indices = init_sphere["particle_index"].astype("int")
>>> trajs = ParticleTrajectories(my_fns, indices, fields=fields)
>>> for t in trajs :
>>> print t["particle_velocity_x"].max(), t["particle_velocity_x"].min()
"""
def __init__(self, outputs, indices, fields=None, suppress_logging=False):
indices.sort() # Just in case the caller wasn't careful
self.field_data = YTFieldData()
if isinstance(outputs, DatasetSeries):
self.data_series = outputs
else:
self.data_series = DatasetSeries(outputs)
self.masks = []
self.sorts = []
self.array_indices = []
self.indices = indices
self.num_indices = len(indices)
self.num_steps = len(outputs)
self.times = []
self.suppress_logging = suppress_logging
# Default fields
if fields is None: fields = []
fields.append("particle_position_x")
fields.append("particle_position_y")
fields.append("particle_position_z")
fields = list(OrderedDict.fromkeys(fields))
if self.suppress_logging:
old_level = int(ytcfg.get("yt","loglevel"))
mylog.setLevel(40)
my_storage = {}
pbar = get_pbar("Constructing trajectory information", len(self.data_series))
for i, (sto, ds) in enumerate(self.data_series.piter(storage=my_storage)):
dd = ds.all_data()
idx_field = dd._determine_fields("particle_index")[0]
newtags = dd[idx_field].ndarray_view().astype("int64")
mask = np.in1d(newtags, indices, assume_unique=True)
sorts = np.argsort(newtags[mask])
self.array_indices.append(np.where(np.in1d(indices, newtags, assume_unique=True))[0])
self.masks.append(mask)
self.sorts.append(sorts)
sto.result_id = ds.parameter_filename
sto.result = ds.current_time
pbar.update(i)
pbar.finish()
if self.suppress_logging:
mylog.setLevel(old_level)
times = []
for fn, time in sorted(my_storage.items()):
times.append(time)
self.times = self.data_series[0].arr([time for time in times], times[0].units)
self.particle_fields = []
# Instantiate fields the caller requested
for field in fields:
self._get_data(field)
def has_key(self, key):
return (key in self.field_data)
def keys(self):
return self.field_data.keys()
def __getitem__(self, key):
"""
Get the field associated with key.
"""
if key == "particle_time":
return self.times
if key not in self.field_data:
self._get_data(key)
return self.field_data[key]
def __setitem__(self, key, val):
"""
Sets a field to be some other value.
"""
self.field_data[key] = val
def __delitem__(self, key):
"""
Delete the field from the trajectory
"""
del self.field_data[key]
def __iter__(self):
"""
This iterates over the trajectories for
the different particles, returning dicts
of fields for each trajectory
"""
for idx in range(self.num_indices):
traj = {}
traj["particle_index"] = self.indices[idx]
traj["particle_time"] = self.times
for field in self.field_data.keys():
traj[field] = self[field][idx,:]
yield traj
def __len__(self):
"""
The number of individual trajectories
"""
return self.num_indices
def add_fields(self, fields):
"""
Add a list of fields to an existing trajectory
Parameters
----------
fields : list of strings
A list of fields to be added to the current trajectory
collection.
Examples
________
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.add_fields(["particle_mass", "particle_gpot"])
"""
for field in fields:
if field not in self.field_data:
self._get_data(field)
def _get_data(self, field):
"""
Get a field to include in the trajectory collection.
The trajectory collection itself is a dict of 2D numpy arrays,
with shape (num_indices, num_steps)
"""
if field not in self.field_data:
if self.suppress_logging:
old_level = int(ytcfg.get("yt","loglevel"))
mylog.setLevel(40)
ds_first = self.data_series[0]
dd_first = ds_first.all_data()
fd = dd_first._determine_fields(field)[0]
if field not in self.particle_fields:
if self.data_series[0].field_info[fd].particle_type:
self.particle_fields.append(field)
particles = np.empty((self.num_indices,self.num_steps))
particles[:] = np.nan
step = int(0)
pbar = get_pbar("Generating field %s in trajectories." % (field), self.num_steps)
my_storage={}
for i, (sto, ds) in enumerate(self.data_series.piter(storage=my_storage)):
mask = self.masks[i]
sort = self.sorts[i]
if field in self.particle_fields:
# This is easy... just get the particle fields
dd = ds.all_data()
pfield = dd[fd].ndarray_view()[mask][sort]
else:
# This is hard... must loop over grids
pfield = np.zeros((self.num_indices))
x = self["particle_position_x"][:,step].ndarray_view()
y = self["particle_position_y"][:,step].ndarray_view()
z = self["particle_position_z"][:,step].ndarray_view()
# This will fail for non-grid index objects
particle_grids, particle_grid_inds = ds.index._find_points(x,y,z)
for grid in particle_grids:
cube = grid.retrieve_ghost_zones(1, [fd])
CICSample_3(x,y,z,pfield,
self.num_indices,
cube[fd],
np.array(grid.LeftEdge).astype(np.float64),
np.array(grid.ActiveDimensions).astype(np.int32),
grid.dds[0])
sto.result_id = ds.parameter_filename
sto.result = (self.array_indices[i], pfield)
pbar.update(step)
step += 1
pbar.finish()
for i, (fn, (indices, pfield)) in enumerate(sorted(my_storage.items())):
particles[indices,i] = pfield
self.field_data[field] = array_like_field(dd_first, particles, fd)
if self.suppress_logging:
mylog.setLevel(old_level)
return self.field_data[field]
def trajectory_from_index(self, index):
"""
Retrieve a single trajectory corresponding to a specific particle
index
Parameters
----------
index : int
This defines which particle trajectory from the
ParticleTrajectories object will be returned.
Returns
-------
A dictionary corresponding to the particle's trajectory and the
fields along that trajectory
Examples
--------
>>> from yt.mods import *
>>> import matplotlib.pylab as pl
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> traj = trajs.trajectory_from_index(indices[0])
>>> pl.plot(traj["particle_time"], traj["particle_position_x"], "-x")
>>> pl.savefig("orbit")
"""
mask = np.in1d(self.indices, (index,), assume_unique=True)
if not np.any(mask):
print("The particle index %d is not in the list!" % (index))
raise IndexError
fields = [field for field in sorted(self.field_data.keys())]
traj = {}
traj["particle_time"] = self.times
traj["particle_index"] = index
for field in fields:
traj[field] = self[field][mask,:][0]
return traj
@parallel_root_only
def write_out(self, filename_base):
"""
Write out particle trajectories to tab-separated ASCII files (one
for each trajectory) with the field names in the file header. Each
file is named with a basename and the index number.
Parameters
----------
filename_base : string
The prefix for the outputted ASCII files.
Examples
--------
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.write_out("orbit_trajectory")
"""
fields = [field for field in sorted(self.field_data.keys())]
num_fields = len(fields)
first_str = "# particle_time\t" + "\t".join(fields)+"\n"
template_str = "%g\t"*num_fields+"%g\n"
for ix in range(self.num_indices):
outlines = [first_str]
for it in range(self.num_steps):
outlines.append(template_str %
tuple([self.times[it]]+[self[field][ix,it] for field in fields]))
fid = open(filename_base + "_%d.dat" % self.indices[ix], "w")
fid.writelines(outlines)
fid.close()
del fid
@parallel_root_only
def write_out_h5(self, filename):
"""
Write out all the particle trajectories to a single HDF5 file
that contains the indices, the times, and the 2D array for each
field individually
Parameters
----------
filename : string
The output filename for the HDF5 file
Examples
--------
>>> from yt.mods import *
>>> trajs = ParticleTrajectories(my_fns, indices)
>>> trajs.write_out_h5("orbit_trajectories")
"""
fid = h5py.File(filename, "w")
fields = [field for field in sorted(self.field_data.keys())]
fid.create_dataset("particle_indices", dtype=np.int32,
data=self.indices)
fid.create_dataset("particle_time", data=self.times)
for field in fields:
fid.create_dataset("%s" % field, data=self[field])
fid.close()
class ProfileND(ParallelAnalysisInterface):
"""The profile object class"""
def __init__(self, data_source, weight_field = None):
self.data_source = data_source
self.ds = data_source.ds
self.field_map = {}
self.field_info = {}
self.field_data = YTFieldData()
if weight_field is not None:
self.variance = YTFieldData()
weight_field = self.data_source._determine_fields(weight_field)[0]
self.weight_field = weight_field
self.field_units = {}
ParallelAnalysisInterface.__init__(self, comm=data_source.comm)
def add_fields(self, fields):
"""Add fields to profile
Parameters
----------
fields : list of field names
A list of fields to create profile histograms for
"""
fields = self.data_source._determine_fields(fields)
for f in fields:
self.field_info[f] = self.data_source.ds.field_info[f]
temp_storage = ProfileFieldAccumulator(len(fields), self.size)
citer = self.data_source.chunks([], "io")
for chunk in parallel_objects(citer):
self._bin_chunk(chunk, fields, temp_storage)
self._finalize_storage(fields, temp_storage)
def set_field_unit(self, field, new_unit):
"""Sets a new unit for the requested field
Parameters
----------
field : string or field tuple
The name of the field that is to be changed.
new_unit : string or Unit object
The name of the new unit.
"""
if field in self.field_units:
self.field_units[field] = \
Unit(new_unit, registry=self.ds.unit_registry)
else:
fd = self.field_map[field]
if fd in self.field_units:
self.field_units[fd] = \
Unit(new_unit, registry=self.ds.unit_registry)
else:
raise KeyError("%s not in profile!" % (field))
def _finalize_storage(self, fields, temp_storage):
# We use our main comm here
# This also will fill _field_data
for i, field in enumerate(fields):
# q values are returned as q * weight but we want just q
temp_storage.qvalues[..., i][temp_storage.used] /= \
temp_storage.weight_values[temp_storage.used]
# get the profile data from all procs
all_store = {self.comm.rank: temp_storage}
all_store = self.comm.par_combine_object(all_store,
"join", datatype="dict")
all_val = np.zeros_like(temp_storage.values)
all_mean = np.zeros_like(temp_storage.mvalues)
all_var = np.zeros_like(temp_storage.qvalues)
all_weight = np.zeros_like(temp_storage.weight_values)
all_used = np.zeros_like(temp_storage.used, dtype="bool")
# Combine the weighted mean and variance from each processor.
# For two samples with total weight, mean, and variance
# given by w, m, and s, their combined mean and variance are:
# m12 = (m1 * w1 + m2 * w2) / (w1 + w2)
# s12 = (m1 * (s1**2 + (m1 - m12)**2) +
# m2 * (s2**2 + (m2 - m12)**2)) / (w1 + w2)
# Here, the mvalues are m and the qvalues are s**2.
for p in sorted(all_store.keys()):
all_used += all_store[p].used
old_mean = all_mean.copy()
old_weight = all_weight.copy()
all_weight[all_store[p].used] += \
all_store[p].weight_values[all_store[p].used]
for i, field in enumerate(fields):
all_val[..., i][all_store[p].used] += \
all_store[p].values[..., i][all_store[p].used]
all_mean[..., i][all_store[p].used] = \
(all_mean[..., i] * old_weight +
all_store[p].mvalues[..., i] *
all_store[p].weight_values)[all_store[p].used] / \
all_weight[all_store[p].used]
all_var[..., i][all_store[p].used] = \
(old_weight * (all_var[..., i] +
(old_mean[..., i] - all_mean[..., i])**2) +
all_store[p].weight_values *
(all_store[p].qvalues[..., i] +
(all_store[p].mvalues[..., i] -
all_mean[..., i])**2))[all_store[p].used] / \
all_weight[all_store[p].used]
all_var = np.sqrt(all_var)
del all_store
self.used = all_used
blank = ~all_used
self.weight = all_weight
self.weight[blank] = 0.0
for i, field in enumerate(fields):
if self.weight_field is None:
self.field_data[field] = \
array_like_field(self.data_source,
all_val[...,i], field)
else:
self.field_data[field] = \
array_like_field(self.data_source,
all_mean[...,i], field)
self.variance[field] = \
array_like_field(self.data_source,
all_var[...,i], field)
self.variance[field][blank] = 0.0
self.field_data[field][blank] = 0.0
self.field_units[field] = self.field_data[field].units
if isinstance(field, tuple):
self.field_map[field[1]] = field
else:
self.field_map[field] = field
def _bin_chunk(self, chunk, fields, storage):
raise NotImplementedError
def _filter(self, bin_fields):
# cut_points is set to be everything initially, but
# we also want to apply a filtering based on min/max
filter = np.ones(bin_fields[0].shape, dtype='bool')
for (mi, ma), data in zip(self.bounds, bin_fields):
filter &= (data > mi)
filter &= (data < ma)
return filter, [data[filter] for data in bin_fields]
def _get_data(self, chunk, fields):
# We are using chunks now, which will manage the field parameters and
# the like.
bin_fields = [chunk[bf] for bf in self.bin_fields]
# We want to make sure that our fields are within the bounds of the
# binning
filter, bin_fields = self._filter(bin_fields)
if not np.any(filter): return None
arr = np.zeros((bin_fields[0].size, len(fields)), dtype="float64")
for i, field in enumerate(fields):
units = chunk.ds.field_info[field].output_units
arr[:,i] = chunk[field][filter].in_units(units)
if self.weight_field is not None:
units = chunk.ds.field_info[self.weight_field].output_units
weight_data = chunk[self.weight_field].in_units(units)
else:
weight_data = np.ones(filter.size, dtype="float64")
weight_data = weight_data[filter]
# So that we can pass these into
return arr, weight_data, bin_fields
def __getitem__(self, field):
fname = self.field_map.get(field, None)
if fname is None:
if isinstance(field, tuple):
fname = self.field_map.get(field[1], None)
elif isinstance(field, DerivedField):
fname = self.field_map.get(field.name[1], None)
if fname is None:
raise KeyError(field)
else:
if getattr(self, 'fractional', False):
return self.field_data[fname]
else:
return self.field_data[fname].in_units(self.field_units[fname])
def items(self):
return [(k,self[k]) for k in self.field_data.keys()]
def keys(self):
return self.field_data.keys()
def __iter__(self):
return sorted(self.items())
def _get_bins(self, mi, ma, n, take_log):
if take_log:
ret = np.logspace(np.log10(mi), np.log10(ma), n+1)
# at this point ret[0] and ret[-1] are not exactly equal to
# mi and ma due to round-off error. Let's force them to be
# mi and ma exactly to avoid incorrectly discarding cells near
# the edges. See Issue #1300.
ret[0], ret[-1] = mi, ma
return ret
else:
return np.linspace(mi, ma, n+1)
def save_as_dataset(self, filename=None):
r"""Export a profile to a reloadable yt dataset.
This function will take a profile and output a dataset
containing all relevant fields. The resulting dataset
can be reloaded as a yt dataset.
Parameters
----------
filename : str, optional
The name of the file to be written. If None, the name
will be a combination of the original dataset plus
the type of object, e.g., Profile1D.
Returns
-------
filename : str
The name of the file that has been created.
Examples
--------
>>> import yt
>>> ds = yt.load("enzo_tiny_cosmology/DD0046/DD0046")
>>> ad = ds.all_data()
>>> profile = yt.create_profile(ad, ["density", "temperature"],
... "cell_mass", weight_field=None,
... n_bins=(128, 128))
>>> fn = profile.save_as_dataset()
>>> prof_ds = yt.load(fn)
>>> print (prof_ds.data["cell_mass"])
(128, 128)
>>> print (prof_ds.data["x"].shape) # x bins as 1D array
(128,)
>>> print (prof_ds.data["density"]) # x bins as 2D array
(128, 128)
>>> p = yt.PhasePlot(prof_ds.data, "density", "temperature",
... "cell_mass", weight_field=None)
>>> p.save()
"""
keyword = "%s_%s" % (str(self.ds), self.__class__.__name__)
filename = get_output_filename(filename, keyword, ".h5")
args = ("field", "log")
extra_attrs = {"data_type": "yt_profile",
"profile_dimensions": self.size,
"weight_field": self.weight_field,
"fractional": self.fractional}
data = {}
data.update(self.field_data)
data["weight"] = self.weight
data["used"] = self.used.astype("float64")
dimensionality = 0
bin_data = []
for ax in "xyz":
if hasattr(self, ax):
dimensionality += 1
data[ax] = getattr(self, ax)
bin_data.append(data[ax])
bin_field_name = "%s_bins" % ax
data[bin_field_name] = getattr(self, bin_field_name)
extra_attrs["%s_range" % ax] = self.ds.arr([data[bin_field_name][0],
data[bin_field_name][-1]])
for arg in args:
key = "%s_%s" % (ax, arg)
extra_attrs[key] = getattr(self, key)
bin_fields = np.meshgrid(*bin_data)
for i, ax in enumerate("xyz"[:dimensionality]):
data[getattr(self, "%s_field" % ax)] = bin_fields[i]
extra_attrs["dimensionality"] = dimensionality
ftypes = dict([(field, "data") for field in data])
save_as_dataset(self.ds, filename, data, field_types=ftypes,
extra_attrs=extra_attrs)
return filename