Solving diffusion convection SDE representation using Milstein method with Boost (c++) - not converging to deterministic solution

Question

I'm working on a research project to simulate the Diffusion advection equation like

I solved the PDE using RK4 and finite differences and I'm comparing the result to solving the stochastic representation from Feynman Kac ($dx=vdt+DdW$) using Milstein's method (Slides for stochastic form slide 64). But when I'm plotting both versions, the stochastic solution doesn't match the deterministic one. Below in red, I made a histogram of the distribution of particles I simulated, and plotted the center of the histogram bar. They don't line up with the blue function solution.

Below is my c++ code and my python script that created the plot. I'm using Boost's custom stepper (Custom stepper doc) as an example. I'm using D(x)=D_0+(D_0x^2), D_0=0.009, $(x) = v_0sin^2(x), v_0=5.

#include <algorithm> // for std::transform and std::accumulate
#include <boost/numeric/odeint.hpp>
#include <cmath>
#include <fstream>
#include <iostream>
#include <numeric>
#include <random>
#include <vector>

using namespace std;
const static size_t N = 1;
typedef vector<vector<double>> ParticlePositions;
const static double v = 0.1; // Convection coefficient
const static double D = 0.009; // Diffusion coefficient
const static double f = 0.0; // Source Term
typedef boost::array< double , N > state_type;

// Function to calculate initial Gaussian distribution
vector<double> initialize_gaussian(double amp, const vector<double>& x)
{
   vector<double> u(x.size());
   double sigma = 0.01;
   for (size_t i = 0; i < x.size(); ++i) {
      u[i] = amp * (1 / sqrt(0.001 * 2 * M_PI)) * exp(-0.5 * (x[i] * x[i] / 0.001));
   }
   return u;
}

double D_eval(double x) {
    return D + (D * x * x);
}
double v_eval(double x) {
    return 50 * v * sin(x) * sin(x);  
}

double dD_dx(double x) {
    return 2 * D * x;
}


// ========================
// === Stochastic Class ===
// === dx=v*dt+D*dW ===
// ========================

// Euler Maruyama: xn+1=xn+dt(w+alpha*dW)
template <size_t N>
class stochastic_euler {
 public:
   typedef boost::array<double, N> state_type;
   typedef boost::array<double, N> deriv_type;
   typedef double value_type;
   typedef double time_type;
   typedef unsigned short order_type;
   typedef boost::numeric::odeint::stepper_tag stepper_category;

   static order_type order(void) { return 1; }

   template <class System>
   void do_step(System system, state_type& x, time_type t, time_type dt) const
   {
      deriv_type det, stoch;
      system.first(x, det); 
      system.second(x, stoch);
      for (size_t i = 0; i < x.size(); ++i) {
         double milCorrect = 4 * D_eval(x[i]) * abs(dD_dx(x[i]));
         x[i] += dt * (det[i] - milCorrect) + \
                     + sqrt(milCorrect * dt * dt) * stoch[i] * stoch[i] + \
                     sqrt(2 * D_eval(x[i]) * dt) * stoch[i]; // sqrt(dt) is standard dev
      }
   }
};

// Pair of functions for the deterministic and stochastic parts of the SDE
struct sde_det {
   void operator()(const state_type& x, state_type& dxdt) const
   {
      for (size_t i = 0; i < x.size(); ++i)
         dxdt[i] = -2*v_eval(x[i]);
   }
};

struct sde_stoch {
   mt19937& m_rng;
   normal_distribution<> m_dist;
   // Constructor
   sde_stoch(mt19937& rng, double sigma) : m_rng(rng), m_dist(0.0, sigma) {}

   void operator()(const state_type& x, state_type& dxdt)
   {
      for (size_t i = 0; i < x.size(); ++i)
         dxdt[i] = m_dist(m_rng);
   }
};


// Used to store particle positions for each time step into vector argument
struct particle_observer {
   ParticlePositions& positions;
   int particle_index;
   int step;

   // Constructor
   particle_observer(ParticlePositions& positions, int particle_index)
       : positions(positions), particle_index(particle_index), step(0) {}

   void operator()(const state_type& x, double t)
   {
      positions[particle_index][step] = x[0];
      step++;
   }
};


