google-deepmind / graphcast

Apache License 2.0
4.36k stars 537 forks source link

evaluation on selected region #46

Closed akhtarvision closed 5 months ago

akhtarvision commented 6 months ago

Hi. Thanks for the great work! One quick question is how one can evaluate the graphcast on the local regions e.g., North America, by adding some lat/lon points and measuring the error on that region? @tewalds

tewalds commented 6 months ago

GraphCast always does whole world prediction. You're welcome to run evaluations you'd like on any region you'd like though. We haven't open sourced our evaluation code, but weatherbench may be helpful there.

Javier-Jimenez99 commented 5 months ago

I tried to call the model with an input_dataset which does not have all lat-lon coordinates. It is predicted over the whole world, can I assume that, at least the coordinates I used as input, are correctly mapped to the output?

tewalds commented 5 months ago

I'm not entirely sure what will happen if you give it a smaller input than expected, but I would not expect useful results. The trunk of the model is a global multimesh. It's possible the encoder/decoder will only connect to a subset of the global mesh, which would lead to the nodes not connected to the inputs would likely be initialized to 0s. Given the multimesh has long range connections, the receptive field of any given output is the entire globe, so I'd expect all outputs to be at least minorly affected by all inputs, which would be 0s instead of something sensible. It's possible, if unlikely, that the model learned to ignore the long range connections leading to a non-global receptive field, but I am certain the receptive field for several variables is large (several thousand of km). If you really want a model that works on patches you'd have to change your internal representation to one without a large receptive field (this ablation is shown in 7.3.1 in the supplementary materials of the science paper while still remaining a global model), and your inputs will need to be as big as the receptive field for your expected output area. This type of model would save memory but would be more computationally expensive, especially if you want to do long time horizon evaluation since the receptive field would become global over a day or two anyway.

MiodragRancic-NOAA commented 5 months ago

What exactly component of Graphcast would need to be modified in order to use it for regional prediction? Thanks!

tewalds commented 5 months ago

Given your region is part of the earth, just generate whole world predictions then crop the output?

If you really wanted to only generate predictions for a subregion (eg to predict at higher resolution), you'll need to modify the model then retrain. The model currently has an encoder in the form of a graphnet from a whole world grid to a icosahedral mesh, then a 16 layer graphnet over a multimesh, then a decoder from the mesh back to the whole world grid. You'd need to replace the icosahedral multimesh with something useful for your patch area. You'll also need to deal with the boundary condition problems, presumably by having a larger input region than output region, though that will make multi-step predictions harder, and updating the loss to only apply to the region of interest. You'd then need to retrain it from scratch.

MiodragRancic-NOAA commented 5 months ago

Thank very much you for your reply!

If there is no way at the moment to use a true regional domain, I am considering to apply a Schmidt transformed cubed sphere, with a high resolution on one cubed face, and a very low elsewhere. We would run the model only over the high resolution domain using previously prescribed b.c., and we would generate values over lower portion using some sort of interpolation/smoothing toward constant values in the interior domain far from the domain of interest. That would be a sort of regional model forced by the prescribed boundary conditions but using formally inputs from the whole sphere. That seems a bit like your suggestion about replacing icosahedral multimesh. And, yes, we would need to retrain the whole system. I think that NN would not care that we use dynamical equations only in one part of the sphere, and prescribe values elsewhere for training.

The main issue in this scenario is how to replace icosahedral multimesh with Schmidt transformed cubed sphere. Hopefully, we will be able to figure out that.

We are actually considering to apply the GraphCast for the regional downscaling from 3 km to 2.5 km and further down for one of our projects. We will have a series of 2.5 km runs for training. It would be nice if we can prove that GraphCast is applicable for the regional, high resolution simulations.

I cc this email to my collaborators on that project,

Thank you once again for your response.

Miodrag Rancic Lynker at NOAA/NWS/NCEP/EMC 5830 University Research Ct, Rm 2786 College Park, MD 20740 301-683-3732

On Wed, Jan 17, 2024 at 11:43 AM Timo Ewalds @.***> wrote:

Given your region is part of the earth, just generate whole world predictions then crop the output?

If you really wanted to only generate predictions for a subregion (eg to predict at higher resolution), you'll need to modify the model then retrain. The model currently has an encoder in the form of a graphnet from a whole world grid to a icosahedral mesh, then a 16 layer graphnet over a multimesh, then a decoder from the mesh back to the whole world grid. You'd need to replace the icosahedral multimesh with something useful for your patch area. You'll also need to deal with the boundary condition problems, presumably by having a larger input region than output region, though that will make multi-step predictions harder, and updating the loss to only apply to the region of interest. You'd then need to retrain it from scratch.

