Calculating long-range particle interactions with numpy

import random
import math
import time
import numpy as np

#import scipy.misc as scm
from multiprocessing import Pool

np.set_printoptions(threshold=np.inf)
random.seed(1)

XX = 0
YY = 1
ZZ = 2

#total particle num
PN = 1000

xyzV = [[0 for i in range(3)] for j in range(PN)]
xyzP = np.zeros((PN,3))
xyzF = np.zeros((PN,3))

def find_pair():
  global PN
  global xyzF
  global xyzP

  tmp_index = np.arange(xyzP.shape[0])

  xx, yy = np.meshgrid(tmp_index, tmp_index)
  distances = np.linalg.norm(xyzP[xx]-xyzP[yy], axis=2)
  dxyz = xyzP[xx] - xyzP[yy]

  distances = distances[:,:,np.newaxis]
  tmp_force = dxyz/distances
  tmp_force = np.nan_to_num(tmp_force,0)

  mmm = tmp_force.sum(axis=1)
  xyzF = xyzF + mmm

def init_lattice():
  global xyzF
  global xyzP

  pnum = 0
  while pnum < PN:
    xyzP[pnum][XX] = random.uniform(-1,1)
    xyzP[pnum][YY] = random.uniform(-1,1)
    xyzP[pnum][ZZ] = random.uniform(-1,1)
    xyzF[pnum][XX] = random.uniform(-1,1)
    xyzF[pnum][YY] = random.uniform(-1,1)
    xyzF[pnum][ZZ] = random.uniform(-1,1)
    pnum += 1


def results_sum():
  #global xyzF
  pnum = 0
  total_F = 0
  while pnum < PN:
    total_F = total_F + xyzF[pnum][XX]
    total_F = total_F + xyzF[pnum][YY]
    total_F = total_F + xyzF[pnum][ZZ]
    pnum += 1

  print (total_F)


if __name__ == "__main__":

  init_lattice()
  find_pair()
  results_sum()

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

CAPTCHA