Dynamic python code loading for hsp2 runtime customization

[rburghol]

Overview (RFC)

This issue provides background information for pull request #119 which hs now been merged.

This PR contains code that:

• Establishes a shared data STATE data structure outlined in Data Model for Shared STATE (supports model modularity) RFC #112
• Integrate loading of dynamic code with develop branch after STATE data PR Devstate #119
• Enables the end user to include a python file that defines functions that can be used inside the model timestep to modify elements of the STATE data model. (see Testing below).
  • Allows modelers to modify runtime state and/or specific parameters of the model with project specific python code.
  • Uses "magic" name import, i.e., given hdf5 file river_model.h5, automatically load code contained in local directory named river_model.py.
• Provides an easy to implement example (see Testing below) of how to modify the withdrawal of a river segment at every time step with a short python function residing in the same directory as the model h5 file..
• Establishes a set of functions that when supported allow several points of custom access during model execution.
  • Ex.: state_step_hydr() can modify STATE variables during the timestep execution of the HYDR/_hydr_ functions. See Dynamic Modifiers below.
• Provides code that integrates modifications to the STATE variable during each timestep in the HYDR/_hydr_ function by calling state_step_hydr() if defined in the model run directory.
• Provides a potential avenue for data modifications by special actions code. (See also SPEC-ACTIONS in hsp2 #90)
• Limitations: The technique that is being used to dynamically import python code for runtime incurs an overhead because it forces recompile for any @njit function that supports the code. For example, using the dynjamic code in_hydr_() shown below adds about 6-8 seconds for each and every model run, regardless of whether or not the dynamic code has been changed. Perhaps there is a better way?

Tasks

• Develop runtime loadable function for Dynamic Loading of Code (below)
• Outline Supported Functions and naming conventions (init, step, etc)
• Outline data model for enabling/disabling
• Integrate dynamic loader into hsp2
• Do pull request

Supported Functions & Special Variables

• General supported functions are:
  • pre_step(): things that happen at the beginning of the code loop, after normal data loading has occurred
  • step(): things that happen during the computational portion of a model time-step
    • Ex: state_step_hydr(state_ix, dict_ix, ts_ix, hydr_ix, step)
  • post_step(): things that happen after a timestep is completed (ex: logging)
  • get_ix: a special function to retrieve local variable names with an implied path to the current domain )
• Naming convention:
  • state_pre_step_[operation]()
  • state_step_[operation]()
  • state_post_step_[operation]()
• Special Variables:
  • state: this is a multi-typed dict (not numba compatible) that is used to pass all the other Dicts related to state
    • state['domain']: the current path to the /STATE/[operation]/[segment], this gets reset each time we switch segment or operation
  • state_paths: this is the path that contains pointers from the full hdf5 paths for state variables to the corresponding unique integer key for the state_ix Dict.
    • state_paths['domain'] = special entry in state_paths that contains the integer key for the current /STATE/[operation]/[segment]. This permits
    • hydr_ix: this is created automatically to hold shortcuts to HYDR variables for the current domain for brevity and fast code exec.
• Minimal Code Support:
  • provide state dictionary in main.py - sets up basic shared numba Dict entries to pass to routines (includes state_paths, state_ix, dict_ix, ts_ix, operation, segment, ...)
  • Add state as argument to operation function,
    • Ex: def hydr(io_manager, siminfo, uci, ts, ftables, state):
  • Add state_ix, dict_ix, ts_ix, state_info, state_paths as arguments to operation numba function
    • These are separated out from state because numba cannot take mixed type Dict
    • Ex: def _hydr_(ui, ts, COLIND, OUTDGT, rowsFT, funct, Olabels, OVOLlabels, state_info, state_paths, state_ix, dict_ix, ts_ix)
  • Pass state_ix, dict_ix, ts_ix, state_info, state_paths as arguments when calling operation numba function
  • Add state_step_[operation] call if defined
  • Set state prior to calling state_step_[operation] (optional, performance) Provide _get_ix function for variables that wish to access state variable changes in the loop

Dynamic Loading of Code

• Code implements methods to customize state variables during time steps:
• Individual functions, if defined will be called as: state_step_hydr(model_exec_list, op_tokens, state_ix, dict_ix, ts_ix, hydr_ix, step)

if 'state_step_hydr' in dir(hsp2_local_py):
        siminfo['state_step_hydr'] = `enabled`

Testing

• Load branch statey
• Run test for model test10
• Modify test10.py, editing the function state_step_hydr() to explore interacting with state.
• Note:
  • Only _hydr_() has support implemented for STATE functions.
  • Only IVOL and OUTDGT modifications are passed from state_ix back into _hydr_
• Current example code below in Code 1

Code 1: Set a dynamic withdrawal in RCHRES R001. See also https://github.com/respec/HSPsquared/blob/99c0e629d88d0665bb45f10250691b2f0679967e/tests/test10/HSP2results/test10.py

import sys
from numba import int8, float32, njit, types, typed # import the types

print("Loaded a set of  HSP2 code!")

@njit
def state_step_hydr(state_info, state_paths, state_ix, dict_ix, ts_ix, hydr_ix, step):
    if (step <= 1):
        print("Custom state_step_hydr() called for ", state_info['segment'])
        print("operation", state_info['operation'])
        print("state at start", state_ix)
        print("domain info", state_info)
        print("state_paths", state_paths)
    # Do a simple withdrawal of 10 MGGD from R001
    if (state_info['segment'] == 'R001') and (state_info['activity'] == 'HYDR'):
        state_ix[hydr_ix['O1']] = 10.0 * 1.547
    # Route point source return from segment R001 demand to R005 inflow (IVOL)
    # For demo purposes this will only use the last state_ix value for R001 demand
    # Realistic approach would run all segments simultaneously or use value from ts_ix (ts_ix loading TBD)
    if (state_info['segment'] == 'R005') and (state_info['activity'] == 'HYDR'):
        state_ix[hydr_ix['IVOL']] += 0.85 * state_ix[state_paths['/STATE/RCHRES_R001/HYDR/O1']]
        if (step <= 1):
            print("IVOL after", state_ix[hydr_ix['IVOL']])
    return

[jdkleiner]

Analysis of testing results:

• first step: run test10 in order to have a .h5 file to analyze
• You can run the following from within tests\test10\HSP2results

import h5py
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# read in results hdf5 file
filename = 'test10.h5'
h5 = h5py.File(filename, 'r')

# extract results table from hdf5 as pandas dataframe
df = pd.DataFrame(np.array(h5['RESULTS']['RCHRES_R001']['HYDR']['table']))

# print all column names
print(list(df.columns.values))

# simple line plot
df.plot(kind='line',
        x='index',
        y=['OVOL1','IVOL'])

plt.show()