Address review comments

Co-Authored-By: Jacob Sznajdman <breakanalysis@gmail.com>
Co-Authored-By: Mats Rydberg <mats.rydberg@neotechnology.com>

Author: 3 people

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

@@ -44,25 +44,25 @@ class BinarizeTask implements Runnable {
     private final double[][] propertyEmbeddings;
 
     private final double threshold;
-    private final BinarizeFeaturesConfig binarizationConfig;
+    private final int dimension;
     private final ProgressTracker progressTracker;
-    private long totalNumFeatures;
+    private long totalFeatureCount;
 
     private double scalarProductSum;
 
     private double scalarProductSumOfSquares;
 
     BinarizeTask(
         Partition partition,
-        HashGNNConfig config,
+        BinarizeFeaturesConfig config,
         HugeObjectArray<HugeAtomicBitSet> truncatedFeatures,
         List<FeatureExtractor> featureExtractors,
         double[][] propertyEmbeddings,
         ProgressTracker progressTracker
     ) {
         this.partition = partition;
-        this.binarizationConfig = config.binarizeFeatures().orElseThrow();
-        this.threshold = binarizationConfig.threshold();
+        this.dimension = config.dimension();
+        this.threshold = config.threshold();
         this.truncatedFeatures = truncatedFeatures;
         this.featureExtractors = featureExtractors;
         this.propertyEmbeddings = propertyEmbeddings;
@@ -76,7 +76,7 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
         SplittableRandom rng,
         ProgressTracker progressTracker,
         TerminationFlag terminationFlag,
-        MutableLong totalNumFeaturesOutput
+        MutableLong totalFeatureCountOutput
     ) {
         progressTracker.beginSubTask("Binarize node property features");
 
@@ -88,14 +88,14 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
         );
 
         var inputDimension = FeatureExtraction.featureCount(featureExtractors);
-        var propertyEmbeddings = embedProperties(config, rng, inputDimension);
+        var propertyEmbeddings = embedProperties(binarizationConfig.dimension(), rng, inputDimension);
 
         var truncatedFeatures = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
 
         var tasks = partition.stream()
             .map(p -> new BinarizeTask(
                 p,
-                config,
+                binarizationConfig,
                 truncatedFeatures,
                 featureExtractors,
                 propertyEmbeddings,
@@ -108,7 +108,7 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
             .terminationFlag(terminationFlag)
             .run();
 
-        totalNumFeaturesOutput.add(tasks.stream().mapToLong(BinarizeTask::totalNumFeatures).sum());
+        totalFeatureCountOutput.add(tasks.stream().mapToLong(BinarizeTask::totalFeatureCount).sum());
 
         var squaredSum = tasks.stream().mapToDouble(BinarizeTask::scalarProductSumOfSquares).sum();
         var sum = tasks.stream().mapToDouble(BinarizeTask::scalarProductSum).sum();
@@ -118,38 +118,47 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
         var variance = (squaredSum - exampleCount * avg * avg) / exampleCount;
         var std = Math.sqrt(variance);
 
-        progressTracker.logInfo(formatWithLocale("Hyperplane scalar products have mean %.4f and standard deviation %.4f. A threshold for binarization may be set to the average plus a few standard deviations.", avg, std));
+        progressTracker.logInfo(formatWithLocale(
+            "Hyperplane scalar products have mean %.4f and standard deviation %.4f. A threshold for binarization may be set to the mean plus a few standard deviations.",
+            avg,
+            std
+        ));
 
         progressTracker.endSubTask("Binarize node property features");
 
         return truncatedFeatures;
     }
+
     // creates a random projection vector for each feature
     // (input features vector for each node is the concatenation of the node's properties)
     // this array is used embed the properties themselves from inputDimension to embeddingDimension dimensions.
-    public static double[][] embedProperties(HashGNNConfig config, SplittableRandom rng, int inputDimension) {
-        var binarizationConfig = config.binarizeFeatures().orElseThrow();
+    public static double[][] embedProperties(int vectorDimension, SplittableRandom rng, int inputDimension) {
         var propertyEmbeddings = new double[inputDimension][];
 
         for (int inputFeature = 0; inputFeature < inputDimension; inputFeature++) {
-            propertyEmbeddings[inputFeature] = new double[binarizationConfig.dimension()];
-            for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
-                // Box-muller transformation to generate gaussian
-                double matrixValue = Math.sqrt(-2*Math.log(rng.nextDouble(0.0, 1.0))) * Math.cos(2*Math.PI * rng.nextDouble(0.0, 1.0));
-                propertyEmbeddings[inputFeature][feature] = matrixValue;
+            propertyEmbeddings[inputFeature] = new double[vectorDimension];
+            for (int feature = 0; feature < vectorDimension; feature++) {
+                propertyEmbeddings[inputFeature][feature] = boxMullerGaussianRandom(rng);
             }
         }
         return propertyEmbeddings;
     }
 
+    private static double boxMullerGaussianRandom(SplittableRandom rng) {
+        return Math.sqrt(-2 * Math.log(rng.nextDouble(
+            0.0,
+            1.0
+        ))) * Math.cos(2 * Math.PI * rng.nextDouble(0.0, 1.0));
+    }
+
     @Override
     public void run() {
         partition.consume(nodeId -> {
-            var featureVector = new float[binarizationConfig.dimension()];
+            var featureVector = new float[dimension];
             FeatureExtraction.extract(nodeId, -1, featureExtractors, new FeatureConsumer() {
                 @Override
                 public void acceptScalar(long nodeOffset, int offset, double value) {
-                    for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
+                    for (int feature = 0; feature < dimension; feature++) {
                         double featureValue = propertyEmbeddings[offset][feature];
                         featureVector[feature] += value * featureValue;
                     }
@@ -159,7 +168,7 @@ public void acceptScalar(long nodeOffset, int offset, double value) {
                 public void acceptArray(long nodeOffset, int offset, double[] values) {
                     for (int inputFeatureOffset = 0; inputFeatureOffset < values.length; inputFeatureOffset++) {
                         double value = values[inputFeatureOffset];
-                        for (int feature = 0; feature < binarizationConfig.dimension(); feature++) {
+                        for (int feature = 0; feature < dimension; feature++) {
                             double featureValue = propertyEmbeddings[offset + inputFeatureOffset][feature];
                             featureVector[feature] += value * featureValue;
                         }
@@ -168,7 +177,7 @@ public void acceptArray(long nodeOffset, int offset, double[] values) {
             });
 
             var featureSet = round(featureVector);
-            totalNumFeatures += featureSet.cardinality();
+            totalFeatureCount += featureSet.cardinality();
             truncatedFeatures.set(nodeId, featureSet);
         });
 
@@ -188,13 +197,14 @@ private HugeAtomicBitSet round(float[] floatVector) {
         return bitset;
     }
 
-    public long totalNumFeatures() {
-        return totalNumFeatures;
+    public long totalFeatureCount() {
+        return totalFeatureCount;
     }
 
     public double scalarProductSum() {
         return scalarProductSum;
     }
+
     public double scalarProductSumOfSquares() {
         return scalarProductSumOfSquares;
     }

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

@@ -40,7 +40,7 @@ class GenerateFeaturesTask implements Runnable {
     private final GenerateFeaturesConfig generateFeaturesConfig;
     private final ProgressTracker progressTracker;
     private final long randomSeed;
-    private long totalNumFeatures = 0;
+    private long totalFeatureCount = 0;
 
     GenerateFeaturesTask(
         Partition partition,
@@ -66,7 +66,7 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
         long randomSeed,
         ProgressTracker progressTracker,
         TerminationFlag terminationFlag,
-        MutableLong totalNumFeaturesOutput
+        MutableLong totalFeatureCountOutput
     ) {
         progressTracker.beginSubTask("Generate base node property features");
 
@@ -77,7 +77,7 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
                 p,
                 graph,
                 randomSeed,
-                config.generateFeatures().get(),
+                config.generateFeatures().orElseThrow(),
                 output,
                 progressTracker
             ))
@@ -88,7 +88,7 @@ static HugeObjectArray<HugeAtomicBitSet> compute(
             .terminationFlag(terminationFlag)
             .run();
 
-        totalNumFeaturesOutput.add(tasks.stream().mapToLong(GenerateFeaturesTask::totalNumFeatures).sum());
+        totalFeatureCountOutput.add(tasks.stream().mapToLong(GenerateFeaturesTask::totalFeatureCount).sum());
 
         progressTracker.endSubTask("Generate base node property features");
 
@@ -104,21 +104,19 @@ public void run() {
             var generatedFeatures = HugeAtomicBitSet.create(dimension);
 
             rng.setSeed(this.randomSeed ^ graph.toOriginalNodeId(nodeId));
-            // without this we get same results for different result, at least on example in doc test
-            rng.setSeed(rng.nextLong());
 
             var randomInts = rng.ints(densityLevel, 0, dimension);
             randomInts.forEach(generatedFeatures::set);
 
-            totalNumFeatures += generatedFeatures.cardinality();
+            totalFeatureCount += generatedFeatures.cardinality();
 
             output.set(nodeId, generatedFeatures);
         });
 
         progressTracker.logProgress(partition.nodeCount());
     }
 
-    public long totalNumFeatures() {
-        return totalNumFeatures;
+    public long totalFeatureCount() {
+        return totalFeatureCount;
     }
 }

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

@@ -19,20 +19,16 @@
  */
 package org.neo4j.gds.embeddings.hashgnn;
 
-import org.apache.commons.lang3.mutable.MutableInt;
 import org.apache.commons.lang3.mutable.MutableLong;
 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;
@@ -51,7 +47,7 @@ public class HashGNN extends Algorithm<HashGNN.HashGNNResult> {
     private final Graph graph;
     private final SplittableRandom rng;
     private final HashGNNConfig config;
-    private final MutableLong totalSetBits = new MutableLong();
+    private final MutableLong currentTotalFeatureCount = new MutableLong();
 
     public HashGNN(Graph graph, HashGNNConfig config, ProgressTracker progressTracker) {
         super(progressTracker);
@@ -93,8 +89,11 @@ public HashGNNResult compute() {
         var embeddingsB = constructInputEmbeddings(rangePartition);
         int embeddingDimension = (int) embeddingsB.get(0).size();
 
-        double avgInputActiveFeatures = totalSetBits.doubleValue() / graph.nodeCount();
-        progressTracker.logInfo(formatWithLocale("Density (number of active features) of binary input features is %.4f.", avgInputActiveFeatures));
+        double avgInputActiveFeatures = currentTotalFeatureCount.doubleValue() / graph.nodeCount();
+        progressTracker.logInfo(formatWithLocale(
+            "Density (number of active features) of binary input features is %.4f.",
+            avgInputActiveFeatures
+        ));
 
         var embeddingsA = HugeObjectArray.newArray(HugeAtomicBitSet.class, graph.nodeCount());
         embeddingsA.setAll(unused -> HugeAtomicBitSet.create(embeddingDimension));
@@ -104,7 +103,8 @@ public HashGNNResult compute() {
             ? 1
             : embeddingDimension * (1 - Math.pow(
                 1 - (1.0 / embeddingDimension),
-                avgDegree)
+                avgDegree
+            )
             );
 
         progressTracker.beginSubTask("Propagate embeddings");
@@ -118,8 +118,8 @@ public HashGNNResult compute() {
                 currentEmbeddings.get(i).clear();
             }
 
-            double scaledNeighborInfluence = graph.relationshipCount() == 0 ? 1.0 : (totalSetBits.doubleValue() / graph.nodeCount()) * config.neighborInfluence() / upperBoundNeighborExpectedBits;
-            totalSetBits.setValue(0);
+            double scaledNeighborInfluence = graph.relationshipCount() == 0 ? 1.0 : (currentTotalFeatureCount.doubleValue() / graph.nodeCount()) * config.neighborInfluence() / upperBoundNeighborExpectedBits;
+            currentTotalFeatureCount.setValue(0);
 
             var hashes = HashTask.compute(
                 embeddingDimension,
@@ -141,11 +141,15 @@ public HashGNNResult compute() {
                 hashes,
                 progressTracker,
                 terminationFlag,
-                totalSetBits
+                currentTotalFeatureCount
             );
 
-            double avgActiveFeatures = totalSetBits.doubleValue() / graph.nodeCount();
-            progressTracker.logInfo(formatWithLocale("After iteration %d average node embedding density (number of active features) is %.4f.", iteration, avgActiveFeatures));
+            double avgActiveFeatures = currentTotalFeatureCount.doubleValue() / graph.nodeCount();
+            progressTracker.logInfo(formatWithLocale(
+                "After iteration %d average node embedding density (number of active features) is %.4f.",
+                iteration,
+                avgActiveFeatures
+            ));
         }
 
         progressTracker.endSubTask("Propagate embeddings");
@@ -209,74 +213,37 @@ 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(
+                return BinarizeTask.compute(
                     graph,
                     partition,
                     config,
                     rng,
                     progressTracker,
                     terminationFlag,
-                    totalSetBits
-                ));
-                bitOffsets.add(embeddingDimension.getValue());
-                embeddingDimension.add(config.binarizeFeatures().get().dimension());
+                    currentTotalFeatureCount
+                );
             } else {
-                inputEmbeddingsList.add(RawFeaturesTask.compute(
+                return RawFeaturesTask.compute(
                     config,
                     progressTracker,
                     graph,
                     partition,
                     terminationFlag,
-                    totalSetBits
-                ));
-                var featureExtractors = FeatureExtraction.propertyExtractors(
-                    graph,
-                    config.featureProperties()
+                    currentTotalFeatureCount
                 );
-                bitOffsets.add(embeddingDimension.getValue());
-                embeddingDimension.add(FeatureExtraction.featureCount(featureExtractors));
             }
+        } else {
+            return GenerateFeaturesTask.compute(
+                graph,
+                partition,
+                config,
+                randomSeed,
+                progressTracker,
+                terminationFlag,
+                currentTotalFeatureCount
+            );
         }
-
-        if (!config.generateFeatures().isPresent()) {
-            return inputEmbeddingsList.get(0);
-        }
-
-        inputEmbeddingsList.add(GenerateFeaturesTask.compute(
-            graph,
-            partition,
-            config,
-            randomSeed,
-            progressTracker,
-            terminationFlag,
-            totalSetBits
-        ));
-        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;
     }
 }