From 34bc10c8ffe49ecb69986cdeeda1fae91702e430 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 10 Jun 2026 14:05:19 +0000 Subject: [PATCH 1/2] Optimize Coordinates construction and de-duplicate coordinate syntax docs Coordinates.__init__ parsed the rich input syntax on every construction (~7 us even for a single point), which dominated single-point evaluations such as eq.B_R(R, Z) (~half of the 25 us call) and the field-line tracer RHS, which builds three Coordinates objects per ODE evaluation. Performance: - Add Coordinates.from_coords() classmethod as the canonical rich-syntax factory; Coordinates.__init__ now takes a fast path for the common case of 1-3 positional scalars/1-D arrays in the default coordinate type and emits a DeprecationWarning for direct rich-syntax construction (named coordinates, (N, dim) arrays, non-default coord_type). - Cache cocos_coefs() with lru_cache and return a read-only mapping (it is a pure function of the COCOS index and was rebuilt per init). - Move the _valid_coordinates* sets to class attributes. - Clean up the slow parser (module-level Iterable import, single np.atleast_1d conversions, set comparison instead of permutation scans; 3D named input is now also converted with np.atleast_1d). - Migrate internal call sites (Equilibrium.coordinates, Surface/ FluxSurface, metis.py, examples, tests) to the warning-free paths. Measured: Coordinates(eq, 0.7, 0.1) 7 -> 2.1 us; kwargs construction 7 -> 3.0 us; eq.B_R(0.7, 0.1) 25 -> 20.7 us. Bulk array calls are spline-bound and unchanged. Documentation: - The coordinate input syntax was fully documented only in Coordinates.__init__ and duplicated in docs/source/coordinates.rst, while the ~44 coordinate-accepting Equilibrium methods documented it not at all (empty :param: stubs). The syntax now lives once in COORDINATES_DOC/COORD_PARAMS_DOC and a new append_to_doc decorator injects the full text into Coordinates.from_coords and Equilibrium.coordinates, and a short cross-referencing note into all coordinate-accepting Equilibrium methods; the empty :param: stubs are removed and coordinates.rst embeds the canonical docstring via automethod instead of a copy. Tests: add fast/slow path parity tests and deprecation tests; full suite 777 passed. Sphinx build warning count drops from 10 to 8 (remaining ones pre-date this change). https://claude.ai/code/session_01KvV66UPfmyCtKH71MthmWo --- docs/source/coordinates.rst | 36 +- examples/example_impact_angle.py | 2 +- ...ample_impact_angle_w_integrated_methods.py | 2 +- pleque/core/cocos.py | 10 +- pleque/core/coordinates.py | 323 +++++++++++++----- pleque/core/equilibrium.py | 201 ++--------- pleque/core/fluxsurface.py | 4 +- pleque/io/metis.py | 2 +- pleque/utils/decorators.py | 27 ++ tests/test_coordinates.py | 66 +++- tests/test_fluxexpansion.py | 6 +- 11 files changed, 374 insertions(+), 305 deletions(-) diff --git a/docs/source/coordinates.rst b/docs/source/coordinates.rst index 3528abe..7e97320 100644 --- a/docs/source/coordinates.rst +++ b/docs/source/coordinates.rst @@ -7,37 +7,13 @@ The tokamak is a curvilinear device which can be described in a number of coordi Accepted coordinates types -------------------------- -**1D - coordinates** +The accepted coordinate types and the full coordinate input syntax (shared by +``Coordinates.from_coords``, ``Equilibrium.coordinates`` and all +coordinate-accepting ``Equilibrium`` methods) are documented in a single +place: -+------------------------+-----------+------------------------------+ -| Coordinate | Code | Note | -+========================+===========+==============================+ -|:math:`\psi_\mathrm{N}` | ``psi_n`` | Default 1D coordinate | -+------------------------+-----------+------------------------------+ -|:math:`\psi` | ``psi`` | | -+------------------------+-----------+------------------------------+ -|:math:`\rho` | ``rho`` | :math:`\rho = \sqrt{\psi_n}` | -+------------------------+-----------+------------------------------+ - -**2D - coordinates** - -+------------------------+--------------+-------------------------------------------------+ -| Coordinate | Code | Note | -+========================+==============+=================================================+ -|:math:`(R, Z)` | ``R, Z`` | Default 2D coordinate | -+------------------------+--------------+-------------------------------------------------+ -|:math:`(r, \theta)` | ``r, theta`` | Polar coordinates with respect to magnetic axis | -+------------------------+--------------+-------------------------------------------------+ - -**3D - coordinates** - -+------------------------+---------------+-------------------------------------------------+ -| Coordinate | Code | Note | -+========================+===============+=================================================+ -|:math:`(R, Z, \phi)` | ``R, Z, phi`` | Default 3D coordinate | -+------------------------+---------------+-------------------------------------------------+ -|:math:`(X, Y, Z)` | ``X, Y, Z`` | | -+------------------------+---------------+-------------------------------------------------+ +.. automethod:: pleque.core.coordinates.Coordinates.from_coords + :noindex: Array shape convention ---------------------- diff --git a/examples/example_impact_angle.py b/examples/example_impact_angle.py index ad4b8d3..31880f6 100644 --- a/examples/example_impact_angle.py +++ b/examples/example_impact_angle.py @@ -27,7 +27,7 @@ axn.plot(first_wall.R,first_wall.Z) -coords=Coordinates(eq,np.vstack((first_wall.R,first_wall.Z)).T) +coords=Coordinates.from_coords(eq,np.vstack((first_wall.R,first_wall.Z)).T) def interpolate_lim(coords, npoints): diff --git a/examples/example_impact_angle_w_integrated_methods.py b/examples/example_impact_angle_w_integrated_methods.py index 36f18f3..cd6f050 100644 --- a/examples/example_impact_angle_w_integrated_methods.py +++ b/examples/example_impact_angle_w_integrated_methods.py @@ -26,7 +26,7 @@ -coords=Coordinates(eq,np.vstack((first_wall.R,first_wall.Z)).T) +coords=Coordinates.from_coords(eq,np.vstack((first_wall.R,first_wall.Z)).T) # Resample using new method coords2=coords.resample2(npoints) diff --git a/pleque/core/cocos.py b/pleque/core/cocos.py index dec7534..f476c05 100644 --- a/pleque/core/cocos.py +++ b/pleque/core/cocos.py @@ -1,3 +1,5 @@ +import functools +import types from typing import TypedDict import numpy as np @@ -15,14 +17,18 @@ CocosInfo = TypedDict('CocosInfo', {'exp_Bp': int, 'sigma_Bp': int, 'sigma_cyl': int, 'sigma_pol': int, 'sign_q': int, 'sign_pprime': int}) +@functools.lru_cache(maxsize=32) def cocos_coefs(cocos_idx) -> CocosInfo: """ Define COCOS coefficients. For more details see Table 1 in O. Sauter, et al., Comp. Phys. Comm. 184 (2013), p. 296 + The result is cached and shared between callers; it is returned as a + read-only mapping. + :param cocos_idx: int, COCOS index - :return: dict with COCOS coefficients + :return: read-only dict with COCOS coefficients """ # Note: sign_q and sign_pprime assumes Ip and B0 to be in direction of a toroidal coordinate. @@ -60,5 +66,5 @@ def cocos_coefs(cocos_idx) -> CocosInfo: cocos["sigma_pol"] = -1 cocos["sign_q"] = -1 - return cocos + return types.MappingProxyType(cocos) diff --git a/pleque/core/coordinates.py b/pleque/core/coordinates.py index 398ed7c..4a93117 100644 --- a/pleque/core/coordinates.py +++ b/pleque/core/coordinates.py @@ -1,83 +1,193 @@ -from collections.abc import Sequence -import itertools +from collections.abc import Iterable +import warnings from typing import Union import numpy as np import xarray -from pleque.utils.decorators import deprecated +from pleque.utils.decorators import append_to_doc, deprecated import pleque.utils.flux_expansions as flux_expansion from .cocos import cocos_coefs from scipy.interpolate import splprep, splev -class Coordinates(object): - - def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, cocos=None, **coords): - r""" - Basic PLEQUE class to handle various coordinate systems in tokamak equilibrium. - - :param equilibrium: - :param *coordinates: * Can be skipped. - * ``array (N, dim)`` - ``N`` points will be generated. - * One, two are three comma separated one dimensional arrays. - :param coord_type: - :param grid: - :param cocos: Define coordinate system cocos. Id `None` equilibrium default cocos is used. - If `equilibrium is None` cocos = 3 (both systems cnt-clockwise) is used. - :param **coords: Lorem ipsum. - - - Default coordinate systems - -------------------------- - - - **1D**: :math:`\psi_\mathrm{N}`, - - **2D**: :math:`(R, Z)`, - - **3D**: :math:`(R, Z, \phi)`. - - Accepted coordinates types - -------------------------- - - **1D - coordinates** - - +------------------------+-----------+------------------------------+ - | Coordinate | Code | Note | - +========================+===========+==============================+ - |:math:`\psi_\mathrm{N}` | ``psi_n`` | Default 1D coordinate | - +------------------------+-----------+------------------------------+ - |:math:`\psi` | ``psi`` | | - +------------------------+-----------+------------------------------+ - |:math:`\rho` | ``rho`` | :math:`\rho = \sqrt{\psi_n}` | - +------------------------+-----------+------------------------------+ +#: Canonical description of the coordinate input syntax accepted across PLEQUE. +#: It is appended (via :func:`pleque.utils.decorators.append_to_doc`) to the +#: docstrings of :meth:`Coordinates.from_coords` and +#: :meth:`pleque.core.Equilibrium.coordinates`, so it is written only once. +COORDINATES_DOC = r""" +**Coordinate input syntax** + +Coordinates may be specified positionally or by name: + +* nothing -- an empty (``dim = 0``) object is created, +* a single ``array (N, dim)`` -- ``N`` points will be generated, +* one, two or three comma separated one dimensional arrays or scalars, +* named coordinates: ``R=, Z=, psi_n=, psi=, rho=, r=, theta=, phi=, X=, Y=``. + +**Default coordinate systems** + +- **1D**: :math:`\psi_\mathrm{N}`, +- **2D**: :math:`(R, Z)`, +- **3D**: :math:`(R, Z, \phi)`. + +**Accepted coordinates types** + +*1D - coordinates* + ++------------------------+-----------+------------------------------+ +| Coordinate | Code | Note | ++========================+===========+==============================+ +|:math:`\psi_\mathrm{N}` | ``psi_n`` | Default 1D coordinate | ++------------------------+-----------+------------------------------+ +|:math:`\psi` | ``psi`` | | ++------------------------+-----------+------------------------------+ +|:math:`\rho` | ``rho`` | :math:`\rho = \sqrt{\psi_n}` | ++------------------------+-----------+------------------------------+ + +*2D - coordinates* + ++------------------------+--------------+-------------------------------------------------+ +| Coordinate | Code | Note | ++========================+==============+=================================================+ +|:math:`(R, Z)` | ``R, Z`` | Default 2D coordinate | ++------------------------+--------------+-------------------------------------------------+ +|:math:`(r, \theta)` | ``r, theta`` | Polar coordinates with respect to magnetic axis | ++------------------------+--------------+-------------------------------------------------+ + +*3D - coordinates* + ++------------------------+---------------+-------------------------------------------------+ +| Coordinate | Code | Note | ++========================+===============+=================================================+ +|:math:`(R, Z, \phi)` | ``R, Z, phi`` | Default 3D coordinate | ++------------------------+---------------+-------------------------------------------------+ +|:math:`(X, Y, Z)` | ``(X, Y, Z)`` | Cartesian coordinates | ++------------------------+---------------+-------------------------------------------------+ +""" + +#: Short note cross-referencing the canonical coordinate syntax description. +#: Appended to the docstrings of all coordinate-accepting ``Equilibrium`` +#: methods. +COORD_PARAMS_DOC = """ +.. note:: This method accepts PLEQUE's rich coordinate input: an existing + :class:`~pleque.core.coordinates.Coordinates` instance, positional + values (``R, Z`` by default for 2D input), or coordinates passed by name + (``R=``, ``Z=``, ``psi_n=``, ``psi=``, ``rho=``, ...), together with the + ``coord_type`` and ``grid`` options. See + :meth:`pleque.core.coordinates.Coordinates.from_coords` for the full + description of the syntax. +""" - **2D - coordintares** - +------------------------+--------------+-------------------------------------------------+ - | Coordinate | Code | Note | - +========================+==============+=================================================+ - |:math:`(R, Z)` | ``R, Z`` | Default 2D coordinate | - +------------------------+--------------+-------------------------------------------------+ - |:math:`(r, \theta)` | ``r, theta`` | Polar coordinates with respect to magnetic axis | - +------------------------+--------------+-------------------------------------------------+ +class Coordinates(object): + r""" + Basic PLEQUE class to handle various coordinate systems in tokamak + equilibrium. - **3D - coordinates** + Preferably created with :meth:`Coordinates.from_coords` or + :meth:`pleque.core.Equilibrium.coordinates`, which accept the full + coordinate input syntax described in :meth:`Coordinates.from_coords`. + """ - +------------------------+---------------+-------------------------------------------------+ - | Coordinate | Code | Note | - +========================+===============+=================================================+ - |:math:`(R, Z, \phi)` | ``R, Z, phi`` | Default 3D coordinate | - +------------------------+---------------+-------------------------------------------------+ - |:math:`(X, Y, Z)` | ``(X, Y, Z)`` | Polar coordinates with respect to magnetic axis | - +------------------------+---------------+-------------------------------------------------+ + # Sets of recognised coordinate names, shared by all instances: + _valid_coordinates = {'R', 'Z', 'psi_n', 'psi', 'rho', 'r', 'theta', 'phi', 'X', 'Y'} + _valid_coordinates_1d = {('psi_n',), ('psi',), ('rho',)} + _valid_coordinates_2d = {('R', 'Z'), ('r', 'theta')} + _valid_coordinates_3d = {('R', 'Z', 'phi'), ('X', 'Y', 'Z')} + # Default coordinate type for a given dimension: + _default_coord_types = {1: ('psi_n',), 2: ('R', 'Z'), 3: ('R', 'Z', 'phi')} - """ + def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, cocos=None, **coords): + r""" + Create Coordinates from positional scalars or 1-D arrays in the + default coordinate type of the given dimension (``psi_n``; ``R, Z``; + ``R, Z, phi``). + + .. deprecated:: 0.0.11 + Constructing ``Coordinates`` directly from the rich coordinate + input (named coordinates such as ``R=``/``psi_n=``, + ``array (N, dim)`` input or a non-default ``coord_type``) is + deprecated; use :meth:`Coordinates.from_coords` or + :meth:`pleque.core.Equilibrium.coordinates` instead. + + :param equilibrium: Instance of ``Equilibrium`` or ``None``. + :param coordinates: One, two or three comma separated scalars or + one dimensional arrays. + :param coord_type: Tuple naming the input coordinates. Only the + default coordinate type of the corresponding + dimension is accepted without deprecation warning. + :param grid: If ``True``, the coordinates span a rectangular grid + (2D only). + :param cocos: Define coordinate system cocos. If `None` equilibrium default cocos is used. + If `equilibrium is None` cocos = 3 (both systems cnt-clockwise) is used. + :param coords: Deprecated here; coordinates passed by name. Use + :meth:`Coordinates.from_coords` instead. + """ + if self._is_rich_input(coordinates, coord_type, coords): + warnings.warn( + 'Constructing Coordinates directly from the rich coordinate input ' + '(named coordinates, array (N, dim) input or a non-default coord_type) ' + 'is deprecated; use Coordinates.from_coords() or ' + 'Equilibrium.coordinates() instead.', + DeprecationWarning, + stacklevel=2, + ) + + self._init_from_input(equilibrium, coordinates, coord_type, grid, cocos, coords) + + @classmethod + @append_to_doc(COORDINATES_DOC) + def from_coords(cls, equilibrium, *coordinates, coord_type=None, grid=False, cocos=None, **coords): + r""" + Create :class:`Coordinates` from any supported coordinate + specification. + + This is the preferred way of constructing :class:`Coordinates` and + the canonical description of the coordinate input syntax accepted + across PLEQUE (for instance by most + :class:`pleque.core.Equilibrium` methods). + + :param equilibrium: Instance of ``Equilibrium`` or ``None``. + :param coordinates: Positional coordinates; see below. + :param coord_type: Tuple naming the input coordinates, + e.g. ``('rho',)`` or ``('Z', 'R')``. + :param grid: If ``True``, the coordinates span a rectangular grid + (2D only). + :param cocos: Define coordinate system cocos. If `None` equilibrium default cocos is used. + If `equilibrium is None` cocos = 3 (both systems cnt-clockwise) is used. + :param coords: Coordinates passed by name; see below. + :return: Instance of :class:`Coordinates`. + """ + if cls is Coordinates: + self = object.__new__(cls) + self._init_from_input(equilibrium, coordinates, coord_type, grid, cocos, coords) + return self + # Subclasses (e.g. Surface) run additional logic in __init__: + return cls(equilibrium, *coordinates, coord_type=coord_type, grid=grid, cocos=cocos, **coords) + + @classmethod + def _is_rich_input(cls, coordinates, coord_type, coords): + """Return True when the input requires the full (deprecated in + ``__init__``) rich-syntax parser instead of positional values in the + default coordinate type.""" + if any(v is not None for v in coords.values()): + return True + n = len(coordinates) + if n == 0: + return False + if coord_type is not None: + if isinstance(coord_type, str): + coord_type = (coord_type,) + if tuple(coord_type) != cls._default_coord_types.get(n): + return True + return n == 1 and np.ndim(coordinates[0]) > 1 + def _init_from_input(self, equilibrium, coordinates, coord_type, grid, cocos, coords): + """Initialise the instance from any supported coordinate + specification (without any deprecation warning).""" self._eq = equilibrium - self._valid_coordinates = {'R', 'Z', 'psi_n', 'psi', 'rho', 'r', 'theta', 'phi', 'X', 'Y'} - self._valid_coordinates_1d = {('psi_n',), ('psi',), ('rho',)} - self._valid_coordinates_2d = {('R', 'Z'), ('r', 'theta')} - self._valid_coordinates_3d = {('R', 'Z', 'phi'), ('X', 'Y', 'Z')} self.dim = -1 # init only self.grid = grid @@ -91,8 +201,56 @@ def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, cocos self.cocos_dict = cocos_coefs(self.cocos) + if self._init_fast_path(coordinates, coord_type, coords): + return + self._evaluate_input(*coordinates, coord_type=coord_type, **coords) + def _init_fast_path(self, coordinates, coord_type, coords): + """ + Set the coordinate attributes directly for the most common input: + one to three positional scalars or 1-D arrays already in the default + coordinate type. Return True on success; on False the caller falls + back to the full input parser. + """ + n = len(coordinates) + if not 1 <= n <= 3: + return False + if coords and not all(v is None for v in coords.values()): + return False + default_type = self._default_coord_types[n] + if coord_type is not None: + if isinstance(coord_type, str): + coord_type = (coord_type,) + if tuple(coord_type) != default_type: + return False + + xs = [] + for c in coordinates: + if isinstance(c, np.ndarray): + if c.ndim != 1 or not np.issubdtype(c.dtype, np.number): + return False + xs.append(c) + elif isinstance(c, (float, int, np.floating, np.integer)): + xs.append(np.array([c])) + else: + return False + + self.dim = n + self._coord_type_input = default_type + self.x1 = self._x1_input = xs[0] + if n >= 2: + self.x2 = self._x2_input = xs[1] + if n == 3: + self.x3 = self._x3_input = xs[2] + + if self.grid and n != 2: + print('WARNING: grid == True is not allowed for dim != 2 (yet).' + 'Turning grid = False.') + self.grid = False + + return True + def __iter__(self): if self.grid: raise TypeError('Grid is not iterable at the moment.') @@ -542,8 +700,6 @@ def length(self): return self._cum_length[-1] def _evaluate_input(self, *coordinates, coord_type=None, **coords): - from collections.abc import Iterable - if len(coordinates) == 0: # todo: self.dim = 0 @@ -582,26 +738,16 @@ def _evaluate_input(self, *coordinates, coord_type=None, **coords): else: raise ValueError('Invalid combination of input coordinates.') elif self.dim == 3: - # if tuple(xy_name) in self._valid_coordinates_3d: - permutations = list(itertools.permutations(xy_name)) - # if any([p in self._valid_coordinates_3d for p in permutations]): - # - # # todo: implement various order of coordinates - # self._x1_input = xy[0] - # self._x2_input = xy[1] - # self._x3_input = xy[2] - # coord_type_ = tuple(xy_name) # todo: make function and use for all - valid = list(self._valid_coordinates_3d) # find the index of the valid coordinate system with known coordinate order ii = [set(item) for item in valid].index(set(xy_name)) actual = valid[ii] - self._x1_input = coords[actual[0]] - self._x2_input = coords[actual[1]] - self._x3_input = coords[actual[2]] + self._x1_input = np.atleast_1d(coords[actual[0]]) + self._x2_input = np.atleast_1d(coords[actual[1]]) + self._x3_input = np.atleast_1d(coords[actual[2]]) coord_type_ = tuple(actual) else: # self._incompatible_dimension_error(self.dim) @@ -656,20 +802,11 @@ def _evaluate_input(self, *coordinates, coord_type=None, **coords): self._x2_input = x2 elif len(coordinates) == 3: self.dim = 3 - x1 = np.atleast_1d(coordinates[0]) - x2 = np.atleast_1d(coordinates[1]) - x3 = np.atleast_1d(coordinates[2]) - # assume _x1_input and _x2_input to be arrays of size (N) - if not isinstance(x1, np.ndarray): - x1 = np.array(x1, ndmin=1) - if not isinstance(x2, np.ndarray): - x2 = np.array(x2, ndmin=1) - if not isinstance(x3, np.ndarray): - x3 = np.array(x3, ndmin=1) - self._x1_input = x1 - self._x2_input = x2 - self._x3_input = x3 + # assume _x1_input, _x2_input and _x3_input to be arrays of size (N) + self._x1_input = np.atleast_1d(coordinates[0]) + self._x2_input = np.atleast_1d(coordinates[1]) + self._x3_input = np.atleast_1d(coordinates[2]) else: self._incompatible_dimension_error(len(coordinates)) @@ -769,13 +906,11 @@ def _convert_to_default_coord_type(self): elif self.dim == 3: # only (R, Z) coordinates are implemented now - # if self._coord_type_input == ('R', 'Z', 'phi'): - if any([p == ('R', 'Z', 'phi') for p in itertools.permutations(self._coord_type_input)]): + if set(self._coord_type_input) == {'R', 'Z', 'phi'}: self.x1 = self._x1_input self.x2 = self._x2_input self.x3 = self._x3_input - # elif self._coord_type_input == ('X', 'Y', 'Z'): - elif any([p == ('X', 'Y', 'Z') for p in itertools.permutations(self._coord_type_input)]): + elif set(self._coord_type_input) == {'X', 'Y', 'Z'}: # todo: COCOS # R(1)**2 = X(1)**2 + Y(2)**2 # Z(2) = Z(3) diff --git a/pleque/core/equilibrium.py b/pleque/core/equilibrium.py index ccadd11..08f8992 100644 --- a/pleque/core/equilibrium.py +++ b/pleque/core/equilibrium.py @@ -8,10 +8,12 @@ from shapely import Polygon, Point import pleque -from pleque.utils.decorators import deprecated, ordered_path_scalar_function, scalar_function, vector_function +from pleque.utils.decorators import append_to_doc, deprecated, ordered_path_scalar_function, scalar_function, \ + vector_function from scipy.interpolate import RectBivariateSpline, UnivariateSpline from pleque.core import Coordinates +from pleque.core.coordinates import COORD_PARAMS_DOC, COORDINATES_DOC from pleque.utils.tools import arglis, xp_sections from pleque.core import FluxFunctions, Surface # , FluxSurface from pleque.core import SurfaceFunctions @@ -410,13 +412,6 @@ def psi(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=Tr """ Psi value - :param psi_n: - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ coord = self.coordinates(*coordinates, R=R, Z=Z, psi_n=psi_n, coord_type=coord_type, grid=grid, **coords) @@ -444,13 +439,6 @@ def diff_psi(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, gr r""" Return the value of :math:`\pm|\nabla \psi|`. It is positive/negative if the :math:`\psi` is increasing/decreasing. - :param coordinates: - :param R: - :param Z: - :param psi_n: - :param coord_type: - :param grid: - :param coords: :return: """ coord = self.coordinates(*coordinates, R=R, Z=Z, psi_n=psi_n, coord_type=coord_type, grid=grid, **coords) @@ -514,13 +502,6 @@ def F(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=True def Fprime(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=True, **coords): ''' - :param coordinates: - :param R: - :param Z: - :param psi_n: - :param coord_type: - :param grid: - :param coords: :return: ''' @@ -549,12 +530,6 @@ def B_abs(self, *coordinates, R=None, Z=None, coord_type=None, grid=True, **coor """ Absolute value of magnetic field in Tesla. - :param grid: - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param coords: :return: Absolute value of magnetic field in Tesla. """ coord = self.coordinates(*coordinates, R=R, Z=Z, coord_type=coord_type, grid=grid, **coords) @@ -569,14 +544,8 @@ def B_abs(self, *coordinates, R=None, Z=None, coord_type=None, grid=True, **coor def Bvec(self, *coordinates, swap_order=False, R=None, Z=None, coord_type=None, grid=True, **coords): """ Magnetic field vector - :param grid: - :param coordinates: :param swap_order: If False, return component-first shape ``(3, ...)``. If True, move the component axis to the end for compatibility. - :param R: - :param Z: - :param coord_type: - :param coords: :return: Magnetic field vector array. Component-first shape is ``(3, n_elements)`` for paired points and ``(3, n_z, n_r)`` for grids. """ @@ -598,14 +567,8 @@ def Bvec(self, *coordinates, swap_order=False, R=None, Z=None, coord_type=None, def Bvec_norm(self, *coordinates, swap_order=False, R=None, Z=None, coord_type=None, grid=True, **coords): """ Magnetic field vector, normalised - :param grid: - :param coordinates: :param swap_order: If False, return component-first shape ``(3, ...)``. If True, move the component axis to the end for compatibility. - :param R: - :param Z: - :param coord_type: - :param coords: :return: Normalised magnetic field vector array. Component-first shape is ``(3, n_elements)`` for paired points and ``(3, n_z, n_r)`` for grids. """ @@ -642,10 +605,6 @@ def _flux_surface(self, *coordinates, resolution=None, dim="step", fluxsurface properties as if it is inside last closed flux surface or if the surface is supposed to be closed are possible. - :param R: - :param Z: - :param psi_n: - :param coord_type: :param coordinates: specifies flux surface to search for (by spatial point or values of psi or psi normalised). If coordinates is spatial point (dim=2) then parameters closed and lcfs are automatically overridden. Coordinates.grid must be False. @@ -739,12 +698,6 @@ def poloidal_mag_flux_exp_coef(self, *coordinates, R=None, Z=None, coord_type=No *Poloidal magnetic flux expansion coefficient* is typically used for :math:`\lambda` scaling in plane perpendicular to the poloidal component of the magnetic field. - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -773,12 +726,6 @@ def effective_poloidal_mag_flux_exp_coef(self, *coordinates, R=None, Z=None, coo .. math:: \lambda^\mathrm{t} = \lambda^\mathrm{u} f_{\mathrm{pol, eff}} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -803,12 +750,6 @@ def poloidal_heat_flux_exp_coef(self, *coordinates, R=None, Z=None, coord_type=N .. math:: q_\theta^\mathrm{t} = \frac{q_\theta^\mathrm{u}}{f_{\mathrm{pol, heat}}} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -837,12 +778,6 @@ def effective_poloidal_heat_flux_exp_coef(self, *coordinates, R=None, Z=None, co .. math:: q_\perp^\mathrm{t} = \frac{q_\theta^\mathrm{u}}{f_{\mathrm{pol, heat, eff}}} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -868,12 +803,6 @@ def parallel_heat_flux_exp_coef(self, *coordinates, R=None, Z=None, coord_type=N .. math:: q_\parallel^\mathrm{t} = \frac{q_\parallel^\mathrm{u}}{f_\parallel} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -902,12 +831,6 @@ def total_heat_flux_exp_coef(self, *coordinates, R=None, Z=None, coord_type=None .. math:: q_\perp^\mathrm{t} = \frac{q_\parallel^\mathrm{u}}{f_{\mathrm{tot}}} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -959,7 +882,7 @@ def _get_surface(self, *coordinates, R=None, Z=None, level=0.5, norm=True, coord contour = find_contour(coordinates.psi, level=level, r=coordinates.R, z=coordinates.Z) for i in range(len(contour)): - contour[i] = Coordinates(self, contour[i]) + contour[i] = Coordinates.from_coords(self, contour[i]) return contour @@ -1260,12 +1183,6 @@ def B_R(self, *coordinates, R=None, Z=None, coord_type=('R', 'Z'), grid=True, ** """ Poloidal value of magnetic field in Tesla. - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: Scalar array with shape ``(n_elements,)`` for paired points and ``(n_z, n_r)`` for grids. """ @@ -1284,12 +1201,6 @@ def B_Z(self, *coordinates, R=None, Z=None, coord_type=None, grid=True, **coords """ Poloidal value of magnetic field in Tesla. - :param grid: - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param coords: :return: Scalar array with shape ``(n_elements,)`` for paired points and ``(n_z, n_r)`` for grids. """ @@ -1309,12 +1220,6 @@ def B_pol(self, *coordinates, R=None, Z=None, coord_type=None, grid=True, **coor """ Absolute value of magnetic field in Tesla. - :param grid: - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param coords: :return: """ coord = self.coordinates(*coordinates, R=R, Z=Z, coord_type=coord_type, grid=grid, **coords) @@ -1328,12 +1233,6 @@ def B_tor(self, *coordinates, R: np.array = None, Z: np.array = None, coord_type """ Toroidal value of magnetic field in Tesla. - :param grid: - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param coords: :return: """ coord = self.coordinates(*coordinates, R=R, Z=Z, coord_type=coord_type, grid=grid, **coords) @@ -1354,12 +1253,6 @@ def abs_q(self, *coordinates, R: np.array = None, Z: np.array = None, coord_type """ Absolute value of q. - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ return np.abs(self.q(*coordinates, R=R, Z=Z, coord_type=coord_type, grid=grid, **coords)) @@ -1376,13 +1269,6 @@ def diff_q(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid """ :param self: - :param coordinates: - :param R: - :param Z: - :param psi_n: - :param coord_type: - :param grid: - :param coords: :return: Derivative of q with respect to psi. """ if not hasattr(self, '_dq_dpsin_spl'): @@ -1420,12 +1306,6 @@ def pol_flux(self, *coordinates, R: np.array = None, Z: np.array = None, coord_t .. math:: \psi_\mathrm{pol} = (2 \pi)^{1 - e_{Bp}} \psi_\mathrm{ref} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -1440,12 +1320,6 @@ def tor_flux(self, *coordinates, R: np.array = None, Z: np.array = None, coord_t .. math:: q = \frac{\mathrm{d \Phi} }{\mathrm{d \psi}} - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ @@ -1495,12 +1369,6 @@ def j_pol(self, *coordinates, R: np.array = None, Z: np.array = None, coord_type [Wesson: Tokamaks, p. 105] - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ from scipy.constants import mu_0 @@ -1518,12 +1386,6 @@ def j_tor(self, *coordinates, R: np.array = None, Z: np.array = None, coord_type .. math:: R p' + \frac{1}{\mu_0 R} ff' - :param coordinates: - :param R: - :param Z: - :param coord_type: - :param grid: - :param coords: :return: """ from scipy.constants import mu_0 @@ -1720,28 +1582,25 @@ def xp_section(self, length: float = 0.15) -> tuple["Coordinates"]: secs_coords = tuple((self.coordinates(sec) for sec in secs)) return secs_coords + @append_to_doc(COORDINATES_DOC) def coordinates(self, *coordinates, coord_type=None, grid=False, **coords): """ Return instance of Coordinates. If instances of coordinates is already on the input, just pass it through. - :param coordinates: - :param coord_type: - :param grid: - :param coords: - :return: + :param coordinates: Positional coordinates; see below. + :param coord_type: Tuple naming the input coordinates, e.g. ``('rho',)`` or ``('Z', 'R')``. + :param grid: If ``True``, the coordinates span a rectangular grid (2D only). + :param coords: Coordinates passed by name; see below. + :return: Instance of :class:`pleque.core.Coordinates`. """ if len(coordinates) >= 1 and isinstance(coordinates[0], Coordinates): return coordinates[0] else: - return Coordinates(self, *coordinates, coord_type=coord_type, grid=grid, **coords) + return Coordinates.from_coords(self, *coordinates, coord_type=coord_type, grid=grid, **coords) def _as_fluxsurface(self, *coordinates, coord_type=None, grid=False, **coords): """ - :param coordinates: - :param coord_type: - :param grid: - :param coords: :return: """ from pleque import FluxSurface @@ -1790,14 +1649,9 @@ def connection_length(self, *coordinates, R: np.array = None, Z: np.array = None Todo: The field line is traced to min/max value of z of first wall, distance is calculated to the last point before first wall. - :param coordinates: - :param R: - :param Z: - :param coord_type: :param direction: if positive trace field line in/cons the direction of magnetic field. :param stopper: (None, 'poloidal', 'z-stopper) force to use stopper. If None stopper is automatically chosen based on psi_n coordinate. - :param coords: :return: """ coords = self.coordinates(*coordinates, R=R, Z=Z, coord_type=coord_type, **coords) @@ -1831,15 +1685,10 @@ def trace_field_line(self, *coordinates, R: np.array = None, Z: np.array = None, One poloidal turn is calculated for field lines inside the separatrix. Outter field lines are limited by z planes given be outermost z coordinates of the first wall. - :param coordinates: - :param R: - :param Z: - :param coord_type: :param direction: if positive trace field line in/cons the direction of magnetic field. :param stopper_method: (None, 'poloidal', 'z-stopper) force to use stopper. If None stopper is automatically chosen based on psi_n coordinate. :param in_first_wall: if True the only inner part of field line is returned. - :param coords: :return: """ @@ -2012,10 +1861,6 @@ def trace_flux_surface(self, *coordinates, s_resolution=1e-3, R=None, TODO support open and/or flux surfaces outise LCFS, needs different stopper - :param R: - :param Z: - :param psi_n: - :param coord_type: :param coordinates: specifies flux surface to search for (by spatial point or values of psi or psi normalised). If coordinates is spatial point (dim=2) then the trace starts at the midplane. Coordinates.grid must be False. @@ -2163,3 +2008,27 @@ def _init_q(self): self._q_spl = UnivariateSpline(psi_n, qs, s=0, k=3) self._dq_dpsin_spl = self._q_spl.derivative() self._q_anideriv_spl = self._q_spl.antiderivative() + + +# All public Equilibrium methods accepting the rich coordinate input +# (see ``Coordinates.from_coords``). A note cross-referencing the canonical +# description of the input syntax is appended to their docstrings, so the +# syntax itself is documented in a single place only. +_COORDINATE_INPUT_METHODS = [ + 'psi', 'nabla_psi', 'diff_psi', 'psi_n', 'r_mid', 'rho', + 'pressure', 'pprime', 'f', 'F', 'Fprime', 'ffprime', 'FFprime', + 'B_abs', 'Bvec', 'Bvec_norm', 'B_R', 'B_Z', 'B_pol', 'B_tor', + 'abs_q', 'q', 'diff_q', 'shear', 'pol_flux', 'tor_flux', + 'j_R', 'j_Z', 'j_pol', 'j_tor', + 'flux_surface', + 'poloidal_mag_flux_exp_coef', 'effective_poloidal_mag_flux_exp_coef', + 'poloidal_heat_flux_exp_coef', 'effective_poloidal_heat_flux_exp_coef', + 'parallel_heat_flux_exp_coef', 'total_heat_flux_exp_coef', + 'outter_parallel_fl_expansion_coef', 'outter_poloidal_fl_expansion_coef', + 'in_first_wall', 'in_lcfs', + 'connection_length', 'trace_field_line', 'trace_flux_surface', +] + +for _name in _COORDINATE_INPUT_METHODS: + append_to_doc(COORD_PARAMS_DOC)(getattr(Equilibrium, _name)) +del _name diff --git a/pleque/core/fluxsurface.py b/pleque/core/fluxsurface.py index 335c0ea..c011186 100644 --- a/pleque/core/fluxsurface.py +++ b/pleque/core/fluxsurface.py @@ -7,7 +7,7 @@ class Surface(Coordinates): - def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, **coords): + def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, cocos=None, **coords): """ Calculates geometrical properties of a specified surface. To make the contour closed, the first and last points in the passed coordinates have to be the same. @@ -16,7 +16,7 @@ def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, **coo :param coords: Instance of coordinate class """ - super().__init__(equilibrium, *coordinates, coord_type=None, grid=False, **coords) + self._init_from_input(equilibrium, coordinates, None, False, cocos, coords) points_RZ = self.as_array(('R', 'Z')) # closed surface has to have identical first and last points and then the shape is polygon diff --git a/pleque/io/metis.py b/pleque/io/metis.py index a476d20..13fa90e 100644 --- a/pleque/io/metis.py +++ b/pleque/io/metis.py @@ -44,7 +44,7 @@ def read(equilibrium, file, time): psi_sep = toexp.psi.values[np.argmax(np.abs(toexp.psi.values))] psi_n = np.abs((toexp.psi.values - psi_axis) / (psi_sep - psi_axis)) - crds = Coordinates(equilibrium=equilibrium, psi_n=psi_n) + crds = Coordinates.from_coords(equilibrium, psi_n=psi_n) for name in list(toexp.variables.keys()): #todo: better handling of complex values diff --git a/pleque/utils/decorators.py b/pleque/utils/decorators.py index 4f0baf7..0f72a4c 100644 --- a/pleque/utils/decorators.py +++ b/pleque/utils/decorators.py @@ -78,6 +78,33 @@ def new_func2(*args, **kwargs): raise TypeError(repr(type(reason))) +def append_to_doc(*snippets): + """ + Append shared documentation snippets to the docstring of the decorated + function (or of a function passed to the returned decorator). + + The original docstring is dedented with :func:`inspect.getdoc`, so the + appended snippets can be written flush-left and the result still renders + correctly with Sphinx autodoc. The function object is modified in place; + no wrapper is created, so there is no runtime overhead. + """ + + text = '\n\n'.join(snippet.strip('\n') for snippet in snippets) + + def decorator(func): + doc = inspect.getdoc(func) + if doc: + func.__doc__ = doc.rstrip() + '\n\n' + text + '\n' + else: + # The leading newline keeps the first snippet line out of the + # docstring-dedent logic of autodoc/inspect.cleandoc, which would + # otherwise strip the indentation of directive continuation lines. + func.__doc__ = '\n' + text + '\n' + return func + + return decorator + + def scalar_function(func): """ Serves to register class functions of Equilibrium class as scalar functions. diff --git a/tests/test_coordinates.py b/tests/test_coordinates.py index c63df3d..40b8874 100644 --- a/tests/test_coordinates.py +++ b/tests/test_coordinates.py @@ -1,4 +1,7 @@ +import warnings + import numpy as np +import pytest from pleque.core import Coordinates from pleque.tests.utils import load_testing_equilibrium @@ -57,12 +60,12 @@ def test_midplane(equilibrium): def test_scalar_point_coordinate(equilibrium): - coord = Coordinates(None, R=1, Z=2) + coord = Coordinates.from_coords(None, R=1, Z=2) assert len(coord) == 1 assert coord.x1[0] == 1 assert coord.x2[0] == 2 - coord = Coordinates(None, R=1, Z=2, phi=0) + coord = Coordinates.from_coords(None, R=1, Z=2, phi=0) assert len(coord) == 1 assert coord.x1[0] == 1 assert coord.x2[0] == 2 @@ -214,13 +217,13 @@ def test_array_input(): rr, zz = np.meshgrid(r, z) - coord = Coordinates(eq, R=rr, Z=zz, grid=False) + coord = Coordinates.from_coords(eq, R=rr, Z=zz, grid=False) assert coord.R.shape == rr.shape assert coord.Z.shape == rr.shape assert coord.r_mid.shape == rr.shape - coord = Coordinates(eq, r=coord.r_mid, theta=np.zeros_like(coord.r_mid), grid=False) + coord = Coordinates.from_coords(eq, r=coord.r_mid, theta=np.zeros_like(coord.r_mid), grid=False) assert coord.R.shape == rr.shape assert coord.Z.shape == rr.shape @@ -240,7 +243,60 @@ def test_distances(): R = 2 N = 52 - coord = Coordinates(None, R=np.ones(N) * R, Z=np.zeros(N), phi=np.linspace(0, 2 * np.pi, N)) + coord = Coordinates.from_coords(None, R=np.ones(N) * R, Z=np.zeros(N), phi=np.linspace(0, 2 * np.pi, N)) calc_length = R * np.pi * 2 assert np.isclose(coord.length, calc_length, atol=1e-2, rtol=1e-2) + + +def test_init_deprecated_rich_input(): + # Rich coordinate input passed directly to the constructor is deprecated: + with pytest.warns(DeprecationWarning): + Coordinates(None, R=1, Z=2) + with pytest.warns(DeprecationWarning): + Coordinates(eq, np.array([[1.0, 0.0], [2.0, 0.5]])) + with pytest.warns(DeprecationWarning): + Coordinates(eq, np.linspace(0, 1, 5), coord_type='rho') + + # ... while `from_coords` and positional input in the default coordinate + # type stay silent: + with warnings.catch_warnings(): + warnings.simplefilter('error') + Coordinates.from_coords(None, R=1, Z=2) + Coordinates.from_coords(eq, np.array([[1.0, 0.0], [2.0, 0.5]])) + Coordinates.from_coords(eq, np.linspace(0, 1, 5), coord_type='rho') + Coordinates(eq) + Coordinates(eq, 0.5) + Coordinates(eq, 1.4, 0.0) + Coordinates(eq, 1.4, 0.0, coord_type=('R', 'Z')) + Coordinates(eq, np.array([1.4]), np.array([0.0]), np.array([0.1])) + + +def test_fast_init_parity(): + # The fast construction path must give the same result as the full + # input parser used by `from_coords`. + cases = [ + (0.5,), # 1d scalar + (np.linspace(0, 1, 5),), # 1d array + (1.4, 0.1), # 2d scalars + (np.linspace(1, 2, 5), np.linspace(-0.2, 0.2, 5)), # 2d arrays + ([1.4, 1.5], [0.0, 0.1]), # 2d lists + (1.4, 0.1, 0.3), # 3d scalars + (np.array([1.4, 1.5]), np.array([0.0, 0.1]), np.array([0.2, 0.3])), + ] + for args in cases: + c1 = Coordinates(eq, *args) + c2 = Coordinates.from_coords(eq, *args) + assert c1.dim == c2.dim + assert c1 == c2 + assert c1._coord_type_input == c2._coord_type_input + assert isinstance(c1._x1_input, np.ndarray) + np.testing.assert_allclose(c1.x1, c2.x1) + if c1.dim >= 2: + np.testing.assert_allclose(c1.x2, c2.x2) + if c1.dim == 3: + np.testing.assert_allclose(c1.x3, c2.x3) + + # grid stays available for 2d input and is demoted otherwise: + assert Coordinates(eq, np.linspace(1, 2, 5), np.linspace(-0.2, 0.2, 5), grid=True).grid + assert not Coordinates(eq, np.linspace(0, 1, 5), grid=True).grid diff --git a/tests/test_fluxexpansion.py b/tests/test_fluxexpansion.py index 901f6bd..6e99db9 100644 --- a/tests/test_fluxexpansion.py +++ b/tests/test_fluxexpansion.py @@ -7,7 +7,7 @@ def test_normal_vectors(): n_points = 30 # Vertical line and its normal - line = Coordinates(None, R=np.linspace(0, 1, n_points), Z=np.zeros(n_points)) + line = Coordinates.from_coords(None, R=np.linspace(0, 1, n_points), Z=np.zeros(n_points)) line_normals = line.normal_vector() assert line_normals.shape == (3, n_points) @@ -16,7 +16,7 @@ def test_normal_vectors(): assert np.allclose(line_normals[2, :], 0) # Diagonal line and its normal - line = Coordinates(None, R=np.linspace(0, 1, n_points), Z=np.linspace(0, 1, n_points)) + line = Coordinates.from_coords(None, R=np.linspace(0, 1, n_points), Z=np.linspace(0, 1, n_points)) line_normals = line.normal_vector() assert np.allclose(line_normals[2, :], 0) @@ -29,7 +29,7 @@ def test_incidence_angle_sin(): n_points = 30 # Vertical line: - line = Coordinates(None, R=np.linspace(0, 1, n_points), Z=np.zeros(n_points)) + line = Coordinates.from_coords(None, R=np.linspace(0, 1, n_points), Z=np.zeros(n_points)) # Generate vectors in various angles (from [0, 2, 0] to [2, 0, 0]) rvec = np.linspace(0, 2, n_points, endpoint=True) From 560afb479ffae0e5f4a12ad5e81ccea4ba9d8bdf Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 12 Jun 2026 08:01:37 +0000 Subject: [PATCH 2/2] fix: resolve ruff lint errors in coordinates and equilibrium - coordinates.py: remove unused itertools import, fix import sort order, remove redundant `object` base class, annotate mutable class-level sets/dict with ClassVar to satisfy RUF012 - equilibrium.py: remove unused `import pleque`, collapse duplicate import blocks (F811 redefinitions introduced by the previous commit), fix import sort order https://claude.ai/code/session_012QryKbWtWjQQrWmHdhe1pR --- pleque/core/coordinates.py | 20 +++++++++----------- pleque/core/equilibrium.py | 23 ++++++++++------------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/pleque/core/coordinates.py b/pleque/core/coordinates.py index 4da3518..6bc80ea 100644 --- a/pleque/core/coordinates.py +++ b/pleque/core/coordinates.py @@ -1,17 +1,15 @@ -from collections.abc import Iterable import warnings -import itertools -from typing import Union +from collections.abc import Iterable +from typing import ClassVar, Union import numpy as np from scipy.interpolate import splev, splprep -from pleque.utils.decorators import append_to_doc, deprecated import pleque.utils.flux_expansions as flux_expansion +from pleque.utils.decorators import append_to_doc, deprecated from .cocos import cocos_coefs - #: Canonical description of the coordinate input syntax accepted across PLEQUE. #: It is appended (via :func:`pleque.utils.decorators.append_to_doc`) to the #: docstrings of :meth:`Coordinates.from_coords` and @@ -81,7 +79,7 @@ """ -class Coordinates(object): +class Coordinates: r""" Basic PLEQUE class to handle various coordinate systems in tokamak equilibrium. @@ -92,13 +90,13 @@ class Coordinates(object): """ # Sets of recognised coordinate names, shared by all instances: - _valid_coordinates = {'R', 'Z', 'psi_n', 'psi', 'rho', 'r', 'theta', 'phi', 'X', 'Y'} - _valid_coordinates_1d = {('psi_n',), ('psi',), ('rho',)} - _valid_coordinates_2d = {('R', 'Z'), ('r', 'theta')} - _valid_coordinates_3d = {('R', 'Z', 'phi'), ('X', 'Y', 'Z')} + _valid_coordinates: ClassVar[set] = {'R', 'Z', 'psi_n', 'psi', 'rho', 'r', 'theta', 'phi', 'X', 'Y'} + _valid_coordinates_1d: ClassVar[set] = {('psi_n',), ('psi',), ('rho',)} + _valid_coordinates_2d: ClassVar[set] = {('R', 'Z'), ('r', 'theta')} + _valid_coordinates_3d: ClassVar[set] = {('R', 'Z', 'phi'), ('X', 'Y', 'Z')} # Default coordinate type for a given dimension: - _default_coord_types = {1: ('psi_n',), 2: ('R', 'Z'), 3: ('R', 'Z', 'phi')} + _default_coord_types: ClassVar[dict] = {1: ('psi_n',), 2: ('R', 'Z'), 3: ('R', 'Z', 'phi')} def __init__(self, equilibrium, *coordinates, coord_type=None, grid=False, cocos=None, **coords): r""" diff --git a/pleque/core/equilibrium.py b/pleque/core/equilibrium.py index 2c95fda..b708c4f 100644 --- a/pleque/core/equilibrium.py +++ b/pleque/core/equilibrium.py @@ -4,26 +4,23 @@ import numpy as np import xarray from scipy.constants import mu_0 -from shapely import Polygon, Point - -import pleque -from pleque.utils.decorators import append_to_doc, deprecated, ordered_path_scalar_function, scalar_function, \ - vector_function - from scipy.interpolate import RectBivariateSpline, UnivariateSpline -from pleque.core import Coordinates -from pleque.core.coordinates import COORD_PARAMS_DOC, COORDINATES_DOC -from pleque.utils.tools import arglis, xp_sections -from pleque.core import FluxFunctions, Surface # , FluxSurface -from pleque.core import SurfaceFunctions -from pleque.core import cocos as cc +from shapely import Point, Polygon + import pleque.utils.equi_tools as eq_tools import pleque.utils.flux_expansions as flux_expansion import pleque.utils.surfaces as surf from pleque.config.settings import get_settings from pleque.core import Coordinates, FluxFunctions, Surface, SurfaceFunctions # , FluxSurface from pleque.core import cocos as cc -from pleque.utils.decorators import deprecated, ordered_path_scalar_function, scalar_function, vector_function +from pleque.core.coordinates import COORD_PARAMS_DOC, COORDINATES_DOC +from pleque.utils.decorators import ( + append_to_doc, + deprecated, + ordered_path_scalar_function, + scalar_function, + vector_function, +) from pleque.utils.surfaces import track_plasma_boundary from pleque.utils.tools import arglis, xp_sections