complete legendre polynomials

Author: yuehhua

src/polynomials.jl

@@ -22,23 +22,24 @@ function legendre_ϕ_ψ(k)
         ϕ_2x_coefs[ki+1, 1:(ki+1)] .= sqrt(2*(2*ki+1)) .* coeffs(l(p2))
     end
 
-    ψ1_coefs .= ϕ_2x_coefs
+    ψ1_coefs = zeros(k, k)
     ψ2_coefs = zeros(k, k)
     for ki in 0:(k-1)
+        ψ1_coefs[ki+1, :] .= ϕ_2x_coefs[ki+1, :]
         for i in 0:(k-1)
             a = ϕ_2x_coefs[ki+1, 1:(ki+1)]
             b = ϕ_coefs[i+1, 1:(i+1)]
             proj_ = proj_factor(a, b)
-            view(ψ1_coefs, ki+1, :) .-= proj_ .* view(ϕ_coefs, i+1, :)
-            view(ψ2_coefs, ki+1, :) .-= proj_ .* view(ϕ_coefs, i+1, :)
+            ψ1_coefs[ki+1, :] .-= proj_ .* view(ϕ_coefs, i+1, :)
+            ψ2_coefs[ki+1, :] .-= proj_ .* view(ϕ_coefs, i+1, :)
         end
 
         for j in 0:(k-1)
             a = ϕ_2x_coefs[ki+1, 1:(ki+1)]
             b = ψ1_coefs[j+1, :]
             proj_ = proj_factor(a, b)
-            view(ψ1_coefs, ki+1, :) .-= proj_ .* view(ψ1_coefs, j+1, :)
-            view(ψ2_coefs, ki+1, :) .-= proj_ .* view(ψ2_coefs, j+1, :)
+            ψ1_coefs[ki+1, :] .-= proj_ .* view(ψ1_coefs, j+1, :)
+            ψ2_coefs[ki+1, :] .-= proj_ .* view(ψ2_coefs, j+1, :)
         end
 
         a = ψ1_coefs[ki+1, :]
@@ -129,16 +130,11 @@ end
 # end
 
 function legendre_filter(k)
-    # x = Symbol('x')
-    # H0 = np.zeros((k,k))
-    # H1 = np.zeros((k,k))
-    # G0 = np.zeros((k,k))
-    # G1 = np.zeros((k,k))
-    # PHI0 = np.zeros((k,k))
-    # PHI1 = np.zeros((k,k))
-    # phi, psi1, psi2 = get_phi_psi(k, base)
-
-    # ----------------------------------------------------------
+    H0 = zeros(k, k)legendre
+    H1 = zeros(k, k)
+    G0 = zeros(k, k)
+    G1 = zeros(k, k)
+    ϕ, ψ1, ψ2 = legendre_ϕ_ψ(k)
 
     # roots = Poly(legendre(k, 2*x-1)).all_roots()
     # x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
@@ -150,29 +146,23 @@ function legendre_filter(k)
     #         G0[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, x_m/2) * phi[kpi](x_m)).sum()
     #         H1[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki]((x_m+1)/2) * phi[kpi](x_m)).sum()
     #         G1[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, (x_m+1)/2) * phi[kpi](x_m)).sum()
-            
-    # PHI0 = np.eye(k)
-    # PHI1 = np.eye(k)
-
-    # ----------------------------------------------------------
 
-    # H0[np.abs(H0)<1e-8] = 0
-    # H1[np.abs(H1)<1e-8] = 0
-    # G0[np.abs(G0)<1e-8] = 0
-    # G1[np.abs(G1)<1e-8] = 0
+    zero_out!(H0)
+    zero_out!(H1)
+    zero_out!(G0)
+    zero_out!(G1)
 
-    # return H0, H1, G0, G1, PHI0, PHI1
+    return H0, H1, G0, G1, I(k), I(k)
 end
 
 function chebyshev_filter(k)
-    # x = Symbol('x')
-    # H0 = np.zeros((k,k))
-    # H1 = np.zeros((k,k))
-    # G0 = np.zeros((k,k))
-    # G1 = np.zeros((k,k))
-    # PHI0 = np.zeros((k,k))
-    # PHI1 = np.zeros((k,k))
-    # phi, psi1, psi2 = get_phi_psi(k, base)
+    H0 = zeros(k, k)
+    H1 = zeros(k, k)
+    G0 = zeros(k, k)
+    G1 = zeros(k, k)
+    Φ0 = zeros(k, k)
+    Φ1 = zeros(k, k)
+    ϕ, ψ1, ψ2 = chebyshev_ϕ_ψ(k)
 
     # ----------------------------------------------------------
 
@@ -193,16 +183,13 @@ function chebyshev_filter(k)
 
     #         PHI0[ki, kpi] = (wm * phi[ki](2*x_m) * phi[kpi](2*x_m)).sum() * 2
     #         PHI1[ki, kpi] = (wm * phi[ki](2*x_m-1) * phi[kpi](2*x_m-1)).sum() * 2
-            
-    # PHI0[np.abs(PHI0)<1e-8] = 0
-    # PHI1[np.abs(PHI1)<1e-8] = 0
-
-    # ----------------------------------------------------------
 
