Cosmo: allow broadcasting in z_at_value

[nstarman]

Signed-off-by: Nathaniel Starkman (@nstarman) nstarkman@protonmail.com

Description

z_at_value only works on scalars. numpy.vectorize cannot be used for broadcasting because it does not preserve units.
This PR renames the current z_at_value to scalar_z_at_value and the new z_at_value correctly broadcasts inputs, calling scalar_z_at_value for each element.

Edit: this does NOT obviate the need for #11361. There is definitely good reason to have an interpolated approach for large arrays. This PR just enables broadcasting on the current function.

Fixes: #11949

[github-actions]

👋 Thank you for your draft pull request! Do you know that you can use [ci skip] or [skip ci] in your commit messages to skip running continuous integration tests until you are ready?

[dhomeier]

I won't have time to provide an actual code review this week, but can you describe how this would be placed in the context of #11361? Specifically, in the scheme of #11361 (comment)
this could constitute the branch taken for moderate input lengths; with larger arrays switched to the interpolation scheme.

Also there is probably some overhead in repeatedly calling scalar_z_at_value that might be removed when using identical method etc., but I am not sure if it has measurable impact.

[nstarman]

I won't have time to provide an actual code review this week, but can you describe how this would be placed in the context of #11361? Specifically, in the scheme of #11361 (comment)
this could constitute the branch taken for moderate input lengths; with larger arrays switched to the interpolation scheme.

I think it's more in the vein of #11361 (comment) and #11361 (comment), where it was suggested that the implementation in #11361 is sufficiently different that it should be in a separate function "z_at_array", and that "z_at_value"� should just be vectorized to a fancy for-loop. No speed increase. Just a lot more convenient.
So #11361 can and should go ahead for doing fast calculations on large arrays. This just makes "z_at_value"� broadcast as expected.

Also there is probably some overhead in repeatedly calling scalar_z_at_value that might be removed when using identical method etc., but I am not sure if it has measurable impact.

I clocked ``z_at_value` as 0.5 ms slower for a single calculation and 4 ms slower for an array of 50
. Surely there are further optimizations to be found, but this seems like it scales reasonably well.

[dhomeier]

I think it's more in the vein of #11361 (comment) and #11361 (comment), where it was suggested that the implementation in #11361 is sufficiently different that it should be in a separate function "z_at_array", and that "z_at_value"� should just be vectorized to a fancy for-loop. No speed increase. Just a lot more convenient.
So #11361 can and should go ahead for doing fast calculations on large arrays. This just makes "z_at_value"� broadcast as expected.

Ah, I think I had missed that part of the discussion or did not realise that it was still pertinent for wrapping up #11361.
But I agree, since the interpolation method makes it functionally different and it will indeed produce slightly different results, it is entirely sensible to put that one into a separate function.

Also there is probably some overhead in repeatedly calling scalar_z_at_value that might be removed when using identical method etc., but I am not sure if it has measurable impact.

I clocked ``z_at_value` as 0.5 ms slower for a single calculation and 4 ms slower for an array of 50

That's actually not measuring the overhead I had in mind – I was rather thinking about doing the setup for the minimiser method only once (when using the same method for all values) and/or vectorising the fval_zmin, fval_zmax, bracket calculation etc. and basically only looping over the minimize_scalar itself.
But when allowing even method as an array input as you are doing here (which I am still unsure if it is a good idea) that setup would become rather tricky, and from my timings for #11080 I still expect by far the most time to be spent in minimize_scalar, so it's probably not worth the effort optimising it here further.

[nstarman]

@dhomeier. The more I think about it, the more I agree that 'method' should not be broadcasted. I've excluded it from the broadcast.

I also add a test to check whether numpy.vectorize works on quantities and this whole PR can be replaced. 🤞 some time in the future it will fail and we can just use numpy.vectorize.

[nstarman]

This PR is premised on the assumption that numpy.vectorize won't soon work on Quantities AND that we should not subclass numpy.vectorize and make the fix ourselves (presumably to be upstreamed). I'm looking at a number of other functions in astropy/cosmology that are not vectorized and I'm wondering about this two assumptions.

