added support for 1D/2D sims

Author: anabiman

MolecularDynamics/dump.py

@@ -30,49 +30,49 @@
 
 import numpy as np
 
-def writeOutput(filename, natoms, timestep, box, **args):
+def writeOutput(filename, natoms, timestep, box, **data):
     """ Writes the output (in dump format) """
 
     axis = ('x', 'y', 'z')
-
+    
     with open(filename, 'a') as fp:
- 
-         fp.write('ITEM: TIMESTEP\n')
-         fp.write('{}\n'.format(timestep))
-         fp.write('ITEM: NUMBER OF ATOMS\n')
-         fp.write('{}\n'.format(natoms))
-         fp.write('ITEM: BOX BOUNDS ff ff ff\n')
-
-         for box_low, box_high in box:
-            fp.write( '{} {}\n'.format(box_low,box_high) )
-  
-         keys = args.keys()
-
-         for key in keys:
-            dims = len(args[key].shape)
+        
+        fp.write('ITEM: TIMESTEP\n')
+        fp.write('{}\n'.format(timestep))
+        
+        fp.write('ITEM: NUMBER OF ATOMS\n')
+        fp.write('{}\n'.format(natoms))
+        
+        fp.write('ITEM: BOX BOUNDS' + ' f' * len(box) + '\n')
+        for box_bounds in box:
+            fp.write('{} {}\n'.format(*box_bounds))
 
-            if dims > 1:
-                for i in range(args[key].shape[1]):
+        for i in range(len(axis) - len(box)):
+            fp.write('0 0\n')
+            
+        keys = list(data.keys())
+        
+        for key in keys:
+            isMatrix = len(data[key].shape) > 1
+            
+            if isMatrix:
+                _, nCols = data[key].shape
+                
+                for i in range(nCols):
                     if key == 'pos':
-                        args[axis[i]] = args[key][:,i]
+                        data['{}'.format(axis[i])] = data[key][:,i]
                     else:
-                        args['{}_{}'.format(key, axis[i])] = args[key][:,i]
-
-                del args[key]
-
-         keys = args.keys()
-         fp.write('ITEM: ATOMS' + (' {}' * len(keys)).format(*keys) + '\n')
-         
-         output = []
-         
-         if args:
-             for key in keys:
-                 data = args[key]
-
-                 if len(output):
-                     output = np.vstack((output, data.T))
-                 else:
-                     output = data
-         
-         if len(output):
-             np.savetxt(fp, output.T)
+                        data['{}_{}'.format(key,axis[i])] = data[key][:,i]
+                        
+                del data[key]
+                
+        keys = data.keys()
+        
+        fp.write('ITEM: ATOMS' + (' {}' * len(data)).format(*data) + '\n')
+        
+        output = []
+        for key in keys:
+            output = np.hstack((output, data[key]))
+            
+        if len(output):
+            np.savetxt(fp, output.reshape((natoms, len(data)), order='F'))

MolecularDynamics/langevin.py

@@ -109,8 +109,8 @@ def run(**args):
 
     natoms, box, dt, temp = args['natoms'], args['box'], args['dt'], args['temp']
     mass, relax, nsteps   = args['mass'], args['relax'], args['steps']
-    ofname, radius, freq  = args['ofname'],  args['radius'], args['freq']
-
+    ofname, freq, radius = args['ofname'], args['freq'], args['radius']
+    
     dim = len(box)
     pos = np.random.rand(natoms,dim)
 
@@ -137,12 +137,12 @@ def run(**args):
         # Check if any particle has collided with the wall
         wallHitCheck(pos,vels,box)
 
-        # Compute output (avg kinetic energy)
-        ins_temp = np.sum(np.dot(mass, (vels - vels.mean(axis=0))**2)) / (Boltzmann * 3 * natoms)
+        # Compute output (temperature)
+        ins_temp = np.sum(np.dot(mass, (vels - vels.mean(axis=0))**2)) / (Boltzmann * dim * natoms)
         output.append([step * dt, ins_temp])
 
         if not step%freq:
-            dump.writeOutput(ofname, natoms, step, box, mass=mass, radius=radius, pos=pos, v=vels)
+            dump.writeOutput(ofname, natoms, step, box, radius=radius, pos=pos, v=vels)
 
     return np.array(output)
 
@@ -158,13 +158,12 @@ def run(**args):
         'steps': 10000,
         'freq': 100,
         'box': ((0, 1e-8), (0, 1e-8), (0, 1e-8)),
-        'ofname': 'traj.dump'
+        'ofname': 'traj-hydrogen-3D.dump'
         }
 
     output = run(**params)
 
     plt.plot(output[:,0] * 1e12, output[:,1])
     plt.xlabel('Time (ps)')
     plt.ylabel('Temp (K)')
-    plt.grid(linestyle=':')
     plt.show()