« back to all changes in this revision
Viewing changes to solving.py
- browse files at revision 85
- compare with another revision
- download diff
- download tarball
- view history from revision 85
- Committer: Simon Funke
- Date: 2012-01-17 19:39:45 UTC
- mfrom: (80.2.10 dolfin_adjoint)
- Revision ID: simon.funke@gmail.com-20120117193945-ng3vty6tixh3quba
Merge from "trunk"
- files added:
- assembly.py
- tests/navier_stokes
- tests/tlm_simple
- files modified:
- __init__.py
added
removed
solving.py
10
10
import libadjoint.exceptions
11
11
12
12
import hashlib
13
import time
14
15
import assembly
13
16
14
17
debugging={}
15
18
77
80
else:
78
81
eq_lhs, eq_rhs = define_nonlinear_equation(eq.lhs, u)
79
82
F = eq.lhs
80
eq_bcs = None
83
eq_bcs = []
81
84
linear = False
82
85
83
# Suppose we are solving for a variable w, and that variable shows up in the
84
# coefficients of eq_lhs/eq_rhs.
85
# Does that mean:
86
# a) the /previous value/ of that variable, and you want to timestep?
87
# b) the /value to be solved for/ in this solve?
88
# i.e. is it timelevel n-1, or n?
89
# if Dolfin is doing a linear solve, we want case a);
90
# if Dolfin is doing a nonlinear solve, we want case b).
91
# so if we are doing a nonlinear solve, we bump the timestep number here
92
# /before/ we map the coefficients -> dependencies,
93
# so that libadjoint records the dependencies with the right timestep number.
94
if not linear:
95
var = adj_variables.next(u)
96
97
# Set up the data associated with the matrix on the left-hand side. This goes on the diagonal
98
# of the 'large' system that incorporates all of the timelevels, which is why it is prefixed
99
# with diag.
100
diag_name = hashlib.md5(str(eq_lhs)).hexdigest() # we don't have a useful human-readable name, so take the md5sum of the string representation of the form
101
diag_deps = [adj_variables[coeff] for coeff in ufl.algorithms.extract_coefficients(eq_lhs) if hasattr(coeff, "function_space")]
102
diag_coeffs = [coeff for coeff in ufl.algorithms.extract_coefficients(eq_lhs) if hasattr(coeff, "function_space")]
103
diag_block = libadjoint.Block(diag_name, dependencies=diag_deps, test_hermitian=debugging["test_hermitian"], test_derivative=debugging["test_derivative"])
104
105
# Similarly, create the object associated with the right-hand side data.
106
if linear:
107
rhs = RHS(eq_rhs)
86
newargs = list(args)
87
88
try:
89
solver_params = kwargs["solver_parameters"]
90
ksp = solver_params["linear_solver"]
91
except KeyError:
92
ksp = None
93
94
try:
95
solver_params = kwargs["solver_parameters"]
96
pc = solver_params["preconditioner"]
97
except KeyError:
98
pc = None
99
100
elif isinstance(args[0], dolfin.cpp.Matrix):
101
linear = True
102
try:
103
eq_lhs = assembly.assembly_cache[args[0]]
104
except KeyError:
105
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("dolfin_adjoint did not assemble your form, and so does not recognise your matrix. Did you from dolfin_adjoint import *?")
106
107
try:
108
eq_rhs = assembly.assembly_cache[args[2]]
109
except KeyError:
110
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("dolfin_adjoint did not assemble your form, and so does not recognise your right-hand side. Did you from dolfin_adjoint import *?")
111
112
u = args[1]
113
if not isinstance(u, dolfin.Function):
114
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("dolfin_adjoint needs the Function you are solving for, rather than its .vector().")
115
newargs = list(args)
116
newargs[1] = u.vector() # for the real solve later
117
118
try:
119
ksp = args[3]
120
except IndexError:
121
ksp = None
122
123
try:
124
pc = args[4]
125
except IndexError:
126
pc = None
127
128
eq_bcs = list(set(assembly.bc_cache[args[0]] + assembly.bc_cache[args[2]]))
129
else:
130
raise libadjoint.exceptions.LibadjointErrorNotImplemented("Don't know how to annotate your equation, sorry!")
131
132
# Suppose we are solving for a variable w, and that variable shows up in the
133
# coefficients of eq_lhs/eq_rhs.
134
# Does that mean:
135
# a) the /previous value/ of that variable, and you want to timestep?
136
# b) the /value to be solved for/ in this solve?
137
# i.e. is it timelevel n-1, or n?
138
# if Dolfin is doing a linear solve, we want case a);
139
# if Dolfin is doing a nonlinear solve, we want case b).
140
# so if we are doing a nonlinear solve, we bump the timestep number here
141
# /before/ we map the coefficients -> dependencies,
142
# so that libadjoint records the dependencies with the right timestep number.
143
if not linear:
144
var = adj_variables.next(u)
145
146
# Set up the data associated with the matrix on the left-hand side. This goes on the diagonal
147
# of the 'large' system that incorporates all of the timelevels, which is why it is prefixed
148
# with diag.
149
diag_name = hashlib.md5(str(eq_lhs) + str(eq_rhs) + str(u) + str(time.time())).hexdigest() # we don't have a useful human-readable name, so take the md5sum of the string representation of the forms
150
diag_deps = [adj_variables[coeff] for coeff in ufl.algorithms.extract_coefficients(eq_lhs) if hasattr(coeff, "function_space")]
151
diag_coeffs = [coeff for coeff in ufl.algorithms.extract_coefficients(eq_lhs) if hasattr(coeff, "function_space")]
152
diag_block = libadjoint.Block(diag_name, dependencies=diag_deps, test_hermitian=debugging["test_hermitian"], test_derivative=debugging["test_derivative"])
153
154
# Similarly, create the object associated with the right-hand side data.
155
if linear:
156
rhs = RHS(eq_rhs)
157
else:
158
rhs = NonlinearRHS(eq_rhs, F, u, bcs, mass=eq_lhs)
159
160
# We need to check if this is the first equation,
161
# so that we can register the appropriate initial conditions.
162
# These equations are necessary so that libadjoint can assemble the
163
# relevant adjoint equations for the adjoint variables associated with
164
# the initial conditions.
165
for coeff, dep in zip(rhs.coefficients(),rhs.dependencies()) + zip(diag_coeffs, diag_deps):
166
# If coeff is not known, it must be an initial condition.
167
if not adjointer.variable_known(dep):
168
169
if not linear: # don't register initial conditions for the first nonlinear solve.
170
if dep == var:
171
continue
172
173
fn_space = coeff.function_space()
174
block_name = "Identity: %s" % str(fn_space)
175
identity_block = libadjoint.Block(block_name)
176
177
init_rhs=Vector(coeff).duplicate()
178
init_rhs.axpy(1.0,Vector(coeff))
179
180
def identity_assembly_cb(variables, dependencies, hermitian, coefficient, context):
181
assert coefficient == 1
182
return (Matrix(ufl.Identity(fn_space.dim())), Vector(dolfin.Function(fn_space)))
183
184
identity_block.assemble=identity_assembly_cb
185
186
if debugging["record_all"]:
187
adjointer.record_variable(dep, libadjoint.MemoryStorage(Vector(coeff)))
188
189
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=RHS(init_rhs))
190
adjointer.register_equation(initial_eq)
191
assert adjointer.variable_known(dep)
192
193
194
# c.f. the discussion above. In the linear case, we want to bump the
195
# timestep number /after/ all of the dependencies' timesteps have been
196
# computed for libadjoint.
197
if linear:
198
var = adj_variables.next(u)
199
200
# With the initial conditions out of the way, let us now define the callbacks that
201
# define the actions of the operator the user has passed in on the lhs of this equation.
202
203
def diag_assembly_cb(dependencies, values, hermitian, coefficient, context):
204
'''This callback must conform to the libadjoint Python block assembly
205
interface. It returns either the form or its transpose, depending on
206
the value of the logical hermitian.'''
207
208
assert coefficient == 1
209
210
value_coeffs=[v.data for v in values]
211
eq_l=dolfin.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
212
213
if hermitian:
214
# Homogenise the adjoint boundary conditions. This creates the adjoint
215
# solution associated with the lifted discrete system that is actually solved.
216
adjoint_bcs = [dolfin.homogenize(bc) for bc in eq_bcs if isinstance(bc, dolfin.DirichletBC)]
217
if len(adjoint_bcs) == 0: adjoint_bcs = None
218
return (Matrix(dolfin.adjoint(eq_l), bcs=adjoint_bcs, pc=pc, ksp=ksp), Vector(None, fn_space=u.function_space()))
108
219
else:
109
rhs = NonlinearRHS(eq_rhs, F, u, bcs, mass=eq_lhs)
110
111
# We need to check if this is the first equation,
112
# so that we can register the appropriate initial conditions.
113
# These equations are necessary so that libadjoint can assemble the
114
# relevant adjoint equations for the adjoint variables associated with
115
# the initial conditions.
116
for coeff, dep in zip(rhs.coefficients(),rhs.dependencies()) + zip(diag_coeffs, diag_deps):
117
# If coeff is not known, it must be an initial condition.
118
if not adjointer.variable_known(dep):
119
120
if not linear: # don't register initial conditions for the first nonlinear solve.
121
if dep == var:
122
continue
123
124
fn_space = coeff.function_space()
125
block_name = "Identity: %s" % str(fn_space)
126
identity_block = libadjoint.Block(block_name)
127
128
init_rhs=Vector(coeff).duplicate()
129
init_rhs.axpy(1.0,Vector(coeff))
130
131
def identity_assembly_cb(variables, dependencies, hermitian, coefficient, context):
132
assert coefficient == 1
133
return (Matrix(ufl.Identity(fn_space.dim())), Vector(dolfin.Function(fn_space)))
134
135
identity_block.assemble=identity_assembly_cb
136
137
if debugging["record_all"]:
138
adjointer.record_variable(dep, libadjoint.MemoryStorage(Vector(coeff)))
139
140
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=RHS(init_rhs))
141
adjointer.register_equation(initial_eq)
142
assert adjointer.variable_known(dep)
143
144
145
# c.f. the discussion above. In the linear case, we want to bump the
146
# timestep number /after/ all of the dependencies' timesteps have been
147
# computed for libadjoint.
148
if linear:
149
var = adj_variables.next(u)
150
151
# With the initial conditions out of the way, let us now define the callbacks that
152
# define the actions of the operator the user has passed in on the lhs of this equation.
153
154
def diag_assembly_cb(dependencies, values, hermitian, coefficient, context):
155
'''This callback must conform to the libadjoint Python block assembly
156
interface. It returns either the form or its transpose, depending on
157
the value of the logical hermitian.'''
158
159
assert coefficient == 1
160
161
value_coeffs=[v.data for v in values]
162
eq_l=dolfin.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
163
164
if hermitian:
165
# Homogenise the adjoint boundary conditions. This creates the adjoint
166
# solution associated with the lifted discrete system that is actually solved.
167
adjoint_bcs = [dolfin.homogenize(bc) for bc in eq_bcs if isinstance(bc, dolfin.DirichletBC)]
168
if len(adjoint_bcs) == 0: adjoint_bcs = None
169
return (Matrix(dolfin.adjoint(eq_l), bcs=adjoint_bcs), Vector(None, fn_space=u.function_space()))
170
else:
171
return (Matrix(eq_l, bcs=eq_bcs), Vector(None, fn_space=u.function_space()))
172
diag_block.assemble = diag_assembly_cb
173
174
def diag_action_cb(dependencies, values, hermitian, coefficient, input, context):
175
value_coeffs = [v.data for v in values]
176
eq_l = dolfin.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
177
178
if hermitian:
179
eq_l = dolfin.adjoint(eq_l)
180
181
output_vec = dolfin.assemble(coefficient * dolfin.action(eq_l, input.data))
182
output_fn = dolfin.Function(input.data.function_space())
183
vec = output_fn.vector()
184
for i in range(len(vec)):
185
vec[i] = output_vec[i]
186
187
return Vector(output_fn)
188
diag_block.action = diag_action_cb
189
190
if len(diag_deps) > 0:
191
# If this block is nonlinear (the entries of the matrix on the LHS
192
# depend on any variable previously computed), then that will induce
193
# derivative terms in the adjoint equations. Here, we define the
194
# callback libadjoint will need to compute such terms.
195
def derivative_action(dependencies, values, variable, contraction_vector, hermitian, input, coefficient, context):
196
dolfin_variable = values[dependencies.index(variable)].data
197
dolfin_values = [val.data for val in values]
198
199
current_form = dolfin.replace(eq_lhs, dict(zip(diag_coeffs, dolfin_values)))
200
201
deriv = dolfin.derivative(current_form, dolfin_variable)
202
args = ufl.algorithms.extract_arguments(deriv)
203
deriv = dolfin.replace(deriv, {args[1]: contraction_vector.data}) # contract over the middle index
204
205
if hermitian:
206
deriv = dolfin.adjoint(deriv)
207
208
action = coefficient * dolfin.action(deriv, input.data)
209
210
return Vector(action)
211
diag_block.derivative_action = derivative_action
212
213
eqn = libadjoint.Equation(var, blocks=[diag_block], targets=[var], rhs=rhs)
214
215
cs = adjointer.register_equation(eqn)
216
if cs == int(libadjoint.constants.adj_constants["ADJ_CHECKPOINT_STORAGE_MEMORY"]):
217
for coeff in adj_variables.coeffs.keys():
218
if adj_variables[coeff] == var: continue
219
adjointer.record_variable(adj_variables[coeff], libadjoint.MemoryStorage(Vector(coeff), cs=True))
220
221
elif cs == int(libadjoint.constants.adj_constants["ADJ_CHECKPOINT_STORAGE_DISK"]):
222
for coeff in adj_variables.coeffs.keys():
223
if adj_variables[coeff] == var: continue
224
adjointer.record_variable(adj_variables[coeff], libadjoint.DiskStorage(Vector(coeff), cs=True))
225
226
return linear
220
return (Matrix(eq_l, bcs=eq_bcs, pc=pc, ksp=ksp), Vector(None, fn_space=u.function_space()))
221
diag_block.assemble = diag_assembly_cb
222
223
def diag_action_cb(dependencies, values, hermitian, coefficient, input, context):
224
value_coeffs = [v.data for v in values]
225
eq_l = dolfin.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
226
227
if hermitian:
228
eq_l = dolfin.adjoint(eq_l)
229
230
output_vec = dolfin.assemble(coefficient * dolfin.action(eq_l, input.data))
231
output_fn = dolfin.Function(input.data.function_space())
232
vec = output_fn.vector()
233
for i in range(len(vec)):
234
vec[i] = output_vec[i]
235
236
return Vector(output_fn)
237
diag_block.action = diag_action_cb
238
239
if len(diag_deps) > 0:
240
# If this block is nonlinear (the entries of the matrix on the LHS
241
# depend on any variable previously computed), then that will induce
242
# derivative terms in the adjoint equations. Here, we define the
243
# callback libadjoint will need to compute such terms.
244
def derivative_action(dependencies, values, variable, contraction_vector, hermitian, input, coefficient, context):
245
dolfin_variable = values[dependencies.index(variable)].data
246
dolfin_values = [val.data for val in values]
247
248
current_form = dolfin.replace(eq_lhs, dict(zip(diag_coeffs, dolfin_values)))
249
250
deriv = dolfin.derivative(current_form, dolfin_variable)
251
args = ufl.algorithms.extract_arguments(deriv)
252
deriv = dolfin.replace(deriv, {args[1]: contraction_vector.data}) # contract over the middle index
253
254
if hermitian:
255
deriv = dolfin.adjoint(deriv)
256
257
action = coefficient * dolfin.action(deriv, input.data)
258
259
return Vector(action)
260
diag_block.derivative_action = derivative_action
261
262
eqn = libadjoint.Equation(var, blocks=[diag_block], targets=[var], rhs=rhs)
263
264
cs = adjointer.register_equation(eqn)
265
if cs == int(libadjoint.constants.adj_constants["ADJ_CHECKPOINT_STORAGE_MEMORY"]):
266
for coeff in adj_variables.coeffs.keys():
267
if adj_variables[coeff] == var: continue
268
adjointer.record_variable(adj_variables[coeff], libadjoint.MemoryStorage(Vector(coeff), cs=True))
269
270
elif cs == int(libadjoint.constants.adj_constants["ADJ_CHECKPOINT_STORAGE_DISK"]):
271
for coeff in adj_variables.coeffs.keys():
272
if adj_variables[coeff] == var: continue
273
adjointer.record_variable(adj_variables[coeff], libadjoint.DiskStorage(Vector(coeff), cs=True))
274
return linear, newargs
227
275
228
276
def solve(*args, **kwargs):
229
277
'''This solve routine comes from the dolfin_adjoint package, and wraps the real Dolfin
242
290
del kwargs["annotate"] # so we don't pass it on to the real solver
243
291
244
292
if to_annotate:
245
linear = annotate(*args, **kwargs)
293
linear, newargs = annotate(*args, **kwargs)
294
else:
295
newargs = args
246
296
247
ret = dolfin.fem.solving.solve(*args, **kwargs)
297
ret = dolfin.fem.solving.solve(*newargs, **kwargs)
248
298
249
299
if to_annotate:
250
300
if not linear and debugging["fussy_replay"]:
257
307
# Finally, if we want to record all of the solutions of the real forward model
258
308
# (for comparison with a libadjoint replay later),
259
309
# then we should record the value of the variable we just solved for.
260
if isinstance(args[0], ufl.classes.Equation):
261
if debugging["record_all"]:
310
if debugging["record_all"]:
311
if isinstance(args[0], ufl.classes.Equation):
262
312
unpacked_args = dolfin.fem.solving._extract_args(*args, **kwargs)
263
313
u = unpacked_args[1]
264
314
adjointer.record_variable(adj_variables[u], libadjoint.MemoryStorage(Vector(u)))
315
elif isinstance(args[0], dolfin.cpp.Matrix):
316
u = args[1]
317
adjointer.record_variable(adj_variables[u], libadjoint.MemoryStorage(Vector(u)))
318
else:
319
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Don't know how to record, sorry")
265
320
266
321
return ret
267
322
446
501
together, duplicating matrices, etc., that occur in the process of constructing
447
502
the adjoint equations.'''
448
503
449
def __init__(self, data, bcs=None):
504
def __init__(self, data, bcs=None, ksp=None, pc=None):
450
505
451
506
if bcs is None:
452
507
self.bcs = []
455
510
456
511
self.data=data
457
512
458
def solve(self, b):
513
self.ksp = ksp
514
self.pc = pc
515
516
def solve(self, var, b):
459
517
460
518
if isinstance(self.data, ufl.Identity):
461
519
x=b.duplicate()
462
520
x.axpy(1.0, b)
463
521
else:
522
if var.type in ['ADJ_TLM', 'ADJ_ADJOINT']:
523
bcs = [dolfin.homogenize(bc) for bc in self.bcs if isinstance(bc, dolfin.DirichletBC)] + [bc for bc in self.bcs if not isinstance(bc, dolfin.DirichletBC)]
524
else:
525
bcs = self.bcs
526
527
solver_parameters = {}
528
if self.pc is not None:
529
solver_parameters["preconditioner"] = self.pc
530
if self.ksp is not None:
531
solver_parameters["linear_solver"] = self.ksp
532
464
533
x=Vector(dolfin.Function(self.test_function().function_space()))
465
dolfin.fem.solving.solve(self.data==b.data, x.data, self.bcs)
534
dolfin.fem.solving.solve(self.data==b.data, x.data, bcs, solver_parameters=solver_parameters)
466
535
467
536
return x
468
537
502
571
503
572
def __init__(self, form):
504
573
505
self.form=form
574
self.form = form
575
self.activated = False
506
576
507
577
def __call__(self, dependencies, values):
508
578
528
598
529
599
def dependencies(self, adjointer, timestep):
530
600
531
if timestep == adjointer.timestep_count-1:
601
if self.activated is False:
532
602
deps = [adj_variables[coeff] for coeff in ufl.algorithms.extract_coefficients(self.form) if hasattr(coeff, "function_space")]
603
self.activated = True
533
604
else:
534
605
deps = []
535
606
765
836
def adjoint_checkpointing(strategy, steps, snaps_on_disk, snaps_in_ram, verbose=False):
766
837
adjointer.set_checkpoint_strategy(strategy)
767
838
adjointer.set_revolve_options(steps, snaps_on_disk, snaps_in_ram, verbose)
839
840
class InitialConditionParameter(libadjoint.Parameter):
841
'''This Parameter is used as input to the tangent linear model (TLM)
842
when one wishes to compute dJ/d(initial condition) in a particular direction (perturbation).'''
843
def __init__(self, coeff, perturbation):
844
'''coeff: the variable whose initial condition you wish to perturb.
845
perturbation: the perturbation direction in which you wish to compute the gradient. Must be a Function.'''
846
self.var = adj_variables[coeff]
847
self.var.c_object.timestep = 0 # we want to put in the source term only at the initial condition.
848
self.perturbation = Vector(perturbation).duplicate()
849
self.perturbation.axpy(1.0, Vector(perturbation))
850
851
def __call__(self, dependencies, values, variable):
852
# The TLM source term only kicks in at the start, for the initial condition:
853
if self.var == variable:
854
return self.perturbation
855
else:
856
return None
857
858
def __str__(self):
859
return self.var.name + ':InitialCondition'
Loggerhead is a web-based interface for Breezy
Version: 2.0.1