Skip to content

Commit

Permalink
Merge pull request #377 from hiddenSymmetries/circularcoil_dofs_update
Browse files Browse the repository at this point in the history 
Updated CircularCoil class to alllow for dofs
  • Loading branch information
@landreman
landreman authored Apr 25, 2024
2 parents d830719 + d82cb17 commit c243322
Show file tree
Hide file tree
Showing 3 changed files with 283 additions and 37 deletions.
125 changes: 95 additions & 30 deletions src/simsopt/field/magneticfieldclasses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -359,14 +359,44 @@ def __init__(self, r0=0.1, center=[0, 0, 0], I=5e5/np.pi, normal=[0, 0]):
self.Inorm = I*4e-7
self.center = center
self.normal = normal
if len(normal) == 2:
theta = normal[0]
phi = normal[1]
else:
theta = np.arctan2(normal[1], normal[0])
phi = np.arctan2(np.sqrt(normal[0]**2+normal[1]**2), normal[2])

self.rotMatrix = np.array([
super().__init__(x0=self.get_dofs(), names=self._make_names(),external_dof_setter=CircularCoil.set_dofs_impl)

def _make_names(self):
if len(self.normal)==2:
normal_names = ['theta', 'phi']
elif len(self.normal)==3:
normal_names = ['x', 'y', 'z']
return ['r0', 'x0', 'y0', 'z0', 'Inorm'] + normal_names

def num_dofs(self):
return 5+len(self.normal)

def get_dofs(self):
return np.concatenate([np.array([self.r0]), np.array(self.center), np.array([self.Inorm]), np.array(self.normal)])

def set_dofs_impl(self, dofs):
self.r0 = dofs[0]
self.center = dofs[1:4].tolist()
self.Inorm = dofs[4]
self.normal = dofs[5:].tolist()

@property
def I(self):
return self.Inorm * 25e5

def _rotmat(self):
if len(self.normal) == 2:
theta = self.get('theta')
phi = self.get('phi')
else:
xn = self.get('x')
yn = self.get('y')
zn = self.get('z')
theta = np.arctan2(yn, xn)
phi = np.arctan2(np.sqrt(xn**2+yn**2), zn)

m = np.array([
[np.cos(phi) * np.cos(theta)**2 + np.sin(theta)**2,
-np.sin(phi / 2)**2 * np.sin(2 * theta),
np.cos(theta) * np.sin(phi)],
Expand All @@ -377,31 +407,31 @@ def __init__(self, r0=0.1, center=[0, 0, 0], I=5e5/np.pi, normal=[0, 0]):
-np.sin(phi) * np.sin(theta),
np.cos(phi)]
])
return m

self.rotMatrixInv = np.array(self.rotMatrix.T)

@property
def I(self):
return self.Inorm * 25e5
def _rotmatinv(self):
m = self._rotmat()
minv = np.array(m.T)
return minv

def _B_impl(self, B):
points = self.get_points_cart_ref()
points = np.array(np.dot(self.rotMatrixInv, np.array(np.subtract(points, self.center)).T).T)
points = np.array(np.dot(self._rotmatinv(), np.array(np.subtract(points, self.center)).T).T)
rho = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]))
r = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]) + np.square(points[:, 2]))
alpha = np.sqrt(self.r0**2 + np.square(r) - 2*self.r0*rho)
beta = np.sqrt(self.r0**2 + np.square(r) + 2*self.r0*rho)
k = np.sqrt(1-np.divide(np.square(alpha), np.square(beta)))
ellipek2 = ellipe(k**2)
ellipkk2 = ellipk(k**2)
B[:] = np.dot(self.rotMatrix, np.array(
B[:] = np.dot(self._rotmat(), np.array(
[self.Inorm*points[:, 0]*points[:, 2]/(2*alpha**2*beta*rho**2+1e-31)*((self.r0**2+r**2)*ellipek2-alpha**2*ellipkk2),
self.Inorm*points[:, 1]*points[:, 2]/(2*alpha**2*beta*rho**2+1e-31)*((self.r0**2+r**2)*ellipek2-alpha**2*ellipkk2),
self.Inorm/(2*alpha**2*beta+1e-31)*((self.r0**2-r**2)*ellipek2+alpha**2*ellipkk2)])).T

def _dB_by_dX_impl(self, dB):
points = self.get_points_cart_ref()
points = np.array(np.dot(self.rotMatrixInv, np.array(np.subtract(points, self.center)).T).T)
points = np.array(np.dot(self._rotmatinv(), np.array(np.subtract(points, self.center)).T).T)
rho = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]))
r = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]) + np.square(points[:, 2]))
alpha = np.sqrt(self.r0**2 + np.square(r) - 2*self.r0*rho)
Expand Down Expand Up @@ -459,19 +489,19 @@ def _dB_by_dX_impl(self, dB):
[dBxdz, dBydz, dBzdz]])

