From c32330478785d9480f8195eb9178d7a3949b1df5 Mon Sep 17 00:00:00 2001 From: Titan Yuan Date: Mon, 24 Feb 2025 11:52:37 -0800 Subject: [PATCH] Add covered assignment solver --- .github/workflows/ci.yaml | 2 +- assignment/BUILD.bazel | 31 ++++++++++ assignment/assignment.cc | 30 +++++++++ assignment/assignment.h | 53 ++++++++++++++++ assignment/covered_assignment.cc | 87 +++++++++++++++++++++++++++ assignment/covered_assignment.h | 30 +++++++++ assignment/covered_assignment_test.cc | 50 +++++++++++++++ 7 files changed, 282 insertions(+), 1 deletion(-) create mode 100644 assignment/BUILD.bazel create mode 100644 assignment/assignment.cc create mode 100644 assignment/assignment.h create mode 100644 assignment/covered_assignment.cc create mode 100644 assignment/covered_assignment.h create mode 100644 assignment/covered_assignment_test.cc diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 8506cdd..f26fa5d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -15,7 +15,7 @@ jobs: - uses: actions/checkout@v4 - uses: bazel-contrib/setup-bazel@0.9.1 - run: bazel build //... - # - run: bazel test --test_output=errors //... + - run: bazel test --test_output=errors //... - uses: actions/upload-artifact@v4 with: name: plugins-${{ matrix.platform }} diff --git a/assignment/BUILD.bazel b/assignment/BUILD.bazel new file mode 100644 index 0000000..051b207 --- /dev/null +++ b/assignment/BUILD.bazel @@ -0,0 +1,31 @@ +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "assignment", + srcs = ["assignment.cc"], + hdrs = ["assignment.h"], + deps = ["@com_google_absl//absl/strings:str_format"], +) + +cc_library( + name = "covered_assignment", + srcs = ["covered_assignment.cc"], + hdrs = ["covered_assignment.h"], + deps = [ + ":assignment", + "//base:logging", + "@com_google_absl//absl/strings:str_format", + "@ortools//ortools/sat:cp_model", + ], +) + +cc_test( + name = "covered_assignment_test", + srcs = ["covered_assignment_test.cc"], + deps = [ + ":assignment", + ":covered_assignment", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/assignment/assignment.cc b/assignment/assignment.cc new file mode 100644 index 0000000..8e328e8 --- /dev/null +++ b/assignment/assignment.cc @@ -0,0 +1,30 @@ +#include "assignment/assignment.h" + +#include +#include + +#include "absl/strings/str_format.h" + +namespace assignment { + +void Assignment::ValidateCosts() const { + // Validate the first dimension of the cost matrix. + if (costs_.size() != num_agents_) { + throw std::invalid_argument( + absl::StrFormat("The assignment cost matrix has an incorrect number of " + "rows: %d vs. %d.", + costs_.size(), num_agents_)); + } + + // Validate the second dimension of the cost matrix. + for (const auto& row : costs_) { + if (row.size() != num_tasks_) { + throw std::invalid_argument( + absl::StrFormat("The assignment cost matrix has an incorrect number " + "of columns: %d vs. %d.", + row.size(), num_tasks_)); + } + } +} + +} // namespace assignment diff --git a/assignment/assignment.h b/assignment/assignment.h new file mode 100644 index 0000000..e166f3d --- /dev/null +++ b/assignment/assignment.h @@ -0,0 +1,53 @@ +// The assignment class is an interface for a cost-based assignment problem. + +#pragma once + +#include +#include + +namespace assignment { + +// Assignment interface. +class Assignment { + public: + // Assignment item type. + using AssignmentItem = std::pair; + + // The assignment cost matrix should be a matrix of dimensions num_agents x + // num_tasks. + Assignment(int num_agents, int num_tasks, + std::vector> costs) + : num_agents_(num_agents), + num_tasks_(num_tasks), + costs_(std::move(costs)) { + Validate(); + ValidateCosts(); + } + + Assignment(const Assignment&) = default; + Assignment& operator=(const Assignment&) = default; + + virtual ~Assignment() = default; + + // Assign the agents to the tasks. + virtual std::vector Assign() const = 0; + + protected: + // Validate the assignment problem. + virtual void Validate() const {}; + + // Number of agents. + int num_agents_ = 0; + + // Number of tasks. + int num_tasks_ = 0; + + // Assignment cost matrix of dimensions num_agents x num_tasks. + std::vector> costs_; + + private: + // Validate the cost matrix. + void ValidateCosts() const; +}; + +} // namespace assignment diff --git a/assignment/covered_assignment.cc b/assignment/covered_assignment.cc new file mode 100644 index 0000000..e631437 --- /dev/null +++ b/assignment/covered_assignment.cc @@ -0,0 +1,87 @@ +#include "assignment/covered_assignment.h" + +#include +#include + +#include "absl/strings/str_format.h" +#include "assignment/assignment.h" +#include "base/logging.h" +#include "ortools/sat/cp_model.h" +#include "ortools/sat/cp_model.pb.h" +#include "ortools/sat/cp_model_solver.h" + +namespace assignment { + +std::vector CoveredAssignment::Assign() const { + // Create the constraint programming model. + operations_research::sat::CpModelBuilder cp_model; + + // Define the variables. + // x[i][j] is an array of boolean variables, such that x[i][j] is true if + // agent i is assigned to task j. + std::vector> x( + num_agents_, std::vector(num_tasks_)); + for (int i = 0; i < num_agents_; ++i) { + for (int j = 0; j < num_tasks_; ++j) { + x[i][j] = cp_model.NewBoolVar(); + } + } + + // Define the constraints. + // Each agent is assigned to one task. + for (int i = 0; i < num_agents_; ++i) { + cp_model.AddExactlyOne(x[i]); + } + // Each task is assigned to at least one agent. + for (int j = 0; j < num_tasks_; ++j) { + std::vector tasks; + for (int i = 0; i < num_agents_; ++i) { + tasks.push_back(x[i][j]); + } + cp_model.AddAtLeastOne(tasks); + } + + // Define the objective function. + operations_research::sat::DoubleLinearExpr total_cost; + for (int i = 0; i < num_agents_; ++i) { + for (int j = 0; j < num_tasks_; ++j) { + total_cost = total_cost.AddTerm(x[i][j], costs_[i][j]); + } + } + cp_model.Minimize(total_cost); + + // Solve the assignment problem. + const operations_research::sat::CpSolverResponse response = + Solve(cp_model.Build()); + + // Check the feasibility of the solution. + if (response.status() == + operations_research::sat::CpSolverStatus::INFEASIBLE) { + LOG(ERROR) << "Covered assignment problem is infeasible."; + return std::vector(); + } + + // Record the assignemnts. + std::vector assignments; + assignments.reserve(num_agents_); + for (int i = 0; i < num_agents_; ++i) { + for (int j = 0; j < num_tasks_; ++j) { + if (SolutionBooleanValue(response, x[i][j])) { + assignments.emplace_back(i, j); + break; + } + } + } + return assignments; +} + +void CoveredAssignment::Validate() const { + // Check that there are at least as many agents as tasks to cover all tasks. + if (num_agents_ < num_tasks_) { + throw std::invalid_argument( + absl::StrFormat("There are fewer agents than tasks: %d vs. %d.", + num_agents_, num_tasks_)); + } +} + +} // namespace assignment diff --git a/assignment/covered_assignment.h b/assignment/covered_assignment.h new file mode 100644 index 0000000..a9c8456 --- /dev/null +++ b/assignment/covered_assignment.h @@ -0,0 +1,30 @@ +// The covered assignment assigns one task to each agent under the condition +// that all tasks are assigned to at least one agent. + +#pragma once + +#include + +#include "assignment/assignment.h" + +namespace assignment { + +// Covered assignment. +class CoveredAssignment : public Assignment { + public: + CoveredAssignment(int num_agents, int num_tasks, + std::vector> costs) + : Assignment(num_agents, num_tasks, std::move(costs)) {} + + CoveredAssignment(const CoveredAssignment&) = default; + CoveredAssignment& operator=(const CoveredAssignment&) = default; + + // Assign the agents to the tasks. + std::vector Assign() const override; + + protected: + // Validate the assignment problem. + void Validate() const override; +}; + +} // namespace assignment diff --git a/assignment/covered_assignment_test.cc b/assignment/covered_assignment_test.cc new file mode 100644 index 0000000..42fe589 --- /dev/null +++ b/assignment/covered_assignment_test.cc @@ -0,0 +1,50 @@ +#include "assignment/covered_assignment.h" + +#include + +#include +#include + +#include "assignment/assignment.h" + +namespace assignment { +namespace { + +TEST(CoveredAssignmentTest, AssignUnique) { + constexpr int kNumAgents = 5; + constexpr int kNumTasks = 4; + const std::vector> costs{ + {90, 80, 75, 70}, {35, 85, 55, 65}, {125, 95, 90, 95}, + {45, 110, 95, 115}, {50, 100, 90, 100}, + }; + CoveredAssignment assignment(kNumAgents, kNumTasks, costs); + const auto assignments = assignment.Assign(); + std::unordered_map expected_assignments{ + {0, 3}, {1, 2}, {2, 1}, {3, 0}, {4, 0}}; + for (const auto& [agent_index, task_index] : assignments) { + EXPECT_EQ(expected_assignments[agent_index], task_index) + << "Agent " << agent_index << " was assigned to task " << task_index + << " but expected task " << expected_assignments[agent_index] << "."; + } +} + +TEST(CoveredAssignmentTest, AssignMultiple) { + constexpr int kNumAgents = 3; + constexpr int kNumTasks = 2; + const std::vector> costs{ + {1, 3}, + {4, 3}, + {2, 2}, + }; + CoveredAssignment assignment(kNumAgents, kNumTasks, costs); + const auto assignments = assignment.Assign(); + std::unordered_map expected_assignments{{0, 0}, {1, 1}, {2, 1}}; + for (const auto& [agent_index, task_index] : assignments) { + EXPECT_EQ(expected_assignments[agent_index], task_index) + << "Agent " << agent_index << " was assigned to task " << task_index + << " but expected task " << expected_assignments[agent_index] << "."; + } +} + +} // namespace +} // namespace assignment