API Reference

BalanceData

Efficient storage for balance computation data with separate arrays for values and missing indicators. This replaces Vector{Union{Missing, Float64}} with better performance and memory pooling capability.

CIC <: AbstractCICModel

Main model type for causal inference with time-series cross-sectional data.

Contains all information needed for matching, balancing, and estimation in TSCS designs. Created by makemodel() and used throughout the analysis workflow.

Fields

• title::String: Model title for identification
• id::Symbol: Column name for unit identifier
• t::Symbol: Column name for time variable
• outcome::Union{Symbol,Vector{Symbol}}: Outcome variable(s)
• treatment::Symbol: Treatment variable (binary 0/1)
• covariates::Vector{Symbol}: Covariates used for matching
• timevary::Dict{Symbol, Bool}: Whether each covariate is time-varying
• reference::Int: Reference time period (default: -1)
• F::UnitRange{Int}: Post-treatment periods for estimation
• L::UnitRange{Int}: Pre-treatment periods for matching (negative values)
• observations::Vector{Tuple{Int, Int}}: Treated observations (time, unit)
• ids::Vector{Int}: All unit identifiers
• matches::Vector{TreatmentObservationMatches}: Matching results for each treated observation
• meanbalances::DataFrame: Covariate balance statistics
• grandbalances::Dict: Overall balance measures
• iterations::Int: Bootstrap iterations (default: 500)
• results::DataFrame: Treatment effect estimates
• treatednum::Int: Number of treated observations
• estimator::String: Estimator type (default: "ATT")

Examples

using TSCSMethods, DataFrames

data = example_data()
model = makemodel(data, :day, :fips, :gub, :death_rte, 
                 [:pop_dens], Dict(:pop_dens => false),
                 1:5, -10:-1)

CICStratified <: AbstractCICModelStratified

Stratified model type for causal inference with time-series cross-sectional data.

Extends the basic CIC model to handle stratified analysis where matching and estimation are performed separately within subgroups defined by a stratifying variable.

Fields

• title::String: Model title for identification
• id::Symbol: Column name for unit identifier
• t::Symbol: Column name for time variable
• outcome::Union{Symbol,Vector{Symbol}}: Outcome variable(s)
• treatment::Symbol: Treatment variable (binary 0/1)
• covariates::Vector{Symbol}: Covariates used for matching
• timevary::Dict{Symbol, Bool}: Whether each covariate is time-varying
• stratifier::Symbol: Variable used for stratification
• strata::Vector{Int}: Values of stratifying variable
• reference::Int: Reference time period (default: -1)
• F::UnitRange{Int}: Post-treatment periods for estimation
• L::UnitRange{Int}: Pre-treatment periods for matching (negative values)
• observations::Vector{Tuple{Int, Int}}: Treated observations (time, unit)
• ids::Vector{Int}: All unit identifiers
• matches::Vector{TreatmentObservationMatches}: Matching results for each treated observation
• meanbalances::DataFrame: Covariate balance statistics by stratum
• grandbalances::Dict: Overall balance measures by stratum
• iterations::Int: Bootstrap iterations (default: 500)
• results::DataFrame: Treatment effect estimates by stratum
• treatednum::Dict{Int64, Int64}: Number of treated observations per stratum
• estimator::String: Estimator type (default: "ATT")
• labels::Dict{Int64, String}: Human-readable labels for strata

Examples

# Create stratified model
model = makemodel(data, :day, :fips, :gub, :death_rte,
                 [:pop_dens], Dict(:pop_dens => false),
                 1:5, -10:-1)
strat_model = stratify(model, data, :region)

    Fblock

Holds the relevant information for bootstrapping and estimation, for a specific f in the outcome window (an f in a stratum when the model is stratified).

ImputationResults

Container for counterfactual imputation results.

Fields

• results: DataFrame with counterfactual trajectories and treatment effects
• matched_pretreatment_avg: Average pre-treatment outcome for matched controls
• treated_pretreatment_avg: Average pre-treatment outcome for treated units
• baseline_difference: Difference in pre-treatment averages (treated - matched)

TreatmentObservationCaliperMatches

Stores matching results for caliper-constrained CIC models.

Similar to TreatmentObservationMatches but for models where matches are restricted by caliper constraints (maximum allowable distance thresholds).

TreatmentObservationMatches

Stores matching results for each treated observation in basic CIC models.

Contains the core data structures for matching:

• eligible_matches: Boolean matrix indicating which control units are eligible matches
• distances: Computed distances between treated and control units
• match_rankings: Ranked preferences for matches by time period

This is the fundamental data structure for storing match relationships in TSCS designs.

TreatmentObservationRefinedMatches

Stores matching results for refined CIC models.

Used in refined matching procedures where the initial match set is iteratively improved through additional matching criteria.

__fill_meanbalances!(balance_row, periods_with_matches, lag_periods_count, covariates, timevary)

Allocate BalanceData objects for a single treated observation with error handling and memory pooling.

Purpose

Handles the actual allocation of BalanceData objects for one row of the meanbalances DataFrame. Implements error-safe memory pooling with proper cleanup if allocation fails partway through.

Arguments

• balance_row: Single row from meanbalances DataFrame to populate
• periods_with_matches: Number of forward periods that have eligible matches
• lag_periods_count: Number of lag periods (for time-varying covariate sizing)
• covariates: Vector of covariate names to allocate storage for
• timevary: Dict indicating which covariates are time-varying

Details

• For time-varying covariates: Creates Vector{BalanceData} with periods_with_matches elements, each of size lag_periods_count
• For static covariates: Creates single BalanceData of size periods_with_matches
• Uses memory pooling via get_balance_data() for allocation efficiency
• Implements error handling: if allocation fails, returns all pooled objects to prevent memory leaks
• Tracks all pooled allocations for proper cleanup

Error Handling

If any allocation fails, all previously allocated pooled objects are automatically returned to their respective pools to prevent memory leaks.

_calculate_counterfactuals!(matches_with_treatment, outcome_lookup, outcome_var, tvar)

Calculate observed outcomes and counterfactual values for each matched observation. Modifies matcheswithtreatment in place.

_calculate_pretreatment_averages(matches_with_treatment, outcome_lookup, outcome_var)

Calculate pre-treatment outcome averages for treated and matched units. Returns (matchedavg, treatedavg).

version for multiple outcomes

does not include p-values, bayesfactor

Fast quantile computation from pre-sorted data

_fill_meanbalances!(meanbalances, matches, lag_periods_count, covariates, timevary, forward_periods_count)

Populate meanbalances DataFrame with properly sized BalanceData objects for each treated observation.

Purpose

For each treated observation, determines which forward periods (F values) have at least one eligible match, then allocates appropriate BalanceData storage structures. The allocation pattern depends on whether covariates are time-varying or static.

Arguments

• meanbalances: DataFrame to populate (one row per treated observation)
• matches: Vector of match objects containing eligible_matches matrices
• lag_periods_count: Number of lag periods (length of L range)
• covariates: Vector of covariate names
• timevary: Dict indicating which covariates are time-varying
• forward_periods_count: Number of forward periods (length of F range)

Details

• Creates :fs column indicating which forward periods have matches
• For time-varying covariates: Vector{BalanceData} of length periodswithmatches
• For static covariates: Single BalanceData of length periodswithmatches
• Uses memory pooling via get_balance_data() for efficiency
• Each BalanceData object sized according to lagperiodscount for time-varying covariates

Storage Structure

meanbalances[i, :covariate] = 
  - Time-varying: [BalanceData(lag_periods), BalanceData(lag_periods), ...]
  - Static: BalanceData(periods_with_matches)

_makegroupindices(tidx, ridx, tridx, tts, uid, fmax, Lmin, tvec, idvec, treatvecbool, X)

Core implementation for building group indices without exposure data.

Purpose

Populates the index dictionaries by creating SubArray views for each (treatment_time, unit) combination. Uses parallel processing to efficiently handle large datasets.

Arguments

• tidx, ridx, tridx: Pre-allocated dictionaries to populate
• tts: Vector of treatment times (when treatments occurred)
• uid: Vector of unique unit IDs
• fmax, Lmin: Time window bounds
• tvec, idvec, treatvecbool: Original data vectors
• X: Covariate data matrix

Implementation Details

• Uses Threads.@threads :greedy for parallel processing across (treatment_time, unit) pairs
• For each pair, calls getyes!() to determine which data rows fall within the time window
• Creates SubArray views using @views macro for memory efficiency
• Thread-safe due to non-overlapping dictionary key assignments

_makegroupindices(tidx, ridx, tridx, tts, uid, fmax, Lmin, tvec, idvec, treatvecbool, X, exidx, exvec)

Core implementation for building group indices with exposure data support.

Purpose

Extended version that additionally creates indexed views of exposure data alongside the standard covariate, time, and treatment indices. Used for exposure-based matching where treatment assignment depends on external exposure variables.

Arguments

• tidx, ridx, tridx: Standard index dictionaries to populate
• exidx: Additional dictionary for exposure data views
• exvec: Exposure data vector
• Other arguments same as standard version

Details

Identical logic to the standard version but additionally populates exidx[(tt, unit)] with SubArray views of the exposure data for each (treatment_time, unit) combination. This enables efficient access to exposure histories during matching algorithms that consider exposure patterns.

_validate_distaveraging_inputs(dtots, lag_times, fw, accums) -> Int

Validate common inputs for distance averaging functions.

