Skip to content

Commit

Permalink
remove _AB fxns for simpler fields
Browse files Browse the repository at this point in the history
  • Loading branch information
@dpanici
dpanici committed Aug 29, 2024
1 parent 0505027 commit 7bb7bd3
Showing 1 changed file with 55 additions and 127 deletions.
182 changes: 55 additions & 127 deletions desc/magnetic_fields/_core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1074,16 +1074,10 @@ def B0(self):
def B0(self, new):
self._B0 = float(np.squeeze(new))

def _compute_A_or_B(
self,
coords,
params=None,
basis="rpz",
source_grid=None,
transforms=None,
compute_A_or_B="B",
def compute_magnetic_field(
self, coords, params=None, basis="rpz", source_grid=None, transforms=None
):
"""Compute magnetic field or vector potential at a set of points.
"""Compute magnetic field at a set of points.
Parameters
----------
Expand All @@ -1098,21 +1092,13 @@ def _compute_A_or_B(
transforms : dict of Transform
Transforms for R, Z, lambda, etc. Default is to build from source_grid
Unused by this MagneticField class.
compute_A_or_B: {"A", "B"}, optional
whether to compute the magnetic vector potential "A" or the magnetic field
"B". Defaults to "B"
Returns
-------
field : ndarray, shape(N,3)
magnetic field or vector potential at specified points
magnetic field at specified points
"""
errorif(
compute_A_or_B not in ["A", "B"],
ValueError,
f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}',
)
params = setdefault(params, {})
B0 = params.get("B0", self.B0)
R0 = params.get("R0", self.R0)
Expand All @@ -1121,48 +1107,13 @@ def _compute_A_or_B(
coords = jnp.atleast_2d(jnp.asarray(coords))
if basis == "xyz":
coords = xyz2rpz(coords)
if compute_A_or_B == "B":
bp = B0 * R0 / coords[:, 0]
brz = jnp.zeros_like(bp)
B = jnp.array([brz, bp, brz]).T
if basis == "xyz":
B = rpz2xyz_vec(B, phi=coords[:, 1])

return B
elif compute_A_or_B == "A":
az = -B0 * R0 * jnp.log(coords[:, 0])
arp = jnp.zeros_like(az)
A = jnp.array([arp, arp, az]).T
# b/c it only has a nonzero z component, no need
# to switch bases back if xyz is given
return A

def compute_magnetic_field(
self, coords, params=None, basis="rpz", source_grid=None, transforms=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
Dict of values for R0 and B0.
basis : {"rpz", "xyz"}
Basis for input coordinates and returned magnetic field.
source_grid : Grid, int or None or array-like, optional
Unused by this MagneticField class.
transforms : dict of Transform
Transforms for R, Z, lambda, etc. Default is to build from source_grid
Unused by this MagneticField class.
Returns
-------
field : ndarray, shape(N,3)
magnetic field at specified points
bp = B0 * R0 / coords[:, 0]
brz = jnp.zeros_like(bp)
B = jnp.array([brz, bp, brz]).T
if basis == "xyz":
B = rpz2xyz_vec(B, phi=coords[:, 1])

"""
return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B")
return B

def compute_magnetic_vector_potential(
self, coords, params=None, basis="rpz", source_grid=None, transforms=None
Expand Down Expand Up @@ -1191,7 +1142,20 @@ def compute_magnetic_vector_potential(
magnetic vector potential at specified points
"""
return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A")
params = setdefault(params, {})
B0 = params.get("B0", self.B0)
R0 = params.get("R0", self.R0)

assert basis.lower() in ["rpz", "xyz"]
coords = jnp.atleast_2d(jnp.asarray(coords))
if basis == "xyz":
coords = xyz2rpz(coords)

Check warning on line 1152 in desc/magnetic_fields/_core.py

View check run for this annotation

Codecov / codecov/patch

desc/magnetic_fields/_core.py#L1152 

Added line #L1152 was not covered by tests
az = -B0 * R0 * jnp.log(coords[:, 0])
arp = jnp.zeros_like(az)
A = jnp.array([arp, arp, az]).T
# b/c it only has a nonzero z component, no need
# to switch bases back if xyz is given
return A


class VerticalMagneticField(_MagneticField, Optimizable):
Expand Down Expand Up @@ -1221,71 +1185,6 @@ def B0(self):
def B0(self, new):
self._B0 = float(np.squeeze(new))

def _compute_A_or_B(
self,
coords,
params=None,
basis="rpz",
source_grid=None,
transforms=None,
compute_A_or_B="B",
):
"""Compute magnetic field or magnetic vector potential 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
Dict of values for B0.
basis : {"rpz", "xyz"}
Basis for input coordinates and returned magnetic field.
source_grid : Grid, int or None or array-like, optional
Unused by this MagneticField class.
compute_A_or_B: {"A", "B"}, optional
whether to compute the magnetic vector potential "A" or the magnetic field
"B". Defaults to "B"
Returns
-------
field : ndarray, shape(N,3)
magnetic field or vector potential at specified points
"""
errorif(
compute_A_or_B not in ["A", "B"],
ValueError,
f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}',
)
params = setdefault(params, {})
B0 = params.get("B0", self.B0)

coords = jnp.atleast_2d(jnp.asarray(coords))
if compute_A_or_B == "B":
bz = B0 * jnp.ones_like(coords[:, 2])
brp = jnp.zeros_like(bz)
B = jnp.array([brp, brp, bz]).T
# b/c it only has a nonzero z component, no need
# to switch bases back if xyz is given

return B
elif compute_A_or_B == "A":
if basis == "xyz":
coords_xyz = coords
coords_rpz = xyz2rpz(coords)
else:
coords_rpz = coords
coords_xyz = rpz2xyz(coords)
ax = B0 / 2 * coords_xyz[:, 1]
ay = -B0 / 2 * coords_xyz[:, 0]

az = jnp.zeros_like(ax)
A = jnp.array([ax, ay, az]).T
if basis == "rpz":
A = xyz2rpz_vec(A, phi=coords_rpz[:, 1])

return A

def compute_magnetic_field(
self, coords, params=None, basis="rpz", source_grid=None, transforms=None
):
Expand All @@ -1311,7 +1210,17 @@ def compute_magnetic_field(
magnetic field at specified points
"""
return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B")
params = setdefault(params, {})
B0 = params.get("B0", self.B0)

coords = jnp.atleast_2d(jnp.asarray(coords))
bz = B0 * jnp.ones_like(coords[:, 2])
brp = jnp.zeros_like(bz)
B = jnp.array([brp, brp, bz]).T
# b/c it only has a nonzero z component, no need
# to switch bases back if xyz is given

return B

def compute_magnetic_vector_potential(
self, coords, params=None, basis="rpz", source_grid=None, transforms=None
Expand Down Expand Up @@ -1340,7 +1249,26 @@ def compute_magnetic_vector_potential(
magnetic vector potential at specified points
"""
return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A")
params = setdefault(params, {})
B0 = params.get("B0", self.B0)

coords = jnp.atleast_2d(jnp.asarray(coords))

if basis == "xyz":
coords_xyz = coords
coords_rpz = xyz2rpz(coords)

Check warning on line 1259 in desc/magnetic_fields/_core.py

View check run for this annotation

Codecov / codecov/patch

desc/magnetic_fields/_core.py#L1258-L1259 

Added lines #L1258 - L1259 were not covered by tests
else:
coords_rpz = coords
coords_xyz = rpz2xyz(coords)
ax = B0 / 2 * coords_xyz[:, 1]
ay = -B0 / 2 * coords_xyz[:, 0]

az = jnp.zeros_like(ax)
A = jnp.array([ax, ay, az]).T
if basis == "rpz":
A = xyz2rpz_vec(A, phi=coords_rpz[:, 1])

return A


class PoloidalMagneticField(_MagneticField, Optimizable):
Expand Down

0 comments on commit 7bb7bd3

Please sign in to comment.