Restructuring heat solver

Author: jussienko

numpy/heat-equation/README.md

@@ -37,6 +37,7 @@ Note: Algorithm is stable only when
 Implement two dimensional heat equation with NumPy using the initial
 temperature field in the file [bottle.dat](bottle.dat) (the file consists of a
 header and 200 x 200 data array). As a boundary condition use fixed values as
-given in the initial field. You can start from the skeleton code in the file
-[skeleton.py](skeleton.py).
+given in the initial field. The main program that is provided in
+[heat_main.py](heat_main.py) can be used as such, implement the required 
+functionality in the module [heat.py](heat.py) (look for **TODO**s).
 

numpy/heat-equation/skeleton.py renamed to numpy/heat-equation/heat.py

@@ -8,23 +8,6 @@
 # Set the colormap
 plt.rcParams['image.cmap'] = 'BrBG'
 
-# Basic variables
-a = 0.5          # Diffusion constant.
-timesteps = 500  # Number of time-steps to evolve system.
-image_interval = 50 # write frequency for png files
-
-# Grid spacings
-dx = 0.01
-dy = 0.01
-dx2 = dx**2
-dy2 = dy**2
-
-# For stability, this is the largest interval possible
-# for the size of the time-step:
-dt = dx2*dy2 / ( 2*a*(dx2+dy2) )
-
-# TODO: Read the initial temperature field from file
-
 
 def evolve(u, u_previous, a, dt):
     """Explicit time evolution.
@@ -37,10 +20,20 @@ def evolve(u, u_previous, a, dt):
     # and the numerical Laplacian according the explicit time evolution method
 
 
-# Write figure of initial field to a file
-plt.imshow(field)
-plt.axis('off')
-plt.savefig('heat_{0:03d}.png'.format(0))
 
+def iterate(field, field0, a, dx, dy, timesteps, image_interval):
+    """Run fixed number of time steps of heat equation"""
 # TODO: Implement the main iteration loop and write the figure 
 # (to a new) file after each 'image_interval' iteration
+
+def init_fields(filename):
+# TODO: Read the initial temperature field from file
+# Create also a copy of the field for the previous time step
+
+    return field, field0
+
+def write_field(field, step):
+    plt.gca().clear()
+    plt.imshow(field)
+    plt.axis('off')
+    plt.savefig('heat_{0:03d}.png'.format(step))

numpy/heat-equation/heat_main.py

@@ -0,0 +1,55 @@
+from __future__ import print_function
+import time
+import argparse
+
+from heat import init_fields, write_field, iterate
+
+
+def main(input_file='bottle.dat', a=0.5, dx=0.1, dy=0.1, 
+         timesteps=200, image_interval=4000):
+
+    # Initialise the temperature field
+    field, field0 = init_fields(input_file)
+
+    print("Heat equation solver")
+    print("Diffusion constant: {}".format(a))
+    print("Input file: {}".format(input_file))
+    print("Parameters")
+    print("----------")
+    print("  nx={} ny={} dx={} dy={}".format(field.shape[0], field.shape[1],
+                                             dx, dy))
+    print("  time steps={}  image interval={}".format(timesteps,
+                                                         image_interval))
+
+    # Plot/save initial field
+    write_field(field, 0)
+    # Iterate
+    t0 = time.time()
+    iterate(field, field0, a, dx, dy, timesteps, image_interval)
+    t1 = time.time()
+    # Plot/save final field
+    write_field(field, timesteps)
+
+    print("Simulation finished in {0} s".format(t1-t0))
+
+if __name__ == '__main__':
+
+    # Process command line arguments
+    parser = argparse.ArgumentParser(description='Heat equation')
+    parser.add_argument('-dx', type=float, default=0.01,
+                        help='grid spacing in x-direction')
+    parser.add_argument('-dy', type=float, default=0.01,
+                        help='grid spacing in y-direction')
+    parser.add_argument('-a', type=float, default=0.5,
+                        help='diffusion constant')
+    parser.add_argument('-n', type=int, default=200,
+                        help='number of time steps')
+    parser.add_argument('-i', type=int, default=4000,
+                        help='image interval')
+    parser.add_argument('-f', type=str, default='bottle.dat', 
+                        help='input file')
+
+    args = parser.parse_args()
+
+    main(args.f, args.a, args.dx, args.dy, args.n, args.i)
+

numpy/heat-equation/solution/heat-serial.py renamed to numpy/heat-equation/solution/heat.py

@@ -1,29 +1,11 @@
-from __future__ import print_function
 import numpy as np
-import time
-
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 
 # Set the colormap
 plt.rcParams['image.cmap'] = 'BrBG'
 
-# Basic parameters
-a = 0.5                # Diffusion constant
-timesteps = 200        # Number of time-steps to evolve system
-image_interval = 4000  # Write frequency for png files
-
-# Grid spacings
-dx = 0.01
-dy = 0.01
-dx2 = dx**2
-dy2 = dy**2
-
-# For stability, this is the largest interval possible
-# for the size of the time-step:
-dt = dx2*dy2 / ( 2*a*(dx2+dy2) )
-
 def evolve(u, u_previous, a, dt, dx2, dy2):
     """Explicit time evolution.
        u:            new temperature field
@@ -40,6 +22,21 @@ def evolve(u, u_previous, a, dt, dx2, dy2):
                  u_previous[1:-1, :-2]) / dy2 )
     u_previous[:] = u[:]
 
+def iterate(field, field0, a, dx, dy, timesteps, image_interval):
+    """Run fixed number of time steps of heat equation"""
+
+    dx2 = dx**2
+    dy2 = dy**2
+
+    # For stability, this is the largest interval possible
+    # for the size of the time-step:
+    dt = dx2*dy2 / ( 2*a*(dx2+dy2) )    
+
+    for m in range(1, timesteps+1):
+        evolve(field, field0, a, dt, dx2, dy2)
+        if m % image_interval == 0:
+            write_field(field, m)
+
 def init_fields(filename):
     # Read the initial temperature field from file
     field = np.loadtxt(filename)
@@ -52,26 +49,4 @@ def write_field(field, step):
     plt.axis('off')
     plt.savefig('heat_{0:03d}.png'.format(step))
 
-def iterate(field, field0, timesteps, image_interval):
-    for m in range(1, timesteps+1):
-        evolve(field, field0, a, dt, dx2, dy2)
-        if m % image_interval == 0:
-            write_field(field, m)
-
-
-def main():
-    # Initialise the temperature field
-    field, field0 = init_fields('bottle.dat')
-    # Plot/save initial field
-    write_field(field, 0)
-    # Iterate
-    t0 = time.time()
-    iterate(field, field0, timesteps, image_interval)
-    t1 = time.time()
-    # Plot/save final field
-    write_field(field, timesteps)
-
-    print("Running time: {0}".format(t1-t0))
 
-if __name__ == '__main__':
-    main()