Embedding initializer is a train parameter

Author: jjaderberg

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

@@ -47,7 +47,6 @@ public class Node2Vec extends Algorithm<Node2VecModel.Result> {
     private final List<Long> sourceNodes;
     private final Optional<Long> maybeRandomSeed;
     private final TrainParameters trainParameters;
-    private final EmbeddingInitializer embeddingInitializer;
 
 
     public static MemoryEstimation memoryEstimation(int walksPerNode, int walkLength, int embeddingDimension) {
@@ -70,8 +69,7 @@ static Node2Vec create(Graph graph, Node2VecBaseConfig config, ProgressTracker p
             config.sourceNodes(),
             config.randomSeed(),
             progressTracker,
-            config.trainParameters(),
-            config.embeddingInitializer()
+            config.trainParameters()
         );
     }
 
@@ -82,8 +80,7 @@ public Node2Vec(
         List<Long> sourceNodes,
         Optional<Long> maybeRandomSeed,
         ProgressTracker progressTracker,
-        TrainParameters trainParameters,
-        EmbeddingInitializer embeddingInitializer
+        TrainParameters trainParameters
     ) {
         super(progressTracker);
         this.graph = graph;
@@ -92,7 +89,6 @@ public Node2Vec(
         this.sourceNodes = sourceNodes;
         this.maybeRandomSeed = maybeRandomSeed;
         this.trainParameters = trainParameters;
-        this.embeddingInitializer = embeddingInitializer;
     }
 
     @Override
@@ -151,7 +147,6 @@ public Node2VecModel.Result compute() {
             graph::toOriginalNodeId,
             graph.nodeCount(),
             trainParameters,
-            embeddingInitializer,
             concurrency,
             maybeRandomSeed,
             walks,

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

@@ -114,7 +114,8 @@ default TrainParameters trainParameters() {
             iterations(),
             windowSize(),
             negativeSamplingRate(),
-            embeddingDimension()
+            embeddingDimension(),
+            embeddingInitializer()
         );
     }
 }

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

@@ -80,7 +80,6 @@ public static MemoryEstimation memoryEstimation(int embeddingDimension) {
         LongUnaryOperator toOriginalId,
         long nodeCount,
         TrainParameters trainParameters,
-        EmbeddingInitializer embeddingInitializer,
         int concurrency,
         Optional<Long> maybeRandomSeed,
         CompressedRandomWalks walks,
@@ -96,7 +95,7 @@ public static MemoryEstimation memoryEstimation(int embeddingDimension) {
             trainParameters.windowSize,
             trainParameters.negativeSamplingRate,
             trainParameters.embeddingDimension,
-            embeddingInitializer,
+            trainParameters.embeddingInitializer,
             concurrency,
             maybeRandomSeed,
             walks,

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

@@ -26,20 +26,23 @@ public class TrainParameters {
     final int windowSize;
     final int negativeSamplingRate;
     final int embeddingDimension;
+    final EmbeddingInitializer embeddingInitializer;
 
     TrainParameters(
         double initialLearningRate,
         double minLearningRate,
         int iterations,
         int windowSize,
         int negativeSamplingRate,
-        int embeddingDimension
+        int embeddingDimension,
+        EmbeddingInitializer embeddingInitializer
     ) {
         this.initialLearningRate = initialLearningRate;
         this.minLearningRate = minLearningRate;
         this.iterations = iterations;
         this.windowSize = windowSize;
         this.negativeSamplingRate = negativeSamplingRate;
         this.embeddingDimension = embeddingDimension;
+        this.embeddingInitializer = embeddingInitializer;
     }
 }