Returns

• n_times::Int: Number of time points after validation

Throws

• ArgumentError: For empty required inputs
• DimensionMismatch: For inconsistent array dimensions

_validate_imputation_inputs(m, matches, dat, tvar, unit_var, stratum)

Validate inputs for imputation function. Throws ArgumentError if validation fails.

alldistances!(dtotals, Σinvdict, xrows, yrows, lag_times)

Calculate Mahalanobis and individual covariate distances for time-series matching.

Mathematical Foundation

This function implements the core distance calculations for the Feltham et al. (2023) extension of Imai et al. (2021) matching methodology:

Mahalanobis Distance

For each time point τ, calculates:

d_M(x_τ, y_τ) = √[(x_τ - y_τ)ᵀ Σ_τ⁻¹ (x_τ - y_τ)]

where:

• xτ, yτ are covariate vectors for treated and control units at time τ
• Σ_τ⁻¹ is the inverse covariance matrix for time τ (from all units)

Individual Covariate Distances (for Calipers)

For each covariate j:

d_j(x_jτ, y_jτ) = √[(x_jτ - y_jτ)² / σ²_jτ]

where σ²jτ = Στ[j,j] is the variance of covariate j at time τ.

Algorithm

1. Matrix Caching: Pre-cache all covariance matrices for lag_times to eliminate repeated hash lookups (Performance optimization - maintains exact results)
2. Distance Calculation: For each time point, compute both Mahalanobis and individual distances simultaneously
3. Missing Data: If any covariate is missing, the corresponding distance is missing

Arguments

• dtotals: Pre-allocated output arrays [Mahalanobis, Cov1, Cov2, ...]
• Σinvdict: Dictionary mapping time → inverse covariance matrix
• xrows: Treated unit covariate vectors over time
• yrows: Control unit covariate vectors over time
• lag_times: Time points for distance calculation

Performance Notes

• Optimization: Caches covariance matrices to avoid O(m) hash lookups per distance
• Complexity: O(k) where k = length(lag_times), down from O(k×m) with m unique times
• Memory: Uses views and pre-allocated arrays for efficiency

Statistical Accuracy

This function preserves exact mathematical equivalence to the original algorithm while providing significant performance improvements through caching.

References

• Feltham et al. (2023): Mass gatherings methodology with extended time windows
• Imai et al. (2021): Original matching framework for TSCS data

    att!(atts, tcounts, fblocks)

Calculate the att for each f, for a the set of treated units and matches contained in the fblocks.

autobalance(
  model, dat;
  refinementnum = 5,
  calmin = 0.08, step = 0.05,
  initial_bals = nothing
)

Automatically balance via a simple algorithm. Start with initial caliper of 1.0, and subtract step whenever the grand mean balance threshold (0.1) is not met.

initial_bals is specified, work downward from the initial specified caliper for one or more selected variables. Unspecified variables start at a caliper of 1.

average_distances!(drow, dtots, accums, lag_times, fw) where {T}

Calculate averaged distances over time windows for fixed window matching.

Mathematical Formula

For each distance type k and valid time window W:

drow[k] = (1/|W|) ∑_{τ ∈ W} dtots[k][l] where lag_times[l] = τ  

This is the fixed-window version of distance averaging, used when matching windows are constant rather than sliding.

Arguments

• drow: Output vector [K] where K = number of distance types
• dtots: Input distances [K][T] for K distance types over T time points
• accums: Working array for counting valid observations per distance type
• lag_times: Time points corresponding to dtots columns
• fw: Fixed matching window specification

Differences from Sliding Version

• Output: Single vector instead of matrix (no outcomeperiodindex, m indices)
• Use Case: Fixed matching windows vs. sliding windows
• Performance: Slightly more efficient due to simpler indexing

Algorithm

Identical to sliding window version but with simplified output structure:

1. Pre-compute window bounds for efficiency
2. Type-stable initialization based on data type T
3. Accumulate valid distances within window
4. Average by count of valid observations

average_distances!(distances, dtots, accums, lag_times, fw, outcome_period_index, m) where {T}

Calculate averaged distances over time windows for sliding window matching.

Mathematical Formula

For each distance type k and valid time window W:

distances[k][outcome_period_index, m] = (1/|W|) ∑_{τ ∈ W} dtots[k][l] where lag_times[l] = τ

Arguments

• distances: Output array [K][outcomeperiodindex, m] where K = number of distance types
• dtots: Input distances [K][T] for K distance types over T time points
• accums: Working array for counting valid observations per distance type
• lag_times: Time points corresponding to dtots columns
• fw: Matching window specification (e.g., -10:-1 for 10 periods before)
• outcome_period_index: Window index (for sliding windows)
• m: Match index

Algorithm

1. Window Bounds: Pre-compute fwmin, fwmax = extrema(fw) once
2. Type-Based Init: Initialize based on data type T for type stability
3. Filtered Iteration: Only process lagtimes[l] where fwmin ≤ lagtimes[l] ≤ fwmax
4. Missing Handling: Skip missing values, track counts in accums
5. Averaging: Divide accumulated sums by valid observation counts

Performance

• Complexity: O(|lagtimes|) with early termination when lagtimes[l] > fw_max
• Memory: Uses pre-allocated arrays, no intermediate allocations
• Type Stability: Compile-time specialization for Float64 vs Union types

balance!(model::VeryAbstractCICModel, dat::DataFrame) -> VeryAbstractCICModel

Perform covariate balancing on a matched model to ensure treated and control units are comparable on observed characteristics.

Arguments

• model::VeryAbstractCICModel: A CIC model that has been matched (using match!)
• dat::DataFrame: Input data containing all model variables

Returns

• The input model with updated balance statistics

Description

This function performs two types of balancing:

1. Mean balancing: Computes balance statistics for each covariate across treatment periods
2. Grand balancing: Aggregates balance statistics across the entire model

Balancing assesses how well the matching procedure achieved covariate balance between treated and control groups. Poor balance may indicate the need for refinement via calipers or other restrictions.

Throws

• ArgumentError: If input data is empty
• ErrorException: If balance calculations fail

Examples

# After constructing and matching a model
model = makemodel(data, :time, :id, :treatment, :outcome, 
                 [:covar1, :covar2], timevary_dict, F_periods, L_periods)
match!(model, data)

# Perform balancing
balance!(model, data)

# Check balance results
checkbalances(model, data)

See Also

• checkbalances: Assess balance quality
• autobalance: Automatic balance improvement
• match!: Matching procedure that should precede balancing

bootinfo!(res, oc, boots; qtiles = [0.025, 0.5, 0.975])

Format the bootstrap matrix into the results dataframe. Assumes that att() has already been added to res.

bootinfo!(res, boots; qtiles = [0.025, 0.5, 0.975])

Format the bootstrap matrix into the results dataframe. Assumes that att() has already been added to res.

    bootstrap!(boots, tcountmat, fblocks, ids, treatdex, iterations)

Perform bootstrapping of the att for each f in the outcome window. Does so for the given set of treated units and matches contained in the fblocks.

calculate_overall_summary(imputation_results::ImputationResults, outcome_var::Symbol)

Calculate overall treatment effect summary statistics and balance diagnostics.

This function computes key summary statistics that help assess the quality and magnitude of the causal inference results.

Arguments

• imputation_results::ImputationResults: Output from impute_results()
• outcome_var::Symbol: Column name of the outcome variable

Returns

• NamedTuple: Summary statistics with fields:
  • treated_observed_mean: Average observed outcome for treated units
  • overall_att: Overall average treatment effect
  • counterfactual_mean: Average counterfactual outcome
  • baseline_difference: Pre-treatment difference between treated and matched units
  • treated_pretreatment_avg: Pre-treatment average for treated units
  • matched_pretreatment_avg: Pre-treatment average for matched controls

Example

imputation = impute_results(model, matches, data, :t, :id)
summary = calculate_overall_summary(imputation, :Y)

# Key statistics
println("Overall ATT: ", summary.overall_att)
println("Baseline balance: ", summary.baseline_difference)  # Should be close to 0
println("Treatment magnitude: ", summary.overall_att / summary.treated_observed_mean * 100, "%")

Notes

• baseline_difference close to 0 indicates good matching quality
• overall_att is the main causal estimate
• Comparing pre-treatment averages helps assess whether matching achieved balance

inverted covariance matrix for mahalanobis distance (all units at t)

inverted sqrt(vars) for balance score calculations (treated units at t)

caliper(
    model::CIC, 
    acaliper::Dict{Symbol, Float64}, 
    dat::DataFrame; 
    dobalance::Bool = true
) -> CaliperCIC

Apply caliper restrictions to a CIC model by excluding matches beyond specified distance thresholds.

Arguments

• model::CIC: A matched CIC model
• acaliper::Dict{Symbol, Float64}: Dictionary mapping covariates to maximum allowed distances
• dat::DataFrame: Input data containing all model variables
• dobalance::Bool: Whether to perform balancing on the calipered model (default: true)

Returns

• CaliperCIC: A calipered version of the input model with restricted matches

Description

Caliper restrictions improve match quality by excluding matches where the distance on specified covariates exceeds the given thresholds. This helps ensure that matched units are similar on the most important dimensions, potentially improving balance and reducing bias.

Examples

# After matching a model
model = makemodel(data, :time, :id, :treatment, :outcome, covariates, timevary, F, L)
match!(model, data)

