add tests

pyrfu/pyrf/autocorr.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-# #rd party imports
+# 3rd party imports
 import numpy as np
 import xarray as xr
 
@@ -22,7 +22,7 @@ def autocorr(inp, maxlags: int = None, normed: bool = True):
     Parameters
     ----------
     inp : xarray.DataArray
-        Input time series.
+        Input time series (scalar of vector).
     maxlags : int, Optional
         Maximum lag in number of points. Default is None (i.e., len(inp) - 1).
     normed : bool, Optional
@@ -35,6 +35,12 @@ def autocorr(inp, maxlags: int = None, normed: bool = True):
 
     """
 
+    # Check input type
+    assert isinstance(inp, xr.DataArray), "inp must be a xarray.DataArray"
+
+    # Check input dimension (scalar or vector)
+    assert inp.ndim < 3, "inp must be a scalar or a vector"
+
     if inp.ndim == 1:
         x = inp.data[:, None]
     else:
@@ -46,19 +52,12 @@ def autocorr(inp, maxlags: int = None, normed: bool = True):
         maxlags = n_t - 1
 
     if maxlags >= n_t or maxlags < 1:
-        raise ValueError(
-            f"maxlags must be None or strictly positive < {n_t:d}",
-        )
+        raise ValueError(f"maxlags must be None or strictly positive < {n_t:d}")
 
-    lags = np.linspace(
-        -float(maxlags),
-        float(maxlags),
-        2 * maxlags + 1,
-        dtype=int,
-    )
+    lags = np.linspace(-float(maxlags), float(maxlags), 2 * maxlags + 1, dtype=int)
     lags = lags * calc_dt(inp)
 
-    out_data = np.zeros_like(x)
+    out_data = np.zeros((maxlags + 1, x.shape[1]))
 
     for i in range(x.shape[1]):
         correls = np.correlate(x[:, i], x[:, i], mode="full")
@@ -75,13 +74,11 @@ def autocorr(inp, maxlags: int = None, normed: bool = True):
             coords=[lags[lags >= 0]],
             dims=["lag"],
         )
-    elif inp.ndim == 2:
+    else:
         out = xr.DataArray(
             out_data,
             coords=[lags[lags >= 0], inp[inp.dims[1]].data],
             dims=["lag", inp.dims[1]],
         )
-    else:
-        raise ValueError("invalid shape!!")
 
     return out

pyrfu/pyrf/calc_ag.py

@@ -3,6 +3,10 @@
 
 # 3rd party imports
 import numpy as np
+import xarray as xr
+
+# Local imports
+from .ts_scalar import ts_scalar
 
 __author__ = "Louis Richard"
 __email__ = "[email protected]"
@@ -67,9 +71,16 @@ def calc_ag(p_xyz):
 
     """
 
+    # Check input type
+    assert isinstance(p_xyz, xr.DataArray), "p_xyz must be a xarray.DataArray"
+
+    # Check import shape
+    message = "p_xyz must be a time series of a tensor"
+    assert p_xyz.data.ndim == 3 and p_xyz.shape[1] == 3 and p_xyz.shape[2] == 3, message
+
     # Diagonal and off-diagonal terms
-    p_11, p_22, _ = [p_xyz[:, 0, 0], p_xyz[:, 1, 1], p_xyz[:, 2, 2]]
-    p_12, p_13, p_23 = [p_xyz[:, 0, 1], p_xyz[:, 0, 2], p_xyz[:, 1, 2]]
+    p_11, p_22, _ = [p_xyz.data[:, 0, 0], p_xyz.data[:, 1, 1], p_xyz.data[:, 2, 2]]
+    p_12, p_13, p_23 = [p_xyz.data[:, 0, 1], p_xyz.data[:, 0, 2], p_xyz.data[:, 1, 2]]
 
     det_p = p_11 * (p_22**2 - p_23**2)
     det_p -= p_12 * (p_12 * p_22 - p_23 * p_13)
@@ -78,5 +89,6 @@ def calc_ag(p_xyz):
     det_g = p_11 * p_22**2
 
     agyrotropy = np.abs(det_p - det_g) / (det_p + det_g)
+    agyrotropy = ts_scalar(p_xyz.time.data, agyrotropy)
 
     return agyrotropy

pyrfu/pyrf/calc_agyro.py

@@ -3,6 +3,10 @@
 
 # 3rd party imports
 import numpy as np
+import xarray as xr
+
+# Local imports
+from .ts_scalar import ts_scalar
 
 __author__ = "Louis Richard"
 __email__ = "[email protected]"
@@ -58,9 +62,17 @@ def calc_agyro(p_xyz):
 
     """
 
+    # Check input type
+    assert isinstance(p_xyz, xr.DataArray), "p_xyz must be a xarray.DataArray"
+
+    # Check import shape
+    message = "p_xyz must be a time series of a tensor"
+    assert p_xyz.data.ndim == 3 and p_xyz.shape[1] == 3 and p_xyz.shape[2] == 3, message
+
     # Parallel and perpendicular components
-    p_perp_1, p_perp_2 = [p_xyz[:, 1, 1], p_xyz[:, 2, 2]]
+    p_perp_1, p_perp_2 = [p_xyz.data[:, 1, 1], p_xyz.data[:, 2, 2]]
 
-    agyro = np.abs(p_perp_1 - p_perp_2) / (p_perp_1 + p_perp_2)
+    agyrotropy = np.abs(p_perp_1 - p_perp_2) / (p_perp_1 + p_perp_2)
+    agyrotropy = ts_scalar(p_xyz.time.data, agyrotropy)
 
-    return agyro
+    return agyrotropy

pyrfu/pyrf/calc_dng.py

@@ -3,6 +3,10 @@
 
 # 3rd party imports
 import numpy as np
+import xarray as xr
+
+# Local imports
+from .ts_scalar import ts_scalar
 
 __author__ = "Louis Richard"
 __email__ = "[email protected]"
@@ -65,13 +69,22 @@ def calc_dng(p_xyz):
 
     """
 
