Source code for mhkit.dolfyn.rotate.base

import numpy as np
import xarray as xr
from numpy.linalg import det, inv
from scipy.spatial.transform import Rotation as R
import warnings


def _make_model(ds):
    """
    The make and model of the instrument that collected the data
    in this data object.
    """
    return '{} {}'.format(ds.attrs['inst_make'],
                          ds.attrs['inst_model']).lower()


def _check_rotmat_det(rotmat, thresh=1e-3):
    """
    Check that the absolute error of the determinant is small.

          abs(det(rotmat) - 1) < thresh

    Returns a boolean array.
    """

    if rotmat.ndim > 2:
        rotmat = np.transpose(rotmat)
    return np.abs(det(rotmat) - 1) < thresh


def _check_rotate_vars(ds, rotate_vars):
    if rotate_vars is None:
        if 'rotate_vars' in ds.attrs:
            rotate_vars = ds.rotate_vars
        else:
            warnings.warn("    'rotate_vars' attribute not found."
                          "Rotating `vel`.")
            rotate_vars = ['vel']

    return rotate_vars


def _set_coords(ds, ref_frame, forced=False):
    """
    Checks the current reference frame and adjusts xarray coords/dims 
    as necessary.
    Makes sure assigned dataarray coordinates match what DOLfYN is reading in.
    """

    make = _make_model(ds)

    XYZ = ['X', 'Y', 'Z']
    ENU = ['E', 'N', 'U']
    beam = ds.beam.values
    principal = ['streamwise', 'x-stream', 'vert']

    # check make/model
    if 'rdi' in make:
        inst = ['X', 'Y', 'Z', 'err']
        earth = ['E', 'N', 'U', 'err']
        princ = ['streamwise', 'x-stream', 'vert', 'err']

    elif 'nortek' in make:
        if 'signature' in make or 'ad2cp' in make:
            inst = ['X', 'Y', 'Z1', 'Z2']
            earth = ['E', 'N', 'U1', 'U2']
            princ = ['streamwise', 'x-stream', 'vert1', 'vert2']

        else:  # AWAC or Vector
            inst = XYZ
            earth = ENU
            princ = principal

    orient = {'beam': beam, 'inst': inst, 'ship': inst, 'earth': earth,
              'principal': princ}
    orientIMU = {'beam': XYZ, 'inst': XYZ, 'ship': XYZ, 'earth': ENU,
                 'principal': principal}

    if forced:
        ref_frame += '-forced'

    # Update 'dir' and 'dirIMU' dimensions
    attrs = ds['dir'].attrs
    attrs.update({'ref_frame': ref_frame})

    ds['dir'] = orient[ref_frame]
    ds['dir'].attrs = attrs
    if hasattr(ds, 'dirIMU'):
        ds['dirIMU'] = orientIMU[ref_frame]
        ds['dirIMU'].attrs = attrs

    ds.attrs['coord_sys'] = ref_frame

    # These are essentially one extra line to scroll through
    tag = ['', '_echo', '_bt']
    for tg in tag:
        if hasattr(ds, 'coord_sys_axes'+tg):
            ds.attrs.pop('coord_sys_axes'+tg)

    return ds


def _beam2inst(dat, reverse=False, force=False):
    """
    Rotate velocities from beam to instrument coordinates.

    Parameters
    ----------
    dat : xarray.Dataset
      The ADCP dataset
    reverse : bool
      If True, this function performs the inverse rotation (inst->beam).
      Default = False
    force : bool, list
      When true do not check which coordinate system the data is in
      prior to performing this rotation. When forced-rotations are
      applied, the string '-forced!' is appended to the
      dat.props['coord_sys'] string. If force is a list, it contains
      a list of variables that should be rotated (rather than the
      default values in adpo.rotate_vars).
      Default = False
    """

    if not force:
        if not reverse and dat.coord_sys.lower() != 'beam':
            raise ValueError('The input must be in beam coordinates.')
        if reverse and dat.coord_sys != 'inst':
            raise ValueError('The input must be in inst coordinates.')

    rotmat = dat['beam2inst_orientmat']

    if isinstance(force, (list, set, tuple)):
        # You can force a distinct set of variables to be rotated by
        # specifying it here.
        rotate_vars = force
    else:
        rotate_vars = [
            ky for ky in dat.rotate_vars if dat[ky].shape[0] == rotmat.shape[0]]

    cs = 'inst'
    if reverse:
        # Can't use transpose because rotation is not between
        # orthogonal coordinate systems
        rotmat = inv(rotmat)
        cs = 'beam'
    for ky in rotate_vars:
        dat[ky].values = np.einsum('ij,j...->i...', rotmat, dat[ky].values)

    if force:
        dat = _set_coords(dat, cs, forced=True)
    else:
        dat = _set_coords(dat, cs)

    return dat


