RaijuDataReader

class PlasmaCalcs.hookups.mage.raiju.raiju_io_tools.RaijuDataReader(ftag, *, dir='.')

Bases: MageicDataReader

handles reading data from raiju file(s).

The most common case is that raiju data is all stored in file like ‘runname.raiju.h5’

ftag: str

file tag, e.g. ‘runname’ from file like ‘runname.raiju.h5’

dir: str

directory where files are located. Stored as self.dirname = os.path.abspath(dir).

While it is technically possible to have more complicated grids,

for now PlasmaCalcs assumes (and checks!) the remix grid is in colat/lon coordinates,

and uses dimensions named (colat, long) corresponding to X and Y in the files.

(X and Y in files are stored as 2D arrays but each actually only varies along 1D.)

This helps to avoid ambiguity with the i, j, k dims of gamera and other mage calculators.

This is okay because, e.g., i=7 would represent very different regions in raiju vs gamera.

Methods

colat(*[, center, cache, squeeze, closeness_as])	return colatitude colat in radians, based on grid coordinates.
directly_loadable_vars()	return tuple of all directly loadable vars from 'Step#0' of file0.
file0_attr(key, *[, decode])	return value of attr, from self.open_file0().attrs[key].
file0_attrs(*[, decode])	return dict of file0.attrs.
fluids_info()	return iterable of dicts of fluid info from Raiju files.
get_var(v, step, *[, slices])	return numpy array of variable v at this step.
get_var0(v, step)	return numpy array of variable v at this step, from file0,
grid([x, center, cache])	return xarray.Dataset of grid coordinates.
grid_coords(*[, center, cache, ij, lat, ...])	return standard grid coords as Dataset with data vars from 'colat', 'lon'.
ijk_coords(*[, center])	return dict of {'i': i_coords, 'j': j_coords, (and 'k': k_coords, if 3D)}
lat(*[, center])	return latitude lat in radians.
lon(*[, center, cache, squeeze, closeness_as])	return longitude lon in radians, based on grid coordinates.
open_file0()	context manager for opening file0, for checking data info like shape, dtype, steps, etc.
steps_info()	return dict of arrays telling 's', 't', 'MJD', and possibly 'timestep' too.
_assert_shape_and_dtype(arr, shape, dtype, ...)	helper for gamera loading routines, ensure var has expected shape and dtype,
_grid_to_center(grid)	returns grid values at cell centers, given grid values (numpy array) at corners.
_resqueeze_slices(array, unsqueezed)	squeeze array along any axes which were unsqueezed.
_unsqueeze_slices(slices)	"un-squeezes slices, returning (unsqueezed_slices, dict telling which slices were changed).

Attributes

Nshape	(Ni, Nj) or (Ni, Nj, Nk) shape of a single full data var.
file0	the file used for checking data info like shape, dtype, steps, timeAttributeCache, etc.
steps	return list of all steps available (as strs).

property Nshape

(Ni, Nj) or (Ni, Nj, Nk) shape of a single full data var.

This is the shape after transposing from fortran order.

static _assert_shape_and_dtype(arr, shape, dtype, *, var)

helper for gamera loading routines, ensure var has expected shape and dtype,

crash if shape or dtype is wrong.

arr can be array or any object with shape and dtype

(e.g. h5py dataset even if didn’t load data into memory yet).

static _grid_to_center(grid)

returns grid values at cell centers, given grid values (numpy array) at corners.

grid: 1D, 2D, or 3D numpy array

grid values at corners, with shape (Ni+1,), (Ni+1, Nj+1), or (Ni+1, Nj+1, Nk+1).

result has shape (Ni,), (Ni, Nj), or (Ni, Nj, Nk).

static _resqueeze_slices(array, unsqueezed)

squeeze array along any axes which were unsqueezed.

assumes array axes will be ‘i’, ‘j’, (and ‘k’ if 3D).

static _unsqueeze_slices(slices)

“un-squeezes slices, returning (unsqueezed_slices, dict telling which slices were changed).

unsqueeze converts integer to iterable of length 1.

The relevant dim should be squeezed away at some point to respect original inputs,

but for internal processing it is nice to not worry about number of dims changing.

colat(*, center=True, cache=True, squeeze='force', closeness_as='coords')

return colatitude colat in radians, based on grid coordinates.

There is a good chance that result will be a 1D array (if squeeze=True),

because colat should only vary with one of the coordinates (i)

(corresponding to ‘X’ in the files).

center: bool

whether to center grid coordinates on cell centers (True) or use the full grid.

if False, result will have size Ni+1 instead of Ni.

cache: bool

whether to use caching, instead of recalculating each time.

squeeze: ‘force’ or bool

whether to squeeze result to 1D array, if possible.

if ‘force’, crash if squeezing isn’t possible.

closeness_as: ‘coords’, ‘attrs’, or None

where to put info about closeness; see xarray_squeeze_close for details.

directly_loadable_vars()

return tuple of all directly loadable vars from ‘Step#0’ of file0.

