Skip to content

Commit 24f7893

Browse files
MikaelSlevinskyMikaelSlevinsky
authored
Merge pull request #38 from MikaelSlevinsky/feat-spherical-vector-fields
Support 2D vector fields on the sphere
2 parents 3de2cec + 9df45a9 commit 24f7893

File tree

4 files changed

+385
-0
lines changed

4 files changed

+385
-0
lines changed

‎src/SphericalHarmonics/SphericalHarmonics.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ include("Butterfly.jl")
1414
include("fastplan.jl")
1515
include("thinplan.jl")
1616
include("synthesisanalysis.jl")
17+
include("vectorfield.jl")
1718

1819
function plan_sph2fourier(A::AbstractMatrix; opts...)
1920
M, N = size(A)

‎src/SphericalHarmonics/vectorfield.jl

Lines changed: 308 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,308 @@
1+
function A_mul_B_vf!(P::RotationPlan, A::AbstractMatrix)
2+
N, M = size(A)
3+
snm = P.snm
4+
cnm = P.cnm
5+
@stepthreads for m = M÷2-1:-1:2
6+
@inbounds for j = m:-2:2
7+
for l = N-j:-1:1
8+
s = snm[l+(j-2)*(2*N+3-j)÷2]
9+
c = cnm[l+(j-2)*(2*N+3-j)÷2]
10+
a1 = A[l+N*(2*m+1)]
11+
a2 = A[l+2+N*(2*m+1)]
12+
a3 = A[l+N*(2*m+2)]
13+
a4 = A[l+2+N*(2*m+2)]
14+
A[l+N*(2*m+1)] = c*a1 + s*a2
15+
A[l+2+N*(2*m+1)] = c*a2 - s*a1
16+
A[l+N*(2*m+2)] = c*a3 + s*a4
17+
A[l+2+N*(2*m+2)] = c*a4 - s*a3
18+
end
19+
end
20+
end
21+
A
22+
end
23+
24+
function At_mul_B_vf!(P::RotationPlan, A::AbstractMatrix)
25+
N, M = size(A)
26+
snm = P.snm
27+
cnm = P.cnm
28+
@stepthreads for m = M÷2-1:-1:2
29+
@inbounds for j = reverse(m:-2:2)
30+
for l = 1:N-j
31+
s = snm[l+(j-2)*(2*N+3-j)÷2]
32+
c = cnm[l+(j-2)*(2*N+3-j)÷2]
33+
a1 = A[l+N*(2*m+1)]
34+
a2 = A[l+2+N*(2*m+1)]
35+
a3 = A[l+N*(2*m+2)]
36+
a4 = A[l+2+N*(2*m+2)]
37+
A[l+N*(2*m+1)] = c*a1 - s*a2
38+
A[l+2+N*(2*m+1)] = c*a2 + s*a1
39+
A[l+N*(2*m+2)] = c*a3 - s*a4
40+
A[l+2+N*(2*m+2)] = c*a4 + s*a3
41+
end
42+
end
43+
end
44+
A
45+
end
46+
47+
48+
function Base.A_mul_B!(Y1::Matrix, Y2::Matrix, SP::SlowSphericalHarmonicPlan, X1::Matrix, X2::Matrix)
49+
RP, p1, p2, B = SP.RP, SP.p1, SP.p2, SP.B
50+
copy!(B, X1)
51+
A_mul_B_vf!(RP, B)
52+
M, N = size(X1)
53+
A_mul_B_col_J!!(Y1, p2, B, 1)
54+
for J = 2:4:N
55+
A_mul_B_col_J!!(Y1, p1, B, J)
56+
J < N && A_mul_B_col_J!!(Y1, p1, B, J+1)
57+
end
58+
for J = 4:4:N
59+
A_mul_B_col_J!!(Y1, p2, B, J)
60+
J < N && A_mul_B_col_J!!(Y1, p2, B, J+1)
61+
end
62+
copy!(B, X2)
63+
A_mul_B_vf!(RP, B)
64+
M, N = size(X2)
65+
A_mul_B_col_J!!(Y2, p2, B, 1)
66+
for J = 2:4:N
67+
A_mul_B_col_J!!(Y2, p1, B, J)
68+
J < N && A_mul_B_col_J!!(Y2, p1, B, J+1)
69+
end
70+
for J = 4:4:N
71+
A_mul_B_col_J!!(Y2, p2, B, J)
72+
J < N && A_mul_B_col_J!!(Y2, p2, B, J+1)
73+
end
74+
Y1
75+
end
76+
77+
function Base.At_mul_B!(Y1::Matrix, Y2::Matrix, SP::SlowSphericalHarmonicPlan, X1::Matrix, X2::Matrix)
78+
RP, p1inv, p2inv, B = SP.RP, SP.p1inv, SP.p2inv, SP.B
79+
copy!(B, X1)
80+
M, N = size(X1)
81+
A_mul_B_col_J!!(Y1, p2inv, B, 1)
82+
for J = 2:4:N
83+
A_mul_B_col_J!!(Y1, p1inv, B, J)
84+
J < N && A_mul_B_col_J!!(Y1, p1inv, B, J+1)
85+
end
86+
for J = 4:4:N
87+
A_mul_B_col_J!!(Y1, p2inv, B, J)
88+
J < N && A_mul_B_col_J!!(Y1, p2inv, B, J+1)
89+
end
90+
sph_zero_spurious_modes_vf!(At_mul_B_vf!(RP, Y1))
91+
copy!(B, X2)
92+
M, N = size(X2)
93+
A_mul_B_col_J!!(Y2, p2inv, B, 1)
94+
for J = 2:4:N
95+
A_mul_B_col_J!!(Y2, p1inv, B, J)
96+
J < N && A_mul_B_col_J!!(Y2, p1inv, B, J+1)
97+
end
98+
for J = 4:4:N
99+
A_mul_B_col_J!!(Y2, p2inv, B, J)
100+
J < N && A_mul_B_col_J!!(Y2, p2inv, B, J+1)
101+
end
102+
sph_zero_spurious_modes_vf!(At_mul_B_vf!(RP, Y2))
103+
Y1
104+
end
105+
106+
Base.Ac_mul_B!(Y1::Matrix, Y2::Matrix, SP::SlowSphericalHarmonicPlan, X1::Matrix, X2::Matrix) = At_mul_B!(Y1, Y2, SP, X1, X2)
107+
108+
109+
function Base.A_mul_B!(Y1::Matrix{T}, Y2::Matrix{T}, P::SynthesisPlan{T}, X1::Matrix{T}, X2::Matrix{T}) where T
110+
M, N = size(X1)
111+
112+
# Column synthesis
113+
PCe = P.planθ[1]
114+
PCo = P.planθ[2]
115+
116+
A_mul_B_col_J!(Y1, PCo, X1, 1)
117+
118+
for J = 2:4:N
119+
X1[1,J] *= two(T)
120+
J < N && (X1[1,J+1] *= two(T))
121+
A_mul_B_col_J!(Y1, PCe, X1, J)
122+
J < N && A_mul_B_col_J!(Y1, PCe, X1, J+1)
123+
X1[1,J] *= half(T)
124+
J < N && (X1[1,J+1] *= half(T))
125+
end
126+
for J = 4:4:N
127+
A_mul_B_col_J!(Y1, PCo, X1, J)
128+
J < N && A_mul_B_col_J!(Y1, PCo, X1, J+1)
129+
end
130+
scale!(half(T), Y1)
131+
132+
# Row synthesis
133+
scale!(inv(sqrt(π)), Y1)
134+
invsqrttwo = inv(sqrt(2))
135+
@inbounds for i = 1:M Y1[i] *= invsqrttwo end
136+
137+
temp = P.temp
138+
planφ = P.planφ
139+
C = P.C
140+
for I = 1:M
141+
copy_row_I!(temp, Y1, I)
142+
row_synthesis!(planφ, C, temp)
143+
copy_row_I!(Y1, temp, I)
144+
end
145+
146+
M, N = size(X2)
147+
148+
# Column synthesis
149+
PCe = P.planθ[1]
150+
PCo = P.planθ[2]
151+
152+
A_mul_B_col_J!(Y2, PCo, X2, 1)
153+
154+
for J = 2:4:N
155+
X2[1,J] *= two(T)
156+
J < N && (X2[1,J+1] *= two(T))
157+
A_mul_B_col_J!(Y2, PCe, X2, J)
158+
J < N && A_mul_B_col_J!(Y2, PCe, X2, J+1)
159+
X2[1,J] *= half(T)
160+
J < N && (X2[1,J+1] *= half(T))
161+
end
162+
for J = 4:4:N
163+
A_mul_B_col_J!(Y2, PCo, X2, J)
164+
J < N && A_mul_B_col_J!(Y2, PCo, X2, J+1)
165+
end
166+
scale!(half(T), Y2)
167+
168+
# Row synthesis
169+
scale!(inv(sqrt(π)), Y2)
170+
invsqrttwo = inv(sqrt(2))
171+
@inbounds for i = 1:M Y2[i] *= invsqrttwo end
172+
173+
temp = P.temp
174+
planφ = P.planφ
175+
C = P.C
176+
for I = 1:M
177+
copy_row_I!(temp, Y2, I)
178+
row_synthesis!(planφ, C, temp)
179+
copy_row_I!(Y2, temp, I)
180+
end
181+
Y1
182+
end
183+
184+
function Base.A_mul_B!(Y1::Matrix{T}, Y2::Matrix{T}, P::AnalysisPlan{T}, X1::Matrix{T}, X2::Matrix{T}) where T
185+
M, N = size(X1)
186+
187+
# Row analysis
188+
temp = P.temp
189+
planφ = P.planφ
190+
C = P.C
191+
for I = 1:M
192+
copy_row_I!(temp, X1, I)
193+
row_analysis!(planφ, C, temp)
194+
copy_row_I!(Y1, temp, I)
195+
end
196+
197+
# Column analysis
198+
PCe = P.planθ[1]
199+
PCo = P.planθ[2]
200+
201+
A_mul_B_col_J!(Y1, PCo, Y1, 1)
202+
for J = 2:4:N
203+
A_mul_B_col_J!(Y1, PCe, Y1, J)
204+
J < N && A_mul_B_col_J!(Y1, PCe, Y1, J+1)
205+
Y1[1,J] *= half(T)
206+
J < N && (Y1[1,J+1] *= half(T))
207+
end
208+
for J = 4:4:N
209+
A_mul_B_col_J!(Y1, PCo, Y1, J)
210+
J < N && A_mul_B_col_J!(Y1, PCo, Y1, J+1)
211+
end
212+
scale!(sqrt(π)*inv(T(M)), Y1)
213+
sqrttwo = sqrt(2)
214+
@inbounds for i = 1:M Y1[i] *= sqrttwo end
215+
216+
M, N = size(X2)
217+
218+
# Row analysis
219+
temp = P.temp
220+
planφ = P.planφ
221+
C = P.C
222+
for I = 1:M
223+
copy_row_I!(temp, X2, I)
224+
row_analysis!(planφ, C, temp)
225+
copy_row_I!(Y2, temp, I)
226+
end
227+
228+
# Column analysis
229+
PCe = P.planθ[1]
230+
PCo = P.planθ[2]
231+
232+
A_mul_B_col_J!(Y2, PCo, Y2, 1)
233+
for J = 2:4:N
234+
A_mul_B_col_J!(Y2, PCe, Y2, J)
235+
J < N && A_mul_B_col_J!(Y2, PCe, Y2, J+1)
236+
Y2[1,J] *= half(T)
237+
J < N && (Y2[1,J+1] *= half(T))
238+
end
239+
for J = 4:4:N
240+
A_mul_B_col_J!(Y2, PCo, Y2, J)
241+
J < N && A_mul_B_col_J!(Y2, PCo, Y2, J+1)
242+
end
243+
scale!(sqrt(π)*inv(T(M)), Y2)
244+
sqrttwo = sqrt(2)
245+
@inbounds for i = 1:M Y2[i] *= sqrttwo end
246+
247+
Y1
248+
end
249+
250+
251+
function sph_zero_spurious_modes_vf!(A::AbstractMatrix)
252+
M, N = size(A)
253+
n = N÷2
254+
A[M, 1] = 0
255+
@inbounds for j = 2:n-1
256+
@simd for i = M-j+2:M
257+
A[i,2j] = 0
258+
A[i,2j+1] = 0
259+
end
260+
end
261+
@inbounds @simd for i = M-n+2:M
262+
A[i,2n] = 0
263+
2n < N && (A[i,2n+1] = 0)
264+
end
265+
A
266+
end
267+
268+
function sphrandvf(::Type{T}, m::Int, n::Int) where T
269+
A = zeros(T, m, 2n-1)
270+
for i = 1:m-1
271+
A[i,1] = rand(T)
272+
end
273+
for j = 1:n-1
274+
for i = 1:m-j+1
275+
A[i,2j] = rand(T)
276+
A[i,2j+1] = rand(T)
277+
end
278+
end
279+
A
280+
end
281+
282+
function sphrandnvf(::Type{T}, m::Int, n::Int) where T
283+
A = zeros(T, m, 2n-1)
284+
for i = 1:m-1
285+
A[i,1] = randn(T)
286+
end
287+
for j = 1:n-1
288+
for i = 1:m-j+1
289+
A[i,2j] = randn(T)
290+
A[i,2j+1] = randn(T)
291+
end
292+
end
293+
A
294+
end
295+
296+
function sphonesvf(::Type{T}, m::Int, n::Int) where T
297+
A = zeros(T, m, 2n-1)
298+
for i = 1:m-1
299+
A[i,1] = one(T)
300+
end
301+
for j = 1:n-1
302+
for i = 1:m-j+1
303+
A[i,2j] = one(T)
304+
A[i,2j+1] = one(T)
305+
end
306+
end
307+
A
308+
end

