added tests (#2)

docs/changelog.rst

@@ -4,6 +4,12 @@
 Change Log
 ==========
 
+Version 0.1.2
+-------------
+
+- Added unit and integration tests
+
+
 Version 0.1.1
 -------------
 

docs/conf.py

@@ -22,7 +22,7 @@
 author = 'Roni Kobrosly'
 
 # The full version, including alpha/beta/rc tags
-release = '0.1.1'
+release = '0.1.2'
 
 
 # -- General configuration ---------------------------------------------------

docs/install.rst

@@ -36,7 +36,7 @@ you can easily install causal-curve using ``pip``::
     pip install causal-curve
 
 
-You can also get the latest version of pyts by cloning the repository::
+You can also get the latest version of causal-curve by cloning the repository::
 
     git clone https://github.com/ronikobrosly/causal-curve.git
     cd causal-curve

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="causal-curve",
-    version="0.1.1",
+    version="0.1.2",
     author="Roni Kobrosly",
     author_email="[email protected]",
     description="A python library with tools to perform causal inference using \

tests/conftest.py

@@ -18,9 +18,12 @@ def full_example_dataset():
 
     treatment = np.random.normal(loc=50.0, scale=10.0, size=n)
     x1 = np.random.normal(loc=50.0, scale=10.0, size=n)
-    outcome = treatment + x1 + np.random.normal(loc=50.0, scale=3.0, size=n)
+    x2 = np.random.normal(loc=0, scale=10.0, size=n)
+    outcome = treatment + x1 + x2 + np.random.normal(loc=50.0, scale=3.0, size=n)
 
-    fixture = pd.DataFrame({"treatment": treatment, "x1": x1, "outcome": outcome})
+    fixture = pd.DataFrame(
+        {"treatment": treatment, "x1": x1, "x2": x2, "outcome": outcome}
+    )
 
     fixture.reset_index(drop=True, inplace=True)
 

tests/integration/test_gps.py

@@ -2,16 +2,16 @@
 
 import pandas as pd
 
-from causal_curve import CDRC
+from causal_curve import GPS
 from tests.test_helpers import assert_df_equal
 
 
-def test_full_cdrc_flow(dataset_fixture):
+def test_full_gps_flow(dataset_fixture):
     """
-    Tests the full flow of the CDRC tool
+    Tests the full flow of the GPS tool
     """
 
-    cdrc = CDRC(
+    gps = GPS(
         treatment_grid_num=10,
         lower_grid_constraint=0.0,
         upper_grid_constraint=1.0,
@@ -21,64 +21,13 @@ def test_full_cdrc_flow(dataset_fixture):
         random_seed=100,
         verbose=False,
     )
-    cdrc.fit(
+    gps.fit(
         T=dataset_fixture["treatment"],
-        X=dataset_fixture["x1"],
+        X=dataset_fixture[["x1", "x2"]],
         y=dataset_fixture["outcome"],
     )
-    cdrc_results = cdrc.calculate_CDRC(0.95)
+    gps_results = gps.calculate_CDRC(0.95)
 
-    expected_df = pd.DataFrame(
-        {
-            "Treatment": [
-                22.104,
-                37.955,
-                42.553,
-                45.957,
-                48.494,
-                51.139,
-                53.994,
-                57.607,
-                61.19,
-                80.168,
-            ],
-            "CDRC": [
-                119.231,
-                138.978,
-                143.009,
-                145.79,
-                148.023,
-                150.729,
-                153.833,
-                157.768,
-                161.83,
-                187.305,
-            ],
-            "Lower_CI": [
-                108.507,
-                136.278,
-                140.41,
-                143.233,
-                145.545,
-                148.375,
-                151.452,
-                155.517,
-                159.391,
-                174.926,
-            ],
-            "Upper_CI": [
-                129.955,
-                141.678,
-                145.609,
-                148.348,
-                150.5,
-                153.083,
-                156.213,
-                160.018,
-                164.27,
-                199.683,
-            ],
-        }
-    )
-
-    assert_df_equal(cdrc_results, expected_df)
+    assert isinstance(gps_results, pd.DataFrame)
+    check = gps_results.columns == ["Treatment", "CDRC", "Lower_CI", "Upper_CI"]
+    assert check.all()

tests/integration/test_tmle.py

@@ -1 +1,28 @@
 """ Integration tests of the tmle.py module """
+
+import pandas as pd
+
+from causal_curve import TMLE
+from tests.test_helpers import assert_df_equal
+
+
+def test_full_tmle_flow(dataset_fixture):
+    """
+    Tests the full flow of the TMLE tool
+    """
+
+    tmle = TMLE(
+        treatment_grid_bins=[22.1, 30, 40, 50, 60, 70, 80.1],
+        random_seed=100,
+        verbose=False,
+    )
+    tmle.fit(
+        T=dataset_fixture["treatment"],
+        X=dataset_fixture[["x1", "x2"]],
+        y=dataset_fixture["outcome"],
+    )
+    tmle_results = tmle.calculate_CDRC(0.95)
+
+    assert isinstance(tmle_results, pd.DataFrame)
+    check = tmle_results.columns == ["Treatment", "CDRC", "Lower_CI", "Upper_CI"]
+    assert check.all()

tests/test_helpers.py

@@ -5,10 +5,11 @@
 from pandas.testing import assert_frame_equal
 
 
-def assert_df_equal(observed_frame, expected_frame):
+def assert_df_equal(observed_frame, expected_frame, check_less_precise):
     """Assert that two pandas dataframes are equal, ignoring ordering of columns."""
     assert_frame_equal(
         observed_frame.sort_index(axis=1),
         expected_frame.sort_index(axis=1),
         check_names=True,
+        check_less_precise=check_less_precise,
     )

tests/unit/test_gps.py

@@ -1 +1,33 @@
 """ Unit tests of the gps.py module """
+
+
+import pandas as pd
+from pygam import LinearGAM
+
+from causal_curve import GPS
+from tests.test_helpers import assert_df_equal
+
+
+def test_GPS_fit(dataset_fixture):
+    """
+    Tests the fit method GPS tool
+    """
+
+    gps = GPS(
+        treatment_grid_num=10,
+        lower_grid_constraint=0.0,
+        upper_grid_constraint=1.0,
+        spline_order=3,
+        n_splines=10,
+        max_iter=100,
+        random_seed=100,
+        verbose=False,
+    )
+    gps.fit(
+        T=dataset_fixture["treatment"],
+        X=dataset_fixture["x1"],
+        y=dataset_fixture["outcome"],
+    )
+
+    assert isinstance(gps.gam_results, LinearGAM)
+    assert gps.gps.shape == (500,)

tests/unit/test_tmle.py

@@ -1 +1,28 @@
 """ Unit tests of the tmle.py module """
+
+
+import pandas as pd
+
+from causal_curve import TMLE
+from tests.test_helpers import assert_df_equal
+
+
+def test_TMLE_fit(dataset_fixture):
+    """
+    Tests the fit method GPS tool
+    """
+
+    tmle = TMLE(
+        treatment_grid_bins=[22.1, 30, 40, 50, 60, 70, 80.1],
+        random_seed=100,
+        verbose=False,
+    )
+    tmle.fit(
+        T=dataset_fixture["treatment"],
+        X=dataset_fixture[["x1", "x2"]],
+        y=dataset_fixture["outcome"],
+    )
+
+    assert tmle.n_obs == 72
+    assert len(tmle.psi_list) == 5
+    assert len(tmle.std_error_ic_list) == 5