// Compute PDE derivative at index i for the current solution u_curr
double compute_rhs(const vector<double>& u_curr, int i, double dx) {
    double x_i = (i * dx) - 2.0; // physical position
    double Dval = D_eval(x_i); 
    double dDval = dD_dx(x_i);
    double w = v_eval(x_i);

    double diff_term = (u_curr[i+1] - 2.0 * u_curr[i] + u_curr[i-1]) / (dx * dx);
    double convection_term = (u_curr[i+1] - u_curr[i-1]) / (2.0 * dx);

    // du/dt
    return Dval * diff_term + dDval * convection_term + w * convection_term;
}

// Perform one RK4 time step, updating sol[j] from sol[j-1]
void update_solution_rk4(vector<vector<double>>& sol, double dx, double dt, int nx, int j) {
    const vector<double>& u_old = sol[j-1];
    vector<double> k1(nx, 0.0), k2(nx, 0.0), k3(nx, 0.0), k4(nx, 0.0);
    vector<double> u_temp1(u_old), u_temp2(u_old), u_temp3(u_old);
    vector<double> u_new(u_old);

    // Compute k1 for interior points
   for(int i = 1; i < nx - 1; ++i) {
      k1[i] = compute_rhs(u_old, i, dx);
      u_temp1[i] = u_old[i] + 0.5 * dt * k1[i];
      k2[i] = compute_rhs(u_temp1, i, dx);
      u_temp2[i] = u_old[i] + 0.5 * dt * k2[i];
      k3[i] = compute_rhs(u_temp2, i, dx);
      u_temp3[i] = u_old[i] + dt * k3[i];
      k4[i] = compute_rhs(u_temp3, i, dx);
      u_new[i] = u_old[i] + (dt / 6.0) * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i]);
    }

    // Apply or maintain boundary conditions if needed (e.g., Dirichlet = 0)
    u_new[0] = 0; //u_old[0];
    u_new[nx-1] = 0; //u_old[nx-1];

    // Store result in sol[j]
    sol[j] = u_new;
}


int main()
{
   cout << "mctest_D_milstein.cpp" << endl;
   typedef boost::array<double, N> state_type;
   using namespace boost::numeric::odeint;
   random_device rd;
   mt19937 gen(rd());
   normal_distribution<> normal_dist(0.0, 1.0);
   uniform_real_distribution<> uniform_dist(-2.0, 2.0);
   
   // Parameters
   double L = 4.0;
   double T = 1.0;
   int nx = 200;
   int nt = 199;

   // ================================
   // === Deterministic Simulation ===
   // ================================

   double dx = L / nx;
   double dt = T / nt;
   

   // Stability condition check
   if (v * dt / dx > 1 || D * dt / (dx * dx) > 0.5) {
      cerr << "Stability condition not met. Reduce dt or dx." << endl;
      return -1;
   }

   // Spatial grid
   vector<double> x(nx);
   for (int i = 0; i < nx; ++i) {
      x[i] = -2.0 + i * dx;
   }

   // Initial condition
   double amp = 1;
   vector<double> u = initialize_gaussian(amp, x);

   // Array to store the solution at each timestep
   vector<vector<double>> sol(nt + 1, vector<double>(nx));
   sol[0] = u;

    // Time-stepping loop
    for (int j = 1; j < nt + 1; ++j) {
        update_solution_rk4(sol, dx, dt, nx, j);
    }

   // save to file

   // =============================
   // === Stochastic Simulation ===
   // =============================

   // nt *= 2;
   // T *= 2;
   vector<int> num_particles_list = {100, 500, 1000, 5000, 10000, 50000, 100000, 500000};

   typedef stochastic_euler<1> euler_maruyama_stepper;

   for (int num_particles : num_particles_list) {
      vector<double> p_init_pos;
      double maxD = *max_element(sol[0].begin(), sol[0].end());
      cout << "MaxD: " << maxD << endl;
      for (int i = 0; i < num_particles; ++i) {
         double prob = uniform_real_distribution<>(0.0, 1.0)(gen);
         double position = uniform_real_distribution<>(-2.0, 2.0)(gen);
         // cout << initialize_gaussian(amp, {position})[0] / maxD << endl;
         if (prob < initialize_gaussian(amp, {position})[0] / maxD) {
            p_init_pos.push_back(position);
         }
      }

      int final_num_particles = p_init_pos.size();
      cout << "Number of particles: "<< final_num_particles << endl;
      // 2D vector (particle# x time steps) initialized to 0.0
      vector<vector<double>> particle_positions(final_num_particles, vector<double>(nt + 1, 0.0));

      for (int i = 0; i < final_num_particles; ++i) {
         particle_positions[i][0] = p_init_pos[i];
      }

      euler_maruyama_stepper stepper;
      mt19937 rng;
      // run do_step for each particle, using observer to store positions into 2D vector
      for (int i = 0; i < final_num_particles; ++i) {
         state_type x = {particle_positions[i][0]};
         particle_observer observer(particle_positions, i);
         integrate_const( stochastic_euler<1>(), make_pair(sde_det(), sde_stoch(rng, 1.0)), 
                           x, 0.0, T, dt, observer );

      }
      // save to file
   }


   return 0;
}

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# ------------- PDE/deterministic -------------
L = 4.0  # Length of the domain
nx = 200  # Number of spatial points

