python source code of ols

Project: whynot (GitHub Link)

whynot-master
- roadmap.txt
- whynot
  - gym
    - utils
      - __init__.py
    - __init__.py
    - spaces
      - __init__.py
    - envs
      - registration.py
      - ode_env.py
      - __init__.py
      - README.md
  - framework.py
  - dynamics.py
  - algorithms
    - causal_suite.py
    - ols.py
    - propensity_score_matching.py
    - __init__.py
    - propensity_weighted_ols.py
  - causal_graphs.py
  - __init__.py
  - utils.py
  - simulators
    - opioid
      - environments.py
      - .experiments.py.swo
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
    - lalonde
      - lalonde.csv
      - experiments.py
      - __init__.py
      - README.md
    - civil_violence
      - experiments.py
      - simulator.py
      - model.py
      - __init__.py
      - README.md
    - world3
      - environments.py
      - experiments.py
      - simulator.py
      - world3_app.js
      - __init__.py
      - README.md
    - dice
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
      - ipopt_bin
        coin-license.txt
        README.md
      - doc
        gams_2013.txt
    - delayed_impact
      - experiments.py
      - simulator.py
      - data
        totals.csv
        transrisk_performance_by_race_ssa.csv
        transrisk_cdf_by_race_ssa.csv
      - __init__.py
      - README.md
      - fico.py
    - __init__.py
    - hiv
      - environments.py
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
    - lotka_volterra
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
    - world2
      - experiments.py
      - simulator.py
      - tables.py
      - __init__.py
      - README.md
    - zika
      - environments.py
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
    - incarceration
      - experiments.py
      - simulator.py
      - __init__.py
      - README.md
    - schelling
      - experiments.py
      - simulator.py
      - model.py
      - __init__.py
      - README.md
    - credit
      - environments.py
      - simulator.py
      - dataloader.py
      - __init__.py
      - README.md
  - traceable_numpy
    - __init__.py
- examples
  - causal_inference
    - unobserved_confounding.ipynb
    - causal_discovery.ipynb
    - regression_discontinuity.ipynb
    - heterogeneous_treatment_effects.ipynb
    - interference.ipynb
    - graphical_methods.ipynb
    - mediation_analysis.ipynb
    - sensitivity_analysis.ipynb
    - instrumental_variables.ipynb
    - scripts
      - causal_search.py
      - mediator_utils.py
      - heterogeneous_utils.py
    - assets
      - lotka_volterra_graph.gml
      - mediators.svg
      - sensitivity_plots
  - reinforcement_learning
    - world3_simulator.ipynb
    - zika_simulator.ipynb
    - scripts
      - utils.py
  - walkthrough
    - walkthrough.ipynb
  - dynamic_decisions
    - scripts
      - utils.py
    - delayed_impact.ipynb
- .pre-commit-config.yaml
- setup.py
- .travis.yml
- README.md
- tests
  - gym_tests
    - test_envs.py
    - spec_list.py
    - test_registration.py
  - test_world2.py
  - test_causal_graphs.py
  - test_opioid.py
  - gamsBaseline.csv
  - test_dice.py
  - test_dynamics.py
  - test_delayed_impact.py
  - test_world3.py
  - test_framework.py
  - __init__.py
  - benchmark.py
  - test_algorithms.py
  - test_simulators.py
- pytest.ini
- requirements.txt
- lint.sh
- .gitignore
- docs
  - Makefile
  - source
    - estimators.rst
    - designing_experiments.rst
    - causal_graph_discovery.rst
    - causal_inference.rst
    - why.rst
    - usage.rst
    - api.rst
    - references.bib
    - examples.rst
    - treatment_effects.rst
    - develop.rst
    - reinforcement_learning.rst
    - quickstart.rst
    - index.rst
    - conf.py
    - simulators.rst
    - simulator_configs
      - civil_violence.rst
      - credit.rst
      - schelling.rst
      - lotka_volterra.rst
      - world2.rst
      - lalonde.rst
      - incarceration.rst
      - hiv.rst
      - zika.rst
      - opioid.rst
      - dice.rst
      - delayed_impact.rst
      - world3.rst
    - _static
      - WhyNot_fullcolor.svg
      - dynamics.svg
    - installation.rst
  - make.bat
  - .gitignore
- LICENSE.txt
- MANIFEST.in

"""Ordinary least-squares based-estimators for causal inference."""
from time import perf_counter

import numpy as np
import statsmodels.api as sm

from whynot.framework import InferenceResult


def estimate_treatment_effect(covariates, treatment, outcome):
    """Run ordinary least squares to estimate causal effect of treatment variable on the outcome Y.

    Parameters
    ----------
        covariates: `np.ndarray`
            Array of shape [num_samples, num_features] of features
        treatment:  `np.ndarray`
            Binary array of shape [num_samples]  indicating treatment status for each
            sample.
        outcome:  `np.ndarray`
            Array of shape [num_samples] containing the observed outcome for each sample.

    Returns
    -------
        result: `whynot.framework.InferenceResult`
            InferenceResult object for this procedure

    """
    features = np.copy(covariates)
    treatment = treatment.reshape(-1, 1)
    features = np.concatenate([treatment, features], axis=1)
    features = sm.add_constant(features, prepend=True, has_constant="add")

    # Only time model fitting, not preprocessing
    start_time = perf_counter()
    model = sm.OLS(outcome, features)
    results = model.fit()
    stop_time = perf_counter()

    # Treatment is the second variable (first is the constant offset)
    ate = results.params[1]
    stderr = results.bse[1]
    conf_int = (ate - 1.96 * stderr, ate + 1.96 * stderr)
    return InferenceResult(
        ate=ate,
        stderr=stderr,
        ci=conf_int,
        individual_effects=None,
        elapsed_time=stop_time - start_time,
    )