From 41f9d6c9cff72e49eef0399f108568a7327b5ead Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Fri, 5 Apr 2024 07:56:48 +0200 Subject: [PATCH] Added mass and size input attributes with docstrings --- networkx/drawing/layout.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/networkx/drawing/layout.py b/networkx/drawing/layout.py index ff4ea874c9c3..bc5056e8e789 100644 --- a/networkx/drawing/layout.py +++ b/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]