From d59124e7e2db8dc25cae022845ef6be4db61802b Mon Sep 17 00:00:00 2001
From: Stephan Hilb <stephan@ecshi.net>
Date: Mon, 5 Jul 2021 17:37:39 +0200
Subject: [PATCH] finish milestone: denoising seems to work

---
 Manifest.toml   |  4 +-
 Project.toml    |  1 +
 src/function.jl |  9 +++--
 src/mesh.jl     |  5 +++
 src/operator.jl | 32 +++++++++-------
 src/run.jl      | 97 ++++++++++++++++++++++++++++++++++++++++---------
 6 files changed, 110 insertions(+), 38 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index 0a76bde..09db26a 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -218,9 +218,9 @@ version = "1.5.1"
 
 [[StaticArrays]]
 deps = ["LinearAlgebra", "Random", "Statistics"]
-git-tree-sha1 = "745914ebcd610da69f3cb6bf76cb7bb83dcb8c9a"
+git-tree-sha1 = "896d55218776ab8f23fb7b222a5a4a946d4aafc2"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.2.4"
+version = "1.2.5"
 
 [[Statistics]]
 deps = ["LinearAlgebra", "SparseArrays"]
diff --git a/Project.toml b/Project.toml
index d5f735f..343cbf9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,5 +7,6 @@ version = "0.1.0"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
diff --git a/src/function.jl b/src/function.jl
index 8a9295c..0e547e0 100644
--- a/src/function.jl
+++ b/src/function.jl
@@ -1,4 +1,5 @@
 using Statistics: mean
+using StaticArrays: SA
 export FeSpace, Mapper, FeFunction, P1, DP0, DP1
 export interpolate!, sample, bind!, evaluate, nabla
 
@@ -9,8 +10,8 @@ struct DP1 end
 
 # FIXME: should be static vectors
 # evaluate all (1-dim) local basis functions against x
-evaluate_basis(::P1, x) = [1 - x[1] - x[2], x[1], x[2]]
-evaluate_basis(::DP0, x) = [1]
+evaluate_basis(::P1, x) = SA[1 - x[1] - x[2], x[1], x[2]]
+evaluate_basis(::DP0, x) = SA[1]
 evaluate_basis(::DP1, x) = evaluate_basis(P1(), x)
 
 # non-mixed
@@ -125,7 +126,7 @@ function interpolate!(dst::FeFunction, ::DP0, expr::Function; params...)
 	    bind!(f, cell)
 	end
 	vertices = mesh.vertices[:, mesh.cells[:, cell]]
-	centroid = mean(vertices, dims = 2)
+	centroid = reshape(mean(vertices, dims = 2), 2)
 	lcentroid = [1/3, 1/3]
 
 	opvalues = map(f -> evaluate(f, lcentroid), params)
@@ -169,7 +170,7 @@ function sample(f::FeFunction)
 	I0 = floor.(Int, vec(minimum(vertices, dims = 2)))
 	I1 = ceil.(Int, vec(maximum(vertices, dims = 2)))
 
-	J = jacobian(x -> vertices * [1 - x[1] - x[2], x[1], x[2]], [0., 0.])
+	J = jacobian(x -> vertices * SA[1 - x[1] - x[2], x[1], x[2]], SA[0., 0.])
 	xloc = J \ (v - vertices[:, 1])
 
 	for I in CartesianIndex(I0[1], I0[2]):CartesianIndex(I1[1], I1[2])
diff --git a/src/mesh.jl b/src/mesh.jl
index f09872c..fb54af3 100644
--- a/src/mesh.jl
+++ b/src/mesh.jl
@@ -43,11 +43,16 @@ init_grid(img::Array{<:Any, 2}, type=:vertex) =
 #    return all(λ .>= 0) && sum(λ) <= 1
 #end
 
+save(filename, f::FeFunction, fs::FeFunction...) =
+    save(filename, f.space.mesh, f, fs...)
+
 function save(filename, mesh::Mesh, fs...)
     cells = [MeshCell(VTKCellTypes.VTK_TRIANGLE, mesh.cells[:, i]) for i in axes(mesh.cells, 2)]
     #vertices = [mesh.vertices[i, j] for i in axes(mesh.vertices, 1), j in axes(mesh.vertices, 2)]
     vtkfile = vtk_grid(filename, mesh.vertices, cells)
     for f in fs
+	f.space.mesh == mesh ||
+	    throw(ArgumentError("meshes do not match"))
         append_data!(vtkfile, f)
     end
     vtk_save(vtkfile)
diff --git a/src/operator.jl b/src/operator.jl
index 16305b1..100f00c 100644
--- a/src/operator.jl
+++ b/src/operator.jl
@@ -1,5 +1,6 @@
 using SparseArrays: sparse
 using LinearAlgebra: det, dot
+using StaticArrays: SA
 using ForwardDiff: jacobian
 
 export Poisson, L2Projection, init_point!, assemble, assemble_rhs
@@ -32,17 +33,18 @@ a(op::Poisson, xloc, u, du, v, dv) = dot(du, dv)
 
 
 # degree 2 quadrature on triangle
