Skip to content
This repository has been archived by the owner on Sep 28, 2024. It is now read-only.

Commit

Permalink
format code
Browse files Browse the repository at this point in the history
  • Loading branch information
yuehhua committed Aug 6, 2022
1 parent ccb4131 commit 0fa7419
Show file tree
Hide file tree
Showing 7 changed files with 225 additions and 201 deletions.
153 changes: 81 additions & 72 deletions src/Transform/polynomials.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,47 +16,47 @@ function legendre_ϕ_ψ(k)
p = Polynomial([-1, 2]) # 2x-1
p2 = Polynomial([-1, 4]) # 4x-1

for ki in 0:(k-1)
for ki in 0:(k - 1)
l = convert(Polynomial, gen_poly(Legendre, ki)) # Legendre of n=ki
ϕ_coefs[ki+1, 1:(ki+1)] .= sqrt(2*ki+1) .* coeffs(l(p))
ϕ_2x_coefs[ki+1, 1:(ki+1)] .= sqrt(2*(2*ki+1)) .* coeffs(l(p2))
ϕ_coefs[ki + 1, 1:(ki + 1)] .= sqrt(2 * ki + 1) .* coeffs(l(p))
ϕ_2x_coefs[ki + 1, 1:(ki + 1)] .= sqrt(2 * (2 * ki + 1)) .* coeffs(l(p2))
end

ψ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)]
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)
ψ1_coefs[ki+1, :] .-= proj_ .* view(ϕ_coefs, i+1, :)
ψ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, :]
for j in 0:(k - 1)
a = ϕ_2x_coefs[ki + 1, 1:(ki + 1)]
b = ψ1_coefs[j + 1, :]
proj_ = proj_factor(a, b)
ψ1_coefs[ki+1, :] .-= proj_ .* view(ψ1_coefs, j+1, :)
ψ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, :]
a = ψ1_coefs[ki + 1, :]
norm1 = proj_factor(a, a)

a = ψ2_coefs[ki+1, :]
norm2 = proj_factor(a, a, complement=true)
a = ψ2_coefs[ki + 1, :]
norm2 = proj_factor(a, a, complement = true)
norm_ = sqrt(norm1 + norm2)
ψ1_coefs[ki+1, :] ./= norm_
ψ2_coefs[ki+1, :] ./= norm_
ψ1_coefs[ki + 1, :] ./= norm_
ψ2_coefs[ki + 1, :] ./= norm_
zero_out!(ψ1_coefs)
zero_out!(ψ2_coefs)
end

ϕ = [Polynomial(ϕ_coefs[i,:]) for i in 1:k]
ψ1 = [Polynomial(ψ1_coefs[i,:]) for i in 1:k]
ψ2 = [Polynomial(ψ2_coefs[i,:]) for i in 1:k]
ϕ = [Polynomial(ϕ_coefs[i, :]) for i in 1:k]
ψ1 = [Polynomial(ψ1_coefs[i, :]) for i in 1:k]
ψ2 = [Polynomial(ψ2_coefs[i, :]) for i in 1:k]

return ϕ, ψ1, ψ2
end
Expand All @@ -68,14 +68,14 @@ function chebyshev_ϕ_ψ(k)
p = Polynomial([-1, 2]) # 2x-1
p2 = Polynomial([-1, 4]) # 4x-1

for ki in 0:(k-1)
for ki in 0:(k - 1)
if ki == 0
ϕ_coefs[ki+1, 1:(ki+1)] .= sqrt(2/π)
ϕ_2x_coefs[ki+1, 1:(ki+1)] .= sqrt(4/π)
ϕ_coefs[ki + 1, 1:(ki + 1)] .= sqrt(2 / π)
ϕ_2x_coefs[ki + 1, 1:(ki + 1)] .= sqrt(4 / π)
else
c = convert(Polynomial, gen_poly(Chebyshev, ki)) # Chebyshev of n=ki
ϕ_coefs[ki+1, 1:(ki+1)] .= 2/sqrt(π) .* coeffs(c(p))
ϕ_2x_coefs[ki+1, 1:(ki+1)] .= sqrt(2) * 2/sqrt(π) .* coeffs(c(p2))
ϕ_coefs[ki + 1, 1:(ki + 1)] .= 2 / sqrt(π) .* coeffs(c(p))
ϕ_2x_coefs[ki + 1, 1:(ki + 1)] .= sqrt(2) * 2 / sqrt(π) .* coeffs(c(p2))
end
end

