How can I use pulp to solve a TPS problem?

Question

I'm working on solving a TPS problem, testing it with different scenarios using Pulp. My Code so far:

from pulp import *

def k_tsp_mtz_encoding(n, k, cost_matrix):
    # Check inputs are OK
    assert 1 <= k < n
    assert len(cost_matrix) == n, f'Cost matrix is not {n}x{n}'
    assert all(len(cj) == n for cj in cost_matrix), f'Cost matrix is not {n}x{n}'

    prob = LpProblem('kTSP', LpMinimize)

    # Decision variables
    x = LpVariable.dicts('x', [(i, j) for i in range(n) for j in range(n) if i != j], 0, 1, LpBinary)
    u = LpVariable.dicts('u', [i for i in range(1, n)], 0, n-1, LpInteger)

    # Objective function
    prob += lpSum(cost_matrix[i][j] * x[(i, j)] for i in range(n) for j in range(n) if i != j)

    # Constraints
    for i in range(n):
        prob += lpSum(x[(i, j)] for j in range(n) if i != j) == 1
        prob += lpSum(x[(j, i)] for j in range(n) if i != j) == 1

    for j in range(1, n):
        prob += lpSum(x[(i, j)] for i in range(n) if i != j) == 1
        prob += lpSum(x[(j, i)] for i in range(n) if i != j) == 1

    # Subtour elimination constraints
    for i in range(1, n):
        for j in range(1, n):  # Avoid accessing cost_matrix[i][i]
            if i != j:
                prob += u[i] - u[j] + n * x[(i, j)] <= n - 1

    # Solve the problem
    prob.solve()

    # Extract tours
    all_tours = []
    for _ in range(k):
        tour = [0]
        j = 0
        while True:
            for i in range(n):
                if value(x[(j, i)]) == 1:
                    tour.append(i)
                    j = i
                    break
            if j == 0:
                break
        all_tours.append(tour)

    return all_tours

I'm testing it with this lines of code but I keep getting errors. I'm doing this in Jupyter Notebooks

cost_matrix=[ [None,3,4,3,5],
             [1, None, 2,4, 1],
             [2, 1, None, 5, 4],
             [1, 1, 5, None, 4],
             [2, 1, 3, 5, None] ]
n=5
k=2
all_tours = k_tsp_mtz_encoding(n, k, cost_matrix)
print(f'Your code returned tours: {all_tours}')
assert len(all_tours) == k, f'k={k} must yield two tours -- your code returns {len(all_tours)} tours instead'

tour_cost = 0
for tour in all_tours:
    assert tour[0] == 0, 'Each salesperson tour must start from vertex 0'
    i = 0
    for j in tour[1:]:
        tour_cost += cost_matrix[i][j]
        i = j
    tour_cost += cost_matrix[i][0]

print(f'Tour cost obtained by your code: {tour_cost}')
assert abs(tour_cost - 12) <= 0.001, f'Expected tour cost is 12, your code returned {tour_cost}'
for i in range(1, n):
    is_in_tour = [ 1 if i in tour else 0 for tour in all_tours]
    assert sum(is_in_tour) == 1, f' vertex {i} is in {sum(is_in_tour)} tours -- this is incorrect'

print('Correct')

This is the most recent error, I was able to get the vertex visited output, but after I modified it I'm just getting this KeyError

KeyError                                  Traceback (most recent call last)
<ipython-input-87-248518380ac2> in <module>
      6 n=5
      7 k=2
----> 8 all_tours = k_tsp_mtz_encoding(n, k, cost_matrix)
      9 print(f'Your code returned tours: {all_tours}')
     10 assert len(all_tours) == k, f'k={k} must yield two tours -- your code returns {len(all_tours)} tours instead'

<ipython-input-86-3fa3b44faee3> in k_tsp_mtz_encoding(n, k, cost_matrix)
     54         while True:
     55             for i in range(n):
---> 56                 if value(x[(j, i)]) == 1:
     57                     tour.append(i)
     58                     j = i

