plasma_plots.arrays

Small xarray metadata helpers used by plasma_plots.

They intentionally live here rather than in Struphy so the plotting package can operate on labeled xarray data from any producer.

Attributes

Name	Description
`DIM_LABELS`	No description.
`LOGICAL_DIMS`	No description.
`SCALARS_EXCLUDE`	No description.

Functions

Name	Description
`angle_period`	The period of an angle coordinate, from its period attribute.
`axis_label`	The axis label of a dimension, with its units.
`close_periodic`	Repeat the first slice at the end of every direction that wraps around in physical space.
`logical_derivative`	The derivative ∂ values / ∂ coordinate along axis.
`logical_dims`	The names of the three logical dimensions of data, radial first.
`map_coordinate`	Replace a coordinate by a function of it, e.g. a logical one by a physical length.
`mapping_jacobian`	The Jacobian J[a, i] = ∂X_a / ∂η_i of the mapping, shape (3, 3, n1, n2, n3).
`periodicity`	How a direction of a (..., 3) array of physical points wraps around, if it does.
`save_scalars`	Save scalar time series to a CSV or NPZ file; under MPI only rank 0 writes it.
`scalar_names`	The names of the scalar time series to use from a collection.
`validate_array`	Check that data is an xarray.DataArray with the required dimensions.
`value_label`	The label of an array's values, with their units.

`DIM_LABELS`attributemodule attribute#

View source

DIM_LABELS = {
  't': '$t$',
  'eta1': '$\\eta_1$',
  'eta2': '$\\eta_2$',
  'eta3': '$\\eta_3$',
  'v1': '$v_1$',
  'v2': '$v_2$',
  'v3': '$v_3$',
  'x': '$x$',
  'y': '$y$',
  'z': '$z$',
  'R': '$R$',
  'Z': '$Z$',
  'component': 'component',
  'marker': 'marker',
  'quantity': 'quantity',
  'rho': '$\\rho$',
  'theta': '$\\theta$',
  'zeta': '$\\zeta$',
  'theta_B': '$\\theta_B$',
  'zeta_B': '$\\zeta_B$',
  'theta_P': '$\\theta_P$'
}

Show moreShow less

`LOGICAL_DIMS`attributemodule attribute#

View source

LOGICAL_DIMS = (
  ('eta1', 'eta2', 'eta3'),
  ('rho', 'theta', 'zeta'),
  ('rho', 'theta_B', 'zeta_B'),
  ('rho', 'theta_P', 'zeta')
)

`SCALARS_EXCLUDE`attributemodule attribute#

View source

SCALARS_EXCLUDE = ('time')

`angle_period`function#

View source

def angle_period(data: xr.DataArray | xr.Dataset, dim: str) -> float | None

The period of an angle coordinate, from its period attribute.

GVEC’s angles carry one (2π for the poloidal angles, 2π/nfp for the toroidal ones, see plasma_plots.gvec.from_gvec()); Struphy’s logical coordinates have none.

Parameters

Name	Type	Description
`data`	`xarray.DataArray or xarray.Dataset`	The field.
`dim`	`str`	One of its coordinates.

Returns

float or None: The period, or None if dim has no period attribute (or is no coordinate).

`axis_label`function#

View source

def axis_label(data: xr.DataArray, dim: str) -> str

The axis label of a dimension, with its units.

The label is the coordinate’s label attribute, else its long_name, else a default for the known dimensions (matplotlib mathtext such as $\eta_1$ for eta1), else the dimension name.

Parameters

Name	Type	Description
`data`	`xarray.DataArray`	The array.
`dim`	`str`	One of its dimensions.

Returns

str: The label, followed by [units] when the coordinate has a units attribute.

Raises

KeyError: If dim is not a dimension of data.

`close_periodic`function#

View source

def close_periodic(data: xr.DataArray, dims=None) -> xr.DataArray

Repeat the first slice at the end of every direction that wraps around in physical space.

Struphy evaluates fields at cell centers, which leave out the seam of a periodic direction (e.g. the poloidal angle), so a surface or pcolormesh drawn through the points has a gap there. A direction counts as periodic when its last slice lies about one grid step from its first one ("open" in periodicity()); a radius, or the ends of a torus sector, do not. The repeated slice gets the logical coordinate one step past the last.

Parameters

Name	Type	Default	Description
`data`	`xarray.DataArray`	required	The field, with the physical `X`, `Y`, `Z` coordinates; without them `data` comes back unchanged.
`dims`	`sequence of str`	`None`	The directions to close, if periodic. Default: those of the logical dimensions (see `logical_dims()`) that `data` has.

Returns

xarray.DataArray: data, one point longer along each periodic direction in dims.

