Christmas lights

lights.py

# Some informal experiments with Christmas Lights for the MOSEK blogpost
from mosek.fusion import *
import numpy as np
import sys

def exRandom():
	D = 2  # Dimension of the tree
	N = 20 # Number of lights
	P = 35 # Number of decorations to shine upon
	T = 10 # Tree top
	W = 3  # Tree width
	L = 1  # Length of segment between lights
	# Strength of each light
	I = np.minimum(1.6, np.maximum(0.4, np.random.normal(size=N) + 1.0))
	# Locations of decorations
	C = [[0,T]]
	while len(C) < P:
		x, y = np.random.uniform(low=-W, high=W), np.random.uniform(low=0, high=T)
		if W*(1-y/T) >= np.abs(x):
			C.append([x,y])
	C = np.array(C)
	return D, N, P, T, W, L, I, C

def exUpDown():
	D = 2  # Dimension of the tree
	N = 25 # Number of lights
	P = 1+4*10 # Number of decorations to shine upon
	T = 10 # Tree top
	W = 3  # Tree width
	L = 1  # Length of segment between lights
	# Strength of each light
	I = np.random.uniform(low=0.8, high=1.2, size=N)
	# Locations of decorations
	C = [[0,T]]
	for i in range(10):
		C.append([(1-i/10)*W, i/10*T])
		C.append([(1-i/10)*(-W), i/10*T])
		C.append([(1-i/10)*W/3, i/10*T])
		C.append([(1-i/10)*(-W/3), i/10*T])
	C = np.array(C)
	return D, N, P, T, W, L, I, C

def exLeftRight():
	D = 2  # Dimension of the tree
	N = 25 # Number of lights
	P = 1+3+5+8 # Number of decorations to shine upon
	T = 10 # Tree top
	W = 3  # Tree width
	L = 0.8  # Length of segment between lights
	# Strength of each light
	I = np.random.uniform(low=0.8, high=1.2, size=N)
	# Locations of decorations
	C = [[0,T]]
	for i in range(1,4):
		C.append([(-W/2)*i/4+W/2*(1-i/4), T*0.7])
	for i in range(1,6):
		C.append([(-W*0.6)*i/6+W*0.6*(1-i/6), T*0.45])
	for i in range(1,9):
		C.append([(-W)*i/9+W*(1-i/9), T*0.1])
	C = np.array(C)
	return D, N, P, T, W, L, I, C

# Defines points in the tree
# W*(1-x[-1]/T) >= ||x[0:-1]||_2
# x[-1] >= 0
def inXmasTree(M, x, W, D, T):
	M.constraint(Expr.vstack(Expr.mul(W, Expr.sub(1, Expr.mul(x.index(D-1), 1/T))),
	                         x.slice(0, D-1)),
	             Domain.inQCone())

	M.constraint(x.index(D-1), Domain.greaterThan(0))

# Returns piecewise linear f(d) 
# (d<=I AND f(d)=I-d) OR (d>=I and f(d)=0)
def approxIllumination(M, d, il, i):
	M.disjunction(DJC.AND(DJC.term(d, Domain.lessThan(i)), DJC.term(Expr.add(il,d), Domain.lessThan(i))),
				  DJC.AND(DJC.term(d, Domain.greaterThan(i)), DJC.term(il, Domain.lessThan(0))))

def modelLights(D, N, P, T, W, L, I, C):
	with Model() as M:
		# Positions of the lights
		x = M.variable("x", [N, D]) 
		# Lights are within the tree
		for i in range(N):
			inXmasTree(M, Var.flatten(x.slice([i,0],[i+1,D])), W, D, T)
		# Lights are not more than L apart
		M.constraint(Expr.hstack(Expr.constTerm(N-1, L),
			                     Expr.sub(x.slice([0,0],[N-1,D]), x.slice([1,0],[N,D]))),
		             Domain.inQCone())

		# (Upper bounds for) distances from i-th decoration to j-th light
		d = M.variable("d", [P, N])
		for i in range(P):
			for j in range(N):
				M.constraint(Expr.vstack(d.index(i,j),
					                     Expr.sub(Var.flatten(x.slice([j,0],[j+1,D])), C[i])),
				             Domain.inQCone())

		# Approximate illumination of i-th decoration by j-th light
		il = M.variable("il", [P, N])
		for i in range(P):
			for j in range(N):
				approxIllumination(M, d.index(i,j), il.index(i,j), I[j])

		# Maximize worst case total illumination of a decoration
		ilmax = M.variable("ilmax")
		M.constraint(Expr.sub(Var.repeat(ilmax, P), Expr.sum(il, 1)), Domain.lessThan(0))
		M.objective(ObjectiveSense.Maximize, ilmax)

		# Startpoint in the corner
		M.constraint(Var.flatten(x.slice([0,0], [1,D])), Domain.equalsTo([W,0]))

		M.setLogHandler(sys.stdout)
		M.setSolverParam("optimizerMaxTime", 200)
		M.setSolverParam("mioTolRelGap", 3)
		M.solve()

		M.acceptedSolutionStatus(AccSolutionStatus.Feasible)
		return x.level().reshape((N,D))

def plot(D, N, P, T, W, L, I, C, x, num):
	import matplotlib.pyplot as plt
	figure, axes = plt.subplots()
	plt.axis('off')
	axes.set_aspect(1)
	plt.fill([-W,W,0,-W], [0,0,T,0], "g")
	alpha = np.ones(P)
	if x is not None:
		for j in range(N):
			axes.add_artist(plt.Circle(x[j,:], I[j], fill=True, color="yellow", alpha=0.3))
		plt.plot(x[:,0], x[:,1], c="black")
		plt.scatter(x[:,0], x[:,1], c="black", s=7)
		for i in range(P):
			alpha[i] = sum(np.maximum(I[j] - np.linalg.norm(x[j,:]-C[i,:]), 0) for j in range(N))
		alpha = np.sqrt(alpha / np.max(alpha))
	for i in range(P):
		plt.scatter([C[i,0]], [C[i,1]], c="red", s=30, alpha=alpha[i], marker=np.random.choice(['h','*','D','d',"8","+"]), edgecolors=None)
	plt.savefig(f'xmas-lights-{num}.png')
	#plt.show()

D, N, P, T, W, L, I, C = exRandom()
plot(D, N, P, T, W, L, I, C, None, 1)
X = modelLights(D, N, P, T, W, L, I, C)
plot(D, N, P, T, W, L, I, C, X, 2)

D, N, P, T, W, L, I, C = exUpDown()
X = modelLights(D, N, P, T, W, L, I, C)
plot(D, N, P, T, W, L, I, C, X, 3)

D, N, P, T, W, L, I, C = exLeftRight()
X = modelLights(D, N, P, T, W, L, I, C)
plot(D, N, P, T, W, L, I, C, X, 4)