def test_set_get_dofs(self):
        """
        If the dofs are set to random numbers, and then you get the dofs,
        you should get what you set.
        """
        for mmax in range(1, 4):
            for nmax in range(4):
                surf = SurfaceHenneberg(nfp=1,
                                        alpha_fac=1,
                                        mmax=mmax,
                                        nmax=nmax)
                self.assertEqual(len(surf.names), len(surf.get_dofs()))

                # R0nH has nmax+1
                # Z0nH has nmax
                # bn has nmax + 1
                # rhomn for m=0 has nmax
                # rhomn for m>0 has 2 * nmax + 1
                nvals = nmax + 1 \
                    + nmax \
                    + nmax + 1 \
                    + nmax \
                    + mmax * (2 * nmax + 1)
                vals = np.random.rand(nvals) - 0.5
                surf.set_dofs(vals)
                np.testing.assert_allclose(surf.get_dofs(), vals)
    def test_indexing(self):
        """
        Check that set_rhomn() sets the expected element of the dofs vector.
        """
        mmax = 3
        nmax = 2
        surf = SurfaceHenneberg(nfp=2, alpha_fac=1, mmax=mmax, nmax=nmax)
        for n in range(nmax + 1):
            surf.R0nH[n] = 1000.0 + n
            surf.bn[n] = 2000.0 + n
        for n in range(1, nmax + 1):
            surf.Z0nH[n] = 3000.0 + n
        for m in range(mmax + 1):
            nmin = -nmax
            if m == 0:
                nmin = 1
            for n in range(nmin, nmax + 1):
                surf.set_rhomn(m, n, m * 100 + n)

        x = surf.get_dofs()
        np.testing.assert_allclose(x, [1.000e+03, 1.001e+03, 1.002e+03, 3.001e+03, 3.002e+03, 2.000e+03, 2.001e+03, \
                                       2.002e+03, 1.000e+00, 2.000e+00, 9.800e+01, 9.900e+01, 1.000e+02, 1.010e+02, \
                                       1.020e+02, 1.980e+02, 1.990e+02, 2.000e+02, 2.010e+02, 2.020e+02, 2.980e+02, \
                                       2.990e+02, 3.000e+02, 3.010e+02, 3.020e+02])
        self.assertListEqual(surf.names, [
            'R0nH(0)', 'R0nH(1)', 'R0nH(2)', 'Z0nH(1)', 'Z0nH(2)', 'bn(0)',
            'bn(1)', 'bn(2)', 'rhomn(0,1)', 'rhomn(0,2)', 'rhomn(1,-2)',
            'rhomn(1,-1)', 'rhomn(1,0)', 'rhomn(1,1)', 'rhomn(1,2)',
            'rhomn(2,-2)', 'rhomn(2,-1)', 'rhomn(2,0)', 'rhomn(2,1)',
            'rhomn(2,2)', 'rhomn(3,-2)', 'rhomn(3,-1)', 'rhomn(3,0)',
            'rhomn(3,1)', 'rhomn(3,2)'
        ])
 def test_axisymm(self):
     """
     Verify to_RZFourier() for an axisymmetric surface with circular
     cross-section.
     """
     R0 = 1.5
     a = 0.3
     for nfp in range(1, 3):
         for alpha_fac in [-1, 0, 1]:
             for mmax in range(1, 3):
                 for nmax in range(3):
                     surfH = SurfaceHenneberg(nfp=nfp,
                                              alpha_fac=alpha_fac,
                                              mmax=mmax,
                                              nmax=nmax)
                     self.assertEqual(surfH.num_dofs(),
                                      len(surfH.get_dofs()))
                     surfH.R0nH[0] = R0
                     surfH.bn[0] = a
                     surfH.set_rhomn(1, 0, a)
                     surf = surfH.to_RZFourier()
                     for m in range(surf.mpol + 1):
                         for n in range(-surf.ntor, surf.ntor + 1):
                             if m == 0 and n == 0:
                                 self.assertAlmostEqual(
                                     surf.get_rc(m, n), R0)
                                 self.assertAlmostEqual(
                                     surf.get_zs(m, n), 0.0)
                             elif m == 1 and n == 0:
                                 self.assertAlmostEqual(
                                     surf.get_rc(m, n), a)
                                 self.assertAlmostEqual(
                                     surf.get_zs(m, n), a)
                             else:
                                 self.assertAlmostEqual(
                                     surf.get_rc(m, n), 0.0)
                                 self.assertAlmostEqual(
                                     surf.get_zs(m, n), 0.0)