lecture

Project3/LyX/Results.lyx

@@ -104,6 +104,199 @@ name "sec:Results"
 \end_inset
 
 
+\end_layout
+
+\begin_layout Standard
+\begin_inset Note Note
+status open
+
+\begin_layout Plain Layout
+\begin_inset Float figure
+placement document
+alignment document
+wide false
+sideways false
+status open
+
+\begin_layout Plain Layout
+\begin_inset Graphics
+	filename Figures/InitialCondition.pdf
+	lyxscale 30
+	width 50text%
+
+\end_inset
+
+
+\begin_inset Caption Standard
+
+\begin_layout Plain Layout
+Initial Condition
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\begin_layout Standard
+\begin_inset Float figure
+placement document
+alignment document
+wide false
+sideways false
+status open
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\begin_layout Plain Layout
+\begin_inset Float figure
+placement document
+alignment document
+wide false
+sideways false
+status open
+
+\begin_layout Plain Layout
+\begin_inset Graphics
+	filename Figures/loss_theta.pdf
+	lyxscale 30
+	width 50text%
+
+\end_inset
+
+
+\begin_inset Caption Standard
+
+\begin_layout Plain Layout
+Loss
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\end_inset
+
+
+\begin_inset Newline newline
+\end_inset
+
+
+\begin_inset Float figure
+placement document
+alignment document
+wide false
+sideways false
+status open
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\begin_layout Plain Layout
+\begin_inset Graphics
+	filename Figures/prediction_kx.pdf
+	lyxscale 30
+	width 60text%
+
+\end_inset
+
+
+\begin_inset Caption Standard
+
+\begin_layout Plain Layout
+Prediction for resistance function 
+\begin_inset Formula $k(x)$
+\end_inset
+
+
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\end_inset
+
+
+\begin_inset Newline newline
+\end_inset
+
+
+\begin_inset Float figure
+placement document
+alignment document
+wide false
+sideways false
+status open
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\begin_layout Plain Layout
+\begin_inset Graphics
+	filename Figures/predicted_vs_true.pdf
+	lyxscale 30
+	width 60text%
+
+\end_inset
+
+
+\begin_inset Caption Standard
+
+\begin_layout Plain Layout
+Prediction vs true data
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\end_inset
+
+
+\begin_inset Caption Standard
+
+\begin_layout Plain Layout
+Solution of MCMC
+\end_layout
+
+\end_inset
+
+
+\end_layout
+
+\begin_layout Plain Layout
+
+\end_layout
+
+\end_inset
+
+
 \end_layout
 
 \end_body

Project3/src/MCMC.py

@@ -4,12 +4,13 @@ import matplotlib.pyplot as plt
 from tqdm import tqdm
 
 # MCMC and solver parameters
-MCMC_STEPS = 500
+MCMC_STEPS = 50#!500
 MCMC_SUBSTEPS = 20
+NUMBER_SAMPLES = 2
 
 x0, x1 = 0, 4
 
-NTime = 2500
+NTime = 500
 T_end = 10
 
 # Upwind viscosity parameter (used in flux computation)
@@ -28,8 +29,7 @@ NT_obs = int(mat['NT_obs'].item())
 times_observed = mat['time_obs'][0]
 n_extraction_points = 4
 x_observed = [float(mat[f'X_{i}'].item()) for i in range(1, n_extraction_points+1)]
-print(f'x_observed: {x_observed}')
-u_true = np.array([mat[f'u_true_X{i}'] for i in range(1, n_extraction_points+1)])[:,:,0]
+u_true = np.array([mat[f'u_true_X{i}'] for i in range(1, n_extraction_points+1)])#[:,:,0]
 u_true_flatten = u_true.flatten()
 
 x = np.linspace(x0, x1, Nx)
@@ -58,7 +58,7 @@ def Upwind(u, M=M):
     F[-1] = f(u[-1])
     return F
 
