CausalEstimate 0.1.0__tar.gz

causalestimate-0.1.0/.github/workflows/format.yml

@@ -0,0 +1,25 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+name: Formatting using black
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+permissions:
+  contents: read
+jobs:
+  format:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+        cache: 'pip' # caching pip dependencies
+    - name: Install dependencies
+      run: |
+        pip install black>=23.10.1
+    - name: Run black
+      run: |
+        black --check CausalEstimate tests

causalestimate-0.1.0/.github/workflows/lint.yml

@@ -0,0 +1,26 @@
+# This workflow will install Python dependencies and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+name: Lint using flake8 
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+permissions:
+  contents: read
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+        cache: 'pip' # caching pip dependencies
+    - name: Install dependencies
+      run: |
+        pip install flake8>=6.1.0
+    - name: Lint with flake8
+      run: |
+        # Check syntax errors
+        flake8 CausalEstimate tests --count --select=E9,F63,F7,F82,U100,E711,E712,E713,E714,E721,F401,F402,F405,F811,F821,F822,F823,F831,F841,F901, --show-source --statistics

causalestimate-0.1.0/.github/workflows/unittest.yml

@@ -0,0 +1,27 @@
+name: 'Unittests'
+
+on:
+    workflow_dispatch:
+    pull_request:
+
+permissions:
+  contents: read
+  pull-requests: read
+
+jobs:
+    tests:
+        runs-on: "ubuntu-latest"
+        steps:
+            - name: Checkout
+              uses: actions/checkout@v4
+
+            - name: Setup dependencies.
+              run: |
+                python -m venv .venv
+                source .venv/bin/activate
+                pip install -r requirements.txt
+
+            - name: Running tests
+              run: |
+                source .venv/bin/activate
+                python -m unittest

causalestimate-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+*.egg-info
+*.pyc
+dist/
+__pycache__/
+notebooks/test*
+CausalEstimate/_version.py

causalestimate-0.1.0/CausalEstimate/api/estimator.py

