Extending functionality with custom parameters
Parameters are a core part of Pywr, allowing you to define how your model behaves. While Pywr comes with a wide range of built-in parameters, you may find that you need to create custom parameters to suit your specific modelling needs. This guide will walk you through the process of creating custom parameters in Pywr.
Currently, Pywr supports custom parameters that are defined in Python. If your parameter is general enough, you may want to consider contributing it to the Pywr project. If you do, please see the Developers Guide for more information on how to do this.
Python functions
The simplest way to create a custom parameter is to define a Python function.
This function should accept at least one argument, which is a ParameterInfo object.
This object contains information from the model, such as the current time step, scenario index,
and any metric values that have been requested.
Additional arguments can also be passed to the function.
Here is an example of a simple custom parameter that returns the current time step:
# custom_parameters.py
from pywr import ParameterInfo
def current_time_step(info: ParameterInfo) -> float:
"""Return the current time step."""
return info.timestep.index
To use this custom parameter in your model it must be defined as a Parameter in your model's JSON file.
Below is an example of how to define the current_time_step parameter in your model's JSON file.
The source field specifies the path to the Python file containing the function,
and the object field specifies the name of the function to call.
{
"parameters": [
{
"meta": {
"name": "current-time-step"
},
"type": "Python",
"source": {
"type": "Path",
"path": "custom_parameters.py"
},
"object": {
"type": "Function",
"class": "current_time_step"
},
"args": [],
"kwargs": {}
}
]
}
Constant arguments
In reality, your function will likely need to accept additional arguments.
These arguments might be constants that change the behaviour of the function, but do not change over time
or are a result of the model's simulation state.
In this case they can be defined as args or kwargs in the parameter definition.
Only simple types that are supported by JSON can be used as arguments, such as strings, numbers, and booleans.
However, by parameterising these values, you can easily change them without modifying the Python code or reuse
the same function with different values in different parts of the model.
# custom_parameters.py
from pywr import ParameterInfo
def current_timestep(info: ParameterInfo, a: float, b: float, some_condition: str = "foo") -> float:
"""Return the current time step."""
match some_condition:
case "foo":
return info.timestep.index + a
case "bar":
return info.timestep.index + b
case _:
raise ValueError(f"Invalid condition: {some_condition}")
To pass these arguments to the function, you can define them in the model's JSON file as follows:
{
"parameters": [
{
"meta": {
"name": "current-time-step"
},
"type": "Python",
"source": {
"type": "Path",
"path": "custom_parameters.py"
},
"object": {
"type": "Function",
"class": "current_timestep"
},
"args": [
1.0,
2.0
],
"kwargs": {
"some_condition": "foo"
}
}
]
}
Metrics from the model
More complex parameters will need information from the model, such as the current volume of a reservoir, or
the value of another parameter, etc.
These values need to be requested in the JSON definition of parameter, and then they can be accessed in the function
using the ParameterInfo object.
# custom_parameters.py
from pywr import ParameterInfo
def factor_volume(info: ParameterInfo, factor: float) -> float:
"""Return the current volume of a reservoir scaled by `factor`."""
volume = info.get_metric("volume")
return factor * volume
The JSON definition of the parameter needs to include a metrics and/or indices field that specifies which model
metrics to request. Both fields are a dictionary where the keys are the keys used to retrieve the values from the
ParameterInfo object, and the values specify the metric to retrieve. Metrics are accessed using get_metric(key),
and indices are accessed using get_index(key).
{
"parameters": [
{
"meta": {
"name": "factor-volume"
},
"type": "Python",
"source": {
"type": "Path",
"path": "custom_parameters.py"
},
"object": {
"type": "Function",
"class": "factor_volume"
},
"args": [
2.0
],
"metrics": {
"volume": {
"type": "Node",
"name": "a-reservoir",
"attribute": "Volume"
}
}
}
]
}
Python classes & stateful parameters
If your parameter needs to maintain state between calls, you can define it as a Python class.
This class should implement an __init__ method that setups up the parameter, including any
initial state.
The __init__ method is passed the args and kwargs defined in the JSON file.
Pywr will create an instance of the class for every scenario in a simulation.
These instances will be reused for each time step in the scenario, allowing you to maintain state across time steps.
Note: Unlike Pywr v1.x a separate instance of the class is created for each scenario. This means you do not have to worry about state being shared between scenarios, and do not need to implement state for each scenario yourself.
The class should also implement before method, which is called for each time step in the scenario.
This method should accept a ParameterInfo object as its only argument.
Finally, the class may also implement an after method, which is called after the resource allocation
has been completed for the time step.
This method can be used to perform any final calculations or updates to the parameter state.
Here is an example of a simple stateful parameter that counts the number of time steps:
# custom_parameters.py
from pywr import ParameterInfo
class TimeStepCounter:
"""A parameter that counts the number of time steps."""
def __init__(self, initial_value: int = 0):
self.count = initial_value
def before(self, _info: ParameterInfo) -> float | None:
"""Return the current time step count."""
# Note that `_info` is not used, but it is required by the interface.
self.count += 1
return self.count
To use this custom parameter in your model, you can define it in the JSON file as follows:
{
"parameters": [
{
"meta": {
"name": "time-step-counter"
},
"type": "Python",
"source": {
"type": "Path",
"path": "custom_parameters.py"
},
"object": {
"type": "Class",
"class": "TimeStepCounter"
},
"args": [
0
],
"kwargs": {}
}
]
}
Using modules instead of files
It might be more convenient to define your custom parameters in a Python module instead of a file. This
allows you to integrate your custom parameters with other Python code, such as unit tests or other utility functions.
To do this, you can use the source field to specify the module name instead of a file path.
Here is an example of how to define a custom parameter in a module (in this case my_model.parameters):
{
"parameters": [
{
"meta": {
"name": "current-time-step"
},
"type": "Python",
"source": {
"type": "Module",
"module": "my_model.parameters"
},
"object": {
"type": "Function",
"class": "current_time_step"
},
"args": [],
"kwargs": {}
}
]
}
Returning integers or multiple values
In the examples above the custom parameter functions return a single floating point value.
However, you can also return integers or multiple values.
In the JSON definition of the parameter, you can specify the return_type field to indicate the type of value
the function will return.
To return an integer, you can set the return_type to "Int".
To return multiple values, you can set the return_type to "Dict" and the function should return a dictionary
where the keys are the names of the values and the values are the values themselves.
An example of a custom parameter that returns multiple values is shown below:
# custom_parameters.py
from pywr import ParameterInfo
def multiple_values(info: ParameterInfo, factor: float) -> dict:
"""Return multiple values."""
return {
"value1": info.timestep.index,
"value2": info.get_metric("volume") * factor
}
The corresponding JSON for this parameter would look like this:
{
"parameters": [
{
"meta": {
"name": "multiple_values"
},
"type": "Python",
"source": {
"type": "Module",
"module": "my_model.parameters"
},
"object": {
"type": "Function",
"class": "multiple_values"
},
"return_type": "Dict",
"args": [
2.0
],
"metrics": {
"volume": {
"type": "Node",
"name": "a-reservoir",
"attribute": "Volume"
}
}
}
]
}
And the returned values can be accessed in the model using the keys defined in the dictionary.
{
"type": "Parameter",
"name": "multiple_values",
"key": "value1"
// or "value2"
}
Before and after methods
The majority of parameters will only need to implement the before method, which is called before the resource
allocation is performed for the time step1. However, in some cases it may be necessary to perform some
calculations after the resource allocation has been completed. Typically, this is the case when the parameter needs to
access the results of the resource allocation, such as the allocated flow or the new volume of a reservoir after the
allocation. In this case, the parameter can implement an after method.
The example below lists a custom parameter that implements both before and after methods. It is a simple crop
water requirement parameter that calculates the water requirement for a crop based on the current month in before,
and then computes a crop yield in after based on the allocated water and the water requirement.
from pathlib import Path
class CropParameter:
"""A simple example of a crop parameter.
It produces an irrigation requirement value based on the month during `before`. This
is intended to be used as a demand (or "max_flow") on an irrigation node. The `after`
method tracks any deficit in irrigation supplied, and at the end of the growing season
returns the yield for that season.
"""
def __init__(self):
self.crop_yield = 0.0
# Example irrigation requirements by month
self.irrigation_requirements = {
1: 0.0, # January
2: 0.0, # February
3: 10.0, # March
4: 20.0, # April
5: 30.0, # May
6: 40.0, # June
7: 30.0, # July
8: 20.0, # August
9: 10.0, # September
10: 0.0, # October
11: 0.0, # November
12: 0.0, # December
}
self.growing_season_months = {3, 4, 5, 6, 7, 8, 9}
def before(self, info) -> float:
"""Return the irrigation requirement for the current month."""
return self.irrigation_requirements.get(info.timestep.month, 0.0)
def after(self, info) -> float:
"""Track the yield based on irrigation supplied."""
irrigation_required = self.irrigation_requirements.get(info.timestep.month, 0.0)
irrigation_supplied = info.get_metric("supplied")
deficit = irrigation_required - irrigation_supplied
if info.timestep.month not in self.growing_season_months:
# Reset yield at the end of the growing season
if info.timestep.month == 10 and info.timestep.day == 1:
final_crop_yield = self.crop_yield
self.crop_yield = 0.0
return final_crop_yield
else:
# Implement a simple crop growth/yield model based on irrigation deficit
self.crop_yield += max(0.0, 1.0 - (deficit / irrigation_required))
return 0.0 # Yield is only returned at the end of the season
def run(model_path: Path):
from pywr import ModelSchema
schema = ModelSchema.from_path(model_path)
model = schema.build(model_path.parent, None)
model.run("clp")
print("Model run complete 🎉")
if __name__ == "__main__":
pth = Path(__file__).parent / "model.json"
run(pth)
When referring to the parameter in the model, the return_value field can be used to specify whether to use the value
returned by the before or after method. By default, the value returned by the before method is used, but if you
want to use the value from the after method you can set return_value to "After". The example below shows how to
use the CropParameter above parameter in a metric set, and specify that the value from the after method should be used.
"metric_sets": [
{
"name": "parameters",
"aggregator": {
"freq": {
"type": "Annual"
},
"func": {
"type": "Sum"
}
},
"metrics": [
{
"type": "Parameter",
"name": "crop1",
"return_value": "After"
}
]
}
],
Cython (and other compiled languages)
Cython functions and classes can be used in Pywr as long as they accessible from Python, and can be imported by Pywr at runtime. In this case using a module for locating the custom parameter is recommended. Otherwise, there is no difference in how you define the custom parameter in the model's JSON file.
Other compiled languages can also be used, but you will need to ensure that the compiled code is accessible from Python.
This can be done by using a Python wrapper around the compiled code, or by using a foreign function interface (FFI)
such as ctypes or cffi.
-
This also is the same as Pywr v1.x where the
beforemethod was the only method that could be implemented. ↩