— Reply to this email directly, view it on GitHub https://github.com/google-deepmind/graphcast/issues/46#issuecomment-1896193940, or unsubscribe https://github.com/notifications/unsubscribe-auth/AJHP7FVFN5DN7SRH6UL223TYO75R5AVCNFSM6AAAAABBDGUPQKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJWGE4TGOJUGA . You are receiving this because you commented.Message ID: @.***>

tewalds commented 5 months ago

If you're willing to put in the effort, you could modify the encoder, decoder and internal mesh to not be uniform. Right now the main computation is done on a uniform 2^6 mesh, but you could do 2^5 in some areas and 2^7 in others. Similarly the encoder/decoder input from a uniform 0.25deg grid, but it could be modified to accept arbitrary points as input, so you don't need to use a cube sphere if you'd prefer to use something that matches your actual geometry better. It may be a bit of work to get that working, but this is definitely the power of graphnets over CNNs or standard image transformers that assume a regular grid.

A 2.5-3km input is much (10^2x) higher than what we currently use, so you're likely to run into memory constraints unless you also put effort into reducing the memory usage. A single 0.25deg grid input with 37 levels and 6 variables per level is almost 1GB (at float32), so 100x the resolution will exhaust any single device's memory, never mind that we take two and need memory for intermediate values, especially at training. Of course only representing 10% of the globe at that resolution would help, but still, there will be work in dealing with this. We already had to use several tricks to run at 0.25deg within 32GB, especially when doing multi-step training. A100/H100s with 80GB will help a bit, but will still struggle. Sharding a graphnet across multiple devices is likely possible, but is likely non-trivial to do performantly. Of course higher resolution will also take longer to train due to additional computation.

MiodragRancic-NOAA commented 5 months ago

Well, it is 2.5 km grid only over North American domain! Thus, assuming that we keep the rest of the globe at a significantly lower resolution, that might work in terms of memory consumption.

Yes, your suggestion to put effort in modifying all three, encoder, decoder and internal mesh is probably the best way to go. However, that would mean getting deeper into GraphCast core. But, essentially, if we do not need to assume standard geographical long-lat grid, that could be additional help.

Thank you for your help!

Miodrag Rancic Lynker at NOAA/NWS/NCEP/EMC 5830 University Research Ct, . Rm 2786 College Park, MD 20740 301-683-3732

On Wed, Jan 17, 2024 at 3:05 PM Timo Ewalds @.***> wrote:

If you're willing to put in the effort, you could modify the encoder, decoder and internal mesh to not be uniform. Right now the main computation is done on a uniform 2^6 mesh, but you could do 2^5 in some areas and 2^7 in others. Similarly the encoder/decoder input from a uniform 0.25deg grid, but it could be modified to accept arbitrary points as input, so you don't need to use a cube sphere if you'd prefer to use something that matches your actual geometry better. It may be a bit of work to get that working, but this is definitely the power of graphnets over CNNs or standard image transformers that assume a regular grid.

A 2.5-3km input is much (10^2x) higher than what we currently use, so you're likely to run into memory constraints unless you also put effort into reducing the memory usage. A single 0.25deg grid input with 37 levels and 6 variables per level is almost 1GB (at float32), so 100x the resolution will exhaust any single device's memory, never mind that we take two and need memory for intermediate values, especially at training. Of course only representing 10% of the globe at that resolution would help, but still, there will be work in dealing with this. We already had to use several tricks to run at 0.25deg within 32GB, especially when doing multi-step training. A100/H100s with 80GB will help a bit, but will still struggle. Sharding a graphnet across multiple devices is likely possible, but is likely non-trivial to do performantly. Of course higher resolution will also take longer to train due to additional computation.

— Reply to this email directly, view it on GitHub https://github.com/google-deepmind/graphcast/issues/46#issuecomment-1896603566, or unsubscribe https://github.com/notifications/unsubscribe-auth/AJHP7FUFQM2BPFLZ2BVM4HTYPAVI7AVCNFSM6AAAAABBDGUPQKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJWGYYDGNJWGY . You are receiving this because you commented.Message ID: @.***>

mjwillson commented 5 months ago

Thanks -- I think we can close this now. To summarise, yes you could potentially adapt GraphCast to be a kind of combination of a high-res local area model + lower-res global model if you adapted the encoder/decoder/internal mesh structures to add local refinements over your local area of interest. It would be a substantial amount of work though and not in scope for this github project.