PartitionFromXarray

class PlasmaCalcs.tools.iterables.PartitionFromXarray(xarray, coord=None, original=None)

Bases: Partition

(ordered) dict of lists partitioned by category.

(assumes python >= 3.7, i.e. dict order is guaranteed to be insertion order)

self assumes it will remain unchanged after originally being created.

self.partition stores {key: list of elements in category}

self.idx stores {key: list of indices of elements in original}

self.flat stores all elements in category order, in a single list.

self.idx_flat concatenates self.idx.values().

self.ridx stores {key: [i such that self.flat[i] == x for x in this category]}

self.ridx_flat stores [i such that self.flat[i] == x for x in original]

iterating through self yields (key, value) pairs.

xarray: xarray.DataArray with ndim <= 1

data array determining partition category for o in original.

self.partition keys will be values of xarray;

partition values will be lists of o from original for which xarray[coord] == o.

E.g. xr.DataArray([‘a’, ‘a’, ‘b’, ‘a’, ‘c’], coords={coord: [‘o1’, ‘o2’, ‘o3’, ‘o4’, ‘o5’]})

–> self.partition = {‘a’: [‘o1’, ‘o2’, ‘o4’], ‘b’: [‘o3’], ‘c’: [‘o5’]}

coord: None or str

coordinate of xarray containing all values from original.

None –> infer as coord = xarray.dims[0] if ndim=1,

else xarray.coords[0] if ndim=0 and len(xarray.coords)==1,

else crash.

original: None or list-like

the original iterable, to be partitioned.

None –> use list(xarray[coord].values).

Methods

__contains__(key)	tells whether key is in self.partition
__iter__()	iterate through (key, val) pairs.
f(x)	return category for x.
get(key[, default])	return self[key].
init_all()	init all the things, based on self.original, self.f, and self._init_keys
init_flat()	create self.flat and self.idx_flat.
init_partition()	creates self.partition and self.idx
init_ridx()	create self.ridx and self.ridx_flat.
items()	return tuple of (key, value) pairs corresponding to self.keys() and self.values().
keys()	return tuple of keys in self.
values()	return tuple of values corresponding to self.keys().

__contains__(key)

tells whether key is in self.partition

__iter__()

iterate through (key, val) pairs.

f(x)

return category for x.

get(key, default=UNSET)

return self[key]. if default is provided and self[key] doesn’t exist, return default.

init_all()

init all the things, based on self.original, self.f, and self._init_keys

init_flat()

create self.flat and self.idx_flat.

self.flat stores all elements in order of categories, in a single list.

self.idx_flat stores the indices of all elements in self.flat.

init_partition()

creates self.partition and self.idx

init_ridx()

create self.ridx and self.ridx_flat.

self.ridx stores {key: [i such that self.flat[i] == x for x in this category]}

self.ridx_flat stores [i such that self.flat[i] == x for x in original]

items()

return tuple of (key, value) pairs corresponding to self.keys() and self.values().

keys()

return tuple of keys in self.

values()

return tuple of values corresponding to self.keys().