Expand All @@ -87,43 +87,43 @@ function chebyshev_ϕ_ψ(k)
# x_m[x_m==0.5] = 0.5 + 1e-8 # add small noise to avoid the case of 0.5 belonging to both phi(2x) and phi(2x-1)
# not needed for our purpose here, we use even k always to avoid
wm = π / k_use / 2

ψ1_coefs = zeros(k, k)
ψ2_coefs = zeros(k, k)

ψ1 = Array{Any,1}(undef, k)
ψ2 = Array{Any,1}(undef, k)
ψ1 = Array{Any, 1}(undef, k)
ψ2 = Array{Any, 1}(undef, k)

for ki in 0:(k-1)
ψ1_coefs[ki+1, :] .= ϕ_2x_coefs[ki+1, :]
for i in 0:(k-1)
proj_ = sum(wm .* ϕ[i+1].(x_m) .* sqrt(2) .* ϕ[ki+1].(2*x_m))
ψ1_coefs[ki+1, :] .-= proj_ .* view(ϕ_coefs, i+1, :)
ψ2_coefs[ki+1, :] .-= proj_ .* view(ϕ_coefs, i+1, :)
for ki in 0:(k - 1)
ψ1_coefs[ki + 1, :] .= ϕ_2x_coefs[ki + 1, :]
for i in 0:(k - 1)
proj_ = sum(wm .* ϕ[i + 1].(x_m) .* sqrt(2) .* ϕ[ki + 1].(2 * x_m))
ψ1_coefs[ki + 1, :] .-= proj_ .* view(ϕ_coefs, i + 1, :)
ψ2_coefs[ki + 1, :] .-= proj_ .* view(ϕ_coefs, i + 1, :)
end

for j in 0:(ki-1)
proj_ = sum(wm .* ψ1[j+1].(x_m) .* sqrt(2) .* ϕ[ki+1].(2*x_m))
ψ1_coefs[ki+1, :] .-= proj_ .* view(ψ1_coefs, j+1, :)
ψ2_coefs[ki+1, :] .-= proj_ .* view(ψ2_coefs, j+1, :)
for j in 0:(ki - 1)
proj_ = sum(wm .* ψ1[j + 1].(x_m) .* sqrt(2) .* ϕ[ki + 1].(2 * x_m))
ψ1_coefs[ki + 1, :] .-= proj_ .* view(ψ1_coefs, j + 1, :)
ψ2_coefs[ki + 1, :] .-= proj_ .* view(ψ2_coefs, j + 1, :)
end

ψ1[ki+1] = ϕ_(ψ1_coefs[ki+1,:]; lb=0., ub=0.5)
ψ2[ki+1] = ϕ_(ψ2_coefs[ki+1,:]; lb=0.5, ub=1.)
norm1 = sum(wm .* ψ1[ki+1].(x_m) .* ψ1[ki+1].(x_m))
norm2 = sum(wm .* ψ2[ki+1].(x_m) .* ψ2[ki+1].(x_m))
ψ1[ki + 1] = ϕ_(ψ1_coefs[ki + 1, :]; lb = 0.0, ub = 0.5)
ψ2[ki + 1] = ϕ_(ψ2_coefs[ki + 1, :]; lb = 0.5, ub = 1.0)

norm1 = sum(wm .* ψ1[ki + 1].(x_m) .* ψ1[ki + 1].(x_m))
norm2 = sum(wm .* ψ2[ki + 1].(x_m) .* ψ2[ki + 1].(x_m))

norm_ = sqrt(norm1 + norm2)
ψ1_coefs[ki+1, :] ./= norm_
ψ2_coefs[ki+1, :] ./= norm_
ψ1_coefs[ki + 1, :] ./= norm_
ψ2_coefs[ki + 1, :] ./= norm_
zero_out!(ψ1_coefs)
zero_out!(ψ2_coefs)
ψ1[ki+1] = ϕ_(ψ1_coefs[ki+1,:]; lb=0., ub=0.5+1e-16)
ψ2[ki+1] = ϕ_(ψ2_coefs[ki+1,:]; lb=0.5+1e-16, ub=1.)

ψ1[ki + 1] = ϕ_(ψ1_coefs[ki + 1, :]; lb = 0.0, ub = 0.5 + 1e-16)
ψ2[ki + 1] = ϕ_(ψ2_coefs[ki + 1, :]; lb = 0.5 + 1e-16, ub = 1.0)
end

return ϕ, ψ1, ψ2
end

Expand All @@ -137,22 +137,26 @@ function legendre_filter(k)
l = convert(Polynomial, gen_poly(Legendre, k))
x_m = roots(l(Polynomial([-1, 2]))) # 2x-1
m = 2 .* x_m .- 1
wm = 1 ./ k ./ legendre_der.(k, m) ./ gen_poly(Legendre, k-1).(m)

for ki in 0:(k-1)
for kpi in 0:(k-1)
H0[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ϕ[ki+1].(x_m/2) .* ϕ[kpi+1].(x_m))
G0[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, x_m/2) .* ϕ[kpi+1].(x_m))
H1[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ϕ[ki+1].((x_m.+1)/2) .* ϕ[kpi+1].(x_m))
G1[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, (x_m.+1)/2) .* ϕ[kpi+1].(x_m))
wm = 1 ./ k ./ legendre_der.(k, m) ./ gen_poly(Legendre, k - 1).(m)

for ki in 0:(k - 1)
for kpi in 0:(k - 1)
H0[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ϕ[ki + 1].(x_m / 2) .* ϕ[kpi + 1].(x_m))
G0[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ψ(ψ1, ψ2, ki, x_m / 2) .* ϕ[kpi + 1].(x_m))
H1[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ϕ[ki + 1].((x_m .+ 1) / 2) .* ϕ[kpi + 1].(x_m))
G1[ki + 1, kpi + 1] = 1 / sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, (x_m .+ 1) / 2) .*
ϕ[kpi + 1].(x_m))
end
end

zero_out!(H0)
zero_out!(H1)
zero_out!(G0)
zero_out!(G1)

return H0, H1, G0, G1, I(k), I(k)
end

Expand All @@ -172,14 +176,19 @@ function chebyshev_filter(k)
# not needed for our purpose here, we use even k always to avoid
wm = π / k_use / 2

for ki in 0:(k-1)
for kpi in 0:(k-1)
H0[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ϕ[ki+1].(x_m/2) .* ϕ[kpi+1].(x_m))
H1[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ϕ[ki+1].((x_m.+1)/2) .* ϕ[kpi+1].(x_m))
G0[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, x_m/2) .* ϕ[kpi+1].(x_m))
G1[ki+1, kpi+1] = 1/sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, (x_m.+1)/2) .* ϕ[kpi+1].(x_m))
Φ0[ki+1, kpi+1] = 2*sum(wm .* ϕ[ki+1].(2x_m) .* ϕ[kpi+1].(2x_m))
Φ1[ki+1, kpi+1] = 2*sum(wm .* ϕ[ki+1].(2 .* x_m .- 1) .* ϕ[kpi+1].(2 .* x_m .- 1))
for ki in 0:(k - 1)
for kpi in 0:(k - 1)
H0[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ϕ[ki + 1].(x_m / 2) .* ϕ[kpi + 1].(x_m))
H1[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ϕ[ki + 1].((x_m .+ 1) / 2) .* ϕ[kpi + 1].(x_m))
G0[ki + 1, kpi + 1] = 1 / sqrt(2) *
sum(wm .* ψ(ψ1, ψ2, ki, x_m / 2) .* ϕ[kpi + 1].(x_m))
G1[ki + 1, kpi + 1] = 1 / sqrt(2) * sum(wm .* ψ(ψ1, ψ2, ki, (x_m .+ 1) / 2) .*
ϕ[kpi + 1].(x_m))
Φ0[ki + 1, kpi + 1] = 2 * sum(wm .* ϕ[ki + 1].(2x_m) .* ϕ[kpi + 1].(2x_m))
Φ1[ki + 1, kpi + 1] = 2 * sum(wm .* ϕ[ki + 1].(2 .* x_m .- 1) .*
ϕ[kpi + 1].(2 .* x_m .- 1))
end
end

