diff --git a/asQ/allatonce/function.py b/asQ/allatonce/function.py index 1f99778d..02b09232 100644 --- a/asQ/allatonce/function.py +++ b/asQ/allatonce/function.py @@ -4,10 +4,11 @@ from functools import reduce from operator import mul import contextlib +from ufl.duals import is_primal, is_dual from asQ.profiling import profiler from asQ.allatonce.mixin import TimePartitionMixin -__all__ = ['time_average', 'AllAtOnceFunction'] +__all__ = ['time_average', 'AllAtOnceFunction', 'AllAtOnceCofunction'] @profiler() @@ -42,11 +43,11 @@ def time_average(aaofunc, uout, uwrk, average='window'): return -class AllAtOnceFunction(TimePartitionMixin): +class AllAtOnceFunctionBase(TimePartitionMixin): @profiler() def __init__(self, ensemble, time_partition, function_space): """ - A function representing multiple timesteps of a time-dependent finite-element problem, + A (co)function representing multiple timesteps of a time-dependent finite-element problem, i.e. the solution to an all-at-once system. :arg ensemble: time-parallel ensemble communicator. The timesteps are partitioned @@ -54,7 +55,8 @@ def __init__(self, ensemble, time_partition, function_space): ensemble.ensemble_comm.size == len(time_partition) must be True. :arg time_partition: a list of integers for the number of timesteps stored on each ensemble rank. - :arg function_space: a FunctionSpace for the solution at a single timestep. + :arg function_space: a Space for the a single timestep. + Either `FunctionSpace` or `DualSpace` depending if the child is AAO(Co)Function. """ self._time_partition_setup(ensemble, time_partition) @@ -67,20 +69,25 @@ def __init__(self, ensemble, time_partition, function_space): self.ncomponents = len(self.field_function_space.subfunctions) - self.function = fd.Function(self.function_space) - self.initial_condition = fd.Function(self.field_function_space) + # this will be renamed either self.function or self.cofunction + self._fbuf = fd.Function(self.function_space) # Functions to view each timestep def field_function(i): - dats = (self.function.subfunctions[j].dat - for j in self._component_indices(i)) + if self.ncomponents == 1: + j = self._component_indices(i)[0] + dat = self._fbuf.subfunctions[j].dat + else: + dat = MixedDat((self._fbuf.subfunctions[j].dat + for j in self._component_indices(i))) + return fd.Function(self.field_function_space, - val=MixedDat(dats)) + val=dat) self._fields = tuple(field_function(i) for i in range(self.nlocal_timesteps)) - # functions containing the last step of the previous + # (co)functions containing the last step of the previous # and current slice for parallel communication self.uprev = fd.Function(self.field_function_space) self.unext = fd.Function(self.field_function_space) @@ -88,7 +95,7 @@ def field_function(i): self.nlocal_dofs = self.function_space.node_set.size self.nglobal_dofs = self.ntimesteps*self.field_function_space.dim() - with self.function.dat.vec as fvec: + with self._fbuf.dat.vec as fvec: sizes = (self.nlocal_dofs, self.nglobal_dofs) self._vec = PETSc.Vec().createWithArray(fvec.array, size=sizes, @@ -204,8 +211,8 @@ def copy(self, copy_values=True): :arg copy_values: If true, the values of the current AllAtOnceFunction will be copied into the new AllAtOnceFunction. """ - new = AllAtOnceFunction(self.ensemble, self.time_partition, - self.field_function_space) + new = type(self)(self.ensemble, self.time_partition, + self.field_function_space) if copy_values: new.assign(self) return new @@ -227,34 +234,34 @@ def assign(self, src, update_halos=True, blocking=True): whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - if isinstance(src, AllAtOnceFunction): - dst_funcs = [self.function, self.initial_condition] - src_funcs = [src.function, src.initial_condition] - # these buffers just will be overwritten if the halos are updated - if not update_halos: - dst_funcs.extend([self.uprev, self.unext]) - src_funcs.extend([src.uprev, src.unext]) - for dst, src in zip(dst_funcs, src_funcs): - dst.assign(src) - + def func_assign(x, y): + return y.assign(x) + + def vec_assign(x, y): + x.copy(y) + + if isinstance(src, type(self)): + return self._vs_op(src, func_assign, vec_assign, + update_ics=True, + update_halos=update_halos, + blocking=blocking) + + # TODO: We should be able to use _vs_op here too but + # test_allatoncesolver:::test_solve_heat_equation + # fails if we do. The only difference is that + # _vs_op accesses the global vec with read/write + # access instead of write only. + # It isn't clear why this makes a difference (it + # shouldn't). elif isinstance(src, PETSc.Vec): with self.global_vec_wo() as gvec: src.copy(gvec) - elif isinstance(src, fd.Function): - if src.function_space() == self.field_function_space: - for i in range(self.nlocal_timesteps): - self[i].assign(src) - self.initial_condition.assign(src) - if not update_halos: - self.uprev.assign(src) - self.unext.assign(src) - elif src.function_space() == self.function_space: - self.function.assign(src) - else: - raise ValueError(f"src must be be in the `function_space` {self.function_space}" - + " or `field_function_space` {self.field_function_space} of the" - + " the AllAtOnceFunction, not in {src.function_space}") + elif isinstance(src, type(self._fbuf)): + return self._vs_op(src, func_assign, vec_assign, + update_ics=True, + update_halos=update_halos, + blocking=blocking) else: raise TypeError(f"src value must be AllAtOnceFunction or PETSc.Vec or field Function, not {type(src)}") @@ -270,7 +277,9 @@ def zero(self, subset=None): :arg subset: pyop2.types.set.Subset indicating the nodes to zero. If None then the whole function is zeroed. """ - funcs = (self.initial_condition, self.function, self.uprev, self.unext) + funcs = [self._fbuf, self.uprev, self.unext] + if hasattr(self, 'initial_condition'): + funcs.append(self.initial_condition) for f in funcs: f.zero(subset=subset) return self @@ -289,12 +298,10 @@ def scale(self, a, update_ics=False, whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - alpha = fd.Constant(a) - - self.function.assign(alpha*self.function) + self._fbuf.assign(a*self._fbuf) - if update_ics: - self.initial_condition.assign(alpha*self.initial_condition) + if update_ics and hasattr(self, 'initial_condition'): + self.initial_condition.assign(a*self.initial_condition) if update_halos: return self.update_time_halos(blocking=blocking) @@ -320,10 +327,8 @@ def axpy(self, a, x, update_ics=False, whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - alpha = fd.Constant(a) - def func_axpy(x, y): - return y.assign(alpha*x + y) + return y.assign(a*x + y) def vec_axpy(x, y): y.axpy(a, x) @@ -354,10 +359,8 @@ def aypx(self, a, x, update_ics=False, whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - alpha = fd.Constant(a) - def func_aypx(x, y): - return y.assign(x + alpha*y) + return y.assign(x + a*y) def vec_aypx(x, y): y.aypx(a, x) @@ -389,11 +392,8 @@ def axpby(self, a, b, x, update_ics=False, whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - alpha = fd.Constant(a) - beta = fd.Constant(b) - def func_axpby(x, y): - return y.assign(alpha*x + beta*y) + return y.assign(a*x + b*y) def vec_axpby(x, y): y.axpby(a, b, x) @@ -425,24 +425,24 @@ def _vs_op(self, x, func_op, vec_op, update_ics=False, whether blocking communication is used. A list of MPI Requests is returned if non-blocking communication is used. """ - if isinstance(x, AllAtOnceFunction): - func_op(x.function, self.function) - if update_ics: + if isinstance(x, type(self)): + func_op(x._fbuf, self._fbuf) + if update_ics and hasattr(self, 'initial_condition'): func_op(x.initial_condition, self.initial_condition) elif isinstance(x, PETSc.Vec): with self.global_vec() as gvec: vec_op(x, gvec) - elif isinstance(x, fd.Function): + elif isinstance(x, type(self._fbuf)): if x.function_space() == self.field_function_space: for i in range(self.nlocal_timesteps): func_op(x, self[i]) - if update_ics: + if update_ics and hasattr(self, 'initial_condition'): func_op(x, self.initial_condition) elif x.function_space() == self.function_space: - func_op(x, self.function) + func_op(x, self._fbuf) else: raise ValueError(f"x must be be in the `function_space` {self.function_space}" @@ -450,7 +450,7 @@ def _vs_op(self, x, func_op, vec_op, update_ics=False, + f" the AllAtOnceFunction, not in {x.function_space}") else: - raise TypeError(f"x value must be AllAtOnceFunction or PETSc.Vec or field Function, not {type(x)}") + raise TypeError(f"x value must be AllAtOnce(Co)Function or PETSc.Vec or field (Co)Function, not {type(x)}") if update_halos: return self.update_time_halos(blocking=blocking) @@ -466,7 +466,7 @@ def global_vec(self): # fvec shares the same storage as _vec, so we need this context # manager to make sure that the data gets copied to/from the # Function.dat storage and _vec. - with self.function.dat.vec: + with self._fbuf.dat.vec: self._vec.stateIncrease() yield self._vec @@ -481,7 +481,7 @@ def global_vec_ro(self): # fvec shares the same storage as _vec, so we need this context # manager to make sure that the data gets copied into _vec from # the Function.dat storage. - with self.function.dat.vec_ro: + with self._fbuf.dat.vec_ro: self._vec.stateIncrease() yield self._vec @@ -496,5 +496,124 @@ def global_vec_wo(self): # fvec shares the same storage as _vec, so we need this context # manager to make sure that the data gets copied back into the # Function.dat storage from _vec. - with self.function.dat.vec_wo: + with self._fbuf.dat.vec_wo: yield self._vec + + +class AllAtOnceFunction(AllAtOnceFunctionBase): + @profiler() + def __init__(self, ensemble, time_partition, function_space): + """ + A function representing multiple timesteps of a time-dependent finite-element problem, + i.e. the solution to an all-at-once system. + + :arg ensemble: time-parallel ensemble communicator. The timesteps are partitioned + over the ensemble members according to time_partition so + ensemble.ensemble_comm.size == len(time_partition) must be True. + :arg time_partition: a list of integers for the number of timesteps stored on each + ensemble rank. + :arg function_space: a FunctionSpace for the solution at a single timestep. + """ + if not is_primal(function_space): + raise TypeError("Cannot only make AllAtOnceFunction from a FunctionSpace") + super().__init__(ensemble, time_partition, function_space) + self.function = self._fbuf + self.initial_condition = fd.Function(self.field_function_space) + + +class AllAtOnceCofunction(AllAtOnceFunctionBase): + @profiler() + def __init__(self, ensemble, time_partition, function_space): + """ + A Cofunction representing multiple timesteps of a time-dependent finite-element problem, + i.e. the solution to an all-at-once system. + + :arg ensemble: time-parallel ensemble communicator. The timesteps are partitioned + over the ensemble members according to time_partition so + ensemble.ensemble_comm.size == len(time_partition) must be True. + :arg time_partition: a list of integers for the number of timesteps stored on each + ensemble rank. + :arg function_space: a FunctionSpace for the solution at a single timestep. + """ + if not is_dual(function_space): + raise TypeError("Can only make an AllAtOnceCofunction from a DualSpace") + super().__init__(ensemble, time_partition, function_space) + self.cofunction = self._fbuf + + @profiler() + def scale(self, a, update_halos=False, blocking=True): + """ + Scale the AllAtOnceCofunction by a scalar. + + :arg a: scalar to multiply the function by. + :arg update_halos: if True then the time-halos will be updated. + :arg blocking: if update_halos is True, then this argument determines + whether blocking communication is used. A list of MPI Requests is returned + if non-blocking communication is used. + """ + return super().scale(a, update_halos=update_halos, blocking=blocking, + update_ics=False) + + @profiler() + def axpy(self, a, x, update_halos=False, blocking=True): + """ + Compute y = a*x + y where y is this AllAtOnceCofunction. + + :arg a: scalar to multiply x. + :arg x: other object for calculation. Can be one of: + - AllAtOnceCofunction: all timesteps are updated, and optionally the ics. + - PETSc Vec: all timesteps are updated. + - firedrake.Cofunction in self.function_space: + all timesteps are updated. + - firedrake.Cofunction in self.field_function_space: + all timesteps are updated, and optionally the ics. + :arg update_halos: if True then the time-halos will be updated. + :arg blocking: if update_halos is True, then this argument determines + whether blocking communication is used. A list of MPI Requests is returned + if non-blocking communication is used. + """ + return super().axpy(a, x, update_halos=update_halos, blocking=blocking, + update_ics=False) + + @profiler() + def aypx(self, a, x, update_halos=False, blocking=True): + """ + Compute y = x + a*y where y is this AllAtOnceCofunction. + + :arg a: scalar to multiply y. + :arg x: other object for calculation. Can be one of: + - AllAtOnceCofunction: all timesteps are updated, and optionally the ics. + - PETSc Vec: all timesteps are updated. + - firedrake.Cofunction in self.function_space: + all timesteps are updated. + - firedrake.Cofunction in self.field_function_space: + all timesteps are updated, and optionally the ics. + :arg update_halos: if True then the time-halos will be updated. + :arg blocking: if update_halos is True, then this argument determines + whether blocking communication is used. A list of MPI Requests is returned + if non-blocking communication is used. + """ + return super().aypx(a, x, update_halos=update_halos, blocking=blocking, + update_ics=False) + + @profiler() + def axpby(self, a, b, x, update_halos=False, blocking=True): + """ + Compute y = a*x + b*y where y is this AllAtOnceCofunction. + + :arg a: scalar to multiply x. + :arg b: scalar to multiply y. + :arg x: other object for calculation. Can be one of: + - AllAtOnceFunction: all timesteps are updated, and optionally the ics. + - PETSc Vec: all timesteps are updated. + - firedrake.Cofunction in self.function_space: + all timesteps are updated. + - firedrake.Cofunction in self.field_function_space: + all timesteps are updated, and optionally the ics. + :arg update_halos: if True then the time-halos will be updated. + :arg blocking: if update_halos is True, then this argument determines + whether blocking communication is used. A list of MPI Requests is returned + if non-blocking communication is used. + """ + return super().axpby(a, b, x, update_halos=update_halos, blocking=blocking, + update_ics=False) diff --git a/asQ/allatonce/solver.py b/asQ/allatonce/solver.py index 990b39ee..0da75a2d 100644 --- a/asQ/allatonce/solver.py +++ b/asQ/allatonce/solver.py @@ -1,7 +1,7 @@ from firedrake.petsc import PETSc, OptionsManager, flatten_parameters from asQ.profiling import profiler -from asQ.allatonce import AllAtOnceJacobian +from asQ.allatonce import AllAtOnceJacobian, AllAtOnceCofunction from asQ.allatonce.mixin import TimePartitionMixin __all__ = ['AllAtOnceSolver'] @@ -135,5 +135,8 @@ def solve(self, rhs=None): if rhs is None: self.snes.solve(None, gvec) else: + if not isinstance(rhs, AllAtOnceCofunction): + msg = f"Right hand side of all-at-once problem must be AllAtOnceCofunction not {type(rhs)}." + raise TypeError(msg) with rhs.global_vec_ro() as rvec: self.snes.solve(rvec, gvec) diff --git a/tests/allatonce/test_allatoncefunction.py b/tests/allatonce/test_allatoncefunction.py index 4098e72f..580680ed 100644 --- a/tests/allatonce/test_allatoncefunction.py +++ b/tests/allatonce/test_allatoncefunction.py @@ -4,6 +4,7 @@ import numpy as np from functools import reduce from operator import mul +from ufl.duals import is_primal, is_dual def random_func(f): @@ -12,16 +13,47 @@ def random_func(f): def random_aaof(aaof): - random_func(aaof.initial_condition) - random_func(aaof.function) + if hasattr(aaof, "initial_condition"): + random_func(aaof.initial_condition) + random_func(function(aaof)) aaof.update_time_halos() +def function(aaof): + if isinstance(aaof, asQ.AllAtOnceFunction): + return aaof.function + elif isinstance(aaof, asQ.AllAtOnceCofunction): + return aaof.cofunction + else: + raise TypeError("This is only meant to be used with AllAtOnce{Function,Cofunction}") + + +def errornorm(u, uh): + # horrible hack just to check if two cofunctions have the same values + if is_primal(u): + return fd.errornorm(u, uh) + elif is_dual(u): + v = fd.Function(u.function_space().dual(), val=u.dat) + vh = fd.Function(uh.function_space().dual(), val=uh.dat) + return fd.errornorm(v, vh) + + +def assign(u, number): + # only use this to assign a numeric value + if is_primal(u): + return u.assign(number) + elif is_dual(u): + for dat in u.dat: + dat.data[:] = number + return u + + nprocs = 4 max_ncpts = 3 ncpts = [pytest.param(i, id=f"{i}component") for i in range(1, max_ncpts+1)] +function_type = ["AllAtOnceFunction", "AllAtOnceCofunction"] @pytest.fixture @@ -80,11 +112,17 @@ def W(request, V): return reduce(mul, [V for _ in range(request.param)]) -@pytest.fixture() -def aaof(ensemble, time_partition, W): +@pytest.fixture(params=function_type) +def aaof(request, ensemble, time_partition, W): if fd.COMM_WORLD.size == 1: return - return asQ.AllAtOnceFunction(ensemble, time_partition, W) + function_type = request.param + if function_type == "AllAtOnceFunction": + return asQ.AllAtOnceFunction(ensemble, time_partition, W) + elif function_type == "AllAtOnceCofunction": + return asQ.AllAtOnceCofunction(ensemble, time_partition, W.dual()) + else: + raise ValueError("Unrecognised all-at-once function type") @pytest.mark.parallel(nprocs=nprocs) @@ -185,8 +223,8 @@ def test_subfunctions(aaof): aaof[i].assign(u) for c in range(aaof.ncomponents): - err = fd.errornorm(u.subfunctions[c], - aaof.function.subfunctions[cpt_idx]) + err = errornorm(u.subfunctions[c], + function(aaof).subfunctions[cpt_idx]) assert (err < 1e-12) cpt_idx += 1 @@ -195,19 +233,16 @@ def test_subfunctions(aaof): cpt_idx = 0 for i in range(aaof.nlocal_timesteps): - norm = fd.norm(aaof[i]) - assert (norm < 1e-12) - for c in range(aaof.ncomponents): - aaof.function.subfunctions[cpt_idx].assign(cpt_idx) + function(aaof).subfunctions[cpt_idx].dat.data_wo[:] = cpt_idx - err = fd.errornorm(aaof[i].subfunctions[c], - aaof.function.subfunctions[cpt_idx]) + err = errornorm(aaof[i].subfunctions[c], + function(aaof).subfunctions[cpt_idx]) assert (err < 1e-12) cpt_idx += 1 if aaof.ncomponents == 1: - err = fd.errornorm(aaof[i], aaof.function.subfunctions[i]) + err = errornorm(aaof[i], function(aaof).subfunctions[i]) assert (err < 1e-12) @@ -222,14 +257,14 @@ def test_bcast_field(aaof): # solution at timestep i is i for slice_index in range(aaof.nlocal_timesteps): window_index = aaof.transform_index(slice_index, from_range='slice', to_range='window') - v.assign(window_index) + assign(v, window_index) aaof[slice_index].assign(v) for i in range(aaof.ntimesteps): - v.assign(-1) - u.assign(i) + assign(v, -1) + assign(u, i) aaof.bcast_field(i, v) - assert (fd.errornorm(v, u) < 1e-12) + assert (errornorm(v, u) < 1e-12) @pytest.mark.parallel(nprocs=nprocs) @@ -253,13 +288,14 @@ def test_copy(aaof): assert aaof.function_space == aaof1.function_space for step in range(aaof.nlocal_timesteps): - err = fd.errornorm(aaof[step], aaof1[step]) + err = errornorm(aaof[step], aaof1[step]) assert (err < 1e-12) - err = fd.errornorm(aaof.initial_condition, aaof1.initial_condition) + if hasattr(aaof, "initial_condition"): + err = errornorm(aaof.initial_condition, aaof1.initial_condition) assert (err < 1e-12) - err = fd.errornorm(aaof.uprev, aaof1.uprev) + err = errornorm(aaof.uprev, aaof1.uprev) assert (err < 1e-12) @@ -280,14 +316,15 @@ def test_assign(aaof): aaof1.assign(aaof) for step in range(aaof.nlocal_timesteps): - err = fd.errornorm(aaof[step], aaof1[step]) + err = errornorm(aaof[step], aaof1[step]) assert (err < 1e-12) - err = fd.errornorm(aaof.initial_condition, - aaof1.initial_condition) + if hasattr(aaof, "initial_condition"): + err = errornorm(aaof.initial_condition, + aaof1.initial_condition) assert (err < 1e-12) - err = fd.errornorm(aaof.uprev, aaof1.uprev) + err = errornorm(aaof.uprev, aaof1.uprev) assert (err < 1e-12) # set from PETSc Vec @@ -296,10 +333,10 @@ def test_assign(aaof): aaof1.assign(gvec0) for step in range(aaof.nlocal_timesteps): - err = fd.errornorm(aaof[step], aaof1[step]) + err = errornorm(aaof[step], aaof1[step]) assert (err < 1e-12) - err = fd.errornorm(aaof.uprev, aaof1.uprev) + err = errornorm(aaof.uprev, aaof1.uprev) assert (err < 1e-12) # set from field function @@ -309,15 +346,15 @@ def test_assign(aaof): aaof.assign(v0) for step in range(aaof.nlocal_timesteps): - err = fd.errornorm(v0, aaof[step]) + err = errornorm(v0, aaof[step]) assert (err < 1e-12) # set from allatonce.function random_aaof(aaof1) - aaof.assign(aaof1.function) + aaof.assign(function(aaof1)) for step in range(aaof.nlocal_timesteps): - err = fd.errornorm(aaof[step], aaof1[step]) + err = errornorm(aaof[step], aaof1[step]) assert (err < 1e-12) @@ -332,9 +369,10 @@ def test_axpy(aaof): aaof1 = aaof.copy(copy_values=False) def check_close(x, y): - err = fd.errornorm(x.function, y.function) + err = errornorm(function(x), function(y)) assert (err < 1e-12) - err = fd.errornorm(x.initial_condition, y.initial_condition) + if hasattr(aaof, "initial_condition"): + err = errornorm(x.initial_condition, y.initial_condition) assert (err < 1e-12) def faxpy(result, a, x, y): @@ -349,8 +387,9 @@ def faxpy(result, a, x, y): # x is aaofunc a = 2 aaof.axpy(a, aaof1) - faxpy(expected.function, a, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpy(function(expected), a, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -358,16 +397,17 @@ def faxpy(result, a, x, y): aaof.assign(orig) # x is aaofunc and update ics - a = 3.5 - aaof.axpy(a, aaof1, update_ics=True) - faxpy(expected.function, a, aaof1.function, orig.function) - faxpy(expected.initial_condition, a, - aaof1.initial_condition, orig.initial_condition) + if hasattr(aaof, "initial_condition"): + a = 3.5 + aaof.axpy(a, aaof1, update_ics=True) + faxpy(function(expected), a, function(aaof1), function(orig)) + faxpy(expected.initial_condition, a, + aaof1.initial_condition, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is PETSc.Vec a = 4.2 @@ -377,8 +417,9 @@ def faxpy(result, a, x, y): aaof.axpy(a, xvec) - faxpy(expected.function, a, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpy(function(expected), a, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -393,7 +434,8 @@ def faxpy(result, a, x, y): aaof.axpy(a, xfunc) for i in range(aaof.nlocal_timesteps): faxpy(expected[i], a, xfunc, orig[i]) - expected.initial_condition.assign(orig.initial_condition) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -401,19 +443,20 @@ def faxpy(result, a, x, y): aaof.assign(orig) # x is a single timestep and update ics - a = 6.1 - random_func(xfunc) + if hasattr(aaof, "initial_condition"): + a = 6.1 + random_func(xfunc) - aaof.axpy(a, xfunc, update_ics=True) - for i in range(aaof.nlocal_timesteps): - faxpy(expected[i], a, xfunc, orig[i]) - faxpy(expected.initial_condition, a, - xfunc, orig.initial_condition) + aaof.axpy(a, xfunc, update_ics=True) + for i in range(aaof.nlocal_timesteps): + faxpy(expected[i], a, xfunc, orig[i]) + faxpy(expected.initial_condition, a, + xfunc, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is a timeseries function a = 7.9 @@ -421,8 +464,9 @@ def faxpy(result, a, x, y): random_func(tsfunc) aaof.axpy(a, tsfunc) - faxpy(expected.function, a, tsfunc, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpy(function(expected), a, tsfunc, function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -438,10 +482,11 @@ def test_aypx(aaof): aaof1 = aaof.copy(copy_values=False) def check_close(x, y): - err = fd.errornorm(x.function, y.function) - assert (err < 1e-12) - err = fd.errornorm(x.initial_condition, y.initial_condition) + err = errornorm(function(x), function(y)) assert (err < 1e-12) + if hasattr(aaof, "initial_condition"): + err = errornorm(x.initial_condition, y.initial_condition) + assert (err < 1e-12) def faypx(result, a, x, y): result.assign(x + a*y) @@ -455,8 +500,9 @@ def faypx(result, a, x, y): # x is aaofunc a = 2 aaof.aypx(a, aaof1) - faypx(expected.function, a, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faypx(function(expected), a, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -464,16 +510,17 @@ def faypx(result, a, x, y): aaof.assign(orig) # x is aaofunc and update ics - a = 3.5 - aaof.aypx(a, aaof1, update_ics=True) - faypx(expected.function, a, aaof1.function, orig.function) - faypx(expected.initial_condition, a, - aaof1.initial_condition, orig.initial_condition) + if hasattr(aaof, "initial_condition"): + a = 3.5 + aaof.aypx(a, aaof1, update_ics=True) + faypx(function(expected), a, function(aaof1), function(orig)) + faypx(expected.initial_condition, a, + aaof1.initial_condition, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is PETSc.Vec a = 4.2 @@ -483,8 +530,9 @@ def faypx(result, a, x, y): aaof.aypx(a, xvec) - faypx(expected.function, a, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faypx(function(expected), a, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -499,7 +547,8 @@ def faypx(result, a, x, y): aaof.aypx(a, xfunc) for i in range(aaof.nlocal_timesteps): faypx(expected[i], a, xfunc, orig[i]) - expected.initial_condition.assign(orig.initial_condition) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -507,19 +556,20 @@ def faypx(result, a, x, y): aaof.assign(orig) # x is a single timestep and update ics - a = 6.1 - random_func(xfunc) + if hasattr(aaof, "initial_condition"): + a = 6.1 + random_func(xfunc) - aaof.aypx(a, xfunc, update_ics=True) - for i in range(aaof.nlocal_timesteps): - faypx(expected[i], a, xfunc, orig[i]) - faypx(expected.initial_condition, a, - xfunc, orig.initial_condition) + aaof.aypx(a, xfunc, update_ics=True) + for i in range(aaof.nlocal_timesteps): + faypx(expected[i], a, xfunc, orig[i]) + faypx(expected.initial_condition, a, + xfunc, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is a timeseries function a = 7.9 @@ -527,8 +577,9 @@ def faypx(result, a, x, y): random_func(tsfunc) aaof.aypx(a, tsfunc) - faypx(expected.function, a, tsfunc, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faypx(function(expected), a, tsfunc, function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -544,10 +595,11 @@ def test_axpby(aaof): aaof1 = aaof.copy(copy_values=False) def check_close(x, y): - err = fd.errornorm(x.function, y.function) - assert (err < 1e-12) - err = fd.errornorm(x.initial_condition, y.initial_condition) + err = errornorm(function(x), function(y)) assert (err < 1e-12) + if hasattr(aaof, "initial_condition"): + err = errornorm(x.initial_condition, y.initial_condition) + assert (err < 1e-12) def faxpby(result, a, b, x, y): result.assign(a*x + b*y) @@ -562,8 +614,9 @@ def faxpby(result, a, b, x, y): a = 2 b = 11.4 aaof.axpby(a, b, aaof1) - faxpby(expected.function, a, b, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpby(function(expected), a, b, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -571,17 +624,18 @@ def faxpby(result, a, b, x, y): aaof.assign(orig) # x is aaofunc and update ics - a = 3.5 - b = 12.1 - aaof.axpby(a, b, aaof1, update_ics=True) - faxpby(expected.function, a, b, aaof1.function, orig.function) - faxpby(expected.initial_condition, a, b, - aaof1.initial_condition, orig.initial_condition) + if hasattr(aaof, "initial_condition"): + a = 3.5 + b = 12.1 + aaof.axpby(a, b, aaof1, update_ics=True) + faxpby(function(expected), a, b, function(aaof1), function(orig)) + faxpby(expected.initial_condition, a, b, + aaof1.initial_condition, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is PETSc.Vec a = 4.2 @@ -592,8 +646,9 @@ def faxpby(result, a, b, x, y): aaof.axpby(a, b, xvec) - faxpby(expected.function, a, b, aaof1.function, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpby(function(expected), a, b, function(aaof1), function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -609,7 +664,8 @@ def faxpby(result, a, b, x, y): aaof.axpby(a, b, xfunc) for i in range(aaof.nlocal_timesteps): faxpby(expected[i], a, b, xfunc, orig[i]) - expected.initial_condition.assign(orig.initial_condition) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -617,20 +673,21 @@ def faxpby(result, a, b, x, y): aaof.assign(orig) # x is a single timestep and update ics - a = 6.1 - b = 15.3 - random_func(xfunc) + if hasattr(aaof, "initial_condition"): + a = 6.1 + b = 15.3 + random_func(xfunc) - aaof.axpby(a, b, xfunc, update_ics=True) - for i in range(aaof.nlocal_timesteps): - faxpby(expected[i], a, b, xfunc, orig[i]) - faxpby(expected.initial_condition, a, b, - xfunc, orig.initial_condition) + aaof.axpby(a, b, xfunc, update_ics=True) + for i in range(aaof.nlocal_timesteps): + faxpby(expected[i], a, b, xfunc, orig[i]) + faxpby(expected.initial_condition, a, b, + xfunc, orig.initial_condition) - check_close(aaof, expected) + check_close(aaof, expected) - # reset - aaof.assign(orig) + # reset + aaof.assign(orig) # x is a timeseries function a = 7.9 @@ -639,8 +696,9 @@ def faxpby(result, a, b, x, y): random_func(tsfunc) aaof.axpby(a, b, tsfunc) - faxpby(expected.function, a, b, tsfunc, orig.function) - expected.initial_condition.assign(orig.initial_condition) + faxpby(function(expected), a, b, tsfunc, function(orig)) + if hasattr(aaof, "initial_condition"): + expected.initial_condition.assign(orig.initial_condition) check_close(aaof, expected) @@ -657,18 +715,19 @@ def test_zero(aaof): random_aaof(aaof) aaof.zero() - norm = fd.norm(aaof.initial_condition) - assert (norm < 1e-12) + if hasattr(aaof, "initial_condition"): + norm = fd.norm(aaof.initial_condition) + assert (norm < 1e-12) - norm = fd.norm(aaof.uprev) - assert (norm < 1e-12) + for dat in aaof.uprev.dat: + assert np.allclose(dat.data, 0) - norm = fd.norm(aaof.unext) - assert (norm < 1e-12) + for dat in aaof.unext.dat: + assert np.allclose(dat.data, 0) for step in range(aaof.nlocal_timesteps): - norm = fd.norm(aaof[step]) - assert (norm < 1e-12) + for dat in aaof[step].dat: + assert np.allclose(dat.data, 0) @pytest.mark.parallel(nprocs=nprocs) @@ -676,17 +735,21 @@ def test_global_vec(aaof): """ test synchronising the global Vec with the local Functions """ - v = fd.Function(aaof.field_function_space, name="v") + vcheck = fd.Function(aaof.field_function_space, name="v") + + u = fd.Function(aaof.field_function_space, name="u") def all_equal(func, val): - v.assign(val) + assign(vcheck, val) for step in range(func.nlocal_timesteps): - err = fd.errornorm(v, func[step]) + err = errornorm(vcheck, func[step]) assert (err < 1e-12) # read only - aaof.initial_condition.assign(10) - aaof.assign(aaof.initial_condition) + assign(u, 10) + if hasattr(aaof, "initial_condition"): + aaof.initial_condition.assign(u) + aaof.assign(u) with aaof.global_vec_ro() as rvec: assert np.allclose(rvec.array, 10) @@ -695,8 +758,10 @@ def all_equal(func, val): all_equal(aaof, 10) # write only - aaof.initial_condition.assign(30) - aaof.assign(aaof.initial_condition) + assign(u, 30) + if hasattr(aaof, "initial_condition"): + aaof.initial_condition.assign(u) + aaof.assign(u) with aaof.global_vec_wo() as wvec: assert np.allclose(wvec.array, 20) @@ -704,8 +769,10 @@ def all_equal(func, val): all_equal(aaof, 40) - aaof.initial_condition.assign(50) - aaof.assign(aaof.initial_condition) + assign(u, 50) + if hasattr(aaof, "initial_condition"): + aaof.initial_condition.assign(u) + aaof.assign(u) with aaof.global_vec() as vec: assert np.allclose(vec.array, 50) @@ -770,13 +837,13 @@ def test_update_time_halos(aaof): size = aaof.ensemble.ensemble_comm.size # solution on this slice - v0.assign(rank) + assign(v0, rank) # solution on previous slice - v1.assign((rank - 1) % size) + assign(v1, (rank - 1) % size) # set last field from each slice aaof[-1].assign(v0) aaof.update_time_halos() - assert (fd.errornorm(aaof.uprev, v1) < 1e-12) + assert (errornorm(aaof.uprev, v1) < 1e-12)