-    # H0[np.abs(H0)<1e-8] = 0
-    # H1[np.abs(H1)<1e-8] = 0
-    # G0[np.abs(G0)<1e-8] = 0
-    # G1[np.abs(G1)<1e-8] = 0
+    zero_out!(H0)
+    zero_out!(H1)
+    zero_out!(G0)
+    zero_out!(G1)
+    zero_out!(Φ0)
+    zero_out!(Φ1)
 
-    # return H0, H1, G0, G1, PHI0, PHI1
+    return H0, H1, G0, G1, Φ0, Φ1
 end

src/utils.jl

@@ -31,7 +31,6 @@ end
 
 function proj_factor(a, b; complement::Bool=false)
     prod_ = convolve(a, b)
-    zero_out!(prod_)
     r = collect(1:length(prod_))
     s = complement ? (1 .- 0.5 .^ r) : (0.5 .^ r)
     proj_ = sum(prod_ ./ r .* s)

src/wavelet.jl

@@ -65,11 +65,11 @@ struct MWT_CZ1d{T,S,R,Q,P}
 end
 
 function MWT_CZ1d(k::Int=3, α::Int=5, L::Int=0, c::Int=1; base::Symbol=:legendre, init=Flux.glorot_uniform)
-    H0, H1, G0, G1, ϕ0, ϕ1 = get_filter(base, k)
-    H0r = zero_out!(H0 * ϕ0)
-    G0r = zero_out!(G0 * ϕ0)
-    H1r = zero_out!(H1 * ϕ1)
-    G1r = zero_out!(G1 * ϕ1)
+    H0, H1, G0, G1, Φ0, Φ1 = get_filter(base, k)
+    H0r = zero_out!(H0 * Φ0)
+    G0r = zero_out!(G0 * Φ0)
+    H1r = zero_out!(H1 * Φ1)
+    G1r = zero_out!(G1 * Φ1)
 
     dim = c*k
     A = SpectralConv(dim=>dim, (α,); init=init)

test/polynomials.jl

@@ -0,0 +1,34 @@
+@testset "polynomials" begin
+    @testset "legendre_ϕ_ψ" begin
+        ϕ, ψ1, ψ2 = NeuralOperators.legendre_ϕ_ψ(10)
+
+        @test all(coeffs(ϕ[1]) .≈ [1.])
+        @test all(coeffs(ϕ[2]) .≈ [-1.7320508075688772, 3.4641016151377544])
+        @test all(coeffs(ϕ[3]) .≈ [2.23606797749979, -13.416407864998739, 13.416407864998739])
+        @test all(coeffs(ϕ[4]) .≈ [-2.6457513110645907, 31.74901573277509, -79.37253933193772, 52.91502622129181])
+        @test all(coeffs(ϕ[5]) .≈ [3.0, -60.0, 270.0, -420.0, 210.0])
+        @test all(coeffs(ϕ[6]) .≈ [-3.3166247903554, 99.498743710662, -696.491205974634, 1857.309882599024,
+                                -2089.4736179239017, 835.7894471695607])
+        @test all(coeffs(ϕ[7]) .≈ [3.605551275463989, -151.43315356948753, 1514.3315356948754, -6057.326142779501,
+                                11357.486517711566, -9994.588135586178, 3331.529378528726])
+        @test all(coeffs(ϕ[8]) .≈ [-3.872983346207417, 216.88706738761536, -2927.9754097328073, 16266.530054071152,
+                                -44732.957648695665, 64415.45901412176, -46522.27595464349, 13292.078844183856])
+        @test all(coeffs(ϕ[9]) .≈ [4.123105625617661, -296.86360504447157, 5195.113088278253, -38097.49598070719,
+                                142865.60992765194, -297160.46864951606, 346687.21342443535, -212257.47760679715,
+                                53064.36940169929])
+        @test all(coeffs(ϕ[10]) .≈ [-4.358898943540674, 392.30090491866065, -8630.619908210534, 80552.45247663166,
+                                -392693.20582357934, 1099540.9763060221, -1832568.2938433702, 1795168.9409077913,
+                                -953683.4998572641, 211929.66663494756])
+        
+        ϕ, ψ1, ψ2 = NeuralOperators.legendre_ϕ_ψ(3)
+        @test coeffs(ϕ[1]) ≈ [1.]
+        @test coeffs(ϕ[2]) ≈ [-1.7320508075688772, 3.4641016151377544]
+        @test coeffs(ϕ[3]) ≈ [2.23606797749979, -13.416407864998739, 13.416407864998739]
+        @test coeffs(ψ1[1]) ≈ [-1.0000000000000122, 6.000000000000073]
+        @test coeffs(ψ1[2]) ≈ [1.7320508075691732, -24.248711305967735, 51.96152422707261]
+        @test coeffs(ψ1[3]) ≈ [2.2360679774995615, -26.832815729994504, 53.665631459989214]
+        @test coeffs(ψ2[1]) ≈ [-5.000000000000066, 6.000000000000073]
+        @test coeffs(ψ2[2]) ≈ [29.44486372867492, -79.67433714817852, 51.96152422707261]
+        @test coeffs(ψ2[3]) ≈ [-29.068883707507286, 80.49844719001908, -53.665631460012115]
+    end
+end

test/runtests.jl

@@ -1,12 +1,14 @@
 using NeuralOperators
 using Test
 using Flux
+using Polynomials
 using Zygote
 using CUDA
 
 CUDA.allowscalar(false)
 
 @testset "NeuralOperators.jl" begin
+    include("polynomials.jl")
     include("fourier.jl")
     include("wavelet.jl")
     include("model.jl")