-def solve_PDE(Theta, T_end, NTime, x, dx, CFL=0.5, time_scheme='Euler'):
+def solve_PDE(Theta, T_end, NTime, x, dx, x_Theta, number_samples=1, CFL=0.5, time_scheme='Euler'):
     """
     Solve the PDE 
         u_t + (k(x) f(u))_x = 0,
@@ -77,7 +77,7 @@ def solve_PDE(Theta, T_end, NTime, x, dx, CFL=0.5, time_scheme='Euler'):
     k = lambda x_val: np.interp(x_val, x_Theta, Theta)
     
     # Initial condition
-    u = np.where(
+    u_init_1 = np.where(
             x <= 0.5, 
             0.2,
             np.where(
@@ -94,55 +94,85 @@ def solve_PDE(Theta, T_end, NTime, x, dx, CFL=0.5, time_scheme='Euler'):
                 )
             )
         )
-    t = 0.0
-    u_vec = [u.copy()]
-    t_vec = [t]
     
-    while t < T_end:
-        # Compute dt using the CFL condition: dt <= CFL*dx / max|k(x) * f'(u)|
-        speeds = np.abs(k(x) * df(u))
-        max_speed = np.max(speeds)
-        dt_CFL = CFL * dx / (max_speed + 1e-6)
-        dt = min(dt_CFL, T_end - t)
-        
-        # Compute fluxes at interfaces
-        F = Upwind(u, M)
-        # Spatial locations of cell interfaces
-        x_half = (x[:-1] + x[1:]) / 2.0
-        k_half = k(x_half)
-        
-        if time_scheme == 'Euler':
-            # Euler update
-            u_new = u.copy()
-            # Update interior cells (i = 1,..., len(u)-2)
-            u_new[1:-1] = u[1:-1] - dt/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
-        elif time_scheme == 'Heun':
-            # Predictor: Euler step
-            u_predictor = u.copy()
-            u_predictor[1:-1] = u[1:-1] - dt/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
-            # Compute fluxes for the predictor
-            F_predictor = Upwind(u_predictor, M)
-            # Average fluxes for a second order update
-            u_new = u.copy()
-            # Compute the slope at u and at u_predictor
-            slope1 = np.zeros_like(u)
-            slope2 = np.zeros_like(u)
-            slope1[1:-1] = -1/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
-            slope2[1:-1] = -1/dx * ( k_half[1:] * F_predictor[1:-1] - k_half[:-1] * F_predictor[:-2] )
-            u_new[1:-1] = u[1:-1] + dt/2.0*(slope1[1:-1] + slope2[1:-1])
-        else:
-            raise ValueError("Unknown time integration scheme.")
-        
-        # Boundary conditions: here, we keep the boundary values fixed (Dirichlet)
-        u_new[0] = u_new[1] # !u[0]
-        u_new[-1] = u_new[-2] # !u[-1]
+    u_init_2 = np.where(
+        x <= 0.5,
+        0.1,
+        np.where(
+            x <= 1.5,
+            0.3,
+            np.where(
+                x <= 2.5,
+                0.7,
+                0.2
+            )
+        )
+    )
+    
+    def solver(u):
+        t = 0.0
+        u_vec = [u.copy()]
+        t_vec = [t]
+        while t < T_end:
+            # Compute dt using the CFL condition: dt <= CFL*dx / max|k(x) * f'(u)|
+            speeds = np.abs(k(x) * df(u))
+            max_speed = np.max(speeds)
+            dt_CFL = CFL * dx / (max_speed + 1e-6)
+            dt = min(dt_CFL, T_end - t)
+            
+            # Compute fluxes at interfaces
+            F = Upwind(u, M)
+            # Spatial locations of cell interfaces
+            x_half = (x[:-1] + x[1:]) / 2.0
+            k_half = k(x_half)
+            
+            if time_scheme == 'Euler':
+                # Euler update
+                u_new = u.copy()
+                # Update interior cells (i = 1,..., len(u)-2)
+                u_new[1:-1] = u[1:-1] - dt/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
+            elif time_scheme == 'Heun':
+                # Predictor: Euler step
+                u_predictor = u.copy()
+                u_predictor[1:-1] = u[1:-1] - dt/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
+                # Compute fluxes for the predictor
+                F_predictor = Upwind(u_predictor, M)
+                # Average fluxes for a second order update
+                u_new = u.copy()
+                # Compute the slope at u and at u_predictor
+                slope1 = np.zeros_like(u)
+                slope2 = np.zeros_like(u)
+                slope1[1:-1] = -1/dx * ( k_half[1:] * F[1:-1] - k_half[:-1] * F[:-2] )
+                slope2[1:-1] = -1/dx * ( k_half[1:] * F_predictor[1:-1] - k_half[:-1] * F_predictor[:-2] )
+                u_new[1:-1] = u[1:-1] + dt/2.0*(slope1[1:-1] + slope2[1:-1])
+            else:
+                raise ValueError("Unknown time integration scheme.")
+            
+            # Boundary conditions: here, we keep the boundary values fixed (Dirichlet)
+            u_new[0] = u_new[1] # !u[0]
+            u_new[-1] = u_new[-2] # !u[-1]
+            
+            u = u_new.copy()
+            t += dt
+            u_vec.append(u.copy())
+            t_vec.append(t)
         
-        u = u_new.copy()
-        t += dt
-        u_vec.append(u.copy())
-        t_vec.append(t)
+        return np.array(t_vec), np.array(u_vec)
+    
+    t_vec_1, u_vec_1 = solver(
+        u=u_init_1
+    )
+    if number_samples == 1:
+        return [t_vec_1], [u_vec_1]
     
-    return np.array(t_vec), np.array(u_vec)
+    # number_samples = 2
+    t_vec_2, u_vec_2 = solver(
+        u=u_init_2
+    )
+    t_vec = [t_vec_1, t_vec_2]
+    u_vec = [u_vec_1, u_vec_2]
+
+    return t_vec, u_vec
 
 def extract_positions(t_vec, h_, x):
     """
@@ -153,54 +183,43 @@ def extract_positions(t_vec, h_, x):
         idx_t = np.argmin(np.abs(t_vec - t_obs))
         h_extracted.append(h_[idx_t, :])
     h_extracted = np.array(h_extracted)[:, [np.argmin(np.abs(x - x_obs)) for x_obs in x_observed]]
-    return h_extracted, times_observed
-
-# Initial guess for Theta (parameters for k(x))
-Theta = np.random.normal(k_mu, k_sigma, k_i+2).clip(0.95, 1.3)
-x_Theta = np.linspace(x0, x1, k_i+2)
-
-# Solve the PDE with the initial Theta
-t_vec, h_Theta = solve_PDE(Theta, T_end, NTime, x, dx, time_scheme='Euler')
-h_Theta, _ = extract_positions(t_vec, h_Theta, x)
-h_Theta = h_Theta.T.flatten()
-dist_Theta = h_Theta - u_true_flatten
-
-# Burn-in phase (without adaptive beta for simplicity)
-for i in range(tau):
-    xi = np.random.normal(0, k_sigma, size=Theta.shape)
-    Phi = np.sqrt(1 - beta**2) * Theta + beta * xi
-    
-    t_vec, h_Phi = solve_PDE(Phi, T_end, NTime, x, dx, time_scheme='Euler')
-    h_Phi, _ = extract_positions(t_vec, h_Phi, x)
-    h_Phi = h_Phi.T.flatten()
-    dist_Phi = h_Phi - u_true_flatten
-    
-    loss_Theta = 0.5 * (dist_Theta @ dist_Theta.T) / gamma**2
-    loss_Phi   = 0.5 * (dist_Phi @ dist_Phi.T) / gamma**2
-    
-    a = min(1, np.exp(loss_Theta - loss_Phi))
-    if np.random.uniform(0, 1) <= a:
-        Theta = Phi
-        h_Theta = h_Phi
-        dist_Theta = dist_Phi
+    return h_extracted
 
