Using xarray with pyrealm?

[davidorme]

Is your feature request related to a problem? Please describe.

So far, pyrealm has been written around using numpy for calculations and expects numpy inputs. However, using xarray might make it easier to implement some functionality, and the ability to label axes could be very useful. Particularly here, being able to identify a time axis and then perform averaging along that axis seems like something that xarray is designed around and provides methods and classes to do. It would avoid having to reimplement (probably poorer) custom solutions using np.ndarrays specifically for pyrealm.

At the moment pyrealm.utilities.check_input_shapes takes a very hard line on only accepting np.ndarray. That is easily changed to allow DataArray inputs, which are otherwise array-like and implement a lot of the numpy API. The problem then is that xarray doesn't implement the whole numpy API, so not all things work. I've tinkered with sticking DataArrays into the P Model implementation and they break on the use of np.xxx.outer in pyrealm.pmodel.functions used in PModelEnvironment - I didn't push the experiment further, but I'd expect similar issues.

It seems like a lot of complexity to change the existing code to be agnostic (or to have numpy/xarray implementations), but equally reinventing the wheel of what xarray does seems like a mistake. Equally, I don't think that making it xarray only is a good idea - it is likely to be slower for one thing.

So, there's a possible advantage but it seems like a bit of a distracting rabbit-hole right now, so I'm going to park it in this issue. Interesting things I saw that might be relevant:

• https://stackoverflow.com/questions/52043669/decorators-to-pass-numpy-arrays-or-xarray-arrays-to-functions
• https://docs.xarray.dev/en/stable/user-guide/computation.html#wrapping-custom-computation

Describe the solution you'd like

Allow seamless use of xarray.DataArray and np.ndarray inputs as array-like inputs.