From b8fc1fdc4f2a7af69ee0c284e4effe33103eb099 Mon Sep 17 00:00:00 2001
From: marcbasquensmunoz <150438807+marcbasquensmunoz@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:38:18 +0200
Subject: [PATCH] Use graphs to model the topology of the network (#22)

* Create new network types

* Update plot_borefield

* Modify equations for graph implementation

* Add reverse graph

* Update examples

* Fix docs

* Update docs

* Fix docs

* Fix docs
---
 BNSPlots/Project.toml                         |   2 +
 BNSPlots/src/BNSPlots.jl                      |   1 +
 BNSPlots/src/borefield.jl                     |  43 +--
 BNSPlots/src/monitor.jl                       |  27 +-
 Project.toml                                  |   1 +
 docs/src/api.md                               |  49 +++-
 docs/src/nonhistory.jl                        |   7 +-
 docs/src/tutorial.jl                          |  13 +-
 examples/Braedstrup/main.jl                   |  31 +-
 examples/g-function/main.jl                   |   4 +-
 examples/tekniska/constant_m.jl               |   6 +-
 examples/tekniska/defs.jl                     |   4 +-
 examples/tekniska/prop_m.jl                   |   4 +-
 examples/tekniska/toggle.jl                   |   4 +-
 src/BoreholeNetworksSimulator.jl              |  10 +-
 .../borefields/EqualBoreholesBorefield.jl     |  19 +-
 src/modular/constraints/HeatLoadConstraint.jl |  13 +-
 .../constraints/InletTempConstraint.jl        |   5 +-
 .../constraints/TotalHeatLoadConstraint.jl    |  17 +-
 src/modular/core/core.jl                      |  37 +--
 src/modular/core/heat_balance.jl              |   4 +-
 src/modular/core/topology.jl                  |  24 +-
 src/modular/interfaces/Constraint.jl          |   2 +-
 src/modular/network/network.jl                | 273 ++++++++++++++++++
 src/modular/operators/SimpleOperator.jl       |  27 +-
 src/modular/simulate.jl                       |  21 +-
 test/api/api.jl                               |   2 +-
 .../test_EqualBoreholesBorefield.jl           |   9 +-
 test/constraints/test_HeatLoadConstraint.jl   |  29 +-
 test/constraints/test_InletTempConstraint.jl  |  31 +-
 .../test_TotalHeatLoadConstraint.jl           |  27 +-
 test/mocks/interfaces/constraint.jl           |   2 +-
 test/test_interfaces.jl                       |   2 +-
 33 files changed, 547 insertions(+), 203 deletions(-)
 create mode 100644 src/modular/network/network.jl

diff --git a/BNSPlots/Project.toml b/BNSPlots/Project.toml
index 95f56fe..7589db5 100644
--- a/BNSPlots/Project.toml
+++ b/BNSPlots/Project.toml
@@ -4,6 +4,8 @@ authors = ["Marc Basquens <marc5@kth.se> and contributors"]
 version = "1.0.0-DEV"
 
 [deps]
+GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 WGLMakie = "276b4fcb-3e11-5398-bf8b-a0c2d153d008"
 
 [compat]
diff --git a/BNSPlots/src/BNSPlots.jl b/BNSPlots/src/BNSPlots.jl
index 9a3593b..188be6d 100644
--- a/BNSPlots/src/BNSPlots.jl
+++ b/BNSPlots/src/BNSPlots.jl
@@ -1,6 +1,7 @@
 module BNSPlots
 
 using WGLMakie
+using GraphMakie
 
 include("results.jl")
 include("borefield.jl")
diff --git a/BNSPlots/src/borefield.jl b/BNSPlots/src/borefield.jl
index 65f6660..53b4fbf 100644
--- a/BNSPlots/src/borefield.jl
+++ b/BNSPlots/src/borefield.jl
@@ -1,8 +1,9 @@
+using Graphs
 
 make_color_range(color_pair, n) = n == 1 ? [color_pair[1]] : range(color_pair[1], stop=color_pair[2], length=n)
 
 """
-    plot_borefield(network, positions; distinguished_branches = [], colors = [])
+    plot_borefield(network, positions; distinguished_boreholes = [])
 
 Makes a plot of the borefield, showing the boreholes numbered and their connections.
 # Arguments
@@ -10,47 +11,33 @@ Makes a plot of the borefield, showing the boreholes numbered and their connecti
 - `positions`: The positions of each borehole.
 
 # Optional arguments
-- `distinguished_branches`: Vector of `Int`. If specified, the branches corresponding to the given values will be highlighted with the colors of `colors`.
-- `colors`: Vector of `Tuples` of `Color`. If specified, uses each pair of colors to define a color gradient to color each of the branches of `distinguished_branches`.
+- `distinguished_boreholes`: Vector of `Tuple{Int, Color}`. If specified, the boreholes corresponding to the given values will be highlighted with each of the colors provided.
 """
-function plot_borefield(network, positions; distinguished_branches = [], colors = [])
-
+function plot_borefield(network, positions; distinguished_boreholes = [])
     scene = Figure()
     axis = scene[1, 1] = Axis(scene, ylabel = "y [m]", xlabel = "x[m]", aspect = DataAspect())
     min_x = minimum(map(x->x[1], positions))
     max_x = maximum(map(x->x[1], positions))
     min_y = minimum(map(x->x[2], positions))
     max_y = maximum(map(x->x[2], positions))
-    
+
     margin = max(max_x-min_x, max_y-min_y) * 0.1
 
     xlims!(axis, min_x-margin, max_x+margin)
     ylims!(axis, min_y-margin, max_y+margin)
 
-    scatter!(axis,  Makie.Point2.(positions), markersize = 15.)
-
-    # Draw boreholes
-    for (i, n_branch) in enumerate(distinguished_branches)
-        branch = network.branches[n_branch]
-        branch_colors = make_color_range(colors[i], length(branch))
-        for (borehole, color) in zip(branch, branch_colors)
-            point =  [ Makie.Point(positions[borehole]) ]
-            scatter!(axis, point, color = color, markersize = 24)
-        end
-    end
-
-    # Write labels with borehole numbers
-    for i in eachindex(1:length(positions))
-        text!(axis, "$i", fontsize = 9, position  = (positions[i] .+  (0.7 , -0.7)) , color = :blue)
-    end
+    borefield = SimpleGraph(network.graph)
+    rem_vertex!(borefield, nv(borefield))
+    rem_vertex!(borefield, nv(borefield))
 
-    # Draw lines representing connections
-    for branch in network.branches
-        for k in 2:length(branch)
-            s = Makie.Point.([ positions[branch][k-1], positions[branch][k] ] )
-            linesegments!(axis, s, color = :red, linewidth = 1.5)
-        end
+    Nb = nv(borefield)
+    borehole_colors = [colorant"black" for i in 1:Nb]
+    borehole_sizes = 12 * ones(Int, Nb)
+    for (bh, color) in distinguished_boreholes
+        borehole_colors[bh] = color
+        borehole_sizes[bh] = 16
     end
+    graphplot!(axis, borefield, layout=positions, nlabels=["$i" for i in 1:Nb], node_color = borehole_colors, node_size = borehole_sizes, nlabels_distance = 5.)
 
     hidespines!(axis)
 
diff --git a/BNSPlots/src/monitor.jl b/BNSPlots/src/monitor.jl
index 3cf53f9..8d9639b 100644
--- a/BNSPlots/src/monitor.jl
+++ b/BNSPlots/src/monitor.jl
@@ -14,7 +14,7 @@ end
 Creates a plot of the result of the simulation.
 # Arguments
 - `containers`: The containers (`SimulationContainers`) containing the result of the simulation through `simulate!`.
-- `branch`: A vector containing the IDs of the boreholes whose data will be displayed. 
+- `boreholes`: A vector containing the IDs of the boreholes whose data will be displayed. 
 - `t`: The times at which the data corresponds. It should normally be `options.t`.
 
 # Optional arguments
@@ -22,9 +22,9 @@ Creates a plot of the result of the simulation.
 - `display`: A vector describing which plots that will be generated. If `:Tfin`, `:Tfout` or `:Tb` are specified, a temperature plot will be created showing the inlet fluid temperature, 
     the outlet fluid temperature, and the borehole wall temperature, respectively. If `:q` is specified, a separate power plot will be created shwoing the heat extracted per meter.
 - `Δt`: The scale of the x-axis in the plot. Possible options: `:year`, `:month`, `:hour`.
-- `color_pair`: A pair of colors used as extrema to generate a range of colors for each borehole. 
+- `colors`: A list of colors used for each borehole. If not specified, the colors used will be between colorant"navajowhite2" and colorant"darkgreen".
 """
-function monitor(containers, branch, t; steps = 1:length(t), display = [:Tfin, :Tfout, :Tb, :q, :mf], Δt = :year, color_pair = (colorant"navajowhite2", colorant"darkgreen"))
+function monitor(containers, boreholes, t; steps = 1:length(t), display = [:Tfin, :Tfout, :Tb, :q, :mf], Δt = :year, colors = [])
     if isempty(display)
         return
     end
@@ -36,8 +36,12 @@ function monitor(containers, branch, t; steps = 1:length(t), display = [:Tfin, :
     grid = scene[1, 1] = GridLayout()
     axes = []
 
-    color_range = make_color_range(color_pair, length(branch)) 
-    
+    if !isempty(colors)
+        color_range = colors
+    else
+        color_range = make_color_range((colorant"navajowhite2", colorant"darkgreen"), length(boreholes)) 
+    end
+
     if anyin([:Tfin, :Tfout, :Tb], display)
         axis_T = Axis(grid[length(axes)+1, 1], ylabel = L"T \, \left[ °C \right]")
         push!(axes, axis_T)
@@ -51,11 +55,11 @@ function monitor(containers, branch, t; steps = 1:length(t), display = [:Tfin, :
         push!(axes, axis_m)
     end
 
-    Tfin = get_Tfin(containers, branch)  
-    Tfout = get_Tfout(containers, branch)  
-    Tb = get_Tb(containers, branch)    
-    q = get_q(containers, branch)
-    mf = get_mf(containers, branch)
+    Tfin = get_Tfin(containers, boreholes)  
+    Tfout = get_Tfout(containers, boreholes)  
+    Tb = get_Tb(containers, boreholes)    
+    q = get_q(containers, boreholes)
+    mf = get_mf(containers, boreholes)
 
     secs_in_year = 8760*3600
     conversion = Dict(:year => 1, :month => 12, :hour => 8760)
@@ -79,10 +83,9 @@ function monitor(containers, branch, t; steps = 1:length(t), display = [:Tfin, :
         end
     end
 
-
     group_color = [PolyElement(color = color, strokecolor = :transparent) for color in color_range]
     group_marker = [LineElement(color = :black), LineElement(color = :black, linestyle = :dash), MarkerElement(marker = :circle, color = :black, strokecolor = :transparent, markersize = 5.)]
-    legend = Legend(scene, [group_color, group_marker], [string.(branch), ["Tfin", "Tfout", "Tb"]], ["Boreholes", "Temperatures"], tellheight = true, tellwidth = false)
+    legend = Legend(scene, [group_color, group_marker], [string.(boreholes), ["Tfin", "Tfout", "Tb"]], ["Boreholes", "Temperatures"], tellheight = true, tellwidth = false)
     legend.titleposition = :top
     legend.orientation = :horizontal
     legend.nbanks = 1
diff --git a/Project.toml b/Project.toml
index 842a42b..3c40715 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,6 +14,7 @@ ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 FiniteLineSource = "64cfb1b1-06f0-48af-b80e-43c410dfe9e8"
 GeometryTypes = "4d00f742-c7ba-57c2-abde-4428a4b178cb"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LegendrePolynomials = "3db4a2ba-fc88-11e8-3e01-49c72059a882"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/docs/src/api.md b/docs/src/api.md
index 3dc1ab2..fbb3392 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -20,10 +20,57 @@ See [Basic tutorial](@ref) for more details.
 BoreholeNetwork
 ```
 
+Some relevant network related functions:
+```@docs
+boreholes_in_branch
+```
+
+```@docs
+first_bhs_in_branch
+```
+
+```@docs
+source
+```
+```@docs
+sink
+```
+```@docs
+connect!
+```
+```@docs
+connect_to_source!
+```
+```@docs
+connect_to_sink!
+```
+```@docs
+connect_in_series!
+```
+```@docs
+connect_in_parallel!
+```
+
 ```@docs
 BoreholeOperation
 ```
 
+```@docs
+Valve
+```
+
+Valve creation:
+```@docs
+valve
+```
+```@docs
+equal_valve
+```
+```@docs
+absolute_valve
+```
+
+
 ```@docs
 Operator
 ```
@@ -31,7 +78,7 @@ Operator
 ### Prewritten operator strategies
 
 ```@docs
-SimpleOperator
+ConstantOperator
 ```
 
 ## Simulation Options
diff --git a/docs/src/nonhistory.jl b/docs/src/nonhistory.jl
index 3023614..72bd8eb 100644
--- a/docs/src/nonhistory.jl
+++ b/docs/src/nonhistory.jl
@@ -15,7 +15,7 @@ using BoreholeNetworksSimulator
 #
 
 Δt = 3600.
-Nt = 2#8760
+Nt = 8760
 
 medium = GroundMedium(α=1e-6, λ=3., T0=10.)
 borehole = SingleUPipeBorehole(H=100., D=10.)
@@ -24,8 +24,9 @@ borefield = EqualBoreholesBorefield(borehole_prototype=borehole, positions=posit
 constraint = constant_HeatLoadConstraint(5 .* ones(BoreholeNetworksSimulator.n_boreholes(borefield)), Nt)
 fluid = Water()
 
-configurations = [BoreholeNetwork([[1], [2]])]
-operator = SimpleOperator(mass_flow = 2., branches = 2)
+network = all_parallel_network(2)
+configurations = [network]
+operator = ConstantOperator(network, mass_flows = 2 * ones(2))
 
 # Now, we define two different options using different `method` parameters, 
 # one with `ConvolutionMethod` corresponding to the convolution,
diff --git a/docs/src/tutorial.jl b/docs/src/tutorial.jl
index aaea2ac..8d59208 100644
--- a/docs/src/tutorial.jl
+++ b/docs/src/tutorial.jl
@@ -72,7 +72,7 @@ borefield = EqualBoreholesBorefield(borehole_prototype=borehole, positions=posit
 # In our example, we want to simulate two independent boreholes, so each of them must be in a separate branch.
 # Also, for the moment, we are only interested in this configuration, so we define:
 
-configurations = [BoreholeNetwork([[1], [2]])]
+network = all_parallel_network(2)
 
 # Even with all these specifications, the evolution of the system is still not fully determined.
 # The missing conditions are referred to as constraints, and are modeled by subtypes of [`Constraint`](@ref).
@@ -80,8 +80,8 @@ configurations = [BoreholeNetwork([[1], [2]])]
 # impose that their inlet temperatures be equal. In our example, since we want out boreholes to be independent, 
 # we will impose the total amount of heat that we want to extract from each borehole. We will impose a constant
 # load, equal for both boreholes. This is specified by
-q1 = 5.
-q2 = 5.
+q1 = H
+q2 = H
 constraint = constant_HeatLoadConstraint([q1, q2], Nt)
 
 # We can finally create the object with all the options:
@@ -94,7 +94,7 @@ options = SimulationOptions(
     fluid = Water(),
     Δt = Δt,
     Nt = Nt,
-    configurations = configurations
+    configurations = [network]
 )
 
 # As we have mentioned, the simulation is designed to allow for a controllable opeartion during its duration.
@@ -105,9 +105,10 @@ options = SimulationOptions(
 # configuration will be used for the next time step. In our case, we only want a static, simple configuration.
 # The second specifies the mass flow rate through each branch of the selected configuration, provided as 
 # a vector. In our example, we will keep this constant through the simulation.
-# For this purpose, there exists the type [`SimpleOperator`](@ref), that implements precisely this strategy.
+# For this purpose, there exists the type [`ConstantOperator`](@ref), that implements precisely this strategy.
 
-operator = SimpleOperator(mass_flow = 2., branches = 2)
+# TODO: Update tutorial with new operation and network
+operator = ConstantOperator(network, mass_flows = 2 * ones(2))
 
 # Before simulating, we first need to call [`initialize`](@ref) to run some precomputations
 # that will be used throught the simulation and to instantiate containers where the result will be written.
diff --git a/examples/Braedstrup/main.jl b/examples/Braedstrup/main.jl
index c796cf6..2a18513 100644
--- a/examples/Braedstrup/main.jl
+++ b/examples/Braedstrup/main.jl
@@ -15,16 +15,18 @@ borehole_locations = "$(@__DIR__)/data/Braedstrup_borehole_coordinates.txt"
 Δt = 8760*3600/12.
 Nt = 120
 
-network = BoreholeNetwork([
-    [35, 36, 29, 37, 30, 22],
-    [43, 48, 42, 41, 40, 34],
-    [47, 46, 45, 39, 32, 33],
-    [44, 38, 31, 24, 25, 26],
-    [18, 12, 11, 17, 16, 23],
-    [19, 13, 7, 6, 5, 10],
-    [20, 14, 8, 3, 2, 1],
-    [27, 28, 21, 15, 9, 4]
-])
+network = BoreholeNetwork(48)
+connect_to_source!(network, [35, 43, 47, 44, 18, 19, 20, 27])
+connect_in_series!(network, [35, 36, 29, 37, 30, 22])
+connect_in_series!(network, [43, 48, 42, 41, 40, 34])
+connect_in_series!(network, [47, 46, 45, 39, 32, 33])
+connect_in_series!(network, [44, 38, 31, 24, 25, 26])
+connect_in_series!(network, [18, 12, 11, 17, 16, 23])
+connect_in_series!(network, [19, 13, 7, 6, 5, 10])
+connect_in_series!(network, [20, 14, 8, 3, 2, 1])
+connect_in_series!(network, [27, 28, 21, 15, 9, 4])
+connect_to_sink!(network, [22, 34, 33, 26, 23, 10, 1, 4])
+
 configurations = [  
     network,            # Heat extraction
     reverse(network)    # Heat injection
@@ -60,7 +62,7 @@ end
 
 function BoreholeNetworksSimulator.operate(operator::SeasonalOperator, i, options, X)
     active_network = options.configurations[operator.seasonal_configuration[i]]
-    BoreholeOperation(active_network, operator.mass_flows)
+    BoreholeOperation(network=active_network, mass_flows=operator.mass_flows)
 end
 
 operator = SeasonalOperator(mass_flows=0.5 .* ones(n_branches(network)), seasonal_configuration=[i%12 in 1:6 ? 2 : 1 for i in 1:Nt])
@@ -74,11 +76,10 @@ containers.X
 ############
 # Draw plots
 
-monitored_branches = [3, 8]
-color_ranges = [(colorant"darkorange", colorant"blue"), (colorant"red", colorant"green")]
+monitoring = [(27, colorant"darkorange"), (28, colorant"green")]
 
-borefiled_plot = plot_borefield(network, borehole_positions, distinguished_branches = monitored_branches, colors = color_ranges)
-branch1 = monitor(containers, network.branches[monitored_branches[1]], options.t, display = [:Tfin], color_pair=color_ranges[1])
+borefield_plot = plot_borefield(network, borehole_positions, distinguished_boreholes = monitoring)
+branch1 = monitor(containers, boreholes_in_branch(network, first_bh=19), options.t, display = [:Tfin])
 
 # save("examples/Braedstrup/plots/Braedstrup_borefield.png", borefiled_plot)
 # save("examples/Braedstrup/plots/branch1.png", branch1)
diff --git a/examples/g-function/main.jl b/examples/g-function/main.jl
index ba5d0c6..c71bd05 100644
--- a/examples/g-function/main.jl
+++ b/examples/g-function/main.jl
@@ -15,7 +15,7 @@ function make_plot(axis, d)
     α = 1e-6
     λ = 3.
 
-    network = BoreholeNetwork([[i] for i in 1:n*m])
+    network = all_parallel_network(n*m)
     configurations = [network]
 
     function create_rectangular_field(n, m, d)
@@ -41,7 +41,7 @@ function make_plot(axis, d)
         configurations = configurations
     )
 
-    operator = SimpleOperator(mass_flow = 1., branches =  n_branches(network))
+    operator = ConstantOperator(network, mass_flows = ones(n*m))
     containers = @time initialize(options)
 
     @time simulate!(operator=operator, options=options, containers=containers)
diff --git a/examples/tekniska/constant_m.jl b/examples/tekniska/constant_m.jl
index e40b241..cd37346 100644
--- a/examples/tekniska/constant_m.jl
+++ b/examples/tekniska/constant_m.jl
@@ -3,13 +3,13 @@ using BNSPlots
 
 include("defs.jl")
 
-operator = SimpleOperator(mass_flow=0.5, branches=n_branches(network))
+operator = ConstantOperator(network, mass_flows=0.5 * ones(10))
 containers = @time initialize(options)
 @time simulate!(operator=operator, options=options, containers=containers)
 
 t_range = (5*8760-24*7):5*8760
-const_m_plot = monitor(containers, [4, 7], options.t, steps = t_range,  color_pair = (colorant"darkgreen", colorant"red")) 
+const_m_plot = monitor(containers, [4, 7], options.t, steps = t_range, colors = [colorant"darkgreen", colorant"red"]) 
 # save("examples/tekniska/plots/const_m.png", const_m_plot)
 
-const_m_plot_5_year = monitor(containers, [4, 7], options.t, color_pair = (colorant"darkgreen", colorant"red")) 
+const_m_plot_5_year = monitor(containers, [4, 7], options.t, colors = [colorant"darkgreen", colorant"red"]) 
 # save("examples/tekniska/plots/const_m_5_years.png", const_m_plot_5_year)
\ No newline at end of file
diff --git a/examples/tekniska/defs.jl b/examples/tekniska/defs.jl
index ab4c30b..caf77e0 100644
--- a/examples/tekniska/defs.jl
+++ b/examples/tekniska/defs.jl
@@ -6,7 +6,7 @@ using WGLMakie
 Δt = 3600.
 Nt = 8760*5
 
-network = BoreholeNetwork([[i] for i in 1:10])
+network = all_parallel_network(10)
 configurations = [  
     network
 ]
@@ -60,7 +60,7 @@ options = SimulationOptions(
 )
 
 group1 = [1, 6, 2, 7, 3, 8]
-borefield_fig = plot_borefield(network, borehole_positions, distinguished_branches=collect(1:10), colors=[i in group1 ? (colorant"red", colorant"red") : (colorant"green", colorant"green") for i in 1:10])
+borefield_fig = plot_borefield(network, borehole_positions, distinguished_boreholes=[(i, i in group1 ? colorant"red" : colorant"green") for i in 1:10])
 
 function plot_weekly_Q()
     fig = Figure(size=(1500, 400))
diff --git a/examples/tekniska/prop_m.jl b/examples/tekniska/prop_m.jl
index 5ca9bc8..f1bdcc6 100644
--- a/examples/tekniska/prop_m.jl
+++ b/examples/tekniska/prop_m.jl
@@ -10,7 +10,7 @@ end
 
 function BoreholeNetworksSimulator.operate(operator::VariableMFOperator, step, options, X)
     operator.mass_flows .= operator.mass_flow_series[step]
-    BoreholeOperation(options.configurations[1], operator.mass_flows)
+    BoreholeOperation(network = options.configurations[1], mass_flows = operator.mass_flows)
 end
 
 operator = VariableMFOperator(0.5 .* Q_tot ./ Q_ref, zeros(n_branches(network)))
@@ -19,6 +19,6 @@ containers = @time initialize(options)
 @time simulate!(operator=operator, options=options, containers=containers)
 
 t_range = (5*8760-24*7):5*8760
-prop_m_plot = monitor(containers, [4, 7], options.t, steps = t_range, color_pair = (colorant"darkgreen", colorant"red")) 
+prop_m_plot = monitor(containers, [4, 7], options.t, steps = t_range, colors = [colorant"darkgreen", colorant"red"]) 
 
 # save("examples/tekniska/plots/prop_m.png", prop_m_plot)
diff --git a/examples/tekniska/toggle.jl b/examples/tekniska/toggle.jl
index 0623b55..fe7da72 100644
--- a/examples/tekniska/toggle.jl
+++ b/examples/tekniska/toggle.jl
@@ -39,7 +39,7 @@ function BoreholeNetworksSimulator.operate(operator::ToggleOperator, step, optio
     mass_flow_containers[branch2] .= mf2
 
     @pack! operator = single_branch, hours_used, toggle, mass_flow_containers
-    BoreholeOperation(options.configurations[1], mass_flow_containers)
+    BoreholeOperation(network = options.configurations[1], mass_flows = mass_flow_containers)
 end
 
 operator = ToggleOperator(Q_threshold = 20000., mass_flow = 0.5, mass_flow_containers = zeros(n_branches(network)))
@@ -47,6 +47,6 @@ containers = @time initialize(options)
 @time simulate!(operator=operator, options=options, containers=containers)
 
 t_range = (5*8760-24*7):5*8760
-toggle_plot = monitor(containers, [4, 7], options.t, steps = t_range, color_pair = (colorant"darkgreen", colorant"red")) 
+toggle_plot = monitor(containers, [4, 7], options.t, steps = t_range, colors = [colorant"darkgreen", colorant"red"]) 
 
 # save("examples/tekniska/plots/toggle.png", toggle_plot)
\ No newline at end of file
diff --git a/src/BoreholeNetworksSimulator.jl b/src/BoreholeNetworksSimulator.jl
index fe64023..edf75f8 100644
--- a/src/BoreholeNetworksSimulator.jl
+++ b/src/BoreholeNetworksSimulator.jl
@@ -17,13 +17,14 @@ using FastGaussQuadrature
 using LegendrePolynomials
 using SpecialFunctions
 using RequiredInterfaces
+using Graphs
 
 using FiniteLineSource
 
 include("utils.jl")
 
 modular = get_all_julia_files_in_dir(joinpath(@__DIR__, "modular"))
-sort_dependencies!(modular, ["interfaces", "core"])
+sort_dependencies!(modular, ["interfaces", "network", "core"])
 include.(modular)
 include("CoolProp/load_properties.jl")
 include("CoolProp/interpolate.jl")
@@ -36,8 +37,11 @@ export Constraint, TotalHeatLoadConstraint, HeatLoadConstraint, InletTempConstra
 export TimeSuperpositionMethod, ConvolutionMethod, NonHistoryMethod
 export BoundaryCondition, NoBoundary, DirichletBoundaryCondition, AdiabaticBoundaryCondition
 export Approximation, MidPointApproximation, MeanApproximation
-export SimulationOptions, SimulationContainers, BoreholeNetwork
-export BoreholeOperation, Operator, SimpleOperator, operate
+export SimulationOptions, SimulationContainers
+export BoreholeNetwork, Valve, BoreholeOperation
+export connect_to_source!, connect_to_sink!, connect!, all_series_network, all_parallel_network, compute_mass_flows!, source, sink, connect_in_series!, equal_valve, valve, absolute_valve, initialize_mass_flows, boreholes_in_branch, first_bhs_in_branch, connect_in_parallel!
+export Operator, SimpleOperator, operate
+export ConstantOperator
 export simulate!, compute_parameters, load_cache!, save_cache, initialize
 export n_branches
 export ThermophysicalProperties
diff --git a/src/modular/borefields/EqualBoreholesBorefield.jl b/src/modular/borefields/EqualBoreholesBorefield.jl
index dd74451..048ee35 100644
--- a/src/modular/borefields/EqualBoreholesBorefield.jl
+++ b/src/modular/borefields/EqualBoreholesBorefield.jl
@@ -33,22 +33,21 @@ function segment_coordinates(bf::EqualBoreholesBorefield, segment)
     (position[1], position[2], z_ref, h)
 end
 
-function internal_model_coeffs!(M, borefield::EqualBoreholesBorefield, medium, operation, T_fluid, fluid)
+function internal_model_coeffs!(M, borefield::EqualBoreholesBorefield, medium, mass_flows, T_fluid, fluid)
     Nb = n_boreholes(borefield)
     λ = get_λ(medium)
 
-    for (i, branch) in enumerate(operation.network.branches)
-        mass_flow = operation.mass_flows[i]
+    # TODO: Optimize for boreholes with same mass_flow
+    for i in 1:Nb
+        mass_flow = mass_flows[i]
 
-        for j in branch
-            Tref = (T_fluid[2*j - 1] + T_fluid[2*j]) / 2
+        Tref = (T_fluid[2i - 1] + T_fluid[2i]) / 2
 
-            k_in, k_out, k_b = uniform_Tb_coeffs(borefield.borehole_prototype, λ, mass_flow, Tref, fluid)
+        k_in, k_out, k_b = uniform_Tb_coeffs(borefield.borehole_prototype, λ, mass_flow, Tref, fluid)
 
-            M[j, j*2 - 1]  = k_in
-            M[j, j*2]      = k_out
-            M[j, Nb*2 + j] = k_b    
-        end
+        M[i, 2i - 1]  = k_in
+        M[i, 2i]      = k_out
+        M[i, 2Nb + i] = k_b    
     end
 end
 
diff --git a/src/modular/constraints/HeatLoadConstraint.jl b/src/modular/constraints/HeatLoadConstraint.jl
index ea7f194..478ff0d 100644
--- a/src/modular/constraints/HeatLoadConstraint.jl
+++ b/src/modular/constraints/HeatLoadConstraint.jl
@@ -42,13 +42,16 @@ function uniform_HeatLoadConstraint(Q_tot::Vector{T}, Nbr) where {T <: Number}
     HeatLoadConstraint(Q)
 end
 
-function constraints_coeffs!(M, ::HeatLoadConstraint, operation, borefield)
+function constraints_coeffs!(M, ::HeatLoadConstraint, borefield::Borefield, network, mass_flows)
     M .= zero(eltype(M))
 
-    for (i, branch) in enumerate(operation.network.branches)
-        Nb = n_boreholes(operation.network)
-        for j in branch
-            M[i, 3Nb + j] = one(eltype(M)) * get_H(borefield, j)
+    Nb = n_boreholes(network)
+    for (i, first_bh) in enumerate(first_bhs_in_branch(network))
+        for bh in neighborhood(network.graph, first_bh, Nb)
+            if bh == source(network) || bh == sink(network) 
+                continue
+            end
+            M[i, 3Nb + bh] = one(eltype(M)) * get_H(borefield, bh)
         end
     end
 end
diff --git a/src/modular/constraints/InletTempConstraint.jl b/src/modular/constraints/InletTempConstraint.jl
index ca7dabc..fdd2bbd 100644
--- a/src/modular/constraints/InletTempConstraint.jl
+++ b/src/modular/constraints/InletTempConstraint.jl
@@ -41,9 +41,10 @@ function uniform_InletTempConstraint(T_in::Vector{N}, Nbr) where {N <: Number}
     InletTempConstraint(T)
 end
 
-function constraints_coeffs!(M, ::InletTempConstraint, operation, borefield)
+function constraints_coeffs!(M, ::InletTempConstraint, borefield::Borefield, network, mass_flows)
     M .= zero(eltype(M))
-    for (i, borehole) in enumerate(first_boreholes(operation.network))
+
+    for (i, borehole) in enumerate(first_bhs_in_branch(network))
         M[i, 2*borehole - 1] = one(eltype(M))
     end
 end
diff --git a/src/modular/constraints/TotalHeatLoadConstraint.jl b/src/modular/constraints/TotalHeatLoadConstraint.jl
index c7b016d..8eb90d8 100644
--- a/src/modular/constraints/TotalHeatLoadConstraint.jl
+++ b/src/modular/constraints/TotalHeatLoadConstraint.jl
@@ -10,19 +10,20 @@ struct TotalHeatLoadConstraint{T <: Number} <: Constraint
 end
 
 
-function constraints_coeffs!(M, ::TotalHeatLoadConstraint, operation, borefield)
+function constraints_coeffs!(M, ::TotalHeatLoadConstraint, borefield::Borefield, network, mass_flows)
     M .= zero(eltype(M))
 
-    Nb = n_boreholes(operation.network)
+    Nb = n_boreholes(network)
 
-    first_bh = operation.network.branches[1][1]
-    for i in 1:(length(operation.network.branches)-1)
-        M[i+1, 2*first_bh-1] = -1.
-        bh_in = 2*operation.network.branches[i+1][1]-1
-        M[i+1, bh_in] = 1.
+    first_bhs = first_bhs_in_branch(network)
+    first_bh = first_bhs[1]
+    for i in 2:n_branches(network)
+        M[i, 2*first_bh-1] = -1.
+        bh_in = 2*first_bhs[i]-1
+        M[i, bh_in] = 1.
     end
 
-    if sum(operation.mass_flows) == 0
+    if sum(mass_flows) == 0
         M[1, 2*first_bh-1] = 1.
         M[1, 2Nb+first_bh] = -1.
         return
diff --git a/src/modular/core/core.jl b/src/modular/core/core.jl
index 71ad33a..79794ea 100644
--- a/src/modular/core/core.jl
+++ b/src/modular/core/core.jl
@@ -1,35 +1,4 @@
 
-"""
-    BoreholeNetwork(branches::Vector{Vector{Int}})
-
-Representation of the hydraulic connections of the boreholes in the network.
-Each element in `branches` should be a vector representing a branch of boreholes connected in series, specified by their identifiers.
-The first borehole of each branch is assumed to be connected in parallel. 
-"""
-struct BoreholeNetwork
-    branches::Vector{Vector{Int}}
-    n_boreholes::Int
-end
-BoreholeNetwork(branches::Vector{Vector{Int}}) = BoreholeNetwork(branches, sum([length(branch) for branch in branches]))
-Base.reverse(network::BoreholeNetwork) = BoreholeNetwork(map(branch -> Base.reverse(branch), network.branches))
-n_branches(network::BoreholeNetwork) = length(network.branches)
-n_boreholes(network::BoreholeNetwork) = network.n_boreholes
-first_boreholes(network::BoreholeNetwork) = map(first, network.branches)
-
-"""
-    BoreholeOperation{Arr <: AbstractArray}(
-        network::BoreholeNetwork         
-        mass_flows::Arr
-    )
-
-Represents a operation state of the network, with `network` representing the hydraulic configuration and `mass_flows` a `Vector` containing the mass flow rate of each branch.
-"""
-@with_kw struct BoreholeOperation{Arr <: AbstractArray}
-    network::BoreholeNetwork         
-    mass_flows::Arr
-end
-BoreholeOperation(::Nothing) = BoreholeOperation(BoreholeNetwork([], 0), @view ones(1)[1:1])
-
 """
     struct SimulationOptions{
                     N <: Number,
@@ -135,11 +104,7 @@ function initialize(options::SimulationOptions)
 end
 function SimulationContainers(options::SimulationOptions) 
     @unpack Nb, Nt = options
-    if Nb > 50
-        SimulationContainers(M = spzeros(4Nb, 4Nb), b = zeros(4Nb), X = zeros(4Nb, Nt), mf = zeros(Nb, Nt))
-    else 
-        SimulationContainers(M = zeros(4Nb, 4Nb), b = zeros(4Nb), X = zeros(4Nb, Nt), mf = zeros(Nb, Nt))
-    end
+    SimulationContainers(M =  Nb > 100 ? spzeros(4Nb, 4Nb) : zeros(4Nb, 4Nb), b = zeros(4Nb), X = zeros(4Nb, Nt), mf = zeros(Nb, Nt))
 end
 
 function branch_of_borehole(operation::BoreholeOperation, borehole)
diff --git a/src/modular/core/heat_balance.jl b/src/modular/core/heat_balance.jl
index a372c44..36593b7 100644
--- a/src/modular/core/heat_balance.jl
+++ b/src/modular/core/heat_balance.jl
@@ -1,10 +1,10 @@
 
-function heat_balance_coeffs!(M, borefield::Borefield, operation::BoreholeOperation, fluid::Fluid)
+function heat_balance_coeffs!(M, borefield::Borefield, mass_flows, fluid::Fluid)
     Nb = n_boreholes(borefield)
     Ns = n_segments(borefield)
 
     for i in 1:Nb
-        heat = cpf(fluid) * operation.mass_flows[branch_of_borehole(operation, i)] 
+        heat = cpf(fluid) * mass_flows[i] 
         M[i, i*2-1] = heat
         M[i, i*2]  = -heat
     end
diff --git a/src/modular/core/topology.jl b/src/modular/core/topology.jl
index aef83a3..7299304 100644
--- a/src/modular/core/topology.jl
+++ b/src/modular/core/topology.jl
@@ -1,13 +1,23 @@
 
-function topology_coeffs!(M, operation::BoreholeOperation)
+function topology_coeffs!(M, network::BoreholeNetwork, mass_flows)
     M .= zero(eltype(M))
     j = 1
-    for branch in operation.network.branches
-        for i in eachindex(@view branch[1:end-1])
-            in = branch[i+1] 
-            out = branch[i] 
-            M[j, 2*in-1] = 1.
-            M[j, 2*out] = -1.
+    Nb = n_boreholes(network)
+
+    for bh_in in 1:Nb
+        if bh_in == sink(network)
+            continue
+        end
+        parents = inneighbors(network.graph, bh_in)
+        filter!(i -> i ≠ source(network), parents)
+        for bh_out in parents
+            if bh_out == source(network)
+                continue
+            end
+            M[j, 2*bh_in-1] = 1.
+            M[j, 2*bh_out] = - mass_flows[bh_out] / mass_flows[bh_in]
+        end
+        if !isempty(parents)
             j += 1
         end
     end
diff --git a/src/modular/interfaces/Constraint.jl b/src/modular/interfaces/Constraint.jl
index 8a2d9ed..848b7a4 100644
--- a/src/modular/interfaces/Constraint.jl
+++ b/src/modular/interfaces/Constraint.jl
@@ -16,6 +16,6 @@ Required functions:
 abstract type Constraint end
 
 @required Constraint begin
-    constraints_coeffs!(M, ::Constraint, operation, borefield)
+    constraints_coeffs!(M, ::Constraint, ::Borefield, network, mass_flows)
     constraints_b!(b, ::Constraint, operation, step)
 end
diff --git a/src/modular/network/network.jl b/src/modular/network/network.jl
new file mode 100644
index 0000000..94a9f27
--- /dev/null
+++ b/src/modular/network/network.jl
@@ -0,0 +1,273 @@
+
+"""
+    BoreholeNetwork(
+        graph::SimpleDiGraph = SimpleDiGraph()
+    )
+
+Represents the network of boreholes with a directed graph `graph`, whose edges represent the flow of fluid.
+The `graph` is expected to have number of vertices `Nb+2`, where `Nb` is the number of boreholes in the borefield.
+Vertexs number `Nb+1` and `Nb+2` represent the mass flow source and sink, respectively. 
+All branches are expected to be connected to the source at the beggining and to the sink at the end.
+
+# Convenience constructors
+ - `BoreholeNetwork(n::Int)`
+Create a network of `n` boreholes (and the source and sink) without any connections.
+
+ - `all_parallel_network(n::Int)`
+Create a network of `n` boreholes connected in parallel.
+
+ - `all_series_network(n::Int)`
+Create a network of `n` boreholes connected in series.
+"""
+@with_kw struct BoreholeNetwork
+    graph::SimpleDiGraph = SimpleDiGraph()
+end
+
+BoreholeNetwork(n::Int) = BoreholeNetwork(graph=SimpleDiGraph(n+2))
+n_boreholes(n::BoreholeNetwork) = nv(n.graph) - 2
+n_branches(n::BoreholeNetwork) = length(neighbors(n.graph, source(n)))
+initialize_mass_flows(network::BoreholeNetwork) = zeros(nv(network.graph))
+"""
+    first_bhs_in_branch(network::BoreholeNetwork)
+
+Return the first borehole in each branch of the borefield.
+"""
+first_bhs_in_branch(network::BoreholeNetwork) = outneighbors(network.graph, source(network))
+
+"""
+    source(network::BoreholeNetwork) 
+
+Return the vertex index of the source.
+"""
+source(network::BoreholeNetwork) = nv(network.graph) - 1
+"""
+    sink(network::BoreholeNetwork) 
+
+Return the vertex index of the sink.
+"""
+sink(network::BoreholeNetwork) = nv(network.graph)
+
+"""
+    connect!(network::BoreholeNetwork, i::Int, j::Int)
+
+Create a flow from borehole `i` to borehole `j`.
+"""
+connect!(network::BoreholeNetwork, i::Int, j::Int) = add_edge!(network.graph, i, j)
+"""
+    connect_to_source!(network::BoreholeNetwork, i::Int)
+
+Create a connection injecting mass flow to borehole `i`.
+"""
+connect_to_source!(network::BoreholeNetwork, i::Int) = add_edge!(network.graph, source(network), i)
+"""
+    connect_to_source!(network::BoreholeNetwork, I::Vector{Int})
+
+Vector version of connect_to_source!.
+"""
+connect_to_source!(network::BoreholeNetwork, I::Vector{Int}) = foreach(i -> add_edge!(network.graph, source(network), i), I)
+"""
+    connect_to_sink!(network::BoreholeNetwork, i::Int)
+
+Create a connection extracting mass flow from borehole `i`.
+"""
+connect_to_sink!(network::BoreholeNetwork, i::Int) = add_edge!(network.graph, i, sink(network))
+"""
+    connect_to_sink!(network::BoreholeNetwork, I::Vector{Int})
+
+Vector version of connect_to_sink!.
+"""
+connect_to_sink!(network::BoreholeNetwork, I::Vector{Int}) = foreach(i -> add_edge!(network.graph, i, sink(network)), I)
+"""
+    connect_in_series!(network::BoreholeNetwork, I::Vector{Int}) 
+
+Connect all the boreholes in `I`, in order.
+"""
+function connect_in_series!(network::BoreholeNetwork, I::Vector{Int}) 
+    for i in 2:length(I)
+        connect!(network, I[i-1], I[i])
+    end
+end
+"""
+    connect_in_parallel!(network::BoreholeNetwork, j::Int, I::Vector{Int}) 
+
+Make a connection from borehole `j` to each of the boreholes in `I`.
+"""
+function connect_in_parallel!(network::BoreholeNetwork, j::Int, I::Vector{Int}) 
+    for i in I
+        connect!(network, j, i)
+    end
+end
+"""
+    boreholes_in_branch(n::BoreholeNetwork; first_bh::Int)
+
+Return all the boreholes in the branch that starts by borehole `first_bh`.
+"""
+function boreholes_in_branch(n::BoreholeNetwork; first_bh::Int)
+    boreholes = neighborhood(n.graph, first_bh, nv(n.graph))
+    filter!(i -> i ≠ sink(n), boreholes)
+    boreholes
+end
+
+function all_parallel_network(n) 
+    network = BoreholeNetwork(n)
+    for i in 1:n
+        connect_to_source!(network, i)
+        connect_to_sink!(network, i)
+    end
+    return network
+end
+
+function all_series_network(n) 
+    network = BoreholeNetwork(n)
+    connect_to_source!(network, 1)
+    for i in 2:n
+        connect!(network, i-1, i)
+    end
+    connect_to_sink!(network, n)
+    return network
+end
+
+"""
+    Valve{T <: Number}(
+        split::Dict{Int, T}
+    )
+
+Models the behaviour of the fluid at pipe divergences. `split` is a dicitonary representing the divergence.
+Its keys represent the boreholes where the flow splits and its values are the proportion of flow going to each borehole. 
+"""
+struct Valve{T <: Number}
+    split::Dict{Int, T}
+end
+act(v::Valve, i, m) = m * v.split[i]
+
+"""
+    equal_valve(nodes::Vector{Int}) 
+
+Return a `Valve` that splits the flow equally to all the boreholes `nodes` in the divergence. 
+"""
+function equal_valve(nodes::Vector{Int}) 
+    valve = Valve(Dict{Int, Float64}())
+    for n in nodes
+        valve.split[n] = 1/length(nodes)
+    end
+    return valve
+end
+"""
+    valve(nodes::Vector{Int}, weights::Vector{T}) 
+
+Return a `Valve` that splits the flow to each borehole `nodes[i]` in the proportion `weights[i]`.
+"""
+function valve(nodes::Vector{Int}, weights::Vector{T}) where {T <: Number}
+    valve = Valve(Dict{Int, T}())
+    for i in 1:length(nodes)
+        valve.split[nodes[i]] = weights[i] 
+    end
+    return valve
+end
+
+"""
+    absolute_valve(nodes::Vector{Int}, mass_flows::Vector{T}) 
+
+Return a `Valve` that splits the flow such that each borehole `nodes[i]` receives an absolute amount 
+of mass flow equal to `mass_flows[i]`.
+"""
+function absolute_valve(nodes::Vector{Int}, mass_flows::Vector{T}) where {T <: Number}
+    valve = Valve(Dict{Int, T}())
+    total_mass_flow = sum(mass_flows)
+    for (i, n) in enumerate(nodes)
+        valve.split[n] = total_mass_flow != 0 ? mass_flows[i] / total_mass_flow : 0.
+    end
+    return valve
+end
+
+"""
+    BoreholeOperation{T <: Number}(
+        valves::Dict{Int, Valve{T}}
+        source_mass_flow::T
+        network::BoreholeNetwork
+    )
+
+Representation of the a state of operation of the borefield. 
+
+# Arguments
+- `valves`: A dicitonary specifying the divergences of the network. Each key is the borehole 
+    where the divergence occurs, and its corresponding value is a `Valve` object, describing the behaviour
+    of the divergence.
+- `source_mass_flow`: The total amount of mass flow injected into all branches.
+- `network`: The representation of the pipe layout of the borefield, including direction of flow.
+
+# Convenience constructors
+`BoreholeOperation(;network::BoreholeNetwork, mass_flows::Vector{T})`
+Builds the operation corresponding to `network` sending mass flow `mass_flows[i]` in each branch `identical`. 
+This method takes care of the initial valve from the source to the start of each branch.
+"""
+struct BoreholeOperation{T <: Number}
+    valves::Dict{Int, Valve{T}}
+    source_mass_flow::T
+    network::BoreholeNetwork
+end
+
+BoreholeOperation(::Nothing) = BoreholeOperation(Dict{Int, Valve{Float64}}(), 0., BoreholeNetwork())
+function BoreholeOperation(;network::BoreholeNetwork, mass_flows::Vector{T}) where {T <: Number}
+    source_valve = absolute_valve(outneighbors(network.graph, source(network)), mass_flows)
+    BoreholeOperation(Dict(source(network) => source_valve), sum(mass_flows), network)
+end
+
+"""
+    Base.reverse(network::BoreholeNetwork) 
+
+Create the reversed flow network from the original `network`.
+"""
+function Base.reverse(network::BoreholeNetwork) 
+    inlets = outneighbors(network.graph, source(network))
+    outlets = inneighbors(network.graph, sink(network))
+    reverse_graph = Base.reverse(network.graph)
+    for inlet in inlets 
+        rem_edge!(reverse_graph, inlet, source(network))
+    end
+    for outlet in outlets 
+        rem_edge!(reverse_graph, sink(network), outlet)
+    end
+    for inlet in inlets 
+        add_edge!(reverse_graph, inlet, sink(network))
+    end
+    for outlet in outlets 
+        add_edge!(reverse_graph, source(network), outlet)
+    end
+    BoreholeNetwork(reverse_graph)
+end
+
+
+"""
+    compute_mass_flows!(mass_flows, network::BoreholeNetwork, operation::BoreholeOperation{T})
+
+Determine the mass flow each borehole receives, as determined by the network topology `network` 
+and the operation `operation` and write the result in `mass_flows`.
+"""
+function compute_mass_flows!(mass_flows, network::BoreholeNetwork, operation::BoreholeOperation{T}) where {T <: Number}
+    mass_flows .= zero(eltype(mass_flows))
+    q = Queue{Tuple{Int, T}}()
+    enqueue!(q, (source(network), operation.source_mass_flow))
+    mass_flows[source(network)] = operation.source_mass_flow
+
+    while !isempty(q)
+        current_node, current_m = dequeue!(q)
+        if current_node in keys(operation.valves)
+            valve = operation.valves[current_node]
+            for n in outneighbors(network.graph, current_node)
+                new_m = act(valve, n, current_m)
+                mass_flows[n] += new_m
+                enqueue!(q, (n, new_m))
+            end
+        elseif length(outneighbors(network.graph, current_node)) > 1
+            message = current_node == source(network) ? "source" : "$(current_node)"
+            throw("Borehole $(message) splits into more than 1 branch; please provide a Valve to specify the flow going to each branch.")
+        else
+            for n in outneighbors(network.graph, current_node)
+                mass_flows[n] += current_m
+                enqueue!(q, (n, current_m))
+            end
+        end
+    end
+end
+
diff --git a/src/modular/operators/SimpleOperator.jl b/src/modular/operators/SimpleOperator.jl
index 6d91114..71d209f 100644
--- a/src/modular/operators/SimpleOperator.jl
+++ b/src/modular/operators/SimpleOperator.jl
@@ -1,17 +1,26 @@
 
 """
-    SimpleOperator{T <: Number} <: Operator (
-        mass_flows::Vector{T}
+    ConstantOperator{T <: Number} <: Operator (
+        valves::Dict{Int, Valve{T}} = Dict{Int, Valve{Float64}}()
+        mass_flow::T
     )
 
-Constant operation strategy, always returning `BoreholeOperation(options.configurations[1], operator.mass_flows)`.
+Constant operation strategy. Its implementation of `operate` is:
+`operate(op::ConstantOperator, step, options, X) = BoreholeOperation(op.valves, op.mass_flow, options.configurations[1])`
+
+# Convenience constructors
+
+    ConstantOperator(network::BoreholeNetwork; mass_flows) 
+
+Builds the initial valve and total source mass flow needed to have mass flow equal to `mass_flows[i]` in branch `i`, according to the network `network`.
 """
-struct SimpleOperator{T <: Number} <: Operator
-    mass_flows::Vector{T}
+@with_kw struct ConstantOperator{T <: Number} <: Operator
+    valves::Dict{Int, Valve{T}} = Dict{Int, Valve{Float64}}()
+    mass_flow::T
 end
-SimpleOperator(;mass_flow, branches) = SimpleOperator(mass_flow .* ones(eltype(mass_flow), branches))
 
-function operate(operator::SimpleOperator, i, options, X)
-    network = options.configurations[1]
-    BoreholeOperation(network, operator.mass_flows)
+function ConstantOperator(network::BoreholeNetwork; mass_flows) 
+    source_valve = absolute_valve(outneighbors(network.graph, source(network)), mass_flows)
+    ConstantOperator(valves = Dict(source(network) => source_valve), mass_flow = sum(mass_flows))
 end
+operate(op::ConstantOperator, step, options, X) = BoreholeOperation(op.valves, op.mass_flow, options.configurations[1])
diff --git a/src/modular/simulate.jl b/src/modular/simulate.jl
index e500c61..ab8fa8a 100644
--- a/src/modular/simulate.jl
+++ b/src/modular/simulate.jl
@@ -30,20 +30,19 @@ function simulate!(;operator, options::SimulationOptions, containers::Simulation
     @unpack M, b, X, mf = containers 
     
     last_operation = BoreholeOperation(nothing)
+    mass_flows = zeros(Nb+2)
     fluid_T = get_T0(medium) .* ones(2Nb)
 
     # Simulation loop
     for i = Ts:Nt
         operation = @views operate(operator, i, options, X[:, 1:i-1])
         operation = unwrap(operation)
+        @unpack network = operation
+        compute_mass_flows!(mass_flows, network, operation)
 
-        for (j, branch) in enumerate(operation.network.branches)
-            for borehole in branch
-                mf[borehole, i] = operation.mass_flows[j]
-            end
-        end
+        @views @. mf[:, i] = mass_flows[1:Nb]
 
-        Nbr = n_branches(operation.network)
+        Nbr = n_branches(network)
 
         internal_model_eqs = 1:Nb
         topology_eqs = Nb+1:2Nb-Nbr
@@ -52,16 +51,14 @@ function simulate!(;operator, options::SimulationOptions, containers::Simulation
         balance_eqs = 3Nb+1:4Nb
 
         # Update M
-        @views internal_model_coeffs!(M[internal_model_eqs, :], borefield, medium, operation, fluid_T, fluid)
-        if last_operation.network != operation.network
-            @views topology_coeffs!(M[topology_eqs, :], operation)
-        end
-        @views constraints_coeffs!(M[constraints_eqs, :], constraint, operation, borefield)
+        @views internal_model_coeffs!(M[internal_model_eqs, :], borefield, medium, mass_flows, fluid_T, fluid)
+        @views topology_coeffs!(M[topology_eqs, :], network, mass_flows)
+        @views constraints_coeffs!(M[constraints_eqs, :], constraint, borefield, network, mass_flows)
         if i == Ts
             @views method_coeffs!(M[method_eqs, :], method, options)
         end
         #if last_operation.mass_flows != operation.mass_flows
-            @views heat_balance_coeffs!(M[balance_eqs, :], borefield, operation, fluid)
+            @views heat_balance_coeffs!(M[balance_eqs, :], borefield, mass_flows, fluid)
         #end
 
         # Update b
diff --git a/test/api/api.jl b/test/api/api.jl
index 65c2669..77b6e64 100644
--- a/test/api/api.jl
+++ b/test/api/api.jl
@@ -1,7 +1,7 @@
 
 @testset "test_Approximations" begin
     @test MeanApproximation() isa Any
-    #@test MidPointApproximation() isa Any
+    @test MidPointApproximation() isa Any
 end
 
 @testset "test_Boreholes" begin
diff --git a/test/borefields/test_EqualBoreholesBorefield.jl b/test/borefields/test_EqualBoreholesBorefield.jl
index 120a1a5..a9033d0 100644
--- a/test/borefields/test_EqualBoreholesBorefield.jl
+++ b/test/borefields/test_EqualBoreholesBorefield.jl
@@ -37,12 +37,13 @@ end
     fluid = FluidMock(cpf = 1000.)
     T_fluid = 10 * ones(2*Nb)
 
-    network = BoreholeNetwork([[1], [2], [3]])
-    mass_flows = ones(Nb)
-    operation = BoreholeOperation(network=network, mass_flows=mass_flows)
+    network = all_parallel_network(Nb)
+    mass_flows = initialize_mass_flows(network)
+    operation = BoreholeOperation(network=network, mass_flows=ones(Nb))
+    compute_mass_flows!(mass_flows, network, operation)
 
     M = zeros(Nb, 4*Nb)
-    internal_model_coeffs!(M, borefield, medium, operation, T_fluid, fluid)
+    internal_model_coeffs!(M, borefield, medium, mass_flows, T_fluid, fluid)
 
     expected = [
         (1, 1, Ci), (1, 2, Co), (1, 2Nb+1, Cb),
diff --git a/test/constraints/test_HeatLoadConstraint.jl b/test/constraints/test_HeatLoadConstraint.jl
index 109e4d0..cd53e86 100644
--- a/test/constraints/test_HeatLoadConstraint.jl
+++ b/test/constraints/test_HeatLoadConstraint.jl
@@ -12,11 +12,14 @@ import BoreholeNetworksSimulator: constraints_coeffs!, constraints_b!
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1], [2], [3]])
+    network = all_parallel_network(Nb)
+    mass_flows = initialize_mass_flows(network)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    compute_mass_flows!(mass_flows, network, operation)
+
     borefield = BorefieldMock(H = H * ones(3))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 3Nb+1, H), (2, 3Nb+2, H), (3, 3Nb+3, H)]
     @test test_sparse_matrix(M, expected)
@@ -34,11 +37,14 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2, 3]])
-    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
     borefield = BorefieldMock(H = H * ones(3))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    network = all_series_network(Nb)
+    mass_flows = initialize_mass_flows(network)
+    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    compute_mass_flows!(mass_flows, network, operation)
+
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 3Nb+1, H), (1, 3Nb+2, H), (1, 3Nb+3, H)]
     @test test_sparse_matrix(M, expected)
@@ -56,11 +62,18 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2], [3, 4], [5]])
+    network = BoreholeNetwork(5)
+    connect_to_source!(network, [1, 3, 5])
+    connect!(network, 1, 2)
+    connect!(network, 3, 4)
+    connect_to_sink!(network, [2, 4, 5])
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    mass_flows = initialize_mass_flows(network)
+    compute_mass_flows!(mass_flows, network, operation)
+
     borefield = BorefieldMock(H = H * ones(5))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 3Nb+1, H), (1, 3Nb+2, H), (2, 3Nb+3, H), (2, 3Nb+4, H), (3, 3Nb+5, H)]
     @test test_sparse_matrix(M, expected)
@@ -79,7 +92,7 @@ end
 
     b = zeros(Nbr)
 
-    network = BoreholeNetwork([[1], [2], [3]])
+    network = all_parallel_network(Nb)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
 
     b_time = zeros(Nbr, Nt)
