metropolis.py

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

NUM_STATES = 4


# Metropolis-Hastings algorithm for Step 3
def metropolis_hastings_ck(
    S, L, c_k_init, prior_mean, prior_std, num_samples=1000, proposal_std=0.1, k=1
):
    """
    Metropolis-Hastings algorithm to sample c_k for a given methylation state k.
    """
    samples = []
    current_ck = c_k_init

    def prior_ck(ck):
        return np.exp(-0.5 * ((ck - prior_mean) / prior_std) ** 2) / (
            np.sqrt(2 * np.pi) * prior_std
        )

    def transition_prob(St, St_prev, Lt_prev, ck):
        exp_ck = np.exp(ck)
        if St == St_prev:
            return (3 / 4) * np.exp(-Lt_prev / exp_ck) + 1 / 4
        else:
            return (1 / 4) * (1 - np.exp(-Lt_prev / exp_ck))

    def likelihood_ck(ck):
        likelihood = 0
        for t in range(1, len(S)):
            St, St_prev, Lt_prev = S[t], S[t - 1], L[t - 1]
            if St_prev == k:
                likelihood += np.log(transition_prob(St, St_prev, Lt_prev, ck))
        return likelihood

    for _ in range(num_samples):
        proposed_ck = np.random.normal(current_ck, proposal_std)

        # prior_current = np.log(prior_ck(current_ck))
        # prior_proposed = np.log(prior_ck(proposed_ck))

        prior_current = stats.norm.logpdf(current_ck, prior_mean, prior_std)
        prior_proposed = stats.norm.logpdf(proposed_ck, prior_mean, prior_std)

        likelihood_current = likelihood_ck(current_ck)
        likelihood_proposed = likelihood_ck(proposed_ck)

        # if likelihood_current == 0:
        #     acceptance_ratio = 1.0 if likelihood_proposed > 0 else 0.0
        # else:
        acceptance_ratio = (prior_proposed + likelihood_proposed) - (
            prior_current + likelihood_current
        )

        if np.random.uniform(0, 1) < min(1, np.exp(acceptance_ratio)):
            current_ck = proposed_ck
            samples.append(current_ck)

    return np.array(samples)


if __name__ == "__main__":
    # Simulated data input
    # Assuming `simulated_data` has been loaded and contains CpG site data
    # Columns: CpG_Site, State, Control_Sample_1, Case_Sample_1, etc.
    simulated_data = np.loadtxt("simulated_cpg_data.csv", delimiter=",", skiprows=1)

    # Extract relevant columns
    S = simulated_data[:, 1].astype(int)  # Hidden states (State column)
    L = np.diff(simulated_data[:, 0])  # Distances between CpG sites

    # Step 3: Sampling c_k for each state
    prior_mean = 0  # Prior mean for c_k
    prior_std = 1  # Prior standard deviation for c_k
    num_samples = 1000  # Number of samples to generate

    # Sample c_k for all states
    c_k_samples = {}
    for k in range(1, NUM_STATES + 1):
        c_k_init = np.random.normal(prior_mean, prior_std)  # Initialize c_k
        c_k_samples[k] = metropolis_hastings_ck(
            S, L, c_k_init, prior_mean, prior_std, num_samples, k=k
        )

    # Visualization of c_k samples
    for k in range(1, NUM_STATES + 1):
        plt.figure(figsize=(8, 5))
        plt.plot(c_k_samples[k], label=f"Samples for c_{k}")
        plt.axhline(np.mean(c_k_samples[k]), color="red", linestyle="--", label="Mean")
        plt.title(f"Metropolis-Hastings Sampling for c_{k}")
        plt.xlabel("Iteration")
        plt.ylabel(f"c_{k}")
        plt.legend()
        plt.grid()
        plt.show()

    # Visualization of final estimated c_k values
    mean_ck_values = [np.mean(c_k_samples[k]) for k in range(1, NUM_STATES + 1)]
    plt.bar(
        range(1, NUM_STATES + 1),
        mean_ck_values,
        tick_label=[f"c_{k}" for k in range(1, NUM_STATES + 1)],
    )
    plt.title("Final Estimated c_k Values")
    plt.xlabel("State")
    plt.ylabel("Mean c_k")
    plt.show()