@@ -0,0 +1,132 @@
+import pandas as pd
+import numpy as np
+from CausalEstimate.estimators.tmle import TMLE
+
+# !TODO: Write test for all functions
+
+
+class Estimator:
+    def __init__(self, methods=None, effect_type="ATE", **kwargs):
+        """
+        Initialize the Estimator class with one or more methods.
+
+        Args:
+            methods (list or str): A list of estimator method names (e.g., ["AIPW", "TMLE"])
+                                   or a single method name (e.g., "AIPW").
+            effect_type (str): The type of effect to estimate (e.g., "ATE", "ATT").
+            **kwargs: Additional keyword arguments for each estimator.
+        """
+        if methods is None:
+            methods = ["AIPW"]  # Default to AIPW if no method is provided.
+
+        # Allow single method or list of methods
+        self.methods = methods if isinstance(methods, list) else [methods]
+        self.effect_type = effect_type
+        self.estimators = self._initialize_estimators(effect_type, **kwargs)
+
+    def _initialize_estimators(self, effect_type, **kwargs):
+        """
+        Initialize the specified estimators based on the methods provided.
+        """
+        estimators = {
+            "TMLE": TMLE,
+            # Add other estimators as needed
+        }
+        initialized_estimators = []
+
+        for method in self.methods:
+            if method in estimators:
+                initialized_estimators.append(
+                    estimators[method](effect_type=effect_type, **kwargs)
+                )
+            else:
+                raise ValueError(f"Method '{method}' is not supported.")
+
+        return initialized_estimators
+
+    def _validate_inputs(self, df, treatment_col, outcome_col):
+        """
+        Validate the input DataFrame and columns for all estimators.
+        """
+        required_columns = [treatment_col, outcome_col]
+        # Check if all required columns exist in the DataFrame
+        for col in required_columns:
+            if col not in df.columns:
+                raise ValueError(f"Column '{col}' is missing from the DataFrame.")
+
+        # Additional validation logic if needed (e.g., check for NaN, etc.)
+        if df[treatment_col].isnull().any():
+            raise ValueError(f"Treatment column '{treatment_col}' contains NaN values.")
+        if df[outcome_col].isnull().any():
+            raise ValueError(f"Outcome column '{outcome_col}' contains NaN values.")
+
+    def _bootstrap_sample(self, df: pd.DataFrame, n_bootstraps: int):
+        """
+        Generate bootstrap samples.
+        """
+        n = len(df)
+        return [df.sample(n=n, replace=True) for _ in range(n_bootstraps)]
+
+    def compute_effect(
+        self,
+        df,
+        treatment_col,
+        outcome_col,
+        bootstrap=False,
+        n_bootstraps=100,
+        **kwargs,
+    ):
+        """
+        Compute treatment effects using the initialized estimators.
+        Can also run bootstrap on all estimators if specified.
+
+        Args:
+            df (pd.DataFrame): The input DataFrame.
+            treatment_col (str): The name of the treatment column.
+            outcome_col (str): The name of the outcome column.
+            bootstrap (bool): Whether to run bootstrapping for the estimators.
+            n_bootstraps (int): Number of bootstrap iterations.
+            sample_size (int): Size of each bootstrap sample.
+            **kwargs: Additional arguments for the estimators.
+
+        Returns:
+            dict: A dictionary where keys are method names and values are computed effects (and optionally standard errors).
+        """
+        # Validate input data and columns
+        self._validate_inputs(df, treatment_col, outcome_col)
+
+        results = {method.__class__.__name__: [] for method in self.estimators}
+
+        if bootstrap:
+            # Perform bootstrapping
+            bootstrap_samples = self._bootstrap_sample(df, n_bootstraps)
+
+            for sample in bootstrap_samples:
+                # For each bootstrap sample, compute the effect using all estimators
+                for estimator in self.estimators:
+                    method_name = type(estimator).__name__
+                    effect = estimator.compute_effect(
+                        sample, treatment_col, outcome_col, **kwargs
+                    )
+                    results[method_name].append(effect)
+
+            # After collecting all bootstrap samples, compute the mean and standard error for each estimator
+            final_results = {}
+            for method_name, effects in results.items():
+                effects_array = np.array(effects)
+                mean_effect = np.mean(effects_array)
+                std_err = np.std(effects_array)
+                final_results[method_name] = (mean_effect, std_err)
+
+        else:
+            # If no bootstrapping, compute the effect directly for each estimator
+            for estimator in self.estimators:
+                method_name = type(estimator).__name__
+                effect = estimator.compute_effect(
+                    df, treatment_col, outcome_col, **kwargs
+                )
+                results[method_name] = effect
+
+            final_results = results
+
+        return final_results

causalestimate-0.1.0/CausalEstimate/estimators/aipw.py

@@ -0,0 +1,31 @@
+from CausalEstimate.estimators.functional.aipw import compute_aipw_ate
+import pandas as pd
+
+
+class IPW:
+    def __init__(self, effect_type="ATE", **kwargs):
+        self.effect_type = effect_type
+        self.kwargs = kwargs
+
+    def compute_effect(
+        self,
+        df: pd.DataFrame,
+        treatment_col: str,
+        outcome_col: str,
+        ps_col: str,
+        predicted_outcome_treated_col: str,
+        predicted_outcome_control_col: str,
+    ) -> float:
+        """
+        Compute the effect using the functional IPW.
+        Available effect types: ATE, ATT, RR, RRT
+        """
+
+        A = df[treatment_col]
+        Y = df[outcome_col]
+        ps = df[ps_col]
+
+        if self.effect_type == "ATE":
+            return compute_aipw_ate(A, Y, ps)
+        else:
+            raise ValueError(f"Effect type '{self.effect_type}' is not supported.")

causalestimate-0.1.0/CausalEstimate/estimators/functional/aipw.py

@@ -0,0 +1,24 @@
+"""
+Augmented Inverse Probability of Treatment Weighting (AIPW)
+References:
+
+ATE:
+        Glynn, Adam N., and Kevin M. Quinn.
+        "An introduction to the augmented inverse propensity weighted estimator." 
+        Political analysis 18.1 (2010): 36-56.
+        note: This also provides a variance estimator for the AIPW estimator.
+"""
+
+from CausalEstimate.estimators.functional.ipw import compute_ipw_ate
+
+
+def compute_aipw_ate(A, Y, ps, Y0_hat, Y1_hat):
+    """
+    Augmented Inverse Probability of Treatment Weighting (AIPW) for ATE.
+    A: treatment assignment, Y: outcome, ps: propensity score
+    Y0_hat: P[Y|A=0], Y1_hat: P[Y|A=1]
+    """
+    ate_ipw = compute_ipw_ate(A, Y, ps)
+    adjustment_factor = (A - ps) / (ps * (1 - ps))
+    ate = ate_ipw - adjustment_factor * ((1 - ps) * Y1_hat + ps * Y0_hat)
+    return ate.mean()

causalestimate-0.1.0/CausalEstimate/estimators/functional/ipw.py

@@ -0,0 +1,106 @@
+"""
+Inverse Probability Weighting (IPW) estimators
+
+References: 
+ATE:
+    Estimation of Average Treatment Effects Honors Thesis Peter Zhang
+    https://lsa.umich.edu/content/dam/econ-assets/Econdocs/HonorsTheses/Estimation%20of%20Average%20Treatment%20Effects.pdf
+
+    Austin, P.C., 2016. Variance estimation when using inverse probability of 
+    treatment weighting (IPTW) with survival analysis. 
+    Statistics in medicine, 35(30), pp.5642-5655.
+
+ATT: 
+    Reifeis et. al. (2022). 
+    On variance of the treatment effect in the treated when estimated by 
+    inverse probability weighting. 
+    American Journal of Epidemiology, 191(6), 1092-1097.
+    https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9271225/
+"""
+
+
+def compute_ipw_risk_ratio(A, Y, ps):
+    """
+    Relative Risk
+    A: treatment assignment, Y: outcome, ps: propensity score
+    """
+    mu_1, mu_0 = compute_mean_potential_outcomes(A, Y, ps)
+    return mu_1 / mu_0
+
+
+def compute_ipw_ate(A, Y, ps):
+    """
+    Average Treatment Effect
+    A: treatment assignment, Y: outcome, ps: propensity score
+    """
+    mu1, mu0 = compute_mean_potential_outcomes(A, Y, ps)
+    return mu1 - mu0
+
+
+def compute_mean_potential_outcomes(A, Y, ps):
+    """
+    Compute E[Y|A=1] and E[Y|A=0] for Y=0/1
+    """
+    mu_1 = (A * Y / ps).mean()
+    mu_0 = ((1 - A) * Y / (1 - ps)).mean()
+    return mu_1, mu_0
+
+
+def compute_ipw_ate_stabilized(A, Y, ps):
+    """
+    Given by Austin (2016)
+    Average Treatment Effect with stabilized weights.
+    A: treatment assignment, Y: outcome, ps: propensity score
+    """
+    W = compute_stabilized_ate_weights(A, ps)
+    Y1_weighed = W * A * Y
+    Y0_weighed = W * (1 - A) * Y
+    return Y1_weighed - Y0_weighed
+
+
+def compute_ipw_att(A, Y, ps):
+    """
+    Average Treatment Effect on the Treated with stabilized weights.
+    Reifeis et. al. (2022).
+    A: treatment assignment, Y: outcome, ps: propensity score
+    """
+    mu_1, mu_0 = compute_mean_potential_outcomes_treated(A, Y, ps)
+    return mu_1 - mu_0
+
+
+def compute_ipw_risk_ratio_treated(A, Y, ps):
+    """
+    Relative Risk of the Treated with stabilized weights. Reifeis et. al. (2022)
+    A: treatment assignment, Y: outcome, ps: propensity score
+    """
+    mu_1, mu_0 = compute_mean_potential_outcomes_treated(A, Y, ps)
+    return mu_1 / mu_0
+
+
+def compute_mean_potential_outcomes_treated(A, Y, ps):
+    """
+    Compute E[Y|A=1] for Y=0/1
+    """
+    W = compute_stabilized_att_weights(A, ps)
+    mu_1 = (W * A * Y).sum() / (W * A).sum()
+    mu_0 = (W * (1 - A) * Y).sum() / (W * (1 - A)).sum()
+    return mu_1, mu_0
+
+
+def compute_stabilized_ate_weights(A, ps):
+    """
+    Compute the (stabilized) weights for the ATE estimator
+    Austin (2016)
+    """
+    weight_treated = A.mean() * A / ps
+    weight_control = (1 - A).mean() * (1 - A) / (1 - ps)
+    return weight_treated + weight_control
+
+
+def compute_stabilized_att_weights(A, ps):
+    """
+    Compute the (stabilized) weights for the ATT estimator
+    As given in the web appendix of Reifeis et. al. (2022)
+    """
+    h = ps / (1 - ps)
+    return A + (1 - A) * h

