Adding node types

Several parts of the code have to be made aware of the new node type. In the rest of this page we shall call our new node type NewNodeType.

1 The Julia core

1.1 Parameters

The parameters object (defined in parameter.jl) passed to the ODE solver must be made aware of the new node type. Therefore define a struct in parameter.jl which holds the data for each node of the new node type:

struct NewNodeType <: AbstractParameterNode
    node_id::Vector{NodeID}
    # Other fields
end

Another abstract type which subtypes from AbstractParameterNode is called AbstractDemandNode. For creating new node type used in allocation, define a struct:

struct NewNodeType <: AbstractDemandNode
    node_id::Vector{NodeID}
    # Other fields
end

These fields do not have to correspond 1:1 with the input tables (see below). The vector with all node IDs that are of the new type in a given model is a mandatory field. Now you can:

  • Add new_node_type::NewNodeType to the Parameters object;
  • Add new_node_type = NewNodeType(db,config) to the function Parameters in read.jl and add new_node_type at the proper location in the Parameters constructor call.

1.2 Reading from configuration

There can be several schemas associated with a single node type. To define a schema for the new node type, add the following to schema.jl:

@schema "ribasim.newnodetype.static" NewNodeTypeStatic

"""
node_id: node ID of the NewNodeType node
"""
@version NewNodeTypeStaticV1 begin
    node_id::Int32
    # Other fields
end

Here static refers to data that does not change over time. For naming conventions of these schemas see Node usage. If a new schema contains a priority column for allocation, it must also be added to the list of all such schemas in the function get_all_priorities in util.jl.

validation.jl deals with checking and applying a specific sorting order for the tabular data (default is sorting by node ID only), see sort_by_function and sorted_table!.

Now we define the function that is called in the second bullet above, in read.jl:

function NewNodeType(db::DB, config::Config)::NewNodeType
    static = load_structvector(db, config, NewNodeTypeStaticV1)
    defaults = (; foo = 1, bar = false)
    # Process potential control states in the static data
    parsed_parameters, valid = parse_static_and_time(db, config, "NewNodeType"; static, defaults)

    if !valid
        error("Errors occurred when parsing NewNodeType data.")
    end

    # Unpack the fields of static as inputs for the NewNodeType constructor
    return NewNodeType(
        NodeID.(NodeType.NewNodeType, parsed_parameters.node_id),
        parsed_parameters.some_property,
        parsed_parameters.control_mapping)
end

1.3 Node behavior

In general if the new node type dictates flow, the behavior of the new node in the Ribasim core is defined in a method of the formulate_flow! function, which is called within the water_balance! (both in solve.jl) function being the right hand side of the system of differential equations solved by Ribasim. Here the details depend highly on the specifics of the node type. An example structure of a formulate_flow! method is given below.

function formulate_flow!(new_node_type::NewNodeType, p::Parameters)::Nothing
    # Retrieve relevant parameters
    (; graph) = p
    (; node_id, param_1, param_2) = new_node_type

    # Loop over nodes of NewNodeType
    for (i, id) in enumerate(node_id)
        # compute e.g. flow based on param_1[i], param_2[i]
    end

    return nothing
end

If the new node type is non-conservative, meaning it either adds or removes water from the model, these boundary flows also need to be recorded. This is done by storing it on the diagonal of the flow[from, to] matrix, e.g. flow[id, id] = q, where q is positive for water added to the model.

1.4 The Jacobian

See Equations for a mathematical description of the Jacobian.

Before the Julia core runs its simulation, the sparsity structure jac_prototype of \(J\) is determined with get_jac_prototype in sparsity.jl. This function runs trough all node types and looks for nodes that create dependencies between states. It creates a sparse matrix of zeros and ones, where the ones denote locations of possible non-zeros in \(J\). Note that only nodes that set flows in the physical layer (or have their own state like PidControl) affect the sparsity structure.

We divide the various node types in groups based on what type of state dependencies they yield, and these groups are discussed below. Each group has its own method update_jac_prototype! in utils.jl for the sparsity structure induced by nodes of that group. NewNodeType should be added to the signature of one these methods, or to the list of node types that do not contribute to the Jacobian in the method of update_jac_prototype! whose signature contains node::AbstractParameterNode. Of course it is also possible that a new method of update_jac_prototype! has to be introduced.

The current dependency groups are:

  • Out-neighbor dependencies: examples are TabulatedRatingCurve, Pump (the latter only in the reduction factor regime and not PID controlled). If the in-neighbor of a node of this group is a basin, then the storage of this basin affects itself and the storage of the outneighbor if that is also a basin;
  • Either-neighbor dependencies: examples are LinearResistance, ManningResistance. If either the in-neighbor or out-neighbor of a node of this group is a basin, the storage of this basin depends on itself. If both the in-neighbor and the out-neighbor are basins, their storages also depend on eachother.
  • The PidControl node is a special case which is discussed in the PID equations.

