brownian_motion.py

import numpy as np

def brownian_particle(n_steps, n_samples, dx, y0):
    """
    Brownian particle simulator
    
    :param n_steps: total steps
    :param n_samples: number of trajectories
    :param dx: maximum step length
    :param y0: starting position
    :return x, y: time, array containing N_samples trajectories with N_steps
    """

    y = np.zeros((n_steps, n_samples))
    x = np.linspace(0, n_steps, n_steps)
    y[0, :] = y0

    for i_sample in range(0, n_samples):
        i_step = 1
        
        while True:

            if i_step >= n_steps:
                break

            random_number = np.random.rand()
            
            #u = (0.5 - random_number) * dx
            
            #print(u, random_number)
            if random_number >= 0.5 and random_number < 0.75: # acceptance rule for moving the particle
                y[i_step, i_sample] = y[i_step - 1, i_sample] + dx
            elif random_number >= 0.75:    
                y[i_step, i_sample] = y[i_step - 1, i_sample] - dx
            else: # the particle remains in the same position
                y[i_step, i_sample] = y[i_step - 1, i_sample] 
                
            i_step += 1
    return x, y


def biased_brownian_particle(n_steps, n_samples, dx, y0):
    """
    Brownian particle simulator
    
    :param n_steps: total steps
    :param n_samples: number of trajectories
    :param dx: maximum step length
    :param y0: starting position
    :return x, y: time, array containing N_samples trajectories with N_steps
    """

    y = np.zeros((n_steps, n_samples))
    x = np.linspace(0, n_steps, n_steps)
    y[0, :] = y0

    for i_sample in range(0, n_samples):
        i_step = 1
        
        while True:

            if i_step >= n_steps:
                break

            random_number = np.random.rand()
            
            u = (0.5 - random_number) * dx
            
            #print(u, random_number)
            if random_number >= 0.5: # acceptance rule for moving the particle
                y[i_step, i_sample] = y[i_step - 1, i_sample] + u

            else: # the particle remains in the same position
                y[i_step, i_sample] = y[i_step - 1, i_sample] 
                
            i_step += 1
    return x, y

if __name__ == "__main__":
	import numpy as np
	import maptlotlib.pyplot as plt
	b = brownian_particle(n_steps = 100, n_samples= 100, dx=1, y0=0 )
	plt.plot(b[0], b[1])
	plt.show()