mosartwmpy tutorial

This tutorial will demonstrate the basic use of mosartwmpy in two scenarios: * First, in standalone mode where all model inputs are provided from files * Second, a contrived example of running the model coupled with runoff input from another model

The use of restart files will also be demonstrated.

from datetime import date, datetime
import numpy as np

from mosartwmpy import Model
from mosartwmpy.plotting.plot import plot_reservoir, plot_variable
from mosartwmpy.utilities.download_data import download_data

/Users/thur961/im3/mosartwmpy/venv/lib/python3.9/site-packages/pandas/compat/__init__.py:124: UserWarning: Could not import the lzma module. Your installed Python is incomplete. Attempting to use lzma compression will result in a RuntimeError.
  warnings.warn(msg)

Download and unpack the tutorial dataset, which covers May of 1981 (this may take a few minutes):

download_data('tutorial')

Downloading example data from https://zenodo.org/record/5597952/files/mosartwmpy-tutorial-1981-05.zip?download=1
https://zenodo.org/record/5597952/files/mosartwmpy-tutorial-1981-05.zip?download=1: 100%|██████████████████████████████████| 16.1M/16.1M [00:05<00:00, 3.22MB/s]
Unzipped: ./input/
Unzipped: ./input/domains/
Unzipped: ./input/domains/land.nc
Unzipped: ./input/domains/mosart.nc
Unzipped: ./input/reservoirs/
Unzipped: ./input/reservoirs/mean_monthly_reservoir_flow.parquet
Unzipped: ./input/reservoirs/mean_monthly_reservoir_demand.parquet
Unzipped: ./input/reservoirs/reservoirs.nc
Unzipped: ./input/reservoirs/dependency_database.parquet
Unzipped: ./input/runoff/
Unzipped: ./input/runoff/runoff_1981_05.nc
Unzipped: ./input/demand/
Unzipped: ./input/demand/demand_1981_05.nc
Download and install complete.

Initialize the model using the provided config.yaml, which has the input paths properly specified for this notebook:

mosart_wm = Model()
mosart_wm.initialize('./config.yaml')

Initalizing model...
Done.

The model is now setup to run in standalone mode, with runoff provided from file.

Let’s run for a couple weeks in this mode (this may take a couple minutes):

mosart_wm.config["simulation.end_date"] = date(1981, 5, 14)
mosart_wm.update_until(mosart_wm.get_end_time())

Beginning simulation for 1981-05-01 through 1981-05-14...
Running mosartwmpy
Simulation completed in 1 minutes and 25 seconds.

Let’s take a look at how the river channels have formed over this week.

mosart_wm.plot_variable('surface_water_amount', log_scale=True)

We can see the dominant river channels beginning to form! Since the initial conditions have no surface water, it can take a few weeks of model time to get to a good baseline, depending on the amount of rainfall.

Now let’s pretend to run the model in coupled mode, as if runoff were being provided by another model such as CLM.

First, we’ll implement dummy functions that create random runoff to take the place of the coupled model:

def get_surface_runoff():
    # provide the surface runoff for each grid cell in mm/s
    data = np.random.normal(0.001, 0.001, mosart_wm.get_grid_shape())
    return np.where(
        data < 0,
        0,
        data
    )

def get_subsurface_runoff():
    # provide the subsurface runoff for each grid cell in mm/s
    data = np.random.normal(0.00015, 0.00015, mosart_wm.get_grid_shape())
    return np.where(
        data < 0,
        0,
        data
    )

In coupled mode, we’ll need to run the model in chunks of time that correspond to the coupling period and update the coupled variables between chunks. For this example, let’s say the coupling period is one day – so we’ll update the runoff variables after each day of model time. The same technique could be extended to run the coupled model between each day rather than generating synthetic data.

# first, disable reading runoff from file
mosart_wm.config["runoff.read_from_file"] = False

# set the end date for two weeks out
mosart_wm.config["simulation.end_date"] = date(1981, 5, 28)

# for each day, set the runoff and simulate for one day
# note that the model expects the input flattened as a 1-d array
while mosart_wm.get_current_time() < mosart_wm.get_end_time():
    # set the surface runoff from a coupled model
    mosart_wm.set_value('surface_runoff_flux', get_surface_runoff().flatten())
    # set the subsurface runoff from a coupled model
    mosart_wm.set_value('subsurface_runoff_flux', get_subsurface_runoff().flatten())
    # simulate for one day
    day = date.fromtimestamp(mosart_wm.get_current_time())
    while day == date.fromtimestamp(mosart_wm.get_current_time()):
        mosart_wm.update()

Let’s see how the random rainfall affected the river channels…

mosart_wm.plot_variable('surface_water_amount', log_scale=True)

Oops, looks like we flooded Canada and Mexico. The river network isn’t very well defined there (the tutorial only includes CONUS data for the most part) so it makes sense. Depending on the random data generation, you’re probably seeing a lot more water across the CONUS too.

And how about the reservoir storage?

mosart_wm.plot_variable('reservoir_water_amount', log_scale=True)

Let’s try restarting the simulation after the original two weeks in May, and run the second two weeks with the runoff from the file. By default, restart files are produced at the end of each model year and also at the end of each simulation. The provided config_with_restart.yaml file will use the restart file from 1981-05-15. You can also use a restart file like an initial conditions file by setting the start date to whatever date you choose. But in this case we’ll restart where we left off.

mosart_wm = Model()
mosart_wm.initialize('./config_with_restart.yaml')
mosart_wm.config["simulation.start_date"] = date(1981, 5, 15)
mosart_wm.config["simulation.end_date"] = date(1981, 5, 28)
mosart_wm.update_until(mosart_wm.get_end_time())

Initalizing model...
Loading restart file from: `./output/tutorial/restart_files/tutorial_restart_1981_05_15.nc`.
Done.
Beginning simulation for 1981-05-15 through 1981-05-28...
Running mosartwmpy
Simulation completed in 1 minutes and 21 seconds.

Let’s look at the surface water and reservoir storage again:

mosart_wm.plot_variable('surface_water_amount', log_scale=True)

mosart_wm.plot_variable('reservoir_water_amount', log_scale=True)

These plots make more sense than the random runoff!

mosartwmpy also includes visualization methods to examine the output over time. For instance, one can plot the behavior of a specific reservoir or create a timeplayer of the water deficit:

plot_reservoir(
    model=mosart_wm,
    grand_id=310,
    start='1981-05-01',
    end='1981-05-31',
)

plot_variable(
    model=mosart_wm,
    variable='WRM_DEFICIT',
    start='1981-05-01',
    end='1981-05-31',
    log_scale=True,
    cmap='autumn_r',
)

Launching server at http://localhost:58014

This concludes the mosartwmpy tutorial. Feel free to open issues with your feedback or start topics on the discussion board!