%23%20Third-party%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.25.1%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20Third-party%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Hello%20World%3A%20Training%20%60neural-lam%60%20on%20DANRA%0A%0A%20%20%20%20This%20notebook%20provides%20a%20**beginner-friendly%2C%20end-to-end%20walkthrough**%20for%20running%20a%20minimal%20model%20training%20pipeline%20in%20%60neural-lam%60%20using%20a%20small%2C%20public%20DANRA%20dataset.%20Every%20step%20uses%20the%20Python%20API%20in%20%60neural_lam.api%60.%0A%0A%20%20%20%20It%20covers%3A%0A%20%20%20%201.%20Environment%20setup%20(CPU-safe)%0A%20%20%20%202.%20Data%20preprocessing%20with%20%60mllam-data-prep%60%0A%20%20%20%203.%20Graph%20generation%20(single-level%2C%20for%20speed)%0A%20%20%20%204.%20Training%20for%201%20epoch%20on%20CPU%0A%20%20%20%205.%20Evaluation%20and%20example%20predictions%0A%20%20%20%206.%20Scaling%20tips%20for%20bigger%20runs%0A%0A%20%20%20%20---%0A%0A%20%20%20%20%23%23%20Prerequisites%20and%20Context%0A%0A%20%20%20%20%3E%20**Important%3A**%20This%20notebook%20is%20designed%20to%20be%20run%20from%20inside%20a%20**local%20clone**%20of%20%5B%60mllam%2Fneural-lam%60%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fneural-lam).%20All%20paths%20are%20relative%20to%20the%20repository%20root.%0A%20%20%20%20%3E%0A%20%20%20%20%3E%20**Paper%20Reference%3A**%20For%20an%20in-depth%20context%20on%20the%20models%20used%20here%2C%20please%20refer%20to%20the%20paper%20%5C%0A%20%20%20%20%20%20%20%20%22Building%20Machine%20Learning%20Limited%20Area%20Models%3A%20Kilometer-Scale%20Weather%20Forecasting%20in%20Realistic%20Settings%22%20%5C%0A%20%20%20%20%20%20%20%20(Adamov%20et%20al.%2C%202025)%20available%20at%20%5BarXiv%3A2504.09340%5D(https%3A%2F%2Farxiv.org%2Fabs%2F2504.09340).%0A%20%20%20%20%3E%0A%20%20%20%20%3E%20**Python%20Version%3A**%20Make%20sure%20you%20are%20using%20a%20Python%20version%20supported%20by%20the%20project%20(3.10-3.14).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20Standard%20library%0A%20%20%20%20import%20os%0A%20%20%20%20from%20pathlib%20import%20Path%0A%0A%20%20%20%20%23%20Run%20on%20CPU%20and%20keep%20W%26B%20metrics%20local%2C%20before%20torch%20and%20wandb%20start%0A%20%20%20%20os.environ%5B%22CUDA_VISIBLE_DEVICES%22%5D%20%3D%20%22%22%0A%20%20%20%20os.environ%5B%22WANDB_MODE%22%5D%20%3D%20%22offline%22%0A%0A%20%20%20%20%23%20Third-party%0A%20%20%20%20import%20fsspec%0A%20%20%20%20import%20matplotlib%0A%20%20%20%20import%20mllam_data_prep%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20xarray%20as%20xr%0A%0A%20%20%20%20%23%20Lightning%20callbacks%20plot%20from%20worker%20threads%2C%20which%20GUI%20backends%0A%20%20%20%20%23%20such%20as%20TkAgg%20do%20not%20support%0A%20%20%20%20matplotlib.use(%22Agg%22)%0A%20%20%20%20%23%20Third-party%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20%23%20First-party%0A%20%20%20%20import%20neural_lam%0A%20%20%20%20from%20neural_lam%20import%20api%2C%20utils%0A%20%20%20%20from%20neural_lam.config%20import%20load_config_and_datastore%0A%20%20%20%20from%20neural_lam.plot_graph%20import%20plot_graph%0A%0A%20%20%20%20%23%20All%20paths%20below%20are%20relative%20to%20the%20repository%20root%0A%20%20%20%20os.chdir(mo.notebook_dir().parents%5B1%5D)%0A%20%20%20%20example_dir%20%3D%20Path(%22tests%2Fdatastore_examples%2Fmdp%2Fdanra_100m_winds%22)%0A%20%20%20%20config_path%20%3D%20str(example_dir%20%2F%20%22config.yaml%22)%0A%20%20%20%20print(%22Working%20directory%3A%22%2C%20Path.cwd())%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20api%2C%0A%20%20%20%20%20%20%20%20config_path%2C%0A%20%20%20%20%20%20%20%20example_dir%2C%0A%20%20%20%20%20%20%20%20fsspec%2C%0A%20%20%20%20%20%20%20%20load_config_and_datastore%2C%0A%20%20%20%20%20%20%20%20mllam_data_prep%2C%0A%20%20%20%20%20%20%20%20neural_lam%2C%0A%20%20%20%20%20%20%20%20np%2C%0A%20%20%20%20%20%20%20%20plot_graph%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20utils%2C%0A%20%20%20%20%20%20%20%20xr%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%201.%20Environment%20Setup%0A%0A%20%20%20%20Set%20up%20the%20environment%20with%20%5B%60uv%60%5D(https%3A%2F%2Fdocs.astral.sh%2Fuv%2F)%20from%20the%20repository%20root%20(see%20%22Installing%20Neural-LAM%22%20in%20the%20%5BREADME%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fneural-lam%2Fblob%2Fmain%2FREADME.md))%20and%20open%20this%20notebook%20with%20marimo%3A%0A%0A%20%20%20%20%60%60%60bash%0A%20%20%20%20uv%20sync%20--extra%20cpu%20--extra%20docs%20--group%20dev%0A%20%20%20%20source%20.venv%2Fbin%2Factivate%0A%20%20%20%20marimo%20edit%20docs%2Fnotebooks%2Fhello_world_danra.py%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20This%20creates%20a%20%60.venv%60%20with%20the%20CPU%20build%20of%20PyTorch%2C%20marimo%20(from%20the%20%60docs%60%20extra)%20and%20the%20development%20dependencies.%20For%20GPU%20runs%20swap%20%60--extra%20cpu%60%20for%20%60--extra%20gpu%60%20(CUDA%2013.0)%20or%20%60--extra%20gpu-cu128%60%20(CUDA%2012.8).%0A%0A%20%20%20%20To%20run%20the%20whole%20tutorial%20non-interactively%20use%20%60python%20docs%2Fnotebooks%2Fhello_world_danra.py%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(neural_lam)%3A%0A%20%20%20%20%23%20Verify%20the%20install%0A%20%20%20%20print(%22neural-lam%20version%3A%22%2C%20neural_lam.__version__)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%202.%20Data%20Configuration%20%26%20Preprocessing%0A%0A%20%20%20%20We%20use%20the%20**DANRA%20100m%20winds%20example**%20already%20in%20this%20repo%20at%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20tests%2Fdatastore_examples%2Fmdp%2Fdanra_100m_winds%2F%0A%20%20%20%20%20%20%E2%94%9C%E2%94%80%E2%94%80%20config.yaml%20%20%20%20%20%20%20%20%20%20%E2%86%90%20neural-lam%20configuration%0A%20%20%20%20%20%20%E2%94%94%E2%94%80%E2%94%80%20danra.datastore.yaml%20%E2%86%90%20mllam-data-prep%20datastore%20config%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20This%20example%20uses%20a%20**cropped%20DANRA%20dataset**%20(~100%C3%97100%20grid%20points%2C%20~10%20days%20in%20April%202022)%20served%20from%20a%20public%20ECMWF%20object%20store%20%E2%80%94%20no%20local%20data%20download%20is%20required.%0A%0A%20%20%20%20%60mllam_data_prep.create_dataset_zarr%60%20reads%20the%20datastore%20config%2C%20fetches%20the%20data%2C%20and%20writes%20a%20processed%20%60.zarr%60%20archive%20next%20to%20it.%0A%0A%20%20%20%20%3E%20**Note%3A**%20This%20will%20download%20and%20process%20approximately%2060%E2%80%93120%20MB%20of%20data.%20It%20may%20take%20a%20few%20minutes%20on%20the%20first%20run.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(example_dir%2C%20fsspec%2C%20mllam_data_prep%2C%20xr)%3A%0A%20%20%20%20_zarr_path%20%3D%20example_dir%20%2F%20%22danra.datastore.zarr%22%0A%20%20%20%20if%20not%20_zarr_path.exists()%3A%0A%20%20%20%20%20%20%20%20%23%20zarr%203%20reads%20the%20remote%20inputs%20through%20an%20async%20https%20filesystem%0A%20%20%20%20%20%20%20%20fsspec.config.conf%5B%22https%22%5D%20%3D%20%7B%22asynchronous%22%3A%20True%7D%0A%20%20%20%20%20%20%20%20mllam_data_prep.create_dataset_zarr(%0A%20%20%20%20%20%20%20%20%20%20%20%20fp_config%3Dexample_dir%20%2F%20%22danra.datastore.yaml%22%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20print(f%22Dataset%20already%20exists%20at%20%7B_zarr_path%7D%2C%20skipping%20preparation.%22)%0A%20%20%20%20ds%20%3D%20xr.open_zarr(_zarr_path)%0A%20%20%20%20ds%0A%20%20%20%20return%20(ds%2C)%0A%0A%0A%40app.cell%0Adef%20_(ds%2C%20np%2C%20plt)%3A%0A%20%20%20%20_da%20%3D%20ds%5B%22state%22%5D.isel(time%3D0%2C%20state_feature%3D0)%0A%20%20%20%20_x_unique%20%3D%20np.unique(ds%5B%22x%22%5D.values)%0A%20%20%20%20_y_unique%20%3D%20np.unique(ds%5B%22y%22%5D.values)%0A%20%20%20%20_grid%20%3D%20np.full((len(_y_unique)%2C%20len(_x_unique))%2C%20np.nan)%0A%20%20%20%20_grid%5B%0A%20%20%20%20%20%20%20%20np.searchsorted(_y_unique%2C%20ds%5B%22y%22%5D.values)%2C%0A%20%20%20%20%20%20%20%20np.searchsorted(_x_unique%2C%20ds%5B%22x%22%5D.values)%2C%0A%20%20%20%20%5D%20%3D%20_da.values%0A%20%20%20%20fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(8%2C%206))%0A%20%20%20%20_pcm%20%3D%20_ax.pcolormesh(%0A%20%20%20%20%20%20%20%20_x_unique%2C%20_y_unique%2C%20_grid%2C%20cmap%3D%22RdBu_r%22%2C%20shading%3D%22auto%22%0A%20%20%20%20)%0A%20%20%20%20_ax.set_xlabel(%22x%20(m)%22)%0A%20%20%20%20_ax.set_ylabel(%22y%20(m)%22)%0A%20%20%20%20_ax.set_title(f%22%7B_da%5B'state_feature'%5D.item()%7D%20at%20t%3D0%22)%0A%20%20%20%20fig.colorbar(_pcm%2C%20ax%3D_ax)%0A%20%20%20%20fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%203.%20Graph%20Generation%0A%0A%20%20%20%20%60neural-lam%60%20uses%20a%20**graph**%20to%20define%20the%20message-passing%20structure%20of%20the%20GNN.%20%60api.create_graph%60%20builds%20it%20with%20%5B%60weather-model-graphs%60%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fweather-model-graphs)%20from%20one%20of%20these%20archetypes%3A%0A%0A%20%20%20%20%7C%20Archetype%20%20%20%20%20%20%7C%20Mesh%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7C%20Model%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7C%0A%20%20%20%20%7C----------------%7C----------------------------%7C------------------------%7C%0A%20%20%20%20%7C%20%60keisler%60%20%20%20%20%20%20%7C%20Single%20level%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%60graph_lam%60%20(this%20notebook)%20%7C%0A%20%20%20%20%7C%20%60graphcast%60%20%20%20%20%7C%20Multi-scale%2C%20flattened%20%20%20%20%20%7C%20%60graph_lam%60%20%20%20%20%20%20%20%20%20%20%20%20%7C%0A%20%20%20%20%7C%20%60hierarchical%60%20%7C%20Hierarchy%20of%20mesh%20levels%20%20%20%7C%20%60hi_lam%60%2C%20%60hi_lam_parallel%60%20%7C%0A%0A%20%20%20%20For%20this%20Hello%20World%20example%20we%20use%20the%20**single-level%20%60keisler%60%20graph**%20%E2%80%94%20the%20lightest%20option%2C%20ideal%20for%20CPU.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(api%2C%20config_path%2C%20example_dir)%3A%0A%20%20%20%20graph_name%20%3D%20%221level%22%0A%20%20%20%20api.create_graph(%0A%20%20%20%20%20%20%20%20config_path%3Dconfig_path%2C%20name%3Dgraph_name%2C%20archetype%3D%22keisler%22%0A%20%20%20%20)%0A%20%20%20%20graph_dir%20%3D%20example_dir%20%2F%20%22graph%22%20%2F%20graph_name%0A%20%20%20%20for%20_f%20in%20sorted(graph_dir.glob(%22*.pt%22))%3A%0A%20%20%20%20%20%20%20%20print(f%22%20%20%7B_f.name%3A40s%7D%20%7B_f.stat().st_size%20%2F%201024%3A.1f%7D%20KB%22)%0A%20%20%20%20return%20graph_dir%2C%20graph_name%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20config_path%2C%0A%20%20%20%20graph_dir%2C%0A%20%20%20%20load_config_and_datastore%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plot_graph%2C%0A%20%20%20%20utils%2C%0A)%3A%0A%20%20%20%20%23%20Interactive%203D%20view%20of%20the%20grid%20and%20mesh%20nodes%20and%20their%20connections%0A%20%20%20%20_%2C%20_datastore%2C%20_%20%3D%20load_config_and_datastore(config_path%3Dconfig_path)%0A%20%20%20%20_xy%20%3D%20_datastore.get_xy(%22state%22%2C%20stacked%3DTrue)%0A%20%20%20%20_grid_xy_max_span%20%3D%20np.max(np.abs(_xy))%0A%20%20%20%20_hierarchical%2C%20_graph_ldict%20%3D%20utils.load_graph(%0A%20%20%20%20%20%20%20%20graph_dir_path%3Dstr(graph_dir)%2C%0A%20%20%20%20%20%20%20%20mesh_node_features_scaling%3D_grid_xy_max_span%2C%0A%20%20%20%20)%0A%20%20%20%20plot_graph(%0A%20%20%20%20%20%20%20%20grid_pos%3D_xy%20%2F%20_grid_xy_max_span%2C%0A%20%20%20%20%20%20%20%20hierarchical%3D_hierarchical%2C%0A%20%20%20%20%20%20%20%20graph_ldict%3D_graph_ldict%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%204.%20Training%20on%20CPU%0A%0A%20%20%20%20%60api.train%60%20takes%20the%20same%20options%20as%20%60python%20-m%20neural_lam.train_model%60%2C%20as%20keyword%20arguments.%20We%20train%20the%20%60graph_lam%60%20model%20for%20**1%20epoch**%20on%20the%20single-level%20graph%3A%0A%0A%20%20%20%20%7C%20Argument%20%7C%20Value%20%7C%20Why%20%7C%0A%20%20%20%20%7C------%7C-------%7C-----%7C%0A%20%20%20%20%7C%20%60model%60%20%7C%20%60%22graph_lam%22%60%20%7C%20Simplest%20model%20%E2%80%94%20compatible%20with%20non-hierarchical%20graphs%20%7C%0A%20%20%20%20%7C%20%60graph%60%20%7C%20%60%221level%22%60%20%7C%20Must%20match%20the%20graph%20name%20used%20in%20the%20previous%20step%20%7C%0A%20%20%20%20%7C%20%60epochs%60%20%7C%20%601%60%20%7C%20Minimal%20run%20%E2%80%94%20just%20to%20verify%20the%20pipeline%20works%20end-to-end%20%7C%0A%20%20%20%20%7C%20%60processor_layers%60%20%7C%20%602%60%20%7C%20Reduced%20from%20default%20(4)%20for%20CPU%20%7C%0A%20%20%20%20%7C%20%60ar_steps_train%60%20%7C%20%601%60%20%7C%20Unroll%201%20time-step%20during%20training%20%E2%80%94%20reduces%20memory%20and%20compute%20%7C%0A%20%20%20%20%7C%20%60ar_steps_eval%60%20%7C%20%601%60%20%7C%20Also%201%20step%20for%20the%20val%20pass%20within%20this%20demo%20run%20%7C%0A%0A%20%20%20%20It%20returns%20a%20%60Run%60%20handle%20with%20the%20run%20directory%20and%20the%20checkpoint%20path.%0A%0A%20%20%20%20**About%20logging%3A**%20By%20default%2C%20%60neural-lam%60%20logs%20via%20%5BWeights%20%26%20Biases%5D(https%3A%2F%2Fwandb.ai).%20To%20suppress%20W%26B%20upload%20during%20this%20demo%2C%20the%20import%20cell%20sets%20%60WANDB_MODE%3Doffline%60%20so%20metrics%20are%20saved%20locally%20without%20requiring%20a%20W%26B%20account%20or%20login.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(api%2C%20config_path%2C%20graph_name)%3A%0A%20%20%20%20train_run%20%3D%20api.train(%0A%20%20%20%20%20%20%20%20config_path%3Dconfig_path%2C%0A%20%20%20%20%20%20%20%20model%3D%22graph_lam%22%2C%0A%20%20%20%20%20%20%20%20graph%3Dgraph_name%2C%0A%20%20%20%20%20%20%20%20epochs%3D1%2C%0A%20%20%20%20%20%20%20%20processor_layers%3D2%2C%0A%20%20%20%20%20%20%20%20ar_steps_train%3D1%2C%0A%20%20%20%20%20%20%20%20ar_steps_eval%3D1%2C%0A%20%20%20%20%20%20%20%20num_workers%3D0%2C%0A%20%20%20%20%20%20%20%20val_steps_to_log%3D%5B1%5D%2C%0A%20%20%20%20)%0A%20%20%20%20print(%22Checkpoint%3A%22%2C%20train_run.checkpoint_path)%0A%20%20%20%20return%20(train_run%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%205.%20Evaluation%20%26%20Visualization%0A%0A%20%20%20%20%60api.evaluate%60%20loads%20the%20checkpoint%20from%20training%20and%20runs%20the%20test%20set.%20The%20returned%20%60Run%60%20points%20to%20the%20plots%20it%20writes%3A%20the%20test%20RMSE%20per%20lead%20time%20(%60rmse_plot%60)%20and%2C%20with%20%60create_gif%3DTrue%60%2C%20a%20GIF%20per%20variable%20for%20each%20of%20the%20%60n_example_pred%60%20example%20forecasts%20(%60example_plots(i)%60).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(api%2C%20config_path%2C%20graph_name%2C%20train_run)%3A%0A%20%20%20%20eval_run%20%3D%20api.evaluate(%0A%20%20%20%20%20%20%20%20config_path%3Dconfig_path%2C%0A%20%20%20%20%20%20%20%20load%3Dstr(train_run.checkpoint_path)%2C%0A%20%20%20%20%20%20%20%20model%3D%22graph_lam%22%2C%0A%20%20%20%20%20%20%20%20graph%3Dgraph_name%2C%0A%20%20%20%20%20%20%20%20processor_layers%3D2%2C%0A%20%20%20%20%20%20%20%20ar_steps_eval%3D4%2C%0A%20%20%20%20%20%20%20%20n_example_pred%3D1%2C%0A%20%20%20%20%20%20%20%20create_gif%3DTrue%2C%0A%20%20%20%20%20%20%20%20num_workers%3D0%2C%0A%20%20%20%20%20%20%20%20val_steps_to_log%3D%5B1%5D%2C%0A%20%20%20%20)%0A%20%20%20%20return%20(eval_run%2C)%0A%0A%0A%40app.cell%0Adef%20_(eval_run%2C%20mo)%3A%0A%20%20%20%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.pdf(src%3Deval_run.rmse_plot%2C%20height%3D%2250vh%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20*%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mo.image(src%3D_gif)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20_gif%20in%20sorted(eval_run.example_plots(1).glob(%22*.gif%22))%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%206.%20Scaling%20Tips%20%26%20Next%20Steps%0A%0A%20%20%20%20%23%23%23%20%F0%9F%94%81%20Use%20a%20larger%20%2F%20hierarchical%20graph%0A%0A%20%20%20%20For%20a%20production-quality%20run%2C%20switch%20to%20the%20hierarchical%20graph%20and%20the%20Hi-LAM%20model%3A%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20api.create_graph(%0A%20%20%20%20%20%20%20%20config_path%3D%22%3Cyour_config.yaml%3E%22%2C%0A%20%20%20%20%20%20%20%20name%3D%22hierarchical%22%2C%0A%20%20%20%20%20%20%20%20archetype%3D%22hierarchical%22%2C%0A%20%20%20%20)%0A%20%20%20%20api.train(%0A%20%20%20%20%20%20%20%20config_path%3D%22%3Cyour_config.yaml%3E%22%2C%0A%20%20%20%20%20%20%20%20model%3D%22hi_lam%22%2C%0A%20%20%20%20%20%20%20%20graph%3D%22hierarchical%22%2C%0A%20%20%20%20%20%20%20%20epochs%3D200%2C%0A%20%20%20%20)%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20same%20steps%20are%20available%20on%20the%20command%20line%20as%20%60python%20-m%20neural_lam.create_graph_with_wmg%60%20and%20%60python%20-m%20neural_lam.train_model%60.%0A%0A%20%20%20%20%23%23%23%20%E2%9A%A1%20Enable%20GPU%20training%0A%0A%20%20%20%20Re-create%20the%20environment%20with%20a%20CUDA%20build%20of%20PyTorch%20and%20remove%20the%20%60CUDA_VISIBLE_DEVICES%60%20line%20from%20the%20import%20cell%2C%20so%20that%20%60devices%60%20(default%20%60auto%60)%20picks%20up%20your%20GPU%3A%0A%20%20%20%20%60%60%60bash%0A%20%20%20%20uv%20sync%20--extra%20gpu%20%20%20%20%20%20%20%20%23%20CUDA%2013.0%0A%20%20%20%20%23%20or%3A%20uv%20sync%20--extra%20gpu-cu128%20%20%20%23%20CUDA%2012.8%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20%23%23%23%20%F0%9F%93%A6%20Use%20larger%20%2F%20full%20DANRA%20data%0A%0A%20%20%20%20Modify%20%60danra.datastore.yaml%60%20to%20extend%20the%20%60coord_ranges.time%60%20window%2C%20or%20point%20%60inputs%5B*%5D.path%60%20at%20a%20larger%20dataset.%20See%20the%20%5Bmllam-data-prep%20README%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fmllam-data-prep)%20for%20full%20configuration%20options%2C%20including%20parallel%20processing%20of%20large%20datasets%20with%20dask.%0A%0A%20%20%20%20%23%23%23%20%F0%9F%93%8A%20Enable%20cloud%20logging%20with%20W%26B%0A%0A%20%20%20%20Remove%20the%20%60WANDB_MODE%60%20line%20in%20the%20import%20cell%20(or%20set%20it%20to%20%60online%60)%20and%20optionally%20pass%20%60logger%3D%22wandb%22%2C%20logger_project%3D%22%3Cyour-project-name%3E%22%60.%0A%0A%20%20%20%20%23%23%23%20%F0%9F%94%97%20Further%20reading%0A%0A%20%20%20%20-%20%5Bneural-lam%20README%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fneural-lam)%0A%20%20%20%20-%20%5Bmllam-data-prep%5D(https%3A%2F%2Fgithub.com%2Fmllam%2Fmllam-data-prep)%0A%20%20%20%20-%20%5BGraph-based%20Neural%20Weather%20Prediction%20for%20Limited%20Area%20Modeling%20(NeurIPS%202023%20CCAI%20workshop)%5D(https%3A%2F%2Farxiv.org%2Fabs%2F2309.17370)%0A%20%20%20%20-%20%5BProbabilistic%20Weather%20Forecasting%20with%20Hierarchical%20GNNs%20(NeurIPS%202024)%5D(https%3A%2F%2Farxiv.org%2Fabs%2F2406.04759)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A
fc4b4229f249f62ec50204f2deb050e5