PlasmaCalcs.hookups.bifrost.bifrost_stagger.StaggerPrePadManager3D

class PlasmaCalcs.hookups.bifrost.bifrost_stagger.StaggerPrePadManager3D(ops, slices=None, *, sizes)

Bases: object

manages pre-padding for stagger operations in ‘x’, ‘y’, ‘z’ dimensions.

(3D stagger operations’ inputs and outputs are always 3D arrays.)

ops: str

string of operations to do, separated by whitespace or ‘_’.

each operation must be one of:

‘xup’, ‘xdn’, ‘yup’, ‘ydn’, ‘zup’, ‘zdn’,

‘ddxup’, ‘ddxdn’, ‘ddyup’, ‘ddydn’, ‘ddzup’, ‘ddzdn’.

slices: dict of indexers or None

dict of ‘x’, ‘y’, ‘z’ to indexers (e.g. from a QuantityLoader).

None –> empty dict.

size: dict of int

lengths of ‘x’, ‘y’, ‘z’ dimensions.

calling self returns padded_slices:

dict of ‘x’, ‘y’, ‘z’ slices to use to pad an array before doing this series of ops,

such that the result will be properly aligned & shaped the same as non-staggered vars.

Note: dims without slice will have result[dim] = None.

dims with slice will have result[dim] =

sum of self.PAD_AMOUNT[‘up’ or ‘dn’] for all ops with dim.

E.g. ‘xup ddxup xdn’ –> result[‘x’] = (2,3) + (2,3) + (3,2) = (7,8).

__init__(ops, slices=None, *, sizes)

Methods

__init__(ops[, slices])
cumulative_pad_amounts()
get_pad_amounts()
get_prepad_dims()
get_prepad_slice(x)
get_prepad_slices()
op_prepad_kws()
stagger_postprocess(array)
unpad1D(x, array)
xarray_copy_kw_from_unstaggered(unstaggered)

Attributes

PAD_AMOUNT

cumulative_pad_amounts()

return amounts of padding for each dim, after each op.

e.g. [{‘x’:(0,0),’y’:(0,0),’z’:(2,3)}, {‘x’:(2,3),’y’:(0,0),’z’:(2,3)}, …]

get_pad_amounts()

return amounts of padding for each dim, based on ops (ignoring self.slices for now).

get_prepad_dims()

return dict of which dims are adjusted by prepadding.

get_prepad_slice(x)

returns padded slice (or other indexer) for this dim (‘x’, ‘y’, or ‘z’)

None if self.pad_amount[x] is None.

None if self.slices[x] == slice(None).

self.slices[x] if self.pad_amount[x] is (0,0).

get_prepad_slices()

return padded_slices to use.

like self.slices but with padding added to provided indexers.

see help(type(self)) for more details.

op_prepad_kws()

return list of dicts to use for stagger.do(array, op, **kw), for op in self.ops.

[TODO] handle left_first=False.

stagger_postprocess(array)

postprocessing of array after stagger ops, to deal with padding details.

Here, does the following:

(1) ensure that array is 3D. (probably a numpy array, but xarray is permitted too.)

(2) if any modes are ‘step_slice’, slice to recover only the relevant values.

(3) squeeze array after stagger ops, to respect any (unpadded) int indexers.

Also asserts that dims to squeeze actually have size 1.

unpad1D(x, array)

remove padding from 1D array corresponding to x dimension.

assumes this is being called after self.get_prepad_slices()

(i.e., expects self.modes and self.to_squeeze_after, to be set appropriately.)

x: ‘x’, ‘y’, or ‘z’.

if self.to_squeeze_after[x], result will be a scalar instead of a 1D array.

xarray_copy_kw_from_unstaggered(unstaggered)

return xarray_copy_kw(unstaggered), updated to remove padding.

E.g.:

staggered = stagger(unstaggered, op, pre_padded=self.get_prepad_slices())

kw_data_array = self.xarray_copy_kw_from_unstaggered(unstaggered)

as_data_array = xr.DataArray(staggered, **kw_data_array)