Merge pull request #20 from bendudson/testing

More tests using pytest

freegs/test_geqdsk.py

@@ -0,0 +1,47 @@
+import numpy
+
+from io import StringIO
+
+from . import _geqdsk
+
+def test_writeread():
+    """
+    Test that data can be written then read back
+    """
+    nx = 65
+    ny = 65
+
+    # Create a dummy dataset
+    data = {"nx":nx, "ny":ny,
+            "rdim":2.0,
+            "zdim":1.5,
+            "rcentr":1.2,
+            "bcentr":2.42,
+            "rleft": 0.5,
+            "zmid":0.1,
+            "rmagx": 1.1,
+            "zmagx": 0.2,
+            "simagx": -2.3,
+            "sibdry": 0.21,
+            "cpasma": 1234521,
+            "fpol":numpy.random.rand(nx),
+            "pres":numpy.random.rand(nx),
+            "qpsi":numpy.random.rand(nx),
+            "psi":numpy.random.rand(nx,ny)
+    }
+
+    output = StringIO()
+
+    # Write to string
+    _geqdsk.write(data, output)
+
+    # Move to the beginning of the buffer
+    output.seek(0)
+
+    # Read from string
+    data2 = _geqdsk.read(output)
+
+    # Check that data and data2 are the same
+    for key in data:
+        numpy.testing.assert_allclose(data2[key], data[key])
+

freegs/test_linearsolve.py

@@ -0,0 +1,73 @@
+"""
+Tests of the linear solver
+"""
+
+from . import multigrid
+
+import numpy as np
+
+###### Laplacian in 2D
+
+def test_direct_laplacian():
+    nx = 65
+    ny = 65
+
+    Lx = 2.
+    Ly = 3.
+
+    dx = Lx/(nx - 1)
+    dy = Ly/(ny - 1)
+
+    xx, yy = np.meshgrid(np.linspace(0,Lx,nx), np.linspace(0,Ly,ny))
+
+    solution = np.exp( - ( (xx - 0.5)**2 + (yy - 0.5)**2 ) / 0.4**2 )
+
+    A = multigrid.LaplacianOp()
+    rhs = A(solution, dx, dy)
+
+    # Copy boundaries
+    rhs[:,0] = solution[:,0]
+    rhs[0,:] = solution[0,:]
+    rhs[:,-1] = solution[:,-1]
+    rhs[-1,:] = solution[-1,:]
+
+    solve = multigrid.MGDirect(multigrid.LaplaceSparse(Lx, Ly)(nx, ny))
+    x = solve(None, rhs)
+
+    assert np.allclose(x, solution)
+
+
+def test_multigrid_laplacian():
+    nx = 65
+    ny = 65
+
+    Lx = 2.
+    Ly = 3.
+
+    dx = Lx/(nx - 1)
+    dy = Ly/(ny - 1)
+
+    xx, yy = np.meshgrid(np.linspace(0,Lx,nx), np.linspace(0,Ly,ny))
+
+    solution = np.exp( - ( (xx - 0.5)**2 + (yy - 0.5)**2 ) / 0.4**2 )
+
+    A = multigrid.LaplacianOp()
+    rhs = A(solution, dx, dy)
+
+    # Copy boundaries
+    rhs[:,0] = solution[:,0]
+    rhs[0,:] = solution[0,:]
+    rhs[:,-1] = solution[:,-1]
+    rhs[-1,:] = solution[-1,:]
+
+    solve = multigrid.createVcycle(nx,ny, multigrid.LaplaceSparse(Lx, Ly),
+                                   ncycle=50, niter=50, nlevels=4, direct=True)
+
+    xinput = np.zeros(rhs.shape)
+    xinput[:,0] = solution[:,0]
+    xinput[0,:] = solution[0,:]
+    xinput[:,-1] = solution[:,-1]
+    xinput[-1,:] = solution[-1,:]
+    x = solve(xinput, rhs)
+
+    assert np.allclose(x, solution, atol=1e-6)