Update

Author: nv-rliu

nx_cugraph/classes/graph.py

@@ -1169,8 +1169,10 @@ def _nodearray_to_set(self, node_ids: cp.ndarray[IndexValue]) -> set[NodeKey]:
         return set(self._nodeiter_to_iter(node_ids.tolist()))
 
     def _nodearray_to_dict(
-        self, values: cp.ndarray[NodeValue]
+        self,
+        values: cp.ndarray[NodeValue],
     ) -> dict[NodeKey, NodeValue]:
+        # values_as_arrays: bool | None = None,
         it = enumerate(values.tolist())
         if (id_to_key := self.id_to_key) is not None:
             return {id_to_key[key]: val for key, val in it}

nx_cugraph/drawing/layout.py

@@ -25,6 +25,7 @@
 from nx_cugraph.utils import (
     _dtype_param,
     _get_float_dtype,
+    _seed_to_int,
     networkx_algorithm,
 )
 
@@ -92,7 +93,7 @@ def forceatlas2_layout(
         # NOTE currently only x & y (dim=2) coordinated are supported by PLC
         #   greater dimensions should be supported in the future to align with nx
         start_pos_arr = G._dict_to_nodearray(
-            pos, default=[np.nan] * 2, dtype=np.dtype(np.float32, dim)
+            pos, default=[np.nan] * dim, dtype=np.dtype((np.float32, 2))
         )
 
         # find, if there exists, the missing position values
@@ -103,11 +104,13 @@ def forceatlas2_layout(
         if num_missing:
             xy_min = cp.nanmin(start_pos_arr, axis=0)
             xy_max = cp.nanmax(start_pos_arr, axis=0)
+            # random state from seed to fill missing coords is different from random
+            # state used for PLC
             seed = create_random_state(seed)
 
             # fill missing gaps with valid random coords
             start_pos_arr[missing_vals] = xy_min + cp.asarray(
-                seed.rand(num_missing, 2), dtype=np.float32
+                seed.rand(num_missing, dim), dtype=np.float32
             ) * (xy_max - xy_min)
 
         x_start = start_pos_arr[:, 0]
@@ -116,6 +119,8 @@ def forceatlas2_layout(
         x_start = None
         y_start = None
 
+    seed = _seed_to_int(seed)
+
     vertices, x_axis, y_axis = plc.force_atlas2(
         plc.ResourceHandle(),
         random_state=seed,
@@ -126,10 +131,15 @@ def forceatlas2_layout(
         outbound_attraction_distribution=outbound_attraction_distribution,
         lin_log_mode=linlog,
         prevent_overlapping=dissuade_hubs,  # this might not be the right usage
+        edge_weight_influence=1,  # default
         jitter_tolerance=jitter_tolerance,
+        barnes_hut_optimize=False,  # default
+        barnes_hut_theta=0,  # default ?
         scaling_ratio=scaling_ratio,
         strong_gravity_mode=strong_gravity,
         gravity=gravity,
+        verbose=False,  # default
+        do_expensive_check=False,  # default
     )
 
     pos_arr = cp.column_stack((x_axis, y_axis))

nx_cugraph/utils/misc.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+# Copyright (c) 2023-2025, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -20,6 +20,7 @@
 
 import cupy as cp
 import numpy as np
+from numpy.random import RandomState
 
 if TYPE_CHECKING:
     import nx_cugraph as nxcg
@@ -124,7 +125,7 @@ def _seed_to_int(seed: int | Random | None) -> int:
     """Handle any valid seed argument and convert it to an int if necessary."""
     if seed is None:
         return
-    if isinstance(seed, Random):
+    if isinstance(seed, (Random, RandomState)):
         return seed.randint(0, sys.maxsize)
     return op.index(seed)  # Ensure seed is integral