Add support for generating node features in HashGNN

Co-Authored-By: Jacob Sznajdman <breakanalysis@gmail.com>

Author: adamnsch

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/BinarizeTask.java

@@ -42,14 +42,14 @@
 
 class BinarizeTask implements Runnable {
     private final Partition partition;
-    private final HashGNNConfig config;
     private final HugeObjectArray<HugeAtomicBitSet> truncatedFeatures;
     private final List<FeatureExtractor> featureExtractors;
     private final int[][] propertyEmbeddings;
     private final List<int[]> hashesList;
     private final HashGNN.MinAndArgmin minAndArgMin;
     private final FeatureBinarizationConfig binarizationConfig;
     private final ProgressTracker progressTracker;
+    private final int sampledBits;
 
     BinarizeTask(
         Partition partition,
@@ -61,14 +61,18 @@ class BinarizeTask implements Runnable {
         ProgressTracker progressTracker
     ) {
         this.partition = partition;
-        this.config = config;
         this.binarizationConfig = config.binarizeFeatures().orElseThrow();
         this.truncatedFeatures = truncatedFeatures;
         this.featureExtractors = featureExtractors;
         this.propertyEmbeddings = propertyEmbeddings;
         this.hashesList = hashesList;
         this.minAndArgMin = new HashGNN.MinAndArgmin();
         this.progressTracker = progressTracker;
+
+        var densityOffset = config.generateFeatures().isPresent()
+            ? config.generateFeatures().get().densityLevel()
+            : 0;
+        this.sampledBits = config.embeddingDensity() - densityOffset;
     }
 
     static HugeObjectArray<HugeAtomicBitSet> compute(
@@ -193,7 +197,7 @@ private HugeAtomicBitSet roundAndSample(BitSet tempBitSet, float[] floatVector)
             }
         }
         var sampledBitset = HugeAtomicBitSet.create(binarizationConfig.dimension());
-        for (int i = 0; i < config.embeddingDensity(); i++) {
+        for (int i = 0; i < sampledBits; i++) {
             hashArgMin(tempBitSet, hashesList.get(i), minAndArgMin);
             if (minAndArgMin.argMin != -1) {
                 sampledBitset.set(minAndArgMin.argMin);

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/GenerateFeaturesTask.java

@@ -0,0 +1,104 @@
+/*
+ * 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.hashgnn;
+
+import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.core.concurrency.RunWithConcurrency;
+import org.neo4j.gds.core.utils.TerminationFlag;
+import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
+import org.neo4j.gds.core.utils.paged.HugeObjectArray;
+import org.neo4j.gds.core.utils.partition.Partition;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+
+import java.util.List;
+import java.util.SplittableRandom;
+import java.util.stream.Collectors;
+
+class GenerateFeaturesTask implements Runnable {
+    private final Partition partition;
+    private final HugeObjectArray<HugeAtomicBitSet> output;
+    private final SplittableRandom rng;
+    private final FeatureBinarizationConfig generateFeaturesConfig;
+    private final ProgressTracker progressTracker;
+
+    GenerateFeaturesTask(
+        Partition partition,
+        SplittableRandom rng,
+        FeatureBinarizationConfig config,
+        HugeObjectArray<HugeAtomicBitSet> output,
+        ProgressTracker progressTracker
+    ) {
+        this.partition = partition;
+        this.rng = rng;
+        this.generateFeaturesConfig = config;
+        this.output = output;
+        this.progressTracker = progressTracker;
+    }
+
+    static HugeObjectArray<HugeAtomicBitSet> compute(
+        Graph graph,
+        List<Partition> partition,
+        HashGNNConfig config,
+        SplittableRandom rng,
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
+    ) {
+        progressTracker.beginSubTask("Generate base node property features");
+
+        var output = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
+
+        var tasks = partition.stream()
+            .map(p -> new GenerateFeaturesTask(
+                p,
+                rng.split(),
+                config.generateFeatures().get(),
+                output,
+                progressTracker
+            ))
+            .collect(Collectors.toList());
+        RunWithConcurrency.builder()
+            .concurrency(config.concurrency())
+            .tasks(tasks)
+            .terminationFlag(terminationFlag)
+            .run();
+
+        progressTracker.endSubTask("Generate base node property features");
+
+        return output;
+    }
+
+    @Override
+    public void run() {
+        int dimension = generateFeaturesConfig.dimension();
+        int densityLevel = generateFeaturesConfig.densityLevel();
+
+        partition.consume(nodeId -> {
+            var generatedFeatures = HugeAtomicBitSet.create(dimension);
+
+            var randomInts = rng.ints(densityLevel, 0, dimension);
+            randomInts.forEach(generatedFeatures::set);
+
+            output.set(nodeId, generatedFeatures);
+        });
+
+        progressTracker.logProgress(partition.nodeCount());
+    }
+
+}

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/HashGNN.java

@@ -19,15 +19,19 @@
  */
 package org.neo4j.gds.embeddings.hashgnn;
 
+import org.apache.commons.lang3.mutable.MutableInt;
 import org.neo4j.gds.Algorithm;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.schema.GraphSchema;
+import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.paged.HugeObjectArray;
+import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.ml.core.features.FeatureExtraction;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -80,16 +84,8 @@ public HashGNNResult compute() {
             .collect(Collectors.toList())
             : List.of(graphCopy);
 
-        var embeddingsB = config.binarizeFeatures().isPresent()
-            ? BinarizeTask.compute(graph, rangePartition, config, rng, progressTracker, terminationFlag)
-            : RawFeaturesTask.compute(config, rng, progressTracker, graph, rangePartition, terminationFlag);
-        int embeddingDimension = config.binarizeFeatures().map(FeatureBinarizationConfig::dimension).orElseGet(() -> {
-            var featureExtractors = FeatureExtraction.propertyExtractors(
-                graph,
-                config.featureProperties()
-            );
-            return FeatureExtraction.featureCount(featureExtractors);
-        });
+        var embeddingsB = constructInputEmbeddings(rangePartition);
+        int embeddingDimension = (int) embeddingsB.get(0).size();
 
         var embeddingsA = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
         embeddingsA.setAll(unused -> HugeAtomicBitSet.create(embeddingDimension));
@@ -201,4 +197,73 @@ public HugeObjectArray<double[]> embeddings() {
         }
     }
 
+    private HugeObjectArray<HugeAtomicBitSet> constructInputEmbeddings(List<Partition> partition) {
+        List<HugeObjectArray<HugeAtomicBitSet>> inputEmbeddingsList = new ArrayList<>();
+        var embeddingDimension = new MutableInt();
+        var bitOffsets = new ArrayList<Integer>();
+
+        if (!config.featureProperties().isEmpty()) {
+            if (config.binarizeFeatures().isPresent()) {
+                inputEmbeddingsList.add(BinarizeTask.compute(
+                    graph,
+                    partition,
+                    config,
+                    rng,
+                    progressTracker,
+                    terminationFlag
+                ));
+                bitOffsets.add(embeddingDimension.getValue());
+                embeddingDimension.add(config.binarizeFeatures().get().dimension());
+            } else {
+                inputEmbeddingsList.add(RawFeaturesTask.compute(
+                    config,
+                    rng,
+                    progressTracker,
+                    graph,
+                    partition,
+                    terminationFlag
+                ));
+                var featureExtractors = FeatureExtraction.propertyExtractors(
+                    graph,
+                    config.featureProperties()
+                );
+                bitOffsets.add(embeddingDimension.getValue());
+                embeddingDimension.add(FeatureExtraction.featureCount(featureExtractors));
+            }
+        }
+
+        if (!config.generateFeatures().isPresent()) {
+            return inputEmbeddingsList.get(0);
+        }
+
+        inputEmbeddingsList.add(GenerateFeaturesTask.compute(
+            graph,
+            partition,
+            config,
+            rng,
+            progressTracker,
+            terminationFlag
+        ));
+        bitOffsets.add(embeddingDimension.getValue());
+        embeddingDimension.add(config.generateFeatures().get().dimension());
+
+        var concatInputEmbeddings = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
+
+        var concatTasks = partition.stream().map(p -> (Runnable) () -> p.consume(nodeId -> {
+            var concatFeatures = HugeAtomicBitSet.create(embeddingDimension.getValue());
+            for (int i = 0; i < inputEmbeddingsList.size(); i++) {
+                var embedding = inputEmbeddingsList.get(i).get(nodeId);
+                int bitOffset = bitOffsets.get(i);
+                embedding.forEachSetBit(bit -> concatFeatures.set(bitOffset + bit));
+            }
+            concatInputEmbeddings.set(nodeId, concatFeatures);
+        })).collect(Collectors.toList());
+
+        RunWithConcurrency.builder()
+            .concurrency(config.concurrency())
+            .tasks(concatTasks)
+            .run();
+
+        return concatInputEmbeddings;
+    }
 }

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/HashGNNConfig.java

@@ -49,6 +49,12 @@ default boolean heterogeneous() {
         return false;
     }
 
+    @Configuration.ToMapValue("org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#toMapBinarizationConfig")
+    @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#parseBinarizationConfig", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
+    default Optional<FeatureBinarizationConfig> generateFeatures() {
+        return Optional.empty();
+    }
+
     @Configuration.ToMapValue("org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#toMapBinarizationConfig")
     @Configuration.ConvertWith(method = "org.neo4j.gds.embeddings.hashgnn.HashGNNConfig#parseBinarizationConfig", inverse = Configuration.ConvertWith.INVERSE_IS_TO_MAP)
     default Optional<FeatureBinarizationConfig> binarizeFeatures() {

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/HashGNNFactory.java

@@ -58,10 +58,16 @@ public HashGNN build(
     public Task progressTask(Graph graph, CONFIG config) {
         var tasks = new ArrayList<Task>();
 
-        if (config.binarizeFeatures().isPresent()) {
-            tasks.add(Tasks.leaf("Binarize node property features", graph.nodeCount()));
-        } else {
-            tasks.add(Tasks.leaf("Extract raw node property features", graph.nodeCount()));
+        if (!config.featureProperties().isEmpty()) {
+            if (config.binarizeFeatures().isPresent()) {
+                tasks.add(Tasks.leaf("Binarize node property features", graph.nodeCount()));
+            } else {
+                tasks.add(Tasks.leaf("Extract raw node property features", graph.nodeCount()));
+            }
+        }
+
+        if (config.generateFeatures().isPresent()) {
+            tasks.add(Tasks.leaf("Generate base node property features", graph.nodeCount()));
         }
 
         int numRelTypes = config.heterogeneous() ? config.relationshipTypes().size() : 1;

algo/src/main/java/org/neo4j/gds/embeddings/hashgnn/RawFeaturesTask.java

@@ -41,12 +41,12 @@
 
 class RawFeaturesTask implements Runnable {
     private final Partition partition;
-    private final HashGNNConfig config;
     private final List<FeatureExtractor> featureExtractors;
     private final int inputDimension;
     private final HugeObjectArray<HugeAtomicBitSet> features;
     private final List<int[]> hashesList;
     private final ProgressTracker progressTracker;
+    private final int sampledBits;
 
     RawFeaturesTask(
         Partition partition,
@@ -58,12 +58,16 @@ class RawFeaturesTask implements Runnable {
         ProgressTracker progressTracker
     ) {
         this.partition = partition;
-        this.config = config;
         this.featureExtractors = featureExtractors;
         this.inputDimension = inputDimension;
         this.features = features;
         this.hashesList = hashesList;
         this.progressTracker = progressTracker;
+
+        var densityOffset = config.generateFeatures().isPresent()
+            ? config.generateFeatures().get().densityLevel()
+            : 0;
+        this.sampledBits = config.embeddingDensity() - densityOffset;
     }
 
     static HugeObjectArray<HugeAtomicBitSet> compute(
@@ -140,7 +144,7 @@ public void acceptArray(long nodeOffset, int offset, double[] values) {
                 return;
             }
             var sampledBitset = HugeAtomicBitSet.create(inputDimension);
-            for (int i = 0; i < config.embeddingDensity(); i++) {
+            for (int i = 0; i < sampledBits; i++) {
                 hashArgMin(nodeFeatures, hashesList.get(i), resMinAndArgMin, tempMinAndArgMin);
                 if (resMinAndArgMin.argMin != -1) {
                     sampledBitset.set(resMinAndArgMin.argMin);