Source code for desc.transform
import numpy as np
import functools
from itertools import permutations, combinations_with_replacement
from desc.backend import jnp, conditional_decorator, jit, use_jax, TextColors, equals
from desc.grid import Grid
from desc.basis import Basis
from desc.equilibrium_io import IOAble
[docs]class Transform(IOAble):
"""Transform
Attributes
----------
grid : Grid
DESCRIPTION
basis : Basis
DESCRIPTION
rcond : float
relative cutoff for singular values in least squares fit
derivatives : ndarray
combinations of derivatives needed
Each row is one set, columns represent the order of derivatives
for [rho, theta, zeta]
matrices : ndarray
DESCRIPTION
pinv : ndarray
DESCRIPTION
"""
_save_attrs_ = ['grid', 'basis', 'derives', 'matrices']
[docs] def __init__(self, grid:Grid=None, basis:Basis=None, derivs=0, rcond=1e-6,
load_from=None, file_format=None, obj_lib=None) -> None:
"""Initializes a Transform
Parameters
----------
grid : Grid
DESCRIPTION
basis : Basis
DESCRIPTION
derivs : int or string
order of derivatives needed, if an int (Default = 0)
OR
type of calculation being performed, if a string
``'force'``: all of the derivatives needed to calculate an
equilibrium from the force balance equations
``'qs'``: all of the derivatives needed to calculate quasi-
symmetry from the triple-product equation
rcond : float
relative cutoff for singular values in least squares fit
Returns
-------
None
"""
if load_from is None:
self.__grid = grid
self.__basis = basis
self.__derivs = derivs
self.__rcond = rcond
self.__matrices = {i: {j: {k: {}
for k in range(4)} for j in range(4)} for i in range(4)}
self.__derivatives = self._get_derivatives_(self.__derivs)
self._sort_derivatives_()
self._build_()
self._build_pinv_()
#self._def_save_attrs_()
else:
self._init_from_file_(load_from=load_from, file_format=file_format, obj_lib=obj_lib)
[docs] def __eq__(self, other) -> bool:
"""Overloads the == operator
Parameters
----------
other : Transform
another Transform object to compare to
Returns
-------
bool
True if other is a Transform with the same attributes as self
False otherwise
"""
if self.__class__ != other.__class__:
return False
return equals(self.__dict__, other.__dict__)
def _get_derivatives_(self, derivs):
"""Get array of derivatives needed for calculating objective function
Parameters
----------
derivs : int or string
order of derivatives needed, if an int (Default = 0)
OR
type of calculation being performed, if a string
``'force'``: all of the derivatives needed to calculate an
equilibrium from the force balance equations
``'qs'``: all of the derivatives needed to calculate quasi-
symmetry from the triple-product equation
Returns
-------
derivatives : ndarray
combinations of derivatives needed
Each row is one set, columns represent the order of derivatives
for [rho, theta, zeta]
"""
if isinstance(derivs, int) and derivs >= 0:
derivatives = np.array([[]])
combos = combinations_with_replacement(range(derivs+1), 3)
for combo in list(combos):
perms = set(permutations(combo))
for perm in list(perms):
if derivatives.shape[1] == 3:
derivatives = np.vstack([derivatives, np.array(perm)])
else:
derivatives = np.array([perm])
elif derivs.lower() == 'force':
derivatives = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0],
[0, 1, 1], [0, 0, 2]])
# FIXME: this assumes the Grid is sorted (which it should be)
if np.all(self.__grid.nodes[:, 0] == np.array([0, 0, 0])):
axis = np.array([[2, 1, 0], [1, 2, 0], [1, 1, 1], [2, 2, 0]])
derivatives = np.vstack([derivatives, axis])
elif derivs.lower() == 'qs':
derivatives = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0],
[0, 1, 1], [0, 0, 2], [3, 0, 0], [2, 1, 0],
[2, 0, 1], [1, 2, 0], [1, 1, 1], [1, 0, 2],
[0, 3, 0], [0, 2, 1], [0, 1, 2], [0, 0, 3],
[2, 2, 0]])
else:
raise NotImplementedError(TextColors.FAIL +
"order options are 'force', 'qs', or a non-negative int"
+ TextColors.ENDC)
return derivatives
def _sort_derivatives_(self) -> None:
"""Sorts derivatives
Returns
-------
None
"""
sort_idx = np.lexsort((self.__derivatives[:, 0],
self.__derivatives[:, 1], self.__derivatives[:, 2]))
self.__derivatives = self.__derivatives[sort_idx]
def _build_(self) -> None:
"""Builds the transform matrices for each derivative order
"""
for d in self.__derivatives:
self.__matrices[d[0]][d[1]][d[2]] = self.__basis.evaluate(
self.__grid.nodes, d)
def _build_pinv_(self) -> None:
"""Builds the transform matrices for each derivative order
"""
# FIXME: this assumes the derivatives are sorted (which they should be)
if np.all(self.__derivatives[0, :] == np.array([0, 0, 0])):
A = self.__matrices[0][0][0]
else:
A = self.__basis.evaluate(self.__grid.nodes, np.array([0, 0, 0]))
self.__pinv = jnp.linalg.pinv(A, rcond=self.__rcond)
def _def_save_attrs_(self) -> None:
"""Defines attributes to save
Returns
-------
None
"""
self._save_attrs_ = ['__grid', '__basis', '__derives', '__matrices']
[docs] def transform(self, c, dr=0, dt=0, dz=0):
"""Transform from spectral domain to physical
Parameters
----------
c : ndarray, shape(N_coeffs,)
spectral coefficients, indexed as (lm,n) flattened in row major order
dr : int
order of radial derivative
dt : int
order of poloidal derivative
dz : int
order of toroidal derivative
Returns
-------
x : ndarray, shape(N_nodes,)
array of values of function at node locations
"""
A = self.__matrices[dr][dt][dz]
if type(A) is dict:
raise ValueError(TextColors.FAIL +
"Derivative orders are out of initialized bounds" +
TextColors.ENDC)
if A.shape[1] != c.size:
raise ValueError(TextColors.FAIL +
"Coefficients dimension ({}) is incompatible with the number of basis modes({})".format(c.size, A.shape[1]) +
TextColors.ENDC)
return jnp.matmul(A, c)
[docs] @conditional_decorator(functools.partial(jit, static_argnums=(0,)), use_jax)
def fit(self, x):
"""Transform from physical domain to spectral using least squares fit
Parameters
----------
x : ndarray, shape(N_nodes,)
values in real space at coordinates specified by self.grid
Returns
-------
c : ndarray, shape(N_coeffs,)
spectral coefficients in self.basis
"""
return jnp.matmul(self.__pinv, x)
[docs] def change_resolution(self, grid:Grid=None, basis:Basis=None) -> None:
"""Re-builds the matrices with a new grid and basis
Parameters
----------
grid : Grid, optional
DESCRIPTION
basis : Basis, optional
DESCRIPTION
Returns
-------
None
"""
if grid is None:
grid = self.__grid
if basis is None:
basis = self.__basis
if self.__grid != grid or self.__basis != basis:
self.__grid = grid
self.__basis = basis
self._build_()
self._build_pinv_()
@property
def grid(self):
return self.__grid
@grid.setter
def grid(self, grid:Grid) -> None:
"""Changes the grid and updates the matrices accordingly
Parameters
----------
grid : Grid
DESCRIPTION
Returns
-------
None
"""
if self.__grid != grid:
self.__grid = grid
self._build_()
self._build_pinv_()
@property
def basis(self):
return self.__basis
@basis.setter
def basis(self, basis:Basis) -> None:
"""Changes the basis and updates the matrices accordingly
Parameters
----------
basis : Basis
DESCRIPTION
Returns
-------
None
"""
if self.__basis != basis:
self.__basis = basis
self._build_()
self._build_pinv_()
@property
def derivs(self):
return self.__derivs
@property
def derivatives(self):
return self.__derivatives
@derivatives.setter
def derivatives(self, derivs) -> None:
"""Changes the order and updates the matrices accordingly
Parameters
----------
derivs : int or string
order of derivatives needed, if an int (Default = 0)
OR
type of calculation being performed, if a string
``'force'``: all of the derivatives needed to calculate an
equilibrium from the force balance equations
``'qs'``: all of the derivatives needed to calculate quasi-
symmetry from the triple-product equation
Returns
-------
None
"""
if derivs != self.__derivs:
self.__derivs = derivs
old_derivatives = self.__derivatives
self.__derivatives = self.get_derivatives(self.__derivs)
self.sort_derivatives()
new_derivatives = self.__derivatives
new_not_in_old = (
new_derivatives[:, None] == old_derivatives).all(-1).any(-1)
derivs_to_add = new_derivatives[~new_not_in_old]
for d in derivs_to_add:
self.__matrices[d[0]][d[1]][d[2]] = self.__basis.evaluate(
self.__grid.nodes, d)
@property
def matrices(self):
return self.__matrices
@property
def num_nodes(self):
return self.__grid.num_nodes
@property
def num_modes(self):
return self.__basis.num_modes
# these functions are currently unused ---------------------------------------
[docs]def zernike_norm(l, m):
"""Norm of a Zernike polynomial with l, m indexing.
Returns the integral (Z^m_l)^2 r dr dt, r=[0,1], t=[0,2*pi]
Parameters
----------
l,m : int
radial and azimuthal mode numbers.
Returns
-------
norm : float
norm of Zernike polynomial over unit disk.
"""
return jnp.sqrt((2 * (l + 1)) / (jnp.pi*(1 + jnp.kronecker(m, 0))))
[docs]def lm_to_fringe(l, m):
"""Convert Zernike (l,m) double index to single Fringe index.
Parameters
----------
l,m : int
radial and azimuthal mode numbers.
Returns
-------
idx : int
Fringe index for l,m
"""
M = (l + np.abs(m)) / 2
return int(M**2 + M + m)
[docs]def fringe_to_lm(idx):
"""Convert single Zernike Fringe index to (l,m) double index.
Parameters
----------
idx : int
Fringe index
Returns
-------
l,m : int
radial and azimuthal mode numbers.
"""
M = (np.ceil(np.sqrt(idx+1)) - 1)
m = idx - M**2 - M
l = 2*M - np.abs(m)
return int(l), int(m)
[docs]def lm_to_ansi(l, m):
"""Convert Zernike (l,m) two term index to ANSI single term index.
Parameters
----------
l,m : int
radial and azimuthal mode numbers.
Returns
-------
idx : int
ANSI index for l,m
"""
return int((l * (l + 2) + m) / 2)
[docs]def ansi_to_lm(idx):
"""Convert Zernike ANSI single term to (l,m) two-term index.
Parameters
----------
idx : int
ANSI index
Returns
-------
l,m : int
radial and azimuthal mode numbers.
"""
l = int(np.ceil((-3 + np.sqrt(9 + 8*idx))/2))
m = 2 * idx - l * (l + 2)
return l, m