Update plotting routines to be usable with all backends and document them

examples/amoc.py

@@ -1,6 +1,4 @@
 import numpy
-import pickle
-
 import matplotlib.pyplot as plt
 
 from transiflow import Continuation
@@ -9,17 +7,10 @@
 from transiflow import utils
 
 
-class Data:
-    def __init__(self):
-        self.mu = []
-        self.value = []
-
-    def append(self, mu, value):
-        self.mu.append(mu)
-        self.value.append(value)
-
-    def callback(self, interface, x, mu):
-        self.append(mu, numpy.max(utils.compute_streamfunction(x, interface)))
+def postprocess(data, interface, x, mu):
+    data['Freshwater Flux'].append(mu)
+    data['Stream Function Maximum'].append(
+        numpy.max(utils.compute_streamfunction(x, interface)))
 
 
 def generate_plots(interface, x, sigma):
@@ -80,30 +71,29 @@ def main():
 
     generate_plots(interface, x1, 0)
 
-    # Enable the lines below to load the solution instead. Same for the ones below
+    # Enable the line below to load the solution instead. Same for the ones below
     # x1 = interface.load_state('x1')
 
     # Perform a continuation to freshwater flux 0.2 without detecting bifurcation points
     # and use this in the bifurcation diagram
-    data2 = Data()
+    data2 = {'Freshwater Flux': [], 'Stream Function Maximum': []}
+    callback = lambda interface, x, mu: postprocess(data2, interface, x, mu)
 
     ds = 0.05
     target = 0.2
     x2, mu2 = continuation.continuation(x1, 'Freshwater Flux', 0, target,
-                                        ds, ds_min=1e-12, callback=data2.callback)
+                                        ds, ds_min=1e-12, callback=callback)
 
     # Write the solution to a file
     interface.save_state('x2', x2)
 
-    generate_plots(interface, x2, mu2)
-
     # Write the data to a file
-    with open('data2', 'wb') as f:
-        pickle.dump(data2, f)
+    interface.save_json('data2.json', data2)
 
-    # Enable the lines below to load the data
-    # with open('data2', 'rb') as f:
-    #     data2 = pickle.load(f)
+    generate_plots(interface, x2, mu2)
+
+    # Enable the line below to load the data
+    # data2 = interface.load_json('data2.json')
 
     # Add asymmetry to the problem
     ds = 0.05
@@ -120,7 +110,7 @@ def main():
     # Go back to the symmetric problem
     ds = -0.05
     target = 0
-    x5, mu5 = continuation.continuation(x4, 'Asymmetry Parameter', mu3, target, ds)
+    x5, mu5 = continuation.continuation(x4, 'Asymmetry Parameter', mu3, target, ds, ds_min=1e-12)
 
     # Write the solution to a file
     interface.save_state('x5', x5)
@@ -132,25 +122,29 @@ def main():
 
     # Now compute the stable branch after the pitchfork bifurcation by going backwards
     # and use this in the bifurcation diagram
-    data6 = Data()
+    data6 = {'Freshwater Flux': [], 'Stream Function Maximum': []}
+    callback = lambda interface, x, mu: postprocess(data6, interface, x, mu)
 
     ds = -0.01
     target = 0.2
     x6, mu6 = continuation.continuation(x5, 'Freshwater Flux', mu4, target,
                                         ds, ds_min=1e-12, ds_max=0.005,
-                                        callback=data6.callback)
+                                        callback=callback)
 
     # Write the solution to a file
     interface.save_state('x6', x6)
 
+    # Write the data to a file
+    interface.save_json('data6.json', data6)
+
     generate_plots(interface, x6, mu6)
 
     # Plot a bifurcation diagram
     plt.title(f'Bifurcation diagram for the AMOC model with $n_x={nx}$, $n_y={ny}$')
     plt.xlabel('$\\sigma$')
     plt.ylabel('Maximum value of the streamfunction')
-    plt.plot(data2.mu, data2.value)
-    plt.plot(data6.mu, data6.value)
+    plt.plot(data2['Freshwater Flux'], data2['Stream Function Maximum'])
+    plt.plot(data6['Freshwater Flux'], data6['Stream Function Maximum'])
     plt.savefig('bifurcation_diagram.eps')
     plt.close()
 

examples/dhc.py

@@ -6,17 +6,11 @@
 from transiflow import utils
 
 
-class Data:
-    def __init__(self):
-        self.mu = []
-        self.value = []
+def postprocess(data, interface, x, mu):
+    data['Rayleigh Number'].append(mu)
+    data['Volume Averaged Kinetic Energy'].append(
+        utils.compute_volume_averaged_kinetic_energy(x, interface))
 
-    def append(self, mu, value):
-        self.mu.append(mu)
-        self.value.append(value)
-
-    def callback(self, interface, x, mu):
-        self.append(mu, utils.compute_volume_averaged_kinetic_energy(x, interface))
 
 def main():
     ''' An example of performing a continuation for a 2D differentially heated cavity and detecting a bifurcation point'''
@@ -37,20 +31,21 @@ def main():
                   'Verbose': False}
 
     interface = Interface(parameters, nx, ny, nz)
