Synthetic Control in Python: Mock Example

import numpy as np

class SyntheticControl:
    def __init__(self, data, treated_index, pre_periods, post_periods):
        self.data = np.array(data, dtype=float)
        self.treated_index = treated_index
        self.pre_periods = pre_periods
        self.post_periods = post_periods
        
    def fit_weights(self):
        pre_data = self.data[:self.pre_periods]
        treated_pre = pre_data[:, self.treated_index]
        donors = np.delete(pre_data, self.treated_index, axis=1)
        
        X = donors.T
        y = treated_pre
        
        # Ridge regression to find optimal donor weights
        lambda_reg = 0.01
        XtX = X @ X.T + lambda_reg * np.eye(X.shape[0])
        XtY = X @ y
        
        weights = np.linalg.solve(XtX, XtY)
        
        # Normalize weights to sum to 1
        weights = weights / np.sum(weights)
        
        return weights
    
    def counterfactual(self):
        weights = self.fit_weights()
        donors = np.delete(self.data[:, self.treated_index:], 0, axis=1) if self.treated_index == 0 else np.delete(self.data, self.treated_index, axis=1)
        
        # Build donor matrix for all periods
        donor_matrix = np.column_stack([self.data[:, j] for j in range(self.data.shape[1]) if j != self.treated_index])
        
        counterfactual = donor_matrix @ weights
        return counterfactual


if __name__ == "__main__":
    # Mock data: 8 periods (5 pre, 3 post), 3 units (0=treated, 1&2=donors)
    mock_data = [
        [10, 12, 11],
        [11, 13, 12],
        [12, 14, 13],
        [13, 15, 14],
        [14, 16, 15],  # end of pre-period
        [15, 17, 16],
        [16, 18, 17],
        [17, 19, 18]
    ]
    
    model = SyntheticControl(mock_data, treated_index=0, pre_periods=5, post_periods=3)
    weights = model.fit_weights()
    counterfactual = model.counterfactual()
    print("Donor weights:", np.round(weights, 4))
    print("Counterfactual values:", np.round(counterfactual, 4))
    print("Observed treated values:", mock_data[5][0], mock_data[6][0], mock_data[7][0])