feat(cosmo): angular_diameter_distance(z1,z2) - #18807

Are you sure you want to arange (pun intented) things in this order? It reminds be a bit of range and np.arange(), where one argument is taken to mean 0, arg0, while two arguments are used as arg0, arg1. @neutrinoceros recently spent much more time than he would have liked getting that right for np.arange(..., like=u.Quantity)...

Might it be better to have something like angular_diameter_distance(z, /, *, z_obs=0.0)? It means more adaptation for those coming from the _z1z2 variant, but would seem much clearer. (Clear enough, in effect, that you could remove the warning for z2 < z1.)

I definitely don't want z_obs as a keyword argument as these aren't compatible with most jit operations, which I very much want to be able support as we adopt the array-api.

IMO func(z) and func(z1, z2) are close enough to how these functions are written mathematically (e.g. $d_A(z)$ and $d_A(z1,z2)$) that it's not confusing. Also, the observer doesn't have to be at z_obs, that's just where the calculation stops. In a ratio of distances, for example, $d_A(z1,z2)/$d_A(z_obs,z2)$ comparisons are being made relative to an observer.
IMO these should be positional-only arguments because it's the order that matters, not the variable name.

It doesn't have to be keyword-only, but effectively having default positional arguments before anything else just does not feel right (for once, typing shows this, as it requires the overload). How about *_distance(z, /, z_ref=0.0) as an alternative? (The mathematical precedent is just as bad; z1 and z2 are just poor choices of names, at least use z_lens, z_src or something like it.)

but effectively having default positional arguments before anything else just does not feel right

Maybe I'm not understanding. There isn't a default positional arguments before anything else. It proceeds z1.

for once, typing shows this, as it requires the overload

It actually doesn't.
Without the overloads it's:

def angular_diameter_distance(self, z: _InputT, z2: _InputT | None = None, /) -> u.Quantity: ...

mypy interprets this as:

• The function always accepts one mandatory argument z: _InputT.
• Optionally accepts a second positional argument z2: _InputT | None.
• The return type is always u.Quantity.

This type-checks fine. The caller can do:

cosmo.angular_diameter_distance(0.5)
cosmo.angular_diameter_distance(0.5, 1.0)

What the @overload does is it says you shouldn't do

cosmo.angular_diameter_distance(0.5, None)
# Mypy error: None not allowed for 2-argument overload

Now that I'm looking into this, maybe I should remove the @overload! None isn't pretty — it's better to omit the argument — but it is valid.

z1 and z2 are just poor choices of names

But they are positional-only parameters. Their names are meaningless. Which appears to be an advantage, since z_lens and z_src are use-case specific. z_ref is fine, but it's not that semantically meaningful, less so than it's argument position.

Edit: I removed the overloads!

these aren't compatible with most jit operations

first time I'm hearing about this. Is this specific to a particular JIT compiler, or is it general ? if not, what JIT do you have in mind ?

About making the signature tighter, have you considered having just one positional-only argument, that would accept either a single Quantity or a pair of quantities ? I find this kind of pattern pretty intuitive to use, and it got so much easier to implement and maintain with structural pattern matching (match/case) that I think they are an attractive idea now.

Hum I'll admit I didn't think of that, but it does seem like a 2 element 1D array doesn't actually match the structure you'd write for a 2-element tuple

>>> a = np.array((1, 2))
>>> match a:
...     case (_, _):
...         print('pair')
...     case _:
...         print('default')
...
default

is this a peculiarity in numpy, or would that be general to all Array API implementers ? if the latter, it looks to me that my idea still works as designed.

(btw, starting with keyword-only arguments and maybe later allow them to be passed positionally too would be backwards compatible, so we don't need to allow positional use from the get go to get the benefit of JIT compilation when we're ready to support it)

but it does seem like a 2 element 1D array doesn't actually match the structure you'd write for a 2-element tuple

But we want to support tuples as valid input as they are ArrayLike. So while 2-elt arrays may be distinguishable (though I have my doubts in general because of e.g. batch axes), the tuple behaviour makes this a non-starter.

starting with keyword-only arguments and maybe later allow them to be passed positionally too would be backwards compatible.

I'm just not sure how this is an improvement. I have yet to see a semantically meaningful variable name I actually like, while the argument position follows standard mathematical convention and thus clearly conveys its meaning. Any confusion should be alleviated by the docstring (if it isn't then we should improve the docstrings).
And keyword arguments are incompatible with mapping ops ("Oh, thanks for catching that. I meant mapping operators" — #18807 (comment)) so IMO it's best to be forward-thinking about this.
A Positional-Or-Keyword argument (so moving the "/") would be fine by me, but only if we had a good semantically meaningful variable name. @mhvk's suggestion of z_ref is definitely the best I've heard so far and IMO it's only meh (is it a reference redshift or the redshift of reference / observer or observee?).
The advantage of starting with Positional-Only is that it's backwards compatible with making an argument Positional-Or-Keyword. So IMO we should start with that and if ever there's a good variable name think about renaming z2 (it's positional and thus has no true name) and making it Positional-Or-Keyword.

@mhvk if you're interested where the overloads originally came from, it's https://cosmology.readthedocs.io/projects/api/latest/api/reference/cosmology.api.HasAngularDiameterDistance.angular_diameter_distance.html#cosmology.api.HasAngularDiameterDistance.angular_diameter_distance,
where we customized Sphinx to read the overloads to build more informative docs.
In that API document we've thought through improvements to the Astropy Cosmology API and I've been slowly working to transition us towards these improvements.

I can remove the overload in this or a followup PR.

I'm thinking about removing this warning in the future. But since this function replaces the deprecated angular_diameter_distance_z1z2 and that had the warning I think it's good to keep for now.

curious, what is the actual issue with z1<z2 ?

z2<z1 makes the distance go negative.

purely stylistic nit, feel free to ignore, though IMO this is easier to understand at a glance

Given that this will appear in a few methods, I think in a followup PR these lines should be shunted off to a helper function. I'd want to do timing tests then. Also, I think this helper function + aszarr are good targets for C code.
Given that, I'm going to push off rewrites here.

curious, what is the actual issue with z1<z2 ?