diff --git a/test/constraints/test_InletTempConstraint.jl b/test/constraints/test_InletTempConstraint.jl
index f2377c3..246289e 100644
--- a/test/constraints/test_InletTempConstraint.jl
+++ b/test/constraints/test_InletTempConstraint.jl
@@ -10,11 +10,14 @@ import BoreholeNetworksSimulator: constraints_coeffs!, constraints_b!
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1], [2], [3]])
-    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
     borefield = BorefieldMock()
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    network = all_parallel_network(Nb)
+    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    mass_flows = initialize_mass_flows(network)
+    compute_mass_flows!(mass_flows, network, operation)
+
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 1, 1.), (2, 3, 1.), (3, 5, 1.)]
     @test test_sparse_matrix(M, expected)
@@ -30,11 +33,14 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2, 3]])
-    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
     borefield = BorefieldMock()
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    network = all_series_network(Nb)
+    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    mass_flows = initialize_mass_flows(network)
+    compute_mass_flows!(mass_flows, network, operation)
+
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 1, 1.)]
     @test test_sparse_matrix(M, expected)
@@ -50,11 +56,18 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2], [3, 4], [5]])
+    network = BoreholeNetwork(5)
+    connect_to_source!(network, [1, 3, 5])
+    connect!(network, 1, 2)
+    connect!(network, 3, 4)
+    connect_to_sink!(network, [2, 4, 5])
+    mass_flows = initialize_mass_flows(network)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    compute_mass_flows!(mass_flows, network, operation)
+    
     borefield = BorefieldMock()
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 1, 1.), (2, 2*2+1, 1.), (3, 4*2+1, 1.)]
     @test test_sparse_matrix(M, expected)