Using jac_prototype the Jacobian of water_balance! is computed automatically using ForwardDiff.jl with memory management provided by PreallocationTools.jl. These computations make use of DiffCache and dual numbers.

2 Python I/O

2.1 Python class

In python/ribasim/ribasim/config.py add

  • the above defined schemas to the imports from ribasim.schemas. This requires code generation to work, see Finishing up;
  • a class of the following form with all schemas associated with the node type:
class NewNodeType(MultiNodeModel):
    static: TableModel[NewNodeTypeStaticSchema] = Field(
        default_factory=TableModel[NewNodeTypeStaticSchema],
        json_schema_extra={"sort_keys": ["node_id"]},
    )

In python/ribasim/ribasim/nodes/__init__.py add

  • NewNodeType to the imports from ribasim.nodes;
  • "NewNodeType" to __all__.

In python/ribasim/ribasim/model.py, add

  • NewNodeType to the imports from ribasim.config;
  • new_node_type as a parameter of the Model class.

In python/ribasim/ribasim/geometry/node.py add a color and shape description in the MARKERS and COLORS dictionaries.

3 QGIS plugin

The script ribasim_qgis/core/nodes.py has to be updated to specify how the new node type is displayed by the QGIS plugin. Specifically:

  • Update the .qml style (using QGIS) in the styles folder for the specific Node.
  • Add an input class per schema, e.g.
class NewNodeTypeStatic:
    input_type = "NewNodeType / static"
    geometry_type = "No Geometry"
    attributes = [
        QgsField("node_id", QVariant.Int)
        # Other fields for properties of this node
    ]

4 Validation

The new node type might have associated restrictions for a model with the new node type so that it behaves properly. Basic node ID and node type validation happens in Model.validate_model in python/ribasim/ribasim/model.py, which automatically considers all node types in the node_types module.

Connectivity validation happens in valid_edges and valid_n_flow_neighbors in core/src/solve.jl. Connectivity rules are specified in core/src/validation.jl. Allowed upstream and downstream neighbor types for new_node_type (the snake case version of NewNodeType) are specified as follows:

# set allowed downstream types
neighbortypes(::Val{:new_node_type}) = Set((:basin,))
# add your newnodetype as acceptable downstream connection of other types
neighbortypes(::Val{:pump}) = Set((:basin, :new_node_type))

The minimum and maximum allowed number of inneighbors and outneighbors for NewNodeType are specified as follows:

# Allowed number of flow/control inneighbors and outneighbors per node type
struct n_neighbor_bounds
    in_min::Int
    in_max::Int
    out_min::Int
    out_max::Int
end

n_neighbor_bounds_flow(::Val{:NewNodeType}) =
    n_neighbor_bounds(0, 0, 1, typemax(Int))

n_neighbor_bounds_control(::Val{:NewNodeType}) =
    n_neighbor_bounds(0, 1, 0, 0)

Here typemax(Int) effectively means unbounded.

5 Tests

Models for the julia tests are generated by running pixi run generate-testmodels, which uses model definitions from the ribasim_testmodels package, see here. These models should also be updated to contain the new node type. Note that certain tests must be updated accordingly when the models used for certain tests are updated, e.g. the final state of the models in core/test/basin.jl. The following function is used to format the array of this final state.

reprf(x) = repr(convert(Vector{Float32}, x))

See here for monitoring of Python test coverage.

If the new node type introduces new (somewhat) complex behavior, a good test is to construct a minimal model containing the new node type in python/ribasim_testmodels/ribasim_testmodels/equations.py and compare the simulation result to the analytical solution (if possible) in core/test/equations.jl.

6 Documentation

There are several parts of the documentation which should be updated with the new node type:

  • If the node has a rol in the physical layer, docs/core/equations should contain a short explanation and if possible an analytical expression for the behavior of the new node;
  • If the node has a role in allocation, docs/core/allocation should make this role clear;
  • docs/reference/node/new-node-type.qmd should contain a short explanation of the node and the possible schemas associated with it;
  • The example models constructed in docs/guide/examples.ipynb should be extended with the new node type or a new example model with the new node type should be made.
  • In _quarto.yml add NewNodeType to the “Node types” contents for the Python API reference.

7 Finishing up

When a new node type is created, one needs to run

pixi run codegen

This will derive all JSON Schemas from the julia code, and write them to the docs folder. From these JSON Schemas the Python modules models.py and config.py are generated.

Since adding a node type touches both the Python and Julia code, it is a good idea to run both the Python test suite and Julia test suite locally before creating a pull request. You can run all tests with:

pixi run tests