When used in forward_multistep (forecast.py) and in sample_rollout (diffusion.py), we assume lead_time_hours is the timestep in between predictions.

Also remove accidentally added lead_time_hours param in forecast.py in commit edfa13

add constant masks for 181x360 data (1deg x 1deg)

(reverting back to old code for now)

Robustify tests:
- check length
- check when lead_time_hours != timedelta_hours

- initialize noise_scheduler outside of jitted function
- fix bug in defining self.state_scaler for state_normalization = delta
- replace non-jittable iteration over timesteps with torch.searchsorted

… is not jittable

…gridpoints)

* Added flexibility for more variables. Added normalization statistics from graphcast as json files, refactored metrics, datasets and tests as well as lightning modules. The latter have to account for nans in the data when including sea-ice-fraction and others
Resolved conflicts during rebase

* Renamed stats to norm, changed variable output for label wrapper, applied changes to configs
Fixed merge conflicts after rebase dev_diffusion

* First changes in stats configs, norm_utils and diffusion module to adapt to more variables
Fixed merge conflicts when rebasing dev_diffusion

* Commiting changes in forecast and norm_utils.py

* Some fixes in the dataloader to use the right time_dim_name

* added small fixes to the code, convert_to_xarray converts with default level and time dim name

* Revert to old era5 convert to xarray

* Updated feature branch with minor changes in the dataloader, added normalization of the forcings to the dataloader
Fixed merge conflicts

* Added description of stats path to dataset

* Fixed problem with debug statement
Fixed merge conflicts

* Added forcings to default configs

* Added fillna in forcings loader

* Using interpolate_na along lat instead of fill

* Save commit
Fixed conflicts

* Restored state before changes to fillna
Fixed merge conflicts

* Make zdim explicit

* Added statistics to diff model
Fixed merge conflicts

* Adapting inference code for predictions

* Changed forecast and dataloader nan behaviour
Fixed merge conflicts

* Reformatted
Fixed merge conflicts

* Masking mechanism fixed in loss calculation, added interpolation for targets and nans as separate arguments to dataloaders, added explicit norm file argument to stats.yaml and NormalizatinModule
Fixed merge conflicts

* Refactored style
Resolved merge conflicts

* Changed position and behaviour of model averaging

* Added test

* Format fixed
Resolved mc

* Fixed formatting

* Fixed ruff after fixing import error in test

* Fixed test problems and formatting

* Fixed ruff-format

* Fixed ruff format

* Fixed lint problem

* Formatting fixes

* Fixed problems

* Format fixes

* Changed naming of prediction difference stats, refactored loss coeffs and state scaler into separate method, removed unnecessary args

* Fixed test error due to missing level dimension indexer

* Moved power statement for scaler and changed mul and div statement in diffusion normalization parts

* Document normalization params

Added documentation for:
residual statistics path, latitude, and WeatherBench latitude coefficients.

---------

Co-authored-by: Renu Singh <[email protected]>