def evalGradientParameter(self,x, mg, misfit_only=False):
self.prior.init_vector(mg,1)
if misfit_only == False:
dm = x[PARAMETER] - self.prior.mean
self.prior.R.mult(dm, mg)
else:
mg.zero()
p0 = dl.Vector()
self.M.init_vector(p0,0)
x[ADJOINT].retrieve(p0, self.simulation_times[1])
mg.axpy(-1., self.Mt_stab*p0)
g = dl.Vector()
self.M.init_vector(g,1)
self.prior.Msolver.solve(g,mg)
grad_norm = g.inner(mg)
return grad_norm
def testRTObservations(self):
Vh = dl.FunctionSpace(self.mesh, "RT", 1)
xvect = dl.interpolate(dl.Expression(("x[0]", "x[1]"), degree=1), Vh).vector()
B = assemblePointwiseObservation(Vh, self.targets, prune_and_sort=False)
out = dl.Vector()
B.init_vector(out,0)
B.mult(xvect, out)
out_np = out.gather_on_zero()
if self.rank == 0:
assert_allclose(self.targets, np.reshape(out_np, (self.ntargets, self.ndim), 'C'))
def solveFwdIncremental(self, sol, rhs):
"""
Solve the 2nd order forward equation
"""
sol.zero()
uold = dl.Vector()
u = dl.Vector()
Muold = dl.Vector()
myrhs = dl.Vector()
self.M.init_vector(uold, 0)
self.M.init_vector(u, 0)
self.M.init_vector(Muold, 0)
self.M.init_vector(myrhs, 0)
for t in self.simulation_times[1::]:
self.M_stab.mult(uold, Muold)
rhs.retrieve(myrhs, t)
myrhs.axpy(1., Muold)
self.solver.solve(u, myrhs)
sol.store(u, t)
uold = u
self.soln_count[2] += 1
def solveAdjIncremental(self, sol, rhs):
"""
Solve the 2nd order adjoint equation
"""
sol.zero()
pold = dl.Vector()
p = dl.Vector()
Mpold = dl.Vector()
myrhs = dl.Vector()
self.M.init_vector(pold, 0)
self.M.init_vector(p, 0)
self.M.init_vector(Mpold, 0)
self.M.init_vector(myrhs, 0)
for t in self.simulation_times[::-1]:
self.Mt_stab.mult(pold, Mpold)
rhs.retrieve(myrhs, t)
Mpold.axpy(1., myrhs)
self.solvert.solve(p, Mpold)
pold = p
sol.store(p, t)
self.soln_count[3] += 1
def proposal(self, current):
#Generate sample from the prior
parRandom.normal(1., self.noise)
w = dl.Vector(self.model.prior.R.mpi_comm())
self.model.prior.init_vector(w, 0)
self.model.prior.sample(self.noise, w, add_mean=False)
# do pCN linear combination with current sample
s = self.parameters["s"]
w *= s
w.axpy(1., self.model.prior.mean)
w.axpy(np.sqrt(1. - s * s), current.m - self.model.prior.mean)
return w
def cost(self, x):
if self.noise_variance is None:
raise ValueError("Noise Variance must be specified")
elif self.noise_variance == 0:
raise ZeroDivisionError(
"Noise Variance must not be 0.0 Set to 1.0 for deterministic inverse problems"
)
r = self.d.copy()
r.axpy(-1., x[STATE])
Wr = dl.Vector(self.W.mpi_comm())
self.W.init_vector(Wr, 0)
self.W.mult(r, Wr)
return r.inner(Wr) / (2. * self.noise_variance)
def test_grad(self):
#Ensure cost gradient is correct
sample = get_prior_sample(self.test_prior)
actual_grad = dl.Vector()
self.test_prior.init_vector(actual_grad, 0)
self.test_prior.grad(sample, actual_grad)
expected_grad_np = np.dot(self.precision,
sample.get_local() - self.means)
self.assertTrue(np.allclose(expected_grad_np, actual_grad.get_local()))
def applyC(self, da, out):
out.zero()
myout = dl.Vector()
self.M.init_vector(myout, 0)
self.M_stab.mult(da, myout)
myout *= -1.
t = self.t_init + self.dt
out.store(myout, t)
myout.zero()
while t < self.t_final:
t += self.dt
out.store(myout, t)
def testScalarObservations(self):
Vh = dl.FunctionSpace(self.mesh, "CG", 1)
xvect = dl.interpolate(dl.Expression("x[0]", degree=1), Vh).vector()
B = assemblePointwiseObservation(Vh, self.targets, prune_and_sort=False)
out = dl.Vector()
B.init_vector(out, 0)
B.mult(xvect, out)
out_np = out.gather_on_zero()
if self.rank == 0:
assert_allclose(self.targets[:,0], out_np)
def __init__(self,
parameter_init,
model,
step_size,
step_num,
alg_name,
adpt_h=False,
**kwargs):
"""
Initialization
"""
# parameters
self.q = parameter_init
self.dim = self.q.size()
self.model = model
self.noise = dl.Vector()
self.model.prior.init_vector(self.noise, "noise")
target_acpt = kwargs.pop('target_acpt', 0.65)
# geometry needed
geom_ord = [0]
if any(s in alg_name for s in ['MALA', 'HMC']): geom_ord.append(1)
if any(s in alg_name for s in ['mMALA', 'mHMC']): geom_ord.append(2)
self.geom = lambda parameter: self.model.get_geom(
parameter, geom_ord=geom_ord, **kwargs)
if '_' in alg_name:
self.ll, self.g, self.HessApply, self.eigs = self.geom(
self.q) # for infmMALA and infmHMC
else:
self.ll, self.g, _, self.eigs = self.geom(self.q)
# sampling setting
self.h = step_size
self.L = step_num
if 'HMC' not in alg_name: self.L = 1
self.alg_name = alg_name
# optional setting for adapting step size
self.adpt_h = adpt_h
if self.adpt_h:
h_adpt = {}
# h_adpt['h']=self._init_h()
h_adpt['h'] = self.h
h_adpt['mu'] = np.log(10 * h_adpt['h'])
h_adpt['loghn'] = 0.
h_adpt['An'] = 0.
h_adpt['gamma'] = 0.05
h_adpt['n0'] = 10
h_adpt['kappa'] = 0.75
h_adpt['a0'] = target_acpt
self.h_adpt = h_adpt
def __init__(self, mesh, Vh, prior):
"""
Construct a model by proving
- the mesh
- the finite element spaces for the STATE/ADJOINT variable and the PARAMETER variable
- the prior information
"""
self.mesh = mesh
self.Vh = Vh
# Initialize Expressions
mtrue_exp = dl.Expression(
'0.2*(15.0 - 5.0*sin(3.1416*(x[1]/20.0 - 0.5)))',
element=Vh[PARAMETER].ufl_element())
self.mtrue = dl.interpolate(mtrue_exp, self.Vh[PARAMETER]).vector()
self.f = dl.Expression(("0.0", "0.0022", "0.0"),
element=Vh[STATE].ufl_element())
self.u_o = dl.Vector()
self.u_bdr = dl.Expression(("0.0", "0.0", "0.0"),
element=Vh[STATE].ufl_element())
self.u_bdr0 = dl.Expression(("0.0", "0.0", "0.0"),
element=Vh[STATE].ufl_element())
self.bc = dl.DirichletBC(self.Vh[STATE], self.u_bdr, Gamma_D)
self.bc0 = dl.DirichletBC(self.Vh[STATE], self.u_bdr0, Gamma_D)
# Assemble constant matrices
self.prior = prior
self.Wuu = self.assembleWuu()
self.computeObservation(self.u_o)
self.A = None
self.At = None
self.C = None
self.Wmm = None
self.Wmu = None
self.gauss_newton_approx = True
self.solver = PETScKrylovSolver(mesh.mpi_comm(), "cg", amg_method())
self.solver_fwd_inc = PETScKrylovSolver(mesh.mpi_comm(), "cg",
amg_method())
self.solver_adj_inc = PETScKrylovSolver(mesh.mpi_comm(), "cg",
amg_method())
self.solver.parameters["relative_tolerance"] = 1e-9
self.solver.parameters["absolute_tolerance"] = 1e-12
self.solver_fwd_inc.parameters = self.solver.parameters
self.solver_adj_inc.parameters = self.solver.parameters
def setUp(self):
mesh = dl.UnitSquareMesh(10, 10)
self.mpi_rank = dl.MPI.rank(mesh.mpi_comm())
self.mpi_size = dl.MPI.size(mesh.mpi_comm())
Vh = dl.FunctionSpace(mesh, 'Lagrange', 1)
uh,vh = dl.TrialFunction(Vh),dl.TestFunction(Vh)
varfM = ufl.inner(uh,vh)*ufl.dx
self.M = dl.assemble(varfM)
x = dl.Vector( mesh.mpi_comm() )
self.M.init_vector(x,0)
k = 121
self.Q = MultiVector(x,k)
def generate_vector(self, component="ALL"):
if component == "ALL":
u = TimeDependentVector(self.sim_times)
u.initialize(self.Q, 0)
a = dl.Vector()
self.Prior.init_vector(a, 0)
p = TimeDependentVector(self.sim_times)
p.initialize(self.Q, 0)
return [u, a, p]
elif component == STATE:
u = TimeDependentVector(self.sim_times)
u.initialize(self.Q, 0)
return u
elif component == PARAMETER:
a = dl.Vector()
self.Prior.init_vector(a, 0)
return a
elif component == ADJOINT:
p = TimeDependentVector(self.sim_times)
p.initialize(self.Q, 0)
return p
else:
raise
def evalFunction(self, ts, t, x, xdot, b):
"""The callback that the TS executes to compute the residual."""
self.update(t)
self.update_x(x)
self.update_xdot(xdot)
b1 = df.Vector()
b2 = df.Vector()
#self.assembler.assemble(self.F_fluid_form, tensor = b1)
df.assemble(self.F_fluid_form, tensor=b1)
[bc.apply(b1) for bc in self.bcs_mesh]
#self.assembler.assemble(self.F_solid_form, tensor = b2)
df.assemble(self.F_solid_form, tensor=b2)
b_wrap = df.PETScVector(b)
# zero all entries
b_wrap.zero()
#df.assemble(b_wrap, self.x)
b_wrap.axpy(1, b1)
b_wrap.axpy(1, b2)
[bc.apply(b_wrap, self.x) for bc in self.bcs]
def cost(self, x):
Rdx = dl.Vector()
self.Prior.init_vector(Rdx, 0)
dx = x[PARAMETER] - self.Prior.mean
self.Prior.R.mult(dx, Rdx)
reg = .5 * Rdx.inner(dx)
u = dl.Vector()
ud = dl.Vector()
self.Q.init_vector(u, 0)
self.Q.init_vector(ud, 0)
misfit = 0
for t in np.arange(self.t_1, self.t_final + (.5 * self.dt), self.dt):
x[STATE].retrieve(u, t)
self.ud.retrieve(ud, t)
diff = u - ud
Qdiff = self.Q * diff
misfit += .5 / self.noise_variance * Qdiff.inner(diff)
c = misfit + reg
return [c, reg, misfit]
def generate_vector(self, component = "ALL"):
if component == "ALL":
u = TimeDependentVector(self.simulation_times)
u.initialize(self.M, 0)
m = dl.Vector()
self.prior.init_vector(m,0)
p = TimeDependentVector(self.simulation_times)
p.initialize(self.M, 0)
return [u, m, p]
elif component == STATE:
u = TimeDependentVector(self.simulation_times)
u.initialize(self.M, 0)
return u
elif component == PARAMETER:
m = dl.Vector()
self.prior.init_vector(m,0)
return m
elif component == ADJOINT:
p = TimeDependentVector(self.simulation_times)
p.initialize(self.M, 0)
return p
else:
raise
def __init__(self,
parameters=CGSolverSteihaug_ParameterList(),
comm=dl.MPI.comm_world):
self.parameters = parameters
self.A = None
self.B_solver = None
self.B_op = None
self.converged = False
self.iter = 0
self.reasonid = 0
self.final_norm = 0
self.TR_radius_2 = None
self.update_x = self.update_x_without_TR
self.r = dl.Vector(comm)
self.z = dl.Vector(comm)
self.Ad = dl.Vector(comm)
self.d = dl.Vector(comm)
self.Bx = dl.Vector(comm)
def gather_vector(u, size=None):
comm = mpi_comm_world()
if size is None:
# size = int(MPI.size(comm) * MPI.sum(comm, u.size()))
size = int(MPI.sum(comm, u.size()))
# From this post: https://fenicsproject.discourse.group/t/gather-function-in-parallel-error/1114/4
u_vec = dolfin.Vector(comm, size)
# Values from everywhere on 0
u_vec = u.gather_on_zero()
# To everywhere from 0
mine = comm.bcast(u_vec)
# Reconstruct
if comm.rank == 0:
x = u_vec
else:
v = dolfin.Vector(MPI.comm_self, size)
v.set_local(mine)
x = v.get_local()
return x
def __init__(self, Vh_STATE, Vhs, weights, geo, bcs0, datafile, variance_u,
variance_g):
if hasattr(geo, "dx"):
self.dx = geo.dx(geo.PHYSICAL)
else:
self.dx = dl.dx
self.ds = geo.ds(geo.AXIS)
x, y, U, V, uu, vv, ww, uv, k = np.loadtxt(datafile,
skiprows=2,
unpack=True)
u_fun_data = VelocityDNS(x=x,
y=y,
U=U,
V=V,
symmetrize=True,
coflow=0.)
u_data = dl.interpolate(u_fun_data, Vhs[0])
if Vh_STATE.num_sub_spaces() == 3:
u_trial, p_trial, g_trial = dl.TrialFunctions(Vh_STATE)
u_test, p_test, g_test = dl.TestFunctions(Vh_STATE)
else:
raise InputError()
Wform = dl.Constant(1./variance_u)*dl.inner(u_trial, u_test)*self.dx +\
dl.Constant(1./variance_g)*g_trial*g_test*self.ds
self.W = dl.assemble(Wform)
dummy = dl.Vector()
self.W.init_vector(dummy, 0)
[bc.zero(self.W) for bc in bcs0]
Wt = Transpose(self.W)
[bc.zero(Wt) for bc in bcs0]
self.W = Transpose(Wt)
xfun = dl.Function(Vh_STATE)
assigner = dl.FunctionAssigner(Vh_STATE, Vhs)
assigner.assign(xfun, [
u_data,
dl.Function(Vhs[1]),
dl.interpolate(dl.Constant(1.), Vhs[2])
])
self.d = xfun.vector()
self.w = (weights * 0.5)
def __init__(self, Vh_STATE, Vhs, bcs0, datafile, dx=dl.dx):
self.dx = dx
x, y, U, V, uu, vv, ww, uv, k = np.loadtxt(datafile,
skiprows=2,
unpack=True)
u_fun_mean = VelocityDNS(x=x,
y=y,
U=U,
V=V,
symmetrize=True,
coflow=0.)
u_fun_data = VelocityDNS(x=x,
y=y,
U=U,
V=V,
symmetrize=False,
coflow=0.)
k_fun_mean = KDNS(x=x, y=y, k=k, symmetrize=True)
k_fun_data = KDNS(x=x, y=y, k=k, symmetrize=False)
u_data = dl.interpolate(u_fun_data, Vhs[0])
k_data = dl.interpolate(k_fun_data, Vhs[2])
noise_var_u = dl.assemble(
dl.inner(u_data - u_fun_mean, u_data - u_fun_mean) * self.dx)
noise_var_k = dl.assemble(
dl.inner(k_data - k_fun_mean, k_data - k_fun_mean) * self.dx)
u_trial, p_trial, k_trial, e_trial = dl.TrialFunctions(Vh_STATE)
u_test, p_test, k_test, e_test = dl.TestFunctions(Vh_STATE)
Wform = dl.Constant(1./noise_var_u)*dl.inner(u_trial, u_test)*self.dx + \
dl.Constant(1./noise_var_k)*dl.inner(k_trial, k_test)*self.dx
self.W = dl.assemble(Wform)
dummy = dl.Vector()
self.W.init_vector(dummy, 0)
[bc.zero(self.W) for bc in bcs0]
Wt = Transpose(self.W)
[bc.zero(Wt) for bc in bcs0]
self.W = Transpose(Wt)
xfun = dl.Function(Vh_STATE)
assigner = dl.FunctionAssigner(Vh_STATE, Vhs)
assigner.assign(
xfun,
[u_data, dl.Function(Vhs[1]), k_data,
dl.Function(Vhs[3])])
self.d = xfun.vector()
def get_diagonal(A, d):
"""
Compute the diagonal of the square operator :math:`A`.
Use :code:`Solver2Operator` if :math:`A^{-1}` is needed.
"""
ej, xj = dl.Vector(d.mpi_comm()), dl.Vector(d.mpi_comm())
A.init_vector(ej, 1)
A.init_vector(xj, 0)
g_size = ej.size()
d.zero()
for gid in range(g_size):
owns_gid = ej.owns_index(gid)
if owns_gid:
SetToOwnedGid(ej, gid, 1.)
ej.apply("insert")
A.mult(ej, xj)
if owns_gid:
val = GetFromOwnedGid(xj, gid)
SetToOwnedGid(d, gid, val)
SetToOwnedGid(ej, gid, 0.)
ej.apply("insert")
d.apply("insert")
def assembleWau(self, x):
"""
Assemble the derivative of the parameter equation with respect to the state
"""
trial = dl.TrialFunction(self.Vh[STATE])
test = dl.TestFunction(self.Vh[PARAMETER])
a = dl.Function(self.Vh[ADJOINT], x[ADJOINT])
c = dl.Function(self.Vh[PARAMETER], x[PARAMETER])
varf = dl.inner(dl.exp(c) * dl.nabla_grad(trial),
dl.nabla_grad(a)) * test * dl.dx
Wau = dl.assemble(varf)
dummy = dl.Vector()
Wau.init_vector(dummy, 0)
self.bc0.zero_columns(Wau, dummy)
return Wau
def estimate_diagonal_inv2(Asolver, k, d, init_vector=None):
"""
An unbiased stochastic estimator for the diagonal of :math:`A^{-1}`.
:math:`d = [ \sum_{j=1}^k v_j .* A^{-1} v_j ] ./ [ \sum_{j=1}^k v_j .* v_j ]`
where
- :math:`v_j` are i.i.d. :math:`\mathcal{N}(0, I)`
- :math:`.*` and :math:`./` represent the element-wise multiplication and division
of vectors, respectively.
Reference:
`Costas Bekas, Effrosyni Kokiopoulou, and Yousef Saad,
An estimator for the diagonal of a matrix,
Applied Numerical Mathematics, 57 (2007), pp. 1214-1229.`
"""
x, b = df.Vector(d.mpi_comm()), df.Vector(d.mpi_comm())
if init_vector:
init_vector(x, 1)
init_vector(b, 0)
else:
if hasattr(Asolver, "init_vector"):
Asolver.init_vector(x, 1)
Asolver.init_vector(b, 0)
else:
Asolver.get_operator().init_vector(x, 1)
Asolver.get_operator().init_vector(b, 0)
d.zero()
for i in range(k):
x.zero()
# parRandom.normal(1., b)
b.set_local(np.random.normal(size=b.size())) # serial hack
Asolver.solve(x, b)
x *= b
d.axpy(1. / float(k), x)
def __init__(self, model, nu):
"""
- model: an object of type Model
- nu: an object of type GaussianLRPosterior
"""
self.model = model
self.nu = nu
self.parameters = {}
self.parameters["inner_rel_tolerance"] = 1e-9
self.noise = dl.Vector()
self.model.prior.init_vector(self.noise, "noise")
self.discard = self.model.generate_vector(PARAMETER)
self.w = self.model.generate_vector(PARAMETER)
def step(self, t0: float, t1: float) -> None:
r"""
Solve on the given time step (t0, t1).
*Arguments*
interval (:py:class:`tuple`)
The time interval (t0, t1) for the step
*Invariants*
Assuming that v\_ is in the correct state for t0, gives
self.vur in correct state at t1.
"""
dt = t1 - t0
theta = self._parameters["theta"]
t = t0 + theta * dt
self.time.assign(t)
# Update matrix and linear solvers etc as needed
if self._timestep is None:
self._timestep = df.Constant(dt)
self._lhs, self._rhs = self.variational_forms(self._timestep)
# Preassemble left-hand side and initialize right-hand side vector
self._lhs_matrix = df.assemble(
self._lhs, keep_diagonal=True) # TODO: Why diagonal?
self._rhs_vector = df.Vector(self._mesh.mpi_comm(),
self._lhs_matrix.size(0))
self._lhs_matrix.init_vector(self._rhs_vector, 0)
# Create linear solver (based on parameter choices)
self._linear_solver = self._create_linear_solver()
else:
self._update_solver(dt)
# Assemble right-hand-side
df.assemble(self._rhs, tensor=self._rhs_vector)
# Apply BCs
for bc in self._bcs:
bc.apply(self._lhs_matrix, self._rhs_vector)
extracellular_indices = np.arange(0, self._rhs_vector.local_size(), 2)
rhs_norm = self._rhs_vector.get_local()[extracellular_indices].sum()
rhs_norm /= extracellular_indices.size
self._rhs_vector.get_local()[extracellular_indices] -= rhs_norm
# Solve problem
self.linear_solver.solve(self.vur.vector(), self._rhs_vector)
def __init__(self, problem, problem2, Q, Qoper, Gamma_z):
"""
Construct the Jacobian operator
"""
self.problem = problem
self.problem2 = problem2
self.misfit = problem.misfit
self.misfit2 = problem2.misfit
self.u_snapshot = dl.Vector()
self.uhat = problem2.generate_vector(STATE)
self.Bv = problem2.generate_vector(ADJOINT)
self.FCprv = problem2.generate_vector(STATE)
self.rhs_fwd = self.problem2.generate_vector(STATE)
self.Q = Q
self.Qoper = Qoper
self.Gamma_z = Gamma_z
def __init__(self, times, tol=1e-10, mpi_comm=dl.MPI.comm_world):
"""
Constructor:
- :code:`times`: time frame at which snapshots are stored.
- :code:`tol` : tolerance to identify the frame of the snapshot.
"""
self.nsteps = len(times)
self.data = []
for i in range(self.nsteps):
self.data.append(dl.Vector(mpi_comm))
self.times = times
self.tol = tol
self.mpi_comm = mpi_comm
def __init__(self, mesh, Vh, prior):
"""
Construct a model by proving
- the mesh
- the finite element spaces for the STATE/ADJOINT variable and the PARAMETER variable
- the prior information
"""
self.mesh = mesh
self.Vh = Vh
# Initialize Expressions
mtrue_exp = dl.Expression(
'std::log(2 + 7*(std::pow(std::pow(x[0] - 0.5,2) + std::pow(x[1] - 0.5,2),0.5) > 0.2))',
element=Vh[PARAMETER].ufl_element())
self.mtrue = dl.interpolate(mtrue_exp, self.Vh[PARAMETER]).vector()
self.f = dl.Constant(1.0)
self.u_o = dl.Vector()
self.u_bdr = dl.Constant(0.0)
self.u_bdr0 = dl.Constant(0.0)
self.bc = dl.DirichletBC(self.Vh[STATE], self.u_bdr, u_boundary)
self.bc0 = dl.DirichletBC(self.Vh[STATE], self.u_bdr0, u_boundary)
# Assemble constant matrices
self.prior = prior
self.Wuu = self.assembleWuu()
self.computeObservation(self.u_o)
self.A = None
self.At = None
self.C = None
self.Wmm = None
self.Wmu = None
self.gauss_newton_approx = False
self.solver = PETScKrylovSolver(mesh.mpi_comm(), "cg", amg_method())
self.solver_fwd_inc = PETScKrylovSolver(mesh.mpi_comm(), "cg",
amg_method())
self.solver_adj_inc = PETScKrylovSolver(mesh.mpi_comm(), "cg",
amg_method())
self.solver.parameters["relative_tolerance"] = 1e-15
self.solver.parameters["absolute_tolerance"] = 1e-20
self.solver_fwd_inc.parameters = self.solver.parameters
self.solver_adj_inc.parameters = self.solver.parameters
def trace(A, mpi_comm=dl.MPI.comm_world):
"""
Compute the trace of a sparse matrix :math:`A`.
"""
v = dl.Vector(mpi_comm)
A.init_vector(v)
nprocs = dl.MPI.size(mpi_comm)
if nprocs > 1:
raise Exception("trace is only serial")
n = A.size(0)
tr = 0.
for i in range(0, n):
[j, val] = A.getrow(i)
tr += val[j == i]
return tr
def pointwise_variance(self, method, k=1000000):
"""
Compute/Estimate the prior pointwise variance.
- If method=="Exact" we compute the diagonal entries of R^{-1} entry by entry.
This requires to solve n linear system in R (not scalable, but ok for illustration purposes).
"""
pw_var = dl.Vector()
self.init_vector(pw_var, 0)
if method == "Exact":
get_diagonal(self.Rsolver, pw_var, solve_mode=True)
elif method == "Estimator":
estimate_diagonal_inv2(self.Rsolver, k, pw_var)
else:
raise NameError("Unknown method")
return pw_var
