much faster Newton solver, still slow

Author: pancetta

pySDC/implementations/problem_classes/GrayScott_MPIFFT.py

@@ -1,4 +1,5 @@
 import numpy as np
+import scipy.sparse as sp
 from mpi4py import MPI
 from mpi4py_fft import PFFT
 
@@ -380,12 +381,23 @@ def solve_system_2(self, rhs, factor, u0, t):
             dg10 = -factor * (tmpv ** 2)
             dg11 = 1 - factor * (2 * tmpu * tmpv - self.params.B)
 
-            for ix, iy in np.ndindex(tmpu.shape):
-                a = np.array([tmpgu[ix, iy], tmpgv[ix, iy]])
-                dg = np.array([[dg00[ix, iy], dg01[ix, iy]], [dg10[ix, iy], dg11[ix, iy]]])
-                b = np.linalg.inv(dg).dot(a)
-                tmpu[ix, iy] -= b[0]
-                tmpv[ix, iy] -= b[1]
+            # interleave and unravel to put into sparse matrix
+            dg00I = np.ravel(np.kron(dg00, np.array([1, 0])))
+            dg01I = np.ravel(np.kron(dg01, np.array([1, 0])))
+            dg10I = np.ravel(np.kron(dg10, np.array([1, 0])))
+            dg11I = np.ravel(np.kron(dg11, np.array([0, 1])))
+
+            # put into sparse matrix
+            dg = sp.diags(dg00I, offsets=0) + sp.diags(dg11I, offsets=0)
+            dg += sp.diags(dg01I, offsets=1, shape=dg.shape) + sp.diags(dg10I, offsets=-1, shape=dg.shape)
+
+            # interleave g terms to apply inverse to it
+            g = np.kron(tmpgu.flatten(), np.array([1, 0])) + np.kron(tmpgv.flatten(), np.array([0, 1]))
+            # invert dg matrix
+            b = sp.linalg.spsolve(dg, g)
+            # update real space vectors
+            tmpu[:] -= b[::2].reshape(self.params.nvars)
+            tmpv[:] -= b[1::2].reshape(self.params.nvars)
 
             # increase iteration count
             n += 1
@@ -535,12 +547,23 @@ def solve_system_2(self, rhs, factor, u0, t):
             dg10 = -factor * (tmpv ** 2)
             dg11 = 1 - factor * (2 * tmpu * tmpv)
 
-            for ix, iy in np.ndindex(tmpu.shape):
-                a = np.array([tmpgu[ix, iy], tmpgv[ix, iy]])
-                dg = np.array([[dg00[ix, iy], dg01[ix, iy]], [dg10[ix, iy], dg11[ix, iy]]])
-                b = np.linalg.inv(dg).dot(a)
-                tmpu[ix, iy] -= b[0]
-                tmpv[ix, iy] -= b[1]
+            # interleave and unravel to put into sparse matrix
+            dg00I = np.ravel(np.kron(dg00, np.array([1, 0])))
+            dg01I = np.ravel(np.kron(dg01, np.array([1, 0])))
+            dg10I = np.ravel(np.kron(dg10, np.array([1, 0])))
+            dg11I = np.ravel(np.kron(dg11, np.array([0, 1])))
+
+            # put into sparse matrix
+            dg = sp.diags(dg00I, offsets=0) + sp.diags(dg11I, offsets=0)
+            dg += sp.diags(dg01I, offsets=1, shape=dg.shape) + sp.diags(dg10I, offsets=-1, shape=dg.shape)
+
+            # interleave g terms to apply inverse to it
+            g = np.kron(tmpgu.flatten(), np.array([1, 0])) + np.kron(tmpgv.flatten(), np.array([0, 1]))
+            # invert dg matrix
+            b = sp.linalg.spsolve(dg, g)
+            # update real-space vectors
+            tmpu[:] -= b[::2].reshape(self.params.nvars)
+            tmpv[:] -= b[1::2].reshape(self.params.nvars)
 
             # increase iteration count
             n += 1

pySDC/playgrounds/mpifft/grayscott.py

@@ -10,7 +10,6 @@
 from pySDC.implementations.sweeper_classes.multi_implicit import multi_implicit
 from pySDC.implementations.problem_classes.GrayScott_MPIFFT import grayscott_imex_diffusion, grayscott_imex_linear, \
     grayscott_mi_diffusion, grayscott_mi_linear