Expand All @@ -189,6 +198,6 @@ function chebyshev_filter(k)
zero_out!(G1)
zero_out!(Φ0)
zero_out!(Φ1)

return H0, H1, G0, G1, Φ0, Φ1
end
20 changes: 10 additions & 10 deletions src/Transform/utils.jl
Original file line number Diff line number Diff line change
@@ -1,46 +1,46 @@
function ϕ_(ϕ_coefs; lb::Real=0., ub::Real=1.)
function ϕ_(ϕ_coefs; lb::Real = 0.0, ub::Real = 1.0)
function partial(x)
mask = (lb x ub) * 1.
mask = (lb x ub) * 1.0
return Polynomial(ϕ_coefs)(x) * mask
end
return partial
end

function ψ(ψ1, ψ2, i, inp)
mask = (inp .> 0.5) .* 1.0
return ψ1[i+1].(inp) .* mask .+ ψ2[i+1].(inp) .* mask
return ψ1[i + 1].(inp) .* mask .+ ψ2[i + 1].(inp) .* mask
end

zero_out!(x; tol=1e-8) = (x[abs.(x) .< tol] .= 0)
zero_out!(x; tol = 1e-8) = (x[abs.(x) .< tol] .= 0)

function gen_poly(poly, n)
x = zeros(n+1)
x = zeros(n + 1)
x[end] = 1
return poly(x)
end

function convolve(a, b)
n = length(b)
y = similar(a, length(a)+n-1)
y = similar(a, length(a) + n - 1)
for i in 1:length(a)
y[i:(i+n-1)] .+= a[i] .* b
y[i:(i + n - 1)] .+= a[i] .* b
end
return y
end

function proj_factor(a, b; complement::Bool=false)
function proj_factor(a, b; complement::Bool = false)
prod_ = convolve(a, b)
r = collect(1:length(prod_))
s = complement ? (1 .- 0.5 .^ r) : (0.5 .^ r)
proj_ = sum(prod_ ./ r .* s)
return proj_
end

_legendre(k, x) = (2k+1) * gen_poly(Legendre, k)(x)
_legendre(k, x) = (2k + 1) * gen_poly(Legendre, k)(x)

function legendre_der(k, x)
out = 0
for i in k-1:-2:-1
for i in (k - 1):-2:-1
out += _legendre(i, x)
end
return out
Expand Down
12 changes: 5 additions & 7 deletions src/Transform/wavelet_transform.jl
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
export WaveletTransform

struct WaveletTransform{N, S}<:AbstractTransform
ec_d
ec_s
struct WaveletTransform{N, S} <: AbstractTransform
ec_d::Any
ec_s::Any
modes::NTuple{N, S} # N == ndims(x)
end

Base.ndims(::WaveletTransform{N}) where {N} = N

function transform(wt::WaveletTransform, 𝐱::AbstractArray)
N = size(X, ndims(wt)-1)
N = size(X, ndims(wt) - 1)
# 1d
Xa = vcat(view(𝐱, :, :, 1:2:N, :), view(𝐱, :, :, 2:2:N, :))
# 2d
Expand All @@ -24,9 +24,7 @@ function transform(wt::WaveletTransform, 𝐱::AbstractArray)
return d, s
end

function inverse(wt::WaveletTransform, 𝐱_fwt::AbstractArray)

end
function inverse(wt::WaveletTransform, 𝐱_fwt::AbstractArray) end

# function truncate_modes(wt::WaveletTransform, 𝐱_fft::AbstractArray)
# return view(𝐱_fft, map(d->1:d, wt.modes)..., :, :) # [ft.modes..., in_chs, batch]
Expand Down
Loading

0 comments on commit 0fa7419

Please sign in to comment.