Add tests, fix displacement method, add notebook

README.rst

@@ -25,6 +25,9 @@ traja |Python-ver| |Travis| |PyPI| |RTD| |Gitter| |Black| |License|
     :target: https://opensource.org/licenses/MIT
     :alt: License: MIT
 
+.. |Binder| image:: https://mybinder.org/badge_logo.svg
+    :target: https://mybinder.org/v2/gh/justinshenk/traja/master?filepath=demo.ipynb
+
 traja is a Python library for trajectory analysis. It extends the capability of
 pandas DataFrame specific for animal trajectory analysis in 2D, and provides
 convenient interfaces to other geometric analysis packages (eg, R and shapely).

environment.yml

@@ -8,8 +8,6 @@ dependencies:
   - pytorch-cpu
   - pip:
     - fastdtw
-    - sphinx-gallery
-    - sphinx_rtd_theme
     - tzlocal
   - pandas
   - numpy

traja/__init__.py

@@ -6,4 +6,4 @@
 from traja import models
 
 __author__ = "justinshenk"
-__version__ = "0.1.0"
+__version__ = "0.1.1"

traja/contrib/dvc.py

@@ -7,20 +7,18 @@
 import re
 from typing import Optional, List
 
+import matplotlib.style as style
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import pandas as pd
 from scipy.stats import ttest_ind, ttest_rel
 import statsmodels.api as sm
 import statsmodels.formula.api as smf
 
-
 import traja
 from traja import TrajaDataFrame
 
-import matplotlib.style as style
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-import numpy as np
-import pandas as pd
-
 try:
     import seaborn as sns
 except ImportError:
@@ -57,7 +55,7 @@ def glm(
     day_range = f"{daily[DAYS_COL].min()}-{daily[DAYS_COL].max()}"
     print(
         f"Fitting generalized linear model on daily data. Days {day_range}\n"
-        f"Model: {response_var} ~ days_from_surgery + diet."
+        f"Model: {response_var} ~ days_from_surgery * diet."
         + (f" Covariance structure ({cov_struct})" if cov_struct else "")
     )
 
@@ -78,7 +76,7 @@ def glm(
         subset = daily[daily.period == period] if period != "Overall" else daily_overall
 
         md = smf.gee(
-            f" {response_var} ~ days_from_surgery + diet",
+            f" {response_var} ~ days_from_surgery * diet",
             data=subset,
             groups="cage",
             **kwargs,
@@ -95,6 +93,34 @@ def glm(
                 sig_pvalues.append((pvalue, factor, period))
         if verbose:
             print(f"END {period} GLM RESULTS\n")
+
+    for diet in ["Control", "Fortasyn"]:
+        for period in [DARK, LIGHT, "Overall"]:
+            subset = (
+                daily.loc[(daily.period == period) & (daily.diet.str.contains(diet))]
+                if period != "Overall"
+                else daily_overall
+            )
+            if subset.empty:
+                import ipdb
+
+                ipdb.set_trace()
+            md = smf.gee(
+                f" {response_var} ~ days_from_surgery",
+                data=subset,
+                groups="cage",
+                **kwargs,
+            )
+            mdf = md.fit()
+            if verbose:
+                print(f"{diet}-only, {period} GLM RESULTS:\n", mdf.summary())
+            pvalues = mdf.pvalues
+            for factor, pvalue in pvalues.items():
+                if pvalue < 0.05 and factor != "Intercept":
+                    print(
+                        f"{factor} effect on {response_var} ({period}) - p-value: {pvalue:.4f}"
+                    )
+                    sig_pvalues.append((pvalue, factor, period))
     if not sig_pvalues:
         print("No significant effect found.\n")
     return sig_pvalues
@@ -1710,7 +1736,7 @@ def main(args):
 
         # Statistical analysis
         if args.glm:
-            glm(daily, response_var="displacment", verbose=args.verbose)
+            glm(daily, response_var="displacement", verbose=args.verbose)
     else:
         # Activity-only
         experiment = DVCExperiment(experiment_name=args.experiment_name)
@@ -1721,12 +1747,11 @@ def main(args):
         experiment.plot_daily_activity(activity, **args.__dict__)
 
     # Statistical analysis
-    if args.glm:
+    if args.glm and args.activity:
         glm(activity, response_var="activity", verbose=args.verbose)
 
 
 def parse_args():
-
     parser = argparse.ArgumentParser(description="Run DVC Experiment")
     parser.add_argument(
         "-c", "--centroids_dir", help="path to directory with centroids CSVs"

traja/tests/test_parsers.py

@@ -0,0 +1,32 @@
+import os
+
+import pandas as pd
+
+import traja
+
+from pandas.util.testing import (
+    assert_frame_equal,
+    assert_index_equal,
+    assert_series_equal,
+)
+
+df = traja.generate(n=20)
+
+
+def test_from_df():
+    df = pd.DataFrame({"x": [1, 2, 3], "y": [2, 3, 4]})
+    trj = traja.parsers.from_df(df)
+    assert_frame_equal(df, trj)
+    assert isinstance(trj, traja.TrajaDataFrame)
+
+
+def test_read_file():
+    datapath = os.path.join(traja.__path__[0], "tests", "data", "3527.csv")
+    trj = traja.parsers.read_file(datapath)
+    assert isinstance(trj, traja.TrajaDataFrame)
+    assert "Frame" in trj
+    assert "Time" in trj
+    assert "TrackId" in trj
+    assert "x" in trj
+    assert "y" in trj
+    assert "ValueChanged" in trj

traja/tests/test_trajadataframe.py

@@ -2,19 +2,12 @@
 import shutil
 import tempfile
 
-import numpy as np
 import pandas as pd
+import pytest
 
 import traja
 from traja import TrajaDataFrame, read_file
 
-import pytest
-from pandas.util.testing import (
-    assert_frame_equal,
-    assert_index_equal,
-    assert_series_equal,
-)
-
 
 class TestDataFrame:
     def setup_method(self):
@@ -28,11 +21,11 @@ def teardown_method(self):
         shutil.rmtree(self.tempdir)
 
     def test_df_init(self):
-        assert type(self.df) is TrajaDataFrame
+        assert isinstance(self.df, TrajaDataFrame)
 
     def test_copy(self):
         df2 = self.df.copy()
-        assert type(df2) is TrajaDataFrame
+        assert (df2, TrajaDataFrame)
         assert df2.xlim == self.df.xlim
 
     def test_dataframe_to_trajadataframe(self):
@@ -43,3 +36,36 @@ def test_dataframe_to_trajadataframe(self):
         tf = TrajaDataFrame(df)
         assert isinstance(df, pd.DataFrame)
         assert isinstance(tf, TrajaDataFrame)
+
+    @pytest.mark.parameterize("method", ["merge", "concat", None])
+    def test_construct_dataframe(self, method):
+        df = traja.TrajaDataFrame(
+            {"x": range(len(self.df)), "y": range(len(self.df))},
+            index=self.df.index,
+            xlim=(0, 2),
+            ylim=(0, 2),
+            spatial_units="m",
+            title="Serious title",
+            fps=2.0,
+            time_units="s",
+            id=42,
+        )
+
+        assert df.title == "Serious title"
+
+        # Test 'merge'
+        df2 = df.copy()
+        assert df2.title == "Serious title"
+
+        assert df._get_time_col() == None
+        assert self.df._get_time_col() == "Time"
+
+        # Modify metavar
+        df.set("title", "New title")
+        assert df.title == "New title"
+
+        # Test __finalize__
+        df_copy = df.copy()
+        df2_copy = df2.copy()
+        df_copy.__finalize__(df2_copy, method=method)
+        assert isinstance(df_copy, traja.TrajaDataFrame)

traja/tests/test_trajectory.py

@@ -1,5 +1,6 @@
 import numpy as np
 import numpy.testing as npt
+import pytest
 from pandas.util.testing import assert_series_equal
 
 import traja
@@ -66,8 +67,29 @@ def test_cartesian_to_polar():
     npt.assert_allclose(theta_actual[:10], theta_expected)
 
 
-def test_expected_sq_displacement():
-    ...
+@pytest.mark.parametrize("eqn1", [True, False])
+def test_expected_sq_displacement(eqn1):
+    disp = traja.expected_sq_displacement(df, eqn1=eqn1)
+
+
+def test_step_lengths():
+    step_lengths = traja.step_lengths(df)
+    assert len(step_lengths == len(df))
+
+
+@pytest.mark.parametrize("w", [None, 2, 3])
+def test_smooth_sg(w):
+    if w == 2:
+        with pytest.raises(Exception):
+            _ = traja.trajectory.smooth_sg(df, w=w)
+    else:
+        trj = traja.trajectory.smooth_sg(df, w=w)
+        assert trj.shape == df.shape
+
+
+@pytest.mark.parametrize("lag", [1, 2])
+def test_angles(lag):
+    angles = traja.angles(df, lag=lag)
 
 
 def test_traj_from_coords():
@@ -79,10 +101,21 @@ def test_traj_from_coords():
     assert_series_equal(trj.time, df.time)
 
 
-def test_distance():
+@pytest.mark.parametrize("method", ["dtw", "hausdorff"])
+def test_distance(method):
     rotated = traja.trajectory.rotate(df, 10).traja.xy[:10]
-    distance = traja.distance(rotated, df.traja.xy)
-    npt.assert_almost_equal(distance, 523.103_701_021_348_1)
+    distance = traja.distance(rotated, df.traja.xy, method=method)
+
+
+@pytest.mark.parametrize("ndarray_type", [True, False])
+def test_grid_coords1D(ndarray_type):
+    xlim, ylim = traja.trajectory._get_xylim(df)
+    bins = traja.trajectory._bins_to_tuple(df, None)
+    grid_indices = traja.grid_coordinates(df, bins=bins, xlim=xlim, ylim=ylim)
+    if ndarray_type:
+        grid_indices = grid_indices.values
+    grid_indices1D = traja._grid_coords1D(grid_indices)
+    assert isinstance(grid_indices, np.ndarray)
 
 
 def test_to_shapely():
@@ -97,7 +130,10 @@ def test_to_shapely():
 
 
 def test_transition_matrix():
-    ...
+    grid_indices = traja.grid_coordinates(df)
+    assert grid_indices.shape[1] == 2
+    grid_indices1D = traja._grid_coords1D(grid_indices)
+    transitions_matrix = traja.transition_matrix(grid_indices1D)
 
 
 def test_calculate_flow_angles():
@@ -107,8 +143,30 @@ def test_calculate_flow_angles():
     npt.assert_approx_equal(U.sum(), expected)
 
 
+def test_resample_time():
+    trj = traja.resample_time(df, "3s")
+    assert isinstance(trj, traja.TrajaDataFrame)
+
+
 def test_transitions():
-    ...
+    transitions = traja.transitions(df)
+    assert isinstance(transitions, np.ndarray)
+
+    # Check when bins set
+    bins = traja._bins_to_tuple(df, bins=None)
+    xmin = df.x.min()
+    xmax = df.x.max()
+    ymin = df.y.min()
+    ymax = df.y.max()
+    xbins = np.linspace(xmin, xmax, bins[0])
+    ybins = np.linspace(ymin, ymax, bins[1])
+    xbin = np.digitize(df.x, xbins)
+    ybin = np.digitize(df.y, ybins)
+
+    df.set("xbin", xbin)
+    df.set("ybin", ybin)
+    transitions = traja.transitions(df)
+    assert isinstance(transitions, np.ndarray)
 
 
 def test_grid_coordinates():