Source code for xicsrt.tools.xicsrt_math

# -*- coding: utf-8 -*-
"""
.. Authors
    Novimir Pablant <npablant@pppl.gov>
    James Kring <jdk0026@tigermail.auburn.edu>
    Yevgeniy Yakusevich <eugenethree@gmail.com>

A set of mathematical utilities and vector convenience functions for XICSRT.
"""

import numpy as np


[docs]def distance_point_to_line(origin, normal, point): """ Find the closest distnace between a point and a line in 3D. """ o = origin n = normal p = point t = np.dot(p - o, n) / np.dot(n, n) d = np.linalg.norm(np.outer(t, n) + o - p, axis=1) return d
[docs]def intersect_ray_plane(ray, plane): """ Find the intersection between a ray and a plane in 3D. """ if ray['origin'].ndim == 1: distance = (np.dot((plane['origin'] - ray['origin']), plane['zaxis']) / np.dot(ray['direction'], plane['zaxis'])) intersect = ray['origin'] + ray['direction'] * distance else: raise NotImplementedError('Only a single ray is currently supported.') return intersect
[docs]def toarray_1d(a): """ Convert the input to a ndarray with at least 1 dimension. This is similar to the numpy function atleast_1d, but has less overhead and is jax compatible. """ a = np.asarray(a) if a.ndim == 0: a = a.reshape(1) return a
[docs]def vector_angle(a, b): """ Find the angle between two vectors. """ if a.ndim == 2 & b.ndim == 2: a_mod = np.linalg.norm(a, axis=1) b_mod = np.linalg.norm(b, axis=1) dot = np.einsum('ij,i,ij,i->i', a, 1/a_mod, b, 1/b_mod, optimize=True) elif a.ndim == 1 & b.ndim == 1: a_mod = np.linalg.norm(a) b_mod = np.linalg.norm(b) dot = np.dot(a/a_mod, b/b_mod) else: raise Exception('Input must have 1 or 2 dimensions.') angle = np.arccos(dot) return angle
[docs]def vector_rotate(a, b, theta): """ Rotate vector a around vector b by an angle theta (radians) Programming Notes: u: parallel projection of a on b_hat. v: perpendicular projection of a on b_hat. w: a vector perpendicular to both a and b. """ a = np.asarray(a) b = np.asarray(b) if a.ndim == 2: b_hat = b / np.linalg.norm(b) dot = np.einsum('ij,j->i', a, b_hat, optimize=True) u = np.einsum('i,j->ij', dot, b_hat, optimize=True) v = a - u w = np.cross(b_hat, v) c = u + v * np.cos(theta) + w * np.sin(theta) elif a.ndim == 1: b_hat = b / np.linalg.norm(b) u = b_hat * np.dot(a, b_hat) v = a - u w = np.cross(b_hat, v) c = u + v * np.cos(theta) + w * np.sin(theta) else: raise Exception('Input array must be 1d (vector) or 2d (array of vectors)') return c
[docs]def magnitude(vector): """ Calculate magnitude of a vector or array of vectors. """ if vector.ndim > 1: mag = np.linalg.norm(vector, axis=1) else: mag = np.linalg.norm(vector) return mag
[docs]def normalize(vector): """ Normalize a vector or an array of vectors. If an array of vectors is given it should have the shape (N,M) where | N: Number of vectors | M: Vector length """ if vector.ndim > 1: norm = np.linalg.norm(vector, axis=1) vector /= np.expand_dims(norm, 1) else: norm = np.linalg.norm(vector) vector /= norm return vector
[docs]def sinusoidal_spiral(phi, b, r0, theta0): r = r0 * (np.sin(theta0 + (b-1)*phi)/np.sin(theta0))**(1/(b-1)) return r
[docs]def rotation_matrix(axis, theta): """ Return the rotation matrix associated with counterclockwise rotation about the given axis by theta radians. """ axis = np.asarray(axis) axis = axis/np.sqrt(np.dot(axis, axis)) a = np.cos(theta/2.0) b, c, d = -axis*np.sin(theta/2.0) aa, bb, cc, dd = a*a, b*b, c*c, d*d bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d matrix = np.array( [[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)] ,[2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)] ,[2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]]) return matrix
[docs]def bragg_angle(wavelength, crystal_spacing): """ The Bragg angle calculation is used so often that it deserves its own function. .. Note:: The crystal_spacing here is the true spacing, not the 2d spacing that is sometimes used in the literature. """ bragg_angle = np.arcsin(wavelength / (2 * crystal_spacing)) return bragg_angle
[docs]def cyl_from_car(point_car): """ Convert from cartesian to cylindrical coordinates. """ if not isinstance(point_car, np.ndarray): point_car = np.array(point_car) point_cyl = np.empty(point_car.shape) if point_car.ndim == 2: point_cyl[:, 0] = np.sqrt(np.sum(point_car[:, 0:2]**2, 1)) point_cyl[:, 1] = np.arctan2(point_car[:, 1], point_car[:, 0]) point_cyl[:, 2] = point_car[:, 2] else: point_cyl[0] = np.sqrt(np.sum(point_car[0:2]**2)) point_cyl[1] = np.arctan2(point_car[1], point_car[0]) point_cyl[2] = point_car[2] return point_cyl
[docs]def car_from_cyl(point_cyl): """ Convert from cylindrical to cartesian coordinates. """ if not isinstance(point_cyl, np.ndarray): point_cyl = np.array(point_cyl) point_car = np.empty(point_cyl.shape) if point_cyl.ndim == 2: point_car[:, 0] = point_cyl[:, 0]*np.cos(point_cyl[:, 1]) point_car[:, 1] = point_cyl[:, 0]*np.sin(point_cyl[:, 1]) point_car[:, 2] = point_cyl[:, 2] else: point_car[0] = point_cyl[0]*np.cos(point_cyl[1]) point_car[1] = point_cyl[0]*np.sin(point_cyl[1]) point_car[2] = point_cyl[2] return point_car
[docs]def tor_from_car(point_car, major_radius): """ Convert from cartesian to toroidal coordinates. Arguments --------- point_car: array [meters] Cartesian coordinates [x,y,z] major_radius: float [meters] Torus Major Radius Returns ------- point_tor: array [meters] Toroidal coordinates [r_min, theta_poloidal, theta_toroidal] """ if not isinstance(point_car, np.ndarray): point_car = np.array(point_car) point_tor = np.empty(point_car.shape) if point_tor.ndim == 2: d = np.linalg.norm(point_car[:, 0:2], axis=1) - major_radius point_tor[:, 2] = np.arctan2(point_car[:, 1], point_car[:, 0]) point_tor[:, 1] = np.arctan2(point_car[:, 2], d) point_tor[:, 0] = np.sqrt(np.power(point_car[:, 2], 2) + np.power(d, 2)) else: d = np.linalg.norm(point_car[0:2]) - major_radius point_tor[2] = np.arctan2(point_car[1], point_car[0]) point_tor[1] = np.arctan2(point_car[2], d) point_tor[0] = np.sqrt(np.power(point_car[2], 2) + np.power(d, 2)) return point_tor
[docs]def car_from_tor(point_tor, major_radius): """ Convert from toroidal to cartesian coordinates. Arguments --------- point_tor: array [meters] Toroidal coordinates [r_min, theta_poloidal, theta_toroidal] major_radius: float [meters] Torus Major Radius Returns ------- point_car: array [meters] Cartesian coordinates [x,y,z] """ if not isinstance(point_tor, np.ndarray): point_tor = np.array(point_tor) point_car = np.empty(point_tor.shape) if point_tor.ndim == 2: point_car[:, 0] = major_radius*np.cos(point_tor[:, 2]) point_car[:, 1] = major_radius*np.sin(point_tor[:, 2]) vector = point_car[:, 0:2]/np.linalg.norm(point_car[:, 0:2], axis=1) point_car[:, 0] = point_car[:, 0] + vector[:, 0]*point_tor[:, 0]*np.cos(point_tor[:, 1]) point_car[:, 1] = point_car[:, 1] + vector[:, 1]*point_tor[:, 0]*np.cos(point_tor[:, 1]) point_car[:, 2] = point_tor[:, 0]*np.sin(point_tor[:, 1]) else: point_car[0] = major_radius*np.cos(point_tor[2]) point_car[1] = major_radius*np.sin(point_tor[2]) vector = point_car[0:2]/np.linalg.norm(point_car[0:2]) point_car[0] = point_car[0] + vector[0]*point_tor[0]*np.cos(point_tor[1]) point_car[1] = point_car[1] + vector[1]*point_tor[0]*np.cos(point_tor[1]) point_car[2] = point_tor[0]*np.sin(point_tor[1]) return point_car
[docs]def point_in_triangle_2d(pt, p0, p1, p2): """ Determine if a point (or set of points) fall within a triangle as specified by three vertices. Calculation is performed in two dimensions (2D). """ pt = np.asarray(pt) if pt.ndim == 2: area = 0.5 * (-p1[1] * p2[0] + p0[1] * (-p1[0] + p2[0]) + p0[0] * (p1[1] - p2[1]) + p1[0] * p2[1]) a = 1 / (2 * area) * (p0[1] * p2[0] - p0[0] * p2[1] + (p2[1] - p0[1]) * pt[:, 0] + (p0[0] - p2[0]) * pt[:, 1]) b = 1 / (2 * area) * (p0[0] * p1[1] - p0[1] * p1[0] + (p0[1] - p1[1]) * pt[:, 0] + (p1[0] - p0[0]) * pt[:, 1]) c = 1 - a - b else: area = 0.5 * (-p1[1] * p2[0] + p0[1] * (-p1[0] + p2[0]) + p0[0] * (p1[1] - p2[1]) + p1[0] * p2[1]) a = 1 / (2 * area) * (p0[1] * p2[0] - p0[0] * p2[1] + (p2[1] - p0[1]) * pt[0] + (p0[0] - p2[0]) * pt[1]) b = 1 / (2 * area) * (p0[0] * p1[1] - p0[1] * p1[0] + (p0[1] - p1[1]) * pt[0] + (p1[0] - p0[0]) * pt[1]) c = 1 - a - b return (a >= 0) & (b >= 0) & (c >= 0)