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="degree"):
"""
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 "deg" in units.lower():
pitch = np.deg2rad(pitch)
roll = np.deg2rad(roll)
heading = np.deg2rad(heading)
elif "rad" in units.lower():
pass
else:
raise ValueError("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, units="degree"):
"""
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 "deg" in units.lower():
pitch = np.deg2rad(pitch)
roll = np.deg2rad(roll)
elif "rad" in units.lower():
pass
else:
raise ValueError("Invalid units")
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)