Add test for work arrays

hysop/tools/tests/test_work_arrays.py

+"""Test different ways to handle internal memory buffers.
+
+"""
+from hysop.problem.simulation import Simulation
+from hysop.tools.parameters import Discretization
+import numpy as np
+from hysop import Field, Box
+from hysop.operators import Stretching, EnergyEnstrophy, Curl
+from hysop.tools.misc import WorkSpaceTools
+import hysop.tools.numpywrappers as npw
+
+cos = np.cos
+sin = np.sin
+
+Nx = 32
+Ny = 106
+Nz = 64
+g = 2
+
+# Domain
+xdom = np.asarray([0., 0., 0.])
+ldom = np.asarray([1., ] * 3)
+box = Box(length=ldom, origin=xdom)
+
+# topo : different discretizations, different distributions
+d1 = Discretization([Nx + 1, Ny + 1, Nz + 1], [g, g, g])
+d2 = Discretization([Ny + 1, Nz + 1, Nx + 14])
+d3 = Discretization([Nz + 1, 2 * Nx + 1, Ny + 1], [g, g, g])
+topo1 = box.create_topology(d1, cutdir=[False, True, True])
+topo3 = box.create_topology(d3, cutdir=[True, True, True])
+
+
+# Create different operators and try to allocate one common workspace
+# for all their internal work arrays.
+def init():
+    """Build, init, setup for operator
+    """
+    # Fields
+    velo = Field(domain=box, is_vector=True, name='Velocity')
+    vorti = Field(domain=box, name='Vorticity', is_vector=True)
+
+    # operators, may work on different topologies
+    op = {}
+    op['stretching'] = Stretching(velo, vorti, discretization=topo1)
+    op['energy'] = EnergyEnstrophy(velo, vorti, discretization=d2)
+    op['curl'] = Curl(invar=velo, outvar=vorti, discretization=topo3)
+
+    for o in op:
+        op[o].discretize()
+    simu = Simulation(nb_iter=4)
+    simu.initialize()
+
+    return op, simu
+
+
+def test_default_work():
+    """Default case : each op has its own work (or no work at all)
+    """
+    op, simu = init()
+    for o in op:
+        op[o].setup()
+        op[o].apply(simu)
+
+    wkop = []
+    for o in op:
+        wkop.append(op[o].discrete_op._rwork)
+    wkref = wkop[0]
+    lref = len(wkref)
+    for w in [ww for ww in wkop[1:] if ww is not None]:
+        lw = len(w)
+        ll = min(lref, lw)
+        for i in xrange(ll):
+            assert not npw.arrays_share_data(wkref[i], w[i])
+
+
+def test_no_common_work():
+    """Each op has its own work (or no work at all), explicitely
+    allocated
+    """
+    op, simu = init()
+    for o in op:
+        wkp = op[o].get_work_properties()['rwork']
+        if wkp is not None:
+            rwork = [npw.zeros(wkp[i]) for i in xrange(len(wkp))]
+        else:
+            rwork = None
+        op[o].setup(rwork=rwork)
+        op[o].apply(simu)
+
+    wkop = []
+    for o in op:
+        wkop.append(op[o].discrete_op._rwork)
+    wkref = wkop[0]
+    lref = len(wkref)
+    for w in [ww for ww in wkop[1:] if ww is not None]:
+        lw = len(w)
+        ll = min(lref, lw)
+        for i in xrange(ll):
+            assert not npw.arrays_share_data(wkref[i], w[i])
+
+
+def test_common_work():
+    """One work for all op
+    """
+    op, simu = init()
+    wkp = []
+    for o in op:
+        wkp.append(op[o].get_work_properties()['rwork'])
+
+    rwork = WorkSpaceTools.allocate_common_workspaces(wkp)
+    for o in op:
+        op[o].setup(rwork=rwork)
+        op[o].apply(simu)
+
+    wkop = []
+    for o in op:
+        wkop.append(op[o].discrete_op._rwork)
+    wkref = wkop[0]
+    lref = len(wkref)
+    for w in [ww for ww in wkop[1:] if ww is not None]:
+        lw = len(w)
+        ll = min(lref, lw)
+        for i in xrange(ll):
+            assert npw.arrays_share_data(wkref[i], w[i])
+
+
+if __name__ == "__main__":
+    test_default_work()
+    test_no_common_work()
+    test_common_work()