DedalusProject / dedalus

A flexible framework for solving PDEs with modern spectral methods.
http://dedalus-project.org/
GNU General Public License v3.0
492 stars 115 forks source link

xarray interface #213

Closed kburns closed 2 years ago

kburns commented 2 years ago

This PR implements an xarray backend for reading in Dedalus datasets.

kburns commented 2 years ago

While loading tasks to DataArrays seems to work great, I think loading full save files to DataSets in general might be bad -- if the variables have different grids for the same coordinates (like a ball and then surrounding shell, or maybe even a ball and a radial slice), my understanding is that xarray will merge/concatenate the two coordinate arrays and expand each dataarray to cover the entire joined coordinate if you try to add them to the same DataSet. I think the quick solution is to just return a dictionary of individual DataArrays -- each with whatever coordinates necessary -- when loading a full output file.

This also somewhat diminishes my motivation for fleshing out lazy loading / virtual merging across time, but that would still be a nice feature to have.

kburns commented 2 years ago

In the last commit I've also added the xarray interface to the post-processing tutorial notebook alongside the current h5py instructions. Maybe keeping them both makes sense until we have more experience / have stabilized the xarray interface a bit more, but it's nice to see how it can cut down on lines of code when plotting.

Unfortunately it's not super clear how to make this an easy replacement for plot bot, since, as far as I can tell, the xarray plotting doesn't offer a simple interface for doing things like automatically matching plot aspect ratios to data, etc., so the same boiler plate around setting up figure panels/axes will essentially need to remain.

kburns commented 2 years ago

Rebased onto latest master, and I think good to merge. This does not include lazy loading, but we could add that after if it's really a problem for big datasets, but until we're doing virtual merges across time, this probably isn't a big problem.

jsoishi commented 2 years ago

Looks good to me!