#! /usr/bin/env python3
#
def brownian_motion_simulation_test ( ):

#*****************************************************************************80
#
## brownian_motion_simulation_test() tests brownian_motion_simulation().
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    22 February 2025
#
#  Author:
#
#    John Burkardt
#
  import matplotlib
  import numpy as np
  import platform

  print ( '' )
  print ( 'brownian_motion_simulation_test():' )
  print ( '  python version:     ' + platform.python_version ( ) )
  print ( '  numpy version:      ' + np.version.version )
  print ( '  matplotlib version: ' + matplotlib.__version__ )
  print ( '  Test brownian_motion_simulation().' )
#
#  Compute the path of a particle undergoing Brownian motion.
#
  for m in range ( 1, 4 ):

    n = 1001
    d = 10.0
    t = 1.0
    x = brownian_motion_simulation ( m, n, d, t )
    brownian_motion_display ( m, n, x )
#
#  Estimate the average displacement of the particle from the origin
#  as a function of time.
#
  for m in range ( 1, 4 ):

    k = 40
    n = 1001
    d = 10.0
    t = 1.0

    dsq = brownian_displacement_simulation ( k, n, m, d, t )
    brownian_displacement_display ( k, n, m, d, t, dsq )
#
#  Plot the location of P particles, taking N steps, in M=2 dimensions,
#  starting from the origin, and
#  each taking K steps.
#
  p = 10000
  n = 100
  m = 1
  d = 10.0
  t = 1.0
  x = brownian_scatter_simulation ( p, n, m, d, t )
  dt = t / float ( n )
  sigma = np.sqrt ( 2.0 * m * d * dt * n )
  brownian_scatter_1d_display ( p, sigma, x )

  p = 1000
  n = 100
  m = 2
  d = 10.0
  t = 1.0
  x = brownian_scatter_simulation ( p, n, m, d, t )
  dt = t / float ( n )
  sigma = np.sqrt ( 2.0 * m * d * dt * n )
  brownian_scatter_2d_display ( p, sigma, x )

  p = 1000
  n = 100
  m = 3
  d = 10.0
  t = 1.0
  x = brownian_scatter_simulation ( p, n, m, d, t )
  brownian_scatter_3d_display ( p, x )
#
#  Terminate.
#
  print ( '' )
  print ( 'brownian_motion_simulation_test():' )
  print ( '  Normal end of execution.' )
  return

def brownian_motion_simulation ( m, n, d, t ):

#*****************************************************************************80
#
## brownian_motion_simulation() simulates Brownian motion.
#
#  Discussion:
#
#    Thanks to Feifei Xu for pointing out a missing factor of 2 in the
#    stepsize calculation, 08 March 2016.
#
#    Thanks to Joerg Peter Pfannmoeller for pointing out a missing factor
#    of M in the stepsize calculation, 23 April 2018.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    02 January 2024
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer M, the spatial dimension.
#
#    integer N, the number of time steps to take, plus 1. 
#
#    real D, the diffusion coefficient.  
#
#    real T, the total time.
#
#  Output:
#
#    real X(N,M), N successive locations of a particle in M dimensional space.
#
  import numpy as np
#
#  Set the time step.
#
  dt = t / float ( n - 1 )
#
#  Make space for the data.
#
  x = np.zeros ( [ n, m ] )
#
#  Set the characteristic stepsize.
#
  sigma = np.sqrt ( 2.0 * m * d * dt )
#
#  Compute the individual steps.
#
  for j in range ( 1, n ):
#
#  S is the stepsize
#
    s = sigma * np.random.standard_normal ( )
#
#  Choose a random unit direction vector.
#
    u = np.random.standard_normal ( m )
    u = u / np.linalg.norm ( u )
#
#  Each position is the sum of the previous steps.
#
    x[j,0:m] = x[j-1,0:m] + s * u[0:m]

  return x

def brownian_displacement_display ( k, n, m, d, t, dsq ):

#*****************************************************************************80
#
## brownian_displacement_display() displays average Brownian motion displacement.
#
#  Discussion:
#
#    Thanks to Feifei Xu for pointing out a missing factor of 2 in the
#    displacement calculation.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    02 January 2024
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer K, the number of repetitions.
#
#    integer N, the number of time steps.  
#
#    integer M, the spatial dimension.
#
#    real D, the diffusion coefficient.
#
#    real T, the total time.
#
#    real DSQ(N,K), the displacements over time for each repetition.
#
  import matplotlib.pyplot as plt
  import numpy as np
#
#  Get the T values.
#
  tvec = np.linspace ( 0, t, n )
#
#  Select 5 random trajectories for display.
#
  for s in range ( 0, 5 ):
    i = int ( k * np.random.random ( ) )
    plt.plot ( tvec, dsq[:,i], 'b-' )
#
#  Display the average displacement.
#
  dsq_ave = np.sum ( dsq, 1 ) / float ( k )
  plt.plot ( tvec, dsq_ave, 'r-', linewidth = 2 )
#
#  Display the ideal displacment.
#
  dsq_ideal = 2.0 * m * d * tvec
  plt.plot ( tvec, dsq_ideal, 'k-', linewidth = 3 )

  plt.grid ( True )
  plt.xlabel ( '<--T-->' )
  plt.ylabel ( '<--D^2-->' )
  plt.title ( 'Squared displacement (Red), Predicted (Black), Samples (Blue)' )

  filename = 'displacement_' + str ( m ) + '.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def brownian_displacement_simulation ( k, n, m, d, t ):

#*****************************************************************************80
#
## brownian_displacement_simulation() simulates Brownian displacement.
#
#  Discussion:
#
#    Thanks to Feifei Xu for pointing out a missing factor of 2 in the
#    stepsize calculation, 08 March 2016.
#
#    This function computes the square of the distance of the Brownian
#    particle from the starting point, repeating this calculation 
#    several times.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    10 June 2018
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer K, the number of repetitions.
#    The default is 20
#
#    integer N, the number of time steps to take, plus 1.
#    This might be 1001.
#
#    integer M, the spatial dimension.  Typically, this is 2.
#
#    real D, the diffusion coefficient.  This might be 10.0.
#    Computationally, this is simply a scale factor between time and space.
#
#    real T, the total time, which defaults to 1.0.
#
#  Output:
#
#    real DSQ(N,K), the displacements over time for each repetition.
#    DSQ(1,:) is 0.0, because we include the displacement at the initial time. 
#
  import numpy as np

  dsq = np.zeros ( [ n, k ] )

  for i in range ( 0, k ):

    x = brownian_motion_simulation ( m, n, d, t )
    dsq[0:n,i] = np.sum ( x[0:n,0:m] ** 2, axis = 1 )
    
  return dsq

def brownian_motion_display ( m, n, x ):

#*****************************************************************************80
#
## brownian_motion_display() displays successive Brownian motion positions.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    26 April 2018
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer M, the spatial dimension.
#    M should be 1, 2 or 3.
#
#    integer N, the number of time steps. 
#
#    real X(N,M), the particle positions.
#
  import matplotlib.pyplot as plt
  import numpy as np
  from mpl_toolkits.mplot3d import Axes3D

  if ( m == 1 ):

    y = np.linspace ( 0, n - 1, n ) / float ( n - 1 )
    plt.plot ( x[:,0], y[:], 'b', linewidth = 2 )
    plt.plot ( x[0,0], y[0], 'g.', markersize = 35 )
    plt.plot ( x[n-1,0], y[n-1], 'r.', markersize = 35 )
    plt.grid ( True )
    plt.xlabel ( '<--X-->' )
    plt.ylabel ( '<--Time-->' )
    plt.title ( 'Brownian motion simulation in 1D' )
    filename = 'motion_' + str ( m ) + 'd.png'
    plt.savefig ( filename )
    print ( '  Graphics saved as "' + filename + '"' )
    plt.close ( )

  elif ( m == 2 ):

    plt.plot ( x[:,0], x[:,1], 'b', linewidth = 2 )
    plt.plot ( x[0,0], x[0,1], 'g.', markersize = 35 )
    plt.plot ( x[n-1,0], x[n-1,1], 'r.', markersize = 35 )
    plt.grid ( True )
    plt.xlabel ( '<--X-->' )
    plt.ylabel ( '<--Y-->' )
    plt.title ( 'Brownian motion simulation in 2D' )
    plt.axis ( 'equal' )
    filename = 'motion_' + str ( m ) + 'd.png'
    plt.savefig ( filename )
    print ( '  Graphics saved as "' + filename + '"' )
    plt.close ( )
#
#  May I just say that I am struck by the inconsistency between 2D and 3D 
#  plotting?
#
  elif ( m == 3 ):

    fig = plt.figure ( )
    ax = fig.add_subplot ( 111, projection = '3d' )
    ax.plot ( x[:,0], x[:,1], x[:,2], 'b', linewidth = 2 )
    ax.scatter ( x[0,0], x[0,1], x[0,2], c = 'g', marker = 'o', s = 100 )
    ax.scatter ( x[n-1,0], x[n-1,1], x[n-1,2], c = 'r', marker = 'o', s = 100 ) 
    ax.grid ( True )
    ax.set_xlabel ( '<--X-->' )
    ax.set_ylabel ( '<--Y-->' )
    ax.set_zlabel ( '<--Z-->' )
    plt.title ( 'Brownian motion simulation in 3D' )
    filename = 'motion_' + str ( m ) + 'd.png'
    plt.savefig ( filename )
    print ( '  Graphics saved as "' + filename + '"' )
    plt.close ( )

  else:

    print ( '' )
    print ( 'brownian_motion_display(): Fatal error!' )
    print ( '  Cannot display data except for M = 1, 2, 3.' )
    raise Exception ( 'brownian_motion_display(): Fatal error!' )

  return

