Thank you for your first document, but how to fix this dimension bug in "Deep-dive into trained models"?

[weatherforecasterwhai]

Hi,

I can't install with "pip install git+https: //...", so I first download the git, then "pip install ." in the neuralGCM directory.

Running "inputs, forcings = model.data_from_xarray(ds_init)" with these bugs:

Traceback (most recent call last): File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/variable.py", line 132, in as_variable obj = Variable(*obj) ^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/variable.py", line 367, in init self._dims = self._parse_dimensions(dims) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/variable.py", line 683, in _parse_dimensions raise ValueError( ValueError: dimensions ('time',) must have the same length as the number of data dimensions, ndim=0

During handling of the above exception, another exception occurred:

Traceback (most recent call last): File "", line 1, in File "/home/whai/neuralgcm/neuralgcm/api.py", line 353, in data_from_xarray inputs = self.inputs_from_xarray(dataset) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/neuralgcm/neuralgcm/api.py", line 340, in inputs_from_xarray return self._data_from_xarray(dataset, self._input_variables) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/neuralgcm/neuralgcm/api.py", line 331, in _data_from_xarray dataset = self._dataset_with_sim_time(dataset) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/neuralgcm/neuralgcm/api.py", line 276, in _dataset_with_sim_time return xarray_utils.ds_with_sim_time( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/dinosaur/xarray_utils.py", line 905, in ds_with_sim_time ds['sim_time'] = ((XR_TIME_NAME,), nondim_time) ~~^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/dataset.py", line 1537, in setitem self.update({key: value}) File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/dataset.py", line 5253, in update merge_result = dataset_update_method(self, other) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/merge.py", line 1104, in dataset_update_method return merge_core( ^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/merge.py", line 755, in merge_core collected = collect_variables_and_indexes(aligned, indexes=indexes) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/merge.py", line 354, in collect_variables_and_indexes variable = as_variable(variable, name=name) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/whai/anaconda3/lib/python3.11/site-packages/xarray/core/variable.py", line 135, in as_variable raise error.class( ValueError: Could not convert tuple of form (dims, data[, attrs, encoding]): (('time',), -92034.607104) to Variable.

[shoyer]

We also made some fixes to the Dinosaur package -- please update that as well!

[weatherforecasterwhai]

Thank you, Sir. After reinstalling dinosaur, it fixes!

I'm really confused about the "hybrid" title of neuralGCM, just like issue #52. Maybe what I asked wasn't clear. Because I can't understand enough of your open codes. Would you please explain fig1 in the paper with the codes and running sequence? How does they "hybrid"? I've read, searched and thought about this, which even make me sleepless.

Really thank you for your attention!

[weatherforecasterwhai]

@shoyer Found a tiny mistake in "Deep-dive into trained models": "pprint.pprint(jax.tree.map(np.shape, encoded))" may be like "pprint.pprint(jax.tree_map(np.shape, encoded))"

And, I want to output vertical velocity after advancing one time like this: " 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 bug? Is the vertical velocity calculated from dinosaur just the large scale omega or the "hybrid" of large scale dynamical and learned parameterization?

Thank you!

[shoyer]

You might find it helpful to look at the gin configuration for our model, starting with the WhirlModel object and working your way down: https://github.com/google-research/neuralgcm/blob/0a18c4b4631d53634fa7f201541c8fa11dfe4e88/neuralgcm/reference_code/paper_configs/stochastic_1_4_deg.gin#L3406

The hybrid model is configured via the StochasticPhysicsParameterizationStep object, with the corrector_module solving the MoistPrimitiveEquationsWithCloudMoisture and the physics parameterization using `DivCurlNeuralParameterization.

[shoyer]

@shoyer Found a tiny mistake in "Deep-dive into trained models": "pprint.pprint(jax.tree.map(np.shape, encoded))" may be like "pprint.pprint(jax.tree_map(np.shape, encoded))"

If you update to the newest version of jax, jax.tree.map should be defined.

[weatherforecasterwhai]

Great help to me, Sir! Tracing from MoistPrimitiveEquationsWithCloudMoisture to time_integration.ImplicitexplitcitODE, the dinosaur part seems to "hybrid" with DivCurlNeuralParameterization through modal_tendencies = self.to_div_curl_fn(nodal_tendencies) by 'u,v,div,curl'. And, would you please train neuralGCM to a higher resolution model, maybe as high as 0.25 degree? Maybe the higher resolutio, the predition is better, just like your paper. It seems neuralGCM using "google_proprietary_code" which others can't use to train neuralGCM.

[shoyer]

We would certainly love to train a 0.25 degree model but haven't figured out how to do it yet, even training at 0.7 degrees with our current approach was extremely expensive and slow.

[weatherforecasterwhai]

We would certainly love to train a 0.25 degree model but haven't figured out how to do it yet, even training at 0.7 degrees with our current approach was extremely expensive and slow.

I've ever asked why not using Transformer instead of MLP in neuralGCM. I don't know MLP even make things so hard. Really thank you for your excellent work!

[weatherforecasterwhai]

@shoyer Would you please help me to output vertical velocity ? Vertical velocity is so important. And I've asked this bug. Running the following after Advancing in time in Deep-dive intor trained models: " 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 bug?

[shoyer]

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

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

[shoyer]

I'm going to close this issue -- please ask new questions in separate issues for each question.