feat: expose the optimisation API in the python wrapper (#148)

The Python API does not enforce a particular optimisation routine. At
the moment, we support:
- LSU
- LUS

Both are, therefore, present in Python as well.

Author: maartenflippo

pumpkin-py/src/lib.rs

@@ -1,5 +1,6 @@
 mod constraints;
 mod model;
+mod optimisation;
 mod result;
 mod variables;
 
@@ -33,6 +34,7 @@ fn pumpkin_py(python: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<result::Solution>()?;
 
     submodule!(constraints, python, m);
+    submodule!(optimisation, python, m);
 
     Ok(())
 }

pumpkin-py/src/model.rs

@@ -2,6 +2,9 @@ use std::num::NonZero;
 use std::path::PathBuf;
 
 use pumpkin_solver::containers::KeyedVec;
+use pumpkin_solver::optimisation::linear_sat_unsat::LinearSatUnsat;
+use pumpkin_solver::optimisation::linear_unsat_sat::LinearUnsatSat;
+use pumpkin_solver::optimisation::OptimisationDirection;
 use pumpkin_solver::options::SolverOptions;
 use pumpkin_solver::predicate;
 use pumpkin_solver::proof::Format;
@@ -14,6 +17,9 @@ use pumpkin_solver::Solver;
 use pyo3::prelude::*;
 
 use crate::constraints::Constraint;
+use crate::optimisation::Direction;
+use crate::optimisation::OptimisationResult;
+use crate::optimisation::Optimiser;
 use crate::result::SatisfactionResult;
 use crate::result::Solution;
 use crate::variables::BoolExpression;
@@ -163,27 +169,9 @@ impl Model {
 
     #[pyo3(signature = (proof=None))]
     fn satisfy(&self, proof: Option<PathBuf>) -> SatisfactionResult {
-        let proof_log = proof
-            .map(|path| ProofLog::cp(&path, Format::Text, true, true))
-            .transpose()
-            .map(|proof| proof.unwrap_or_default())
-            .expect("failed to create proof file");
-
-        let options = SolverOptions {
-            proof_log,
-            ..Default::default()
-        };
-
-        let mut solver = Solver::with_options(options);
-
-        let solver_setup = self
-            .create_variable_map(&mut solver)
-            .and_then(|variable_map| {
-                self.post_constraints(&mut solver, &variable_map)?;
-                Ok(variable_map)
-            });
+        let solver_setup = self.create_solver(proof);
 
-        let Ok(variable_map) = solver_setup else {
+        let Ok((mut solver, variable_map)) = solver_setup else {
             return SatisfactionResult::Unsatisfiable();
         };
 
@@ -202,6 +190,62 @@ impl Model {
             pumpkin_solver::results::SatisfactionResult::Unknown => SatisfactionResult::Unknown(),
         }
     }
+
+    #[pyo3(signature = (objective, optimiser=Optimiser::LinearSatUnsat, direction=Direction::Minimise, proof=None))]
+    fn optimise(
+        &self,
+        objective: IntExpression,
+        optimiser: Optimiser,
+        direction: Direction,
+        proof: Option<PathBuf>,
+    ) -> OptimisationResult {
+        let solver_setup = self.create_solver(proof);
+
+        let Ok((mut solver, variable_map)) = solver_setup else {
+            return OptimisationResult::Unsatisfiable();
+        };
+
+        let mut brancher = solver.default_brancher();
+
+        let direction = match direction {
+            Direction::Minimise => OptimisationDirection::Minimise,
+            Direction::Maximise => OptimisationDirection::Maximise,
+        };
+
+        let objective = objective.to_affine_view(&variable_map);
+
+        let result = match optimiser {
+            Optimiser::LinearSatUnsat => solver.optimise(
+                &mut brancher,
+                &mut Indefinite,
+                LinearSatUnsat::new(direction, objective, |_, _| {}),
+            ),
+            Optimiser::LinearUnsatSat => solver.optimise(
+                &mut brancher,
+                &mut Indefinite,
+                LinearUnsatSat::new(direction, objective, |_, _| {}),
+            ),
+        };
+
+        match result {
+            pumpkin_solver::results::OptimisationResult::Satisfiable(solution) => {
+                OptimisationResult::Satisfiable(Solution {
+                    solver_solution: solution,
+                    variable_map,
+                })
+            }
+            pumpkin_solver::results::OptimisationResult::Optimal(solution) => {
+                OptimisationResult::Optimal(Solution {
+                    solver_solution: solution,
+                    variable_map,
+                })
+            }
+            pumpkin_solver::results::OptimisationResult::Unsatisfiable => {
+                OptimisationResult::Unsatisfiable()
+            }
+            pumpkin_solver::results::OptimisationResult::Unknown => OptimisationResult::Unknown(),
+        }
+    }
 }
 
 impl Model {
@@ -252,6 +296,33 @@ impl Model {
 
         Ok(())
     }
+
+    fn create_solver(
+        &self,
+        proof: Option<PathBuf>,
+    ) -> Result<(Solver, VariableMap), ConstraintOperationError> {
+        let proof_log = proof
+            .map(|path| ProofLog::cp(&path, Format::Text, true, true))
+            .transpose()
+            .map(|proof| proof.unwrap_or_default())
+            .expect("failed to create proof file");
+
+        let options = SolverOptions {
+            proof_log,
+            ..Default::default()
+        };
+
+        let mut solver = Solver::with_options(options);
+
+        let variable_map = self
+            .create_variable_map(&mut solver)
+            .and_then(|variable_map| {
+                self.post_constraints(&mut solver, &variable_map)?;
+                Ok(variable_map)
+            })?;
+
+        Ok((solver, variable_map))
+    }
 }
 
 #[derive(Clone)]

pumpkin-py/src/optimisation.rs

@@ -0,0 +1,36 @@
+use pyo3::prelude::*;
+
+use crate::result::Solution;
+
+#[pyclass]
+pub enum OptimisationResult {
+    /// The problem was solved to optimality, and the solution is an optimal one.
+    Optimal(Solution),
+    /// At least one solution was identified, and the solution is the best one.
+    Satisfiable(Solution),
+    /// The problem was unsatisfiable.
+    Unsatisfiable(),
+    /// None of the other variants were concluded.
+    Unknown(),
+}
+
+#[pyclass(eq, eq_int)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum Optimiser {
+    LinearSatUnsat,
+    LinearUnsatSat,
+}
+
+#[pyclass(eq, eq_int)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum Direction {
+    Minimise,
+    Maximise,
+}
+
+pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<Optimiser>()?;
+    m.add_class::<Direction>()?;
+    m.add_class::<OptimisationResult>()?;
+    Ok(())
+}

