Water Allocation

Version:

Experience level: Beginner
Reasoning types: Prescriptive
Industry: Utilities
Tags: AllocationLPNetwork Flow
Experience level: Beginner
Reasoning types: Prescriptive
Industry: Utilities
Tags: AllocationLPNetwork Flow
Experience level: Intermediate
Reasoning types: Prescriptive
Industry: Utilities & Resources
Tags: Resource AllocationNetwork FlowNonlinear ProgrammingIpopt

What this template is for

Water utilities (and any operator of a distribution network) must decide how to route limited supply to satisfy demand at the lowest possible cost. This template models a small water distribution network: sources (e.g., reservoirs and groundwater) send water to users (municipal, industrial, agricultural) over a set of connections with capacity limits and transmission losses.

This template uses RelationalAI’s prescriptive reasoning (optimization) capabilities to choose flows on each connection that meet all user demands while minimizing total sourcing cost.

Who this is for

You want a small, end-to-end example of prescriptive reasoning (optimization) with RelationalAI.
You’re comfortable with basic Python and the idea of constraints + objectives.

What you’ll build

A semantic model of water Source, demand User, and Connection entities using concepts and typed properties.
A linear program (LP) that chooses a non-negative flow on each source-to-user connection.
Constraints for source capacity, per-connection max flow, and demand satisfaction after accounting for loss rates.
A solver run (HiGHS backend) that prints a readable table of non-trivial flow allocations.

What’s included

Model + solve script: water_allocation.py
Sample data: data/sources.csv, data/users.csv, data/connections.csv
Outputs: Printed termination status, objective value, and a flow allocation table

Prerequisites

Access

A Snowflake account that has the RAI Native App installed.
A Snowflake user with permissions to access the RAI Native App.

Tools

Python >= 3.10

Quickstart

Follow these steps to run the template with the included sample data.

Download the ZIP file for this template and extract it:
Terminal window
```
curl -O https://private.relational.ai/templates/zips/v0.13/water_allocation.zip
unzip water_allocation.zip
cd water_allocation
```
You can also download the template ZIP using the “Download ZIP” button at the top of this page.

Create and activate a virtual environment

python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip

Install dependencies

From this folder:
Terminal window
```
python -m pip install .
```
Configure Snowflake connection and RAI profile
Terminal window
```
rai init
```
Run the template
Terminal window
```
python water_allocation.py
```

Expected output

Status: OPTIMAL
Total cost: $874.28

Flow allocations:
      source        user       flow
 Reservoir_A   Municipal 317.934783
 Reservoir_A  Industrial 182.065217
 Reservoir_A Agricultural 500.000000
 Reservoir_B   Municipal 316.980758
 Reservoir_B  Industrial 250.000000
 Reservoir_B Agricultural  56.818182

Template structure

.
├─ README.md
├─ pyproject.toml
├─ water_allocation.py         # main runner / entrypoint
└─ data/                       # sample input data
   ├─ sources.csv
   ├─ users.csv
   └─ connections.csv

Start here: water_allocation.py

Sample data

Data files are in data/.

`sources.csv`

Defines water sources, each with a capacity (maximum total outflow) and a per-unit cost.

Column	Meaning
`id`	Unique source identifier
`name`	Source name (e.g., `Reservoir_A`)
`capacity`	Maximum total flow out of the source
`cost_per_unit`	Cost per unit of water sourced

`users.csv`

Defines demand points (users). The priority column is included in the sample data but is not used in the optimization model in this template.

Column	Meaning
`id`	Unique user identifier
`name`	User name (e.g., `Municipal`)
`demand`	Required delivered volume (after losses)
`priority`	Integer priority (not used in the model)

`connections.csv`

Defines feasible source-to-user connections with a maximum allowed flow and a fractional loss rate.

Column	Meaning
`source_id`	Foreign key to `sources.csv.id`
`user_id`	Foreign key to `users.csv.id`
`max_flow`	Upper bound on flow along this connection
`loss_rate`	Fraction of flow lost in transmission (e.g., `0.08` means 8% loss)

Model overview

The semantic model for this template is built around three concepts.

`Source`

A water supply point with limited capacity and a per-unit sourcing cost.

Property	Type	Identifying?	Notes
`id`	int	Yes	Loaded as the key from `data/sources.csv`
`name`	string	No	Human-readable label used in output
`capacity`	float	No	Upper bound on total outflow
`cost_per_unit`	float	No	Multiplies flow in the objective