-quadrature() = [1/6, 1/6, 1/6], [1/6 4/6 1/6; 1/6 1/6 4/6]
+quadrature() = SA[1/6, 1/6, 1/6], SA[1/6 4/6 1/6; 1/6 1/6 4/6]
 
 
 elmap(mesh, cell, x) =
-    mesh.vertices[:, mesh.cells[:, cell]] * [1 - x[1] - x[2], x[1], x[2]]
+    mesh.vertices[:, mesh.cells[:, cell]] * SA[1 - x[1] - x[2], x[1], x[2]]
 
+assemble(op::Operator) = assemble(
+    op.space, (x...; y...) -> a(op, x...; y...); params(op)...)
 
-function assemble(op::Operator)
-    space = op.space
+function assemble(space::FeSpace, a; params...)
     mesh = space.mesh
-    opparams = params(op)
+    opparams = NamedTuple(params)
 
     # precompute basis at quadrature points
 
@@ -58,7 +60,7 @@ function assemble(op::Operator)
     for r in 1:nrdims
 	for k in axes(qx, 2)
 	    qphi[r, k, r, :] .= evaluate_basis(space.element, qx[:, k])
-	    dqphi[r, :, k, r, :] .= transpose(jacobian(x -> evaluate_basis(space.element, x), qx[:, k])::Array{Float64,2})
+	    dqphi[r, :, k, r, :] .= transpose(jacobian(x -> evaluate_basis(space.element, x), qx[:, k])::AbstractArray{Float64,2})
 	end
     end
 
@@ -75,7 +77,7 @@ function assemble(op::Operator)
 	for f in opparams
 	    bind!(f, cell)
 	end
-	delmap = jacobian(x -> elmap(mesh, cell, x), [0., 0.])::Array{Float64,2} # constant on element
+	delmap = jacobian(x -> elmap(mesh, cell, x), SA[0., 0.])::Array{Float64,2} # constant on element
 	delmapinv = inv(delmap) # constant on element
 	intel = abs(det(delmap))
 
@@ -95,7 +97,7 @@ function assemble(op::Operator)
 		    phij .= reshape(view(qphi, :, k, jdim, ldofj), space.size)
 		    dphij .= reshape(view(dqphi, :, :, k, jdim, ldofj) * delmapinv, (space.size..., :))
 
-		    gdofv = qw[k] * a(op, xhat, phii, dphii, phij, dphij; opvalues...) * intel
+		    gdofv = qw[k] * a(xhat, phii, dphii, phij, dphij; opvalues...) * intel
 
 		    push!(I, gdofi)
 		    push!(J, gdofj)
@@ -111,10 +113,12 @@ function assemble(op::Operator)
     return A
 end
 
-function assemble_rhs(op::Operator)
-    space = op.space
+assemble_rhs(op::Operator) = assemble_rhs(
+    op.space, (x...; y...) -> l(op, x...; y...); params(op)...)
+
+function assemble_rhs(space::FeSpace, l; params...)
     mesh = space.mesh
-    opparams = params(op)
+    opparams = NamedTuple(params)
 
     # precompute basis at quadrature points
 
@@ -130,7 +134,7 @@ function assemble_rhs(op::Operator)
     for r in 1:nrdims
 	for k in axes(qx, 2)
 	    qphi[r, k, r, :] .= evaluate_basis(space.element, qx[:, k])
-	    dqphi[r, :, k, r, :] .= transpose(jacobian(x -> evaluate_basis(space.element, x), qx[:, k])::Array{Float64,2})
+	    dqphi[r, :, k, r, :] .= transpose(jacobian(x -> evaluate_basis(space.element, x), qx[:, k])::AbstractArray{Float64,2})
 	end
     end
 
@@ -143,7 +147,7 @@ function assemble_rhs(op::Operator)
 	for f in opparams
 	    bind!(f, cell)
 	end
-	delmap = jacobian(x -> elmap(mesh, cell, x), [0., 0.])::Array{Float64,2} # constant on element
+	delmap = jacobian(x -> elmap(mesh, cell, x), SA[0., 0.])::Array{Float64,2} # constant on element
 	delmapinv = inv(delmap) # constant on element
 	intel = abs(det(delmap))
 
@@ -159,7 +163,7 @@ function assemble_rhs(op::Operator)
 		phij .= reshape(qphi[:, k, jdim, ldofj], space.size)
 		dphij .= reshape(dqphi[:, :, k, jdim, ldofj] * delmapinv, (space.size..., :))
 
-		gdofv = qw[k] * l(op, xhat, phij, dphij; opvalues...) * intel
+		gdofv = qw[k] * l(xhat, phij, dphij; opvalues...) * intel
 
 		b[gdofj] += gdofv
 	    end
diff --git a/src/run.jl b/src/run.jl
index 54f6e0b..37a20ab 100644
--- a/src/run.jl
+++ b/src/run.jl
@@ -14,53 +14,114 @@ function myrun()
     Vp1 = FeSpace(mesh, DP0(), (1,))
     Vp2 = FeSpace(mesh, DP1(), (m, d))
 
