feat: implement predicate and assumption solving in python wrapper (#150)

Assumption solving in the Rust layer currently performs semantic
minimisation of the core. E.g. if the assumptions contain `[y <= 1]` and
`[y != 0]`, and the domain of `y` starts at 0, then the core may contain
`[y == 1]` rather than the original predicates in the assumptions.

Author: maartenflippo

pumpkin-py/src/lib.rs

@@ -28,9 +28,11 @@ macro_rules! submodule {
 fn pumpkin_py(python: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<variables::IntExpression>()?;
     m.add_class::<variables::BoolExpression>()?;
-    m.add_class::<model::Comparator>()?;
+    m.add_class::<variables::Comparator>()?;
+    m.add_class::<variables::Predicate>()?;
     m.add_class::<model::Model>()?;
     m.add_class::<result::SatisfactionResult>()?;
+    m.add_class::<result::SatisfactionUnderAssumptionsResult>()?;
     m.add_class::<result::Solution>()?;
 
     submodule!(constraints, python, m);

pumpkin-py/src/model.rs

@@ -21,11 +21,13 @@ use crate::optimisation::Direction;
 use crate::optimisation::OptimisationResult;
 use crate::optimisation::Optimiser;
 use crate::result::SatisfactionResult;
+use crate::result::SatisfactionUnderAssumptionsResult;
 use crate::result::Solution;
 use crate::variables::BoolExpression;
 use crate::variables::BoolVariable;
 use crate::variables::IntExpression;
 use crate::variables::IntVariable;
+use crate::variables::Predicate;
 use crate::variables::VariableMap;
 
 #[pyclass]
@@ -36,15 +38,6 @@ pub struct Model {
     constraints: Vec<ModelConstraint>,
 }
 
-#[pyclass(eq, eq_int)]
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub enum Comparator {
-    NotEqual,
-    Equal,
-    LessThanOrEqual,
-    GreaterThanOrEqual,
-}
-
 #[pymethods]
 impl Model {
     #[new]
@@ -121,23 +114,13 @@ impl Model {
         }
     }
 
-    #[pyo3(signature = (integer, comparator, value, name=None))]
-    fn predicate_as_boolean(
-        &mut self,
-        integer: IntExpression,
-        comparator: Comparator,
-        value: i32,
-        name: Option<&str>,
-    ) -> BoolExpression {
+    #[pyo3(signature = (predicate, name=None))]
+    fn predicate_as_boolean(&mut self, predicate: Predicate, name: Option<&str>) -> BoolExpression {
         self.boolean_variables
             .push(ModelBoolVar {
                 name: name.map(|n| n.to_owned()),
                 integer_equivalent: None,
-                predicate: Some(Predicate {
-                    integer,
-                    comparator,
-                    value,
-                }),
+                predicate: Some(predicate),
             })
             .into()
     }
@@ -191,6 +174,70 @@ impl Model {
         }
     }
 
+    #[pyo3(signature = (assumptions))]
+    fn satisfy_under_assumptions(
+        &self,
+        assumptions: Vec<Predicate>,
+    ) -> SatisfactionUnderAssumptionsResult {
+        let solver_setup = self.create_solver(None);
+
+        let Ok((mut solver, variable_map)) = solver_setup else {
+            return SatisfactionUnderAssumptionsResult::Unsatisfiable();
+        };
+
+        let mut brancher = solver.default_brancher();
+
+        let solver_assumptions = assumptions
+            .iter()
+            .map(|pred| pred.to_solver_predicate(&variable_map))
+            .collect::<Vec<_>>();
+
+        // Maarten: I do not understand why it is necessary, but we have to create a local variable
+        // here     that is the result of the `match` statement. Otherwise the compiler
+        // complains that     `solver` and `brancher` potentially do not live long enough.
+        //
+        //     Ideally this would not be necessary, but perhaps it is unavoidable with the setup we
+        //     currently have. Either way, we take the suggestion by the compiler.
+        let result = match solver.satisfy_under_assumptions(&mut brancher, &mut Indefinite, &solver_assumptions) {
+            pumpkin_solver::results::SatisfactionResultUnderAssumptions::Satisfiable(solution) => {
+                SatisfactionUnderAssumptionsResult::Satisfiable(Solution {
+                    solver_solution: solution,
+                    variable_map,
+                })
+            }
+            pumpkin_solver::results::SatisfactionResultUnderAssumptions::UnsatisfiableUnderAssumptions(mut result) => {
+                // Maarten: For now we assume that the core _must_ consist of the predicates that
+                //     were the input to the solve call. In general this is not the case, e.g. when
+                //     the assumptions can be semantically minized (the assumptions [y <= 1],
+                //     [y >= 0] and [y != 0] will be compressed to [y == 1] which would end up in
+                //     the core).
+                //
+                //     In the future, perhaps we should make the distinction between predicates and
+                //     literals in the python wrapper as well. For now, this is the simplest way
+                //     forward. I expect that the situation above almost never happens in practice.
+                let core = result
+                    .extract_core()
+                    .iter()
+                    .map(|predicate| assumptions
+                         .iter()
+                         .find(|pred| pred.to_solver_predicate(&variable_map) == *predicate)
+                         .copied()
+                         .expect("predicates in core must be part of the assumptions"))
+                    .collect();
+
+                SatisfactionUnderAssumptionsResult::UnsatisfiableUnderAssumptions(core)
+            }
+            pumpkin_solver::results::SatisfactionResultUnderAssumptions::Unsatisfiable => {
+                SatisfactionUnderAssumptionsResult::Unsatisfiable()
+            }
+            pumpkin_solver::results::SatisfactionResultUnderAssumptions::Unknown => {
+                SatisfactionUnderAssumptionsResult::Unknown()
+            }
+        };
+
+        result
+    }
+
     #[pyo3(signature = (objective, optimiser=Optimiser::LinearSatUnsat, direction=Direction::Minimise, proof=None))]
     fn optimise(
         &self,
@@ -411,26 +458,3 @@ impl ModelBoolVar {
         Ok(literal)
     }
 }
-
-struct Predicate {
-    integer: IntExpression,
-    comparator: Comparator,
-    value: i32,
-}
-
-impl Predicate {
-    /// Convert the predicate in the model domain to a predicate in the solver domain.
-    fn to_solver_predicate(
-        &self,
-        variable_map: &VariableMap,
-    ) -> pumpkin_solver::predicates::Predicate {
-        let affine_view = self.integer.to_affine_view(variable_map);
-
-        match self.comparator {
-            Comparator::NotEqual => predicate![affine_view != self.value],
-            Comparator::Equal => predicate![affine_view == self.value],
-            Comparator::LessThanOrEqual => predicate![affine_view <= self.value],
-            Comparator::GreaterThanOrEqual => predicate![affine_view >= self.value],
-        }
-    }
-}

pumpkin-py/src/result.rs

@@ -3,6 +3,7 @@ use pyo3::prelude::*;
 
 use crate::variables::BoolExpression;
 use crate::variables::IntExpression;
+use crate::variables::Predicate;
 use crate::variables::VariableMap;
 
 #[pyclass]
@@ -13,6 +14,15 @@ pub enum SatisfactionResult {
     Unknown(),
 }
 
+#[pyclass]
+#[allow(clippy::large_enum_variant)]
+pub enum SatisfactionUnderAssumptionsResult {
+    Satisfiable(Solution),
+    UnsatisfiableUnderAssumptions(Vec<Predicate>),
+    Unsatisfiable(),
+    Unknown(),
+}
+
 #[pyclass]
 #[derive(Clone)]
 pub struct Solution {
@@ -32,3 +42,7 @@ impl Solution {
             .get_literal_value(variable.to_literal(&self.variable_map))
     }
 }
+
+#[pyclass]
+#[derive(Clone)]
+pub struct CoreExtractor {}

pumpkin-py/src/variables.rs

@@ -1,12 +1,13 @@
 use pumpkin_solver::containers::KeyedVec;
 use pumpkin_solver::containers::StorageKey;
+use pumpkin_solver::predicate;
 use pumpkin_solver::variables::AffineView;
 use pumpkin_solver::variables::DomainId;
 use pumpkin_solver::variables::Literal;
 use pumpkin_solver::variables::TransformableVariable;
 use pyo3::prelude::*;
 
-#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct IntVariable(usize);
 
 impl StorageKey for IntVariable {
@@ -19,8 +20,8 @@ impl StorageKey for IntVariable {
     }
 }
 
-#[pyclass]
-#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
+#[pyclass(eq, hash, frozen)]
+#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct IntExpression {
     pub variable: IntVariable,
     pub offset: i32,
@@ -83,6 +84,52 @@ impl IntExpression {
     }
 }
 
+#[pyclass(eq, eq_int, hash, frozen)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum Comparator {
+    NotEqual,
+    Equal,
+    LessThanOrEqual,
+    GreaterThanOrEqual,
+}
+
+#[pyclass(eq, get_all, hash, frozen)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct Predicate {
+    pub variable: IntExpression,
+    pub comparator: Comparator,
+    pub value: i32,
+}
+
+#[pymethods]
+impl Predicate {
+    #[new]
+    fn new(variable: IntExpression, comparator: Comparator, value: i32) -> Self {
+        Self {
+            variable,
+            comparator,
+            value,
+        }
+    }
+}
+
+impl Predicate {
+    /// Convert the predicate in the model domain to a predicate in the solver domain.
+    pub(crate) fn to_solver_predicate(
+        self,
+        variable_map: &VariableMap,
+    ) -> pumpkin_solver::predicates::Predicate {
+        let affine_view = self.variable.to_affine_view(variable_map);
+
+        match self.comparator {
+            Comparator::NotEqual => predicate![affine_view != self.value],
+            Comparator::Equal => predicate![affine_view == self.value],
+            Comparator::LessThanOrEqual => predicate![affine_view <= self.value],
+            Comparator::GreaterThanOrEqual => predicate![affine_view >= self.value],
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
 pub struct BoolVariable(usize);
 
@@ -96,7 +143,7 @@ impl StorageKey for BoolVariable {
     }
 }
 
-#[pyclass]
+#[pyclass(eq)]
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
 pub struct BoolExpression(BoolVariable, bool);
 

pumpkin-py/tests/test_assumptions.py

@@ -0,0 +1,36 @@
+from pumpkin_py import Comparator, Model, Predicate, SatisfactionUnderAssumptionsResult
+from pumpkin_py.constraints import LessThanOrEquals
+
+
+def test_assumptions_are_respected():
+    model = Model()
+
+    x = model.new_integer_variable(1, 5, name="x")
+
+    assumption = Predicate(x, Comparator.LessThanOrEqual, 3)
+
+    result = model.satisfy_under_assumptions([assumption])
+    assert isinstance(result, SatisfactionUnderAssumptionsResult.Satisfiable)
+
+    solution = result._0
+    x_value = solution.int_value(x)
+    assert x_value <= 3
+
+
+def test_core_extraction():
+    model = Model()
+
+    x = model.new_integer_variable(1, 5, name="x")
+    y = model.new_integer_variable(1, 5, name="x")
+
+    x_ge_3 = Predicate(x, Comparator.GreaterThanOrEqual, 3)
+    y_ge_3 = Predicate(y, Comparator.GreaterThanOrEqual, 3)
+
+    model.add_constraint(LessThanOrEquals([x, y], 5))
+
+    result = model.satisfy_under_assumptions([x_ge_3, y_ge_3])
+    assert isinstance(result, SatisfactionUnderAssumptionsResult.UnsatisfiableUnderAssumptions)
+
+    core = set(result._0)
+    assert set([x_ge_3, y_ge_3]) == core
+