Example of how to use a model with heiplanet-models¶

In [ ]:

import heiplanet_models as mb
from pathlib import Path
import xarray as xr
import rioxarray as rioxr
import numpy as np
from copy import deepcopy
from datetime import datetime
import json

import heiplanet_models as mb from pathlib import Path import xarray as xr import rioxarray as rioxr import numpy as np from copy import deepcopy from datetime import datetime import json

Set up the variables to run the model¶

data is assumed to live in a directory:

data/in

relative to the base directory of the package

In [ ]:

path = Path.cwd().parent.parent / "data" / "in"
path

path = Path.cwd().parent.parent / "data" / "in" path

In [ ]:

data_file = "era5_data_2016-2017_allm_2t_tp_monthly_unicoords_adjlon_celsius_mm_05deg_trim_ts20250923-065745_hssc-laptop01.nc"

data_file = "era5_data_2016-2017_allm_2t_tp_monthly_unicoords_adjlon_celsius_mm_05deg_trim_ts20250923-065745_hssc-laptop01.nc"

In [ ]:

ro_file = "R0_pip_stats.csv"

ro_file = "R0_pip_stats.csv"

make output directory¶

It helps to add the date and time to the output name to keep track of multiple runs

In [ ]:

outpath = Path.cwd().parent.parent / "data" / "out" / datetime.today().strftime("%Y-%m-%d")
outpath.mkdir(parents=True, exist_ok=True)

outpath = Path.cwd().parent.parent / "data" / "out" / datetime.today().strftime("%Y-%m-%d") outpath.mkdir(parents=True, exist_ok=True)

define the model config¶

do the following:

• read the default config

In [ ]:

with open(Path.cwd() / "config_jmodel.json", "r") as f:
    config = json.load(f)

with open(Path.cwd() / "config_jmodel.json", "r") as f: config = json.load(f)

In [ ]:

config

config

update the config with anything you might want

In [ ]:

config["graph"]["setup_modeldata"]["kwargs"]["input"] = str(path / data_file)
config["graph"]["setup_modeldata"]["kwargs"]["r0_path"] = str(path / ro_file)
config["graph"]["setup_modeldata"]["kwargs"]["output"] = str(
    outpath / f"output_jmodel_europe.nc"
)

config["graph"]["setup_modeldata"]["kwargs"]["input"] = str(path / data_file) config["graph"]["setup_modeldata"]["kwargs"]["r0_path"] = str(path / ro_file) config["graph"]["setup_modeldata"]["kwargs"]["output"] = str( outpath / f"output_jmodel_europe.nc" )

store the config file to the output folder

In [ ]:

with open(outpath / "config_europe.json", "w") as f:
    json.dump(config, f)

with open(outpath / "config_europe.json", "w") as f: json.dump(config, f)

Make a computation_graph object

In [ ]:

computation = mb.computation_graph.ComputationGraph(config)

computation = mb.computation_graph.ComputationGraph(config)

Visualize the computation graph to check the model is put together correctly. This is a good way to check whether the system is modeled in software the correct way.

In [ ]:

computation.visualize(outpath / "computation_jmodel.png")

computation.visualize(outpath / "computation_jmodel.png")

Run model¶

We now execute the build computation graph, load the created data back in to inspect it.

In [ ]:

computation.execute()

computation.execute()

In [ ]:

data = xr.open_dataset(path / data_file, chunks=None)
data

data = xr.open_dataset(path / data_file, chunks=None) data

We can also access individual nodes in the computation graph by name, e.g., for debugging

In [ ]:

input = computation.task_graph["read_input_data"].compute()

input = computation.task_graph["read_input_data"].compute()

In [ ]:

input["t2m"].sel(time="2016-08-01").plot(cmap="viridis")

input["t2m"].sel(time="2016-08-01").plot(cmap="viridis")

In [ ]:

input["t2m"].mean(dim="time").plot(cmap="viridis")

input["t2m"].mean(dim="time").plot(cmap="viridis")

Have a look at the results of the model run¶

In [ ]:

data = xr.open_dataset(outpath / "output_jmodel_europe.nc", engine="rasterio")

data = xr.open_dataset(outpath / "output_jmodel_europe.nc", engine="rasterio")

In [ ]:

data

data

In [ ]:

data["R0"].mean(dim="time").plot(cmap="viridis")

data["R0"].mean(dim="time").plot(cmap="viridis")

In [ ]:

data["R0"].sel(time="2016-08-01").plot(cmap="viridis")

data["R0"].sel(time="2016-08-01").plot(cmap="viridis")

Run model in parallel mode on a global map¶

In [ ]:

output = outpath / "output_JModel_global.nc"
run_mode = "parallelized"
grid_data_baseurl = None
nuts_level = None
resolution = None
year = None

output = outpath / "output_JModel_global.nc" run_mode = "parallelized" grid_data_baseurl = None nuts_level = None resolution = None year = None

When we pass None for the nuts_level and resolution, the global map data 'as is' will be used for the Jmodel

update the config file

In [ ]:

global_config = deepcopy(config)
global_config["graph"]["setup_modeldata"]["kwargs"]["input"] = str(path / data_file)
global_config["graph"]["setup_modeldata"]["kwargs"]["r0_path"] = str(path / ro_file)
global_config["graph"]["setup_modeldata"]["kwargs"]["output"] = str(output)
global_config["graph"]["setup_modeldata"]["kwargs"]["run_mode"] = run_mode
global_config["graph"]["setup_modeldata"]["kwargs"]["grid_data_baseurl"] = (
    grid_data_baseurl
)
global_config["graph"]["setup_modeldata"]["kwargs"]["nuts_level"] = nuts_level
global_config["graph"]["setup_modeldata"]["kwargs"]["resolution"] = resolution
global_config["graph"]["setup_modeldata"]["kwargs"]["year"] = year

global_config = deepcopy(config) global_config["graph"]["setup_modeldata"]["kwargs"]["input"] = str(path / data_file) global_config["graph"]["setup_modeldata"]["kwargs"]["r0_path"] = str(path / ro_file) global_config["graph"]["setup_modeldata"]["kwargs"]["output"] = str(output) global_config["graph"]["setup_modeldata"]["kwargs"]["run_mode"] = run_mode global_config["graph"]["setup_modeldata"]["kwargs"]["grid_data_baseurl"] = ( grid_data_baseurl ) global_config["graph"]["setup_modeldata"]["kwargs"]["nuts_level"] = nuts_level global_config["graph"]["setup_modeldata"]["kwargs"]["resolution"] = resolution global_config["graph"]["setup_modeldata"]["kwargs"]["year"] = year

In [ ]:

global_config

global_config

again, save the config with the data, build a computation graph and run it. We leave out the visualization here, because we have not changed anything in that regard. We made no new output directory, so the data will land in the same output we had before.

In [ ]:

with open(outpath / "config_global.json", "w") as f:
    json.dump(config, f)

with open(outpath / "config_global.json", "w") as f: json.dump(config, f)

In [ ]:

computation_global = mb.computation_graph.ComputationGraph(global_config)

computation_global = mb.computation_graph.ComputationGraph(global_config)

In [ ]:

computation_global.execute()

computation_global.execute()

Check out the input map

In [ ]:

input = computation_global.task_graph["read_input_data"].compute()

input = computation_global.task_graph["read_input_data"].compute()

In [ ]:

input["t2m"].mean(dim="time").plot(cmap="viridis")

input["t2m"].mean(dim="time").plot(cmap="viridis")

Look at the output data

In [ ]:

data = xr.open_dataset(output, engine="rasterio", chunks=None).compute()

data = xr.open_dataset(output, engine="rasterio", chunks=None).compute()

In [ ]:

data["R0"].mean(dim="time").plot(cmap="viridis")

data["R0"].mean(dim="time").plot(cmap="viridis")

In [ ]:

data["R0"].sel(time="2016-08-01").plot(cmap="viridis")

data["R0"].sel(time="2016-08-01").plot(cmap="viridis")

Concluding remarks¶

In this notebook, the three steps of using a heiplanet model have been shown:

• create a configuration dictionary and modify it in accordance with your aims
• instantiate a computation graph and inspect it visually if desired
• run the computation graph with the desired scheduler.

This notebook also shows some of the recommended best practices for working with heiplanet-models. In particular:

• save the configuration file alongside the data
• separation of concerns between configuraion and code
• how to use visualization to verify model composition

In the future, we aim to automate much of the configuration saving and updating.

In [ ]: