From 04196f41d3774ab693ed2bd1af616189ea959b03 Mon Sep 17 00:00:00 2001 From: Guillaume Pinot Date: Wed, 8 Sep 2021 21:57:12 +0200 Subject: [PATCH] Linear constraint support in the builder --- .gitignore | 1 + src/builder.rs | 305 ++++++++++++++++++++++++++++++++-------- tests/tutorial_tests.rs | 59 ++++++++ 3 files changed, 308 insertions(+), 57 deletions(-) create mode 100644 tests/tutorial_tests.rs diff --git a/.gitignore b/.gitignore index 96ef6c0..80aca69 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /target Cargo.lock +*~ diff --git a/src/builder.rs b/src/builder.rs index 773a4e1..360650a 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -1,4 +1,5 @@ use crate::proto; +use proto::constraint_proto::Constraint as CstEnum; #[derive(Default, Debug)] pub struct CpModelBuilder { @@ -43,76 +44,149 @@ impl CpModelBuilder { } pub fn add_or(&mut self, vars: impl IntoIterator) -> Constraint { - let index = self.proto.constraints.len(); - self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::BoolOr( - proto::BoolArgumentProto { - literals: vars.into_iter().map(|v| v.0).collect(), - }, - )), - ..Default::default() - }); - Constraint(index) + self.add_cst(CstEnum::BoolOr(proto::BoolArgumentProto { + literals: vars.into_iter().map(|v| v.0).collect(), + })) } pub fn add_and(&mut self, vars: impl IntoIterator) -> Constraint { - let index = self.proto.constraints.len(); - self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::BoolAnd( - proto::BoolArgumentProto { - literals: vars.into_iter().map(|v| v.0).collect(), - }, - )), - ..Default::default() - }); - Constraint(index) + self.add_cst(CstEnum::BoolAnd(proto::BoolArgumentProto { + literals: vars.into_iter().map(|v| v.0).collect(), + })) } pub fn add_at_most_one(&mut self, vars: impl IntoIterator) -> Constraint { - let index = self.proto.constraints.len(); - self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::AtMostOne( - proto::BoolArgumentProto { - literals: vars.into_iter().map(|v| v.0).collect(), - }, - )), - ..Default::default() - }); - Constraint(index) + self.add_cst(CstEnum::AtMostOne(proto::BoolArgumentProto { + literals: vars.into_iter().map(|v| v.0).collect(), + })) } pub fn add_exactly_one(&mut self, vars: impl IntoIterator) -> Constraint { - let index = self.proto.constraints.len(); - self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::ExactlyOne( - proto::BoolArgumentProto { - literals: vars.into_iter().map(|v| v.0).collect(), - }, - )), - ..Default::default() - }); - Constraint(index) + self.add_cst(CstEnum::ExactlyOne(proto::BoolArgumentProto { + literals: vars.into_iter().map(|v| v.0).collect(), + })) } pub fn add_xor(&mut self, vars: impl IntoIterator) -> Constraint { + self.add_cst(CstEnum::BoolXor(proto::BoolArgumentProto { + literals: vars.into_iter().map(|v| v.0).collect(), + })) + } + pub fn add_all_different(&mut self, vars: impl IntoIterator) -> Constraint { + self.add_cst(CstEnum::AllDiff(proto::AllDifferentConstraintProto { + vars: vars.into_iter().map(|v| v.0).collect(), + })) + } + pub fn add_linear_constraint( + &mut self, + expr: &LinearExpr, + (begin, end): (i64, i64), + ) -> Constraint { + self.add_cst(CstEnum::Linear(proto::LinearConstraintProto { + vars: expr.vars.clone(), + coeffs: expr.coeffs.clone(), + domain: vec![begin - expr.constant, end - expr.constant], + })) + } + pub fn add_eq, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_eq_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_eq_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([val, val]); + self.add_cst(CstEnum::Linear(cst)) + } + pub fn add_ge, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_ge_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_ge_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([val, i64::MAX]); + self.add_cst(CstEnum::Linear(cst)) + } + pub fn add_le, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_le_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_le_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([i64::MIN, val]); + self.add_cst(CstEnum::Linear(cst)) + } + pub fn add_gt, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_gt_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_gt_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([val + 1, i64::MAX]); + self.add_cst(CstEnum::Linear(cst)) + } + pub fn add_lt, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_lt_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_lt_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([i64::MIN, val - 1]); + self.add_cst(CstEnum::Linear(cst)) + } + pub fn add_ne, U: Into>( + &mut self, + lhs: T, + rhs: U, + ) -> Constraint { + self.add_ne_by_ref(&lhs.into(), &rhs.into()) + } + pub fn add_ne_by_ref(&mut self, lhs: &LinearExpr, rhs: &LinearExpr) -> Constraint { + let (mut cst, val) = create_linear_proto(lhs, rhs); + cst.domain.extend([i64::MIN, val - 1, val + 1, i64::MAX]); + self.add_cst(CstEnum::Linear(cst)) + } + fn add_cst(&mut self, cst: CstEnum) -> Constraint { let index = self.proto.constraints.len(); self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::BoolXor( - proto::BoolArgumentProto { - literals: vars.into_iter().map(|v| v.0).collect(), - }, - )), + constraint: Some(cst), ..Default::default() }); Constraint(index) } - pub fn add_all_different(&mut self, vars: impl IntoIterator) -> Constraint { - let index = self.proto.constraints.len(); - self.proto.constraints.push(proto::ConstraintProto { - constraint: Some(proto::constraint_proto::Constraint::AllDiff( - proto::AllDifferentConstraintProto { - vars: vars.into_iter().map(|v| v.0).collect(), - }, - )), - ..Default::default() + + pub fn minimize>(&mut self, expr: T) { + let expr = expr.into(); + self.proto.objective = Some(proto::CpObjectiveProto { + vars: expr.vars, + coeffs: expr.coeffs, + offset: expr.constant as f64, + scaling_factor: 1., + domain: vec![], + }); + } + pub fn maximize>(&mut self, expr: T) { + let mut expr = expr.into(); + for coeff in &mut expr.coeffs { + *coeff *= -1; + } + self.proto.objective = Some(proto::CpObjectiveProto { + vars: expr.vars, + coeffs: expr.coeffs, + offset: -expr.constant as f64, + scaling_factor: -1., + domain: vec![], }); - Constraint(index) } pub fn stats(&self) -> String { @@ -161,13 +235,130 @@ impl From for IntVar { impl IntVar { pub fn solution_value(self, response: &proto::CpSolverResponse) -> i64 { if self.0 < 0 { - use std::ops::Not; - 1 - IntVar::from(BoolVar(self.0).not()).solution_value(response) + 1 - self.not().solution_value(response) } else { response.solution[self.0 as usize] } } + fn not(self) -> Self { + IntVar::from(!BoolVar(self.0)) + } } #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Constraint(usize); + +#[derive(Clone, Default, Debug)] +pub struct LinearExpr { + vars: Vec, + coeffs: Vec, + constant: i64, +} +impl std::ops::AddAssign for LinearExpr { + fn add_assign(&mut self, rhs: i64) { + self.constant += rhs; + } +} +impl> std::ops::AddAssign<(i64, V)> for LinearExpr { + fn add_assign(&mut self, (coeff, var): (i64, V)) { + let var = var.into(); + if var.0 < 0 { + self.vars.push(var.not().0); + self.coeffs.push(-coeff); + self.constant += coeff; + } else { + self.vars.push(var.0); + self.coeffs.push(coeff); + } + } +} +impl> std::ops::AddAssign for LinearExpr { + fn add_assign(&mut self, var: V) { + *self += (1, var); + } +} +impl> std::iter::Extend<(i64, V)> for LinearExpr { + fn extend>(&mut self, iter: I) { + for (coeff, var) in iter.into_iter() { + *self += (coeff, var.into()); + } + } +} +impl> std::iter::FromIterator<(i64, V)> for LinearExpr { + fn from_iter>(iter: I) -> Self { + let mut res = Self::default(); + res.extend(iter); + res + } +} +impl> std::iter::Extend for LinearExpr { + fn extend>(&mut self, iter: I) { + for var in iter.into_iter() { + *self += var.into(); + } + } +} +impl> std::iter::FromIterator for LinearExpr { + fn from_iter>(iter: I) -> Self { + let mut res = Self::default(); + res.extend(iter); + res + } +} +impl> From for LinearExpr { + fn from(var: V) -> Self { + let mut res = Self::default(); + res += var; + res + } +} +impl From for LinearExpr { + fn from(constant: i64) -> Self { + let mut res = Self::default(); + res += constant; + res + } +} +impl, const L: usize> From<[(i64, V); L]> for LinearExpr { + fn from(expr: [(i64, V); L]) -> Self { + let mut res = Self::default(); + for term in expr { + res += term; + } + res + } +} +impl std::ops::Add for LinearExpr +where + LinearExpr: std::ops::AddAssign, +{ + type Output = LinearExpr; + fn add(mut self, rhs: T) -> Self::Output { + self += rhs; + self + } +} + +fn create_linear_proto( + left: &LinearExpr, + right: &LinearExpr, +) -> (proto::LinearConstraintProto, i64) { + ( + proto::LinearConstraintProto { + vars: left + .vars + .iter() + .copied() + .chain(right.vars.iter().copied()) + .collect(), + coeffs: left + .coeffs + .iter() + .copied() + .chain(right.coeffs.iter().map(|&c| -c)) + .collect(), + domain: vec![], + }, + right.constant - left.constant, + ) +} diff --git a/tests/tutorial_tests.rs b/tests/tutorial_tests.rs new file mode 100644 index 0000000..d4384b0 --- /dev/null +++ b/tests/tutorial_tests.rs @@ -0,0 +1,59 @@ +use cp_sat::builder::CpModelBuilder; +use cp_sat::proto::CpSolverStatus; + +#[test] +fn presentation_problem() { + let mut model = CpModelBuilder::default(); + let domain = [(0, 2)]; + let x = model.new_int_var_with_name(domain, "x"); + let y = model.new_int_var_with_name(domain, "y"); + let z = model.new_int_var_with_name(domain, "z"); + + model.add_ne(x, y); + + let response = model.solve(); + println!("{:?}", response); + + assert_eq!(response.status(), CpSolverStatus::Optimal); + let x_val = x.solution_value(&response); + let y_val = y.solution_value(&response); + let z_val = z.solution_value(&response); + println!("x = {}", x_val); + println!("y = {}", y_val); + println!("z = {}", z_val); + + assert!(x_val != y_val); + assert!(0 <= x_val && x_val <= 2); + assert!(0 <= y_val && y_val <= 2); + assert!(0 <= z_val && z_val <= 2); +} + +#[test] +fn solving_a_cp_problem() { + let mut model = CpModelBuilder::default(); + let var_upper_bound = vec![50, 45, 37].into_iter().max().unwrap(); + let x = model.new_int_var_with_name([(0, var_upper_bound)], "x"); + let y = model.new_int_var_with_name([(0, var_upper_bound)], "y"); + let z = model.new_int_var_with_name([(0, var_upper_bound)], "z"); + + model.add_le([(2, x), (7, y), (3, z)], 50); + model.add_le([(3, x), (-5, y), (7, z)], 45); + model.add_le([(5, x), (2, y), (-6, z)], 37); + + model.maximize([(2, x), (2, y), (3, z)]); + + let response = model.solve(); + assert_eq!(response.status(), CpSolverStatus::Optimal); + let x_val = x.solution_value(&response); + let y_val = y.solution_value(&response); + let z_val = z.solution_value(&response); + println!("objective: {}", response.objective_value); + println!("x = {}", x_val); + println!("y = {}", y_val); + println!("z = {}", z_val); + + assert_eq!(35., response.objective_value); + assert!(2 * x_val + 7 * y_val + 3 * z_val <= 50); + assert!(3 * x_val - 5 * y_val + 7 * z_val <= 45); + assert!(5 * x_val + 2 * y_val - 6 * z_val <= 37); +}