def test_is_zero_simple_vector_expressions(): mesh = UnitSquareMesh(4, 4) V = FunctionSpace(mesh, 'CG', 1) v = TestFunction(V) u = TrialFunction(V) check_is_zero(dot(as_vector([Zero(), u]), as_vector([Zero(), v])), 1) check_is_zero(dot(as_vector([Zero(), u]), as_vector([v, Zero()])), 0)
def test_is_zero_simple_scalar_expressions(): mesh = UnitSquareMesh(4, 4) V = FunctionSpace(mesh, 'CG', 1) v = TestFunction(V) u = TrialFunction(V) check_is_zero(Zero() * u * v, 0) check_is_zero(Constant(0) * u * v, 1) check_is_zero(Zero() * v, 0) check_is_zero(Constant(0) * v, 1) check_is_zero(0 * u * v, 0) check_is_zero(0 * v, 0) check_is_zero(ufl.sin(0) * v, 0) check_is_zero(ufl.cos(0) * v, 1)
def argument(self, arg): """ Argument is UFL lingo for test (num=0) and trial (num=1) functions """ # Do not make several new args for the same original arg if arg in self._cache: return self._cache[arg] # Get some data about the argument and get our target index V = arg.function_space() N = V.num_sub_spaces() num = arg.number() idx_wanted = [self._index_v, self._index_u][num] new_arg = [] for idx in range(N): # Construct non-coupled argument Vi = V.sub(idx).collapse() a = dolfin.function.argument.Argument(Vi, num, part=arg.part()) indices = numpy.ndindex(a.ufl_shape) if idx == idx_wanted: # This is a wanted index new_arg.extend(a[I] for I in indices) else: # This index should be pruned new_arg.extend(Zero() for _ in indices) new_arg = as_vector(new_arg) self._cache[arg] = new_arg return new_arg
def construct_modified_terminal(mt, terminal): """Construct a modified terminal given terminal modifiers from an analysed modified terminal and a terminal.""" expr = terminal if mt.reference_value: expr = ReferenceValue(expr) dim = expr.ufl_domain().topological_dimension() for n in range(mt.local_derivatives): # Return zero if expression is trivially constant. This has to # happen here because ReferenceGrad has no access to the # topological dimension of a literal zero. if is_cellwise_constant(expr): expr = Zero(expr.ufl_shape + (dim, ), expr.ufl_free_indices, expr.ufl_index_dimensions) else: expr = ReferenceGrad(expr) # No need to apply restrictions to ConstantValue terminals if not isinstance(expr, ConstantValue): if mt.restriction == '+': expr = PositiveRestricted(expr) elif mt.restriction == '-': expr = NegativeRestricted(expr) return expr
def argument(self, obj): Q = obj.ufl_function_space() dom = Q.ufl_domain() sub_elements = obj.ufl_element().sub_elements() # If not a mixed element, do nothing if not isinstance(obj.ufl_element(), MixedElement): return obj # Split into sub-elements, creating appropriate space for each args = [] for i, sub_elem in enumerate(sub_elements): Q_i = FunctionSpace(dom, sub_elem) a = Argument(Q_i, obj.number(), part=obj.part()) indices = [()] for m in a.ufl_shape: indices = [(k + (j, )) for k in indices for j in range(m)] if (i == self.idx[obj.number()]): args += [a[j] for j in indices] else: args += [Zero() for j in indices] return as_vector(args)
def zero(self, o): fi = o.ufl_free_indices fid = o.ufl_index_dimensions mapped_fi = tuple(self.index(Index(count=i)) for i in fi) paired_fid = [(mapped_fi[pos], fid[pos]) for pos, a in enumerate(fi)] new_fi, new_fid = zip(*tuple(sorted(paired_fid))) return Zero(o.ufl_shape, new_fi, new_fid)
def modified_terminal(self, o): mt = analyse_modified_terminal(o) t = mt.terminal r = mt.restriction if isinstance(t, Argument) and r != self.restrictions[t.number()]: return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions) else: return o
def coefficient(self, o): # Define dw/dw := d/ds [w + s v] = v # Return corresponding argument if we can find o among w do = self._w2v.get(o) if do is not None: return do # Look for o among coefficient derivatives dos = self._cd.get(o) if dos is None: # If o is not among coefficient derivatives, return # do/dw=0 do = Zero(o.ufl_shape) return do else: # Compute do/dw_j = do/dw_h : v. # Since we may actually have a tuple of oprimes and vs in a # 'mixed' space, sum over them all to get the complete inner # product. Using indices to define a non-compound inner product. # Example: # (f:g) -> (dfdu:v):g + f:(dgdu:v) # shape(dfdu) == shape(f) + shape(v) # shape(f) == shape(g) == shape(dfdu : v) # Make sure we have a tuple to match the self._v tuple if not isinstance(dos, tuple): dos = (dos, ) if len(dos) != len(self._v): error( "Got a tuple of arguments, expecting a matching tuple of coefficient derivatives." ) dosum = Zero(o.ufl_shape) for do, v in zip(dos, self._v): so, oi = as_scalar(do) rv = len(v.ufl_shape) oi1 = oi[:-rv] oi2 = oi[-rv:] prod = so * v[oi2] if oi1: dosum += as_tensor(prod, oi1) else: dosum += prod return dosum
def conditional(self, f): # TODO: Is this right? What about non-scalars? c, a, b = f.operands() s = f.shape() ufl_assert(s == (), "TODO: Assuming scalar valued expressions.") _0 = Zero() _1 = IntValue(1) da = conditional(c, _1, _0) db = conditional(c, _0, _1) return (None, da, db)
def coefficient_derivative(self, o, expr, coefficients, arguments, cds): # If we're only taking a derivative wrt part of an argument in # a mixed space other bits might come back as zero. We want to # propagate a zero in that case. argument, = arguments if all(isinstance(a, Zero) for a in argument.ufl_operands): return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions) else: return self.reuse_if_untouched(o, expr, coefficients, arguments, cds)
def argument(self, o): from ufl import split from firedrake import MixedFunctionSpace, FunctionSpace V = o.function_space() if len(V) == 1: # Not on a mixed space, just return ourselves. return o if o in self._arg_cache: return self._arg_cache[o] V_is = V.split() indices = self.blocks[o.number()] try: indices = tuple(indices) nidx = len(indices) except TypeError: # Only one index provided. indices = (indices, ) nidx = 1 if nidx == 1: W = V_is[indices[0]] W = FunctionSpace(W.mesh(), W.ufl_element()) a = (Argument(W, o.number(), part=o.part()), ) else: W = MixedFunctionSpace([V_is[i] for i in indices]) a = split(Argument(W, o.number(), part=o.part())) args = [] for i in range(len(V_is)): if i in indices: c = indices.index(i) a_ = a[c] if len(a_.ufl_shape) == 0: args += [a_] else: args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)] else: args += [ Zero() for j in numpy.ndindex(V_is[i].ufl_element().value_shape()) ] return self._arg_cache.setdefault(o, as_vector(args))
def getFormStage(F, butch, u0, t, dt, bcs=None, splitting=None, nullspace=None): """Given a time-dependent variational form and a :class:`ButcherTableau`, produce UFL for the s-stage RK method. :arg F: UFL form for the semidiscrete ODE/DAE :arg butch: the :class:`ButcherTableau` for the RK method being used to advance in time. :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg t: a :class:`Constant` referring to the current time level. Any explicit time-dependence in F is included :arg dt: a :class:`Constant` referring to the size of the current time step. :arg splitting: a callable that maps the (floating point) Butcher matrix a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used to vary between the classical RK formulation and Butcher's reformulation that leads to a denser mass matrix with block-diagonal stiffness. Only `AI` and `IA` are currently supported. :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possible time-dependent) boundary conditions imposed on the system. :arg nullspace: A list of tuples of the form (index, VSB) where index is an index into the function space associated with `u` and VSB is a :class: `firedrake.VectorSpaceBasis` instance to be passed to a `firedrake.MixedVectorSpaceBasis` over the larger space associated with the Runge-Kutta method On output, we return a tuple consisting of several parts: - Fnew, the :class:`Form` - possibly a 4-tuple containing information needed to solve a mass matrix to update the solution (this is empty for RadauIIA methods for which there is a trivial update function. - UU, the :class:`firedrake.Function` holding all the stage time values. It lives in a :class:`firedrake.FunctionSpace` corresponding to the s-way tensor product of the space on which the semidiscrete form lives. - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the stages, - 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object that represents the nullspace of the coupled system - `gblah`, a list of tuples of the form (f, expr, method), where f is a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`. At each time step, each expr needs to be re-interpolated/projected onto the corresponding f in order for Firedrake to pick up that time-dependent boundary conditions need to be re-applied. The interpolation/projection is encoded in method, which is either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending on whether the function space for f supports interpolation or not. """ v = F.arguments()[0] V = v.function_space() assert V == u0.function_space() num_stages = butch.num_stages num_fields = len(V) # s-way product space for the stage variables Vbig = reduce(mul, (V for _ in range(num_stages))) VV = TestFunction(Vbig) UU = Function(Vbig) vecconst = np.vectorize(Constant) C = vecconst(butch.c) A = vecconst(butch.A) # set up the pieces we need to work with to do our substitutions nsxnf = (num_stages, num_fields) if num_fields == 1: u0bits = np.array([u0], dtype="O") vbits = np.array([v], dtype="O") if num_stages == 1: # single-stage method VVbits = np.array([[VV]], dtype="O") UUbits = np.array([[UU]], dtype="O") else: # multi-stage methods VVbits = np.reshape(split(VV), nsxnf) UUbits = np.reshape(split(UU), nsxnf) else: u0bits = np.array(list(split(u0)), dtype="O") vbits = np.array(list(split(v)), dtype="O") VVbits = np.reshape(split(VV), nsxnf) UUbits = np.reshape(split(UU), nsxnf) split_form = extract_terms(F) Fnew = Zero() # first, process terms with a time derivative. I'm # assuming we have something of the form inner(Dt(g(u0)), v)*dx # For each stage i, this gets replaced with # inner((g(stages[i]) - g(u0))/dt, v)*dx # but we have to carefully handle the cases where g indexes into # pieces of u dtless = strip_dt_form(split_form.time) if splitting is None or splitting == AI: for i in range(num_stages): repl = {t: t + C[i] * dt} for j in range(num_fields): repl[u0bits[j]] = UUbits[i][j] - u0bits[j] repl[vbits[j]] = VVbits[i][j] # Also get replacements right for indexing. for j in range(num_fields): for ii in np.ndindex(u0bits[j].ufl_shape): repl[u0bits[j][ii]] = UUbits[i][j][ii] - u0bits[j][ii] repl[vbits[j][ii]] = VVbits[i][j][ii] Fnew += replace(dtless, repl) # Now for the non-time derivative parts for i in range(num_stages): # replace test function repl = {} for k in range(num_fields): repl[vbits[k]] = VVbits[i][k] for ii in np.ndindex(vbits[k].ufl_shape): repl[vbits[k][ii]] = VVbits[i][k][ii] Ftmp = replace(split_form.remainder, repl) # replace the solution with stage values for j in range(num_stages): repl = {t: t + C[j] * dt} for k in range(num_fields): repl[u0bits[k]] = UUbits[j][k] for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[j][k][ii] # and sum the contribution Fnew += A[i, j] * dt * replace(Ftmp, repl) elif splitting == IA: Ainv = np.vectorize(Constant)(np.linalg.inv(butch.A)) # time derivative part gets inverse of Butcher matrix. for i in range(num_stages): repl = {} for k in range(num_fields): repl[vbits[k]] = VVbits[i][k] for ii in np.ndindex(vbits[k].ufl_shape): repl[vbits[k][ii]] = VVbits[i][k][ii] Ftmp = replace(dtless, repl) for j in range(num_stages): repl = {t: t + C[j] * dt} for k in range(num_fields): repl[u0bits[k]] = (UUbits[j][k] - u0bits[k]) for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[j][k][ii] - u0bits[k][ii] Fnew += Ainv[i, j] * replace(Ftmp, repl) # rest of the operator: just diagonal! for i in range(num_stages): repl = {t: t + C[i] * dt} for j in range(num_fields): repl[u0bits[j]] = UUbits[i][j] repl[vbits[j]] = VVbits[i][j] # Also get replacements right for indexing. for j in range(num_fields): for ii in np.ndindex(u0bits[j].ufl_shape): repl[u0bits[j][ii]] = UUbits[i][j][ii] repl[vbits[j][ii]] = VVbits[i][j][ii] Fnew += dt * replace(split_form.remainder, repl) else: raise NotImplementedError("Can't handle that splitting type") if bcs is None: bcs = [] bcsnew = [] gblah = [] # For each BC, we need a new BC for each stage # so we need to figure out how the function is indexed (mixed + vec) # and then set it to have the value of the original argument at # time t+C[i]*dt. for bc in bcs: if num_fields == 1: # not mixed space comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(comp) Vbigi = lambda i: Vbig[i].sub(comp) else: Vsp = V Vbigi = lambda i: Vbig[i] else: # mixed space sub = bc.function_space_index() comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(sub).sub(comp) Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp) else: Vsp = V.sub(sub) Vbigi = lambda i: Vbig[sub + num_fields * i] bcarg = bc._original_arg for i in range(num_stages): try: gdat = interpolate(bcarg, Vsp) gmethod = lambda gd, gc: gd.interpolate(gc) except: # noqa: E722 gdat = project(bcarg, Vsp) gmethod = lambda gd, gc: gd.project(gc) gcur = replace(bcarg, {t: t + C[i] * dt}) bcsnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain)) gblah.append((gdat, gcur, gmethod)) nspacenew = getNullspace(V, Vbig, butch, nullspace) # For RIIA, we have an optimized update rule and don't need to # build the variational form for doing updates. # But something's broken with null spaces, so that's a TO-DO. unew = Function(V) Fupdate = inner(unew - u0, v) * dx B = vectorize(Constant)(butch.b) C = vectorize(Constant)(butch.c) for i in range(num_stages): repl = {t: t + C[i] * dt} for k in range(num_fields): repl[u0bits[k]] = UUbits[i][k] for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[i][k][ii] eFFi = replace(split_form.remainder, repl) Fupdate += dt * B[i] * eFFi # And the BC's for the update -- just the original BC at t+dt update_bcs = [] update_bcs_gblah = [] for bc in bcs: if num_fields == 1: # not mixed space comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(comp) else: Vsp = V else: # mixed space sub = bc.function_space_index() comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(sub).sub(comp) else: Vsp = V.sub(sub) bcarg = bc._original_arg try: gdat = interpolate(bcarg, Vsp) gmethod = lambda gd, gc: gd.interpolate(gc) except: # noqa: E722 gdat = project(bcarg, Vsp) gmethod = lambda gd, gc: gd.project(gc) gcur = replace(bcarg, {t: t + dt}) update_bcs.append(DirichletBC(Vsp, gdat, bc.sub_domain)) update_bcs_gblah.append((gdat, gcur, gmethod)) return (Fnew, (unew, Fupdate, update_bcs, update_bcs_gblah), UU, bcsnew, gblah, nspacenew)
def independent_operator(self, o): "Return a zero with the right shape and indices for operators independent of differentiation variable." return Zero(o.ufl_shape + self._var_shape, o.ufl_free_indices, o.ufl_index_dimensions)
def getForm(F, butch, t, dt, u0, bcs=None): """Given a time-dependent variational form and a :class:`ButcherTableau`, produce UFL for the s-stage RK method. :arg F: UFL form for the semidiscrete ODE/DAE :arg butch: the :class:`ButcherTableau` for the RK method being used to advance in time. :arg t: a :class:`Constant` referring to the current time level. Any explicit time-dependence in F is included :arg dt: a :class:`Constant` referring to the size of the current time step. :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possible time-dependent) boundary conditions imposed on the system. On output, we return a tuple consisting of four parts: - Fnew, the :class:`Form` - k, the :class:`firedrake.Function` holding all the stages. It lives in a :class:`firedrake.FunctionSpace` corresponding to the s-way tensor product of the space on which the semidiscrete form lives. - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the stages, - `gblah`, a list of pairs of the form (f, expr), where f is a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`. at each time step, each expr needs to be re-interpolated/projected onto the corresponding f in order for Firedrake to pick up that time-dependent boundary conditions need to be re-applied. """ v = F.arguments()[0] V = v.function_space() assert V == u0.function_space() A = numpy.array([[Constant(aa) for aa in arow] for arow in butch.A]) c = numpy.array([Constant(ci) for ci in butch.c]) num_stages = len(c) num_fields = len(V) Vbig = numpy.prod([V for i in range(num_stages)]) # Silence a warning about transfer managers when we # coarsen coefficients in V push_parent(V.dm, Vbig.dm) vnew = TestFunction(Vbig) k = Function(Vbig) if len(V) == 1: u0bits = [u0] vbits = [v] if num_stages == 1: vbigbits = [vnew] kbits = [k] else: vbigbits = split(vnew) kbits = split(k) else: u0bits = split(u0) vbits = split(v) vbigbits = split(vnew) kbits = split(k) kbits_np = numpy.zeros((num_stages, num_fields), dtype="object") for i in range(num_stages): for j in range(num_fields): kbits_np[i, j] = kbits[i * num_fields + j] Ak = A @ kbits_np Fnew = Zero() for i in range(num_stages): repl = {t: t + c[i] * dt} for j, (ubit, vbit, kbit) in enumerate(zip(u0bits, vbits, kbits)): repl[ubit] = ubit + dt * Ak[i, j] repl[vbit] = vbigbits[num_fields * i + j] repl[TimeDerivative(ubit)] = kbits_np[i, j] if (len(ubit.ufl_shape) == 1): for kk, kbitbit in enumerate(kbits_np[i, j]): repl[TimeDerivative(ubit[kk])] = kbitbit repl[ubit[kk]] = repl[ubit][kk] repl[vbit[kk]] = repl[vbit][kk] Fnew += replace(F, repl) bcnew = [] gblah = [] if bcs is None: bcs = [] for bc in bcs: if isinstance(bc.domain_args[0], str): boundary = bc.domain_args[0] else: boundary = () try: for j in bc.sub_domain: boundary += j except TypeError: boundary = (bc.sub_domain, ) gfoo = expand_derivatives(diff(bc._original_arg, t)) if len(V) == 1: for i in range(num_stages): gcur = replace(gfoo, {t: t + c[i] * dt}) try: gdat = interpolate(gcur, V) except NotImplementedError: gdat = project(gcur, V) gblah.append((gdat, gcur)) bcnew.append(DirichletBC(Vbig[i], gdat, boundary)) else: sub = bc.function_space_index() for i in range(num_stages): gcur = replace(gfoo, {t: t + butch.c[i] * dt}) try: gdat = interpolate(gcur, V.sub(sub)) except NotImplementedError: gdat = project(gcur, V) gblah.append((gdat, gcur)) bcnew.append( DirichletBC(Vbig[sub + num_fields * i], gdat, boundary)) return Fnew, k, bcnew, gblah
def independent_terminal(self, o): "Return a zero with the right shape for terminals independent of differentiation variable." return Zero(o.ufl_shape + self._var_shape)
def grad(self, g): # If we hit this type, it has already been propagated to a # coefficient (or grad of a coefficient), # FIXME: Assert # this! so we need to take the gradient of the variation or # return zero. Complications occur when dealing with # derivatives w.r.t. single components... # Figure out how many gradients are around the inner terminal ngrads = 0 o = g while isinstance(o, Grad): o, = o.ufl_operands ngrads += 1 if not isinstance(o, FormArgument): error("Expecting gradient of a FormArgument, not %s" % ufl_err_str(o)) def apply_grads(f): for i in range(ngrads): f = Grad(f) return f # Find o among all w without any indexing, which makes this # easy for (w, v) in zip(self._w, self._v): if o == w and isinstance(v, FormArgument): # Case: d/dt [w + t v] return apply_grads(v) # If o is not among coefficient derivatives, return do/dw=0 gprimesum = Zero(g.ufl_shape) def analyse_variation_argument(v): # Analyse variation argument if isinstance(v, FormArgument): # Case: d/dt [w[...] + t v] vval, vcomp = v, () elif isinstance(v, Indexed): # Case: d/dt [w + t v[...]] # Case: d/dt [w[...] + t v[...]] vval, vcomp = v.ufl_operands vcomp = tuple(vcomp) else: error("Expecting argument or component of argument.") if not all(isinstance(k, FixedIndex) for k in vcomp): error("Expecting only fixed indices in variation.") return vval, vcomp def compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp): # Apply gradients directly to argument vval, and get the # right indexed scalar component(s) kk = indices(ngrads) Dvkk = apply_grads(vval)[vcomp + kk] # Place scalar component(s) Dvkk into the right tensor # positions if wshape: Ejj, jj = unit_indexed_tensor(wshape, wcomp) else: Ejj, jj = 1, () gprimeterm = as_tensor(Ejj * Dvkk, jj + kk) return gprimeterm # Accumulate contributions from variations in different # components for (w, v) in zip(self._w, self._v): # Analyse differentiation variable coefficient if isinstance(w, FormArgument): if not w == o: continue wshape = w.ufl_shape if isinstance(v, FormArgument): # Case: d/dt [w + t v] return apply_grads(v) elif isinstance(v, ListTensor): # Case: d/dt [w + t <...,v,...>] for wcomp, vsub in unwrap_list_tensor(v): if not isinstance(vsub, Zero): vval, vcomp = analyse_variation_argument(vsub) gprimesum = gprimesum + compute_gprimeterm( ngrads, vval, vcomp, wshape, wcomp) else: if wshape != (): error("Expecting scalar coefficient in this branch.") # Case: d/dt [w + t v[...]] wval, wcomp = w, () vval, vcomp = analyse_variation_argument(v) gprimesum = gprimesum + compute_gprimeterm( ngrads, vval, vcomp, wshape, wcomp) elif isinstance( w, Indexed ): # This path is tested in unit tests, but not actually used? # Case: d/dt [w[...] + t v[...]] # Case: d/dt [w[...] + t v] wval, wcomp = w.ufl_operands if not wval == o: continue assert isinstance(wval, FormArgument) if not all(isinstance(k, FixedIndex) for k in wcomp): error( "Expecting only fixed indices in differentiation variable." ) wshape = wval.ufl_shape vval, vcomp = analyse_variation_argument(v) gprimesum = gprimesum + compute_gprimeterm( ngrads, vval, vcomp, wshape, wcomp) else: error("Expecting coefficient or component of coefficient.") # FIXME: Handle other coefficient derivatives: oprimes = # self._cd.get(o) if 0: oprimes = self._cd.get(o) if oprimes is None: if self._cd: # TODO: Make it possible to silence this message # in particular? It may be good to have for # debugging... warning("Assuming d{%s}/d{%s} = 0." % (o, self._w)) else: # Make sure we have a tuple to match the self._v tuple if not isinstance(oprimes, tuple): oprimes = (oprimes, ) if len(oprimes) != len(self._v): error("Got a tuple of arguments, expecting a" " matching tuple of coefficient derivatives.") # Compute dg/dw_j = dg/dw_h : v. # Since we may actually have a tuple of oprimes and vs # in a 'mixed' space, sum over them all to get the # complete inner product. Using indices to define a # non-compound inner product. for (oprime, v) in zip(oprimes, self._v): error("FIXME: Figure out how to do this with ngrads") so, oi = as_scalar(oprime) rv = len(v.ufl_shape) oi1 = oi[:-rv] oi2 = oi[-rv:] prod = so * v[oi2] if oi1: gprimesum += as_tensor(prod, oi1) else: gprimesum += prod return gprimesum
def ConstantOrZero(x): return Zero() if abs(complex(x)) < 1.e-10 else Constant(x)
def getForm(F, butch, t, dt, u0, bcs=None, bc_type="DAE", splitting=AI, nullspace=None): """Given a time-dependent variational form and a :class:`ButcherTableau`, produce UFL for the s-stage RK method. :arg F: UFL form for the semidiscrete ODE/DAE :arg butch: the :class:`ButcherTableau` for the RK method being used to advance in time. :arg t: a :class:`Constant` referring to the current time level. Any explicit time-dependence in F is included :arg dt: a :class:`Constant` referring to the size of the current time step. :arg splitting: a callable that maps the (floating point) Butcher matrix a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used to vary between the classical RK formulation and Butcher's reformulation that leads to a denser mass matrix with block-diagonal stiffness. Some choices of function will assume that `butch.A` is invertible. :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possible time-dependent) boundary conditions imposed on the system. :arg bc_type: How to manipulate the strongly-enforced boundary conditions to derive the stage boundary conditions. Should be a string, either "DAE", which implements BCs as constraints in the style of a differential-algebraic equation, or "ODE", which takes the time derivative of the boundary data and evaluates this for the stage values :arg nullspace: A list of tuples of the form (index, VSB) where index is an index into the function space associated with `u` and VSB is a :class: `firedrake.VectorSpaceBasis` instance to be passed to a `firedrake.MixedVectorSpaceBasis` over the larger space associated with the Runge-Kutta method On output, we return a tuple consisting of four parts: - Fnew, the :class:`Form` - k, the :class:`firedrake.Function` holding all the stages. It lives in a :class:`firedrake.FunctionSpace` corresponding to the s-way tensor product of the space on which the semidiscrete form lives. - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the stages, - 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object that represents the nullspace of the coupled system - `gblah`, a list of tuples of the form (f, expr, method), where f is a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`. At each time step, each expr needs to be re-interpolated/projected onto the corresponding f in order for Firedrake to pick up that time-dependent boundary conditions need to be re-applied. The interpolation/projection is encoded in method, which is either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending on whether the function space for f supports interpolation or not. """ v = F.arguments()[0] V = v.function_space() assert V == u0.function_space() c = numpy.array([Constant(ci) for ci in butch.c], dtype=object) bA1, bA2 = splitting(butch.A) try: bA1inv = numpy.linalg.inv(bA1) except numpy.linalg.LinAlgError: bA1inv = None try: bA2inv = numpy.linalg.inv(bA2) A2inv = numpy.array([[ConstantOrZero(aa) for aa in arow] for arow in bA2inv], dtype=object) except numpy.linalg.LinAlgError: raise NotImplementedError("We require A = A1 A2 with A2 invertible") A1 = numpy.array([[ConstantOrZero(aa) for aa in arow] for arow in bA1], dtype=object) if bA1inv is not None: A1inv = numpy.array([[ConstantOrZero(aa) for aa in arow] for arow in bA1inv], dtype=object) else: A1inv = None num_stages = butch.num_stages num_fields = len(V) Vbig = reduce(mul, (V for _ in range(num_stages))) vnew = TestFunction(Vbig) w = Function(Vbig) if len(V) == 1: u0bits = [u0] vbits = [v] if num_stages == 1: vbigbits = [vnew] wbits = [w] else: vbigbits = split(vnew) wbits = split(w) else: u0bits = split(u0) vbits = split(v) vbigbits = split(vnew) wbits = split(w) wbits_np = numpy.zeros((num_stages, num_fields), dtype=object) for i in range(num_stages): for j in range(num_fields): wbits_np[i, j] = wbits[i * num_fields + j] A1w = A1 @ wbits_np A2invw = A2inv @ wbits_np Fnew = Zero() for i in range(num_stages): repl = {t: t + c[i] * dt} for j, (ubit, vbit) in enumerate(zip(u0bits, vbits)): repl[ubit] = ubit + dt * A1w[i, j] repl[vbit] = vbigbits[num_fields * i + j] repl[TimeDerivative(ubit)] = A2invw[i, j] if (len(ubit.ufl_shape) == 1): for kk in range(len(A1w[i, j])): repl[TimeDerivative(ubit[kk])] = A2invw[i, j][kk] repl[ubit[kk]] = repl[ubit][kk] repl[vbit[kk]] = repl[vbit][kk] Fnew += replace(F, repl) bcnew = [] gblah = [] if bcs is None: bcs = [] if bc_type == "ODE": assert splitting == AI, "ODE-type BC aren't implemented for this splitting strategy" u0_mult_np = numpy.divide(1.0, butch.c, out=numpy.zeros_like(butch.c), where=butch.c != 0) u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np], dtype=object) def bc2gcur(bc, i): gorig = bc._original_arg gfoo = expand_derivatives(diff(gorig, t)) return replace(gfoo, {t: t + c[i] * dt}) + u0_mult[i] * gorig elif bc_type == "DAE": if bA1inv is None: raise NotImplementedError( "Cannot have DAE BCs for this Butcher Tableau/splitting") u0_mult_np = A1inv @ numpy.ones_like(butch.c) u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np], dtype=object) def bc2gcur(bc, i): gorig = as_ufl(bc._original_arg) gcur = 0 for j in range(num_stages): gcur += ConstantOrZero(bA1inv[i, j]) / dt * replace( gorig, {t: t + c[j] * dt}) return gcur else: raise ValueError("Unrecognised bc_type: %s", bc_type) # This logic uses information set up in the previous section to # set up the new BCs for either method for bc in bcs: if num_fields == 1: # not mixed space comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(comp) Vbigi = lambda i: Vbig[i].sub(comp) else: Vsp = V Vbigi = lambda i: Vbig[i] else: # mixed space sub = bc.function_space_index() comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(sub).sub(comp) Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp) else: Vsp = V.sub(sub) Vbigi = lambda i: Vbig[sub + num_fields * i] for i in range(num_stages): gcur = bc2gcur(bc, i) blah = BCStageData(Vsp, gcur, u0, u0_mult, i, t, dt) gdat, gcr, gmethod = blah.gstuff gblah.append((gdat, gcr, gmethod)) bcnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain)) nspnew = getNullspace(V, Vbig, butch, nullspace) return Fnew, w, bcnew, nspnew, gblah
def negative_restricted(self, o): return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions)