<jupyter_start><jupyter_text>Toy Datasets<jupyter_code># For Base cohort
age = 45
lab1 = 3.4
lab2 = 5000
pcond1 = 0.3
pcat1 = 0.5

mean = [age, lab1, lab2]
cov = [[14, 3, 100], 
       [3, 2, 50],
       [100, 50, 25000]]


np.random.seed(42)
x = np.random.multivariate_normal(mean, cov, config['nbase'])
x[x< 0] = 0

lab3 = np.random.beta(1, 2, size=config['nbase'])
x = np.concatenate((x, lab3.reshape(config['nbase'], -1)), axis=1)

cond1 = np.random.binomial(1, pcond1, size=config['nbase']).reshape(config['nbase'], -1)
x = np.concatenate((x, cond1), axis=1)

cat1 = np.random.binomial(3, pcat1, size=config['nbase']).reshape(config['nbase'], -1)
x = np.concatenate((x, cat1), axis=1)

data_base = pd.DataFrame(x, columns=['age', 'lab1', 'lab2', 'lab3', 'cond1', 'cat1'])
# For External cohort
age = 50
lab1 = 5.5
lab2 = 4800
pcond1 = 0.5
pcat1 = 0.7
factor = 1.5

mean = [age, lab1, lab2]
cov = [[20, 5, 150], 
       [5, 4, 100],
       [150, 100, 55000]]


x = np.random.multivariate_normal(mean, cov, config['neca'])
x[x< 0] = 0

lab3 = factor*np.random.beta(1, 2, size=config['neca'])
x = np.concatenate((x, lab3.reshape(config['neca'], -1)), axis=1)

cond1 = np.random.binomial(1, pcond1, size=config['neca']).reshape(config['neca'], -1)
x = np.concatenate((x, cond1), axis=1)

cat1 = np.random.binomial(3, pcat1, size=config['neca']).reshape(config['neca'], -1)
x = np.concatenate((x, cat1), axis=1)

data_eca = pd.DataFrame(x, columns=['age', 'lab1', 'lab2', 'lab3', 'cond1', 'cat1'])
# Define Cohen's d: Standardized Mean Difference
def cohensd(d1, d2):
    num = np.mean(d1) - np.mean(d2)
    den = np.sqrt((np.std(d1) ** 2 + np.std(d2) ** 2) / 2)
    cohens_d = num / den
    return cohens_d
# Plot data distribution
def plot_dist(data1, data2, col, ax=None):
    if not ax:
        _, ax = plt.subplots(1,1)
        
    ax.hist(data1[col], density=True, fill='black', histtype='stepfilled', 
             edgecolor='black', bins= 20, linewidth=1.2, label=data1.name)
    ax.hist(data2[col], density=True, bins= 20, fill='blue', histtype='stepfilled',
             edgecolor='blue', linewidth=1.2, label=data2.name, alpha=0.7)
    ax.legend(loc='best')
    ax.set_title(col)
    return 
# View the Starting distribution

_, ax = plt.subplots(3,2, figsize=(15, 10))

data_base.name = 'RCT Cohort'
data_eca.name = 'ECA Cohort'
plot_dist(data_base, data_eca, 'age', ax[0][0])
plot_dist(data_base, data_eca, 'lab1', ax[0][1])
plot_dist(data_base, data_eca, 'lab2', ax[1][0])
plot_dist(data_base, data_eca, 'lab3', ax[1][1])
plot_dist(data_base, data_eca, 'cond1', ax[2][0])
plot_dist(data_base, data_eca, 'cat1', ax[2][1])
# Precompute parameters required for KDE estimation
dist_density, waypoints, barwidth, kde_base = {}, {}, {}, {}

for c in data_base.columns:
    a = np.quantile(pd.concat([data_base[c], data_eca[c]]), 0.05)
    b = np.quantile(pd.concat([data_base[c], data_eca[c]]), 0.95)
    waypoints[c] = torch.tensor(np.linspace(a, b, config['nwaypoints']))
    barwidth[c] = torch.diff(waypoints[c])[0]
            
v, b = [], []
for c in data_base.columns:
    v.append(waypoints[c])
    b.append(barwidth[c])

waypoints = torch.stack(v,1)
barwidths = torch.stack(b,0).unsqueeze(0).repeat(config['nwaypoints'], 1)
import torch
from torch.distributions import Normal
from torch.distributions.distribution import Distribution

