vector_analysis.py

import numpy as np
import itertools

def vf_gradient(field, *, axis=None):
    '''
    Gradient of a vector field.

    :param field: (n, c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Gradient of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if axis == None:
        axis = tuple([i for i in range(len(field))])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif len(field) != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    num_dims = len(field)
    ret = []
    for i in axis:
        ret.append([np.gradient(field[j], axis=i) for j in range(num_dims)])
    return np.array(ret)

def gradient(field, rank, *, axis=None):
    '''
    Gradient of a field.

    :param field: (c1, c2, ... , cn) or (n, c1, c2, ... , cn)
    :param rank: {0, 1}
    :param axis: None or int or tuple of ints, optional.
    :return: Gradient of the field.
    :rtype: ndarray
    :raises ValueError:
    '''
    if rank == 0: # for a scalar field.
        return np.array(np.gradient(field, axis=axis))
    elif rank == 1: # for a vector field.
        return vf_gradient(field, axis=axis)
    else:
        raise ValueError('\'axis\' entry is out of bounds.')


def vf_divergence(field, *, axis=None):
    '''
    Divergence of a vector field.

    :param field: (n, c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Divergence of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if axis == None:
        axis = tuple([i for i in range(len(field))])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif len(field) != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    return np.ufunc.reduce(np.add, [np.gradient(field[i], axis=j) for i, j in enumerate(axis)])


def tf_divergence(field, *, axis=None):
    '''
    Divergence of a tensor field.

    :param field: (n, n, c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Divergence of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if axis == None:
        axis = tuple([i for i in range(len(field))])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif len(field) != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    num_dims = len(field)
    ret = []
    for i in range(num_dims):
        ret.append(np.ufunc.reduce(np.add, [np.gradient(field[j, i], axis=k) for j, k in enumerate(axis)]))
    return np.array(ret)


def divergence(field, rank, *, axis=None):
    '''
    Divergence of a field.

    :param field: (n, c1, c2, ... , cn) or (n, n, c1, c2, ... , cn)
    :param rank: {1, 2}
    :param axis: None or tuple of ints, optional.
    :return: Divergence of the field.
    :rtype: ndarray
    :raises ValueError:
    '''
    if rank == 1:  # for a vector field.
        return vf_divergence(field, axis=axis)
    elif rank == 2:  # for a tensor field.
        return tf_divergence(field, axis=axis)
    else:
        raise ValueError('\'axis\' entry is out of bounds')


def vf_curl(field, *, axis=None):
    '''
    Curl of a vector field.

    :param field: (n, c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Curl of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if len(field) != 3:
        raise ValueError('')
    elif axis == None:
        axis = tuple([i for i in range(len(field))])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif len(field) != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    num_dims = len(field)
    return np.array([np.gradient(field[(i + 2) % num_dims], axis=axis[(i + 1) % num_dims]) - np.gradient(
        field[(i + 1) % num_dims], axis=axis[(i + 2) % num_dims]) for i in range(num_dims)])


def curl(field, rank, *, axis=None):
    '''
    Curl of a vector field.

    :param field: (n, c1, c2, ... , cn)
    :param rank: {1}
    :param axis: None or tuple of ints, optional.
    :return: Curl of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if rank == 1:  # for a vector field.
        return vf_curl(field, axis=axis)
    else:
        raise ValueError('\'axis\' entry is out of bounds')


def sf_laplacian(field, *, axis=None):
    '''
    Laplacian of a scalar field.

    :param field: (c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Laplacian of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if axis == None:
        axis = tuple([i for i in range(field.ndim)])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif field.ndim != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    grad = np.gradient(field, axis=axis)
    return np.ufunc.reduce(np.add, [np.gradient(grad[i], axis=j) for i, j in enumerate(axis)])


def vf_laplacian(field, *, axis=None):
    '''
    Laplacian of a vector field.

    :param field: (n, c1, c2, ... , cn)
    :param axis: None or tuple of ints, optional.
    :return: Laplacian of the field.
    :rtype: ndarray
    :raises TypeError:
    :raises ValueError:
    '''
    if axis == None:
        axis = tuple([i for i in range(len(field))])
    elif not isinstance(axis, tuple):
        raise TypeError('\'axis\' entry must be a tuple.')
    elif len(field) != len(axis):
        raise ValueError('\'axis\' and \'field\' must have the same length.')

    ret = []
    for f in field:
        g = np.gradient(f, axis=axis)
        ret.append(np.ufunc.reduce(np.add, [np.gradient(g[j], axis=k) for j, k in enumerate(axis)]))
    return np.array(ret)


def laplacian(field, rank, *, axis=None):
    '''
    Laplacian of a field.

    :param field: (c1, c2, ... , cn) or (n, c1, c2, ... , cn)
    :param rank: {0, 1}
    :param axis: None or tuple of ints, optional.
    :return: Laplacian of the field.
    :rtype: ndarray
    :raises ValueError:
    '''
    if rank == 0:  # for a scalar field.
        return sf_laplacian(field, axis=axis)
    elif rank == 1:  # for a vector field.
        return vf_laplacian(field, axis=axis)
    else:
        raise ValueError('\'axis\' entry is out of bounds')