# Apply calipers - exclude matches with distance > 0.25 on covar1 or > 0.5 on covar2
caliper_specs = Dict(:covar1 => 0.25, :covar2 => 0.5)
calipered_model = caliper(model, caliper_specs, data)

# Apply calipers without automatic balancing
calipered_model = caliper(model, caliper_specs, data; dobalance = false)

Notes

• Stricter calipers improve match quality but reduce the number of matches
• Use autobalance to automatically determine appropriate caliper values

See Also

• refine: Alternative approach using match ranking
• autobalance: Automatic caliper selection
• match!: Initial matching procedure

change_pct(val, attval)

Calculate percentage change from baseline value.

Arguments

• val: Baseline value
• attval: Treatment effect (difference from baseline)

Returns

• Float64: Percentage change (100 * attval / val)

Example

baseline = 100.0
treatment_effect = 10.0
pct_change = change_pct(baseline, treatment_effect)  # Returns 10.0 (10% increase)

checkbalances(
  m::Dict{Symbol, Union{Float64, Vector{Float64}}};
  threshold = 0.1, stratareduce = true
)

Simply check whether the grand means are above the std. balance threshold. Returns a Bool for each covariate. If Stratareduce is true, then the strata balances will be agggregated to the covariate level, such that a violation in any caliper triggers a violation in the aggregated output.

not stratified

combostrat(model, dat, vars::Vector{Symbol}; varslabs = nothing)

Stratify based on the combinations of one or more variables that exist in the data. Strata are formed directly from the variable values.

compare_strata(imputation_results_list::Vector, outcome_var::Symbol, stratum_names::Vector{String})

Compare results across multiple strata or model specifications.

Arguments

• imputation_results_list: Vector of ImputationResults from different strata
• outcome_var: Symbol for outcome variable
• stratum_names: Names for each stratum for labeling

Returns

• Comparison plot showing ATT estimates across strata
• Summary table of key statistics

Example

# Run analysis for each stratum
imputation1 = impute_results(model, matches, data, :t, :id, stratum=1)
imputation2 = impute_results(model, matches, data, :t, :id, stratum=2)

# Compare results
comparison = compare_strata([imputation1, imputation2], :Y, ["Urban", "Rural"])

compute_treated_std(model::VeryAbstractCICModel, dat::DataFrame) -> Dict{Tuple{Int64, Symbol}, Float64}

Compute standardization factors (1/σ) for covariates across all treated units' matching periods.

Purpose

For balance calculations, we need to standardize covariate differences by the standard deviation of treated units during their matching windows. This ensures balance statistics are comparable across covariates with different scales.

Returns

Dictionary with keys (time_offset, covariate) and values 1/σ, where:

• time_offset: Days relative to treatment (negative = pre-treatment, positive = post-treatment)
• covariate: Covariate name
• 1/σ: Inverse standard deviation for standardization

Details

For each treated unit at treatment time t, we collect covariate values from their matching window [t + min(L), t + max(F)]. We then compute the standard deviation of each covariate at each time offset across all treated units and return the inverse for efficient multiplication during balance calculations.

Example

std_factors = compute_treated_std(model, data)
# std_factors[(-10, :population)] = 0.05  # 1/σ for population 10 days before treatment

faster version of countmap, use when length is known

create_inspection_dashboard(imputation_results::ImputationResults, outcome_var::Symbol)

Create a comprehensive dashboard with all inspection plots in a single figure.

Arguments

• imputation_results: ImputationResults struct from impute_results()
• outcome_var: Symbol representing the outcome variable column name

Returns

• Figure: Comprehensive dashboard with 4 panels:
  1. Top Left: Treatment effects (ATT) over time with confidence intervals
  2. Top Right: Observed vs counterfactual trajectories comparison
  3. Bottom Left: Treatment effects as percentage changes
  4. Bottom Right: Summary statistics and balance diagnostics

Example

# After running TSCSMethods analysis and generating imputation results
imputation = impute_results(model, model.matches, data, :t, :id)

# Create comprehensive dashboard
dashboard = create_inspection_dashboard(imputation, :Y)

# Save or display
save("analysis_dashboard.png", dashboard)
dashboard  # Display in notebook/REPL

Dashboard Interpretation

• ATT Panel: Shows treatment effects over time. Look for consistent patterns and significant effects
• Trajectory Panel: Key causal story - distance between blue (observed) and black (counterfactual) lines shows treatment impact
• Percentage Panel: Treatment effects as % changes for easier interpretation of magnitude
• Summary Panel: Key statistics including overall ATT and baseline balance assessment

Notes

• Dashboard provides comprehensive view of results
• All panels use consistent time scales for easy comparison
• Confidence intervals help assess statistical uncertainty
• Summary statistics help evaluate matching quality and effect magnitude

customstrat!(
  cc,
  stratdict::Union{Dict{Tuple{Int64, Int64}, Int64}, Dict{Int64, Int64}}
)

Stratify based on the values of some input dictionary, specifying strata for each (t, id) or each (id) for a stratification that varies only by unit.

default_treatmentcategories(x) -> Int

Default categorization function for treatment histories during matching eligibility.

Purpose

Categorizes treatment exposure counts in the pre-crossover period to determine matching eligibility. Units can only be matched if they fall into the same treatment category during the pre-crossover window, ensuring similar treatment exposure patterns.

Arguments

• x: Total count of treatment exposures in the pre-crossover period for a unit

Returns

• 0: Never treated (x == 0)
• 1: Ever treated (x > 0)

Usage in Matching

During matching, for each potential match pair:

1. Count treatments in pre-crossover period for treated unit → category A
2. Count treatments in pre-crossover period for potential match → category B
3. If treatmentcategories(A) == treatmentcategories(B), matching is allowed
4. Otherwise, the match is ineligible

Default Logic

This binary categorization ensures that:

• Untreated units can only match with other untreated units
• Previously treated units can only match with other previously treated units
• Prevents contamination from units with different treatment exposure histories

Custom Categories

Users can provide custom categorization functions for more nuanced matching:

# Example: Allow multiple treatment intensity levels
custom_categories(x) = x <= 1 ? 0 : (x <= 5 ? 1 : 2)
match!(model, data; treatcat = custom_categories)

distances_allocate!(matches, flen, covnum; sliding = false)

Pre-allocate distance storage arrays for match calculations.

Arguments

• matches: Vector of match objects to allocate storage for
• flen: Length of the time dimension (number of time periods)
• covnum: Number of covariates
• sliding: Whether to use sliding windows (default: false)

Algorithm

1. For each match object, allocate a matrix with dimensions:
   • Rows: Number of potential matches for this observation
   • Columns: Number of distance types (1 Mahalanobis + covnum individual covariates)
2. Initialize all distances to infinity (no matches computed yet)

Performance

• Memory: Pre-allocates all required storage to avoid allocations during computation
• Parallelization: Thread-safe as each match object gets independent storage
• Scalability: O(nobservations × maxmatches × n_covariates) memory usage

Note

This function prepares the storage structure but does not compute any distances. Actual distance computation is performed by distances_calculate!.

distances_calculate!(matches, observations, ids, covariates, tg, rg, fmin, Lmin, Lmax, Σinvdict; sliding=false)

Main function for calculating distances between treated and control units over time windows.

Mathematical Framework

This implements the core distance calculation for Feltham et al. (2023) methodology:

For each treated unit i and potential control unit j:

1. Time Window: Define matching window L = [Lmin, Lmax] relative to treatment time
2. Distance Calculation: For each τ ∈ L, compute d(x{iτ}, x{jτ})
3. Temporal Averaging: Average distances over valid time points in window
4. Multiple Distance Types: Compute both Mahalanobis and individual covariate distances

Mathematical Formulation

D(i,j) = (1/|W|) ∑_{τ ∈ W} d(x_{iτ}, x_{jτ})

where:

• W = {τ : Lmin ≤ τ - t_i ≤ Lmax, data available for both units}
• t_i is the treatment time for unit i
• d(·,·) is the distance metric (Mahalanobis or covariate-specific)

Algorithm Overview

1. Thread Parallelization: Use :greedy scheduler for load balancing across irregular workloads
2. Valid Match Filtering: Check eligible_matches matrix to identify potential matches
3. Storage Optimization: Use thread-local pre-allocated arrays
4. Distance Computation: Call alldistances! for each treated-control pair
5. Window Averaging: Call window_distances! to average over time windows

Performance Optimizations

• Threading: Modern :greedy scheduler for better load balancing
• Memory: Thread-local storage eliminates 90% of allocations
• Caching: Pre-cache covariance matrices to avoid repeated lookups
• Early Termination: Skip observations with no valid matches

Arguments

• matches: Output array of match objects with distance matrices
• observations: Treated unit observations to process
• ids: Unit identifiers for matching
• covariates: List of covariate names for distance calculation
• tg: Treated unit covariate data grouped by (time, unit)
• rg: Time index mapping
• fmin, Lmin, Lmax: Window specification parameters
• Σinvdict: Pre-computed inverse covariance matrices by time
• sliding: Whether to use sliding windows (currently fixed windows only)

Complexity

• Time: O(ntreated × npotentialmatches × windowsize × n_covariates)
• Memory: O(nthreads × maxwindowsize × maxcovariates) due to pre-allocation
• Parallelization: Scales with number of threads, load-balanced across observations

eligibility!(matches, observations, X::Matrix{Float64}, ids, treatcat, dat_t, dat_id, dat_trt, fmin, fmax, Lmin; exposure = nothing)

