Diffusion3d/Diffusion3D.py at master · dloman/Diffusion3d · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
#!/usr/bin/python
#
# Author: Daniel Loman
# Date: July 12, 2012
# Filename: Diffusion3D.py
# Description:
#   This program creates 3d matrix then simulates diffusion through the matrix
# and then it visualizes this diffusion using matplotlib
#


import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sys import argv
from os import system, remove
from string import atoi
################################################################################
################################################################################
def Usage():
  print 'Diffuse.py [n t Initialize BoundryCondition](optional settings)'

################################################################################
def Plot(Grid, ax):
  n=len(Grid)+1
  x, y, z = np.mgrid[1:n,1:n,1:n]
  color = ax.scatter(x, y, z, c=Grid, marker='o')
  return color

################################################################################
def GetNeighborhoodDifference(GridValue,BigGrid,x,y,z):
  Diff=0
  for i in range(3):
    for j in range(3):
      for k in range(3):
         Diff+=(GridValue-BigGrid[x+i,y+j,z+k])/20
  if Diff > GridValue:
    return 0
  else:
   return Diff

################################################################################
def GetNextGrid(Grid,BoundryCondition):
  n=len(Grid)
  NextGrid=np.zeros(shape=(n,n,n))
  BigGrid=np.ones(shape=(n+2,n+2,n+2))*BoundryCondition
  BigGrid[1:n+1,1:n+1,1:n+1]=Grid
  for i in range(n):
    for j in range(n):
      for k in range(n):
        Diff =GetNeighborhoodDifference(Grid[i,j,k],BigGrid,i,j,k)
        if Diff >=0:
           NextGrid[i,j,k] = Grid[i,j,k] - Diff
        else:
           NextGrid[i,j,k] = 0


  return NextGrid


###############################################
################################################################################
def Diffuse(n,t,Initialize=0,BoundryCondition=0):
  if Initialize == 0:
    Grid = np.floor(np.ones(shape=(n,n,n))*10)
  elif Initialize ==1:
    Grid = np.floor(np.random.rand(n,n,n)*10)
  else:
    print 'add different initializations besides initilize =0'
    return
  print 'Running %dx%dx%d Simulation with %d TimeSteps---This Might Take a While' % (n,n,n,t)
  fig = plt.figure()
  ax = fig.add_subplot(111, projection='3d')
  for i in range(t):
    ax.cla()
    color = Plot(Grid,ax)
    if i == 0:
      color.set_clim(0, 10)
      plt.colorbar(color)
    fname = '_tmp%05d.png'%i
  #  print 'Saving frame', fname
    fig.savefig(fname)
    Grid=GetNextGrid(Grid,BoundryCondition)
  #plt.show()
  print 'Making movie animation.mpg - this make take a while'
  system("ffmpeg -y -qscale 1 -f image2 -r 5 -i _tmp%05d.png Animation.mkv")
  system("rm _tmp* ")

################################################################################
################################################################################
if __name__ == '__main__':
  if len(argv)<2:
    Diffuse(10,100)
  elif len(argv) == 2:
    if argv[1]=='-h' or argv[1]=='--help':
      Usage()
    else:
      Diffuse(atoi(argv[1]),100)
  elif len(argv) == 3:
    Diffuse(atoi(argv[1]),atoi(argv[2]))
  elif len(argv) == 4:
    Diffuse(atoi(argv[1]),atoi(argv[2]),atoi(argv[3]))
  elif len(argv) == 5:
    Diffuse(atoi(argv[1]),atoi(argv[2]),atoi(argv[3]),atoi(argv[4]))
  else:
    Usage()