Sudoku solver using integer programming

sudoku.py

import numpy as np
from scipy import optimize, sparse

board_coords = np.array([
    [0, 1, 2],
    [0, 4, 3],
    [0, 7, 4],
    [1, 0, 6],
    [1, 8, 3],
    [2, 2, 4],
    [2, 6, 5],
    [3, 3, 8],
    [3, 5, 6],
    [4, 0, 8],
    [4, 4, 1],
    [4, 8, 6],
    [5, 3, 7],
    [5, 5, 5],
    [6, 2, 7],
    [6, 6, 6],
    [7, 0, 4],
    [7, 8, 8],
    [8, 1, 3],
    [8, 4, 4],
    [8, 7, 2]
])
# convert the digits to be zero-indexed
board_coords_zero = board_coords.copy()
board_coords_zero[:, 2] -= 1

board = np.zeros((9, 9))

# initial board
for row in board_coords_zero:
    board[row[0], row[1]] = row[2]

# solution for checking constraints and final solution
# also helpful for fixing intuition about the ordering
# of the constraints - our solution will be in the
# same order as the raveled grid
solution = np.array([
    [9, 2, 5, 6, 3, 1, 8, 4, 7],
    [6, 1, 8, 5, 7, 4, 2, 9, 3],
    [3, 7, 4, 9, 8, 2, 5, 6, 1],
    [7, 4, 9, 8, 2, 6, 1, 3, 5],
    [8, 5, 2, 4, 1, 3, 9, 7, 6],
    [1, 6, 3, 7, 9, 5, 4, 8, 2],
    [2, 8, 7, 3, 5, 9, 6, 1, 4],
    [4, 9, 1, 2, 6, 7, 3, 5, 8],
    [5, 3, 6, 1, 4, 8, 7, 2, 9],
])

grid_solution = np.zeros((9, 9, 9))
for i in range(9):
    for j in range(9):
        grid_solution[i, j, solution[i, j] - 1] = 1
grid_solution = grid_solution.ravel(order='F')

# unstack in column-major and then put back in 2d row major to get
# the grid shape i, j, k
grid = np.mgrid[0:9, 0:9, 0:9].ravel(order='F').reshape(-1, 3, order='C')

# the first set of constraints we need to sum over k,
# so we want the raveled index for each i, j pair over k

constraints_row = np.repeat(np.arange(9 ** 2), 9)
row_constraints_col = np.arange(9 ** 3)

A_row = sparse.coo_array((
    np.ones(9 ** 3),
    (constraints_row, row_constraints_col)),
    shape=(9 ** 2, 9 ** 3)
)
assert (A_row.sum(1) == 9).all()
assert (A_row @ grid_solution == 1).all()

# then for every i, k pair over j
# finally for every k, j pair over i

# Since every number should be in every row, we'll have a constraint
# that every row should sum to 1, after unraveling the 9 x 9 x 9 board
# that means every 9 rows should sum to 1
# that's 81 row constraints, 9 for each of the 9 numbers

# let j be the columns
# for them every 9th row should sum to 1

col_constraints_col = (
    np.tile(np.arange(0, 9 ** 3, 9), 9)
    # offset by 1 more every 9 elements
    + np.repeat(np.arange(9), 81)
)
A_col = sparse.coo_array((
    np.ones(9 ** 3),
    (constraints_row, col_constraints_col)),
    shape=(9 ** 2, 9 ** 3)
)
assert (A_col.sum(1) == 9).all()
assert (A_col @ grid_solution == 1).all()


# let k be the numbers 0 through 8
# we have to use every number once so we also need these to sum to 1
# over the kth dimension
# the kth dimension moves the slowest so for this one we need it to move
# every 27th row
digit_constraints_col = (
    np.tile(np.arange(0, 9 ** 3, 9 ** 2), 81)
    # offset by 1 more every 9 elements
    + constraints_row
)
A_digits = sparse.coo_array((
    np.ones(9 ** 3),
    (constraints_row, digit_constraints_col)
), shape=(9 ** 2, 9 ** 3))
assert (A_digits.sum(1) == 9).all()
assert (A_col @ grid_solution == 1).all()

# subgrid constraints
# there are 9 sub grids along the rows and columns <3, 3-5, and 6-8
# each of these subgrids must also have each of the digits
grid_constraints = []
row_grid = np.array([0, 3, 6, 9])
col_grid = np.array([0, 3, 6, 9])
for i in range(3):
    for j in range(3):
        grid_constraints.append(
            np.ravel_multi_index(

                grid[(row_grid[i] <= grid[:, 0])
                     & (row_grid[i + 1] > grid[:, 0])
                     & (col_grid[j] <= grid[:, 1])
                     & (col_grid[j + 1] > grid[:, 1])
                     ].T,
                (9, 9, 9),
                order='F')
        )

grid_rows = np.repeat(np.arange(81), 9)

A_grid = sparse.coo_array((
    np.ones(9 ** 3),
    (grid_rows, np.hstack((grid_constraints)))
), shape=(9 ** 2, 9 ** 3))
assert (A_grid.sum(1) == 9).all()
assert (A_grid @ grid_solution == 1).all()


# Stack together all of the constraints
A = sparse.vstack((A_row, A_col, A_digits, A_grid))

assert (A @ grid_solution == 1).all()

c = np.cumsum(np.ones(9 ** 3)) / 9 ** 3
integrality = np.ones_like(c)
b_l = np.ones(A.shape[0])
b_u = np.ones(A.shape[0])

x_b_l = np.zeros(9 ** 3)
# starting values
x_b_l[np.ravel_multi_index(board_coords_zero.T, (9, 9, 9), order='F')] = 1
x_b_u = np.ones(9 ** 3)
bounds = optimize.Bounds(lb=x_b_l, ub=x_b_u)

constraints = optimize.LinearConstraint(A, b_l, b_u)
res = optimize.milp(c=c, constraints=constraints, bounds=bounds,
                    integrality=integrality,
                    options=dict(disp=True)
                    )
assert (res.x == grid_solution).all()

found_grid_solution = res.x.reshape(9, 9, 9, order='F')
found_solution = np.zeros((9, 9))

for i, j, value in zip(*np.where(found_grid_solution)):
    found_solution[i, j] = value