`User`

A demand point that must receive enough effective inflow (flow adjusted for losses).

Property	Type	Identifying?	Notes
`id`	int	Yes	Loaded as the key from `data/users.csv`
`name`	string	No	Human-readable label used in output
`demand`	float	No	Required delivered volume
`priority`	int	No	Included in sample data; not used in this model

`Connection`

A feasible link from a Source to a User that can carry flow up to max_flow and loses a fraction loss_rate in transit.

Property	Type	Identifying?	Notes
`source`	`Source`	Part of compound key	Joined from `data/connections.csv.source_id`
`user`	`User`	Part of compound key	Joined from `data/connections.csv.user_id`
`max_flow`	float	No	Upper bound on the decision variable
`loss_rate`	float	No	Used to compute delivered inflow: `flow * (1 - loss_rate)`
`flow`	float	No	Continuous decision variable ()

How it works

This section walks through the highlights in water_allocation.py.

Import libraries and configure inputs

First, the script imports the Semantics APIs (Model, data, where, require, sum, select) and configures the local DATA_DIR it will read CSVs from:

from pathlib import Path

import pandas
from pandas import read_csv

from relationalai.semantics import Model, data, require, select, sum, where
from relationalai.semantics.reasoners.optimization import Solver, SolverModel

# --------------------------------------------------
# Configure inputs
# --------------------------------------------------

DATA_DIR = Path(__file__).parent / "data"

# Disable pandas inference of string types. This ensures that string columns
# in the CSVs are loaded as object dtype. This is only required when using
# relationalai versions prior to v1.0.
pandas.options.future.infer_string = False

Define concepts and load CSV data

Next, it declares Source and User concepts and loads sources.csv and users.csv into those concepts with data(...).into(...):

# Create a Semantics model container.
model = Model("water_allocation", config=globals().get("config", None), use_lqp=False)

# Source concept: water supply points with capacity and per-unit cost.
Source = model.Concept("Source")
Source.id = model.Property("{Source} has {id:int}")
Source.name = model.Property("{Source} has {name:string}")
Source.capacity = model.Property("{Source} has {capacity:float}")
Source.cost_per_unit = model.Property("{Source} has {cost_per_unit:float}")

# Load source data from CSV.
source_csv = read_csv(DATA_DIR / "sources.csv")
data(source_csv).into(Source, keys=["id"])

# User concept: demand points with required volume (and a priority field).
User = model.Concept("User")
User.id = model.Property("{User} has {id:int}")
User.name = model.Property("{User} has {name:string}")
User.demand = model.Property("{User} has {demand:float}")
User.priority = model.Property("{User} has {priority:int}")

# Load user data from CSV.
user_csv = read_csv(DATA_DIR / "users.csv")
data(user_csv).into(User, keys=["id"])

Then it declares a Connection concept and uses where(...).define(...) to join each connections.csv row to the corresponding Source and User:

# Connection concept: links a Source to a User with transmission parameters.
Connection = model.Concept("Connection")
Connection.source = model.Property("{Connection} from {source:Source}")
Connection.user = model.Property("{Connection} to {user:User}")
Connection.max_flow = model.Property("{Connection} has {max_flow:float}")
Connection.loss_rate = model.Property("{Connection} has {loss_rate:float}")
Connection.x_flow = model.Property("{Connection} has {flow:float}")

# Load connection data from CSV.
conn_data = data(read_csv(DATA_DIR / "connections.csv"))

# Define Connection entities by joining the CSV data with Source and User.
where(
    Source.id == conn_data.source_id,
    User.id == conn_data.user_id
).define(
    Connection.new(
        source=Source,
        user=User,
        max_flow=conn_data.max_flow,
        loss_rate=conn_data.loss_rate,
    )
)

Define decision variables, constraints, and objective

With the network data in place, the script creates a continuous SolverModel and declares Connection.x_flow as a non-negative decision variable with an upper bound from Connection.max_flow:

Conn = Connection.ref()

# Create a continuous optimization model.
s = SolverModel(model, "cont")

# Decision variable: flow on each connection (continuous, non-negative).
s.solve_for(
    Connection.x_flow,
    name=["flow", Connection.source.name, Connection.user.name],
    lower=0,
    upper=Connection.max_flow,
)