+    # Check input type
+    assert isinstance(p_xyz, xr.DataArray), "p_xyz must be a xarray.DataArray"
+
+    # Check import shape
+    message = "p_xyz must be a time series of a tensor"
+    assert p_xyz.data.ndim == 3 and p_xyz.shape[1] == 3 and p_xyz.shape[2] == 3, message
+
     # Parallel and perpendicular components
-    p_para, p_perp = [p_xyz[:, 0, 0], (p_xyz[:, 1, 1] + p_xyz[:, 2, 2]) / 2]
+    p_para = p_xyz.data[:, 0, 0]
+    p_perp = (p_xyz.data[:, 1, 1] + p_xyz.data[:, 2, 2]) / 2
 
     # Off-diagonal terms
-    p_12, p_13, p_23 = [p_xyz[:, 0, 1], p_xyz[:, 0, 2], p_xyz[:, 1, 2]]
+    p_12, p_13, p_23 = [p_xyz.data[:, 0, 1], p_xyz.data[:, 0, 2], p_xyz.data[:, 1, 2]]
 
     d_ng = np.sqrt(8 * (p_12**2 + p_13**2 + p_23**2))
     d_ng /= p_para + 2 * p_perp
+    d_ng = ts_scalar(p_xyz.time.data, d_ng)
 
     return d_ng

pyrfu/pyrf/calc_sqrtq.py

@@ -3,6 +3,10 @@
 
 # 3rd party imports
 import numpy as np
+import xarray as xr
+
+# Local imports
+from .ts_scalar import ts_scalar
 
 __author__ = "Louis Richard"
 __email__ = "[email protected]"
@@ -64,13 +68,22 @@ def calc_sqrtq(p_xyz):
 
     """
 
+    # Check input type
+    assert isinstance(p_xyz, xr.DataArray), "p_xyz must be a xarray.DataArray"
+
+    # Check import shape
+    message = "p_xyz must be a time series of a tensor"
+    assert p_xyz.data.ndim == 3 and p_xyz.shape[1] == 3 and p_xyz.shape[2] == 3, message
+
     # Parallel and perpendicular components
-    p_para, p_perp = [p_xyz[:, 0, 0], (p_xyz[:, 1, 1] + p_xyz[:, 2, 2]) / 2]
+    p_para = p_xyz.data[:, 0, 0]
+    p_perp = (p_xyz.data[:, 1, 1] + p_xyz.data[:, 2, 2]) / 2
 
     # Off-diagonal terms
-    p_12, p_13, p_23 = [p_xyz[:, 0, 1], p_xyz[:, 0, 2], p_xyz[:, 1, 2]]
+    p_12, p_13, p_23 = [p_xyz.data[:, 0, 1], p_xyz.data[:, 0, 2], p_xyz.data[:, 1, 2]]
 
     sqrt_q = np.sqrt(p_12**2 + p_13**2 + p_23**2)
     sqrt_q /= np.sqrt(p_perp**2 + 2 * p_perp * p_para)
+    sqrt_q = ts_scalar(p_xyz.time.data, sqrt_q)
 
     return sqrt_q

pyrfu/pyrf/cdfepoch2datetime64.py

@@ -65,7 +65,7 @@ def cdfepoch2datetime64(epochs):
 
     Parameters
     ----------
-    epochs : array_like
+    epochs : float or int or array_like
         CDF epochs to convert.
 
     Returns
@@ -75,6 +75,10 @@ def cdfepoch2datetime64(epochs):
 
     """
 
+    # Check input type
+    message = "epochs must be array_like"
+    assert isinstance(epochs, (float, int, list, np.ndarray)), message
+
     times = cdfepoch.breakdown(epochs)
     times = np.transpose(np.atleast_2d(times))
 

pyrfu/pyrf/datetime2iso8601.py

@@ -2,6 +2,9 @@
 # -*- coding: utf-8 -*-
 
 # 3rd party imports
+import datetime
+
+import numpy as np
 import pandas as pd
 
 __author__ = "Louis Richard"
@@ -27,14 +30,20 @@ def datetime2iso8601(time):
 
     """
 
-    if isinstance(time, list):
+    # Check input type
+    message = "time must be array_like or datetime.datetime"
+    assert isinstance(time, (list, np.ndarray, datetime.datetime)), message
+
+    if isinstance(time, (np.ndarray, list)):
         return list(map(datetime2iso8601, time))
 
+    assert isinstance(time, datetime.datetime), "time datetime.datetime"
+
     time_datetime = pd.Timestamp(time)
 
     # Convert to string
     datetime_str = time_datetime.strftime("%Y-%m-%dT%H:%M:%S.%f")
 
-    tt2000 = f"{datetime_str}{time_datetime.nanosecond:03d}"
+    time_iso8601 = f"{datetime_str}{time_datetime.nanosecond:03d}"
 
-    return tt2000
+    return time_iso8601