Welcome to lightweight-mcnnm’s documentation!

lightweight-mcnnm is a Python package that provides a lightweight and performant implementation of the Matrix Completion with Nuclear Norm Minimization (MC-NNM) estimator for causal inference in panel data settings.

Contents:

• Installation
  • Requirements
  • Installing from PyPI
  • Upgrading
  • Installing for Development
• Usage
  • JIT Compilation
  • Comprehensive Example
  • Basic Usage
  • Input Data
  • Generating Synthetic Data
  • Advanced Usage
  • Interpreting Results
• Validation Methods
  • 1. Cross-Validation
  • 2. Holdout Validation
  • Customizing Validation in estimate()
• API Reference
  • Main Functions
  • Classes
  • Internal Functions
• Examples
  • Basic Usage
  • Example 1: Staggered Adoption with Cross-Validation
  • Example 2: Block Assignment with Cross-Validation
  • Example 3: Single Treated Unit with Covariates
  • Example 4: Estimation with Holdout Validation
  • Example 5: Estimation with Pre-specified Lambda Values
  • Example 6: Estimation with Autocorrelation
  • Matrix Completion
  • Covariates
• Development
  • Setting up the Development Environment
  • Running Tests
  • Code Style
  • Pre-commit Hooks
  • Contributing
  • Building Documentation
• References
  • Academic References
  • Acknowledgments

What is lightweight-mcnnm?

lightweight-mcnnm implements the MC-NNM estimator exactly as described in Matrix Completion Methods for Causal Panel Data Models by Susan Athey, Mohsen Bayati, Nikolay Doudchenko, Guido Imbens, and Khashayar Khosravi (2021). This estimator provides a powerful tool for estimating causal effects in panel data settings, particularly when dealing with complex treatment patterns and potential confounders.

The implementation focuses on performance and minimal dependencies, making it suitable for use in various environments, including GPUs and cloud clusters.

Features

• Lightweight implementation with minimal dependencies

• Utilizes JAX for improved performance and GPU compatibility

• Faithful to the original MC-NNM algorithm as described in Athey et al. (2021)

• Suitable for large-scale panel data analysis

• Supports various treatment assignment mechanisms (staggered adoption, block assignment, single treated unit)

• Includes unit-specific, time-specific, and unit-time specific covariates

• Offers flexible validation methods for parameter selection (cross-validation and holdout)

Comparison to Other Implementations

lightweight-mcnnm is designed to be lightweight and easy to use, with a focus on performance and minimal dependencies. The other two main implementations of the MC-NNM estimator are CausalTensor and fect . Both packages implement MC-NNM as part of a broader set of causal inference methods. Both implement covariates and cross-validation differently from this package. For a detailed comparison, see this notebook.

Quick Start

Here’s a simple example of how to use lightweight-mcnnm:

from mcnnm import estimate, generate_data

Y, W, X, Z, V, true_params = generate_data(
     nobs=50,
     nperiods=10,
     unit_fe=True,
     time_fe=True,
     X_cov=True,
     Z_cov=True,
     V_cov=True,
     seed=2024,
     noise_scale=0.1,
     autocorrelation=0.0,
     assignment_mechanism="staggered",
     treatment_probability=0.1,
 )

# Run estimation
results = estimate(
    Y=Y,
    Mask=W,
    X=X,
    Z=Z,
    V=V,
    Omega=None,
    use_unit_fe=True,
    use_time_fe=True,
    lambda_L=None,
    lambda_H=None,
    validation_method='cv',
    K=3,
    n_lambda=30,
)

print(f"\nTrue effect: {true_params['treatment_effect']}, Estimated effect: {results.tau:.3f}")
print(f"Chosen lambda_L: {results.lambda_L:.4f}, lambda_H: {results.lambda_H:.4f}")

For more detailed information and examples, check out the Usage and Examples pages.

Indices and tables

• Index

• Module Index

• Search Page