CP-SAT: performance tips for the enumeration of all solutions

[tanlin2013]

What version of OR-Tools and what language are you using?
Version: 9.0.9048
Language: Python

Which solver are you using (e.g. CP-SAT, Routing Solver, GLOP, BOP, Gurobi)
CP-SAT

What operating system (Linux, Windows, ...) and version?
macOS Sonoma 14.4.1

What did you do?
Hi there, I'm working on a sudoku-like constraint problem that requires the enumeration of all possible solutions. I wonder are there any general performance tips for this task? It might be problem-specific indeed. At least parallelism doesn't seem to be an option, then what are the other things I could have tried from the broader aspect?

The problem I consider is a 2d grid of size (Lx, Ly), the variables are defined on the edges, being either 0 or 1 (boolean), drawn as the arrows in dark gray. Additionally, periodic boundaries are assumed.

There are 2 constraints,

• The overall fluxes at each vertex node (in orange), must sum up to some given integer (called charge) ranging from -2 to 2. In the example here, it is always 2-in-2-out on every node, corresponding to the charge 0. So charge = up_edge + right_edge - down_edge - left_edge.

• The sum of edges in every row or column, must equal to a given tuple of integers (W_row, W_col). In this example, it will be (0, 0) for both rows and columns. Remember the periodic boundaries.

What did you expect to see

Of course, the number of solutions grows really fast in an exponential manner. So is the time required...
(This table assumes all charges are 0, and (W_row, W_col) = (0, 0).)

\size	(4, 4)	(6, 4)	(8, 4)	(6, 6)
n_sols	990	32,810	1,159,166	5,482,716
cpu time	< 1s	13s	20m	6h

Code
Below is the essential part of the implementation, or the full implementation here.

(i, j, k) is the coordinate system I use to label each edge. We can basically assign every up_edge and right_edge to the vertex on position (i, j), and k for these 2 edges.

@dataclass(slots=True)
class CpModel:
    shape: Tuple[int, int]
    charge_distri: npt.NDArray[np.int64]
    flux_sector: Optional[Tuple[int, int]] = field(default=None)
    _model: cp_model.CpModel = field(init=False, repr=False)
    _solver: cp_model.CpSolver = field(init=False, repr=False)
    _links: Dict = field(default_factory=dict, init=False, repr=False)
    _callback: SolutionCallback = field(init=False, repr=False)

    def __post_init__(self):
        self._model = cp_model.CpModel()
        self._solver = cp_model.CpSolver()
        for j, i, k in product(range(self.shape[1]), range(self.shape[0]), range(2)):
            self._links[(i, j, k)] = self._model.NewIntVar(0, 1, f"link_{i}_{j}_{k}")
        self.constraint_gauss_law()
        self.constraint_flux_sector()
        self._callback = SolutionCallback(self._links)

    def constraint_gauss_law(self) -> None:
        for i, j in product(range(self.shape[0]), range(self.shape[1])):
            self._model.Add(
                self._links[(i, j, 0)]
                + self._links[(i, j, 1)]
                - self._links[((i - 1) % self.shape[0], j, 0)]
                - self._links[(i, (j - 1) % self.shape[1], 1)]
                == self.charge_distri[i, j]
            )

    def constraint_flux_sector(self) -> None:
        if self.flux_sector is not None:
            if (np.array(self.shape) % 2 != 0).any():
                raise ValueError("The shape of the lattice must be even.")
            for i in range(self.shape[0]):
                self._model.Add(
                    sum(self._links[(i, j, 0)] for j in range(self.shape[1])) - self.shape[1] // 2
                    == self.flux_sector[0]
                )
            for j in range(self.shape[1]):
                self._model.Add(
                    sum(self._links[(i, j, 1)] for i in range(self.shape[0])) - self.shape[0] // 2
                    == self.flux_sector[1]
                )

    @property
    def n_solutions(self) -> int:
        return self._callback.n_solutions

    def solve(self, all_solutions: bool = True, log_search_progress: bool = False) -> None:
        self._solver.parameters.enumerate_all_solutions = all_solutions
        self._solver.parameters.log_search_progress = log_search_progress
        status = self._solver.solve(self._model, self._callback)

[lperron]

Look for embarrassingly parallel search.

Basically select k variables, create all possible assignments, enumerate all solutions from one possible assignment in parallel.

Bonus point if you only generate valid assignments.