‎test/sphericalharmonics/sphericalharmonictests.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,6 @@ include("pointwisetests.jl")
1717

1818
include("synthesisanalysistests.jl")
1919

20+
include("vectorfieldtests.jl")
21+
2022
include("apitests.jl")

‎test/sphericalharmonics/vectorfieldtests.jl

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
using FastTransforms, Compat
2+
using Compat.Test
3+
4+
@testset "Test vector field transforms" begin
5+
# f = (θ,φ) -> cospi(θ) + sinpi(θ)*(1+cospi(2θ))*sinpi(φ) + sinpi(θ)^5*(cospi(5φ)-sinpi(5φ))
6+
∇θf = (θ,φ) -> π*(-sinpi(θ) + (cospi(θ)*(1+cospi(2θ)) - 2*sinpi(θ)*sinpi(2θ))*sinpi(φ) + 5*sinpi(θ)^4*cospi(θ)*(cospi(5φ)-sinpi(5φ)))
7+
∇φf = (θ,φ) -> π*((1+cospi(2θ))*cospi(φ) - 5*sinpi(θ)^4*(sinpi(5φ)+cospi(5φ)))
8+
9+
n = 6
10+
θ = (0.5:n-0.5)/n
11+
φ = (0:2n-2)*2/(2n-1)
12+
∇θF = [∇θf(θ,φ) for θ in θ, φ in φ]
13+
∇φF = [∇φf(θ,φ) for θ in θ, φ in φ]
14+
V1 = zero(∇θF)
15+
V2 = zero(∇φF)
16+
Pa = FastTransforms.plan_analysis(∇θF)
17+
A_mul_B!(V1, V2, Pa, ∇θF, ∇φF)
18+
P = SlowSphericalHarmonicPlan(V1)
19+
20+
U1 = zero(V1)
21+
U2 = zero(V2)
22+
At_mul_B!(U1, U2, P, V1, V2)
23+
24+
W1 = zero(U1)
25+
W2 = zero(U2)
26+
27+
A_mul_B!(W1, W2, P, U1, U2)
28+
29+
Ps = FastTransforms.plan_synthesis(W1)
30+
31+
G1 = zero(∇θF)
32+
G2 = zero(∇φF)
33+
34+
A_mul_B!(G1, G2, Ps, W1, W2)
35+
36+
@test vecnorm(∇θF - G1)/vecnorm(∇θF) < n*eps()
37+
@test vecnorm(∇φF - G2)/vecnorm(∇φF) < n*eps()
38+
39+
y = (1.0, 2.0, 3.0)
40+
for k in (10, 20, 40)
41+
∇θf = (θ,φ) -> -2π*k*sin(k*((sinpi(θ)*cospi(φ) - y[1])^2 + (sinpi(θ)*sinpi(φ) - y[2])^2 + (cospi(θ) - y[3])^2))*( (sinpi(θ)*cospi(φ) - y[1])*(cospi(θ)*cospi(φ)) + (sinpi(θ)*sinpi(φ) - y[2])*(cospi(θ)*sinpi(φ)) - (cospi(θ) - y[3])*sinpi(θ) )
42+
∇φf = (θ,φ) -> -2π*k*sin(k*((sinpi(θ)*cospi(φ) - y[1])^2 + (sinpi(θ)*sinpi(φ) - y[2])^2 + (cospi(θ) - y[3])^2))*( (sinpi(θ)*cospi(φ) - y[1])*(-sinpi(φ)) + (sinpi(θ)*sinpi(φ) - y[2])*(cospi(φ)) )
43+
n = 12k
44+
45+
θ = (0.5:n-0.5)/n
46+
φ = (0:2n-2)*2/(2n-1)
47+
∇θF = [∇θf(θ,φ) for θ in θ, φ in φ]
48+
∇φF = [∇φf(θ,φ) for θ in θ, φ in φ]
49+
V1 = zero(∇θF)
50+
V2 = zero(∇φF)
51+
Pa = FastTransforms.plan_analysis(∇θF)
52+
A_mul_B!(V1, V2, Pa, ∇θF, ∇φF)
53+
P = SlowSphericalHarmonicPlan(V1)
54+
55+
U1 = zero(V1)
56+
U2 = zero(V2)
57+
At_mul_B!(U1, U2, P, V1, V2)
58+
59+
W1 = zero(U1)
60+
W2 = zero(U2)
61+
62+
A_mul_B!(W1, W2, P, U1, U2)
63+
64+
Ps = FastTransforms.plan_synthesis(W1)
65+
66+
G1 = zero(∇θF)
67+
G2 = zero(∇φF)
68+
69+
A_mul_B!(G1, G2, Ps, W1, W2)
70+
71+
@test vecnorm(∇θF - G1)/vecnorm(∇θF) < n*eps()
72+
@test vecnorm(∇φF - G2)/vecnorm(∇φF) < n*eps()
73+
end
74+
end

0 commit comments

Comments
 (0)