Skip to main content

otspot_model/
expression.rs

1//! Expression type and operator overloads for the modeling API
2
3use std::collections::HashMap;
4use std::ops::{Add, Mul, Neg, Sub};
5
6use super::constraint::{Constraint, ConstraintSense};
7use super::variable::Variable;
8
9/// A linear expression: sum of (coefficient * variable) + constant.
10///
11/// Internally stored as a flat `HashMap<Variable, f64>` plus a scalar constant.
12/// This representation automatically merges like terms: `x + 2*x` → `{x: 3.0}`.
13#[derive(Debug, Clone, Default)]
14pub struct Expression {
15    pub(crate) coefficients: HashMap<Variable, f64>,
16    pub(crate) constant: f64,
17}
18
19impl Expression {
20    /// Create an expression representing a single scalar constant.
21    pub fn from_constant(c: f64) -> Self {
22        Expression {
23            coefficients: HashMap::new(),
24            constant: c,
25        }
26    }
27
28    /// Retrieve the coefficient for a given variable (0.0 if not present).
29    pub(crate) fn coefficient(&self, var: Variable) -> f64 {
30        self.coefficients.get(&var).copied().unwrap_or(0.0)
31    }
32
33    /// Merge another expression into this one (in-place add).
34    fn merge_add(&mut self, rhs: Expression) {
35        for (var, coeff) in rhs.coefficients {
36            *self.coefficients.entry(var).or_insert(0.0) += coeff;
37        }
38        self.constant += rhs.constant;
39    }
40
41    // --- Constraint builders ---
42
43    fn constraint(self, rhs: impl Into<Expression>, sense: ConstraintSense) -> Constraint {
44        let mut lhs = self;
45        let mut rhs_expr = rhs.into();
46        for (var, coeff) in rhs_expr.coefficients.drain() {
47            *lhs.coefficients.entry(var).or_insert(0.0) -= coeff;
48        }
49        let rhs_val = rhs_expr.constant - lhs.constant;
50        lhs.constant = 0.0;
51        Constraint { lhs, rhs: rhs_val, sense }
52    }
53
54    /// Create a `<=` constraint: `self <= rhs`.
55    pub fn leq(self, rhs: impl Into<Expression>) -> Constraint {
56        self.constraint(rhs, ConstraintSense::Le)
57    }
58
59    /// Create a `>=` constraint: `self >= rhs`.
60    pub fn geq(self, rhs: impl Into<Expression>) -> Constraint {
61        self.constraint(rhs, ConstraintSense::Ge)
62    }
63
64    /// Create an `==` constraint: `self == rhs`.
65    pub fn eq_constraint(self, rhs: impl Into<Expression>) -> Constraint {
66        self.constraint(rhs, ConstraintSense::Eq)
67    }
68}
69
70// --- From conversions ---
71
72impl From<Variable> for Expression {
73    fn from(var: Variable) -> Self {
74        let mut coefficients = HashMap::new();
75        coefficients.insert(var, 1.0);
76        Expression {
77            coefficients,
78            constant: 0.0,
79        }
80    }
81}
82
83impl From<f64> for Expression {
84    fn from(c: f64) -> Self {
85        Expression::from_constant(c)
86    }
87}
88
89impl From<i32> for Expression {
90    fn from(c: i32) -> Self {
91        Expression::from_constant(c as f64)
92    }
93}
94
95// --- Negation ---
96
97impl Neg for Expression {
98    type Output = Expression;
99    fn neg(mut self) -> Expression {
100        for coeff in self.coefficients.values_mut() {
101            *coeff = -*coeff;
102        }
103        self.constant = -self.constant;
104        self
105    }
106}
107
108impl Neg for Variable {
109    type Output = Expression;
110    fn neg(self) -> Expression {
111        -Expression::from(self)
112    }
113}
114
115// --- Expression + Expression ---
116
117impl Add for Expression {
118    type Output = Expression;
119    fn add(mut self, rhs: Expression) -> Expression {
120        self.merge_add(rhs);
121        self
122    }
123}
124
125// --- Expression - Expression ---
126
127impl Sub for Expression {
128    type Output = Expression;
129    fn sub(self, rhs: Expression) -> Expression {
130        self + (-rhs)
131    }
132}
133
134// --- f64 * Expression ---
135
136impl Mul<Expression> for f64 {
137    type Output = Expression;
138    fn mul(self, mut rhs: Expression) -> Expression {
139        for coeff in rhs.coefficients.values_mut() {
140            *coeff *= self;
141        }
142        rhs.constant *= self;
143        rhs
144    }
145}
146
147impl Mul<f64> for Expression {
148    type Output = Expression;
149    fn mul(self, rhs: f64) -> Expression {
150        rhs * self
151    }
152}
153
154// --- f64 * Variable → Expression ---
155
156impl Mul<Variable> for f64 {
157    type Output = Expression;
158    fn mul(self, rhs: Variable) -> Expression {
159        let mut coefficients = HashMap::new();
160        coefficients.insert(rhs, self);
161        Expression {
162            coefficients,
163            constant: 0.0,
164        }
165    }
166}
167
168impl Mul<f64> for Variable {
169    type Output = Expression;
170    fn mul(self, rhs: f64) -> Expression {
171        rhs * self
172    }
173}
174
175// --- Variable + Variable → Expression ---
176
177impl Add<Variable> for Variable {
178    type Output = Expression;
179    fn add(self, rhs: Variable) -> Expression {
180        Expression::from(self) + Expression::from(rhs)
181    }
182}
183
184impl Sub<Variable> for Variable {
185    type Output = Expression;
186    fn sub(self, rhs: Variable) -> Expression {
187        Expression::from(self) - Expression::from(rhs)
188    }
189}
190
191// --- Variable + Expression / Expression + Variable ---
192
193impl Add<Expression> for Variable {
194    type Output = Expression;
195    fn add(self, rhs: Expression) -> Expression {
196        Expression::from(self) + rhs
197    }
198}
199
200impl Add<Variable> for Expression {
201    type Output = Expression;
202    fn add(mut self, rhs: Variable) -> Expression {
203        *self.coefficients.entry(rhs).or_insert(0.0) += 1.0;
204        self
205    }
206}
207
208impl Sub<Expression> for Variable {
209    type Output = Expression;
210    fn sub(self, rhs: Expression) -> Expression {
211        Expression::from(self) - rhs
212    }
213}
214
215impl Sub<Variable> for Expression {
216    type Output = Expression;
217    fn sub(mut self, rhs: Variable) -> Expression {
218        *self.coefficients.entry(rhs).or_insert(0.0) -= 1.0;
219        self
220    }
221}
222
223// --- f64 + Expression / Expression + f64 ---
224
225impl Add<f64> for Expression {
226    type Output = Expression;
227    fn add(mut self, rhs: f64) -> Expression {
228        self.constant += rhs;
229        self
230    }
231}
232
233impl Add<Expression> for f64 {
234    type Output = Expression;
235    fn add(self, mut rhs: Expression) -> Expression {
236        rhs.constant += self;
237        rhs
238    }
239}
240
241impl Sub<f64> for Expression {
242    type Output = Expression;
243    fn sub(mut self, rhs: f64) -> Expression {
244        self.constant -= rhs;
245        self
246    }
247}
248
249impl Sub<Expression> for f64 {
250    type Output = Expression;
251    fn sub(self, rhs: Expression) -> Expression {
252        self + (-rhs)
253    }
254}
255
256// --- f64 + Variable / Variable + f64 ---
257
258impl Add<f64> for Variable {
259    type Output = Expression;
260    fn add(self, rhs: f64) -> Expression {
261        Expression::from(self) + rhs
262    }
263}
264
265impl Add<Variable> for f64 {
266    type Output = Expression;
267    fn add(self, rhs: Variable) -> Expression {
268        self + Expression::from(rhs)
269    }
270}
271
272impl Sub<f64> for Variable {
273    type Output = Expression;
274    fn sub(self, rhs: f64) -> Expression {
275        Expression::from(self) - rhs
276    }
277}
278
279impl Sub<Variable> for f64 {
280    type Output = Expression;
281    fn sub(self, rhs: Variable) -> Expression {
282        self - Expression::from(rhs)
283    }
284}