Skip to content

Commit c3d534b

Browse files
dlfivefiftydlfivefifty
committed
2 parents f76c4ff + 625f342 commit c3d534b

File tree

3 files changed

+424
-0
lines changed

3 files changed

+424
-0
lines changed

‎src/gradient.jl

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# This file calculates the surface gradient of a scalar field.
2+
3+
function gradient!(U::Matrix{T}, ∇θU::Matrix{T}, ∇φU::Matrix{T}) where T
4+
@assert size(U) == size(∇θU) == size(∇φU)
5+
N, M = size(U)
6+
7+
# The first column is easy.
8+
@inbounds @simd for ℓ = 1:N-1
9+
∇θU[ℓ, 1] = -sqrt(T(ℓ*(ℓ+1)))*U[ℓ+1, 1]
10+
∇φU[ℓ, 1] = 0
11+
end
12+
∇θU[N, 1] = 0
13+
∇φU[N, 1] = 0
14+
15+
# Next, we differentiate with respect to φ, which preserves the order. It swaps sines and cosines in longitude, though.
16+
@inbounds for m = 1:M÷2
17+
@simd for ℓ = 1:N+1-m
18+
∇φU[ℓ, 2m] = -m*U[ℓ, 2m+1]
19+
∇φU[ℓ, 2m+1] = m*U[ℓ, 2m]
20+
end
21+
end
22+
23+
# Then, we differentiate with respect to θ, which preserves the order but divides by sin(θ).
24+
25+
@inbounds for m = 1:M÷2
26+
ℓ = 1
27+
bℓ = -(ℓ+m+1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
28+
∇θU[ℓ, 2m] = bℓ*U[ℓ+1, 2m]
29+
∇θU[ℓ, 2m+1] = bℓ*U[ℓ+1, 2m+1]
30+
@simd for ℓ = 2:N-m
31+
aℓ = (ℓ+m-2)*sqrt(T((ℓ-1)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
32+
bℓ = -(ℓ+m+1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
33+
∇θU[ℓ, 2m] = aℓ*U[ℓ-1, 2m] + bℓ*U[ℓ+1, 2m]
34+
∇θU[ℓ, 2m+1] = aℓ*U[ℓ-1, 2m+1] + bℓ*U[ℓ+1, 2m+1]
35+
end
36+
ℓ = N-m+1
37+
aℓ = (ℓ+m-2)*sqrt(T((ℓ-1)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
38+
∇θU[ℓ, 2m] = aℓ*U[ℓ-1, 2m]
39+
∇θU[ℓ, 2m+1] = aℓ*U[ℓ-1, 2m+1]
40+
end
41+
42+
# Finally, we divide by sin(θ), which can be done by decrementing the order P_ℓ^m ↘ P_ℓ^{m-1}.
43+
@inbounds for m = 1:M÷2
44+
ℓ = N+1-m
45+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
46+
∇θU[ℓ, 2m] = ∇θU[ℓ, 2m]/aℓ
47+
∇θU[ℓ, 2m+1] = ∇θU[ℓ, 2m+1]/aℓ
48+
∇φU[ℓ, 2m] = ∇φU[ℓ, 2m]/aℓ
49+
∇φU[ℓ, 2m+1] = ∇φU[ℓ, 2m+1]/aℓ
50+
ℓ = N+1-m-1
51+
if ℓ > 0
52+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
53+
∇θU[ℓ, 2m] = ∇θU[ℓ, 2m]/aℓ
54+
∇θU[ℓ, 2m+1] = ∇θU[ℓ, 2m+1]/aℓ
55+
∇φU[ℓ, 2m] = ∇φU[ℓ, 2m]/aℓ
56+
∇φU[ℓ, 2m+1] = ∇φU[ℓ, 2m+1]/aℓ
57+
end
58+
@simd for ℓ = N+1-m-2:-1:1
59+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
60+
bℓ = -sqrt(T(ℓ*(ℓ+1))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
61+
∇θU[ℓ, 2m] = (∇θU[ℓ, 2m] - bℓ*∇θU[ℓ+2, 2m])/aℓ
62+
∇θU[ℓ, 2m+1] = (∇θU[ℓ, 2m+1] - bℓ*∇θU[ℓ+2, 2m+1])/aℓ
63+
∇φU[ℓ, 2m] = (∇φU[ℓ, 2m] - bℓ*∇φU[ℓ+2, 2m])/aℓ
64+
∇φU[ℓ, 2m+1] = (∇φU[ℓ, 2m+1] - bℓ*∇φU[ℓ+2, 2m+1])/aℓ
65+
end
66+
end
67+
68+
∇θU
69+
end
70+
71+
function curl!(U::Matrix, U1::Matrix, U2::Matrix)
72+
gradient!(U, U2, U1)
73+
N, M = size(U)
74+
@inbounds for j = 1:M
75+
for i = 1:N
76+
U1[i,j] = -U1[i,j]
77+
end
78+
end
79+
U1
80+
end
81+
82+
83+
function partial_gradient!(U::Matrix{T}, ∇θU::Matrix{T}, ∇φU::Matrix{T}) where T
84+
@assert size(U) == size(∇θU) == size(∇φU)
85+
N, M = size(U)
86+
87+
# The first column is easy.
88+
@inbounds @simd for ℓ = 1:N-1
89+
∇θU[ℓ, 1] = -sqrt(T(ℓ*(ℓ+1)))*U[ℓ+1, 1]
90+
∇φU[ℓ, 1] = 0
91+
end
92+
∇θU[N, 1] = 0
93+
∇φU[N, 1] = 0
94+
95+
# Next, we differentiate with respect to φ, which preserves the order. It swaps sines and cosines in longitude, though.
96+
@inbounds for m = 1:M÷2
97+
@simd for ℓ = 1:N+1-m
98+
∇φU[ℓ, 2m] = -m*U[ℓ, 2m+1]
99+
∇φU[ℓ, 2m+1] = m*U[ℓ, 2m]
100+
end
101+
end
102+
103+
# Then, we differentiate with respect to θ, which preserves the order but divides by sin(θ).
104+
105+
@inbounds for m = 1:M÷2
106+
ℓ = 1
107+
bℓ = -(ℓ+m+1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
108+
∇θU[ℓ, 2m] = bℓ*U[ℓ+1, 2m]
109+
∇θU[ℓ, 2m+1] = bℓ*U[ℓ+1, 2m+1]
110+
@simd for ℓ = 2:N-m
111+
aℓ = (ℓ+m-2)*sqrt(T((ℓ-1)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
112+
bℓ = -(ℓ+m+1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
113+
∇θU[ℓ, 2m] = aℓ*U[ℓ-1, 2m] + bℓ*U[ℓ+1, 2m]
114+
∇θU[ℓ, 2m+1] = aℓ*U[ℓ-1, 2m+1] + bℓ*U[ℓ+1, 2m+1]
115+
end
116+
ℓ = N-m+1
117+
aℓ = (ℓ+m-2)*sqrt(T((ℓ-1)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
118+
∇θU[ℓ, 2m] = aℓ*U[ℓ-1, 2m]
119+
∇θU[ℓ, 2m+1] = aℓ*U[ℓ-1, 2m+1]
120+
end
121+
122+
∇θU
123+
end

‎src/helmholtzhodge.jl

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
include("gradient.jl")
2+
3+
using BandedMatrices
4+
5+
import BandedMatrices: BandedQ
6+
7+
# Store QR factorizations required to apply the Helmholtz-Hodge decomposition.
8+
struct HelmholtzHodge{T}
9+
Q::Vector{BandedQ{T}}
10+
R::Vector{BandedMatrix{T, Matrix{T}}}
11+
X::Vector{T}
12+
end
13+
14+
function HelmholtzHodge(::Type{T}, N::Int) where T
15+
Q = Vector{BandedQ{T}}(undef, N)
16+
R = Vector{BandedMatrix{T, Matrix{T}}}(undef, N)
17+
for m = 1:N
18+
Q[m], R[m] = qr(helmholtzhodgeconversion(T, N, m))
19+
end
20+
HelmholtzHodge(Q, R, zeros(T, 2N+2))
21+
end
22+
23+
function helmholtzhodgeconversion(::Type{T}, N::Int, m::Int) where T
24+
A = BandedMatrix(Zeros{T}(2N+4-2m, 2N+2-2m), (2, 2))
25+
for ℓ = 1:N+1-m
26+
A[2ℓ, 2ℓ-1] = m
27+
A[2ℓ-1, 2ℓ] = m
28+
cst = (m+ℓ-1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
29+
A[2ℓ+2, 2ℓ] = cst
30+
A[2ℓ+1, 2ℓ-1] = cst
31+
end
32+
for ℓ = 1:N-m
33+
cst = -(m+ℓ+1)*sqrt(T(ℓ*(ℓ+2m))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
34+
A[2ℓ-1, 2ℓ+1] = cst
35+
A[2ℓ, 2ℓ+2] = cst
36+
end
37+
A
38+
end
39+
40+
# This function works in-place on the data stored in the factorization.
41+
function solvehelmholtzhodge!(HH::HelmholtzHodge{T}, m::Int) where T
42+
Q = HH.Q[m]
43+
R = HH.R[m]
44+
X = HH.X
45+
46+
# Step 1: apply Q'
47+
48+
H=Q.H
49+
m=Q.m
50+
51+
M=size(H,1)
52+
x=pointer(X)
53+
h=pointer(H)
54+
st=stride(H,2)
55+
sz=sizeof(T)
56+
57+
for k=1:min(size(H,2),m-M+1)
58+
wp=h+sz*st*(k-1)
59+
xp=x+sz*(k-1)
60+
61+
dt=BandedMatrices.dot(M,wp,1,xp,1)
62+
BandedMatrices.axpy!(M,-2*dt,wp,1,xp,1)
63+
end
64+
65+
for k=m-M+2:size(H,2)
66+
p=k-m+M-1
67+
68+
wp=h+sz*st*(k-1)
69+
xp=x+sz*(k-1)
70+
71+
dt=BandedMatrices.dot(M-p,wp,1,xp,1)
72+
BandedMatrices.axpy!(M-p,-2*dt,wp,1,xp,1)
73+
end
74+
75+
# Step 2: backsolve with (square) R
76+
77+
hhdtbsv!('U', 'N', 'N', size(R.data, 2), R.u, pointer(R.data), size(R.data, 1), pointer(X), 1)
78+
79+
X
80+
end
81+
82+
83+
function helmholtzhodge!(HH::HelmholtzHodge{T}, U1, U2, V1, V2) where T
84+
N, M = size(V1)
85+
86+
# U1 is for e_theta and U2 is for e_phi.
87+
# The first columns are easy.
88+
U1[1, 1] = 0
89+
U2[1, 1] = 0
90+
@inbounds @simd for ℓ = 1:N-1
91+
U1[ℓ+1, 1] = -V1[ℓ, 1]/sqrt(T(ℓ*(ℓ+1)))
92+
U2[ℓ+1, 1] = -V2[ℓ, 1]/sqrt(T(ℓ*(ℓ+1)))
93+
end
94+
95+
# First, we multiply by sin(θ), which can be done by incrementing the order P_ℓ^{m-1} ↗ P_ℓ^m.
96+
@inbounds for m = 1:M÷2
97+
@simd for ℓ = 1:N-1-m
98+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
99+
bℓ = -sqrt(T(ℓ*(ℓ+1))/T((2ℓ+2m-1)*(2ℓ+2m+1)))
100+
V1[ℓ, 2m] = aℓ*V1[ℓ, 2m] + bℓ*V1[ℓ+2, 2m]
101+
V1[ℓ, 2m+1] = aℓ*V1[ℓ, 2m+1] + bℓ*V1[ℓ+2, 2m+1]
102+
V2[ℓ, 2m] = aℓ*V2[ℓ, 2m] + bℓ*V2[ℓ+2, 2m]
103+
V2[ℓ, 2m+1] = aℓ*V2[ℓ, 2m+1] + bℓ*V2[ℓ+2, 2m+1]
104+
end
105+
ℓ = N-m
106+
if ℓ > 0
107+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
108+
V1[ℓ, 2m] = aℓ*V1[ℓ, 2m]
109+
V1[ℓ, 2m+1] = aℓ*V1[ℓ, 2m+1]
110+
V2[ℓ, 2m] = aℓ*V2[ℓ, 2m]
111+
V2[ℓ, 2m+1] = aℓ*V2[ℓ, 2m+1]
112+
end
113+
ℓ = N+1-m
114+
aℓ = sqrt(T((ℓ+2m-2)*(ℓ+2m-1))/T((2ℓ+2m-3)*(2ℓ+2m-1)))
115+
V1[ℓ, 2m] = aℓ*V1[ℓ, 2m]
116+
V1[ℓ, 2m+1] = aℓ*V1[ℓ, 2m+1]
117+
V2[ℓ, 2m] = aℓ*V2[ℓ, 2m]
118+
V2[ℓ, 2m+1] = aℓ*V2[ℓ, 2m+1]
119+
end
120+
121+
# Next, we solve the banded linear systems.
122+
for m = 1:M÷2
123+
readin1!(HH.X, V1, V2, N, m)
124+
solvehelmholtzhodge!(HH, m)
125+
writeout1!(HH.X, U1, U2, N, m)
126+
readin2!(HH.X, V1, V2, N, m)
127+
solvehelmholtzhodge!(HH, m)
128+
writeout2!(HH.X, U1, U2, N, m)
129+
end
130+
131+
U1
132+
end
133+
134+
function readin1!(X, V1, V2, N, m)
135+
X[1] = V1[1, 2m]
136+
X[2N+2-2m] = V2[N+1-m, 2m+1]
137+
@inbounds for ℓ = 1:N-m
138+
X[2ℓ] = V2[ℓ, 2m+1]
139+
X[2ℓ+1] = V1[ℓ+1, 2m]
140+
end
141+
X[2N+3-2m] = X[2N+4-2m] = 0
142+
end
143+
144+
function readin2!(X, V1, V2, N, m)
145+
X[1] = V1[1, 2m+1]
146+
X[2N+2-2m] = -V2[N+1-m, 2m]
147+
@inbounds for ℓ = 1:N-m
148+
X[2ℓ] = -V2[ℓ, 2m]
149+
X[2ℓ+1] = V1[ℓ+1, 2m+1]
150+
end
151+
X[2N+3-2m] = X[2N+4-2m] = 0
152+
end
153+
154+
function writeout1!(X, U1, U2, N, m)
155+
U1[1, 2m] = X[1]
156+
U2[1, 2m+1] = X[N+1-m]
157+
@inbounds for ℓ = 1:N-1-m
158+
U1[ℓ+1, 2m] = X[2ℓ+1]
159+
U2[ℓ, 2m+1] = X[2ℓ]
160+
end
161+
U2[N-m, 2m+1] = X[2N-2m]
162+
end
163+
164+
function writeout2!(X, U1, U2, N, m)
165+
U1[1, 2m+1] = X[1]
166+
U2[1, 2m] = X[N+1-m]
167+
@inbounds for ℓ = 1:N-1-m
168+
U1[ℓ+1, 2m+1] = X[2ℓ+1]
169+
U2[ℓ, 2m] = -X[2ℓ]
170+
end
171+
U2[N-m, 2m] = -X[2N-2m]
172+
end
173+
174+
using LinearAlgebra
175+
import LinearAlgebra: BlasInt
176+
import LinearAlgebra.BLAS: libblas, @blasfunc
177+
178+
for (fname, elty) in ((:dtbsv_,:Float64),
179+
(:stbsv_,:Float32),
180+
(:ztbsv_,:ComplexF64),
181+
(:ctbsv_,:ComplexF32))
182+
@eval begin
183+
function hhdtbsv!(uplo::Char, trans::Char, diag::Char,
184+
n::Int, k::Int, A::Ptr{$elty}, lda::Int,
185+
x::Ptr{$elty}, incx::Int)
186+
ccall((@blasfunc($fname), libblas), Nothing,
187+
(Ref{UInt8}, Ref{UInt8}, Ref{UInt8},
188+
Ref{BlasInt}, Ref{BlasInt},
189+
Ptr{$elty}, Ref{BlasInt},
190+
Ptr{$elty}, Ref{BlasInt}),
191+
uplo, trans, diag,
192+
n, k,
193+
A, lda,
194+
x, incx)
195+
x
196+
end
197+
end
198+
end

0 commit comments

Comments
 (0)