git_tutorial_mcmc/fit_line_emcee.py at main · PhysicistSouravDas/git_tutorial_mcmc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
"""
A sample script to fit a simulated data with a straight line.

This script is a part of demonstration for students
to collaborate via github to complete the tasks and push to
github
"""

import numpy as np
import emcee
import matplotlib.pyplot as plt

# 1. Load the data generated by the first script
data = np.loadtxt("mock_data.txt", comments="#")
x, y, yerr = data[:, 0], data[:, 1], data[:, 2]

# CLASS EXERCISE: FILL IN THE MISSING FUNCTIONS BELOW
# PUSH YOUR CHANGES

def log_likelihood(theta, x, y, yerr):
    """
    Calculates the log-likelihood of the data given the parameters.
    """
    a, b = theta

    # TODO: Calculate the model predicted y values for the line y = a*x + b
    # model = ...

    # TODO: Calculate the log-likelihood assuming Gaussian errors.
    # Hint: The formula is -0.5 * sum(((y - model) / yerr)^2)
    # return ...
    pass # Remove this pass when implementing

def log_prior(theta):
    """
    Defines the prior probability for the parameters 'a' and 'b'.
    """
    a, b = theta

    # TODO: Define flat (uniform) priors for 'a' and 'b'.
    # Example: 'a' should be between -10 and 10, 'b' between -10 and 10.
    # If the parameters are inside this range, return 0.0.
    # If they fall outside, return -np.inf
    # return ...
    pass # Remove this pass when implementing

def log_probability(theta, x, y, yerr):
    """
    Combines the prior and likelihood to get the full posterior probability.
    """
    # TODO: Calculate the log_prior.
    # If it is not finite (-np.inf), immediately return -np.inf

    # TODO: If the prior is valid, return the sum of the log_prior and the log_likelihood
    # return ...
    pass # Remove this pass when implementing

# ==============================================================================
# EMCEE SETUP AND EXECUTION (DO NOT MODIFY UNLESS YOU WANT TO EXPERIMENT)
# ==============================================================================

# Set up the MCMC sampler parameters
ndim = 2          # Number of parameters we are fitting (a and b)
nwalkers = 32     # Number of MCMC walkers

# Initialize the starting positions of the walkers in a tiny Gaussian ball
# around a random guess (e.g., a=1.0, b=0.0)
initial_guess = np.array([1.0, 0.0])
pos = initial_guess + 1e-4 * np.random.randn(nwalkers, ndim)

# Initialize the emcee EnsembleSampler
# We pass the log_probability function and the extra arguments (x, y, yerr)
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, args=(x, y, yerr))

# Run the MCMC sampler
print("Running MCMC...")
sampler.run_mcmc(pos, 2000, progress=True)
print("MCMC completed!")

# Discard the first 500 steps as "burn-in" and flatten the chain
flat_samples = sampler.get_chain(discard=500, thin=15, flat=True)

# Print the results
a_fit = np.percentile(flat_samples[:, 0], [16, 50, 84])
b_fit = np.percentile(flat_samples[:, 1], [16, 50, 84])
print(f"Fitted a = {a_fit[1]:.3f} (True: 2.5)")
print(f"Fitted b = {b_fit[1]:.3f} (True: -1.0)")

# (Optional) plot the trace or use the 'corner' package if installed!
# import corner
# fig = corner.corner(flat_samples, labels=["a", "b"], truths=[2.5, -1.0])
# plt.show()