-# from pySDC.implementations.problem_classes.GrayScott_FFT import grayscott_imex_linear
 from pySDC.implementations.transfer_classes.TransferMesh_MPIFFT import fft_to_fft
 
 
@@ -30,7 +29,7 @@ def run_simulation(spectral=None, splitting_type=None, ml=None, num_procs=None):
     # initialize level parameters
     level_params = dict()
     level_params['restol'] = 1E-12
-    level_params['dt'] = 1E-00
+    level_params['dt'] = 8E-00
     level_params['nsweeps'] = [1]
     level_params['residual_type'] = 'last_abs'
 
@@ -57,12 +56,12 @@ def run_simulation(spectral=None, splitting_type=None, ml=None, num_procs=None):
     problem_params['Dv'] = 0.00001
     problem_params['A'] = 0.04
     problem_params['B'] = 0.1
-    problem_params['newton_maxiter'] = 100
+    problem_params['newton_maxiter'] = 50
     problem_params['newton_tol'] = 1E-11
 
     # initialize step parameters
     step_params = dict()
-    step_params['maxiter'] = 500
+    step_params['maxiter'] = 100
     step_params['errtol'] = 1E-09
 
     # initialize controller parameters
@@ -95,7 +94,7 @@ def run_simulation(spectral=None, splitting_type=None, ml=None, num_procs=None):
 
     # set time parameters
     t0 = 0.0
-    Tend = 4
+    Tend = 32
 
     f = None
     if rank == 0:
@@ -127,21 +126,21 @@ def run_simulation(spectral=None, splitting_type=None, ml=None, num_procs=None):
     # call main function to get things done...
     uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
 
-    plt.figure()
-    plt.imshow(P.fft.backward(uend[..., 0]))#, vmin=0, vmax=1)
-    # plt.imshow(np.fft.irfft2(uend.values[..., 0]))#, vmin=0, vmax=1)
-    plt.title('u')
-    plt.colorbar()
-    plt.figure()
-    plt.imshow(P.fft.backward(uend[..., 1]))#, vmin=0, vmax=1)
-    # plt.imshow(np.fft.irfft2(uend.values[..., 1]))#, vmin=0, vmax=1)
-    plt.title('v')
-    plt.colorbar()
     # plt.figure()
-    # plt.imshow(uend[..., 0] + uend[..., 1])
-    # plt.title('sum')
+    # plt.imshow(P.fft.backward(uend[..., 0]))#, vmin=0, vmax=1)
+    # # plt.imshow(np.fft.irfft2(uend.values[..., 0]))#, vmin=0, vmax=1)
+    # plt.title('u')
+    # plt.colorbar()
+    # plt.figure()
+    # plt.imshow(P.fft.backward(uend[..., 1]))#, vmin=0, vmax=1)
+    # # plt.imshow(np.fft.irfft2(uend.values[..., 1]))#, vmin=0, vmax=1)
+    # plt.title('v')
     # plt.colorbar()
-    plt.show()
+    # # plt.figure()
+    # # plt.imshow(uend[..., 0] + uend[..., 1])
+    # # plt.title('sum')
+    # # plt.colorbar()
+    # plt.show()
     # # exit()
 
     if rank == 0:
@@ -185,14 +184,14 @@ def main():
     """
     # run_simulation(spectral=False, splitting_type='diffusion', ml=False, num_procs=1)
     # run_simulation(spectral=True, splitting_type='diffusion', ml=False, num_procs=1)
-    run_simulation(spectral=True, splitting_type='linear', ml=False, num_procs=1)
+    # run_simulation(spectral=True, splitting_type='linear', ml=False, num_procs=1)
     # run_simulation(spectral=False, splitting_type='diffusion', ml=True, num_procs=1)
     # run_simulation(spectral=True, splitting_type='diffusion', ml=True, num_procs=1)
     # run_simulation(spectral=False, splitting_type='diffusion', ml=True, num_procs=10)
     # run_simulation(spectral=True, splitting_type='diffusion', ml=True, num_procs=10)
 
     # run_simulation(spectral=False, splitting_type='mi_diffusion', ml=False, num_procs=1)
-    # run_simulation(spectral=True, splitting_type='mi_diffusion', ml=False, num_procs=1)
+    run_simulation(spectral=True, splitting_type='mi_diffusion', ml=False, num_procs=1)
     # run_simulation(spectral=False, splitting_type='mi_linear', ml=False, num_procs=1)
     # run_simulation(spectral=True, splitting_type='mi_linear', ml=False, num_procs=1)