diff --git a/docs/source/simsopt.geo.rst b/docs/source/simsopt.geo.rst index 0a1bf5c32..92859f7c7 100644 --- a/docs/source/simsopt.geo.rst +++ b/docs/source/simsopt.geo.rst @@ -76,6 +76,14 @@ simsopt.geo.curvexyzfourier module :undoc-members: :show-inheritance: +simsopt.geo.curvexyzfouriersymmetries module +-------------------------------------------- + +.. automodule:: simsopt.geo.curvexyzfouriersymmetries + :members: + :undoc-members: + :show-inheritance: + simsopt.geo.finitebuild module ------------------------------ diff --git a/src/simsopt/geo/__init__.py b/src/simsopt/geo/__init__.py index 2270653fb..25cf7ea42 100644 --- a/src/simsopt/geo/__init__.py +++ b/src/simsopt/geo/__init__.py @@ -6,6 +6,7 @@ from .curvehelical import * from .curverzfourier import * from .curvexyzfourier import * +from .curvexyzfouriersymmetries import * from .curveperturbed import * from .curveobjectives import * from .curveplanarfourier import * @@ -28,6 +29,7 @@ __all__ = (curve.__all__ + curvehelical.__all__ + curverzfourier.__all__ + curvexyzfourier.__all__ + + curvexyzfouriersymmetries.__all__ + curveperturbed.__all__ + curveobjectives.__all__ + curveplanarfourier.__all__ + finitebuild.__all__ + plotting.__all__ + diff --git a/src/simsopt/geo/curvexyzfouriersymmetries.py b/src/simsopt/geo/curvexyzfouriersymmetries.py new file mode 100644 index 000000000..b537e94f3 --- /dev/null +++ b/src/simsopt/geo/curvexyzfouriersymmetries.py @@ -0,0 +1,142 @@ +import jax.numpy as jnp +import numpy as np +from .curve import JaxCurve +from math import gcd + +__all__ = ['CurveXYZFourierSymmetries'] + + +def jaxXYZFourierSymmetriescurve_pure(dofs, quadpoints, order, nfp, stellsym, ntor): + + theta, m = jnp.meshgrid(quadpoints, jnp.arange(order + 1), indexing='ij') + + if stellsym: + xc = dofs[:order+1] + ys = dofs[order+1:2*order+1] + zs = dofs[2*order+1:] + + xhat = np.sum(xc[None, :] * jnp.cos(2 * jnp.pi * nfp*m*theta), axis=1) + yhat = np.sum(ys[None, :] * jnp.sin(2 * jnp.pi * nfp*m[:, 1:]*theta[:, 1:]), axis=1) + + z = jnp.sum(zs[None, :] * jnp.sin(2*jnp.pi*nfp * m[:, 1:]*theta[:, 1:]), axis=1) + else: + xc = dofs[0 : order+1] + xs = dofs[order+1 : 2*order+1] + yc = dofs[2*order+1: 3*order+2] + ys = dofs[3*order+2: 4*order+2] + zc = dofs[4*order+2: 5*order+3] + zs = dofs[5*order+3: ] + + xhat = np.sum(xc[None, :] * jnp.cos(2*jnp.pi*nfp*m*theta), axis=1) + np.sum(xs[None, :] * jnp.sin(2*jnp.pi*nfp*m[:, 1:]*theta[:, 1:]), axis=1) + yhat = np.sum(yc[None, :] * jnp.cos(2*jnp.pi*nfp*m*theta), axis=1) + np.sum(ys[None, :] * jnp.sin(2*jnp.pi*nfp*m[:, 1:]*theta[:, 1:]), axis=1) + + z = np.sum(zc[None, :] * jnp.cos(2*jnp.pi*nfp*m*theta), axis=1) + np.sum(zs[None, :] * jnp.sin(2*jnp.pi*nfp*m[:, 1:]*theta[:, 1:]), axis=1) + + angle = 2 * jnp.pi * quadpoints * ntor + x = jnp.cos(angle) * xhat - jnp.sin(angle) * yhat + y = jnp.sin(angle) * xhat + jnp.cos(angle) * yhat + + gamma = jnp.zeros((len(quadpoints),3)) + gamma = gamma.at[:, 0].add(x) + gamma = gamma.at[:, 1].add(y) + gamma = gamma.at[:, 2].add(z) + return gamma + + +class CurveXYZFourierSymmetries(JaxCurve): + r'''A curve representation that allows for stellarator and discrete rotational symmetries. This class can be used to + represent a helical coil that does not lie on a torus. The coordinates of the curve are given by: + + .. math:: + x(\theta) &= \hat x(\theta) \cos(2 \pi \theta n_{\text{tor}}) - \hat y(\theta) \sin(2 \pi \theta n_{\text{tor}})\\ + y(\theta) &= \hat x(\theta) \sin(2 \pi \theta n_{\text{tor}}) + \hat y(\theta) \cos(2 \pi \theta n_{\text{tor}})\\ + z(\theta) &= \sum_{m=1}^{\text{order}} z_{s,m} \sin(2 \pi n_{\text{fp}} m \theta) + + where + + .. math:: + \hat x(\theta) &= x_{c, 0} + \sum_{m=1}^{\text{order}} x_{c,m} \cos(2 \pi n_{\text{fp}} m \theta)\\ + \hat y(\theta) &= \sum_{m=1}^{\text{order}} y_{s,m} \sin(2 \pi n_{\text{fp}} m \theta)\\ + + + if the coil is stellarator symmetric. When the coil is not stellarator symmetric, the formulas above + become + + .. math:: + x(\theta) &= \hat x(\theta) \cos(2 \pi \theta n_{\text{tor}}) - \hat y(\theta) \sin(2 \pi \theta n_{\text{tor}})\\ + y(\theta) &= \hat x(\theta) \sin(2 \pi \theta n_{\text{tor}}) + \hat y(\theta) \cos(2 \pi \theta n_{\text{tor}})\\ + z(\theta) &= z_{c, 0} + \sum_{m=1}^{\text{order}} \left[ z_{c, m} \cos(2 \pi n_{\text{fp}} m \theta) + z_{s, m} \sin(2 \pi n_{\text{fp}} m \theta) \right] + + where + + .. math:: + \hat x(\theta) &= x_{c, 0} + \sum_{m=1}^{\text{order}} \left[ x_{c, m} \cos(2 \pi n_{\text{fp}} m \theta) + x_{s, m} \sin(2 \pi n_{\text{fp}} m \theta) \right] \\ + \hat y(\theta) &= y_{c, 0} + \sum_{m=1}^{\text{order}} \left[ y_{c, m} \cos(2 \pi n_{\text{fp}} m \theta) + y_{s, m} \sin(2 \pi n_{\text{fp}} m \theta) \right] \\ + + Args: + quadpoints: number of grid points/resolution along the curve, + order: how many Fourier harmonics to include in the Fourier representation, + nfp: discrete rotational symmetry number, + stellsym: stellaratory symmetry if True, not stellarator symmetric otherwise, + ntor: the number of times the curve wraps toroidally before biting its tail. Note, + it is assumed that nfp and ntor are coprime. If they are not coprime, + then then the curve actually has nfp_new:=nfp // gcd(nfp, ntor), + and ntor_new:=ntor // gcd(nfp, ntor). The operator `//` is integer division. + To avoid confusion, we assert that ntor and nfp are coprime at instantiation. + ''' + + def __init__(self, quadpoints, order, nfp, stellsym, ntor=1, **kwargs): + if isinstance(quadpoints, int): + quadpoints = np.linspace(0, 1, quadpoints, endpoint=False) + pure = lambda dofs, points: jaxXYZFourierSymmetriescurve_pure( + dofs, points, order, nfp, stellsym, ntor) + + if gcd(ntor, nfp) != 1: + raise Exception('nfp and ntor must be coprime') + + self.order = order + self.nfp = nfp + self.stellsym = stellsym + self.ntor = ntor + self.coefficients = np.zeros(self.num_dofs()) + if "dofs" not in kwargs: + if "x0" not in kwargs: + kwargs["x0"] = self.coefficients + else: + self.set_dofs_impl(kwargs["x0"]) + + super().__init__(quadpoints, pure, names=self._make_names(order), **kwargs) + + def _make_names(self, order): + if self.stellsym: + x_cos_names = [f'xc({i})' for i in range(0, order + 1)] + x_names = x_cos_names + y_sin_names = [f'ys({i})' for i in range(1, order + 1)] + y_names = y_sin_names + z_sin_names = [f'zs({i})' for i in range(1, order + 1)] + z_names = z_sin_names + else: + x_names = ['xc(0)'] + x_cos_names = [f'xc({i})' for i in range(1, order + 1)] + x_sin_names = [f'xs({i})' for i in range(1, order + 1)] + x_names += x_cos_names + x_sin_names + y_names = ['yc(0)'] + y_cos_names = [f'yc({i})' for i in range(1, order + 1)] + y_sin_names = [f'ys({i})' for i in range(1, order + 1)] + y_names += y_cos_names + y_sin_names + z_names = ['zc(0)'] + z_cos_names = [f'zc({i})' for i in range(1, order + 1)] + z_sin_names = [f'zs({i})' for i in range(1, order + 1)] + z_names += z_cos_names + z_sin_names + + return x_names + y_names + z_names + + def num_dofs(self): + return (self.order+1) + self.order + self.order if self.stellsym else 3*(2*self.order+1) + + def get_dofs(self): + return self.coefficients + + def set_dofs_impl(self, dofs): + self.coefficients[:] = dofs[:] + diff --git a/tests/geo/test_curve.py b/tests/geo/test_curve.py index acf3314bf..a997a6e84 100644 --- a/tests/geo/test_curve.py +++ b/tests/geo/test_curve.py @@ -11,12 +11,14 @@ from simsopt.geo.curverzfourier import CurveRZFourier from simsopt.geo.curveplanarfourier import CurvePlanarFourier from simsopt.geo.curvehelical import CurveHelical +from simsopt.geo.curvexyzfouriersymmetries import CurveXYZFourierSymmetries from simsopt.geo.curve import RotatedCurve, curves_to_vtk from simsopt.geo import parameters from simsopt.configs.zoo import get_ncsx_data, get_w7x_data +from simsopt.field import BiotSavart, Current, coils_via_symmetries, Coil from simsopt.field.coil import coils_to_makegrid from simsopt.geo import CurveLength, CurveCurveDistance - +from math import gcd try: import pyevtk @@ -36,7 +38,7 @@ def taylor_test(f, df, x, epsilons=None, direction=None): dfx = df(x)@direction if epsilons is None: epsilons = np.power(2., -np.asarray(range(7, 20))) - # print("################################################################################") + print("################################################################################") err_old = 1e9 counter = 0 for eps in epsilons: @@ -46,6 +48,8 @@ def taylor_test(f, df, x, epsilons=None, direction=None): fminuseps = f(x - eps * direction) dfest = (fpluseps-fminuseps)/(2*eps) err = np.linalg.norm(dfest - dfx) + print(err, err/err_old) + assert err < 1e-9 or err < 0.3 * err_old if err < 1e-9: break @@ -53,7 +57,7 @@ def taylor_test(f, df, x, epsilons=None, direction=None): counter += 1 if err > 1e-10: assert counter > 3 - # print("################################################################################") + print("################################################################################") def get_curve(curvetype, rotated, x=np.asarray([0.5])): @@ -73,9 +77,15 @@ def get_curve(curvetype, rotated, x=np.asarray([0.5])): curve = CurveHelical(x, order, 5, 2, 1.0, 0.3, x0=np.ones((2*order,))) elif curvetype == "CurvePlanarFourier": curve = CurvePlanarFourier(x, order, 2, True) + elif curvetype == "CurveXYZFourierSymmetries1": + curve = CurveXYZFourierSymmetries(x, order, 2, True) + elif curvetype == "CurveXYZFourierSymmetries2": + curve = CurveXYZFourierSymmetries(x, order, 2, False) + elif curvetype == "CurveXYZFourierSymmetries3": + curve = CurveXYZFourierSymmetries(x, order, 2, False, ntor=3) else: assert False - + dofs = np.zeros((curve.dof_size, )) if curvetype in ["CurveXYZFourier", "JaxCurveXYZFourier"]: dofs[1] = 1. @@ -87,10 +97,38 @@ def get_curve(curvetype, rotated, x=np.asarray([0.5])): dofs[order+1] = 0.1 elif curvetype in ["CurveHelical", "CurveHelicalInitx0"]: dofs[0] = np.pi/2 + elif curvetype == "CurveXYZFourierSymmetries1": + R = 1 + r = 0.5 + curve.set('xc(0)', R) + curve.set('xc(1)', -r) + curve.set('zs(1)', -r) + dofs = curve.get_dofs() + elif curvetype == "CurveXYZFourierSymmetries2": + R = 1 + r = 0.5 + curve.set('xc(0)', R) + curve.set('xs(1)', -0.1*r) + curve.set('xc(1)', -r) + curve.set('zs(1)', -r) + curve.set('zc(0)', 1) + curve.set('zs(1)', r) + dofs = curve.get_dofs() + elif curvetype == "CurveXYZFourierSymmetries3": + R = 1 + r = 0.5 + curve.set('xc(0)', R) + curve.set('xs(1)', -0.1*r) + curve.set('xc(1)', -r) + curve.set('zs(1)', -r) + curve.set('zc(0)', 1) + curve.set('zs(1)', r) + dofs = curve.get_dofs() else: assert False curve.x = dofs + rand_scale * np.random.rand(len(dofs)).reshape(dofs.shape) + if rotated: curve = RotatedCurve(curve, 0.5, flip=False) return curve @@ -98,7 +136,239 @@ def get_curve(curvetype, rotated, x=np.asarray([0.5])): class Testing(unittest.TestCase): - curvetypes = ["CurveXYZFourier", "JaxCurveXYZFourier", "CurveRZFourier", "CurvePlanarFourier", "CurveHelical", "CurveHelicalInitx0"] + curvetypes = ["CurveXYZFourier", "JaxCurveXYZFourier", "CurveRZFourier", "CurvePlanarFourier", "CurveHelical", "CurveXYZFourierSymmetries1","CurveXYZFourierSymmetries2", "CurveXYZFourierSymmetries3", "CurveHelicalInitx0"] + + def get_curvexyzfouriersymmetries(self, stellsym=True, x=None, nfp=None, ntor=1): + # returns a CurveXYZFourierSymmetries that is randomly perturbed + + np.random.seed(1) + rand_scale = 1e-2 + + if nfp is None: + nfp = 3 + if x is None: + x = np.linspace(0, 1, 200, endpoint=False) + + order = 2 + curve = CurveXYZFourierSymmetries(x, order, nfp, stellsym, ntor=ntor) + R = 1 + r = 0.25 + curve.set('xc(0)', R) + curve.set('xc(2)', r) + curve.set('ys(2)', -r) + curve.set('zs(1)', -2*r) + curve.set('zs(2)', r) + dofs = curve.x.copy() + curve.x = dofs + rand_scale * np.random.rand(len(dofs)).reshape(dofs.shape) + + return curve + + def test_curvexyzsymmetries_raisesexception(self): + # test ensures that an exception is raised when you try and create a curvexyzfouriersymmetries + # where gcd(ntor, nfp) != 1. + + order = 1 + nfp = 1 + ntor = 2 + # nfp = 1 and ntor = 2 here, so it should work + curve = CurveXYZFourierSymmetries(100, order, nfp, True, ntor=ntor, x0=np.ones(3*order+1)) + print(curve.x) + + with self.assertRaises(Exception): + order = 1 + nfp = 2 + ntor = 2 + # nfp = 2 and ntor = 2 here, so an exception should be raised + _ = CurveXYZFourierSymmetries(100, order, nfp, True, ntor=ntor, x0=np.ones(3*order+1)) + + def test_curvehelical_is_curvexyzfouriersymmetries(self): + # this test checks that both helical coil representations can produce the same helical curve on a torus + order = 1 + nfp = 2 + curve1 = CurveXYZFourierSymmetries(100, order, nfp, True) + R = 1 + r = 0.5 + curve1.set('xc(0)', R) + curve1.set('xc(1)', r) + curve1.set('zs(1)', -r) + curve2 = CurveHelical(np.linspace(0, 1, 100, endpoint=False), order, nfp, 1, R, r, x0=np.zeros((2*order,))) + np.testing.assert_allclose(curve1.gamma(), curve2.gamma(), atol=1e-14) + + def test_trefoil_nonstellsym(self): + r''' This test checks that a CurveXYZFourierSymmetries can represent a non-stellarator symmetric + trefoil knot. A parametric representation of a trefoil knot is given by: + + x(t) = sin(t) + 2sin(t) + y(t) = cos(t) - 2cos(t) + z(t) = -sin(3t) + + The above can be rewritten the CurveXYZFourierSymmetries representation, with + order = 1, nfp = 3, and ntor = 2: + + x(t) = xhat(t) * cos(ntor*t) - yhat(t) * sin(ntor*t), + y(t) = xhat(t) * sin(ntor*t) + yhat(t) * cos(ntor*t), + z(t) = -sin(nfp*t), + + where + xhat(t) = sin(nfp*t), + yhat(t) = -2 + cos(nfp*t), + + i.e., The dofs are xs(1)=1, yc(0)=-2, yc(1)=1, zs(1)=-1, and zero otherwise. + ''' + + order = 1 + nfp = 3 + ntor = 2 + x = np.linspace(0, 1, 500, endpoint=False) + curve = CurveXYZFourierSymmetries(x, order, nfp, False, ntor=ntor) + + X = np.sin(2*np.pi * x) + 2 * np.sin(2*2*np.pi * x) + Y = np.cos(2*np.pi * x) - 2 * np.cos(2*2*np.pi * x) + Z = -np.sin(3*2*np.pi*x) + XYZ = np.concatenate([X[:, None], Y[:, None], Z[:, None]], axis=1) + + curve.set('xs(1)', 1.) + curve.set('yc(0)', -2.) + curve.set('yc(1)', 1.) + curve.set('zs(1)', -1.) + np.testing.assert_allclose(curve.gamma(), XYZ, atol=1e-14) + + def test_trefoil_stellsym(self): + r''' This test checks that a CurveXYZFourierSymmetries can represent a stellarator symmetric + trefoil knot. A parametric representation of a trefoil knot is given by: + + x(t) = cos(t) - 2cos(t), + y(t) =-sin(t) - 2sin(t), + z(t) = -sin(3t). + + The above can be rewritten the CurveXYZFourierSymmetries representation, with + order = 1, nfp = 3, and ntor = 2: + + x(t) = xhat(t) * cos(ntor*t) - yhat(t) * sin(ntor*t), + y(t) = xhat(t) * sin(ntor*t) + yhat(t) * cos(ntor*t), + z(t) = -sin(nfp*t), + + where + + xhat(t) = -2 + cos(nfp*t) + yhat(t) = -sin(nfp*t) + + i.e., xc(0)=-2, xc(1)=1, ys(1)=-1, zs(1)=-1. + ''' + + order = 1 + nfp = 3 + ntor = 2 + x = np.linspace(0, 1, 500, endpoint=False) + curve = CurveXYZFourierSymmetries(x, order, nfp, True, ntor=ntor) + + X = np.cos(2*np.pi * x) - 2 * np.cos(2*2*np.pi * x) + Y =-np.sin(2*np.pi * x) - 2 * np.sin(2*2*np.pi * x) + Z = -np.sin(3*2*np.pi*x) + XYZ = np.concatenate([X[:, None], Y[:, None], Z[:, None]], axis=1) + + curve.set('xc(0)', -2) + curve.set('xc(1)', 1) + curve.set('ys(1)', -1) + curve.set('zs(1)', -1) + np.testing.assert_allclose(curve.gamma(), XYZ, atol=1e-14) + + + def test_nonstellsym(self): + # this test checks that you can obtain a stellarator symmetric magnetic field from two non-stellarator symmetric + # CurveXYZFourierSymmetries curves. + for nfp in [1, 2, 3, 4, 5, 6]: + for ntor in [1, 2, 3, 4, 5, 6]: + with self.subTest(nfp=nfp, ntor=ntor): + self.subtest_nonstellsym(nfp, ntor) + + def subtest_nonstellsym(self, nfp, ntor): + if gcd(ntor, nfp) != 1: + return + + # this test checks that you can obtain a stellarator symmetric magnetic field from two non-stellarator symmetric + # CurveXYZFourierSymmetries curves. + curve = self.get_curvexyzfouriersymmetries(stellsym=False, nfp=nfp, ntor=ntor) + current = Current(1e5) + coils = coils_via_symmetries([curve], [current], 1, True) + bs = BiotSavart(coils) + bs.set_points([[1, 1, 1], [1, -1, -1]]) + B=bs.B_cyl() + np.testing.assert_allclose(B[0, 0], -B[1, 0], atol=1e-14) + np.testing.assert_allclose(B[0, 1], B[1, 1], atol=1e-14) + np.testing.assert_allclose(B[0, 2], B[1, 2], atol=1e-14) + + def test_xyzhelical_symmetries(self): + # checking various symmetries of the CurveXYZFourierSymmetries representation + for nfp in [1, 2, 3, 4, 5, 6]: + for ntor in [1, 2, 3, 4, 5, 6]: + with self.subTest(nfp=nfp, ntor=ntor): + self.subtest_xyzhelical_symmetries(nfp, ntor) + + def subtest_xyzhelical_symmetries(self, nfp, ntor): + if gcd(nfp, ntor) != 1 : + return + + # does the stellarator symmetric curve have rotational symmetry? + curve = self.get_curvexyzfouriersymmetries(stellsym=True, nfp=nfp, x=np.array([0.123, 0.123+1/nfp]), ntor=ntor) + out = curve.gamma() + + # NOTE: the point at angle t+1/nfp is the point at angle t, but rotated by 2pi *(ntor/nfp) radians. + alpha = 2*np.pi*ntor/nfp + R = np.array([ + [np.cos(alpha), -np.sin(alpha), 0], + [np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1] + ]) + + print(R@out[0], out[1]) + np.testing.assert_allclose(out[1], R@out[0], atol=1e-14) + + # does the stellarator symmetric curve indeed pass through (x0, 0, 0)? + curve = self.get_curvexyzfouriersymmetries(stellsym=True, nfp=nfp, x = np.array([0]), ntor=ntor) + out = curve.gamma() + assert np.abs(out[0, 0]) > 1e-3 + np.testing.assert_allclose(out[0, 1], 0, atol=1e-14) + np.testing.assert_allclose(out[0, 2], 0, atol=1e-14) + + + # does the non-stellarator symmetric curve not pass through (x0, 0, 0)? + curve = self.get_curvexyzfouriersymmetries(stellsym=False, nfp=nfp, x = np.array([0]), ntor=ntor) + out = curve.gamma() + assert np.abs(out[0, 0]) > 1e-3 + assert np.abs(out[0, 1]) > 1e-3 + assert np.abs(out[0, 2]) > 1e-3 + + # is the stellarator symmetric curve actually stellarator symmetric? + curve = self.get_curvexyzfouriersymmetries(stellsym=True, nfp=nfp, x = np.array([0.123, -0.123]), ntor=ntor) + pts = curve.gamma() + np.testing.assert_allclose(pts[0, 0], pts[1, 0], atol=1e-14) + np.testing.assert_allclose(pts[0, 1], -pts[1, 1], atol= 1e-14) + np.testing.assert_allclose(pts[0, 2], -pts[1, 2], atol= 1e-14) + + # is the field from the stellarator symmetric curve actually stellarator symmetric? + curve = self.get_curvexyzfouriersymmetries(stellsym=True, nfp=nfp, x=np.linspace(0, 1, 200, endpoint=False), ntor=ntor) + current = Current(1e5) + coil = Coil(curve, current) + bs = BiotSavart([coil]) + bs.set_points([[1, 1, 1], [1, -1, -1]]) + B=bs.B_cyl() + np.testing.assert_allclose(B[0, 0],-B[1, 0], atol=1e-14) + np.testing.assert_allclose(B[0, 1], B[1, 1], atol=1e-14) + np.testing.assert_allclose(B[0, 2], B[1, 2], atol=1e-14) + + # does the non-stellarator symmetric curve have rotational symmetry still? + # NOTE: the point at angle t+1/nfp is the point at angle t, but rotated by 2pi *(ntor/nfp) radians. + curve = self.get_curvexyzfouriersymmetries(stellsym=False, nfp=nfp, x = np.array([0.123, 0.123+1./nfp]), ntor=ntor) + out = curve.gamma() + alpha = 2*np.pi*ntor/nfp + R = np.array([ + [np.cos(alpha), -np.sin(alpha), 0], + [np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1] + ]) + print(R@out[0], out[1]) + np.testing.assert_allclose(out[1], R@out[0], atol=1e-14) def test_curve_helical_xyzfourier(self): x = np.asarray([0.6])