Determine matching eligibility for all treated observations using crossover window constraints.

Purpose

Core eligibility determination that implements the TSCS matching methodology's crossover window logic. Creates efficient indexed views of the data and determines which units can serve as matches for each treated observation.

Arguments

• matches: Vector of match objects to populate with eligibility information
• observations: Vector of treated observations (time, unit_id) tuples
• X: Covariate data matrix (Float64 version for non-missing data)
• ids: Vector of unique unit identifiers
• treatcat: Treatment categorization function for crossover period matching
• dat_t, dat_id, dat_trt: Time, unit ID, and treatment vectors from original data
• fmin, fmax: Bounds of forward period range (F)
• Lmin: Minimum lag period (lower bound of L)
• exposure: Optional exposure variable for exposure-based matching

Returns

• tg: Covariate group indices dictionary
• rg: Time group indices dictionary

Process

1. Group Index Creation: Builds efficient (treatmenttime, unitid) indexed views
2. Eligibility Determination: Calls eligiblematches!() to apply crossover window logic
3. Exposure Handling: Supports both standard and exposure-based matching

Crossover Window Logic

For each treated observation, potential matches are evaluated using:

• Pre-crossover treatment similarity (via treatcat function)
• Post-crossover treatment exclusion (no treatment allowed in specific windows)
• Time window constraints ensuring proper temporal alignment

eligibility!(matches, observations, X::Matrix{Union{Missing, Float64}}, ids, treatcat, dat_t, dat_id, dat_trt, fmin, fmax, Lmin; exposure = nothing)

Determine matching eligibility with support for missing covariate data.

Purpose

Missing data version of the core eligibility function. Handles datasets where covariates may contain missing values, using appropriate handling in the groupindices creation and subsequent matching algorithms.

Arguments

Same as standard eligibility!() but with:

• X: Covariate matrix allowing missing values (Matrix{Union{Missing, Float64}})

Missing Data Handling

• Creates appropriate SubArray types that can handle missing values
• Propagates missing data handling through the matching pipeline
• Distance calculations will appropriately handle missing values in downstream functions

Implementation Notes

Identical logic to the Float64 version but uses different type dispatch to ensure proper handling of missing values throughout the matching process.

eligibility!(matches, observations, X::Matrix{Union{}}, ids, treatcat, dat_t, dat_id, dat_trt, fmin, fmax, Lmin; exposure = nothing)

Handle matching eligibility when no covariates are specified.

Purpose

Special case handler for models with no covariates (empty covariates array). Converts the empty matrix to an appropriate Float64 matrix structure and delegates to the standard eligibility function.

Arguments

Same as standard eligibility!() but with:

• X: Empty covariate matrix (Matrix{Union{}})

Behavior

• Creates a properly sized Matrix{Float64} with 0 columns but correct row count
• Delegates to the standard Float64 version of eligibility!()
• Enables matching based purely on treatment history patterns without covariate distance

Use Cases

• Treatment effect estimation based only on temporal patterns
• Matching when covariates are unavailable or not needed
• Sensitivity analysis excluding covariate matching

eligibility(model::VeryAbstractCICModel)

Calculate unit eligibility across all treated observations.

Arguments

• model: Fitted TSCSMethods model

Returns

• Matrix{Int}: Eligibility matrix (units × time periods F)

Description

For each potential control unit, calculates how many times it is eligible to serve as a match across all treated observations and time periods. Higher eligibility indicates units that are consistently good potential matches for multiple treated units.

Useful for identifying:

• Units that are consistently good matches across multiple treated observations
• Potential issues with match quality or overlap
• Whether certain units dominate the matching process
• Balance in the matching pool

Examples

eligibility_matrix = eligibility(model)

# Find highly eligible units
total_eligibility = sum(eligibility_matrix, dims=2)
highly_eligible_indices = findall(x -> x > 5, vec(total_eligibility))
highly_eligible_ids = model.ids[highly_eligible_indices]

println("Units eligible for >5 matches: $highly_eligible_ids")

# Check eligibility by time period
for (f_idx, f) in enumerate(model.F)
    eligible_count = sum(eligibility_matrix[:, f_idx])
    println("F=$f: $eligible_count total eligible matches")
end

estimate!(
    ccr::AbstractCICModel, dat;
    iterations = nothing,
    percentiles = [0.025, 0.5, 0.975],
    overallestimate = false,
    bayesfactor = true
)

Perform ATT estimation, with bootstrapped CIs.

estimate!(ccr::AbstractCICModelStratified, dat; iterations = nothing)

Perform ATT stratified estimation, with bootstrapped CIs.

example_data()

Load pre-existing example data for testing and demonstration.

Returns

• DataFrame: Panel data with columns: date, fips, popdens, cumuldeathrate, deathrte, gub, day

Description

Loads example panel data from the package's vignette directory. This is real data from the package examples, providing a consistent dataset for testing and tutorials.

For generating synthetic data with custom parameters, use example_data_generated(), policy_data(), or economic_data() instead.

not stratified

export_results_csv(inspection_results, output_path::String)

Export the results DataFrame to CSV for further analysis in other software.

Arguments

• inspection_results: Output from inspect_results()
• output_path: Path for the CSV file

Example

inspection = inspect_results(imputation, :Y)
export_results_csv(inspection, "analysis_results.csv")

filterunits!(m, omap)

Remove treated units and matches that do not have values defined in the outcome window. This could be modified to account for missingness easily. This should probably be integrated into the match! procedure. (keep separate for now). (Updates ranks as well.)

find_periods_with_matches!(has_matches_by_period, eligible_matches_matrix)

Determine which forward periods have at least one eligible match across all potential match units.

Purpose

For a given treated observation, creates a boolean vector indicating which forward periods (F values) have at least one unit eligible to serve as a match. This is used to determine which periods need BalanceData storage allocation.

Arguments

• has_matches_by_period: Output boolean vector to populate (length = number of F periods)
• eligible_matches_matrix: Boolean matrix where [unit, period] indicates if unit is eligible match for that period

Details

• Performs column-wise any() operation across the eligiblematchesmatrix
• has_matches_by_period[i] = true if any unit can match in forward period i
• has_matches_by_period[i] = false if no units are eligible for forward period i
• Used to optimize storage: only allocate BalanceData for periods with potential matches

Example

# eligible_matches_matrix: 3 units × 4 periods
# [true  false true  true ]  # unit 1
# [false true  false false]  # unit 2  
# [false false false true ]  # unit 3
# 
# Result: [true, true, true, true] - all periods have ≥1 eligible match

check possibility of the match for each f

check possibility of the match for each f

version with exposure

check possibility of the match for each f

version with exposure

check possibility of the match for each f

generate_normal_effects(peak_effect::Float64, F_range::UnitRange{Int64})

Generate F-period effects following a normal distribution centered on middle F period.

Example

effects = generate_normal_effects(1.0, 1:10)  # Peak effect of 1.0 at F=5-6

generate_realistic_tscs(; kwargs...)

Generate synthetic data that mimics the WORKING example_data() treatment pattern but with known, controllable treatment effects for validation.

Key Design

• Event-based treatment: Treatment indicator = 1 only on treatment day
• Sparse treatment: ~2% of observations treated (like real example)
• Staggered timing: Different units treated on different days
• Known ATT: Specified treatment effect we can validate against

Parameters

• true_att::Float64: Known treatment effect to recover
• n_units::Int=100: Number of units (like example data)
• n_days::Int=90: Number of time periods (like example data)
• n_treated::Int=20: Number of treated units (like example data)
• treatment_start_day::Int=40: Earliest treatment day (like example data)
• outcome_scale::Float64=0.5: Base outcome level (like death rates)
• seed::Int=1234: Random seed

Returns

DataFrame with same structure as example_data():

• date, fips, pop_dens, cumul_death_rate, death_rte, gub, day
• Treatment effect built into death_rte outcome

generate_simple_tscs(; kwargs...)

Generate synthetic time-series cross-sectional data with known treatment effect.

Mathematical Model

Yit = αi + λt + τ*Dit + ε_it

Where:

• Y_it: outcome for unit i at time t
• α_i: unit fixed effects (time-invariant differences between units)
• λ_t: time fixed effects (common time trends)
• D_it: treatment indicator (1 if unit i treated at time t, 0 otherwise)
• τ: treatment effect (THIS IS WHAT WE WANT TO RECOVER)
• ε_it: idiosyncratic error term

Key Properties

• Parallel trends satisfied by construction (common λ_t)
• No confounding (treatment assignment independent of potential outcomes)
• Known counterfactual (Yit without treatment = αi + λt + εit)

Parameters

• true_att::Float64: The true average treatment effect (τ)
• n_units::Int=100: Number of cross-sectional units
• n_periods::Int=50: Number of time periods
• treatment_period::Int=25: When treatment begins (must be > max(abs(L)))
• n_treated::Int=20: Number of units that receive treatment
• unit_fe_var::Float64=2.0: Variance of unit fixed effects (α_i)
• time_fe_var::Float64=1.0: Variance of time fixed effects (λ_t)
• error_var::Float64=1.0: Variance of idiosyncratic errors (ε_it)
• seed::Int=1234: Random seed for reproducibility

Returns

DataFrame with columns:

• unit_id: Unit identifier (1, 2, ..., n_units)
• time_period: Time identifier (1, 2, ..., n_periods)
• treatment: Treatment indicator (0/1)
• outcome: Generated outcome variable Y_it
• unit_fe: Unit fixed effect α_i (for validation)
• time_fe: Time fixed effect λ_t (for validation)
• true_control_outcome: Counterfactual outcome without treatment

Example

# Generate data where true ATT = 2.5
data = generate_simple_tscs(true_att=2.5, n_units=50, n_periods=40, seed=42)

# The package should recover ATT ≈ 2.5
model = makemodel(data, :time_period, :unit_id, :treatment, :outcome, [], Dict(), 1:5, -10:-1)
match!(model, data)
estimate!(model, data)
@test abs(model.overall.ATT - 2.5) < 0.2  # Should be close to true effect

generate_tscs_with_covariates(; kwargs...)

Generate synthetic TSCS data with covariates that affect both treatment and outcome.

Mathematical Model

Yit = αi + λt + β*Xit + τ*Dit + εit Dit = treatment assignment potentially related to Xit

Where X_it are time-varying covariates that may confound the relationship.

Additional Parameters

• covariate_effects::Vector{Float64}: Effects of each covariate on outcome (β)
• covariate_names::Vector{Symbol}: Names for the covariates
• confounding_strength::Float64=0.0: How much covariates affect treatment assignment

get_balance_data(size::Int, fill_missing::Bool = true) -> (BalanceData, Bool)

Efficiently allocate BalanceData objects using memory pooling for performance optimization.

Purpose

Creates BalanceData objects while reusing pre-allocated memory buffers when possible. This reduces allocation overhead during intensive balance calculations across many treated observations and matching periods.

Arguments

• size: Number of elements needed in the BalanceData object
• fill_missing: Whether to initialize all values as missing (default: true)

Returns

Tuple of:

• BalanceData: The allocated/reused balance data object
• Bool: Whether the object came from the pool (true) or was newly allocated (false)

Details

• For small objects (≤100 elements): attempts to reuse pooled memory
• For large objects or when pool is empty: allocates new memory directly
• Pooled objects should be returned via return_balance_data() when no longer needed
• Thread-safe through channel-based pooling system

Example

balance_data, is_pooled = get_balance_data(50, true)
# ... use balance_data ...
return_balance_data(balance_data, is_pooled)  # Return to pool when done

get_thread_storage(n_times::Int, n_covariates::Int) -> ThreadLocalDistanceStorage

Get or create thread-local storage for distance calculations.

Algorithm

1. Task Identification: Get current task (modern alternative to deprecated threadid())
2. Storage Check: Check if storage exists and is large enough for this task
3. Resize/Create: If needed, create new storage with sufficient capacity
4. Return: Provide storage for immediate use

Parameters

• n_times: Number of time periods needed
• n_covariates: Number of covariates needed

Returns

• ThreadLocalDistanceStorage: Pre-allocated arrays ready for use

Task Safety

Uses modern task-local storage instead of deprecated Threads.threadid(). Each task maintains independent storage, preventing race conditions while maximizing memory reuse and providing better composability.

getoutcomemap(outcome, data, t, id)

Create a fast lookup dictionary mapping (time, unit) tuples to outcome values.

Purpose

This function creates an optimized lookup structure for outcome data that will be accessed repeatedly during estimation. Instead of searching through the DataFrame for each outcome lookup, we pre-build a hash table for O(1) access.

Algorithm

1. Dictionary Creation: Build Dict{Tuple{Int, Int}, Float64} where keys are (time, unit_id)
2. Data Population: Iterate through all rows, storing non-missing outcomes
3. Missing Data Handling: Skip any rows where the outcome is missing

Performance Benefits

• Fast Lookups: O(1) average case vs O(n) DataFrame search
• Memory Efficient: Only stores non-missing outcomes
• Estimation Optimization: Critical for bootstrap performance where outcomes are accessed thousands of times

Arguments

• outcome: Symbol indicating the outcome column name
• data: Input DataFrame containing the panel data
• t: Symbol for the time variable column
• id: Symbol for the unit identifier column

Returns

Dict{Tuple{Int, Int}, Float64}: Dictionary mapping (time, unit_id) → outcome value

Usage in Estimation

During ATT estimation, we need outcome values for:

• Treated units at outcome periods: outcomemap[(outcome_time, treated_unit)]
• Matched units at outcome periods: outcomemap[(outcome_time, matched_unit)]
• All units at reference periods: outcomemap[(reference_time, unit)]

Example

# If data has columns :time, :unit_id, :outcome
outcome_lookup = getoutcomemap(:outcome, data, :time, :unit_id)
# Later: quickly get outcome for unit 5 at time 10
value = outcome_lookup[(10, 5)]

not stratified

getyes!(yesrows, tvec, idvec, tt, fmax, Lmin, unit)

Determine which data rows fall within the matching time window for a specific unit and treatment time.

Purpose

Populates a boolean mask indicating which rows of the original data should be included in the indexed view for a given (treatment_time, unit) combination. This implements the core logic for defining matching windows in TSCS designs.

Arguments

• yesrows: Output boolean vector to populate (same length as tvec)
• tvec: Time vector from original data
• idvec: Unit ID vector from original data
• tt: Treatment time (when treatment occurred)
• fmax: Maximum forward period (upper bound of F range)
• Lmin: Minimum lag period (lower bound of L range, typically negative)
• unit: Unit ID to extract data for

Time Window Logic

Selects rows where:

1. tvec[k] >= tt + Lmin (at or after start of lag period)
2. tvec[k] < tt + fmax (before end of forward period)
3. idvec[k] == unit (belongs to the specified unit)

Example Time Windows

Treatment at tt=100, Lmin=-30, fmax=40:
- Window: [70, 139] (from 100-30 to 100+40-1)
- Includes: lag periods for covariate measurement + forward periods for outcomes
- Excludes: data outside this window or from other units

Performance Notes

• Operates in-place on yesrows for memory efficiency
• Used within threaded loops, so must be thread-safe (which it is)

grandbalance!(model::AbstractCICModelStratified)

Calculate the overall mean covariate balance for a model.

grandbalance!(model::AbstractCICModel)

Calculate the overall mean covariate balance for a model.

impute_results(m, matches, dat, tvar, unit_var; stratum = 1)

Generate counterfactual outcomes Y(0) using matched control units.

Arguments

• m: Fitted TSCSMethods model
• matches: DataFrame of matched units from matching procedure
• dat: Original panel data
• tvar: Time variable name in data
• unit_var: Unit ID variable name in data
• stratum: Stratum number for stratified models (default: 1)

Returns

• ImputationResults: Structured container with counterfactual analysis

The results DataFrame includes:

• counterfactual_trajectory: What treated outcomes would have been without treatment
• counterfactual_values: Average of matched control outcomes at each time
• counterfactual_lower/counterfactual_upper: Confidence intervals for counterfactuals
• pct_change/pct_change_lower/pct_change_upper: Treatment effects as percentage changes
• daily_counterfactual_change: Day-to-day variation in counterfactual outcomes

Creates counterfactual predictions by averaging matched control outcomes at each time period, enabling visualization of treated vs counterfactual trajectories.

initialize_balance_storage!(model)

Prepare the meanbalances DataFrame storage for balance calculations.

Purpose

Allocates and initializes the meanbalances DataFrame which stores balance statistics for each treated observation. The DataFrame structure depends on whether covariates are time-varying or static.

Details

• Creates one row per treated observation
• For time-varying covariates: Vector{Vector{BalanceData}} (periods × time points)
• For static covariates: Vector{BalanceData} (just periods)
• Uses object pooling for memory efficiency

inspect_results(imputation_results::ImputationResults, outcome_var::Symbol; 
               plot_percentage_changes::Bool = false)

Comprehensive inspection of TSCSMethods results with clean separation of data processing and plotting.

Arguments

• imputation_results: ImputationResults struct from impute_results()
• outcome_var: Symbol representing the outcome variable column name
• plot_percentage_changes: Whether to include percentage change plots (default: false)

Returns

• NamedTuple with fields:
  • treatment_effects_plot: Figure showing ATT over time with confidence intervals
  • counterfactual_plot: Figure showing observed vs counterfactual trajectories
  • summary: NamedTuple with summary statistics and balance diagnostics
  • data: Prepared data structures for custom plotting or further analysis
  • percentage_plot: (optional) Figure showing treatment effects as percentage changes

Example

using TSCSMethods

# Complete workflow
data = example_data()
model = makemodel(data, :t, :id, :gub, :Y, [:X1, :X2], 
                  Dict(:X1 => false, :X2 => false), -15:-10, 1:5)
match!(model, wids=50)
balance!(model)
estimate!(model, dobayesfactor=false)

# Generate imputation results  
imputation = impute_results(model, model.matches, data, :t, :id)

# Comprehensive inspection
inspection = inspect_results(imputation, :Y, plot_percentage_changes=true)

# View plots
inspection.treatment_effects_plot    # Shows ATT over time
inspection.counterfactual_plot       # Shows causal story: observed vs counterfactual  
inspection.percentage_plot           # Shows effects as % changes

# Access summary statistics
summary = inspection.summary
println("Average treatment effect: ", round(summary.overall_att, digits=3))
println("Baseline balance: ", round(summary.baseline_difference, digits=3))

# Access underlying data for custom analysis  
results_df = inspection.results_df
custom_plot = plot_treatment_effects(results_df, title="My Custom ATT Plot")

