Skip to main content

epanet_rs/solver/
state.rs

1//! Per-step mutable solver state (flows, heads, demands, statuses, settings, resistances).
2
3use simplelog::debug;
4
5use crate::model::control::ControlCondition;
6use crate::model::link::{LinkStatus, LinkTrait};
7use crate::model::network::Network;
8use crate::model::node::NodeType;
9use crate::model::options::DemandModel;
10use crate::model::units::{Cfs, Ft3};
11use crate::utils::time::seconds_to_hhmmss;
12
13/// The solver state is the initial/final state of the solver for a single step
14#[derive(Debug, Clone)]
15pub struct SolverState {
16    pub flows: Vec<f64>,
17    pub heads: Vec<f64>,
18    pub demands: Vec<f64>,
19    pub emitter_flows: Vec<f64>,
20    pub demand_flows: Vec<f64>,
21    pub statuses: Vec<LinkStatus>,
22    pub settings: Vec<f64>,
23    pub resistances: Vec<f64>,
24    /// version of the topology of the network on which the state was created
25    pub topology_version: u32,
26    /// version of the properties of the network on which the state was created
27    pub properties_version: u32,
28}
29
30impl SolverState {
31    /// Create a new solver state with the initial values for the flows, heads, demands and statuses and calculate resistances
32    pub fn new_with_initial_values(network: &Network) -> Self {
33        // get the initial emitter flows
34        let initial_emitter_flows = network
35            .nodes
36            .iter()
37            .map(|n| {
38                if let NodeType::Junction(junction) = &n.node_type {
39                    if junction.emitter_coefficient > 0.0 {
40                        1.0
41                    } else {
42                        0.0
43                    }
44                } else {
45                    0.0
46                }
47            })
48            .collect::<Vec<f64>>();
49
50        Self {
51            flows: network
52                .links
53                .iter()
54                .map(|l| l.initial_flow())
55                .collect::<Vec<f64>>(),
56            heads: network
57                .nodes
58                .iter()
59                .map(|n| n.initial_head())
60                .collect::<Vec<f64>>(),
61            emitter_flows: initial_emitter_flows,
62            demands: vec![0.0; network.nodes.len()],
63            demand_flows: vec![0.0; network.nodes.len()],
64            settings: network
65                .links
66                .iter()
67                .map(|l| l.initial_setting())
68                .collect::<Vec<f64>>(),
69            statuses: network
70                .links
71                .iter()
72                .map(|l| l.initial_status)
73                .collect::<Vec<LinkStatus>>(),
74            resistances: network
75                .links
76                .iter()
77                .map(|l| l.resistance())
78                .collect::<Vec<f64>>(),
79            topology_version: network.topology_version,
80            properties_version: network.properties_version,
81        }
82    }
83
84    /// Applies demand and head patterns to the state at the given time.
85    /// TODO: Add support for multiple patterns
86    pub fn apply_patterns(&mut self, network: &Network, time: usize) {
87        let time_options = &network.options.time_options;
88        let pattern_time = time_options.pattern_start + time;
89        let pattern_index = pattern_time / time_options.pattern_timestep;
90
91        for (i, node) in network.nodes.iter().enumerate() {
92            let NodeType::Reservoir(reservoir) = &node.node_type else {
93                continue;
94            };
95            let Some(pat_idx) = reservoir.head_pattern_index else {
96                continue;
97            };
98            let pattern = &network.patterns[pat_idx];
99            self.heads[i] =
100                node.elevation * pattern.multipliers[pattern_index % pattern.multipliers.len()];
101        }
102
103        let default_pattern_idx = network
104            .options
105            .pattern
106            .as_ref()
107            .or(Some(&Box::from("1")))
108            .and_then(|id| network.pattern_map.get(id).copied());
109
110        self.demands = network
111            .nodes
112            .iter()
113            .map(|n| {
114                let NodeType::Junction(junction) = &n.node_type else {
115                    return 0.0;
116                };
117                // loop over all demand categories to calculate the total demand
118                let mut total_demand = 0.0;
119                for demand in junction.demands.iter() {
120                    let pat_idx = demand.pattern_index.or(default_pattern_idx);
121                    let pattern = pat_idx.map(|idx| &network.patterns[idx]);
122                    let multiplier = match pattern {
123                        Some(p) => p.multipliers[pattern_index % p.multipliers.len()],
124                        None => 1.0,
125                    };
126                    total_demand +=
127                        demand.basedemand * multiplier * network.options.demand_multiplier
128                }
129                total_demand
130            })
131            .collect::<Vec<Cfs>>();
132
133        if network.options.demand_model == DemandModel::PDA {
134            self.demand_flows = self.demands.clone();
135        }
136    }
137
138    /// Applies controls to the state at the given time.
139    pub fn apply_controls(&mut self, network: &Network, time: usize) {
140        let clocktime = (time + network.options.time_options.start_clocktime) % (24 * 3600);
141
142        for control in &network.controls {
143            // skip controls that are not pressure or level controls
144            if matches!(
145                control.condition,
146                ControlCondition::LowPressure { .. } | ControlCondition::HighPressure { .. }
147            ) {
148                continue;
149            }
150            if matches!(
151                control.condition,
152                ControlCondition::LowLevel { .. } | ControlCondition::HighLevel { .. }
153            ) && time == 0
154            {
155                continue;
156            }
157            if control.is_active(self, network, time, clocktime) {
158                debug!(
159                    "<yellow>Activating control: {:?} at time {}</>",
160                    control,
161                    seconds_to_hhmmss(clocktime)
162                );
163                control.activate(self, network);
164            }
165        }
166    }
167
168    /// Updates the tank levels for the given time step.
169    pub fn update_tanks(&mut self, network: &Network, timestep: usize) {
170        for (tank_index, node) in network.nodes.iter().enumerate() {
171            if let NodeType::Tank(tank) = &node.node_type {
172                // calculate the delta volume and new head for the tank
173                let delta_volume = self.demands[tank_index] * timestep as Ft3;
174                let new_head = tank.new_head(delta_volume, self.heads[tank_index]);
175                self.heads[tank_index] = new_head;
176            }
177        }
178    }
179
180    /// Updates the state with the changes to the network.
181    pub fn update_with_network_changes(&mut self, network: &mut Network) {
182        // if the topology has changed, reset the state to the initial values
183        if network.topology_version != self.topology_version {
184            *self = SolverState::new_with_initial_values(network);
185            // clear the updated nodes and links
186            network.updated_nodes.clear();
187            network.updated_links.clear();
188        } else {
189            for node_index in network.updated_nodes.iter() {
190                let node = &network.nodes[*node_index];
191
192                match &node.node_type {
193                    NodeType::Junction(junction) => {
194                        // set the emitter flow to 1.0 if the emitter coefficient is greater than 0.0, otherwise set it to 0.0
195                        self.emitter_flows[*node_index] = if junction.emitter_coefficient > 0.0 {
196                            1.0
197                        } else {
198                            0.0
199                        };
200                    }
201                    NodeType::Tank(_) | NodeType::Reservoir(_) => {
202                        // update the head
203                        self.heads[*node_index] = node.initial_head();
204                    }
205                }
206            }
207
208            for link_index in network.updated_links.iter() {
209                let link = &network.links[*link_index];
210                // update the resistance of the link
211                self.resistances[*link_index] = link.resistance();
212                // reset the flow of the link to the initial flow
213                self.flows[*link_index] = link.initial_flow();
214                // reset the status of the link to the initial status
215                // TODO: investigate whether this leads to inconsistent results when changing non-status parameters of the link
216                self.statuses[*link_index] = link.initial_status;
217            }
218        }
219
220        // clear the updated nodes and links
221        network.updated_nodes.clear();
222        network.updated_links.clear();
223
224        self.properties_version = network.properties_version;
225    }
226}