dB[:] = np.array([
[np.dot(self.rotMatrixInv[:, 0], np.dot(self.rotMatrix[0, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 1], np.dot(self.rotMatrix[0, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 2], np.dot(self.rotMatrix[0, :], dB_by_dXm))],
[np.dot(self.rotMatrixInv[:, 0], np.dot(self.rotMatrix[1, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 1], np.dot(self.rotMatrix[1, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 2], np.dot(self.rotMatrix[1, :], dB_by_dXm))],
[np.dot(self.rotMatrixInv[:, 0], np.dot(self.rotMatrix[2, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 1], np.dot(self.rotMatrix[2, :], dB_by_dXm)),
np.dot(self.rotMatrixInv[:, 2], np.dot(self.rotMatrix[2, :], dB_by_dXm))]]).T
[np.dot(self._rotmatinv()[:, 0], np.dot(self._rotmat()[0, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 1], np.dot(self._rotmat()[0, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 2], np.dot(self._rotmat()[0, :], dB_by_dXm))],
[np.dot(self._rotmatinv()[:, 0], np.dot(self._rotmat()[1, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 1], np.dot(self._rotmat()[1, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 2], np.dot(self._rotmat()[1, :], dB_by_dXm))],
[np.dot(self._rotmatinv()[:, 0], np.dot(self._rotmat()[2, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 1], np.dot(self._rotmat()[2, :], dB_by_dXm)),
np.dot(self._rotmatinv()[:, 2], np.dot(self._rotmat()[2, :], dB_by_dXm))]]).T

def _A_impl(self, A):
points = self.get_points_cart_ref()
points = np.array(np.dot(self.rotMatrixInv, np.array(np.subtract(points, self.center)).T).T)
points = np.array(np.dot(self._rotmatinv(), np.array(np.subtract(points, self.center)).T).T)
rho = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]))
r = np.sqrt(np.square(points[:, 0]) + np.square(points[:, 1]) + np.square(points[:, 2]))
alpha = np.sqrt(self.r0**2 + np.square(r) - 2*self.r0*rho)
Expand All @@ -484,7 +514,7 @@ def _A_impl(self, A):
denom = ((points[:, 0]**2+points[:, 1]**2+1e-31)*np.sqrt(self.r0**2+points[:, 0]**2+points[:, 1]**2+2*self.r0*np.sqrt(points[:, 0]**2+points[:, 1]**2)+points[:, 2]**2+1e-31))
fak = num/denom
pts = fak[:, None]*np.concatenate((-points[:, 1][:, None], points[:, 0][:, None], np.zeros((points.shape[0], 1))), axis=-1)
A[:] = -self.Inorm/2*np.dot(self.rotMatrix, pts.T).T
A[:] = -self.Inorm/2*np.dot(self._rotmat(), pts.T).T

def as_dict(self, serial_objs_dict):
d = super().as_dict(serial_objs_dict=serial_objs_dict)
Expand All @@ -499,6 +529,44 @@ def from_dict(cls, d, serial_objs_dict, recon_objs):
field.set_points_cart(xyz)
return field

def gamma(self, points=64):
"""Export points of the coil."""

angle_points = np.linspace(0, 2*np.pi, points+1)[:-1]

x = self.r0 * np.cos(angle_points)
y = self.r0 * np.sin(angle_points)
z = 0 * angle_points

coords = np.add(np.dot(self._rotmat(), np.column_stack([x,y,z]).T).T, self.center)
return coords

def to_vtk(self, filename, close=False):
"""
Export circular coil to VTK format
Args:
filename: Name of the file to write.
close: Whether to draw the segment from the last quadrature point back to the first.
"""
from pyevtk.hl import polyLinesToVTK

def wrap(data):
return np.concatenate([data, [data[0]]])

# get the coordinates
if close:
x = wrap(self.gamma()[:, 0])
y = wrap(self.gamma()[:, 1])
z = wrap(self.gamma()[:, 2])
ppl = np.asarray([self.gamma().shape[0]+1])
else:
x = self.gamma()[:, 0]
y = self.gamma()[:, 1]
z = self.gamma()[:, 2]
ppl = np.asarray([self.gamma().shape[0]])

polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl)

class DipoleField(MagneticField):
r"""
Expand Down Expand Up @@ -675,10 +743,7 @@ def _toVTK(self, vtkname):
mtheta_normalized = np.ascontiguousarray(mtheta / self.m_maxima)

# Save all the data to a vtk file which can be visualized nicely with ParaView
data = {"m": (mx, my, mz), "m_normalized": (mx_normalized, my_normalized, mz_normalized),
"m_rphiz": (mr, mphi, mz), "m_rphiz_normalized": (mr_normalized, mphi_normalized, mz_normalized),
"m_rphitheta": (mrminor, mphi, mtheta),
"m_rphitheta_normalized": (mrminor_normalized, mphi_normalized, mtheta_normalized)}
data = {"m": (mx, my, mz), "m_normalized": (mx_normalized, my_normalized, mz_normalized), "m_rphiz": (mr, mphi, mz), "m_rphiz_normalized": (mr_normalized, mphi_normalized, mz_normalized), "m_rphitheta": (mrminor, mphi, mtheta), "m_rphitheta_normalized": (mrminor_normalized, mphi_normalized, mtheta_normalized)}
from pyevtk.hl import pointsToVTK
pointsToVTK(str(vtkname), ox, oy, oz, data=data)

Expand Down
53 changes: 46 additions & 7 deletions tests/field/test_magneticfields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,20 @@ def test_circularcoil_Bfield(self):
Bfield_regen = json.loads(field_json_str, cls=GSONDecoder)
self.assertTrue(np.allclose(Bfield.B(), Bfield_regen.B()))

def compare_gammas(circular_coil, general_coil):
# Verify that the gamma values are the same, up to a shift in the
# array index.
gamma1 = general_coil.curve.gamma()
gamma2 = circular_coil.gamma(len(curve.quadpoints))
if general_coil.current.get_value() * circular_coil.I < 0:
# Currents are opposite sign, so the direction of the points
# will be reversed.
gamma1 = np.flipud(gamma1)

index = np.argmin(np.linalg.norm(gamma1[0, None] - gamma2, axis=1))
gamma3 = np.roll(gamma2, -index, axis=0)
np.testing.assert_allclose(gamma1, gamma3, atol=1e-14)

# Verify that divergence is zero
dB1_by_dX = Bfield.dB_by_dX()
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints)))
Expand All @@ -232,11 +246,13 @@ def test_circularcoil_Bfield(self):
points = np.asarray(npoints * [[-1.41513202e-03, 8.99999382e-01, -3.14473221e-04]])
np.random.seed(0)
points += pointVar * (np.random.rand(*points.shape)-0.5)

## verify with a x^2+z^2=radius^2 circular coil
normal = [np.pi/2, np.pi/2]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], radius, 0., center[1], 0., 0., center[2], 0., radius])
Bcircular = BiotSavart([Coil(curve, Current(current))])
general_coil = Coil(curve, Current(current))
Bcircular = BiotSavart([general_coil])
Bfield = CircularCoil(I=current, r0=radius, normal=normal, center=center)
Bfield.set_points(points)
Bcircular.set_points(points)
Expand All @@ -246,11 +262,14 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints)))
assert np.allclose(dB1_by_dX, transpGradB1)
compare_gammas(Bfield, general_coil)

# use normal = [0, 1, 0]
normal = [0, 1, 0]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], radius, 0., center[1], 0., 0., center[2], 0., radius])
Bcircular = BiotSavart([Coil(curve, Current(current))])
general_coil = Coil(curve, Current(current))
Bcircular = BiotSavart([general_coil])
Bfield = CircularCoil(I=current, r0=radius, normal=normal, center=center)
Bfield.set_points(points)
Bcircular.set_points(points)
Expand All @@ -260,11 +279,14 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints)))
assert np.allclose(dB1_by_dX, transpGradB1)
compare_gammas(Bfield, general_coil)

## verify with a y^2+z^2=radius^2 circular coil
normal = [0, np.pi/2]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], 0, 0., center[1], radius, 0., center[2], 0., radius])
Bcircular = BiotSavart([Coil(curve, Current(-current))])
general_coil = Coil(curve, Current(-current))
Bcircular = BiotSavart([general_coil])
Bfield = CircularCoil(I=current, r0=radius, normal=normal, center=center)
Bfield.set_points(points)
Bcircular.set_points(points)
Expand All @@ -274,6 +296,7 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints))) # divergence
assert np.allclose(dB1_by_dX, transpGradB1) # symmetry of the gradient
compare_gammas(Bfield, general_coil)

# one points
Bfield.set_points(np.asarray([[0.1, 0.2, 0.3]]))
Expand All @@ -294,7 +317,8 @@ def test_circularcoil_Bfield(self):
normal = [1, 0, 0]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], 0, 0., center[1], radius, 0., center[2], 0., radius])
Bcircular = BiotSavart([Coil(curve, Current(-current))])
general_coil = Coil(curve, Current(-current))
Bcircular = BiotSavart([general_coil])
Bfield = CircularCoil(I=current, r0=radius, normal=normal, center=center)
Bfield.set_points(points)
Bcircular.set_points(points)
Expand All @@ -304,12 +328,15 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints))) # divergence
assert np.allclose(dB1_by_dX, transpGradB1) # symmetry of the gradient
compare_gammas(Bfield, general_coil)

## verify with a x^2+y^2=radius^2 circular coil
center = [0, 0, 0]
normal = [0, 0]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], 0, radius, center[1], radius, 0., center[2], 0., 0.])
Bcircular = BiotSavart([Coil(curve, Current(current))])
general_coil = Coil(curve, Current(current))
Bcircular = BiotSavart([general_coil])
curve2 = CurveRZFourier(300, 1, 1, True)
curve2.set_dofs([radius, 0, 0])
Bcircular2 = BiotSavart([Coil(curve2, Current(current))])
Expand All @@ -325,12 +352,15 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular2.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints))) # divergence
assert np.allclose(dB1_by_dX, transpGradB1) # symmetry of the gradient
compare_gammas(Bfield, general_coil)

# use normal = [0, 0, 1]
center = [0, 0, 0]
normal = [0, 0, 1]
curve = CurveXYZFourier(300, 1)
curve.set_dofs([center[0], 0, radius, center[1], radius, 0., center[2], 0., 0.])
Bcircular = BiotSavart([Coil(curve, Current(current))])
general_coil = Coil(curve, Current(current))
Bcircular = BiotSavart([general_coil])
curve2 = CurveRZFourier(300, 1, 1, True)
curve2.set_dofs([radius, 0, 0])
Bcircular2 = BiotSavart([Coil(curve, Current(current))])
Expand All @@ -346,6 +376,8 @@ def test_circularcoil_Bfield(self):
assert np.allclose(Bfield.dB_by_dX(), Bcircular2.dB_by_dX())
assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints))) # divergence
assert np.allclose(dB1_by_dX, transpGradB1) # symmetry of the gradient
compare_gammas(Bfield, general_coil)

## Test with results from coilpy
radius = 1.2345
center = np.array([0.123, 1.456, 2.789])
Expand All @@ -372,7 +404,14 @@ def test_circularcoil_Bfield(self):

field = CircularCoil(r0=radius, center=center, I=current, normal=[np.pi/2, -angle])
field.set_points(points)
np.allclose(field.B(), [[0.01016974, 0.00629875, -0.00220838]])
np.testing.assert_allclose(field.B(), [[0.01016974, 0.00629875, -0.00220838]], rtol=1e-6)
# test coil location
np.testing.assert_allclose(field.gamma(points=4), [[1.3575,1.456,2.789],[0.123,center[1]+radius*np.cos(-angle),center[2]-radius*np.sin(-angle)],
[-1.1115,1.456,2.789],[0.123,center[1]-radius*np.cos(-angle),center[2]+radius*np.sin(-angle)]])
with ScratchDir("."):
for close in [True, False]:
field.to_vtk('test', close=close)


def test_circularcoil_Bfield_toroidal_arrangement(self):
# This makes N_coils with centered at major radius R_m
Expand Down
Loading

0 comments on commit c243322

Please sign in to comment.