"""
AMR index container class
"""
from __future__ import print_function
#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------
import h5py
import numpy as np
import string, re, gc, time
from yt.extern.six.moves import cPickle
from yt.extern.six.moves import zip as izip
import weakref
from itertools import chain
from yt.funcs import *
from yt.utilities.logger import ytLogger as mylog
from yt.arraytypes import blankRecordArray
from yt.config import ytcfg
from yt.fields.field_info_container import NullFunc
from yt.geometry.geometry_handler import \
Index, YTDataChunk, ChunkDataCache
from yt.utilities.definitions import MAXLEVEL
from yt.utilities.physical_constants import sec_per_year
from yt.utilities.io_handler import io_registry
from yt.utilities.parallel_tools.parallel_analysis_interface import \
ParallelAnalysisInterface
from .grid_container import \
GridTree, MatchPointsToGrids
from yt.data_objects.data_containers import data_object_registry
[docs]class GridIndex(Index):
"""The index class for patch and block AMR datasets. """
float_type = 'float64'
_preload_implemented = False
_index_properties = ("grid_left_edge", "grid_right_edge",
"grid_levels", "grid_particle_count",
"grid_dimensions")
def _setup_geometry(self):
mylog.debug("Counting grids.")
self._count_grids()
mylog.debug("Initializing grid arrays.")
self._initialize_grid_arrays()
mylog.debug("Parsing index.")
self._parse_index()
mylog.debug("Constructing grid objects.")
self._populate_grid_objects()
mylog.debug("Re-examining index")
self._initialize_level_stats()
def __del__(self):
del self.grid_dimensions
del self.grid_left_edge
del self.grid_right_edge
del self.grid_levels
del self.grid_particle_count
del self.grids
@property
def parameters(self):
return self.dataset.parameters
def _detect_output_fields_backup(self):
# grab fields from backup file as well, if present
return
try:
backup_filename = self.dataset.backup_filename
f = h5py.File(backup_filename, 'r')
g = f["data"]
grid = self.grids[0] # simply check one of the grids
grid_group = g["grid_%010i" % (grid.id - grid._id_offset)]
for field_name in grid_group:
if field_name != 'particles':
self.field_list.append(field_name)
except KeyError:
return
except IOError:
return
[docs] def select_grids(self, level):
"""
Returns an array of grids at *level*.
"""
return self.grids[self.grid_levels.flat == level]
[docs] def get_levels(self):
for level in range(self.max_level+1):
yield self.select_grids(level)
def _initialize_grid_arrays(self):
mylog.debug("Allocating arrays for %s grids", self.num_grids)
self.grid_dimensions = np.ones((self.num_grids,3), 'int32')
self.grid_left_edge = self.ds.arr(np.zeros((self.num_grids,3),
self.float_type), 'code_length')
self.grid_right_edge = self.ds.arr(np.ones((self.num_grids,3),
self.float_type), 'code_length')
self.grid_levels = np.zeros((self.num_grids,1), 'int32')
self.grid_particle_count = np.zeros((self.num_grids,1), 'int32')
[docs] def clear_all_data(self):
"""
This routine clears all the data currently being held onto by the grids
and the data io handler.
"""
for g in self.grids: g.clear_data()
self.io.queue.clear()
[docs] def get_smallest_dx(self):
"""
Returns (in code units) the smallest cell size in the simulation.
"""
return self.select_grids(self.grid_levels.max())[0].dds[:].min()
def _initialize_level_stats(self):
# Now some statistics:
# 0 = number of grids
# 1 = number of cells
# 2 = blank
desc = {'names': ['numgrids','numcells','level'],
'formats':['Int64']*3}
self.level_stats = blankRecordArray(desc, MAXLEVEL)
self.level_stats['level'] = [i for i in range(MAXLEVEL)]
self.level_stats['numgrids'] = [0 for i in range(MAXLEVEL)]
self.level_stats['numcells'] = [0 for i in range(MAXLEVEL)]
for level in range(self.max_level+1):
self.level_stats[level]['numgrids'] = np.sum(self.grid_levels == level)
li = (self.grid_levels[:,0] == level)
self.level_stats[level]['numcells'] = self.grid_dimensions[li,:].prod(axis=1).sum()
@property
def grid_corners(self):
return np.array([
[self.grid_left_edge[:,0], self.grid_left_edge[:,1], self.grid_left_edge[:,2]],
[self.grid_right_edge[:,0], self.grid_left_edge[:,1], self.grid_left_edge[:,2]],
[self.grid_right_edge[:,0], self.grid_right_edge[:,1], self.grid_left_edge[:,2]],
[self.grid_left_edge[:,0], self.grid_right_edge[:,1], self.grid_left_edge[:,2]],
[self.grid_left_edge[:,0], self.grid_left_edge[:,1], self.grid_right_edge[:,2]],
[self.grid_right_edge[:,0], self.grid_left_edge[:,1], self.grid_right_edge[:,2]],
[self.grid_right_edge[:,0], self.grid_right_edge[:,1], self.grid_right_edge[:,2]],
[self.grid_left_edge[:,0], self.grid_right_edge[:,1], self.grid_right_edge[:,2]],
], dtype='float64')
[docs] def lock_grids_to_parents(self):
r"""This function locks grid edges to their parents.
This is useful in cases where the grid structure may be somewhat
irregular, or where setting the left and right edges is a lossy
process. It is designed to correct situations where left/right edges
may be set slightly incorrectly, resulting in discontinuities in images
and the like.
"""
mylog.info("Locking grids to parents.")
for i, g in enumerate(self.grids):
si = g.get_global_startindex()
g.LeftEdge = self.ds.domain_left_edge + g.dds * si
g.RightEdge = g.LeftEdge + g.ActiveDimensions * g.dds
self.grid_left_edge[i,:] = g.LeftEdge
self.grid_right_edge[i,:] = g.RightEdge
[docs] def print_stats(self):
"""
Prints out (stdout) relevant information about the simulation
"""
header = "%3s\t%6s\t%14s\t%14s" % ("level","# grids", "# cells",
"# cells^3")
print(header)
print("%s" % (len(header.expandtabs())*"-"))
for level in range(MAXLEVEL):
if (self.level_stats['numgrids'][level]) == 0:
break
print("% 3i\t% 6i\t% 14i\t% 14i" % \
(level, self.level_stats['numgrids'][level],
self.level_stats['numcells'][level],
np.ceil(self.level_stats['numcells'][level]**(1./3))))
dx = self.select_grids(level)[0].dds[0]
print("-" * 46)
print(" \t% 6i\t% 14i" % (self.level_stats['numgrids'].sum(), self.level_stats['numcells'].sum()))
print("\n")
try:
print("z = %0.8f" % (self["CosmologyCurrentRedshift"]))
except:
pass
print("t = %0.8e = %0.8e s = %0.8e years" % \
(self.ds.current_time.in_units("code_time"),
self.ds.current_time.in_units("s"),
self.ds.current_time.in_units("yr")))
print("\nSmallest Cell:")
u=[]
for item in ("Mpc", "pc", "AU", "cm"):
print("\tWidth: %0.3e %s" % (dx.in_units(item), item))
def _find_field_values_at_points(self, fields, coords):
r"""Find the value of fields at a set of coordinates.
Returns the values [field1, field2,...] of the fields at the given
(x, y, z) points. Returns a numpy array of field values cross coords
"""
coords = YTArray(ensure_numpy_array(coords),'code_length', registry=self.ds.unit_registry)
grids = self._find_points(coords[:,0], coords[:,1], coords[:,2])[0]
fields = ensure_list(fields)
mark = np.zeros(3, dtype=np.int)
out = []
# create point -> grid mapping
grid_index = {}
for coord_index, grid in enumerate(grids):
if grid not in grid_index:
grid_index[grid] = []
grid_index[grid].append(coord_index)
out = np.zeros((len(fields),len(coords)), dtype=np.float64)
for grid in grid_index:
cellwidth = (grid.RightEdge - grid.LeftEdge) / grid.ActiveDimensions
for field in fields:
for coord_index in grid_index[grid]:
mark = ((coords[coord_index,:] - grid.LeftEdge) / cellwidth).astype('int')
out[:,coord_index] = grid[field][mark[0],mark[1],mark[2]]
return out
def _find_points(self, x, y, z) :
"""
Returns the (objects, indices) of leaf grids containing a number of (x,y,z) points
"""
x = ensure_numpy_array(x)
y = ensure_numpy_array(y)
z = ensure_numpy_array(z)
if not len(x) == len(y) == len(z):
raise AssertionError("Arrays of indices must be of the same size")
grid_tree = self._get_grid_tree()
pts = MatchPointsToGrids(grid_tree, len(x), x, y, z)
ind = pts.find_points_in_tree()
return self.grids[ind], ind
def _get_grid_tree(self):
left_edge = self.ds.arr(np.zeros((self.num_grids, 3)),
'code_length')
right_edge = self.ds.arr(np.zeros((self.num_grids, 3)),
'code_length')
level = np.zeros((self.num_grids), dtype='int64')
parent_ind = np.zeros((self.num_grids), dtype='int64')
num_children = np.zeros((self.num_grids), dtype='int64')
dimensions = np.zeros((self.num_grids, 3), dtype="int32")
for i, grid in enumerate(self.grids) :
left_edge[i,:] = grid.LeftEdge
right_edge[i,:] = grid.RightEdge
level[i] = grid.Level
if grid.Parent is None :
parent_ind[i] = -1
else :
parent_ind[i] = grid.Parent.id - grid.Parent._id_offset
num_children[i] = np.int64(len(grid.Children))
dimensions[i,:] = grid.ActiveDimensions
return GridTree(self.num_grids, left_edge, right_edge, dimensions,
parent_ind, level, num_children)
[docs] def convert(self, unit):
return self.dataset.conversion_factors[unit]
def _identify_base_chunk(self, dobj):
fast_index = None
def _gsort(g):
if g.filename is None:
return g.id
return g.filename
if dobj._type_name == "grid":
dobj._chunk_info = np.empty(1, dtype='object')
dobj._chunk_info[0] = weakref.proxy(dobj)
elif getattr(dobj, "_grids", None) is None:
gi = dobj.selector.select_grids(self.grid_left_edge,
self.grid_right_edge,
self.grid_levels)
grids = list(sorted(self.grids[gi], key = _gsort))
dobj._chunk_info = np.empty(len(grids), dtype='object')
for i, g in enumerate(grids):
dobj._chunk_info[i] = g
# These next two lines, when uncommented, turn "on" the fast index.
#if dobj._type_name != "grid":
# fast_index = self._get_grid_tree()
if getattr(dobj, "size", None) is None:
dobj.size = self._count_selection(dobj, fast_index = fast_index)
if getattr(dobj, "shape", None) is None:
dobj.shape = (dobj.size,)
dobj._current_chunk = list(self._chunk_all(dobj, cache = False,
fast_index = fast_index))[0]
def _count_selection(self, dobj, grids = None, fast_index = None):
if fast_index is not None:
return fast_index.count(dobj.selector)
if grids is None: grids = dobj._chunk_info
count = sum((g.count(dobj.selector) for g in grids))
return count
def _chunk_all(self, dobj, cache = True, fast_index = None):
gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
fast_index = fast_index or getattr(dobj._current_chunk, "_fast_index",
None)
yield YTDataChunk(dobj, "all", gobjs, dobj.size,
cache, fast_index = fast_index)
def _chunk_spatial(self, dobj, ngz, sort = None, preload_fields = None):
gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
if sort in ("+level", "level"):
giter = sorted(gobjs, key = g.Level)
elif sort == "-level":
giter = sorted(gobjs, key = -g.Level)
elif sort is None:
giter = gobjs
if preload_fields is None: preload_fields = []
preload_fields, _ = self._split_fields(preload_fields)
if self._preload_implemented and len(preload_fields) > 0 and ngz == 0:
giter = ChunkDataCache(list(giter), preload_fields, self)
for i, og in enumerate(giter):
if ngz > 0:
g = og.retrieve_ghost_zones(ngz, [], smoothed=True)
else:
g = og
size = self._count_selection(dobj, [og])
if size == 0: continue
# We don't want to cache any of the masks or icoords or fcoords for
# individual grids.
yield YTDataChunk(dobj, "spatial", [g], size, cache = False)
_grid_chunksize = 1000
def _chunk_io(self, dobj, cache=True, local_only=False,
preload_fields=None, chunk_sizing="auto"):
# local_only is only useful for inline datasets and requires
# implementation by subclasses.
if preload_fields is None:
preload_fields = []
preload_fields, _ = self._split_fields(preload_fields)
gfiles = defaultdict(list)
gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
fast_index = dobj._current_chunk._fast_index
for g in gobjs:
gfiles[g.filename].append(g)
# We can apply a heuristic here to make sure we aren't loading too
# many grids all at once.
if chunk_sizing == "auto":
chunk_ngrids = len(gobjs)
if chunk_ngrids > 0:
nproc = np.float(ytcfg.getint("yt", "__global_parallel_size"))
chunking_factor = np.ceil(self._grid_chunksize*nproc/chunk_ngrids).astype("int")
size = max(self._grid_chunksize//chunking_factor, 1)
else:
size = self._grid_chunksize
elif chunk_sizing == "config_file":
size = ytcfg.getint("yt", "chunk_size")
elif chunk_sizing == "just_one":
size = 1
elif chunk_sizing == "old":
size = self._grid_chunksize
else:
raise RuntimeError("%s is an invalid value for the 'chunk_sizing' argument." % chunk_sizing)
for fn in sorted(gfiles):
gs = gfiles[fn]
for grids in (gs[pos:pos + size] for pos
in range(0, len(gs), size)):
dc = YTDataChunk(dobj, "io", grids,
self._count_selection(dobj, grids),
cache = cache, fast_index = fast_index)
# We allow four full chunks to be included.
with self.io.preload(dc, preload_fields,
4.0 * size):
yield dc