Complete cosmo_array numpy support and add cosmo_quantity

[kyleaoman]

Unyt has a unyt_quantity for scalars with units attached (a subclass of unyt_array). It makes sense to have an analogue cosmo_quantity subclass of cosmo_array, mainly so that various numpy functions don't silently cast to unyt_quantity. While I'm at it am also reviewing testing of numpy function return types and cleaning up helper functions as needed.

Closes #102
Closes #123

So far I have new tests implemented, and the new cosmo_quantity class implemented. As collateral to all of this I'm implementing wrapping all of the numpy functions supported by unyt. Those are now all wrapped. I'd previously wrapped all of the numpy ufuncs so most of the logic existed already, just repackaged existing helper functions and created a couple of new ones, and now cribbed off of the unyt wrappers to make sure I don't miss any edge cases (unless they missed any!).

After a bit of tidying up and a couple of squashed bugs in unyt the test suite now passes (with the not-yet-released unyt version with my patches applied, they say hopefully released by the end of the year). Will need to pin dependencies to unyt >=3.0.4, and probably numpy >=2.0, too.

Left to do:

• Think about more test coverage.
• Tidy up warnings in tests.
• Consider refactoring all of the "ca_cfs" logic by using cosmo_factor(None, None) instead. That's a big job, though.
• Look at a few places where to_comoving might be better as convert_to_comoving (in-place).
• Look at behaviour of defaulting to comoving when mixed arguments are given, there's a kwarg now that can control this, should I use it elsewhere?
• Tidy up definitions, probably move a bunch of helpers over from objects.py to _array_functions.py.
• Make a PR for BUG: numpy function take strips units with scalar index input yt-project/unyt#549 and get it merged, then implement a wrapper for np.take.
• Document _array_functions.py.
• Review documentation of objects.py
• Review/write test docstrings.
• Fix comoving == None prints as (Physical).
• Look at recursion in cosmo_array.__new__ for lists of lists of cosmo_arrays, etc. (shouldn't recurse here because unyt doesn't, but __new__ needs some cleanup if possible anyway).
• Implement bypass_validation?
• Look in visualisation routines to address Projection/slicing returns unyt_array instead of cosmo_array #123 ?
• Move the cosmo+unit wrapper for vis return images out of the backends and into the high-level render functions + update tests.
• Review & expand narrative docs.
• Review visualisation narrative docs.
• Recruit some beta testers.
• Check that everything is merged on unyt side and that they make a release. Will probably have to rename some stuff when they fix the fft->ftt typo.
• Pin dependencies (numpy >= 2.0.0 for sure, maybe >= 2.1.0; unyt >= 3.0.4).

[kyleaoman]

@JBorrow @robjmcgibbon Everything that both of you have raised is now addressed. I've checked and in some cases revised the visualisation examples from the narrative docs, they now all run out of the box.

There are a small number of API changes from the user perspective - a few necessary given the changes, and a few where it seemed like this would be a good time to make them. Notably the removal of boxsize as an argument to a couple of visualisation functions where an object with metadata is already being passed - better to just get the box size internally in that case.

There are two minor (non-blocking, imo) outstanding issues that need changed in unyt to be rectified. These have xfailing tests in place, so we'll catch them when unyt fixes them (I already contributed one fix that's merged to their main, just depends when they release again...). The two issues are that using np.block will crash, and that if someone ever decided to do U.Mpc * cosmo_array(...) they'd get back a unyt_array, whereas cosmo_array(...) * U.Mpc gives a cosmo_array as you'd expect it to. The way that __mul__ and __rmul__ work is kind of awful, had never really dug deep into that bit of python before.

One gotcha that I have no good way of addressing is that unyt only traverses one level of nested containers on initialisation. This means that cosmo_array([[x, x], [x, x]]) where x is a cosmo_array or unyt_array completely ignores any extra attributes of x and just treats it as a raw numpy array. A full-depth traversal would be a significant design change on the unyt side, not to mention being really really expensive in a lot of common use cases, so I don't think there's a strong argument to change this behaviour. It does tend to come up when you define say a mask region as [[xmin, xmax], [ymin, ymax], [zmin, zmax]] where all of those are cosmo_quantity objects... just need to get in the habit of coding around that, I guess. The docs do cover this, for all of the good that will do. Old code will usually complain with a warning here, so that's something.

Aaaaanyway, dare I say, ready to merge?

[kyleaoman]

Ah and I tested that numpy>=2.0.0 isn't enough now, so require >=2.1.0. And unyt>=3.0.4 now (just released with a bunch of fixes that I contributed throughout all of this).

[robjmcgibbon]

Fantastic, thanks Kyle, looks good to me. I will let Josh have the final say since he had more comments than me.

[JBorrow]

Are you ready to merge this? You just made a new commit 😅

[kyleaoman]

I'm tinkering with the leftover non-blocking stuff and discovered a dodgy test, figured I'd sneak it in...

[JBorrow]

Do you want a release to be minted or is it not ready for that?

[kyleaoman]

I think it's ready for a release. There might be a follow-up minor release, but that would first involve wrangling some changes in unyt and getting them to release them, so rather go ahead with this. As far as I know there's nothing serious outstanding.