diff --git a/src/simsopt/field/magneticfieldclasses.py b/src/simsopt/field/magneticfieldclasses.py index 101fe80e7..915ced98b 100644 --- a/src/simsopt/field/magneticfieldclasses.py +++ b/src/simsopt/field/magneticfieldclasses.py @@ -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)], @@ -377,16 +407,16 @@ 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) @@ -394,14 +424,14 @@ def _B_impl(self, B): 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) @@ -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) @@ -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) @@ -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""" @@ -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) diff --git a/tests/field/test_magneticfields.py b/tests/field/test_magneticfields.py index ed13a4d60..df46b91ca 100644 --- a/tests/field/test_magneticfields.py +++ b/tests/field/test_magneticfields.py @@ -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))) @@ -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) @@ -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) @@ -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) @@ -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]])) @@ -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) @@ -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))]) @@ -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))]) @@ -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]) @@ -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 diff --git a/tests/field/test_magneticfields_optimization.py b/tests/field/test_magneticfields_optimization.py new file mode 100644 index 000000000..da8249641 --- /dev/null +++ b/tests/field/test_magneticfields_optimization.py @@ -0,0 +1,142 @@ +import unittest +import numpy as np +from pathlib import Path + +from monty.tempfile import ScratchDir + +from simsopt.field import CircularCoil +from simsopt.objectives import LeastSquaresProblem +from simsopt.solve.serial import least_squares_serial_solve +from simsopt import make_optimizable + +TEST_DIR = (Path(__file__).parent / ".." / "test_files").resolve() + +class Testing(unittest.TestCase): + + def test_circularcoil_current_optimization(self): + + # test optimization of I and r0 + coil = CircularCoil() + x0 = np.random.rand(coil.dof_size) + x0[1] = 0 + x0[2] = 0 + x0[3] = 0 + x0[5] = 0 + x0[6] = 0 + coil.x = x0 + coil.fix([1, 2, 3, 5, 6]) + + print('Initial coil radius: ', coil.x[0]) + print('Initial coil current: ', coil.x[1]) + + points = np.array([[0, 0, 0]]) + coil.set_points(points) + + def get_I(c): + return c.I + + def get_B(c): + return c.B()[0][2] + + I_coil = make_optimizable(get_I, coil) + B_coil = make_optimizable(get_B, coil) + + Bmag = 1.2 * 2 * np.pi / 1.12345 + prob = LeastSquaresProblem.from_tuples([(I_coil.J, 1.2e7, 2e6), (B_coil.J, Bmag, 1.0)]) + + with ScratchDir("."): + least_squares_serial_solve(prob) + + print(' Final coil radius: ', coil.x[0]) + print(' Final coil current: ', coil.x[1]) + assert np.allclose(coil.x, [1.12345, 4.8], atol=1e-6) + + def test_circularcoil_position_optimization(self): + + # test center optimization + coil = CircularCoil() + x0 = np.random.rand(coil.dof_size) + x0[0] = 1.12345 + x0[4] = 4.8 + x0[5] = 0 + x0[6] = 0 + coil.x = x0 + coil.fix([0, 4, 5, 6]) + + print('Initial coil position: ', coil.x) + + def Bx(c): + return c.B()[0][0] + + def By(c): + return c.B()[0][1] + + def Bz(c): + return c.B()[0][2] + + Bx_coil = make_optimizable(Bx, coil) + By_coil = make_optimizable(By, coil) + Bz_coil = make_optimizable(Bz, coil) + + Bmag = 1.2 * 2 * np.pi / 1.12345 + prob = LeastSquaresProblem.from_tuples( + [(Bx_coil.J, 0, 1.0), + (By_coil.J, 0, 1.0), + (Bz_coil.J, Bmag, 1.0)] + ) + + points = np.array([[1, 1, 1]]) + coil.set_points(points) + + with ScratchDir("."): + least_squares_serial_solve(prob, ftol=1e-15, xtol=1e-15, gtol=1e-15) + + coil.unfix('y0') + coil.unfix('z0') + + print(' Final coil position: ', coil.x) + assert np.allclose(coil.x, np.ones(3), atol=1e-6) + + def test_circularcoil_orientation_optimization(self): + + # test normal optimization + coil = CircularCoil() + x0 = np.random.rand(coil.dof_size) + x0[0] = 1.12345 + x0[1] = 0 + x0[2] = 0 + x0[3] = 0 + x0[4] = 4.8 + coil.x = x0 + coil.fix([0, 1, 2, 3, 4]) + + print('Initial coil normal: ', coil.x) + + points = np.array([[0, 0, 0]]) + coil.set_points(points) + + def Bx(c): + return c.B()[0][0] + + def By(c): + return c.B()[0][1] + + def Bz(c): + return c.B()[0][2] + + Bx_coil = make_optimizable(Bx, coil) + By_coil = make_optimizable(By, coil) + Bz_coil = make_optimizable(Bz, coil) + + Bmag = np.sqrt(2) / 2 * 1.2 * 2 * np.pi / 1.12345 + prob = LeastSquaresProblem.from_tuples( + [(Bx_coil.J, Bmag, 1.0), + (By_coil.J, 0, 1.0), + (Bz_coil.J, Bmag, 1.0)] + ) + + with ScratchDir("."): + least_squares_serial_solve(prob) + + print(' Final coil normal: ', coil.x) + assert np.allclose(coil.x, [0, np.pi / 4], atol=1e-6)