Notes

• This function coordinates the entire inspection workflow
• Data processing and plotting are separated for flexibility
• All plots contain proper labels and confidence intervals
• Use plot_percentage_changes=true for additional percentage change visualization

make_groupindices(tvec, treatvec, idvec, uid, fmax, Lmin, X; exvec = nothing)

Create efficient lookup dictionaries for accessing unit-specific data during matching windows.

Purpose

Builds indexed views of the data organized by (treatmenttime, unitid) pairs, covering the relevant time windows for matching. This enables fast access to covariate values, treatment histories, and time vectors for each unit during balance calculations.

Arguments

• tvec: Time vector from the data
• treatvec: Treatment indicator vector (0/1 or boolean)
• idvec: Unit identifier vector
• uid: Vector of unique unit IDs to process
• fmax: Maximum forward period (upper bound of F range)
• Lmin: Minimum lag period (lower bound of L range)
• X: Covariate data matrix
• exvec: Optional exposure vector for exposure-based matching

Returns

Tuple of dictionaries with keys (treatment_time, unit_id):

• tidx: Covariate values for unit during time window [tt + Lmin, tt + fmax)
• ridx: Time values for unit during the window
• tridx: Treatment indicators for unit during the window
• exidx: Exposure values (only if exvec provided)

Details

For each treated observation at time tt, creates indexed views for all units covering the time window from tt + Lmin to tt + fmax - 1. This spans both the lag periods (for baseline covariate measurement) and forward periods (for outcome measurement).

Performance Notes

• Uses threaded processing for parallel index construction
• Returns memory-efficient SubArray views rather than copying data
• Pre-allocates dictionaries with appropriate size hints

Example

tidx, ridx, tridx = make_groupindices(data.time, data.treatment, data.unit_id, 
                                      unique_units, 40, -30, covariate_matrix)
# Access unit 1's covariates during treatment time 100's window:
unit1_covariates = tidx[(100, 1)]  # SubArray view of relevant data

make_timeunit_lookup(dat, variable, time_col, unit_var)

Create lookup dictionary mapping (time, unit) tuples to variable values. Returns Dict{Tuple, Union{Float64, Missing}} for fast value retrieval.

Arguments

• dat: DataFrame with panel data
• variable: Column name for the variable to look up
• time_col: Column name for time variable
• unit_col: Column name for unit ID variable

    makefblocks(subTus, subMus, subWos, subWrs, subFs)

Populate a set of fblocks to estimate and bootstrap the ATT.

makemodel(
    dat::DataFrame, 
    t::Symbol, 
    id::Symbol, 
    treatment::Symbol, 
    outcome::Union{Symbol, Vector{Symbol}},
    covariates::Vector{Symbol}, 
    timevary::Dict{Symbol, Bool},
    F::UnitRange{Int}, 
    L::UnitRange{Int};
    title::String = "model",
    estimator::String = "ATT"
) -> CIC

Construct a CIC (Changes-in-Changes) model for causal inference analysis.

Arguments

• dat::DataFrame: Input data containing all required variables
• t::Symbol: Column name for time variable
• id::Symbol: Column name for unit identifier
• treatment::Symbol: Column name for treatment variable (0/1 coding expected)
• outcome::Union{Symbol, Vector{Symbol}}: Column name(s) for outcome variable(s)
• covariates::Vector{Symbol}: Vector of column names for matching covariates
• timevary::Dict{Symbol, Bool}: Dictionary mapping each covariate to whether it's time-varying
• F::UnitRange{Int}: Time periods for treatment effect estimation (post-treatment)
• L::UnitRange{Int}: Time periods for pre-treatment matching window
• title::String: Optional title for the model (default: "model")
• estimator::String: Estimator type (default: "ATT" for Average Treatment Effect on Treated)

Returns

• CIC: A constructed CIC model ready for matching, balancing, and estimation

Throws

• ArgumentError: If input data is empty, required columns are missing, or parameters are invalid

Examples

using TSCSMethods, DataFrames

# Create sample data
data = DataFrame(
    time = repeat(1:100, 50),
    unit_id = repeat(1:50, inner=100),
    treated = repeat([fill(1, 10); fill(0, 40)], inner=100),
    outcome = randn(5000),
    covar1 = randn(5000),
    covar2 = randn(5000)
)

# Construct model
model = makemodel(
    data, :time, :unit_id, :treated, :outcome,
    [:covar1, :covar2],
    Dict(:covar1 => false, :covar2 => true),
    51:60, 40:49  # F and L periods
)

References

Based on the methodology from Imai, Kim, and Wang (2021) for matching methods with time-series cross-sectional data.

match!(model::AbstractCICModel, dat::DataFrame; treatcat = default_treatmentcategories, exposure = nothing, variancesonly = true)

Perform matching for treatment events using Mahalanobis distance with crossover window restrictions.

Purpose

Core matching algorithm that identifies eligible control units for each treated observation, calculates distances using Mahalanobis metric, and ranks potential matches. Implements the extended TSCS matching methodology with crossover window constraints to ensure proper counterfactual identification.

Arguments

• model: CIC model containing treated observations and match storage structures
• dat: DataFrame with time-series cross-sectional data
• treatcat: Function categorizing treatment histories (default: default_treatmentcategories)
• exposure: Optional exposure variable name for exposure-based matching
• variancesonly: Use diagonal covariance matrix (variances only) vs full covariance (default: true)

Returns

Modified model with populated matches containing:

• Eligible match indicators for each treated observation
• Mahalanobis distances to potential matches
• Ranked match preferences for each forward period

Matching Process

1. Eligibility Determination: Uses crossover window logic to determine which units can serve as matches
2. Distance Calculation: Computes Mahalanobis distances using sample covariance matrix
3. Ranking: Orders potential matches by distance for each forward period (F)

Crossover Window Logic

For treatment at time t with forward period f:

• Pre-crossover period: [t + f - max(F), t + f - min(F)]
• Post-crossover period: [t + f - min(F), t + f - max(F)]
• Matches must have similar treatment patterns in pre-crossover period
• Matches cannot be treated in post-crossover period

Performance Notes

• Uses threaded processing for parallel distance calculations
• Memory-efficient SubArray views via groupindices system
• Handles missing data through multiple dispatch

Example

model = makemodel(data, :time, :id, :treatment, :outcome, covariates, timevary, F, L)
match!(model, data)  # Standard matching
match!(model, data; exposure = :policy_exposure)  # With exposure matching

matchinfo(refined_model, original_model; maxrank = 5)

Generate match information DataFrame for refined/calipered models.

Purpose

Creates match information for models that have undergone refinement or caliper operations, showing how refined matches relate to the original full model's ranking system. Useful for understanding how refinement affects match selection and quality.

Arguments

• refined_model: Refined or calipered CIC model (result of refine() or caliper())
• original_model: Original full CIC model before refinement
• maxrank: Maximum number of matches to include per F period (default: 5)

Returns

DataFrame with columns:

• timetreated: Time when treatment occurred
• treatedunit: ID of the treated unit
• f: Forward period within the outcome window
• matchunits: Vector of matched control unit IDs from refined model
• ranks: Vector of rank positions from original model's full ranking

Description

For each treatment observation in the refined model:

1. Finds the corresponding observation in the original model
2. Extracts eligible matches from the refined model
3. Determines their ranks in the original model's full ranking
4. Returns matches sorted by their original ranks

This allows comparison between refined and original matching, showing whether refinement preserves the best matches or introduces changes.

Examples

# Compare refined vs original matches
original_matches = matchinfo(original_model; maxrank = 10)
refined_matches = matchinfo(refined_model, original_model; maxrank = 10)

# Check if refinement preserved top matches
for (orig_row, ref_row) in zip(eachrow(original_matches), eachrow(refined_matches))
    if orig_row.treatedunit == ref_row.treatedunit && orig_row.f == ref_row.f
        original_top3 = orig_row.matchunits[1:min(3, length(orig_row.matchunits))]
        refined_units = ref_row.matchunits
        preserved = intersect(original_top3, refined_units)
        println("Unit $(orig_row.treatedunit), F=$(orig_row.f): $(length(preserved))/3 top matches preserved")
    end
end

matchinfo(model::Union{CIC, CICStratified}; maxrank = 5)

Generate match information DataFrame for standard CIC models.

Purpose

Creates a detailed summary of matching results showing which control units were matched to each treated observation across all forward periods (F). Provides match unit IDs and their ranking within each time period.

Arguments

• model: Fitted CIC or CICStratified model with completed matching
• maxrank: Maximum number of top-ranked matches to include per F period (default: 5)

Returns

DataFrame with columns:

• timetreated: Time when treatment occurred
• treatedunit: ID of the treated unit
• f: Forward period within the outcome window
• matchunits: Vector of matched control unit IDs (ranked best to worst)
• ranks: Vector of rank positions (1 = best match, 2 = second best, etc.)

Description

For each treatment observation and each forward period F, this function:

1. Extracts the top maxrank matched control units
2. Converts internal match indices to actual unit IDs
3. Filters out periods with no available matches
4. Returns results sorted by treatment time and unit ID

This is useful for:

• Inspecting match quality and consistency across time periods
• Understanding which units serve as controls for specific treatments
• Debugging matching algorithm performance
• Creating match summaries for reporting

Examples

# Get match information with top 3 matches per period
matches_df = matchinfo(model; maxrank = 3)

