def test_ngroups_scan(self):
"""
Verify that all quantities make sense when ngroups >= nprocs_world
and ngroups = None.
"""
rank_world = MPI.COMM_WORLD.Get_rank()
nprocs = MPI.COMM_WORLD.Get_size()
for ngroups in range(-1, nprocs + 3):
m = MpiPartition(ngroups=ngroups)
self.assertGreaterEqual(m.ngroups, 1)
self.assertLessEqual(m.ngroups, nprocs)
self.assertEqual(m.rank_world, rank_world)
self.assertGreaterEqual(m.rank_groups, 0)
self.assertLess(m.rank_groups, nprocs)
self.assertEqual(m.nprocs_world, nprocs)
self.assertGreaterEqual(m.nprocs_groups, 1)
self.assertLessEqual(m.nprocs_groups, nprocs)
self.assertEqual(m.proc0_world, rank_world == 0)
if m.proc0_groups:
self.assertGreaterEqual(m.rank_leaders, 0)
self.assertLessEqual(m.rank_leaders, nprocs)
self.assertGreaterEqual(m.nprocs_leaders, 1)
self.assertLessEqual(m.nprocs_leaders, nprocs)
else:
self.assertEqual(m.rank_leaders, -1)
self.assertEqual(m.nprocs_leaders, -1)
# The sizes of the worker groups should be relatively
# even, with a difference of no more than 1 between the
# largest and the smallest.
if m.proc0_world:
group_sizes = np.zeros(nprocs, dtype='i')
group_sizes[0] = m.nprocs_groups
for j in range(1, nprocs):
group_sizes[j] = m.comm_world.recv(tag=j)
print('group_sizes:', group_sizes)
self.assertLessEqual(
np.max(group_sizes) - np.min(group_sizes), 1)
else:
m.comm_world.send(m.nprocs_groups, 0, tag=m.rank_world)
m.write()
def test_ngroups1(self):
"""
Verify that all quantities make sense when ngroups = 1.
"""
rank_world = MPI.COMM_WORLD.Get_rank()
nprocs = MPI.COMM_WORLD.Get_size()
m = MpiPartition(ngroups=1)
self.assertEqual(m.ngroups, 1)
self.assertEqual(m.rank_world, rank_world)
self.assertEqual(m.rank_groups, rank_world)
self.assertEqual(m.rank_leaders, 0 if rank_world == 0 else -1)
self.assertEqual(m.nprocs_world, nprocs)
self.assertEqual(m.nprocs_groups, nprocs)
self.assertEqual(m.nprocs_leaders, 1 if rank_world == 0 else -1)
self.assertEqual(m.proc0_world, rank_world == 0)
self.assertEqual(m.proc0_groups, rank_world == 0)
m.write()
from simsopt.objectives.least_squares import LeastSquaresProblem
from simsopt.solve.mpi import least_squares_mpi_solve
"""
In this example, we simultaneously optimize for quasisymmetry and
the elimination of magnetic islands, with both VMEC and SPEC called in
the objective function.
Below, the argument max_nfev=1 in least_squares_mpi_solve causes the
optimization to stop after only a single iteration, so this example
does not take too long to run. For a real optimization, that argument
should be removed.
"""
log()
mpi = MpiPartition()
mpi.write()
vmec_filename = os.path.join(os.path.dirname(__file__), 'inputs',
'input.nfp2_QA_iota0.4_withIslands')
vmec = Vmec(vmec_filename, mpi=mpi)
surf = vmec.boundary
spec_filename = os.path.join(os.path.dirname(__file__), 'inputs',
'nfp2_QA_iota0.4_withIslands.sp')
spec = Spec(spec_filename, mpi=mpi)
# This next line is where the boundary surface objects of VMEC and
# SPEC are linked:
spec.boundary = surf
# Define parameter space:
def test_stellopt_scenarios_1DOF_circularCrossSection_varyR0_targetVolume(self):
"""
This script implements the "1DOF_circularCrossSection_varyR0_targetVolume"
example from
https://github.com/landreman/stellopt_scenarios
This optimization problem has one independent variable, representing
the mean major radius. The problem also has one objective: the plasma
volume. There is not actually any need to run an equilibrium code like
VMEC since the objective function can be computed directly from the
boundary shape. But this problem is a fast way to test the
optimization infrastructure with VMEC.
Details of the optimum and a plot of the objective function landscape
can be found here:
https://github.com/landreman/stellopt_scenarios/tree/master/1DOF_circularCrossSection_varyR0_targetVolume
"""
# logging.basicConfig(level=logging.DEBUG)
# logger = logging.getLogger('[{}]'.format(MPI.COMM_WORLD.Get_rank()) + __name__)
logger = logging.getLogger(__name__)
for ngroups in range(1, 1 + MPI.COMM_WORLD.Get_size()):
for grad in [False, True]:
# In the next line, we can adjust how many groups the pool of MPI
# processes is split into.
mpi = MpiPartition(ngroups=ngroups)
mpi.write()
# Start with a default surface, which is axisymmetric with major
# radius 1 and minor radius 0.1.
equil = Vmec(mpi=mpi)
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)
# VMEC parameters are all fixed by default, while surface parameters are all non-fixed by default.
# You can choose which parameters are optimized by setting their 'fixed' attributes.
surf.all_fixed()
surf.set_fixed('rc(0,0)', False)
# Each Target is then equipped with a shift and weight, to become a
# term in a least-squares objective function
desired_volume = 0.15
prob = LeastSquaresProblem([(equil.volume, desired_volume, 1)])
# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_mpi_solve(prob, mpi=mpi, grad=grad)
# Make sure all procs call VMEC:
objective = prob.objective()
if mpi.proc0_world:
print("At the optimum,")
print(" rc(m=0,n=0) = ", surf.get_rc(0, 0))
print(" volume, according to VMEC = ", equil.volume())
print(" volume, according to Surface = ", surf.volume())
print(" objective function = ", objective)
assert np.abs(surf.get_rc(0, 0) - 0.7599088773175) < 1.0e-5
assert np.abs(equil.volume() - 0.15) < 1.0e-6
assert np.abs(surf.volume() - 0.15) < 1.0e-6
assert prob.objective() < 1.0e-15