def test_cmscr1d_weak_solution_zero_Dirichlet(self):
print("Running test 'test_cmscr1d_weak_solution_zero_Dirichlet'")
# Define temporal and spatial sample points.
m, n = 10, 20
# Define mesh and function space.
mesh = UnitSquareMesh(m - 1, n - 1)
V = dh.create_function_space(mesh, 'default')
W = dh.create_vector_function_space(mesh, 'default')
# Define boundary conditions for velocity.
bc = DirichletBC(W.sub(0), Constant(0), dh.DirichletBoundary())
# Create zero function.
f = Function(V)
# Compute velocity.
v, k, res, fun, converged = cmscr1d_weak_solution(W, f,
f.dx(0), f.dx(1),
1.0, 1.0,
1.0, 1.0, 1.0,
bcs=bc)
v = v.vector().get_local()
k = k.vector().get_local()
np.testing.assert_allclose(v.shape, m * n)
np.testing.assert_allclose(v, np.zeros_like(v))
np.testing.assert_allclose(k.shape, m * n)
np.testing.assert_allclose(k, np.zeros_like(k))
def test_cms1d_weak_solution_given_velocity_default(self):
# Define temporal and spatial sample points.
m, n = 10, 20
# Define mesh and function space.
mesh = UnitSquareMesh(m - 1, n - 1)
V = dh.create_function_space(mesh, 'default')
# Create zero function.
f = Function(V)
# Create velocity.
v = Function(V)
vx = Function(V)
# Compute source.
k, res, fun = cms1d_weak_solution_given_velocity(
V, f, f.dx(0), f.dx(1), v, vx, 1.0, 1.0)
k = k.vector().get_local()
np.testing.assert_allclose(k.shape, m * n)
np.testing.assert_allclose(k, np.zeros_like(k))
V = dh.create_function_space(mesh, 'periodic')
# Create zero function.
f = Function(V)
# Compute source.
k, res, fun = cms1d_weak_solution_given_velocity(
V, f, f.dx(0), f.dx(1), v, vx, 1.0, 1.0)
k = k.vector().get_local()
np.testing.assert_allclose(k.shape, m * (n - 1))
np.testing.assert_allclose(k, np.zeros_like(k))
def test_cm1d_weak_solution_default(self):
# Define temporal and spatial sample points.
m, n = 10, 20
# Define mesh and function space.
mesh = UnitSquareMesh(m - 1, n - 1)
V = dh.create_function_space(mesh, 'default')
# Create zero function.
f = Function(V)
# Compute velocity.
v, res, fun = cm1d_weak_solution(V, f, f.dx(0), f.dx(1), 1.0, 1.0)
v = v.vector().get_local()
np.testing.assert_allclose(v.shape, m * n)
np.testing.assert_allclose(v, np.zeros_like(v))
np.testing.assert_allclose(res, 0)
np.testing.assert_allclose(fun, 0)
V = dh.create_function_space(mesh, 'periodic')
# Create zero function.
f = Function(V)
# Compute velocity.
v, res, fun = cm1d_weak_solution(V, f, f.dx(0), f.dx(1), 1.0, 1.0)
v = v.vector().get_local()
np.testing.assert_allclose(v.shape, m * (n - 1))
np.testing.assert_allclose(v, np.zeros_like(v))
np.testing.assert_allclose(res, 0)
np.testing.assert_allclose(fun, 0)
def cmscr1d_img_pb(img: np.array,
alpha0: float, alpha1: float,
alpha2: float, alpha3: float,
beta: float, deriv='mesh') \
-> (np.array, np.array, float, float, bool):
"""Computes the L2-H1 mass conserving flow with source for a 1D image
sequence with spatio-temporal and convective regularisation with periodic
spatial boundary.
Args:
img (np.array): 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
alpha0 (float): The spatial regularisation parameter for v.
alpha1 (float): The temporal regularisation parameter for v.
alpha2 (float): The spatial regularisation parameter for k.
alpha3 (float): The temporal regularisation parameter for k.
beta (float): The convective regularisation parameter.
deriv (str): Specifies how to approximate pertial derivatives.
When set to 'mesh' it uses FEniCS built in function.
Returns:
v (np.array): A velocity array of shape (m, n).
k (np.array): A source array of shape (m, n).
res (float): The residual.
fun (float): The function value.
converged (bool): True if Newton's method converged.
"""
# Check for valid arguments.
valid = {'mesh'}
if deriv not in valid:
raise ValueError("Argument 'deriv' must be one of %r." % valid)
# Create mesh.
m, n = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space.
V = dh.create_function_space(mesh, 'periodic')
W = dh.create_vector_function_space(mesh, 'periodic')
# Convert array to function.
f = Function(V)
f.vector()[:] = dh.img2funvec_pb(img)
# Compute partial derivatives.
ft, fx = f.dx(0), f.dx(1)
# Compute velocity.
v, k, res, fun, converged = cmscr1d_weak_solution(W, f, ft, fx, alpha0,
alpha1, alpha2, alpha3,
beta)
# Convert back to array and return.
v = dh.funvec2img(v.vector().get_local(), m, n)
k = dh.funvec2img(k.vector().get_local(), m, n)
return v, k, res, fun, converged
def of1d_img(img: np.array, alpha0: float, alpha1: float, deriv) \
-> (np.array, float, float):
"""Computes the L2-H1 optical flow for a 1D image sequence.
Takes a one-dimensional image sequence and returns a minimiser of the
Horn-Schunck functional with spatio-temporal regularisation.
Allows to specify how to approximate partial derivatives of f numerically.
Args:
img (np.array): 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
alpha0 (float): Spatial regularisation parameter.
alpha1 (float): Temporal regularisation parameter.
deriv (str): Specifies how to approximate pertial derivatives.
When set to 'mesh' it uses FEniCS built in function.
When set to 'fd' it uses finite differences.
Returns:
v (np.array): A velocity array of shape (m, n).
res (float): The residual.
fun (float): The function value.
"""
# Check for valid arguments.
valid = {'mesh', 'fd'}
if deriv not in valid:
raise ValueError("Argument 'deriv' must be one of %r." % valid)
# Create mesh.
m, n = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space.
V = dh.create_function_space(mesh, 'default')
# Convert array to function.
f = Function(V)
f.vector()[:] = dh.img2funvec(img)
# Compute partial derivatives.
if deriv is 'mesh':
ft, fx = f.dx(0), f.dx(1)
if deriv is 'fd':
imgt, imgx = nh.partial_derivatives(img)
ft, fx = Function(V), Function(V)
ft.vector()[:] = dh.img2funvec(imgt)
fx.vector()[:] = dh.img2funvec(imgx)
# Compute velocity.
v, res, fun = of1d_weak_solution(V, f, ft, fx, alpha0, alpha1)
# Convert to array and return.
return dh.funvec2img(v.vector().get_local(), m, n), res, fun
def plot_deltas_2D(vertex_vector, mesh_path, save_dir):
mesh = Mesh()
with XDMFFile(mesh_path) as f:
f.read(mesh)
V = FunctionSpace(mesh, 'CG', 1)
value = Function(V)
value.vector()[:] = vertex_vector[dof_to_vertex_map(V)]
delta_S = project(value.dx(0), V)
delta_file = XDMFFile(save_dir)
delta_file.write_checkpoint(delta_S, 'delta_S', 0, XDMFFile.Encoding.HDF5,
True)
delta_v = project(value.dx(1), V)
delta_file.write_checkpoint(delta_v, 'delta_v', 0, XDMFFile.Encoding.HDF5,
True)
def cm1d_img_pb(img: np.array,
alpha0: float,
alpha1: float,
deriv='mesh') -> np.array:
"""Computes the L2-H1 mass conserving flow for a 1D image sequence with
periodic spatial boundary.
Allows to specify how to approximate partial derivatives of f numerically.
Note that the last column of img is ignored.
Args:
img (np.array): 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
alpha0 (float): Spatial regularisation parameter.
alpha1 (float): Temporal regularisation parameter.
deriv (str): Specifies how to approximate pertial derivatives.
When set to 'mesh' it uses FEniCS built in function.
Returns:
v (np.array): A velocity array of shape (m, n).
res (float): The residual.
func (float): The value of the functional.
"""
# Check for valid arguments.
valid = {'mesh'}
if deriv not in valid:
raise ValueError("Argument 'deriv' must be one of %r." % valid)
# Create mesh.
m, n = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space.
V = dh.create_function_space(mesh, 'periodic')
# Convert array to function.
f = Function(V)
f.vector()[:] = dh.img2funvec_pb(img)
# Compute partial derivatives.
ft, fx = f.dx(0), f.dx(1)
# Compute velocity.
v, res, fun = cm1d_weak_solution(V, f, ft, fx, alpha0, alpha1)
# Convert to array and return.
return dh.funvec2img_pb(v.vector().get_local(), m, n), res, fun
def compute_velocity_correction(
ui, p0, p1, u_bcs, rho, mu, dt, rotational_form, my_dx, tol, verbose
):
"""Compute the velocity correction according to
.. math::
U = u_0 - \\frac{dt}{\\rho} \\nabla (p_1-p_0).
"""
W = ui.function_space()
P = p1.function_space()
u = TrialFunction(W)
v = TestFunction(W)
a3 = dot(u, v) * my_dx
phi = Function(P)
phi.assign(p1)
if p0:
phi -= p0
if rotational_form:
r = SpatialCoordinate(W.mesh())[0]
div_ui = 1 / r * (r * ui[0]).dx(0) + ui[1].dx(1)
phi += mu * div_ui
L3 = dot(ui, v) * my_dx - dt / rho * (phi.dx(0) * v[0] + phi.dx(1) * v[1]) * my_dx
u1 = Function(W)
solve(
a3 == L3,
u1,
bcs=u_bcs,
solver_parameters={
"linear_solver": "iterative",
"symmetric": True,
"preconditioner": "hypre_amg",
"krylov_solver": {
"relative_tolerance": tol,
"absolute_tolerance": 0.0,
"maximum_iterations": 100,
"monitor_convergence": verbose,
},
},
)
# u = project(ui - k/rho * grad(phi), V)
# div_u = 1/r * div(r*u)
r = SpatialCoordinate(W.mesh())[0]
div_u1 = 1.0 / r * (r * u1[0]).dx(0) + u1[1].dx(1)
info("||u||_div = {:e}".format(sqrt(assemble(div_u1 * div_u1 * my_dx))))
return u1
def cm1dvelocity(img: np.array, vel: np.array, alpha0: float,
alpha1: float) -> np.array:
"""Computes the source for a L2-H1 mass conserving flow for a 1D image
sequence and a given velocity.
Takes a one-dimensional image sequence and a velocity, and returns a
minimiser of the L2-H1 mass conservation functional with spatio-temporal
regularisation.
Args:
img (np.array): A 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
vel (np.array): A 1D image sequence of shape (m, n).
alpha0 (float): The spatial regularisation parameter.
alpha1 (float): The temporal regularisation parameter.
Returns:
k: A source array of shape (m, n).
"""
# Create mesh.
[m, n] = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space and functions.
V = FunctionSpace(mesh, 'CG', 1)
k = TrialFunction(V)
w = TestFunction(V)
# Convert image to function.
f = Function(V)
f.vector()[:] = dh.img2funvec(img)
# Convert velocity to function.
v = Function(V)
v.vector()[:] = dh.img2funvec(vel)
# Define derivatives of data.
ft = Function(V)
ftv = np.diff(img, axis=0) * (m - 1)
ftv = np.concatenate((ftv, ftv[-1, :].reshape(1, n)), axis=0)
ft.vector()[:] = dh.img2funvec(ftv)
fx = Function(V)
fxv = np.gradient(img, 1 / (n - 1), axis=1)
fx.vector()[:] = dh.img2funvec(fxv)
# Define weak formulation.
A = k * w * dx + alpha0 * k.dx(1) * w.dx(1) * dx + alpha1 * k.dx(0) * w.dx(
0) * dx
b = (ft + v.dx(1) * f + v * fx) * w * dx
# Compute solution.
k = Function(V)
solve(A == b, k)
# Convert back to array.
k = dh.funvec2img(k.vector().get_local(), m, n)
return k
def test_cmscr1d_weak_solution_default(self):
print("Running test 'test_cmscr1d_weak_solution_default'")
# Define temporal and spatial sample points.
m, n = 10, 20
# Define mesh and function space.
mesh = UnitSquareMesh(m - 1, n - 1)
V = dh.create_function_space(mesh, 'default')
W = dh.create_vector_function_space(mesh, 'default')
# Create zero function.
f = Function(V)
# Compute velocity.
v, k, res, fun, converged = cmscr1d_weak_solution(W, f,
f.dx(0), f.dx(1),
1.0, 1.0,
1.0, 1.0, 1.0)
v = v.vector().get_local()
k = k.vector().get_local()
np.testing.assert_allclose(v.shape, m * n)
np.testing.assert_allclose(v, np.zeros_like(v))
np.testing.assert_allclose(k.shape, m * n)
np.testing.assert_allclose(k, np.zeros_like(k))
V = dh.create_function_space(mesh, 'periodic')
# Create zero function.
f = Function(V)
# Compute velocity.
v, k, res, fun, converged = cmscr1d_weak_solution(W, f,
f.dx(0), f.dx(1),
1.0, 1.0, 1.0,
1.0, 1.0)
v = v.vector().get_local()
k = k.vector().get_local()
np.testing.assert_allclose(v.shape, m * n)
np.testing.assert_allclose(v, np.zeros_like(v))
np.testing.assert_allclose(k.shape, m * n)
np.testing.assert_allclose(k, np.zeros_like(k))
def cm1dsource(img: np.array, k: np.array, alpha0: float,
alpha1: float) -> np.array:
"""Computes the L2-H1 mass conserving flow for a 1D image sequence and a
given source.
Takes a one-dimensional image sequence and a source, and returns a
minimiser of the L2-H1 mass conservation functional with spatio-temporal
regularisation.
Args:
img (np.array): A 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
k (np.array): A 1D image sequence of shape (m, n).
alpha0 (float): The spatial regularisation parameter.
alpha1 (float): The temporal regularisation parameter.
Returns:
v: A velocity array of shape (m, n).
"""
# Create mesh.
[m, n] = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space and functions.
V = FunctionSpace(mesh, 'CG', 1)
v = TrialFunction(V)
w = TestFunction(V)
# Convert image to function.
f = Function(V)
f.vector()[:] = dh.img2funvec(img)
# Convert source to function.
g = Function(V)
g.vector()[:] = dh.img2funvec(k)
# Define derivatives of data.
ft = Function(V)
ftv = np.diff(img, axis=0) * (m - 1)
ftv = np.concatenate((ftv, ftv[-1, :].reshape(1, n)), axis=0)
ft.vector()[:] = dh.img2funvec(ftv)
fx = Function(V)
fxv = np.gradient(img, 1 / (n - 1), axis=1)
fx.vector()[:] = dh.img2funvec(fxv)
ft = f.dx(0)
fx = f.dx(1)
# Define weak formulation.
A = - (fx*v + f*v.dx(1)) * (fx*w + f*w.dx(1))*dx \
- alpha0*v.dx(1)*w.dx(1)*dx - alpha1*v.dx(0)*w.dx(0)*dx
b = ft * (fx * w + f * w.dx(1)) * dx - g * (fx * w + f * w.dx(1)) * dx
# Compute solution.
v = Function(V)
solve(A == b, v)
# Convert back to array.
vel = dh.funvec2img(v.vector().get_local(), m, n)
return vel
def of2dmcs(img1: np.array, img2: np.array, alpha0: float, alpha1: float,
beta0: float, beta1: float) -> (np.array, np.array):
"""Computes the L2-H1 optical flow for a 2D two-channel image sequence and
source for the second channel (optical flow 2d multi-channel with source).
Takes a two-dimensional image sequence and returns a minimiser of the
Horn-Schunck functional with source and with spatio-temporal
regularisation for both velocity and source.
Args:
img1 (np.array): A 2D image sequence of shape (t, m, n), where t is the
number of time steps and (n, n) is the number of
pixels.
img2 (np.array): A 2D image sequence of shape (t, m, n), where t is the
number of time steps and (n, n) is the number of
pixels.
alpha0 (float): The spatial regularisation parameter.
alpha1 (float): The temporal regularisation parameter.
Returns:
v (np.array): A velocity array of shape (t, m, n, 2).
k (np.array): A source array of shape (t, m, n).
"""
# Create mesh.
[t, m, n] = img1.shape
mesh = UnitCubeMesh(t-2, m-1, n-1)
# Define function space and functions.
V = FunctionSpace(mesh, 'CG', 1)
W = VectorFunctionSpace(mesh, 'CG', 1, dim=3)
v1, v2, k = TrialFunctions(W)
w1, w2, w3 = TestFunctions(W)
# Convert image to function.
f1, f2 = Function(V), Function(V)
f1.vector()[:] = dh.imgseq2funvec(img1[0:-1])
f2.vector()[:] = dh.imgseq2funvec(img2[0:-1])
# Define function to compute temporal derivative.
def time_deriv(img: np.array) -> Function:
# Evaluate function at vertices.
mc = mesh.coordinates().reshape((-1, 3))
hx, hy, hz = 1./(t-2), 1./(m-1), 1./(n-1)
x = np.array(mc[:, 0]/hx, dtype=int)
y = np.array(mc[:, 1]/hy, dtype=int)
z = np.array(mc[:, 2]/hz, dtype=int)
# Map pixel values to vertices.
d2v = dof_to_vertex_map(V)
# Compute derivative wrt. time.
imgt = img[1:] - img[0:-1]
ftv = imgt[x, y, z]
# Create function.
ft = Function(V)
ft.vector()[:] = ftv[d2v]
return ft
# Compute temporal derivatives.
f1t = time_deriv(img1)
f2t = time_deriv(img2)
# Define derivatives of data.
f1x, f1y = f1.dx(1), f1.dx(2)
f2x, f2y = f2.dx(1), f2.dx(2)
# Define weak formulation.
A = f1x*(f1x*v1 + f1y*v2)*w1*dx + f2x*(f2x*v1 + f2y*v2 - k)*w1*dx \
+ f1y*(f1x*v1 + f1y*v2)*w2*dx + f2y*(f2x*v1 + f2y*v2 - k)*w2*dx \
- (f2x*v1 + f2y*v2 - k)*w3*dx \
+ alpha0*v1.dx(1)*w1.dx(1)*dx + alpha0*v1.dx(2)*w1.dx(2)*dx \
+ alpha1*v1.dx(0)*w1.dx(0)*dx \
+ alpha0*v2.dx(1)*w2.dx(1)*dx + alpha0*v2.dx(2)*w2.dx(2)*dx \
+ alpha1*v2.dx(0)*w2.dx(0)*dx \
+ beta0*k.dx(1)*w3.dx(1)*dx + beta0*k.dx(2)*w3.dx(2)*dx \
+ beta1*k.dx(0)*w3.dx(0)*dx
b = - f1x*f1t*w1*dx - f2x*f2t*w1*dx \
- f1y*f1t*w2*dx - f2y*f2t*w2*dx \
+ f2t*w3*dx
# Compute solution.
v = Function(W)
solve(A == b, v, [], solver_parameters={"linear_solver": "cg"})
# Split solution into functions.
v1, v2, k = v.split(deepcopy=True)
# Convert back to arrays.
v1 = dh.funvec2imgseq(v1.vector().get_local(), t-1, m, n)
v2 = dh.funvec2imgseq(v2.vector().get_local(), t-1, m, n)
k = dh.funvec2imgseq(k.vector().get_local(), t-1, m, n)
return (np.stack((v1, v2), axis=3), k)
def step(self,
dt,
u1, p1,
u, p0,
u_bcs, p_bcs,
f0=None, f1=None,
verbose=True,
tol=1.0e-10
):
'''General pressure projection scheme as described in section 3.4 of
:cite:`GMS06`.
'''
# Some initial sanity checkups.
assert dt > 0.0
# Define coefficients
k = Constant(dt)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Compute tentative velocity step.
# rho (u[-1] + (u.\nabla)u) = mu 1/r \div(r \nabla u) + rho g.
with Message('Computing tentative velocity'):
solver = NewtonSolver()
solver.parameters['maximum_iterations'] = 5
solver.parameters['absolute_tolerance'] = tol
solver.parameters['relative_tolerance'] = 0.0
solver.parameters['report'] = True
# The nonlinear term makes the problem generally nonsymmetric.
solver.parameters['linear_solver'] = 'gmres'
# If the nonsymmetry is too strong, e.g., if u[-1] is large, then
# AMG preconditioning might not work very well.
# Use HYPRE-Euclid instead of ILU for parallel computation.
solver.parameters['preconditioner'] = 'hypre_euclid'
solver.parameters['krylov_solver']['relative_tolerance'] = tol
solver.parameters['krylov_solver']['absolute_tolerance'] = 0.0
solver.parameters['krylov_solver']['maximum_iterations'] = 1000
solver.parameters['krylov_solver']['monitor_convergence'] = verbose
ui = Function(self.W)
step_problem = \
TentativeVelocityProblem(ui,
self.theta,
self.rho, self.mu,
u, p0, k,
u_bcs,
f0, f1,
stabilization=self.stabilization,
dx=self.dx
)
# Take u[-1] as initial guess.
ui.assign(u[-1])
solver.solve(step_problem, ui.vector())
#div_u = 1/r * div(r*ui)
#plot(div_u)
#interactive()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
with Message('Computing pressure correction'):
# The pressure correction is based on the update formula
#
# ( dphi/dr )
# rho/dt (u_{n+1}-u*) + ( dphi/dz ) = 0
# (1/r dphi/dtheta)
#
# with
#
# phi = p_{n+1} - p*
#
# and
#
# 1/r d/dr (r ur_{n+1})
# + d/dz ( uz_{n+1})
# + 1/r d/dtheta( utheta_{n+1}) = 0
#
# With the assumption that u does not change in the direction
# theta, one derives
#
# - 1/r * div(r * \nabla phi) = 1/r * rho/dt div(r*(u_{n+1} - u*))
# - 1/r * n. (r * \nabla phi) = 1/r * rho/dt n.(r*(u_{n+1} - u*))
#
# In its weak form, this is
#
# \int r * \grad(phi).\grad(q) *2*pi =
# - rho/dt \int div(r*u*) q *2*p
# - rho/dt \int_Gamma n.(r*(u_{n+1}-u*)) q *2*pi.
#
# (The terms 1/r cancel with the volume elements 2*pi*r.)
# If Dirichlet boundary conditions are applied to both u* and u_n
# (the latter in the final step), the boundary integral vanishes.
#
alpha = 1.0
#alpha = 3.0 / 2.0
rotational_form = False
self._pressure_poisson(p1, p0,
self.mu, ui,
self.rho * alpha * ui / k,
p_bcs=p_bcs,
rotational_form=rotational_form,
tol=tol,
verbose=verbose
)
#plot(ui, title='u intermediate')
##plot(f, title='f')
##plot(ui[1], title='u intermediate[1]')
#plot(div(ui), title='div(u intermediate)')
#plot(p1, title='p1')
#interactive()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Velocity correction.
# U = u[-1] - dt/rho \nabla (p1-p0).
with Message('Computing velocity correction'):
u = TrialFunction(self.W)
v = TestFunction(self.W)
a3 = dot(u, v) * dx
phi = Function(self.P)
phi.assign(p1)
if p0:
phi -= p0
if rotational_form:
r = Expression('x[0]', degree=1, domain=self.W.mesh())
div_ui = 1/r * (r * ui[0]).dx(0) + ui[1].dx(1)
phi += self.mu * div_ui
L3 = dot(ui, v) * dx \
- k / self.rho * (phi.dx(0) * v[0] + phi.dx(1) * v[1]) * dx
# Jacobi preconditioning for mass matrix is order-optimal.
solve(
a3 == L3, u1,
bcs=u_bcs,
solver_parameters={
'linear_solver': 'iterative',
'symmetric': True,
'preconditioner': 'jacobi',
'krylov_solver': {'relative_tolerance': tol,
'absolute_tolerance': 0.0,
'maximum_iterations': 100,
'monitor_convergence': verbose}
}
)
#u = project(ui - k/rho * grad(phi), V)
# div_u = 1/r * div(r*u)
r = Expression('x[0]', degree=1, domain=self.W.mesh())
div_u1 = 1.0 / r * (r * u1[0]).dx(0) + u1[1].dx(1)
info('||u||_div = %e' % sqrt(assemble(div_u1 * div_u1 * dx)))
#plot(div_u)
#interactive()
## Ha! TODO remove
#u1.vector()[:] = ui.vector()
return
def heston_transform(parameters,
get_error=False,
file_path=None,
plot_deltas=False):
experiment = parameters.experiment
domain_name = parameters.domain
mesh_name = parameters.mesh
mesh = Mesh()
mesh_path = parameters.get_mesh_path()
with XDMFFile(mesh_path) as f:
f.read(mesh)
# Create functions in function spaces defined by both original and
# transformed meshes and copy value function data from one to another
t_mesh_path = f'meshes/{domain_name}/t_{domain_name}_{mesh_name}.xdmf'
if not os.path.exists(t_mesh_path):
mesh2 = get_original_mesh(t_mesh_path, mesh, parameters)
else:
mesh2 = Mesh()
with XDMFFile(t_mesh_path) as f:
f.read(mesh2)
V = FunctionSpace(mesh, 'CG', 1)
t_V = FunctionSpace(mesh2, 'CG', 1)
value = Function(V)
control = Function(V)
t_value = Function(t_V)
t_control = Function(t_V)
if not file_path:
f = XDMFFile(f'out/{experiment}/{domain_name}/{mesh_name}/v.xdmf')
t_f = XDMFFile(f'out/{experiment}/{domain_name}/{mesh_name}/t_v.xdmf')
else:
f = XDMFFile(file_path)
path, file_name = os.path.split(file_path)
t_f = XDMFFile(os.path.join(path, f't_{file_name}'))
# Save Final Time Condition
f.read_checkpoint(value, 'value_func', 0)
t_value.vector()[:] = value.vector()[:]
t_f.write_checkpoint(t_value, 'value_func', parameters.T)
# Iterate through remaining timesteps (only some of them are
# saved to the file)
i = 1
parameters.calculate_save_interval()
dt = parameters.T / parameters.M
for k in range(parameters.M - 1, -1, -1):
if k % parameters.save_interval != 0:
continue
f.read_checkpoint(value, 'value_func', i)
f.read_checkpoint(control, 'control', i - 1)
t_value.vector()[:] = value.vector()[:]
t_control.vector()[:] = control.vector()[:]
t_f.write_checkpoint(t_value, 'value_func', k * dt,
XDMFFile.Encoding.HDF5, True)
t_f.write_checkpoint(t_control, 'control', k * dt,
XDMFFile.Encoding.HDF5, True)
if plot_deltas:
delta_S = project(t_value.dx(0), t_V)
t_f.write_checkpoint(delta_S, 'delta_S', k * dt,
XDMFFile.Encoding.HDF5, True)
delta_v = project(t_value.dx(1), t_V)
t_f.write_checkpoint(delta_v, 'delta_v', k * dt,
XDMFFile.Encoding.HDF5, True)
i += 1
if k == 0 and get_error:
error_calc(t_value, parameters.t_v_e, mesh2, f't_{mesh_name}',
f'out/{experiment}/{domain_name}/errors.json')
def cmscr1d_img(img: np.array,
alpha0: float, alpha1: float,
alpha2: float, alpha3: float,
beta: float, deriv, mesh=None, bc='natural') \
-> (np.array, np.array, float, float, bool):
"""Computes the L2-H1 mass conserving flow with source for a 1D image
sequence with spatio-temporal and convective regularisation.
Allows to specify how to approximate partial derivatives of f numerically.
Args:
img (np.array): 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
alpha0 (float): The spatial regularisation parameter for v.
alpha1 (float): The temporal regularisation parameter for v.
alpha2 (float): The spatial regularisation parameter for k.
alpha3 (float): The temporal regularisation parameter for k.
beta (float): The convective regularisation parameter.
deriv (str): Specifies how to approximate pertial derivatives.
When set to 'mesh' it uses FEniCS built in function.
When set to 'fd' it uses finite differences.
mesh: A custom mesh (optional). Must have (m - 1, n - 1) cells.
bc (str): One of {'natural', 'zero', 'zerospace'} for boundary
conditions for the velocity v (optional).
Returns:
v (np.array): A velocity array of shape (m, n).
k (np.array): A source array of shape (m, n).
res (float): The residual.
fun (float): The function value.
converged (bool): True if Newton's method converged.
"""
# Check for valid arguments.
valid = {'mesh', 'fd'}
if deriv not in valid:
raise ValueError("Argument 'deriv' must be one of %r." % valid)
# Create mesh.
m, n = img.shape
if mesh is None:
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function spaces.
V = dh.create_function_space(mesh, 'default')
W = dh.create_vector_function_space(mesh, 'default')
# Convert array to function.
f = Function(V)
f.vector()[:] = dh.img2funvec(img)
# Compute partial derivatives.
if deriv is 'mesh':
ft, fx = f.dx(0), f.dx(1)
if deriv is 'fd':
imgt, imgx = nh.partial_derivatives(img)
ft, fx = Function(V), Function(V)
ft.vector()[:] = dh.img2funvec(imgt)
fx.vector()[:] = dh.img2funvec(imgx)
# Check for valid arguments.
valid = {'natural', 'zero', 'zerospace'}
if bc not in valid:
raise ValueError("Argument 'bc' must be one of %r." % valid)
# Define boundary conditions for velocity.
if bc is 'natural':
bc = []
if bc is 'zero':
bc = DirichletBC(W.sub(0), Constant(0), dh.DirichletBoundary())
if bc is 'zerospace':
bc = DirichletBC(W.sub(0), Constant(0), dh.DirichletBoundarySpace())
# Compute velocity.
v, k, res, fun, converged = cmscr1d_weak_solution(W,
f,
ft,
fx,
alpha0,
alpha1,
alpha2,
alpha3,
beta,
bcs=bc)
# Convert back to array and return.
v = dh.funvec2img(v.vector().get_local(), m, n)
k = dh.funvec2img(k.vector().get_local(), m, n)
return v, k, res, fun, converged
class TVPD():
""" Total variation using primal-dual Newton """
def __init__(self, parameters):
"""
TV regularization in primal-dual format
Input parameters:
* k = regularization parameter
* eps = regularization constant (see above)
* rescaledradiusdual = radius of dual set
* exact = use full TV (bool)
* PCGN = use GN Hessian to precondition (bool);
not recommended for performance but can help avoid num.instability
* print (bool)
"""
self.parameters = {}
self.parameters['k'] = 1.0
self.parameters['eps'] = 1e-2
self.parameters['rescaledradiusdual'] = 1.0
self.parameters['exact'] = False
self.parameters['PCGN'] = False
self.parameters['print'] = False
self.parameters['correctcost'] = True
self.parameters['amg'] = 'default'
assert parameters.has_key('Vm')
self.parameters.update(parameters)
self.Vm = self.parameters['Vm']
k = self.parameters['k']
eps = self.parameters['eps']
exact = self.parameters['exact']
amg = self.parameters['amg']
self.m = Function(self.Vm)
testm = TestFunction(self.Vm)
trialm = TrialFunction(self.Vm)
# WARNING: should not be changed.
# As it is, code only works with DG0
if self.parameters.has_key('Vw'):
Vw = self.parameters['Vw']
else:
Vw = FunctionSpace(self.Vm.mesh(), 'DG', 0)
self.wx = Function(Vw)
self.wxrs = Function(Vw) # re-scaled dual variable
self.wxhat = Function(Vw)
self.gwx = Function(Vw)
self.wy = Function(Vw)
self.wyrs = Function(Vw) # re-scaled dual variable
self.wyhat = Function(Vw)
self.gwy = Function(Vw)
self.wxsq = Vector()
self.wysq = Vector()
self.normw = Vector()
self.factorw = Vector()
testw = TestFunction(Vw)
trialw = TrialFunction(Vw)
normm = inner(nabla_grad(self.m), nabla_grad(self.m))
TVnormsq = normm + Constant(eps)
TVnorm = sqrt(TVnormsq)
if self.parameters['correctcost']:
meshtmp = UnitSquareMesh(self.Vm.mesh().mpi_comm(), 10, 10)
Vtmp = FunctionSpace(meshtmp, 'CG', 1)
x = SpatialCoordinate(meshtmp)
correctioncost = 1. / assemble(sqrt(4.0 * x[0] * x[0]) * dx)
else:
correctioncost = 1.0
self.wkformcost = Constant(k * correctioncost) * TVnorm * dx
if exact:
sys.exit(1)
# self.w = nabla_grad(self.m)/TVnorm # full Hessian
# self.Htvw = inner(Constant(k) * nabla_grad(testm), self.w) * dx
self.misfitwx = inner(testw, self.wx * TVnorm - self.m.dx(0)) * dx
self.misfitwy = inner(testw, self.wy * TVnorm - self.m.dx(1)) * dx
self.Htvx = assemble(inner(Constant(k) * testm.dx(0), trialw) * dx)
self.Htvy = assemble(inner(Constant(k) * testm.dx(1), trialw) * dx)
self.massw = inner(TVnorm * testw, trialw) * dx
mpicomm = self.Vm.mesh().mpi_comm()
invMwMat, VDM, VDM = setupPETScmatrix(Vw, Vw, 'aij', mpicomm)
for ii in VDM.dofs():
invMwMat[ii, ii] = 1.0
invMwMat.assemblyBegin()
invMwMat.assemblyEnd()
self.invMwMat = PETScMatrix(invMwMat)
self.invMwd = Vector()
self.invMwMat.init_vector(self.invMwd, 0)
self.invMwMat.init_vector(self.wxsq, 0)
self.invMwMat.init_vector(self.wysq, 0)
self.invMwMat.init_vector(self.normw, 0)
self.invMwMat.init_vector(self.factorw, 0)
u = Function(Vw)
uflrank = len(u.ufl_shape)
if uflrank == 0:
ones = ("1.0")
elif uflrank == 1:
ones = (("1.0", "1.0"))
else:
sys.exit(1)
u = interpolate(Constant(ones), Vw)
self.one = u.vector()
self.wkformAx = inner(testw, trialm.dx(0) - \
self.wx * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAxrs = inner(testw, trialm.dx(0) - \
self.wxrs * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAy = inner(testw, trialm.dx(1) - \
self.wy * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
self.wkformAyrs = inner(testw, trialm.dx(1) - \
self.wyrs * inner(nabla_grad(self.m), nabla_grad(trialm)) / TVnorm) * dx
kovsq = Constant(k) / TVnorm
self.wkformGNhess = kovsq * inner(nabla_grad(trialm),
nabla_grad(testm)) * dx
factM = 1e-2 * k
M = assemble(inner(testm, trialm) * dx)
self.sMass = M * factM
self.Msolver = PETScKrylovSolver('cg', 'jacobi')
self.Msolver.parameters["maximum_iterations"] = 2000
self.Msolver.parameters["relative_tolerance"] = 1e-24
self.Msolver.parameters["absolute_tolerance"] = 1e-24
self.Msolver.parameters["error_on_nonconvergence"] = True
self.Msolver.parameters["nonzero_initial_guess"] = False
self.Msolver.set_operator(M)
if amg == 'default': self.amgprecond = amg_solver()
else: self.amgprecond = amg
if self.parameters['print']:
print '[TVPD] TV regularization -- primal-dual method',
if self.parameters['PCGN']:
print ' -- PCGN',
print ' -- k={}, eps={}'.format(k, eps)
print '[TVPD] preconditioner = {}'.format(self.amgprecond)
print '[TVPD] correction cost with factor={}'.format(
correctioncost)
def isTV(self):
return True
def isPD(self):
return True
def cost(self, m):
""" evaluate the cost functional at m """
setfct(self.m, m)
return assemble(self.wkformcost)
def costvect(self, m_in):
return self.cost(m_in)
def _assemble_invMw(self):
""" Assemble inverse of matrix Mw,
weighted mass matrix in dual space """
# WARNING: only works if Mw is diagonal (e.g, DG0)
Mw = assemble(self.massw)
Mwd = get_diagonal(Mw)
as_backend_type(self.invMwd).vec().pointwiseDivide(\
as_backend_type(self.one).vec(),\
as_backend_type(Mwd).vec())
self.invMwMat.set_diagonal(self.invMwd)
def grad(self, m):
""" compute the gradient at m """
setfct(self.m, m)
self._assemble_invMw()
self.gwx.vector().zero()
self.gwx.vector().axpy(1.0, assemble(self.misfitwx))
normgwx = norm(self.gwx.vector())
self.gwy.vector().zero()
self.gwy.vector().axpy(1.0, assemble(self.misfitwy))
normgwy = norm(self.gwy.vector())
if self.parameters['print']:
print '[TVPD] |gw|={}'.format(np.sqrt(normgwx**2 + normgwy**2))
return self.Htvx*(self.wx.vector() - self.invMwd*self.gwx.vector()) \
+ self.Htvy*(self.wy.vector() - self.invMwd*self.gwy.vector())
#return assemble(self.Htvw) - self.Htv*(self.invMwd*self.gw.vector())
def gradvect(self, m_in):
return self.grad(m_in)
def assemble_hessian(self, m):
""" build Hessian matrix at given point m """
setfct(self.m, m)
self._assemble_invMw()
self.Ax = assemble(self.wkformAx)
Hxasym = MatMatMult(self.Htvx, MatMatMult(self.invMwMat, self.Ax))
Hx = (Hxasym + Transpose(Hxasym)) * 0.5
Axrs = assemble(self.wkformAxrs)
Hxrsasym = MatMatMult(self.Htvx, MatMatMult(self.invMwMat, Axrs))
Hxrs = (Hxrsasym + Transpose(Hxrsasym)) * 0.5
self.Ay = assemble(self.wkformAy)
Hyasym = MatMatMult(self.Htvy, MatMatMult(self.invMwMat, self.Ay))
Hy = (Hyasym + Transpose(Hyasym)) * 0.5
Ayrs = assemble(self.wkformAyrs)
Hyrsasym = MatMatMult(self.Htvy, MatMatMult(self.invMwMat, Ayrs))
Hyrs = (Hyrsasym + Transpose(Hyrsasym)) * 0.5
self.H = Hx + Hy
self.Hrs = Hxrs + Hyrs
PCGN = self.parameters['PCGN']
if PCGN:
HGN = assemble(self.wkformGNhess)
self.precond = HGN + self.sMass
else:
self.precond = self.Hrs + self.sMass
def hessian(self, mhat):
return self.Hrs * mhat
def compute_what(self, mhat):
""" Compute update direction for what, given mhat """
self.wxhat.vector().zero()
self.wxhat.vector().axpy(
1.0, self.invMwd * (self.Ax * mhat - self.gwx.vector()))
normwxhat = norm(self.wxhat.vector())
self.wyhat.vector().zero()
self.wyhat.vector().axpy(
1.0, self.invMwd * (self.Ay * mhat - self.gwy.vector()))
normwyhat = norm(self.wyhat.vector())
if self.parameters['print']:
print '[TVPD] |what|={}'.format(
np.sqrt(normwxhat**2 + normwyhat**2))
def update_w(self, mhat, alphaLS, compute_what=True):
""" update dual variable in direction what
and update re-scaled version """
if compute_what: self.compute_what(mhat)
self.wx.vector().axpy(alphaLS, self.wxhat.vector())
self.wy.vector().axpy(alphaLS, self.wyhat.vector())
# rescaledradiusdual=1.0: checked empirically to be max radius acceptable
rescaledradiusdual = self.parameters['rescaledradiusdual']
# wx**2
as_backend_type(self.wxsq).vec().pointwiseMult(\
as_backend_type(self.wx.vector()).vec(),\
as_backend_type(self.wx.vector()).vec())
# wy**2
as_backend_type(self.wysq).vec().pointwiseMult(\
as_backend_type(self.wy.vector()).vec(),\
as_backend_type(self.wy.vector()).vec())
# |w|
self.normw = self.wxsq + self.wysq
as_backend_type(self.normw).vec().sqrtabs()
# |w|/r
as_backend_type(self.normw).vec().pointwiseDivide(\
as_backend_type(self.normw).vec(),\
as_backend_type(self.one*rescaledradiusdual).vec())
# max(1.0, |w|/r)
# as_backend_type(self.factorw).vec().pointwiseMax(\
# as_backend_type(self.one).vec(),\
# as_backend_type(self.normw).vec())
count = pointwiseMaxCount(self.factorw, self.normw, 1.0)
# rescale wx and wy
as_backend_type(self.wxrs.vector()).vec().pointwiseDivide(\
as_backend_type(self.wx.vector()).vec(),\
as_backend_type(self.factorw).vec())
as_backend_type(self.wyrs.vector()).vec().pointwiseDivide(\
as_backend_type(self.wy.vector()).vec(),\
as_backend_type(self.factorw).vec())
minf = self.factorw.min()
maxf = self.factorw.max()
if self.parameters['print']:
# print 'min(factorw)={}, max(factorw)={}'.format(minf, maxf)
print '[TVPD] perc. dual entries rescaled={:.2f} %, min(factorw)={}, max(factorw)={}'.format(\
100.*float(count)/self.factorw.size(), minf, maxf)
def getprecond(self):
""" Precondition by inverting the TV Hessian """
solver = PETScKrylovSolver('cg', self.amgprecond)
solver.parameters["maximum_iterations"] = 2000
solver.parameters["relative_tolerance"] = 1e-24
solver.parameters["absolute_tolerance"] = 1e-24
solver.parameters["error_on_nonconvergence"] = True
solver.parameters["nonzero_initial_guess"] = False
# used to compare iterative application of preconditioner
# with exact application of preconditioner:
#solver = PETScLUSolver("petsc")
#solver.parameters['symmetric'] = True
#solver.parameters['reuse_factorization'] = True
solver.set_operator(self.precond)
return solver
def init_vector(self, u, dim):
self.sMass.init_vector(u, dim)
ds = Measure("ds", subdomain_data=boundary_parts)
# Define forms (dont redefine functions used here)
# step 1
u0 = Function(V)
u1 = Function(V)
p0 = Function(Q)
u_tent = TrialFunction(V)
v = TestFunction(V)
# U_ = 1.5*u0 - 0.5*u1
# nonlinearity = inner(dot(0.5 * (u_tent.dx(0) + u0.dx(0)), U_), v) * dx
# F_tent = (1./dt)*inner(u_tent - u0, v) * dx + nonlinearity\
# + nu*inner(0.5 * (u_tent.dx(0) + u0.dx(0)), v.dx(0)) * dx + inner(p0.dx(0), v) * dx\
# - inner(f, v)*dx # solve to u_
# using explicite scheme: so LHS has interpretation as heat equation, RHS are sources
F_tent = (1./dt)*inner(u_tent - u0, v)*dx + inner(dot(u0.dx(0), u0), v)*dx + nu*inner((u_tent.dx(0) + u0.dx(0)), v.dx(0)) * dx + inner(p0.dx(0), v) * dx\
- inner(f, v)*dx # solve to u_
a_tent, L_tent = system(F_tent)
# step 2
u_tent_computed = Function(V)
p = TrialFunction(Q)
q = TestFunction(Q)
F_p = inner(grad(p-p0), grad(q))*dx + (1./dt)*u_tent_computed.dx(0)*q*dx # + 2*p.dx(0)*q*ds(1) # tried to force dp/dn=0 on inflow
# TEST: prescribe Neumann outflow BC
# F_p = inner(grad(p-p0), grad(q))*dx + (1./dt)*u_tent_computed.dx(0)*q*dx + (1./dt)*(v_in_expr-u_tent_computed)*q*ds(2)
a_p, L_p = system(F_p)
A_p = assemble(a_p)
as_backend_type(A_p).set_nullspace(null_space)
print(A_p.array())
# step 2 rotation
# F_p_rot = inner(grad(p), grad(q))*dx + (1./dt)*u_tent_computed.dx(0)*q*dx + (1./dt)*(v_in_expr-u_tent_computed)*q*ds(2)
ds = Measure("ds", subdomain_data=boundary_parts)
# Define forms (dont redefine functions used here)
# step 1
u0 = Function(V)
u1 = Function(V)
p0 = Function(Q)
u_tent = TrialFunction(V)
v = TestFunction(V)
# U_ = 1.5*u0 - 0.5*u1
# nonlinearity = inner(dot(0.5 * (u_tent.dx(0) + u0.dx(0)), U_), v) * dx
# F_tent = (1./dt)*inner(u_tent - u0, v) * dx + nonlinearity\
# + nu*inner(0.5 * (u_tent.dx(0) + u0.dx(0)), v.dx(0)) * dx + inner(p0.dx(0), v) * dx\
# - inner(f, v)*dx # solve to u_
# using explicite scheme: so LHS has interpretation as heat equation, RHS are sources
F_tent = (1./dt)*inner(u_tent - u0, v)*dx + inner(dot(u0.dx(0), u0), v)*dx + nu*inner((u_tent.dx(0) + u0.dx(0)), v.dx(0)) * dx + inner(p0.dx(0), v) * dx\
- inner(f, v)*dx # solve to u_
a_tent, L_tent = system(F_tent)
# step 2
u_tent_computed = Function(V)
p = TrialFunction(Q)
q = TestFunction(Q)
F_p = inner(grad(p - p0), grad(q)) * dx + (1. / dt) * u_tent_computed.dx(
0) * q * dx # + 2*p.dx(0)*q*ds(1) # tried to force dp/dn=0 on inflow
# TEST: prescribe Neumann outflow BC
# F_p = inner(grad(p-p0), grad(q))*dx + (1./dt)*u_tent_computed.dx(0)*q*dx + (1./dt)*(v_in_expr-u_tent_computed)*q*ds(2)
a_p, L_p = system(F_p)
A_p = assemble(a_p)
as_backend_type(A_p).set_nullspace(null_space)
print(A_p.array())
# step 2 rotation