causalestimate-0.1.0/CausalEstimate/estimators/functional/matching.py

@@ -0,0 +1,42 @@
+import pandas as pd
+from CausalEstimate.utils.helpers import check_required_columns
+
+
+def compute_matching_ate(
+    Y: pd.Series,
+    matching_df: pd.DataFrame,
+    treated_col: str = "treated_pid",
+    control_col: str = "control_pid",
+) -> float:
+    """
+    Compute the effect using matching with vectorized Pandas operations.
+
+    Args:
+        Y (pd.Series): Outcomes for both treated and control units.
+        matching_df (pd.DataFrame): DataFrame containing matching results with columns
+                                    treated_col and control_col indicating the
+                                    matched treated and control unit IDs.
+
+    Returns:
+        float: The estimated ATE using the matched data.
+    """
+    check_required_columns(matching_df, [treated_col, control_col])
+    # Merge the treated outcomes with the control outcomes
+    merged_df = matching_df.merge(
+        Y.rename("treated_outcome"), left_on=treated_col, right_index=True
+    )
+    merged_df = merged_df.merge(
+        Y.rename("control_outcome"), left_on=control_col, right_index=True
+    )
+
+    # Compute the average control outcome for each treated unit using groupby
+    avg_control_outcomes = merged_df.groupby(treated_col)["control_outcome"].mean()
+
+    # Compute the difference between treated outcome and control outcome for each treated unit
+    treated_outcomes = Y[
+        avg_control_outcomes.index
+    ]  # Ensure we're aligning with grouped treated pids
+    diffs = treated_outcomes - avg_control_outcomes
+
+    # Return the average difference (ATE)
+    return diffs.mean()

causalestimate-0.1.0/CausalEstimate/estimators/functional/tmle.py

@@ -0,0 +1,41 @@
+from scipy.special import expit, logit
+from statsmodels.api import add_constant
+from statsmodels.genmod.families import Binomial
+from statsmodels.genmod.generalized_linear_model import GLM
+
+
+def compute_tmle_ate(A, Y, ps, Y0_hat, Y1_hat, Yhat):
+    """
+    Estimate the average treatment effect using the targeted maximum likelihood estimation (TMLE) method.
+    A: treatment assignment, Y: outcome, ps: propensity score,
+    Y0_hat: P[Y|A=0], Y1_hat: P[Y|A=1], Yhat: P[Y]
+    """
+    epsilon = estimate_fluctuation_parameter(A, Y, ps, Yhat)
+    return update_ate_estimate(ps, Y0_hat, Y1_hat, epsilon)
+
+
+def update_ate_estimate(ps, Y0_hat, Y1_hat, epsilon) -> tuple:
+    """Update the Q_star values using the fluctuation parameter epsilon."""
+    H_1 = 1 / ps
+    Q_star_1 = expit(logit(Y1_hat) + epsilon * H_1)
+
+    H_0 = 1 / (1 - ps)
+    Q_star_0 = expit(logit(Y0_hat) - epsilon * H_0)
+
+    return (Q_star_1 - Q_star_0).mean()
+
+
+def estimate_fluctuation_parameter(A, Y, ps, Yhat) -> float:
+    """
+    Estimate the fluctuation parameter epsilon using a logistic regression model.
+    Returns the estimated epsilon.
+    """
+    # compute the clever covariate H
+    H = A / ps - (1 - A) / (1 - ps)
+
+    # Use logit of the current outcome as offset
+    offset = logit(Y)
+
+    # Fit the model with offset
+    model = GLM(Yhat, add_constant(H), family=Binomial(), offset=offset).fit()
+    return model.params[0]

causalestimate-0.1.0/CausalEstimate/estimators/ipw.py