-
     continuation = Continuation(interface, newton_tolerance=1e-7)
 
     # Compute an initial guess
     x0 = interface.vector()
 
     # Store data for computing the bifurcation diagram using postprocessing
-    data = Data()
+    data = {'Rayleigh Number': [], 'Volume Averaged Kinetic Energy': []}
+    callback = lambda interface, x, mu: postprocess(data, interface, x, mu)
 
-    # Perform an initial continuation to Rayleigh number 1e8 without detecting bifurcation points
+    # Perform an initial continuation to Rayleigh number 1e8 without
+    # detecting bifurcation points
     ds = 100
     target = 1e9
     x, mu = continuation.continuation(x0, 'Rayleigh Number', 0, target,
-                                      ds, ds_max=1e8, callback=data.callback)
+                                      ds, ds_max=1e8, callback=callback)
 
     parameters['Eigenvalue Solver'] = {}
     parameters['Eigenvalue Solver']['Target'] = 0
@@ -61,18 +56,18 @@ def main():
     ds = 1e8
     target = 3.5e9
     x2, mu2 = continuation.continuation(x, 'Rayleigh Number', mu, target, ds, ds_max=1e8,
-                                        detect_bifurcations=True, callback=data.callback)
+                                        detect_bifurcations=True, callback=callback)
 
     ke = utils.compute_volume_averaged_kinetic_energy(x2, interface)
 
     # Compute the unstable branch after the bifurcation
     continuation = Continuation(interface, newton_tolerance=1e-4)
     x3, mu3 = continuation.continuation(x2, 'Rayleigh Number', mu2, target,
-                                        ds, ds_max=1e8, callback=data.callback)
+                                        ds, ds_max=1e8, callback=callback)
 
     # Plot a bifurcation diagram
     bif = plt.scatter(mu2, ke, marker='^')
-    plt.plot(data.mu, data.value)
+    plt.plot(data['Rayleigh Number'], data['Volume Averaged Kinetic Energy'])
 
     plt.title('Bifurcation diagram for the differentially heated cavity with $n_x=n_z={}$'.format(nx))
     plt.xlabel('Rayleigh number')

examples/ldc.py

@@ -6,17 +6,10 @@
 from transiflow import utils
 
 
-class Data:
-    def __init__(self):
-        self.mu = []
-        self.value = []
-
-    def append(self, mu, value):
-        self.mu.append(mu)
-        self.value.append(value)
-
-    def callback(self, interface, x, mu):
-        self.append(mu, utils.compute_volume_averaged_kinetic_energy(x, interface))
+def postprocess(data, interface, x, mu):
+    data['Reynolds Number'].append(mu)
+    data['Volume Averaged Kinetic Energy'].append(
+        utils.compute_volume_averaged_kinetic_energy(x, interface))
 
 
 def main():
@@ -43,13 +36,14 @@ def main():
     x0 = continuation.continuation(x0, 'Lid Velocity', 0, 1, 1)[0]
 
     # Store data for computing the bifurcation diagram using postprocessing
-    data = Data()
+    data = {'Reynolds Number': [], 'Volume Averaged Kinetic Energy': []}
+    callback = lambda interface, x, mu: postprocess(data, interface, x, mu)
 
     # Perform an initial continuation to Reynolds number 7000 without detecting bifurcation points
     ds = 100
     target = 6000
     x, mu = continuation.continuation(x0, 'Reynolds Number', 0, target, ds,
-                                      callback=data.callback)
+                                      callback=callback)
 
     # Now detect the bifurcation point
     parameters['Eigenvalue Solver'] = {}
@@ -62,18 +56,18 @@ def main():
     target = 10000
     x2, mu2 = bifurcation_continuation.continuation(x, 'Reynolds Number', mu, target,
                                                     ds, ds_max=100, detect_bifurcations=True,
-                                                    callback=data.callback)
+                                                    callback=callback)
 
     ke = utils.compute_volume_averaged_kinetic_energy(x2, interface)
 
     # Compute the unstable branch after the bifurcation
     target = 10000
     x3, mu3 = continuation.continuation(x2, 'Reynolds Number', mu2, target, ds,
-                                        callback=data.callback)
+                                        callback=callback)
 
     # Plot a bifurcation diagram
     bif = plt.scatter(mu2, ke, marker='^')
-    plt.plot(data.mu, data.value)
+    plt.plot(data['Reynolds Number'], data['Volume Averaged Kinetic Energy'])
 
     plt.title('Bifurcation diagram for the lid-driven cavity with $n_x=n_z={}$'.format(nx))
     plt.xlabel('Reynolds number')

examples/ldc3.py

@@ -7,17 +7,10 @@
 from transiflow import utils
 
 
-class Data:
-    def __init__(self):
-        self.t = []
-        self.value = []
-
-    def append(self, t, value):
-        self.t.append(t)
-        self.value.append(value)
-
-    def callback(self, interface, x, t):
-        self.append(t, utils.compute_volume_averaged_kinetic_energy(x, interface))
+def postprocess(data, interface, x, t):
+    data['t'].append(t)
+    data['Volume Averaged Kinetic Energy'].append(
+        utils.compute_volume_averaged_kinetic_energy(x, interface))
 
 
 def main():
