From 82ab086b399bcd5a699750f74cacaa87762a49e2 Mon Sep 17 00:00:00 2001 From: louis-richard Date: Thu, 3 Aug 2023 22:43:37 +0200 Subject: [PATCH] add tests and improve coverage --- pyrfu/pyrf/cotrans.py | 108 ++++++-------- pyrfu/pyrf/ebsp.py | 25 ++-- pyrfu/pyrf/int_sph_dist.py | 83 +++++------ pyrfu/tests/test_pyrf.py | 298 +++++++++++++++++++++++++++++++++---- 4 files changed, 361 insertions(+), 153 deletions(-) diff --git a/pyrfu/pyrf/cotrans.py b/pyrfu/pyrf/cotrans.py index 8630626..4d64edb 100644 --- a/pyrfu/pyrf/cotrans.py +++ b/pyrfu/pyrf/cotrans.py @@ -12,7 +12,9 @@ # Local imports from ..models import igrf +from .ts_tensor_xyz import ts_tensor_xyz from .ts_vec_xyz import ts_vec_xyz +from .unix2datetime64 import unix2datetime64 __author__ = "Louis Richard" __email__ = "louisr@irfu.se" @@ -49,9 +51,8 @@ def _dipole_direction_gse(time, flag: str = "dipole"): np.sin(np.deg2rad(phi)), ], ).T - dipole_direction_gse_ = cotrans( - np.hstack([time[:, None], dipole_direction_geo_]), + ts_vec_xyz(unix2datetime64(time), dipole_direction_geo_), "geo>gse", ) @@ -67,6 +68,8 @@ def _transformation_matrix(t, tind, hapgood, *args): transf_mat_out[:, 2, 2] = np.ones(len(t)) for j, t_num in enumerate(tind[::-1]): + assert abs(t_num) in list(range(1, 6)), "t_num must be +/- 1, 2, 3, 4, 5" + if t_num in [-1, 1]: if hapgood: theta = 100.461 + 36000.770 * t_zero + 15.04107 * ut @@ -116,8 +119,8 @@ def _transformation_matrix(t, tind, hapgood, *args): elif t_num in [-3, 3]: dipole_direction_gse_ = _dipole_direction_gse(t, "dipole") - y_e = dipole_direction_gse_[:, 2] # 1st col is time - z_e = dipole_direction_gse_[:, 3] + y_e = dipole_direction_gse_[:, 1] # 1st col is time + z_e = dipole_direction_gse_[:, 2] psi = np.rad2deg(np.arctan(y_e / z_e)) transf_mat = _triang(-psi * np.sign(t_num), 0) # inverse if -3 @@ -126,24 +129,21 @@ def _transformation_matrix(t, tind, hapgood, *args): dipole_direction_gse_ = _dipole_direction_gse(t, "dipole") mu = np.arctan( - dipole_direction_gse_[:, 1] - / np.sqrt(np.sum(dipole_direction_gse_[:, 2:] ** 2, axis=1)), + dipole_direction_gse_[:, 0] + / np.sqrt(np.sum(dipole_direction_gse_[:, 1:] ** 2, axis=1)), ) mu = np.rad2deg(mu) transf_mat = _triang(-mu * np.sign(t_num), 1) - elif t_num in [-5, 5]: + else: lambda_, phi = igrf(t, "dipole") transf_mat = np.matmul(_triang(phi - 90, 1), _triang(lambda_, 2)) if t_num == -5: transf_mat = np.transpose(transf_mat, [0, 2, 1]) - else: - raise ValueError - - if j == len(tind): + if j == 0: transf_mat_out = transf_mat else: transf_mat_out = np.matmul(transf_mat, transf_mat_out) @@ -205,27 +205,31 @@ def cotrans(inp, flag, hapgood: bool = True): """ + assert isinstance(inp, xr.DataArray), "inp must be a xarray.DataArray" + assert inp.ndim < 3, "inp must be scalar or vector" + if ">" in flag: ref_syst_in, ref_syst_out = flag.split(">") else: ref_syst_in, ref_syst_out = [None, flag.lower()] - if isinstance(inp, xr.DataArray): - if "COORDINATE_SYSTEM" in inp.attrs: - ref_syst_internal = inp.attrs["COORDINATE_SYSTEM"].lower() - ref_syst_internal = ref_syst_internal.split(">")[0] - else: - ref_syst_internal = None - - if ref_syst_in is not None and ref_syst_internal is not None: - message = "input ref. frame in variable and input flag differs" - assert ref_syst_internal == ref_syst_in, message - elif ref_syst_in is None and ref_syst_internal is not None: - ref_syst_in = ref_syst_internal.lower() - elif ref_syst_in is None and ref_syst_internal is None: - raise ValueError("input reference frame undefined") + if "COORDINATE_SYSTEM" in inp.attrs: + ref_syst_internal = inp.attrs["COORDINATE_SYSTEM"].lower() + ref_syst_internal = ref_syst_internal.split(">")[0] + else: + ref_syst_internal = None + if ref_syst_in is not None and ref_syst_internal is not None: + message = "input ref. frame in variable and input flag differs" + assert ref_syst_internal == ref_syst_in, message + flag = f"{ref_syst_in}>{ref_syst_out}" + elif ref_syst_in is None and ref_syst_internal is not None: + ref_syst_in = ref_syst_internal.lower() flag = f"{ref_syst_in}>{ref_syst_out}" + elif flag.lower() == "dipoledirectiongse": + flag = flag.lower() + elif ref_syst_in is None and ref_syst_internal is None: + raise ValueError(f"Transformation {flag} is unknown!") if ref_syst_in == ref_syst_out: return inp @@ -234,24 +238,13 @@ def cotrans(inp, flag, hapgood: bool = True): j2000 = 946727930.8160001 # j2000 = Time("J2000", format="jyear_str").unix - if isinstance(inp, xr.DataArray): - time = inp.time.data - t = (time.astype(np.int64) * 1e-9).astype(np.float64) - - # Terrestial Time (seconds since J2000) - tts = t - j2000 - inp_ts = inp - inp = inp.data - - elif isinstance(inp, np.ndarray): - time = (inp[:, 0] * 1e9).astype("datetime64[ns]") - t = inp[:, 0] - # Terrestial Time (seconds since J2000) - tts = t - j2000 - inp_ts = None - inp = inp[:, 1:] - else: - raise TypeError("invalid input") + time = inp.time.data + t = (time.astype(np.int64) * 1e-9).astype(np.float64) + + # Terrestial Time (seconds since J2000) + tts = t - j2000 + inp_ts = inp + inp = inp.data if hapgood: day_start_epoch = time.astype("datetime64[D]") @@ -302,29 +295,18 @@ def cotrans(inp, flag, hapgood: bool = True): tind = transformation_dict[flag] - elif flag == "dipoledirectiongse": - out_data = _dipole_direction_gse(t) - return ts_vec_xyz(inp.time.data, out_data) + transf_mat = _transformation_matrix(t, tind, hapgood, *args_trans_mat) - else: - raise ValueError(f"Transformation {flag} is unknown!") + if inp.ndim == 1: + out = ts_tensor_xyz(inp_ts.time.data, transf_mat) - transf_mat = _transformation_matrix(t, tind, hapgood, *args_trans_mat) - - if inp.ndim == 2: - out = np.einsum("kji,ki->kj", transf_mat, inp) - elif inp.ndim == 1: - out = transf_mat - else: - raise ValueError - - if inp_ts is not None: - out_data = out - out = inp_ts.copy() - out.data = out_data - out.attrs["COORDINATE_SYSTEM"] = ref_syst_out.upper() + else: + out_data = np.einsum("kji,ki->kj", transf_mat, inp) + out = inp_ts.copy() + out.data = out_data + out.attrs["COORDINATE_SYSTEM"] = ref_syst_out.upper() else: - out = np.hstack([t[:, None], out]) + out = _dipole_direction_gse(t) return out diff --git a/pyrfu/pyrf/ebsp.py b/pyrfu/pyrf/ebsp.py index 0ba19fc..2a8e0af 100644 --- a/pyrfu/pyrf/ebsp.py +++ b/pyrfu/pyrf/ebsp.py @@ -113,21 +113,17 @@ def _freq_int(freq_int, delta_b): pc12_range, other_range = [False, False] if isinstance(freq_int, str): - if freq_int.lower() == "pc12": + if freq_int.lower() == "pc35": + freq_int = [0.002, 0.1] + + delta_t = 60 # local + else: pc12_range = True freq_int = [0.1, 5.0] delta_t = 1 # local - elif freq_int.lower() == "pc35": - freq_int = [0.002, 0.1] - - delta_t = 60 # local - - else: - raise ValueError("Invalid format of interval") - fs_out = 1 / delta_t else: if freq_int[1] >= freq_int[0]: @@ -443,9 +439,7 @@ def ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int, **kwargs): if fac_matrix is None: xyz = xyz[:-1, :] else: - fac_matrix["t"] = fac_matrix["t"][:-1, :] - - fac_matrix["rotMatrix"] = fac_matrix["rotMatrix"][:-1, :, :] + fac_matrix = fac_matrix[:-1, ...] if want_ee: e_xyz = e_xyz[:-1, :] @@ -480,9 +474,6 @@ def ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int, **kwargs): idx_nan_e = np.isnan(e_xyz.data) idx_nan_eisr2 = np.isnan(eisr2.data) - if e_xyz.shape[1] < 3: - raise IndexError("E must be a 3D vector to be rotated to FAC") - if fac_matrix is None: e_xyz = convert_fac(e_xyz, b_bgd, xyz) else: @@ -650,7 +641,9 @@ def ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int, **kwargs): arg_ = ts_vec_xyz(time_b0, np.transpose(tmp)) we = convert_fac(arg_, fac_matrix) else: - we = np.hstack([we[:, :2], we_z]) + we = np.transpose( + np.vstack([np.transpose(we[:, :2]), np.transpose(we_z)]) + ) power_e = 2 * np.pi * (we * np.conj(we)) / new_freq_mat power_e = np.vstack([power_e.T, np.sum(power_e, axis=1)]).T diff --git a/pyrfu/pyrf/int_sph_dist.py b/pyrfu/pyrf/int_sph_dist.py index 16a0983..453271f 100644 --- a/pyrfu/pyrf/int_sph_dist.py +++ b/pyrfu/pyrf/int_sph_dist.py @@ -75,15 +75,11 @@ def int_sph_dist(vdf, speed, phi, theta, speed_grid, **kwargs): # Overwrite projection dimension if azimuthal angle of projection # plane is not provided. Set the azimuthal angle grid width. - if phi_grid is None or projection_dim == "1d": - projection_dim = "1d" - d_phi_grid = 1.0 - elif phi_grid is not None and projection_dim.lower() in ["2d", "3d"]: + if phi_grid is not None and projection_dim.lower() in ["2d", "3d"]: d_phi_grid = np.median(np.diff(phi_grid)) else: - raise RuntimeError( - "1d projection with phi_grid provided doesn't make sense!!", - ) + projection_dim = "1d" + d_phi_grid = 1.0 # Make sure the transformation matrix is orthonormal. x_phat = xyz[:, 0] / np.linalg.norm(xyz[:, 0]) # re-normalize @@ -135,36 +131,9 @@ def int_sph_dist(vdf, speed, phi, theta, speed_grid, **kwargs): n_mc_mat = n_mc_mat.astype(int) - if projection_base == "pol": - # Area or line element (primed) - d_a_grid = speed_grid ** (int(projection_dim[0]) - 1) * d_phi_grid * d_v_grid - d_a_grid = d_a_grid.astype(np.float64) - - if projection_dim == "1d": - f_g = mc_pol_1d( - vdf, - speed, - phi, - theta, - d_v, - d_v_m, - d_phi, - d_theta, - speed_grid_edges, - d_a_grid, - v_lim, - a_lim, - n_mc_mat, - r_mat, - ) - else: - raise NotImplementedError( - "2d projection on polar grid is not ready yet!!", - ) - - elif projection_base == "cart" and projection_dim == "2d": + if projection_base == "cart" and projection_dim == "2d": d_a_grid = d_v_grid**2 - f_g = mc_cart_2d( + f_g = _mc_cart_2d( vdf, speed, phi, @@ -182,7 +151,7 @@ def int_sph_dist(vdf, speed, phi, theta, speed_grid, **kwargs): ) elif projection_base == "cart" and projection_dim == "3d": d_a_grid = d_v_grid**3 - f_g = mc_cart_3d( + f_g = _mc_cart_3d( vdf, speed, phi, @@ -199,11 +168,33 @@ def int_sph_dist(vdf, speed, phi, theta, speed_grid, **kwargs): r_mat, ) else: - raise ValueError("Invalid base!!") + # Area or line element (primed) + d_a_grid = speed_grid ** (int(projection_dim[0]) - 1) * d_phi_grid * d_v_grid + d_a_grid = d_a_grid.astype(np.float64) - if projection_dim == "1d": - pst = {"f": f_g, "vx": speed_grid, "vx_edges": speed_grid_edges} - elif projection_dim == "2d" and projection_base == "cart": + if projection_dim == "1d": + f_g = _mc_pol_1d( + vdf, + speed, + phi, + theta, + d_v, + d_v_m, + d_phi, + d_theta, + speed_grid_edges, + d_a_grid, + v_lim, + a_lim, + n_mc_mat, + r_mat, + ) + else: + raise NotImplementedError( + "2d projection on polar grid is not ready yet!!", + ) + + if projection_dim == "2d" and projection_base == "cart": pst = { "f": f_g, "vx": speed_grid, @@ -222,15 +213,13 @@ def int_sph_dist(vdf, speed, phi, theta, speed_grid, **kwargs): "vz_edges": speed_grid_edges, } else: - raise NotImplementedError( - "2d projection on polar grid is not ready yet!!", - ) + pst = {"f": f_g, "vx": speed_grid, "vx_edges": speed_grid_edges} return pst @numba.jit(cache=True, nogil=True, parallel=True, nopython=True) -def mc_pol_1d( +def _mc_pol_1d( vdf, v, phi, @@ -341,7 +330,7 @@ def mc_pol_1d( @numba.jit(cache=True, nogil=True, parallel=True, nopython=True) -def mc_cart_3d( +def _mc_cart_3d( vdf, v, phi, @@ -454,7 +443,7 @@ def mc_cart_3d( @numba.jit(cache=True, nogil=True, parallel=True, nopython=True) -def mc_cart_2d( +def _mc_cart_2d( vdf, v, phi, diff --git a/pyrfu/tests/test_pyrf.py b/pyrfu/tests/test_pyrf.py index bdb0e90..c7130e3 100644 --- a/pyrfu/tests/test_pyrf.py +++ b/pyrfu/tests/test_pyrf.py @@ -13,6 +13,10 @@ from ddt import data, ddt, idata, unpack from pyrfu import pyrf +from pyrfu.pyrf.compress_cwt import _compress_cwt_1d +from pyrfu.pyrf.ebsp import _average_data, _censure_plot, _freq_int +from pyrfu.pyrf.int_sph_dist import _mc_cart_2d, _mc_cart_3d, _mc_pol_1d +from pyrfu.pyrf.wavelet import _power_c, _power_r __author__ = "Louis Richard" __email__ = "louisr@irfu.se" @@ -22,10 +26,11 @@ __status__ = "Prototype" -def generate_timeline(f_s, n_pts: int = 10000): - ref_time = np.datetime64("2019-01-01T00:00:00.000") +def generate_timeline(f_s, n_pts: int = 10000, dtype="datetime64[ns]"): + ref_time = np.datetime64("2019-01-01T00:00:00.000000000") times = [ref_time + np.timedelta64(int(i * 1e9 / f_s), "ns") for i in range(n_pts)] - return np.array(times) + times = np.array(times).astype(dtype) + return times def generate_data(n_pts, kind: str = "scalar"): @@ -41,14 +46,18 @@ def generate_data(n_pts, kind: str = "scalar"): return data -def generate_ts(f_s, n_pts, kind: str = "scalar"): +def generate_ts(f_s, n_pts, kind: str = "scalar", attrs: dict = {}): if kind == "scalar": - out = pyrf.ts_scalar(generate_timeline(f_s, n_pts), generate_data(n_pts, kind)) + out = pyrf.ts_scalar( + generate_timeline(f_s, n_pts), generate_data(n_pts, kind), attrs=attrs + ) elif kind == "vector": - out = pyrf.ts_vec_xyz(generate_timeline(f_s, n_pts), generate_data(n_pts, kind)) + out = pyrf.ts_vec_xyz( + generate_timeline(f_s, n_pts), generate_data(n_pts, kind), attrs=attrs + ) elif kind == "tensor": out = pyrf.ts_tensor_xyz( - generate_timeline(f_s, n_pts), generate_data(n_pts, kind) + generate_timeline(f_s, n_pts), generate_data(n_pts, kind), attrs=attrs ) else: raise ValueError("Invalid kind of data!!") @@ -561,6 +570,12 @@ def test_compress_cwt_output(self): self.assertIsInstance(result[1], np.ndarray) self.assertIsInstance(result[2], np.ndarray) + def test_compress_cwt_1d(self): + result = _compress_cwt_1d.__wrapped__( + np.random.random((1000, 100)), random.randint(2, 100) + ) + self.assertIsInstance(result, np.ndarray) + class ConvertFACTestCase(unittest.TestCase): def test_convert_fac_input(self): @@ -626,39 +641,55 @@ def test_convert_fac_output(self): @ddt class CotransTestCase(unittest.TestCase): - def test_cotrans_input(self): - with self.assertRaises(TypeError): - pyrf.cotrans(0.0, "gse>gsm") - - with self.assertRaises(IndexError): - pyrf.cotrans(generate_data(100), "gse>gsm") - - with self.assertRaises(ValueError): - pyrf.cotrans(generate_ts(64.0, 100, "vector"), "gsm") + @data( + (0.0, "gse>gsm", True), + (generate_data(100), "gse>gsm", True), + (generate_ts(64.0, 100, "vector"), "gsm", True), + (generate_ts(64.0, 100, "tensor"), "gse>gsm", True), + ( + generate_ts(64.0, 100, "vector", {"COORDINATE_SYSTEM": "gse"}), + "gsm>sm", + True, + ), + ) + @unpack + def test_cotrans_input(self, inp, flag, hapgood): + with self.assertRaises((TypeError, IndexError, ValueError, AssertionError)): + pyrf.cotrans(inp, flag, hapgood) - @idata(itertools.permutations(["gei", "geo", "gse", "gsm", "mag", "sm"], 2)) - def test_cotrans_output(self, value): + @idata(itertools.product(["gei", "geo", "gse", "gsm", "mag", "sm"], repeat=2)) + def test_cotrans_output_trans(self, value): transf = f"{value[0]}>{value[1]}" - result = pyrf.cotrans(generate_ts(64.0, 100, "vector"), transf) + + inp = generate_ts(64.0, 100, "vector") + result = pyrf.cotrans(inp, transf, hapgood=True) self.assertIsInstance(result, xr.DataArray) self.assertListEqual(list(result.shape), [100, 3]) - result = pyrf.cotrans(generate_ts(64.0, 100, "vector"), transf, hapgood=False) + result = pyrf.cotrans(inp, transf, hapgood=False) self.assertIsInstance(result, xr.DataArray) self.assertListEqual(list(result.shape), [100, 3]) - inp = generate_ts(64.0, 100, "vector") inp.attrs["COORDINATE_SYSTEM"] = value[0] + result = pyrf.cotrans(inp, transf, hapgood=False) + self.assertIsInstance(result, xr.DataArray) + self.assertListEqual(list(result.shape), [100, 3]) + result = pyrf.cotrans(inp, value[1], hapgood=False) self.assertIsInstance(result, xr.DataArray) self.assertListEqual(list(result.shape), [100, 3]) - inp = generate_ts(64.0, 100, "vector") - inp.attrs["COORDINATE_SYSTEM"] = value[0] result = pyrf.cotrans(inp, value[1], hapgood=True) self.assertIsInstance(result, xr.DataArray) self.assertListEqual(list(result.shape), [100, 3]) + def test_cotrans_output_exot(self): + inp = generate_ts(64.0, 100, "scalar") + result = pyrf.cotrans(inp, "gse>gsm", hapgood=True) + self.assertIsInstance(result, xr.DataArray) + result = pyrf.cotrans(inp, "dipoledirectiongse", hapgood=True) + self.assertIsInstance(result, xr.DataArray) + class CrossTestCase(unittest.TestCase): def test_cross_input(self): @@ -895,6 +926,7 @@ class EbspTestCase(unittest.TestCase): generate_ts(64.0, 100, "vector"), generate_ts(64.0, 100, "vector"), [1e0, 1e1], + {}, ), ( generate_ts(64.0, 97, "vector"), @@ -903,6 +935,34 @@ class EbspTestCase(unittest.TestCase): generate_ts(64.0, 100, "vector"), generate_ts(64.0, 120, "vector"), [1e0, 1e1], + {}, + ), + ( + generate_ts(99.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + [1e0, 1e1], + {}, + ), + ( + generate_ts(40.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(40.0, 100, "vector"), + generate_ts(40.0, 100, "vector"), + generate_ts(40.0, 100, "vector"), + [1e0, 1e1], + {}, + ), + ( + generate_ts(64.0, 97, "vector"), + generate_ts(64.0, 97, "vector"), + generate_ts(64.0, 97, "vector"), + generate_ts(64.0, 97, "vector"), + generate_ts(64.0, 97, "vector"), + [1e0, 1e1], + {}, ), ( generate_ts(64.0, 97, "vector"), @@ -911,6 +971,7 @@ class EbspTestCase(unittest.TestCase): generate_ts(64.0, 97, "vector"), generate_ts(64.0, 97, "vector"), [1e0, 1e1], + {"fac_matrix": generate_ts(64.0, 100, "tensor")}, ), ( generate_ts(64.0, 100, "vector"), @@ -919,11 +980,13 @@ class EbspTestCase(unittest.TestCase): generate_ts(64.0, 100, "vector"), None, [1e0, 1e1], + {}, ), ) @unpack - def test_ebsp_input_pass(self, e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int): - self.assertIsNotNone(pyrf.ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int)) + def test_ebsp_input_pass(self, e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int, options): + result = pyrf.ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int, **options) + self.assertIsInstance(result, dict) @data( ( @@ -958,10 +1021,26 @@ def test_ebsp_input_pass(self, e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int): generate_ts(64.0, 100, "vector"), [1e0, 1e1], ), + ( + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + [1e1, 1e0], + ), + ( + generate_ts(64.0, 100, "vector")[:, :2], + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + generate_ts(64.0, 100, "vector"), + [1e0, 1e1], + ), ) @unpack def test_ebsp_input_fail(self, e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int): - with self.assertRaises((AssertionError, TypeError, IndexError)): + with self.assertRaises((AssertionError, TypeError, IndexError, ValueError)): pyrf.ebsp(e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int) @data( @@ -1005,6 +1084,16 @@ def test_ebsp_input_fail(self, e_xyz, db_xyz, b_xyz, b_bgd, xyz, freq_int): "fac_matrix": None, "m_width_coeff": 1, }, + { + "polarization": False, + "no_resample": False, + "fac": False, + "de_dot_b0": True, + "full_b_db": False, + "nav": 8, + "fac_matrix": None, + "m_width_coeff": 1, + }, { "polarization": False, "no_resample": False, @@ -1105,6 +1194,34 @@ def test_ebsp_freq_int_fail(self, value): value, ) + @data(([0.32, 3.2], generate_ts(64.0, 100000, "vector"))) + @unpack + def test_average_data(self, freq_int, data): + _, _, _, out_time = _freq_int(freq_int, data) + in_time = data.time.data.astype(np.float64) / 1e9 + + result = _average_data.__wrapped__(data.data, in_time, out_time, None) + self.assertIsInstance(result, np.ndarray) + self.assertListEqual(list(result.shape), [len(out_time), 3]) + + @data(([0.32, 3.2], generate_ts(64.0, 100000, "vector"))) + @unpack + def test_censure_plot(self, freq_int, data): + _, _, out_sampling, out_time = _freq_int(freq_int, data) + a_ = np.logspace(1, 2, 12) + idx_nan = np.full(len(data), False) + idx_nan[np.random.randint(100000, size=(100))] = True + censure = np.floor(2 * a_ * out_sampling / 64.0 * 8) + result = _censure_plot.__wrapped__( + np.random.random((len(out_time), len(a_))), + idx_nan, + censure, + len(data), + a_, + ) + self.assertIsInstance(result, np.ndarray) + self.assertListEqual(list(result.shape), [len(out_time), len(a_)]) + class EndTestCase(unittest.TestCase): def test_end_input(self): @@ -1339,11 +1456,118 @@ class IncrementsTestCase(unittest.TestCase): generate_ts(64.0, 100, "tensor"), ) def test_increments_output(self, value): - result = pyrf.increments(value, random.randint(1, 99)) + result = pyrf.increments(value, random.randint(1, 50)) self.assertIsInstance(result[0], np.ndarray) self.assertIsInstance(result[1], xr.DataArray) +@ddt +class IntSphDistTestCase(unittest.TestCase): + @data( + {"projection_base": "pol", "projection_dim": "2d"}, + ) + def test_int_sph_dist_input(self, value): + vdf = np.random.random((51, 32, 16)) + speed = np.linspace(0, 1, 51) + phi = np.arange(32) + theta = np.arange(16) + speed_grid = np.linspace(-1, 1, 101) + + with self.assertRaises((RuntimeError, NotImplementedError)): + pyrf.int_sph_dist(vdf, speed, phi, theta, speed_grid, **value) + + @data( + {}, + {"weight": "lin"}, + {"weight": "log"}, + {"speed_edges": np.linspace(-0.01, 1.01, 52)}, + {"speed_grid_edges": np.linspace(-1.01, 1.01, 102)}, + { + "phi_grid": np.arange(0, 32), + "projection_base": "cart", + "projection_dim": "2d", + }, + {"projection_base": "cart", "projection_dim": "3d"}, + ) + def test_int_sph_dist_output(self, value): + vdf = np.random.random((51, 32, 16)) + speed = np.linspace(0, 1, 51) + phi = np.arange(32) + theta = np.arange(16) + speed_grid = np.linspace(-1, 1, 101) + result = pyrf.int_sph_dist(vdf, speed, phi, theta, speed_grid, **value) + self.assertIsInstance(result, dict) + + @data( + ( + np.random.random((51, 32, 16)), + np.linspace(0, 1, 51), + np.arange(32), + np.arange(16), + np.ones(51) * 0.02, + np.ones(51) * 0.01, + np.ones(32), + np.ones(16), + np.linspace(-1.01, 1.01, 102), + np.ones(101) * 0.02 * np.pi / 16, + np.array([-np.inf, np.inf]), + np.array([-np.pi, np.pi]), + np.ones((51, 32, 16), dtype=np.int64) * 10, + np.eye(3), + ) + ) + def test_mc_pol_1d(self, value): + vdf, *args = value + vdf[vdf < 1e-2] = 0 + self.assertIsInstance(_mc_pol_1d.__wrapped__(vdf, *args), np.ndarray) + + @data( + ( + np.random.random((51, 32, 16)), + np.linspace(0, 1, 51), + np.arange(32), + np.arange(16), + np.ones(51) * 0.02, + np.ones(51) * 0.01, + np.ones(32), + np.ones(16), + np.linspace(-1.01, 1.01, 102), + 0.02**2, + np.array([-np.inf, np.inf]), + np.array([-np.pi, np.pi]), + (np.ones((51, 32, 16), dtype=np.int64) * 10).astype(int), + np.eye(3), + ) + ) + def test_mc_cart_2d(self, value): + vdf, *args = value + vdf[vdf < 1e-2] = 0 + self.assertIsInstance(_mc_cart_2d.__wrapped__(vdf, *args), np.ndarray) + + @data( + ( + np.random.random((51, 32, 16)), + np.linspace(0, 1, 51), + np.arange(32), + np.arange(16), + np.ones(51) * 0.02, + np.ones(51) * 0.01, + np.ones(32), + np.ones(16), + np.linspace(-1.01, 1.01, 102), + 0.02**2, + np.array([-np.inf, np.inf]), + np.array([-np.pi, np.pi]), + (np.ones((51, 32, 16), dtype=np.int64) * 10).astype(int), + np.eye(3), + ) + ) + def test_mc_cart_3d(self, value): + vdf, *args = value + vdf[vdf < 1e-2] = 0 + self.assertIsInstance(_mc_cart_3d.__wrapped__(vdf, *args), np.ndarray) + + @ddt class Iso86012Unix(unittest.TestCase): @data( @@ -1774,6 +1998,26 @@ def test_wavelet_input(self, inp, options): def test_wavelet_output(self, inp, options): self.assertIsNotNone(pyrf.wavelet(inp, **options)) + @data( + ( + np.random.random((100, 3)) + np.random.random((100, 3)) * 1j, + np.random.random((100, 3)), + ) + ) + @unpack + def test_power_r(self, power, new_freq_mat): + self.assertIsNotNone(_power_r.__wrapped__(power, new_freq_mat)) + + @data( + ( + np.random.random((100, 3)) + np.random.random((100, 3)) * 1j, + np.random.random((100, 3)), + ) + ) + @unpack + def test_power_c(self, power, new_freq_mat): + self.assertIsNotNone(_power_c.__wrapped__(power, new_freq_mat)) + if __name__ == "__main__": unittest.main()