diff --git a/tests/test_order_constit.py b/tests/test_order_constit.py new file mode 100644 index 0000000..8c87e45 --- /dev/null +++ b/tests/test_order_constit.py @@ -0,0 +1,77 @@ +import numpy as np +import pytest + +from utide import reconstruct, solve +from utide._ut_constants import constit_index_dict, ut_constants + + +ts = 735604 +duration = 35 + +time = np.linspace(ts, ts + duration, 842) +tref = (time[-1] + time[0]) / 2 + +const = ut_constants.const + +amps = [1.0, 0.5, 0.6, 0.1] +names = ["M2", "S2", "K1", "O1"] +cpds = [24 * const.freq[constit_index_dict[name]] for name in names] +sinusoids = [] +for amp, cpd in zip(amps, cpds): + arg = 2 * np.pi * (time - tref) * cpd + sinusoids.append(amp * np.cos(arg)) +tide = np.hstack(tuple(sinusoids)).sum(axis=0) + +np.random.seed(1234) +noise = 1e-3 * np.random.randn(len(time)) + +time_series = tide + noise + +opts0 = { + "constit": ["K1", "M2", "O1", "S2"], + "order_constit": "frequency", + "phase": "raw", + "nodal": False, + "trend": False, + "method": "ols", + "conf_int": "MC", + "Rayleigh_min": 0.95, +} + + +@pytest.mark.parametrize("conf_int", ["none", "linear", "MC"]) +def test_order(conf_int): + + orders = [None, "PE", "frequency", opts0["constit"]] + if conf_int != "none": + orders.append("SNR") + elevs = [] + ts_elevs = [] + vels = [] + ts_vels = [] + for order in orders: + opts = opts0.copy() + opts["order_constit"] = order + opts["conf_int"] = conf_int + elevs.append(solve(time, time_series, lat=45, **opts)) + vels.append(solve(time, time_series, time_series, lat=45, **opts)) + ts_elevs.append(reconstruct(time, elevs[-1], min_SNR=0)) + ts_vels.append(reconstruct(time, vels[-1], min_SNR=0)) + + # Are the reconstructions all the same? + for i in range(1, len(elevs)): + assert (ts_elevs[i].h == ts_elevs[0].h).all() + assert (ts_vels[i].u == ts_vels[0].u).all() + assert (ts_vels[i].v == ts_vels[0].v).all() + + # Is None equivalent to "PE"? (Just a spot check.) + assert (elevs[0].name == elevs[1].name).all() + assert (elevs[0].A == elevs[1].A).all() + + +def test_invalid_snr(): + opts = opts0.copy() + opts["conf_int"] = "none" + opts["order_constit"] = "SNR" + with pytest.raises(ValueError): + solve(time, time_series, lat=45, **opts) diff --git a/utide/_solve.py b/utide/_solve.py index 675f3d5..20ed0e8 100644 --- a/utide/_solve.py +++ b/utide/_solve.py @@ -5,10 +5,9 @@ import numpy as np from ._time_conversion import _normalize_time -from ._ut_constants import constit_index_dict from .confidence import _confidence from .constituent_selection import ut_cnstitsel -from .diagnostics import ut_diagn +from .diagnostics import _PE, _SNR, ut_diagn from .ellipse_params import ut_cs2cep from .harmonics import ut_E from .robustfit import robustfit @@ -17,6 +16,7 @@ default_opts = { "constit": "auto", + "order_constit": None, "conf_int": "linear", "method": "ols", "trend": True, @@ -37,6 +37,8 @@ def _process_opts(opts, is_2D): newopts.update_values(strict=True, **opts) # TODO: add more validations. newopts.infer = validate_infer(newopts.infer, is_2D) + snr = newopts.conf_int != "none" + newopts.order_constit = validate_order_constit(newopts.order_constit, snr) compat_opts = _translate_opts(newopts) @@ -48,6 +50,7 @@ def _translate_opts(opts): # Here or elsewhere, proper validation remains to be added. oldopts = Bunch() oldopts.cnstit = opts.constit + oldopts.ordercnstit = opts.order_constit oldopts.infer = opts.infer # we will not use the matlab names, though oldopts.conf_int = True @@ -101,6 +104,22 @@ def validate_infer(infer, is_2D): return infer +def validate_order_constit(arg, have_snr): + available = ["PE", "frequency"] + if have_snr: + available.append("SNR") + if arg is None: + return "PE" + if isinstance(arg, str) and arg in available: + return arg + if not isinstance(arg, str) and np.iterable(arg): + return arg # TODO: add checking of its elements + raise ValueError( + f"order_constit must be one of {available} or" + f" a sequence of constituents, not '{arg}'", + ) + + def solve(t, u, v=None, lat=None, **opts): """ Calculate amplitude, phase, confidence intervals of tidal constituents. @@ -122,7 +141,7 @@ def solve(t, u, v=None, lat=None, **opts): standard library `datetime` proleptic Gregorian calendar, starting with 1 at 00:00 on January 1 of year 1; this is the 'datenum' used by Matplotlib. - constit : {'auto', array_like}, optional + constit : {'auto', sequence}, optional List of strings with standard letter abbreviations of tidal constituents; or 'auto' to let the list be determined based on the time span. @@ -165,6 +184,14 @@ def solve(t, u, v=None, lat=None, **opts): amp_ratios and phase_offsets have length N for a scalar time series, or 2N for a vector series. + order_constit : {'PE', 'SNR', 'frequency', sequence}, optional + The default is 'PE' (percent energy) order, returning results ordered from + high energy to low. + The 'SNR' order is from high signal-to-noise ratio to low, and is + available only if `conf_int` is not 'none'. The + 'frequency' order is from low to high frequency. Alternatively, a + sequence of constituent names may be supplied, typically the same list as + given in the *constit* option. MC_n : integer, optional Not yet implemented. robust_kw : dict, optional @@ -370,7 +397,7 @@ def _solv1(tin, uin, vin, lat, **opts): coef.theta = np.hstack((coef.theta, theta)) coef.g = np.hstack((coef.g, g)) - if opt["conf_int"] is True: + if opt["conf_int"]: coef = _confidence( coef, cnstit, @@ -392,63 +419,50 @@ def _solv1(tin, uin, vin, lat, **opts): # Diagnostics. if not opt["nodiagn"]: - coef, indPE = ut_diagn(coef, opt) + coef = ut_diagn(coef) + # Adds a diagn dictionary, always sorted by energy. + # This doesn't seem very useful. Let's directly add the variables + # to the base coef structure. Then they can be sorted with everything + # else. + coef["PE"] = _PE(coef) + coef["SNR"] = _SNR(coef) # Re-order constituents. - if opt["ordercnstit"] is not None: + coef = _reorder(coef, opt) + # This might have added PE if it was not already present. - if opt["ordercnstit"] == "frq": - ind = coef["aux"]["frq"].argsort() + if opt["RunTimeDisp"]: + print("done.") - elif opt["ordercnstit"] == "snr": - if not opt["nodiagn"]: - ind = coef["diagn"]["SNR"].argsort()[::-1] - else: - if opt["twodim"]: - SNR = (coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) / ( - (coef["Lsmaj_ci"] / 1.96) ** 2 + (coef["Lsmin_ci"] / 1.96) ** 2 - ) + return coef - else: - SNR = (coef["A"] ** 2) / (coef["A_ci"] / 1.96) ** 2 - ind = SNR.argsort()[::-1] +def _reorder(coef, opt): + if opt["ordercnstit"] == "PE": + # Default: order by decreasing energy. + if "PE" not in coef: + coef["PE"] = _PE(coef) + ind = coef["PE"].argsort()[::-1] - else: - ilist = [constit_index_dict[name] for name in opt["ordercnstit"]] - ind = np.array(ilist, dtype=int) + elif opt["ordercnstit"] == "frequency": + ind = coef["aux"]["frq"].argsort() - else: # Default: order by decreasing energy. - if not opt["nodiagn"]: - ind = indPE - else: - if opt["twodim"]: - PE = np.sum(coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) - PE = 100 * (coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) / PE - else: - PE = 100 * coef["A"] ** 2 / np.sum(coef["A"] ** 2) - - ind = PE.argsort()[::-1] - - reorderlist = ["g", "name"] - if opt.twodim: - reorderlist.extend(["Lsmaj", "Lsmin", "theta"]) - if opt.conf_int: - reorderlist.extend(["Lsmaj_ci", "Lsmin_ci", "theta_ci", "g_ci"]) + elif opt["ordercnstit"] == "SNR": + # If we are here, we should be guaranteed to have SNR already. + ind = coef["SNR"].argsort()[::-1] else: - reorderlist.append("A") - if opt.conf_int: - reorderlist.extend(["A_ci", "g_ci"]) + namelist = list(coef["name"]) + ilist = [namelist.index(name) for name in opt["ordercnstit"]] + ind = np.array(ilist, dtype=int) + + arrays = "name PE SNR A A_ci g g_ci Lsmaj Lsmaj_ci Lsmin Lsmin_ci theta theta_ci" + reorderlist = [a for a in arrays.split() if a in coef] for key in reorderlist: coef[key] = coef[key][ind] coef["aux"]["frq"] = coef["aux"]["frq"][ind] coef["aux"]["lind"] = coef["aux"]["lind"][ind] - - if opt["RunTimeDisp"]: - print("done.") - return coef @@ -532,7 +546,7 @@ def _slvinit(tin, uin, vin, lat, **opts): opt["rmin"] = 1 opt["method"] = "ols" opt["tunrdn"] = 1 - opt["linci"] = 0 + opt["linci"] = False opt["white"] = 0 opt["nrlzn"] = 200 opt["lsfrqosmp"] = 1 diff --git a/utide/diagnostics.py b/utide/diagnostics.py index bae848b..b6a250a 100644 --- a/utide/diagnostics.py +++ b/utide/diagnostics.py @@ -1,58 +1,44 @@ import numpy as np -def ut_diagn(coef, opt): - - if opt["RunTimeDisp"]: - print("diagnostics ... ", end="") - coef["diagn"] = {} +def _PE(coef): + """ + Return the energy percentage for each constituent. + """ + if "Lsmaj" in coef: + E = coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2 + PE = (100 / np.sum(E)) * E + else: + PE = 100 * coef["A"] ** 2 / np.sum(coef["A"] ** 2) + return PE - if opt["twodim"]: - PE = np.sum(coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) - PE = 100 * (coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) / PE +def _SNR(coef): + """ + Return the signal-to-noise ratio for each constituent. + """ + if "Lsmaj" in coef: SNR = (coef["Lsmaj"] ** 2 + coef["Lsmin"] ** 2) / ( (coef["Lsmaj_ci"] / 1.96) ** 2 + (coef["Lsmin_ci"] / 1.96) ** 2 ) - else: - PE = 100 * coef["A"] ** 2 / np.sum(coef["A"] ** 2) SNR = (coef["A"] ** 2) / (coef["A_ci"] / 1.96) ** 2 + return SNR + +def ut_diagn(coef): + """ + Add to coef the names, PE, and SNR, *always* sorted by energy. + + To be eliminated... + """ + coef["diagn"] = {} + PE = _PE(coef) + SNR = _SNR(coef) indPE = PE.argsort()[::-1] coef["diagn"]["name"] = coef["name"][indPE] coef["diagn"]["PE"] = PE[indPE] coef["diagn"]["SNR"] = SNR[indPE] - return coef, indPE - - -# [~,indPE] = sort(PE,'descend'); -# coef.diagn.name = coef.name(indPE); -# coef.diagn.PE = PE(indPE); -# coef.diagn.SNR = SNR; % used in ut_diagntable; ordered by PE there -# if opt.twodim -# [coef.diagn,usnrc,vsnrc] = ut_diagntable(coef,cnstit,... -# t,u,v,xmod,m,B,W,varMSM,Gall,Hall,elor,varcov_mCw,indPE); -# else -# [coef.diagn,usnrc,~] = ut_diagntable(coef,cnstit,... -# t,u,[],xmod,m,B,W,varMSM,Gall,Hall,elor,varcov_mCw,indPE); -# end -# if opt.diagnplots -# tmp = nan*ones(size(uin)); -# tmp(uvgd) = usnrc; -# usnrc = tmp; -# tmp = nan*ones(size(uin)); -# tmp(uvgd) = e; -# e = tmp; -# if opt.twodim -# tmp = nan*ones(size(uin)); -# tmp(uvgd) = vsnrc; -# vsnrc = tmp; -# ut_diagnfigs(coef,indPE,tin,uin,vin,usnrc,vsnrc,e); -# else -# ut_diagnfigs(coef,indPE,tin,uin,[],usnrc,[],e); -# end -# end -# end + return coef