# plasma_plots.arrays

*module*

Small xarray metadata helpers used by [`plasma_plots`][plasma_plots].

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

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L1-L1)

## plasma_plots.arrays.DIM_LABELS

*attribute* · *module attribute*

```python
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$'}
```

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L15-L37)

## plasma_plots.arrays.LOGICAL_DIMS

*attribute* · *module attribute*

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

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L43-L48)

## plasma_plots.arrays.SCALARS_EXCLUDE

*attribute* · *module attribute*

```python
SCALARS_EXCLUDE = ('time')
```

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L38-L38)

## plasma_plots.arrays.angle_period

*function*

```python
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()`][plasma_plots.gvec.from_gvec]); Struphy's logical coordinates have none.

**Parameters**

- `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).

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L87-L107)

## plasma_plots.arrays.axis_label

*function*

```python
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**

- `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``.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L144-L174)

## plasma_plots.arrays.close_periodic

*function*

```python
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()`][plasma_plots.arrays.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**

- `data` (`xarray.DataArray`) — The field, with the physical ``X``, ``Y``, ``Z`` coordinates; without them ``data`` comes back unchanged.
- `dims` (`sequence of str`) (default: `None`) — The directions to close, if periodic. Default: those of the logical dimensions (see [`logical_dims()`][plasma_plots.arrays.logical_dims]) that ``data`` has.

**Returns**

- (`xarray.DataArray`) — ``data``, one point longer along each periodic direction in ``dims``.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L413-L467)

## plasma_plots.arrays.logical_derivative

*function*

```python
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()`][plasma_plots.arrays.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**

- `values` (`numpy.ndarray`) — The values to differentiate.
- `coordinate` (`numpy.ndarray`) — The coordinate along ``axis``, one value per point.
- `axis` (`int`) — The axis of ``values`` to differentiate along.
- `kind` (`('open', 'closed', None)`) (default: `"open"`) — How the direction wraps around, from [`periodicity()`][plasma_plots.arrays.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``.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L470-L522)

## plasma_plots.arrays.logical_dims

*function*

```python
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`][plasma_plots.arrays.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**

- `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**

```pycon
>>> 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')
```

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L51-L84)

## plasma_plots.arrays.map_coordinate

*function*

```python
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**

- `data` (`xarray.DataArray`) — The array.
- `dim` (`str`) — The dimension whose coordinate is mapped, e.g. ``"eta1"``.
- `mapping` (`callable or float`) — A function of the coordinate's values (``lambda eta1: 0.1 + 0.9 * eta1``), or a factor (a length, for ``length * eta1``).
- `name` (`str`) (default: `None`) — The name of the new dimension, e.g. ``"r"``. Default: keep ``dim``.
- `units` (`str`) (default: `None`) — The new coordinate's units, e.g. ``"m"``. Default: none.
- `label` (`str`) (default: `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**

```pycon
>>> 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")
```

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L177-L239)

## plasma_plots.arrays.mapping_jacobian

*function*

```python
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()`][plasma_plots.arrays.logical_dims]): spectrally around periodic directions (see [`periodicity()`][plasma_plots.arrays.periodicity]),
second order elsewhere (see [`logical_derivative()`][plasma_plots.arrays.logical_derivative]).

**Parameters**

- `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).

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L525-L569)

## plasma_plots.arrays.periodicity

*function*

```python
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**

- `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).

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L376-L410)

## plasma_plots.arrays.save_scalars

*function*

```python
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**

- `scalars` (`xarray.Dataset or mapping`) — The scalars, e.g. ``out.scalars``. Each selected one must have ``t`` as its only dimension, with identical time coordinates.
- `path` (`str`) — The file to write.
- `names` (`sequence of str`) (default: `None`) — The scalars to save. Default: every one not in ``exclude``; see [`scalar_names()`][plasma_plots.arrays.scalar_names].
- `exclude` (`sequence of str`) (default: `SCALARS_EXCLUDE`) — Names left out when ``names`` is not given. Default: ``("time",)``.
- `fmt` (`('csv', 'npz')`) (default: `"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**

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

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L297-L373)

## plasma_plots.arrays.scalar_names

*function*

```python
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**

- `scalars` (`xarray.Dataset or mapping`) — The scalars, e.g. ``out.scalars``.
- `names` (`sequence of str`) (default: `None`) — Explicit names, which must all be present; ``exclude`` then does not apply. Default: every variable not in ``exclude``.
- `exclude` (`sequence of str`) (default: `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``.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L261-L294)

## plasma_plots.arrays.validate_array

*function*

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

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

**Parameters**

- `data` (`xarray.DataArray`) — The array to check.
- `required_dims` (`sequence of str`) (default: `()`) — 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.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L110-L141)

## plasma_plots.arrays.value_label

*function*

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

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

**Parameters**

- `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.

[View source](https://github.com/max-models/plasma-plots/blob/devel/src/plasma_plots/arrays.py#L242-L258)