# Read the CSV file
pdeDf = pd.read_csv('deterministic_positions.csv', header=None)

pde_data = pdeDf.values.flatten()
dx = L / nx
x = np.arange(-2, 2, dx)

# ------------- SDE/stochastic -------------
num_particles_list = [100, 500, 1000, 5000, 10000, 50000, 100000, 500000]
average_errors = []
for num_particles in num_particles_list:
   df = pd.read_csv(f"particle_positions_{num_particles}.csv")

   # Get the last timestep
   sdeDf = df['TimeStep'].max()
   # Filter the DataFrame to include only the last timestep
   last_timestep_data = df[df['TimeStep'] == sdeDf]
   # Calculate the number of particles
   num_particles = len(last_timestep_data)

   counts, bin_edges = np.histogram(last_timestep_data['Position'], bins=num_particles//50+4)

   # Plot the particle positions at the last timestep
   # using histogram (frequency of values per bin) get center of each bar and y value
   bins = (bin_edges[:-1] + bin_edges[1:]) / 2

   counts = counts / np.max(counts) * np.max(pde_data)

   errors = []
   for i, bin_center in enumerate(bins):
         closest_x_index = np.argmin(np.abs(x - bin_center))
         error = (np.abs(counts[i] - pde_data[closest_x_index]))**2 # RSS
         errors.append(error) # TODO: check collisional paper summary pdf to verify similar error formula

   errors = np.array(errors)
   # print(errors)
   average_errors.append(np.mean(errors))
   print(f'Average error for {num_particles} particles: {np.mean(errors)}')

   # ------------- Plotting -------------
   plt.figure()
   
   color = 'tab:blue'
   plt.plot(x, pde_data, label='Deterministic Distribution Func', color=color)
   # ax1.set_ylabel('Deterministic Density', color=color)
   plt.tick_params(axis='y', labelcolor=color)
   plt.xlabel('Final Position')
   plt.title(f'Deterministic Distribution & Stochastic Frequency for {num_particles} particles')
   # plt.xlim(-0.75, 0.75)
   color = 'tab:red'
   plt.scatter(bins, counts, color=color, label='Stochastic Binned Frequency')
   # plt.legend()


#Plot all particle positions
if(True):
   num_particles = 50000
   df = pd.read_csv(f"particle_positions_{num_particles}.csv")
   plt.figure()
   for particle_id in df['Particle'].unique():
      particle_data = df[df['Particle'] == particle_id]
      plt.plot(particle_data['TimeStep'], particle_data['Position'], alpha=0.5)
   plt.xlabel('Time Step')
   plt.ylabel('Position')
   plt.title(f'Particle Trajectories for {num_particles} particles')
   plt.grid(True)
   plt.show()

What I have tried so far

I've tried doubling the diffusion and convection functions for just the stochastic version, increasing the number of particles simulated, and increasing my simulation time, but the solutions never converge. I looked at the Feynman Kac formula, and it looks like the stochastic form of the diffusion convection equation is correct. I also tried varying the time step and as I decrease dt, the deterministic/PDE solution does get closer (the decay on the left side gets steeper), but the solution becomes unstable before it converges. I also noticed that the peak of my histogram (highest red dot) doesn't line up with the peak of the PDE solution.

Long story short: why does my deterministic and stochastic solutions to (I think) the same equation not line up?