`logical_derivative`function#

View source

def logical_derivative(values: np.ndarray, coordinate: np.ndarray, axis: int, kind: str | None) -> np.ndarray

The derivative ∂ values / ∂ coordinate along axis.

Around a periodic direction on a uniform grid (kind "open" or "closed", see periodicity()) the derivative is spectral (an FFT), which is exact for the smooth, periodic angles of Struphy’s mappings; the Nyquist mode of an even number of points is dropped. Elsewhere numpy’s second-order differences (one-sided at the ends; first order with only two points).

Parameters

Name	Type	Default	Description
`values`	`numpy.ndarray`	required	The values to differentiate.
`coordinate`	`numpy.ndarray`	required	The coordinate along `axis`, one value per point.
`axis`	`int`	required	The axis of `values` to differentiate along.
`kind`	`('open', 'closed', None)`	`"open"`	How the direction wraps around, from `periodicity()`: `"open"` (the seam is left out), `"closed"` (the last point repeats the first) or `None` (not periodic).

Returns

numpy.ndarray: The derivative, of the same shape as values.

`logical_dims`function#

View source

def logical_dims(data: xr.DataArray | xr.Dataset) -> tuple[str, str, str]

The names of the three logical dimensions of data, radial first.

One of LOGICAL_DIMS: the one with the most of its names among the dimensions of data (then among its coordinates, for a slice that selected some away), Struphy’s ("eta1", "eta2", "eta3") when none match.

Parameters

Name	Type	Description
`data`	`xarray.DataArray or xarray.Dataset`	The field.

Returns

tuple of str: Three dimension names, e.g. ("eta1", "eta2", "eta3") or ("rho", "theta_B", "zeta_B"); data need not have all three.

Examples

>>> logical_dims(phi)  # Struphy
('eta1', 'eta2', 'eta3')
>>> # GVEC, Boozer
>>> logical_dims(ev.mod_B.plasma.data.slice(x="zeta_B", y="theta_B", rho=0.5))
('rho', 'theta_B', 'zeta_B')

`map_coordinate`function#

View source

def map_coordinate(data: xr.DataArray, dim: str, mapping, *, name: str | None = None, units: str | None = None, label: str | None = None) -> xr.DataArray

Replace a coordinate by a function of it, e.g. a logical one by a physical length.

Every plot then draws and labels the new coordinate, e.g. the minor radius r = a1 + (a2 - a1) * eta1 in meters instead of eta1. With name, the dimension is renamed too, so r is what selections (r=0.5) and plots (x="r") name; without it, dim keeps its name and only its values, label and units change.

Parameters

Name	Type	Default	Description
`data`	`xarray.DataArray`	required	The array.
`dim`	`str`	required	The dimension whose coordinate is mapped, e.g. `"eta1"`.
`mapping`	`callable or float`	required	A function of the coordinate’s values (`lambda eta1: 0.1 + 0.9 * eta1`), or a factor (a length, for `length * eta1`).
`name`	`str`	`None`	The name of the new dimension, e.g. `"r"`. Default: keep `dim`.
`units`	`str`	`None`	The new coordinate’s units, e.g. `"m"`. Default: none.
`label`	`str`	`None`	The new coordinate’s axis label (its `long_name`), mathtext allowed. Default: `name`.

Returns

xarray.DataArray: The array over the new coordinate; the data itself is unchanged.

Raises

KeyError: If dim is not a dimension of data.

Examples

>>> map_coordinate(
...     T, "eta1", lambda eta1: 0.1 + 0.9 * eta1, name="r", units="m"
... ).plasma.plot.lineout(x="r", t=-1)
>>> # a length of 2π along eta1
>>> map_coordinate(n, "eta1", 2 * np.pi, units="m")

`mapping_jacobian`function#

View source

def mapping_jacobian(data: xr.DataArray) -> np.ndarray

The Jacobian J[a, i] = ∂X_a / ∂η_i of the mapping, shape (3, 3, n1, n2, n3).

Differentiated numerically from the X, Y, Z coordinates of data on its logical grid ((eta1, eta2, eta3), or GVEC’s, see logical_dims()): spectrally around periodic directions (see periodicity()), second order elsewhere (see logical_derivative()).

Parameters

Name	Type	Description
`data`	`xarray.DataArray`	An array with the three logical dimensions, each with at least two points, and the `X`, `Y`, `Z` coordinates over them.

Returns

numpy.ndarray: J of shape (3, 3, n1, n2, n3): the Cartesian component a first, the logical direction i second.

Raises

ValueError: If a logical dimension is missing or has fewer than two points (pass a struphy domain to the calling function instead).