Next, it adds two constraints with require(...) and s.satisfy(...): source capacity (total outflow per source) and demand satisfaction (effective inflow per user, accounting for losses):

# Constraint: total outflow from each source must not exceed its capacity.
outflow = sum(Conn.flow).where(Conn.source == Source).per(Source)
source_limit = require(outflow <= Source.capacity)
s.satisfy(source_limit)

# Constraint: effective inflow to each user must meet demand (accounting for losses).
effective_inflow = (
    sum(Conn.flow * (1 - Conn.loss_rate)).where(Conn.user == User).per(User)
)
meet_demand = require(effective_inflow >= User.demand)
s.satisfy(meet_demand)

Finally, it minimizes total sourcing cost by summing flow times the per-unit cost of the corresponding source:

# Objective: minimize total cost.
total_cost = sum(Connection.x_flow * Connection.source.cost_per_unit)
s.minimize(total_cost)

Solve and print results

The script solves with the HiGHS backend, prints status and objective value, and then uses select(...).where(...).to_df() to display only non-trivial flows (> 0.001):

solver = Solver("highs")
s.solve(solver, time_limit_sec=60)

print(f"Status: {s.termination_status}")
print(f"Total cost: ${s.objective_value:.2f}")

allocations = select(
    Connection.source.name.alias("source"),
    Connection.user.name.alias("user"),
    Connection.x_flow,
).where(
    Connection.x_flow > 0.001
).to_df()

print("\nFlow allocations:")
print(allocations.to_string(index=False))

Customize this template

Use your own data

Replace the CSVs in data/ with your own, keeping the same column names (or update the loading logic in water_allocation.py).
Make sure your connections.csv only references valid source_id and user_id values that exist in sources.csv and users.csv.
Remember that users.csv.demand is the delivered requirement; because of losses, the required flow may be higher than demand.

Tune parameters

Adjust cost tradeoffs by changing sources.csv.cost_per_unit.
Model higher leakage/evaporation by changing connections.csv.loss_rate.
Tighten or relax infrastructure constraints by changing connections.csv.max_flow and sources.csv.capacity.

Extend the model

Add intermediate transfer nodes (source → junction → user) by introducing a node concept and splitting Connection into edge-to-edge connections.
Incorporate user priorities (currently loaded but unused) by changing the objective to penalize unmet demand with priority weights, or by adding staged solves.

Troubleshooting

rai init fails or the script cannot authenticate

Run rai init again and confirm you can see/select the right Snowflake account, role, warehouse, and the RAI Native App.
If you use multiple profiles, set RAI_PROFILE to the intended profile before running the script.
Ensure your Snowflake user has privileges to use the RAI Native App in the selected context.

ModuleNotFoundError for relationalai or pandas

Confirm your virtual environment is active (source .venv/bin/activate).
From the template folder, reinstall dependencies with python -m pip install ..
Verify you’re using Python 3.10+ with python —version.

Missing or incorrect CSV columns (for example: KeyError or AttributeError)

Verify the CSV headers match the expected schemas:
- sources.csv: id, name, capacity, cost_per_unit
- users.csv: id, name, demand, priority
- connections.csv: source_id, user_id, max_flow, loss_rate
Ensure numeric columns parse as numbers (no stray text like “800 gallons”).

Status: INFEASIBLE

Check feasibility limits: for each user, the sum of max_flow * (1 - loss_rate) across incoming connections must be at least demand.
Check that total source capacity is sufficient (and remember that losses may require total flow to exceed total delivered demand).
If you increased loss_rate, you may need to increase max_flow or capacity to compensate.

Flow allocation table is empty

The output filters out very small flows with Connection.x_flow > 0.001. If your solution uses tiny flows, lower the threshold.
If the model is infeasible, no decision variables will have meaningful values; check the printed termination status first.

What this template is for

This template uses RelationalAI’s prescriptive reasoning (optimization) capabilities to choose flows on each connection that meet all user demands while minimizing total sourcing cost.

Who this is for

You want a small, end-to-end example of prescriptive reasoning (optimization) with RelationalAI.
You’re comfortable with basic Python and the idea of constraints + objectives.

What you’ll build

A semantic model of water Source, demand User, and Connection entities using concepts and typed properties.
A linear program (LP) that chooses a non-negative flow on each source-to-user connection.
Constraints for source capacity, per-connection max flow, and demand satisfaction after accounting for loss rates.
A solver run (HiGHS backend) that prints a readable table of non-trivial flow allocations.

