`BoundaryError` Jacobian can be `NaN` with some kinetic profiles

[ddudt]

The BoundaryError objective requires computing the pressure "p" at the boundary for the pressure balance equation. Only the pressure is used, not the pressure gradient. But if the pressure gradient is undefined at the boundary, then this results in NaN values for the Jacobian of the pressure balance equations.

A common use case is when creating kinetic profiles from an existing pressure profile:

# equilibrium with pressure profile
eq = get("DSHAPE")
print(f"p(1) = {eq.compute('p', grid=LinearGrid(rho=1.0))['p'][0]}")
print(f"p'(1) = {eq.compute('p_r', grid=LinearGrid(rho=1.0))['p_r'][0]}")

# change to equivalent kinetic profiles
density = PowerProfile(0.5, ScaledProfile(0.5 / elementary_charge, eq.pressure))
temperature = PowerProfile(0.5, ScaledProfile(0.5 / elementary_charge, eq.pressure))
eq.pressure = None
eq.atomic_number = PowerSeriesProfile([1], [0])
eq.electron_density = density
eq.electron_temperature = temperature
eq.ion_temperature = temperature
print(f"p(1) = {eq.compute('p', grid=LinearGrid(rho=1.0))['p'][0]}")
print(f"p'(1) = {eq.compute('p_r', grid=LinearGrid(rho=1.0))['p_r'][0]}")

# BoundaryError objective
field = ToroidalMagneticField(B0=0.2, R0=3.5)
objective = ObjectiveFunction(BoundaryError(eq=eq, field=field, field_fixed=True))
objective.build()

f = objective.compute_unscaled(objective.x(eq))
J = objective.jac_unscaled(objective.x(eq))
print(np.any(np.isnan(f)))
print(np.any(np.isnan(J)))

[f0uriest]

This isn't really an issue with BoundaryError, since force balance is NaN as well at rho=1. The issue is the non-integer power profile where the profile goes to zero, since $\frac{d\sqrt{x}}{dx} \rvert_{x = 0} = inf$

we might want to add some sort of safepow function with a custom jvp? or we may be able to rewrite the compute functions for ne_r etc to enforce 0*inf == 0

[dpanici]

Try second way, using 0*inf==0 and only use for the kinetic profile derivatives, make some sortof safemultiply function possibly

[dpanici]

Two issues

• first is computing things and getting nan
  • custom rule to p_r calculation can handle (make sure 0*inf = 0)
• second is is that BoundaryError's jacobian contains nan
  • harder bc this is derivatives w.r.t p_l etc.

Which elements of J are nan? wrt p0, p2 etc?

[ddudt]

PR #1712 is a partial fix that resolves this issue when using forward mode AD, but this remains an issue with reverse mode.