-LOSS_THETA = []
-LOSS_PHI = []
+def MCMC(
+    MCMC_STEPS=MCMC_STEPS,
+    MCMC_SUBSTEPS=MCMC_SUBSTEPS,
+    tau=tau,
+    beta=beta,
+    gamma=gamma,
+    k_i=k_i,
+    k_mu=k_mu,
+    k_sigma=k_sigma,
+    u_true=u_true,
+    number_samples=1
+):
+    if number_samples == 1:
+        u_true = u_true[:,:,0]
+        u_true_flatten = u_true.flatten()
+    else:
+        u_true_flatten = u_true.flatten()
+    # Initial guess for Theta (parameters for k(x))
+    Theta = np.random.normal(k_mu, k_sigma, k_i+2).clip(0.95, 1.3)
+    x_Theta = np.linspace(x0, x1, k_i+2)
 
-print('Starting MCMC with adaptive beta')
-# For adaptive beta, track the number of accepted proposals in each block.
-target_accept_low = 0.2
-target_accept_high = 0.5
+    # Solve the PDE with the initial Theta
+    t_vec, h_Theta = solve_PDE(Theta, T_end, NTime, x, dx, x_Theta, number_samples)
+    h_Theta = np.array([extract_positions(t_vec[i], h_Theta[i], x) for i in range(number_samples)]).T
+    h_Theta = h_Theta.flatten()
+    dist_Theta = h_Theta - u_true_flatten
 
-for mcmc_iter in tqdm(range(MCMC_STEPS)):
-    accepted = 0
-    for _ in range(MCMC_SUBSTEPS):
+    # Burn-in phase (without adaptive beta for simplicity)
+    for i in range(tau):
         xi = np.random.normal(0, k_sigma, size=Theta.shape)
         Phi = np.sqrt(1 - beta**2) * Theta + beta * xi
         
-        t_vec, h_Phi = solve_PDE(Phi, T_end, NTime, x, dx, time_scheme='Euler')
-        h_Phi, _ = extract_positions(t_vec, h_Phi, x)
-        h_Phi = h_Phi.T.flatten()
+        t_vec, h_Phi = solve_PDE(Phi, T_end, NTime, x, dx, x_Theta, number_samples)
+        h_Phi = np.array([extract_positions(t_vec[i], h_Phi[i], x) for i in range(number_samples)]).T
+        h_Phi = h_Phi.flatten()
         dist_Phi = h_Phi - u_true_flatten
         
         loss_Theta = 0.5 * (dist_Theta @ dist_Theta.T) / gamma**2
@@ -211,71 +230,123 @@ for mcmc_iter in tqdm(range(MCMC_STEPS)):
             Theta = Phi
             h_Theta = h_Phi
             dist_Theta = dist_Phi
-            accepted += 1
 
-    # Save losses at the end of each outer step
-    LOSS_THETA.append(loss_Theta)
-    LOSS_PHI.append(loss_Phi)
-    
-    # Compute acceptance ratio for this block
-    acc_ratio = accepted / MCMC_SUBSTEPS
-    acceptance_history.append(acc_ratio)
-    beta_history.append(beta)
-    
-    # Adapt beta based on acceptance ratio:
-    if acc_ratio < target_accept_low:
-        beta *= 0.9
-    elif acc_ratio > target_accept_high:
-        beta *= 1.1
-# Combine both loss functions
-plt.figure()
-plt.title('Losses')
-plt.plot(LOSS_THETA, '--o', color='red', label='$\Phi(\Theta)$')
-plt.plot(LOSS_PHI, '--o', color='green', label='$\Phi(\phi$')
-plt.xlabel('Iteration')
-plt.ylabel('Loss')
-plt.legend()
-plt.grid()
-plt.savefig('Figures/losses.pdf')
-plt.show()
+    LOSS_THETA = []
+    LOSS_PHI = []
 
