from . import vector as r_vec
from . import awac as r_awac
from . import signature as r_sig
from . import rdi as r_rdi
from .base import _make_model
import numpy as np
import xarray as xr
import warnings
# The 'rotation chain'
rc = ["beam", "inst", "earth", "principal"]
rot_module_dict = {
# Nortek instruments
"vector": r_vec,
"awac": r_awac,
"signature": r_sig,
"ad2cp": r_sig,
# TRDI instruments
"rdi": r_rdi,
}
[docs]
def rotate2(ds, out_frame="earth", inplace=True):
"""
Rotate a dataset to a new coordinate system.
Parameters
----------
ds : xr.Dataset
The dolfyn dataset (ADV or ADCP) to rotate.
out_frame : string {'beam', 'inst', 'earth', 'principal'}
The coordinate system to rotate the data into.
inplace : bool
When True ``ds`` is modified. When False a copy is returned.
Default = True
Returns
-------
ds : xarray.Dataset or None
Returns a new rotated dataset **when ``inplace=False``**, otherwise
returns None.
Notes
-----
- This function rotates all variables in ``ds.attrs['rotate_vars']``.
- In order to rotate to the 'principal' frame, a value should exist for
``ds.attrs['principal_heading']``. The function
:func:`calc_principal_heading <dolfyn.calc_principal_heading>`
is recommended for this purpose, e.g.:
ds.attrs['principal_heading'] = dolfyn.calc_principal_heading(ds['vel'].mean(range))
where here we are using the depth-averaged velocity to calculate
the principal direction.
"""
# Create and return deep copy if not writing "in place"
if not inplace:
ds = ds.copy(deep=True)
csin = ds.coord_sys.lower()
if csin == "ship":
csin = "inst"
# Returns True/False if head2inst_rotmat has been set/not-set.
r_vec._check_inst2head_rotmat(ds)
if out_frame == "principal" and csin != "earth":
warnings.warn(
"You are attempting to rotate into the 'principal' "
"coordinate system, but the dataset is in the {} "
"coordinate system. Be sure that 'principal_heading' is "
"defined based on the earth coordinate system.".format(csin)
)
rmod = None
for ky in rot_module_dict:
if ky in _make_model(ds):
rmod = rot_module_dict[ky]
break
if rmod is None:
raise ValueError(
"Rotations are not defined for " "instrument '{}'.".format(_make_model(ds))
)
# Get the 'indices' of the rotation chain
try:
iframe_in = rc.index(csin)
except ValueError:
raise Exception(
"The coordinate system of the input "
"dataset, '{}', is invalid.".format(ds.coord_sys)
)
try:
iframe_out = rc.index(out_frame.lower())
except ValueError:
raise Exception(
"The specified output coordinate system "
"is invalid, please select one of: 'beam', 'inst', "
"'earth', 'principal'."
)
if iframe_out == iframe_in:
print("Data is already in the {} coordinate system".format(out_frame))
if iframe_out > iframe_in:
reverse = False
else:
reverse = True
while ds.coord_sys.lower() != out_frame.lower():
csin = ds.coord_sys
if csin == "ship":
csin = "inst"
inow = rc.index(csin)
if reverse:
func = getattr(rmod, "_" + rc[inow - 1] + "2" + rc[inow])
else:
func = getattr(rmod, "_" + rc[inow] + "2" + rc[inow + 1])
ds = func(ds, reverse=reverse)
if not inplace:
return ds
[docs]
def calc_principal_heading(vel, tidal_mode=True):
"""
Compute the principal angle of the horizontal velocity.
Parameters
----------
vel : np.ndarray (2,...,Nt), or (3,...,Nt)
The 2D or 3D velocity array (3rd-dim is ignored in this calculation)
tidal_mode : bool
If true, range is set from 0 to +/-180 degrees. If false, range is 0 to
360 degrees. Default = True
Returns
-------
p_heading : float or ndarray
The principal heading in degrees clockwise from North.
Notes
-----
When tidal_mode=True, this tool calculates the heading that is
aligned with the bidirectional flow. It does so following these
steps:
1. rotates vectors with negative velocity by 180 degrees
2. then doubles those angles to make a complete circle again
3. computes a mean direction from this, and halves that angle
(to undo the doubled-angles in step 2)
4. The returned angle is forced to be between 0 and 180. So, you
may need to add 180 to this if you want your positive
direction to be in the western-half of the plane.
Otherwise, this function simply computes the average direction
using a vector method.
"""
dt = vel[0] + vel[1] * 1j
if tidal_mode:
# Flip all vectors that are below the x-axis
dt[dt.imag <= 0] *= -1
# Now double the angle, so that angles near pi and 0 get averaged
# together correctly:
dt *= np.exp(1j * np.angle(dt))
dt = np.ma.masked_invalid(dt)
# Divide the angle by 2 to remove the doubling done on the previous
# line.
pang = np.angle(np.nanmean(dt, -1, dtype=np.complex128)) / 2
else:
pang = np.angle(np.nanmean(dt, -1))
return np.round((90 - np.rad2deg(pang)), decimals=4)
[docs]
def set_declination(ds, declin, inplace=True):
"""
Set the magnetic declination
Parameters
----------
ds : xarray.Dataset or :class:`dolfyn.velocity.Velocity`
The input dataset or velocity class
declination : float
The value of the magnetic declination in degrees (positive
values specify that Magnetic North is clockwise from True North)
inplace : bool
When True ``ds`` is modified. When False a copy is returned.
Default = True
Returns
-------
ds : xarray.Dataset or None
Returns a new dataset with declination set **when
``inplace=False``**, otherwise returns None.
Notes
-----
This function modifies the data object in the following ways:
- If the dataset is in the *earth* reference frame at the time of
setting declination, it will be rotated into the "*True-East*,
*True-North*, Up" (hereafter, ETU) coordinate system
- ``dat['orientmat']`` is modified to be an ETU to
instrument (XYZ) rotation matrix (rather than the magnetic-ENU to
XYZ rotation matrix). Therefore, all rotations to/from the 'earth'
frame will now be to/from this ETU coordinate system.
- The value of the specified declination will be stored in
``dat.attrs['declination']``
- ``dat['heading']`` is adjusted for declination
(i.e., it is relative to True North).
- If ``dat.attrs['principal_heading']`` is set, it is
adjusted to account for the orientation of the new 'True'
earth coordinate system (i.e., calling set_declination on a
data object in the principal coordinate system, then calling
dat.rotate2('earth') will yield a data object in the new
'True' earth coordinate system)
"""
# Create and return deep copy if not writing "in place"
if not inplace:
ds = ds.copy(deep=True)
if "declination" in ds.attrs:
angle = declin - ds.attrs.pop("declination")
else:
angle = declin
cd = np.cos(-np.deg2rad(angle))
sd = np.sin(-np.deg2rad(angle))
# The ordering is funny here because orientmat is the
# transpose of the inst->earth rotation matrix:
Rdec = np.array([[cd, -sd, 0], [sd, cd, 0], [0, 0, 1]])
if ds.coord_sys == "earth":
rotate2earth = True
rotate2(ds, "inst", inplace=True)
else:
rotate2earth = False
ds["orientmat"].values = np.einsum(
"kj...,ij->ki...",
ds["orientmat"].values,
Rdec,
).astype(np.float32)
if "heading" in ds:
heading = ds["heading"] + angle
heading[heading > 180] -= 360
ds["heading"].values = heading
if rotate2earth:
rotate2(ds, "earth", inplace=True)
if "principal_heading" in ds.attrs:
ds.attrs["principal_heading"] += angle
ds.attrs["declination"] = declin
ds.attrs["declination_in_orientmat"] = 1 # logical
if not inplace:
return ds
[docs]
def set_inst2head_rotmat(ds, rotmat, inplace=True):
"""
Set the instrument to head rotation matrix for the Nortek ADV if it
hasn't already been set through a '.userdata.json' file.
Parameters
----------
ds : xarray.Dataset
The data set to assign inst2head_rotmat
rotmat : float
3x3 rotation matrix
inplace : bool
When True ``ds`` is modified. When False a copy is returned.
Default = True
Returns
-------
ds : xarray.Dataset or None
Returns a new dataset with inst2head_rotmat set **when
``inplace=False``**, otherwise returns None.
Notes
-----
If the data object is in earth or principal coords, it is first
rotated to 'inst' before assigning inst2head_rotmat, it is then
rotated back to the coordinate system in which it was input. This
way the inst2head_rotmat gets applied correctly (in inst
coordinate system).
"""
# Create and return deep copy if not writing "in place"
if not inplace:
ds = ds.copy(deep=True)
if not ds.inst_model.lower() == "vector":
raise Exception(
"Setting 'inst2head_rotmat' is only supported " "for Nortek Vector ADVs."
)
if ds.get("inst2head_rotmat", None) is not None:
raise Exception(
"You are setting 'inst2head_rotmat' after it has already "
"been set. You can only set it once."
)
csin = ds.coord_sys
if csin not in ["inst", "beam"]:
rotate2(ds, "inst", inplace=True)
ds["inst2head_rotmat"] = xr.DataArray(
np.array(rotmat), dims=["x1", "x2"], coords={"x1": [1, 2, 3], "x2": [1, 2, 3]}
)
ds.attrs["inst2head_rotmat_was_set"] = 1 # logical
# Note that there is no validation that the user doesn't
# change `ds.attrs['inst2head_rotmat']` after calling this
# function.
if not csin == "beam": # csin not 'beam', then we're in inst
ds = r_vec._rotate_inst2head(ds)
if csin not in ["inst", "beam"]:
rotate2(ds, csin, inplace=True)
if not inplace:
return ds