Use matrix multiplication operator (#4390)

Co-authored-by: Ross Barnowski <rossbar@berkeley.edu>

1 files changed, 21 insertions, 26 deletions

diff --git a/networkx/drawing/tests/test_layout.py b/networkx/drawing/tests/test_layout.py
index b94d1b87..9527f363 100644
--- a/networkx/drawing/tests/test_layout.py
+++ b/networkx/drawing/tests/test_layout.py
@@ -4,7 +4,7 @@ from networkx.testing import almost_equal
 
 import pytest
 
-numpy = pytest.importorskip("numpy")
+np = pytest.importorskip("numpy")
 test_smoke_empty_graphscipy = pytest.importorskip("scipy")
 
 
@@ -23,7 +23,7 @@ class TestLayout:
         for k in positions:
             if prev_val is not None:
                 diff = positions[k] - prev_val
-                distances.append(numpy.dot(diff, diff) ** 0.5)
+                distances.append((diff @ diff) ** 0.5)
             prev_val = positions[k]
         return distances
 
@@ -94,10 +94,10 @@ class TestLayout:
         nx.kamada_kawai_layout(G, dim=3)
 
     def check_scale_and_center(self, pos, scale, center):
-        center = numpy.array(center)
+        center = np.array(center)
         low = center - scale
         hi = center + scale
-        vpos = numpy.array(list(pos.values()))
+        vpos = np.array(list(pos.values()))
         length = vpos.max(0) - vpos.min(0)
         assert (length <= 2 * scale).all()
         assert (vpos >= low).all()
@@ -179,9 +179,9 @@ class TestLayout:
         vpos = nx.shell_layout(G, [[0], [1, 2], [3]])
         assert not vpos[0].any()
         assert vpos[3].any()  # ensure node 3 not at origin (#3188)
-        assert numpy.linalg.norm(vpos[3]) <= 1  # ensure node 3 fits (#3753)
+        assert np.linalg.norm(vpos[3]) <= 1  # ensure node 3 fits (#3753)
         vpos = nx.shell_layout(G, [[0], [1, 2], [3]], rotate=0)
-        assert numpy.linalg.norm(vpos[3]) <= 1  # ensure node 3 fits (#3753)
+        assert np.linalg.norm(vpos[3]) <= 1  # ensure node 3 fits (#3753)
 
     def test_smoke_initial_pos_fruchterman_reingold(self):
         pos = nx.circular_layout(self.Gi)
@@ -310,10 +310,10 @@ class TestLayout:
     def test_kamada_kawai_costfn_1d(self):
         costfn = nx.drawing.layout._kamada_kawai_costfn
 
-        pos = numpy.array([4.0, 7.0])
-        invdist = 1 / numpy.array([[0.1, 2.0], [2.0, 0.3]])
+        pos = np.array([4.0, 7.0])
+        invdist = 1 / np.array([[0.1, 2.0], [2.0, 0.3]])
 
-        cost, grad = costfn(pos, numpy, invdist, meanweight=0, dim=1)
+        cost, grad = costfn(pos, np, invdist, meanweight=0, dim=1)
 
         assert almost_equal(cost, ((3 / 2.0 - 1) ** 2))
         assert almost_equal(grad[0], -0.5)
@@ -322,12 +322,12 @@ class TestLayout:
     def check_kamada_kawai_costfn(self, pos, invdist, meanwt, dim):
         costfn = nx.drawing.layout._kamada_kawai_costfn
 
-        cost, grad = costfn(pos.ravel(), numpy, invdist, meanweight=meanwt, dim=dim)
+        cost, grad = costfn(pos.ravel(), np, invdist, meanweight=meanwt, dim=dim)
 
-        expected_cost = 0.5 * meanwt * numpy.sum(numpy.sum(pos, axis=0) ** 2)
+        expected_cost = 0.5 * meanwt * np.sum(np.sum(pos, axis=0) ** 2)
         for i in range(pos.shape[0]):
             for j in range(i + 1, pos.shape[0]):
-                diff = numpy.linalg.norm(pos[i] - pos[j])
+                diff = np.linalg.norm(pos[i] - pos[j])
                 expected_cost += (diff * invdist[i][j] - 1.0) ** 2
 
         assert almost_equal(cost, expected_cost)
@@ -336,31 +336,27 @@ class TestLayout:
         for nd in range(pos.shape[0]):
             for dm in range(pos.shape[1]):
                 idx = nd * pos.shape[1] + dm
-                pos0 = pos.flatten()
+                ps = pos.flatten()
 
-                pos0[idx] += dx
-                cplus = costfn(
-                    pos0, numpy, invdist, meanweight=meanwt, dim=pos.shape[1]
-                )[0]
+                ps[idx] += dx
+                cplus = costfn(ps, np, invdist, meanweight=meanwt, dim=pos.shape[1])[0]
 
-                pos0[idx] -= 2 * dx
-                cminus = costfn(
-                    pos0, numpy, invdist, meanweight=meanwt, dim=pos.shape[1]
-                )[0]
+                ps[idx] -= 2 * dx
+                cminus = costfn(ps, np, invdist, meanweight=meanwt, dim=pos.shape[1])[0]
 
                 assert almost_equal(grad[idx], (cplus - cminus) / (2 * dx), places=5)
 
     def test_kamada_kawai_costfn(self):
-        invdist = 1 / numpy.array([[0.1, 2.1, 1.7], [2.1, 0.2, 0.6], [1.7, 0.6, 0.3]])
+        invdist = 1 / np.array([[0.1, 2.1, 1.7], [2.1, 0.2, 0.6], [1.7, 0.6, 0.3]])
         meanwt = 0.3
 
         # 2d
-        pos = numpy.array([[1.3, -3.2], [2.7, -0.3], [5.1, 2.5]])
+        pos = np.array([[1.3, -3.2], [2.7, -0.3], [5.1, 2.5]])
 
         self.check_kamada_kawai_costfn(pos, invdist, meanwt, 2)
 
         # 3d
-        pos = numpy.array([[0.9, 8.6, -8.7], [-10, -0.5, -7.1], [9.1, -8.1, 1.6]])
+        pos = np.array([[0.9, 8.6, -8.7], [-10, -0.5, -7.1], [9.1, -8.1, 1.6]])
 
         self.check_kamada_kawai_costfn(pos, invdist, meanwt, 3)
 
@@ -395,8 +391,7 @@ class TestLayout:
         G = nx.empty_graph(2)
         vpos = {0: (0.0, 0.0), 1: (1.0, 1.0)}
         s_vpos = nx.rescale_layout_dict(vpos)
-        norm = numpy.linalg.norm
-        assert norm([sum(x) for x in zip(*s_vpos.values())]) < 1e-6
+        assert np.linalg.norm([sum(x) for x in zip(*s_vpos.values())]) < 1e-6
 
         G = nx.empty_graph(3)
         vpos = {0: (0, 0), 1: (1, 1), 2: (0.5, 0.5)}