8000 Replace _prod_vectorized by @-multiplication. by anntzer · Pull Request #15374 · matplotlib/matplotlib · GitHub
[go: up one dir, main page]
Skip to content

Replace _prod_vectorized by @-multiplication. #15374

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 10, 2020
Merged

Replace _prod_vectorized by @-multiplication. #15374

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 41 additions & 72 deletions lib/matplotlib/tri/triinterpolate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,15 +514,14 @@ def _get_alpha_vec(x, y, tris_pts):
ab = _transpose_vectorized(abT)
OM = np.stack([x, y], axis=1) - tris_pts[:, 0, :]

metric = _prod_vectorized(ab, abT)
metric = ab @ abT
# Here we try to deal with the colinear cases.
# metric_inv is in this case set to the Moore-Penrose pseudo-inverse
# meaning that we will still return a set of valid barycentric
# coordinates.
metric_inv = _pseudo_inv22sym_vectorized(metric)
Covar = _prod_vectorized(ab, _transpose_vectorized(
np.expand_dims(OM, ndim)))
ksi = _prod_vectorized(metric_inv, Covar)
Covar = ab @ _transpose_vectorized(np.expand_dims(OM, ndim))
ksi = metric_inv @ Covar
alpha = _to_matrix_vectorized([
[1-ksi[:, 0, 0]-ksi[:, 1, 0]], [ksi[:, 0, 0]], [ksi[:, 1, 0]]])
return alpha
Expand Down Expand Up @@ -576,9 +575,9 @@ def _compute_tri_eccentricities(tris_pts):
c = np.expand_dims(tris_pts[:, 1, :] - tris_pts[:, 0, :], axis=2)
# Do not use np.squeeze, this is dangerous if only one triangle
# in the triangulation...
dot_a = _prod_vectorized(_transpose_vectorized(a), a)[:, 0, 0]
dot_b = _prod_vectorized(_transpose_vectorized(b), b)[:, 0, 0]
dot_c = _prod_vectorized(_transpose_vectorized(c), c)[:, 0, 0]
dot_a = (_transpose_vectorized(a) @ a)[:, 0, 0]
dot_b = (_transpose_vectorized(b) @ b)[:, 0, 0]
dot_c = (_transpose_vectorized(c) @ c)[:, 0, 0]
# Note that this line will raise a warning for dot_a, dot_b or dot_c
# zeros, but we choose not to support triangles with duplicate points.
return _to_matrix_vectorized([[(dot_c-dot_b) / dot_a],
Expand Down Expand Up @@ -711,15 +710,12 @@ def get_function_values(self, alpha, ecc, dofs):
V = _to_matrix_vectorized([
[x_sq*x], [y_sq*y], [z_sq*z], [x_sq*z], [x_sq*y], [y_sq*x],
[y_sq*z], [z_sq*y], [z_sq*x], [x*y*z]])
prod = _prod_vectorized(self.M, V)
prod += _scalar_vectorized(E[:, 0, 0],
_prod_vectorized(self.M0, V))
prod += _scalar_vectorized(E[:, 1, 0],
_prod_vectorized(self.M1, V))
prod += _scalar_vectorized(E[:, 2, 0],
_prod_vectorized(self.M2, V))
prod = self.M @ V
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ V)
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ V)
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ V)
s = _roll_vectorized(prod, 3*subtri, axis=0)
return _prod_vectorized(dofs, s)[:, 0, 0]
return (dofs @ s)[:, 0, 0]

def get_function_derivatives(self, alpha, J, ecc, dofs):
"""
Expand Down Expand Up @@ -761,23 +757,20 @@ def get_function_derivatives(self, alpha, J, ecc, dofs):
[ -z_sq, 2.*x*z-z_sq],
[ x*z-y*z, x*y-y*z]])
# Puts back dV in first apex basis
dV = _prod_vectorized(dV, _extract_submatrices(
self.rotate_dV, subtri, block_size=2, axis=0))

prod = _prod_vectorized(self.M, dV)
prod += _scalar_vectorized(E[:, 0, 0],
_prod_vectorized(self.M0, dV))
prod += _scalar_vectorized(E[:, 1, 0],
_prod_vectorized(self.M1, dV))
prod += _scalar_vectorized(E[:, 2, 0],
_prod_vectorized(self.M2, dV))
dV = dV @ _extract_submatrices(
self.rotate_dV, subtri, block_size=2, axis=0)

prod = self.M @ dV
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ dV)
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ dV)
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ dV)
dsdksi = _roll_vectorized(prod, 3*subtri, axis=0)
dfdksi = _prod_vectorized(dofs, dsdksi)
dfdksi = dofs @ dsdksi
# In global coordinates:
# Here we try to deal with the simplest colinear cases, returning a
# null matrix.
J_inv = _safe_inv22_vectorized(J)
dfdx = _prod_vectorized(J_inv, _transpose_vectorized(dfdksi))
dfdx = J_inv @ _transpose_vectorized(dfdksi)
return dfdx

def get_function_hessians(self, alpha, J, ecc, dofs):
Expand All @@ -800,9 +793,9 @@ def get_function_hessians(self, alpha, J, ecc, dofs):
as a column-matrices of shape (N x 3 x 1).
"""
d2sdksi2 = self.get_d2Sidksij2(alpha, ecc)
d2fdksi2 = _prod_vectorized(dofs, d2sdksi2)
d2fdksi2 = dofs @ d2sdksi2
H_rot = self.get_Hrot_from_J(J)
d2fdx2 = _prod_vectorized(d2fdksi2, H_rot)
d2fdx2 = d2fdksi2 @ H_rot
return _transpose_vectorized(d2fdx2)

