def meshX(coord=2, filename="visdump_mesh.h5", directory=".", key=None):
return mesh.meshElemCentroids(filename,directory,key)[:,coord]
def transect_data(varnames,
keys='all',
directory=".",
filename="visdump_data.h5",
mesh_filename="visdump_mesh.h5",
coord_order=None,
deformable=False,
return_map=False):
"""Pulls simulation output into structured 2D arrays for transect-based, (i,j) indexing.
Input:
varnames | A list of variable names to pull, e.g.
| ['saturation_liquid', 'saturation_ice'], or a single variable
| name, e.g. 'saturation_liquid'
keys | Indices of timesteps to pull. Either an int (i.e. 0, -1, etc)
| for the kth timestep, or a list of ints, or 'all'.
directory | Directory of the run. Defaults to '.'
filename | Filename of the run. Defaults to 'visdump_data.h5'
mesh_filename | Filename of the mesh. Defaults to 'visdump_mesh.h5'
coord_order | Order of the transect coordinates. Defaults to ['x','z']. The
| mesh is sorted in this order.
deformable | Is the mesh deforming?
return_map | See return value below.
Output:
Output is an array of shape:
( len(varnames+2), len(keys), n_cells_coord_order[0], n_cells_coord_order[1] )
data[0,0,:,:] is the coord_order[0] centroid
data[1,0,:,:] is the coord_order[1] centroid
data[i+2,k,:,:] is the ith varname data at the kth requested timestep, sorted in
the same way as the centroids.
Note that the data is re-ordered in INCREASING coordinate, i.e. bottom to top in z.
If return_map is True, then returns a tuple, (data, map) where
map is a (NX,NZ) array of integers specifying which global id
corresponds to the (i,j) cell. This is useful for mapping input
data back INTO the unstructured mesh.
Example usage:
Calculate and plot the thaw depth at step 5.
// Pull saturation ice -- TD is where sat ice = 0."
data = transect_data(['saturation_ice', 5)
// x coordinate for plotting
x = data[0,0,:,0]
// for each column, find highest z where sat_ice > 0.
td_i = np.array([np.where(data[2,0,i,:] > 0.)[0][-1] for i in range(data.shape[2])])
// now that we have an index into the highest cell with ice, determine td as the
// mean of the highest cell with ice and the one above that. Note this assumes
// all columns have some thawing.
td_z = np.array( [ (dat[1,0,i,td_i[i]] + dat[1,0,i,td_i[i+1]]) / 2.
for i in range(len(td_i)) ] )
plt.plot(x, td_z)
"""
if coord_order is None:
coord_order = ['x', 'z']
if type(varnames) is str:
varnames = [
varnames,
]
# get centroids
xyz = mesh.meshElemCentroids(mesh_filename, directory)
# round to avoid issues
xyz = np.round(xyz, decimals=5)
# get ordering of centroids
dtype = [(coord_order[0], float), (coord_order[1], float)]
num_order = []
for i in coord_order:
if i == 'x':
num_order.append(0)
elif i == 'y':
num_order.append(1)
elif i == 'z':
num_order.append(2)
xyz_sort_order = np.array(
[tuple([xyz[i, x] for x in num_order]) for i in range(len(xyz))],
dtype=dtype)
xyz_sorting = xyz_sort_order.argsort(order=coord_order)
with h5py.File(os.path.join(directory, filename), 'r') as dat:
keys_avail = dat[fullname(varnames[0])].keys()
keys_avail.sort(lambda a, b: int.__cmp__(int(a), int(b)))
if keys == 'all':
keys = keys_avail
elif type(keys) is str:
keys = [
keys,
]
elif type(keys) is int:
keys = [
keys_avail[keys],
]
elif type(keys) is slice:
keys = keys_avail[keys]
elif type(keys) is list:
if all(type(k) is int for k in keys):
keys = [keys_avail[k] for k in keys]
elif all(type(k) is str for k in keys):
pass
else:
raise RuntimeError(
"Keys requested cannot be processed -- should be 'all', int, or str key, or list of ints or strs."
)
# get data
vals = np.zeros((len(varnames) + 2, len(keys), len(xyz)), 'd')
for i, key in enumerate(keys):
if deformable:
xyz = mesh.meshElemCentroids(mesh_filename, directory)
vals[0, i, :] = xyz[xyz_sorting, num_order[0]]
vals[1, i, :] = xyz[xyz_sorting, num_order[1]]
for j, varname in enumerate(varnames):
vals[j + 2, i, :] = dat[fullname(varname)][key][:,
0][xyz_sorting]
# reshape the data
# determine nx
nx = len(set(vals[0, 0, :]))
nz = vals.shape[2] / nx
if (nx * nz != vals.shape[2]):
raise RuntimeError(
"Assumption about first coordinate being cleanly binnable is falling apart -- ask Ethan to rethink this algorithm!"
)
shp = vals.shape
if not return_map:
return vals.reshape(shp[0], shp[1], nx, nz)
else:
return vals.reshape(shp[0], shp[1], nx,
nz), xyz_sorting.reshape(nx, nz)
def meshX(coord=2, filename="visdump_mesh.h5", directory=".", key=None):
return mesh.meshElemCentroids(filename, directory, key)[:, coord]