feat: MN3 and MilkyWayPotential2022

[adrn]

This implements the MN3 and MilkyWayPotential2022 potential models.

In writing the implementations, I have a few questions:

1. Subclassing AbstractPotential seems to require specifying several boilerplate class attributes:

_: KW_ONLY
units: AbstractUnitSystem = eqx.field(converter=unitsystem, static=True)
constants: ImmutableMap[str, Quantity] = eqx.field(
    default=default_constants, converter=ImmutableMap
)

Can you say more about these? Is there no way we can move default behavior to the __check_init__ of the base class?

2. Why do the potential methods all have to be decorated with @partial(jax.jit)? Isn't that handled by equinox?
3. The tests are currently failing because I need to update the expected values for the two potentials I implemented. How did you compute these and how do you know they are "correct" or actually expected? In Gala, I did some tests against sympy -- do we want something like that in galax?

[nstarman]

Subclassing AbstractPotential seems to require specifying several boilerplate class attributes:

Yeah. Technically units and constants are inherited and don't need to be specified.
This is a stylistic thing, where the norm is to have all fields specified at class declaration.
So in galax IMO it's good that when users code-peak they see everything. When they write a potential, 🤷.

Why do the potential methods all have to be decorated with @partial(jax.jit)? Isn't that handled by equinox?

It's not handled by equinox. That handles PyTree'ing self and other stuff.
Decorating everything with jit is something we need to AB test! I think most of these should be converted to @partial(jax.jit, inline=True), which should allow JAX to intelligently fuse functions so that only the top-level function appears jitted. ATM I think we are slowing stuff down by jitting sub-functions.

How did you compute these and how do you know they are "correct" or actually expected?

Manually. This isn't for correctness, only to ensure nothing changes.
To check for correctness we are testing against gala.

In Gala, I did some tests against sympy -- do we want something like that in galax?

Yes!

It would be great to have 3 types of tests:

1. Correctness tests against Sympy (instead of against gala)
2. API drift tests (what we currently have, but let's do it more sensibly with griffe for static API analysis and then pytest-arraydiff for I/O)
3. Ecosystem checks against gala, galpy, agama, etc.

[nstarman]

😎

[codecov]

Codecov Report

Attention: Patch coverage is 99.20000% with 1 line in your changes missing coverage. Please review.

Project coverage is 96.75%. Comparing base (6f057d0) to head (0949a07).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/galax/potential/_src/builtin/special.py	93.75%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #418      +/-   ##
==========================================
+ Coverage   96.63%   96.75%   +0.12%     
==========================================
  Files          83       84       +1     
  Lines        3090     3207     +117     
==========================================
+ Hits         2986     3103     +117     
  Misses        104      104              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[adrn]

I sloppily let a merge commit sneak in but maybe that is ok if you squash and merge?

[nstarman]

Nice! Seeing the separated classes convinced me that this is the right approach.

[adrn]

I think the only remaining issue here is to address the "low" test coverage of the added potentials.

[adrn]

@nstarman I think this is ready for another look!

[nstarman]

🎉