1import math 2 3import numpy as np 4 5from yt.units.yt_array import YTArray 6 7prec_accum = { 8 int: np.int64, 9 np.int8: np.int64, 10 np.int16: np.int64, 11 np.int32: np.int64, 12 np.int64: np.int64, 13 np.uint8: np.uint64, 14 np.uint16: np.uint64, 15 np.uint32: np.uint64, 16 np.uint64: np.uint64, 17 float: np.float64, 18 np.float16: np.float64, 19 np.float32: np.float64, 20 np.float64: np.float64, 21 complex: np.complex128, 22 np.complex64: np.complex128, 23 np.complex128: np.complex128, 24 np.dtype("int"): np.int64, 25 np.dtype("int8"): np.int64, 26 np.dtype("int16"): np.int64, 27 np.dtype("int32"): np.int64, 28 np.dtype("int64"): np.int64, 29 np.dtype("uint8"): np.uint64, 30 np.dtype("uint16"): np.uint64, 31 np.dtype("uint32"): np.uint64, 32 np.dtype("uint64"): np.uint64, 33 np.dtype("float"): np.float64, 34 np.dtype("float16"): np.float64, 35 np.dtype("float32"): np.float64, 36 np.dtype("float64"): np.float64, 37 np.dtype("complex"): np.complex128, 38 np.dtype("complex64"): np.complex128, 39 np.dtype("complex128"): np.complex128, 40} 41 42 43def periodic_position(pos, ds): 44 r"""Assuming periodicity, find the periodic position within the domain. 45 46 Parameters 47 ---------- 48 pos : array 49 An array of floats. 50 51 ds : ~yt.data_objects.static_output.Dataset 52 A simulation static output. 53 54 Examples 55 -------- 56 >>> a = np.array([1.1, 0.5, 0.5]) 57 >>> data = {"Density": np.ones([32, 32, 32])} 58 >>> ds = load_uniform_grid(data, [32, 32, 32], 1.0) 59 >>> ppos = periodic_position(a, ds) 60 >>> ppos 61 array([ 0.1, 0.5, 0.5]) 62 """ 63 64 off = (pos - ds.domain_left_edge) % ds.domain_width 65 return ds.domain_left_edge + off 66 67 68def periodic_dist(a, b, period, periodicity=(True, True, True)): 69 r"""Find the Euclidean periodic distance between two sets of points. 70 71 Parameters 72 ---------- 73 a : array or list 74 Either an ndim long list of coordinates corresponding to a single point 75 or an (ndim, npoints) list of coordinates for many points in space. 76 77 b : array of list 78 Either an ndim long list of coordinates corresponding to a single point 79 or an (ndim, npoints) list of coordinates for many points in space. 80 81 period : float or array or list 82 If the volume is symmetrically periodic, this can be a single float, 83 otherwise an array or list of floats giving the periodic size of the 84 volume for each dimension. 85 86 periodicity : An ndim-element tuple of booleans 87 If an entry is true, the domain is assumed to be periodic along 88 that direction. 89 90 Examples 91 -------- 92 >>> a = [0.1, 0.1, 0.1] 93 >>> b = [0.9, 0, 9, 0.9] 94 >>> period = 1.0 95 >>> dist = periodic_dist(a, b, 1.0) 96 >>> dist 97 0.346410161514 98 """ 99 a = np.array(a) 100 b = np.array(b) 101 period = np.array(period) 102 103 if period.size == 1: 104 period = np.array([period, period, period]) 105 106 if a.shape != b.shape: 107 raise RuntimeError("Arrays must be the same shape.") 108 109 if period.shape != a.shape and len(a.shape) > 1: 110 n_tup = tuple(1 for i in range(a.ndim - 1)) 111 period = np.tile(np.reshape(period, (a.shape[0],) + n_tup), (1,) + a.shape[1:]) 112 elif len(a.shape) == 1: 113 a = np.reshape(a, (a.shape[0],) + (1, 1)) 114 b = np.reshape(b, (a.shape[0],) + (1, 1)) 115 period = np.reshape(period, (a.shape[0],) + (1, 1)) 116 117 c = np.empty((2,) + a.shape, dtype="float64") 118 c[0, :] = np.abs(a - b) 119 120 p_directions = [i for i, p in enumerate(periodicity) if p] 121 np_directions = [i for i, p in enumerate(periodicity) if not p] 122 for d in p_directions: 123 c[1, d, :] = period[d, :] - np.abs(a - b)[d, :] 124 for d in np_directions: 125 c[1, d, :] = c[0, d, :] 126 127 d = np.amin(c, axis=0) ** 2 128 r2 = d.sum(axis=0) 129 if r2.size == 1: 130 return np.sqrt(r2[0, 0]) 131 return np.sqrt(r2) 132 133 134def periodic_ray(start, end, left=None, right=None): 135 """ 136 periodic_ray(start, end, left=None, right=None) 137 138 Break up periodic ray into non-periodic segments. 139 Accepts start and end points of periodic ray as YTArrays. 140 Accepts optional left and right edges of periodic volume as YTArrays. 141 Returns a list of lists of coordinates, where each element of the 142 top-most list is a 2-list of start coords and end coords of the 143 non-periodic ray: 144 145 [[[x0start,y0start,z0start], [x0end, y0end, z0end]], 146 [[x1start,y1start,z1start], [x1end, y1end, z1end]], 147 ...,] 148 149 Parameters 150 ---------- 151 start : array 152 The starting coordinate of the ray. 153 end : array 154 The ending coordinate of the ray. 155 left : optional, array 156 The left corner of the periodic domain. If not given, an array 157 of zeros with same size as the starting coordinate us used. 158 right : optional, array 159 The right corner of the periodic domain. If not given, an array 160 of ones with same size as the starting coordinate us used. 161 162 Examples 163 -------- 164 >>> import yt 165 >>> start = yt.YTArray([0.5, 0.5, 0.5]) 166 >>> end = yt.YTArray([1.25, 1.25, 1.25]) 167 >>> periodic_ray(start, end) 168 [ 169 [ 170 YTArray([0.5, 0.5, 0.5]) (dimensionless), 171 YTArray([1., 1., 1.]) (dimensionless) 172 ], 173 [ 174 YTArray([0., 0., 0.]) (dimensionless), 175 YTArray([0.25, 0.25, 0.25]) (dimensionless) 176 ] 177 ] 178 """ 179 180 if left is None: 181 left = np.zeros(start.shape) 182 if right is None: 183 right = np.ones(start.shape) 184 dim = right - left 185 186 vector = end - start 187 wall = np.zeros_like(start) 188 close = np.zeros(start.shape, dtype=object) 189 190 left_bound = vector < 0 191 right_bound = vector > 0 192 no_bound = vector == 0.0 193 bound = vector != 0.0 194 195 wall[left_bound] = left[left_bound] 196 close[left_bound] = np.max 197 wall[right_bound] = right[right_bound] 198 close[right_bound] = np.min 199 wall[no_bound] = np.inf 200 close[no_bound] = np.min 201 202 segments = [] 203 this_start = start.copy() 204 this_end = end.copy() 205 t = 0.0 206 tolerance = 1e-6 207 while t < 1.0 - tolerance: 208 hit_left = (this_start <= left) & (vector < 0) 209 if (hit_left).any(): 210 this_start[hit_left] += dim[hit_left] 211 this_end[hit_left] += dim[hit_left] 212 hit_right = (this_start >= right) & (vector > 0) 213 if (hit_right).any(): 214 this_start[hit_right] -= dim[hit_right] 215 this_end[hit_right] -= dim[hit_right] 216 217 nearest = vector.unit_array * np.array( 218 [close[q]([this_end[q], wall[q]]) for q in range(start.size)] 219 ) 220 dt = ((nearest - this_start) / vector)[bound].min() 221 now = this_start + vector * dt 222 close_enough = np.abs(now - nearest) / np.abs(vector.max()) < 1e-10 223 now[close_enough] = nearest[close_enough] 224 segments.append([this_start.copy(), now.copy()]) 225 this_start = now.copy() 226 t += dt 227 228 return segments 229 230 231def euclidean_dist(a, b): 232 r"""Find the Euclidean distance between two points. 233 234 Parameters 235 ---------- 236 a : array or list 237 Either an ndim long list of coordinates corresponding to a single point 238 or an (ndim, npoints) list of coordinates for many points in space. 239 240 b : array or list 241 Either an ndim long list of coordinates corresponding to a single point 242 or an (ndim, npoints) list of coordinates for many points in space. 243 244 Examples 245 -------- 246 >>> a = [0.1, 0.1, 0.1] 247 >>> b = [0.9, 0, 9, 0.9] 248 >>> period = 1.0 249 >>> dist = euclidean_dist(a, b) 250 >>> dist 251 1.38564064606 252 253 """ 254 a = np.array(a) 255 b = np.array(b) 256 if a.shape != b.shape: 257 RuntimeError("Arrays must be the same shape.") 258 c = a.copy() 259 np.subtract(c, b, c) 260 np.power(c, 2, c) 261 c = c.sum(axis=0) 262 if isinstance(c, np.ndarray): 263 np.sqrt(c, c) 264 else: 265 # This happens if a and b only have one entry. 266 c = math.sqrt(c) 267 return c 268 269 270def rotate_vector_3D(a, dim, angle): 271 r"""Rotates the elements of an array around an axis by some angle. 272 273 Given an array of 3D vectors a, this rotates them around a coordinate axis 274 by a clockwise angle. An alternative way to think about it is the 275 coordinate axes are rotated counterclockwise, which changes the directions 276 of the vectors accordingly. 277 278 Parameters 279 ---------- 280 a : array 281 An array of 3D vectors with dimension Nx3. 282 283 dim : integer 284 A integer giving the axis around which the vectors will be rotated. 285 (x, y, z) = (0, 1, 2). 286 287 angle : float 288 The angle in radians through which the vectors will be rotated 289 clockwise. 290 291 Examples 292 -------- 293 >>> a = np.array([[1, 1, 0], [1, 0, 1], [0, 1, 1], [1, 1, 1], [3, 4, 5]]) 294 >>> b = rotate_vector_3D(a, 2, np.pi / 2) 295 >>> print(b) 296 [[ 1.00000000e+00 -1.00000000e+00 0.00000000e+00] 297 [ 6.12323400e-17 -1.00000000e+00 1.00000000e+00] 298 [ 1.00000000e+00 6.12323400e-17 1.00000000e+00] 299 [ 1.00000000e+00 -1.00000000e+00 1.00000000e+00] 300 [ 4.00000000e+00 -3.00000000e+00 5.00000000e+00]] 301 302 """ 303 mod = False 304 if len(a.shape) == 1: 305 mod = True 306 a = np.array([a]) 307 if a.shape[1] != 3: 308 raise ValueError("The second dimension of the array a must be == 3!") 309 if dim == 0: 310 R = np.array( 311 [ 312 [1, 0, 0], 313 [0, np.cos(angle), np.sin(angle)], 314 [0, -np.sin(angle), np.cos(angle)], 315 ] 316 ) 317 elif dim == 1: 318 R = np.array( 319 [ 320 [np.cos(angle), 0, -np.sin(angle)], 321 [0, 1, 0], 322 [np.sin(angle), 0, np.cos(angle)], 323 ] 324 ) 325 elif dim == 2: 326 R = np.array( 327 [ 328 [np.cos(angle), np.sin(angle), 0], 329 [-np.sin(angle), np.cos(angle), 0], 330 [0, 0, 1], 331 ] 332 ) 333 else: 334 raise ValueError("dim must be 0, 1, or 2!") 335 if mod: 336 return np.dot(R, a.T).T[0] 337 else: 338 return np.dot(R, a.T).T 339 340 341def modify_reference_frame(CoM, L, P=None, V=None): 342 r"""Rotates and translates data into a new reference frame to make 343 calculations easier. 344 345 This is primarily useful for calculations of halo data. 346 The data is translated into the center of mass frame. 347 Next, it is rotated such that the angular momentum vector for the data 348 is aligned with the z-axis. Put another way, if one calculates the angular 349 momentum vector on the data that comes out of this function, it will 350 always be along the positive z-axis. 351 If the center of mass is re-calculated, it will be at the origin. 352 353 Parameters 354 ---------- 355 CoM : array 356 The center of mass in 3D. 357 358 L : array 359 The angular momentum vector. 360 361 Optional 362 -------- 363 364 P : array 365 The positions of the data to be modified (i.e. particle or grid cell 366 positions). The array should be Nx3. 367 368 V : array 369 The velocities of the data to be modified (i.e. particle or grid cell 370 velocities). The array should be Nx3. 371 372 Returns 373 ------- 374 L : array 375 The angular momentum vector equal to [0, 0, 1] modulo machine error. 376 377 P : array 378 The modified positional data. Only returned if P is not None 379 380 V : array 381 The modified velocity data. Only returned if V is not None 382 383 Examples 384 -------- 385 >>> CoM = np.array([0.5, 0.5, 0.5]) 386 >>> L = np.array([1, 0, 0]) 387 >>> P = np.array([[1, 0.5, 0.5], [0, 0.5, 0.5], [0.5, 0.5, 0.5], [0, 0, 0]]) 388 >>> V = p.copy() 389 >>> LL, PP, VV = modify_reference_frame(CoM, L, P, V) 390 >>> LL 391 array([ 6.12323400e-17, 0.00000000e+00, 1.00000000e+00]) 392 >>> PP 393 array([[ 3.06161700e-17, 0.00000000e+00, 5.00000000e-01], 394 [ -3.06161700e-17, 0.00000000e+00, -5.00000000e-01], 395 [ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00], 396 [ 5.00000000e-01, -5.00000000e-01, -5.00000000e-01]]) 397 >>> VV 398 array([[ -5.00000000e-01, 5.00000000e-01, 1.00000000e+00], 399 [ -5.00000000e-01, 5.00000000e-01, 3.06161700e-17], 400 [ -5.00000000e-01, 5.00000000e-01, 5.00000000e-01], 401 [ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]]) 402 403 """ 404 # First translate the positions to center of mass reference frame. 405 if P is not None: 406 P = P - CoM 407 408 # is the L vector pointing in the Z direction? 409 if np.inner(L[:2], L[:2]) == 0.0: 410 # the reason why we need to explicitly check for the above 411 # is that formula is used in denominator 412 # this just checks if we need to flip the z axis or not 413 if L[2] < 0.0: 414 # this is just a simple flip in direction of the z axis 415 if P is not None: 416 P = -P 417 if V is not None: 418 V = -V 419 420 # return the values 421 if V is None and P is not None: 422 return L, P 423 elif P is None and V is not None: 424 return L, V 425 else: 426 return L, P, V 427 428 # Normal vector is not aligned with simulation Z axis 429 # Therefore we are going to have to apply a rotation 430 # Now find the angle between modified L and the x-axis. 431 LL = L.copy() 432 LL[2] = 0.0 433 theta = np.arccos(np.inner(LL, [1.0, 0.0, 0.0]) / np.inner(LL, LL) ** 0.5) 434 if L[1] < 0.0: 435 theta = -theta 436 # Now rotate all the position, velocity, and L vectors by this much around 437 # the z axis. 438 if P is not None: 439 P = rotate_vector_3D(P, 2, theta) 440 if V is not None: 441 V = rotate_vector_3D(V, 2, theta) 442 L = rotate_vector_3D(L, 2, theta) 443 # Now find the angle between L and the z-axis. 444 theta = np.arccos(np.inner(L, [0.0, 0.0, 1.0]) / np.inner(L, L) ** 0.5) 445 # This time we rotate around the y axis. 446 if P is not None: 447 P = rotate_vector_3D(P, 1, theta) 448 if V is not None: 449 V = rotate_vector_3D(V, 1, theta) 450 L = rotate_vector_3D(L, 1, theta) 451 452 # return the values 453 if V is None and P is not None: 454 return L, P 455 elif P is None and V is not None: 456 return L, V 457 else: 458 return L, P, V 459 460 461def compute_rotational_velocity(CoM, L, P, V): 462 r"""Computes the rotational velocity for some data around an axis. 463 464 This is primarily for halo computations. 465 Given some data, this computes the circular rotational velocity of each 466 point (particle) in reference to the axis defined by the angular momentum 467 vector. 468 This is accomplished by converting the reference frame of the center of 469 mass of the halo. 470 471 Parameters 472 ---------- 473 CoM : array 474 The center of mass in 3D. 475 476 L : array 477 The angular momentum vector. 478 479 P : array 480 The positions of the data to be modified (i.e. particle or grid cell 481 positions). The array should be Nx3. 482 483 V : array 484 The velocities of the data to be modified (i.e. particle or grid cell 485 velocities). The array should be Nx3. 486 487 Returns 488 ------- 489 v : array 490 An array N elements long that gives the circular rotational velocity 491 for each datum (particle). 492 493 Examples 494 -------- 495 >>> CoM = np.array([0, 0, 0]) 496 >>> L = np.array([0, 0, 1]) 497 >>> P = np.array([[1, 0, 0], [1, 1, 1], [0, 0, 1], [1, 1, 0]]) 498 >>> V = np.array([[0, 1, 10], [-1, -1, -1], [1, 1, 1], [1, -1, -1]]) 499 >>> circV = compute_rotational_velocity(CoM, L, P, V) 500 >>> circV 501 array([ 1. , 0. , 0. , 1.41421356]) 502 503 """ 504 # First we translate into the simple coordinates. 505 L, P, V = modify_reference_frame(CoM, L, P, V) 506 # Find the vector in the plane of the galaxy for each position point 507 # that is perpendicular to the radial vector. 508 radperp = np.cross([0, 0, 1], P) 509 # Find the component of the velocity along the radperp vector. 510 # Unf., I don't think there's a better way to do this. 511 res = np.empty(V.shape[0], dtype="float64") 512 for i, rp in enumerate(radperp): 513 temp = np.dot(rp, V[i]) / np.dot(rp, rp) * rp 514 res[i] = np.dot(temp, temp) ** 0.5 515 return res 516 517 518def compute_parallel_velocity(CoM, L, P, V): 519 r"""Computes the parallel velocity for some data around an axis. 520 521 This is primarily for halo computations. 522 Given some data, this computes the velocity component along the angular 523 momentum vector. 524 This is accomplished by converting the reference frame of the center of 525 mass of the halo. 526 527 Parameters 528 ---------- 529 CoM : array 530 The center of mass in 3D. 531 532 L : array 533 The angular momentum vector. 534 535 P : array 536 The positions of the data to be modified (i.e. particle or grid cell 537 positions). The array should be Nx3. 538 539 V : array 540 The velocities of the data to be modified (i.e. particle or grid cell 541 velocities). The array should be Nx3. 542 543 Returns 544 ------- 545 v : array 546 An array N elements long that gives the parallel velocity for 547 each datum (particle). 548 549 Examples 550 -------- 551 >>> CoM = np.array([0, 0, 0]) 552 >>> L = np.array([0, 0, 1]) 553 >>> P = np.array([[1, 0, 0], [1, 1, 1], [0, 0, 1], [1, 1, 0]]) 554 >>> V = np.array([[0, 1, 10], [-1, -1, -1], [1, 1, 1], [1, -1, -1]]) 555 >>> paraV = compute_parallel_velocity(CoM, L, P, V) 556 >>> paraV 557 array([10, -1, 1, -1]) 558 559 """ 560 # First we translate into the simple coordinates. 561 L, P, V = modify_reference_frame(CoM, L, P, V) 562 # And return just the z-axis velocities. 563 return V[:, 2] 564 565 566def compute_radial_velocity(CoM, L, P, V): 567 r"""Computes the radial velocity for some data around an axis. 568 569 This is primarily for halo computations. 570 Given some data, this computes the radial velocity component for the data. 571 This is accomplished by converting the reference frame of the center of 572 mass of the halo. 573 574 Parameters 575 ---------- 576 CoM : array 577 The center of mass in 3D. 578 579 L : array 580 The angular momentum vector. 581 582 P : array 583 The positions of the data to be modified (i.e. particle or grid cell 584 positions). The array should be Nx3. 585 586 V : array 587 The velocities of the data to be modified (i.e. particle or grid cell 588 velocities). The array should be Nx3. 589 590 Returns 591 ------- 592 v : array 593 An array N elements long that gives the radial velocity for 594 each datum (particle). 595 596 Examples 597 -------- 598 >>> CoM = np.array([0, 0, 0]) 599 >>> L = np.array([0, 0, 1]) 600 >>> P = np.array([[1, 0, 0], [1, 1, 1], [0, 0, 1], [1, 1, 0]]) 601 >>> V = np.array([[0, 1, 10], [-1, -1, -1], [1, 1, 1], [1, -1, -1]]) 602 >>> radV = compute_radial_velocity(CoM, L, P, V) 603 >>> radV 604 array([ 1. , 1.41421356 , 0. , 0.]) 605 606 """ 607 # First we translate into the simple coordinates. 608 L, P, V = modify_reference_frame(CoM, L, P, V) 609 # We find the tangential velocity by dotting the velocity vector 610 # with the cylindrical radial vector for this point. 611 # Unf., I don't think there's a better way to do this. 612 P[:, 2] = 0 613 res = np.empty(V.shape[0], dtype="float64") 614 for i, rad in enumerate(P): 615 temp = np.dot(rad, V[i]) / np.dot(rad, rad) * rad 616 res[i] = np.dot(temp, temp) ** 0.5 617 return res 618 619 620def compute_cylindrical_radius(CoM, L, P, V): 621 r"""Compute the radius for some data around an axis in cylindrical 622 coordinates. 623 624 This is primarily for halo computations. 625 Given some data, this computes the cylindrical radius for each point. 626 This is accomplished by converting the reference frame of the center of 627 mass of the halo. 628 629 Parameters 630 ---------- 631 CoM : array 632 The center of mass in 3D. 633 634 L : array 635 The angular momentum vector. 636 637 P : array 638 The positions of the data to be modified (i.e. particle or grid cell 639 positions). The array should be Nx3. 640 641 V : array 642 The velocities of the data to be modified (i.e. particle or grid cell 643 velocities). The array should be Nx3. 644 645 Returns 646 ------- 647 cyl_r : array 648 An array N elements long that gives the radial velocity for 649 each datum (particle). 650 651 Examples 652 -------- 653 >>> CoM = np.array([0, 0, 0]) 654 >>> L = np.array([0, 0, 1]) 655 >>> P = np.array([[1, 0, 0], [1, 1, 1], [0, 0, 1], [1, 1, 0]]) 656 >>> V = np.array([[0, 1, 10], [-1, -1, -1], [1, 1, 1], [1, -1, -1]]) 657 >>> cyl_r = compute_cylindrical_radius(CoM, L, P, V) 658 >>> cyl_r 659 array([ 1. , 1.41421356, 0. , 1.41421356]) 660 661 """ 662 # First we translate into the simple coordinates. 663 L, P, V = modify_reference_frame(CoM, L, P, V) 664 # Demote all the positions to the z=0 plane, which makes the distance 665 # calculation very easy. 666 P[:, 2] = 0 667 return np.sqrt((P * P).sum(axis=1)) 668 669 670def ortho_find(vec1): 671 r"""Find two complementary orthonormal vectors to a given vector. 672 673 For any given non-zero vector, there are infinite pairs of vectors 674 orthonormal to it. This function gives you one arbitrary pair from 675 that set along with the normalized version of the original vector. 676 677 Parameters 678 ---------- 679 vec1 : array_like 680 An array or list to represent a 3-vector. 681 682 Returns 683 ------- 684 vec1 : array 685 The original 3-vector array after having been normalized. 686 687 vec2 : array 688 A 3-vector array which is orthonormal to vec1. 689 690 vec3 : array 691 A 3-vector array which is orthonormal to vec1 and vec2. 692 693 Raises 694 ------ 695 ValueError 696 If input vector is the zero vector. 697 698 Notes 699 ----- 700 Our initial vector is `vec1` which consists of 3 components: `x1`, `y1`, 701 and `z1`. ortho_find determines a vector, `vec2`, which is orthonormal 702 to `vec1` by finding a vector which has a zero-value dot-product with 703 `vec1`. 704 705 .. math:: 706 707 vec1 \cdot vec2 = x_1 x_2 + y_1 y_2 + z_1 z_2 = 0 708 709 As a starting point, we arbitrarily choose `vec2` to have `x2` = 1, 710 `y2` = 0: 711 712 .. math:: 713 714 vec1 \cdot vec2 = x_1 + (z_1 z_2) = 0 715 716 \rightarrow z_2 = -(x_1 / z_1) 717 718 Of course, this will fail if `z1` = 0, in which case, let's say use 719 `z2` = 1 and `x2` = 0: 720 721 .. math:: 722 723 \rightarrow y_2 = -(z_1 / y_1) 724 725 Similarly, if `y1` = 0, this case will fail, in which case we use 726 `y2` = 1 and `z2` = 0: 727 728 .. math:: 729 730 \rightarrow x_2 = -(y_1 / x_1) 731 732 Since we don't allow `vec1` to be zero, all cases are accounted for. 733 734 Producing `vec3`, the complementary orthonormal vector to `vec1` and `vec2` 735 is accomplished by simply taking the cross product of `vec1` and `vec2`. 736 737 Examples 738 -------- 739 >>> a = [1.0, 2.0, 3.0] 740 >>> a, b, c = ortho_find(a) 741 >>> a 742 array([ 0.26726124, 0.53452248, 0.80178373]) 743 >>> b 744 array([ 0.9486833 , 0. , -0.31622777]) 745 >>> c 746 array([-0.16903085, 0.84515425, -0.50709255]) 747 """ 748 vec1 = np.array(vec1, dtype=np.float64) 749 # Normalize 750 norm = np.sqrt(np.vdot(vec1, vec1)) 751 if norm == 0: 752 raise ValueError("Zero vector used as input.") 753 vec1 /= norm 754 x1 = vec1[0] 755 y1 = vec1[1] 756 z1 = vec1[2] 757 if z1 != 0: 758 x2 = 1.0 759 y2 = 0.0 760 z2 = -(x1 / z1) 761 norm2 = (1.0 + z2 ** 2.0) ** (0.5) 762 elif y1 != 0: 763 x2 = 0.0 764 z2 = 1.0 765 y2 = -(z1 / y1) 766 norm2 = (1.0 + y2 ** 2.0) ** (0.5) 767 else: 768 y2 = 1.0 769 z2 = 0.0 770 x2 = -(y1 / x1) 771 norm2 = (1.0 + z2 ** 2.0) ** (0.5) 772 vec2 = np.array([x2, y2, z2]) 773 vec2 /= norm2 774 vec3 = np.cross(vec1, vec2) 775 return vec1, vec2, vec3 776 777 778def quartiles(a, axis=None, out=None, overwrite_input=False): 779 """ 780 Compute the quartile values (25% and 75%) along the specified axis 781 in the same way that the numpy.median calculates the median (50%) value 782 alone a specified axis. Check numpy.median for details, as it is 783 virtually the same algorithm. 784 785 Returns an array of the quartiles of the array elements [lower quartile, 786 upper quartile]. 787 788 Parameters 789 ---------- 790 a : array_like 791 Input array or object that can be converted to an array. 792 axis : {None, int}, optional 793 Axis along which the quartiles are computed. The default (axis=None) 794 is to compute the quartiles along a flattened version of the array. 795 out : ndarray, optional 796 Alternative output array in which to place the result. It must 797 have the same shape and buffer length as the expected output, 798 but the type (of the output) will be cast if necessary. 799 overwrite_input : {False, True}, optional 800 If True, then allow use of memory of input array (a) for 801 calculations. The input array will be modified by the call to 802 quartiles. This will save memory when you do not need to preserve 803 the contents of the input array. Treat the input as undefined, 804 but it will probably be fully or partially sorted. Default is 805 False. Note that, if `overwrite_input` is True and the input 806 is not already an ndarray, an error will be raised. 807 808 Returns 809 ------- 810 quartiles : ndarray 811 A new 2D array holding the result (unless `out` is specified, in 812 which case that array is returned instead). If the input contains 813 integers, or floats of smaller precision than 64, then the output 814 data-type is float64. Otherwise, the output data-type is the same 815 as that of the input. 816 817 See Also 818 -------- 819 numpy.median, numpy.mean, numpy.percentile 820 821 Notes 822 ----- 823 Given a vector V of length N, the quartiles of V are the 25% and 75% values 824 of a sorted copy of V, ``V_sorted`` - i.e., ``V_sorted[(N-1)/4]`` and 825 ``3*V_sorted[(N-1)/4]``, when N is odd. When N is even, it is the average 826 of the two values bounding these values of ``V_sorted``. 827 828 Examples 829 -------- 830 >>> a = np.arange(100).reshape(10, 10) 831 >>> a 832 array([[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 833 [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], 834 [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], 835 [30, 31, 32, 33, 34, 35, 36, 37, 38, 39], 836 [40, 41, 42, 43, 44, 45, 46, 47, 48, 49], 837 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59], 838 [60, 61, 62, 63, 64, 65, 66, 67, 68, 69], 839 [70, 71, 72, 73, 74, 75, 76, 77, 78, 79], 840 [80, 81, 82, 83, 84, 85, 86, 87, 88, 89], 841 [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]]) 842 >>> mu.quartiles(a) 843 array([ 24.5, 74.5]) 844 >>> mu.quartiles(a, axis=0) 845 array([[ 15., 16., 17., 18., 19., 20., 21., 22., 23., 24.], 846 [ 65., 66., 67., 68., 69., 70., 71., 72., 73., 74.]]) 847 >>> mu.quartiles(a, axis=1) 848 array([[ 1.5, 11.5, 21.5, 31.5, 41.5, 51.5, 61.5, 71.5, 81.5, 849 91.5], 850 [ 6.5, 16.5, 26.5, 36.5, 46.5, 56.5, 66.5, 76.5, 86.5, 851 96.5]]) 852 """ 853 if overwrite_input: 854 if axis is None: 855 a_sorted = sorted(a.ravel()) 856 else: 857 a.sort(axis=axis) 858 a_sorted = a 859 else: 860 a_sorted = np.sort(a, axis=axis) 861 if axis is None: 862 axis = 0 863 indexer = [slice(None)] * a_sorted.ndim 864 indices = [int(a_sorted.shape[axis] / 4), int(a_sorted.shape[axis] * 0.75)] 865 result = [] 866 for index in indices: 867 if a_sorted.shape[axis] % 2 == 1: 868 # index with slice to allow mean (below) to work 869 indexer[axis] = slice(index, index + 1) 870 else: 871 indexer[axis] = slice(index - 1, index + 1) 872 # special cases for small arrays 873 if a_sorted.shape[axis] == 2: 874 # index with slice to allow mean (below) to work 875 indexer[axis] = slice(index, index + 1) 876 # Use mean in odd and even case to coerce data type 877 # and check, use out array. 878 result.append(np.mean(a_sorted[tuple(indexer)], axis=axis, out=out)) 879 return np.array(result) 880 881 882def get_perspective_matrix(fovy, aspect, z_near, z_far): 883 """ 884 Given a field of view in radians, an aspect ratio, and a near 885 and far plane distance, this routine computes the transformation matrix 886 corresponding to perspective projection using homogenous coordinates. 887 888 Parameters 889 ---------- 890 fovy : scalar 891 The angle in degrees of the field of view. 892 893 aspect : scalar 894 The aspect ratio of width / height for the projection. 895 896 z_near : scalar 897 The distance of the near plane from the camera. 898 899 z_far : scalar 900 The distance of the far plane from the camera. 901 902 Returns 903 ------- 904 persp_matrix : ndarray 905 A new 4x4 2D array. Represents a perspective transformation 906 in homogeneous coordinates. Note that this matrix does not 907 actually perform the projection. After multiplying a 4D 908 vector of the form (x_0, y_0, z_0, 1.0), the point will be 909 transformed to some (x_1, y_1, z_1, w). The final projection 910 is applied by performing a divide by w, that is 911 (x_1/w, y_1/w, z_1/w, w/w). The matrix uses a row-major 912 ordering, rather than the column major ordering typically 913 used by OpenGL. 914 915 Notes 916 ----- 917 The usage of 4D homogeneous coordinates is for OpenGL and GPU 918 hardware that automatically performs the divide by w operation. 919 See the following for more details about the OpenGL perspective matrices. 920 921 https://www.tomdalling.com/blog/modern-opengl/explaining-homogenous-coordinates-and-projective-geometry/ 922 http://www.songho.ca/opengl/gl_projectionmatrix.html 923 924 """ 925 926 tan_half_fovy = np.tan(np.radians(fovy) / 2) 927 928 result = np.zeros((4, 4), dtype="float32", order="C") 929 # result[0][0] = 1 / (aspect * tan_half_fovy) 930 # result[1][1] = 1 / tan_half_fovy 931 # result[2][2] = - (z_far + z_near) / (z_far - z_near) 932 # result[3][2] = -1 933 # result[2][3] = -(2 * z_far * z_near) / (z_far - z_near) 934 935 f = z_far 936 n = z_near 937 938 t = tan_half_fovy * n 939 b = -t * aspect 940 r = t * aspect 941 l = -t * aspect 942 943 result[0][0] = (2 * n) / (r - l) 944 result[2][0] = (r + l) / (r - l) 945 result[1][1] = (2 * n) / (t - b) 946 result[1][2] = (t + b) / (t - b) 947 result[2][2] = -(f + n) / (f - n) 948 result[2][3] = -2 * f * n / (f - n) 949 result[3][2] = -1 950 951 return result 952 953 954def get_orthographic_matrix(maxr, aspect, z_near, z_far): 955 """ 956 Given a field of view in radians, an aspect ratio, and a near 957 and far plane distance, this routine computes the transformation matrix 958 corresponding to perspective projection using homogenous coordinates. 959 960 Parameters 961 ---------- 962 maxr : scalar 963 should be ``max(|x|, |y|)`` 964 965 aspect : scalar 966 The aspect ratio of width / height for the projection. 967 968 z_near : scalar 969 The distance of the near plane from the camera. 970 971 z_far : scalar 972 The distance of the far plane from the camera. 973 974 Returns 975 ------- 976 persp_matrix : ndarray 977 A new 4x4 2D array. Represents a perspective transformation 978 in homogeneous coordinates. Note that this matrix does not 979 actually perform the projection. After multiplying a 4D 980 vector of the form (x_0, y_0, z_0, 1.0), the point will be 981 transformed to some (x_1, y_1, z_1, w). The final projection 982 is applied by performing a divide by w, that is 983 (x_1/w, y_1/w, z_1/w, w/w). The matrix uses a row-major 984 ordering, rather than the column major ordering typically 985 used by OpenGL. 986 987 Notes 988 ----- 989 The usage of 4D homogeneous coordinates is for OpenGL and GPU 990 hardware that automatically performs the divide by w operation. 991 See the following for more details about the OpenGL perspective matrices. 992 993 http://www.scratchapixel.com/lessons/3d-basic-rendering/perspective-and-orthographic-projection-matrix/orthographic-projection-matrix 994 https://www.tomdalling.com/blog/modern-opengl/explaining-homogenous-coordinates-and-projective-geometry/ 995 http://www.songho.ca/opengl/gl_projectionmatrix.html 996 997 """ 998 999 r = maxr * aspect 1000 t = maxr 1001 l = -r 1002 b = -t 1003 1004 result = np.zeros((4, 4), dtype="float32", order="C") 1005 result[0][0] = 2.0 / (r - l) 1006 result[1][1] = 2.0 / (t - b) 1007 result[2][2] = -2.0 / (z_far - z_near) 1008 result[3][3] = 1 1009 1010 result[3][0] = -(r + l) / (r - l) 1011 result[3][1] = -(t + b) / (t - b) 1012 result[3][2] = -(z_far + z_near) / (z_far - z_near) 1013 1014 return result 1015 1016 1017def get_lookat_matrix(eye, center, up): 1018 """ 1019 Given the position of a camera, the point it is looking at, and 1020 an up-direction. Computes the lookat matrix that moves all vectors 1021 such that the camera is at the origin of the coordinate system, 1022 looking down the z-axis. 1023 1024 Parameters 1025 ---------- 1026 eye : array_like 1027 The position of the camera. Must be 3D. 1028 1029 center : array_like 1030 The location that the camera is looking at. Must be 3D. 1031 1032 up : array_like 1033 The direction that is considered up for the camera. Must be 1034 3D. 1035 1036 Returns 1037 ------- 1038 lookat_matrix : ndarray 1039 A new 4x4 2D array in homogeneous coordinates. This matrix 1040 moves all vectors in the same way required to move the camera 1041 to the origin of the coordinate system, with it pointing down 1042 the negative z-axis. 1043 1044 """ 1045 1046 eye = np.array(eye) 1047 center = np.array(center) 1048 up = np.array(up) 1049 1050 f = (center - eye) / np.linalg.norm(center - eye) 1051 s = np.cross(f, up) / np.linalg.norm(np.cross(f, up)) 1052 u = np.cross(s, f) 1053 1054 result = np.zeros((4, 4), dtype="float32", order="C") 1055 1056 result[0][0] = s[0] 1057 result[0][1] = s[1] 1058 result[0][2] = s[2] 1059 result[1][0] = u[0] 1060 result[1][1] = u[1] 1061 result[1][2] = u[2] 1062 result[2][0] = -f[0] 1063 result[2][1] = -f[1] 1064 result[2][2] = -f[2] 1065 result[0][3] = -np.dot(s, eye) 1066 result[1][3] = -np.dot(u, eye) 1067 result[2][3] = np.dot(f, eye) 1068 result[3][3] = 1.0 1069 return result 1070 1071 1072def get_translate_matrix(dx, dy, dz): 1073 """ 1074 Given a movement amount for each coordinate, creates a translation 1075 matrix that moves the vector by each amount. 1076 1077 Parameters 1078 ---------- 1079 dx : scalar 1080 A translation amount for the x-coordinate 1081 1082 dy : scalar 1083 A translation amount for the y-coordinate 1084 1085 dz : scalar 1086 A translation amount for the z-coordinate 1087 1088 Returns 1089 ------- 1090 trans_matrix : ndarray 1091 A new 4x4 2D array. Represents a translation by dx, dy 1092 and dz in each coordinate respectively. 1093 """ 1094 result = np.zeros((4, 4), dtype="float32", order="C") 1095 1096 result[0][0] = 1.0 1097 result[1][1] = 1.0 1098 result[2][2] = 1.0 1099 result[3][3] = 1.0 1100 1101 result[0][3] = dx 1102 result[1][3] = dy 1103 result[2][3] = dz 1104 1105 return result 1106 1107 1108def get_scale_matrix(dx, dy, dz): 1109 """ 1110 Given a scaling factor for each coordinate, returns a matrix that 1111 corresponds to the given scaling amounts. 1112 1113 Parameters 1114 ---------- 1115 dx : scalar 1116 A scaling factor for the x-coordinate. 1117 1118 dy : scalar 1119 A scaling factor for the y-coordinate. 1120 1121 dz : scalar 1122 A scaling factor for the z-coordinate. 1123 1124 Returns 1125 ------- 1126 scale_matrix : ndarray 1127 A new 4x4 2D array. Represents a scaling by dx, dy, and dz 1128 in each coordinate respectively. 1129 """ 1130 result = np.zeros((4, 4), dtype="float32", order="C") 1131 1132 result[0][0] = dx 1133 result[1][1] = dy 1134 result[2][2] = dz 1135 result[3][3] = 1 1136 1137 return result 1138 1139 1140def get_rotation_matrix(theta, rot_vector): 1141 """ 1142 Given an angle theta and a 3D vector rot_vector, this routine 1143 computes the rotation matrix corresponding to rotating theta 1144 radians about rot_vector. 1145 1146 Parameters 1147 ---------- 1148 theta : scalar 1149 The angle in radians. 1150 1151 rot_vector : array_like 1152 The axis of rotation. Must be 3D. 1153 1154 Returns 1155 ------- 1156 rot_matrix : ndarray 1157 A new 3x3 2D array. This is the representation of a 1158 rotation of theta radians about rot_vector in the simulation 1159 box coordinate frame 1160 1161 See Also 1162 -------- 1163 ortho_find 1164 1165 Examples 1166 -------- 1167 >>> a = [0, 1, 0] 1168 >>> theta = 0.785398163 # pi/4 1169 >>> rot = mu.get_rotation_matrix(theta, a) 1170 >>> rot 1171 array([[ 0.70710678, 0. , 0.70710678], 1172 [ 0. , 1. , 0. ], 1173 [-0.70710678, 0. , 0.70710678]]) 1174 >>> np.dot(rot, a) 1175 array([ 0., 1., 0.]) 1176 # since a is an eigenvector by construction 1177 >>> np.dot(rot, [1, 0, 0]) 1178 array([ 0.70710678, 0. , -0.70710678]) 1179 """ 1180 1181 ux = rot_vector[0] 1182 uy = rot_vector[1] 1183 uz = rot_vector[2] 1184 cost = np.cos(theta) 1185 sint = np.sin(theta) 1186 1187 R = np.array( 1188 [ 1189 [ 1190 cost + ux ** 2 * (1 - cost), 1191 ux * uy * (1 - cost) - uz * sint, 1192 ux * uz * (1 - cost) + uy * sint, 1193 ], 1194 [ 1195 uy * ux * (1 - cost) + uz * sint, 1196 cost + uy ** 2 * (1 - cost), 1197 uy * uz * (1 - cost) - ux * sint, 1198 ], 1199 [ 1200 uz * ux * (1 - cost) - uy * sint, 1201 uz * uy * (1 - cost) + ux * sint, 1202 cost + uz ** 2 * (1 - cost), 1203 ], 1204 ] 1205 ) 1206 1207 return R 1208 1209 1210def quaternion_mult(q1, q2): 1211 """ 1212 1213 Multiply two quaternions. The inputs are 4-component numpy arrays 1214 in the order [w, x, y, z]. 1215 1216 """ 1217 w = q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] - q1[3] * q2[3] 1218 x = q1[0] * q2[1] + q1[1] * q2[0] + q1[2] * q2[3] - q1[3] * q2[2] 1219 y = q1[0] * q2[2] + q1[2] * q2[0] + q1[3] * q2[1] - q1[1] * q2[3] 1220 z = q1[0] * q2[3] + q1[3] * q2[0] + q1[1] * q2[2] - q1[2] * q2[1] 1221 return np.array([w, x, y, z]) 1222 1223 1224def quaternion_to_rotation_matrix(quaternion): 1225 """ 1226 1227 This converts a quaternion representation of on orientation to 1228 a rotation matrix. The input is a 4-component numpy array in 1229 the order [w, x, y, z], and the output is a 3x3 matrix stored 1230 as a 2D numpy array. We follow the approach in 1231 "3D Math Primer for Graphics and Game Development" by 1232 Dunn and Parberry. 1233 1234 """ 1235 1236 w = quaternion[0] 1237 x = quaternion[1] 1238 y = quaternion[2] 1239 z = quaternion[3] 1240 1241 R = np.empty((3, 3), dtype=np.float64) 1242 1243 R[0][0] = 1.0 - 2.0 * y ** 2 - 2.0 * z ** 2 1244 R[0][1] = 2.0 * x * y + 2.0 * w * z 1245 R[0][2] = 2.0 * x * z - 2.0 * w * y 1246 1247 R[1][0] = 2.0 * x * y - 2.0 * w * z 1248 R[1][1] = 1.0 - 2.0 * x ** 2 - 2.0 * z ** 2 1249 R[1][2] = 2.0 * y * z + 2.0 * w * x 1250 1251 R[2][0] = 2.0 * x * z + 2.0 * w * y 1252 R[2][1] = 2.0 * y * z - 2.0 * w * x 1253 R[2][2] = 1.0 - 2.0 * x ** 2 - 2.0 * y ** 2 1254 1255 return R 1256 1257 1258def rotation_matrix_to_quaternion(rot_matrix): 1259 """ 1260 1261 Convert a rotation matrix-based representation of an 1262 orientation to a quaternion. The input should be a 1263 3x3 rotation matrix, while the output will be a 1264 4-component numpy array. We follow the approach in 1265 "3D Math Primer for Graphics and Game Development" by 1266 Dunn and Parberry. 1267 1268 """ 1269 m11 = rot_matrix[0][0] 1270 m12 = rot_matrix[0][1] 1271 m13 = rot_matrix[0][2] 1272 m21 = rot_matrix[1][0] 1273 m22 = rot_matrix[1][1] 1274 m23 = rot_matrix[1][2] 1275 m31 = rot_matrix[2][0] 1276 m32 = rot_matrix[2][1] 1277 m33 = rot_matrix[2][2] 1278 1279 four_w_squared_minus_1 = m11 + m22 + m33 1280 four_x_squared_minus_1 = m11 - m22 - m33 1281 four_y_squared_minus_1 = m22 - m11 - m33 1282 four_z_squared_minus_1 = m33 - m11 - m22 1283 max_index = 0 1284 four_max_squared_minus_1 = four_w_squared_minus_1 1285 if four_x_squared_minus_1 > four_max_squared_minus_1: 1286 four_max_squared_minus_1 = four_x_squared_minus_1 1287 max_index = 1 1288 if four_y_squared_minus_1 > four_max_squared_minus_1: 1289 four_max_squared_minus_1 = four_y_squared_minus_1 1290 max_index = 2 1291 if four_z_squared_minus_1 > four_max_squared_minus_1: 1292 four_max_squared_minus_1 = four_z_squared_minus_1 1293 max_index = 3 1294 1295 max_val = 0.5 * np.sqrt(four_max_squared_minus_1 + 1.0) 1296 mult = 0.25 / max_val 1297 1298 if max_index == 0: 1299 w = max_val 1300 x = (m23 - m32) * mult 1301 y = (m31 - m13) * mult 1302 z = (m12 - m21) * mult 1303 elif max_index == 1: 1304 x = max_val 1305 w = (m23 - m32) * mult 1306 y = (m12 + m21) * mult 1307 z = (m31 + m13) * mult 1308 elif max_index == 2: 1309 y = max_val 1310 w = (m31 - m13) * mult 1311 x = (m12 + m21) * mult 1312 z = (m23 + m32) * mult 1313 elif max_index == 3: 1314 z = max_val 1315 w = (m12 - m21) * mult 1316 x = (m31 + m13) * mult 1317 y = (m23 + m32) * mult 1318 1319 return np.array([w, x, y, z]) 1320 1321 1322def get_sph_r(coords): 1323 # The spherical coordinates radius is simply the magnitude of the 1324 # coordinate vector. 1325 1326 return np.sqrt(np.sum(coords ** 2, axis=0)) 1327 1328 1329def resize_vector(vector, vector_array): 1330 if len(vector_array.shape) == 4: 1331 res_vector = np.resize(vector, (3, 1, 1, 1)) 1332 else: 1333 res_vector = np.resize(vector, (3, 1)) 1334 return res_vector 1335 1336 1337def normalize_vector(vector): 1338 # this function normalizes 1339 # an input vector 1340 1341 L2 = np.atleast_1d(np.linalg.norm(vector)) 1342 L2[L2 == 0] = 1.0 1343 vector = vector / L2 1344 return vector 1345 1346 1347def get_sph_theta(coords, normal): 1348 # The angle (theta) with respect to the normal (J), is the arccos 1349 # of the dot product of the normal with the normalized coordinate 1350 # vector. 1351 1352 res_normal = resize_vector(normal, coords) 1353 1354 # check if the normal vector is normalized 1355 # since arccos requires the vector to be normalised 1356 res_normal = normalize_vector(res_normal) 1357 1358 tile_shape = [1] + list(coords.shape)[1:] 1359 1360 J = np.tile(res_normal, tile_shape) 1361 1362 JdotCoords = np.sum(J * coords, axis=0) 1363 1364 with np.errstate(invalid="ignore"): 1365 ret = np.arccos(JdotCoords / np.sqrt(np.sum(coords ** 2, axis=0))) 1366 1367 ret[np.isnan(ret)] = 0 1368 1369 return ret 1370 1371 1372def get_sph_phi(coords, normal): 1373 # We have freedom with respect to what axis (xprime) to define 1374 # the disk angle. Here I've chosen to use the axis that is 1375 # perpendicular to the normal and the y-axis. When normal == 1376 # y-hat, then set xprime = z-hat. With this definition, when 1377 # normal == z-hat (as is typical), then xprime == x-hat. 1378 # 1379 # The angle is then given by the arctan of the ratio of the 1380 # yprime-component and the xprime-component of the coordinate 1381 # vector. 1382 1383 normal = normalize_vector(normal) 1384 (zprime, xprime, yprime) = ortho_find(normal) 1385 1386 res_xprime = resize_vector(xprime, coords) 1387 res_yprime = resize_vector(yprime, coords) 1388 1389 tile_shape = [1] + list(coords.shape)[1:] 1390 Jx = np.tile(res_xprime, tile_shape) 1391 Jy = np.tile(res_yprime, tile_shape) 1392 1393 Px = np.sum(Jx * coords, axis=0) 1394 Py = np.sum(Jy * coords, axis=0) 1395 1396 return np.arctan2(Py, Px) 1397 1398 1399def get_cyl_r(coords, normal): 1400 # The cross product of the normal (J) with a coordinate vector 1401 # gives a vector of magnitude equal to the cylindrical radius. 1402 1403 res_normal = resize_vector(normal, coords) 1404 res_normal = normalize_vector(res_normal) 1405 1406 tile_shape = [1] + list(coords.shape)[1:] 1407 J = np.tile(res_normal, tile_shape) 1408 1409 JcrossCoords = np.cross(J, coords, axisa=0, axisb=0, axisc=0) 1410 return np.sqrt(np.sum(JcrossCoords ** 2, axis=0)) 1411 1412 1413def get_cyl_z(coords, normal): 1414 # The dot product of the normal (J) with the coordinate vector 1415 # gives the cylindrical height. 1416 1417 res_normal = resize_vector(normal, coords) 1418 res_normal = normalize_vector(res_normal) 1419 1420 tile_shape = [1] + list(coords.shape)[1:] 1421 J = np.tile(res_normal, tile_shape) 1422 1423 return np.sum(J * coords, axis=0) 1424 1425 1426def get_cyl_theta(coords, normal): 1427 # This is identical to the spherical phi component 1428 1429 return get_sph_phi(coords, normal) 1430 1431 1432def get_cyl_r_component(vectors, theta, normal): 1433 # The r of a vector is the vector dotted with rhat 1434 1435 normal = normalize_vector(normal) 1436 (zprime, xprime, yprime) = ortho_find(normal) 1437 1438 res_xprime = resize_vector(xprime, vectors) 1439 res_yprime = resize_vector(yprime, vectors) 1440 1441 tile_shape = [1] + list(vectors.shape)[1:] 1442 Jx = np.tile(res_xprime, tile_shape) 1443 Jy = np.tile(res_yprime, tile_shape) 1444 1445 rhat = Jx * np.cos(theta) + Jy * np.sin(theta) 1446 1447 return np.sum(vectors * rhat, axis=0) 1448 1449 1450def get_cyl_theta_component(vectors, theta, normal): 1451 # The theta component of a vector is the vector dotted with thetahat 1452 normal = normalize_vector(normal) 1453 (zprime, xprime, yprime) = ortho_find(normal) 1454 1455 res_xprime = resize_vector(xprime, vectors) 1456 res_yprime = resize_vector(yprime, vectors) 1457 1458 tile_shape = [1] + list(vectors.shape)[1:] 1459 Jx = np.tile(res_xprime, tile_shape) 1460 Jy = np.tile(res_yprime, tile_shape) 1461 1462 thetahat = -Jx * np.sin(theta) + Jy * np.cos(theta) 1463 1464 return np.sum(vectors * thetahat, axis=0) 1465 1466 1467def get_cyl_z_component(vectors, normal): 1468 # The z component of a vector is the vector dotted with zhat 1469 normal = normalize_vector(normal) 1470 (zprime, xprime, yprime) = ortho_find(normal) 1471 1472 res_zprime = resize_vector(zprime, vectors) 1473 1474 tile_shape = [1] + list(vectors.shape)[1:] 1475 zhat = np.tile(res_zprime, tile_shape) 1476 1477 return np.sum(vectors * zhat, axis=0) 1478 1479 1480def get_sph_r_component(vectors, theta, phi, normal): 1481 # The r component of a vector is the vector dotted with rhat 1482 normal = normalize_vector(normal) 1483 (zprime, xprime, yprime) = ortho_find(normal) 1484 1485 res_xprime = resize_vector(xprime, vectors) 1486 res_yprime = resize_vector(yprime, vectors) 1487 res_zprime = resize_vector(zprime, vectors) 1488 1489 tile_shape = [1] + list(vectors.shape)[1:] 1490 1491 Jx, Jy, Jz = ( 1492 YTArray(np.tile(rprime, tile_shape), "") 1493 for rprime in (res_xprime, res_yprime, res_zprime) 1494 ) 1495 1496 rhat = ( 1497 Jx * np.sin(theta) * np.cos(phi) 1498 + Jy * np.sin(theta) * np.sin(phi) 1499 + Jz * np.cos(theta) 1500 ) 1501 1502 return np.sum(vectors * rhat, axis=0) 1503 1504 1505def get_sph_phi_component(vectors, phi, normal): 1506 # The phi component of a vector is the vector dotted with phihat 1507 normal = normalize_vector(normal) 1508 (zprime, xprime, yprime) = ortho_find(normal) 1509 1510 res_xprime = resize_vector(xprime, vectors) 1511 res_yprime = resize_vector(yprime, vectors) 1512 1513 tile_shape = [1] + list(vectors.shape)[1:] 1514 Jx = YTArray(np.tile(res_xprime, tile_shape), "") 1515 Jy = YTArray(np.tile(res_yprime, tile_shape), "") 1516 1517 phihat = -Jx * np.sin(phi) + Jy * np.cos(phi) 1518 1519 return np.sum(vectors * phihat, axis=0) 1520 1521 1522def get_sph_theta_component(vectors, theta, phi, normal): 1523 # The theta component of a vector is the vector dotted with thetahat 1524 normal = normalize_vector(normal) 1525 (zprime, xprime, yprime) = ortho_find(normal) 1526 1527 res_xprime = resize_vector(xprime, vectors) 1528 res_yprime = resize_vector(yprime, vectors) 1529 res_zprime = resize_vector(zprime, vectors) 1530 1531 tile_shape = [1] + list(vectors.shape)[1:] 1532 Jx, Jy, Jz = ( 1533 YTArray(np.tile(rprime, tile_shape), "") 1534 for rprime in (res_xprime, res_yprime, res_zprime) 1535 ) 1536 1537 thetahat = ( 1538 Jx * np.cos(theta) * np.cos(phi) 1539 + Jy * np.cos(theta) * np.sin(phi) 1540 - Jz * np.sin(theta) 1541 ) 1542 1543 return np.sum(vectors * thetahat, axis=0) 1544