-    g = FeFunction(Vu)
-    u = FeFunction(Vu)
-    p1 = FeFunction(Vp1)
-    p2 = FeFunction(Vp2)
+    g = FeFunction(Vu, name="g")
+    u = FeFunction(Vu, name="u")
+    p1 = FeFunction(Vp1, name="p1")
+    p2 = FeFunction(Vp2, name="p2")
     du = FeFunction(Vu)
     dp1 = FeFunction(Vp1)
     dp2 = FeFunction(Vp2)
     nablau = nabla(u)
     nabladu = nabla(du)
 
-    g.data .= 1
+    g.data .= 0
     u.data .= 0
-    p1.data .= -1.
+    p1.data .= 0
     p2.data .= 0
     du.data .= 0
     dp1.data .= 0
     dp2.data .= 0
 
-    alpha1 = 1.
-    alpha2 = 1.
-    beta = 1e-3
-    lambda = 1.
-
+    alpha1 = 0.
+    alpha2 = 10.
+    beta = 1e-2
+    lambda = 0.01
     gamma1 = 1e-3
     gamma2 = 1e-3
 
+    interpolate!(g, x -> norm(x - SA[0.5, 0.5]) < 0.3)
+
+    # TODO: fixme
+    #T(u) = [u[1] + u[2]]
     T(u) = u
+    S(u) = u
 
-    function dp1_update(xloc; g, u, p1, du)
+    function dp1_update(x; g, u, p1, du)
 	m1 = max(gamma1, norm(T(u) - g))
-	cond = norm(T(u) - g) >= gamma1 ?
+	cond = norm(T(u) - g) > gamma1 ?
 	    dot(T(u) - g, T(du)) / norm(T(u) - g)^2 * p1 :
 	    zeros(size(p1))
 	return -p1 + alpha1 / m1 * (T(u + du) - g) - cond
     end
 
-    interpolate!(p1, dp1_update; g, u, p1, du)
 
-    function dp2_update(xloc; u, nablau, p2, du, nabladu)
+    function dp2_update(x; u, nablau, p2, du, nabladu)
 	m2 = max(gamma2, norm(nablau))
-	cond = norm(nablau) >= gamma2 ?
+	cond = norm(nablau) > gamma2 ?
 	    dot(nablau, nabladu) / norm(nablau)^2 * p2 :
 	    zeros(size(p2))
 	return -p2 + lambda / m2 * (nablau + nabladu) - cond
     end
 
-    interpolate!(p2, dp2_update; u, nablau, p2, du, nabladu)
 
-    return p1, p2
+    @inline function du_a(x, du, nabladu, phi, nablaphi; g, u, nablau, p1, p2)
+	m1 = max(gamma1, norm(T(u) - g))
+	cond1 = norm(T(u) - g) > gamma1 ?
+	    dot(T(u) - g, T(du)) / norm(T(u) - g)^2 * p1 :
+	    zeros(size(p1))
+	a1 = alpha1 / m1 * dot(T(du), T(phi)) -
+	    dot(cond1, T(phi))
+
+	m2 = max(gamma2, norm(nablau))
+	cond2 = norm(nablau) > gamma2 ?
+	    dot(nablau, nabladu) / norm(nablau)^2 * p2 :
+	    zeros(size(p2))
+	a2 = lambda / m2 * dot(nabladu, nablaphi) -
+	    dot(cond2, nablaphi)
+
+	aB = alpha2 * dot(T(du), T(phi)) +
+	    beta * dot(S(du), S(phi))
+
+	return a1 + a2 + aB
+    end
+
+    @inline function du_l(x, phi, nablaphi; g, u, nablau)
+	aB = alpha2 * dot(T(u), T(phi)) +
+	    beta * dot(S(u), S(phi))
+	m1 = max(gamma1, norm(T(u) - g))
+	p1part = alpha1 / m1 * dot(T(u) - g, T(phi))
+	m2 = max(gamma2, norm(nablau))
+	p2part = lambda / m2 * dot(nablau, nablaphi)
+	gpart = alpha2 * dot(g, T(phi))
+
+	return -aB - p1part - p2part + gpart
+    end
+
+    for i = 1:6
+	A = assemble(Vu, du_a; g, u, nablau, p1, p2)
+	b = assemble_rhs(Vu, du_l; g, u, nablau)
+
+	du.data .= A \ b
+
+	interpolate!(dp1, dp1_update; g, u, p1, du)
+	interpolate!(dp2, dp2_update; u, nablau, p2, du, nabladu)
+
+	u.data .+= du.data
+	p1.data .+= dp1.data
+	p2.data .+= dp2.data
+
+	# reproject
+	p1.space.element::DP0
+	p1.data .= clamp.(p1.data, -alpha1, alpha1)
+	p2.space.element::DP1
+	p2d = reshape(p2.data, m * d * 3, :) # no copy
+	for i in axes(p2d, 2)
+	    p2in = norm(p2d[:, i])
+	    if p2in > lambda
+		p2d[:, i] .*= lambda ./ p2in
+	    end
+	end
+    end
+
+    save("test.vtu", g, u, p1)
 end
-- 
GitLab