cbps 0.2.0__tar.gz → 0.2.1__tar.gz

{cbps-0.2.0 → cbps-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cbps
-Version: 0.2.0
+Version: 0.2.1
 Summary: Covariate Balancing Propensity Score (CBPS) for causal inference in Python
 Author: Cai Xuanyu, Xu Wenli
 Maintainer: Cai Xuanyu, Xu Wenli

{cbps-0.2.0 → cbps-0.2.1}/cbps/__init__.py

@@ -110,7 +110,7 @@ import warnings
 import pandas as pd
 import numpy as np
 
-__version__ = "0.1.0"
+__version__ = "0.2.1"
 
 from cbps.core.results import CBPSResults, CBPSSummary
 from cbps.core.cbps_binary import cbps_binary_fit
@@ -725,10 +725,12 @@ def CBPS(
     covariates : np.ndarray, optional
         Covariate matrix. Required when using the array interface. Should not
         include an intercept column (it will be added automatically).
-    att : int, default 1
-        Target estimand. 0 for ATE (average treatment effect), 1 for ATT
-        with the second level as treated, 2 for ATT with the first level as
-        treated. For non-binary treatments, only ATE is available.
+    att : int or str, default 1
+        Target estimand. Accepts integer (0, 1, 2) or string ('ate', 'att', 'atc').
+        0 / 'ate': ATE (average treatment effect for entire population).
+        1 / 'att': ATT (average treatment effect on the treated, T=1 as treated).
+        2 / 'atc': ATC (average treatment effect on the control, T=0 as treated).
+        For non-binary treatments, only ATE is available.
     ATT : int, optional
         Deprecated. Use lowercase ``att`` instead.
     method : {'over', 'exact'}, default 'over'
@@ -833,20 +835,29 @@ def CBPS(
         two_step = twostep
 
     # Parameter validation
-    # att must be 0, 1, or 2
-    # att=0: ATE, att=1: ATT (T=1 as treated), att=2: ATT (T=0 as treated)
-    # Check type first, then value range (TypeError before ValueError)
-    if not isinstance(att, (int, np.integer)):
+    # att accepts int (0,1,2) or str ('ate','att','atc')
+    _ATT_STR_MAP = {'ate': 0, 'att': 1, 'atc': 2}
+    if isinstance(att, str):
+        att_lower = att.lower().strip()
+        if att_lower not in _ATT_STR_MAP:
+            raise ValueError(
+                f"Invalid att='{att}'. Accepted string values: 'ate', 'att', 'atc'.\n"
+                f"  'ate' (=0): Average Treatment Effect\n"
+                f"  'att' (=1): Average Treatment effect on the Treated\n"
+                f"  'atc' (=2): Average Treatment effect on the Control"
+            )
+        att = _ATT_STR_MAP[att_lower]
+    elif not isinstance(att, (int, np.integer)):
         raise TypeError(
-            f"att must be an integer (0, 1, or 2), got type {type(att).__name__}: {att}"
+            f"att must be int (0,1,2) or str ('ate','att','atc'), got {type(att).__name__}: {att}"
         )
     if att not in [0, 1, 2]:
         raise ValueError(
             f"Invalid att parameter: {att}\n\n"
             f"att must be 0, 1, or 2:\n"
-            f"  att=0: ATE (Average Treatment Effect) for entire population\n"
-            f"  att=1: ATT (Average Treatment effect on the Treated, T=1 as treated)\n"
-            f"  att=2: ATT (Average Treatment effect on the Treated, T=0 as treated)\n\n"
+            f"  att=0 / 'ate': ATE (Average Treatment Effect) for entire population\n"
+            f"  att=1 / 'att': ATT (Average Treatment effect on the Treated, T=1 as treated)\n"
+            f"  att=2 / 'atc': ATT (Average Treatment effect on the Treated, T=0 as treated)\n\n"
             f"You provided: att={att}"
         )
     
@@ -938,19 +949,22 @@ def CBPS(
         )
     
     # Validate att parameter
-    if not isinstance(att, (int, np.integer)):
+    _ATT_STR_MAP = {'ate': 0, 'att': 1, 'atc': 2}
+    if isinstance(att, str):
+        att_lower = att.lower().strip()
+        if att_lower not in _ATT_STR_MAP:
+            raise ValueError(
+                f"Invalid att='{att}'. Accepted: 'ate', 'att', 'atc' or 0, 1, 2."
+            )
+        att = _ATT_STR_MAP[att_lower]
+    elif not isinstance(att, (int, np.integer)):
         raise TypeError(
-            f"att must be an integer (0, 1, or 2), got {type(att).__name__}. "
-            f"Received: att={att}"
+            f"att must be int (0,1,2) or str ('ate','att','atc'), got {type(att).__name__}: {att}"
         )
     if att not in (0, 1, 2):
         raise ValueError(
             f"att must be 0 (ATE), 1 (ATT treated=level2), or 2 (ATT treated=level1). "
-            f"Received: att={att}\n\n"
-            f"Explanation:\n"
-            f"  att=0: Average Treatment Effect (ATE) for entire population\n"
-            f"  att=1: Average Treatment effect on the Treated (ATT), second level as treated\n"
-            f"  att=2: ATT with first level as treated"
+            f"Received: att={att}"
         )
     
     # Validate method parameter

{cbps-0.2.0 → cbps-0.2.1}/cbps.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cbps
-Version: 0.2.0
+Version: 0.2.1
 Summary: Covariate Balancing Propensity Score (CBPS) for causal inference in Python
 Author: Cai Xuanyu, Xu Wenli
 Maintainer: Cai Xuanyu, Xu Wenli

{cbps-0.2.0 → cbps-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cbps"
-version = "0.2.0"
+version = "0.2.1"
 description = "Covariate Balancing Propensity Score (CBPS) for causal inference in Python"
 readme = "README.md"
 license = "AGPL-3.0-only"

{cbps-0.2.0 → cbps-0.2.1}/tests/test_api.py

@@ -881,3 +881,52 @@ class TestAPIErrorHandling:
 
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])
+
+
+class TestAttStringParameter:
+    """Test that CBPS() high-level API accepts string att values."""
+
+    def test_ate_string(self):
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        fit = CBPS('treat ~ age + educ + re74', data=data, att='ate')
+        assert fit.converged
+
+    def test_att_string(self):
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        fit = CBPS('treat ~ age + educ + re74', data=data, att='att')
+        assert fit.converged
+
+    def test_atc_string(self):
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        fit = CBPS('treat ~ age + educ + re74', data=data, att='atc')
+        assert fit.converged
+
+    def test_case_insensitive(self):
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        fit = CBPS('treat ~ age + educ + re74', data=data, att='ATE')
+        assert fit.converged
+
+    def test_invalid_string_raises(self):
+        import pytest
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        with pytest.raises(ValueError, match="Invalid att"):
+            CBPS('treat ~ age + educ + re74', data=data, att='invalid')
+
+    def test_string_equivalent_to_int(self):
+        import numpy as np
+        from cbps import CBPS
+        from cbps.datasets import load_lalonde
+        data = load_lalonde()
+        fit_int = CBPS('treat ~ age + educ + re74', data=data, att=0)
+        fit_str = CBPS('treat ~ age + educ + re74', data=data, att='ate')
+        assert np.allclose(fit_int.coefficients, fit_str.coefficients)