KeyError: (0, 0)```

Answer 1

Break this up into multiple functions with typed parameters. This will reduce the risk of programming error. Also, .dicts() is inappropriate for this application; use .matrix().

The following does solve your key error but does not attempt to produce a meaningful linear problem. It successfully solves but then your assertion fails; I leave that to you to debug.

from typing import Iterator, Sequence

import pulp


def tsp_make_vars(n: int) -> tuple[
    list[list[pulp.LpVariable]],
    list[pulp.LpVariable],
]:
    x = pulp.LpVariable.matrix(
        name='x', indices=(range(n), range(n)), cat=pulp.LpBinary,
    )
    for i in range(n):
        x[i][i] = 0
    u = pulp.LpVariable.matrix(
        name='u', indices=range(n), lowBound=0, upBound=n-1, cat=pulp.LpInteger,
    )
    u[0] = None
    return x, u


def tsp_make_prob(
    cost_matrix: Sequence[Sequence[float]],
    x: list[list[pulp.LpVariable]],
) -> pulp.LpProblem:
    n = len(cost_matrix)
    assert all(len(cj) == n for cj in cost_matrix), f'Cost matrix is not {n}x{n}'

    prob = pulp.LpProblem(name='kTSP', sense=pulp.LpMinimize)
    prob.setObjective(pulp.lpDot(x, cost_matrix))
    return prob


def tsp_add_constraints(
    prob: pulp.LpProblem,
    x: list[list[pulp.LpVariable]],
    u: list[pulp.LpVariable],
) -> None:
    n = len(x)

    for i, row in enumerate(x):
        prob.addConstraint(
            name=f'excl_{i}j', constraint=pulp.lpSum(row) == 1,
        )
    for j in range(n):
        prob.addConstraint(
            name=f'excl_i{j}',
            constraint=pulp.lpSum(row[j] for row in x) == 1,
        )

    for i, (xi, ui) in enumerate(zip(x[1:], u[1:]), start=1):
        for j, (xij, uj) in enumerate(zip(xi[1:], u[1:]), start=1):
            if i != j:
                prob.addConstraint(
                    name=f'subtour_{i},{j}',
                    constraint=ui - uj + n*xij <= n - 1,
                )


def tsp_solve(prob: pulp.LpProblem) -> None:
    prob.solve()
    if prob.status != pulp.LpStatusOptimal:
        raise ValueError(prob.status)


def tsp_extract_tours(
    x: list[list[pulp.LpVariable]],
    k: int,
) -> Iterator[list[int]]:
    n = len(x)
    assert 1 <= k < n

    for _ in range(k):
        tour = [0]
        j = 0
        while True:
            for i in range(n):
                if pulp.value(x[j][i]) > 0.5:
                    tour.append(i)
                    j = i
                    break
            if j == 0:
                break
        yield tour


def test() -> None:
    cost_matrix = (
        (0, 3, 4, 3, 5),
        (1, 0, 2, 4, 1),
        (2, 1, 0, 5, 4),
        (1, 1, 5, 0, 4),
        (2, 1, 3, 5, 0),
    )
    n = len(cost_matrix)
    k = 2

    x, u = tsp_make_vars(n=n)
    prob = tsp_make_prob(cost_matrix=cost_matrix, x=x)
    tsp_add_constraints(prob=prob, x=x, u=u)
    print(prob)
    tsp_solve(prob=prob)

    all_tours = tuple(tsp_extract_tours(x=x, k=k))
    print('Tours:', all_tours)
    assert len(all_tours) == k, f'k={k} must yield two tours -- code returned {len(all_tours)} tours instead'

    tour_cost = 0
    for tour in all_tours:
        assert tour[0] == 0, 'Each salesperson tour must start from vertex 0'
        i = 0
        for j in tour[1:]:
            tour_cost += cost_matrix[i][j]
            i = j
        tour_cost += cost_matrix[i][0]

    print(f'Tour cost: {tour_cost}')
    assert abs(tour_cost - 12) <= 0.001, f'Expected tour cost is 12, not {tour_cost}'
    for i in range(1, n):
        is_in_tour = [1 if i in tour else 0 for tour in all_tours]
        assert sum(is_in_tour) == 1, f' vertex {i} is in {sum(is_in_tour)} tours -- this is incorrect'


if __name__ == '__main__':
    test()

kTSP:
MINIMIZE
3*x_0_1 + 4*x_0_2 + 3*x_0_3 + 5*x_0_4 + 1*x_1_0 + 2*x_1_2 + 4*x_1_3 + 1*x_1_4 + 2*x_2_0 + 1*x_2_1 + 5*x_2_3 + 4*x_2_4 + 1*x_3_0 + 1*x_3_1 + 5*x_3_2 + 4*x_3_4 + 2*x_4_0 + 1*x_4_1 + 3*x_4_2 + 5*x_4_3 + 0
SUBJECT TO
excl_0j: x_0_1 + x_0_2 + x_0_3 + x_0_4 = 1

excl_1j: x_1_0 + x_1_2 + x_1_3 + x_1_4 = 1

excl_2j: x_2_0 + x_2_1 + x_2_3 + x_2_4 = 1

excl_3j: x_3_0 + x_3_1 + x_3_2 + x_3_4 = 1

excl_4j: x_4_0 + x_4_1 + x_4_2 + x_4_3 = 1

excl_i0: x_1_0 + x_2_0 + x_3_0 + x_4_0 = 1

excl_i1: x_0_1 + x_2_1 + x_3_1 + x_4_1 = 1

excl_i2: x_0_2 + x_1_2 + x_3_2 + x_4_2 = 1

excl_i3: x_0_3 + x_1_3 + x_2_3 + x_4_3 = 1

excl_i4: x_0_4 + x_1_4 + x_2_4 + x_3_4 = 1

subtour_1,2: u_1 - u_2 + 5 x_1_2 <= 4

subtour_1,3: u_1 - u_3 + 5 x_1_3 <= 4

subtour_1,4: u_1 - u_4 + 5 x_1_4 <= 4

subtour_2,1: - u_1 + u_2 + 5 x_2_1 <= 4

subtour_2,3: u_2 - u_3 + 5 x_2_3 <= 4

subtour_2,4: u_2 - u_4 + 5 x_2_4 <= 4

subtour_3,1: - u_1 + u_3 + 5 x_3_1 <= 4

subtour_3,2: - u_2 + u_3 + 5 x_3_2 <= 4

subtour_3,4: u_3 - u_4 + 5 x_3_4 <= 4

subtour_4,1: - u_1 + u_4 + 5 x_4_1 <= 4

subtour_4,2: - u_2 + u_4 + 5 x_4_2 <= 4

subtour_4,3: - u_3 + u_4 + 5 x_4_3 <= 4

VARIABLES
0 <= u_1 <= 4 Integer
0 <= u_2 <= 4 Integer
0 <= u_3 <= 4 Integer
0 <= u_4 <= 4 Integer
0 <= x_0_1 <= 1 Integer
0 <= x_0_2 <= 1 Integer
0 <= x_0_3 <= 1 Integer
0 <= x_0_4 <= 1 Integer
0 <= x_1_0 <= 1 Integer
0 <= x_1_2 <= 1 Integer
0 <= x_1_3 <= 1 Integer
0 <= x_1_4 <= 1 Integer
0 <= x_2_0 <= 1 Integer
0 <= x_2_1 <= 1 Integer
0 <= x_2_3 <= 1 Integer
0 <= x_2_4 <= 1 Integer
0 <= x_3_0 <= 1 Integer
0 <= x_3_1 <= 1 Integer
0 <= x_3_2 <= 1 Integer
0 <= x_3_4 <= 1 Integer
0 <= x_4_0 <= 1 Integer
0 <= x_4_1 <= 1 Integer
0 <= x_4_2 <= 1 Integer
0 <= x_4_3 <= 1 Integer

Welcome to the CBC MILP Solver 
Version: 2.10.3 
Build Date: Dec 15 2019 
...

Result - Optimal solution found

Objective value:                10.00000000
Enumerated nodes:               0
Total iterations:               0
Time (CPU seconds):             0.00
Time (Wallclock seconds):       0.01

Option for printingOptions changed from normal to all
Total time (CPU seconds):       0.00   (Wallclock seconds):       0.01

Tours: ([0, 3, 1, 4, 2, 0], [0, 3, 1, 4, 2, 0])
Tour cost: 20