class GaussianKDE(Distribution):
    def __init__(self, X, weights=None):
        """
        X : tensor (n, d)
          `n` points with `d` dimensions to which KDE will be fit
        bw : numeric
          bandwidth for Gaussian kernel
        """
        self.X = X
        self.ncols = X.shape[1]
        self.n = X.shape[0]
        self.bw = self.scotts_factor()
        self.mvn = Normal(loc=torch.zeros(1), scale=1)
        
        if weights is not None:
            self.weights = weights/sum(weights)
        else:
            self.weights = (1/len(X))*torch.ones(len(X))
        self.neff = 1/sum(self.weights**2)
    
    def scotts_factor(self):
        """Compute Scott's factor.
        Returns
        -------
        s : float
            Scott's factor.
        """
        factor = np.power(self.n, -1./5)
        sigma = self.X.std(axis=0)
        return 1.06*sigma*factor

    def __call__(self, Y, X=None):
        """Returns the kernel density estimates of each point in `Y`.

        Parameters
        ----------
        Y : tensor (m, d)
          `m` points with `d` dimensions for which the probability density will
          be calculated
        X : tensor (n, d), optional
          `n` points with `d` dimensions to which KDE will be fit. Provided to
          allow batch calculations in `log_prob`. By default, `X` is None and
          all points used to initialize KernelDensityEstimator are included.


        Returns
        -------
        probs : tensor (m)
          probability densities for each of the queried points in `Y`
        """
           
        if X is not None:
            self.X = X
        
        logprob = (self.mvn.log_prob((self.X.unsqueeze(1) - Y) / self.bw))
        
#         probs = ((1/bw) * torch.exp(logprob).sum(dim=0))/self.n
            
        # experiment
        b1 = self.weights.repeat(Y.shape[0],1).T
        b22 = b1.unsqueeze(2).repeat(1, 1, self.ncols)
        w = (b22 * torch.exp(logprob))/sum(self.weights)
        probs = (1/self.bw)* w.sum(dim=0)
        
        return probs*barwidths
class KDEoptim(nn.Module):
    """Custom Pytorch model for gradient optimization.
    """
    def __init__(self, X, Y, wp):
        
        super().__init__()
        base, eca = [], []
        for i in X.columns:
            base.append(torch.tensor(X[i]))
            eca.append(torch.tensor(Y[i]))
            
        X = torch.stack(base, dim=1)
        Y = torch.stack(eca, dim=1)
        
        weights = torch.ones(Y.shape[0], requires_grad=True)/len(Y)  
        self.weights = nn.Parameter(weights)
        self.waypoints = wp
        
        self.Y = Y
        self.kde_target = GaussianKDE(X)(self.waypoints)
        
    def forward(self):
        """
        forward fn
        """
        self.kde2_estimated = GaussianKDE(self.Y, 
                                     weights=self.weights)(self.waypoints)
        return sum(torch.square(self.kde_target - self.kde2_estimated)).sum()<jupyter_output><empty_output><jupyter_text>Optimized weights<jupyter_code>def training_loop(model, optimizer, n=config['iterations']):
    "Training loop for torch model."
    losses = []
    red = None
    for i in range(n):
        preds = model()
        preds.backward()
        optimizer.step()
        optimizer.zero_grad()
        loss = (preds.detach().numpy()).item()
        losses.append(loss) 
        
        with torch.no_grad():
            for param in model.parameters():
                param.clamp_(0, 100)
                
        if i % config['prn'] == 0:
            print(f'Training {i} epochs:', loss)  

    return losses
# instantiate model
m = KDEoptim(data_base, data_eca, waypoints)

# # Instantiate optimizer
opt = torch.optim.Adam(m.parameters(), lr=0.001)
losses = training_loop(m, opt)
optweights = m.weights.detach().numpy()
plt.hist(optweights/sum(optweights), bins=len(optweights)//100)
plt.axvline(x=1/len(optweights),color='black', linestyle='dashed')
plt.yscale('log')
plt.show()
data_eca_study = data_eca.sample(config['nbase'], replace=True, weights=optweights, 
                             random_state=42)
data_eca_study.name = 'Optimized ECA Cohort'

f, ax = plt.subplots(len(data_base.columns),2, figsize=(15, 20))

for i, col in enumerate(data_base.columns):
    plot_dist(data_base, data_eca, col, ax[i][0])
    plot_dist(data_base, data_eca_study, col, ax[i][1])

f.suptitle('Before (left) and after(right) Matching') 
for col in data_base.columns:
    print(f'Distribution of {col}')
    print('Before balancing', ks_2samp(data_base[col], data_eca[col]))
    print('After balancing', ks_2samp(data_base[col], data_eca_study[col]))
    print('\n')
for col in data_base.columns:
    print(f'Distribution of {col}')
    print('Before balancing', cohensd(data_base[col], data_eca[col]))
    print('After balancing', cohensd(data_base[col], data_eca_study[col]))
    print('\n')
end = datetime.now()
print('Total time:', end-start)<jupyter_output>Total time: 0:03:12.895979
{'fraction_non_alphanumeric': 0.15988511249401627, 'fraction_numerical': 0.03135471517472475, 'mean_word_length': 2.822964318389753, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '216438', 'license': 'MIT', 'n_tokens_mistral': 3158, 'n_tokens_neox': 2991, 'n_words': 653, 'path': 'Pytorch-towards-wp.ipynb', 'repo_name': 'csrinesh/balance_cohort', 'size': '8356'}