Pretrained model configs diverged from code

[robert-DL]

The available models that are pretrained diverged from the codebase.

In particular

1. the val and the metric entries are missing in the archesweathergen config of the pretrained diffusion model. This leads to issues when one tries to load the model.
2. Loading models for mode=test requires a 'pytorch-lightning_version' key in the state dicts ( triggered when calling trainer.test).

Solutions:

1. can be resolved by updating the configs accordingly.
2. can be resolved by either modifying the state dicts or loading the state dicts before trainer.test()

[thaondc-mbzuai]

Hello @robert-DL, I’ve encountered the same issue and was wondering if you have found out the correct parameters for the .yaml config file of archesweathergen?

[robert-DL]

Hi @thaondc-mbzuai!

Indeed there is a solution to the problem that works for now. You can adapt the "old" configs with the new ones you can find in the config folder of the repository - everything missing (val key) should be there.

For the pytorch_lightning_version problem you can either add a surrogate version (0.0.0) to the state dict manually or load the state dict manually before trainer.test in main_hydra.py and set checkpoint_pth=None in trainer.test()

An update will be planned accordingly.

[thaondc-mbzuai]

Thank you, @robert-DL. I’ve resolved my issue by updating the .yaml config file for archesweathergen using the configs folder, and archesweathergen now runs smoothly.

However, I noticed that several evaluation metrics mentioned in the publication (Continuous Ranked Probability Score (CRPS), Energy Score, spread-skill ratio, and activity) are not currently implemented in the source code. It would be greatly appreciated if these could be included in a future update.

Additionally, the rank histogram metric appears to be non-functional at the moment. I encountered the following error during its execution:

File "/l/users/thao.nguyen/geoarches/geoarches/evaluation/eval_multistep.py", line 302, in 
    main()
File "/l/users/thao.nguyen/geoarches/geoarches/evaluation/eval_multistep.py", line 271, in main
    labelled_metric_output = metric.compute()
                             ^^^^^^^^^^^^^^^^
File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/torchmetrics/metric.py", line 700, in wrapped_func
    value = _squeeze_if_scalar(compute(*args, **kwargs))
                               ^^^^^^^^^^^^^^^^^^^^^^^^
File "/l/users/thao.nguyen/geoarches/geoarches/metrics/metric_base.py", line 153, in compute
    outputs = metric.compute()
              ^^^^^^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/torchmetrics/metric.py", line 700, in wrapped_func
    value = _squeeze_if_scalar(compute(*args, **kwargs))
                               ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/l/users/thao.nguyen/geoarches/geoarches/metrics/label_wrapper.py", line 232, in compute
    return self._convert(self.metric.compute())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/l/users/thao.nguyen/geoarches/geoarches/metrics/label_wrapper.py", line 220, in _convert
    ds = xr.Dataset(
         ^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/xarray/core/dataset.py", line 747, in __init__
    variables, coord_names, dims, indexes, _ = merge_data_and_coords(
                                               ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/xarray/core/dataset.py", line 460, in merge_data_and_coords
    return merge_core(
           ^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/xarray/core/merge.py", line 699, in merge_core
    collected = collect_variables_and_indexes(aligned, indexes=indexes)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/xarray/core/merge.py", line 358, in collect_variables_and_indexes
    variable = as_variable(variable, name=name, auto_convert=False)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/thao.nguyen/.conda/envs/weather/lib/python3.11/site-packages/xarray/core/variable.py", line 155, in as_variable
    raise error.__class__(
TypeError: Variable 'rankhist': Could not convert tuple of form (dims, data[, attrs, encoding]): (['prediction_timedelta', 'variable', 'rank'], tensor([[[1.0628, 0.9656, 0.9889,  ..., 0.9742, 0.9807, 1.1022],
         [1.1722, 0.9815, 0.9830,  ..., 0.9595, 0.9747, 0.9964],
         [1.7516, 1.1756, 1.0709,  ..., 1.0399, 1.1551, 1.5471],
         [1.1685, 1.0270, 1.0112,  ..., 1.0120, 1.0605, 1.2805]],

        [[1.0632, 0.9951, 0.9821,  ..., 1.0186, 1.0266, 1.2518],
         [1.2571, 1.0110, 1.0100,  ..., 0.9732, 0.9874, 1.0573],
         [1.7254, 1.2061, 1.0798,  ..., 1.0834, 1.1817, 1.5354],
         [1.3438, 1.0524, 1.0009,  ..., 1.0625, 1.1245, 1.3739]],

        [[1.1502, 1.0263, 1.0097,  ..., 1.0110, 1.0302, 1.2134],
         [1.1845, 1.0389, 1.0200,  ..., 1.0043, 1.0243, 1.0851],
         [1.7430, 1.1792, 1.0583,  ..., 1.0735, 1.1606, 1.4901],
         [1.2588, 1.0403, 1.0060,  ..., 1.1167, 1.1772, 1.4015]],

        ...,

        [[1.0115, 0.9871, 0.9944,  ..., 1.0057, 1.0473, 1.0792],
         [1.0453, 1.0188, 1.0167,  ..., 0.9922, 0.9881, 1.0192],
         [1.5685, 1.1647, 1.0673,  ..., 1.0481, 1.0894, 1.2054],
         [1.0533, 1.0980, 0.9972,  ..., 0.9405, 0.9311, 1.0248]],

        [[1.0039, 1.0114, 0.9692,  ..., 1.0232, 1.0116, 1.0603],
         [1.0201, 1.0004, 1.0142,  ..., 0.9969, 0.9949, 1.0255],
         [1.5140, 1.1822, 1.0756,  ..., 1.0156, 1.0717, 1.2108],
         [0.9944, 1.0199, 1.0299,  ..., 0.9442, 1.0108, 1.1811]],

        [[0.9736, 0.9727, 0.9738,  ..., 1.0189, 1.0098, 1.0126],
         [1.0080, 0.9990, 0.9880,  ..., 0.9677, 1.0194, 1.0420],
         [1.5378, 1.2083, 1.0986,  ..., 1.0014, 1.0555, 1.1953],
         [0.8734, 0.9773, 1.0110,  ..., 1.0126, 0.9782, 1.1674]]],
       device='cuda:0')) to Variable.

[14renus]

Hi @thaondc-mbzuai ,

CRPS, energy score and spread skill score (spskr) are all implemented in ensemble metrics: https://github.com/INRIA/geoarches/blob/main/geoarches/metrics/ensemble_metrics.py#L147

As for rank histogram errors, can you please open a separate issue giving the context of how you are running the code and the arguments you pass in?

Thanks
Renu

[14renus]

@robert-DL Hey Robert,
Do you mind sharing the changes you needed to make to the pretrained diffusion config so we can update huggingface accordingly?

Thanks
Renu

[robert-DL]

I just copied the config part that needs to be used for the pretrained ArchesWeatherGen version. The val dict has to be placed under module (similar to inference).

Regarding the torch lightning version, we either have to modify the state dicts by adding a placeholder version for lightning and reupload the file or we have to code around that by first asking if the mode equals "test" and if yes, check if the state dict contains the key. If not, load the state dict by a hand and set ckpt_path=None

val: num_members: 5 rollout_iterations: 2 # number of rollouts/multistep metrics: - {_target_: geoarches.metrics.ensemble_metrics.Era5EnsembleMetrics, lead_time_hours: 24} - {_target_: geoarches.metrics.brier_skill_score.Era5BrierSkillScore, lead_time_hours: 24}