Coverage for /usr/share/miniconda3/envs/dolfin/lib/python3.8/site-packages/block/splitting.py: 89%
74 statements
« prev ^ index » next coverage.py v7.2.1, created at 2023-03-20 13:03 +0000
1from builtins import map
2from builtins import range
3from ufl.corealg.map_dag import MultiFunction
4import collections
5SplitForm = collections.namedtuple("SplitForm", ["indices", "form"])
7class FormSplitter(MultiFunction):
9 """Split a form in a list of subtrees for each component of the
10 mixed space it is built on. See :meth:`split` for a usage
11 description."""
13 def split(self, form):
14 """Split the form.
16 :arg form: the form to split.
18 This is a no-op if none of the arguments in the form are
19 defined on :class:`~.MixedFunctionSpace`\s.
21 The return-value is a tuple for which each entry is.
23 .. code-block:: python
25 (argument_indices, form)
27 Where ``argument_indices`` is a tuple indicating which part of
28 the mixed space the form belongs to, it has length equal to
29 the number of arguments in the form. Hence functionals have
30 a 0-tuple, 1-forms have a 1-tuple and 2-forms a 2-tuple
31 of indices.
33 For example, consider the following code:
35 .. code-block:: python
37 V = FunctionSpace(m, 'CG', 1)
38 W = V*V*V
39 u, v, w = TrialFunctions(W)
40 p, q, r = TestFunctions(W)
41 a = q*u*dx + p*w*dx
43 Then splitting the form returns a tuple of two forms.
45 .. code-block:: python
47 ((0, 2), w*p*dx),
48 (1, 0), q*u*dx))
50 """
51 from ufl.algorithms.map_integrands import map_integrand_dags
52 from numpy import ndindex
53 args = form.arguments()
54 if all(a.function_space().num_sub_spaces() == 0 for a in args): 54 ↛ exit, 54 ↛ 562 missed branches: 1) line 54 didn't finish the generator expression on line 54, 2) line 54 didn't jump to line 56, because the condition on line 54 was never true
55 # No mixed spaces, just return the form directly.
56 idx = tuple([0]*len(form.arguments()))
57 return (SplitForm(indices=idx, form=form), )
58 forms = []
59 # How many subspaces do we have for each argument?
60 shape = tuple(max(1, a.function_space().num_sub_spaces()) for a in args)
61 # Walk over all the indices of the spaces
62 for idx in ndindex(shape):
63 # Which subspace are we currently interested in?
64 self.idx = dict(enumerate(idx))
65 # Cache for the arguments we construct
66 self._args = {}
67 # Visit the form
68 f = map_integrand_dags(self, form)
69 # Zero-simplification may result in an empty form, only
70 # collect those that are non-zero.
71 if len(f.integrals()) > 0: 71 ↛ 62line 71 didn't jump to line 62, because the condition on line 71 was never false
72 forms.append(SplitForm(indices=idx, form=f))
73 return tuple(forms)
75 expr = MultiFunction.reuse_if_untouched
77 def multi_index(self, o):
78 return o
80 def argument(self, o):
81 from ufl import as_vector
82 from ufl.constantvalue import Zero
83 from dolfin.function.argument import Argument
84 from numpy import ndindex
85 V = o.function_space()
86 if V.num_sub_spaces() == 0: 86 ↛ 88line 86 didn't jump to line 88, because the condition on line 86 was never true
87 # Not on a mixed space, just return ourselves.
88 return o
89 # Already seen this argument, return the cached version.
90 if o in self._args:
91 return self._args[o]
92 args = []
93 for i, V_i_sub in enumerate(V.split()):
94 # Walk over the subspaces and build a vector that is zero
95 # where the argument does not match the one we're looking
96 # for and is just the non-mixed argument when we do want
97 # it.
98 V_i = V_i_sub.collapse()
99 a = Argument(V_i, o.number(), part=o.part())
100 indices = ndindex(a.ufl_shape)
101 if self.idx[o.number()] == i:
102 args += [a[j] for j in indices]
103 else:
104 args += [Zero() for j in indices]
105 self._args[o] = as_vector(args)
106 return self._args[o]
108def split_form(form):
109 # Takes a mixed space form and returns a numpy array of forms for use with cbc.block
110 # Array entries are the forms on each collapsed subspace of the mixed space.
111 from numpy import zeros
113 args = form.arguments()
114 assert len(args) <= 2
116 shape = tuple(max(1, a.function_space().num_sub_spaces()) for a in args)
118 forms = zeros(shape, dtype = object)
120 for ij, form_ij in FormSplitter().split(form):
121 forms[ij] = form_ij
123 return forms
126def _collapse_bc(bc):
127 from dolfin import DirichletBC, FunctionSpace, Constant
128 sub_space = FunctionSpace(bc.function_space())
130 assert len(sub_space.component()) == 1
132 sub_id = sub_space.component()[0]
133 if hasattr(bc, 'domain_args'): 133 ↛ 134line 133 didn't jump to line 134, because the condition on line 133 was never true
134 sub_bc = DirichletBC(sub_space.collapse(), bc.value(), *bc.domain_args)
135 else:
136 if hasattr(bc.value(), 'values'): 136 ↛ 139line 136 didn't jump to line 139, because the condition on line 136 was never false
137 value = Constant(bc.value().values())
138 else:
139 value = bc.value()
140 sub_bc = DirichletBC(sub_space.collapse(), value, bc.user_sub_domain())
141 return (int(sub_id), sub_bc)
144def split_bcs(bcs, m):
145 # return a list of lists of DirichletBC for use with cbc.block
146 # we need the number of blocks m to ensure correct length of output list;
147 # if m is None, the list will stop at the last active BC
148 bc_map = [_collapse_bc(bc) for bc in bcs]
149 if m is None: 149 ↛ 152line 149 didn't jump to line 152, because the condition on line 149 was never false
150 m = max(i for i,_ in bc_map)+1
152 collapsed_bcs = [[] for i in range(m)]
153 for i, bc_i in bc_map:
154 collapsed_bcs[i].append(bc_i)
156 return collapsed_bcs