Added mass and size input attributes with docstrings

networkx/drawing/layout.py

@@ -1133,6 +1133,8 @@ def forceatlas2_layout(
     distributed_action=False,
     strong_gravity=False,
     adjust_sizes=False,
+    mass="mass",
+    size="size",
     dissuade_hubs=False,
     weight=None,
     linlog=False,
@@ -1173,6 +1175,12 @@ def forceatlas2_layout(
         Normalizes the attraction forces according to the degree of the nodes
     strong_gravity : bool (default: false)
         Controls the "pull" to the center of mass of the plot
+    mass: str (default: "mass")
+        Looks up node attribute mass. The mass increases the attractive pull towards other nodes, similar
+    to a planetary body. The value defaults to 1.
+    size: str (default: "size")
+        Looks up node attribute size. .Allows for node sizes to have a larger size, preventing a "crowding effect".
+    The value defaults to 1
     adjust_sizes: bool (default: False)
         Prevent node overlapping in the layout
     dissuade_hubs : bool (default: False)
@@ -1248,13 +1256,13 @@ def estimate_factor(n, swing, traction, speed, speed_efficiency, jitter_toleranc
     # center locs around 0
     pos_arr = (np.random.rand(len(G), dim) - 0.5) * (max_pos_range - min_pos_range)
 
-    mass = np.zeros(len(G))
-    size = np.zeros(len(G))
+    masses = np.zeros(len(G))
+    sizes = np.zeros(len(G))
     for idx, node in enumerate(G.nodes()):
         if node in pos:
             pos_arr[idx] = pos[node].copy()
-        mass[idx] = G.nodes[node].get("mass", G.degree(node) + 1)
-        size[idx] = G.nodes[node].get("size", 1)
+        masses[idx] = G.nodes[node].get(masses, G.degree(node) + 1)
+        sizes[idx] = G.nodes[node].get(sizes, 1)
 
     n = G.number_of_nodes()
     gravities = np.zeros((n, dim))
@@ -1283,12 +1291,12 @@ def estimate_factor(n, swing, traction, speed, speed_efficiency, jitter_toleranc
             attraction = -np.einsum("ijk, ij -> ik", diff, A)
 
         if distributed_action:
-            attraction /= mass[:, None]
+            attraction /= masses[:, None]
 
         # repulsion
-        tmp = mass[:, None] @ mass[None]
+        tmp = masses[:, None] @ masses[None]
         if adjust_sizes:
-            distance += -size[:, None] - size[None]
+            distance += -sizes[:, None] - sizes[None]
             # multiply negative distance * 100 (squared below)
             distance[distance < 0] = 1 / 10
 
@@ -1303,7 +1311,7 @@ def estimate_factor(n, swing, traction, speed, speed_efficiency, jitter_toleranc
         # gravity
         gravities = (
             -gravity
-            * mass[:, None]
+            * masses[:, None]
             * pos_arr
             / np.linalg.norm(pos_arr, axis=-1)[:, None]
         )
@@ -1314,8 +1322,8 @@ def estimate_factor(n, swing, traction, speed, speed_efficiency, jitter_toleranc
         update = attraction + repulsion + gravities
 
         # compute total swing and traction
-        swing += (mass * np.linalg.norm(pos_arr - update, axis=-1)).sum()
-        traction += (0.5 * mass * np.linalg.norm(pos_arr + update, axis=-1)).sum()
+        swing += (masses * np.linalg.norm(pos_arr - update, axis=-1)).sum()
+        traction += (0.5 * masses * np.linalg.norm(pos_arr + update, axis=-1)).sum()
 
         speed, speed_efficiency = estimate_factor(
             n,
@@ -1328,12 +1336,12 @@ def estimate_factor(n, swing, traction, speed, speed_efficiency, jitter_toleranc
 
         # update pos
         if adjust_sizes:
-            swinging = mass * np.linalg.norm(update, axis=-1)
+            swinging = masses * np.linalg.norm(update, axis=-1)
             factor = 0.1 * speed / (1 + np.sqrt(speed * swinging))
             df = np.linalg.norm(update, axis=-1)
             factor = np.minimum(factor * df, 10.0 * np.ones(df.shape)) / df
         else:
-            swinging = mass * np.linalg.norm(update, axis=-1)
+            swinging = masses * np.linalg.norm(update, axis=-1)
             factor = speed / (1 + np.sqrt(speed * swinging))
 
         pos_arr += update * factor[:, None]