How to output vertical velocity?

[weatherforecasterwhai]

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]

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