Benchmark and expose disjoint-set islands
Signed-off-by: teerthsharma <teerths57@gmail.com>
This commit is contained in:
committed by
teerthsharma
parent
ba2782140f
commit
5d91d878c2
@@ -79,6 +79,12 @@ mujoco_test(
|
||||
ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
|
||||
)
|
||||
|
||||
mujoco_test(
|
||||
island_benchmark_test
|
||||
MAIN_TARGET benchmark::benchmark_main
|
||||
ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
|
||||
)
|
||||
|
||||
mujoco_test(
|
||||
engine_util_sparse_benchmark_test
|
||||
MAIN_TARGET benchmark::benchmark_main
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Benchmarks island discovery on deterministic geodesic Rips graphs. The
|
||||
// corpus spans the connectivity transition of points sampled on S^2 and adds
|
||||
// MuJoCo-relevant static and repeated incidences. Corpus construction and
|
||||
// validation are deliberately outside the timed region.
|
||||
|
||||
#include <benchmark/benchmark.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "src/engine/engine_island.h"
|
||||
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
struct Edge {
|
||||
int first;
|
||||
int second;
|
||||
};
|
||||
|
||||
struct GraphCase {
|
||||
std::string name;
|
||||
int node_count;
|
||||
std::vector<Edge> incidences;
|
||||
std::vector<int> expected_partition;
|
||||
int active_nodes;
|
||||
int expected_components;
|
||||
int pre_bridge_components;
|
||||
bool bridge_added;
|
||||
std::uint64_t expected_checksum;
|
||||
};
|
||||
|
||||
struct Point {
|
||||
double x;
|
||||
double y;
|
||||
double z;
|
||||
};
|
||||
|
||||
std::uint64_t SplitMix64(std::uint64_t& state) {
|
||||
state += 0x9e3779b97f4a7c15ULL;
|
||||
std::uint64_t value = state;
|
||||
value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL;
|
||||
value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL;
|
||||
return value ^ (value >> 31);
|
||||
}
|
||||
|
||||
double Uniform01(std::uint64_t& state) {
|
||||
return static_cast<double>(SplitMix64(state) >> 11) * 0x1.0p-53;
|
||||
}
|
||||
|
||||
std::vector<Point> SampleSphere(int count, std::uint64_t seed) {
|
||||
constexpr double kTwoPi = 6.283185307179586476925286766559;
|
||||
std::vector<Point> points;
|
||||
points.reserve(count);
|
||||
for (int i = 0; i < count; ++i) {
|
||||
const double z = 2.0 * Uniform01(seed) - 1.0;
|
||||
const double angle = kTwoPi * Uniform01(seed);
|
||||
const double radial = std::sqrt(std::max(0.0, 1.0 - z*z));
|
||||
points.push_back({radial * std::cos(angle), radial * std::sin(angle), z});
|
||||
}
|
||||
return points;
|
||||
}
|
||||
|
||||
double Dot(const Point& a, const Point& b) {
|
||||
return a.x*b.x + a.y*b.y + a.z*b.z;
|
||||
}
|
||||
|
||||
std::vector<Edge> RipsEdges(const std::vector<Point>& points, double target_degree) {
|
||||
const double probability = target_degree / (points.size() - 1);
|
||||
const double radius = 2.0 * std::asin(std::sqrt(probability));
|
||||
const double minimum_dot = std::cos(radius);
|
||||
std::vector<Edge> edges;
|
||||
for (int i = 0; i < static_cast<int>(points.size()); ++i) {
|
||||
for (int j = i + 1; j < static_cast<int>(points.size()); ++j) {
|
||||
if (Dot(points[i], points[j]) >= minimum_dot) {
|
||||
edges.push_back({i, j});
|
||||
}
|
||||
}
|
||||
}
|
||||
return edges;
|
||||
}
|
||||
|
||||
std::vector<int> CanonicalPartition(int node_count, const std::vector<Edge>& edges) {
|
||||
std::vector<std::vector<int>> adjacency(node_count);
|
||||
std::vector<bool> active(node_count, false);
|
||||
for (const Edge& edge : edges) {
|
||||
active[edge.first] = true;
|
||||
active[edge.second] = true;
|
||||
if (edge.first != edge.second) {
|
||||
adjacency[edge.first].push_back(edge.second);
|
||||
adjacency[edge.second].push_back(edge.first);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> partition(node_count, -1);
|
||||
std::queue<int> pending;
|
||||
for (int start = 0; start < node_count; ++start) {
|
||||
if (!active[start] || partition[start] != -1) {
|
||||
continue;
|
||||
}
|
||||
partition[start] = start;
|
||||
pending.push(start);
|
||||
while (!pending.empty()) {
|
||||
const int node = pending.front();
|
||||
pending.pop();
|
||||
for (int neighbor : adjacency[node]) {
|
||||
if (partition[neighbor] == -1) {
|
||||
partition[neighbor] = start;
|
||||
pending.push(neighbor);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return partition;
|
||||
}
|
||||
|
||||
int CountComponents(const std::vector<int>& partition) {
|
||||
int count = 0;
|
||||
for (int node = 0; node < static_cast<int>(partition.size()); ++node) {
|
||||
count += partition[node] == node;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
std::uint64_t PartitionChecksum(const std::vector<int>& partition) {
|
||||
std::uint64_t hash = 1469598103934665603ULL;
|
||||
for (int value : partition) {
|
||||
hash ^= static_cast<std::uint32_t>(value);
|
||||
hash *= 1099511628211ULL;
|
||||
}
|
||||
return hash;
|
||||
}
|
||||
|
||||
std::vector<int> CanonicalizeLabels(const std::vector<int>& labels) {
|
||||
std::vector<int> minimum(labels.size(), std::numeric_limits<int>::max());
|
||||
for (int node = 0; node < static_cast<int>(labels.size()); ++node) {
|
||||
if (labels[node] >= 0) {
|
||||
minimum[labels[node]] = std::min(minimum[labels[node]], node);
|
||||
}
|
||||
}
|
||||
std::vector<int> canonical(labels.size(), -1);
|
||||
for (int node = 0; node < static_cast<int>(labels.size()); ++node) {
|
||||
if (labels[node] >= 0) {
|
||||
canonical[node] = minimum[labels[node]];
|
||||
}
|
||||
}
|
||||
return canonical;
|
||||
}
|
||||
|
||||
void DeterministicShuffle(std::vector<Edge>& edges, std::uint64_t seed) {
|
||||
for (std::size_t i = edges.size(); i > 1; --i) {
|
||||
const std::size_t j = SplitMix64(seed) % i;
|
||||
std::swap(edges[i - 1], edges[j]);
|
||||
}
|
||||
}
|
||||
|
||||
bool AddCriticalBridge(const std::vector<Point>& points, std::vector<Edge>& edges) {
|
||||
const std::vector<int> partition = CanonicalPartition(points.size(), edges);
|
||||
if (CountComponents(partition) < 2) {
|
||||
return false;
|
||||
}
|
||||
|
||||
double best_dot = -2.0;
|
||||
Edge bridge{-1, -1};
|
||||
for (int i = 0; i < static_cast<int>(points.size()); ++i) {
|
||||
for (int j = i + 1; j < static_cast<int>(points.size()); ++j) {
|
||||
if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] &&
|
||||
Dot(points[i], points[j]) > best_dot) {
|
||||
best_dot = Dot(points[i], points[j]);
|
||||
bridge = {i, j};
|
||||
}
|
||||
}
|
||||
}
|
||||
if (bridge.first >= 0) {
|
||||
edges.push_back(bridge);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
GraphCase MakeCase(std::string name, int node_count, double target_degree,
|
||||
std::uint64_t seed, bool critical_bridge, bool static_rows,
|
||||
bool repeated_rows) {
|
||||
const std::vector<Point> points = SampleSphere(node_count, seed);
|
||||
std::vector<Edge> edges = RipsEdges(points, target_degree);
|
||||
const int pre_bridge_components =
|
||||
critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1;
|
||||
bool bridge_added = false;
|
||||
if (critical_bridge) {
|
||||
bridge_added = AddCriticalBridge(points, edges);
|
||||
}
|
||||
|
||||
const std::vector<Edge> unique_edges = edges;
|
||||
if (static_rows) {
|
||||
for (int node = 0; node < node_count; node += 17) {
|
||||
edges.push_back({node, node});
|
||||
}
|
||||
}
|
||||
if (repeated_rows) {
|
||||
for (std::size_t i = 0; i < unique_edges.size(); i += 11) {
|
||||
edges.push_back(unique_edges[i]);
|
||||
edges.push_back({unique_edges[i].second, unique_edges[i].first});
|
||||
}
|
||||
}
|
||||
DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL);
|
||||
|
||||
std::vector<int> expected = CanonicalPartition(node_count, edges);
|
||||
const int active_nodes = std::count_if(expected.begin(), expected.end(),
|
||||
[](int component) { return component >= 0; });
|
||||
const int components = CountComponents(expected);
|
||||
const std::uint64_t checksum = PartitionChecksum(expected);
|
||||
return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes,
|
||||
components, pre_bridge_components, bridge_added, checksum};
|
||||
}
|
||||
|
||||
const std::vector<GraphCase>& Corpus() {
|
||||
static const std::vector<GraphCase> corpus = {
|
||||
MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false),
|
||||
MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0),
|
||||
0x33960002ULL, true, false, false),
|
||||
MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)),
|
||||
0x33960003ULL, false, false, false),
|
||||
MakeCase("GroundedStaticRepeated_S2Rips_256", 256,
|
||||
2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true),
|
||||
MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL,
|
||||
false, false, true),
|
||||
MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)),
|
||||
0x33960006ULL, true, false, false),
|
||||
};
|
||||
return corpus;
|
||||
}
|
||||
|
||||
struct FloodFillWorkspace {
|
||||
explicit FloodFillWorkspace(int node_count)
|
||||
: adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count),
|
||||
colind(node_count * node_count), stack(node_count * node_count + node_count),
|
||||
island(node_count) {}
|
||||
|
||||
std::vector<unsigned char> adjacency;
|
||||
std::vector<int> rownnz;
|
||||
std::vector<int> rowadr;
|
||||
std::vector<int> colind;
|
||||
std::vector<int> stack;
|
||||
std::vector<int> island;
|
||||
};
|
||||
|
||||
int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) {
|
||||
const int n = graph.node_count;
|
||||
std::fill(work.adjacency.begin(), work.adjacency.end(), 0);
|
||||
std::fill(work.rownnz.begin(), work.rownnz.end(), 0);
|
||||
for (const Edge& edge : graph.incidences) {
|
||||
work.adjacency[edge.first*n + edge.second] = 1;
|
||||
work.adjacency[edge.second*n + edge.first] = 1;
|
||||
}
|
||||
|
||||
int address = 0;
|
||||
for (int row = 0; row < n; ++row) {
|
||||
work.rowadr[row] = address;
|
||||
for (int column = 0; column < n; ++column) {
|
||||
if (work.adjacency[row*n + column]) {
|
||||
work.colind[address++] = column;
|
||||
++work.rownnz[row];
|
||||
}
|
||||
}
|
||||
}
|
||||
return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(),
|
||||
work.colind.data(), work.stack.data());
|
||||
}
|
||||
|
||||
struct DsuWorkspace {
|
||||
explicit DsuWorkspace(int node_count)
|
||||
: parent(node_count), island(node_count), dof_count(node_count, 1) {}
|
||||
|
||||
std::vector<int> parent;
|
||||
std::vector<int> island;
|
||||
std::vector<int> dof_count;
|
||||
};
|
||||
|
||||
int RunDsu(const GraphCase& graph, DsuWorkspace& work) {
|
||||
_mjPRIVATE_dsuInit(work.parent.data(), graph.node_count);
|
||||
for (const Edge& edge : graph.incidences) {
|
||||
_mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second);
|
||||
}
|
||||
int dof_count = 0;
|
||||
return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(),
|
||||
graph.node_count, &dof_count);
|
||||
}
|
||||
|
||||
bool Validate(const GraphCase& graph) {
|
||||
FloodFillWorkspace flood(graph.node_count);
|
||||
DsuWorkspace dsu(graph.node_count);
|
||||
const int flood_components = RunFloodFill(graph, flood);
|
||||
const int dsu_components = RunDsu(graph, dsu);
|
||||
const bool bridge_valid = graph.pre_bridge_components < 0 ||
|
||||
(graph.bridge_added &&
|
||||
graph.pre_bridge_components == graph.expected_components + 1);
|
||||
return bridge_valid && flood_components == graph.expected_components &&
|
||||
dsu_components == graph.expected_components &&
|
||||
CanonicalizeLabels(flood.island) == graph.expected_partition &&
|
||||
CanonicalizeLabels(dsu.island) == graph.expected_partition &&
|
||||
PartitionChecksum(graph.expected_partition) == graph.expected_checksum;
|
||||
}
|
||||
|
||||
void BM_FloodFill(benchmark::State& state, const GraphCase* graph) {
|
||||
if (!Validate(*graph)) {
|
||||
state.SkipWithError("invalid S2-Rips graph fixture");
|
||||
return;
|
||||
}
|
||||
FloodFillWorkspace work(graph->node_count);
|
||||
state.SetLabel("edges=" + std::to_string(graph->incidences.size()) +
|
||||
" active=" + std::to_string(graph->active_nodes) +
|
||||
" components=" + std::to_string(graph->expected_components) +
|
||||
" pre_bridge=" + std::to_string(graph->pre_bridge_components) +
|
||||
" bridge_added=" + std::to_string(graph->bridge_added) +
|
||||
" checksum=" + std::to_string(graph->expected_checksum));
|
||||
for (auto _ : state) {
|
||||
int components = RunFloodFill(*graph, work);
|
||||
benchmark::DoNotOptimize(components);
|
||||
benchmark::ClobberMemory();
|
||||
}
|
||||
state.SetItemsProcessed(state.iterations() * graph->incidences.size());
|
||||
}
|
||||
|
||||
void BM_Dsu(benchmark::State& state, const GraphCase* graph) {
|
||||
if (!Validate(*graph)) {
|
||||
state.SkipWithError("invalid S2-Rips graph fixture");
|
||||
return;
|
||||
}
|
||||
DsuWorkspace work(graph->node_count);
|
||||
state.SetLabel("edges=" + std::to_string(graph->incidences.size()) +
|
||||
" active=" + std::to_string(graph->active_nodes) +
|
||||
" components=" + std::to_string(graph->expected_components) +
|
||||
" pre_bridge=" + std::to_string(graph->pre_bridge_components) +
|
||||
" bridge_added=" + std::to_string(graph->bridge_added) +
|
||||
" checksum=" + std::to_string(graph->expected_checksum));
|
||||
for (auto _ : state) {
|
||||
int components = RunDsu(*graph, work);
|
||||
benchmark::DoNotOptimize(components);
|
||||
benchmark::ClobberMemory();
|
||||
}
|
||||
state.SetItemsProcessed(state.iterations() * graph->incidences.size());
|
||||
}
|
||||
|
||||
const bool kRegistered = [] {
|
||||
for (const GraphCase& graph : Corpus()) {
|
||||
benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph);
|
||||
benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph);
|
||||
}
|
||||
return true;
|
||||
}();
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
@@ -16,6 +16,8 @@
|
||||
|
||||
#include "src/engine/engine_island.h"
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
@@ -35,6 +37,310 @@ using ::testing::NotNull;
|
||||
using ::testing::Pointwise;
|
||||
using IslandTest = MujocoTest;
|
||||
|
||||
TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) {
|
||||
int parent[] = {8, 6, 7, 5};
|
||||
|
||||
_mjPRIVATE_dsuInit(parent, 0);
|
||||
EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5));
|
||||
|
||||
_mjPRIVATE_dsuInit(parent, 4);
|
||||
EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) {
|
||||
int parent[] = {0, 0, 1, 2, 3};
|
||||
|
||||
EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0);
|
||||
EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0);
|
||||
EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) {
|
||||
int parent[] = {-1, -1, -1, -1, -1, -1};
|
||||
|
||||
_mjPRIVATE_dsuUnion(parent, -1, 4);
|
||||
_mjPRIVATE_dsuUnion(parent, 3, -1);
|
||||
_mjPRIVATE_dsuUnion(parent, 5, 2);
|
||||
_mjPRIVATE_dsuUnion(parent, 4, 5);
|
||||
_mjPRIVATE_dsuUnion(parent, 3, 4);
|
||||
|
||||
EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
|
||||
for (int tree = 2; tree < 6; ++tree) {
|
||||
EXPECT_EQ(_mjPRIVATE_dsuFind(parent, tree), 2);
|
||||
}
|
||||
EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) {
|
||||
int parent[] = {-1, -1, -1, -1};
|
||||
_mjPRIVATE_dsuUnion(parent, 3, 1);
|
||||
_mjPRIVATE_dsuUnion(parent, 2, 1);
|
||||
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
|
||||
|
||||
_mjPRIVATE_dsuUnion(parent, 1, 3);
|
||||
_mjPRIVATE_dsuUnion(parent, 3, 1);
|
||||
_mjPRIVATE_dsuUnion(parent, 2, 2);
|
||||
_mjPRIVATE_dsuUnion(parent, -1, 2);
|
||||
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) {
|
||||
int self_parent[] = {-1, -1, -1};
|
||||
_mjPRIVATE_dsuUnion(self_parent, 1, 1);
|
||||
EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1));
|
||||
|
||||
int static_first[] = {-1, -1, -1};
|
||||
_mjPRIVATE_dsuUnion(static_first, -1, 2);
|
||||
EXPECT_THAT(static_first, ElementsAre(-1, -1, 2));
|
||||
|
||||
int static_second[] = {-1, -1, -1};
|
||||
_mjPRIVATE_dsuUnion(static_second, 0, -1);
|
||||
EXPECT_THAT(static_second, ElementsAre(0, -1, -1));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) {
|
||||
int distinct_parent[] = {0, 0, 2, 2};
|
||||
_mjPRIVATE_dsuUnion(distinct_parent, 1, 3);
|
||||
EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2));
|
||||
|
||||
int shared_parent[] = {0, 0, 0, 3};
|
||||
_mjPRIVATE_dsuUnion(shared_parent, 1, 2);
|
||||
EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3));
|
||||
|
||||
int long_paths[] = {0, 0, 1, 3, 3, 4};
|
||||
_mjPRIVATE_dsuUnion(long_paths, 2, 5);
|
||||
EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) {
|
||||
int parent[] = {-1, -1, -1, -1, -1, -1};
|
||||
_mjPRIVATE_dsuUnion(parent, 0, 1);
|
||||
_mjPRIVATE_dsuUnion(parent, 1, 2);
|
||||
_mjPRIVATE_dsuUnion(parent, 2, 0);
|
||||
_mjPRIVATE_dsuUnion(parent, 0, 2);
|
||||
_mjPRIVATE_dsuUnion(parent, 3, 4);
|
||||
_mjPRIVATE_dsuUnion(parent, 4, 5);
|
||||
EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3));
|
||||
|
||||
_mjPRIVATE_dsuUnion(parent, 5, 0);
|
||||
EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) {
|
||||
int parent[] = {-1, 1, 1, 3};
|
||||
|
||||
EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1),
|
||||
"self-incidence of the static tree");
|
||||
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) {
|
||||
int island[] = {71};
|
||||
int parent[] = {72};
|
||||
const int tree_dofnum[] = {73};
|
||||
int nidof = -1;
|
||||
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0);
|
||||
EXPECT_EQ(nidof, 0);
|
||||
EXPECT_THAT(island, ElementsAre(71));
|
||||
EXPECT_THAT(parent, ElementsAre(72));
|
||||
|
||||
parent[0] = -1;
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0);
|
||||
EXPECT_EQ(nidof, 0);
|
||||
EXPECT_THAT(island, ElementsAre(-1));
|
||||
EXPECT_THAT(parent, ElementsAre(-1));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) {
|
||||
int parent[] = {-1, 1, 1, 2, 4, 4, 6};
|
||||
const int tree_dofnum[] = {1000, 0, 3, 5, 7, 11, 13};
|
||||
int island[] = {9, 9, 9, 9, 9, 9, 9};
|
||||
int nidof = -1;
|
||||
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3);
|
||||
EXPECT_EQ(nidof, 39);
|
||||
EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2));
|
||||
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) {
|
||||
int parent[] = {-1, 1, 1, 2, 4, 4, 5, 7, 7, 8};
|
||||
const int tree_dofnum[] = {99, 0, 2, 3, 0, 5, 7, 11, 0, 13};
|
||||
int island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9};
|
||||
int nidof = -1;
|
||||
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3);
|
||||
EXPECT_EQ(nidof, 41);
|
||||
EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2));
|
||||
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 4, 7, 7, 7));
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) {
|
||||
constexpr int kTreeCount = 4096;
|
||||
std::vector<int> parent(kTreeCount);
|
||||
std::vector<int> island(kTreeCount, -2);
|
||||
std::vector<int> tree_dofnum(kTreeCount);
|
||||
parent[0] = 0;
|
||||
int expected_nidof = 0;
|
||||
for (int tree = 1; tree < kTreeCount; ++tree) {
|
||||
parent[tree] = tree - 1;
|
||||
tree_dofnum[tree] = tree % 5;
|
||||
expected_nidof += tree_dofnum[tree];
|
||||
}
|
||||
|
||||
int nidof = -1;
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(),
|
||||
kTreeCount, &nidof),
|
||||
1);
|
||||
EXPECT_EQ(nidof, expected_nidof);
|
||||
for (int tree = 0; tree < kTreeCount; ++tree) {
|
||||
EXPECT_EQ(island[tree], 0);
|
||||
EXPECT_EQ(parent[tree], 0);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) {
|
||||
constexpr int kTreeCount = 4096;
|
||||
std::vector<int> parent(kTreeCount);
|
||||
std::vector<int> island(kTreeCount, -2);
|
||||
std::vector<int> tree_dofnum(kTreeCount);
|
||||
_mjPRIVATE_dsuInit(parent.data(), kTreeCount);
|
||||
|
||||
int expected_nidof = 0;
|
||||
for (int tree = kTreeCount - 1; tree > 0; --tree) {
|
||||
_mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1);
|
||||
}
|
||||
for (int tree = 0; tree < kTreeCount; ++tree) {
|
||||
tree_dofnum[tree] = tree % 7;
|
||||
expected_nidof += tree_dofnum[tree];
|
||||
}
|
||||
|
||||
int nidof = -1;
|
||||
EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(),
|
||||
kTreeCount, &nidof),
|
||||
1);
|
||||
EXPECT_EQ(nidof, expected_nidof);
|
||||
for (int tree = 0; tree < kTreeCount; ++tree) {
|
||||
EXPECT_EQ(island[tree], 0);
|
||||
EXPECT_EQ(parent[tree], 0);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) {
|
||||
constexpr uint32_t kSeed = 0x5eed3396u;
|
||||
constexpr int kTrials = 2000;
|
||||
uint32_t state = kSeed;
|
||||
auto next = [&state]() {
|
||||
state = state * 1664525u + 1013904223u;
|
||||
return state;
|
||||
};
|
||||
|
||||
for (int trial = 0; trial < kTrials; ++trial) {
|
||||
const int ntree = 1 + next() % 64;
|
||||
const int nedge = next() % 192;
|
||||
std::vector<std::array<int, 2>> edges;
|
||||
edges.reserve(nedge);
|
||||
for (int edge = 0; edge < nedge; ++edge) {
|
||||
int tree1;
|
||||
int tree2;
|
||||
switch (next() % 8) {
|
||||
case 0:
|
||||
tree1 = -1;
|
||||
tree2 = next() % ntree;
|
||||
break;
|
||||
case 1:
|
||||
tree1 = next() % ntree;
|
||||
tree2 = -1;
|
||||
break;
|
||||
case 2:
|
||||
tree1 = next() % ntree;
|
||||
tree2 = tree1;
|
||||
break;
|
||||
case 3:
|
||||
if (!edges.empty()) {
|
||||
const auto& previous = edges[next() % edges.size()];
|
||||
tree1 = previous[0];
|
||||
tree2 = previous[1];
|
||||
break;
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 4:
|
||||
if (!edges.empty()) {
|
||||
const auto& previous = edges[next() % edges.size()];
|
||||
tree1 = previous[1];
|
||||
tree2 = previous[0];
|
||||
break;
|
||||
}
|
||||
[[fallthrough]];
|
||||
default:
|
||||
tree1 = next() % ntree;
|
||||
tree2 = next() % ntree;
|
||||
break;
|
||||
}
|
||||
edges.push_back({tree1, tree2});
|
||||
}
|
||||
|
||||
std::vector<int> parent(ntree);
|
||||
_mjPRIVATE_dsuInit(parent.data(), ntree);
|
||||
for (const auto& edge : edges) {
|
||||
_mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]);
|
||||
}
|
||||
|
||||
std::vector<int> active(ntree);
|
||||
std::vector<std::vector<int>> adjacent(ntree);
|
||||
for (const auto& edge : edges) {
|
||||
if (edge[0] >= 0) active[edge[0]] = 1;
|
||||
if (edge[1] >= 0) active[edge[1]] = 1;
|
||||
if (edge[0] >= 0 && edge[1] >= 0) {
|
||||
adjacent[edge[0]].push_back(edge[1]);
|
||||
adjacent[edge[1]].push_back(edge[0]);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> expected_island(ntree, -1);
|
||||
std::vector<int> expected_parent(ntree, -1);
|
||||
int expected_nisland = 0;
|
||||
for (int start = 0; start < ntree; ++start) {
|
||||
if (!active[start] || expected_island[start] != -1) continue;
|
||||
std::vector<int> pending = {start};
|
||||
std::vector<int> component;
|
||||
expected_island[start] = expected_nisland;
|
||||
while (!pending.empty()) {
|
||||
const int tree = pending.back();
|
||||
pending.pop_back();
|
||||
component.push_back(tree);
|
||||
for (int neighbor : adjacent[tree]) {
|
||||
if (expected_island[neighbor] == -1) {
|
||||
expected_island[neighbor] = expected_nisland;
|
||||
pending.push_back(neighbor);
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int tree : component) expected_parent[tree] = start;
|
||||
++expected_nisland;
|
||||
}
|
||||
|
||||
std::vector<int> tree_dofnum(ntree);
|
||||
int expected_nidof = 0;
|
||||
for (int tree = 0; tree < ntree; ++tree) {
|
||||
tree_dofnum[tree] = next() % 8;
|
||||
if (active[tree]) expected_nidof += tree_dofnum[tree];
|
||||
}
|
||||
std::vector<int> island(ntree, -2);
|
||||
int nidof = -1;
|
||||
const int nisland = _mjPRIVATE_dsuAssign(
|
||||
island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof);
|
||||
|
||||
SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial
|
||||
<< " ntree=" << ntree << " nedge=" << nedge);
|
||||
EXPECT_EQ(nisland, expected_nisland);
|
||||
EXPECT_EQ(nidof, expected_nidof);
|
||||
EXPECT_EQ(island, expected_island);
|
||||
EXPECT_EQ(parent, expected_parent);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(IslandTest, FloodFillSingleton) {
|
||||
// adjacency matrix for the graph 0 1 2
|
||||
// U U
|
||||
@@ -532,6 +838,13 @@ TEST_F(IslandTest, IslandEfc) {
|
||||
EXPECT_EQ(data->nf, 2);
|
||||
EXPECT_EQ(data->nl, 1);
|
||||
EXPECT_EQ(data->nefc, 30);
|
||||
EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(1, 0, 0, 6));
|
||||
EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 1, 1, 0));
|
||||
EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(6, 17, 1, 6));
|
||||
EXPECT_THAT(AsVector(data->efc_island, data->nefc),
|
||||
ElementsAre(0, 3, 3, 3, 3, 3, 3, 1, 2, 0,
|
||||
0, 0, 0, 0, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
|
||||
Reference in New Issue
Block a user