-# Only the phi(theta) loss is plotted
-plt.figure()
-plt.title('Evolution of entropy loss')
-plt.loglog(LOSS_THETA, '--o', color='red', label='$\Phi(\Theta)$')
-plt.xlabel('Iteration')
-plt.ylabel('$\Phi(\Theta)$')
-plt.grid()
-plt.savefig('Figures/loss_theta.pdf')
-plt.show()
+    print('Starting MCMC with adaptive beta')
+    # For adaptive beta, track the number of accepted proposals in each block.
+    target_accept_low = 0.2
+    target_accept_high = 0.5
 
-# Solve with the final Theta
-t_vec, h_final = solve_PDE(Theta, T_end, NTime, x, dx)
-h_final, _ = extract_positions(t_vec, h_final, x)
-h_final = h_final.T
+    for mcmc_iter in tqdm(range(MCMC_STEPS)):
+        accepted = 0
+        for _ in range(MCMC_SUBSTEPS):
+            xi = np.random.normal(0, k_sigma, size=Theta.shape)
+            Phi = np.sqrt(1 - beta**2) * Theta + beta * xi
+            
+            t_vec, h_Phi = solve_PDE(Phi, T_end, NTime, x, dx, x_Theta, number_samples)
+            h_Phi = np.array([extract_positions(t_vec[i], h_Phi[i], x) for i in range(number_samples)]).T
+            h_Phi = h_Phi.flatten()
+            dist_Phi = h_Phi - u_true_flatten
+            
+            loss_Theta = 0.5 * (dist_Theta @ dist_Theta.T) / gamma**2
+            loss_Phi   = 0.5 * (dist_Phi @ dist_Phi.T) / gamma**2
+            
+            a = min(1, np.exp(loss_Theta - loss_Phi))
+            if np.random.uniform(0, 1) <= a:
+                Theta = Phi
+                h_Theta = h_Phi
+                dist_Theta = dist_Phi
+                accepted += 1
 