pumpkin-py/tests/test_optimisation.py

@@ -0,0 +1,30 @@
+from pumpkin_py import Model
+from pumpkin_py.optimisation import Direction, OptimisationResult
+
+
+def test_linear_sat_unsat_minimisation():
+    model = Model()
+
+    objective = model.new_integer_variable(1, 5, name="objective")
+
+    result = model.optimise(objective, direction=Direction.Minimise)
+
+    assert isinstance(result, OptimisationResult.Optimal)
+
+    solution = result._0
+    assert solution.int_value(objective) == 1
+
+
+def test_linear_sat_unsat_maximisation():
+    model = Model()
+
+    objective = model.new_integer_variable(1, 5, name="objective")
+
+    result = model.optimise(objective, direction=Direction.Maximise)
+
+    assert isinstance(result, OptimisationResult.Optimal)
+
+    solution = result._0
+    assert solution.int_value(objective) == 5
+
+

pumpkin-solver/src/bin/pumpkin-solver/flatzinc/mod.rs

@@ -18,7 +18,6 @@ use pumpkin_solver::constraints::cumulative;
 use pumpkin_solver::optimisation::linear_sat_unsat::LinearSatUnsat;
 use pumpkin_solver::optimisation::linear_unsat_sat::LinearUnsatSat;
 use pumpkin_solver::optimisation::OptimisationDirection;
-use pumpkin_solver::optimisation::OptimisationProcedure;
 use pumpkin_solver::optimisation::OptimisationStrategy;
 use pumpkin_solver::options::CumulativeOptions;
 use pumpkin_solver::results::solution_iterator::IteratedSolution;

pumpkin-solver/src/optimisation/linear_sat_unsat.rs

