def test_init(self):
s = SurfaceRZFourier(nfp=2, mpol=3, ntor=2)
self.assertEqual(s.rc.shape, (4, 5))
self.assertEqual(s.zs.shape, (4, 5))
s = SurfaceRZFourier(nfp=10, mpol=1, ntor=3, stellsym=False)
self.assertEqual(s.rc.shape, (2, 7))
self.assertEqual(s.zs.shape, (2, 7))
self.assertEqual(s.rs.shape, (2, 7))
self.assertEqual(s.zc.shape, (2, 7))
def test_2dof_surface_opt(self):
"""
Optimize the minor radius and elongation of an axisymmetric torus to
obtain a desired volume and area.
"""
for solver in solvers:
desired_volume = 0.6
desired_area = 8.0
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
surf = SurfaceRZFourier(quadpoints_phi=62, quadpoints_theta=63)
# Set initial surface shape. It helps to make zs(1,0) larger
# than rc(1,0) since there are two solutions to this
# optimization problem, and for testing we want to find one
# rather than the other.
surf.set_zs(1, 0, 0.2)
# Parameters are all non-fixed by default, meaning they will be
# optimized. You can choose to exclude any subset of the variables
# from the space of independent variables by setting their 'fixed'
# property to True.
surf.set_fixed('rc(0,0)')
# Each function you want in the objective function is then
# equipped with a shift and weight, to become a term in a
# least-squares objective function. A list of terms are
# combined to form a nonlinear-least-squares problem.
prob = LeastSquaresProblem([(surf.volume, desired_volume, 1),
(surf.area, desired_area, 1)])
# Verify the state vector and names are what we expect
np.testing.assert_allclose(prob.x, [0.1, 0.2])
self.assertEqual(prob.dofs.names[0][:28],
'rc(1,0) of SurfaceRZFourier ')
self.assertEqual(prob.dofs.names[1][:28],
'zs(1,0) of SurfaceRZFourier ')
# Solve the minimization problem:
solver(prob)
# Check results
self.assertAlmostEqual(surf.get_rc(0, 0), 1.0, places=13)
self.assertAlmostEqual(surf.get_rc(1, 0),
0.10962565115956417,
places=13)
self.assertAlmostEqual(surf.get_zs(0, 0), 0.0, places=13)
self.assertAlmostEqual(surf.get_zs(1, 0),
0.27727411213693337,
places=13)
self.assertAlmostEqual(surf.volume(), desired_volume, places=8)
self.assertAlmostEqual(surf.area(), desired_area, places=8)
self.assertLess(np.abs(prob.objective()), 1.0e-15)
def test_guidingcenterphihits(self):
bsh = self.bsh
ma = self.ma
nparticles = 2
m = PROTON_MASS
q = ELEMENTARY_CHARGE
Ekin = 9000 * ONE_EV
nphis = 10
phis = np.linspace(0, 2 * np.pi, nphis, endpoint=False)
mpol = 5
ntor = 5
nfp = 3
s = SurfaceRZFourier(mpol=mpol,
ntor=ntor,
stellsym=True,
nfp=nfp,
quadpoints_phi=np.linspace(0,
1,
nfp * 2 * ntor + 1,
endpoint=False),
quadpoints_theta=np.linspace(0,
1,
2 * mpol + 1,
endpoint=False))
s.fit_to_curve(ma, 0.10, flip_theta=False)
sc = SurfaceClassifier(s, h=0.1, p=2)
if with_evtk:
sc.to_vtk('/tmp/classifier')
# check that the axis is classified as inside the domain
assert sc.evaluate(ma.gamma()[:1, :]) > 0
assert sc.evaluate(2 * ma.gamma()[:1, :]) < 0
np.random.seed(1)
gc_tys, gc_phi_hits = trace_particles_starting_on_curve(
ma,
bsh,
nparticles,
tmax=1e-4,
seed=1,
mass=m,
charge=q,
Ekin=Ekin,
umin=-0.1,
umax=+0.1,
phis=phis,
mode='gc_vac',
stopping_criteria=[LevelsetStoppingCriterion(sc)])
if with_evtk:
particles_to_vtk(gc_tys, '/tmp/particles_gc')
for i in range(nparticles):
assert validate_phi_hits(gc_phi_hits[i], bsh, nphis)
def test_from_focus(self):
"""
Try reading in a focus-format file.
"""
filename = TEST_DIR / 'tf_only_half_tesla.plasma'
s = SurfaceRZFourier.from_focus(filename)
self.assertEqual(s.nfp, 3)
self.assertTrue(s.stellsym)
self.assertEqual(s.rc.shape, (11, 13))
self.assertEqual(s.zs.shape, (11, 13))
self.assertAlmostEqual(s.rc[0, 6], 1.408922E+00)
self.assertAlmostEqual(s.rc[0, 7], 2.794370E-02)
self.assertAlmostEqual(s.zs[0, 7], -1.909220E-02)
self.assertAlmostEqual(s.rc[10, 12], -6.047097E-05)
self.assertAlmostEqual(s.zs[10, 12], 3.663233E-05)
self.assertAlmostEqual(s.get_rc(0, 0), 1.408922E+00)
self.assertAlmostEqual(s.get_rc(0, 1), 2.794370E-02)
self.assertAlmostEqual(s.get_zs(0, 1), -1.909220E-02)
self.assertAlmostEqual(s.get_rc(10, 6), -6.047097E-05)
self.assertAlmostEqual(s.get_zs(10, 6), 3.663233E-05)
true_area = 24.5871075268402
true_volume = 2.96201898538042
#print("computed area: ", area, ", correct value: ", true_area, \
# " , difference: ", area - true_area)
#print("computed volume: ", volume, ", correct value: ", \
# true_volume, ", difference:", volume - true_volume)
self.assertAlmostEqual(s.area(), true_area, places=4)
self.assertAlmostEqual(s.volume(), true_volume, places=3)
def get_surface(surfacetype, stellsym, phis=None, thetas=None, ntor=5, mpol=5):
nfp = 3
nphi = 11 if surfacetype == "SurfaceXYZTensorFourier" else 15
ntheta = 11 if surfacetype == "SurfaceXYZTensorFourier" else 15
if phis is None:
phis = np.linspace(0, 1/nfp, nphi, endpoint=False)
if thetas is None:
if surfacetype == "SurfaceXYZTensorFourier":
thetas = np.linspace(0, 1, ntheta, endpoint=False)
else:
thetas = np.linspace(0, 1/(1. + int(stellsym)), ntheta, endpoint=False)
if surfacetype == "SurfaceXYZFourier":
s = SurfaceXYZFourier(mpol=mpol, ntor=ntor, nfp=nfp, stellsym=stellsym,
quadpoints_phi=phis, quadpoints_theta=thetas)
elif surfacetype == "SurfaceRZFourier":
s = SurfaceRZFourier(mpol=mpol, ntor=ntor, nfp=nfp, stellsym=stellsym,
quadpoints_phi=phis, quadpoints_theta=thetas)
elif surfacetype == "SurfaceXYZTensorFourier":
s = SurfaceXYZTensorFourier(mpol=mpol, ntor=ntor, nfp=nfp, stellsym=stellsym,
clamped_dims=[False, False, False],
quadpoints_phi=phis, quadpoints_theta=thetas
)
else:
raise Exception("surface type not implemented")
return s
def test_surface_sampling(self):
np.random.seed(1)
nquadpoints = int(4e2)
surface = SurfaceRZFourier(nfp=1, stellsym=True, mpol=1, ntor=0, quadpoints_phi=nquadpoints, quadpoints_theta=nquadpoints)
dofs = surface.get_dofs()
dofs[0] = 1
dofs[1] = 0.8
surface.set_dofs(dofs)
n = np.linalg.norm(surface.normal(), axis=2)
print(np.min(n), np.max(n))
start = int(0.2*nquadpoints)
stop = int(0.5*nquadpoints)
from scipy.integrate import simpson
quadpoints_phi = surface.quadpoints_phi
quadpoints_theta = surface.quadpoints_theta
lineintegrals = [simpson(y=n[i, start:stop], x=quadpoints_theta[start:stop]) for i in range(start, stop)]
area_of_subset = simpson(y=lineintegrals, x=quadpoints_phi[start:stop])
total_area = surface.area()
print("area_of_subset/total_area", area_of_subset/total_area)
nsamples = int(1e6)
xyz, idxs = draw_uniform_on_surface(surface, nsamples, safetyfactor=10)
samples_in_range = np.sum((idxs[0] >= start) * (idxs[0] < stop)*(idxs[1] >= start) * (idxs[1] < stop))
print("samples_in_range/nsamples", samples_in_range/nsamples)
print("fraction of samples if uniform", (stop-start)**2/(nquadpoints**2))
assert abs(samples_in_range/nsamples - area_of_subset/total_area) < 1e-2
def test_tracing_on_surface_runs(self):
bsh = self.bsh
ma = self.ma
nparticles = 1
m = PROTON_MASS
q = ELEMENTARY_CHARGE
tmax = 1e-3
Ekin = 9000 * ONE_EV
np.random.seed(1)
ntor = 1
mpol = 1
stellsym = True
nfp = 3
phis = np.linspace(0, 1, 50, endpoint=False)
thetas = np.linspace(0, 1, 50, endpoint=False)
s = SurfaceRZFourier(mpol=mpol,
ntor=ntor,
stellsym=stellsym,
nfp=nfp,
quadpoints_phi=phis,
quadpoints_theta=thetas)
s.fit_to_curve(ma, 0.03, flip_theta=False)
gc_tys, gc_phi_hits = trace_particles_starting_on_surface(
s,
bsh,
nparticles,
tmax=tmax,
seed=1,
mass=m,
charge=q,
Ekin=Ekin,
umin=-0.80,
umax=-0.70,
phis=[],
mode='gc_vac',
tol=1e-11,
stopping_criteria=[IterationStoppingCriterion(10)])
assert len(gc_tys[0]) == 11
def test_area_volume(self):
"""
Test the calculation of area and volume for an axisymmetric surface
"""
s = SurfaceRZFourier()
s.rc[0, 0] = 1.3
s.rc[1, 0] = 0.4
s.zs[1, 0] = 0.2
true_area = 15.827322032265993
true_volume = 2.0528777154265874
self.assertAlmostEqual(s.area(), true_area, places=4)
self.assertAlmostEqual(s.volume(), true_volume, places=3)
def test_derivatives(self):
"""
Check the automatic differentiation for area and volume.
"""
for mpol in range(1, 3):
for ntor in range(2):
for nfp in range(1, 4):
s = SurfaceRZFourier(nfp=nfp, mpol=mpol, ntor=ntor)
x0 = s.get_dofs()
x = np.random.rand(len(x0)) - 0.5
x[0] = np.random.rand() + 2
# This surface will probably self-intersect, but I
# don't think this actually matters here.
s.set_dofs(x)
dofs = Dofs([s.area, s.volume])
jac = dofs.jac()
fd_jac = dofs.fd_jac()
print('difference for surface test_derivatives:',
jac - fd_jac)
np.testing.assert_allclose(jac,
fd_jac,
rtol=1e-4,
atol=1e-4)
def test_distance(self):
c = CurveRZFourier(100, 1, 1, False)
dofs = c.get_dofs()
dofs[0] = 1.
c.set_dofs(dofs)
s = SurfaceRZFourier(mpol=1, ntor=1)
s.fit_to_curve(c, 0.2, flip_theta=True)
xyz = np.asarray([[0, 0, 0], [1., 0, 0], [2., 0., 0]])
d = signed_distance_from_surface(xyz, s)
assert np.allclose(d, [-0.8, 0.2, -0.8])
s.fit_to_curve(c, 0.2, flip_theta=False)
d = signed_distance_from_surface(xyz, s)
assert np.allclose(d, [-0.8, 0.2, -0.8])
def test_set_dofs(self):
"""
Test that we can set the shape from a 1D vector
"""
# First try an axisymmetric surface for simplicity:
s = SurfaceRZFourier()
s.set_dofs([2.9, -1.1, 0.7])
self.assertAlmostEqual(s.rc[0, 0], 2.9)
self.assertAlmostEqual(s.rc[1, 0], -1.1)
self.assertAlmostEqual(s.zs[1, 0], 0.7)
# Now try a nonaxisymmetric shape:
s = SurfaceRZFourier(mpol=3, ntor=1)
s.set_dofs(np.array(list(range(21))) + 1)
self.assertAlmostEqual(s.rc[0, 0], 0)
self.assertAlmostEqual(s.rc[0, 1], 1)
self.assertAlmostEqual(s.rc[0, 2], 2)
self.assertAlmostEqual(s.rc[1, 0], 3)
self.assertAlmostEqual(s.rc[1, 1], 4)
self.assertAlmostEqual(s.rc[1, 2], 5)
self.assertAlmostEqual(s.rc[2, 0], 6)
self.assertAlmostEqual(s.rc[2, 1], 7)
self.assertAlmostEqual(s.rc[2, 2], 8)
self.assertAlmostEqual(s.rc[3, 0], 9)
self.assertAlmostEqual(s.rc[3, 1], 10)
self.assertAlmostEqual(s.rc[3, 2], 11)
self.assertAlmostEqual(s.zs[0, 0], 0)
self.assertAlmostEqual(s.zs[0, 1], 0)
self.assertAlmostEqual(s.zs[0, 2], 12)
self.assertAlmostEqual(s.zs[1, 0], 13)
self.assertAlmostEqual(s.zs[1, 1], 14)
self.assertAlmostEqual(s.zs[1, 2], 15)
self.assertAlmostEqual(s.zs[2, 0], 16)
self.assertAlmostEqual(s.zs[2, 1], 17)
self.assertAlmostEqual(s.zs[2, 2], 18)
self.assertAlmostEqual(s.zs[3, 0], 19)
self.assertAlmostEqual(s.zs[3, 1], 20)
self.assertAlmostEqual(s.zs[3, 2], 21)
def test_convert_back(self):
"""
If we start with a SurfaceRZFourier, convert to Garabedian, and
convert back to SurfaceFourier, we should get back what we
started with.
"""
for mpol in range(1, 4):
for ntor in range(5):
for nfp in range(1, 4):
sf1 = SurfaceRZFourier(nfp=nfp, mpol=mpol, ntor=ntor)
# Set all dofs to random numbers in [-2, 2]:
sf1.set_dofs(
(np.random.rand(len(sf1.get_dofs())) - 0.5) * 4)
sg = sf1.to_Garabedian()
sf2 = sg.to_RZFourier()
np.testing.assert_allclose(sf1.rc, sf2.rc)
np.testing.assert_allclose(sf1.zs, sf2.zs)
https://github.com/landreman/stellopt_scenarios/tree/master/1DOF_circularCrossSection_varyR0_targetVolume
This example uses the serial solver instead of the MPI solver, since
the MPI case is covered by the example
stellopt_scenarios_1DOF_circularCrossSection_varyAxis_targetIota_spec.
"""
# Print detailed logging info. This could be commented out if desired.
logging.basicConfig(level=logging.DEBUG)
# Create a Spec object:
equil = Spec()
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
equil.boundary = SurfaceRZFourier(nfp=5, mpol=1, ntor=1)
surf = equil.boundary
# Set the initial boundary shape. Here is one syntax:
surf.set('rc(0,0)', 1.0)
# Here is another syntax:
surf.set_rc(0, 1, 0.1)
surf.set_zs(0, 1, 0.1)
surf.set_rc(1, 0, 0.1)
surf.set_zs(1, 0, 0.1)
print('rc:', surf.rc)
print('zs:', surf.zs)
# Surface parameters are all non-fixed by default. You can choose
def get_surface(surfacetype, stellsym, phis=None, thetas=None):
np.random.seed(2)
mpol = 4
ntor = 3
nfp = 2
phis = phis if phis is not None else np.linspace(0, 1, 31, endpoint=False)
thetas = thetas if thetas is not None else np.linspace(0, 1, 31, endpoint=False)
if surfacetype == "SurfaceRZFourier":
from simsopt.geo.surfacerzfourier import SurfaceRZFourier
s = SurfaceRZFourier(nfp=nfp, stellsym=stellsym, mpol=mpol, ntor=ntor, quadpoints_phi=phis, quadpoints_theta=thetas)
s.set_dofs(s.get_dofs()*0.)
s.rc[0, ntor + 0] = 1
s.rc[1, ntor + 0] = 0.3
s.zs[1, ntor + 0] = 0.3
elif surfacetype == "SurfaceXYZFourier":
from simsopt.geo.surfacexyzfourier import SurfaceXYZFourier
s = SurfaceXYZFourier(nfp=nfp, stellsym=stellsym, mpol=mpol, ntor=ntor, quadpoints_phi=phis, quadpoints_theta=thetas)
s.set_dofs(s.get_dofs()*0.)
s.xc[0, ntor + 1] = 1.
s.xc[1, ntor + 1] = 0.1
s.ys[0, ntor + 1] = 1.
s.ys[1, ntor + 1] = 0.1
s.zs[1, ntor] = 0.1
elif surfacetype == "SurfaceXYZTensorFourier":
from simsopt.geo.surfacexyztensorfourier import SurfaceXYZTensorFourier
s = SurfaceXYZTensorFourier(
nfp=nfp, stellsym=stellsym, mpol=mpol, ntor=ntor,
clamped_dims=[False, not stellsym, True],
quadpoints_phi=phis, quadpoints_theta=thetas)
s.set_dofs(s.get_dofs()*0.)
s.x[0, 0] = 1.0
s.x[1, 0] = 0.1
s.z[mpol+1, 0] = 0.1
else:
assert False
dofs = np.asarray(s.get_dofs())
np.random.seed(2)
rand_scale = 0.01
s.set_dofs(dofs + rand_scale * np.random.rand(len(dofs)).reshape(dofs.shape))
return s
def test_get_dofs(self):
"""
Test that we can convert the degrees of freedom into a 1D vector
"""
# First try an axisymmetric surface for simplicity:
s = SurfaceRZFourier()
s.rc[0, 0] = 1.3
s.rc[1, 0] = 0.4
s.zs[0, 0] = 0.3
s.zs[1, 0] = 0.2
dofs = s.get_dofs()
self.assertEqual(dofs.shape, (3, ))
self.assertAlmostEqual(dofs[0], 1.3)
self.assertAlmostEqual(dofs[1], 0.4)
self.assertAlmostEqual(dofs[2], 0.2)
# Now try a nonaxisymmetric shape:
s = SurfaceRZFourier(mpol=3, ntor=1)
s.rc[:, :] = [[100, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
s.zs[:, :] = [[101, 102, 13], [14, 15, 16], [17, 18, 19], [20, 21, 22]]
dofs = s.get_dofs()
self.assertEqual(dofs.shape, (21, ))
for j in range(21):
self.assertAlmostEqual(dofs[j], j + 2)
def test_stopping_criteria(self):
bsh = self.bsh
ma = self.ma
nparticles = 1
m = PROTON_MASS
q = ELEMENTARY_CHARGE
tmax = 1e-3
Ekin = 9000 * ONE_EV
np.random.seed(1)
gc_tys, gc_phi_hits = trace_particles_starting_on_curve(
ma,
bsh,
nparticles,
tmax=tmax,
seed=1,
mass=m,
charge=q,
Ekin=Ekin,
umin=-0.80,
umax=-0.70,
phis=[],
mode='gc_vac',
tol=1e-11,
stopping_criteria=[IterationStoppingCriterion(10)])
assert len(gc_tys[0]) == 11
# consider a particle with mostly perpendicular velocity so that it get's lost
ntor = 1
mpol = 1
stellsym = True
nfp = 3
phis = np.linspace(0, 1, 50, endpoint=False)
thetas = np.linspace(0, 1, 50, endpoint=False)
s = SurfaceRZFourier(mpol=mpol,
ntor=ntor,
stellsym=stellsym,
nfp=nfp,
quadpoints_phi=phis,
quadpoints_theta=thetas)
s.fit_to_curve(ma, 0.03, flip_theta=False)
sc = SurfaceClassifier(s, h=0.1, p=2)
gc_tys, gc_phi_hits = trace_particles_starting_on_curve(
ma,
bsh,
nparticles,
tmax=tmax,
seed=1,
mass=m,
charge=q,
Ekin=Ekin,
umin=-0.01,
umax=+0.01,
phis=[],
mode='gc_vac',
tol=1e-11,
stopping_criteria=[LevelsetStoppingCriterion(sc)])
if with_evtk:
particles_to_vtk(gc_tys, '/tmp/particles_gc')
assert gc_phi_hits[0][-1][1] == -1
assert np.all(sc.evaluate(gc_tys[0][:, 1:4]) > 0)
#!/usr/bin/env python3
from simsopt.geo.surfacerzfourier import SurfaceRZFourier
from simsopt import LeastSquaresProblem
from simsopt import least_squares_serial_solve
"""
Optimize the minor radius and elongation of an axisymmetric torus to
obtain a desired volume and area.
"""
desired_volume = 0.6
desired_area = 8.0
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
surf = SurfaceRZFourier()
# Parameters are all non-fixed by default, meaning they will be
# optimized. You can choose to exclude any subset of the variables
# from the space of independent variables by setting their 'fixed'
# property to True.
surf.set_fixed('rc(0,0)')
# Each target function is then equipped with a shift and weight, to
# become a term in a least-squares objective function
term1 = (surf.volume, desired_volume, 1)
term2 = (surf.area, desired_area, 1)
# A list of terms are combined to form a nonlinear-least-squares
# problem.
prob = LeastSquaresProblem([term1, term2])
def test_change_resolution(self):
"""
Check that we can change mpol and ntor.
"""
for mpol in [1, 2]:
for ntor in [0, 1]:
s = SurfaceRZFourier(mpol=mpol, ntor=ntor)
n = len(s.get_dofs())
s.set_dofs((np.random.rand(n) - 0.5) * 0.01)
s.set_rc(0, 0, 1.0)
s.set_rc(1, 0, 0.1)
s.set_zs(1, 0, 0.13)
v1 = s.volume()
a1 = s.area()
s.change_resolution(mpol + 1, ntor)
s.recalculate = True
v2 = s.volume()
a2 = s.area()
self.assertAlmostEqual(v1, v2)
self.assertAlmostEqual(a1, a2)
s.change_resolution(mpol, ntor + 1)
s.recalculate = True
v2 = s.volume()
a2 = s.area()
self.assertAlmostEqual(v1, v2)
self.assertAlmostEqual(a1, a2)
s.change_resolution(mpol + 1, ntor + 1)
s.recalculate = True
v2 = s.volume()
a2 = s.area()
self.assertAlmostEqual(v1, v2)
self.assertAlmostEqual(a1, a2)
coils = stellarator.coils
currents = stellarator.currents
bs = BiotSavart(coils, currents)
bs_tf = BiotSavart(coils, currents)
mpol = 5
ntor = 5
stellsym = True
nfp = 3
constraint_weight = 1e0
phis = np.linspace(0, 1 / nfp, 25, endpoint=False)
thetas = np.linspace(0, 1, 25, endpoint=False)
s = SurfaceRZFourier(mpol=mpol,
ntor=ntor,
stellsym=stellsym,
nfp=nfp,
quadpoints_phi=phis,
quadpoints_theta=thetas)
s.fit_to_curve(ma, 0.2, flip_theta=True)
# First optimize at fixed volume
qfm = QfmResidual(s, bs)
qfm.J()
vol = Volume(s)
vol_target = vol.J()
qfm_surface = QfmSurface(bs, s, vol, vol_target)
res = qfm_surface.minimize_qfm_penalty_constraints_LBFGS(
def test_aspect_ratio(self):
"""
Test that the aspect ratio of a torus with random minor and major radius 0.1 <= minor_R <= major_R
is properly computed to be major_R/minor_R.
"""
s = SurfaceRZFourier(nfp=2, mpol=3, ntor=2)
s.rc = s.rc * 0
s.rs = s.rs * 0
s.zc = s.zc * 0
s.zs = s.zs * 0
r1 = np.random.random_sample() + 0.1
r2 = np.random.random_sample() + 0.1
major_R = np.max([r1, r2])
minor_R = np.min([r1, r2])
s.rc[0, 2] = major_R
s.rc[1, 2] = minor_R
s.zs[1, 2] = minor_R
print("AR approx: ", s.aspect_ratio(), "Exact: ", major_R / minor_R)
self.assertAlmostEqual(s.aspect_ratio(), major_R / minor_R)
