Add prolate spheroidal coordinates

[adrn]

I'm gearing up to add the Staeckel Fudge to Galax, so it would be useful to have a representation of prolate spheroidal coordinates. Transformations to this coordinate system need to specify a focal length parameter, Delta, so it's a bit different from other d3 representations. Thoughts on this approach?

[adrn]

Tests are failing for two reasons that might be related:

• How can I tell coordinax that Delta is not a coordinate component and is instead a parameter that helps specify the coordinate system?
• I can't get the self-transform to work (without explicitly specifying a value for Delta in represent_as) because something is inspecting the call signature and knows that that is a missing field in the dataclass.

One approach could be to allow a default Delta=None (because writing the coordinates themselves doesn't require a focal length, only transformations)?

[nstarman]

How can I tell coordinax that Delta is not a coordinate component and is instead a parameter that helps specify the coordinate system?

That's a good question! We don't have a mechanism for that. This is complicated, essentially each value of $\Delta$ should result in a different representation. The non-viable but technically correct solution is to dynamically create a new class for each $Delta$. Not doing that, we could make $\Delta$ just be a property of the transformation, e.g. a mandatory kwarg to represent_as. However we still want to ensure math happens correctly, e.g. when 2 coords with different $\Delta$ are added it errors (or does some transformation then adds correctly). Therefore __add__ needs to know about $\Delta$, which can only happen if it's stored on the object. There are two options, both are a fair bit of work, and I'm not sure which is best.

1. $\Delta$ is a field on the representation class. We would then need to change the fundamental assumption on representations, which is that all fields are coordinate axes. coordinax uses field_items() and related functions liberally to perform generalized operations on representations. We would need to define new functions like axis_field_items to replace the other functions everywhere (or leverage multiple dispatch as I discuss below).
2. We define a new class (name TBD) RepDeltaWrapper that accepts the representation and the delta. The representation is still "pure" in that all fields are coordinates. RepDeltaWrapper would need to be made to work with the existing representation machinery. This might end up involving all the changes described for (1) + registering all the primitives.

So writing this out, (2) might be technically / mathematically better, but (1) seems best to do now.

Order of events: 😓

1. Define an Axis field descriptor, like how I've done in astropy.cosmology.Parameter. This is used to annotate all the axis fields in coordinax (all of them currently). It should be parametric wrt the type it sets, so Axis[Distance] or Axis[Angle].
2. Open a PR in dataclassish to define a filter flag type base class, like we did in unxt. The Flag ABC should error, the 'default' filter flag should be the same as if it weren't present. This is actually technically not needed because the multiple dispatch only cares about step 4. So maybe we just open an Issue on dataclassish for now.
3. In coordinax define an AxesFilterFlag and register its behaviour for the field_items() etc functions. It can filter the dataclass' fields by whether they are annotated as Axis types in the class definition (they would still be Quantity on the instance).
   At this point all the existing representations should work completely unchanged since all their fields are axes.
4. Define the prolate spheroidal coordinate classes, with Delta not as an axis. Register in all the behaviours for checking Delta when doing math operations, etc. Random extra thought: Delta should be a scalar only.

[nstarman]

@adrn I started steps 1-3 in #225. The tests on using filtered axes is pretty minimal — just on FourVector — so there might be a couple functions I forgot to refactor. But this PR should now be able to define the Delta attribute!

[adrn]

@nstarman I think this is ready for a look! Tests are failing with a gnarly recursion error in test_jax_ops.py that I haven't figured out yet.

[adrn]

OK it turns out the recursion error was because one of the tests was calling the fallback represent_as(CartesianPos3D, ProlateSpheroidalPos), which shouldn't be possible because this transformation has no defined Delta, and (unrelated) was triggering a recursion. Fixing in the next batch of commits

[codecov]

Codecov Report

Attention: Patch coverage is 97.41935% with 4 lines in your changes missing coverage. Please review.

Project coverage is 97.03%. Comparing base (c2e561e) to head (26bf681).

Files with missing lines	Patch %	Lines
src/coordinax/_src/vectors/d3/spheroidal.py	93.65%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #220      +/-   ##
==========================================
+ Coverage   96.98%   97.03%   +0.04%     
==========================================
  Files         113      114       +1     
  Lines        3455     3605     +150     
==========================================
+ Hits         3351     3498     +147     
- Misses        104      107       +3     

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

[adrn]

@nstarman OK tests are passing, so take a look at how I handled Delta and the kwargs and let me know what you think!

[nstarman]

Looks really good. Small comments only. Then should be ready 🚀

[nstarman]

Can this be more concise while retaining clarity?

[nstarman]

ChatGPT suggests:

• ref_name (short for “reference name”)
• limit_name (indicating the name of the limiting value)
• max_name (if it’s always describing the maximum)
• comp_name (short for “comparison name”)

Which do you prefer / have another suggestion?

[adrn]

🤷 - is it worth it for a few characters on an internal function? I can, but personally don't think it's worth it!

[adrn]

Hm, not sure about the failing test. Adding the converter to VectorAttribute causes filter_jit to fail here:

@eqx.filter_jit
def func(q):
    return q.represent_as(cx.ProlateSpheroidalPos, Delta=Quantity(1.0, "kpc"))

q = cx.CartesianPos3D.from_([1, 2, 3], "kpc")
func(q)

[nstarman]

Can you try removing the parametrization in the converter? Maybe just Quantity.from_. I think the dimensions will be checked against the field annotation when it's being set.

[nstarman]

But also Happy Thanksgiving 🦃! For after the holidays :)