JuMP kmeans with constraints

I am solving k-means using JuMP.


X = [
	2.0 2.0;
	2.5 2.1;
	7.0 7.0;
	2.2 2.3;
	6.8 7.0;
	7.2 7.5;]
k=2
m, n = size(X)
W = exp.(-Distances.pairwise(Distances.Euclidean(), X'))
I = Matrix(1.0 * LinearAlgebra.I, m, m)
model = JuMP.Model(SCS.Optimizer)
JuMP.unset_silent(model)
@JuMP.variable(model, Z[1:m, 1:m], PSD)
@JuMP.constraint(model, Z .>= 0)
@JuMP.constraint(model, sum(Z; dims=2) .== ones(m))
@JuMP.constraint(model, LinearAlgebra.tr(Z) == k)
@JuMP.objective(model, Min, LinearAlgebra.tr(W * (I - Z)))
JuMP.optimize!(model)
Z_value = JuMP.value.(Z)

Everything works great.

However, I want to impose constraints based on Z value.

I wanted to add something like this:

JuMP.register(model, :proc_matrix, m * m, proc_matrix; autodiff = true)
@JuMP.variable(model, distances[1:m, 1:m] <= 0.1)
@JuMP.constraint(model, distances .== proc_matrix(Z...))

function proc_matrix(a...)
	m = round(Int, floor(sqrt(length(a))))
	M = reshape([a...], m, m)
	Distances.pairwise(Distances.Euclidean(), M)
end

However, even this does not work.

Any suggestions?!

I also want to make proc_matrix and the constraint more complicated but first I want to figure out how to process a variable matrix and return external output back to JuMP. Thank you!

A couple of quick points:

  • JuMP.register is for registering user-defined nonlinear functions. Currently, you must use them in @NL... macros (but we’re working on fixing this)
  • SCS is a conic optimizer; it does not support arbitrary nonlinear functions.
  • User-defined functions must return a scalar; they cannot return a matrix

If you want to keep using a conic optimizer, you’ll have to reformulate the problem into a conic program. If you want to impose arbitrary nonlinear restrictions on the distances, you’ll need a nonlinear solver that also supports PSD constraints, and then I think you’re out of luck?

Assuming I didn’t make any mistakes, I think you’re after something like this:

# JuMP.register(model, :proc_matrix, m * m, proc_matrix; autodiff = true)
# @JuMP.variable(model, distances[1:m, 1:m] <= 0.1)
# @JuMP.constraint(model, distances .== proc_matrix(Z...))
# function proc_matrix(a...)
# 	m = round(Int, floor(sqrt(length(a))))
# 	M = reshape([a...], m, m)
# 	Distances.pairwise(Distances.Euclidean(), M)
# end

# is equivalent to
# @constraint(
# 	model, 
# 	[row in 1:m, col in 1:m], 
# 	Distances.evaluate(Distances.Euclidean(), Z[row, :], Z[: col]) <= 0.1,
# )

# is equivalent to
# JuMP.@constraint(
# 	model,
# 	[row in 1:m, col in 1:m],
# 	sqrt(sum((Z[row, i] - Z[i, col])^2 for i in 1:m)) <= 0.1
# )

# is equivalent to
JuMP.@constraint(
	model,
	[row in 1:m, col in 1:m],
	[0.1; Z[row, :] .- Z[:, col]] in JuMP.SecondOrderCone(),
)

@odow this is a great example. Now I understand better how it works.

Ultimately I want to get based on the distances an estimate of the probability of each point being associated with a given cluster. There are 2 clusters in this case. I want to create a 6x2 matrix with the probability of each point belonging to each cluster. I expect the sum of the rows of this matrix to add to 1.

I know how to code this outside jump based on the Z values. However, I do not know how to implement this inside jump. i want to add penalties based on the values of this 6x2 matrix.

Are there tricks in SCS or JuMP to do this? I checked the documentation and the source code of SCS, and I could find a clear solution. Thank you!

I want to be able to do something like this:

X = [
	2.0 2.0;
	2.5 2.1;
	7.0 7.0;
	2.2 2.3;
	6.8 7.0;
	7.2 7.5;]
C = [10, 5, 5, 10, 1, 1]
k=2
m, n = size(X)
W = exp.(-Distances.pairwise(Distances.Euclidean(), X'))
I = Matrix(1.0 * LinearAlgebra.I, m, m)
model = JuMP.Model(SCS.Optimizer)
JuMP.unset_silent(model)
@JuMP.variable(model, Z[1:m, 1:m], PSD)
@JuMP.constraint(model, Z .>= 0)
@JuMP.constraint(model, sum(Z; dims=2) .== ones(m))
@JuMP.constraint(model, LinearAlgebra.tr(Z) == k)
@JuMP.variable(model, capacity[1:k] == 10)
@JuMP.constraint(model, sum(unique(abs.(trunc.(Z; digits = 1)) .* C; dims=2); dims=1) .== capacity)
@JuMP.objective(model, Min, LinearAlgebra.tr(W * (I - Z)))
JuMP.optimize!(model)

I don’t know if you’re on the right track. There are some reformulations that could achieve what you’re looking for, but they get complicated. And they introduce binary and integer variables, so now you’ll need a mixed-integer conic solver, and the only mixed-integer PSD solver in JuMP’s list of supported solvers is Pajarito.

Stepping back, why solve K-means with a conic solver? What are the alternatives?

unique and trunc give you trouble. The key for trunc with digits = 1 is to see that Z[i,j] is between 0 and 1, so that the trunc will be an integer between 0 and 9 (followed a division by 10).
So that gives:

JuMP.@variable(model, Z_trunc[1:m, 1:m] >= 0, Symmetric, Int)
JuMP.@constraint(model. Z_trunc .<= 10 * Z .<= Z_trunc + 1)
JuMP.@constraint(model, sum(unique(abs.(Z_trunc / 10) .* C; dims=2); dims=1) .== capacity)

then we can drop the abs (why is it needed?) because Z_trunc cannot be negative.

JuMP.@constraint(model, sum(unique(Z_trunc / 10 .* C; dims=2); dims=1) .== capacity)

I’m a but confused by the usage of unique. I invented some fake data:

julia> Z_trunc
6Ă—6 Matrix{Float64}:
 0.2  0.1  0.0  0.0  0.1  0.2
 0.2  0.0  0.2  0.2  0.1  0.0
 0.0  0.2  0.1  0.1  0.1  0.2
 0.1  0.1  0.2  0.1  0.0  0.2
 0.2  0.1  0.1  0.0  0.2  0.1
 0.2  0.0  0.2  0.2  0.1  0.1

julia> unique(Z_trunc, dims = 2)
6Ă—6 Matrix{Float64}:
 0.2  0.1  0.0  0.0  0.1  0.2
 0.2  0.0  0.2  0.2  0.1  0.0
 0.0  0.2  0.1  0.1  0.1  0.2
 0.1  0.1  0.2  0.1  0.0  0.2
 0.2  0.1  0.1  0.0  0.2  0.1
 0.2  0.0  0.2  0.2  0.1  0.1

Did you want to count the number of unique entries in each row?

If we omit the unique, then the last part is:

Y = Z_trunc .* C / 10
JuMP.@constraint(model, sum(Y; dims=1) .== 10)

Puttig it all together

JuMP.@variable(model, Z_trunc[1:m, 1:m] >= 0, Symmetric, Int)
JuMP.@constraint(model. Z_trunc .<= 10 * Z .<= Z_trunc + 1)
JuMP.@expression(model, Y, Z_trunc .* C / 10)
JuMP.@constraint(model, sum(Y; dims=1) .== 10)
1 Like