Evaluating Hypothesis Tests

"""
Module for executing a given Monte Carlo scenario.

This module contains the SimulationRunner class, which executes a scenario
by combining a DGP and a test. It handles data generation, testing,
and result collection.

Classes:
    SimulationRunner: Runs simulations and summarizes results.
"""

import numpy as np

from sim_infrastructure.protocols import DGPProtocol, TestProtocol


class SimulationRunner:
    """Runs Monte Carlo simulations for a given DGP and test.

    Attributes:
        dgp (DGPProtocol): data-generating process with a sample() method and
            attributes that describe the DGP (covar_corr and common_coef_val).
        test (TestProtocol): test with a test() method and attributes that
            describe test name and test decision after seeing the data.
        test_decisions (np.ndarray): array of decisions made by the test in
            each Monte Carlo replication.
    """

    def __init__(
        self,
        dgp: DGPProtocol,
        test: TestProtocol,
    ) -> None:
        self.dgp: DGPProtocol = dgp
        self.test: TestProtocol = test
        self.test_decisions: np.ndarray = np.empty(0)

    def simulate(self, n_sim: int, n_obs: int, first_seed: int | None = None) -> None:
        """Runs simulations and stores test decisions.

        Args:
            n_sim (int): number of simulations to run.
            n_obs (int): number of sample points in each Monte Carlo dataset.
            first_seed (int | None): starting random seed for reproducibility.
                Defaults to None.
        """
        # Preallocate array to hold test decisions
        self.test_decisions = np.empty(n_sim)

        # Run simulation
        for sim_id in range(n_sim):
            # Draw data
            x, y = self.dgp.sample(
                n_obs, seed=first_seed + sim_id if first_seed else None
            )
            # Test null of zero coefficients
            self.test.test(x, y)
            # Store error
            self.test_decisions[sim_id] = self.test.decision

    def summarize_results(self) -> dict:
        """Return a summary of results: test name, DGP characteristics, power.

        Returns:
            dict: power of current test at current DGP
        """
        return {
            "test": self.test.name,
            "rho": self.dgp.covar_corr,
            "common_coef_val": self.dgp.common_coef_val,
            "power": self.test_decisions.mean(),
        }

"""
Module for joint tests for joint hypothesie.

Classes:
    WaldWithOLS: OLS-based Wald test that all non-intercept coefficients are zero.
"""

import numpy as np
import pandas as pd
from statsmodels.regression.linear_model import OLS


class WaldWithOLS:
    """Class for applying Wald test to test null of zero coefficients.

    Tailored to linear model
        Y_i = theta0 + theta1 X_{i1} + ... + thetap X_{ip} + U_i
        
    Tests the null:
        H0: theta1 = ... = thetap = 0

    Attributes:
        name (str): test name, "Wald"
        decision (np.bool): whether null is rejected
    """

    def __init__(self) -> None:
        self.name: str = "Wald"
        self.decision: np.bool

    def test(self, x: pd.DataFrame, y: pd.DataFrame) -> None:
        """Carry out the Wald test with given data.

        Args:
            x (pd.DataFrame): covariates, including a leading constant column.
            y (pd.DataFrame): outcomes.
        """
        # Fit models
        lin_reg = OLS(y, x)
        lin_reg_fit = lin_reg.fit()

        # Perform Wald test
        num_covars = x.shape[1]
        wald_restriction_matrix = np.concatenate(
            (np.zeros((num_covars - 1, 1)), np.eye(num_covars - 1)), axis=1
        )

        wald_test = lin_reg_fit.wald_test(
            wald_restriction_matrix,
            use_f=False,
            scalar=True,
        )

        self.decision = wald_test.pvalue <= 0.05

"""
Module for multiple tests for testing joint hypotheses.

Classes:
    BonferronigMultipleTWithOLS: OLS-based multiple t-test that all non-intercept
        coefficients are zero. Uses Bonferroni correction.
"""

import numpy as np
import pandas as pd
from statsmodels.regression.linear_model import OLS
from statsmodels.stats.multitest import multipletests


