Saving and loading pyqg models

[asross]

Currently, it's pretty difficult to save and load pyqg models. Trying to do it with pickle.dump and pickle.load doesn't work because of Cython issues, and trying to do it with netcdf files is also a bit tricky. This is currently the shortest way I know how to save and load a pyqg.Model:

import json
import pyqg
import numpy as np

# Initialize a model
m1 = pyqg.QGModel(**some_kwargs)
m1.run()

# Save it to disk, along with its parameters
ds = m1.to_dataset() # convert to xarray dataset
ds = ds.assign_attrs(pyqg_kwargs=json.dumps(m1._config)) # save the parameters used to generate the model
ds = ds.drop_vars([k for k,v in ds.variables.items() if np.iscomplexobj(v)]) # netcdf doesn't support complex numbers
ds.to_netcdf('my_model.nc') # finally dump it to netcdf

# Load it
ds = xr.open_dataset('my_model.nc')
m2 = pyqg.QGModel(**json.loads(ds.attrs['pyqg_kwargs']))
m2.q = ds.q.isel(time=-1).data # overwrite its main state variable
m2._invert() # regenerate other variables and fields
m2._calc_derived_fields()
for k in m2.diagnostics.keys(): # reset diagnostics
    m2.diagnostics[k]['value'] = ds[k].data
    m2.diagnostics[k]['count'] = 1

# Now m1 and m2 are essentially identical

This is a lot of work, and I would love to be able to just do:

# Initialize a model
m1 = pyqg.QGModel(**some_kwargs)
m1.run()

# Save it
m1.save('my_model.nc')

# Load it
m2 = pyqg.QGModel.load('my_model.nc')

Would this feature be useful, and is there a better way we can implement it other than the hacky code I just sketched out?

[rabernat]

Yes this is a good idea. We have done similar stuff when we needed to restart a model.

Ideally we could define a netCDF schema that contains:

• The model type (e.g. QGModel)
• All of the necessary parameters to initialize the model
• All of the state variables
• [maybe] the state of the diagnostic accumulators. Alternatively, we punt on this and just say that the diagnostics have to be reinitialized from scratch when you restart from a file. (That's what most GCMs do.)

To implement this, I would write two functions.

• Model.checkpoint('checkpoint.nc') - the challenge here is discovering all of the necessary attributes to save from each subclass 🤔
• from_checkpoint('checkpoint.nc') (top level function)
  • First this figures out what model class to initialize and initializes a fresh model
  • Next it updates the state from the file