@@ -12,13 +12,34 @@ use crate::variables::IntegerVariable;
 use crate::ConstraintOperationError;
 use crate::Solver;
 
+/// Implements the linear SAT-UNSAT (LSU) optimisation procedure.
 #[derive(Debug, Clone, Copy)]
-pub struct LinearSatUnsat<Var: IntegerVariable, Callback> {
+pub struct LinearSatUnsat<Var, Callback> {
     direction: OptimisationDirection,
     objective: Var,
     solution_callback: Callback,
 }
 
+impl<Var, Callback> LinearSatUnsat<Var, Callback>
+where
+    // The trait bound here is not common; see
+    // linear_unsat_sat for more info.
+    Callback: Fn(&Solver, SolutionReference),
+{
+    /// Create a new instance of [`LinearSatUnsat`].
+    pub fn new(
+        direction: OptimisationDirection,
+        objective: Var,
+        solution_callback: Callback,
+    ) -> Self {
+        Self {
+            direction,
+            objective,
+            solution_callback,
+        }
+    }
+}
+
 impl<Var: IntegerVariable, Callback> LinearSatUnsat<Var, Callback> {
     /// Given the current objective value `best_objective_value`, it adds a constraint specifying
     /// that the objective value should be at most `best_objective_value - 1`. Note that it is
@@ -64,14 +85,6 @@ where
     Var: IntegerVariable,
     Callback: Fn(&Solver, SolutionReference),
 {
-    fn new(direction: OptimisationDirection, objective: Var, solution_callback: Callback) -> Self {
-        Self {
-            direction,
-            objective,
-            solution_callback,
-        }
-    }
-
     fn optimise(
         &mut self,
         brancher: &mut impl Brancher,

pumpkin-solver/src/optimisation/linear_unsat_sat.rs

@@ -12,24 +12,46 @@ use crate::termination::TerminationCondition;
 use crate::variables::IntegerVariable;
 use crate::Solver;
 
+/// Implements the linear UNSAT-SAT (LUS) optimisation procedure.
 #[derive(Debug, Clone, Copy)]
-pub struct LinearUnsatSat<Var: IntegerVariable, Callback> {
+pub struct LinearUnsatSat<Var, Callback> {
     direction: OptimisationDirection,
     objective: Var,
     solution_callback: Callback,
 }
 
-impl<Var: IntegerVariable, Callback: Fn(&Solver, SolutionReference)>
-    OptimisationProcedure<Var, Callback> for LinearUnsatSat<Var, Callback>
+impl<Var, Callback> LinearUnsatSat<Var, Callback>
+where
+    // The trait bound here is contrary to common
+    // practice; typically the bounds are only enforced
+    // where they are required (in this case, in the
+    // implementation of OptimisationProcedure).
+    //
+    // However, if we don't have the trait bound here,
+    // the compiler may implement `FnOnce` for the
+    // empty closure, which causes problems. So, we
+    // have the hint here.
+    //
+    // Similar is also the case in linear SAT-UNSAT.
+    Callback: Fn(&Solver, SolutionReference),
 {
-    fn new(direction: OptimisationDirection, objective: Var, solution_callback: Callback) -> Self {
+    /// Create a new instance of [`LinearUnsatSat`].
+    pub fn new(
+        direction: OptimisationDirection,
+        objective: Var,
+        solution_callback: Callback,
+    ) -> Self {
         Self {
             direction,
             objective,
             solution_callback,
         }
     }
+}
 
+impl<Var: IntegerVariable, Callback: Fn(&Solver, SolutionReference)>
+    OptimisationProcedure<Var, Callback> for LinearUnsatSat<Var, Callback>
+{
     fn optimise(
         &mut self,
         brancher: &mut impl Brancher,

pumpkin-solver/src/optimisation/mod.rs

@@ -15,8 +15,6 @@ pub mod linear_sat_unsat;
 pub mod linear_unsat_sat;
 
 pub trait OptimisationProcedure<Var: IntegerVariable, Callback: Fn(&Solver, SolutionReference)> {
-    fn new(direction: OptimisationDirection, objective: Var, solution_callback: Callback) -> Self;
-
     fn optimise(
         &mut self,
         brancher: &mut impl Brancher,