@@ -0,0 +1,37 @@
+from CausalEstimate.estimators.functional.ipw import (
+    compute_ipw_ate,
+    compute_ipw_ate_stabilized,
+    compute_ipw_att,
+    compute_ipw_risk_ratio_treated,
+    compute_ipw_risk_ratio,
+)
+
+
+class IPW:
+    def __init__(self, effect_type="ATE", **kwargs):
+        self.effect_type = effect_type
+        self.kwargs = kwargs
+
+    def compute_effect(self, df, treatment_col, outcome_col, ps_col) -> float:
+        """
+        Compute the effect using the functional IPW.
+        Available effect types: ATE, ATT, RR, RRT
+        """
+
+        A = df[treatment_col]
+        Y = df[outcome_col]
+        ps = df[ps_col]
+
+        if self.effect_type == "ATE":
+            if self.kwargs.get("stabilized", False):
+                return compute_ipw_ate(A, Y, ps)
+            else:
+                return compute_ipw_ate_stabilized(A, Y, ps)
+        elif self.effect_type == "ATT":
+            return compute_ipw_att(A, Y, ps)
+        elif self.effect_type == "RR":
+            return compute_ipw_risk_ratio(A, Y, ps)
+        elif self.effect_type == "RRT":
+            return compute_ipw_risk_ratio_treated(A, Y, ps)
+        else:
+            raise ValueError(f"Effect type '{self.effect_type}' is not supported.")

causalestimate-0.1.0/CausalEstimate/estimators/matching.py

@@ -0,0 +1,29 @@
+from CausalEstimate.matching.matching import match_optimal
+from CausalEstimate.estimators.functional.matching import compute_matching_ate
+
+
+class MATCHING:
+    def __init__(self, effect_type="ATE", **kwargs):
+        self.effect_type = effect_type
+        self.kwargs = kwargs
+
+    def compute_effect(self, df, treatment_col, outcome_col, ps_col) -> float:
+        """
+        Compute the effect using the functional IPW.
+        Available effect types: ATE, RR, OR
+        """
+
+        Y = df[outcome_col].values
+
+        df["index"] = df.index  # temporary index column
+        matched = match_optimal(
+            df,
+            treatment_col=treatment_col,
+            ps_col=ps_col,
+            pid_col="index",
+            **self.kwargs,
+        )
+        if self.effect_type == "ATE":
+            return compute_matching_ate(Y, matched)
+        else:
+            raise ValueError(f"Effect type '{self.effect_type}' is not supported.")

causalestimate-0.1.0/CausalEstimate/estimators/tmle.py

@@ -0,0 +1,35 @@
+from CausalEstimate.estimators.functional.tmle import compute_tmle_ate
+import pandas as pd
+
+
+class TMLE:
+    def __init__(self, effect_type="ATE", **kwargs):
+        self.effect_type = effect_type
+        self.kwargs = kwargs
+
+    def compute_effect(
+        self,
+        df: pd.DataFrame,
+        treatment_col: str,
+        outcome_col: str,
+        ps_col: str,
+        predicted_outcome_col: str,
+        predicted_outcome_treated_col: str,
+        predicted_outcome_control_col: str,
+    ) -> float:
+        """
+        Compute the effect using the functional IPW.
+        Available effect types: ATE
+        """
+
+        A = df[treatment_col]
+        Y = df[outcome_col]
+        ps = df[ps_col]
+        Y0_hat = df[predicted_outcome_control_col]
+        Y1_hat = df[predicted_outcome_treated_col]
+        Yhat = df[predicted_outcome_col]
+
+        if self.effect_type == "ATE":
+            return compute_tmle_ate(A, Y, ps, Y0_hat, Y1_hat, Yhat)
+        else:
+            raise ValueError(f"Effect type '{self.effect_type}' is not supported.")

causalestimate-0.1.0/CausalEstimate/filter/filter.py

@@ -0,0 +1,10 @@
+import pandas as pd
+
+
+def filter_by_column(df: pd.DataFrame, column: str, value: int) -> pd.DataFrame:
+    """
+    Filters the DataFrame by a specifeid column and value.
+    """
+    if column not in df.columns:
+        raise ValueError(f"Column '{column}' not found in the DataFrame.")
+    return df[df[column] == value].reset_index(drop=True)

