Hello,
I am attaching a simple example that uses TSAdjoint to calculate the
sensitivity of a 1 degree of freedom ODE. When using theta methods, it returns
the right sensitivity as given by the analytical solution. When I set
`ts.setThetaEndpoint(True)`, it does not work. Is there a theoretical reason
why? I was using this option because using the options
ts.setType(ts.Type.THETA)
ts.setTheta(0.5)
ts.setThetaEndpoint(True)
is equivalent to using ts.Type.CN
Thanks
Miguel
Miguel A. Salazar de Troya
Postdoctoral Researcher, Lawrence Livermore National Laboratory
B141
Rm: 1085-5
Ph: 1(925) 422-6411
# Testing TSAdjoint and matrix-free Jacobian
import sys, petsc4py
petsc4py.init(sys.argv)
import numpy as np
from petsc4py import PETSc
class VDP(object):
n = 1
comm = PETSc.COMM_SELF
def __init__(self, mu_=1.0, mf_=False):
self.mu_ = mu_
self.mf_ = mf_
self.init_cond = 10.0
if self.mf_:
self.J_ = PETSc.Mat().createDense([self.n, self.n], comm=self.comm)
self.J_.setUp()
self.Jp_ = PETSc.Mat().createDense([self.n, 1], comm=self.comm)
self.Jp_.setUp()
def initialCondition(self, u):
u[0] = self.init_cond
u.assemble()
def evalIFunction(self, ts, t, u, udot, f):
mu = self.mu_
f[0] = udot[0] - mu * u[0]
f.assemble()
def evalIJacobian(self, ts, t, u, udot, sigma, A, B):
if not self.mf_:
J = A
else:
J = self.J_
mu = self.mu_
J[0, 0] = sigma * 1.0 - mu
J.assemble()
if A != B:
B.assemble()
return True # same nonzero pattern
def evalFunction(self, ts, t, u, f):
mu = self.mu_
f[0] = u[0] * mu
f.assemble()
def evalJacobian(self, ts, t, u, A, B):
if not self.mf_:
J = A
else:
J = self.J_
mu = self.mu_
J[0, 0] = mu
J.assemble()
if A != B:
B.assemble()
return True # same nonzero pattern
def evalJacobianP(self, ts, t, u, C):
if not self.mf_:
Jp = C
else:
Jp = self.Jp_
Jp[0, 0] = u[0]
Jp.assemble()
return True
def integrand(self, ts, t, u, f):
f[0] = u[0] * self.mu_ * self.mu_
f.assemble()
def formRHSintegrandJacobian(self, ts, t, u, J, P):
J[0, 0] = 1.0 * self.mu_ * self.mu_
J.assemble()
return True
def formRHSintegrandJacobianP(self, ts, t, u, J):
J[0, 0] = u[0] * self.mu_ * 2.0
J.assemble()
return True
class JacShell:
def __init__(self, ode):
self.ode_ = ode
def mult(self, A, x, y):
"y <- A * x"
self.ode_.J_.mult(x, y)
def multTranspose(self, A, x, y):
"y <- A' * x"
self.ode_.J_.multTranspose(x, y)
class JacPShell:
def __init__(self, ode):
self.ode_ = ode
def multTranspose(self, A, x, y):
"y <- A' * x"
self.ode_.Jp_.multTranspose(x, y)
OptDB = PETSc.Options()
mu_ = OptDB.getScalar("mu", 2.0e0)
mf_ = OptDB.getBool("mf", False)
ode = VDP(mu_, mf_)
if not mf_:
J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm)
J.setUp()
Jp = PETSc.Mat().createDense([ode.n, 1], comm=ode.comm)
Jp.setUp()
else:
J = PETSc.Mat().create()
J.setSizes([ode.n, ode.n])
J.setType("python")
shell = JacShell(ode)
J.setPythonContext(shell)
J.setUp()
J.assemble()
Jp = PETSc.Mat().create()
Jp.setSizes([ode.n, 1])
Jp.setType("python")
shell = JacPShell(ode)
Jp.setPythonContext(shell)
Jp.setUp()
Jp.assemble()
u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
f = u.duplicate()
adj_u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
adj_p = PETSc.Vec().createSeq(1, comm=ode.comm)
ts = PETSc.TS().create(comm=ode.comm)
ts.setType(ts.Type.THETA)
ts.setTheta(0.5) # adjoint for 1.0 (backward Euler) is currently broken in PETSc
#ts.setThetaEndpoint(True)
ts.setProblemType(False.ProblemType.LINEAR)
ts.setProblemType(ts.EquationType.IMPLICIT)
ts.setIFunction(ode.evalIFunction, f)
ts.setIJacobian(ode.evalIJacobian, J)
ts.setRHSJacobianP(ode.evalJacobianP, Jp)
integral = True
if integral:
ts.createQuadratureTS()
tsq = ts.getQuadratureTS()[1]
Jint = J.duplicate()
Jintp = J.duplicate()
tsq.setRHSFunction(ode.integrand)
tsq.setRHSJacobian(ode.formRHSintegrandJacobian, Jint)
tsq.setRHSJacobianP(ode.formRHSintegrandJacobianP, Jintp)
ts.setSaveTrajectory()
ts.setTime(0.0)
ts.setTimeStep(0.001)
final_T = 1.0
ts.setMaxTime(final_T)
ts.setMaxSteps(1000)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
ts.setFromOptions()
ode.initialCondition(u)
ts.solve(u)
print("Cost function")
a, b = ode.init_cond, ode.mu_
if integral:
ts.getCostIntegral().view()
print("Exact value: {0:.5f}".format(a * b * (np.exp(b * final_T) - 1.0)))
else:
print(u.array[0])
print("Exact value: {0:.5f}".format(a * (np.exp(b * final_T))))
# I don't understand why I need to define it as 1.0 for the sensitivities of the final state to work
# The reason is because of the way PETSc distingushes between time integrals
# and just scalars functions of the final state
if integral:
adj_u[0] = 0.0
else:
adj_u[0] = 1.0
adj_u.assemble()
adj_p[0] = 0
adj_p.assemble()
ts.setCostGradients(adj_u, adj_p)
ts.adjointSolve()
# Real sensitivity of the final solution
if integral:
sensitivity_a = 1.0 * b * (-1.0 + np.exp(b * final_T))
sensitivity_b = a * b * final_T * np.exp(b * final_T) + a * (np.exp(b * final_T) - 1)
else:
sensitivity_a = np.exp(b * final_T)
sensitivity_b = a * final_T * np.exp(b * final_T)
print(" ")
print("Numerical sensitivity w.r.t. a")
print(adj_u.array[0])
print("Real sensitivity w.r.t a")
print(sensitivity_a)
print(" ")
print("Numerical sensitivity w.r.t. b")
if integral:
print(adj_p.array[0])
else:
print(adj_p.array[0])
print("Real sensitivity w.r.t b")
print(sensitivity_b)