@mhvk as a resident expert on numpy...

[mhvk]

@nstarman - it would definitely be good to get np.vectorize to work - just never has risen to high enough priority to think what is blocking it. Probably best, though, to discuss in a separate issue: Could you make a minimal example that shows that a function that itself works with Quantity cannot be used with np.vectorize?

[nstarman]

👍 See #11893 for a semi-functional example implementation.

[nstarman]

Ok, if all this passes, I still need to do a full docstring for z_at_value. I think it now merits documentation. The bracket rules in particular need explication.

[mhvk]

Somewhat random comments. Does look good generally.

[nstarman]

If it's approved, I'm going to squash some of my commit history. So don't merge, thanks!

[pllim]

So don't merge, thanks!

Might be safer to turn this into draft then, but up to you. Thanks!

[mhvk]

Looks good! Two very small comments that are easy to implement, and a question whether you even want to expose z_at_scalar_value.

[mhvk]

I thought this was now OK!?

[nstarman]

It's fixed in that I manually check for a scalar and correct the output. I think there's some way to get 'reduce_ok' to do this automatically, but I can't get the examples in the numpy documentation to work here.

[nstarman]

Hence the note for myself or an enterprising contributor. Low priority, but it would be nice.

[mhvk]

Ah, I thought it was the problem that the code returned a 1-D array rather than an array scalar. I think reduce_ok is really meant for dimensions that are reduced over (i.e., it tells the iterator that it is OK for a dimension to be missing or be unity in the output).

[nstarman]

In the first example in https://numpy.org/doc/stable/reference/arrays.nditer.html#reduction-iteration it seems to work for degrading to a scalar.

[mhvk]

Indeed, and I did not know that. But I think it is not directly relevant, since you are not actually reducing over any dimension..

[mhvk]

The overhead from setting up the iterator is likely small compared to the actual function call, so might it make sense to just state that z_at_value can now handle arrays? I.e., not mention z_at_scalar_value at all?

[dhomeier]

The overhead from setting up the iterator is likely small compared to the actual function call,

That was the one clocked in #11778 (comment), right?

[nstarman]

That was a previous implementation. I'll reclock.

[nstarman]

It's a pretty negligible difference. And scales as we'd expect.

[dhomeier]

The extra overhead from calling z_at_value on an array is expected? Not a serious difference, just wondering where that additional time comes in...

[nstarman]

Looking at the above, I'm inclined to agree.
The point of z_at_scalar_value public is that it's more performant. Except it basically isn't. I'm going to make it a private function. The interpolated case (#11361) should introduce a new function, perhaps z_at_interp_value. If I can make z_at_scalar_value more performant relative to z_at_value, I'll make it public.

[mhvk]

@dhomeier - list comprehension is really quite fast, hard to np.nditer to match, with all the broadcasting etc.

[dhomeier]

The point of z_at_scalar_value public is that it's more performant. Except it basically isn't. I'm going to

Well, 12-15 % faster is still faster, even if perhaps not enough to justify keeping it public. What was more surprising to me was that even looping over it 100 times is still 15 % faster than calling z_at_value directly on the value. Unless you are seeing some caching artefacts in timeit, but I can't really see where they'd come in.

[dhomeier]

@dhomeier - list comprehension is really quite fast, hard to np.nditer to match, with all the broadcasting etc.

Sounds sensible – I had not looked at the internals since my early incomplete review, so I thought there was still some code that the vectorised version did not execute on every single element of fval; but seeing now that there is also a bunch of extras included for accepting sequences of bracket etc.

[nstarman]

That's probably from the fact that z_at_value needs to consider how to broadcast all the values. I think if I were to do array nesting and so forth we would see that the manual implementation is equivalent to z_at_value. z_at_value just doesn't know when it can simplify. If python allowed for multiple dispatch..

[nstarman]

Ok. Tests should pass and I've cleaned the commit history.
Edit: and some spacing in the docstrings.

[mhvk]

Nice!