Add weights option to event_study

.github/workflows/ci-tests.yaml

@@ -52,6 +52,9 @@ jobs:
         run: |
           pixi run tests-regular
 
+      - name: Build coverage.xml
+        run: pixi run coverage-report
+
       - name: Upload coverage to Codecov (partial)
         uses: codecov/codecov-action@v4
         with:
@@ -101,6 +104,9 @@ jobs:
           pixi run tests-against-r-core
           pixi run tests-against-r-extended
 
+      - name: Build coverage.xml
+        run: pixi run coverage-report
+
       - name: Upload coverage to Codecov (partial)
         uses: codecov/codecov-action@v4
         with:
@@ -109,18 +115,6 @@ jobs:
           flags: tests-vs-r
           files: coverage.xml
 
-  merge_coverage:
-    name: "Merge Coverage"
-    runs-on: ubuntu-latest
-    needs: [test, test_slow]
-    steps:
-      - name: Final coverage merge
-        uses: codecov/codecov-action@v4
-        with:
-          token: ${{ secrets.CODECOV_TOKEN }}
-          partial: false
-          flags: final
-
   build-docs:
     name: "Build Docs"
     runs-on: ubuntu-latest

pyfixest/did/did.py

@@ -26,13 +26,17 @@ class DID(ABC):
         YYYYMMDDHHMMSS, i.e. it must be possible to compare two dates via '>'.
         Datetime variables are currently not accepted. Never treated units must
         have a value of 0.
-    xfml : str
+    cluster : Optional[str]
+        The name of the cluster variable.
+    weights : Optional[str]
+        Default is None. Weights for WLS estimation. If None, all observations
+        are weighted equally. If a string, the name of the column in `data` that
+        contains the weights. Must be analytic weights for now.
+    xfml : Optional[str]
         The formula for the covariates.
-    att : str
+    att : Optional[bool], default=True
         Whether to estimate the average treatment effect on the treated (ATT) or
         the canonical event study design with all leads and lags. Default is True.