What’s included

Model + solve script: water_allocation.py
Sample data: data/sources.csv, data/users.csv, data/connections.csv
Outputs: Printed termination status, objective value, and a flow allocation table

Prerequisites

Access

A Snowflake account that has the RAI Native App installed.
A Snowflake user with permissions to access the RAI Native App.

Tools

Python >= 3.10

Quickstart

Follow these steps to run the template with the included sample data.

Download the ZIP file for this template and extract it:
Terminal window
```
curl -O https://private.relational.ai/templates/zips/v0.14/water_allocation.zip
unzip water_allocation.zip
cd water_allocation
```
You can also download the template ZIP using the “Download ZIP” button at the top of this page.

Create and activate a virtual environment

python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip

Install dependencies

From this folder:
Terminal window
```
python -m pip install .
```
Configure Snowflake connection and RAI profile
Terminal window
```
rai init
```
Run the template
Terminal window
```
python water_allocation.py
```

Expected output

Status: OPTIMAL
Total cost: $874.28

Flow allocations:
      source        user       flow
 Reservoir_A   Municipal 317.934783
 Reservoir_A  Industrial 182.065217
 Reservoir_A Agricultural 500.000000
 Reservoir_B   Municipal 316.980758
 Reservoir_B  Industrial 250.000000
 Reservoir_B Agricultural  56.818182

Template structure

.
├─ README.md
├─ pyproject.toml
├─ water_allocation.py         # main runner / entrypoint
└─ data/                       # sample input data
   ├─ sources.csv
   ├─ users.csv
   └─ connections.csv

Start here: water_allocation.py

Sample data

Data files are in data/.

`sources.csv`

Defines water sources, each with a capacity (maximum total outflow) and a per-unit cost.

Column	Meaning
`id`	Unique source identifier
`name`	Source name (e.g., `Reservoir_A`)
`capacity`	Maximum total flow out of the source
`cost_per_unit`	Cost per unit of water sourced

`users.csv`

Defines demand points (users). The priority column is included in the sample data but is not used in the optimization model in this template.

Column	Meaning
`id`	Unique user identifier
`name`	User name (e.g., `Municipal`)
`demand`	Required delivered volume (after losses)
`priority`	Integer priority (not used in the model)

`connections.csv`

Defines feasible source-to-user connections with a maximum allowed flow and a fractional loss rate.

Column	Meaning
`source_id`	Foreign key to `sources.csv.id`
`user_id`	Foreign key to `users.csv.id`
`max_flow`	Upper bound on flow along this connection
`loss_rate`	Fraction of flow lost in transmission (e.g., `0.08` means 8% loss)

Model overview

The semantic model for this template is built around three concepts.

`Source`

A water supply point with limited capacity and a per-unit sourcing cost.

Property	Type	Identifying?	Notes
`id`	int	Yes	Loaded as the key from `data/sources.csv`
`name`	string	No	Human-readable label used in output
`capacity`	float	No	Upper bound on total outflow
`cost_per_unit`	float	No	Multiplies flow in the objective

`User`

A demand point that must receive enough effective inflow (flow adjusted for losses).

Property	Type	Identifying?	Notes
`id`	int	Yes	Loaded as the key from `data/users.csv`
`name`	string	No	Human-readable label used in output
`demand`	float	No	Required delivered volume
`priority`	int	No	Included in sample data; not used in this model

`Connection`

A feasible link from a Source to a User that can carry flow up to max_flow and loses a fraction loss_rate in transit.

Property	Type	Identifying?	Notes
`source`	`Source`	Part of compound key	Joined from `data/connections.csv.source_id`
`user`	`User`	Part of compound key	Joined from `data/connections.csv.user_id`
`max_flow`	float	No	Upper bound on the decision variable
`loss_rate`	float	No	Used to compute delivered inflow: `flow * (1 - loss_rate)`
`flow`	float	No	Continuous decision variable ()

How it works

This section walks through the highlights in water_allocation.py.

Import libraries and configure inputs

First, the script imports the Semantics APIs (Model, Relationship, data, where, require, sum, select) and configures the local DATA_DIR it will read CSVs from:

from pathlib import Path

import pandas
from pandas import read_csv