class BonferronigMultipleTWithOLS:
    """Class for applying a multiple t-test to test a null of zero coefficients.

    Tailored to linear model
        Y_i = theta0 + theta1 X_{i1} + ... + thetap X_{ip} + U_i

    Tests the null:
        H0: theta1 = ... = thetap = 0

    Attributes:
        name (str): test name, "Bonferroni t"
        decision (np.bool): whether null is rejected
    """

    def __init__(self) -> None:
        self.name: str = "Bonferroni t"
        self.decision: np.bool

    def test(self, x: pd.DataFrame, y: pd.DataFrame) -> None:
        """Carry out the multiple t-test with given data.

        Args:
            x (pd.DataFrame): covariates, including a leading constant column.
            y (pd.DataFrame): outcomes.
        """
        # Fit models
        lin_reg = OLS(y, x)
        lin_reg_fit = lin_reg.fit()

        # Perform multiple t-tests
        p_vals_t = lin_reg_fit.pvalues.iloc[1:]
        t_test_corrected_bonf = multipletests(
            p_vals_t,
            method="bonferroni", 
        )
        self.decision = t_test_corrected_bonf[0].sum() > 0

"""
Module for linear data-generating processes (DGPs).

This module contains classes for generating data from simple linear models.

Classes:
    BivariateLinearModel: A DGP for static linear models with normal variables,
        adjustable correlation between covariates, and two variables + constant.
"""

import numpy as np
import pandas as pd

class BivariateLinearModel:
    """A data-generating process (DGP) for a static linear model

    Attributes:
        common_coef_val (float): Common values for theta1 and theta2
        covar_corr (float): Correlation coefficient between covariates
    """

    def __init__(self, common_coef_val: float, covar_corr: float) -> None:
        self.common_coef_val: float = common_coef_val
        self.covar_corr: float = covar_corr

    def sample(
        self, n_obs: int, seed: int | None = None
    ) -> tuple[pd.DataFrame, pd.DataFrame]:
        """Samples data from the static DGP.

        Args:
            n_obs (int): Number of observations to sample.
            seed (int, optional): Random seed for reproducibility. Defaults to None.

        Returns:
            tuple: (x, y) DataFrames, each of length n_obs.
        """
        # Initialize RNG
        rng = np.random.default_rng(seed)

        # Construct mean and covariance for coefficients
        x_mean = np.array([1, 0, 0])
        x_covar = np.array([[0, 0, 0], [0, 1, self.covar_corr], [0, self.covar_corr, 1]])

        # Construct vector of coefficients
        thetas = np.array([1, self.common_coef_val, self.common_coef_val]
                          )
        # Draw covariates and residuals, combine into output 
        covariates = rng.multivariate_normal(
            x_mean,
            x_covar,
            size=n_obs,
        )
        resids = rng.normal(0, np.sqrt(1), size=n_obs)
        y = (covariates @ thetas) + resids

        # Convert output into pandas dataframes with dynamic variable names for X
        covariates_df = pd.DataFrame(
            covariates,
            columns=[f"X{i}" for i in range(covariates.shape[1])],
        )
        y_df = pd.DataFrame(
            y,
            columns=["y"],
        )
        return covariates_df, y_df

"""
Module for defining simulation scenarios.

This module contains scenarios for comparing power functions of tests under model:
    Y = theta0  + theta1*X1 + theta2*X2 + U

The null being tested is
    H0: theta1=theta2=0

Classes:
    SimulationScenario: data class for scenarios

Variables:
    scenarios (list): list of scenarios to un.
"""

from dataclasses import dataclass
from itertools import product

import numpy as np

from dgps.linear import BivariateLinearModel
from sim_infrastructure.protocols import DGPProtocol, TestProtocol
from tests.joint import WaldWithOLS
from tests.multiple import BonferronigMultipleTWithOLS


@dataclass(frozen=True)
class SimulationScenario:
    """A single simulation scenario: DGP, test, and sample size."""

    name: str  # For readability
    dgp: type[DGPProtocol]
    dgp_params: dict  # E.g. thetas go here
    test: type[TestProtocol]
    test_params: dict  #
    sample_size: int
    n_simulations: int = 500
    first_seed: int = 1


