google-research / neuralgcm

Hybrid ML + physics model of the Earth's atmosphere
https://neuralgcm.readthedocs.io
Apache License 2.0
594 stars 60 forks source link

How to output vertical velocity? #68

Open weatherforecasterwhai opened 3 months ago

weatherforecasterwhai commented 3 months ago

Would you please help me to output vertical velocity ? Vertical velocity is so important. Running the following after "Advancing in time" in "Deep-dive intor trained models" in the document: " from dinosaur import primitive_equations omega=primitive_equations.compute_vertical_velocity(advanced,model.data_coords) " But it runs in following errors:

Traceback (most recent call last): File "", line 1, in File "/home/whai/anaconda3/lib/python3.11/site-packages/dinosaur/primitive_equations.py", line 234, in compute_vertical_velocity sigma_dot_boundaries = compute_diagnostic_state(state, coords).sigma_dot_full ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/contextlib.py", line 81, in inner return func(*args, **kwds) ^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/dinosaur/primitive_equations.py", line 175, in compute_diagnostic_state nodal_vorticity = to_nodal_fn(state.vorticity) ^^^^^^^^^^^^^^^ AttributeError: 'ModelState' object has no attribute 'vorticity'

The "advaned" state "pprint" has vorticity. Would you please help me to fix this ?

shoyer commented 3 months ago

There is a nested State object -- try passing advanced.state into compute_diagnostic_state().