[docs]def euler2orient(heading, pitch, roll, units='degrees'): """ Calculate the orientation matrix from DOLfYN-defined euler angles. This function is not likely to be called during data processing since it requires DOLfYN-defined euler angles. It is intended for testing DOLfYN. The matrices H, P, R are the transpose of the matrices for rotation about z, y, x as shown here https://en.wikipedia.org/wiki/Rotation_matrix. The transpose is used because in DOLfYN the orientation matrix is organized for rotation from EARTH --> INST, while the wiki's matrices are organized for rotation from INST --> EARTH. Parameters ---------- heading : numpy.ndarray The heading angle. pitch : numpy.ndarray The pitch angle. roll : numpy.ndarray The roll angle. units : str Units in degrees or radians. is 'degrees' Returns ======= omat : numpy.ndarray (3x3xtime) The orientation matrix of the data. The returned orientation matrix obeys the following conventions: - a "ZYX" rotation order. That is, these variables are computed assuming that rotation from the earth -> instrument frame happens by rotating around the z-axis first (heading), then rotating around the y-axis (pitch), then rotating around the x-axis (roll). Note this requires matrix multiplication in the reverse order. - heading is defined as the direction the x-axis points, positive clockwise from North (this is *opposite* the right-hand-rule around the Z-axis), range 0-360 degrees. - pitch is positive when the x-axis pitches up (this is *opposite* the right-hand-rule around the Y-axis) - roll is positive according to the right-hand-rule around the instrument's x-axis """ if units.lower() == 'degrees': pitch = np.deg2rad(pitch) roll = np.deg2rad(roll) heading = np.deg2rad(heading) elif units.lower() == 'radians': pass else: raise Exception("Invalid units") # Converts the DOLfYN-defined heading to one that follows the right-hand-rule # reports heading as rotation of the y-axis positive counterclockwise from North heading = np.pi / 2 - heading # Converts the DOLfYN-defined pitch to one that follows the right-hand-rule. pitch = -pitch ch = np.cos(heading) sh = np.sin(heading) cp = np.cos(pitch) sp = np.sin(pitch) cr = np.cos(roll) sr = np.sin(roll) zero = np.zeros_like(sr) one = np.ones_like(sr) H = np.array( [[ch, sh, zero], [-sh, ch, zero], [zero, zero, one], ]) P = np.array( [[cp, zero, -sp], [zero, one, zero], [sp, zero, cp], ]) R = np.array( [[one, zero, zero], [zero, cr, sr], [zero, -sr, cr], ]) return np.einsum('ij...,jk...,kl...->il...', R, P, H)
[docs]def orient2euler(omat): """ Calculate DOLfYN-defined euler angles from the orientation matrix. Parameters ---------- omat : numpy.ndarray The orientation matrix Returns ------- heading : numpy.ndarray The heading angle. Heading is defined as the direction the x-axis points, positive clockwise from North (this is *opposite* the right-hand-rule around the Z-axis), range 0-360 degrees. pitch : np.ndarray The pitch angle (degrees). Pitch is positive when the x-axis pitches up (this is *opposite* the right-hand-rule around the Y-axis). roll : np.ndarray The roll angle (degrees). Roll is positive according to the right-hand-rule around the instrument's x-axis. """ if isinstance(omat, np.ndarray) and \ omat.shape[:2] == (3, 3): pass elif hasattr(omat, 'orientmat'): omat = omat['orientmat'].values # Note: orientation matrix is earth->inst unless supplied by an external IMU hh = np.rad2deg(np.arctan2(omat[0, 0], omat[0, 1])) hh %= 360 return ( # heading hh, # pitch np.rad2deg(np.arcsin(omat[0, 2])), # roll np.rad2deg(np.arctan2(omat[1, 2], omat[2, 2])), )
[docs]def quaternion2orient(quaternions): """ Calculate orientation from Nortek AHRS quaternions, where q = [W, X, Y, Z] instead of the standard q = [X, Y, Z, W] = [q1, q2, q3, q4] Parameters ---------- quaternions : xarray.DataArray Quaternion dataArray from the raw dataset Returns ------- orientmat : numpy.ndarray The earth2inst rotation maxtrix as calculated from the quaternions See Also -------- scipy.spatial.transform.Rotation """ omat = type(quaternions)(np.empty((3, 3, quaternions.time.size))) omat = omat.rename({'dim_0': 'earth', 'dim_1': 'inst', 'dim_2': 'time'}) for i in range(quaternions.time.size): r = R.from_quat([quaternions.isel(q=1, time=i), quaternions.isel(q=2, time=i), quaternions.isel(q=3, time=i), quaternions.isel(q=0, time=i)]) omat[..., i] = r.as_matrix() # quaternions in inst2earth reference frame, need to rotate to earth2inst omat.values = np.rollaxis(omat.values, 1) earth = xr.DataArray(['E', 'N', 'U'], dims=['earth'], name='earth', attrs={ 'units': '1', 'long_name': 'Earth Reference Frame', 'coverage_content_type': 'coordinate'}) inst = xr.DataArray(['X', 'Y', 'Z'], dims=['inst'], name='inst', attrs={ 'units': '1', 'long_name': 'Instrument Reference Frame', 'coverage_content_type': 'coordinate'}) return omat.assign_coords({'earth': earth, 'inst': inst, 'time': quaternions.time})
[docs]def calc_tilt(pitch, roll): """ Calculate "tilt", the vertical inclination, from pitch and roll. Parameters ---------- roll : numpy.ndarray or xarray.DataArray Instrument roll in degrees pitch : numpy.ndarray or xarray.DataArray Instrument pitch in degrees Returns ------- tilt : numpy.ndarray Vertical inclination of the instrument in degrees """ if 'xarray' in type(pitch).__module__: pitch = pitch.values if 'xarray' in type(roll).__module__: roll = roll.values tilt = np.arctan( np.sqrt(np.tan(np.deg2rad(roll)) ** 2 + np.tan(np.deg2rad(pitch)) ** 2) ) return np.rad2deg(tilt)