causalestimate-0.1.0/CausalEstimate/matching/assignment.py

@@ -0,0 +1,50 @@
+from typing import Tuple
+import numpy as np
+import pandas as pd
+from scipy.optimize import linear_sum_assignment
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import min_weight_full_bipartite_matching
+from CausalEstimate.matching.helpers import is_sparse
+
+
+def assign_controls(distance_matrix: np.array) -> Tuple[np.array, np.array]:
+    """
+    Assigns controls to treated individuals using either sparse or dense assignment.
+    Uses sparse assignment if the distance matrix is sparse; otherwise, falls back to dense assignment.
+    """
+    if is_sparse(distance_matrix):
+        try:
+            row_indices, col_indices = sparse_assignment(distance_matrix)
+        except ValueError:
+            raise ValueError("Cannot assign unique controls to treated.")
+    else:
+        distance_matrix[distance_matrix == 0] = np.inf
+        row_indices, col_indices = linear_sum_assignment(distance_matrix)
+
+    return row_indices, col_indices
+
+
+def sparse_assignment(dist_mat: np.array) -> Tuple[np.array, np.array]:
+    """
+    Performs assignment using sparse matching.
+    """
+    sparse_d_mat = csr_matrix(dist_mat)
+    row_ind, col_ind = min_weight_full_bipartite_matching(sparse_d_mat)
+    return row_ind, col_ind
+
+
+def validate_control_availability(
+    treated: pd.DataFrame, controls: pd.DataFrame, n_controls: int
+) -> None:
+    """
+    Validates whether there are enough controls to match the treated individuals.
+    Raises an error if there aren't enough controls.
+    """
+    if len(treated) == 0:
+        raise ValueError("No treated units have sufficient controls")
+    if len(controls) < n_controls * len(treated):
+        raise ValueError(
+            f"Not enough controls to match.\nN_controls: {len(controls)}\
+            N_treated: {len(treated)}\
+            Required N_controls: {n_controls*len(treated)} (n_controls x n_treated)"
+        )

causalestimate-0.1.0/CausalEstimate/matching/distance.py

@@ -0,0 +1,44 @@
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+
+
+def compute_distance_matrix(
+    treated: pd.DataFrame, control: pd.DataFrame, ps_col: str
+) -> np.array:
+    """
+    Computes the distance matrix (absolute differences) between treated and control individuals based on propensity scores.
+    Args:
+        treated_df (pd.DataFrame): DataFrame of treated individuals.
+        control_df (pd.DataFrame): DataFrame of control individuals.
+        ps_col (str): Column name for propensity scores.
+    Returns:
+        np.array: Distance matrix with treated individuals as rows and controls as columns.
+    """
+    treated_ps = treated[ps_col].to_numpy()
+    control_ps = control[ps_col].to_numpy()
+    dist_mat = np.abs(treated_ps.reshape(-1, 1) - control_ps.reshape(1, -1))
+    return dist_mat
+
+
+def filter_treated_w_insufficient_controls(
+    dist_mat: np.array, treated_df: pd.DataFrame, n_controls: int
+) -> Tuple[np.array, pd.DataFrame]:
+    """
+    Filters out treated individuals who do not have enough valid controls based on the distance matrix.
+    Args:
+        distance_matrix (np.array): Distance matrix with treated individuals as rows and controls as columns.
+                            zero entries indicate invalid matches.
+        treated_df (pd.DataFrame): DataFrame of treated individuals.
+        n_controls (int): Minimum number of controls required for each treated individual.
+    Returns:
+        Tuple[np.array, pd.DataFrame]: Updated distance matrix and treated DataFrame after filtering.
+    """
+    valid_controls_per_treated = np.count_nonzero(dist_mat, axis=1)
+    sufficient_controls_mask = valid_controls_per_treated >= n_controls
+
+    if np.any(~sufficient_controls_mask):
+        dist_mat = dist_mat[sufficient_controls_mask].reshape(-1, dist_mat.shape[1])
+        treated_df = treated_df[sufficient_controls_mask]
+    return dist_mat, treated_df