Skip to content

Instantly share code, notes, and snippets.

@matthewelmer
Last active October 9, 2023 08:04
Show Gist options
  • Save matthewelmer/8feddff81a2ddeba85712fb2a47fa53a to your computer and use it in GitHub Desktop.
Save matthewelmer/8feddff81a2ddeba85712fb2a47fa53a to your computer and use it in GitHub Desktop.
Exemplary Julia Code (p3.jl, the others are included just in case you want to try running it)
module CGL;
export cgl;
using Revise; # you must do this before loading any revisable packages
function cgl(n::Integer)::Vector{Real}
# Returns n cgl points on the closed interval [-1, 1].
return @. cos(pi * ((1:n) - n) / (n - 1));
end
function cgl(n::Integer, zL::Real, zR::Real)::Vector{Real}
# Returns n cgl points between zL and zR i.e. the open sub-interval (zL, zR)
# Check input
if zL < -1 || zR > 1
return []; # TODO(m elmer): Exception
end
# Modify formula for CGL point generation to exclude boundaries and map it
# to the sub-interval
n += 2
modified_cos_arg = @. pi * ((2:(n-1)) - n) / (n - 1);
scale_to_sub_interval = (zR - zL) / 2;
shift_to_sub_interval = (zR - zL) / 2 + zL
return @. cos(modified_cos_arg) * scale_to_sub_interval + shift_to_sub_interval;
end
# For some reason, trying to specify the type of z_poi as `::Vector{Real}`
# causes the following error: `MethodError: no method matching cgl(::Int64,
# ::Vector{Float64})` (TODO: file bug report??)
function cgl(n::Integer, z_poi::Vector{T})::Vector{Real} where T <: Real
if minimum(z_poi) < -1 || maximum(z_poi) > 1
return []; # TODO(m elmer): Exception
end
# Initialize bins for the CGL points between each point of interest, scaling
# the number of points in each bin roughly to how much space it takes on the
# interval.
nk = Vector{Integer}(undef, (size(z_poi, 1) + 2) - 1);
nk[1] = round(Integer, n * (z_poi[1] - (-1)) / 2) - 1;
for i in 2:(size(nk)[1] - 1)
nk[i] = round(Integer, n * (z_poi[i] - z_poi[i - 1]) / 2) - 1;
end
nk[end] = round(Integer, n * (1 - z_poi[end]) / 2) - 2;
# If the rough distribution of points resulted in more points than the user
# asked for, remove points from the bins with the most points.
while sum(nk) + size(z_poi, 1) + 2 > n
nk[argmax(nk)] -= 1;
end
# If the rough distribution of points resulted in fewer points than the user
# asked for, add points to the bins with the fewest points.
while sum(nk) + size(z_poi, 1) + 2 < n
nk[argmin(nk)] += 1;
end
# Now that we've sorted out exactly how many points we want in each bin,
# generate the CGL points within each bin, put them between the points of
# interest, and return the result.
result = [-1; cgl(nk[1], -1, z_poi[1])];
for i in 1:(size(z_poi, 1) - 1)
result = [result; z_poi[i]; cgl(nk[i + 1], z_poi[i], z_poi[i + 1])];
end
result = [result; z_poi[end]; cgl(nk[end], z_poi[end], 1); 1];
return result;
end
end;
module LeastSquares;
export sqrls;
using Revise; # you must do this before loading any revisable packages
using LinearAlgebra;
function sqrls(A::Matrix{T}, b::Vector{T})::Tuple{Vector{T}, T} where T <: Number
# Solves Ax = b for x using QR least squares with scaling.
# Returns x along with the condition number of the R matrix.
S = 1 ./ sqrt.(sum(A .* A, dims=1));
S = reshape(S, (size(A)[2],));
QR_decomp = qr(A * diagm(S));
x = S .* (QR_decomp \ b); # This does QR least squares
cn = cond(QR_decomp.R);
return x, cn
end
end;
module OrthoPoly;
export legendre_mat, legendre_mat_1, legendre_mat_2,
chebyshev_mat, chebyshev_mat_1, chebyshev_mat_2;
using Revise; # you must do this before loading any revisable packages
# TODO(melmer): Add types (maybe?)
function legendre_mat(z, d)
# Return a matrix containing every Legendre polynomial up to degree d
# evaluated at each point in z. Degree increases column-by-column, and the
# value of z increases row-by-row.
if d < 0
throw(ArgumentError(:d))
end
L = Matrix{Float64}(undef, size(z, 1), d+1)
L[:, 1] .= 1
L[:, 2] = z
for k = 3:(d+1)
# NOTE(melmer): In my testing, vectorizing this inner for loop increased
# runtime and memory allocation.
for j in 1:size(L, 1)
L[j, k] = (2 * (k-2) + 1) / ((k-2) + 1) *
z[j] * L[j, k-1] -
(k-2) / ((k-2) + 1) * L[j, k-2];
end
end
return L
end
function legendre_mat_1(L)
# Return a matrix containing the first derivative of every Legendre
# polynomial up to degree d evaluated at each point in z. Degree increases
# column-by-column, and the value of z increases row-by-row.
z = L[:, 2]
LD = similar(L)
LD[:, 1] .= 0
LD[:, 2] .= 1
for k in 3:size(L, 2)
# NOTE(melmer): In my testing, vectorizing this inner for loop increased
# runtime and memory allocation.
for j in 1:size(L, 1)
LD[j, k] = (2 * (k-2) + 1) / ((k-2) + 1) *
(L[j, k-1] + z[j] * LD[j, k-1]) -
(k-2) / ((k-2) + 1) * LD[j, k-2];
end
end
return LD
end
function legendre_mat_2(L, LD)
# Return a matrix containing the second derivative of every Legendre
# polynomial up to degree d evaluated at each point in z. Degree increases
# column-by-column, and the value of z increases row-by-row.
z = L[:, 2]
LDD = similar(LD)
LDD[:, 1:2] .= 0
for k in 3:size(LD, 2)
# NOTE(melmer): In my testing, vectorizing this inner for loop increased
# runtime and memory allocation.
for j in 1:size(L, 1)
LDD[j, k] = (2 * (k-2) + 1) / ((k-2) + 1) *
(2 * LD[j, k-1] + z[j] * LDD[j, k-1]) -
(k-2) / ((k-2) + 1) * LDD[j, k-2];
end
end
return LDD
end
function chebyshev_mat(z, d)
# TODO(melmer): Bring in line with Legendre polynomial generators
# Return a matrix containing every Chebyshev polynomial up to degree d
# evaluated at each point in z. Degree increases column-by-column, and the
# value of z increases row-by-row.
if d < 0
throw(ArgumentError(:d))
end
n = size(z, 1)
T = Matrix{Float64}(undef, n, d+1)
T[:, 1] .= 1
T[:, 2] = z
for k = 3:(d+1)
km2 = k - 2
T[:, k] = 2 * z .* T[:, k-1] .- km2 / (km2 + 1) * T[:, k-2]
end
return T
end
function chebyshev_mat_1(T)
# TODO(melmer): Bring in line with Legendre polynomial generators
# Return a matrix containing the first derivative of every Chebyshev
# polynomial up to degree d evaluated at each point in z. Degree increases
# column-by-column, and the value of z increases row-by-row.
x = T[:, 2];
TD = similar(T);
TD[:, 1] = zeros(size(T, 1), 1);
TD[:, 2] = ones(size(T, 1), 1);
for i = 3:size(T, 2)
for j = 1:size(T, 1)
TD[j, i] = 2 * (1 * T[j, i-1] + x[j] * TD[j, i-1]) - TD[j, i-2];
end
end
return TD;
end
function chebyshev_mat_2(T, TD)
# TODO(melmer): Bring in line with Legendre polynomial generators
# Return a matrix containing the second derivative of every Chebyshev
# polynomial up to degree d evaluated at each point in z. Degree increases
# column-by-column, and the value of z increases row-by-row.
x = T[:, 2];
TDD = similar(TD);
TDD[:, 1:2] = zeros(size(TD, 1), 2);
for i = 3:size(TD, 2)
for j = 1:size(TD, 1)
TDD[j, i] = 2 * (2 * TD[j, i-1] + x[j] * TDD[j, i-1]) - TDD[j, i-2];
end
end
return TDD;
end
end;
cn = 583.288
Computation: 0.000609 seconds (124 allocations: 2.100 MiB)
Plotting: 0.158417 seconds (13.15 k allocations: 1.157 MiB)
Total: 0.159164 seconds (13.32 k allocations: 3.261 MiB)
module P3 # Make the script a module for a deterministic workflow
using Revise
using Plots
using .Main: @showp, @showc, nearly_equal
include("ortho_poly.jl")
using .OrthoPoly
include("least_squares.jl")
using .LeastSquares
include("cgl.jl")
using .CGL
const n = 250
const Nbf = 35 + 2 # We're stripping a couple off, so add a couple
const _1 = ones(n)
const z(t) = 2 / 7 * t - 1
const t = 7 / 2 * (cgl(n, [z(1); z(2); z(3); z(5)]) .+ 1)
const c = 2 / 7
function computation()
# Need to keep a non-truncated version while generating derivatives
H = legendre_mat(z.(t), Nbf - 1)
HD = legendre_mat_1(H)
HDD = legendre_mat_2(H, HD)
H = H[:, 3:end]
HD = HD[:, 3:end]
HDD = HDD[:, 3:end]
H_teq1 = legendre_mat(z.(_1), Nbf - 1)
H_teq2 = legendre_mat(z.(2 * _1), Nbf - 1)
H_teq3 = legendre_mat(z.(3 * _1), Nbf - 1)
H_teq5 = legendre_mat(z.(5 * _1), Nbf - 1)
HD_teq1 = legendre_mat_1(H_teq1)
H_teq1 = H_teq1[:, 3:end]
H_teq2 = H_teq2[:, 3:end]
H_teq3 = H_teq3[:, 3:end]
H_teq5 = H_teq5[:, 3:end]
HD_teq1 = HD_teq1[:, 3:end]
A = c^2 * HDD .+
3 * c * t .* HD .-
6 * H .+
(π / (1 - 3 * π) * t .+ 2) .* (5 * H_teq2 .- 2 * H_teq5) .-
3 * t * 1 / (1 - 3 * π) .* (π * H_teq3 .- c * HD_teq1)
b = 2 .+ 12 / (1 - 3 * π) * t
ξ, cn = sqrls(A, b)
@showc cn
λ = 1 / 3 .* (2 * H_teq5 .- 5 * H_teq2) * ξ
y = H * ξ .+
λ .+
1 / (1 - 3 * π) * (π * (λ .+ H_teq3 * ξ) .+ 4 .- c * HD_teq1 * ξ) .* t
residual = A * ξ .- b
return (y=y, residual=residual)
end
function plotting(results)
default_fig_width, default_fig_height = (600, 400)
size = (default_fig_width - 100, default_fig_height)
p = plot(
t,
results.y,
label=nothing,
xlabel="\$t\$",
ylabel="\$y\$",
title="TFC Solution",
dpi=300,
linewidth=2,
size=size
)
savefig("p3_tfc_soln.png")
p = plot(
t,
results.residual,
label=nothing,
title="TFC Residual",
xlabel="\$t\$",
ylabel="\$y\$",
dpi=300,
linewidth=2,
markershape=:circle,
markersize=3,
size=size,
yrotation=55,
ytickfontvalign=:bottom,
ytickfonthalign=:left
)
savefig("p3_tfc_residual.png")
end
function main()
results = @time "Computation" computation()
@time "Plotting" plotting(results)
end
@time "Total" main()
end;
module StartupUtils
export @showc, @showp, nearly_equal
using Revise;
using LinearAlgebra;
using Plots;
# Compact show
macro showc(exs...)
blk = Expr(:block)
for ex in exs
push!(blk.args, :(println($(sprint(Base.show_unquoted,ex)*" = "),
repr(begin local value = $(esc(ex)) end, context = :compact=>true))))
end
isempty(exs) || push!(blk.args, :value)
return blk
end
# Pretty show
macro showp(exs...)
blk = Expr(:block)
for ex in exs
push!(blk.args, :(println($(sprint(Base.show_unquoted,ex)*" = "),
repr(MIME("text/plain"), begin local value = $(esc(ex)) end, context = :limit=>true))))
end
isempty(exs) || push!(blk.args, :value)
return blk
end
# Relatively robust floating-point comparisons
function nearly_equal(
a::T, b::T, rel_tol::T = 128 * eps(T), abs_tol::T = eps(T)
)::Bool where T <: AbstractFloat
# Some handy constants for determining whether the defaults are suitable:
# eps(Float64) = 2.220446049250313e-16
# eps(Float32) = 1.1920929f-7
# eps(Float16) = 0.000977
# floatmin(Float64) = 2.2250738585072014e-308
# floatmin(Float32) = 1.1754944f-38
# floatmin(Float16) = 6.104e-5
# See also: nextfloat e.g. nextfloat(1000.0::Float64) - 1000.0::Float64 =
# 1.1368683772161603e-13
if a == b return true; end
diff = abs(a - b);
norm = min(abs(a) + abs(b), floatmax(T))
return diff < max(abs_tol, rel_tol * norm);
end
# Now make it work on vectors
# TODO(m elmer): SIMD this hoe
function nearly_equal(
a::Vector{T}, b::Vector{T}, rel_tol::T = 128 * eps(T),
abs_tol::T = 128 * floatmin(T)
)::Bool where T <: AbstractFloat
if size(a) != size(b) throw(DimensionMismatch("Dimension mismatch.")); end
for i in 1:size(a, 1)
if !nearly_equal(a[i], b[i], rel_tol, abs_tol) return false; end
end
return true;
end
end
using .StartupUtils;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment