Source code for geoana.spatial

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
from .utils import mkvc


[docs]def cylindrical_2_cartesian(grid, vec=None): """ Take a grid or vector (if provided) defined in cylindrical coordinates :math:`(r, \\theta, z)` and transform it to cartesian coordinates, :math:`(x, y, z)`. **Required** :param numpy.ndarray grid: grid in cylindrical coordinates :math:`(r, \\theta, z)` **Optional** :param numpy.ndarray vec: (optional) vector defined in cylindrical coordinates **Returns** :return: grid or vector (if provided) in cartesian coordinates :math:`(x, y, z)` :rtype: numpy.ndarray """ grid = np.atleast_2d(grid) if vec is None: return np.hstack([ mkvc(grid[:, 0]*np.cos(grid[:, 1]), 2), mkvc(grid[:, 0]*np.sin(grid[:, 1]), 2), mkvc(grid[:, 2], 2) ]) if len(vec.shape) == 1 or vec.shape[1] == 1: vec = vec.reshape(grid.shape, order='F') x = vec[:, 0] * np.cos(grid[:, 1]) - vec[:, 1] * np.sin(grid[:, 1]) y = vec[:, 0] * np.sin(grid[:, 1]) + vec[:, 1] * np.cos(grid[:, 1]) newvec = [x, y] if grid.shape[1] == 3: z = vec[:, 2] newvec += [z] return np.vstack(newvec).T
[docs]def cartesian_2_cylindrical(grid, vec=None): """ Takes a grid or vector (if provided) defined in cartesian coordinates :math:`(x, y, z)` and transform it to cylindrical coordinates, :math:`(r, \\theta, z)`. **Required** :param numpy.ndarray grid: grid in cartesian coordinates :math:`(x, y, z)` **Optional** :param numpy.ndarray vec: (optional) vector defined in cartesian coordinates **Returns** :return: grid or vector (if provided) in cylindrical coordinates :math:`(r, \\theta, z)` :rtype: numpy.ndarray """ grid = np.atleast_2d(grid) if vec is None: return np.hstack([ mkvc(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), 2), mkvc(np.arctan2(grid[:, 1], grid[:, 0]), 2), mkvc(grid[:, 2], 2) ]) if len(vec.shape) == 1 or vec.shape[1] == 1: vec = vec.reshape(grid.shape, order='F') theta = np.arctan2(grid[:, 1], grid[:, 0]) return np.hstack([ mkvc(np.cos(theta)*vec[:, 0] + np.sin(theta)*vec[:, 1], 2), mkvc(-np.sin(theta)*vec[:, 0] + np.cos(theta)*vec[:, 1], 2), mkvc(vec[:, 2], 2) ])
[docs]def spherical_2_cartesian(grid, vec=None): """ Take a grid or vector (if provided) defined in spherical coordinates :math:`(r, \\theta, \\phi)` and transform it to cartesian coordinates, :math:`(x, y, z)`. **Required** :param numpy.ndarray grid: grid in spherical coordinates :math:`(r, \\theta, \\phi)` **Optional** :param numpy.ndarray vec: (optional) vector defined in spherical coordinates **Returns** :return: grid or vector (if provided) in cartesian coordinates :math:`(x, y, z)` :rtype: numpy.ndarray """ grid = np.atleast_2d(grid) if vec is None: return np.hstack([ mkvc(grid[:, 0] * np.sin(grid[:, 2]) * np.cos(grid[:, 1]), 2), mkvc(grid[:, 0] * np.sin(grid[:, 2]) * np.sin(grid[:, 1]), 2), mkvc(grid[:, 0] * np.cos(grid[:, 2]), 2) ]) if len(vec.shape) == 1 or vec.shape[1] == 1: vec = vec.reshape(grid.shape, order='F') x = ( vec[:, 0] * np.sin(grid[:, 2]) * np.cos(grid[:, 1]) + vec[:, 2] * np.cos(grid[:, 2]) * np.cos(grid[:, 1]) - vec[:, 1] * np.sin(grid[:, 1]) ) y = ( vec[:, 0] * np.sin(grid[:, 2]) * np.sin(grid[:, 1]) + vec[:, 2] * np.cos(grid[:, 2]) * np.sin(grid[:, 1]) - vec[:, 1] * np.cos(grid[:, 1]) ) z = ( vec[:, 0] * np.cos(grid[:, 2]) - vec[:, 2] * np.sin(grid[:, 2]) ) newvec = [x, y, z] return np.vstack(newvec).T
[docs]def cartesian_2_spherical(grid, vec=None): """ Takes a grid or vector (if provided) defined in cartesian coordinates :math:`(x, y, z)` and transform it to spherical coordinates, :math:`(r, \\theta, \\phi)`. **Required** :param numpy.ndarray grid: grid in cartesian coordinates :math:`(x, y, z)` **Optional** :param numpy.ndarray vec: (optional) vector defined in cartesian coordinates **Returns** :return: grid or vector (if provided) in spherical coordinates :math:`(r, \\theta, \\phi)` :rtype: numpy.ndarray """ grid = np.atleast_2d(grid) if vec is None: return np.hstack([ mkvc(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2 + grid[:, 2]**2), 2), mkvc(np.arctan2(grid[:, 1], grid[:, 0]), 2), mkvc( np.arctan2(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), grid[:, 2]), 2 ), ]) if len(vec.shape) == 1 or vec.shape[1] == 1: vec = vec.reshape(grid.shape, order='F') theta = np.arctan2(grid[:, 1], grid[:, 0]) phi = np.arctan2(np.sqrt(grid[:, 0]**2 + grid[:, 1]**2), grid[:, 2]) r = ( vec[:, 0] * np.sin(phi) * np.cos(theta) + vec[:, 1] * np.sin(phi) * np.sin(theta) + vec[:, 2] * np.cos(phi) ) theta = - vec[:, 0] * np.sin(theta) + vec[:, 1] * np.cos(theta) phi = ( vec[:, 0] * np.cos(phi) * np.cos(theta) + vec[:, 1] * np.cos(phi) * np.sin(theta) - vec[:, 2] * np.sin(phi) ) newvec = [r, theta, phi] return np.vstack(newvec).T
[docs]def vector_magnitude(v): """ Amplitude of a vector, v. **Required** :param numpy.ndarray v: vector array **Returns** :returns: magnitude of a vector (n, 1) :rtype: numpy.ndarray """ v = np.atleast_2d(v) return np.sqrt((v**2).sum(axis=1))
[docs]def vector_distance(xyz, origin=np.r_[0., 0., 0.]): """ Vector distance of a grid, xyz from an origin origin. **Required** :param numpy.ndarray xyz: grid (npoints x 3) **Optional** :param numpy.ndarray origin: origin (default: [0., 0., 0.]) **Returns** :returns: vector distance from a grid of points from the origin (npoints x 3) :rtype: numpy.ndarray """ assert(xyz.shape[1] == 3), ( "the xyz grid should be npoints by 3, the shape provided is {}".format( xyz.shape ) ) if len(origin) != 3: raise Exception( "the origin must be length 3, the length provided is {}".format( len(origin) ) ) dx = xyz[:, 0] - origin[0] dy = xyz[:, 1] - origin[1] dz = xyz[:, 2] - origin[2] return np.c_[dx, dy, dz]
[docs]def distance(xyz, origin=np.r_[0., 0., 0.]): """ Radial distance from an grid of points to the origin **Required** :param numpy.ndarray xyz: grid (npoints x 3) **Optional** :param numpy.ndarray origin: origin (default: [0., 0., 0.]) **Returns** :returns: distance between each point and the origin (npoints x 1) :rtype: numpy.ndarray """ dxyz = vector_distance(xyz, origin) return vector_magnitude(dxyz)
[docs]def vector_dot(xyz, vector): """ Take a dot product between an array of vectors, xyz and a vector [x, y, z] **Required** :param numpy.ndarray xyz: grid (npoints x 3) :param numpy.ndarray vector: vector (1 x 3) **Returns** :returns: dot product between the grid and the (1 x 3) vector, returns an (npoints x 1) array :rtype: numpy.ndarray """ if len(vector) != 3: raise Exception( "vector should be length 3, the provided length is {}".format( len(vector) ) ) return vector[0]*xyz[:, 0] + vector[1]*xyz[:, 1] + vector[2]*xyz[:, 2]
[docs]def repeat_scalar(scalar, dim=3): """ Repeat a spatially distributed scalar value dim times to simplify multiplication with a vector. **Required** :param numpy.ndarray scalar: (n x 1) array of scalars **Optional** :param int dim: dimension of the second axis for the output (default = 3) **Returns** :returns: (n x dim) array of the repeated vector :rtype: numpy.ndarray """ assert len(scalar) in scalar.shape, ( "input must be a scalar. The shape you provided is {}".format( scalar.shape ) ) return np.kron(np.ones((1, dim)), np.atleast_2d(scalar).T)
[docs]def rotation_matrix_from_normals(v0, v1, tol=1e-20): """ Performs the minimum number of rotations to define a rotation from the direction indicated by the vector n0 to the direction indicated by n1. The axis of rotation is n0 x n1 https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula :param numpy.ndarray v0: vector of length 3 :param numpy.ndarray v1: vector of length 3 :param float tol: tolerance. If the norm of the cross product between the two vectors is below this, no rotation is performed default = 1e-20 :rtype: numpy.ndarray :return: 3 x 3 rotation matrix which rotates the frame so that n0 is aligned with n1 """ v0 = mkvc(v0) v1 = mkvc(v1) # ensure both n0, n1 are vectors of length 1 assert len(v0) == 3, "Length of n0 should be 3" assert len(v1) == 3, "Length of n1 should be 3" # ensure both are true normals n0 = v0*1./np.linalg.norm(v0) n1 = v1*1./np.linalg.norm(v1) n0dotn1 = n0.dot(n1) # define the rotation axis, which is the cross product of the two vectors rotAx = np.cross(n0, n1) if np.linalg.norm(rotAx) < tol: return np.eye(3, dtype=float) rotAx *= 1./np.linalg.norm(rotAx) cosT = n0dotn1/(np.linalg.norm(n0)*np.linalg.norm(n1)) sinT = np.sqrt(1.-n0dotn1**2) ux = np.array( [ [0., -rotAx[2], rotAx[1]], [rotAx[2], 0., -rotAx[0]], [-rotAx[1], rotAx[0], 0.] ], dtype=float ) return np.eye(3, dtype=float) + sinT*ux + (1.-cosT)*(ux.dot(ux))
[docs]def rotate_points_from_normals(xyz, n0, n1, x0=np.r_[0., 0., 0.]): """ rotates a grid so that the vector n0 is aligned with the vector n1 **Required** :param numpy.ndarray xyz: :param numpy.ndarray n0: vector of length 3, should have norm 1 :param numpy.ndarray n1: vector of length 3, should have norm 1 **Optional** :param numpy.ndarray x0: vector of length 3, point about which we perform the rotation **Returns** :rtype: numpy.ndarray :return: (3x3) rotation matrix which rotates the frame so that n0 is aligned with n1 """ R = rotation_matrix_from_normals(n0, n1) if xyz.shape[1] != 3: raise AssertionError("Grid xyz should be 3 wide") if len(x0) != 3: raise AssertionError("x0 should have length 3") x0 = np.array(x0.flatten()) # ensure it is an array not a vector3 X0 = np.ones([xyz.shape[0], 1])*mkvc(x0) return np.dot(xyz - X0, R.T) + X0 # equivalent to (R*(XYZ - X0)).T + X0