Decouple embedding initializer

Author: jjaderberg

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.embeddings.node2vec;
+
+import org.neo4j.gds.utils.StringJoining;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
+import static org.neo4j.gds.utils.StringFormatting.toUpperCaseWithLocale;
+
+public enum EmbeddingInitializer {
+    UNIFORM,
+    NORMALIZED;
+
+    private static final List<String> VALUES = Arrays
+        .stream(EmbeddingInitializer.values())
+        .map(EmbeddingInitializer::name)
+        .collect(Collectors.toList());
+
+    public static EmbeddingInitializer parse(Object input) {
+        if (input instanceof String) {
+            var inputString = toUpperCaseWithLocale((String) input);
+
+            if (!VALUES.contains(inputString)) {
+                throw new IllegalArgumentException(formatWithLocale(
+                    "EmbeddingInitializer `%s` is not supported. Must be one of: %s.",
+                    input,
+                    StringJoining.join(VALUES)
+                ));
+            }
+
+            return valueOf(toUpperCaseWithLocale(inputString));
+        } else if (input instanceof EmbeddingInitializer) {
+            return (EmbeddingInitializer) input;
+        }
+
+        throw new IllegalArgumentException(formatWithLocale(
+            "Expected EmbeddingInitializer or String. Got %s.",
+            input.getClass().getSimpleName()
+        ));
+    }
+
+    public static String toString(EmbeddingInitializer embeddingInitializer) {
+        return embeddingInitializer.toString();
+    }
+}

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

@@ -47,7 +47,7 @@ public class Node2Vec extends Algorithm<Node2VecModel.Result> {
     private final List<Long> sourceNodes;
     private final Optional<Long> maybeRandomSeed;
     private final TrainParameters trainParameters;
-    private final Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer;
+    private final EmbeddingInitializer embeddingInitializer;
 
 
     public static MemoryEstimation memoryEstimation(int walksPerNode, int walkLength, int embeddingDimension) {
@@ -83,7 +83,7 @@ public Node2Vec(
         Optional<Long> maybeRandomSeed,
         ProgressTracker progressTracker,
         TrainParameters trainParameters,
-        Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer
+        EmbeddingInitializer embeddingInitializer
     ) {
         super(progressTracker);
         this.graph = graph;

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

@@ -24,53 +24,12 @@
 import org.neo4j.gds.config.AlgoBaseConfig;
 import org.neo4j.gds.config.EmbeddingDimensionConfig;
 import org.neo4j.gds.traversal.RandomWalkBaseConfig;
-import org.neo4j.gds.utils.StringJoining;
 
-import java.util.Arrays;
 import java.util.List;
-import java.util.stream.Collectors;
 
-import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
-import static org.neo4j.gds.utils.StringFormatting.toUpperCaseWithLocale;
 
 public interface Node2VecBaseConfig extends AlgoBaseConfig, EmbeddingDimensionConfig, RandomWalkBaseConfig {
 
-    enum EmbeddingInitializer {
-        UNIFORM,
-        NORMALIZED;
-
-        private static final List<String> VALUES = Arrays
-            .stream(Node2VecBaseConfig.EmbeddingInitializer.values())
-            .map(Node2VecBaseConfig.EmbeddingInitializer::name)
-            .collect(Collectors.toList());
-        public static Node2VecBaseConfig.EmbeddingInitializer parse(Object input) {
-            if (input instanceof String) {
-                var inputString = toUpperCaseWithLocale((String) input);
-
-                if (!VALUES.contains(inputString)) {
-                    throw new IllegalArgumentException(formatWithLocale(
-                        "EmbeddingInitializer `%s` is not supported. Must be one of: %s.",
-                        input,
-                        StringJoining.join(VALUES)
-                    ));
-                }
-
-                return valueOf(toUpperCaseWithLocale(inputString));
-            } else if (input instanceof Node2VecBaseConfig.EmbeddingInitializer) {
-                return (Node2VecBaseConfig.EmbeddingInitializer) input;
-            }
-
-            throw new IllegalArgumentException(formatWithLocale(
-                "Expected EmbeddingInitializer or String. Got %s.",
-                input.getClass().getSimpleName()
-            ));
-        }
-
-        public static String toString(Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer) {
-            return embeddingInitializer.toString();
-        }
-    }
-
     @Value.Default
     @Configuration.IntegerRange(min = 2)
     default int windowSize() {
@@ -103,8 +62,8 @@ default int embeddingDimension() {
     }
 
     @Value.Default
-    @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.node2vec.Node2VecBaseConfig.EmbeddingInitializer#parse", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
-    @Configuration.ToMapValue("org.neo4j.gds.embeddings.node2vec.Node2VecBaseConfig.EmbeddingInitializer#toString")
+    @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.node2vec.EmbeddingInitializer#parse", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
+    @Configuration.ToMapValue("org.neo4j.gds.embeddings.node2vec.EmbeddingInitializer#toString")
     default EmbeddingInitializer embeddingInitializer() {
         return EmbeddingInitializer.NORMALIZED;
     }

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

@@ -54,7 +54,7 @@ public class Node2VecModel {
     private final int embeddingDimension;
     private final int windowSize;
     private final int negativeSamplingRate;
-    private final Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer;
+    private final EmbeddingInitializer embeddingInitializer;
     private final int concurrency;
     private final CompressedRandomWalks walks;
     private final RandomWalkProbabilities randomWalkProbabilities;
@@ -80,7 +80,7 @@ public static MemoryEstimation memoryEstimation(int embeddingDimension) {
         LongUnaryOperator toOriginalId,
         long nodeCount,
         TrainParameters trainParameters,
-        Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer,
+        EmbeddingInitializer embeddingInitializer,
         int concurrency,
         Optional<Long> maybeRandomSeed,
         CompressedRandomWalks walks,
@@ -114,7 +114,7 @@ public static MemoryEstimation memoryEstimation(int embeddingDimension) {
         int windowSize,
         int negativeSamplingRate,
         int embeddingDimension,
-        Node2VecBaseConfig.EmbeddingInitializer embeddingInitializer,
+        EmbeddingInitializer embeddingInitializer,
         int concurrency,
         Optional<Long> maybeRandomSeed,
         CompressedRandomWalks walks,

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

@@ -40,6 +40,8 @@
 import org.neo4j.gds.StoreLoaderBuilder;
 import org.neo4j.gds.TestProgressTracker;
 import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.collections.ha.HugeLongArray;
+import org.neo4j.gds.collections.ha.HugeObjectArray;
 import org.neo4j.gds.collections.hsa.HugeSparseLongArray;
 import org.neo4j.gds.compat.Neo4jProxy;
 import org.neo4j.gds.compat.TestLog;
@@ -50,12 +52,9 @@
 import org.neo4j.gds.core.loading.construction.GraphFactory;
 import org.neo4j.gds.core.loading.construction.RelationshipsBuilder;
 import org.neo4j.gds.core.utils.Intersections;
-import org.neo4j.gds.collections.ha.HugeLongArray;
-import org.neo4j.gds.collections.ha.HugeObjectArray;
 import org.neo4j.gds.core.utils.progress.EmptyTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.shuffle.ShuffleUtil;
-import org.neo4j.gds.embeddings.node2vec.Node2VecBaseConfig.EmbeddingInitializer;
 import org.neo4j.gds.gdl.GdlFactory;
 import org.neo4j.gds.ml.core.tensor.FloatVector;