-    cluster : str
-        The name of the cluster variable.
     """
 
     @abstractmethod
@@ -44,8 +48,9 @@ def __init__(
         tname: str,
         gname: str,
         cluster: Optional[str] = None,
+        weights: Optional[str] = None,
         xfml: Optional[str] = None,
-        att: bool = True,
+        att: Optional[bool] = True,
     ):
         # do some checks here
 
@@ -57,9 +62,10 @@ def __init__(
         self._xfml = xfml
         self._att = att
         self._cluster = cluster
+        self._weights = weights
+        self._weights_type = "aweights"
 
         # check if tname and gname are of type int (either int 64, 32, 8)
-
         for var in [self._tname, self._gname]:
             if self._data[var].dtype not in [
                 "int64",

pyfixest/did/did2s.py

@@ -36,18 +36,18 @@ class DID2S(DID):
         YYYYMMDDHHMMSS, i.e. it must be possible to compare two dates via '>'.
         Datetime variables are currently not accepted. Never treated units
         must have a value of 0.
-    xfml : str
-        The formula for the covariates.
-    att : str
-        Whether to estimate the pooled average treatment effect on the treated
-        (ATT) or the canonical event study design with all leads and lags / the
-        ATT for each period. Default is True.
     cluster : str
         The name of the cluster variable.
-    weights : Optional[str].
+    weights : Optional[str]
         Default is None. Weights for WLS estimation. If None, all observations
         are weighted equally. If a string, the name of the column in `data` that
-        contains the weights.
+        contains the weights. Must be analytic weights for now.
+    xfml : Optional[str]
+        The formula for the covariates.
+    att : Optional[bool], default=True
+        Whether to estimate the pooled average treatment effect on the treated
+        (ATT) or the canonical event study design with all leads and lags / the
+        ATT for each period. Default is True.
     """
 
     def __init__(
@@ -59,8 +59,8 @@ def __init__(
         gname: str,
         cluster: str,
         weights: Optional[str] = None,
-        att: bool = True,
         xfml: Optional[str] = None,
+        att: Optional[bool] = True,
     ):
         super().__init__(
             data=data,
@@ -71,6 +71,7 @@ def __init__(
             xfml=xfml,
             att=att,
             cluster=cluster,
+            weights=weights,
         )
 
         self._estimator = "did2s"
@@ -86,17 +87,14 @@ def __init__(
         self._first_u = np.array([])
         self._second_u = np.array([])
 
-        # column name with weights. None by default
-        self._weights_name = weights
-
     def estimate(self):
         """Estimate the two-step DID2S model."""
         return _did2s_estimate(
             data=self._data,
             yname=self._yname,
             _first_stage=self._fml1,
             _second_stage=self._fml2,
-            weights=self._weights_name,
+            weights=self._weights,
             treatment="is_treated",
         )  # returns triple Feols, first_u, second_u
 
@@ -121,7 +119,7 @@ def vcov(self):
             first_u=self._first_u,
             second_u=self._second_u,
             cluster=self._cluster,
-            weights=self._weights_name,
+            weights=self._weights,
         )
 
     def iplot(
@@ -175,10 +173,10 @@ def _did2s_estimate(
         The formula for the second stage.
     treatment: str
         The name of the treatment variable. Must be boolean.
-    weights : Optional[str].
+    weights : Optional[str]
         Default is None. Weights for WLS estimation. If None, all observations
         are weighted equally. If a string, the name of the column in `data` that
-        contains the weights.
+        contains the weights. Must be analytic weights for now.
 
     Returns
     -------
@@ -283,10 +281,10 @@ def _did2s_vcov(
         The second stage residuals.
     cluster: str
         The name of the cluster variable.
-    weights : Optional[str].
+    weights : Optional[str]
         Default is None. Weights for WLS estimation. If None, all observations
         are weighted equally. If a string, the name of the column in `data` that
-        contains the weights.
+        contains the weights. Must be analytic weights for now.
 
     Returns
     -------

pyfixest/did/estimation.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 import pandas as pd
 
@@ -15,17 +15,19 @@ def event_study(
     idname: str,
     tname: str,
     gname: str,
-    xfml: Optional[str] = None,
     cluster: Optional[str] = None,
-    estimator: Optional[str] = "twfe",
+    weights: Optional[str] = None,
+    xfml: Optional[str] = None,
     att: Optional[bool] = True,
+    estimator: Optional[Literal["did2s", "twfe", "saturated"]] = "twfe",
 ):
     """
     Estimate Event Study Model.
 
     This function allows for the estimation of treatment effects using different
     estimators. Currently, it supports "twfe" for the two-way fixed effects
-    estimator and "did2s" for Gardner's two-step DID2S estimator. Other estimators
+    estimator, "did2s" for Gardner's two-step DID2S estimator, and "saturated" for
+    a Sun & Abraham staggered event study estimator. Other estimators
     are in development.
 
     Parameters
@@ -42,14 +44,18 @@ def event_study(
         Unit-specific time of initial treatment.
     cluster: Optional[str]
         The name of the cluster variable. If None, defaults to idname.
-    xfml : str
+    weights : Optional[str]
+        Default is None. Weights for WLS estimation. If None, all observations
+        are weighted equally. If a string, the name of the column in `data` that
+        contains the weights. Must be analytic weights for now.
+    xfml : Optional[str]
         The formula for the covariates.
-    estimator : str
-        The estimator to use. Options are "did2s", "twfe", and "saturated".
-    att : bool, optional
+    att : Optional[bool]
         If True, estimates the average treatment effect on the treated (ATT).
         If False, estimates the canonical event study design with all leads and
         lags. Default is True.
+    estimator : Optional[str], default="twfe"
+        The estimator to use. Options are "did2s", "twfe", and "saturated".
 
     Returns
     -------
@@ -96,10 +102,11 @@ def event_study(
     assert isinstance(idname, str), "idname must be a string"
     assert isinstance(tname, str), "tname must be a string"
     assert isinstance(gname, str), "gname must be a string"
+    assert isinstance(cluster, str) or cluster is None, "cluster must be a string"
+    assert isinstance(weights, str) or weights is None, "weights must be a string"
     assert isinstance(xfml, str) or xfml is None, "xfml must be a string or None"
-    assert isinstance(estimator, str), "estimator must be a string"
     assert isinstance(att, bool), "att must be a boolean"
-    assert isinstance(cluster, str) or cluster is None, "cluster must be a string"
+    assert isinstance(estimator, str), "estimator must be a string"
 
     cluster = idname if cluster is None else cluster
 
@@ -110,9 +117,10 @@ def event_study(
             idname=idname,
             tname=tname,
             gname=gname,
+            cluster=cluster,
+            weights=weights,
             xfml=xfml,
             att=att,
-            cluster=cluster,
         )
 
         fit, did2s._first_u, did2s._second_u = did2s.estimate()
@@ -130,9 +138,10 @@ def event_study(
             idname=idname,
             tname=tname,
             gname=gname,
+            cluster=cluster,
+            weights=weights,
             xfml=xfml,
             att=att,
-            cluster=cluster,
         )
         fit = twfe.estimate()
         fit._yname = twfe._yname
@@ -151,9 +160,10 @@ def event_study(
             idname=idname,
             tname=tname,
             gname=gname,
+            cluster=cluster,
+            weights=weights,
             xfml=xfml,
             att=att,
-            cluster=cluster,
         )
         fit = saturated.estimate()
         vcov = fit.vcov(vcov={"CRV1": cluster})