# View matches for a specific treatment
unit_matches = filter(row -> row.treatedunit == 1001 && row.timetreated == 15, matches_df)
for row in eachrow(unit_matches)
    println("F=$(row.f): matched to units $(row.matchunits) with ranks $(row.ranks)")
end

# Check if any treatment lacks matches
no_matches = filter(row -> isempty(row.matchunits), matches_df)
if !isempty(no_matches)
    println("Warning: Some treatments have no matches")
end

meanbalance!(model, dat, tg, rg)

Calculate the mean balances, for each treated observation from the full set of balances. This will limit the calculations to include those present in model.matches, e.g. in case a caliper has been applied.

meanbalance!(model)

Calculate the mean balances, for each treated observation from the full set of balances. This will limit the calculations to include those present in model.matches, e.g. in case a caliper has been applied.

name_model(model::VeryAbstractCICModel)

Generate the filename for a set of models.

treated observation weights

obsinfo(
    match_info_df, data, variables;
    full_model_observations = nothing,
    time_column = :t, id_column = :id
)

Extract covariate information for treated units at their treatment times.

Purpose

Gathers covariate values for all treated units at the exact time of their treatment events. For time-varying covariates, this ensures values are captured at the moment of treatment rather than at arbitrary time points. Useful for creating treatment group summaries and balance tables.

Arguments

• match_info_df: DataFrame from matchinfo() containing treatment observations
• data: Original time-series cross-sectional dataset
• variables: Vector of column names to extract (e.g., [:population, :gdp, :policy_index])
• full_model_observations: Optional vector of all treatment observations for comparison (default: nothing)
• time_column: Name of time variable in data (default: :t)
• id_column: Name of unit identifier in data (default: :id)

Returns

DataFrame with columns:

• timetreated: Time when treatment occurred
• treatedunit: ID of the treated unit
• removed: Boolean indicating if observation was removed in refinement (if full_model_observations provided)
• Additional columns for each variable in variables

Description

For each unique treatment observation:

1. Identifies the corresponding row in the original data
2. Extracts covariate values at the time of treatment
3. Optionally compares against full model to identify removed observations
4. Returns sorted results with all requested covariate information

This is essential for:

• Creating balance tables showing pre-treatment characteristics
• Comparing refined vs full model treatment groups
• Understanding treatment group composition
• Generating summary statistics for reporting

Examples

# Get basic covariate info for treated units
match_df = matchinfo(model)
covariates = [:population, :gdp_per_capita, :unemployment_rate]
treated_info = obsinfo(match_df, data, covariates; 
                      time_column = :year, id_column = :state_id)

# Compare full vs refined model
full_obs = model.observations
refined_df = matchinfo(refined_model)
comparison = obsinfo(refined_df, data, covariates;
                    full_model_observations = full_obs,
                    time_column = :year, id_column = :state_id)

# View removed observations
removed_units = filter(row -> row.removed, comparison)
println("Removed $(nrow(removed_units)) treatment observations in refinement")

# Create balance table
using Statistics
for var in covariates
    mean_val = mean(skipmissing(treated_info[!, var]))
    println("$(var): mean = $(round(mean_val, digits=3))")
end

Notes

• Uses exact time-unit matching to ensure temporal accuracy
• Handles missing values gracefully through type system
• Maintains sort order for consistent output
• Column names are parameterized to avoid hardcoded dependencies

plot_counterfactual_comparison(results_df, outcome_var; kwargs...)

Plot observed vs counterfactual trajectories. Takes DataFrame from preparecounterfactualcomparison_data().

Arguments

• results_df: DataFrame with columns f, [outcome_var], counterfactual_trajectory, etc.
• outcome_var: Symbol for outcome variable column name
• title: Plot title (default: "Observed vs Counterfactual Trajectories")
• xlabel: X-axis label (default: "Time Period")
• ylabel: Y-axis label (default: string of outcome_var)

plot_percentage_changes(results_df; kwargs...)

Plot treatment effects as percentage changes. Takes DataFrame from preparepercentagechange_data().

Arguments

• results_df: DataFrame with columns f, pct_change, pct_change_lower, pct_change_upper
• title: Plot title (default: "Treatment Effects (% Change)")
• xlabel: X-axis label (default: "Time Period")
• ylabel: Y-axis label (default: "Percent Change")
• markersize: Size of scatter points (default: 8)

prepare_counterfactual_comparison_data(imputation_results::ImputationResults, outcome_var::Symbol)

Extract and sort data for observed vs counterfactual trajectory plotting.

This function prepares the DataFrame needed to visualize the key causal inference comparison: what actually happened (observed) vs. what would have happened without treatment (counterfactual).

Arguments

• imputation_results::ImputationResults: Output from impute_results()
• outcome_var::Symbol: Column name of the outcome variable (e.g., :Y, :deaths, :GDP)

Returns

• DataFrame: Sorted by time period, containing all original columns plus outcome variable. Key columns for plotting:
  • f: Time periods
  • [outcome_var]: Actual observed outcomes for treated units
  • counterfactual_trajectory: Estimated counterfactual outcomes
  • counterfactual_lower: Lower bounds of counterfactual confidence intervals
  • counterfactual_upper: Upper bounds of counterfactual confidence intervals
  • att: Treatment effects (observed - counterfactual)

Example

imputation = impute_results(model, matches, data, :t, :id)  
results_df = prepare_counterfactual_comparison_data(imputation, :Y)

# The key causal story (now just DataFrame columns)
observed_trajectory = results_df[!, :Y]                    # What actually happened
counterfactual_trajectory = results_df.counterfactual_trajectory  # What would have happened
treatment_effect = results_df.att                          # The causal impact

Notes

The counterfactual trajectory represents the estimated outcome path that treated units would have followed if they had not received treatment, constructed by averaging outcomes from matched control units.

prepare_percentage_change_data(imputation_results::ImputationResults)

Extract and sort percentage change data for plotting.

Arguments

• imputation_results::ImputationResults: Output from impute_results()

Returns

• DataFrame: Sorted by time period, containing columns:
  • f: Time periods
  • pct_change: Treatment effects as percentage changes
  • pct_change_lower: Lower bounds of percentage change confidence intervals
  • pct_change_upper: Upper bounds of percentage change confidence intervals

Example

imputation = impute_results(model, matches, data, :t, :id)
results_df = prepare_percentage_change_data(imputation)

# Access percentage data (now just DataFrame columns)
pct_effects = results_df.pct_change        # [2.1, 4.2, 1.5, 0.5, ...]
pct_lower = results_df.pct_change_lower    # [0.1, 2.0, -0.5, -1.5, ...]

prepare_treatment_effect_data(imputation_results::ImputationResults)

Extract and sort treatment effect data for plotting.

This function processes ImputationResults to create a clean DataFrame containing treatment effects (ATT) and confidence intervals organized by time period.

Arguments

• imputation_results::ImputationResults: Output from impute_results()

Returns

• DataFrame: Sorted by time period, containing columns:
  • f: Time periods
  • att: Average treatment effects for each time period
  • 2.5%: Lower bounds of 95% confidence intervals
  • 97.5%: Upper bounds of 95% confidence intervals

Example

imputation = impute_results(model, matches, data, :t, :id)
results_df = prepare_treatment_effect_data(imputation)

# Access the data (now just DataFrame columns)
plot_times = results_df.f           # [-4, -3, -2, -1, 1, 2, 3, 4, 5]
effects = results_df.att            # [0.1, 0.2, 0.15, 0.05, 0.3, 0.4, ...]
lower_ci = results_df[!, Symbol("2.5%")]  # [-0.1, 0.0, -0.05, -0.15, 0.1, ...]

processunits(matches, observations, outcome, F, ids, reference, t, id, data)

Pre-process matched units and outcomes for estimation and bootstrapping.

Algorithm Overview

This function transforms the match structure into arrays optimized for ATT estimation:

1. Data Size Calculation: For each treated unit, calculate total data points needed:

   • sum(eligible_matches): Total number of matched control units across all F periods
   • valid_outcome_periods_count: Number of F periods with at least one match (treated unit appears once per valid F period)
   • Total = matches + treated unit observations

2. Memory Pre-allocation: Create vectors to store:

   • Wos: Weighted outcomes (positive for treated, negative for matches)
   • Wrs: Weighted reference period outcomes
   • Tus: Treated unit IDs
   • Mus: Match unit IDs (includes treated unit ID for treated observations)
   • Fs: F period indicators

3. Data Population: Use unitstore! to populate arrays with outcome data and weights

Mathematical Foundation

For ATT estimation, we need paired (treated, control) observations for each F period. The weighting scheme is:

• Treated unit: weight = +1.0 for outcome, -1.0 for reference
• Matched units: weight = -1/nmatches for outcome, +1/nmatches for reference

Arguments

• matches: Vector of match objects with eligible_matches matrices
• observations: Vector of (time, unit) tuples for treated units
• outcome: Outcome variable symbol
• F: Post-treatment periods for effect estimation
• ids: Unit identifier mapping
• reference: Reference period offset (typically negative)
• t, id: Time and unit variable symbols
• data: Input DataFrame

Returns

Tuple of vectors: (Tus, Mus, Wos, Wrs, Fs) where each element contains data for all treated units, structured for efficient estimation.

Performance Notes

• Uses threaded allocation and population for scalability
• Pre-calculates data sizes to avoid dynamic resizing
• Structured for efficient bootstrapping operations

