Commit for SOR solver with test cases by group III.

projects/project_5/sor.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+
+def poisson_solver_sor(rho, boxsize, omega=None, tol=1e-6, maxiter=2000, epsilon0=1.):
+    '''Solve the Poisson equation on 2D grid with periodic boundary conditions by SOR method
+    within a given tolerance. (Red-Black-Grid version)
+
+    Inputs:
+    rho (float matrix, shape (m, n)): charge distribution; m, n > 1
+    boxsize (float vector, > 0, len = 2): box lengths for x and y direction
+    omega (None or float between 0 and 2): relaxation factor, leave as None for auto selection
+    tol (float >= 0): tolerance between nearest two iteration to stop; 0 for reaching maxiter
+    maxiter (int > 0): maximum iteration number
+    epsilon0 (float > 0): electric permittivity
+
+    Outputs:
+    phi (float matrix, shape (m, n)): electrostatic potential on grid
+    error (float): final error
+    '''
+
+    # check the inputs
+    rho = np.array(rho)
+    if rho.ndim != 2 or np.shape(rho)[0] < 2 or np.shape(rho)[1] < 2:
+        raise ValueError('We need at least two grid points on each direction')
+    m, n = np.shape(rho)
+    if np.shape(boxsize) != (2,) or boxsize[0] <= 0 or boxsize[1] <= 0:
+        raise ValueError('We need correct boxsize')
+    if omega is not None and (omega <= 0 or omega >= 2):
+        raise ValueError('We need correct relaxation factor')
+    if tol < 0:
+        raise ValueError('We need correct tolerance')
+    if maxiter <= 0:
+        raise ValueError('We need correct iteration number')
+    if epsilon0 <= 0:
+        raise ValueError('We need correct electric permittivity')
+
+    # optimal omega
+    # ref: http://www.math.umbc.edu/~kogan/technical_papers/2007/Yang_Gobbert.pdf
+    if omega is None:
+        omega = 2 / (1 + np.sin(np.pi / max(m, n)))
+
+    # preparing the red and black grids
+    row_even, row_odd = np.arange(0, m, 2), np.arange(1, m, 2)
+    col_even, col_odd = np.arange(0, n, 2), np.arange(1, n, 2)
+    red_row = np.concatenate((np.tile(row_even, len(col_even)), np.tile(row_odd, len(col_odd))))
+    red_col = np.concatenate((np.repeat(col_even, len(row_even)), np.repeat(col_odd, len(row_odd))))
+    blk_row = np.concatenate((np.tile(row_even, len(col_odd)), np.tile(row_odd, len(col_even))))
+    blk_col = np.concatenate((np.repeat(col_odd, len(row_even)), np.repeat(col_even, len(row_odd))))
+
+    # calculate the coefficients
+    one_minus_omega = 1. - omega
+    hx2 = (boxsize[0] / m) ** 2
+    hy2 = (boxsize[1] / n) ** 2
+    coeff_x = hy2 / (2 * (hx2 + hy2))
+    coeff_y = hx2 / (2 * (hx2 + hy2))
+    coeff_rho = hx2 * hy2 / (2 * epsilon0 * (hx2 + hy2))
+    rho = rho * coeff_rho
+
+    # initiate variables for iteration
+    phi = np.zeros((m, n))
+    errors = np.zeros((m, n))
+    error = 1000.
+    newvalue = 0.
+    iter_num = 0
+
+    # define update function (vectorized)
+    def _update(i, j):
+        newvalue = one_minus_omega * phi[i, j] + \
+                    omega * (coeff_x * (phi[(i - 1) % m, j] + phi[(i + 1) % m, j]) + \
+                        coeff_y * (phi[i, (j - 1) % n] + phi[i, (j + 1) % n]) + \
+                            rho[i, j])
+        errors[i, j] = newvalue - phi[i, j]
+        phi[i, j] = newvalue
+    update = np.vectorize(_update, signature='(n),(m)->()')
+
+    # iteration
+    while error > tol and iter_num < maxiter:
+        # errors.fill(0.)
+        update(red_row, red_col)
+        update(blk_row, blk_col)
+        error = np.linalg.norm(errors)
+        # print("one iter finished.")
+        iter_num += 1
+
+    return phi, error

projects/project_5/test/test_sor.py

@@ -0,0 +1,33 @@
+import pytest
+import numpy as np
+from itertools import product
+from project_5.sor import poisson_solver_sor
+
+
+@pytest.mark.parametrize('rho, boxsize, omega, exception', [
+    ([[1, 2], [3, 4]], (1), 1.2, ValueError),
+    ([[-1, 2, 3], [1]], (1, 2), 0.5, ValueError),
+    ([[1]], (1, 2), 1.2, ValueError),
+    ([[-1, "hello"], [1, 2]], (1, 2), 0.5, TypeError),
+    ([[2], [3, 4]], (1, 2), 3, ValueError),
+    ([[-1, 2, 3], [1], [4, 5, 6]], (1, 2), 0.5, ValueError)])
+def test_poisson_solver_sor_exceptions(rho, boxsize, omega, exception):
+    with pytest.raises(exception):
+        poisson_solver_sor(rho, boxsize, omega)
+
+
+@pytest.mark.parametrize(
+    'nx,ny', [(nx, ny) for nx, ny in product(
+        [50, 100],
+        [50, 100])])
+def test_cossin(nx, ny):
+    """Test a known solution of Poisson's equation.
+    """
+    gx = np.linspace(-np.pi, np.pi, nx, endpoint=False)
+    gy = np.linspace(-np.pi, np.pi, ny, endpoint=False)
+    x, y = np.meshgrid(gx, gy, indexing='ij')
+    phi = np.cos(x) + np.sin(y)
+    np.testing.assert_allclose(
+        poisson_solver_sor(phi, (2 * np.pi, 2 * np.pi))[0],
+        phi,
+        rtol=1e-2, atol=5e-2)