# Create DGP combinations indexed by correlation and coefficient values
common_coef_vals = np.linspace(-1.75, 1.75, 231)
covar_corr_vals = np.linspace(-0.99, 0.99, 51)
dgps = [
    (
        BivariateLinearModel,
        {"common_coef_val": common_coef_val, "covar_corr": covar_corr},
    )
    for common_coef_val, covar_corr in product(common_coef_vals, covar_corr_vals)
]

# Create list of tests
tests = [
    (WaldWithOLS, {}),
    (BonferronigMultipleTWithOLS, {}),
]

sample_sizes = [200]

# Generate all combinations
scenarios = [
    SimulationScenario(
        name="",
        dgp=dgp_class,
        dgp_params=dgp_params,
        test=test_class,
        test_params=test_params,
        sample_size=size,
    )
    for (dgp_class, dgp_params), (
        test_class,
        test_params,
    ), size in product(dgps, tests, sample_sizes)
]

from concurrent.futures import ThreadPoolExecutor, as_completed

from tqdm import tqdm

from sim_infrastructure.runner import SimulationRunner
from sim_infrastructure.scenarios import SimulationScenario


class SimulationOrchestratorParallel:
    """Parallelized simulation orchestrator class.

    Attributes:
        scenarios (list[SimulationScenario]): information about simulation
            scenario: DGP and estimator type, associated parameters, number of
            simulations to run, sample size, first seed.
        summary_results (list): simulation results for each scenario, as
            returned by corresponding simulation runner.
    """

    def __init__(self, scenarios: list[SimulationScenario]) -> None:
        self.scenarios = scenarios
        self.summary_results = []

    def _run_single_scenario(self, scenario: SimulationScenario):
        """Run a single scenario and return the results."""

        dgp = scenario.dgp(**scenario.dgp_params)
        estimator = scenario.test(**scenario.test_params)

        runner = SimulationRunner(dgp, estimator)
        runner.simulate(
            n_sim=scenario.n_simulations,
            n_obs=scenario.sample_size,
            first_seed=scenario.first_seed,
        )

        return runner.summarize_results()

    def run_all(self, max_workers: int | None = None):
        """Run all scenarios with thread-based parallelism.
        Args:
            max_workers (int | None, optional): number of threads to use in 
                execution. Defaults to None.
        """
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit all scenarios to the executor
            futures = {
                executor.submit(self._run_single_scenario, scenario): scenario
                for scenario in self.scenarios
            }

            # Collect results as they complete, with a progress bar
            for future in tqdm(
                as_completed(futures), total=len(futures), desc="Running simulations"
            ):
                self.summary_results.append(future.result())

"""
Entry point for running simulations for comparing power of a joint Wald test
vs. a multiple t-test.

This module contains scenarios for comparing power functions of tests under model:
    Y = theta0  + theta1*X1 + theta2*X2 + U

The null being tested is
    H0: theta1 = theta2 = 0

Scenarios considered: normal DGP with varying values of coefficients and varying
values of correlation between x1 and x2.

Usage:
    python -X gil=0 main.py

Output:
    Power surface comparison plots saved under PLOT_FOLDER
"""

from pathlib import Path

import pandas as pd

from sim_infrastructure.orchestrators import SimulationOrchestratorParallel
from sim_infrastructure.results_processor import ResultsProcessor
from sim_infrastructure.scenarios import scenarios

SIM_RESULTS_PATH = Path() / "results" / "sim_results.csv"
PLOT_FOLDER = Path() / "results" / "plots"


if __name__ == "__main__":
    # Create and execute simulations
    orchestrator = SimulationOrchestratorParallel(scenarios)
    orchestrator.run_all()

    # Export results
    pd.DataFrame(orchestrator.summary_results).to_csv(SIM_RESULTS_PATH)

    # Export plots
    results_processor = ResultsProcessor(SIM_RESULTS_PATH, PLOT_FOLDER)
    results_processor.export_all_plots()