Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 9a476dd

Browse files
anntzeranntzer
committed
Replace _prod_vectorized by @-multiplication.
@-multiplication has the exact same broadcasting behavior as _prod_vectorized (`...ij,...jk->...ik` in einsum notation).
1 parent 327df3b commit 9a476dd

1 file changed

Lines changed: 41 additions & 72 deletions

File tree

‎lib/matplotlib/tri/triinterpolate.py‎

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

517-
metric = _prod_vectorized(ab, abT)
517+
metric = ab @ abT
518518
# Here we try to deal with the colinear cases.
519519
# metric_inv is in this case set to the Moore-Penrose pseudo-inverse
520520
# meaning that we will still return a set of valid barycentric
521521
# coordinates.
522522
metric_inv = _pseudo_inv22sym_vectorized(metric)
523-
Covar = _prod_vectorized(ab, _transpose_vectorized(
524-
np.expand_dims(OM, ndim)))
525-
ksi = _prod_vectorized(metric_inv, Covar)
523+
Covar = ab @ _transpose_vectorized(np.expand_dims(OM, ndim))
524+
ksi = metric_inv @ Covar
526525
alpha = _to_matrix_vectorized([
527526
[1-ksi[:, 0, 0]-ksi[:, 1, 0]], [ksi[:, 0, 0]], [ksi[:, 1, 0]]])
528527
return alpha
@@ -576,9 +575,9 @@ def _compute_tri_eccentricities(tris_pts):
576575
c = np.expand_dims(tris_pts[:, 1, :] - tris_pts[:, 0, :], axis=2)
577576
# Do not use np.squeeze, this is dangerous if only one triangle
578577
# in the triangulation...
579-
dot_a = _prod_vectorized(_transpose_vectorized(a), a)[:, 0, 0]
580-
dot_b = _prod_vectorized(_transpose_vectorized(b), b)[:, 0, 0]
581-
dot_c = _prod_vectorized(_transpose_vectorized(c), c)[:, 0, 0]
578+
dot_a = (_transpose_vectorized(a) @ a)[:, 0, 0]
579+
dot_b = (_transpose_vectorized(b) @ b)[:, 0, 0]
580+
dot_c = (_transpose_vectorized(c) @ c)[:, 0, 0]
582581
# Note that this line will raise a warning for dot_a, dot_b or dot_c
583582
# zeros, but we choose not to support triangles with duplicate points.
584583
return _to_matrix_vectorized([[(dot_c-dot_b) / dot_a],
@@ -711,15 +710,12 @@ def get_function_values(self, alpha, ecc, dofs):
711710
V = _to_matrix_vectorized([
712711
[x_sq*x], [y_sq*y], [z_sq*z], [x_sq*z], [x_sq*y], [y_sq*x],
713712
[y_sq*z], [z_sq*y], [z_sq*x], [x*y*z]])
714-
prod = _prod_vectorized(self.M, V)
715-
prod += _scalar_vectorized(E[:, 0, 0],
716-
_prod_vectorized(self.M0, V))
717-
prod += _scalar_vectorized(E[:, 1, 0],
718-
_prod_vectorized(self.M1, V))
719-
prod += _scalar_vectorized(E[:, 2, 0],
720-
_prod_vectorized(self.M2, V))
713+
prod = self.M @ V
714+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ V)
715+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ V)
716+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ V)
721717
s = _roll_vectorized(prod, 3*subtri, axis=0)
722-
return _prod_vectorized(dofs, s)[:, 0, 0]
718+
return (dofs @ s)[:, 0, 0]
723719