def get_d2Sidksij2(self, alpha, ecc):
Expand Down Expand Up @@ -837,15 +830,12 @@ def get_d2Sidksij2(self, alpha, ecc):
[ 0., 2.*x-4.*z, -2.*z],
[ -2.*z, -2.*y, x-y-z]])
# Puts back d2V in first apex basis
d2V = _prod_vectorized(d2V, _extract_submatrices(
self.rotate_d2V, subtri, block_size=3, axis=0))
prod = _prod_vectorized(self.M, d2V)
prod += _scalar_vectorized(E[:, 0, 0],
_prod_vectorized(self.M0, d2V))
prod += _scalar_vectorized(E[:, 1, 0],
_prod_vectorized(self.M1, d2V))
prod += _scalar_vectorized(E[:, 2, 0],
_prod_vectorized(self.M2, d2V))
d2V = d2V @ _extract_submatrices(
self.rotate_d2V, subtri, block_size=3, axis=0)
prod = self.M @ d2V
prod += _scalar_vectorized(E[:, 0, 0], self.M0 @ d2V)
prod += _scalar_vectorized(E[:, 1, 0], self.M1 @ d2V)
prod += _scalar_vectorized(E[:, 2, 0], self.M2 @ d2V)
d2sdksi2 = _roll_vectorized(prod, 3*subtri, axis=0)
return d2sdksi2

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

# 1) matrix to rotate dofs in global coordinates
J1 = _prod_vectorized(self.J0_to_J1, J)
J2 = _prod_vectorized(self.J0_to_J2, J)
J1 = self.J0_to_J1 @ J
J2 = self.J0_to_J2 @ J
DOF_rot = np.zeros([n, 9, 9], dtype=np.float64)
DOF_rot[:, 0, 0] = 1
DOF_rot[:, 3, 3] = 1
Expand All @@ -891,13 +881,11 @@ def get_bending_matrices(self, J, ecc):
alpha = np.expand_dims(alpha, 2)
weight = weights[igauss]
d2Skdksi2 = self.get_d2Sidksij2(alpha, ecc)
d2Skdx2 = _prod_vectorized(d2Skdksi2, H_rot)
K += weight * _prod_vectorized(_prod_vectorized(d2Skdx2, self.E),
_transpose_vectorized(d2Skdx2))
d2Skdx2 = d2Skdksi2 @ H_rot
K += weight * (d2Skdx2 @ self.E @ _transpose_vectorized(d2Skdx2))

# 4) With nodal (not elem) dofs
K = _prod_vectorized(_prod_vectorized(_transpose_vectorized(DOF_rot),
K), DOF_rot)
K = _transpose_vectorized(DOF_rot) @ K @ DOF_rot

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

expand_indices = np.ones([ntri, 9, 1], dtype=np.int32)
f_row_indices = _prod_vectorized(_transpose_vectorized(f_dof_indices),
_transpose_vectorized(expand_indices))
f_col_indices = _prod_vectorized(expand_indices, f_dof_indices)
f_row_indices = _transpose_vectorized(expand_indices @ f_dof_indices)
f_col_indices = expand_indices @ f_dof_indices
K_elem = self.get_bending_matrices(J, ecc)

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

# Extracting Ff force vector in dense format
Expand Down Expand Up @@ -1058,12 +1045,12 @@ def get_dof_vec(tri_z, tri_dz, J):
"""
npt = tri_z.shape[0]
dof = np.zeros([npt, 9], dtype=np.float64)
J1 = _prod_vectorized(_ReducedHCT_Element.J0_to_J1, J)
J2 = _prod_vectorized(_ReducedHCT_Element.J0_to_J2, J)
J1 = _ReducedHCT_Element.J0_to_J1 @ J
J2 = _ReducedHCT_Element.J0_to_J2 @ J

col0 = _prod_vectorized(J, np.expand_dims(tri_dz[:, 0, :], axis=2))
col1 = _prod_vectorized(J1, np.expand_dims(tri_dz[:, 1, :], axis=2))
col2 = _prod_vectorized(J2, np.expand_dims(tri_dz[:, 2, :], axis=2))
col0 = J @ np.expand_dims(tri_dz[:, 0, :], axis=2)
col1 = J1 @ np.expand_dims(tri_dz[:, 1, :], axis=2)
col2 = J2 @ np.expand_dims(tri_dz[:, 2, :], axis=2)

dfdksi = _to_matrix_vectorized([
[col0[:, 0, 0], col1[:, 0, 0], col2[:, 0, 0]],
Expand Down Expand Up @@ -1377,7 +1364,6 @@ def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000):
# The following private functions:
# :func:`_safe_inv22_vectorized`
# :func:`_pseudo_inv22sym_vectorized`
# :func:`_prod_vectorized`
# :func:`_scalar_vectorized`
# :func:`_transpose_vectorized`
# :func:`_roll_vectorized`
Expand Down Expand Up @@ -1499,23 +1485,6 @@ def _pseudo_inv22sym_vectorized(M):
return M_inv


def _prod_vectorized(M1, M2):
"""
Matrix product between arrays of matrices, or a matrix and an array of
matrices (*M1* and *M2*)
"""
sh1 = M1.shape
sh2 = M2.shape
assert len(sh1) >= 2
assert len(sh2) >= 2
assert sh1[-1] == sh2[-2]

ndim1 = len(sh1)
t1_index = [*range(ndim1-2), ndim1-1, ndim1-2]
return np.sum(np.transpose(M1, t1_index)[..., np.newaxis] *
M2[..., np.newaxis, :], -3)


def _scalar_vectorized(scalar, M):
"""
Scalar product between scalars and matrices.
Expand Down
0