from relationalai.semantics import Model, Relationship, data, require, select, sum, where
from relationalai.semantics.reasoners.optimization import Solver, SolverModel

# --------------------------------------------------
# Configure inputs
# --------------------------------------------------

DATA_DIR = Path(__file__).parent / "data"

# Disable pandas inference of string types. This ensures that string columns
# in the CSVs are loaded as object dtype. This is only required when using
# relationalai versions prior to v1.0.
pandas.options.future.infer_string = False

Define concepts and load CSV data

Next, it declares Source and User concepts and loads sources.csv and users.csv into those concepts with data(...).into(...):

# Create a Semantics model container.
model = Model("water_allocation", config=globals().get("config", None))

# Source concept: water supply points with capacity and per-unit cost.
Source = model.Concept("Source")
Source.id = model.Property("{Source} has {id:int}")
Source.name = model.Property("{Source} has {name:string}")
Source.capacity = model.Property("{Source} has {capacity:float}")
Source.cost_per_unit = model.Property("{Source} has {cost_per_unit:float}")

# Load source data from CSV.
source_csv = read_csv(DATA_DIR / "sources.csv")
data(source_csv).into(Source, keys=["id"])

# User concept: demand points with required volume (and a priority field).
User = model.Concept("User")
User.id = model.Property("{User} has {id:int}")
User.name = model.Property("{User} has {name:string}")
User.demand = model.Property("{User} has {demand:float}")
User.priority = model.Property("{User} has {priority:int}")

# Load user data from CSV.
user_csv = read_csv(DATA_DIR / "users.csv")
data(user_csv).into(User, keys=["id"])

Then it declares a Connection concept and uses where(...).define(...) to join each connections.csv row to the corresponding Source and User:

# Connection concept: links a Source to a User with transmission parameters.
Connection = model.Concept("Connection")
Connection.source = model.Relationship("{Connection} from {source:Source}")
Connection.user = model.Relationship("{Connection} to {user:User}")
Connection.max_flow = model.Property("{Connection} has {max_flow:float}")
Connection.loss_rate = model.Property("{Connection} has {loss_rate:float}")
Connection.x_flow = model.Property("{Connection} has {flow:float}")

# Load connection data from CSV.
conn_data = data(read_csv(DATA_DIR / "connections.csv"))

# Define Connection entities by joining the CSV data with Source and User.
where(
    Source.id == conn_data.source_id,
    User.id == conn_data.user_id
).define(
    Connection.new(
        source=Source,
        user=User,
        max_flow=conn_data.max_flow,
        loss_rate=conn_data.loss_rate,
    )
)

Define decision variables, constraints, and objective

With the network data in place, the script creates a continuous SolverModel and declares Connection.x_flow as a non-negative decision variable with an upper bound from Connection.max_flow:

ConnectionRef = Connection.ref()

# Create a continuous optimization model.
s = SolverModel(model, "cont")

# Decision variable: flow on each connection (continuous, non-negative).
s.solve_for(
    Connection.x_flow,
    name=["flow", Connection.source.name, Connection.user.name],
    lower=0,
    upper=Connection.max_flow,
)

Next, it adds two constraints with require(...) and s.satisfy(...): source capacity (total outflow per source) and demand satisfaction (effective inflow per user, accounting for losses):

# Constraint: total outflow from each source must not exceed its capacity.
outflow = sum(ConnectionRef.x_flow).where(ConnectionRef.source == Source).per(Source)
source_limit = require(outflow <= Source.capacity)
s.satisfy(source_limit)

# Constraint: effective inflow to each user must meet demand (accounting for losses).
effective_inflow = (
    sum(ConnectionRef.x_flow * (1 - ConnectionRef.loss_rate)).where(ConnectionRef.user == User).per(User)
)
meet_demand = require(effective_inflow >= User.demand)
s.satisfy(meet_demand)

Finally, it minimizes total sourcing cost by summing flow times the per-unit cost of the corresponding source:

# Objective: minimize total cost.
total_cost = sum(Connection.x_flow * Connection.source.cost_per_unit)
s.minimize(total_cost)

Solve and print results

The script solves with the HiGHS backend, prints status and objective value, and then uses select(...).where(...).to_df() to display only non-trivial flows (> 0.001):

solver = Solver("highs")
s.solve(solver, time_limit_sec=60)