724720
def get_function_derivatives(self, alpha, J, ecc, dofs):
725721
"""
@@ -761,23 +757,20 @@ def get_function_derivatives(self, alpha, J, ecc, dofs):
761757
[ -z_sq, 2.*x*z-z_sq],
762758
[ x*z-y*z, x*y-y*z]])
763759
# Puts back dV in first apex basis
764-
dV = _prod_vectorized(dV, _extract_submatrices(
765-
self.rotate_dV, subtri, block_size=2, axis=0))
766-
767-
prod = _prod_vectorized(self.M, dV)
768-
prod += _scalar_vectorized(E[:, 0, 0],
769-
_prod_vectorized(self.M0, dV))
770-
prod += _scalar_vectorized(E[:, 1, 0],
771-
_prod_vectorized(self.M1, dV))
772-
prod += _scalar_vectorized(E[:, 2, 0],
773-
_prod_vectorized(self.M2, dV))
760+
dV = dV @ _extract_submatrices(
761+
self.rotate_dV, subtri, block_size=2, axis=0)
762+
763+
prod = self.M @ dV
764+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ dV)
765+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ dV)
766+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ dV)
774767
dsdksi = _roll_vectorized(prod, 3*subtri, axis=0)
775-
dfdksi = _prod_vectorized(dofs, dsdksi)
768+
dfdksi = dofs @ dsdksi
776769
# In global coordinates:
777770
# Here we try to deal with the simplest colinear cases, returning a
778771
# null matrix.
779772
J_inv = _safe_inv22_vectorized(J)
780-
dfdx = _prod_vectorized(J_inv, _transpose_vectorized(dfdksi))
773+
dfdx = J_inv @ _transpose_vectorized(dfdksi)
781774
return dfdx
782775

783776
def get_function_hessians(self, alpha, J, ecc, dofs):
@@ -800,9 +793,9 @@ def get_function_hessians(self, alpha, J, ecc, dofs):
800793
as a column-matrices of shape (N x 3 x 1).
801794
"""
802795
d2sdksi2 = self.get_d2Sidksij2(alpha, ecc)
803-
d2fdksi2 = _prod_vectorized(dofs, d2sdksi2)
796+
d2fdksi2 = dofs @ d2sdksi2
804797
H_rot = self.get_Hrot_from_J(J)
805-
d2fdx2 = _prod_vectorized(d2fdksi2, H_rot)
798+
d2fdx2 = d2fdksi2 @ H_rot
806799
return _transpose_vectorized(d2fdx2)
807800

808801
def get_d2Sidksij2(self, alpha, ecc):
@@ -837,15 +830,12 @@ def get_d2Sidksij2(self, alpha, ecc):
837830
[ 0., 2.*x-4.*z, -2.*z],
838831
[ -2.*z, -2.*y, x-y-z]])
839832
# Puts back d2V in first apex basis
840-
d2V = _prod_vectorized(d2V, _extract_submatrices(
841-
self.rotate_d2V, subtri, block_size=3, axis=0))
842-
prod = _prod_vectorized(self.M, d2V)
843-
prod += _scalar_vectorized(E[:, 0, 0],
844-
_prod_vectorized(self.M0, d2V))
845-
prod += _scalar_vectorized(E[:, 1, 0],
846-
_prod_vectorized(self.M1, d2V))
847-
prod += _scalar_vectorized(E[:, 2, 0],
848-
_prod_vectorized(self.M2, d2V))
833+
d2V = d2V @ _extract_submatrices(
834+
self.rotate_d2V, subtri, block_size=3, axis=0)
835+
prod = self.M @ d2V
836+
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ d2V)
837+
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ d2V)
838+
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ d2V)
849839
d2sdksi2 = _roll_vectorized(prod, 3*subtri, axis=0)
850840
return d2sdksi2
851841

@@ -868,8 +858,8 @@ def get_bending_matrices(self, J, ecc):
868858
n = np.size(ecc, 0)
869859

870860
# 1) matrix to rotate dofs in global coordinates
871-
J1 = _prod_vectorized(self.J0_to_J1, J)
872-
J2 = _prod_vectorized(self.J0_to_J2, J)
861+
J1 = self.J0_to_J1 @ J
862+
J2 = self.J0_to_J2 @ J
873863
DOF_rot = np.zeros([n, 9, 9], dtype=np.float64)
874864
DOF_rot[:, 0, 0] = 1
875865
DOF_rot[:, 3, 3] = 1
@@ -891,13 +881,11 @@ def get_bending_matrices(self, J, ecc):
891881
alpha = np.expand_dims(alpha, 2)
892882
weight = weights[igauss]
893883
d2Skdksi2 = self.get_d2Sidksij2(alpha, ecc)
894-
d2Skdx2 = _prod_vectorized(d2Skdksi2, H_rot)
895-
K += weight * _prod_vectorized(_prod_vectorized(d2Skdx2, self.E),
896-
_transpose_vectorized(d2Skdx2))
884+
d2Skdx2 = d2Skdksi2 @ H_rot
885+
K += weight * (d2Skdx2 @ self.E @ _transpose_vectorized(d2Skdx2))
897886

