~libadjoint/libadjoint/dolfin_predictability

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
import dolfin
import solving
import assembly
import libadjoint

lu_solvers = {}
adj_lu_solvers = {}

def make_LUSolverMatrix(form, reuse_factorization):
  class LUSolverMatrix(solving.Matrix):
    def solve(self, var, b):

      if reuse_factorization is False:
        return solving.Matrix.solve(self, var, b)

      if var.type in ['ADJ_TLM', 'ADJ_ADJOINT']:
        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)]
      else:
        bcs = self.bcs

      if var.type in ['ADJ_FORWARD', 'ADJ_TLM']:
        solver = lu_solvers[form]
      else:
        if adj_lu_solvers[form] is None:
          A = assembly.assemble(self.data); [bc.apply(A) for bc in bcs]
          adj_lu_solvers[form] = LUSolver(A)
          adj_lu_solvers[form].parameters["reuse_factorization"] = True

        solver = adj_lu_solvers[form]

      x = solving.Vector(dolfin.Function(self.test_function().function_space()))

      if b.data is None:
        # This means we didn't get any contribution on the RHS of the adjoint system. This could be that the
        # simulation ran further ahead than when the functional was evaluated, or it could be that the
        # functional is set up incorrectly.
        dolfin.info_red("Warning: got zero RHS for the solve associated with variable %s" % var)
      else:
        b_vec = dolfin.assemble(b.data); [bc.apply(b_vec) for bc in bcs]
        solver.solve(x.data.vector(), b_vec, annotate=False)

      return x
  return LUSolverMatrix

class LUSolver(dolfin.LUSolver):
  def __init__(self, *args):
    try:
      self.operator = args[0].form
    except AttributeError:
      raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your matrix A has to have the .form attribute: was it assembled after from dolfin_adjoint import *?")

    dolfin.LUSolver.__init__(self, *args)

  def solve(self, *args, **kwargs):

    annotate = True
    if "annotate" in kwargs:
      annotate = kwargs["annotate"]
      del kwargs["annotate"]

    if annotate:
      if len(args) != 2:
        raise libadjoint.exceptions.LibadjointErrorInvalidInputs("The annotated LUSolver.solve must be called like solve(x, b).")

      A = self.operator

      try:
        x = args[0].function
      except AttributeError:
        raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your solution x has to have a .function attribute; is it the .vector() of a Function?")

      try:
        b = args[1].form
      except AttributeError:
        raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your RHS b has to have the .form attribute: was it assembled after from dolfin_adjoint import *?")

      try:
        eq_bcs = list(set(A.bcs + b.bcs))
      except AttributeError:
        assert not hasattr(A, 'bcs') and not hasattr(b, 'bcs')
        eq_bcs = []

      if self.parameters["reuse_factorization"]:
        lu_solvers[A] = self
        adj_lu_solvers[A] = None

      solving.annotate(A == b, x, eq_bcs, solver_parameters={"linear_solver": "lu"}, matrix_class=make_LUSolverMatrix(A, self.parameters["reuse_factorization"]))

    out = dolfin.LUSolver.solve(self, *args, **kwargs)

    if annotate:
      if solving.debugging["record_all"]:
        solving.adjointer.record_variable(solving.adj_variables[x], libadjoint.MemoryStorage(solving.Vector(x)))

    return out