Skip to content

Commit ab3b998

Browse files
shsmsshsms
authored
Return Formulas from formula generators instead of Strings (#10)
The `Formula` type implements the `Add` and `Sub` traits, making it easy to combine formulas.
2 parents 93bb81c + 8f60f27 commit ab3b998

File tree

12 files changed

+209
-119
lines changed

12 files changed

+209
-119
lines changed

‎src/graph.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ mod formulas;
1313
pub mod iterators;
1414

1515
use crate::{ComponentGraphConfig, Edge, Node};
16+
pub use formulas::Formula;
1617
use petgraph::graph::{DiGraph, NodeIndex};
1718
use std::collections::HashMap;
1819

‎src/graph/formulas.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,58 +12,61 @@ use crate::Node;
1212

1313
mod expr;
1414
mod fallback;
15+
mod formula;
1516
mod generators;
1617
mod traversal;
1718

19+
pub use formula::Formula;
20+
1821
/// Formulas for various microgrid metrics.
1922
impl<N, E> ComponentGraph<N, E>
2023
where
2124
N: Node,
2225
E: Edge,
2326
{
2427
/// Returns a string representing the consumer formula for the graph.
25-
pub fn consumer_formula(&self) -> Result<String, Error> {
28+
pub fn consumer_formula(&self) -> Result<Formula, Error> {
2629
generators::consumer::ConsumerFormulaBuilder::try_new(self)?.build()
2730
}
2831

2932
/// Returns a string representing the grid formula for the graph.
30-
pub fn grid_formula(&self) -> Result<String, Error> {
33+
pub fn grid_formula(&self) -> Result<Formula, Error> {
3134
generators::grid::GridFormulaBuilder::try_new(self)?.build()
3235
}
3336

3437
/// Returns a string representing the producer formula for the graph.
35-
pub fn producer_formula(&self) -> Result<String, Error> {
38+
pub fn producer_formula(&self) -> Result<Formula, Error> {
3639
generators::producer::ProducerFormulaBuilder::try_new(self)?.build()
3740
}
3841

3942
/// Returns a string representing the battery formula for the graph.
40-
pub fn battery_formula(&self, battery_ids: Option<BTreeSet<u64>>) -> Result<String, Error> {
43+
pub fn battery_formula(&self, battery_ids: Option<BTreeSet<u64>>) -> Result<Formula, Error> {
4144
generators::battery::BatteryFormulaBuilder::try_new(self, battery_ids)?.build()
4245
}
4346

4447
/// Returns a string representing the CHP formula for the graph.
45-
pub fn chp_formula(&self, chp_ids: Option<BTreeSet<u64>>) -> Result<String, Error> {
48+
pub fn chp_formula(&self, chp_ids: Option<BTreeSet<u64>>) -> Result<Formula, Error> {
4649
generators::chp::CHPFormulaBuilder::try_new(self, chp_ids)?.build()
4750
}
4851

4952
/// Returns a string representing the PV formula for the graph.
50-
pub fn pv_formula(&self, pv_inverter_ids: Option<BTreeSet<u64>>) -> Result<String, Error> {
53+
pub fn pv_formula(&self, pv_inverter_ids: Option<BTreeSet<u64>>) -> Result<Formula, Error> {
5154
generators::pv::PVFormulaBuilder::try_new(self, pv_inverter_ids)?.build()
5255
}
5356

5457
/// Returns a string with the coalesce formula for the given component IDs.
5558
///
5659
/// This formula uses the `COALESCE` function to return the first non-null
5760
/// value from the components with the provided IDs.
58-
pub fn coalesce(&self, component_ids: BTreeSet<u64>) -> Result<String, Error> {
61+
pub fn coalesce(&self, component_ids: BTreeSet<u64>) -> Result<Formula, Error> {
5962
generators::generic::CoalesceFormulaBuilder::try_new(self, component_ids)?.build()
6063
}
6164

6265
/// Returns a string representing the EV charger formula for the graph.
6366
pub fn ev_charger_formula(
6467
&self,
6568
ev_charger_ids: Option<BTreeSet<u64>>,
66-
) -> Result<String, Error> {
69+
) -> Result<Formula, Error> {
6770
generators::ev_charger::EVChargerFormulaBuilder::try_new(self, ev_charger_ids)?.build()
6871
}
6972
}

‎src/graph/formulas/formula.rs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// License: MIT
2+
// Copyright © 2025 Frequenz Energy-as-a-Service GmbH
3+
4+
//! This module defines the `Formula` struct and its operations.
5+
6+
use super::expr::Expr;
7+
8+
// This struct represents the microgrid metric formulas that are generated by
9+
// traversing the component graph.
10+
//
11+
// `Formula` objects can be added or subtracted from each other, and they can
12+
// converted to a string representation, before they are passed to an evaluator.
13+
#[derive(Debug, Clone)]
14+
pub struct Formula {
15+
expr: Expr,
16+
}
17+
18+
impl Formula {
19+
pub(crate) fn new(expr: Expr) -> Self {
20+
Formula { expr }
21+
}
22+
}
23+
24+
impl std::ops::Add for Formula {
25+
type Output = Self;
26+
27+
fn add(self, rhs: Self) -> Self::Output {
28+
Formula {
29+
expr: self.expr + rhs.expr,
30+
}
31+
}
32+
}
33+
34+
impl std::ops::Sub for Formula {
35+
type Output = Self;
36+
37+
fn sub(self, rhs: Self) -> Self::Output {
38+
Formula {
39+
expr: self.expr - rhs.expr,
40+
}
41+
}
42+
}
43+
44+
impl std::fmt::Display for Formula {
45+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46+
self.expr.fmt(f)
47+
}
48+
}
49+
50+
impl From<Formula> for String {
51+
fn from(formula: Formula) -> Self {
52+
formula.expr.to_string()
53+
}
54+
}

‎src/graph/formulas/generators/battery.rs‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
use std::collections::BTreeSet;
77

88
use crate::component_category::CategoryPredicates;
9+
use crate::graph::formulas::expr::Expr;
10+
use crate::graph::formulas::Formula;
911
use crate::{ComponentGraph, Edge, Error, Node};
1012

1113
pub(crate) struct BatteryFormulaBuilder<'a, N, E>
@@ -47,14 +49,14 @@ where
4749
/// This is the sum of all battery_inverters in the graph. If the
4850
/// battery_ids are provided, only the batteries with the given ids are
4951
/// included in the formula.
50-
pub fn build(self) -> Result<String, Error> {
52+
pub fn build(self) -> Result<Formula, Error> {
5153
if self.inverter_ids.is_empty() {
52-
return Ok("0.0".to_string());
54+
return Ok(Formula::new(Expr::number(0.0)));
5355
}
5456

5557
self.graph
5658
.fallback_expr(self.inverter_ids, false)
57-
.map(|expr| expr.to_string())
59+
.map(Formula::new)
5860
}
5961

6062
fn find_inverter_ids(
@@ -104,7 +106,7 @@ mod tests {
104106
builder.connect(grid, grid_meter);
105107

106108
let graph = builder.build(None)?;
107-
let formula = graph.battery_formula(None)?;
109+
let formula = graph.battery_formula(None)?.to_string();
108110
assert_eq!(formula, "0.0");
109111

110112
// Add a battery meter with one inverter and one battery.
@@ -115,7 +117,7 @@ mod tests {
115117
assert_eq!(meter_bat_chain.component_id(), 2);
116118

117119
let graph = builder.build(None)?;
118-
let formula = graph.battery_formula(None)?;
120+
let formula = graph.battery_formula(None)?.to_string();
119121
assert_eq!(formula, "COALESCE(#3, #2, 0.0)");
120122

121123
// Add a second battery meter with one inverter and two batteries.
@@ -125,18 +127,22 @@ mod tests {
125127
assert_eq!(meter_bat_chain.component_id(), 5);
126128

127129
let graph = builder.build(None)?;
128-
let formula = graph.battery_formula(None)?;
130+
let formula = graph.battery_formula(None)?.to_string();
129131
assert_eq!(formula, "COALESCE(#3, #2, 0.0) + COALESCE(#6, #5, 0.0)");
130132

131-
let formula = graph.battery_formula(Some(BTreeSet::from([4])))?;
133+
let formula = graph
134+
.battery_formula(Some(BTreeSet::from([4])))?
135+
.to_string();
132136
assert_eq!(formula, "COALESCE(#3, #2, 0.0)");
133137

134-
let formula = graph.battery_formula(Some(BTreeSet::from([7, 8])))?;
138+
let formula = graph
139+
.battery_formula(Some(BTreeSet::from([7, 8])))?
140+
.to_string();
135141
assert_eq!(formula, "COALESCE(#6, #5, 0.0)");
136142

137143
let formula = graph
138-
.battery_formula(Some(BTreeSet::from([4, 8, 7])))
139-
.unwrap();
144+
.battery_formula(Some(BTreeSet::from([4, 8, 7])))?
145+
.to_string();
140146
assert_eq!(formula, "COALESCE(#3, #2, 0.0) + COALESCE(#6, #5, 0.0)");
141147

142148
// Add a third battery meter with two inverters with two connected batteries.
@@ -146,7 +152,7 @@ mod tests {
146152
assert_eq!(meter_bat_chain.component_id(), 9);
147153

148154
let graph = builder.build(None)?;
149-
let formula = graph.battery_formula(None)?;
155+
let formula = graph.battery_formula(None)?.to_string();
150156
assert_eq!(
151157
formula,
152158
concat!(
@@ -157,8 +163,8 @@ mod tests {
157163
);
158164

159165
let formula = graph
160-
.battery_formula(Some(BTreeSet::from([12, 13])))
161-
.unwrap();
166+
.battery_formula(Some(BTreeSet::from([12, 13])))?
167+
.to_string();
162168
assert_eq!(
163169
formula,
164170
"COALESCE(#11 + #10, #9, COALESCE(#11, 0.0) + COALESCE(#10, 0.0))"
@@ -171,7 +177,7 @@ mod tests {
171177
assert_eq!(meter_pv_chain.component_id(), 14);
172178

173179
let graph = builder.build(None)?;
174-
let formula = graph.battery_formula(None)?;
180+
let formula = graph.battery_formula(None)?.to_string();
175181
assert_eq!(
176182
formula,
177183
concat!(
@@ -208,7 +214,7 @@ mod tests {
208214
allow_unspecified_inverters: true,
209215
..Default::default()
210216
}))?;
211-
let formula = graph.battery_formula(None)?;
217+
let formula = graph.battery_formula(None)?.to_string();
212218
assert_eq!(
213219
formula,
214220
concat!(
@@ -220,22 +226,26 @@ mod tests {
220226
);
221227

222228
let formula = graph
223-
.battery_formula(Some(BTreeSet::from([19, 21])))
224-
.unwrap();
229+
.battery_formula(Some(BTreeSet::from([19, 21])))?
230+
.to_string();
225231
assert_eq!(
226232
formula,
227233
"COALESCE(#20 + #18, #17, COALESCE(#20, 0.0) + COALESCE(#18, 0.0))"
228234
);
229235

230-
let formula = graph.battery_formula(Some(BTreeSet::from([19]))).unwrap();
236+
let formula = graph
237+
.battery_formula(Some(BTreeSet::from([19])))?
238+
.to_string();
231239
assert_eq!(formula, "COALESCE(#18, 0.0)");
232240

233-
let formula = graph.battery_formula(Some(BTreeSet::from([21]))).unwrap();
241+
let formula = graph
242+
.battery_formula(Some(BTreeSet::from([21])))?
243+
.to_string();
234244
assert_eq!(formula, "COALESCE(#20, 0.0)");
235245

236246
let formula = graph
237-
.battery_formula(Some(BTreeSet::from([4, 12, 13, 19])))
238-
.unwrap();
247+
.battery_formula(Some(BTreeSet::from([4, 12, 13, 19])))?
248+
.to_string();
239249
assert_eq!(
240250
formula,
241251
concat!(

‎src/graph/formulas/generators/chp.rs‎

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
use std::collections::BTreeSet;
77

88
use crate::component_category::CategoryPredicates;
9+
use crate::graph::formulas::expr::Expr;
10+
use crate::graph::formulas::Formula;
911
use crate::{ComponentGraph, Edge, Error, Node};
1012

1113
pub(crate) struct CHPFormulaBuilder<'a, N, E>
@@ -43,9 +45,9 @@ where
4345
///
4446
/// This is the sum of all CHPs in the graph. If the chp_ids are provided,
4547
/// only the CHPs with the given ids are included in the formula.
46-
pub fn build(self) -> Result<String, Error> {
48+
pub fn build(self) -> Result<Formula, Error> {
4749
if self.chp_ids.is_empty() {
48-
return Ok("0.0".to_string());
50+
return Ok(Formula::new(Expr::number(0.0)));
4951
}
5052

5153
for id in &self.chp_ids {
@@ -59,7 +61,7 @@ where
5961

6062
self.graph
6163
.fallback_expr(self.chp_ids, false)
62-
.map(|expr| expr.to_string())
64+
.map(Formula::new)
6365
}
6466
}
6567

@@ -78,7 +80,7 @@ mod tests {
7880
builder.connect(grid, grid_meter);
7981

8082
let graph = builder.build(None)?;
81-
let formula = graph.chp_formula(None)?;
83+
let formula = graph.chp_formula(None)?.to_string();
8284
assert_eq!(formula, "0.0");
8385

8486
// Add a chp meter with one chp
@@ -89,7 +91,7 @@ mod tests {
8991
assert_eq!(meter_chp_chain.component_id(), 2);
9092

9193
let graph = builder.build(None)?;
92-
let formula = graph.chp_formula(None)?;
94+
let formula = graph.chp_formula(None)?.to_string();
9395
assert_eq!(formula, "COALESCE(#3, #2, 0.0)");
9496

9597
// Add a battery meter with one inverter and two batteries.
@@ -99,7 +101,7 @@ mod tests {
99101
assert_eq!(meter_bat_chain.component_id(), 4);
100102

101103
let graph = builder.build(None)?;
102-
let formula = graph.chp_formula(None)?;
104+
let formula = graph.chp_formula(None)?.to_string();
103105
assert_eq!(formula, "COALESCE(#3, #2, 0.0)");
104106

105107
// Add a chp meter with two CHPs.
@@ -109,7 +111,7 @@ mod tests {
109111
assert_eq!(meter_chp_chain.component_id(), 8);
110112

111113
let graph = builder.build(None)?;
112-
let formula = graph.chp_formula(None)?;
114+
let formula = graph.chp_formula(None)?.to_string();
113115
assert_eq!(
114116
formula,
115117
concat!(
@@ -118,7 +120,9 @@ mod tests {
118120
)
119121
);
120122

121-
let formula = graph.chp_formula(Some(BTreeSet::from([10, 3]))).unwrap();
123+
let formula = graph
124+
.chp_formula(Some(BTreeSet::from([10, 3])))?
125+
.to_string();
122126
assert_eq!(formula, "COALESCE(#3, #2, 0.0) + COALESCE(#10, 0.0)");
123127

124128
// add a meter direct to the grid with three CHPs
@@ -128,7 +132,7 @@ mod tests {
128132
assert_eq!(meter_chp_chain.component_id(), 11);
129133

130134
let graph = builder.build(None)?;
131-
let formula = graph.chp_formula(None)?;
135+
let formula = graph.chp_formula(None)?.to_string();
132136
assert_eq!(
133137
formula,
134138
concat!(
@@ -143,8 +147,8 @@ mod tests {
143147
);
144148

145149
let formula = graph
146-
.chp_formula(Some(BTreeSet::from([3, 9, 10, 12, 13])))
147-
.unwrap();
150+
.chp_formula(Some(BTreeSet::from([3, 9, 10, 12, 13])))?
151+
.to_string();
148152
assert_eq!(
149153
formula,
150154
concat!(
@@ -155,8 +159,8 @@ mod tests {
155159
);
156160

157161
let formula = graph
158-
.chp_formula(Some(BTreeSet::from([3, 9, 10, 12, 13, 14])))
159-
.unwrap();
162+
.chp_formula(Some(BTreeSet::from([3, 9, 10, 12, 13, 14])))?
163+
.to_string();
160164
assert_eq!(
161165
formula,
162166
concat!(
@@ -170,7 +174,9 @@ mod tests {
170174
),
171175
);
172176

173-
let formula = graph.chp_formula(Some(BTreeSet::from([10, 14]))).unwrap();
177+
let formula = graph
178+
.chp_formula(Some(BTreeSet::from([10, 14])))?
179+
.to_string();
174180
assert_eq!(formula, "COALESCE(#10, 0.0) + COALESCE(#14, 0.0)");
175181

176182
// Failure cases:

0 commit comments

Comments
 (0)