1""" 2Interpolation inside triangular grids. 3""" 4from __future__ import (absolute_import, division, print_function, 5 unicode_literals) 6 7import six 8from six.moves import xrange 9 10from matplotlib.tri import Triangulation 11from matplotlib.tri.trifinder import TriFinder 12from matplotlib.tri.tritools import TriAnalyzer 13import numpy as np 14import warnings 15 16__all__ = ('TriInterpolator', 'LinearTriInterpolator', 'CubicTriInterpolator') 17 18 19class TriInterpolator(object): 20 """ 21 Abstract base class for classes used to perform interpolation on 22 triangular grids. 23 24 Derived classes implement the following methods: 25 26 - ``__call__(x, y)`` , 27 where x, y are array_like point coordinates of the same shape, and 28 that returns a masked array of the same shape containing the 29 interpolated z-values. 30 31 - ``gradient(x, y)`` , 32 where x, y are array_like point coordinates of the same 33 shape, and that returns a list of 2 masked arrays of the same shape 34 containing the 2 derivatives of the interpolator (derivatives of 35 interpolated z values with respect to x and y). 36 37 """ 38 def __init__(self, triangulation, z, trifinder=None): 39 if not isinstance(triangulation, Triangulation): 40 raise ValueError("Expected a Triangulation object") 41 self._triangulation = triangulation 42 43 self._z = np.asarray(z) 44 if self._z.shape != self._triangulation.x.shape: 45 raise ValueError("z array must have same length as triangulation x" 46 " and y arrays") 47 48 if trifinder is not None and not isinstance(trifinder, TriFinder): 49 raise ValueError("Expected a TriFinder object") 50 self._trifinder = trifinder or self._triangulation.get_trifinder() 51 52 # Default scaling factors : 1.0 (= no scaling) 53 # Scaling may be used for interpolations for which the order of 54 # magnitude of x, y has an impact on the interpolant definition. 55 # Please refer to :meth:`_interpolate_multikeys` for details. 56 self._unit_x = 1.0 57 self._unit_y = 1.0 58 59 # Default triangle renumbering: None (= no renumbering) 60 # Renumbering may be used to avoid unnecessary computations 61 # if complex calculations are done inside the Interpolator. 62 # Please refer to :meth:`_interpolate_multikeys` for details. 63 self._tri_renum = None 64 65 # __call__ and gradient docstrings are shared by all subclasses 66 # (except, if needed, relevant additions). 67 # However these methods are only implemented in subclasses to avoid 68 # confusion in the documentation. 69 _docstring__call__ = """ 70 Returns a masked array containing interpolated values at the specified 71 x,y points. 72 73 Parameters 74 ---------- 75 x, y : array-like 76 x and y coordinates of the same shape and any number of 77 dimensions. 78 79 Returns 80 ------- 81 z : np.ma.array 82 Masked array of the same shape as *x* and *y* ; values 83 corresponding to (*x*, *y*) points outside of the triangulation 84 are masked out. 85 86 """ 87 88 _docstringgradient = """ 89 Returns a list of 2 masked arrays containing interpolated derivatives 90 at the specified x,y points. 91 92 Parameters 93 ---------- 94 x, y : array-like 95 x and y coordinates of the same shape and any number of 96 dimensions. 97 98 Returns 99 ------- 100 dzdx, dzdy : np.ma.array 101 2 masked arrays of the same shape as *x* and *y* ; values 102 corresponding to (x,y) points outside of the triangulation 103 are masked out. 104 The first returned array contains the values of 105 :math:`\\frac{\\partial z}{\\partial x}` and the second those of 106 :math:`\\frac{\\partial z}{\\partial y}`. 107 108 """ 109 110 def _interpolate_multikeys(self, x, y, tri_index=None, 111 return_keys=('z',)): 112 """ 113 Versatile (private) method defined for all TriInterpolators. 114 115 :meth:`_interpolate_multikeys` is a wrapper around method 116 :meth:`_interpolate_single_key` (to be defined in the child 117 subclasses). 118 :meth:`_interpolate_single_key actually performs the interpolation, 119 but only for 1-dimensional inputs and at valid locations (inside 120 unmasked triangles of the triangulation). 121 122 The purpose of :meth:`_interpolate_multikeys` is to implement the 123 following common tasks needed in all subclasses implementations: 124 125 - calculation of containing triangles 126 - dealing with more than one interpolation request at the same 127 location (e.g., if the 2 derivatives are requested, it is 128 unnecessary to compute the containing triangles twice) 129 - scaling according to self._unit_x, self._unit_y 130 - dealing with points outside of the grid (with fill value np.nan) 131 - dealing with multi-dimensionnal *x*, *y* arrays: flattening for 132 :meth:`_interpolate_params` call and final reshaping. 133 134 (Note that np.vectorize could do most of those things very well for 135 you, but it does it by function evaluations over successive tuples of 136 the input arrays. Therefore, this tends to be more time consuming than 137 using optimized numpy functions - e.g., np.dot - which can be used 138 easily on the flattened inputs, in the child-subclass methods 139 :meth:`_interpolate_single_key`.) 140 141 It is guaranteed that the calls to :meth:`_interpolate_single_key` 142 will be done with flattened (1-d) array_like input parameters `x`, `y` 143 and with flattened, valid `tri_index` arrays (no -1 index allowed). 144 145 Parameters 146 ---------- 147 x, y : array_like 148 x and y coordinates indicating where interpolated values are 149 requested. 150 tri_index : integer array_like, optional 151 Array of the containing triangle indices, same shape as 152 *x* and *y*. Defaults to None. If None, these indices 153 will be computed by a TriFinder instance. 154 (Note: For point outside the grid, tri_index[ipt] shall be -1). 155 return_keys : tuple of keys from {'z', 'dzdx', 'dzdy'} 156 Defines the interpolation arrays to return, and in which order. 157 158 Returns 159 ------- 160 ret : list of arrays 161 Each array-like contains the expected interpolated values in the 162 order defined by *return_keys* parameter. 163 """ 164 # Flattening and rescaling inputs arrays x, y 165 # (initial shape is stored for output) 166 x = np.asarray(x, dtype=np.float64) 167 y = np.asarray(y, dtype=np.float64) 168 sh_ret = x.shape 169 if x.shape != y.shape: 170 raise ValueError("x and y shall have same shapes." 171 " Given: {0} and {1}".format(x.shape, y.shape)) 172 x = np.ravel(x) 173 y = np.ravel(y) 174 x_scaled = x/self._unit_x 175 y_scaled = y/self._unit_y 176 size_ret = np.size(x_scaled) 177 178 # Computes & ravels the element indexes, extract the valid ones. 179 if tri_index is None: 180 tri_index = self._trifinder(x, y) 181 else: 182 if (tri_index.shape != sh_ret): 183 raise ValueError( 184 "tri_index array is provided and shall" 185 " have same shape as x and y. Given: " 186 "{0} and {1}".format(tri_index.shape, sh_ret)) 187 tri_index = np.ravel(tri_index) 188 189 mask_in = (tri_index != -1) 190 if self._tri_renum is None: 191 valid_tri_index = tri_index[mask_in] 192 else: 193 valid_tri_index = self._tri_renum[tri_index[mask_in]] 194 valid_x = x_scaled[mask_in] 195 valid_y = y_scaled[mask_in] 196 197 ret = [] 198 for return_key in return_keys: 199 # Find the return index associated with the key. 200 try: 201 return_index = {'z': 0, 'dzdx': 1, 'dzdy': 2}[return_key] 202 except KeyError: 203 raise ValueError("return_keys items shall take values in" 204 " {'z', 'dzdx', 'dzdy'}") 205 206 # Sets the scale factor for f & df components 207 scale = [1., 1./self._unit_x, 1./self._unit_y][return_index] 208 209 # Computes the interpolation 210 ret_loc = np.empty(size_ret, dtype=np.float64) 211 ret_loc[~mask_in] = np.nan 212 ret_loc[mask_in] = self._interpolate_single_key( 213 return_key, valid_tri_index, valid_x, valid_y) * scale 214 ret += [np.ma.masked_invalid(ret_loc.reshape(sh_ret), copy=False)] 215 216 return ret 217 218 def _interpolate_single_key(self, return_key, tri_index, x, y): 219 """ 220 Performs the interpolation at points belonging to the triangulation 221 (inside an unmasked triangles). 222 223 Parameters 224 ---------- 225 return_index : string key from {'z', 'dzdx', 'dzdy'} 226 Identifies the requested values (z or its derivatives) 227 tri_index : 1d integer array 228 Valid triangle index (-1 prohibited) 229 x, y : 1d arrays, same shape as `tri_index` 230 Valid locations where interpolation is requested. 231 232 Returns 233 ------- 234 ret : 1-d array 235 Returned array of the same size as *tri_index* 236 """ 237 raise NotImplementedError("TriInterpolator subclasses" + 238 "should implement _interpolate_single_key!") 239 240 241class LinearTriInterpolator(TriInterpolator): 242 """ 243 A LinearTriInterpolator performs linear interpolation on a triangular grid. 244 245 Each triangle is represented by a plane so that an interpolated value at 246 point (x,y) lies on the plane of the triangle containing (x,y). 247 Interpolated values are therefore continuous across the triangulation, but 248 their first derivatives are discontinuous at edges between triangles. 249 250 Parameters 251 ---------- 252 triangulation : :class:`~matplotlib.tri.Triangulation` object 253 The triangulation to interpolate over. 254 z : array_like of shape (npoints,) 255 Array of values, defined at grid points, to interpolate between. 256 trifinder : :class:`~matplotlib.tri.TriFinder` object, optional 257 If this is not specified, the Triangulation's default TriFinder will 258 be used by calling 259 :func:`matplotlib.tri.Triangulation.get_trifinder`. 260 261 Methods 262 ------- 263 `__call__` (x, y) : Returns interpolated values at x,y points 264 `gradient` (x, y) : Returns interpolated derivatives at x,y points 265 266 """ 267 def __init__(self, triangulation, z, trifinder=None): 268 TriInterpolator.__init__(self, triangulation, z, trifinder) 269 270 # Store plane coefficients for fast interpolation calculations. 271 self._plane_coefficients = \ 272 self._triangulation.calculate_plane_coefficients(self._z) 273 274 def __call__(self, x, y): 275 return self._interpolate_multikeys(x, y, tri_index=None, 276 return_keys=('z',))[0] 277 __call__.__doc__ = TriInterpolator._docstring__call__ 278 279 def gradient(self, x, y): 280 return self._interpolate_multikeys(x, y, tri_index=None, 281 return_keys=('dzdx', 'dzdy')) 282 gradient.__doc__ = TriInterpolator._docstringgradient 283 284 def _interpolate_single_key(self, return_key, tri_index, x, y): 285 if return_key == 'z': 286 return (self._plane_coefficients[tri_index, 0]*x + 287 self._plane_coefficients[tri_index, 1]*y + 288 self._plane_coefficients[tri_index, 2]) 289 elif return_key == 'dzdx': 290 return self._plane_coefficients[tri_index, 0] 291 elif return_key == 'dzdy': 292 return self._plane_coefficients[tri_index, 1] 293 else: 294 raise ValueError("Invalid return_key: " + return_key) 295 296 297class CubicTriInterpolator(TriInterpolator): 298 """ 299 A CubicTriInterpolator performs cubic interpolation on triangular grids. 300 301 In one-dimension - on a segment - a cubic interpolating function is 302 defined by the values of the function and its derivative at both ends. 303 This is almost the same in 2-d inside a triangle, except that the values 304 of the function and its 2 derivatives have to be defined at each triangle 305 node. 306 307 The CubicTriInterpolator takes the value of the function at each node - 308 provided by the user - and internally computes the value of the 309 derivatives, resulting in a smooth interpolation. 310 (As a special feature, the user can also impose the value of the 311 derivatives at each node, but this is not supposed to be the common 312 usage.) 313 314 Parameters 315 ---------- 316 triangulation : :class:`~matplotlib.tri.Triangulation` object 317 The triangulation to interpolate over. 318 z : array_like of shape (npoints,) 319 Array of values, defined at grid points, to interpolate between. 320 kind : {'min_E', 'geom', 'user'}, optional 321 Choice of the smoothing algorithm, in order to compute 322 the interpolant derivatives (defaults to 'min_E'): 323 324 - if 'min_E': (default) The derivatives at each node is computed 325 to minimize a bending energy. 326 - if 'geom': The derivatives at each node is computed as a 327 weighted average of relevant triangle normals. To be used for 328 speed optimization (large grids). 329 - if 'user': The user provides the argument `dz`, no computation 330 is hence needed. 331 332 trifinder : :class:`~matplotlib.tri.TriFinder` object, optional 333 If not specified, the Triangulation's default TriFinder will 334 be used by calling 335 :func:`matplotlib.tri.Triangulation.get_trifinder`. 336 dz : tuple of array_likes (dzdx, dzdy), optional 337 Used only if *kind* ='user'. In this case *dz* must be provided as 338 (dzdx, dzdy) where dzdx, dzdy are arrays of the same shape as *z* and 339 are the interpolant first derivatives at the *triangulation* points. 340 341 Methods 342 ------- 343 `__call__` (x, y) : Returns interpolated values at x,y points 344 `gradient` (x, y) : Returns interpolated derivatives at x,y points 345 346 Notes 347 ----- 348 This note is a bit technical and details the way a 349 :class:`~matplotlib.tri.CubicTriInterpolator` computes a cubic 350 interpolation. 351 352 The interpolation is based on a Clough-Tocher subdivision scheme of 353 the *triangulation* mesh (to make it clearer, each triangle of the 354 grid will be divided in 3 child-triangles, and on each child triangle 355 the interpolated function is a cubic polynomial of the 2 coordinates). 356 This technique originates from FEM (Finite Element Method) analysis; 357 the element used is a reduced Hsieh-Clough-Tocher (HCT) 358 element. Its shape functions are described in [1]_. 359 The assembled function is guaranteed to be C1-smooth, i.e. it is 360 continuous and its first derivatives are also continuous (this 361 is easy to show inside the triangles but is also true when crossing the 362 edges). 363 364 In the default case (*kind* ='min_E'), the interpolant minimizes a 365 curvature energy on the functional space generated by the HCT element 366 shape functions - with imposed values but arbitrary derivatives at each 367 node. The minimized functional is the integral of the so-called total 368 curvature (implementation based on an algorithm from [2]_ - PCG sparse 369 solver): 370 371 .. math:: 372 373 E(z) = \\ \\frac{1}{2} \\int_{\\Omega} \\left( 374 \\left( \\frac{\\partial^2{z}}{\\partial{x}^2} \\right)^2 + 375 \\left( \\frac{\\partial^2{z}}{\\partial{y}^2} \\right)^2 + 376 2\\left( \\frac{\\partial^2{z}}{\\partial{y}\\partial{x}} 377 \\right)^2 \\right) dx\\,dy 378 379 If the case *kind* ='geom' is chosen by the user, a simple geometric 380 approximation is used (weighted average of the triangle normal 381 vectors), which could improve speed on very large grids. 382 383 References 384 ---------- 385 .. [1] Michel Bernadou, Kamal Hassan, "Basis functions for general 386 Hsieh-Clough-Tocher triangles, complete or reduced.", 387 International Journal for Numerical Methods in Engineering, 388 17(5):784 - 789. 2.01. 389 .. [2] C.T. Kelley, "Iterative Methods for Optimization". 390 391 """ 392 def __init__(self, triangulation, z, kind='min_E', trifinder=None, 393 dz=None): 394 TriInterpolator.__init__(self, triangulation, z, trifinder) 395 396 # Loads the underlying c++ _triangulation. 397 # (During loading, reordering of triangulation._triangles may occur so 398 # that all final triangles are now anti-clockwise) 399 self._triangulation.get_cpp_triangulation() 400 401 # To build the stiffness matrix and avoid zero-energy spurious modes 402 # we will only store internally the valid (unmasked) triangles and 403 # the necessary (used) points coordinates. 404 # 2 renumbering tables need to be computed and stored: 405 # - a triangle renum table in order to translate the result from a 406 # TriFinder instance into the internal stored triangle number. 407 # - a node renum table to overwrite the self._z values into the new 408 # (used) node numbering. 409 tri_analyzer = TriAnalyzer(self._triangulation) 410 (compressed_triangles, compressed_x, compressed_y, tri_renum, 411 node_renum) = tri_analyzer._get_compressed_triangulation(True, True) 412 self._triangles = compressed_triangles 413 self._tri_renum = tri_renum 414 # Taking into account the node renumbering in self._z: 415 node_mask = (node_renum == -1) 416 self._z[node_renum[~node_mask]] = self._z 417 self._z = self._z[~node_mask] 418 419 # Computing scale factors 420 self._unit_x = np.ptp(compressed_x) 421 self._unit_y = np.ptp(compressed_y) 422 self._pts = np.column_stack([compressed_x / self._unit_x, 423 compressed_y / self._unit_y]) 424 # Computing triangle points 425 self._tris_pts = self._pts[self._triangles] 426 # Computing eccentricities 427 self._eccs = self._compute_tri_eccentricities(self._tris_pts) 428 # Computing dof estimations for HCT triangle shape function 429 self._dof = self._compute_dof(kind, dz=dz) 430 # Loading HCT element 431 self._ReferenceElement = _ReducedHCT_Element() 432 433 def __call__(self, x, y): 434 return self._interpolate_multikeys(x, y, tri_index=None, 435 return_keys=('z',))[0] 436 __call__.__doc__ = TriInterpolator._docstring__call__ 437 438 def gradient(self, x, y): 439 return self._interpolate_multikeys(x, y, tri_index=None, 440 return_keys=('dzdx', 'dzdy')) 441 gradient.__doc__ = TriInterpolator._docstringgradient 442 443 def _interpolate_single_key(self, return_key, tri_index, x, y): 444 tris_pts = self._tris_pts[tri_index] 445 alpha = self._get_alpha_vec(x, y, tris_pts) 446 ecc = self._eccs[tri_index] 447 dof = np.expand_dims(self._dof[tri_index], axis=1) 448 if return_key == 'z': 449 return self._ReferenceElement.get_function_values( 450 alpha, ecc, dof) 451 elif return_key in ['dzdx', 'dzdy']: 452 J = self._get_jacobian(tris_pts) 453 dzdx = self._ReferenceElement.get_function_derivatives( 454 alpha, J, ecc, dof) 455 if return_key == 'dzdx': 456 return dzdx[:, 0, 0] 457 else: 458 return dzdx[:, 1, 0] 459 else: 460 raise ValueError("Invalid return_key: " + return_key) 461 462 def _compute_dof(self, kind, dz=None): 463 """ 464 Computes and returns nodal dofs according to kind 465 466 Parameters 467 ---------- 468 kind: {'min_E', 'geom', 'user'} 469 Choice of the _DOF_estimator subclass to perform the gradient 470 estimation. 471 dz: tuple of array_likes (dzdx, dzdy), optional 472 Used only if *kind=user ; in this case passed to the 473 :class:`_DOF_estimator_user`. 474 475 Returns 476 ------- 477 dof : array_like, shape (npts,2) 478 Estimation of the gradient at triangulation nodes (stored as 479 degree of freedoms of reduced-HCT triangle elements). 480 """ 481 if kind == 'user': 482 if dz is None: 483 raise ValueError("For a CubicTriInterpolator with " 484 "*kind*='user', a valid *dz* " 485 "argument is expected.") 486 TE = _DOF_estimator_user(self, dz=dz) 487 elif kind == 'geom': 488 TE = _DOF_estimator_geom(self) 489 elif kind == 'min_E': 490 TE = _DOF_estimator_min_E(self) 491 else: 492 raise ValueError("CubicTriInterpolator *kind* proposed: {0} ; " 493 "should be one of: " 494 "'user', 'geom', 'min_E'".format(kind)) 495 return TE.compute_dof_from_df() 496 497 @staticmethod 498 def _get_alpha_vec(x, y, tris_pts): 499 """ 500 Fast (vectorized) function to compute barycentric coordinates alpha. 501 502 Parameters 503 ---------- 504 x, y : array-like of dim 1 (shape (nx,)) 505 Coordinates of the points whose points barycentric 506 coordinates are requested 507 tris_pts : array like of dim 3 (shape: (nx,3,2)) 508 Coordinates of the containing triangles apexes. 509 510 Returns 511 ------- 512 alpha : array of dim 2 (shape (nx,3)) 513 Barycentric coordinates of the points inside the containing 514 triangles. 515 """ 516 ndim = tris_pts.ndim-2 517 518 a = tris_pts[:, 1, :] - tris_pts[:, 0, :] 519 b = tris_pts[:, 2, :] - tris_pts[:, 0, :] 520 abT = np.concatenate([np.expand_dims(a, ndim+1), 521 np.expand_dims(b, ndim+1)], ndim+1) 522 ab = _transpose_vectorized(abT) 523 x = np.expand_dims(x, ndim) 524 y = np.expand_dims(y, ndim) 525 OM = np.concatenate([x, y], ndim) - tris_pts[:, 0, :] 526 527 metric = _prod_vectorized(ab, abT) 528 # Here we try to deal with the colinear cases. 529 # metric_inv is in this case set to the Moore-Penrose pseudo-inverse 530 # meaning that we will still return a set of valid barycentric 531 # coordinates. 532 metric_inv = _pseudo_inv22sym_vectorized(metric) 533 Covar = _prod_vectorized(ab, _transpose_vectorized( 534 np.expand_dims(OM, ndim))) 535 ksi = _prod_vectorized(metric_inv, Covar) 536 alpha = _to_matrix_vectorized([ 537 [1-ksi[:, 0, 0]-ksi[:, 1, 0]], [ksi[:, 0, 0]], [ksi[:, 1, 0]]]) 538 return alpha 539 540 @staticmethod 541 def _get_jacobian(tris_pts): 542 """ 543 Fast (vectorized) function to compute triangle jacobian matrix. 544 545 Parameters 546 ---------- 547 tris_pts : array like of dim 3 (shape: (nx,3,2)) 548 Coordinates of the containing triangles apexes. 549 550 Returns 551 ------- 552 J : array of dim 3 (shape (nx,2,2)) 553 Barycentric coordinates of the points inside the containing 554 triangles. 555 J[itri,:,:] is the jacobian matrix at apex 0 of the triangle 556 itri, so that the following (matrix) relationship holds: 557 [dz/dksi] = [J] x [dz/dx] 558 with x: global coordinates 559 ksi: element parametric coordinates in triangle first apex 560 local basis. 561 """ 562 a = np.array(tris_pts[:, 1, :] - tris_pts[:, 0, :]) 563 b = np.array(tris_pts[:, 2, :] - tris_pts[:, 0, :]) 564 J = _to_matrix_vectorized([[a[:, 0], a[:, 1]], 565 [b[:, 0], b[:, 1]]]) 566 return J 567 568 @staticmethod 569 def _compute_tri_eccentricities(tris_pts): 570 """ 571 Computes triangle eccentricities 572 573 Parameters 574 ---------- 575 tris_pts : array like of dim 3 (shape: (nx,3,2)) 576 Coordinates of the triangles apexes. 577 578 Returns 579 ------- 580 ecc : array like of dim 2 (shape: (nx,3)) 581 The so-called eccentricity parameters [1] needed for 582 HCT triangular element. 583 """ 584 a = np.expand_dims(tris_pts[:, 2, :] - tris_pts[:, 1, :], axis=2) 585 b = np.expand_dims(tris_pts[:, 0, :] - tris_pts[:, 2, :], axis=2) 586 c = np.expand_dims(tris_pts[:, 1, :] - tris_pts[:, 0, :], axis=2) 587 # Do not use np.squeeze, this is dangerous if only one triangle 588 # in the triangulation... 589 dot_a = _prod_vectorized(_transpose_vectorized(a), a)[:, 0, 0] 590 dot_b = _prod_vectorized(_transpose_vectorized(b), b)[:, 0, 0] 591 dot_c = _prod_vectorized(_transpose_vectorized(c), c)[:, 0, 0] 592 # Note that this line will raise a warning for dot_a, dot_b or dot_c 593 # zeros, but we choose not to support triangles with duplicate points. 594 return _to_matrix_vectorized([[(dot_c-dot_b) / dot_a], 595 [(dot_a-dot_c) / dot_b], 596 [(dot_b-dot_a) / dot_c]]) 597 598 599# FEM element used for interpolation and for solving minimisation 600# problem (Reduced HCT element) 601class _ReducedHCT_Element(): 602 """ 603 Implementation of reduced HCT triangular element with explicit shape 604 functions. 605 606 Computes z, dz, d2z and the element stiffness matrix for bending energy: 607 E(f) = integral( (d2z/dx2 + d2z/dy2)**2 dA) 608 609 *** Reference for the shape functions: *** 610 [1] Basis functions for general Hsieh-Clough-Tocher _triangles, complete or 611 reduced. 612 Michel Bernadou, Kamal Hassan 613 International Journal for Numerical Methods in Engineering. 614 17(5):784 - 789. 2.01 615 616 *** Element description: *** 617 9 dofs: z and dz given at 3 apex 618 C1 (conform) 619 620 """ 621 # 1) Loads matrices to generate shape functions as a function of 622 # triangle eccentricities - based on [1] p.11 ''' 623 M = np.array([ 624 [ 0.00, 0.00, 0.00, 4.50, 4.50, 0.00, 0.00, 0.00, 0.00, 0.00], 625 [-0.25, 0.00, 0.00, 0.50, 1.25, 0.00, 0.00, 0.00, 0.00, 0.00], 626 [-0.25, 0.00, 0.00, 1.25, 0.50, 0.00, 0.00, 0.00, 0.00, 0.00], 627 [ 0.50, 1.00, 0.00, -1.50, 0.00, 3.00, 3.00, 0.00, 0.00, 3.00], 628 [ 0.00, 0.00, 0.00, -0.25, 0.25, 0.00, 1.00, 0.00, 0.00, 0.50], 629 [ 0.25, 0.00, 0.00, -0.50, -0.25, 1.00, 0.00, 0.00, 0.00, 1.00], 630 [ 0.50, 0.00, 1.00, 0.00, -1.50, 0.00, 0.00, 3.00, 3.00, 3.00], 631 [ 0.25, 0.00, 0.00, -0.25, -0.50, 0.00, 0.00, 0.00, 1.00, 1.00], 632 [ 0.00, 0.00, 0.00, 0.25, -0.25, 0.00, 0.00, 1.00, 0.00, 0.50]]) 633 M0 = np.array([ 634 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 635 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 636 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 637 [-1.00, 0.00, 0.00, 1.50, 1.50, 0.00, 0.00, 0.00, 0.00, -3.00], 638 [-0.50, 0.00, 0.00, 0.75, 0.75, 0.00, 0.00, 0.00, 0.00, -1.50], 639 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 640 [ 1.00, 0.00, 0.00, -1.50, -1.50, 0.00, 0.00, 0.00, 0.00, 3.00], 641 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 642 [ 0.50, 0.00, 0.00, -0.75, -0.75, 0.00, 0.00, 0.00, 0.00, 1.50]]) 643 M1 = np.array([ 644 [-0.50, 0.00, 0.00, 1.50, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 645 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 646 [-0.25, 0.00, 0.00, 0.75, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 647 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 648 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 649 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 650 [ 0.50, 0.00, 0.00, -1.50, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 651 [ 0.25, 0.00, 0.00, -0.75, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 652 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]]) 653 M2 = np.array([ 654 [ 0.50, 0.00, 0.00, 0.00, -1.50, 0.00, 0.00, 0.00, 0.00, 0.00], 655 [ 0.25, 0.00, 0.00, 0.00, -0.75, 0.00, 0.00, 0.00, 0.00, 0.00], 656 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 657 [-0.50, 0.00, 0.00, 0.00, 1.50, 0.00, 0.00, 0.00, 0.00, 0.00], 658 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 659 [-0.25, 0.00, 0.00, 0.00, 0.75, 0.00, 0.00, 0.00, 0.00, 0.00], 660 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 661 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], 662 [ 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]]) 663 664 # 2) Loads matrices to rotate components of gradient & Hessian 665 # vectors in the reference basis of triangle first apex (a0) 666 rotate_dV = np.array([[ 1., 0.], [ 0., 1.], 667 [ 0., 1.], [-1., -1.], 668 [-1., -1.], [ 1., 0.]]) 669 670 rotate_d2V = np.array([[1., 0., 0.], [0., 1., 0.], [ 0., 0., 1.], 671 [0., 1., 0.], [1., 1., 1.], [ 0., -2., -1.], 672 [1., 1., 1.], [1., 0., 0.], [-2., 0., -1.]]) 673 674 # 3) Loads Gauss points & weights on the 3 sub-_triangles for P2 675 # exact integral - 3 points on each subtriangles. 676 # NOTE: as the 2nd derivative is discontinuous , we really need those 9 677 # points! 678 n_gauss = 9 679 gauss_pts = np.array([[13./18., 4./18., 1./18.], 680 [ 4./18., 13./18., 1./18.], 681 [ 7./18., 7./18., 4./18.], 682 [ 1./18., 13./18., 4./18.], 683 [ 1./18., 4./18., 13./18.], 684 [ 4./18., 7./18., 7./18.], 685 [ 4./18., 1./18., 13./18.], 686 [13./18., 1./18., 4./18.], 687 [ 7./18., 4./18., 7./18.]], dtype=np.float64) 688 gauss_w = np.ones([9], dtype=np.float64) / 9. 689 690 # 4) Stiffness matrix for curvature energy 691 E = np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 2.]]) 692 693 # 5) Loads the matrix to compute DOF_rot from tri_J at apex 0 694 J0_to_J1 = np.array([[-1., 1.], [-1., 0.]]) 695 J0_to_J2 = np.array([[ 0., -1.], [ 1., -1.]]) 696 697 def get_function_values(self, alpha, ecc, dofs): 698 """ 699 Parameters 700 ---------- 701 alpha : is a (N x 3 x 1) array (array of column-matrices) of 702 barycentric coordinates, 703 ecc : is a (N x 3 x 1) array (array of column-matrices) of triangle 704 eccentricities, 705 dofs : is a (N x 1 x 9) arrays (arrays of row-matrices) of computed 706 degrees of freedom. 707 708 Returns 709 ------- 710 Returns the N-array of interpolated function values. 711 """ 712 subtri = np.argmin(alpha, axis=1)[:, 0] 713 ksi = _roll_vectorized(alpha, -subtri, axis=0) 714 E = _roll_vectorized(ecc, -subtri, axis=0) 715 x = ksi[:, 0, 0] 716 y = ksi[:, 1, 0] 717 z = ksi[:, 2, 0] 718 x_sq = x*x 719 y_sq = y*y 720 z_sq = z*z 721 V = _to_matrix_vectorized([ 722 [x_sq*x], [y_sq*y], [z_sq*z], [x_sq*z], [x_sq*y], [y_sq*x], 723 [y_sq*z], [z_sq*y], [z_sq*x], [x*y*z]]) 724 prod = _prod_vectorized(self.M, V) 725 prod += _scalar_vectorized(E[:, 0, 0], 726 _prod_vectorized(self.M0, V)) 727 prod += _scalar_vectorized(E[:, 1, 0], 728 _prod_vectorized(self.M1, V)) 729 prod += _scalar_vectorized(E[:, 2, 0], 730 _prod_vectorized(self.M2, V)) 731 s = _roll_vectorized(prod, 3*subtri, axis=0) 732 return _prod_vectorized(dofs, s)[:, 0, 0] 733 734 def get_function_derivatives(self, alpha, J, ecc, dofs): 735 """ 736 Parameters 737 ---------- 738 *alpha* is a (N x 3 x 1) array (array of column-matrices of 739 barycentric coordinates) 740 *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at 741 triangle first apex) 742 *ecc* is a (N x 3 x 1) array (array of column-matrices of triangle 743 eccentricities) 744 *dofs* is a (N x 1 x 9) arrays (arrays of row-matrices) of computed 745 degrees of freedom. 746 747 Returns 748 ------- 749 Returns the values of interpolated function derivatives [dz/dx, dz/dy] 750 in global coordinates at locations alpha, as a column-matrices of 751 shape (N x 2 x 1). 752 """ 753 subtri = np.argmin(alpha, axis=1)[:, 0] 754 ksi = _roll_vectorized(alpha, -subtri, axis=0) 755 E = _roll_vectorized(ecc, -subtri, axis=0) 756 x = ksi[:, 0, 0] 757 y = ksi[:, 1, 0] 758 z = ksi[:, 2, 0] 759 x_sq = x*x 760 y_sq = y*y 761 z_sq = z*z 762 dV = _to_matrix_vectorized([ 763 [ -3.*x_sq, -3.*x_sq], 764 [ 3.*y_sq, 0.], 765 [ 0., 3.*z_sq], 766 [ -2.*x*z, -2.*x*z+x_sq], 767 [-2.*x*y+x_sq, -2.*x*y], 768 [ 2.*x*y-y_sq, -y_sq], 769 [ 2.*y*z, y_sq], 770 [ z_sq, 2.*y*z], 771 [ -z_sq, 2.*x*z-z_sq], 772 [ x*z-y*z, x*y-y*z]]) 773 # Puts back dV in first apex basis 774 dV = _prod_vectorized(dV, _extract_submatrices( 775 self.rotate_dV, subtri, block_size=2, axis=0)) 776 777 prod = _prod_vectorized(self.M, dV) 778 prod += _scalar_vectorized(E[:, 0, 0], 779 _prod_vectorized(self.M0, dV)) 780 prod += _scalar_vectorized(E[:, 1, 0], 781 _prod_vectorized(self.M1, dV)) 782 prod += _scalar_vectorized(E[:, 2, 0], 783 _prod_vectorized(self.M2, dV)) 784 dsdksi = _roll_vectorized(prod, 3*subtri, axis=0) 785 dfdksi = _prod_vectorized(dofs, dsdksi) 786 # In global coordinates: 787 # Here we try to deal with the simplest colinear cases, returning a 788 # null matrix. 789 J_inv = _safe_inv22_vectorized(J) 790 dfdx = _prod_vectorized(J_inv, _transpose_vectorized(dfdksi)) 791 return dfdx 792 793 def get_function_hessians(self, alpha, J, ecc, dofs): 794 """ 795 Parameters 796 ---------- 797 *alpha* is a (N x 3 x 1) array (array of column-matrices) of 798 barycentric coordinates 799 *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at 800 triangle first apex) 801 *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle 802 eccentricities 803 *dofs* is a (N x 1 x 9) arrays (arrays of row-matrices) of computed 804 degrees of freedom. 805 806 Returns 807 ------- 808 Returns the values of interpolated function 2nd-derivatives 809 [d2z/dx2, d2z/dy2, d2z/dxdy] in global coordinates at locations alpha, 810 as a column-matrices of shape (N x 3 x 1). 811 """ 812 d2sdksi2 = self.get_d2Sidksij2(alpha, ecc) 813 d2fdksi2 = _prod_vectorized(dofs, d2sdksi2) 814 H_rot = self.get_Hrot_from_J(J) 815 d2fdx2 = _prod_vectorized(d2fdksi2, H_rot) 816 return _transpose_vectorized(d2fdx2) 817 818 def get_d2Sidksij2(self, alpha, ecc): 819 """ 820 Parameters 821 ---------- 822 *alpha* is a (N x 3 x 1) array (array of column-matrices) of 823 barycentric coordinates 824 *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle 825 eccentricities 826 827 Returns 828 ------- 829 Returns the arrays d2sdksi2 (N x 3 x 1) Hessian of shape functions 830 expressed in covariante coordinates in first apex basis. 831 """ 832 subtri = np.argmin(alpha, axis=1)[:, 0] 833 ksi = _roll_vectorized(alpha, -subtri, axis=0) 834 E = _roll_vectorized(ecc, -subtri, axis=0) 835 x = ksi[:, 0, 0] 836 y = ksi[:, 1, 0] 837 z = ksi[:, 2, 0] 838 d2V = _to_matrix_vectorized([ 839 [ 6.*x, 6.*x, 6.*x], 840 [ 6.*y, 0., 0.], 841 [ 0., 6.*z, 0.], 842 [ 2.*z, 2.*z-4.*x, 2.*z-2.*x], 843 [2.*y-4.*x, 2.*y, 2.*y-2.*x], 844 [2.*x-4.*y, 0., -2.*y], 845 [ 2.*z, 0., 2.*y], 846 [ 0., 2.*y, 2.*z], 847 [ 0., 2.*x-4.*z, -2.*z], 848 [ -2.*z, -2.*y, x-y-z]]) 849 # Puts back d2V in first apex basis 850 d2V = _prod_vectorized(d2V, _extract_submatrices( 851 self.rotate_d2V, subtri, block_size=3, axis=0)) 852 prod = _prod_vectorized(self.M, d2V) 853 prod += _scalar_vectorized(E[:, 0, 0], 854 _prod_vectorized(self.M0, d2V)) 855 prod += _scalar_vectorized(E[:, 1, 0], 856 _prod_vectorized(self.M1, d2V)) 857 prod += _scalar_vectorized(E[:, 2, 0], 858 _prod_vectorized(self.M2, d2V)) 859 d2sdksi2 = _roll_vectorized(prod, 3*subtri, axis=0) 860 return d2sdksi2 861 862 def get_bending_matrices(self, J, ecc): 863 """ 864 Parameters 865 ---------- 866 *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at 867 triangle first apex) 868 *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle 869 eccentricities 870 871 Returns 872 ------- 873 Returns the element K matrices for bending energy expressed in 874 GLOBAL nodal coordinates. 875 K_ij = integral [ (d2zi/dx2 + d2zi/dy2) * (d2zj/dx2 + d2zj/dy2) dA] 876 tri_J is needed to rotate dofs from local basis to global basis 877 """ 878 n = np.size(ecc, 0) 879 880 # 1) matrix to rotate dofs in global coordinates 881 J1 = _prod_vectorized(self.J0_to_J1, J) 882 J2 = _prod_vectorized(self.J0_to_J2, J) 883 DOF_rot = np.zeros([n, 9, 9], dtype=np.float64) 884 DOF_rot[:, 0, 0] = 1 885 DOF_rot[:, 3, 3] = 1 886 DOF_rot[:, 6, 6] = 1 887 DOF_rot[:, 1:3, 1:3] = J 888 DOF_rot[:, 4:6, 4:6] = J1 889 DOF_rot[:, 7:9, 7:9] = J2 890 891 # 2) matrix to rotate Hessian in global coordinates. 892 H_rot, area = self.get_Hrot_from_J(J, return_area=True) 893 894 # 3) Computes stiffness matrix 895 # Gauss quadrature. 896 K = np.zeros([n, 9, 9], dtype=np.float64) 897 weights = self.gauss_w 898 pts = self.gauss_pts 899 for igauss in range(self.n_gauss): 900 alpha = np.tile(pts[igauss, :], n).reshape(n, 3) 901 alpha = np.expand_dims(alpha, 2) 902 weight = weights[igauss] 903 d2Skdksi2 = self.get_d2Sidksij2(alpha, ecc) 904 d2Skdx2 = _prod_vectorized(d2Skdksi2, H_rot) 905 K += weight * _prod_vectorized(_prod_vectorized(d2Skdx2, self.E), 906 _transpose_vectorized(d2Skdx2)) 907 908 # 4) With nodal (not elem) dofs 909 K = _prod_vectorized(_prod_vectorized(_transpose_vectorized(DOF_rot), 910 K), DOF_rot) 911 912 # 5) Need the area to compute total element energy 913 return _scalar_vectorized(area, K) 914 915 def get_Hrot_from_J(self, J, return_area=False): 916 """ 917 Parameters 918 ---------- 919 *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at 920 triangle first apex) 921 922 Returns 923 ------- 924 Returns H_rot used to rotate Hessian from local basis of first apex, 925 to global coordinates. 926 if *return_area* is True, returns also the triangle area (0.5*det(J)) 927 """ 928 # Here we try to deal with the simplest colinear cases ; a null 929 # energy and area is imposed. 930 J_inv = _safe_inv22_vectorized(J) 931 Ji00 = J_inv[:, 0, 0] 932 Ji11 = J_inv[:, 1, 1] 933 Ji10 = J_inv[:, 1, 0] 934 Ji01 = J_inv[:, 0, 1] 935 H_rot = _to_matrix_vectorized([ 936 [Ji00*Ji00, Ji10*Ji10, Ji00*Ji10], 937 [Ji01*Ji01, Ji11*Ji11, Ji01*Ji11], 938 [2*Ji00*Ji01, 2*Ji11*Ji10, Ji00*Ji11+Ji10*Ji01]]) 939 if not return_area: 940 return H_rot 941 else: 942 area = 0.5 * (J[:, 0, 0]*J[:, 1, 1] - J[:, 0, 1]*J[:, 1, 0]) 943 return H_rot, area 944 945 def get_Kff_and_Ff(self, J, ecc, triangles, Uc): 946 """ 947 Builds K and F for the following elliptic formulation: 948 minimization of curvature energy with value of function at node 949 imposed and derivatives 'free'. 950 Builds the global Kff matrix in cco format. 951 Builds the full Ff vec Ff = - Kfc x Uc 952 953 Parameters 954 ---------- 955 *J* is a (N x 2 x 2) array of jacobian matrices (jacobian matrix at 956 triangle first apex) 957 *ecc* is a (N x 3 x 1) array (array of column-matrices) of triangle 958 eccentricities 959 *triangles* is a (N x 3) array of nodes indexes. 960 *Uc* is (N x 3) array of imposed displacements at nodes 961 962 Returns 963 ------- 964 (Kff_rows, Kff_cols, Kff_vals) Kff matrix in coo format - Duplicate 965 (row, col) entries must be summed. 966 Ff: force vector - dim npts * 3 967 """ 968 ntri = np.size(ecc, 0) 969 vec_range = np.arange(ntri, dtype=np.int32) 970 c_indices = -np.ones(ntri, dtype=np.int32) # for unused dofs, -1 971 f_dof = [1, 2, 4, 5, 7, 8] 972 c_dof = [0, 3, 6] 973 974 # vals, rows and cols indices in global dof numbering 975 f_dof_indices = _to_matrix_vectorized([[ 976 c_indices, triangles[:, 0]*2, triangles[:, 0]*2+1, 977 c_indices, triangles[:, 1]*2, triangles[:, 1]*2+1, 978 c_indices, triangles[:, 2]*2, triangles[:, 2]*2+1]]) 979 980 expand_indices = np.ones([ntri, 9, 1], dtype=np.int32) 981 f_row_indices = _prod_vectorized(_transpose_vectorized(f_dof_indices), 982 _transpose_vectorized(expand_indices)) 983 f_col_indices = _prod_vectorized(expand_indices, f_dof_indices) 984 K_elem = self.get_bending_matrices(J, ecc) 985 986 # Extracting sub-matrices 987 # Explanation & notations: 988 # * Subscript f denotes 'free' degrees of freedom (i.e. dz/dx, dz/dx) 989 # * Subscript c denotes 'condensated' (imposed) degrees of freedom 990 # (i.e. z at all nodes) 991 # * F = [Ff, Fc] is the force vector 992 # * U = [Uf, Uc] is the imposed dof vector 993 # [ Kff Kfc ] 994 # * K = [ ] is the laplacian stiffness matrix 995 # [ Kcf Kff ] 996 # * As F = K x U one gets straightforwardly: Ff = - Kfc x Uc 997 998 # Computing Kff stiffness matrix in sparse coo format 999 Kff_vals = np.ravel(K_elem[np.ix_(vec_range, f_dof, f_dof)]) 1000 Kff_rows = np.ravel(f_row_indices[np.ix_(vec_range, f_dof, f_dof)]) 1001 Kff_cols = np.ravel(f_col_indices[np.ix_(vec_range, f_dof, f_dof)]) 1002 1003 # Computing Ff force vector in sparse coo format 1004 Kfc_elem = K_elem[np.ix_(vec_range, f_dof, c_dof)] 1005 Uc_elem = np.expand_dims(Uc, axis=2) 1006 Ff_elem = - _prod_vectorized(Kfc_elem, Uc_elem)[:, :, 0] 1007 Ff_indices = f_dof_indices[np.ix_(vec_range, [0], f_dof)][:, 0, :] 1008 1009 # Extracting Ff force vector in dense format 1010 # We have to sum duplicate indices - using bincount 1011 Ff = np.bincount(np.ravel(Ff_indices), weights=np.ravel(Ff_elem)) 1012 return Kff_rows, Kff_cols, Kff_vals, Ff 1013 1014 1015# :class:_DOF_estimator, _DOF_estimator_user, _DOF_estimator_geom, 1016# _DOF_estimator_min_E 1017# Private classes used to compute the degree of freedom of each triangular 1018# element for the TriCubicInterpolator. 1019class _DOF_estimator(): 1020 """ 1021 Abstract base class for classes used to perform estimation of a function 1022 first derivatives, and deduce the dofs for a CubicTriInterpolator using a 1023 reduced HCT element formulation. 1024 Derived classes implement compute_df(self,**kwargs), returning 1025 np.vstack([dfx,dfy]).T where : dfx, dfy are the estimation of the 2 1026 gradient coordinates. 1027 """ 1028 def __init__(self, interpolator, **kwargs): 1029 if not isinstance(interpolator, CubicTriInterpolator): 1030 raise ValueError("Expected a CubicTriInterpolator object") 1031 self._pts = interpolator._pts 1032 self._tris_pts = interpolator._tris_pts 1033 self.z = interpolator._z 1034 self._triangles = interpolator._triangles 1035 (self._unit_x, self._unit_y) = (interpolator._unit_x, 1036 interpolator._unit_y) 1037 self.dz = self.compute_dz(**kwargs) 1038 self.compute_dof_from_df() 1039 1040 def compute_dz(self, **kwargs): 1041 raise NotImplementedError 1042 1043 def compute_dof_from_df(self): 1044 """ 1045 Computes reduced-HCT elements degrees of freedom, knowing the 1046 gradient. 1047 """ 1048 J = CubicTriInterpolator._get_jacobian(self._tris_pts) 1049 tri_z = self.z[self._triangles] 1050 tri_dz = self.dz[self._triangles] 1051 tri_dof = self.get_dof_vec(tri_z, tri_dz, J) 1052 return tri_dof 1053 1054 @staticmethod 1055 def get_dof_vec(tri_z, tri_dz, J): 1056 """ 1057 Computes the dof vector of a triangle, knowing the value of f, df and 1058 of the local Jacobian at each node. 1059 1060 *tri_z*: array of shape (3,) of f nodal values 1061 *tri_dz*: array of shape (3,2) of df/dx, df/dy nodal values 1062 *J*: Jacobian matrix in local basis of apex 0 1063 1064 Returns dof array of shape (9,) so that for each apex iapex: 1065 dof[iapex*3+0] = f(Ai) 1066 dof[iapex*3+1] = df(Ai).(AiAi+) 1067 dof[iapex*3+2] = df(Ai).(AiAi-)] 1068 """ 1069 npt = tri_z.shape[0] 1070 dof = np.zeros([npt, 9], dtype=np.float64) 1071 J1 = _prod_vectorized(_ReducedHCT_Element.J0_to_J1, J) 1072 J2 = _prod_vectorized(_ReducedHCT_Element.J0_to_J2, J) 1073 1074 col0 = _prod_vectorized(J, np.expand_dims(tri_dz[:, 0, :], axis=2)) 1075 col1 = _prod_vectorized(J1, np.expand_dims(tri_dz[:, 1, :], axis=2)) 1076 col2 = _prod_vectorized(J2, np.expand_dims(tri_dz[:, 2, :], axis=2)) 1077 1078 dfdksi = _to_matrix_vectorized([ 1079 [col0[:, 0, 0], col1[:, 0, 0], col2[:, 0, 0]], 1080 [col0[:, 1, 0], col1[:, 1, 0], col2[:, 1, 0]]]) 1081 dof[:, 0:7:3] = tri_z 1082 dof[:, 1:8:3] = dfdksi[:, 0] 1083 dof[:, 2:9:3] = dfdksi[:, 1] 1084 return dof 1085 1086 1087class _DOF_estimator_user(_DOF_estimator): 1088 """ dz is imposed by user / Accounts for scaling if any """ 1089 def compute_dz(self, dz): 1090 (dzdx, dzdy) = dz 1091 dzdx = dzdx * self._unit_x 1092 dzdy = dzdy * self._unit_y 1093 return np.vstack([dzdx, dzdy]).T 1094 1095 1096class _DOF_estimator_geom(_DOF_estimator): 1097 """ Fast 'geometric' approximation, recommended for large arrays. """ 1098 def compute_dz(self): 1099 """ 1100 self.df is computed as weighted average of _triangles sharing a common 1101 node. On each triangle itri f is first assumed linear (= ~f), which 1102 allows to compute d~f[itri] 1103 Then the following approximation of df nodal values is then proposed: 1104 f[ipt] = SUM ( w[itri] x d~f[itri] , for itri sharing apex ipt) 1105 The weighted coeff. w[itri] are proportional to the angle of the 1106 triangle itri at apex ipt 1107 """ 1108 el_geom_w = self.compute_geom_weights() 1109 el_geom_grad = self.compute_geom_grads() 1110 1111 # Sum of weights coeffs 1112 w_node_sum = np.bincount(np.ravel(self._triangles), 1113 weights=np.ravel(el_geom_w)) 1114 1115 # Sum of weighted df = (dfx, dfy) 1116 dfx_el_w = np.empty_like(el_geom_w) 1117 dfy_el_w = np.empty_like(el_geom_w) 1118 for iapex in range(3): 1119 dfx_el_w[:, iapex] = el_geom_w[:, iapex]*el_geom_grad[:, 0] 1120 dfy_el_w[:, iapex] = el_geom_w[:, iapex]*el_geom_grad[:, 1] 1121 dfx_node_sum = np.bincount(np.ravel(self._triangles), 1122 weights=np.ravel(dfx_el_w)) 1123 dfy_node_sum = np.bincount(np.ravel(self._triangles), 1124 weights=np.ravel(dfy_el_w)) 1125 1126 # Estimation of df 1127 dfx_estim = dfx_node_sum/w_node_sum 1128 dfy_estim = dfy_node_sum/w_node_sum 1129 return np.vstack([dfx_estim, dfy_estim]).T 1130 1131 def compute_geom_weights(self): 1132 """ 1133 Builds the (nelems x 3) weights coeffs of _triangles angles, 1134 renormalized so that np.sum(weights, axis=1) == np.ones(nelems) 1135 """ 1136 weights = np.zeros([np.size(self._triangles, 0), 3]) 1137 tris_pts = self._tris_pts 1138 for ipt in range(3): 1139 p0 = tris_pts[:, (ipt) % 3, :] 1140 p1 = tris_pts[:, (ipt+1) % 3, :] 1141 p2 = tris_pts[:, (ipt-1) % 3, :] 1142 alpha1 = np.arctan2(p1[:, 1]-p0[:, 1], p1[:, 0]-p0[:, 0]) 1143 alpha2 = np.arctan2(p2[:, 1]-p0[:, 1], p2[:, 0]-p0[:, 0]) 1144 # In the below formula we could take modulo 2. but 1145 # modulo 1. is safer regarding round-off errors (flat triangles). 1146 angle = np.abs(np.mod((alpha2-alpha1) / np.pi, 1.)) 1147 # Weight proportional to angle up np.pi/2 ; null weight for 1148 # degenerated cases 0. and np.pi (Note that `angle` is normalized 1149 # by np.pi) 1150 weights[:, ipt] = 0.5 - np.abs(angle-0.5) 1151 return weights 1152 1153 def compute_geom_grads(self): 1154 """ 1155 Compute the (global) gradient component of f assumed linear (~f). 1156 returns array df of shape (nelems,2) 1157 df[ielem].dM[ielem] = dz[ielem] i.e. df = dz x dM = dM.T^-1 x dz 1158 """ 1159 tris_pts = self._tris_pts 1160 tris_f = self.z[self._triangles] 1161 1162 dM1 = tris_pts[:, 1, :] - tris_pts[:, 0, :] 1163 dM2 = tris_pts[:, 2, :] - tris_pts[:, 0, :] 1164 dM = np.dstack([dM1, dM2]) 1165 # Here we try to deal with the simplest colinear cases: a null 1166 # gradient is assumed in this case. 1167 dM_inv = _safe_inv22_vectorized(dM) 1168 1169 dZ1 = tris_f[:, 1] - tris_f[:, 0] 1170 dZ2 = tris_f[:, 2] - tris_f[:, 0] 1171 dZ = np.vstack([dZ1, dZ2]).T 1172 df = np.empty_like(dZ) 1173 1174 # With np.einsum : could be ej,eji -> ej 1175 df[:, 0] = dZ[:, 0]*dM_inv[:, 0, 0] + dZ[:, 1]*dM_inv[:, 1, 0] 1176 df[:, 1] = dZ[:, 0]*dM_inv[:, 0, 1] + dZ[:, 1]*dM_inv[:, 1, 1] 1177 return df 1178 1179 1180class _DOF_estimator_min_E(_DOF_estimator_geom): 1181 """ 1182 The 'smoothest' approximation, df is computed through global minimization 1183 of the bending energy: 1184 E(f) = integral[(d2z/dx2 + d2z/dy2 + 2 d2z/dxdy)**2 dA] 1185 """ 1186 def __init__(self, Interpolator): 1187 self._eccs = Interpolator._eccs 1188 _DOF_estimator_geom.__init__(self, Interpolator) 1189 1190 def compute_dz(self): 1191 """ 1192 Elliptic solver for bending energy minimization. 1193 Uses a dedicated 'toy' sparse Jacobi PCG solver. 1194 """ 1195 # Initial guess for iterative PCG solver. 1196 dz_init = _DOF_estimator_geom.compute_dz(self) 1197 Uf0 = np.ravel(dz_init) 1198 1199 reference_element = _ReducedHCT_Element() 1200 J = CubicTriInterpolator._get_jacobian(self._tris_pts) 1201 eccs = self._eccs 1202 triangles = self._triangles 1203 Uc = self.z[self._triangles] 1204 1205 # Building stiffness matrix and force vector in coo format 1206 Kff_rows, Kff_cols, Kff_vals, Ff = reference_element.get_Kff_and_Ff( 1207 J, eccs, triangles, Uc) 1208 1209 # Building sparse matrix and solving minimization problem 1210 # We could use scipy.sparse direct solver ; however to avoid this 1211 # external dependency an implementation of a simple PCG solver with 1212 # a simplendiagonal Jocabi preconditioner is implemented. 1213 tol = 1.e-10 1214 n_dof = Ff.shape[0] 1215 Kff_coo = _Sparse_Matrix_coo(Kff_vals, Kff_rows, Kff_cols, 1216 shape=(n_dof, n_dof)) 1217 Kff_coo.compress_csc() 1218 Uf, err = _cg(A=Kff_coo, b=Ff, x0=Uf0, tol=tol) 1219 # If the PCG did not converge, we return the best guess between Uf0 1220 # and Uf. 1221 err0 = np.linalg.norm(Kff_coo.dot(Uf0) - Ff) 1222 if err0 < err: 1223 # Maybe a good occasion to raise a warning here ? 1224 warnings.warn("In TriCubicInterpolator initialization, PCG sparse" 1225 " solver did not converge after 1000 iterations. " 1226 "`geom` approximation is used instead of `min_E`") 1227 Uf = Uf0 1228 1229 # Building dz from Uf 1230 dz = np.empty([self._pts.shape[0], 2], dtype=np.float64) 1231 dz[:, 0] = Uf[::2] 1232 dz[:, 1] = Uf[1::2] 1233 return dz 1234 1235 1236# The following private :class:_Sparse_Matrix_coo and :func:_cg provide 1237# a PCG sparse solver for (symmetric) elliptic problems. 1238class _Sparse_Matrix_coo(object): 1239 def __init__(self, vals, rows, cols, shape): 1240 """ 1241 Creates a sparse matrix in coo format 1242 *vals*: arrays of values of non-null entries of the matrix 1243 *rows*: int arrays of rows of non-null entries of the matrix 1244 *cols*: int arrays of cols of non-null entries of the matrix 1245 *shape*: 2-tuple (n,m) of matrix shape 1246 1247 """ 1248 self.n, self.m = shape 1249 self.vals = np.asarray(vals, dtype=np.float64) 1250 self.rows = np.asarray(rows, dtype=np.int32) 1251 self.cols = np.asarray(cols, dtype=np.int32) 1252 1253 def dot(self, V): 1254 """ 1255 Dot product of self by a vector *V* in sparse-dense to dense format 1256 *V* dense vector of shape (self.m,) 1257 """ 1258 assert V.shape == (self.m,) 1259 return np.bincount(self.rows, 1260 weights=self.vals*V[self.cols], 1261 minlength=self.m) 1262 1263 def compress_csc(self): 1264 """ 1265 Compress rows, cols, vals / summing duplicates. Sort for csc format. 1266 """ 1267 _, unique, indices = np.unique( 1268 self.rows + self.n*self.cols, 1269 return_index=True, return_inverse=True) 1270 self.rows = self.rows[unique] 1271 self.cols = self.cols[unique] 1272 self.vals = np.bincount(indices, weights=self.vals) 1273 1274 def compress_csr(self): 1275 """ 1276 Compress rows, cols, vals / summing duplicates. Sort for csr format. 1277 """ 1278 _, unique, indices = np.unique( 1279 self.m*self.rows + self.cols, 1280 return_index=True, return_inverse=True) 1281 self.rows = self.rows[unique] 1282 self.cols = self.cols[unique] 1283 self.vals = np.bincount(indices, weights=self.vals) 1284 1285 def to_dense(self): 1286 """ 1287 Returns a dense matrix representing self. 1288 Mainly for debugging purposes. 1289 """ 1290 ret = np.zeros([self.n, self.m], dtype=np.float64) 1291 nvals = self.vals.size 1292 for i in range(nvals): 1293 ret[self.rows[i], self.cols[i]] += self.vals[i] 1294 return ret 1295 1296 def __str__(self): 1297 return self.to_dense().__str__() 1298 1299 @property 1300 def diag(self): 1301 """ 1302 Returns the (dense) vector of the diagonal elements. 1303 """ 1304 in_diag = (self.rows == self.cols) 1305 diag = np.zeros(min(self.n, self.n), dtype=np.float64) # default 0. 1306 diag[self.rows[in_diag]] = self.vals[in_diag] 1307 return diag 1308 1309 1310def _cg(A, b, x0=None, tol=1.e-10, maxiter=1000): 1311 """ 1312 Use Preconditioned Conjugate Gradient iteration to solve A x = b 1313 A simple Jacobi (diagonal) preconditionner is used. 1314 1315 Parameters 1316 ---------- 1317 A: _Sparse_Matrix_coo 1318 *A* must have been compressed before by compress_csc or 1319 compress_csr method. 1320 1321 b: array 1322 Right hand side of the linear system. 1323 1324 Returns 1325 ------- 1326 x: array. 1327 The converged solution. 1328 err: float 1329 The absolute error np.linalg.norm(A.dot(x) - b) 1330 1331 Other parameters 1332 ---------------- 1333 x0: array. 1334 Starting guess for the solution. 1335 tol: float. 1336 Tolerance to achieve. The algorithm terminates when the relative 1337 residual is below tol. 1338 maxiter: integer. 1339 Maximum number of iterations. Iteration will stop 1340 after maxiter steps even if the specified tolerance has not 1341 been achieved. 1342 """ 1343 n = b.size 1344 assert A.n == n 1345 assert A.m == n 1346 b_norm = np.linalg.norm(b) 1347 1348 # Jacobi pre-conditioner 1349 kvec = A.diag 1350 # For diag elem < 1e-6 we keep 1e-6. 1351 kvec = np.where(kvec > 1.e-6, kvec, 1.e-6) 1352 1353 # Initial guess 1354 if x0 is None: 1355 x = np.zeros(n) 1356 else: 1357 x = x0 1358 1359 r = b - A.dot(x) 1360 w = r/kvec 1361 1362 p = np.zeros(n) 1363 beta = 0.0 1364 rho = np.dot(r, w) 1365 k = 0 1366 1367 # Following C. T. Kelley 1368 while (np.sqrt(abs(rho)) > tol*b_norm) and (k < maxiter): 1369 p = w + beta*p 1370 z = A.dot(p) 1371 alpha = rho/np.dot(p, z) 1372 r = r - alpha*z 1373 w = r/kvec 1374 rhoold = rho 1375 rho = np.dot(r, w) 1376 x = x + alpha*p 1377 beta = rho/rhoold 1378 #err = np.linalg.norm(A.dot(x) - b) # absolute accuracy - not used 1379 k += 1 1380 err = np.linalg.norm(A.dot(x) - b) 1381 return x, err 1382 1383 1384# The following private functions: 1385# :func:`_inv22_vectorized` 1386# :func:`_safe_inv22_vectorized` 1387# :func:`_pseudo_inv22sym_vectorized` 1388# :func:`_prod_vectorized` 1389# :func:`_scalar_vectorized` 1390# :func:`_transpose_vectorized` 1391# :func:`_roll_vectorized` 1392# :func:`_to_matrix_vectorized` 1393# :func:`_extract_submatrices` 1394# provide fast numpy implementation of some standard operations on arrays of 1395# matrices - stored as (:, n_rows, n_cols)-shaped np.arrays. 1396def _inv22_vectorized(M): 1397 """ 1398 Inversion of arrays of (2,2) matrices. 1399 """ 1400 assert (M.ndim == 3) 1401 assert (M.shape[-2:] == (2, 2)) 1402 M_inv = np.empty_like(M) 1403 delta_inv = np.reciprocal(M[:, 0, 0]*M[:, 1, 1] - M[:, 0, 1]*M[:, 1, 0]) 1404 M_inv[:, 0, 0] = M[:, 1, 1]*delta_inv 1405 M_inv[:, 0, 1] = -M[:, 0, 1]*delta_inv 1406 M_inv[:, 1, 0] = -M[:, 1, 0]*delta_inv 1407 M_inv[:, 1, 1] = M[:, 0, 0]*delta_inv 1408 return M_inv 1409 1410 1411# Development note: Dealing with pathologic 'flat' triangles in the 1412# CubicTriInterpolator code and impact on (2,2)-matrix inversion functions 1413# :func:`_safe_inv22_vectorized` and :func:`_pseudo_inv22sym_vectorized`. 1414# 1415# Goals: 1416# 1) The CubicTriInterpolator should be able to handle flat or almost flat 1417# triangles without raising an error, 1418# 2) These degenerated triangles should have no impact on the automatic dof 1419# calculation (associated with null weight for the _DOF_estimator_geom and 1420# with null energy for the _DOF_estimator_min_E), 1421# 3) Linear patch test should be passed exactly on degenerated meshes, 1422# 4) Interpolation (with :meth:`_interpolate_single_key` or 1423# :meth:`_interpolate_multi_key`) shall be correctly handled even *inside* 1424# the pathologic triangles, to interact correctly with a TriRefiner class. 1425# 1426# Difficulties: 1427# Flat triangles have rank-deficient *J* (so-called jacobian matrix) and 1428# *metric* (the metric tensor = J x J.T). Computation of the local 1429# tangent plane is also problematic. 1430# 1431# Implementation: 1432# Most of the time, when computing the inverse of a rank-deficient matrix it 1433# is safe to simply return the null matrix (which is the implementation in 1434# :func:`_safe_inv22_vectorized`). This is because of point 2), itself 1435# enforced by: 1436# - null area hence null energy in :class:`_DOF_estimator_min_E` 1437# - angles close or equal to 0 or np.pi hence null weight in 1438# :class:`_DOF_estimator_geom`. 1439# Note that the function angle -> weight is continuous and maximum for an 1440# angle np.pi/2 (refer to :meth:`compute_geom_weights`) 1441# The exception is the computation of barycentric coordinates, which is done 1442# by inversion of the *metric* matrix. In this case, we need to compute a set 1443# of valid coordinates (1 among numerous possibilities), to ensure point 4). 1444# We benefit here from the symmetry of metric = J x J.T, which makes it easier 1445# to compute a pseudo-inverse in :func:`_pseudo_inv22sym_vectorized` 1446def _safe_inv22_vectorized(M): 1447 """ 1448 Inversion of arrays of (2,2) matrices, returns 0 for rank-deficient 1449 matrices. 1450 1451 *M* : array of (2,2) matrices to inverse, shape (n,2,2) 1452 """ 1453 assert M.ndim == 3 1454 assert M.shape[-2:] == (2, 2) 1455 M_inv = np.empty_like(M) 1456 prod1 = M[:, 0, 0]*M[:, 1, 1] 1457 delta = prod1 - M[:, 0, 1]*M[:, 1, 0] 1458 1459 # We set delta_inv to 0. in case of a rank deficient matrix ; a 1460 # rank-deficient input matrix *M* will lead to a null matrix in output 1461 rank2 = (np.abs(delta) > 1e-8*np.abs(prod1)) 1462 if np.all(rank2): 1463 # Normal 'optimized' flow. 1464 delta_inv = 1./delta 1465 else: 1466 # 'Pathologic' flow. 1467 delta_inv = np.zeros(M.shape[0]) 1468 delta_inv[rank2] = 1./delta[rank2] 1469 1470 M_inv[:, 0, 0] = M[:, 1, 1]*delta_inv 1471 M_inv[:, 0, 1] = -M[:, 0, 1]*delta_inv 1472 M_inv[:, 1, 0] = -M[:, 1, 0]*delta_inv 1473 M_inv[:, 1, 1] = M[:, 0, 0]*delta_inv 1474 return M_inv 1475 1476 1477def _pseudo_inv22sym_vectorized(M): 1478 """ 1479 Inversion of arrays of (2,2) SYMMETRIC matrices ; returns the 1480 (Moore-Penrose) pseudo-inverse for rank-deficient matrices. 1481 1482 In case M is of rank 1, we have M = trace(M) x P where P is the orthogonal 1483 projection on Im(M), and we return trace(M)^-1 x P == M / trace(M)**2 1484 In case M is of rank 0, we return the null matrix. 1485 1486 *M* : array of (2,2) matrices to inverse, shape (n,2,2) 1487 """ 1488 assert M.ndim == 3 1489 assert M.shape[-2:] == (2, 2) 1490 M_inv = np.empty_like(M) 1491 prod1 = M[:, 0, 0]*M[:, 1, 1] 1492 delta = prod1 - M[:, 0, 1]*M[:, 1, 0] 1493 rank2 = (np.abs(delta) > 1e-8*np.abs(prod1)) 1494 1495 if np.all(rank2): 1496 # Normal 'optimized' flow. 1497 M_inv[:, 0, 0] = M[:, 1, 1] / delta 1498 M_inv[:, 0, 1] = -M[:, 0, 1] / delta 1499 M_inv[:, 1, 0] = -M[:, 1, 0] / delta 1500 M_inv[:, 1, 1] = M[:, 0, 0] / delta 1501 else: 1502 # 'Pathologic' flow. 1503 # Here we have to deal with 2 sub-cases 1504 # 1) First sub-case: matrices of rank 2: 1505 delta = delta[rank2] 1506 M_inv[rank2, 0, 0] = M[rank2, 1, 1] / delta 1507 M_inv[rank2, 0, 1] = -M[rank2, 0, 1] / delta 1508 M_inv[rank2, 1, 0] = -M[rank2, 1, 0] / delta 1509 M_inv[rank2, 1, 1] = M[rank2, 0, 0] / delta 1510 # 2) Second sub-case: rank-deficient matrices of rank 0 and 1: 1511 rank01 = ~rank2 1512 tr = M[rank01, 0, 0] + M[rank01, 1, 1] 1513 tr_zeros = (np.abs(tr) < 1.e-8) 1514 sq_tr_inv = (1.-tr_zeros) / (tr**2+tr_zeros) 1515 #sq_tr_inv = 1. / tr**2 1516 M_inv[rank01, 0, 0] = M[rank01, 0, 0] * sq_tr_inv 1517 M_inv[rank01, 0, 1] = M[rank01, 0, 1] * sq_tr_inv 1518 M_inv[rank01, 1, 0] = M[rank01, 1, 0] * sq_tr_inv 1519 M_inv[rank01, 1, 1] = M[rank01, 1, 1] * sq_tr_inv 1520 1521 return M_inv 1522 1523 1524def _prod_vectorized(M1, M2): 1525 """ 1526 Matrix product between arrays of matrices, or a matrix and an array of 1527 matrices (*M1* and *M2*) 1528 """ 1529 sh1 = M1.shape 1530 sh2 = M2.shape 1531 assert len(sh1) >= 2 1532 assert len(sh2) >= 2 1533 assert sh1[-1] == sh2[-2] 1534 1535 ndim1 = len(sh1) 1536 t1_index = list(xrange(ndim1-2)) + [ndim1-1, ndim1-2] 1537 return np.sum(np.transpose(M1, t1_index)[..., np.newaxis] * 1538 M2[..., np.newaxis, :], -3) 1539 1540 1541def _scalar_vectorized(scalar, M): 1542 """ 1543 Scalar product between scalars and matrices. 1544 """ 1545 return scalar[:, np.newaxis, np.newaxis]*M 1546 1547 1548def _transpose_vectorized(M): 1549 """ 1550 Transposition of an array of matrices *M*. 1551 """ 1552 ndim = M.ndim 1553 assert ndim == 3 1554 return np.transpose(M, [0, ndim-1, ndim-2]) 1555 1556 1557def _roll_vectorized(M, roll_indices, axis): 1558 """ 1559 Rolls an array of matrices along an axis according to an array of indices 1560 *roll_indices* 1561 *axis* can be either 0 (rolls rows) or 1 (rolls columns). 1562 """ 1563 assert axis in [0, 1] 1564 ndim = M.ndim 1565 assert ndim == 3 1566 ndim_roll = roll_indices.ndim 1567 assert ndim_roll == 1 1568 sh = M.shape 1569 r, c = sh[-2:] 1570 assert sh[0] == roll_indices.shape[0] 1571 vec_indices = np.arange(sh[0], dtype=np.int32) 1572 1573 # Builds the rolled matrix 1574 M_roll = np.empty_like(M) 1575 if axis == 0: 1576 for ir in range(r): 1577 for ic in range(c): 1578 M_roll[:, ir, ic] = M[vec_indices, (-roll_indices+ir) % r, ic] 1579 elif axis == 1: 1580 for ir in range(r): 1581 for ic in range(c): 1582 M_roll[:, ir, ic] = M[vec_indices, ir, (-roll_indices+ic) % c] 1583 return M_roll 1584 1585 1586def _to_matrix_vectorized(M): 1587 """ 1588 Builds an array of matrices from individuals np.arrays of identical 1589 shapes. 1590 *M*: ncols-list of nrows-lists of shape sh. 1591 1592 Returns M_res np.array of shape (sh, nrow, ncols) so that: 1593 M_res[...,i,j] = M[i][j] 1594 """ 1595 assert isinstance(M, (tuple, list)) 1596 assert all([isinstance(item, (tuple, list)) for item in M]) 1597 c_vec = np.asarray([len(item) for item in M]) 1598 assert np.all(c_vec-c_vec[0] == 0) 1599 r = len(M) 1600 c = c_vec[0] 1601 M00 = np.asarray(M[0][0]) 1602 dt = M00.dtype 1603 sh = [M00.shape[0], r, c] 1604 M_ret = np.empty(sh, dtype=dt) 1605 for irow in range(r): 1606 for icol in range(c): 1607 M_ret[:, irow, icol] = np.asarray(M[irow][icol]) 1608 return M_ret 1609 1610 1611def _extract_submatrices(M, block_indices, block_size, axis): 1612 """ 1613 Extracts selected blocks of a matrices *M* depending on parameters 1614 *block_indices* and *block_size*. 1615 1616 Returns the array of extracted matrices *Mres* so that: 1617 M_res[...,ir,:] = M[(block_indices*block_size+ir), :] 1618 """ 1619 assert block_indices.ndim == 1 1620 assert axis in [0, 1] 1621 1622 r, c = M.shape 1623 if axis == 0: 1624 sh = [block_indices.shape[0], block_size, c] 1625 elif axis == 1: 1626 sh = [block_indices.shape[0], r, block_size] 1627 1628 dt = M.dtype 1629 M_res = np.empty(sh, dtype=dt) 1630 if axis == 0: 1631 for ir in range(block_size): 1632 M_res[:, ir, :] = M[(block_indices*block_size+ir), :] 1633 elif axis == 1: 1634 for ic in range(block_size): 1635 M_res[:, :, ic] = M[:, (block_indices*block_size+ic)] 1636 1637 return M_res 1638