diff --git a/CHANGELOG.md b/CHANGELOG.md index 56a93d27b0..b69ca99d02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,11 +15,14 @@ in the MAKEGRID format for use with other codes. least squares fit is now weighted inversely with the distance from the axis to improve the accuracy for low aspect ratio. - Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z` -keyword arguments, which form a hollow cylindrical bounding box. If the field line -trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying -function of the distance to the box to stop the trajectory and prevent tracing the field -line out to infinity, which is both costly and unnecessary when making a Poincare plot, +keyword arguments, which form a hollow cylindrical bounding box. If the field line +trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying +function of the distance to the box to stop the trajectory and prevent tracing the field +line out to infinity, which is both costly and unnecessary when making a Poincare plot, the principle purpose of the function. +- Adds a new class ``DommaschkPotentialField`` which allows creation of magnetic fields based +off of the vacuum potentials detailed in Representations for Vacuum Potentials in Stellarators +https://doi.org/10.1016/0010-4655(86)90109-8. Speed Improvements diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 2b87cab2bf..a65cbbaa0a 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -626,7 +626,9 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi"] = ( @@ -649,7 +651,9 @@ def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_t"] = ( @@ -670,7 +674,9 @@ def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_z"] = ( @@ -691,7 +697,7 @@ def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi"] = transforms["potential"]( @@ -714,7 +720,7 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_t"] = transforms["potential_dtheta"]( @@ -737,7 +743,7 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_z"] = transforms["potential_dzeta"]( @@ -762,8 +768,8 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs): coordinates="tz", data=["Phi_t", "Phi_z", "e_theta", "e_zeta", "|e_theta x e_zeta|"], parameterization=[ - "desc.magnetic_fields.CurrentPotentialField", - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.CurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", ], ) def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs): diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index cc91e326ff..285d851d77 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -211,15 +211,15 @@ def _decorator(func): "desc.geometry.curve.FourierPlanarCurve", "desc.geometry.core.Curve", ], - "desc.magnetic_fields.CurrentPotentialField": [ + "desc.magnetic_fields._current_potential.CurrentPotentialField": [ "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", - "desc.magnetic_fields.MagneticField", + "desc.magnetic_fields._core.MagneticField", ], - "desc.magnetic_fields.FourierCurrentPotentialField": [ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": [ "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", - "desc.magnetic_fields.MagneticField", + "desc.magnetic_fields._core.MagneticField", ], "desc.coils.SplineXYZCoil": [ "desc.geometry.curve.SplineXYZCurve", diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py new file mode 100644 index 0000000000..722f64adda --- /dev/null +++ b/desc/magnetic_fields/__init__.py @@ -0,0 +1,16 @@ +"""Classes for Magnetic Fields.""" + +from ._core import ( + PoloidalMagneticField, + ScalarPotentialField, + ScaledMagneticField, + SplineMagneticField, + SumMagneticField, + ToroidalMagneticField, + VerticalMagneticField, + _MagneticField, + field_line_integrate, + read_BNORM_file, +) +from ._current_potential import CurrentPotentialField, FourierCurrentPotentialField +from ._dommaschk import DommaschkPotentialField, dommaschk_potential diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields/_core.py similarity index 69% rename from desc/magnetic_fields.py rename to desc/magnetic_fields/_core.py index f6372dadbe..727a3f14f1 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields/_core.py @@ -9,15 +9,14 @@ from desc.backend import fori_loop, jit, jnp, odeint from desc.basis import DoubleFourierSeries -from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.compute import rpz2xyz_vec, xyz2rpz from desc.derivatives import Derivative from desc.equilibrium import EquilibriaFamily, Equilibrium -from desc.geometry import FourierRZToroidalSurface from desc.grid import LinearGrid from desc.io import IOAble from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter from desc.transform import Transform -from desc.utils import copy_coeffs, errorif, flatten_list, setdefault, warnif +from desc.utils import copy_coeffs, flatten_list, setdefault from desc.vmec_utils import ptolemy_identity_fwd, ptolemy_identity_rev @@ -1342,9 +1341,9 @@ def compute_magnetic_field( """ assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) + coords = coords.astype(float) if basis == "xyz": coords = xyz2rpz(coords) - Rq, phiq, Zq = coords.T if params is None: params = self._params @@ -1479,664 +1478,3 @@ def odefun(rpz, s): r = x[:, :, 0].T.reshape((len(phis), *rshape)) z = x[:, :, 2].T.reshape((len(phis), *rshape)) return r, z - - -class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): - """Magnetic field due to a surface current potential on a toroidal surface. - - Surface current K is assumed given by K = n x ∇ Φ - where: - - - n is the winding surface unit normal. - - Phi is the current potential function, which is a function of theta and zeta. - - This function then uses biot-savart to find the B field from this current - density K on the surface. - - Parameters - ---------- - potential : callable - function to compute the current potential. Should have a signature of - the form potential(theta,zeta,**params) -> ndarray. - theta,zeta are poloidal and toroidal angles on the surface - potential_dtheta: callable - function to compute the theta derivative of the current potential - potential_dzeta: callable - function to compute the zeta derivative of the current potential - params : dict, optional - default parameters to pass to potential function (and its derivatives) - R_lmn, Z_lmn : array-like, shape(k,) - Fourier coefficients for winding surface R and Z in cylindrical coordinates - modes_R : array-like, shape(k,2) - poloidal and toroidal mode numbers [m,n] for R_lmn. - modes_Z : array-like, shape(k,2) - mode numbers associated with Z_lmn, defaults to modes_R - NFP : int - number of field periods - sym : bool - whether to enforce stellarator symmetry for the surface geometry. - Default is "auto" which enforces if modes are symmetric. If True, - non-symmetric modes will be truncated. - M, N: int or None - Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R - and modes_Z. - name : str - name for this field - check_orientation : bool - ensure that this surface has a right handed orientation. Do not set to False - unless you are sure the parameterization you have given is right handed - (ie, e_theta x e_zeta points outward from the surface). - - """ - - _io_attrs_ = ( - _MagneticField._io_attrs_ - + FourierRZToroidalSurface._io_attrs_ - + [ - "_params", - ] - ) - - def __init__( - self, - potential, - potential_dtheta, - potential_dzeta, - params=None, - R_lmn=None, - Z_lmn=None, - modes_R=None, - modes_Z=None, - NFP=1, - sym="auto", - M=None, - N=None, - name="", - check_orientation=True, - ): - assert callable(potential), "Potential must be callable!" - assert callable(potential_dtheta), "Potential derivative must be callable!" - assert callable(potential_dzeta), "Potential derivative must be callable!" - - self._potential = potential - self._potential_dtheta = potential_dtheta - self._potential_dzeta = potential_dzeta - self._params = params - - super().__init__( - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - M=M, - N=N, - name=name, - check_orientation=check_orientation, - ) - - @property - def params(self): - """Dict of parameters to pass to potential function and its derivatives.""" - return self._params - - @params.setter - def params(self, new): - warnif( - len(new) != len(self._params), - UserWarning, - "Length of new params is different from length of current params! " - "May cause errors unless potential function is also changed.", - ) - self._params = new - - @property - def potential(self): - """Potential function, signature (theta,zeta,**params) -> potential value.""" - return self._potential - - @potential.setter - def potential(self, new): - if new != self._potential: - assert callable(new), "Potential must be callable!" - self._potential = new - - @property - def potential_dtheta(self): - """Phi poloidal deriv. function, signature (theta,zeta,**params) -> value.""" - return self._potential_dtheta - - @potential_dtheta.setter - def potential_dtheta(self, new): - if new != self._potential_dtheta: - assert callable(new), "Potential derivative must be callable!" - self._potential_dtheta = new - - @property - def potential_dzeta(self): - """Phi toroidal deriv. function, signature (theta,zeta,**params) -> value.""" - return self._potential_dzeta - - @potential_dzeta.setter - def potential_dzeta(self, new): - if new != self._potential_dzeta: - assert callable(new), "Potential derivative must be callable!" - self._potential_dzeta = new - - def save(self, file_name, file_format=None, file_mode="w"): - """Save the object. - - **Not supported for this object!** - - Parameters - ---------- - file_name : str file path OR file instance - location to save object - file_format : str (Default hdf5) - format of save file. Only used if file_name is a file path - file_mode : str (Default w - overwrite) - mode for save file. Only used if file_name is a file path - - """ - raise OSError( - "Saving CurrentPotentialField is not supported," - " as the potential function cannot be serialized." - ) - - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None - ): - """Compute magnetic field at a set of points. - - Parameters - ---------- - coords : array-like shape(n,3) - Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. - params : dict or array-like of dict, optional - Dictionary of optimizable parameters, eg field.params_dict. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - source_grid : Grid, int or None or array-like, optional - Source grid upon which to evaluate the surface current density K. - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * self.M, - N=30 + 2 * self.N, - NFP=self.NFP, - ) - return _compute_magnetic_field_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - ) - - @classmethod - def from_surface( - cls, - surface, - potential, - potential_dtheta, - potential_dzeta, - params=None, - ): - """Create CurrentPotentialField using geometry provided by given surface. - - Parameters - ---------- - surface: FourierRZToroidalSurface, optional, default None - Existing FourierRZToroidalSurface object to create a - CurrentPotentialField with. - potential : callable - function to compute the current potential. Should have a signature of - the form potential(theta,zeta,**params) -> ndarray. - theta,zeta are poloidal and toroidal angles on the surface - potential_dtheta: callable - function to compute the theta derivative of the current potential - potential_dzeta: callable - function to compute the zeta derivative of the current potential - params : dict, optional - default parameters to pass to potential function (and its derivatives) - - """ - errorif( - not isinstance(surface, FourierRZToroidalSurface), - TypeError, - "Expected type FourierRZToroidalSurface for argument surface, " - f"instead got type {type(surface)}", - ) - - R_lmn = surface.R_lmn - Z_lmn = surface.Z_lmn - modes_R = surface._R_basis.modes[:, 1:] - modes_Z = surface._Z_basis.modes[:, 1:] - NFP = surface.NFP - sym = surface.sym - name = surface.name - - return cls( - potential, - potential_dtheta, - potential_dzeta, - params, - R_lmn, - Z_lmn, - modes_R, - modes_Z, - NFP, - sym, - name=name, - check_orientation=False, - ) - - -class FourierCurrentPotentialField( - _MagneticField, FourierRZToroidalSurface, Optimizable -): - """Magnetic field due to a surface current potential on a toroidal surface. - - Surface current K is assumed given by - - K = n x ∇ Φ - - Φ(θ,ζ) = Φₛᵥ(θ,ζ) + Gζ/2π + Iθ/2π - - where: - - n is the winding surface unit normal. - - Phi is the current potential function, which is a function of theta and zeta, - and is given as a secular linear term in theta/zeta and a double Fourier - series in theta/zeta. - - This function then uses biot-savart to find the B field from this current - density K on the surface. - - Parameters - ---------- - Phi_mn : ndarray - Fourier coefficients of the double FourierSeries part of the current potential. - modes_Phi : array-like, shape(k,2) - Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn - coefficients. - I : float - Net current linking the plasma and the surface toroidally - Denoted I in the algorithm - G : float - Net current linking the plasma and the surface poloidally - Denoted G in the algorithm - NOTE: a negative G will tend to produce a positive toroidal magnetic field - B in DESC, as in DESC the poloidal angle is taken to be positive - and increasing when going in the clockwise direction, which with the - convention n x grad(phi) will result in a toroidal field in the negative - toroidal direction. - sym_Phi : {False,"cos","sin"} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. - M_Phi, N_Phi: int or None - Maximum poloidal and toroidal mode numbers for the single valued part of the - current potential. - R_lmn, Z_lmn : array-like, shape(k,) - Fourier coefficients for winding surface R and Z in cylindrical coordinates - modes_R : array-like, shape(k,2) - poloidal and toroidal mode numbers [m,n] for R_lmn. - modes_Z : array-like, shape(k,2) - mode numbers associated with Z_lmn, defaults to modes_R - NFP : int - number of field periods - sym : bool - whether to enforce stellarator symmetry for the surface geometry. - Default is "auto" which enforces if modes are symmetric. If True, - non-symmetric modes will be truncated. - M, N: int or None - Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R - and modes_Z. - name : str - name for this field - check_orientation : bool - ensure that this surface has a right handed orientation. Do not set to False - unless you are sure the parameterization you have given is right handed - (ie, e_theta x e_zeta points outward from the surface). - - """ - - _io_attrs_ = ( - _MagneticField._io_attrs_ - + FourierRZToroidalSurface._io_attrs_ - + ["_Phi_mn", "_I", "_G", "_Phi_basis", "_M_Phi", "_N_Phi", "_sym_Phi"] - ) - - def __init__( - self, - Phi_mn=np.array([0.0]), - modes_Phi=np.array([[0, 0]]), - I=0, - G=0, - sym_Phi=False, - M_Phi=None, - N_Phi=None, - R_lmn=None, - Z_lmn=None, - modes_R=None, - modes_Z=None, - NFP=1, - sym="auto", - M=None, - N=None, - name="", - check_orientation=True, - ): - Phi_mn, modes_Phi = map(np.asarray, (Phi_mn, modes_Phi)) - assert ( - Phi_mn.size == modes_Phi.shape[0] - ), "Phi_mn size and modes_Phi.shape[0] must be the same size!" - - assert np.issubdtype(modes_Phi.dtype, np.integer) - - M_Phi = setdefault(M_Phi, np.max(abs(modes_Phi[:, 0]))) - N_Phi = setdefault(N_Phi, np.max(abs(modes_Phi[:, 1]))) - - self._M_Phi = M_Phi - self._N_Phi = N_Phi - - self._sym_Phi = sym_Phi - self._Phi_basis = DoubleFourierSeries(M=M_Phi, N=N_Phi, NFP=NFP, sym=sym_Phi) - self._Phi_mn = copy_coeffs(Phi_mn, modes_Phi, self._Phi_basis.modes[:, 1:]) - - self._I = float(np.squeeze(I)) - self._G = float(np.squeeze(G)) - - super().__init__( - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - M=M, - N=N, - name=name, - check_orientation=check_orientation, - ) - - @optimizable_parameter - @property - def I(self): # noqa: E743 - """Net current linking the plasma and the surface toroidally.""" - return self._I - - @I.setter - def I(self, new): # noqa: E743 - self._I = float(np.squeeze(new)) - - @optimizable_parameter - @property - def G(self): - """Net current linking the plasma and the surface poloidally.""" - return self._G - - @G.setter - def G(self, new): - self._G = float(np.squeeze(new)) - - @optimizable_parameter - @property - def Phi_mn(self): - """Fourier coefficients describing single-valued part of potential.""" - return self._Phi_mn - - @Phi_mn.setter - def Phi_mn(self, new): - if len(new) == self.Phi_basis.num_modes: - self._Phi_mn = jnp.asarray(new) - else: - raise ValueError( - f"Phi_mn should have the same size as the basis, got {len(new)} for " - + f"basis with {self.Phi_basis.num_modes} modes." - ) - - @property - def Phi_basis(self): - """DoubleFourierSeries: Spectral basis for Phi.""" - return self._Phi_basis - - @property - def sym_Phi(self): - """str: Type of symmetry of periodic part of Phi (no symmetry if False).""" - return self._sym_Phi - - @property - def M_Phi(self): - """int: Poloidal resolution of periodic part of Phi.""" - return self._M_Phi - - @property - def N_Phi(self): - """int: Toroidal resolution of periodic part of Phi.""" - return self._N_Phi - - def change_Phi_resolution(self, M=None, N=None, NFP=None, sym_Phi=None): - """Change the maximum poloidal and toroidal resolution for Phi. - - Parameters - ---------- - M : int - Poloidal resolution to change Phi basis to. - If None, defaults to current self.Phi_basis poloidal resolution - N : int - Toroidal resolution to change Phi basis to. - If None, defaults to current self.Phi_basis toroidal resolution - NFP : int - Number of field periods for surface and Phi basis. - If None, defaults to current NFP. - Note: will change the NFP of the surface geometry as well as the - Phi basis. - sym_Phi : {"auto","cos","sin",False} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. Default is "auto" which enforces if modes are symmetric. - If True, non-symmetric modes will be truncated. - - """ - M = M or self._M_Phi - N = N or self._M_Phi - NFP = NFP or self.NFP - sym_Phi = sym_Phi or self.sym_Phi - - Phi_modes_old = self.Phi_basis.modes - self.Phi_basis.change_resolution(M=M, N=N, NFP=self.NFP, sym=sym_Phi) - - self._Phi_mn = copy_coeffs(self.Phi_mn, Phi_modes_old, self.Phi_basis.modes) - self._M_Phi = M - self._N_Phi = N - self._sym_Phi = sym_Phi - self.change_resolution( - NFP=NFP - ) # make sure surface and Phi basis NFP are the same - - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None - ): - """Compute magnetic field at a set of points. - - Parameters - ---------- - coords : array-like shape(n,3) - Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. - params : dict or array-like of dict, optional - Dictionary of optimizable parameters, eg field.params_dict. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - source_grid : Grid, int or None or array-like, optional - Source grid upon which to evaluate the surface current density K. - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * max(self.M, self.M_Phi), - N=30 + 2 * max(self.N, self.N_Phi), - NFP=self.NFP, - ) - return _compute_magnetic_field_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - ) - - @classmethod - def from_surface( - cls, - surface, - Phi_mn=np.array([0.0]), - modes_Phi=np.array([[0, 0]]), - I=0, - G=0, - sym_Phi=False, - M_Phi=None, - N_Phi=None, - ): - """Create FourierCurrentPotentialField using geometry of given surface. - - Parameters - ---------- - surface: FourierRZToroidalSurface, optional, default None - Existing FourierRZToroidalSurface object to create a - CurrentPotentialField with. - Phi_mn : ndarray - Fourier coefficients of the double FourierSeries of the current potential. - Should correspond to the given DoubleFourierSeries basis object passed in. - modes_Phi : array-like, shape(k,2) - Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn - coefficients - I : float - Net current linking the plasma and the surface toroidally - Denoted I in the algorithm - G : float - Net current linking the plasma and the surface poloidally - Denoted G in the algorithm - NOTE: a negative G will tend to produce a positive toroidal magnetic field - B in DESC, as in DESC the poloidal angle is taken to be positive - and increasing when going in the clockwise direction, which with the - convention n x grad(phi) will result in a toroidal field in the negative - toroidal direction. - sym_Phi : {False,"cos","sin"} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. - M_Phi, N_Phi: int or None - Maximum poloidal and toroidal mode numbers for the single valued part of the - current potential. - - """ - if not isinstance(surface, FourierRZToroidalSurface): - raise TypeError( - "Expected type FourierRZToroidalSurface for argument surface, " - f"instead got type {type(surface)}" - ) - R_lmn = surface.R_lmn - Z_lmn = surface.Z_lmn - modes_R = surface._R_basis.modes[:, 1:] - modes_Z = surface._Z_basis.modes[:, 1:] - NFP = surface.NFP - sym = surface.sym - name = surface.name - - return cls( - Phi_mn=Phi_mn, - modes_Phi=modes_Phi, - I=I, - G=G, - sym_Phi=sym_Phi, - M_Phi=M_Phi, - N_Phi=N_Phi, - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - name=name, - check_orientation=False, - ) - - -def _compute_magnetic_field_from_CurrentPotentialField( - field, - coords, - source_grid, - params=None, - basis="rpz", -): - """Compute magnetic field at a set of points. - - Parameters - ---------- - field : CurrentPotentialField or FourierCurrentPotentialField - current potential field object from which to compute magnetic field. - coords : array-like shape(N,3) - cylindrical or cartesian coordinates - source_grid : Grid, - source grid upon which to evaluate the surface current density K - params : dict, optional - parameters to pass to compute function - should include the potential - basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - assert basis.lower() in ["rpz", "xyz"] - coords = jnp.atleast_2d(jnp.asarray(coords)) - if basis == "rpz": - coords = rpz2xyz(coords) - - # compute surface current, and store grid quantities - # needed for integration in class - # TODO: does this have to be xyz, or can it be computed in rpz as well? - data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) - - _rs = xyz2rpz(data["x"]) - _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) - - # surface element, must divide by NFP to remove the NFP multiple on the - # surface grid weights, as we account for that when doing the for loop - # over NFP - _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP - - def nfp_loop(j, f): - # calculate (by rotating) rs, rs_t, rz_t - phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( - 2 * jnp.pi - ) - # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) - rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T - rs = rpz2xyz(rs) - K = rpz2xyz_vec(_K, phi=phi) - fj = biot_savart_general( - coords, - rs, - K, - _dV, - ) - f += fj - return f - - B = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) - if basis == "rpz": - B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) - return B diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py new file mode 100644 index 0000000000..d728cd50cd --- /dev/null +++ b/desc/magnetic_fields/_current_potential.py @@ -0,0 +1,675 @@ +"""Magnetic field due to sheet current on a winding surface.""" + +import numpy as np + +from desc.backend import fori_loop, jnp +from desc.basis import DoubleFourierSeries +from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.geometry import FourierRZToroidalSurface +from desc.grid import LinearGrid +from desc.optimizable import Optimizable, optimizable_parameter +from desc.utils import copy_coeffs, errorif, setdefault, warnif + +from ._core import _MagneticField, biot_savart_general + + +class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): + """Magnetic field due to a surface current potential on a toroidal surface. + + Surface current K is assumed given by K = n x ∇ Φ + where: + + - n is the winding surface unit normal. + - Phi is the current potential function, which is a function of theta and zeta. + + This function then uses biot-savart to find the B field from this current + density K on the surface. + + Parameters + ---------- + potential : callable + function to compute the current potential. Should have a signature of + the form potential(theta,zeta,**params) -> ndarray. + theta,zeta are poloidal and toroidal angles on the surface + potential_dtheta: callable + function to compute the theta derivative of the current potential + potential_dzeta: callable + function to compute the zeta derivative of the current potential + params : dict, optional + default parameters to pass to potential function (and its derivatives) + R_lmn, Z_lmn : array-like, shape(k,) + Fourier coefficients for winding surface R and Z in cylindrical coordinates + modes_R : array-like, shape(k,2) + poloidal and toroidal mode numbers [m,n] for R_lmn. + modes_Z : array-like, shape(k,2) + mode numbers associated with Z_lmn, defaults to modes_R + NFP : int + number of field periods + sym : bool + whether to enforce stellarator symmetry for the surface geometry. + Default is "auto" which enforces if modes are symmetric. If True, + non-symmetric modes will be truncated. + M, N: int or None + Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R + and modes_Z. + name : str + name for this field + check_orientation : bool + ensure that this surface has a right handed orientation. Do not set to False + unless you are sure the parameterization you have given is right handed + (ie, e_theta x e_zeta points outward from the surface). + + """ + + _io_attrs_ = ( + _MagneticField._io_attrs_ + + FourierRZToroidalSurface._io_attrs_ + + [ + "_params", + ] + ) + + def __init__( + self, + potential, + potential_dtheta, + potential_dzeta, + params=None, + R_lmn=None, + Z_lmn=None, + modes_R=None, + modes_Z=None, + NFP=1, + sym="auto", + M=None, + N=None, + name="", + check_orientation=True, + ): + assert callable(potential), "Potential must be callable!" + assert callable(potential_dtheta), "Potential derivative must be callable!" + assert callable(potential_dzeta), "Potential derivative must be callable!" + + self._potential = potential + self._potential_dtheta = potential_dtheta + self._potential_dzeta = potential_dzeta + self._params = params + + super().__init__( + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + M=M, + N=N, + name=name, + check_orientation=check_orientation, + ) + + @property + def params(self): + """Dict of parameters to pass to potential function and its derivatives.""" + return self._params + + @params.setter + def params(self, new): + warnif( + len(new) != len(self._params), + UserWarning, + "Length of new params is different from length of current params! " + "May cause errors unless potential function is also changed.", + ) + self._params = new + + @property + def potential(self): + """Potential function, signature (theta,zeta,**params) -> potential value.""" + return self._potential + + @potential.setter + def potential(self, new): + if new != self._potential: + assert callable(new), "Potential must be callable!" + self._potential = new + + @property + def potential_dtheta(self): + """Phi poloidal deriv. function, signature (theta,zeta,**params) -> value.""" + return self._potential_dtheta + + @potential_dtheta.setter + def potential_dtheta(self, new): + if new != self._potential_dtheta: + assert callable(new), "Potential derivative must be callable!" + self._potential_dtheta = new + + @property + def potential_dzeta(self): + """Phi toroidal deriv. function, signature (theta,zeta,**params) -> value.""" + return self._potential_dzeta + + @potential_dzeta.setter + def potential_dzeta(self, new): + if new != self._potential_dzeta: + assert callable(new), "Potential derivative must be callable!" + self._potential_dzeta = new + + def save(self, file_name, file_format=None, file_mode="w"): + """Save the object. + + **Not supported for this object!** + + Parameters + ---------- + file_name : str file path OR file instance + location to save object + file_format : str (Default hdf5) + format of save file. Only used if file_name is a file path + file_mode : str (Default w - overwrite) + mode for save file. Only used if file_name is a file path + + """ + raise OSError( + "Saving CurrentPotentialField is not supported," + " as the potential function cannot be serialized." + ) + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * self.M, + N=30 + 2 * self.N, + NFP=self.NFP, + ) + return _compute_magnetic_field_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + + @classmethod + def from_surface( + cls, + surface, + potential, + potential_dtheta, + potential_dzeta, + params=None, + ): + """Create CurrentPotentialField using geometry provided by given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface, optional, default None + Existing FourierRZToroidalSurface object to create a + CurrentPotentialField with. + potential : callable + function to compute the current potential. Should have a signature of + the form potential(theta,zeta,**params) -> ndarray. + theta,zeta are poloidal and toroidal angles on the surface + potential_dtheta: callable + function to compute the theta derivative of the current potential + potential_dzeta: callable + function to compute the zeta derivative of the current potential + params : dict, optional + default parameters to pass to potential function (and its derivatives) + + """ + errorif( + not isinstance(surface, FourierRZToroidalSurface), + TypeError, + "Expected type FourierRZToroidalSurface for argument surface, " + f"instead got type {type(surface)}", + ) + + R_lmn = surface.R_lmn + Z_lmn = surface.Z_lmn + modes_R = surface._R_basis.modes[:, 1:] + modes_Z = surface._Z_basis.modes[:, 1:] + NFP = surface.NFP + sym = surface.sym + name = surface.name + + return cls( + potential, + potential_dtheta, + potential_dzeta, + params, + R_lmn, + Z_lmn, + modes_R, + modes_Z, + NFP, + sym, + name=name, + check_orientation=False, + ) + + +class FourierCurrentPotentialField( + _MagneticField, FourierRZToroidalSurface, Optimizable +): + """Magnetic field due to a surface current potential on a toroidal surface. + + Surface current K is assumed given by + + K = n x ∇ Φ + + Φ(θ,ζ) = Φₛᵥ(θ,ζ) + Gζ/2π + Iθ/2π + + where: + + - n is the winding surface unit normal. + - Phi is the current potential function, which is a function of theta and zeta, + and is given as a secular linear term in theta/zeta and a double Fourier + series in theta/zeta. + + This function then uses biot-savart to find the B field from this current + density K on the surface. + + Parameters + ---------- + Phi_mn : ndarray + Fourier coefficients of the double FourierSeries part of the current potential. + modes_Phi : array-like, shape(k,2) + Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn + coefficients. + I : float + Net current linking the plasma and the surface toroidally + Denoted I in the algorithm + G : float + Net current linking the plasma and the surface poloidally + Denoted G in the algorithm + NOTE: a negative G will tend to produce a positive toroidal magnetic field + B in DESC, as in DESC the poloidal angle is taken to be positive + and increasing when going in the clockwise direction, which with the + convention n x grad(phi) will result in a toroidal field in the negative + toroidal direction. + sym_Phi : {False,"cos","sin"} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. + M_Phi, N_Phi: int or None + Maximum poloidal and toroidal mode numbers for the single valued part of the + current potential. + R_lmn, Z_lmn : array-like, shape(k,) + Fourier coefficients for winding surface R and Z in cylindrical coordinates + modes_R : array-like, shape(k,2) + poloidal and toroidal mode numbers [m,n] for R_lmn. + modes_Z : array-like, shape(k,2) + mode numbers associated with Z_lmn, defaults to modes_R + NFP : int + number of field periods + sym : bool + whether to enforce stellarator symmetry for the surface geometry. + Default is "auto" which enforces if modes are symmetric. If True, + non-symmetric modes will be truncated. + M, N: int or None + Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R + and modes_Z. + name : str + name for this field + check_orientation : bool + ensure that this surface has a right handed orientation. Do not set to False + unless you are sure the parameterization you have given is right handed + (ie, e_theta x e_zeta points outward from the surface). + + """ + + _io_attrs_ = ( + _MagneticField._io_attrs_ + + FourierRZToroidalSurface._io_attrs_ + + ["_Phi_mn", "_I", "_G", "_Phi_basis", "_M_Phi", "_N_Phi", "_sym_Phi"] + ) + + def __init__( + self, + Phi_mn=np.array([0.0]), + modes_Phi=np.array([[0, 0]]), + I=0, + G=0, + sym_Phi=False, + M_Phi=None, + N_Phi=None, + R_lmn=None, + Z_lmn=None, + modes_R=None, + modes_Z=None, + NFP=1, + sym="auto", + M=None, + N=None, + name="", + check_orientation=True, + ): + Phi_mn, modes_Phi = map(np.asarray, (Phi_mn, modes_Phi)) + assert ( + Phi_mn.size == modes_Phi.shape[0] + ), "Phi_mn size and modes_Phi.shape[0] must be the same size!" + + assert np.issubdtype(modes_Phi.dtype, np.integer) + + M_Phi = setdefault(M_Phi, np.max(abs(modes_Phi[:, 0]))) + N_Phi = setdefault(N_Phi, np.max(abs(modes_Phi[:, 1]))) + + self._M_Phi = M_Phi + self._N_Phi = N_Phi + + self._sym_Phi = sym_Phi + self._Phi_basis = DoubleFourierSeries(M=M_Phi, N=N_Phi, NFP=NFP, sym=sym_Phi) + self._Phi_mn = copy_coeffs(Phi_mn, modes_Phi, self._Phi_basis.modes[:, 1:]) + + self._I = float(np.squeeze(I)) + self._G = float(np.squeeze(G)) + + super().__init__( + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + M=M, + N=N, + name=name, + check_orientation=check_orientation, + ) + + @optimizable_parameter + @property + def I(self): # noqa: E743 + """Net current linking the plasma and the surface toroidally.""" + return self._I + + @I.setter + def I(self, new): # noqa: E743 + self._I = float(np.squeeze(new)) + + @optimizable_parameter + @property + def G(self): + """Net current linking the plasma and the surface poloidally.""" + return self._G + + @G.setter + def G(self, new): + self._G = float(np.squeeze(new)) + + @optimizable_parameter + @property + def Phi_mn(self): + """Fourier coefficients describing single-valued part of potential.""" + return self._Phi_mn + + @Phi_mn.setter + def Phi_mn(self, new): + if len(new) == self.Phi_basis.num_modes: + self._Phi_mn = jnp.asarray(new) + else: + raise ValueError( + f"Phi_mn should have the same size as the basis, got {len(new)} for " + + f"basis with {self.Phi_basis.num_modes} modes." + ) + + @property + def Phi_basis(self): + """DoubleFourierSeries: Spectral basis for Phi.""" + return self._Phi_basis + + @property + def sym_Phi(self): + """str: Type of symmetry of periodic part of Phi (no symmetry if False).""" + return self._sym_Phi + + @property + def M_Phi(self): + """int: Poloidal resolution of periodic part of Phi.""" + return self._M_Phi + + @property + def N_Phi(self): + """int: Toroidal resolution of periodic part of Phi.""" + return self._N_Phi + + def change_Phi_resolution(self, M=None, N=None, NFP=None, sym_Phi=None): + """Change the maximum poloidal and toroidal resolution for Phi. + + Parameters + ---------- + M : int + Poloidal resolution to change Phi basis to. + If None, defaults to current self.Phi_basis poloidal resolution + N : int + Toroidal resolution to change Phi basis to. + If None, defaults to current self.Phi_basis toroidal resolution + NFP : int + Number of field periods for surface and Phi basis. + If None, defaults to current NFP. + Note: will change the NFP of the surface geometry as well as the + Phi basis. + sym_Phi : {"auto","cos","sin",False} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. Default is "auto" which enforces if modes are symmetric. + If True, non-symmetric modes will be truncated. + + """ + M = M or self._M_Phi + N = N or self._M_Phi + NFP = NFP or self.NFP + sym_Phi = sym_Phi or self.sym_Phi + + Phi_modes_old = self.Phi_basis.modes + self.Phi_basis.change_resolution(M=M, N=N, NFP=self.NFP, sym=sym_Phi) + + self._Phi_mn = copy_coeffs(self.Phi_mn, Phi_modes_old, self.Phi_basis.modes) + self._M_Phi = M + self._N_Phi = N + self._sym_Phi = sym_Phi + self.change_resolution( + NFP=NFP + ) # make sure surface and Phi basis NFP are the same + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * max(self.M, self.M_Phi), + N=30 + 2 * max(self.N, self.N_Phi), + NFP=self.NFP, + ) + return _compute_magnetic_field_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + + @classmethod + def from_surface( + cls, + surface, + Phi_mn=np.array([0.0]), + modes_Phi=np.array([[0, 0]]), + I=0, + G=0, + sym_Phi=False, + M_Phi=None, + N_Phi=None, + ): + """Create FourierCurrentPotentialField using geometry of given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface, optional, default None + Existing FourierRZToroidalSurface object to create a + CurrentPotentialField with. + Phi_mn : ndarray + Fourier coefficients of the double FourierSeries of the current potential. + Should correspond to the given DoubleFourierSeries basis object passed in. + modes_Phi : array-like, shape(k,2) + Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn + coefficients + I : float + Net current linking the plasma and the surface toroidally + Denoted I in the algorithm + G : float + Net current linking the plasma and the surface poloidally + Denoted G in the algorithm + NOTE: a negative G will tend to produce a positive toroidal magnetic field + B in DESC, as in DESC the poloidal angle is taken to be positive + and increasing when going in the clockwise direction, which with the + convention n x grad(phi) will result in a toroidal field in the negative + toroidal direction. + sym_Phi : {False,"cos","sin"} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. + M_Phi, N_Phi: int or None + Maximum poloidal and toroidal mode numbers for the single valued part of the + current potential. + + """ + if not isinstance(surface, FourierRZToroidalSurface): + raise TypeError( + "Expected type FourierRZToroidalSurface for argument surface, " + f"instead got type {type(surface)}" + ) + R_lmn = surface.R_lmn + Z_lmn = surface.Z_lmn + modes_R = surface._R_basis.modes[:, 1:] + modes_Z = surface._Z_basis.modes[:, 1:] + NFP = surface.NFP + sym = surface.sym + name = surface.name + + return cls( + Phi_mn=Phi_mn, + modes_Phi=modes_Phi, + I=I, + G=G, + sym_Phi=sym_Phi, + M_Phi=M_Phi, + N_Phi=N_Phi, + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + name=name, + check_orientation=False, + ) + + +def _compute_magnetic_field_from_CurrentPotentialField( + field, + coords, + source_grid, + params=None, + basis="rpz", +): + """Compute magnetic field at a set of points. + + Parameters + ---------- + field : CurrentPotentialField or FourierCurrentPotentialField + current potential field object from which to compute magnetic field. + coords : array-like shape(N,3) + cylindrical or cartesian coordinates + source_grid : Grid, + source grid upon which to evaluate the surface current density K + params : dict, optional + parameters to pass to compute function + should include the potential + basis : {"rpz", "xyz"} + basis for input coordinates and returned magnetic field + + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + if basis == "rpz": + coords = rpz2xyz(coords) + + # compute surface current, and store grid quantities + # needed for integration in class + # TODO: does this have to be xyz, or can it be computed in rpz as well? + data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) + + _rs = xyz2rpz(data["x"]) + _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) + + # surface element, must divide by NFP to remove the NFP multiple on the + # surface grid weights, as we account for that when doing the for loop + # over NFP + _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP + + def nfp_loop(j, f): + # calculate (by rotating) rs, rs_t, rz_t + phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( + 2 * jnp.pi + ) + # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) + rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T + rs = rpz2xyz(rs) + K = rpz2xyz_vec(_K, phi=phi) + fj = biot_savart_general( + coords, + rs, + K, + _dV, + ) + f += fj + return f + + B = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) + if basis == "rpz": + B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) + return B diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py new file mode 100644 index 0000000000..ac5720224a --- /dev/null +++ b/desc/magnetic_fields/_dommaschk.py @@ -0,0 +1,520 @@ +"""Dommaschk potential utility functions. + +based off Representations for vacuum potentials in stellarators +https://doi.org/10.1016/0010-4655(86)90109-8 + +""" + +from desc.backend import cond, fori_loop, gammaln, jit, jnp +from desc.derivatives import Derivative + +from ._core import ScalarPotentialField, _MagneticField + + +class DommaschkPotentialField(ScalarPotentialField): + """Magnetic field due to a Dommaschk scalar magnetic potential in rpz coordinates. + + From Dommaschk 1986 paper https://doi.org/10.1016/0010-4655(86)90109-8 + + this is the field due to the dommaschk potential (eq. 1) for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + ms : 1D array-like of int + first indices of V_m_l terms (eq. 12 of reference) + ls : 1D array-like of int + second indices of V_m_l terms (eq. 12 of reference) + a_arr : 1D array-like of float + a_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*D_m_l terms + b_arr : 1D array-like of float + b_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*D_m_l terms + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*N_m_l-1 term + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms + B0: float + scale strength of the magnetic field's 1/R portion + + """ + + def __init__( + self, + ms=jnp.array([0]), + ls=jnp.array([0]), + a_arr=jnp.array([0.0]), + b_arr=jnp.array([0.0]), + c_arr=jnp.array([0.0]), + d_arr=jnp.array([0.0]), + B0=1.0, + ): + ms = jnp.atleast_1d(jnp.asarray(ms)) + ls = jnp.atleast_1d(jnp.asarray(ls)) + a_arr = jnp.atleast_1d(jnp.asarray(a_arr)) + b_arr = jnp.atleast_1d(jnp.asarray(b_arr)) + c_arr = jnp.atleast_1d(jnp.asarray(c_arr)) + d_arr = jnp.atleast_1d(jnp.asarray(d_arr)) + + assert ( + ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size + ), "Passed in arrays must all be of the same size!" + assert not jnp.any( + jnp.logical_or(ms < 0, ls < 0) + ), "m and l mode numbers must be >= 0!" + assert ( + jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1 + ), "B0 should be a scalar value!" + + params = {} + params["ms"] = ms + params["ls"] = ls + params["a_arr"] = a_arr + params["b_arr"] = b_arr + params["c_arr"] = c_arr + params["d_arr"] = d_arr + params["B0"] = B0 + + super().__init__(dommaschk_potential, params) + + @classmethod + def fit_magnetic_field( # noqa: C901 - FIXME - simplify + cls, field, coords, max_m, max_l, sym=False, verbose=1 + ): + """Fit a vacuum magnetic field with a Dommaschk Potential field. + + Parameters + ---------- + field (MagneticField or callable or ndarray): magnetic field to fit + if callable, must accept (num_nodes,3) array of rpz coords as argument + and output (num_nodes,3) as the B field in cylindrical rpz basis. + if ndarray, must be an ndarray of the magnetic field in rpz, + of shape (num_nodes,3) with the columns being (B_R, B_phi, B_Z) + coords (ndarray): shape (num_nodes,3) of R,phi,Z points to fit field at + max_m (int): maximum m to use for Dommaschk Potentials + max_l (int): maximum l to use for Dommaschk Potentials + sym (bool): if field is stellarator symmetric or not. + if True, only stellarator-symmetric modes will + be included in the fitting + verbose (int): verbosity level of fitting routine, > 0 prints residuals + """ + # We seek c in Ac = b + # A will be the BR, Bphi and BZ from each individual + # dommaschk potential basis function evaluated at each node + # c is the dommaschk potential coefficients + # c will be [B0, a_00, a_10, a_01, a_11... etc] + # b is the magnetic field at each node which we are fitting + if isinstance(field, _MagneticField): + B = field.compute_magnetic_field(coords) + elif callable(field): + B = field(coords) + else: # it must be the field evaluated at the passed-in coords + B = field + # TODO: add basis argument for if passed-in field or callable + # evaluates rpz or xyz basis magnetic field vector, + # and what basis coords is + + ######### + # make b + ######### + # we will have the rhs be 3*num_nodes in length (bc of vector B) + + rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") + + ##################### + # b is made, now do A + ##################### + num_modes = 1 + (max_l + 1) * (max_m + 1) * 4 + # TODO: if symmetric, technically only need half the modes + # however, the field and functions are setup to accept equal + # length arrays for a,b,c,d, so we will just zero out the + # modes that don't fit symmetry, but in future + # should refactor code to have a 3rd index so that + # we have a = V_ml0, b = V_ml1, c = V_ml2, d = V_ml3 + # and the modes array can then be [m,l,x] where x is 0,1,2,3 + # and we dont need to keep track of a,b,c,d separately + + # TODO: technically we can drop some modes + # since if max_l=0, there are only ever nonzero terms for a and b + # and if max_m=0, there are only ever nonzero terms for a and c + # but since we are only fitting in a least squares sense, + # and max_l and max_m should probably be both nonzero anyways, + # this is not an issue right now + + # mode numbers + ms = [] + ls = [] + + # order of coeffs in the vector c are B0, a_ml, b_ml, c_ml, d_ml + a_s = [] + b_s = [] + c_s = [] + d_s = [] + zero_due_to_sym_inds = [] + abcd_zero_due_to_sym_inds = [ + [], + [], + [], + [], + ] # indices that should be 0 due to symmetry + # if sym is True, when l is even then we need a=d=0 + # and if l is odd then b=c=0 + + for l in range(max_l + 1): + for m in range(max_m + 1): + if not sym: + pass # no sym, use all coefs + elif l // 2 == 0: + zero_due_to_sym_inds = [0, 3] # a=d=0 for even l with sym + elif l // 2 == 1: + zero_due_to_sym_inds = [1, 2] # b=c=0 for odd l with sym + for which_coef in range(4): + if which_coef == 0: + a_s.append(1) + elif which_coef == 1: + b_s.append(1) + elif which_coef == 2: + c_s.append(1) + elif which_coef == 3: + d_s.append(1) + if which_coef in zero_due_to_sym_inds: + abcd_zero_due_to_sym_inds[which_coef].append(0) + else: + abcd_zero_due_to_sym_inds[which_coef].append(1) + + ms.append(m) + ls.append(l) + for i in range(4): + abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) + assert (len(a_s) + len(b_s) + len(c_s) + len(d_s)) == num_modes - 1 + + params = { + "ms": ms, + "ls": ls, + "a_arr": a_s, + "b_arr": b_s, + "c_arr": c_s, + "d_arr": d_s, + "B0": 0.0, + } + n = ( + round(num_modes - 1) / 4 + ) # how many l-m mode pairs there are, also is len(a_s) + n = int(n) + domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) + + def get_B_dom(coords, X, ms, ls): + """Fxn wrapper to find jacobian of dommaschk B wrt coefs a,b,c,d.""" + # zero out any terms that should be zero due to symmetry, which + # we cataloged earlier for each a_arr,b_arr,c_arr,d_arr + # that way the resulting modes after pinv don't contain them either + return domm_field.compute_magnetic_field( + coords, + params={ + "ms": jnp.asarray(ms), + "ls": jnp.asarray(ls), + "a_arr": jnp.asarray(X[1 : n + 1]) * abcd_zero_due_to_sym_inds[0], + "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]) + * abcd_zero_due_to_sym_inds[1], + "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]) + * abcd_zero_due_to_sym_inds[2], + "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]) + * abcd_zero_due_to_sym_inds[3], + "B0": X[0], + }, + ) + + X = [] + for key in ["B0", "a_arr", "b_arr", "c_arr", "d_arr"]: + obj = params[key] + if isinstance(obj, list): + X += obj + else: + X += [obj] + X = jnp.asarray(X) + + jac = jit(Derivative(get_B_dom, argnum=1))( + coords, X, params["ms"], params["ls"] + ) + + A = jac.reshape((rhs.size, len(X)), order="F") + + # now solve Ac=b for the coefficients c + + # TODO: use min singular value to give sense of cond number? + c, res, _, _ = jnp.linalg.lstsq(A, rhs) + + if verbose > 0: + # res is a list of len(1) so index into it + print(f"Sum of Squares Residual of fit: {res[0]:1.4e} T") + + # recover the params from the c coefficient vector + B0 = c[0] + + # we zero out the terms that should be zero due to symmetry here + # TODO: should also just not return any zeroed-out modes, but + # the way the modes are cataloged here with the ls and ms arrays, + # it is not straightforward to do that. By that I mean + # say ls = [1,2] ms = [1,1] + a_arr = c[1 : n + 1] * abcd_zero_due_to_sym_inds[0] + b_arr = c[n + 1 : 2 * n + 1] * abcd_zero_due_to_sym_inds[1] + c_arr = c[2 * n + 1 : 3 * n + 1] * abcd_zero_due_to_sym_inds[2] + d_arr = c[3 * n + 1 : 4 * n + 1] * abcd_zero_due_to_sym_inds[3] + + return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) + + +true_fun = lambda m_n: 0.0 # used for returning 0 when conditionals evaluate to True + + +@jit +def gamma(n): + """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" + return jnp.exp(gammaln(n)) + + +@jit +def alpha(m, n): + """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + # modified for eqns 31 and 32 + + def false_fun(m_n): + m, n = m_n + return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) + + def bool_fun(n): + return n < 0 + + return cond( + bool_fun(n), + true_fun, + false_fun, + ( + m, + n, + ), + ) + + +@jit +def alphastar(m, n): + """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + + def false_fun(m_n): # modified for eqns 31 and 32 + m, n = m_n + return (2 * n + m) * alpha(m, n) + + return cond(n < 0, true_fun, false_fun, (m, n)) + + +@jit +def beta(m, n): + """Beta of eq 28, modified for eqns 31 and 32.""" + + def false_fun(m_n): + m, n = m_n + return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) + + +@jit +def betastar(m, n): + """Beta* of eq 28, modified for eqns 31 and 32.""" + + def false_fun(m_n): + m, n = m_n + return (2 * n - m) * beta(m, n) + + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) + + +@jit +def gamma_n(m, n): + """gamma_n of eq 33.""" + + def body_fun(i, val): + return val + 1 / i + 1 / (m + i) + + def false_fun(m_n): + m, n = m_n + return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) + + return cond(n <= 0, true_fun, false_fun, (m, n)) + + +@jit +def gamma_nstar(m, n): + """gamma_n star of eq 33.""" + + def false_fun(m_n): + m, n = m_n + return (2 * n + m) * gamma_n(m, n) + + return cond(n <= 0, true_fun, false_fun, (m, n)) + + +@jit +def CD_m_k(R, m, k): + """Eq 31 of Dommaschk paper.""" + + def body_fun(j, val): + result = ( + val + + ( + -( + alpha(m, j) + * ( + alphastar(m, k - m - j) * jnp.log(R) + + gamma_nstar(m, k - m - j) + - alpha(m, k - m - j) + ) + - gamma_n(m, j) * alphastar(m, k - m - j) + + alpha(m, j) * betastar(m, k - j) + ) + * R ** (2 * j + m) + ) + + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def CN_m_k(R, m, k): + """Eq 32 of Dommaschk paper.""" + + def body_fun(j, val): + result = ( + val + + ( + ( + alpha(m, j) + * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) + - gamma_n(m, j) * alpha(m, k - m - j) + + alpha(m, j) * beta(m, k - j) + ) + * R ** (2 * j + m) + ) + - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def D_m_n(R, Z, m, n): + """D_m_n term in eqn 8 of Dommaschk paper.""" + # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper + + def body_fun(k, val): + coef = CD_m_k(R, m, k) / gamma(n - 2 * k + 1) + exp = n - 2 * k + # derivative of 0**0 is ill defined, so we do this to enforce it being 0 + exp = jnp.where((Z == 0) & (exp == 0), 1, exp) + return val + coef * Z**exp + + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def N_m_n(R, Z, m, n): + """N_m_n term in eqn 9 of Dommaschk paper.""" + # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper + + def body_fun(k, val): + coef = CN_m_k(R, m, k) / gamma(n - 2 * k + 1) + exp = n - 2 * k + # derivative of 0**0 is ill defined, so we do this to enforce it being 0 + exp = jnp.where((Z == 0) & (exp == 0), 1, exp) + return val + coef * Z**exp + + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def V_m_l(R, phi, Z, m, l, a, b, c, d): + """Eq 12 of Dommaschk paper. + + Parameters + ---------- + R,phi,Z : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + m : int + first index of V_m_l term + l : int + second index of V_m_l term + a : float + a_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*D_m_l + b : float + b_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*D_m_l + c : float + c_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*N_m_l-1 + d : float + d_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*N_m_l-1 + + Returns + ------- + value : array-like + Value of this V_m_l term evaluated at the given R,phi,Z points + (same size as the size of the given R,phi, or Z arrays). + + """ + return (a * jnp.cos(m * phi) + b * jnp.sin(m * phi)) * D_m_n(R, Z, m, l) + ( + c * jnp.cos(m * phi) + d * jnp.sin(m * phi) + ) * N_m_n(R, Z, m, l - 1) + + +@jit +def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): + """Eq 1 of Dommaschk paper. + + this is the total dommaschk potential for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + R,phi,Z : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + ms : 1D array-like of int + first indices of V_m_l terms + ls : 1D array-like of int + second indices of V_m_l terms + a_arr : 1D array-like of float + a_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*D_m_l + b_arr : 1D array-like of float + b_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*D_m_l + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*N_m_l-1 + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*N_m_l-1 + B0: float, toroidal magnetic field strength scale, this is the strength of the + 1/R part of the magnetic field and is the Bphi at R=1. + + Returns + ------- + value : array-like + Value of the total dommaschk potential evaluated + at the given R,phi,Z points + (same size as the size of the given R,phi, Z arrays). + """ + ms, ls, a_arr, b_arr, c_arr, d_arr = map( + jnp.atleast_1d, (ms, ls, a_arr, b_arr, c_arr, d_arr) + ) + R, phi, Z = map(jnp.atleast_1d, (R, phi, Z)) + R, phi, Z = jnp.broadcast_arrays(R, phi, Z) + ms, ls, a_arr, b_arr, c_arr, d_arr = jnp.broadcast_arrays( + ms, ls, a_arr, b_arr, c_arr, d_arr + ) + value = B0 * phi # phi term + + def body(i, val): + val += V_m_l(R, phi, Z, ms[i], ls[i], a_arr[i], b_arr[i], c_arr[i], d_arr[i]) + return val + + return fori_loop(0, len(ms), body, value) diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index 6fbc8bd8c4..d6f7ad6c19 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -636,14 +636,14 @@ def compute(self, eq_params, constants=None): "Phi_mn": eq_params["Phi_mn"], } sheet_source_data = compute_fun( - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", self._sheet_data_keys, params=sheet_params, transforms=constants["sheet_source_transforms"], profiles={}, ) sheet_eval_data = compute_fun( - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", self._sheet_data_keys, params=sheet_params, transforms=constants["sheet_eval_transforms"], diff --git a/docs/api.rst b/docs/api.rst index f2d791224d..1375a4c307 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -142,6 +142,7 @@ Magnetic Fields desc.magnetic_fields.PoloidalMagneticField desc.magnetic_fields.SplineMagneticField desc.magnetic_fields.ScalarPotentialField + desc.magnetic_fields.DommaschkPotentialField desc.magnetic_fields.field_line_integrate desc.magnetic_fields.read_BNORM_file diff --git a/docs/api_fields.rst b/docs/api_fields.rst index d3ba0614f4..7e183117e8 100644 --- a/docs/api_fields.rst +++ b/docs/api_fields.rst @@ -22,6 +22,7 @@ a surface and computes the normal field strength on that surface. :template: class.rst desc.magnetic_fields.SplineMagneticField + desc.magnetic_fields.DommaschkPotentialField desc.magnetic_fields.ScalarPotentialField desc.magnetic_fields.ToroidalMagneticField desc.magnetic_fields.VerticalMagneticField diff --git a/setup.cfg b/setup.cfg index 033cbe2a0b..20e90c9dde 100644 --- a/setup.cfg +++ b/setup.cfg @@ -71,12 +71,7 @@ ignore = D105, per-file-ignores = # need to import things to top level even if they aren't used there - desc/compute/__init__.py: F401 - desc/equilibrium/__init__.py: F401 - desc/geometry/__init__.py: F401 - desc/io/__init__.py: F401 - desc/objectives/__init__.py: F401 - desc/optimize/__init__.py: F401 + desc/*/__init__.py: F401 # too many long lines to deal with now desc/compute/data_index.py: E501 # need imports in weird order for selecting device before benchmarks diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl index 702deaeb77..8ff60241da 100644 Binary files a/tests/inputs/master_compute_data.pkl and b/tests/inputs/master_compute_data.pkl differ diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 7d65b6b5ac..bb9c6dd167 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1239,14 +1239,14 @@ def test_compute_everything(): **elliptic_cross_section_with_torsion ), # magnetic fields - "desc.magnetic_fields.CurrentPotentialField": CurrentPotentialField( + "desc.magnetic_fields._current_potential.CurrentPotentialField": CurrentPotentialField( # noqa:E501 **elliptic_cross_section_with_torsion, potential=lambda theta, zeta, G: G * zeta / 2 / np.pi, potential_dtheta=lambda theta, zeta, G: np.zeros_like(theta), potential_dzeta=lambda theta, zeta, G: G * np.ones_like(theta) / 2 / np.pi, params={"G": 1e7}, ), - "desc.magnetic_fields.FourierCurrentPotentialField": ( + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": ( FourierCurrentPotentialField( **elliptic_cross_section_with_torsion, I=0, G=1e7 ) @@ -1270,7 +1270,7 @@ def test_compute_everything(): ), } assert things.keys() == data_index.keys(), ( - f"Missing the parameterizations {data_index.keys() - things.keys()}" + f"Missing the parameterization {data_index.keys() - things.keys()}" f" to test against master." ) # use this low resolution grid for equilibria to reduce file size diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index c1bbb89417..11e584f3af 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -13,6 +13,7 @@ from desc.io import load from desc.magnetic_fields import ( CurrentPotentialField, + DommaschkPotentialField, FourierCurrentPotentialField, PoloidalMagneticField, ScalarPotentialField, @@ -22,6 +23,7 @@ field_line_integrate, read_BNORM_file, ) +from desc.magnetic_fields._dommaschk import CD_m_k, CN_m_k def phi_lm(R, phi, Z, a, m): @@ -707,7 +709,7 @@ def test_Bnormal_calculation(self): surface = get("DSHAPE").surface Bnorm, _ = tfield.compute_Bnormal(surface) # should have 0 Bnormal because surface is axisymmetric - np.testing.assert_allclose(Bnorm, 0, atol=3e-15) + np.testing.assert_allclose(Bnorm, 0, atol=1e-14) @pytest.mark.unit def test_Bnormal_save_and_load_DSHAPE(self, tmpdir_factory): @@ -770,6 +772,170 @@ def test_Bnormal_save_and_load_HELIOTRON(self, tmpdir_factory): with pytest.raises(AssertionError, match="sym"): Bnorm_from_file = read_BNORM_file(path, asym_surf, grid) + +@pytest.mark.unit +def test_dommaschk_CN_CD_m_0(): + """Test of CD_m_k and CN_m_k when k=0.""" + # based off eqn 8 and 9 of Dommaschk paper + # https://doi.org/10.1016/0010-4655(86)90109-8 + for m in range(1, 6): + # test of CD_m_k based off eqn 8 + R = np.linspace(0.1, 1, 100) + res1 = CD_m_k(R, m, 0) + res2 = 0.5 * (R**m + R ** (-m)) + np.testing.assert_allclose(res1, res2) + + # test of CN_m_k based off eqn 9 + res1 = CN_m_k(R, m, 0) + res2 = 0.5 * (R**m - R ** (-m)) / m + np.testing.assert_allclose(res1, res2, atol=1e-15) + + +@pytest.mark.unit +def test_dommaschk_field_errors(): + """Test the assert statements of the DommaschkField function.""" + ms = [1] + ls = [1] + a_arr = [1] + b_arr = [1] + c_arr = [1] + d_arr = [1, 1] # length is not equal to the rest + with pytest.raises(AssertionError): + DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + d_arr = [1] + ms = [-1] # negative mode number + with pytest.raises(AssertionError): + DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + + +@pytest.mark.unit +def test_dommaschk_radial_field(): + """Test the Dommaschk potential for a pure toroidal (Bphi~1/R) field.""" + phi = np.linspace(0, 2 * np.pi, 10) + R = np.linspace(0.1, 1.5, 50) + Z = np.linspace(-0.05, 0.5, 50) + R, phi, Z = np.meshgrid(R, phi, Z) + coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + ms = [0] + ls = [0] + a_arr = [0] + b_arr = [0] + c_arr = [0] + d_arr = [0] + B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0) + np.testing.assert_array_equal(B_dom[:, 1], 1 / R.flatten()) + np.testing.assert_allclose(B_dom[:, 2], 0) + + +@pytest.mark.unit +def test_dommaschk_vertical_field(): + """Test the Dommaschk potential for a 1/R toroidal + pure vertical field.""" + phi = jnp.linspace(0, 2 * jnp.pi, 10) + R = jnp.linspace(0.1, 1.5, 50) + Z = jnp.linspace(-0.5, 0.5, 50) + R, phi, Z = jnp.meshgrid(R, phi, Z) + coords = jnp.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + ms = [0] + ls = [1] + a_arr = [1] + b_arr = [0] + c_arr = [0] + d_arr = [0] + B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + B_dom = B.compute_magnetic_field(coords) + ones = jnp.ones_like(B_dom[:, 0]) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-14) + np.testing.assert_allclose(B_dom[:, 2], ones, atol=5e-15) + + +@pytest.mark.unit +def test_dommaschk_fit_toroidal_field(): + """Test the Dommaschk potential fit for a 1/R toroidal scaled to 2 T.""" + phi = np.linspace(0, 2 * jnp.pi, 3) + R = np.linspace(0.1, 1.5, 3) + Z = np.linspace(-0.5, 0.5, 3) + R, phi, Z = np.meshgrid(R, phi, Z) + coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + max_l = 2 + max_m = 2 + B0 = 2 # scale strength for to 1/R field to fit + + def B0_over_R(coord): + B_R = np.zeros_like(coord[:, 0]) + B_phi = B0 / coord[:, 0] + B_Z = np.zeros_like(coord[:, 0]) + return np.vstack((B_R, B_phi, B_Z)).T + + B = DommaschkPotentialField.fit_magnetic_field( + B0_over_R, coords, max_m, max_l, sym=True + ) + + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=2e-14) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=2e-14) + np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=2e-14) + + # only nonzero coefficient of the field should be the B0 + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14) + for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]: + np.testing.assert_allclose(B._params[coef], 0, atol=1e-14) + + # test fit from values + B = DommaschkPotentialField.fit_magnetic_field( + B0_over_R(coords), coords, max_m, max_l, sym=True + ) + + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=2e-14) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=2e-14) + np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=2e-14) + + # only nonzero coefficient of the field should be the B0 + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14) + for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]: + np.testing.assert_allclose(B._params[coef], 0, atol=1e-14) + + +@pytest.mark.unit +def test_dommaschk_fit_vertical_and_toroidal_field(): + """Test the Dommaschk potential fit for a toroidal and a vertical field.""" + phi = np.linspace(0, 2 * np.pi, 3) + R = np.linspace(0.1, 1.5, 3) + Z = np.linspace(-0.5, 0.5, 3) + R, phi, Z = np.meshgrid(R, phi, Z) + coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + max_l = 1 + max_m = 1 + B0 = 2 # scale strength for to 1/R field to fit + B0_Z = 1 # scale strength for to uniform vertical field to fit + field = ToroidalMagneticField(B0=B0, R0=1) + VerticalMagneticField(B0=B0_Z) + + B = DommaschkPotentialField.fit_magnetic_field(field, coords, max_m, max_l) + + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-14) + np.testing.assert_allclose(B_dom[:, 2], B0_Z, atol=1e-14) + + np.testing.assert_allclose(B._params["B0"], B0) + + # only nonzero coefficient of the field should be the B0 and a_ml = a_01 + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14) + for coef, m, l in zip(B._params["a_arr"], B._params["ms"], B._params["ls"]): + if m == 0 and l == 1: + np.testing.assert_allclose(coef, B0_Z) + else: + np.testing.assert_allclose(coef, 0, atol=1e-14) + for name in ["b_arr", "c_arr", "d_arr"]: + np.testing.assert_allclose(B._params[name], 0, atol=1e-14) + @pytest.mark.unit def test_spline_field_jit(self): """Test that the spline field can be passed to a jitted function."""