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

#*****************************************************************************80
#
## scatter_2d() follows multiple Brownian particles in 2D.
#
#  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

  print ( '' )
  print ( 'scatter_2d():' )
  print ( '  Display multiple Brownian particles in 2D.' )

  p = 1000
  n = 100
  m = 2
  d = 10.0
  t = 1.0
  x = scatter_simulation ( p, n, m, d, t )
  dt = t / float ( n )
  sigma = np.sqrt ( 2.0 * m * d * dt * n )
#
#  Display particles.
#
  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.show ( )
  plt.close ( )

  return

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

#*****************************************************************************80
#
## 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

if ( __name__ == '__main__' ):
  scatter_2d ( )