@@ -73,7 +86,7 @@ end
 
     b = zeros(Nbr)
 
-    network = BoreholeNetwork([[1], [2], [3]])
+    network = all_parallel_network(Nb)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
 
     b_time = zeros(Nbr, Nt)
diff --git a/test/constraints/test_TotalHeatLoadConstraint.jl b/test/constraints/test_TotalHeatLoadConstraint.jl
index e273483..1aff5fe 100644
--- a/test/constraints/test_TotalHeatLoadConstraint.jl
+++ b/test/constraints/test_TotalHeatLoadConstraint.jl
@@ -12,11 +12,13 @@ import BoreholeNetworksSimulator: constraints_coeffs!, constraints_b!
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1], [2], [3]])
+    network = all_parallel_network(3)
+    mass_flows = initialize_mass_flows(network)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
     borefield = BorefieldMock(H = H * ones(3))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    compute_mass_flows!(mass_flows, network, operation)
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [
         (1, 3Nb+1, H), (1, 3Nb+2, H), (1, 3Nb+3, H),
@@ -38,11 +40,13 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2, 3]])
+    network = all_series_network(Nb)
+    mass_flows = initialize_mass_flows(network)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    compute_mass_flows!(mass_flows, network, operation)
     borefield = BorefieldMock(H = H * ones(3))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [(1, 3Nb+1, H), (1, 3Nb+2, H), (1, 3Nb+3, H)]
     @test test_sparse_matrix(M, expected)