print(f"Status: {s.termination_status}")
print(f"Total cost: ${s.objective_value:.2f}")

allocations = select(
    Connection.source.name.alias("source"),
    Connection.user.name.alias("user"),
    Connection.x_flow,
).where(
    Connection.x_flow > 0.001
).to_df()

print("\nFlow allocations:")
print(allocations.to_string(index=False))

Customize this template

Use your own data

Replace the CSVs in data/ with your own, keeping the same column names (or update the loading logic in water_allocation.py).
Make sure your connections.csv only references valid source_id and user_id values that exist in sources.csv and users.csv.
Remember that users.csv.demand is the delivered requirement; because of losses, the required flow may be higher than demand.

Tune parameters

Adjust cost tradeoffs by changing sources.csv.cost_per_unit.
Model higher leakage/evaporation by changing connections.csv.loss_rate.
Tighten or relax infrastructure constraints by changing connections.csv.max_flow and sources.csv.capacity.

Extend the model

Add intermediate transfer nodes (source → junction → user) by introducing a node concept and splitting Connection into edge-to-edge connections.
Incorporate user priorities (currently loaded but unused) by changing the objective to penalize unmet demand with priority weights, or by adding staged solves.

Troubleshooting

rai init fails or the script cannot authenticate

Run rai init again and confirm you can see/select the right Snowflake account, role, warehouse, and the RAI Native App.
If you use multiple profiles, set RAI_PROFILE to the intended profile before running the script.
Ensure your Snowflake user has privileges to use the RAI Native App in the selected context.

ModuleNotFoundError for relationalai or pandas

Confirm your virtual environment is active (source .venv/bin/activate).
From the template folder, reinstall dependencies with python -m pip install ..
Verify you’re using Python 3.10+ with python —version.

Missing or incorrect CSV columns (for example: KeyError or AttributeError)

Verify the CSV headers match the expected schemas:
- sources.csv: id, name, capacity, cost_per_unit
- users.csv: id, name, demand, priority
- connections.csv: source_id, user_id, max_flow, loss_rate
Ensure numeric columns parse as numbers (no stray text like “800 gallons”).

Status: INFEASIBLE

Check feasibility limits: for each user, the sum of max_flow * (1 - loss_rate) across incoming connections must be at least demand.
Check that total source capacity is sufficient (and remember that losses may require total flow to exceed total delivered demand).
If you increased loss_rate, you may need to increase max_flow or capacity to compensate.

Flow allocation table is empty

The output filters out very small flows with Connection.x_flow > 0.001. If your solution uses tiny flows, lower the threshold.
If the model is infeasible, no decision variables will have meaningful values; check the printed termination status first.

What this template is for

This template uses prescriptive reasoning (optimization) to minimize the cost of distributing water from sources to users with nonlinear transmission losses.

Water utilities must distribute water from multiple sources (reservoirs, groundwater) to multiple user groups (municipal, industrial, agricultural). Each source has a limited capacity and a different extraction cost. Each connection in the distribution network has a maximum flow rate and a transmission loss rate that reduces the effective amount delivered.

This template uses prescriptive reasoning to find the minimum-cost allocation that satisfies every user’s demand. It models the distribution network as a flow problem with source capacity constraints, demand satisfaction constraints with nonlinear transmission losses, and connection flow limits.

The key feature is nonlinear loss modeling: transmission losses increase with utilization (effective delivery = flow * (1 - loss_rate * flow / max_flow)), creating a quadratic constraint that requires the Ipopt nonlinear solver. This is more realistic than constant-rate losses — at low flow the loss is small, but at capacity the full loss rate applies.

Who this is for

Water resource planners and utility operations analysts
Engineers modeling distribution networks with capacity and loss
Developers learning network flow optimization with RelationalAI

What you’ll build

A nonlinear optimization model for minimum-cost water distribution solved with Ipopt
Source capacity constraints limiting total outflow per source
Demand constraints with nonlinear (utilization-dependent) transmission losses
Flow upper bounds on individual connections

What’s included

water_allocation.py — Main script defining the network model, constraints, and solver call
data/sources.csv — Water sources with capacity and cost per unit
data/users.csv — User groups with demand and priority
data/connections.csv — Network connections with max flow and loss rate
pyproject.toml — Python package configuration

Prerequisites

Access

A Snowflake account that has the RAI Native App installed.
A Snowflake user with permissions to access the RAI Native App.

Tools

Python >= 3.10

Quickstart

Download ZIP:
Terminal window
```
curl -O https://private.relational.ai/templates/zips/v1/water_allocation.zip
unzip water_allocation.zip
cd water_allocation
```
You can also download the template ZIP using the “Download ZIP” button at the top of this page.

Create venv:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

Install:
Terminal window
```
python -m pip install .
```
Configure:
Terminal window
```
rai init
```
Run:
Terminal window
```
python water_allocation.py
```

Expected output:

Status: LOCALLY_SOLVED
Total cost: $853.39

Flow allocations:
  Reservoir_A  Agricultural  352.77
  Reservoir_A    Industrial  247.23
  Reservoir_A     Municipal  400.00
  Reservoir_B  Agricultural  182.06
  Reservoir_B    Industrial  177.93
  Reservoir_B     Municipal  228.99

Groundwater is not used — the solver routes all flow through the cheaper reservoirs. Reservoir A supplies its full municipal capacity (400 units) while splitting the remainder between agricultural and industrial users. Reservoir B covers the remaining demand.

Template structure

.
├── README.md
├── pyproject.toml
├── water_allocation.py
└── data/
    ├── connections.csv
    ├── sources.csv
    └── users.csv

How it works

1. Define sources, users, and connections

The model loads three concepts from CSV. Sources have capacity and cost. Users have demand and priority. Connections link sources to users with max flow and loss rate:

Source = Concept("Source", identify_by={"id": Integer})
Source.capacity = Property(f"{Source} has {Float:capacity}")
Source.cost_per_unit = Property(f"{Source} has {Float:cost_per_unit}")

User = Concept("User", identify_by={"id": Integer})
User.demand = Property(f"{User} has {Float:demand}")

Connection = Concept("Connection")
Connection.source = Property(f"{Connection} from {Source}", short_name="source")
Connection.user = Property(f"{Connection} to {User}", short_name="user")
Connection.max_flow = Property(f"{Connection} has {Float:max_flow}")
Connection.loss_rate = Property(f"{Connection} has {Float:loss_rate}")

2. Define the flow variable

Each connection gets a continuous flow variable bounded between zero and its maximum flow:

problem.solve_for(
    Connection.x_flow,
    name=["flow", Connection.source.name, Connection.user.name],
    lower=0,
    upper=Connection.max_flow
)

3. Add capacity and demand constraints

Source capacity limits total outflow. Demand constraints use nonlinear losses — loss increases with utilization, so effective delivery per connection is flow * (1 - loss_rate * flow / max_flow):

outflow = sum(ConnectionRef.x_flow).where(ConnectionRef.source == Source).per(Source)
problem.satisfy(model.require(outflow <= Source.capacity))

effective_inflow = sum(
    ConnectionRef.x_flow * (1 - ConnectionRef.loss_rate * ConnectionRef.x_flow / ConnectionRef.max_flow)
).where(ConnectionRef.user == User).per(User)
problem.satisfy(model.require(effective_inflow >= User.demand))

This quadratic constraint makes the problem nonlinear, requiring the Ipopt solver.

4. Minimize cost

The objective minimizes total extraction cost across all active flows:

total_cost = sum(Connection.x_flow * Connection.source.cost_per_unit)
problem.minimize(total_cost)

Customize this template

Add seasonal variation by introducing time periods with different source capacities and user demands.
Include priority-based allocation using the priority field to penalize unmet demand differently for each user group.
Add minimum flow requirements on certain connections to model contractual obligations.
Extend the network with intermediate nodes (pumping stations, treatment plants) that add processing costs or additional capacity constraints.

Troubleshooting

Status: INFEASIBLE

Total source capacity (after losses) is insufficient to meet all user demands. Check that the sum of source capacities minus worst-case losses covers total demand. You can increase source capacity in sources.csv, reduce demands in users.csv, or add new connections in connections.csv.

Some connections show zero flow

The solver avoids expensive routes when cheaper alternatives exist. If a source has a high cost per unit, its connections may carry zero flow. This is expected behavior for a cost-minimizing solution.

ModuleNotFoundError

Make sure you activated the virtual environment and ran python -m pip install . to install all dependencies listed in pyproject.toml.

Connection or authentication errors

Run rai init to configure your Snowflake connection. Verify that your account has the RAI Native App installed and that your user has the required permissions.