[amss] Rewrite lecture using JAX

[HumphreyYang]

Rewrite amss using JAX will also resolve #223 related to removing the outdated interpolation package.

We can use the interpolater in jax.scipy https://docs.jax.dev/en/latest/_autosummary/jax.scipy.interpolate.RegularGridInterpolator.html in JAX as suggested by @jstac.

The lecture will also benefit from autodifferentiation.

[HumphreyYang]

I have been trying to rewrite this lecture in JAX, but I failed to generate the same plot for the complete vs. incomplete market case under the same setup, even though the code produces the same results for the first few plots.

I think the code in this lecture is sensitive to numerical precision and diverges with slight changes in parameters or precision, as noted in the code comment:

class AMSS:
    # WARNING: THE CODE IS EXTREMELY SENSITIVE TO CHOICES OF PARAMETERS.
    # DO NOT CHANGE THE PARAMETERS AND EXPECT IT TO WORK

I will try to rewrite the code without aiming to replicate the original plot and instead conduct sanity checks to see whether the results align with the theoretical predictions.

I will work on it once we push towards the advanced series.

[jstac]

Many thanks @HumphreyYang !

@thomassargent30 will be the best guide concerning theory.

It's not great that we have that comment in our code so it will be nice if we can get a little more robustness :-)

(BTW, did you try with 64 bit floats? That might make a difference.)

[HumphreyYang]

Many thanks, @jstac for the suggestion!

Yes, I used high-precision mode, but the result is still off. I’ll try to make it more robust and see if we can add an exercise where we can tweak the parameters and perform sensitivity analysis.

[jstac]

Thanks @HumphreyYang, it's a big challenge. Please touch base with @thomassargent30 regarding changes.