`periodicity`function#

View source

def periodicity(points: np.ndarray, axis: int) -> str | None

How a direction of a (..., 3) array of physical points wraps around, if it does.

The direction is "closed" when its last slice coincides with its first (to 1e-9 of the median step), and "open" when the seam between them is at most 1.5 times the last grid step everywhere.

Parameters

Name	Type	Description
`points`	`numpy.ndarray`	Physical points, with the Cartesian coordinates `(X, Y, Z)` along the last axis.
`axis`	`int`	The axis of `points` to examine; it needs at least three points to be periodic.

Returns

{'closed', 'open', None}: "closed" (the last slice repeats the first), "open" (it stops one step short of the seam, as on Struphy’s cell centers), or None (not periodic, e.g. a radius or the ends of a torus sector).

`save_scalars`function#

View source

def save_scalars(scalars: xr.Dataset | Mapping, path: str, *, names=None, exclude=SCALARS_EXCLUDE, fmt=None) -> str

Save scalar time series to a CSV or NPZ file; under MPI only rank 0 writes it.

Parameters

Name	Type	Default	Description
`scalars`	`xarray.Dataset or mapping`	required	The scalars, e.g. `out.scalars`. Each selected one must have `t` as its only dimension, with identical time coordinates.
`path`	`str`	required	The file to write.
`names`	`sequence of str`	`None`	The scalars to save. Default: every one not in `exclude`; see `scalar_names()`.
`exclude`	`sequence of str`	`SCALARS_EXCLUDE`	Names left out when `names` is not given. Default: `("time",)`.
`fmt`	`('csv', 'npz')`	`"csv"`	The file format. Default: from the extension of `path`, else `"csv"`. A CSV has a header line `t,<names>` and one row per time; an NPZ has the arrays `t` and one per name.

Returns

str: path, on every rank.

Raises

ValueError: If a scalar has dimensions other than t, the time coordinates differ, or the format is unknown.
KeyError: If one of names is not in scalars.

Examples

>>> save_scalars(out.scalars, "scalars.csv", names=["en_E", "en_tot"])

`scalar_names`function#

View source

def scalar_names(scalars: xr.Dataset | Mapping, *, names=None, exclude=SCALARS_EXCLUDE) -> list[str]

The names of the scalar time series to use from a collection.

Parameters

Name	Type	Default	Description
`scalars`	`xarray.Dataset or mapping`	required	The scalars, e.g. `out.scalars`.
`names`	`sequence of str`	`None`	Explicit names, which must all be present; `exclude` then does not apply. Default: every variable not in `exclude`.
`exclude`	`sequence of str`	`SCALARS_EXCLUDE`	Names left out when `names` is not given. Default: `("time",)`.

Returns

list of str: The selected names, in order.

Raises

KeyError: If one of names is not in scalars.

`validate_array`function#

View source

def validate_array(data: xr.DataArray, *, required_dims: Sequence[str] = ()) -> xr.DataArray

Check that data is an xarray.DataArray with the required dimensions.

Parameters

Name	Type	Default	Description
`data`	`xarray.DataArray`	required	The array to check.
`required_dims`	`sequence of str`	`()`	Dimensions `data` must have. Default: none.

Returns

xarray.DataArray: data itself, unchanged.

Raises

TypeError: If data is not an xarray.DataArray.
ValueError: If one of required_dims is missing.

`value_label`function#

View source

def value_label(data: xr.DataArray) -> str

The label of an array’s values, with their units.

Parameters

Name	Type	Description
`data`	`xarray.DataArray`	The array. The label is its `label` attribute, else `long_name`, else its name.

Returns

str: "label [units]", with the units attribute or a.u.; only "[units]" when there is no label.

plasma_plots.arrays

DIM_LABELSattributemodule attribute#

LOGICAL_DIMSattributemodule attribute#

SCALARS_EXCLUDEattributemodule attribute#

angle_periodfunction#

axis_labelfunction#

close_periodicfunction#

logical_derivativefunction#

logical_dimsfunction#

map_coordinatefunction#

mapping_jacobianfunction#

periodicityfunction#

save_scalarsfunction#

scalar_namesfunction#

validate_arrayfunction#

value_labelfunction#

`DIM_LABELS`attributemodule attribute#

`LOGICAL_DIMS`attributemodule attribute#

`SCALARS_EXCLUDE`attributemodule attribute#

`angle_period`function#

`axis_label`function#

`close_periodic`function#

`logical_derivative`function#

`logical_dims`function#

`map_coordinate`function#

`mapping_jacobian`function#

`periodicity`function#

`save_scalars`function#

`scalar_names`function#

`validate_array`function#

`value_label`function#