Fix right candidate selection for large cluster refinement

build.gradle.kts

@@ -81,6 +81,7 @@ subprojects {
 
     dependencies {
         "testImplementation"("org.junit.jupiter:junit-jupiter-api:5.3.0")
+        "testImplementation"("org.junit.jupiter:junit-jupiter-params:5.3.0")
         "testRuntimeOnly"("org.junit.jupiter:junit-jupiter-engine:5.3.0")
         "testImplementation"(group = "org.assertj", name = "assertj-core", version = "3.11.1")
 

common/src/main/java/com/bakdata/dedupe/clustering/RefineCluster.java

@@ -55,6 +55,7 @@
 import lombok.Value;
 import lombok.experimental.FieldDefaults;
 import lombok.experimental.Wither;
+import org.apache.commons.lang3.tuple.Pair;
 
 
 /**
@@ -131,19 +132,25 @@ private static double scoreClustering(final byte[] partitions, final double[][]
         return score;
     }
 
-    private static List<WeightedEdge> getRandomEdges(final int potentialNumEdges, final int desiredNumEdges) {
-        final List<WeightedEdge> weightedEdges;
-        weightedEdges = RANDOM.ints(0, potentialNumEdges)
-                .distinct()
-                .limit(desiredNumEdges)
-                .mapToObj(i -> {
-                    // reverse of Gaussian
-                    int leftIndex = (int) (Math.sqrt(i + 0.25) - 0.5);
-                    int rightIndex = i - getNumEdges(leftIndex) + leftIndex;
-                    return WeightedEdge.of(leftIndex, rightIndex, Double.NaN);
-                })
-                .collect(Collectors.toList());
-        return weightedEdges;
+    static List<WeightedEdge> getRandomEdges(final int potentialNumEdges, final int desiredNumEdges) {
+        return RANDOM.ints(0, potentialNumEdges)
+            .distinct()
+            .mapToObj(RefineCluster::createGaussPair)
+            .filter(RefineCluster::isNotSelfPair)
+            .map(p -> WeightedEdge.of(p.getLeft(), p.getRight(), Double.NaN))
+            .limit(desiredNumEdges)
+            .collect(Collectors.toList());
+    }
+
+    private static <T> boolean isNotSelfPair(final Pair<T, T> pair) {
+        return !pair.getLeft().equals(pair.getRight());
+    }
+
+    static Pair<Integer, Integer> createGaussPair(final int i) {
+        // reverse of Gaussian
+        final int leftIndex = (int) (Math.sqrt(2 * i + 0.25) - 0.5);
+        final int rightIndex = i - getNumEdges(leftIndex + 1);
+        return Pair.of(leftIndex, rightIndex);
     }
 
     private List<ClassifiedCandidate<T>> getRelevantClassifications(final Cluster<C, ? super T> cluster,

common/src/test/java/com/bakdata/dedupe/clustering/RefineClusterTest.java

@@ -0,0 +1,47 @@
+package com.bakdata.dedupe.clustering;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static com.bakdata.dedupe.clustering.RefineCluster.createGaussPair;
+import static com.bakdata.dedupe.clustering.RefineCluster.getRandomEdges;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class RefineClusterTest {
+
+	static Stream<Arguments> generateGaussPairs() {
+		final int n = 7;
+		final AtomicInteger i = new AtomicInteger();
+		return IntStream.range(0, n)
+			.boxed()
+			.flatMap(leftIndex -> IntStream.rangeClosed(0, leftIndex)
+				.boxed()
+				.map(rightIndex -> Pair.of(leftIndex, rightIndex)))
+			.map(p -> Arguments.of(i.getAndIncrement(), p.getLeft(), p.getRight()));
+	}
+
+	@ParameterizedTest
+	@MethodSource("generateGaussPairs")
+	void shouldCreateCorrectGaussPair(final int i, final int leftIndex, final int rightIndex) {
+		assertThat(createGaussPair(i))
+			.as("%d should generate %d and %d", i, leftIndex, rightIndex)
+			.satisfies(edge -> assertThat(edge.getLeft()).isEqualTo(leftIndex))
+			.satisfies(edge -> assertThat(edge.getRight()).isEqualTo(rightIndex));
+	}
+
+	@Test
+	void shouldCreateCorrectNumberOfRandomEdges() {
+		final int potentialNumEdges = 55; // gaussian sum of 11
+		final int desiredNumEdges = 45; // gaussian sum of 10
+		assertThat(getRandomEdges(potentialNumEdges, desiredNumEdges))
+			.hasSize(desiredNumEdges);
+	}
+
+}