@@ -60,11 +64,18 @@ end
 
     M = zeros(Nbr, 4*Nb)
 
-    network = BoreholeNetwork([[1, 2], [3, 4], [5]])
-    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
     borefield = BorefieldMock(H = H * ones(5))
 
-    constraints_coeffs!(M, constraint, operation, borefield)
+    network = BoreholeNetwork(5)
+    connect_to_source!(network, [1, 3, 5])
+    connect!(network, 1, 2)
+    connect!(network, 3, 4)
+    connect_to_sink!(network, [2, 4, 5])
+    mass_flows = initialize_mass_flows(network)
+    operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
+    compute_mass_flows!(mass_flows, network, operation)
+
+    constraints_coeffs!(M, constraint, borefield, network, mass_flows)
 
     expected = [
         (1, 3Nb+1, H), (1, 3Nb+2, H), (1, 3Nb+3, H), (1, 3Nb+4, H), (1, 3Nb+5, H),
@@ -84,7 +95,7 @@ end
 
     b = zeros(Nbr)
 
-    network = BoreholeNetwork([[1], [2], [3]])
+    network = all_parallel_network(Nb)
     operation = BoreholeOperation(network=network, mass_flows=ones(Nbr))
 
     b_time = zeros(Nbr, Nt)
diff --git a/test/mocks/interfaces/constraint.jl b/test/mocks/interfaces/constraint.jl
index 47b866c..9c0824f 100644
--- a/test/mocks/interfaces/constraint.jl
+++ b/test/mocks/interfaces/constraint.jl
@@ -8,5 +8,5 @@ Mock for testing purposes.
     M = []
     b = []
 end
-BoreholeNetworksSimulator.constraints_coeffs!(M, c::ConstraintMock, operation, borefield) = c.M .= M
+BoreholeNetworksSimulator.constraints_coeffs!(M, c::ConstraintMock, ::Borefield, network, mass_flows) = c.M .= M
 BoreholeNetworksSimulator.constraints_b!(b, c::ConstraintMock, operation, step) = c.b .= b
diff --git a/test/test_interfaces.jl b/test/test_interfaces.jl
index 8c1ed8a..641a573 100644
--- a/test/test_interfaces.jl
+++ b/test/test_interfaces.jl
@@ -30,5 +30,5 @@ end
 end
 
 @testset "test_Operator_interfaces" begin
-    @test RI.check_interface_implemented(Operator, SimpleOperator)
+    @test RI.check_interface_implemented(Operator, ConstantOperator)
 end