desired_iota = 0.41
iota_weight = 1
term2 = (vmec.iota_axis, desired_iota, iota_weight)

prob = LeastSquaresProblem([term1, term2])

# Solve the minimization problem:
least_squares_mpi_solve(prob, mpi=mpi, grad=True)

# Evaluate quantities on all processes, in case communication is
# required:
final_objective = prob.objective()
vmec_volume = vmec.volume()
spec_volume = spec.volume()
surf_volume = surf.volume()
vmec_iota = vmec.iota_axis()
spec_iota = spec.iota()

if mpi.proc0_world:
    logging.info("At the optimum,")
    logging.info(f" objective function = {final_objective}")
    logging.info(f" rc(m=1,n=1) = {surf.get_rc(1, 1)}")
    logging.info(f" zs(m=1,n=1) = {surf.get_zs(1, 1)}")
    logging.info(f" volume, according to VMEC    = {vmec_volume}")
    logging.info(f" volume, according to SPEC    = {spec_volume}")
    logging.info(f" volume, according to Surface = {surf_volume}")
    logging.info(f" iota on axis, from VMEC       = {vmec_iota}")
    logging.info(f" iota at mid-radius, from SPEC = {spec_iota}")

    assert np.abs(surf.get_rc(1, 1) - 0.0313066948) < 1.0e-3
    assert np.abs(surf.get_zs(1, 1) - (-0.031232391)) < 1.0e-3
# combined to form a nonlinear-least-squares problem.
desired_volume = 0.15
volume_weight = 1
term1 = (equil.volume, desired_volume, volume_weight)

desired_iota = 0.41
iota_weight = 1
term2 = (equil.iota_axis, desired_iota, iota_weight)

prob = LeastSquaresProblem([term1, term2])

# Solve the minimization problem:
least_squares_mpi_solve(prob, mpi, grad=True)

objective = prob.objective()
if mpi.proc0_world:
    print("At the optimum,")
    print(" rc(m=1,n=1) = ", surf.get_rc(1, 1))
    print(" zs(m=1,n=1) = ", surf.get_zs(1, 1))
    print(" volume, according to VMEC    = ", equil.volume())
    print(" volume, according to Surface = ", surf.volume())
    print(" iota on axis = ", equil.iota_axis())
    print(" objective function = ", objective)

    assert np.abs(surf.get_rc(1, 1) - 0.0313066948) < 1.0e-3
    assert np.abs(surf.get_zs(1, 1) - (-0.031232391)) < 1.0e-3
    assert np.abs(equil.volume() - 0.178091) < 1.0e-3
    assert np.abs(surf.volume() - 0.178091) < 1.0e-3
    assert np.abs(equil.iota_axis() - 0.4114567) < 1.0e-4
    assert prob.objective() < 1.0e-2
https://github.com/landreman/stellopt_scenarios/tree/master/2DOF_circularCrossSection_varyAxis_targetIotaAndQuasisymmetry
See that website for a detailed description of the problem and plots
of the objective function landscape.
"""

vmec = Vmec(os.path.join(os.path.dirname(__file__), 'inputs', 'input.2DOF_circularCrossSection_varyAxis_targetIotaAndQuasisymmetry'))

# Define parameter space:
vmec.boundary.all_fixed()
vmec.boundary.set_fixed("rc(0,1)", False)
vmec.boundary.set_fixed("zs(0,1)", False)

# Define objective function:
boozer = Boozer(vmec, mpol=32, ntor=16)
qs = Quasisymmetry(boozer,
                   1.0,  # Radius to target
                   1, 0,  # (M, N) you want in |B|
                   normalization="symmetric",
                   weight="stellopt_ornl")

# Objective function is 100 * (iota - (-0.41))^2 + 1 * (qs - 0)^2
prob = LeastSquaresProblem([(vmec.iota_axis, -0.41, 100),
                            (qs, 0, 1)])

least_squares_serial_solve(prob)

print("Final values before shifting and scaling:", prob.dofs.f())
print("Final residuals:", prob.f())
print("Final state vector:", prob.x)
print("Final iota on axis:", vmec.iota_axis())
equil.boundary = surf

# 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('Delta(1,-1)', False)

# Each function we 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.
desired_iota = -0.41
prob = LeastSquaresProblem([(equil.iota_axis, desired_iota, 1)])

# Solve the minimization problem. We can choose whether to use a
# derivative-free or derivative-based algorithm.
least_squares_mpi_solve(prob, mpi, grad=False)

# Make sure all procs call VMEC:
objective = prob.objective()
if mpi.proc0_world:
    print("At the optimum,")
    print(" Delta(m=1,n=-1) = ", surf.get_Delta(1, -1))
    print(" iota on axis = ", equil.iota_axis())
    print(" objective function = ", objective)

assert np.abs(surf.get_Delta(1, -1) - 0.08575) < 1.0e-4
assert np.abs(equil.iota_axis() - desired_iota) < 1.0e-5
assert prob.objective() < 1.0e-15