@@ -39,7 +32,8 @@ def main():
     interface = Interface(parameters, nx, ny)
 
     # Store data for computing the bifurcation diagram using postprocessing
-    data = Data()
+    data = {'t': [], 'Volume Averaged Kinetic Energy': []}
+    callback = lambda interface, x, t: postprocess(data, interface, x, t)
 
     n = interface.discretization.dof * nx * ny
     x = numpy.random.random(n)
@@ -50,14 +44,14 @@ def main():
     for mu in range(0, 100, 10):
         interface.set_parameter('Reynolds Number', mu)
         time_integration = TimeIntegration(interface)
-        x, t = time_integration.integration(x, 1, 10, data.callback)
+        x, t = time_integration.integration(x, 1, 10, callback)
 
         # Plot the traced value during the time integration
         # plt.plot(data.t, data.value)
         # plt.show()
 
         mu_list.append(mu)
-        value_list.append(data.value[-1])
+        value_list.append(data['Volume Averaged Kinetic Energy'][-1])
 
     # Plot a bifurcation diagram
     plt.plot(mu_list, value_list)

examples/ldc_3d.py

@@ -1,45 +1,22 @@
-import pickle
-
 from transiflow import Continuation
 from transiflow import Interface
 from transiflow import utils
 
 
-class Data:
-    def __init__(self):
-        self.mu = []
-        self.value = []
-
-    def append(self, mu, value):
-        self.mu.append(mu)
-        self.value.append(value)
-
-
-def compute_volume_averaged_kinetic_energy(x_local, interface):
-    # Because HYMLS uses local subdomains, we need a different interface for postprocessing purposes
-    # that operates on the global domain
-    postprocess_interface = Interface(interface.parameters,
-                                      interface.nx_global, interface.ny_global, interface.nz_global,
-                                      interface.discretization.dim, interface.discretization.dof)
-
-    return utils.compute_volume_averaged_kinetic_energy(x_local.array, postprocess_interface)
-
-
 def postprocess(data, interface, x, mu, enable_output):
     if not enable_output:
         return
 
-    nx = interface.nx_global
-    ny = interface.ny_global
-    nz = interface.nz_global
+    nx = interface.nx
+    ny = interface.ny
+    nz = interface.nz
 
-    x_local = x.gather()
-    if interface.comm.MyPID() == 0:
-        # Store data for a bifurcation diagram at every continuation step
-        data.append(mu, utils.compute_volume_averaged_kinetic_energy(x_local, interface))
+    # Store data for a bifurcation diagram at every continuation step
+    data['Reynolds Number'].append(mu)
+    data['Volume Averaged Kinetic Energy'].append(
+        utils.compute_volume_averaged_kinetic_energy(x, interface))
 
-        with open('ldc_bif_' + str(nx) + '_' + str(ny) + '_' + str(nz) + '.obj', 'wb') as f:
-            pickle.dump(data, f)
+    interface.save_json('ldc_bif_' + str(nx) + '_' + str(ny) + '_' + str(nz) + '.json', data)
 
     # Store the solution at every continuation step
     interface.save_state(str(mu), x)
@@ -75,25 +52,27 @@ def main():
     parameters['Solver']['Iterative Solver']['Num Blocks'] = 100
     parameters['Solver']['Iterative Solver']['Convergence Tolerance'] = 1e-6
 
-    # Use one level in the HYMLS preconditioner. More levels means a less accurate factorization,
-    # but more parallelism and a smaller linear system at the coarsest level.
+    # Use one level in the HYMLS preconditioner. More levels means a
+    # less accurate factorization, but more parallelism and a smaller
+    # linear system at the coarsest level.
     parameters['Preconditioner'] = {}
     parameters['Preconditioner']['Number of Levels'] = 1
 
-    # Define a HYMLS interface that handles everything that is different when using HYMLS+Trilinos
-    # instead of NumPy as computational backend
+    # Define a HYMLS interface that handles everything that is
+    # different when using HYMLS+Trilinos instead of SciPy as
+    # computational backend
     interface = Interface(parameters, nx, ny, nz, backend='HYMLS')
-
-    data = Data()
-    callback = lambda interface, x, mu: postprocess(data, interface, x, mu, enable_output)
-
     continuation = Continuation(interface)
 
     # Compute an initial guess
     x0 = interface.vector()
-    x = continuation.continuation(x0, 'Lid Velocity', 0, 1, 1, callback=callback)[0]
+    x = continuation.continuation(x0, 'Lid Velocity', 0, 1, 1)[0]
+
+    # Perform an initial continuation to Reynolds number 1700 without
+    # detecting bifurcation points
+    data = {'Reynolds Number': [], 'Volume Averaged Kinetic Energy': []}
+    callback = lambda interface, x, mu: postprocess(data, interface, x, mu, enable_output)
 
-    # Perform an initial continuation to Reynolds number 1700 without detecting bifurcation points
     ds = 100
     target = 1800
     x, mu = continuation.continuation(x, 'Reynolds Number', 0, target,