property file0

the file used for checking data info like shape, dtype, steps, timeAttributeCache, etc.

Here, just returns abspath to the {self.ftag}.raiju.h5 file.

file0_attr(key, *, decode=True)

return value of attr, from self.open_file0().attrs[key].

decode: bool, whether to result.decode(‘utf-8’) if result is a bytes string.

file0_attrs(*, decode=True)

return dict of file0.attrs.

E.g. result[‘UnitsID’] tells the UnitsID for the data.

decode: bool, whether to s.decode(‘utf-8’) for bytes strings s.

fluids_info()

return iterable of dicts of fluid info from Raiju files.

info includes ‘i’, ‘m’, ‘q’, and ‘params’.

m is in amu (e.g. ~1 for H) and q is in elementary charge units (e.g. -1 for electrons).

get_var(v, step, *, slices=None)

return numpy array of variable v at this step.

(Transposed from fortran order to python order.)

slices: None or dict

if provided, only load the pieces needed by slices; keys can be i, j, k.

Not implemented here; subclass should implement. (See also: self.get_var0)

get_var0(v, step)

return numpy array of variable v at this step, from file0,

assuming file0 data shape matches full Nshape (else will crash).

(Transposed from fortran order to python order.)

grid(x=None, *, center=True, cache=True)

return xarray.Dataset of grid coordinates.

Will return DataArray instead if x is a single str.

x: None, ‘x’, ‘y’, ‘z’, or list of strs.

tells which coordinates to return. (‘X’, ‘Y’, ‘Z’ in the files themselves.)

‘x’, ‘y’, or ‘z’ –> return DataArray of that coordinate only.

list of strs –> return Dataset but with those coordinates only.

None –> (‘x’, ‘y’) if ‘Z’ missing from file0, else (‘x’, ‘y’, ‘z’)

center: bool

whether to return coordinates at cell centers.

When False, the grid has shape (Ni+1, Nj+1) or (Ni+1, Nj+1, Nk+1).

telling the coordinates at grid corners.

When True, the grid has shape (Ni, Nj) or (Ni, Nj, Nk),

telling the coordinates at grid centers.

cache: bool

whether to cache result (for x and center).

if True, cache result in case it is reloaded.

if False, don’t check or alter cache.

grid_coords(*, center=True, cache=True, ij=False, lat=False, as_coords=False)

return standard grid coords as Dataset with data vars from ‘colat’, ‘lon’.

Some kwargs control which data vars are included, and what will be the underlying dims.

‘colat’ and ‘lon’ are always provided in the result, as 1D coords,

and will crash if that is not possible.

center, cache:

passed directly to self.grid(); see help(self.grid) for details.

ij: bool

whether to include ‘i’ and ‘j’ coords and dims in the result.

if False, will use ‘colat’ and ‘lon’ as dims instead of ‘i’ and ‘j’.

lat: bool

whether to include ‘lat’ coord in result. lat = np.pi/2 - colat.

as_coords: bool

whether to set_coords for all data vars in result.

if False, data_vars will not be included in result.coords.

if True, will do result.set_coords(result.data_vars).

(there is no option to include ‘x’ and ‘y’ coords in result,

because the file data for ‘X’ and ‘Y’ actually corresponds to ‘colat’ and ‘lon’,

not any sort of rectilinear x and y values.)

ijk_coords(*, center=False)

return dict of {‘i’: i_coords, ‘j’: j_coords, (and ‘k’: k_coords, if 3D)}

where coords depend on center:

if center: i_coords = np.arange(Ni)

else: i_coords = -0.5+np.arange(1+Ni)

and similar for j_coords and k_coords.

lat(*, center=True, **kw_colat)

return latitude lat in radians. lat = np.pi/2 - colat. kwargs go to self.colat().

center: bool, whether to center (full shape Ni) or use full grid (full shape Ni+1).

lon(*, center=True, cache=True, squeeze='force', closeness_as='coords')

return longitude lon in radians, based on grid coordinates.

There is a good chance that result will be a 1D array (if squeeze=True),

because lon should only vary with one of the coordinates (j)

(corresponding to ‘Y’ in the files).

center: bool

whether to center grid coordinates on cell centers (True) or use the full grid.

if False, result will have size Nj+1 instead of Nj.

cache: bool

whether to use caching, instead of recalculating each time.

squeeze: ‘force’ or bool

whether to squeeze result to 1D array, if possible.

if ‘force’, crash if squeezing isn’t possible.

closeness_as: ‘coords’, ‘attrs’, or None

where to put info about closeness; see xarray_squeeze_close for details.

open_file0()

context manager for opening file0, for checking data info like shape, dtype, steps, etc.

property steps

return list of all steps available (as strs).

Sorts steps in increasing order of step number, i.e. step 7 comes before step 10.

steps_info()

return dict of arrays telling ‘s’, ‘t’, ‘MJD’, and possibly ‘timestep’ too.

‘s’ is step number (as str); ‘t’ is time; ‘MJD’ is Modified Julian Date.