-
Notifications
You must be signed in to change notification settings - Fork 157
/
variational_solver.py
236 lines (188 loc) · 8.71 KB
/
variational_solver.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
from __future__ import absolute_import
import weakref
import ufl
from pyop2.profiling import timed_function, profile
from firedrake import solving_utils
from firedrake import ufl_expr
from firedrake.petsc import PETSc
__all__ = ["LinearVariationalProblem",
"LinearVariationalSolver",
"NonlinearVariationalProblem",
"NonlinearVariationalSolver"]
class NonlinearVariationalProblem(object):
"""Nonlinear variational problem F(u; v) = 0."""
def __init__(self, F, u, bcs=None, J=None,
Jp=None,
form_compiler_parameters=None,
nest=None):
"""
:param F: the nonlinear form
:param u: the :class:`.Function` to solve for
:param bcs: the boundary conditions (optional)
:param J: the Jacobian J = dF/du (optional)
:param Jp: a form used for preconditioning the linear system,
optional, if not supplied then the Jacobian itself
will be used.
:param dict form_compiler_parameters: parameters to pass to the form
compiler (optional)
:param nest: indicate if matrices on mixed spaces should be
built as monolithic operators (suitable for direct
solves), or as nested blocks (suitable for fieldsplit
preconditioning). If not provided, uses the default
given by ``parameters["matnest"]``.
"""
from firedrake import solving
# Extract and check arguments
u = solving._extract_u(u)
bcs = solving._extract_bcs(bcs)
# Store input UFL forms and solution Function
self.F = F
# Use the user-provided Jacobian. If none is provided, derive
# the Jacobian from the residual.
self.J = J or ufl_expr.derivative(F, u)
self.Jp = Jp
self.u = u
self.bcs = bcs
self._nest = nest
# Store form compiler parameters
self.form_compiler_parameters = form_compiler_parameters
self._constant_jacobian = False
@property
def dm(self):
return self.u.function_space()._dm
class NonlinearVariationalSolver(object):
"""Solves a :class:`NonlinearVariationalProblem`."""
_id = 0
def __init__(self, problem, **kwargs):
"""
:arg problem: A :class:`NonlinearVariationalProblem` to solve.
:kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or
:class:`.MixedVectorSpaceBasis`) spanning the null
space of the operator.
:kwarg solver_parameters: Solver parameters to pass to PETSc.
This should be a dict mapping PETSc options to values. For
example, to set the nonlinear solver type to just use a linear
solver:
:kwarg options_prefix: an optional prefix used to distinguish
PETSc options. If not provided a unique prefix will be
created. Use this option if you want to pass options
to the solver from the command line in addition to
through the ``solver_parameters`` dict.
.. code-block:: python
{'snes_type': 'ksponly'}
PETSc flag options should be specified with `bool` values. For example:
.. code-block:: python
{'snes_monitor': True}
"""
parameters, nullspace, options_prefix = solving_utils._extract_kwargs(**kwargs)
# Do this first so __del__ doesn't barf horribly if we get an
# error in __init__
if options_prefix is not None:
self._opt_prefix = options_prefix
self._auto_prefix = False
else:
self._opt_prefix = 'firedrake_snes_%d_' % NonlinearVariationalSolver._id
self._auto_prefix = True
NonlinearVariationalSolver._id += 1
assert isinstance(problem, NonlinearVariationalProblem)
ctx = solving_utils._SNESContext(problem)
self.snes = PETSc.SNES().create()
self.snes.setOptionsPrefix(self._opt_prefix)
# Mixed problem, use jacobi pc if user has not supplied one.
if ctx.is_mixed:
parameters.setdefault('pc_type', 'jacobi')
# Allow command-line arguments to override dict parameters
opts = PETSc.Options()
for k, v in opts.getAll().iteritems():
if k.startswith(self._opt_prefix):
parameters[k[len(self._opt_prefix):]] = v
self._problem = problem
self._ctx = ctx
self.snes.setDM(problem.dm)
ctx.set_function(self.snes)
ctx.set_jacobian(self.snes)
ctx.set_nullspace(nullspace, problem.J.arguments()[0].function_space()._ises)
self.parameters = parameters
def __del__(self):
# Remove stuff from the options database
# It's fixed size, so if we don't it gets too big.
if self._auto_prefix and hasattr(self, '_opt_prefix'):
opts = PETSc.Options()
for k in self.parameters.iterkeys():
del opts[self._opt_prefix + k]
delattr(self, '_opt_prefix')
@property
def parameters(self):
return self._parameters
@parameters.setter
def parameters(self, val):
assert isinstance(val, dict), 'Must pass a dict to set parameters'
self._parameters = val
solving_utils.update_parameters(self, self.snes)
@timed_function("SNES solver execution")
@profile
def solve(self):
dm = self.snes.getDM()
dm.setAppCtx(weakref.proxy(self._ctx))
dm.setCreateMatrix(self._ctx.create_matrix)
# Apply the boundary conditions to the initial guess.
for bc in self._problem.bcs:
bc.apply(self._problem.u)
# User might have updated parameters dict before calling
# solve, ensure these are passed through to the snes.
solving_utils.update_parameters(self, self.snes)
with self._problem.u.dat.vec as v:
self.snes.solve(None, v)
solving_utils.check_snes_convergence(self.snes)
class LinearVariationalProblem(NonlinearVariationalProblem):
"""Linear variational problem a(u, v) = L(v)."""
def __init__(self, a, L, u, bcs=None, aP=None,
form_compiler_parameters=None,
nest=None,
constant_jacobian=True):
"""
:param a: the bilinear form
:param L: the linear form
:param u: the :class:`.Function` to solve for
:param bcs: the boundary conditions (optional)
:param aP: an optional operator to assemble to precondition
the system (if not provided a preconditioner may be
computed from ``a``)
:param dict form_compiler_parameters: parameters to pass to the form
compiler (optional)
:param nest: indicate if matrices on mixed spaces should be
built as monolithic operators (suitable for direct
solves), or as nested blocks (suitable for fieldsplit
preconditioning). If not provided, uses the default
given by ``parameters["matnest"]``.
:param constant_jacobian: (optional) flag indicating that the
Jacobian is constant (i.e. does not depend on
varying fields). If your Jacobian can change, set
this flag to ``False``.
"""
# In the linear case, the Jacobian is the equation LHS.
J = a
F = ufl.action(J, u) - L
super(LinearVariationalProblem, self).__init__(F, u, bcs, J, aP,
form_compiler_parameters=form_compiler_parameters, nest=nest)
self._constant_jacobian = constant_jacobian
class LinearVariationalSolver(NonlinearVariationalSolver):
"""Solves a :class:`LinearVariationalProblem`."""
def __init__(self, *args, **kwargs):
"""
:arg problem: A :class:`LinearVariationalProblem` to solve.
:kwarg solver_parameters: Solver parameters to pass to PETSc.
This should be a dict mapping PETSc options to values.
:kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or
:class:`.MixedVectorSpaceBasis`) spanning the null
space of the operator.
:kwarg options_prefix: an optional prefix used to distinguish
PETSc options. If not provided a unique prefix will be
created. Use this option if you want to pass options
to the solver from the command line in addition to
through the ``solver_parameters`` dict.
"""
super(LinearVariationalSolver, self).__init__(*args, **kwargs)
self.parameters.setdefault('snes_type', 'ksponly')
self.parameters.setdefault('ksp_rtol', 1.0e-7)
solving_utils.update_parameters(self, self.snes)