8000 Merge pull request #15374 from anntzer/cleanup-tri · QuLogic/matplotlib@ec45ec7 · GitHub
[go: up one dir, main page]
Skip to content

Commit ec45ec7

Browse files
timhoffmtimhoffm
authored
Merge pull request matplotlib#15374 from anntzer/cleanup-tri
Replace _prod_vectorized by @-multiplication.
2 parents bd765e0 + 9a476dd commit ec45ec7

File tree

1 file changed

+41
-72
lines changed

1 file changed

+41
-72
lines changed

‎lib/matplotlib/tri/triinterpolate.py

Lines changed: 41 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -505,15 +505,14 @@ def _get_alpha_vec(x, y, tris_pts):
505505
ab = _transpose_vectorized(abT)
506506
OM = np.stack([x, y], axis=1) - tris_pts[:, 0, :]
507507

508-
metric = _prod_vectorized(ab, abT)
508+
metric = ab @ abT
509509
# Here we try to deal with the colinear cases.
510510
# metric_inv is in this case set to the Moore-Penrose pseudo-inverse
511511
# meaning that we will still return a set of valid barycentric
512512
# coordinates.
513513
metric_inv = _pseudo_inv22sym_vectorized(metric)
514-
Covar = _prod_vectorized(ab, _transpose_vectorized(
515-
np.expand_dims(OM, ndim)))
516-
ksi = _prod_vectorized(metric_inv, Covar)
514+
Covar = ab @ _transpose_vectorized(np.expand_dims(OM, ndim))
515+
ksi = metric_inv @ Covar
517516
alpha = _to_matrix_vectorized([
518517
[1-ksi[:, 0, 0]-ksi[:, 1, 0]], [ksi[:, 0, 0]], [ksi[:, 1, 0]]])
519518
return alpha
@@ -567,9 +566,9 @@ def _compute_tri_eccentricities(tris_pts):
567566
c = np.expand_dims(tris_pts[:, 1, :] - tris_pts[:, 0, :], axis=2)
568567
# Do not use np.squeeze, this is dangerous if only one triangle
569568
# in the triangulation...
570-
dot_a = _prod_vectorized(_transpose_vectorized(a), a)[:, 0, 0]
571-
dot_b = _prod_vectorized(_transpose_vectorized(b), b)[:, 0, 0]
572-
dot_c = _prod_vectorized(_transpose_vectorized(c), c)[:, 0, 0]
569+
dot_a = (_transpose_vectorized(a) @ a)[:, 0, 0]
570+
dot_b = (_transpose_vectorized(b) @ b)[:, 0, 0]
571+
dot_c = (_transpose_vectorized(c) @ c)[:, 0, 0]
573572
# Note that this line will raise a warning for dot_a, dot_b or dot_c
574573
# zeros, but we choose not to support triangles with duplicate points.
575574
return _to_matrix_vectorized([[(dot_c-dot_b) / dot_a],
@@ -702,15 +701,12 @@ def get_function_values(self, alpha, ecc, dofs):
702701
V = _to_matrix_vectorized([
703702
[x_sq*x], [y_sq*y], [z_sq*z], [x_sq*z], [x_sq*y], [y_sq*x],
704703
[y_sq*z], [z_sq*y], [z_sq*x], [x*y*z]])
705-
prod = _prod_vectorized(self.M, V)
706-
prod += _scalar_vectorized(E[:, 0, 0],
707-
_prod_vectorized(self.M0, V))
708-
prod += _scalar_vectorized(E[:, 1, 0],
709-
_prod_vectorized(self.M1, V))
710-
prod += _scalar_vectorized(E[:, 2, 0],
711-
_prod_vectorized(self.M2, V))
704+
prod = self.M @ V
705+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ V)
706+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ V)
707+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ V)
712708
s = _roll_vectorized(prod, 3*subtri, axis=0)
713-
return _prod_vectorized(dofs, s)[:, 0, 0]
709+
return (dofs @ s)[:, 0, 0]
714710

715711
def get_function_derivatives(self, alpha, J, ecc, dofs):
716712
"""
@@ -752,23 +748,20 @@ def get_function_derivatives(self, alpha, J, ecc, dofs):
752748
[ -z_sq, 2.*x*z-z_sq],
753749
[ x*z-y*z, x*y-y*z]])
754750
# Puts back dV in first apex basis
755-
dV = _prod_vectorized(dV, _extract_submatrices(
756-
self.rotate_dV, subtri, block_size=2, axis=0))
757-
758-
prod = _prod_vectorized(self.M, dV)
759-
prod += _scalar_vectorized(E[:, 0, 0],
760-
_prod_vectorized(self.M0, dV))
761-
prod += _scalar_vectorized(E[:, 1, 0],
762-
_prod_vectorized(self.M1, dV))
763-
prod += _scalar_vectorized(E[:, 2, 0],
764-
_prod_vectorized(self.M2, dV))
751+
dV = dV @ _extract_submatrices(
752+
self.rotate_dV, subtri, block_size=2, axis=0)
753+
754+
prod = self.M @ dV
755+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ dV)
756+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ dV)
757+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ dV)
765758
dsdksi = _roll_vectorized(prod, 3*subtri, axis=0)
766-
dfdksi = _prod_vectorized(dofs, dsdksi)
759+
dfdksi = dofs @ dsdksi
767760
# In global coordinates:
768761
# Here we try to deal with the simplest colinear cases, returning a
769762
# null matrix.
770763
J_inv = _safe_inv22_vectorized(J)
771-
dfdx = _prod_vectorized(J_inv, _transpose_vectorized(dfdksi))
764+
dfdx = J_inv @ _transpose_vectorized(dfdksi)
772765
return dfdx
773766

774767
def get_function_hessians(self, alpha, J, ecc, dofs):
@@ -791,9 +784,9 @@ def get_function_hessians(self, alpha, J, ecc, dofs):
791784
as a column-matrices of shape (N x 3 x 1).
792785
"""
793786
d2sdksi2 = self.get_d2Sidksij2(alpha, ecc)
794-
d2fdksi2 = _prod_vectorized(dofs, d2sdksi2)
787+
d2fdksi2 = dofs @ d2sdksi2
795788
H_rot = self.get_Hrot_from_J(J)
796-
d2fdx2 = _prod_vectorized(d2fdksi2, H_rot)
789+
d2fdx2 = d2fdksi2 @ H_rot
797790
return _transpose_vectorized(d2fdx2)
798791

799792
def get_d2Sidksij2(self, alpha, ecc):
@@ -828,15 +821,12 @@ def get_d2Sidksij2(self, alpha, ecc):
828821
[ 0., 2.*x-4.*z, -2.*z],
829822
[ -2.*z, -2.*y, x-y-z]])
830823
# Puts back d2V in first apex basis
831-
d2V = _prod_vectorized(d2V, _extract_submatrices(
832-
self.rotate_d2V, subtri, block_size=3, axis=0))
833-
prod = _prod_vectorized(self.M, d2V)
834-
prod += _scalar_vectorized(E[:, 0, 0],
835-
_prod_vectorized(self.M0, d2V))
836-
prod += _scalar_vectorized(E[:, 1, 0],
837-
_prod_vectorized(self.M1, d2V))
838-
prod += _scalar_vectorized(E[:, 2, 0],
839-
_prod_vectorized(self.M2, d2V))
824+
d2V = d2V @ _extract_submatrices(
825+
self.rotate_d2V, subtri, block_size=3, axis=0)
826+
prod = self.M @ d2V
827+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ d2V)
828+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ d2V)
829+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ d2V)
840830
d2sdksi2 = _roll_vectorized(prod, 3*subtri, axis=0)
841831
return d2sdksi2
842832

@@ -859,8 +849,8 @@ def get_bending_matrices(self, J, ecc):
859849
n = np.size(ecc, 0)
860850

861851
# 1) matrix to rotate dofs in global coordinates
862-
J1 = _prod_vectorized(self.J0_to_J1, J)
863-
J2 = _prod_vectorized(self.J0_to_J2, J)
852+
J1 = self.J0_to_J1 @ J
853+
J2 = self.J0_to_J2 @ J
864854
DOF_rot = np.zeros([n, 9, 9], dtype=np.float64)
865855
DOF_rot[:, 0, 0] = 1
866856
DOF_rot[:, 3, 3] = 1
@@ -882,13 +872,11 @@ def get_bending_matrices(self, J, ecc):
882872
alpha = np.expand_dims(alpha, 2)
883873
weight = weights[igauss]
884874
d2Skdksi2 = self.get_d2Sidksij2(alpha, ecc)
885-
d2Skdx2 = _prod_vectorized(d2Skdksi2, H_rot)
886-
K += weight * _prod_vectorized(_prod_vectorized(d2Skdx2, self.E),
887-
_transpose_vectorized(d2Skdx2))
875+
d2Skdx2 = d2Skdksi2 @ H_rot
876+
K += weight * (d2Skdx2 @ self.E @ _transpose_vectorized(d2Skdx2))
888877

889878
# 4) With nodal (not elem) dofs
890-
K = _prod_vectorized(_prod_vectorized(_transpose_vectorized(DOF_rot),
891-
K), DOF_rot)
879+
K = _transpose_vectorized(DOF_rot) @ K @ DOF_rot
892880

893881
# 5) Need the area to compute total element energy
894882
return _scalar_vectorized(area, K)
@@ -960,9 +948,8 @@ def get_Kff_and_Ff(self, J, ecc, triangles, Uc):
960948
c_indices, triangles[:, 2]*2, triangles[:, 2]*2+1]])
961949

962950
expand_indices = np.ones([ntri, 9, 1], dtype=np.int32)
963-
f_row_indices = _prod_vectorized(_transpose_vectorized(f_dof_indices),
964-
_transpose_vectorized(expand_indices))
965-
f_col_indices = _prod_vectorized(expand_indices, f_dof_indices)
951+
f_row_indices = _transpose_vectorized(expand_indices @ f_dof_indices)
952+
f_col_indices = expand_indices @ f_dof_indices
966953
K_elem = self.get_bending_matrices(J, ecc)
967954

968955
# Extracting sub-matrices
@@ -985,7 +972,7 @@ def get_Kff_and_Ff(self, J, ecc, triangles, Uc):
985972
# Computing Ff force vector in sparse coo format
986973
Kfc_elem = K_elem[np.ix_(vec_range, f_dof, c_dof)]
987974
Uc_elem = np.expand_dims(Uc, axis=2)
988-
Ff_elem = - _prod_vectorized(Kfc_elem, Uc_elem)[:, :, 0]
975+
Ff_elem = -(Kfc_elem @ Uc_elem)[:, :, 0]
989976
Ff_indices = f_dof_indices[np.ix_(vec_range, [0], f_dof)][:, 0, :]
990977

991978
# Extracting Ff force vector in dense format
@@ -1059,12 +1046,12 @@ def get_dof_vec(tri_z, tri_dz, J):
10591046
"""
10601047
npt = tri_z.shape[0]
10611048
dof = np.zeros([npt, 9], dtype=np.float64)
1062-
J1 = _prod_vectorized(_ReducedHCT_Element.J0_to_J1, J)
1063-
J2 = _prod_vectorized(_ReducedHCT_Element.J0_to_J2, J)
1049+
J1 = _ReducedHCT_Element.J0_to_J1 @ J
1050+
J2 = _ReducedHCT_Element.J0_to_J2 @ J
10641051

1065-
col0 = _prod_vectorized(J, np.expand_dims(tri_dz[:, 0, :], axis=2))
1066-
col1 = _prod_vectorized(J1, np.expand_dims(tri_dz[:, 1, :], axis=2))
1067-
col2 = _prod_vectorized(J2, np.expand_dims(tri_dz[:, 2, :], axis=2))
1052+
col0 = J @ np.expand_dims(tri_dz[:, 0, :], axis=2)
1053+
col1 = J1 @ np.expand_dims(tri_dz[:, 1, :], axis=2)
1054+
col2 = J2 @ np.expand_dims(tri_dz[:, 2, :], axis=2)
10681055

10691056
dfdksi = _to_matrix_vectorized([
10701057
[col0[:, 0, 0], col1[:, 0, 0], col2[:, 0, 0]],
@@ -1370,7 +1357,6 @@ def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000):
13701357
# The following private functions:
13711358
# :func:`_safe_inv22_vectorized`
13721359
# :func:`_pseudo_inv22sym_vectorized`
1373-
# :func:`_prod_vectorized`
13741360
# :func:`_scalar_vectorized`
13751361
# :func:`_transpose_vectorized`
13761362
# :func:`_roll_vectorized`
@@ -1492,23 +1478,6 @@ def _pseudo_inv22sym_vectorized(M):
14921478
return M_inv
14931479

14941480

1495-
def _prod_vectorized(M1, M2):
1496-
"""
1497-
Matrix product between arrays of matrices, or a matrix and an array of
1498-
matrices (*M1* and *M2*)
1499-
"""
1500-
sh1 = M1.shape
1501-
sh2 = M2.shape
1502-
assert len(sh1) >= 2
1503-
assert len(sh2) >= 2
1504-
assert sh1[-1] == sh2[-2]
1505-
1506-
ndim1 = len(sh1)
1507-
t1_index = [*range(ndim1-2), ndim1-1, ndim1-2]
1508-
return np.sum(np.transpose(M1, t1_index)[..., np.newaxis] *
1509-
M2[..., np.newaxis, :], -3)
1510-
1511-
15121481
def _scalar_vectorized(scalar, M):
15131482
"""
15141483
Scalar product between scalars and matrices.

0 commit comments

Comments
 (0)
0