n-queens in Fusion

In a nice post on LinkedIn by Alireza Soroudi the author presents a short GAMS integer model for the classical n-queens puzzle: place the maximal number of non-attacking queens on an $n\times n$ chessboard.

We couldn't resist writing it up in Python Fusion as well.

	from mosek.fusion import *
	import numpy as np
	import sys
	n = int(sys.argv[1])
	M = Model()
	x = M.variable("x", [n,n], Domain.binary())
	M.objective(ObjectiveSense.Maximize, Expr.sum(x))
	M.constraint(Expr.sum(x, 0), Domain.lessThan(1))
	M.constraint(Expr.sum(x, 1), Domain.lessThan(1))
	for k in range(-n+1,n):
	M.constraint(Expr.dot(x, np.diag([1]*(n-abs(k)), k)), Domain.lessThan(1))
	M.constraint(Expr.dot(x, np.fliplr(np.diag([1]*(n-abs(k)), k))), Domain.lessThan(1))
	M.solve()
	print(np.array2string(np.fromiter(
	map(lambda u: '\u2655' if u>0.5 else '_', x.level()), dtype='<U1').reshape((n,n)),
	formatter={'str_kind': lambda u: u}))

view raw nqueens.py hosted with ❤ by GitHub

And here is the solution we get for $n=8$:

[[_ ♕ _ _ _ _ _ _]

 [_ _ _ _ _ ♕ _ _]

 [♕ _ _ _ _ _ _ _]

 [_ _ _ _ _ _ ♕ _]

 [_ _ _ ♕ _ _ _ _]

 [_ _ _ _ _ _ _ ♕]

 [_ _ ♕ _ _ _ _ _]

  [_ _ _ _ ♕ _ _ _]]

Grouping in Fusion

We sometimes get asked how to efficiently perform a "group by" operation on a variable. That is, we have a variable $x=(x_0,\ldots,x_{n-1})$ naturally divided into groups and we want to add constraints for each group or constraints on aggregate values within groups. This arises for instance in portfolio optimization where the groups are assets, each consisting of a (varying) number of tax lots.

To keep the discussion more concrete, suppose we have a variable $x=(x_0,x_1,x_2,x_3,x_4,x_5)$ which consists of 3 groups $(x_0,x_1,x_2)$, $(x_3)$ and $(x_4,x_5)$. Let's say we want to express:

• an upper bound $b_i$ on the total value within each group,
• a joint volatility constraint such as $$\gamma\geq\left\|G\cdot \left[\begin{array}{c} x_0+x_1+x_2-i_0 \\ x_3-i_1 \\ x_4+x_5-i_2\end{array}\right]\right\|_2$$ for some matrix $G$ and constant index weights $i_0,i_1,i_2$
• the constraint that all values within one group have the same sign.

Pick, slice

A straightforward solution is to pick (Expr.pick) the content of each group into a separate view and add relevant constraints in a loop over the groups. One could also take slices (Expr.slice) when the groups form contiguous subsequences. This approach is implemented below in Python Fusion.

	from mosek.fusion import *
	import numpy as np
	M = Model()
	x = M.variable(6)
	groups = [ [0,1,2], [3], [4,5] ]
	# Some data
	bounds = [ 1.0, 2.0, 3.0 ]
	ind_weights = [ 1.1, 2.2, 3.3 ]
	gamma = 0.01
	G = np.random.uniform(size=(3,3))
	# Sum of each group bounded above
	for i in range(len(groups)):
	M.constraint(Expr.sum(Expr.pick(x, groups[i])), Domain.lessThan(bounds[i]))
	# Volatility bound
	# grpSumVec containins the sums in each group
	grpSumVec = Expr.vstack( [Expr.sum(Expr.pick(x, grp)) for grp in groups] )
	M.constraint( Expr.vstack(gamma, Expr.mul(G, Expr.sub(grpSumVec, ind_weights))), Domain.inQCone() )

view raw group-pick.py hosted with ❤ by GitHub

Loop-free

The previous solution requires picking and stacking expressions in a loop over all groups. This can sometimes be slow. A nicer and more efficient solution uses a bit of linear algebra to perform the grouping. We first encode the groups via a sparse membership matrix $\mathrm{Memb}$, where the rows are groups, columns are entries of $x$, and there is a $1$ whenever an entry belongs to a group. In our example $$\mathrm{Memb}=\left[\begin{array}{cccccc}1&1&1&0&0&0\\0&0&0&1&0&0\\0&0&0&0&1&1\end{array}\right] .$$

This matrix is easy to construct in Fusion. 

	# Membership matrix
	Mi = [0, 0, 0, 1, 2, 2] # Rows, Group number for entries of x
	Mj = [0, 1, 2, 3, 4, 5] # Columns
	Mv = [1.0] * 6 # Values
	Memb = Matrix.sparse(3, 6, Mi, Mj, Mv)

view raw group-membdef.py hosted with ❤ by GitHub

Note that $\mathrm{Memb}\cdot x$ is the vector of all group sums, in our case $$\mathrm{Memb}\cdot x = \left[\begin{array}{c} x_0+x_1+x_2 \\ x_3 \\ x_4+x_5\end{array}\right].$$

That means we can express both of the previous models in a single call without looping. Since $\mathrm{Memb}$ is very sparse, this becomes a very efficient representation. Of course it is important to keep $\mathrm{Memb}$ as a sparse matrix.

	# Bounded sums
	M.constraint( Expr.mul(Memb, x), Domain.lessThan(bounds) )
	# Volatility bound
	M.constraint( Expr.vstack(gamma,
	Expr.mul(G, Expr.sub( Expr.mul(Memb, x), ind_weights))),
	Domain.inQCone() )

view raw group-loopfree.py hosted with ❤ by GitHub

Loop-free same sign

We can now use the same matrix to model the last problem from the introduction: all entries within each group must have the same sign. We introduce a sequence of binary variables $z=(z_0,\ldots,z_{g-1})$, one for each of the $g$ groups. The $j$-th variable will determine the sign of elements in the $j$-th group. That is imposed by constraints $$-M(1-z_j)\leq x_i\leq Mz_j,$$ whenever $x_i$ belongs to $j$-th group. We can use the pick/slice strategy per group, as before, or observe that $\mathrm{Memb}^T\cdot z$ produces the vector with the correct binary variable for each entry in $x$. I our case, if $z=(z_0,z_1,z_2)$ then $$\mathrm{Memb}^T\cdot z = (z_0,z_0,z_0,z_1,z_2,z_2)^T.$$ Now each inequality in (4) can be written as a single constraint:

	z = M.variable(3, Domain.binary()) # Indicates sign for each group
	bigM = 100
	# All items in one group have the same sign
	zMatch = Expr.mul(Memb.transpose(), z)
	M.constraint( Expr.sub(x, Expr.mul(bigM, zMatch)), Domain.lessThan(0) )
	M.constraint( Expr.sub(Expr.mul(-bigM, Expr.sub(1, zMatch)), x), Domain.lessThan(0) )

view raw group-sign.py hosted with ❤ by GitHub