Fix NaN in Omnigenity

[dpanici]

Resolves #2120

[dpanici]

This fails, but with M_B=2 it passes

It also fails when NFP=1. So issue seems to be when M_B=3+ here?

    @pytest.mark.unit
    @pytest.mark.parametrize(
        "helicity", [(1, 0)]
    )  # @pytest.mark.parametrize("helicity", [(1, 0), (1, 1), (0, 1)])
    def test_objective_no_nangrad_omnigenity(self, helicity):
        """Omnigenity."""
        surf = FourierRZToroidalSurface.from_qp_model(
            major_radius=1,
            aspect_ratio=20,
            elongation=6,
            mirror_ratio=0.2,
            torsion=0.1,
            NFP=2,
            sym=True,
        )

        eq = Equilibrium(Psi=6e-3, M=4, N=4, surface=surf)
        field = OmnigenousField(
            L_B=1,
            M_B=3, # works for 2
            L_x=1,
            M_x=1,
            N_x=1,
            NFP=eq.NFP,
            helicity=helicity,
            # B_lm=np.array([0.8, 1.2]),
        )


        obj = ObjectiveFunction(Omnigenity(eq=eq, field=field))
        obj.build()
        g = obj.grad(obj.x())
        assert not np.any(np.isnan(g)), str(helicity)

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -1.80 %    |     3.991e+03      |     3.920e+03      |    -71.79    |       39.34        |       36.64        |
  test_proximal_jac_w7x_with_eq_update   |    1.89 %    |     6.470e+03      |     6.593e+03      |    122.30    |       161.60       |       162.19       |
  test_proximal_freeb_jac                |    0.41 %    |     1.319e+04      |     1.324e+04      |    53.52     |       84.33        |       83.71        |
  test_proximal_freeb_jac_blocked        |    0.13 %    |     7.508e+03      |     7.517e+03      |     9.72     |       74.48        |       73.37        |
  test_proximal_freeb_jac_batched        |    0.29 %    |     7.487e+03      |     7.508e+03      |    21.48     |       73.17        |       73.28        |
  test_proximal_jac_ripple               |    1.62 %    |     3.473e+03      |     3.529e+03      |    56.20     |       66.07        |       66.17        |
  test_proximal_jac_ripple_bounce1d      |    2.78 %    |     3.509e+03      |     3.606e+03      |    97.71     |       77.12        |       77.42        |
  test_eq_solve                          |    0.93 %    |     1.984e+03      |     2.002e+03      |    18.41     |       94.91        |       94.62        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[dpanici]

Further info: passing B_lm to be a well (instead of the default constant |B|=1 T) avoids the NaN. So the issue maybe appears to be when we have more than 2 spline knots and they are all a constant value.

i.e. passing

            B_lm=np.concatenate(
                (
                    np.array([0.999, 1.0, 1.001]),
                    np.zeros((1 * 3,)),  # same field on all flux surfaces
                )
            ),
        )

into the Omnigenity objective above seems to avoid the NaN completely. So the issue must be the splining of a flat well