From fa92c7061af08f6a8b2a8d38775bd7a3891c8271 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 10:52:43 +0100 Subject: [PATCH 01/28] fix .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 9b1a5785..4c6a7142 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ _git2_* docs/build *.info .vscode/spellright.dict +deps/build.log From 173ee665347cecc1d3883a27fbc226a7ccaaa938 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 10:54:41 +0100 Subject: [PATCH 02/28] add new Propagator module --- src/Luna.jl | 3 ++- src/Propagator.jl | 9 +++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) create mode 100644 src/Propagator.jl diff --git a/src/Luna.jl b/src/Luna.jl index 24de3188..4270b2bb 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -57,6 +57,7 @@ include("Grid.jl") include("Modes.jl") include("Fields.jl") include("RK45.jl") +include("Propagator.jl") include("LinearOps.jl") include("Capillary.jl") include("Antiresonant.jl") @@ -88,7 +89,7 @@ export Utils, Scans, Output, Maths, PhysData, Grid, RK45, Modes, Capillary, Rect Nonlinear, Ionisation, NonlinearRHS, LinearOps, Stats, Polarisation, Tools, Plotting, Raman, Antiresonant, Fields, Processing, Interface, SFA, prop_capillary, prop_gnlse, Pulses, Scan, runscan, makefilename, addvariable!, - StepIndexFibre, SimpleFibre + StepIndexFibre, SimpleFibre, Propagator # for a tuple of TimeFields we assume all inputs are for mode 1 function doinput_sm(grid, inputs::Tuple{Vararg{T} where T <: Fields.TimeField}, FT) diff --git a/src/Propagator.jl b/src/Propagator.jl new file mode 100644 index 00000000..fd527579 --- /dev/null +++ b/src/Propagator.jl @@ -0,0 +1,9 @@ +module Propagator +import Dates +import Logging +import Printf: @sprintf +import Luna.Utils: format_elapsed + + + +end # module From d1655655df72155434da31bfc73230599df15fff Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 11:07:26 +0100 Subject: [PATCH 03/28] Fix deps --- Project.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Project.toml b/Project.toml index a1fb9c8f..4e471255 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826" NumericalIntegration = "e7bfaba1-d571-5449-8927-abc22e82249b" Optim = "429524aa-4258-5aef-a3af-852621145aeb" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Peaks = "18e31ff7-3703-566c-8e60-38913d67486b" PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" @@ -82,6 +83,7 @@ Logging = "1.9" Memoize = "0.4.4" NumericalIntegration = "0.3.3" Optim = "1.4" +OrdinaryDiffEq = "6.102.0" Peaks = "0.3.2, 0.4, 0.5" PhysicalConstants = "0.2" Pkg = "1.9" From 741be490b46363b0028dd673356b4d8dbd0a067a Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 19:57:56 +0100 Subject: [PATCH 04/28] first attempt --- src/Luna.jl | 12 ++--- src/Propagator.jl | 109 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 5 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 4270b2bb..f737008f 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -384,11 +384,13 @@ function run(Eω, grid, save_modeinfo_maybe(output, transform) flush(stderr) # flush std error once before starting to show setup steps - RK45.solve_precon( - transform, linop, Eω, z0, init_dz, grid.zmax, stepfun=stepfun, - max_dt=max_dz, min_dt=min_dz, - rtol=rtol, atol=atol, safety=safety, norm=norm, - status_period=status_period) + #RK45.solve_precon( + # transform, linop, Eω, z0, init_dz, grid.zmax, stepfun=stepfun, + # max_dt=max_dz, min_dt=min_dz, + # rtol=rtol, atol=atol, safety=safety, norm=norm, + # status_period=status_period) + Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; + rtol, atol, max_dz, min_dz, status_period) end # run some code for precompilation diff --git a/src/Propagator.jl b/src/Propagator.jl index fd527579..505012cf 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -3,7 +3,116 @@ import Dates import Logging import Printf: @sprintf import Luna.Utils: format_elapsed +import OrdinaryDiffEq as ODE +mutable struct Printer{DT} + status_period::Int + start::DT + tic::DT + steps::Int + zmax::Float64 +end +function Printer(status_period, zmax) + Logging.@info "Starting propagation" + Printer(status_period, Dates.now(), Dates.now(), 0, zmax) +end + +function printstep(p::Printer, z, dz) + p.steps += 1 + if Dates.value(Dates.now() - p.tic) > 1000*p.status_period + speed = z/(Dates.value(Dates.now() - p.start)/1000) + eta_in_s = (p.zmax - z)/speed + if eta_in_s > 356400 + Logging.@info @sprintf("Progress: %.2f %%, ETA: XX:XX:XX, stepsize %.2e", + z/p.zmax*100, dz) + else + eta_in_ms = Dates.Millisecond(ceil(eta_in_s*1000)) + etad = Dates.DateTime(Dates.UTInstant(eta_in_ms)) + Logging.@info @sprintf("Progress: %.2f %%, ETA: %s, stepsize %.2e", + z/p.zmax*100, Dates.format(etad, "HH:MM:SS"), dz) + end + flush(stderr) + p.tic = Dates.now() + end +end + +function printstop(p::Printer) + totaltime = Dates.now() - p.start + dtstring = format_elapsed(totaltime) + Logging.@info @sprintf("Propagation finished in %s, %d steps", + dtstring, p.steps) +end + +abstract type AbstractPropagator end + +struct ConstPropagator{NLT} <: AbstractPropagator + L::Vector{ComplexF64} + nonlinop!::NLT + Eωtmp::Vector{ComplexF64} + Pωtmp::Vector{ComplexF64} +end + +#function make_propagator(f!, linop::Vector{ComplexF64}, u0::Vector{ComplexF64}) +# prop = ConstPropagator(linop, f!, similar(u0), similar(u0)) +#end + +function fcl!(du,u,p,z) + @. p.Eωtmp = u * exp(p.L * z) + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) + @. du = p.Pωtmp * exp(-p.L * z) +end + +struct NonConstPropagator{LT, NLT, SFT, PT} <: AbstractPropagator + linop!::LT + nonlinop!::NLT + stepfun::SFT + n::Int + Eωtmp::Vector{ComplexF64} + Pωtmp::Vector{ComplexF64} + p::PT +end + +function fncl!(du,u,p,z) + Eω = @views u[1:p.n] + L = @views u[p.n+1:end] + dEω = @views du[1:p.n] + dL = @views du[p.n+1:end] + @. p.Eωtmp = Eω * exp(L) + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) + @. dEω = p.Pωtmp * exp(-L) + p.linop!(dL, z) +end + +function callbackncl(integrator) + n = integrator.p.n + Eω = @views integrator.p.u[1:n] + L = @views integrator.p.u[n+1:end] + @. integrator.p.Eωtmp = Eω * exp(L) + interp = let integrator=integrator, n=n + function interp(z) + u = integrator(z) + Eω = @views u[1:n] + L = @views u[n+1:end] + @. Eω * exp(L) + end + end + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) + printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) + u_modified!(integrator, false) +end + +function propagate(f!, linop!, Eω0, z, zmax, stepfun; + rtol=1e-6, atol=1e-10, max_dz=Inf, min_dz=0, status_period=1) + p = Printer(status_period, zmax) + prop = NonConstPropagator(linop!, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), p) + u0 = vcat(Eω0, zero(Eω0)) + prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) + cb = ODE.DiscreteCallback((u,t,integrator) -> true, callbackncl, save_positions=(false,false)) + integrator = ODE.init(prob, ODE.Tsit5(); adaptive=true, reltol=rtol, abstol=atol, + dtmin=min_dz, dtmax=max_dz, callback=cb) + ODE.solve!(integrator) + printstop(p) +end end # module From 6149e4e679561da07a19c8e1ddc64cb5f7a37a29 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 20:11:44 +0100 Subject: [PATCH 05/28] working prop --- src/Luna.jl | 14 +++++++------- src/Propagator.jl | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index f737008f..88e1cd58 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -395,13 +395,13 @@ end # run some code for precompilation Logging.with_logger(Logging.NullLogger()) do - prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, - τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) + #prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, + # τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) prop_capillary(125e-6, 0.3, :He, (1.0, 0); λ0=800e-9, energy=1e-9, τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) - prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, - τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11, - modes=4) + #prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, + # τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11, + # modes=4) p = Tools.capillary_params(120e-6, 10e-15, 800e-9, 125e-6, :He, P=1.0) # gnlse_sol.jl example but with N=1 and 100th of the fibre length @@ -416,8 +416,8 @@ Logging.with_logger(Logging.NullLogger()) do λ0 = 835e-9 λlims = [450e-9, 8000e-9] trange = 4e-12 - output = prop_gnlse(γ, flength, βs; λ0, τfwhm, power=P0, pulseshape=:sech, λlims, trange, - raman=true, shock=true, fr, shotnoise=true, saveN=11) + #output = prop_gnlse(γ, flength, βs; λ0, τfwhm, power=P0, pulseshape=:sech, λlims, trange, + # raman=true, shock=true, fr, shotnoise=true, saveN=11) end end # module diff --git a/src/Propagator.jl b/src/Propagator.jl index 505012cf..5e097cdc 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -86,8 +86,8 @@ end function callbackncl(integrator) n = integrator.p.n - Eω = @views integrator.p.u[1:n] - L = @views integrator.p.u[n+1:end] + Eω = @views integrator.u[1:n] + L = @views integrator.u[n+1:end] @. integrator.p.Eωtmp = Eω * exp(L) interp = let integrator=integrator, n=n function interp(z) @@ -99,7 +99,7 @@ function callbackncl(integrator) end integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - u_modified!(integrator, false) + ODE.u_modified!(integrator, false) end function propagate(f!, linop!, Eω0, z, zmax, stepfun; From f1e5cb40ce40e1ace816f2a5fd7fe7e10617f464 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 20:51:16 +0100 Subject: [PATCH 06/28] make multimode work --- src/Luna.jl | 18 ++++++------- src/Propagator.jl | 64 ++++++++++++++++++++++++++++++++++------------- 2 files changed, 55 insertions(+), 27 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 88e1cd58..02c0de86 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -359,7 +359,7 @@ save_modeinfo_maybe(output, t) = nothing function run(Eω, grid, linop, transform, FT, output; min_dz=0, max_dz=grid.zmax/2, init_dz=1e-4, z0=0.0, - rtol=1e-6, atol=1e-10, safety=0.9, norm=RK45.weaknorm, + rtol=1e-3, atol=1e-6, safety=0.9, norm=RK45.weaknorm, status_period=1) Et = FT \ Eω @@ -395,13 +395,13 @@ end # run some code for precompilation Logging.with_logger(Logging.NullLogger()) do - #prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, - # τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) - prop_capillary(125e-6, 0.3, :He, (1.0, 0); λ0=800e-9, energy=1e-9, + prop_capillary(125e-6, 0.03, :He, 1.0; λ0=800e-9, energy=1e-9, τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) - #prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9, - # τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11, - # modes=4) + prop_capillary(125e-6, 0.03, :He, (1.0, 0); λ0=800e-9, energy=1e-9, + τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11) + prop_capillary(125e-6, 0.03, :He, 1.0; λ0=800e-9, energy=1e-9, + τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11, + modes=4) p = Tools.capillary_params(120e-6, 10e-15, 800e-9, 125e-6, :He, P=1.0) # gnlse_sol.jl example but with N=1 and 100th of the fibre length @@ -416,8 +416,8 @@ Logging.with_logger(Logging.NullLogger()) do λ0 = 835e-9 λlims = [450e-9, 8000e-9] trange = 4e-12 - #output = prop_gnlse(γ, flength, βs; λ0, τfwhm, power=P0, pulseshape=:sech, λlims, trange, - # raman=true, shock=true, fr, shotnoise=true, saveN=11) + output = prop_gnlse(γ, flength, βs; λ0, τfwhm, power=P0, pulseshape=:sech, λlims, trange, + raman=true, shock=true, fr, shotnoise=true, saveN=11) end end # module diff --git a/src/Propagator.jl b/src/Propagator.jl index 5e097cdc..42f7011d 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -14,10 +14,15 @@ mutable struct Printer{DT} end function Printer(status_period, zmax) - Logging.@info "Starting propagation" Printer(status_period, Dates.now(), Dates.now(), 0, zmax) end +function printstart(p::Printer) + p.start = Dates.now() + p.tic = Dates.now() + Logging.@info "Starting propagation" +end + function printstep(p::Printer, z, dz) p.steps += 1 if Dates.value(Dates.now() - p.tic) > 1000*p.status_period @@ -46,30 +51,41 @@ end abstract type AbstractPropagator end -struct ConstPropagator{NLT} <: AbstractPropagator - L::Vector{ComplexF64} +struct ConstPropagator{NLT, SFT, PT, AT} <: AbstractPropagator + L::AT nonlinop!::NLT - Eωtmp::Vector{ComplexF64} - Pωtmp::Vector{ComplexF64} + stepfun::SFT + Eωtmp::AT + Pωtmp::AT + p::PT end -#function make_propagator(f!, linop::Vector{ComplexF64}, u0::Vector{ComplexF64}) -# prop = ConstPropagator(linop, f!, similar(u0), similar(u0)) -#end - function fcl!(du,u,p,z) @. p.Eωtmp = u * exp(p.L * z) p.nonlinop!(p.Pωtmp, p.Eωtmp, z) @. du = p.Pωtmp * exp(-p.L * z) end -struct NonConstPropagator{LT, NLT, SFT, PT} <: AbstractPropagator +function callbackcl(integrator) + @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.L * integrator.t) + interp = let integrator=integrator + function interp(z) + u = integrator(z) + @. u * exp(integrator.p.L * z) + end + end + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) + printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) + ODE.u_modified!(integrator, false) +end + +struct NonConstPropagator{LT, NLT, SFT, PT, AT} <: AbstractPropagator linop!::LT nonlinop!::NLT stepfun::SFT n::Int - Eωtmp::Vector{ComplexF64} - Pωtmp::Vector{ComplexF64} + Eωtmp::AT + Pωtmp::AT p::PT end @@ -102,17 +118,29 @@ function callbackncl(integrator) ODE.u_modified!(integrator, false) end -function propagate(f!, linop!, Eω0, z, zmax, stepfun; - rtol=1e-6, atol=1e-10, max_dz=Inf, min_dz=0, status_period=1) - p = Printer(status_period, zmax) - prop = NonConstPropagator(linop!, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), p) +function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer) where N + prop = ConstPropagator(linop, f!, stepfun, similar(Eω0), similar(Eω0), printer) + prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) + prob, callbackcl +end + +function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) + prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) u0 = vcat(Eω0, zero(Eω0)) prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) - cb = ODE.DiscreteCallback((u,t,integrator) -> true, callbackncl, save_positions=(false,false)) + prob, callbackncl +end + +function propagate(f!, linop, Eω0, z, zmax, stepfun; + rtol=1e-3, atol=1e-6, max_dz=Inf, min_dz=0, status_period=1) + printer = Printer(status_period, zmax) + prob, cbfunc = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) + cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) integrator = ODE.init(prob, ODE.Tsit5(); adaptive=true, reltol=rtol, abstol=atol, dtmin=min_dz, dtmax=max_dz, callback=cb) + printstart(printer) ODE.solve!(integrator) - printstop(p) + printstop(printer) end end # module From 848a16c40ad17764a45f870161db4e728e50069e Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 21:58:49 +0100 Subject: [PATCH 07/28] fix init --- src/Luna.jl | 2 +- src/Propagator.jl | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 02c0de86..91bb9265 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -390,7 +390,7 @@ function run(Eω, grid, # rtol=rtol, atol=atol, safety=safety, norm=norm, # status_period=status_period) Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, max_dz, min_dz, status_period) + rtol, atol, init_dz, max_dz, min_dz, status_period) end # run some code for precompilation diff --git a/src/Propagator.jl b/src/Propagator.jl index 42f7011d..d5a13990 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -51,7 +51,7 @@ end abstract type AbstractPropagator end -struct ConstPropagator{NLT, SFT, PT, AT} <: AbstractPropagator +struct ConstPropagator{NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator L::AT nonlinop!::NLT stepfun::SFT @@ -79,7 +79,7 @@ function callbackcl(integrator) ODE.u_modified!(integrator, false) end -struct NonConstPropagator{LT, NLT, SFT, PT, AT} <: AbstractPropagator +struct NonConstPropagator{LT, NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator linop!::LT nonlinop!::NLT stepfun::SFT @@ -132,12 +132,12 @@ function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) end function propagate(f!, linop, Eω0, z, zmax, stepfun; - rtol=1e-3, atol=1e-6, max_dz=Inf, min_dz=0, status_period=1) + rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, status_period=1) printer = Printer(status_period, zmax) prob, cbfunc = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) integrator = ODE.init(prob, ODE.Tsit5(); adaptive=true, reltol=rtol, abstol=atol, - dtmin=min_dz, dtmax=max_dz, callback=cb) + dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) printstart(printer) ODE.solve!(integrator) printstop(printer) From 93e9b39f9269a32f2fe2c39b6f60ede67c7c712e Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 22:11:10 +0100 Subject: [PATCH 08/28] Add interface --- src/Interface.jl | 24 ++++++++++++++++++++---- src/Luna.jl | 31 ++++++++++++++++++++++--------- 2 files changed, 42 insertions(+), 13 deletions(-) diff --git a/src/Interface.jl b/src/Interface.jl index 7495088a..42cfff0b 100644 --- a/src/Interface.jl +++ b/src/Interface.jl @@ -339,11 +339,19 @@ If `raman` is `true`, then the following options apply: - `scanidx`: Current scan index within a scan being run. Only used when `scan` is passed. - `filename`: Can be used to to overwrite the scan name when running a parameter scan. The running `scanidx` will be appended to this filename. Ignored if no `scan` is given. + +# Solver options +- `rtol::Number`: Relative tolerance for the solver. Defaults to an optimal solver dependent value. +- `atol::Number`: Absolute tolerance for the solver. Defaults to an optimal solver dependent value. +- `solver::Symbol`: The solver to use. Defaults to `:Tsit5` from DifferentialEquations.jl. + Other possible choices can be found in the non-stiff ODE documentation of DifferentialEquations.jl. + For the original Luna RK45 solver, use `:OrigRK45`. - `status_period::Number`: Interval (in seconds) between printed status updates. """ -function prop_capillary(args...; status_period=5, kwargs...) +function prop_capillary(args...; rtol=nothing, atol=nothing, solver=:Tsit5, + status_period=5, kwargs...) Eω, grid, linop, transform, FT, output = prop_capillary_args(args...; kwargs...) - Luna.run(Eω, grid, linop, transform, FT, output; status_period) + Luna.run(Eω, grid, linop, transform, FT, output; rtol, atol, solver, status_period) output end @@ -897,11 +905,19 @@ Note that the current GNLSE model is single mode only. - `scanidx`: Current scan index within a scan being run. Only used when `scan` is passed. - `filename`: Can be used to to overwrite the scan name when running a parameter scan. The running `scanidx` will be appended to this filename. Ignored if no `scan` is given. + +# Solver options +- `rtol::Number`: Relative tolerance for the solver. Defaults to an optimal solver dependent value. +- `atol::Number`: Absolute tolerance for the solver. Defaults to an optimal solver dependent value. +- `solver::Symbol`: The solver to use. Defaults to `:Tsit5` from DifferentialEquations.jl. + Other possible choices can be found in the non-stiff ODE documentation of DifferentialEquations.jl. + For the original Luna RK45 solver, use `:OrigRK45`. - `status_period::Number`: Interval (in seconds) between printed status updates. """ -function prop_gnlse(args...; status_period=5, kwargs...) +function prop_gnlse(args...; rtol=nothing, atol=nothing, solver=:Tsit5, + status_period=5, kwargs...) Eω, grid, linop, transform, FT, output = prop_gnlse_args(args...; kwargs...) - Luna.run(Eω, grid, linop, transform, FT, output; status_period) + Luna.run(Eω, grid, linop, transform, FT, output; rtol, atol, solver, status_period) output end diff --git a/src/Luna.jl b/src/Luna.jl index 91bb9265..d84efdaa 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -359,8 +359,8 @@ save_modeinfo_maybe(output, t) = nothing function run(Eω, grid, linop, transform, FT, output; min_dz=0, max_dz=grid.zmax/2, init_dz=1e-4, z0=0.0, - rtol=1e-3, atol=1e-6, safety=0.9, norm=RK45.weaknorm, - status_period=1) + rtol=nothing, atol=nothing, safety=0.9, norm=RK45.weaknorm, + status_period=1, solver=:Tsit5) Et = FT \ Eω @@ -384,13 +384,26 @@ function run(Eω, grid, save_modeinfo_maybe(output, transform) flush(stderr) # flush std error once before starting to show setup steps - #RK45.solve_precon( - # transform, linop, Eω, z0, init_dz, grid.zmax, stepfun=stepfun, - # max_dt=max_dz, min_dt=min_dz, - # rtol=rtol, atol=atol, safety=safety, norm=norm, - # status_period=status_period) - Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period) + + if solver == :OrigRK45 + Logging.@info("Using original Luna RK45 solver.") + rtol = isnothing(rtol) ? 1e-3 : rtol + atol = isnothing(atol) ? 1e-6 : atol + Logging.@info("Using rtol = $rtol, atol = $atol") + RK45.solve_precon( + transform, linop, Eω, z0, init_dz, grid.zmax, stepfun=stepfun, + max_dt=max_dz, min_dt=min_dz, + rtol=rtol, atol=atol, safety=safety, norm=norm, + status_period=status_period) + return + else + Logging.@info("Using $solver solver") + rtol = isnothing(rtol) ? 1e-3 : rtol + atol = isnothing(atol) ? 1e-6 : atol + Logging.@info("Using rtol = $rtol, atol = $atol") + Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; + rtol, atol, init_dz, max_dz, min_dz, status_period) + end end # run some code for precompilation From 5b58db41ad4cc8e57c3452952040b03942d94cb6 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 22:28:38 +0100 Subject: [PATCH 09/28] Match orig precision --- src/Luna.jl | 8 ++++---- src/Propagator.jl | 13 ++++++------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index d84efdaa..51e845d3 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -387,8 +387,8 @@ function run(Eω, grid, if solver == :OrigRK45 Logging.@info("Using original Luna RK45 solver.") - rtol = isnothing(rtol) ? 1e-3 : rtol - atol = isnothing(atol) ? 1e-6 : atol + rtol = isnothing(rtol) ? 1e-6 : rtol + atol = isnothing(atol) ? 1e-10 : atol Logging.@info("Using rtol = $rtol, atol = $atol") RK45.solve_precon( transform, linop, Eω, z0, init_dz, grid.zmax, stepfun=stepfun, @@ -398,8 +398,8 @@ function run(Eω, grid, return else Logging.@info("Using $solver solver") - rtol = isnothing(rtol) ? 1e-3 : rtol - atol = isnothing(atol) ? 1e-6 : atol + rtol = isnothing(rtol) ? 1e-2 : rtol + atol = isnothing(atol) ? 1e-5 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; rtol, atol, init_dz, max_dz, min_dz, status_period) diff --git a/src/Propagator.jl b/src/Propagator.jl index d5a13990..ab334d03 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -9,12 +9,11 @@ mutable struct Printer{DT} status_period::Int start::DT tic::DT - steps::Int zmax::Float64 end function Printer(status_period, zmax) - Printer(status_period, Dates.now(), Dates.now(), 0, zmax) + Printer(status_period, Dates.now(), Dates.now(), zmax) end function printstart(p::Printer) @@ -24,7 +23,6 @@ function printstart(p::Printer) end function printstep(p::Printer, z, dz) - p.steps += 1 if Dates.value(Dates.now() - p.tic) > 1000*p.status_period speed = z/(Dates.value(Dates.now() - p.start)/1000) eta_in_s = (p.zmax - z)/speed @@ -42,11 +40,12 @@ function printstep(p::Printer, z, dz) end end -function printstop(p::Printer) +function printstop(p::Printer, integrator) totaltime = Dates.now() - p.start dtstring = format_elapsed(totaltime) - Logging.@info @sprintf("Propagation finished in %s, %d steps", - dtstring, p.steps) + Logging.@info @sprintf("Propagation finished in %s", dtstring) + Logging.@info @sprintf("Steps accepted: %d; rejected: %d", integrator.stats.naccept, integrator.stats.nreject) + Logging.@info @sprintf("Nonlinear function calls: %d", integrator.stats.nf) end abstract type AbstractPropagator end @@ -140,7 +139,7 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) printstart(printer) ODE.solve!(integrator) - printstop(printer) + printstop(printer, integrator) end end # module From 4d0963e73838a2e5c59d2bcfc9f2bcc53e666c21 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 22:39:12 +0100 Subject: [PATCH 10/28] Actually allow solver choice --- src/Luna.jl | 2 +- src/Propagator.jl | 17 +++++++++++------ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 51e845d3..9cf0761f 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -402,7 +402,7 @@ function run(Eω, grid, atol = isnothing(atol) ? 1e-5 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period) + rtol, atol, init_dz, max_dz, min_dz, status_period, solver) end end diff --git a/src/Propagator.jl b/src/Propagator.jl index ab334d03..b2ec0f3e 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -44,7 +44,8 @@ function printstop(p::Printer, integrator) totaltime = Dates.now() - p.start dtstring = format_elapsed(totaltime) Logging.@info @sprintf("Propagation finished in %s", dtstring) - Logging.@info @sprintf("Steps accepted: %d; rejected: %d", integrator.stats.naccept, integrator.stats.nreject) + Logging.@info @sprintf("Steps accepted: %d; rejected: %d", + integrator.stats.naccept, integrator.stats.nreject) Logging.@info @sprintf("Nonlinear function calls: %d", integrator.stats.nf) end @@ -73,7 +74,8 @@ function callbackcl(integrator) @. u * exp(integrator.p.L * z) end end - integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, + ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) ODE.u_modified!(integrator, false) end @@ -112,7 +114,8 @@ function callbackncl(integrator) @. Eω * exp(L) end end - integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, + ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) ODE.u_modified!(integrator, false) end @@ -124,18 +127,20 @@ function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printe end function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) - prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) + prop = NonConstPropagator(linop, f!, stepfun, + length(Eω0), similar(Eω0), similar(Eω0), printer) u0 = vcat(Eω0, zero(Eω0)) prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) prob, callbackncl end function propagate(f!, linop, Eω0, z, zmax, stepfun; - rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, status_period=1) + rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, + status_period=1, solver=:Tsit5) printer = Printer(status_period, zmax) prob, cbfunc = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) - integrator = ODE.init(prob, ODE.Tsit5(); adaptive=true, reltol=rtol, abstol=atol, + integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) printstart(printer) ODE.solve!(integrator) From 6c645fc9a27814efe83971bded3950f86269cf18 Mon Sep 17 00:00:00 2001 From: John Travers Date: Mon, 1 Sep 2025 22:45:46 +0100 Subject: [PATCH 11/28] Add some code comments --- src/Propagator.jl | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/src/Propagator.jl b/src/Propagator.jl index b2ec0f3e..34b51700 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -51,6 +51,7 @@ end abstract type AbstractPropagator end +# For a constant linear operator, we can integrate L(z) analytically struct ConstPropagator{NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator L::AT nonlinop!::NLT @@ -61,12 +62,13 @@ struct ConstPropagator{NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator end function fcl!(du,u,p,z) - @. p.Eωtmp = u * exp(p.L * z) - p.nonlinop!(p.Pωtmp, p.Eωtmp, z) - @. du = p.Pωtmp * exp(-p.L * z) + @. p.Eωtmp = u * exp(p.L * z) # Transform back from interaction picture + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator + @. du = p.Pωtmp * exp(-p.L * z) # Transform to interaction picture end function callbackcl(integrator) + # The output we want must be transformed back from the interaction picture @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.L * integrator.t) interp = let integrator=integrator function interp(z) @@ -77,9 +79,11 @@ function callbackcl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - ODE.u_modified!(integrator, false) + ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end +# For a non-constant linear operator, we need to integrate L(z) numerically along with +# the solution. We do this by simply including the linear operator in the state vector. struct NonConstPropagator{LT, NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator linop!::LT nonlinop!::NLT @@ -95,16 +99,17 @@ function fncl!(du,u,p,z) L = @views u[p.n+1:end] dEω = @views du[1:p.n] dL = @views du[p.n+1:end] - @. p.Eωtmp = Eω * exp(L) - p.nonlinop!(p.Pωtmp, p.Eωtmp, z) - @. dEω = p.Pωtmp * exp(-L) - p.linop!(dL, z) + @. p.Eωtmp = Eω * exp(L) # Transform back from interaction picture + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator + @. dEω = p.Pωtmp * exp(-L) # Transform to interaction picture + p.linop!(dL, z) # Integrate linear operator end function callbackncl(integrator) n = integrator.p.n Eω = @views integrator.u[1:n] L = @views integrator.u[n+1:end] + # The output we want must be transformed back from the interaction picture @. integrator.p.Eωtmp = Eω * exp(L) interp = let integrator=integrator, n=n function interp(z) @@ -117,7 +122,7 @@ function callbackncl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - ODE.u_modified!(integrator, false) + ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer) where N @@ -129,7 +134,7 @@ end function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) - u0 = vcat(Eω0, zero(Eω0)) + u0 = vcat(Eω0, zero(Eω0)) # Initial linear operator is zero prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) prob, callbackncl end @@ -139,6 +144,7 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; status_period=1, solver=:Tsit5) printer = Printer(status_period, zmax) prob, cbfunc = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) + # We do all saving and stats in a callback called at every step cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) From 6d2c099364fd9b9bd2cdd973b8e5b7ede5a5035b Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 11:29:29 +0100 Subject: [PATCH 12/28] test componentwise tols --- src/Propagator.jl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/Propagator.jl b/src/Propagator.jl index 34b51700..09ff427d 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -125,25 +125,27 @@ function callbackncl(integrator) ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end -function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer) where N +function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer, rtol, atol) where N prop = ConstPropagator(linop, f!, stepfun, similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) - prob, callbackcl + prob, callbackcl, rtol, atol end -function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) +function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) u0 = vcat(Eω0, zero(Eω0)) # Initial linear operator is zero + #rtol = vcat(ones(length(Eω0))*rtol, ones(length(Eω0))*1e-10) + #atol = vcat(ones(length(Eω0))*atol, ones(length(Eω0))*1e-15) prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) - prob, callbackncl + prob, callbackncl, rtol, atol end function propagate(f!, linop, Eω0, z, zmax, stepfun; rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, status_period=1, solver=:Tsit5) printer = Printer(status_period, zmax) - prob, cbfunc = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer) + prob, cbfunc, rtol, atol = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) # We do all saving and stats in a callback called at every step cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, From 5aad2eb922b0470642d293f4455c3025160d7ac0 Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 11:41:18 +0100 Subject: [PATCH 13/28] Look at fixing tstops --- src/Propagator.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Propagator.jl b/src/Propagator.jl index 09ff427d..5e1f8464 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -143,13 +143,13 @@ end function propagate(f!, linop, Eω0, z, zmax, stepfun; rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, - status_period=1, solver=:Tsit5) + status_period=1, solver=:Tsit5, zstops=nothing) printer = Printer(status_period, zmax) prob, cbfunc, rtol, atol = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) # We do all saving and stats in a callback called at every step cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, - dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) + dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb, tstops=zstops) printstart(printer) ODE.solve!(integrator) printstop(printer, integrator) From 2f5da0df9bd4c93257461027a2cffff26c0872a0 Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 15:40:27 +0100 Subject: [PATCH 14/28] relax gradient const comparison tests --- test/test_gradient.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/test_gradient.jl b/test/test_gradient.jl index 54502d49..f475ddf8 100644 --- a/test/test_gradient.jl +++ b/test/test_gradient.jl @@ -75,8 +75,8 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) -@test all(output_grad.data["Eω"][grid.sidx, :] .≈ output_const.data["Eω"][grid.sidx, :]) -@test all(output_grad_array.data["Eω"][grid.sidx, :] .≈ output_const.data["Eω"][grid.sidx, :]) +@test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) +@test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end @testset "envelope" begin @@ -134,6 +134,6 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) -@test all(output_grad.data["Eω"][grid.sidx, :] .≈ output_const.data["Eω"][grid.sidx, :]) -@test all(output_grad_array.data["Eω"][grid.sidx, :] .≈ output_const.data["Eω"][grid.sidx, :]) +@test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) +@test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end From a024354f3ae74ba38de4fa909887c5f5aa27544d Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 15:42:34 +0100 Subject: [PATCH 15/28] comment on test relaxation --- test/test_gradient.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/test_gradient.jl b/test/test_gradient.jl index f475ddf8..5c8553fd 100644 --- a/test/test_gradient.jl +++ b/test/test_gradient.jl @@ -75,6 +75,8 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) +# TODO: tolerances here are quite high, not because of inherent errors, but due to different ODE solvers +# it would be good to investigate this further @test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) @test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end @@ -134,6 +136,8 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) +# TODO: tolerances here are quite high, not because of inherent errors, but due to different ODE solvers +# it would be good to investigate this further @test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) @test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end From 7e724b4fb2b50b9ed4236cb74df3a09bcc7a3006 Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 17:10:51 +0100 Subject: [PATCH 16/28] Add analytical gradient case test --- src/Luna.jl | 2 +- src/Propagator.jl | 73 ++++++++++++++++++++++++++++--------------- test/test_gradient.jl | 67 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 111 insertions(+), 31 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 9cf0761f..70792618 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -402,7 +402,7 @@ function run(Eω, grid, atol = isnothing(atol) ? 1e-5 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period, solver) + rtol, atol, init_dz, max_dz, min_dz, status_period, solver)#, zstops=output.save_cond.grid) end end diff --git a/src/Propagator.jl b/src/Propagator.jl index 5e1f8464..43bf9d80 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -51,29 +51,33 @@ end abstract type AbstractPropagator end -# For a constant linear operator, we can integrate L(z) analytically -struct ConstPropagator{NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator - L::AT +# For the cases where we can integrate L(z) analytically +struct AnalyticalPropagator{LT, NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator + Li!::LT # function to get integrated linear operator at z nonlinop!::NLT stepfun::SFT + Litmp::AT Eωtmp::AT Pωtmp::AT p::PT end function fcl!(du,u,p,z) - @. p.Eωtmp = u * exp(p.L * z) # Transform back from interaction picture - p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator - @. du = p.Pωtmp * exp(-p.L * z) # Transform to interaction picture + p.Li!(p.Litmp, z) # Get integrated linear operator at z + @. p.Eωtmp = u * exp(p.Litmp) # Transform back from interaction picture + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator + @. du = p.Pωtmp * exp(-p.Litmp) # Transform to interaction picture end function callbackcl(integrator) # The output we want must be transformed back from the interaction picture - @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.L * integrator.t) + integrator.p.Li!(integrator.p.Litmp, integrator.t) + @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.Litmp) interp = let integrator=integrator function interp(z) - u = integrator(z) - @. u * exp(integrator.p.L * z) + u = integrator.sol(z) + integrator.p.Li!(integrator.p.Litmp, z) + @. u * exp(integrator.p.Litmp) end end integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, @@ -83,9 +87,9 @@ function callbackcl(integrator) end # For a non-constant linear operator, we need to integrate L(z) numerically along with -# the solution. We do this by simply including the linear operator in the state vector. +# the solution. We do this by simply including the integral of linear operator in the state vector. struct NonConstPropagator{LT, NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropagator - linop!::LT + linop!::LT # function to get linear operator at z nonlinop!::NLT stepfun::SFT n::Int @@ -95,28 +99,28 @@ struct NonConstPropagator{LT, NLT, SFT, PT, AT<:AbstractArray} <: AbstractPropag end function fncl!(du,u,p,z) - Eω = @views u[1:p.n] - L = @views u[p.n+1:end] + Eω = @views u[1:p.n] # Actual Eω + Li = @views u[p.n+1:end] # Cumulatively integrated linear operator dEω = @views du[1:p.n] - dL = @views du[p.n+1:end] - @. p.Eωtmp = Eω * exp(L) # Transform back from interaction picture - p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator - @. dEω = p.Pωtmp * exp(-L) # Transform to interaction picture - p.linop!(dL, z) # Integrate linear operator + dLi = @views du[p.n+1:end] + @. p.Eωtmp = Eω * exp(Li) # Transform back from interaction picture + p.nonlinop!(p.Pωtmp, p.Eωtmp, z) # Apply nonlinear operator + @. dEω = p.Pωtmp * exp(-Li) # Transform to interaction picture + p.linop!(dLi, z) # Integrate linear operator end function callbackncl(integrator) n = integrator.p.n - Eω = @views integrator.u[1:n] - L = @views integrator.u[n+1:end] + Eω = @views integrator.u[1:n] # Actual Eω + Li = @views integrator.u[n+1:end] # Cumulatively integrated linear operator # The output we want must be transformed back from the interaction picture - @. integrator.p.Eωtmp = Eω * exp(L) + @. integrator.p.Eωtmp = Eω * exp(Li) interp = let integrator=integrator, n=n function interp(z) - u = integrator(z) + u = integrator.sol(z) Eω = @views u[1:n] - L = @views u[n+1:end] - @. Eω * exp(L) + Li = @views u[n+1:end] + @. Eω * exp(Li) end end integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, @@ -125,18 +129,35 @@ function callbackncl(integrator) ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end +# Constant linear operator case--linop is an array function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer, rtol, atol) where N - prop = ConstPropagator(linop, f!, stepfun, similar(Eω0), similar(Eω0), printer) + # For a constant linear operator L, the integral is just L*z + Li! = let linop=linop + function Li!(out, z) + @. out = linop * z + end + end + prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) + prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) + prob, callbackcl, rtol, atol +end + +# For a linop tuple we expect a pair of functions (linop, ilinop) where the second function provides the +# cumulatively integrated linear operator. This is mostly for testing. +function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) where N + Li! = linop[2] + prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) prob, callbackcl, rtol, atol end +# General linop, we integrate numerically along with the solution function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) u0 = vcat(Eω0, zero(Eω0)) # Initial linear operator is zero #rtol = vcat(ones(length(Eω0))*rtol, ones(length(Eω0))*1e-10) - #atol = vcat(ones(length(Eω0))*atol, ones(length(Eω0))*1e-15) + #atol = vcat(ones(length(Eω0))*atol, ones(length(Eω0))*1e-10) prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) prob, callbackncl, rtol, atol end diff --git a/test/test_gradient.jl b/test/test_gradient.jl index 5c8553fd..299c2413 100644 --- a/test/test_gradient.jl +++ b/test/test_gradient.jl @@ -75,8 +75,11 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) -# TODO: tolerances here are quite high, not because of inherent errors, but due to different ODE solvers -# it would be good to investigate this further +# TODO: tolerances here are quite high, because the const propagator and non-const propagator handle the +# integration of the linear part differently. The absolute error is small (e.g. here 6 significant digits) +# and is easily reduced by tightening tolerances. This does not arise with the piecewise integrator originally +# used, because the const and piecewise operators are identical for a constant gradient. But see the +# analytical gradient test below for an example of where the piecewise approximation breaks down @test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) @test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end @@ -136,8 +139,64 @@ statsfun = Stats.collect_stats(grid, Eω, output_grad_array = Output.MemoryOutput(0, grid.zmax, 201, statsfun) Luna.run(Eω, grid, linop, transform, FT, output_grad_array, status_period=10) -# TODO: tolerances here are quite high, not because of inherent errors, but due to different ODE solvers -# it would be good to investigate this further +# See comment on line 78 @test isapprox(output_grad.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) @test isapprox(output_grad_array.data["Eω"][grid.sidx, :], output_const.data["Eω"][grid.sidx, :], rtol=3e-6) end + +@testset "analytical gradient" begin +a = 13e-6 +gas = :Ar +pres = 5 +τ = 30e-15 +λ0 = 800e-9 +L = 5e-2 + +# Common setup +grid = Grid.RealGrid(L, λ0, (160e-9, 3000e-9), 0.5e-12) +inputs = Fields.GaussField(λ0=λ0, τfwhm=τ, energy=1e-6) +responses = (Nonlinear.Kerr_field(PhysData.γ3_gas(gas)),) + +# Constant +dens0 = PhysData.density(gas, pres) +dens(z) = dens0 +m = Capillary.MarcatiliMode(a, gas, pres, loss=false) +aeff(z) = Modes.Aeff(m, z=z) +energyfun, energyfunω = Fields.energyfuncs(grid) +linop, βfun!, frame_vel, αfun = LinearOps.make_const_linop(grid, m, λ0) + +# Test a linearly increasing linop for analytical verification of propagator +linopfun = let linop=linop, L=L + function linopfun(out, z) + @. out = linop * z / L # linear increase from 0 to linop over distance L + end +end +linopifun = let linop=linop, L=L + function linopifun(out, z) + @. out = linop * z^2 / (2 * L) # integral of linopfun + end +end + +Eω, transform, FT = Luna.setup( + grid, dens, responses, inputs, βfun!, aeff) +statsfun = Stats.collect_stats(grid, Eω, + Stats.ω0(grid), + Stats.energy(grid, energyfunω)) + + +# Test analytical propagator with linopfun and linopifun +output_ana = Output.MemoryOutput(0, grid.zmax, 201, statsfun) +Luna.run(Eω, grid, (linopfun, linopifun), transform, FT, output_ana, status_period=10) + +# Test numerically integrated propagator with linopfun +output_num = Output.MemoryOutput(0, grid.zmax, 201, statsfun) +Luna.run(Eω, grid, linopfun, transform, FT, output_num, status_period=10) + +# Test piecewise version +output_piece = Output.MemoryOutput(0, grid.zmax, 201, statsfun) +Luna.run(Eω, grid, linopfun, transform, FT, output_piece, status_period=10, solver=:OrigRK45) + +@test isapprox(output_num.data["Eω"][grid.sidx, :], output_ana.data["Eω"][grid.sidx, :], rtol=2e-3) +# Piecewise solver struggles with the high frequency part: +@test isapprox(output_piece.data["Eω"][grid.sidx, :], output_ana.data["Eω"][grid.sidx, :], rtol=0.4) +end From 357b8d888c05b2864761036c4f7d3a6cffa055ad Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 21:01:15 +0100 Subject: [PATCH 17/28] Try to fix absorbing boundaries --- src/Propagator.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/Propagator.jl b/src/Propagator.jl index 43bf9d80..2175b882 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -83,7 +83,8 @@ function callbackcl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + @. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) # copy back as we modify u in stepfun (absorbing boundaries) + #ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end # For a non-constant linear operator, we need to integrate L(z) numerically along with @@ -126,7 +127,8 @@ function callbackncl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + @. Eω = integrator.p.Eωtmp * exp(-Li) # copy back as we modify u in stepfun (absorbing boundaries) + # ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end # Constant linear operator case--linop is an array From 6e2f269d5018c86cd0220b5d745d38fae7954057 Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 23:53:45 +0100 Subject: [PATCH 18/28] Add work-precision script --- scripts/solver_work_precision.jl | 260 +++++ scripts/solver_work_precision.svg | 1548 +++++++++++++++++++++++++++++ src/Propagator.jl | 11 +- 3 files changed, 1814 insertions(+), 5 deletions(-) create mode 100644 scripts/solver_work_precision.jl create mode 100644 scripts/solver_work_precision.svg diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl new file mode 100644 index 00000000..4f6b7ca0 --- /dev/null +++ b/scripts/solver_work_precision.jl @@ -0,0 +1,260 @@ +# Calculate work-precision plots for various NLSE solvers + +using DifferentialEquations, SciMLOperators +import FFTW +import LinearAlgebra: inv, mul!, ldiv!, norm, Diagonal +using PyPlot +import Luna + +# NLSE grid and temporary storage +mutable struct NLSE{TFT} + n::Int + dt::Float64 + dΩ::Float64 + T::Vector{Float64} + Ω::Vector{Float64} + ut::Vector{ComplexF64} + utmp::Vector{ComplexF64} + utmp2::Vector{ComplexF64} + dutmp::Vector{ComplexF64} + L::Vector{ComplexF64} + FT::TFT + cz::Float64 + nfunc::Int +end + +function NLSE(dt, trange) + n = nextpow(2, ceil(Int, trange/dt)) + dt = trange/n + dΩ = 2π/trange + T = collect((-n//2:n//2-1)*dt) + Ω = FFTW.fftshift((-n//2:n//2-1)*dΩ) + ut = @. complex(5*sech(T)) + utmp = similar(ut) + utmp2 = similar(ut) + dutmp = similar(ut) + L = @. 1im*Ω^2/2 + FT = FFTW.plan_fft(ut) + inv(FT) + cz = 0.0 + NLSE(n, dt, dΩ, T, Ω, ut, utmp, utmp2, dutmp, L, FT, cz, 0) +end + +function reset!(nlse::NLSE) + nlse.cz = 0.0 + nlse.nfunc = 0 +end + +# explicit linear operator +function f1!(du,u,p,z) + @. du = p.L*u +end + +# nonlinear operator (Kerr effect) +function f2!(du,u,p,z) + p.nfunc += 1 + ldiv!(p.ut, p.FT, u) + @. p.utmp = -1im*abs2(p.ut)*p.ut + mul!(du, p.FT, p.utmp) +end + +# full interaction picture, constant L, analytically integrated +function fpre!(du,u,p,z) + @. p.utmp = u*exp(p.L*z) + f2!(du,p.utmp,p,z) + @. du *= exp(-p.L*z) +end + +# piecewise interaction picture, constant L, analytically integrated +function fpre2!(du,u,p,z) + @. p.utmp = u*exp(p.L*(z - p.cz)) + f2!(du,p.utmp,p,z) + @. du *= exp(-p.L*(z - p.cz)) +end + +# full interaction picture, L numerically integrated +function fdbl!(du,u,p,z) + uu = @view u[1:length(u)÷2] + ll = @view u[length(u)÷2+1:end] + duu = @view du[1:length(u)÷2] + dll = @view du[length(u)÷2+1:end] + @. p.utmp = uu*exp(ll) + f2!(p.utmp2,p.utmp,p,z) + @. duu = p.utmp2*exp(-ll) + @. dll = p.L +end + +# explicitly call both linear and nonlinear terms, this is stiff +function fall!(du,u,p,z) + f1!(p.dutmp,u,p,z) + f2!(du,u,p,z) + du .+= p.dutmp +end + +# 5th order soliton initial condition +function getinit(nlse::NLSE) + ut0 = @. complex.(5*sech(nlse.T)) + FFTW.fft(ut0) +end + +# reset u and cz at each step for piecewise interaction picture solver +function resetaffect!(integrator) + integrator.u .= integrator.u .* exp.(integrator.p.L .* (integrator.t - integrator.p.cz)) + integrator.p.cz = integrator.t +end + +function geterror(nlse, u) + ana = @. 5*sech(nlse.T)*exp(-1im*π/4) + norm(FFTW.ifft(u) .- ana)/norm(ana) +end + +function run(prob, solver, adaptive, dt, reltol, abstol; cb=nothing) + zs = range(0.0, π/2, length=201) + println("building") + @time integrator = init(prob, solver; dt, adaptive, reltol, abstol, saveat=zs, callback=cb) + println("starting") + @time u = solve!(integrator) + zs, u +end + +# run full interaction picture, constant L, analytically integrated +function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + prob = ODEProblem(fpre!, getinit(nlse), (0.0, π/2), nlse) + zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) + for (i,z) in enumerate(zs) + @. res[:,i] = u[:,i] * exp(nlse.L * z) + end + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +# run full interaction picture, L numerically integrated +function run_numfullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + u0 = vcat(getinit(nlse), zero(nlse.L)) + prob = ODEProblem(fdbl!, u0, (0.0, π/2), nlse) + zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) + for i in 1:length(zs) + @. res[:,i] = u[1:nlse.n,i] * exp(u[nlse.n + 1:end,i]) + end + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +# piecewise interaction picture, constant L, analytically integrated +function run_pieceip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + prob = ODEProblem(fpre2!, getinit(nlse), (0.0, π/2), nlse) + cb = DiscreteCallback((u,t,integrator) -> true, resetaffect!, save_positions=(false,true)) + _, u = run(prob, solver, adaptive, dt, reltol, abstol; cb) + res = Array(u) + zs = u.t + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +function run_stiff(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + prob = ODEProblem(fall!, getinit(nlse), (0.0, π/2), nlse) + zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + res = Array(u) + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +# Exponential RK integrator +function run_splitlin(nlse::NLSE; solver=ETDRK4(), adaptive=false, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + op = DiagonalOperator(nlse.L) + f = SplitFunction(op, f2!) + prob = SplitODEProblem(f, getinit(nlse), (0.0, π/2), nlse) + zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + res = Array(u) + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +# Luna original RK45 solver +function run_Luna(nlse::NLSE; solver=nothing, adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + z, u, steps = Luna.RK45.solve_precon((du, u, z) -> f2!(du, u, nlse, z), nlse.L, getinit(nlse), 0.0, dt, π/2; + rtol=reltol, atol=abstol, output=true, locextrap=true)#, norm=Luna.RK45.normnorm) + err = geterror(nlse, u[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + z, u, nlse.nfunc, err +end + +# Luna new solver +function run_newLuna(nlse::NLSE; solver=:Tsit5, adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + zs = range(0.0, π/2, length=201) + iz = 2 + res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) + res[:,1] = getinit(nlse) + function stepfun(u, z, dz, interpolant) + while iz <= length(zs) && z >= zs[iz] + res[:,iz] = interpolant(zs[iz]) + iz += 1 + end + end + sol = Luna.Propagator.propagate((du, u, z) -> f2!(du, u, nlse, z), nlse.L, res[:,1], 0, π/2, stepfun; + rtol=reltol, atol=abstol, init_dz=dt, max_dz=Inf, min_dz=0, + status_period=10, solver) + err = geterror(nlse, res[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + zs, res, nlse.nfunc, err +end + +function workprecision(nlse::NLSE, solvers) + errs = [] + nfs = [] + for (i,solverset) in enumerate(solvers) + solver, dts, reltols, abstols = solverset + errsi = zeros(length(dts)) + nfsi = zeros(length(dts)) + for j in 1:length(dts) + z, u, nfuncs, err = solver(nlse; reltol=reltols[j], abstol=abstols[j], dt=dts[j]) + errsi[j] = err + nfsi[j] = nfuncs + end + loglog(errsi, nfsi, label=string(solver)) + push!(errs, errsi) + push!(nfs, nfsi) + end + legend() + PyPlot.grid() + xlabel("Error") + ylabel("Function Evaluations") + ylim(2e3,4e4) + xlim(1e-6,1e-1) + errs, nfs +end + + +nlse = NLSE(0.016, 48.0); + +errs, nfs = workprecision(nlse, ( + (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), + (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), + (run_numfullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), + (run_splitlin, collect(logrange(1e-4, 1e-2, 30)), 1e-6 .* ones(30), 1e-6 .* ones(30)), + (run_Luna, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30)), + (run_newLuna, 0.0002 .* ones(40), collect(logrange(5e-5, 1.2e-1, 40)), 1e-6 .* ones(40)) +)) + +savefig(joinpath(pkgdir(Luna), "scripts/solver_work_precision.svg")) diff --git a/scripts/solver_work_precision.svg b/scripts/solver_work_precision.svg new file mode 100644 index 00000000..89503d60 --- /dev/null +++ b/scripts/solver_work_precision.svg @@ -0,0 +1,1548 @@ + + + + + + + + 2025-09-02T23:52:06.914015 + image/svg+xml + + + Matplotlib v3.10.1, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/Propagator.jl b/src/Propagator.jl index 2175b882..78d12382 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -83,8 +83,8 @@ function callbackcl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - @. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) # copy back as we modify u in stepfun (absorbing boundaries) - #ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + #@. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) # copy back as we modify u in stepfun (absorbing boundaries) + ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end # For a non-constant linear operator, we need to integrate L(z) numerically along with @@ -127,8 +127,8 @@ function callbackncl(integrator) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - @. Eω = integrator.p.Eωtmp * exp(-Li) # copy back as we modify u in stepfun (absorbing boundaries) - # ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + #@. Eω = integrator.p.Eωtmp * exp(-Li) # copy back as we modify u in stepfun (absorbing boundaries) + ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal end # Constant linear operator case--linop is an array @@ -174,8 +174,9 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb, tstops=zstops) printstart(printer) - ODE.solve!(integrator) + sol = ODE.solve!(integrator) printstop(printer, integrator) + sol end end # module From 155a91cd7efb22e0546da6cea8c74ee23cf479b3 Mon Sep 17 00:00:00 2001 From: John Travers Date: Tue, 2 Sep 2025 23:56:58 +0100 Subject: [PATCH 19/28] add plotting to work-precision --- scripts/solver_work_precision.jl | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index 4f6b7ca0..40f5d1ef 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -245,6 +245,23 @@ function workprecision(nlse::NLSE, solvers) errs, nfs end +function plot_nlse(nlse::NLSE, z, u) + IT = 10log10.(abs2.(FFTW.ifft(u,1))) + IW = 10log10.(abs2.(FFTW.fftshift(u,1))) + IT .-= maximum(IT) + IW .-= maximum(IW) + figure() + subplot(121) + pcolormesh(z, nlse.T, IT, clim=(-160,0)) + xlabel("Position") + ylabel("Time") + subplot(122) + pcolormesh(z, FFTW.fftshift(nlse.Ω), IW, clim=(-160,0)) + colorbar() + xlabel("Position") + ylabel("Frequency") + tight_layout() +end nlse = NLSE(0.016, 48.0); From 94ed472452b29843a3d6ea298f1495d5207ea95d Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 3 Sep 2025 10:45:19 +0100 Subject: [PATCH 20/28] add work-precision script --- .gitignore | 1 + scripts/solver_work_precision.jl | 65 +- scripts/solver_work_precision.svg | 336 ++- scripts/solver_work_precision_cmp.svg | 3575 +++++++++++++++++++++++++ src/Propagator.jl | 4 +- 5 files changed, 3832 insertions(+), 149 deletions(-) create mode 100644 scripts/solver_work_precision_cmp.svg diff --git a/.gitignore b/.gitignore index 4c6a7142..8e456c27 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ docs/build *.info .vscode/spellright.dict deps/build.log +.DS_Store diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index 40f5d1ef..fec7fec2 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -5,6 +5,7 @@ import FFTW import LinearAlgebra: inv, mul!, ldiv!, norm, Diagonal using PyPlot import Luna +import Printf: @sprintf # NLSE grid and temporary storage mutable struct NLSE{TFT} @@ -188,10 +189,21 @@ function run_splitlin(nlse::NLSE; solver=ETDRK4(), adaptive=false, dt=0.0002, re end # Luna original RK45 solver -function run_Luna(nlse::NLSE; solver=nothing, adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) +function run_Luna_weak(nlse::NLSE; solver=nothing, adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) reset!(nlse) z, u, steps = Luna.RK45.solve_precon((du, u, z) -> f2!(du, u, nlse, z), nlse.L, getinit(nlse), 0.0, dt, π/2; - rtol=reltol, atol=abstol, output=true, locextrap=true)#, norm=Luna.RK45.normnorm) + rtol=reltol, atol=abstol, output=true, locextrap=true, norm=Luna.RK45.weaknorm) + err = geterror(nlse, u[:,end]) + println("nfunc: $(nlse.nfunc)") + println("error: $err") + z, u, nlse.nfunc, err +end + +# Luna original RK45 solver with better norm +function run_Luna_norm(nlse::NLSE; solver=nothing, adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) + reset!(nlse) + z, u, steps = Luna.RK45.solve_precon((du, u, z) -> f2!(du, u, nlse, z), nlse.L, getinit(nlse), 0.0, dt, π/2; + rtol=reltol, atol=abstol, output=true, locextrap=true, norm=Luna.RK45.normnorm) err = geterror(nlse, u[:,end]) println("nfunc: $(nlse.nfunc)") println("error: $err") @@ -245,33 +257,54 @@ function workprecision(nlse::NLSE, solvers) errs, nfs end -function plot_nlse(nlse::NLSE, z, u) +function plot_nlse(nlse::NLSE, z, u; axs=nothing) IT = 10log10.(abs2.(FFTW.ifft(u,1))) IW = 10log10.(abs2.(FFTW.fftshift(u,1))) IT .-= maximum(IT) IW .-= maximum(IW) - figure() - subplot(121) - pcolormesh(z, nlse.T, IT, clim=(-160,0)) - xlabel("Position") - ylabel("Time") - subplot(122) - pcolormesh(z, FFTW.fftshift(nlse.Ω), IW, clim=(-160,0)) - colorbar() - xlabel("Position") - ylabel("Frequency") - tight_layout() + if isnothing(axs) + fig = PyPlot.plt.figure(constrained_layout=true, figsize=(10, 6)) + axd = fig.subplot_mosaic( + """ + ab + """) + axs = (axd["a"], axd["b"]) + end + axs[1].pcolormesh(z, nlse.T, IT, clim=(-200,0), rasterized=true) + axs[1].set_xlabel("Position") + axs[1].set_ylabel("Time") + img = axs[2].pcolormesh(z, FFTW.fftshift(nlse.Ω), IW, clim=(-200,0), rasterized=true) + axs[2].set_xlabel("Position") + axs[2].set_ylabel("Frequency") + colorbar(img, ax=axs, fraction=0.05, pad=0.1, label="dB") +end + +function plot_nlse_cmp(nlse::NLSE, data) + fig = PyPlot.plt.figure(constrained_layout=true, figsize=(10, 3*length(data))) + ax_array = fig.subplots(length(data), 2) + for (i, (z, u, nfs, err)) in enumerate(data) + axs = (ax_array[i,1], ax_array[i,2]) + plot_nlse(nlse, z, u; axs=axs) + errs = @sprintf("%.2e", err) + axs[1].set_title("nfs=$(nfs), err=$(errs)") + end end nlse = NLSE(0.016, 48.0); +# run at plot work-precision errs, nfs = workprecision(nlse, ( (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), (run_numfullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), (run_splitlin, collect(logrange(1e-4, 1e-2, 30)), 1e-6 .* ones(30), 1e-6 .* ones(30)), - (run_Luna, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30)), + (run_Luna_weak, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30)), + (run_Luna_norm, 0.0002 .* ones(30), collect(logrange(1e-7, 1e-1, 30)), 1e-6 .* ones(30)), (run_newLuna, 0.0002 .* ones(40), collect(logrange(5e-5, 1.2e-1, 40)), 1e-6 .* ones(40)) )) - savefig(joinpath(pkgdir(Luna), "scripts/solver_work_precision.svg")) + +# run a comparison to visualise the error +data = [run_newLuna(nlse; reltol=rtol, abstol=1e-6) for rtol in (1e-1, 6.9e-2, 1.2e-2, 5e-4)] +plot_nlse_cmp(nlse, data) +savefig(joinpath(pkgdir(Luna), "scripts/solver_work_precision_cmp.svg"), dpi=600) diff --git a/scripts/solver_work_precision.svg b/scripts/solver_work_precision.svg index 89503d60..3e1018dd 100644 --- a/scripts/solver_work_precision.svg +++ b/scripts/solver_work_precision.svg @@ -6,7 +6,7 @@ - 2025-09-02T23:52:06.914015 + 2025-09-03T09:48:58.613292 image/svg+xml @@ -42,16 +42,16 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - - + @@ -142,11 +142,11 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -190,11 +190,11 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -232,11 +232,11 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -287,11 +287,11 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -334,11 +334,11 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -354,285 +354,285 @@ L 414.72 41.472 - - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + @@ -707,16 +707,16 @@ z +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - - + @@ -731,82 +731,82 @@ L -3.5 0 - - + - + - + - + - + - + - + - + - + - + - + @@ -1060,7 +1060,7 @@ L 230.260186 192.099599 L 244.720153 202.979358 L 289.512597 220.129617 L 348.824101 234.235934 -" clip-path="url(#pf8be939810)" style="fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square"/> +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square"/> +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #ff7f0e; stroke-width: 1.5; stroke-linecap: square"/> +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #2ca02c; stroke-width: 1.5; stroke-linecap: square"/> +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #d62728; stroke-width: 1.5; stroke-linecap: square"/> +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #9467bd; stroke-width: 1.5; stroke-linecap: square"/> + + + +" clip-path="url(#p37c236bc63)" style="fill: none; stroke: #e377c2; stroke-width: 1.5; stroke-linecap: square"/> - - - + - + - - + - + - - + - + - - + - + @@ -1467,15 +1500,15 @@ L 328.977813 99.439188 - - + - - + + - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -1541,7 +1615,7 @@ z - + diff --git a/scripts/solver_work_precision_cmp.svg b/scripts/solver_work_precision_cmp.svg new file mode 100644 index 00000000..1b973b5b --- /dev/null +++ b/scripts/solver_work_precision_cmp.svg @@ -0,0 +1,3575 @@ + + + + + + + + 2025-09-03T10:43:14.527977 + image/svg+xml + + + Matplotlib v3.10.1, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/Propagator.jl b/src/Propagator.jl index 78d12382..5ce9c00d 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -132,7 +132,7 @@ function callbackncl(integrator) end # Constant linear operator case--linop is an array -function makeprop(f!, linop::Array{ComplexF64,N}, Eω0, z, zmax, stepfun, printer, rtol, atol) where N +function makeprop(f!, linop::Array{ComplexF64}, Eω0, z, zmax, stepfun, printer, rtol, atol) # For a constant linear operator L, the integral is just L*z Li! = let linop=linop function Li!(out, z) @@ -146,7 +146,7 @@ end # For a linop tuple we expect a pair of functions (linop, ilinop) where the second function provides the # cumulatively integrated linear operator. This is mostly for testing. -function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) where N +function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) Li! = linop[2] prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) From 8d4296ea1930f695ecbf0d9b913de2e584fe6b72 Mon Sep 17 00:00:00 2001 From: John Travers Date: Thu, 4 Sep 2025 21:39:30 +0100 Subject: [PATCH 21/28] weird issue hunted down --- scripts/solver_work_precision.jl | 66 +++++++++++++++++++++++--------- src/Propagator.jl | 29 ++++++++++---- 2 files changed, 69 insertions(+), 26 deletions(-) diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index fec7fec2..938d878a 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -104,6 +104,10 @@ function resetaffect!(integrator) integrator.p.cz = integrator.t end +function noaffect!(integrator) + # do nothing +end + function geterror(nlse, u) ana = @. 5*sech(nlse.T)*exp(-1im*π/4) norm(FFTW.ifft(u) .- ana)/norm(ana) @@ -119,7 +123,7 @@ function run(prob, solver, adaptive, dt, reltol, abstol; cb=nothing) end # run full interaction picture, constant L, analytically integrated -function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) +function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6, fullret=false) reset!(nlse) prob = ODEProblem(fpre!, getinit(nlse), (0.0, π/2), nlse) zs, u = run(prob, solver, adaptive, dt, reltol, abstol) @@ -130,6 +134,9 @@ function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol err = geterror(nlse, res[:,end]) println("nfunc: $(nlse.nfunc)") println("error: $err") + if fullret + return zs, res, nlse.nfunc, err, u + end zs, res, nlse.nfunc, err end @@ -138,7 +145,9 @@ function run_numfullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, rel reset!(nlse) u0 = vcat(getinit(nlse), zero(nlse.L)) prob = ODEProblem(fdbl!, u0, (0.0, π/2), nlse) - zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + cb = DiscreteCallback((u,t,integrator) -> true, noaffect!, save_positions=(true,true)) + zs, u = run(prob, solver, adaptive, dt, reltol, abstol; cb=cb) + zs = Array(u.t) res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) for i in 1:length(zs) @. res[:,i] = u[1:nlse.n,i] * exp(u[nlse.n + 1:end,i]) @@ -236,7 +245,7 @@ function workprecision(nlse::NLSE, solvers) errs = [] nfs = [] for (i,solverset) in enumerate(solvers) - solver, dts, reltols, abstols = solverset + solver, dts, reltols, abstols, label = solverset errsi = zeros(length(dts)) nfsi = zeros(length(dts)) for j in 1:length(dts) @@ -244,7 +253,10 @@ function workprecision(nlse::NLSE, solvers) errsi[j] = err nfsi[j] = nfuncs end - loglog(errsi, nfsi, label=string(solver)) + if isnothing(label) + label = string(solver) + end + loglog(errsi, nfsi, label=label) push!(errs, errsi) push!(nfs, nfsi) end @@ -293,18 +305,36 @@ end nlse = NLSE(0.016, 48.0); # run at plot work-precision -errs, nfs = workprecision(nlse, ( - (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), - (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), - (run_numfullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30)), - (run_splitlin, collect(logrange(1e-4, 1e-2, 30)), 1e-6 .* ones(30), 1e-6 .* ones(30)), - (run_Luna_weak, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30)), - (run_Luna_norm, 0.0002 .* ones(30), collect(logrange(1e-7, 1e-1, 30)), 1e-6 .* ones(30)), - (run_newLuna, 0.0002 .* ones(40), collect(logrange(5e-5, 1.2e-1, 40)), 1e-6 .* ones(40)) -)) -savefig(joinpath(pkgdir(Luna), "scripts/solver_work_precision.svg")) +# errs, nfs = workprecision(nlse, ( +# (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_numfullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_splitlin, collect(logrange(1e-4, 1e-2, 30)), 1e-6 .* ones(30), 1e-6 .* ones(30), nothing), +# (run_Luna_weak, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30), nothing), +# (run_Luna_norm, 0.0002 .* ones(30), collect(logrange(1e-7, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(5e-5, 1.2e-1, 40)), 1e-6 .* ones(40), nothing) +# )) +# savefig("scripts/solver_work_precision_nofsal.svg")) + +# errs, nfs = workprecision(nlse, ( +# (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_Luna_weak, 0.0002 .* ones(30), collect(logrange(1e-10, 1e-3, 30)), 1e-10 .* ones(30), nothing), +# (run_Luna_norm, 0.0002 .* ones(30), collect(logrange(1e-7, 1e-1, 30)), 1e-6 .* ones(30), nothing), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(5e-5, 1.2e-1, 40)), 1e-6 .* ones(40), nothing) +# )) +# savefig(solver_work_precision_nabsbound.svg")) + +# errs, nfs = workprecision(nlse, ( +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-4 .* ones(40), "1e-4"), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-6 .* ones(40), "1e-6"), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-8 .* ones(40), "1e-8"), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-10 .* ones(40), "1e-10"), +# )) +# savefig(solver_work_precision_atolscan.svg")) + -# run a comparison to visualise the error -data = [run_newLuna(nlse; reltol=rtol, abstol=1e-6) for rtol in (1e-1, 6.9e-2, 1.2e-2, 5e-4)] -plot_nlse_cmp(nlse, data) -savefig(joinpath(pkgdir(Luna), "scripts/solver_work_precision_cmp.svg"), dpi=600) +# # run a comparison to visualise the error +# data = [run_newLuna(nlse; reltol=rtol, abstol=1e-6) for rtol in (1e-1, 6.9e-2, 1.2e-2, 5e-4)] +# plot_nlse_cmp(nlse, data) +# savefig("solver_work_precision_cmp.svg"), dpi=600) diff --git a/src/Propagator.jl b/src/Propagator.jl index 5ce9c00d..934490e2 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -73,6 +73,8 @@ function callbackcl(integrator) # The output we want must be transformed back from the interaction picture integrator.p.Li!(integrator.p.Litmp, integrator.t) @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.Litmp) + + # define interp function to pass to output interp = let integrator=integrator function interp(z) u = integrator.sol(z) @@ -80,11 +82,16 @@ function callbackcl(integrator) @. u * exp(integrator.p.Litmp) end end + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) + printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - #@. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) # copy back as we modify u in stepfun (absorbing boundaries) - ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + + # copy back as we (might) modify u in stepfun (absorbing boundaries) + integrator.p.Li!(integrator.p.Litmp, integrator.t) + @. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) + #ODE.u_modified!(integrator, false) # We did (possibly) mutate the solution, so cannot keep fsal end # For a non-constant linear operator, we need to integrate L(z) numerically along with @@ -116,6 +123,8 @@ function callbackncl(integrator) Li = @views integrator.u[n+1:end] # Cumulatively integrated linear operator # The output we want must be transformed back from the interaction picture @. integrator.p.Eωtmp = Eω * exp(Li) + + # define interp function to pass to output interp = let integrator=integrator, n=n function interp(z) u = integrator.sol(z) @@ -124,11 +133,14 @@ function callbackncl(integrator) @. Eω * exp(Li) end end + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, ODE.get_proposed_dt(integrator), interp) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) - #@. Eω = integrator.p.Eωtmp * exp(-Li) # copy back as we modify u in stepfun (absorbing boundaries) - ODE.u_modified!(integrator, false) # We didn't mutate the solution, so can keep fsal + + # copy back as we (might) modify u in stepfun (absorbing boundaries) + @. Eω = integrator.p.Eωtmp * exp(-Li) + #ODE.u_modified!(integrator, false) # We did (possibly) mutate the solution, so cannot keep fsal end # Constant linear operator case--linop is an array @@ -145,7 +157,7 @@ function makeprop(f!, linop::Array{ComplexF64}, Eω0, z, zmax, stepfun, printer, end # For a linop tuple we expect a pair of functions (linop, ilinop) where the second function provides the -# cumulatively integrated linear operator. This is mostly for testing. +# cumulatively integrated linear operator. This is mostly for testing with analytically integrable operators. function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) Li! = linop[2] prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) @@ -170,9 +182,10 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; printer = Printer(status_period, zmax) prob, cbfunc, rtol, atol = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) # We do all saving and stats in a callback called at every step - cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(false,false)) - integrator = ODE.init(prob, getproperty(ODE, solver)(); adaptive=true, reltol=rtol, abstol=atol, - dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb, tstops=zstops) + cb = ODE.DiscreteCallback((u,t,integrator) -> true, cbfunc, save_positions=(true,true)) + solveri = getproperty(ODE, solver)() + integrator = ODE.init(prob, solveri; adaptive=true, reltol=rtol, abstol=atol, + dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) printstart(printer) sol = ODE.solve!(integrator) printstop(printer, integrator) From 7b9ccc8406959259e2902f75c86cadd4f4245304 Mon Sep 17 00:00:00 2001 From: John Travers Date: Thu, 4 Sep 2025 21:40:44 +0100 Subject: [PATCH 22/28] remove figures --- scripts/solver_work_precision.svg | 1622 ----------- scripts/solver_work_precision_cmp.svg | 3575 ------------------------- 2 files changed, 5197 deletions(-) delete mode 100644 scripts/solver_work_precision.svg delete mode 100644 scripts/solver_work_precision_cmp.svg diff --git a/scripts/solver_work_precision.svg b/scripts/solver_work_precision.svg deleted file mode 100644 index 3e1018dd..00000000 --- a/scripts/solver_work_precision.svg +++ /dev/null @@ -1,1622 +0,0 @@ - - - - - - - - 2025-09-03T09:48:58.613292 - image/svg+xml - - - Matplotlib v3.10.1, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/scripts/solver_work_precision_cmp.svg b/scripts/solver_work_precision_cmp.svg deleted file mode 100644 index 1b973b5b..00000000 --- a/scripts/solver_work_precision_cmp.svg +++ /dev/null @@ -1,3575 +0,0 @@ - - - - - - - - 2025-09-03T10:43:14.527977 - image/svg+xml - - - Matplotlib v3.10.1, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - From daba0464208d0ff02913fcf647514423d2fc9cc4 Mon Sep 17 00:00:00 2001 From: John Travers Date: Thu, 4 Sep 2025 21:54:23 +0100 Subject: [PATCH 23/28] fix analytic gradient test --- scripts/solver_work_precision.jl | 5 +++-- src/Luna.jl | 2 +- src/Propagator.jl | 2 -- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index 938d878a..a9cdea70 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -304,7 +304,7 @@ end nlse = NLSE(0.016, 48.0); -# run at plot work-precision +# run work-precision plots for various solvers # errs, nfs = workprecision(nlse, ( # (run_fullip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), # (run_pieceip, 0.0002 .* ones(30), collect(logrange(1e-5, 1e-1, 30)), 1e-6 .* ones(30), nothing), @@ -325,15 +325,16 @@ nlse = NLSE(0.016, 48.0); # )) # savefig(solver_work_precision_nabsbound.svg")) +# # wor-precision curves for multiple atol values for new Luna solver # errs, nfs = workprecision(nlse, ( # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-4 .* ones(40), "1e-4"), +# (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-5 .* ones(40), "1e-5"), # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-6 .* ones(40), "1e-6"), # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-8 .* ones(40), "1e-8"), # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-10 .* ones(40), "1e-10"), # )) # savefig(solver_work_precision_atolscan.svg")) - # # run a comparison to visualise the error # data = [run_newLuna(nlse; reltol=rtol, abstol=1e-6) for rtol in (1e-1, 6.9e-2, 1.2e-2, 5e-4)] # plot_nlse_cmp(nlse, data) diff --git a/src/Luna.jl b/src/Luna.jl index 70792618..0b54b453 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -365,7 +365,7 @@ function run(Eω, grid, Et = FT \ Eω function stepfun(Eω, z, dz, interpolant) - Eω .*= grid.ωwin + #Eω .*= grid.ωwin ldiv!(Et, FT, Eω) Et .*= grid.twin mul!(Eω, FT, Et) diff --git a/src/Propagator.jl b/src/Propagator.jl index 934490e2..0f94f844 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -91,7 +91,6 @@ function callbackcl(integrator) # copy back as we (might) modify u in stepfun (absorbing boundaries) integrator.p.Li!(integrator.p.Litmp, integrator.t) @. integrator.u = integrator.p.Eωtmp * exp(-integrator.p.Litmp) - #ODE.u_modified!(integrator, false) # We did (possibly) mutate the solution, so cannot keep fsal end # For a non-constant linear operator, we need to integrate L(z) numerically along with @@ -140,7 +139,6 @@ function callbackncl(integrator) # copy back as we (might) modify u in stepfun (absorbing boundaries) @. Eω = integrator.p.Eωtmp * exp(-Li) - #ODE.u_modified!(integrator, false) # We did (possibly) mutate the solution, so cannot keep fsal end # Constant linear operator case--linop is an array From c23b126bcc66c46f39115d76289ff11ac8604edb Mon Sep 17 00:00:00 2001 From: John Travers Date: Thu, 4 Sep 2025 23:19:19 +0100 Subject: [PATCH 24/28] fix continuing? --- src/Luna.jl | 2 +- src/Propagator.jl | 16 +++++++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 0b54b453..e48813f8 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -402,7 +402,7 @@ function run(Eω, grid, atol = isnothing(atol) ? 1e-5 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period, solver)#, zstops=output.save_cond.grid) + rtol, atol, init_dz, max_dz, min_dz, status_period, solver) end end diff --git a/src/Propagator.jl b/src/Propagator.jl index 0f94f844..54d7485e 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -149,6 +149,9 @@ function makeprop(f!, linop::Array{ComplexF64}, Eω0, z, zmax, stepfun, printer, @. out = linop * z end end + if z > 0 # Handle continuing + Eω0 = @. Eω0 * exp(-linop * z) + end prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) prob, callbackcl, rtol, atol @@ -159,6 +162,10 @@ end function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) Li! = linop[2] prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) + if z > 0 # Handle continuing + Li!(prop.Litmp, z) + Eω0 = @. Eω0 * exp(-prop.Litmp) + end prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) prob, callbackcl, rtol, atol end @@ -167,9 +174,12 @@ end function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) - u0 = vcat(Eω0, zero(Eω0)) # Initial linear operator is zero - #rtol = vcat(ones(length(Eω0))*rtol, ones(length(Eω0))*1e-10) - #atol = vcat(ones(length(Eω0))*atol, ones(length(Eω0))*1e-10) + Li = zeros(ComplexF64, size(Eω0)) + if z > 0 # Handle continuing + quadgk!(linop, Li, 0.0, z) + Eω0 = @. Eω0 * exp(-Li) + end + u0 = vcat(Eω0, Li) # state vector includes integrated linear operator prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) prob, callbackncl, rtol, atol end From b88ba1ff3b06260233aae80fee7eddd215040628 Mon Sep 17 00:00:00 2001 From: John Travers Date: Fri, 5 Sep 2025 00:04:11 +0100 Subject: [PATCH 25/28] tweak default precision --- scripts/solver_work_precision.jl | 4 ++-- src/Luna.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index a9cdea70..26a36de0 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -325,7 +325,7 @@ nlse = NLSE(0.016, 48.0); # )) # savefig(solver_work_precision_nabsbound.svg")) -# # wor-precision curves for multiple atol values for new Luna solver +# # work-precision curves for multiple atol values for new Luna solver # errs, nfs = workprecision(nlse, ( # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-4 .* ones(40), "1e-4"), # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-5 .* ones(40), "1e-5"), @@ -333,7 +333,7 @@ nlse = NLSE(0.016, 48.0); # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-8 .* ones(40), "1e-8"), # (run_newLuna, 0.0002 .* ones(40), collect(logrange(1e-10, 1.2e-1, 40)), 1e-10 .* ones(40), "1e-10"), # )) -# savefig(solver_work_precision_atolscan.svg")) +# savefig("solver_work_precision_atolscan.svg") # # run a comparison to visualise the error # data = [run_newLuna(nlse; reltol=rtol, abstol=1e-6) for rtol in (1e-1, 6.9e-2, 1.2e-2, 5e-4)] diff --git a/src/Luna.jl b/src/Luna.jl index e48813f8..9130e39b 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -398,8 +398,8 @@ function run(Eω, grid, return else Logging.@info("Using $solver solver") - rtol = isnothing(rtol) ? 1e-2 : rtol - atol = isnothing(atol) ? 1e-5 : atol + rtol = isnothing(rtol) ? 7e-2 : rtol + atol = isnothing(atol) ? 1e-6 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; rtol, atol, init_dz, max_dz, min_dz, status_period, solver) From c146cb60507166d595ab3871853fc8f77099ccc7 Mon Sep 17 00:00:00 2001 From: John Travers Date: Fri, 5 Sep 2025 09:33:46 +0100 Subject: [PATCH 26/28] smarter continuing --- src/Propagator.jl | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/src/Propagator.jl b/src/Propagator.jl index 54d7485e..90947a14 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -144,14 +144,11 @@ end # Constant linear operator case--linop is an array function makeprop(f!, linop::Array{ComplexF64}, Eω0, z, zmax, stepfun, printer, rtol, atol) # For a constant linear operator L, the integral is just L*z - Li! = let linop=linop + Li! = let linop=linop, z0=z function Li!(out, z) - @. out = linop * z + @. out = linop * (z - z0) # difference from z0 to handle continuing end end - if z > 0 # Handle continuing - Eω0 = @. Eω0 * exp(-linop * z) - end prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) prob, callbackcl, rtol, atol @@ -160,12 +157,16 @@ end # For a linop tuple we expect a pair of functions (linop, ilinop) where the second function provides the # cumulatively integrated linear operator. This is mostly for testing with analytically integrable operators. function makeprop(f!, linop::Tuple, Eω0, z, zmax, stepfun, printer, rtol, atol) - Li! = linop[2] - prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) - if z > 0 # Handle continuing - Li!(prop.Litmp, z) - Eω0 = @. Eω0 * exp(-prop.Litmp) + Li0 = zero(Eω0) + ilinop! = linop[2] + ilinop!(Li0, z) + Li! = let ilinop! = ilinop!, Li0=Li0 + function Li!(out, z) + ilinop!(out, z) + out .-= Li0 # handle continuing + end end + prop = AnalyticalPropagator(Li!, f!, stepfun, similar(Eω0), similar(Eω0), similar(Eω0), printer) prob = ODE.ODEProblem(fcl!, Eω0, (z, zmax), prop) prob, callbackcl, rtol, atol end @@ -174,12 +175,8 @@ end function makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) prop = NonConstPropagator(linop, f!, stepfun, length(Eω0), similar(Eω0), similar(Eω0), printer) - Li = zeros(ComplexF64, size(Eω0)) - if z > 0 # Handle continuing - quadgk!(linop, Li, 0.0, z) - Eω0 = @. Eω0 * exp(-Li) - end - u0 = vcat(Eω0, Li) # state vector includes integrated linear operator + # continuing is handled implicilty + u0 = vcat(Eω0, zero(Eω0)) # state vector includes integrated linear operator prob = ODE.ODEProblem(fncl!, u0, (z, zmax), prop) prob, callbackncl, rtol, atol end From 3227a8739437019068577ac2b72e83901a8c03ec Mon Sep 17 00:00:00 2001 From: John Travers Date: Sun, 7 Sep 2025 13:53:46 +0100 Subject: [PATCH 27/28] fix step size continuation --- scripts/solver_work_precision.jl | 14 +++---- src/Luna.jl | 8 ++-- src/Output.jl | 67 +++++++++++++++++++++++++------- src/Propagator.jl | 23 +++++++++-- test/test_output.jl | 4 +- 5 files changed, 86 insertions(+), 30 deletions(-) diff --git a/scripts/solver_work_precision.jl b/scripts/solver_work_precision.jl index 26a36de0..6b318bcc 100644 --- a/scripts/solver_work_precision.jl +++ b/scripts/solver_work_precision.jl @@ -119,14 +119,14 @@ function run(prob, solver, adaptive, dt, reltol, abstol; cb=nothing) @time integrator = init(prob, solver; dt, adaptive, reltol, abstol, saveat=zs, callback=cb) println("starting") @time u = solve!(integrator) - zs, u + zs, u, integrator end # run full interaction picture, constant L, analytically integrated function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6, fullret=false) reset!(nlse) prob = ODEProblem(fpre!, getinit(nlse), (0.0, π/2), nlse) - zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + zs, u, integrator = run(prob, solver, adaptive, dt, reltol, abstol) res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) for (i,z) in enumerate(zs) @. res[:,i] = u[:,i] * exp(nlse.L * z) @@ -135,7 +135,7 @@ function run_fullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol println("nfunc: $(nlse.nfunc)") println("error: $err") if fullret - return zs, res, nlse.nfunc, err, u + return zs, res, nlse.nfunc, err, u, integrator end zs, res, nlse.nfunc, err end @@ -146,7 +146,7 @@ function run_numfullip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, rel u0 = vcat(getinit(nlse), zero(nlse.L)) prob = ODEProblem(fdbl!, u0, (0.0, π/2), nlse) cb = DiscreteCallback((u,t,integrator) -> true, noaffect!, save_positions=(true,true)) - zs, u = run(prob, solver, adaptive, dt, reltol, abstol; cb=cb) + zs, u, integrator = run(prob, solver, adaptive, dt, reltol, abstol; cb=cb) zs = Array(u.t) res = Array{Complex{Float64}}(undef, nlse.n, length(zs)) for i in 1:length(zs) @@ -163,7 +163,7 @@ function run_pieceip(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, relto reset!(nlse) prob = ODEProblem(fpre2!, getinit(nlse), (0.0, π/2), nlse) cb = DiscreteCallback((u,t,integrator) -> true, resetaffect!, save_positions=(false,true)) - _, u = run(prob, solver, adaptive, dt, reltol, abstol; cb) + _, u, integrator = run(prob, solver, adaptive, dt, reltol, abstol; cb) res = Array(u) zs = u.t err = geterror(nlse, res[:,end]) @@ -175,7 +175,7 @@ end function run_stiff(nlse::NLSE; solver=Tsit5(), adaptive=true, dt=0.0002, reltol=1e-2, abstol=1e-6) reset!(nlse) prob = ODEProblem(fall!, getinit(nlse), (0.0, π/2), nlse) - zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + zs, u, integrator = run(prob, solver, adaptive, dt, reltol, abstol) res = Array(u) err = geterror(nlse, res[:,end]) println("nfunc: $(nlse.nfunc)") @@ -189,7 +189,7 @@ function run_splitlin(nlse::NLSE; solver=ETDRK4(), adaptive=false, dt=0.0002, re op = DiagonalOperator(nlse.L) f = SplitFunction(op, f2!) prob = SplitODEProblem(f, getinit(nlse), (0.0, π/2), nlse) - zs, u = run(prob, solver, adaptive, dt, reltol, abstol) + zs, u, integrator = run(prob, solver, adaptive, dt, reltol, abstol) res = Array(u) err = geterror(nlse, res[:,end]) println("nfunc: $(nlse.nfunc)") diff --git a/src/Luna.jl b/src/Luna.jl index 9130e39b..89692cf0 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -364,16 +364,16 @@ function run(Eω, grid, Et = FT \ Eω - function stepfun(Eω, z, dz, interpolant) + function stepfun(Eω, z, dz, interpolant; cache=nothing) #Eω .*= grid.ωwin ldiv!(Et, FT, Eω) Et .*= grid.twin mul!(Eω, FT, Et) - output(Eω, z, dz, interpolant) + output(Eω, z, dz, interpolant; cache) end # check_cache does nothing except for HDF5Outputs - Eωc, zc, dzc = Output.check_cache(output, Eω, z0, init_dz) + Eωc, zc, dzc, stepcache = Output.check_cache(output, Eω, z0, init_dz) if zc > z0 Logging.@info("Found cached propagation. Resuming...") Eω, z0, init_dz = Eωc, zc, dzc @@ -402,7 +402,7 @@ function run(Eω, grid, atol = isnothing(atol) ? 1e-6 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period, solver) + rtol, atol, init_dz, max_dz, min_dz, status_period, solver, stepcache) end end diff --git a/src/Output.jl b/src/Output.jl index af2bf0fd..cc137087 100644 --- a/src/Output.jl +++ b/src/Output.jl @@ -39,10 +39,14 @@ function MemoryOutput(save_cond, yname, tname, statsfun=nostats, script=nothing) MemoryOutput(save_cond, yname, tname, 0, data, statsfun) end -function initialise(o::MemoryOutput, y) +function initialise(o::MemoryOutput, t, dt, y) dims = init_dims(size(y), o.save_cond) o.data[o.yname] = Array{ComplexF64}(undef, dims) o.data[o.tname] = Array{Float64}(undef, (dims[end],)) + o.data["cache"] = Dict{String, Any}() + o.data["cache"]["y"] = copy(y) + o.data["cache"]["t"] = t + o.data["cache"]["dt"] = dt end "getindex works interchangeably so when switching from one Output to @@ -62,10 +66,10 @@ haskey(o::MemoryOutput, key) = haskey(o.data, key) yfun: callable which returns interpolated function value at different t Note that from RK45.jl, this will be called with yn and tn as arguments. """ -function (o::MemoryOutput)(y, t, dt, yfun) +function (o::MemoryOutput)(y, t, dt, yfun; kwargs...) save, ts = o.save_cond(y, t, dt, o.saved) append_stats!(o, o.statsfun(y, t, dt)) - !haskey(o.data, o.yname) && initialise(o, y) + !haskey(o.data, o.yname) && initialise(o, t, dt, y) while save s = size(o.data[o.yname]) if s[end] < o.saved+1 @@ -79,6 +83,9 @@ function (o::MemoryOutput)(y, t, dt, yfun) o.saved += 1 save, ts = o.save_cond(y, t, dt, o.saved) end + o.data["cache"]["y"] .= y + o.data["cache"]["t"] = t + o.data["cache"]["dt"] = dt end function append_stats!(o::MemoryOutput, d) @@ -212,7 +219,7 @@ function HDF5Output(fpath::AbstractString) HDF5Output(fpath, 0, 0, 1; readonly=true) end -function initialise(o::HDF5Output, y) +function initialise(o::HDF5Output, t, dt, y) ydims = size(y) idims = init_dims(ydims, o.save_cond) cdims = collect(idims) @@ -235,10 +242,15 @@ function initialise(o::HDF5Output, y) o.cachehash = hash((statsnames, size(y))) file["meta"]["cachehash"] = o.cachehash if o.cache - file["meta"]["cache"]["t"] = typemin(0.0) - file["meta"]["cache"]["dt"] = typemin(0.0) + file["meta"]["cache"]["t"] = t + file["meta"]["cache"]["dt"] = dt file["meta"]["cache"]["y"] = y file["meta"]["cache"]["saved"] = 0 + file["meta"]["cache"]["dtpropose"] = 0.0 + file["meta"]["cache"]["dtcache"] = 0.0 + file["meta"]["cache"]["qold"] = 0.0 + file["meta"]["cache"]["erracc"] = 0.0 + file["meta"]["cache"]["dtacc"] = 0.0 end end end @@ -310,13 +322,13 @@ end yfun: callable which returns interpolated function value at different t Note that from RK45.jl, this will be called with yn and tn as arguments. """ -function (o::HDF5Output)(y, t, dt, yfun) +function (o::HDF5Output)(y, t, dt, yfun; cache=nothing) o.readonly && error("Cannot add data to read-only output!") save, ts = o.save_cond(y, t, dt, o.saved) push!(o.stats_tmp, o.statsfun(y, t, dt)) if save HDF5.h5open(o.fpath, "r+") do file - !HDF5.haskey(file, o.yname) && initialise(o, y) + !HDF5.haskey(file, o.yname) && initialise(o, t, dt, y) statsnames = sort(collect(keys(o.stats_tmp[end]))) cachehash = hash((statsnames, size(y))) cachehash == o.cachehash || error( @@ -345,6 +357,11 @@ function (o::HDF5Output)(y, t, dt, yfun) write(file["meta"]["cache"]["dt"], dt) write(file["meta"]["cache"]["y"], y) write(file["meta"]["cache"]["saved"], o.saved) + write(file["meta"]["cache"]["dtpropose"], cache.dtpropose) + write(file["meta"]["cache"]["dtcache"], cache.dtcache) + write(file["meta"]["cache"]["qold"], cache.qold) + write(file["meta"]["cache"]["erracc"], cache.erracc) + write(file["meta"]["cache"]["dtacc"], cache.dtacc) end end end @@ -464,19 +481,43 @@ Check for an existing cached propagation in the output `o` and return this cache """ function check_cache(o::HDF5Output, y, t, dt) if !o.cache || !haskey(o["meta"]["cache"], "t") - return y, t, dt + return y, t, dt, nothing end tc = o["meta"]["cache"]["t"] if tc < t - return y, t, dt + return y, t, dt, nothing end yc = o["meta"]["cache"]["y"] dtc = o["meta"]["cache"]["dt"] - return yc, tc, dtc + dtacc = o["meta"]["cache"]["dtacc"] + dtpropose = o["meta"]["cache"]["dtpropose"] + dtcache = o["meta"]["cache"]["dtcache"] + qold = o["meta"]["cache"]["qold"] + erracc = o["meta"]["cache"]["erracc"] + dtacc = o["meta"]["cache"]["dtacc"] + return yc, tc, dtc, (;qold, erracc, dtacc, dtpropose, dtcache) end -# For other outputs (e.g. MemoryOutput or another function), checking the cache does nothing. -check_cache(o, y, t, dt) = y, t, dt +""" + check_cache(o::MemoryOutput, y, t, dt) + +Check for an existing cached propagation in the output `o` and return this cache if present. +""" +function check_cache(o::MemoryOutput, y, t, dt) + if !haskey(o, "cache") + return y, t, dt, nothing + end + tc = o.data["cache"]["t"] + if tc < t + return y, t, dt, nothing + end + yc = o.data["cache"]["y"] + dtc = o.data["cache"]["dt"] + return yc, tc, dtc, nothing +end + +# For other outputs, checking the cache does nothing. +check_cache(o, y, t, dt) = y, t, dt, nothing "Condition callable that distributes save points evenly on a grid" struct GridCondition diff --git a/src/Propagator.jl b/src/Propagator.jl index 90947a14..4cb7a04a 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -70,6 +70,7 @@ function fcl!(du,u,p,z) end function callbackcl(integrator) + println(integrator.t, " ", ODE.get_proposed_dt(integrator), " ", integrator.t - integrator.tprev, " ", integrator.dt) # The output we want must be transformed back from the interaction picture integrator.p.Li!(integrator.p.Litmp, integrator.t) @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.Litmp) @@ -77,14 +78,20 @@ function callbackcl(integrator) # define interp function to pass to output interp = let integrator=integrator function interp(z) - u = integrator.sol(z) + u = integrator(z) integrator.p.Li!(integrator.p.Litmp, z) @. u * exp(integrator.p.Litmp) end end + cache = (dtpropose = integrator.dtpropose, + dtcache = integrator.dtcache, + qold = integrator.qold, + erracc = integrator.erracc, + dtacc = integrator.dtacc) + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, - ODE.get_proposed_dt(integrator), interp) + integrator.dt, interp; cache) printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) @@ -126,7 +133,7 @@ function callbackncl(integrator) # define interp function to pass to output interp = let integrator=integrator, n=n function interp(z) - u = integrator.sol(z) + u = integrator(z) Eω = @views u[1:n] Li = @views u[n+1:end] @. Eω * exp(Li) @@ -183,7 +190,7 @@ end function propagate(f!, linop, Eω0, z, zmax, stepfun; rtol=1e-3, atol=1e-6, init_dz=1e-4, max_dz=Inf, min_dz=0, - status_period=1, solver=:Tsit5, zstops=nothing) + status_period=1, solver=:Tsit5, zstops=nothing, stepcache=nothing) printer = Printer(status_period, zmax) prob, cbfunc, rtol, atol = makeprop(f!, linop, Eω0, z, zmax, stepfun, printer, rtol, atol) # We do all saving and stats in a callback called at every step @@ -191,7 +198,15 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; solveri = getproperty(ODE, solver)() integrator = ODE.init(prob, solveri; adaptive=true, reltol=rtol, abstol=atol, dt=init_dz, dtmin=min_dz, dtmax=max_dz, callback=cb) + if !isnothing(stepcache) + integrator.qold = stepcache.qold + integrator.erracc = stepcache.erracc + integrator.dtacc = stepcache.dtacc + integrator.dtpropose = stepcache.dtpropose + integrator.dtcache = stepcache.dtcache + end printstart(printer) + println("size of integrator: ", Base.summarysize(integrator)) sol = ODE.solve!(integrator) printstop(printer, integrator) sol diff --git a/test/test_output.jl b/test/test_output.jl index d8db385c..e5921f65 100644 --- a/test/test_output.jl +++ b/test/test_output.jl @@ -225,8 +225,8 @@ fpath = joinpath(homedir(), ".luna", "output_test", "test.h5") Stats.ω0(grid), Stats.energy(grid, energyfunω)) output = Output.HDF5Output(fpath, 0, grid.zmax, 51, statsfun) - function stepfun(Eω, z, dz, interpolant) - output(Eω, z, dz, interpolant) + function stepfun(Eω, z, dz, interpolant; cache) + output(Eω, z, dz, interpolant; cache) if z > 3e-2 error("Oh no!") end From dd922b8bd380af62a9635734a9e9017a2eed9abf Mon Sep 17 00:00:00 2001 From: John Travers Date: Sun, 7 Sep 2025 14:25:01 +0100 Subject: [PATCH 28/28] fix and clean continuing --- src/Luna.jl | 7 +++--- src/Output.jl | 61 ++++++++++++++++++++++++++++++++++----------- src/Propagator.jl | 25 +++++++++++-------- test/test_output.jl | 10 ++++---- 4 files changed, 70 insertions(+), 33 deletions(-) diff --git a/src/Luna.jl b/src/Luna.jl index 89692cf0..a6203255 100644 --- a/src/Luna.jl +++ b/src/Luna.jl @@ -364,12 +364,12 @@ function run(Eω, grid, Et = FT \ Eω - function stepfun(Eω, z, dz, interpolant; cache=nothing) + function stepfun(Eω, z, dz, interpolant; stepcache=nothing) #Eω .*= grid.ωwin ldiv!(Et, FT, Eω) Et .*= grid.twin mul!(Eω, FT, Et) - output(Eω, z, dz, interpolant; cache) + output(Eω, z, dz, interpolant; stepcache) end # check_cache does nothing except for HDF5Outputs @@ -402,7 +402,8 @@ function run(Eω, grid, atol = isnothing(atol) ? 1e-6 : atol Logging.@info("Using rtol = $rtol, atol = $atol") Propagator.propagate(transform, linop, Eω, z0, grid.zmax, stepfun; - rtol, atol, init_dz, max_dz, min_dz, status_period, solver, stepcache) + rtol, atol, init_dz, max_dz, min_dz, status_period, + solver, stepcache) end end diff --git a/src/Output.jl b/src/Output.jl index cc137087..3231ea22 100644 --- a/src/Output.jl +++ b/src/Output.jl @@ -66,7 +66,7 @@ haskey(o::MemoryOutput, key) = haskey(o.data, key) yfun: callable which returns interpolated function value at different t Note that from RK45.jl, this will be called with yn and tn as arguments. """ -function (o::MemoryOutput)(y, t, dt, yfun; kwargs...) +function (o::MemoryOutput)(y, t, dt, yfun; stepcache=nothing) save, ts = o.save_cond(y, t, dt, o.saved) append_stats!(o, o.statsfun(y, t, dt)) !haskey(o.data, o.yname) && initialise(o, t, dt, y) @@ -86,6 +86,13 @@ function (o::MemoryOutput)(y, t, dt, yfun; kwargs...) o.data["cache"]["y"] .= y o.data["cache"]["t"] = t o.data["cache"]["dt"] = dt + if !isnothing(stepcache) + o.data["cache"]["dtpropose"] = stepcache.dtpropose + o.data["cache"]["dtcache"] = stepcache.dtcache + o.data["cache"]["qold"] = stepcache.qold + o.data["cache"]["erracc"] = stepcache.erracc + o.data["cache"]["dtacc"] = stepcache.dtacc + end end function append_stats!(o::MemoryOutput, d) @@ -322,7 +329,7 @@ end yfun: callable which returns interpolated function value at different t Note that from RK45.jl, this will be called with yn and tn as arguments. """ -function (o::HDF5Output)(y, t, dt, yfun; cache=nothing) +function (o::HDF5Output)(y, t, dt, yfun; stepcache=nothing) o.readonly && error("Cannot add data to read-only output!") save, ts = o.save_cond(y, t, dt, o.saved) push!(o.stats_tmp, o.statsfun(y, t, dt)) @@ -357,11 +364,13 @@ function (o::HDF5Output)(y, t, dt, yfun; cache=nothing) write(file["meta"]["cache"]["dt"], dt) write(file["meta"]["cache"]["y"], y) write(file["meta"]["cache"]["saved"], o.saved) - write(file["meta"]["cache"]["dtpropose"], cache.dtpropose) - write(file["meta"]["cache"]["dtcache"], cache.dtcache) - write(file["meta"]["cache"]["qold"], cache.qold) - write(file["meta"]["cache"]["erracc"], cache.erracc) - write(file["meta"]["cache"]["dtacc"], cache.dtacc) + if !isnothing(stepcache) + write(file["meta"]["cache"]["dtpropose"], stepcache.dtpropose) + write(file["meta"]["cache"]["dtcache"], stepcache.dtcache) + write(file["meta"]["cache"]["qold"], stepcache.qold) + write(file["meta"]["cache"]["erracc"], stepcache.erracc) + write(file["meta"]["cache"]["dtacc"], stepcache.dtacc) + end end end end @@ -489,13 +498,21 @@ function check_cache(o::HDF5Output, y, t, dt) end yc = o["meta"]["cache"]["y"] dtc = o["meta"]["cache"]["dt"] - dtacc = o["meta"]["cache"]["dtacc"] - dtpropose = o["meta"]["cache"]["dtpropose"] - dtcache = o["meta"]["cache"]["dtcache"] - qold = o["meta"]["cache"]["qold"] - erracc = o["meta"]["cache"]["erracc"] - dtacc = o["meta"]["cache"]["dtacc"] - return yc, tc, dtc, (;qold, erracc, dtacc, dtpropose, dtcache) + if haskey(o["meta"]["cache"], "dtacc") + dtacc = o["meta"]["cache"]["dtacc"] + dtpropose = o["meta"]["cache"]["dtpropose"] + dtcache = o["meta"]["cache"]["dtcache"] + qold = o["meta"]["cache"]["qold"] + erracc = o["meta"]["cache"]["erracc"] + dtacc = o["meta"]["cache"]["dtacc"] + stepcache = (; dtacc, dtpropose, dtcache, qold, erracc) + if all(iszero, stepcache) + stepcache = nothing + end + else + stepcache = nothing + end + return yc, tc, dtc, stepcache end """ @@ -513,7 +530,21 @@ function check_cache(o::MemoryOutput, y, t, dt) end yc = o.data["cache"]["y"] dtc = o.data["cache"]["dt"] - return yc, tc, dtc, nothing + if haskey(o["cache"], "dtacc") + dtacc = o["cache"]["dtacc"] + dtpropose = o["cache"]["dtpropose"] + dtcache = o["cache"]["dtcache"] + qold = o["cache"]["qold"] + erracc = o["cache"]["erracc"] + dtacc = o["cache"]["dtacc"] + stepcache = (; dtacc, dtpropose, dtcache, qold, erracc) + if all(iszero, stepcache) + stepcache = nothing + end + else + stepcache = nothing + end + return yc, tc, dtc, stepcache end # For other outputs, checking the cache does nothing. diff --git a/src/Propagator.jl b/src/Propagator.jl index 4cb7a04a..08c8e14f 100644 --- a/src/Propagator.jl +++ b/src/Propagator.jl @@ -70,7 +70,6 @@ function fcl!(du,u,p,z) end function callbackcl(integrator) - println(integrator.t, " ", ODE.get_proposed_dt(integrator), " ", integrator.t - integrator.tprev, " ", integrator.dt) # The output we want must be transformed back from the interaction picture integrator.p.Li!(integrator.p.Litmp, integrator.t) @. integrator.p.Eωtmp = integrator.u * exp(integrator.p.Litmp) @@ -84,15 +83,15 @@ function callbackcl(integrator) end end - cache = (dtpropose = integrator.dtpropose, - dtcache = integrator.dtcache, - qold = integrator.qold, - erracc = integrator.erracc, - dtacc = integrator.dtacc) + stepcache = (dtpropose = integrator.dtpropose, + dtcache = integrator.dtcache, + qold = integrator.qold, + erracc = integrator.erracc, + dtacc = integrator.dtacc) integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, - integrator.dt, interp; cache) - + integrator.dt, interp; stepcache) + printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) # copy back as we (might) modify u in stepfun (absorbing boundaries) @@ -140,8 +139,15 @@ function callbackncl(integrator) end end + stepcache = (dtpropose = integrator.dtpropose, + dtcache = integrator.dtcache, + qold = integrator.qold, + erracc = integrator.erracc, + dtacc = integrator.dtacc) + integrator.p.stepfun(integrator.p.Eωtmp, integrator.t, - ODE.get_proposed_dt(integrator), interp) + integrator.dt, interp; stepcache) + printstep(integrator.p.p, integrator.t, ODE.get_proposed_dt(integrator)) # copy back as we (might) modify u in stepfun (absorbing boundaries) @@ -206,7 +212,6 @@ function propagate(f!, linop, Eω0, z, zmax, stepfun; integrator.dtcache = stepcache.dtcache end printstart(printer) - println("size of integrator: ", Base.summarysize(integrator)) sol = ODE.solve!(integrator) printstop(printer, integrator) sol diff --git a/test/test_output.jl b/test/test_output.jl index e5921f65..be639302 100644 --- a/test/test_output.jl +++ b/test/test_output.jl @@ -225,8 +225,8 @@ fpath = joinpath(homedir(), ".luna", "output_test", "test.h5") Stats.ω0(grid), Stats.energy(grid, energyfunω)) output = Output.HDF5Output(fpath, 0, grid.zmax, 51, statsfun) - function stepfun(Eω, z, dz, interpolant; cache) - output(Eω, z, dz, interpolant; cache) + function stepfun(Eω, z, dz, interpolant; stepcache) + output(Eω, z, dz, interpolant; stepcache) if z > 3e-2 error("Oh no!") end @@ -263,9 +263,9 @@ fpath = joinpath(homedir(), ".luna", "output_test", "test.h5") Eω = output["Eω"][idx1:idx2, :] Iω = abs2.(Eω) Iωm = abs2.(Eωm) - @test norm(Iω - Iωm)/norm(Iω) < 1e-7 - @test all(isapprox.(Eωm, Eω, atol=1e-4*maximum(abs.(Eωm)))) - @test all(isapprox.(output["stats"]["ω0"], mem.data["stats"]["ω0"], rtol=1e-6)) + @test norm(Iω - Iωm)/norm(Iω) < 6e-9 + @test all(isapprox.(Eωm, Eω, atol=1e-5*maximum(abs.(Eωm)))) + @test all(isapprox.(output["stats"]["ω0"], mem.data["stats"]["ω0"], rtol=3e-6)) @test all(output["stats"]["energy"] .≈ mem.data["stats"]["energy"]) @test output["z"] == mem.data["z"] @test output["grid"] == Grid.to_dict(grid)