<jupyter_start><jupyter_text>Iterative methods for solving GEE<jupyter_code>import numpy as np
import pandas as pd
from scipy import linalg
from typing import List, NamedTuple

class CovarianceParameters(NamedTuple('CovarianceParameters', [
    ('alpha', np.float64),
    ('sigma2', np.float64),
])):
    def make_correlation_matrix(self, size):
        correlation_matrix = np.eye(size)
        for i in range(size - 1):
            for j in range(i + 1, size):
                correlation_matrix[i, j] = correlation_matrix[j, i] = np.power(self.alpha, j - i)
        return correlation_matrix
    
    def make_covariance_matrix(self, size):
        return self.make_correlation_matrix(size)*self.sigma2
    
class Cluster(NamedTuple('Cluster', [
    ('X', np.array),
    ('y', np.array),
    ('covariance', np.array),
])):
    """Cluster covariates, response, and covariance structure."""
def estimate_covariance_parameters(clusters, beta):    
    def _estimate_covariance_parameters(X, y, beta):
        epsilon = y - X.dot(beta)        
        covariance_matrix = np.outer(epsilon, epsilon)
        sigma2 = np.diag(covariance_matrix)
        rho = [covariance_matrix[i, i + 1] for i in range(len(covariance_matrix) - 1)]
        return sigma2, rho
    
    sigma2 = [] 
    rho = []
    for cluster in clusters:
        cluster_sigma2, cluster_rho = _estimate_covariance_parameters(
            cluster.X, cluster.y, beta)
        sigma2.extend(cluster_sigma2)
        rho.extend(cluster_rho)  
    sigma2 = np.mean(sigma2)
    return CovarianceParameters(alpha=np.mean(rho)/sigma2, sigma2=sigma2)

def estimate_beta(clusters: List[Cluster]):
    """Estimate beta under the assumption that clusters."""
    cluster_weights = [
        linalg.cho_solve(linalg.cho_factor(cluster.covariance), np.eye(len(cluster.y))) for cluster in clusters
    ]    
    projected_X = np.sum([
        cluster.X.T.dot(weights).dot(cluster.X)
        for cluster, weights in zip(clusters, cluster_weights)
    ], 0)
    projected_y = np.sum([
        cluster.X.T.dot(weights).dot(cluster.y)
        for cluster, weights in zip(clusters, cluster_weights)
    ], 0)    
    beta = linalg.cho_solve(linalg.cho_factor(projected_X), projected_y)
    
    bread = linalg.cho_solve(linalg.cho_factor(projected_X), np.eye(len(projected_X)))
    residuals = [cluster.y - cluster.X.dot(beta) for cluster in clusters]
    sigma = np.mean([np.outer(residual, residual) for residual in residuals], 0)
    meat = np.sum([
        cluster.X.T.dot(weights).dot(sigma).dot(weights).dot(cluster.X)
        for cluster, residual, weights in zip(clusters, residuals, cluster_weights)
    ], 0)
    return beta, bread.dot(meat).dot(bread)

def update_clusters(clusters, covariance_parameters):
    return [
        Cluster(X=cluster.X, y=cluster.y,
                covariance=covariance_parameters.make_covariance_matrix(len(cluster.y)))
        for cluster in clusters]
orthodont_data = pd.read_csv('orthodont.csv')
orthodont_data = orthodont_data.set_index('Subject')
orthodont_data.head(8)
def make_covariates(data_frame):    
    age = (data_frame['age'] - 8).values
    is_female = (data_frame['Sex'] == 'Female').values.astype(np.float64)
    return np.column_stack((
        np.ones(len(data_frame)),
        age,
        is_female,
        age*is_female,        
    ))

def make_response(data_frame):
    return data_frame['distance'].values
def make_table(beta, beta_covariance):
    std_errors = np.sqrt(np.diag(beta_covariance))    
    return pd.DataFrame(list(zip(beta, std_errors)),
                        columns=['Estimate', 'Standard error'],
                        index=['$\\beta_0$', '$\\beta_1$', '$\\beta_2$', '$\\beta_3$'])<jupyter_output><empty_output><jupyter_text>Initially, we assume no covariance.<jupyter_code>clusters = [
    Cluster(
        X=make_covariates(orthodont_data.loc[i]),
        y=make_response(orthodont_data.loc[i]),
        covariance=np.eye(len(orthodont_data.loc[i])),
    )   
    for i in np.unique(orthodont_data.index)
]

beta, beta_covariance = estimate_beta(clusters)
with open('beta0_estimate.tex', 'w') as f:
    f.write(make_table(beta, beta_covariance).to_latex(escape=False))<jupyter_output><empty_output><jupyter_text>In each iteration, we estimate covariance parameters with $\hat{\beta}$. We compute a new correlation structure for each cluster, and use it to get a better estimate for $\beta$.<jupyter_code>def _update_estimates(clusters, beta):
    covariance_parameters = estimate_covariance_parameters(clusters, beta)
    clusters = update_clusters(clusters, covariance_parameters)
    return covariance_parameters, clusters, estimate_beta(clusters)

previous_beta, (covariance_parameters, clusters, (beta, beta_covariance)) = (
    beta, _update_estimates(clusters, beta))

with open('beta1_estimate.tex', 'w') as f:
    f.write(make_table(beta, beta_covariance).to_latex(escape=False))
covariance_parameters
while np.sum(np.abs(previous_beta - beta)) > 1e-12:
    previous_beta, (covariance_parameters, clusters, (beta, beta_covariance)) = (
        beta, _update_estimates(clusters, beta))
    print(beta, covariance_parameters)
    print(np.sqrt(np.diag(beta_covariance)))<jupyter_output>[22.75026233  0.76945687 -1.55885764 -0.28569215] CovarianceParameters(alpha=0.6135188339895254, sigma2=4.91059980525123)
[0.53558315 0.08739713 0.83909753 0.13692499]
[22.7502655   0.76945666 -1.55886113 -0.28569188] CovarianceParameters(alpha=0.6135308146270501, sigma2=4.910651850771101)
[0.53558381 0.08739724 0.83909856 0.13692515]
[22.75026552  0.76945666 -1.55886115 -0.28569188] CovarianceParameters(alpha=0.61353088133117, sigma2=4.910652141021707)
[0.53558381 0.08739724 0.83909857 0.13692515]
[22.75026552  0.76945666 -1.55886115 -0.28569188] CovarianceParameters(alpha=0.6135308817025603, sigma2=4.9106521426377565)
[0.53558381 0.08739724 0.83909857 0.13692515]
[22.75026552  0.76945666 -1.55886115 -0.28569188] CovarianceParameters(alpha=0.6135308817046287, sigma2=4.910652142646756)
[0.53558381 0.08739724 0.83909857 0.13692515]
[22.75026552  0.76945666 -1.55886115 -0.28569188] CovarianceParameters(alpha=0.6135308817046394, sigma2=4.910652142646804)
[0.53558381 0.08739724 0.83909857 [...]<jupyter_text>The final, converged estimates are below.<jupyter_code>with open('beta_final_estimate.tex', 'w') as f:
    f.write(make_table(beta, beta_covariance).to_latex(escape=False))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1250968842039994, 'fraction_numerical': 0.10354983723453728, 'mean_word_length': 4.116574147501982, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '69646', 'license': 'MIT', 'n_tokens_mistral': 2703, 'n_tokens_neox': 2239, 'n_words': 401, 'path': 'hw4/iterative_gee_dental.ipynb', 'repo_name': 'ppham27/stat571', 'size': '6451'}