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 demand_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 numerical considerations for a mathematical description of the Jacobian and how it is evaluated, and numerical implementation for the implementation.

The sparsity structure of the Jacobian is detected automatically with SparseConnectivityTracer.jl, and the Jacobian is computed with automatic differentiation using ForwardDiff.jl. Hence no Jacobian specific code is needed for a new node type, as long as formulate_flow! for NewNodeType adheres to the following:

  • The code must be compatible with dual numbers (e.g. no hardcoded Float64 buffers for intermediate results that depend on the state);
  • Basin storages and levels may only be obtained from the storage_uplink and storage_downlink arguments, i.e. the flow of a node may only depend on the storages of its direct upstream and downstream Basin;
  • Flows that are controlled by PidControl or ContinuousControl must obtain the controlled value from the PID integral or the ContinuousControl compound variable respectively.

Otherwise the dependency of the flow on the state is missed in the reduced linear solve (build_J_inner!), or an error is thrown.

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", QMetaType.Type.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_links 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/getting-started/examples.qmd 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