Enhance Node2Vec TrainTask creation tests

Co-authored-by: Ioannis Panagiotas <[email protected]>

Author: vnickolov

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2VecModel.java

@@ -19,10 +19,10 @@
  */
 package org.neo4j.gds.embeddings.node2vec;
 
-import org.neo4j.gds.collections.ha.HugeLongArray;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
+import org.neo4j.gds.core.utils.partition.DegreePartition;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.BitUtil;
@@ -154,10 +154,15 @@ Node2VecResult train() {
         }
         progressTracker.endSubTask();
 
-        return  new Node2VecResult(centerEmbeddings, lossPerIteration);
+        return new Node2VecResult(centerEmbeddings, lossPerIteration);
     }
 
-    private HugeObjectArray<FloatVector> initializeEmbeddings(LongUnaryOperator toOriginalNodeId, long nodeCount, int embeddingDimensions, Random random) {
+    private HugeObjectArray<FloatVector> initializeEmbeddings(
+        LongUnaryOperator toOriginalNodeId,
+        long nodeCount,
+        int embeddingDimensions,
+        Random random
+    ) {
         HugeObjectArray<FloatVector> embeddings = HugeObjectArray.newArray(
             FloatVector.class,
             nodeCount
@@ -206,18 +211,16 @@ private TrainingTask(
             HugeObjectArray<FloatVector> centerEmbeddings,
             HugeObjectArray<FloatVector> contextEmbeddings,
             PositiveSampleProducer positiveSampleProducer,
-            HugeLongArray negativeSamples,
+            NegativeSampleProducer negativeSampleProducer,
             float learningRate,
             int negativeSamplingRate,
             int embeddingDimensions,
-            ProgressTracker progressTracker,
-            long randomSeed,
-            int taskId
+            ProgressTracker progressTracker
         ) {
             this.centerEmbeddings = centerEmbeddings;
             this.contextEmbeddings = contextEmbeddings;
             this.positiveSampleProducer = positiveSampleProducer;
-            this.negativeSampleProducer = new NegativeSampleProducer(negativeSamples, randomSeed + taskId);
+            this.negativeSampleProducer = negativeSampleProducer;
             this.learningRate = learningRate;
             this.negativeSamplingRate = negativeSamplingRate;
 
@@ -254,7 +257,7 @@ private void trainSample(long center, long context, boolean positive) {
             double positiveSigmoid = Sigmoid.sigmoid(affinity);
             double negativeSigmoid = 1 - positiveSigmoid;
 
-            lossSum -= positive ? Math.log(positiveSigmoid+EPSILON) : Math.log(negativeSigmoid+EPSILON);
+            lossSum -= positive ? Math.log(positiveSigmoid + EPSILON) : Math.log(negativeSigmoid + EPSILON);
 
             float gradient = positive ? (float) -negativeSigmoid : (float) positiveSigmoid;
             // we are doing gradient descent, so we go in the negative direction of the gradient here
@@ -290,37 +293,47 @@ void addAll(FloatConsumer other) {
         }
     }
 
-    List<TrainingTask> createTrainingTasks(float learningRate, AtomicInteger taskIndex){
+    List<TrainingTask> createTrainingTasks(float learningRate, AtomicInteger taskIndex) {
         return PartitionUtils.degreePartitionWithBatchSize(
             walks.size(),
             walks::walkLength,
             BitUtil.ceilDiv(randomWalkProbabilities.sampleCount(), concurrency.value()),
             partition -> {
-
                 var taskId = taskIndex.getAndIncrement();
-                var positiveSampleProducer = new PositiveSampleProducer(
-                    walks.iterator(partition.startNode(), partition.nodeCount()),
-                    randomWalkProbabilities.positiveSamplingProbabilities(),
-                    windowSize,
-                    Optional.of(randomSeed),
-                    taskId
-                );
-
+                var taskRandomSeed = randomSeed + taskId;
+                var positiveSampleProducer = createPositiveSampleProducer(partition, taskRandomSeed);
+                var negativeSampleProducer = createNegativeSampleProducer(taskRandomSeed);
                 return new TrainingTask(
                     centerEmbeddings,
                     contextEmbeddings,
                     positiveSampleProducer,
-                    randomWalkProbabilities.negativeSamplingDistribution(),
+                    negativeSampleProducer,
                     learningRate,
                     negativeSamplingRate,
                     embeddingDimension,
-                    progressTracker,
-                    randomSeed,
-                    taskId
+                    progressTracker
                 );
             }
         );
+    }
 
+    NegativeSampleProducer createNegativeSampleProducer(long randomSeed) {
+        return new NegativeSampleProducer(
+            randomWalkProbabilities.negativeSamplingDistribution(),
+            randomSeed
+        );
+    }
+
+    PositiveSampleProducer createPositiveSampleProducer(
+        DegreePartition partition,
+        long randomSeed
+    ) {
+        return new PositiveSampleProducer(
+            walks.iterator(partition.startNode(), partition.nodeCount()),
+            randomWalkProbabilities.positiveSamplingProbabilities(),
+            windowSize,
+            randomSeed
+        );
     }
 
 }

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/PositiveSampleProducer.java

@@ -22,9 +22,7 @@
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
 
 import java.util.Iterator;
-import java.util.Optional;
 import java.util.SplittableRandom;
-import java.util.concurrent.ThreadLocalRandom;
 
 import static org.neo4j.gds.mem.BitUtil.ceilDiv;
 
@@ -42,14 +40,13 @@ public class PositiveSampleProducer {
     private int contextWordIndex;
     private int currentWindowStart;
     private int currentWindowEnd;
-    private SplittableRandom probabilitySupplier;
+    private final SplittableRandom probabilitySupplier;
 
     PositiveSampleProducer(
         Iterator<long[]> walks,
         HugeDoubleArray samplingProbabilities,
         int windowSize,
-        Optional<Long> maybeRandomSeed,
-        int taskId
+        long randomSeed
     ) {
         this.walks = walks;
         this.samplingProbabilities = samplingProbabilities;
@@ -60,10 +57,7 @@ public class PositiveSampleProducer {
         this.currentWalk = new long[0];
         this.centerWordIndex = -1;
         this.contextWordIndex = 1;
-        probabilitySupplier = maybeRandomSeed
-            .map(seed -> new SplittableRandom(taskId + seed))
-            .orElseGet(() -> new SplittableRandom(ThreadLocalRandom.current().nextLong()));
-
+        probabilitySupplier = new SplittableRandom(randomSeed);
     }
 
     public boolean next(long[] buffer) {

algo/src/test/java/org/neo4j/gds/embeddings/node2vec/Node2VecModelTest.java

@@ -19,19 +19,31 @@
  */
 package org.neo4j.gds.embeddings.node2vec;
 
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.neo4j.gds.collections.ha.HugeLongArray;
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.Intersections;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.Optional;
 import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.LongUnaryOperator;
 import java.util.stream.LongStream;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 class Node2VecModelTest {
 
@@ -196,6 +208,53 @@ void randomSeed(int iterations) {
         }
     }
 
+    @Test
+    @DisplayName("When creating multiple tasks with random seed the actual seed for the task should be `randomSeed + taskId`.")
+    void shouldCreateTrainingTasksWithCorrectRandomSeed() {
+        var randomWalksMock = mock(CompressedRandomWalks.class);
+        when(randomWalksMock.size()).thenReturn(10L);
+        when(randomWalksMock.walkLength(anyLong())).thenReturn(3);
+
+        var randomWalkProbabilitiesMock = mock(RandomWalkProbabilities.class);
+        when(randomWalkProbabilitiesMock.sampleCount()).thenReturn(30L);
+        when(randomWalkProbabilitiesMock.negativeSamplingDistribution()).thenReturn(HugeLongArray.newArray(10));
+
+        var trainParametersMock = mock(TrainParameters.class);
+        when(trainParametersMock.embeddingInitializer()).thenReturn(EmbeddingInitializer.UNIFORM);
+
+        var node2VecModel = spy(
+            new Node2VecModel(
+                LongUnaryOperator.identity(),
+                1000,
+                trainParametersMock,
+                new Concurrency(4),
+                Optional.of(1L), // Random Seed
+                randomWalksMock,
+                randomWalkProbabilitiesMock,
+                ProgressTracker.NULL_TRACKER
+            )
+        );
+
+        var taskIdTracker = new AtomicInteger(0);
+        var trainingTasks = node2VecModel.createTrainingTasks(0.2f, taskIdTracker);
+
+        assertThat(trainingTasks).hasSize(5);
+
+        verify(node2VecModel, times(5)).createPositiveSampleProducer(any(), anyLong());
+        verify(node2VecModel, times(1)).createPositiveSampleProducer(any(), eq(1L));
+        verify(node2VecModel, times(1)).createPositiveSampleProducer(any(), eq(2L));
+        verify(node2VecModel, times(1)).createPositiveSampleProducer(any(), eq(3L));
+        verify(node2VecModel, times(1)).createPositiveSampleProducer(any(), eq(4L));
+        verify(node2VecModel, times(1)).createPositiveSampleProducer(any(), eq(5L));
+
+        verify(node2VecModel, times(5)).createNegativeSampleProducer(anyLong());
+        verify(node2VecModel, times(1)).createNegativeSampleProducer(1L);
+        verify(node2VecModel, times(1)).createNegativeSampleProducer(2L);
+        verify(node2VecModel, times(1)).createNegativeSampleProducer(3L);
+        verify(node2VecModel, times(1)).createNegativeSampleProducer(4L);
+        verify(node2VecModel, times(1)).createNegativeSampleProducer(5L);
+    }
+
     private static CompressedRandomWalks generateRandomWalks(
         RandomWalkProbabilities.Builder probabilitiesBuilder,
         long numberOfClusters,

algo/src/test/java/org/neo4j/gds/embeddings/node2vec/PositiveSampleProducerTest.java

@@ -29,7 +29,6 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.Optional;
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
 
@@ -61,7 +60,6 @@ void doesNotCauseStackOverflow() {
             walks.iterator(0, nbrOfWalks),
             HugeDoubleArray.of(LongStream.range(0, nbrOfWalks).mapToDouble((l) -> 1.0).toArray()),
             10,
-            Optional.empty(),
             0
         );
 
@@ -90,7 +88,6 @@ void doesNotCauseStackOverflowDueToBadLuck() {
             walks.iterator(0, nbrOfWalks),
             probabilities,
             10,
-            Optional.empty(),
             0
         );
         // does not overflow the stack = passes test
@@ -115,7 +112,6 @@ void doesNotAttemptToFetchOutsideBatch() {
             walks.iterator(0, nbrOfWalks / 2),
             HugeDoubleArray.of(LongStream.range(0, nbrOfWalks).mapToDouble((l) -> 1.0).toArray()),
             10,
-            Optional.empty(),
             0
         );
 
@@ -141,7 +137,6 @@ void shouldProducePairsWith(
             walks.iterator(0, walks.size()),
             centerNodeProbabilities,
             windowSize,
-            Optional.empty(),
             0
         );
         while (producer.next(buffer)) {
@@ -165,7 +160,6 @@ void shouldProducePairsWithBounds() {
             walks.iterator(0, 2),
             centerNodeProbabilities,
             3,
-            Optional.empty(),
             0
         );
         while (producer.next(buffer)) {
@@ -212,7 +206,6 @@ void shouldRemoveDownsampledWordFromWalk() {
             walks.iterator(0, walks.size()),
             centerNodeProbabilities,
             3,
-            Optional.empty(),
             0
         );