898887
# 4) With nodal (not elem) dofs
899-
K = _prod_vectorized(_prod_vectorized(_transpose_vectorized(DOF_rot),
900-
K), DOF_rot)
888+
K = _transpose_vectorized(DOF_rot) @ K @ DOF_rot
901889

902890
# 5) Need the area to compute total element energy
903891
return _scalar_vectorized(area, K)
@@ -968,9 +956,8 @@ def get_Kff_and_Ff(self, J, ecc, triangles, Uc):
968956
c_indices, triangles[:, 2]*2, triangles[:, 2]*2+1]])
969957

970958
expand_indices = np.ones([ntri, 9, 1], dtype=np.int32)
971-
f_row_indices = _prod_vectorized(_transpose_vectorized(f_dof_indices),
972-
_transpose_vectorized(expand_indices))
973-
f_col_indices = _prod_vectorized(expand_indices, f_dof_indices)
959+
f_row_indices = _transpose_vectorized(expand_indices @ f_dof_indices)
960+
f_col_indices = expand_indices @ f_dof_indices
974961
K_elem = self.get_bending_matrices(J, ecc)
975962

976963
# Extracting sub-matrices
@@ -993,7 +980,7 @@ def get_Kff_and_Ff(self, J, ecc, triangles, Uc):
993980
# Computing Ff force vector in sparse coo format
994981
Kfc_elem = K_elem[np.ix_(vec_range, f_dof, c_dof)]
995982
Uc_elem = np.expand_dims(Uc, axis=2)
996-
Ff_elem = - _prod_vectorized(Kfc_elem, Uc_elem)[:, :, 0]
983+
Ff_elem = -(Kfc_elem @ Uc_elem)[:, :, 0]
997984
Ff_indices = f_dof_indices[np.ix_(vec_range, [0], f_dof)][:, 0, :]
998985

999986
# Extracting Ff force vector in dense format
@@ -1058,12 +1045,12 @@ def get_dof_vec(tri_z, tri_dz, J):
10581045
"""
10591046
npt = tri_z.shape[0]
10601047
dof = np.zeros([npt, 9], dtype=np.float64)
1061-
J1 = _prod_vectorized(_ReducedHCT_Element.J0_to_J1, J)
1062-
J2 = _prod_vectorized(_ReducedHCT_Element.J0_to_J2, J)
1048+
J1 = _ReducedHCT_Element.J0_to_J1 @ J
1049+
J2 = _ReducedHCT_Element.J0_to_J2 @ J
10631050

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

10681055
dfdksi = _to_matrix_vectorized([
10691056
[col0[:, 0, 0], col1[:, 0, 0], col2[:, 0, 0]],
@@ -1377,7 +1364,6 @@ def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000):
13771364
# The following private functions:
13781365
# :func:`_safe_inv22_vectorized`
13791366
# :func:`_pseudo_inv22sym_vectorized`
1380-
# :func:`_prod_vectorized`
13811367
# :func:`_scalar_vectorized`
13821368
# :func:`_transpose_vectorized`
13831369
# :func:`_roll_vectorized`
@@ -1499,23 +1485,6 @@ def _pseudo_inv22sym_vectorized(M):
14991485
return M_inv
15001486

15011487

1502-
def _prod_vectorized(M1, M2):
1503-
"""
1504-
Matrix product between arrays of matrices, or a matrix and an array of
1505-
matrices (*M1* and *M2*)
1506-
"""
1507-
sh1 = M1.shape
1508-
sh2 = M2.shape
1509-
assert len(sh1) >= 2
1510-
assert len(sh2) >= 2
1511-
assert sh1[-1] == sh2[-2]
1512-
1513-
ndim1 = len(sh1)
1514-
t1_index = [*range(ndim1-2), ndim1-1, ndim1-2]
1515-
return np.sum(np.transpose(M1, t1_index)[..., np.newaxis] *
1516-
M2[..., np.newaxis, :], -3)
1517-
1518-
15191488
def _scalar_vectorized(scalar, M):
15201489
"""
15211490
Scalar product between scalars and matrices.

0 commit comments

Comments
 (0)