quickatt(matchseries; outcomes = [:deathrte, :case_rte, :deaths, :cases], F = 10:40, ttt = false, tm1 = 30 )

Quickly calculate the att for specified outcomes based on the output of the matchprocess(). Does not calculate confidence intervals.

Calculates both the individual-level unit effects, and the average effects.

quick_inspection(imputation_results::ImputationResults, outcome_var::Symbol)

Quick inspection function that prints summary statistics and optionally shows plots.

Arguments

• imputation_results: Output from impute_results()
• outcome_var: Symbol for outcome variable

Returns

• Summary statistics printed to console
• Basic diagnostic information about matching quality

Example

imputation = impute_results(model, matches, data, :t, :id)
quick_inspection(imputation, :Y)

refine(
    model::CIC, 
    dat::DataFrame;
    refinementnum::Int = 5, 
    dobalance::Bool = true,
    doestimate::Bool = true
) -> RefinedCIC

Refine a CIC model by keeping only the best refinementnum matches for each treated unit.

Arguments

• model::CIC: A matched CIC model
• dat::DataFrame: Input data containing all model variables
• refinementnum::Int: Number of best matches to retain per treated unit (default: 5)
• dobalance::Bool: Whether to perform balancing on the refined model (default: true)
• doestimate::Bool: Whether to perform estimation on the refined model (default: true)

Returns

• RefinedCIC: A refined version of the input model with reduced matches

Description

Refinement improves match quality by keeping only the closest matches for each treated unit, based on the distance metrics computed during the matching stage. This typically improves covariate balance and estimation precision at the cost of statistical power.

Examples

# After matching a model
model = makemodel(data, :time, :id, :treatment, :outcome, covariates, timevary, F, L)
match!(model, data)

# Refine to top 3 matches per treated unit
refined_model = refine(model, data; refinementnum = 3)

# Refine without automatic balancing and estimation
refined_model = refine(model, data; refinementnum = 5, dobalance = false, doestimate = false)

See Also

• caliper: Alternative approach using distance thresholds
• match!: Initial matching procedure
• balance!: Balance assessment

return_balance_data(balance_data::BalanceData, is_pooled::Bool)

Return a BalanceData object to the memory pool for reuse, if it came from the pool.

Purpose

Completes the memory pooling cycle by returning used BalanceData objects back to the pool for future reuse. This prevents memory fragmentation and reduces allocation overhead in subsequent balance calculations.

Arguments

• balance_data: The BalanceData object to potentially return to pool
• is_pooled: Whether this object originally came from the pool (from get_balance_data)

Details

• Only returns objects to pool if they originally came from it (is_pooled == true)
• Only pools objects with ≤100 elements (larger ones are discarded)
• Thread-safe operation using channel-based pools
• Should be called when BalanceData objects are no longer needed

Usage Pattern

balance_data, is_pooled = get_balance_data(50, true)
# ... use balance_data for calculations ...
return_balance_data(balance_data, is_pooled)  # Return when done

save_inspection_plots(inspection_results, output_dir::String = ".", prefix::String = "analysis")

Save all plots from inspect_results() to files with publication-quality settings.

Arguments

• inspection_results: Output from inspect_results()
• output_dir: Directory to save plots (default: current directory)
• prefix: File prefix for saved plots (default: "analysis")

Example

# Run inspection
inspection = inspect_results(imputation, :Y, plot_percentage_changes=true)

# Save all plots
save_inspection_plots(inspection, "results/", "covid_analysis")
# Creates: covid_analysis_treatment_effects.png, covid_analysis_counterfactual.png, etc.

showmatches(model::VeryAbstractCICModel, treatment_observation::Tuple{Int, Int})

Show ranked matches for a specific treatment observation.

Arguments

• model: Fitted TSCSMethods model
• treatment_observation: Tuple of (treatmenttime, unitid)

Returns

• Vector{Vector{Int}}: Ranked matches for each time period F (sorted best to worst)
• "Treatment observation not found": If observation doesn't exist in model

Description

For a specific treatment event, shows the ranked list of matched control unit indices for each time period in the post-treatment window F. Units are ranked from best match (rank 1) to worst match based on the matching algorithm's distance calculations and balancing constraints.

The returned structure is a vector where each element corresponds to a time period F, containing the ranked list of control unit indices (internal to model.matches). To get actual unit IDs, use model.ids[index].

This function is useful for:

• Debugging matching quality for specific treatment events
• Understanding which units serve as matches across different time periods
• Assessing match consistency over time
• Manual inspection of matching results

Examples

# Show matches for unit 1001 treated at time 15
matches = showmatches(model, (15, 1001))

if matches isa String
    println(matches)  # "Treatment observation not found"
else
    println("Matches across F periods:")
    for (f_idx, f) in enumerate(model.F)
        period_matches = matches[f_idx]
        if !isempty(period_matches)
            # Convert to actual unit IDs
            match_ids = [model.ids[idx] for idx in period_matches[1:min(3, length(period_matches))]]
            println("F=$f: top 3 matches = $match_ids")
        else
            println("F=$f: no matches available")
        end
    end
end

# Check consistency of top match across periods
if matches isa Vector && all(length.(matches) .> 0)
    top_matches = [model.ids[period_matches[1]] for period_matches in matches]
    println("Top match by period: $top_matches")
end

stratify(stratfunc::Function, args...; kwargs...)

Apply a stratification function, its arguments, to apply the stratification, calculate the stratified grandbalances, the treated observations in each group, and return the updated model with plot labels.

    stratifyinputs(X, s, strata)

Stratify the output from processunits().

function treatedinfo(
  model, variables, dat;
)

Gives variable values for treated observations present in the model, for the chosen set of variables. Order is the same as model.observations.

assign treatment numbers in each category

trim_model(model)

Remove treated observations that do not have any valid matches. This copies!

unitcounts(m)

count the total number of treated and the number of matched units for each F.

Use the mu matrix to generate vectors of weighted outcomes and unit information, as a preprocessing step for estimation.

validate_dgp(data::DataFrame, true_att::Float64)

Validate that the data generating process is working correctly. This is a helper function for debugging DGPs.

variable_filter(
  model, variable, dat;
  mn = nothing, mx = nothing
)

Remove treated observations according to some variable values. The filtered out treated observations will still left as treated in the data, and will not be allowed to be control units for matching.

Generic function to stratify on a covariate present in the dataframe.

Assumes that matching, balancing, and meanbalancing has ocurred.

whentreated(unit_id::Int, model::VeryAbstractCICModel)

Find when a specific unit received treatment.

Arguments

• unit_id: ID of the unit to look up
• model: Fitted TSCSMethods model

Returns

• Vector{Tuple{Int, Int}}: Treatment events for this unit as (treatmenttime, unitid)
• "Unit not found": If unit was never treated

Description

Searches through all treatment observations to find when (if ever) a specific unit received treatment. Returns all treatment events for units that may have multiple treatments over time.

This function is useful for:

• Verifying treatment timing for specific units
• Checking if a unit received treatment multiple times
• Debugging treatment assignment issues
• Understanding the treatment history of particular units

Examples

# Check when unit 1001 was treated
treatment_events = whentreated(1001, model)

if treatment_events isa String
    println("Unit 1001 was never treated")
else
    treatment_times = [event[1] for event in treatment_events]
    println("Unit 1001 treated at times: $treatment_times")
    
    if length(treatment_events) > 1
        println("Note: Unit had multiple treatments!")
    end
end

# Check multiple units
for unit_id in [1001, 1002, 1003]
    events = whentreated(unit_id, model)
    if events isa String
        println("Unit $unit_id: never treated")
    else
        println("Unit $unit_id: treated $(length(events)) times")
    end
end

window_distances!(distances, dtots, accums, valideligible_matchescols, validunits, tt, Σinvdict, treated_covariate_rows, lag_times, tg, fmin, Lmin, Lmax; sliding = false)

Assign distance calculations for temporal matching windows.

Arguments

• distances: Output matrix to store computed distances
• dtots: Pre-allocated arrays for distance calculations per covariate
• accums: Accumulator arrays for counting valid observations
• valideligible_matchescols: Valid match columns for each potential match
• validunits: Unit identifiers for valid matches
• tt: Treatment time indicator
• Σinvdict: Dictionary of inverse covariance matrices by time period
• treated_covariate_rows: Treated unit covariate rows over time
• lag_times: Time periods for the treated unit
• tg: Treatment group data structure
• fmin: Minimum relative time offset for matching window
• Lmin: Minimum absolute time for matching
• Lmax: Maximum absolute time for matching
• sliding: Whether to use sliding windows (default: false, not implemented)

Algorithm

1. Window Definition: Create temporal matching windows based on treatment timing
2. Distance Computation: For each valid match and time window:
   • Extract control unit data for the same time periods
   • Calculate Mahalanobis and covariate-specific distances using alldistances!
   • Average distances over the matching window using average_distances!
3. Storage: Store final averaged distances in the distances matrix

Mathematical Foundation

Implements the temporal windowing approach from Feltham et al. (2023), where matches are formed based on distance averages over specified pre-treatment periods.

Performance Notes

• Complexity: O(nmatches × windowsize × n_covariates)
• Threading: Called within threaded loops, uses thread-local storage
• Memory: Reuses pre-allocated arrays to minimize allocations

Note

Currently only supports fixed windows. Sliding window implementation would allow the matching window to vary by treatment time.