-colors = ['blue', 'red', 'green', 'purple']
-plt.figure()
-plt.title('Predicted vs True')
-for i in range(4):
-    plt.plot([], [], color=colors[i], label=f'x = {x_observed[i]}')
-    plt.plot(h_final[i, :], color=colors[i])
-    plt.plot(u_true[i, :], '--o', color=colors[i])
-plt.plot([], [], color='black', label='Predicted')
-plt.plot([], [], '--o', color='black', label='True')
-plt.xlabel('Time')
-plt.ylabel('u(x, t)')
-plt.legend(ncol=4, loc='upper center', bbox_to_anchor=(0.5, -0.1))
-plt.grid()
-plt.savefig('Figures/predicted_vs_true.pdf')
-plt.show()
+        # Save losses at the end of each outer step
+        LOSS_THETA.append(loss_Theta)
+        LOSS_PHI.append(loss_Phi)
+        
+        # Compute acceptance ratio for this block
+        acc_ratio = accepted / MCMC_SUBSTEPS
+        acceptance_history.append(acc_ratio)
+        beta_history.append(beta)
+        
+        # Adapt beta based on acceptance ratio:
+        if acc_ratio < target_accept_low:
+            beta *= 0.9
+        elif acc_ratio > target_accept_high:
+            beta *= 1.1
+
+    return Theta, x_Theta, LOSS_THETA, LOSS_PHI
+
+if __name__ == '__main__':
+    Theta, x_Theta, LOSS_THETA, LOSS_PHI = MCMC(
+        MCMC_STEPS=MCMC_STEPS,
+        MCMC_SUBSTEPS=MCMC_SUBSTEPS,
+        tau=tau,
+        beta=beta,
+        gamma=gamma,
+        k_i=k_i,
+        k_mu=k_mu,
+        k_sigma=k_sigma,
+        u_true=u_true,
+        number_samples=NUMBER_SAMPLES
+    )
+    if NUMBER_SAMPLES == 1:
+        FigureFolder = 'Figures/SingleSample'
+    else:
+        FigureFolder = 'Figures/DoubleSample'
+    # Combine both loss functions
+    plt.figure()
+    plt.title('Losses')
+    [plt.plot(LOSS_THETA[:][i], '--o', color='red', label='$\Phi(\Theta)$') for i in range(NUMBER_SAMPLES)]
+    [plt.plot(LOSS_PHI[:][i], '--o', color='green', label='$\Phi(\phi$)') for i in range(NUMBER_SAMPLES)]
+    plt.xlabel('Iteration')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.grid()
+    plt.savefig(f'{FigureFolder}/losses.pdf', bbox_inches='tight')
+    plt.show()
+
+    # Only the phi(theta) loss is plotted
+    plt.figure()
+    plt.title('Evolution of entropy loss')
+    [plt.loglog(LOSS_THETA[:][i], '--o', color='red', label='$\Phi(\Theta)$') for i in range(NUMBER_SAMPLES)]
+    plt.xlabel('Iteration')
+    plt.ylabel('$\Phi(\Theta)$')
+    plt.grid()
+    plt.savefig(f'{FigureFolder}/loss_theta.pdf', bbox_inches='tight')
+    plt.show()
+
+    # Solve with the final Theta
+    t_vec, h_final = solve_PDE(Theta, T_end, NTime, x, dx, x_Theta, NUMBER_SAMPLES)
+    h_final = [extract_positions(t_vec[i], h_final[i], x) for i in range(NUMBER_SAMPLES)]
+    h_final = np.array([h_final[i].T for i in range(NUMBER_SAMPLES)])
+
+    colors = ['blue', 'red', 'green', 'purple']
+    for sample in range(NUMBER_SAMPLES):
+        plt.figure()
+        plt.title(f'Predicted vs True - sample {sample+1}')
+        for i in range(4):
+            plt.plot([], [], color=colors[i], label=f'x = {x_observed[i]}')
+            plt.plot(h_final[sample, i, :], color=colors[i])
+            plt.plot(u_true[i, :, sample], '--o', color=colors[i])
+        plt.plot([], [], color='black', label='Predicted')
+        plt.plot([], [], '--o', color='black', label='True')
+        plt.xlabel('Time')
+        plt.ylabel('u(x,t)')
+        plt.legend(ncol=4, loc='upper center', bbox_to_anchor=(0.5, -0.1))
+        plt.grid()
+        if NUMBER_SAMPLES == 1:
+            plt.savefig(f'{FigureFolder}/predicted_vs_true_single_sample.pdf', bbox_inches='tight')
+        else:
+            plt.savefig(f'{FigureFolder}/predicted_vs_true_both_samples_sample_{sample+1}.pdf', bbox_inches='tight')
+        plt.show()
 
-# Plot the final theta -> prediction for k(x)
-plt.figure()
-plt.title('Prediction for k(x)')
-plt.plot(x_Theta, Theta, '--o', color='blue')
-plt.xlabel('$x$')
-plt.ylabel('$k(x)$')
-plt.grid()
-plt.savefig('Figures/prediction_kx.pdf')
-plt.show()
+    # Plot the final theta -> prediction for k(x)
+    plt.figure()
+    plt.title('Prediction for k(x)')
+    plt.plot(x_Theta, Theta, '--o', color='blue')
+    plt.xlabel('$x$')
+    plt.ylabel('$k(x)$')
+    plt.grid()
+    plt.savefig(f'{FigureFolder}/prediction_kx.pdf', bbox_inches='tight')
+    plt.show()