Source code for torx.vector.vector_operations_m

"""Functions to operator on vectors."""
import numpy as np
import xarray as xr
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def vector_magnitude(input_array: xr.DataArray):
    """Return the vector magnitude."""
    return np.sqrt(vector_dot(input_array, input_array))




@autodoc_function
def poloidal_vector_magnitude(input_array: xr.DataArray):
    """Return the vector magnitude of the poloidal component."""
    return vector_magnitude(input_array.sel(vector=["eR", "eZ"]))




@autodoc_function
def vector_dot(input_a: xr.DataArray, input_b: xr.DataArray):
    """Return the dot product of two vectors."""
    assert ("vector" in input_a.dims) and ("vector" in input_b.dims)
    assert input_a.vector.size == input_b.vector.size

    dims = list(input_a.dims)
    dims.remove("vector")

    a_attrs = input_a.attrs.copy()
    b_attrs = input_b.attrs.copy()

    # Handle arbitrary vector length
    result = input_a.isel(vector=0) * input_b.isel(vector=0)
    for i in np.arange(1, input_a.vector.size):
        result += input_a.isel(vector = i) * input_b.isel(vector = i)
    
    if "norm" in a_attrs and "norm" in b_attrs:
        result.attrs["norm"] = a_attrs["norm"] * b_attrs["norm"]

    return result.transpose(*dims).drop_vars("vector")




@autodoc_function
def vector_cross(input_a, input_b):
    """
    Return the cross product of two vectors.

    vec(a) cross vec(b) = mag(a) mag(b) sin(theta) vec(n)
    where theta is the angle between vec(a) and vec(b) and
    where vec(n) is a unit vector perpendicular to both vec(a) and vec(b)
    The sign of vec(b) is given by right-hand-rule (i.e. if 'a' is index
    finger, 'b' is middle finger, 'n' is thumb, or alternatively if 'a'
    is thumb, 'b' is index finger, 'n' is middle finger).
    """
    dask_gufunc_kwargs_dict = dict(output_sizes=dict(vector=3))

    return xr.apply_ufunc(
        np.cross,
        input_a,
        input_b,
        input_core_dims=[["vector"], ["vector"]],
        output_core_dims=[["vector"]],
        dask="parallelized",
        output_dtypes=[input_a.dtype],
        dask_gufunc_kwargs=dask_gufunc_kwargs_dict
    )




@autodoc_function
def unit_vector(input_array: xr.DataArray):
    """Return the unit vector in the direction of the input array."""
    return input_array / vector_magnitude(input_array)




@autodoc_function
def scalar_projection(input_a: xr.DataArray, input_b: xr.DataArray):
    """
    Return the magnitude of input_a which is parallel to input_b.

    See https://en.wikipedia.org/wiki/Vector_projection for notation.
    scalar_projection is 'a_1'.
    """
    return vector_dot(input_a, unit_vector(input_b))




@autodoc_function
def vector_projection(input_a: xr.DataArray, input_b: xr.DataArray):
    """
    Return a projected vector.

    The result is the orthogonal projection of input_a onto a
    straight line parallel to input_b. It is parallel to input_b.

    See https://en.wikipedia.org/wiki/Vector_projection for notation.
    vector_rejection is 'vec(a)_1'.
    """
    return scalar_projection(input_a, input_b) * unit_vector(input_b)




@autodoc_function
def vector_rejection(input_a: xr.DataArray, input_b: xr.DataArray):
    """
    Return a rejected vector.

    The result is the projection of input_a onto a straight
    line orthogonal to input_b.

    See https://en.wikipedia.org/wiki/Vector_projection for notation.
    vector_rejection is 'vec(a)_2'.
    """
    return input_a - vector_projection(input_a, input_b)