def brownian_scatter_1d_display ( p, sigma, x ):

#*****************************************************************************80
#
## brownian_scatter_1d_display() histograms P independent particles after N steps.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    22 February 2025
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  plt.clf ( )
  plt.hist ( x, bins = 31, density = True )
  x2 = np.linspace ( -10, +10, 101 )
  y2 = 1.0 / np.sqrt ( 2.0 * np.pi * sigma**2 ) * np.exp ( - 0.5 * x2**2 / sigma**2 ) 
  plt.plot ( x2, y2, 'r-', linewidth = 3 )
  plt.legend ( [ 'Theory', 'Data' ] )
  plt.grid ( True )
  plt.xlabel ( '<--X-->' )
  plt.ylabel ( '<--Frequency-->' )
  plt.title ( 'P particles in 1 dimension after N steps' )
  filename = 'scatter_1d.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def brownian_scatter_2d_display ( p, sigma, x ):

#*****************************************************************************80
#
## brownian_scatter_2d_display() histograms P independent particles after N steps.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    22 February 2025
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  plt.clf ( )
  plt.plot ( x[:,0], x[:,1], 'b.', markersize = 5 )
  theta = np.linspace ( 0.0, 2.0 * np.pi, 51 )
  x2 = np.sqrt ( 2.0 ) * sigma * np.cos ( theta )
  y2 = np.sqrt ( 2.0 ) * sigma * np.sin ( theta )
  plt.plot ( x2, y2, 'r-', linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<--X-->' )
  plt.ylabel ( '<--Y-->' )
  plt.title ( 'P particles in 2 dimensions after N steps' )
  plt.axis ( 'equal' )
  filename = 'scatter_2d.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def brownian_scatter_3d_display ( p, x ):

#*****************************************************************************80
#
## brownian_scatter_3d_display() histograms P independent particles after N steps.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    23 February 2025
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np
  from mpl_toolkits.mplot3d import Axes3D

  fig = plt.figure ( )
  ax = fig.add_subplot ( 111, projection = '3d' )
  ax.scatter ( x[:,0], x[:,1], x[:,2], c = 'b', marker = 'o', s = 5 )
  ax.grid ( True )
  ax.set_xlabel ( '<--X-->' )
  ax.set_ylabel ( '<--Y-->' )
  ax.set_zlabel ( '<--Z-->' )
  plt.title ( 'P particles in 3 dimensions after N steps' )
  filename = 'scatter_3d.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def brownian_scatter_simulation ( p, n, m, d, t ):

#*****************************************************************************80
#
## brownian_scatter_simulation() lets P particles take N Brownian steps each.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    22 February 2025
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    integer p: the number of particles.
#
#    integer m: the spatial dimension.
#
#    integer n: the number of time steps to take 
#
#    real d: the diffusion coefficient.  
#
#    real t: the total time.
#
#  Output:
#
#    real x(p,m): the locations of P particles in M dimensional space
#    after each has taken N steps.
#
  import numpy as np
#
#  Set the time step.
#
  dt = t / n
#
#  Make space for the data.
#
  x = np.zeros ( [ p, m ] )
#
#  Compute the individual steps.
#
  for j in range ( 0, n ):
#
#  Choose direction at random.
#
    direct = np.random.standard_normal ( [ p, m ] )
    norm_direct = np.sqrt ( np.sum ( direct ** 2, axis = 1 ) )
    for i in range ( 0, m ):
      direct[i] = direct[i] / norm_direct[i]
#
#  Choose stepsize at random.
#
    step = np.sqrt ( 2.0 * m * d * dt ) * np.random.standard_normal ( p )
#
#  Increment the position.
#
    for j in range ( 0, m ):
      x[:,j] = x[:,j] + step[:] * direct[:,j]

  return x

def timestamp ( ):

#*****************************************************************************80
#
## timestamp() prints the date as a timestamp.
#
#  Licensing:
#
#    This code is distributed under the MIT license. 
#
#  Modified:
#
#    06 April 2013
#
#  Author:
#
#    John Burkardt
#
  import time

  t = time.time ( )
  print ( time.ctime ( t ) )

  return

if ( __name__ == '__main__' ):
  timestamp ( )
  brownian_motion_simulation_test ( )
  timestamp ( )