Improve tasks example

hysop_examples/tasks/tasks.py

@@ -15,6 +15,10 @@ def compute(args):
     from hysop.constants import HYSOP_DEFAULT_TASK_ID, Implementation
     from hysop.operators import CustomOperator, Dummy
     has_tasks = not args.proc_tasks is None
+    try:
+        has_tasks_overlapping = any(len(_) > 1 for _ in args.proc_tasks)
+    except TypeError:
+        has_tasks_overlapping = False
 
     # Define domain
     npts = args.npts
@@ -26,7 +30,10 @@ def compute(args):
     A = Field(domain=box, name='A', dtype=args.dtype)
     B = Field(domain=box, name='B', dtype=args.dtype)
     C = Field(domain=box, name='C', dtype=args.dtype)
-
+    if has_tasks_overlapping:
+        Ap = Field(domain=box, name='Ap', dtype=args.dtype)
+    else:
+        Ap = A
     # Build the mpi_params for each task (keep in mind tha some procs can be defined on both tasks)
     if has_tasks:
         mpi_params = {TASK_A: None, TASK_B: None}
@@ -89,13 +96,17 @@ def compute(args):
 #         `--B--,              // Inter-task communication step
 #            |  op3 B,C->A
 #         ,--A--'              // Inter-task communication step
-# endA       |
+# endA  A->A |
 
 # Note : without endA operator, the second communication is not
 # automatically inserted because A field is an output for task A
 # and cannot be invalidated by op3 on othre task. Inter-task
 # invalidation in graph building is not yet implemented
 
+# Note this algorithm is not working with overlapping tasks
+# This is due to the fact that A is both input and output in taskA.
+# A workaround is to use an other field A' and endA is then a copy
+
     custom1 = CustomOperator(name='custom_op1',
                              func=custom_func1, invars=(A, ), outvars=(A,),
                              variables={A: npts, }, implementation=impl,
@@ -105,12 +116,21 @@ def compute(args):
                              variables={A: npts, B: npts}, implementation=impl,
                              **_get_op_kwds(TASK_A))
     custom3 = CustomOperator(name='custom_op3',
-                             func=custom_func3, invars=(B, C), outvars=(A,),
-                             variables={A: npts, B: npts, C: npts}, implementation=impl,
+                             func=custom_func3, invars=(B, C), outvars=(Ap,),
+                             variables={Ap: npts, B: npts, C: npts}, implementation=impl,
                              **_get_op_kwds(TASK_B))
-    endA = Dummy(name='endA',
-                 variables={A: npts, }, implementation=impl,
-                 **_get_op_kwds(TASK_A))
+    if has_tasks_overlapping:
+        def copy(_Ap, _A):
+            _A.data[0][...] = _Ap.data[0]
+            print(np.min(_A.data[0]), np.max(_A.data[0]))
+        endA = CustomOperator(name='endA',
+                              func=copy, invars=(Ap,), outvars=(A,),
+                              variables={Ap: npts, A: npts}, implementation=impl,
+                              **_get_op_kwds(TASK_A))
+    else:
+        endA = Dummy(name='endA',
+                     variables={A: npts, }, implementation=impl,
+                     **_get_op_kwds(TASK_A))
 
     # Create the problem we want to solve and insert our operator
     problem = Problem()
@@ -135,7 +155,7 @@ def compute(args):
         problem.initialize_field(C, formula=init_C)
     problem.solve(simu, dry_run=args.dry_run)
 
-    for dA in A.discrete_fields.values():
+    for dA in A.discrete_fields.values()+Ap.discrete_fields.values():
         if not np.allclose(dA.data[0], 1.):
             print np.min(dA.data[0]), np.max(dA.data[0])
         else: