From 396ba6c348f5c386a048b0043e44909afd594839 Mon Sep 17 00:00:00 2001 From: Mathieu Roule <72818492+MathieuRoule@users.noreply.github.com> Date: Wed, 27 Mar 2024 16:25:52 +0100 Subject: [PATCH] WMat refactoring Independant Fourier transform w.r.t. angle from basis FT computations remove the basisFT structures moved the angle_gradients and v to vp mappings to OrbitalElements --- src/GFunc.jl | 8 +- src/Isochrone/WMatIsochrone.jl | 419 --------------------------------- src/LinearResponse.jl | 6 - src/TimeResponse.jl | 207 ---------------- src/WMat.jl | 392 +++++++++++++++--------------- 5 files changed, 188 insertions(+), 844 deletions(-) delete mode 100644 src/Isochrone/WMatIsochrone.jl delete mode 100644 src/TimeResponse.jl diff --git a/src/GFunc.jl b/src/GFunc.jl index 7f7fd0b..20383c9 100644 --- a/src/GFunc.jl +++ b/src/GFunc.jl @@ -60,7 +60,7 @@ function MakeGu(ndFdJ::Function, # vp -> v vp = δvp*(kvval-0.5) - vval = vFromvp(vp,vmin,vmax,params.VMAPN) + vval = v_from_vp(vp, vmin, vmax, n=params.VMAPN) #### # (u,v) -> (a,e) @@ -76,12 +76,12 @@ function MakeGu(ndFdJ::Function, # compute Jacobians ##### # vp -> v - Jacv = DvDvp(vp,vmin,vmax,params.VMAPN) + _, Jacv = v_from_vp_derivative(vp, vmin, vmax, n=params.VMAPN) # (u,v) -> (α,β). # Renormalized. (2/(ωmax-ωmin) * |∂(α,β)/∂(u,v)|) - resonance = OrbitalElements.Resonance(n1,n2,Wdata.ωmin,Wdata.ωmax) - RenormalizedJacαβ = (2/(Wdata.ωmax-Wdata.ωmin))*OrbitalElements.uv_to_αβ_jacobian(uval,vval,resonance) + resonance = Resonance(n1,n2,Wdata.ωmin,Wdata.ωmax) + RenormalizedJacαβ = (2/(Wdata.ωmax-Wdata.ωmin)) * uv_to_αβ_jacobian(uval,vval,resonance) # (α,β) -> (E,L): this is the most expensive function here, # so we have pre-tabulated it diff --git a/src/Isochrone/WMatIsochrone.jl b/src/Isochrone/WMatIsochrone.jl deleted file mode 100644 index ae0af33..0000000 --- a/src/Isochrone/WMatIsochrone.jl +++ /dev/null @@ -1,419 +0,0 @@ - -######################################################################## -# -# Data structure -# -######################################################################## - -""" -@TO DESCRIBE -""" -function WMatdataCreateIsochrone(n1::Int64,n2::Int64, - params::LinearParameters) - - # compute the frequency scaling factors for this resonance - ωmin, ωmax = OrbitalElements.FindWminWmaxIsochrone(n1,n2) - - # Useful parameters - nradial = params.nradial - Ku, Kv = params.Ku, params.Kv - - return WMatdataType(ωmin,ωmax,zeros(Float64,2,Ku), - zeros(Float64,nradial,Kv,Ku), - zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku), # Orbital mappings - zeros(Float64,Kv,Ku)) -end - - - -######################################################################## -# -# Fourier Transform at one location -# -######################################################################## - -""" - WintegrandIsochrone(u,a,e) - -Integrand computation/update for FT of basis elements -""" -function WintegrandIsochrone(w::Float64, - a::Float64,e::Float64,L::Float64, - Ω1::Float64,Ω2::Float64, - bc::Float64,M::Float64,G::Float64, - basis::AstroBasis.AbstractAstroBasis, - params::LinearParameters)::Tuple{Float64,Float64} - - # Current location of the radius, r=r(w) - rval = OrbitalElements.ru(w,a,e) - - # Current value of the radial frequency integrand (almost dθ/dw) - rp,ra = OrbitalElements.RpRaFromAE(a,e) - gval = OrbitalElements.ThetaRpRaIsochrone(rp,ra,u,bc=bc,Ω₀=params.Orbitalparams.Ω₀) - - # collect the basis elements (in place!) - AstroBasis.tabUl!(basis,params.lharmonic,rval) - - # the velocity for integration (dθ1dw, dθ2dw) - return Ω1*gval, (Ω2 - L/(rval^(2)))*gval -end - -""" - WBasisFTIsochrone(a,e,Ω₁,Ω₂,n1,n2,bc,M,G,basisFT,restab,params) - -Fourier Transform of basis elements using RK4 scheme -result stored in place -""" -function WBasisFTIsochrone(a::Float64,e::Float64, - Ω1::Float64,Ω2::Float64, - n1::Int64,n2::Int64, - bc::Float64,M::Float64,G::Float64, - basis::AstroBasis.AbstractAstroBasis, - restab::Vector{Float64}, - params::LinearParameters) - - @assert length(restab) == basis.nradial "LinearResponse.WBasisFT: FT array not of the same size as the basis" - - # Integration step - Kwp = (params.ADAPTIVEKW) ? ceil(Int64,params.Kw/(0.1+(1-e))) : params.Kw - - # Caution : Reverse integration (lower error at apocenter than pericenter) - # -> Result to multiply by -1 - dw = (2.0)/(Kwp) - - # get (rp,ra) - rp,ra = OrbitalElements.RpRaFromAE(a,e) - - # need angular momentum - Lval = OrbitalElements.isochroneLfromrpra(rp,ra,bc,M,G) - - # Initialise the state vectors: u, theta1, (theta2-psi) - u, theta1, theta2 = -1.0, 0.0, 0.0 - - # launch the integration from the left boundary - gval = OrbitalElements.ThetaRpRaIsochrone(rp,ra,u,bc=bc,Ω₀=params.Orbitalparams.Ω₀) - - # Current location of the radius, r=r(u) - rval = OrbitalElements.ru(u,a,e) - - # the velocity for integration - dt1du, dt2du = Ω1*gval, (Ω2 - Lval/(rval^(2)))*gval - - # collect the basis elements (in place!) - AstroBasis.tabUl!(basis,params.lharmonic,rval) - - fill!(restab,0.0) - - # start the integration loop now that we are initialised - # at each step, we are performing an RK4-like calculation - for istep=1:Kwp - - # RK4 step 1 - # compute the first prefactor - pref1 = (1.0/6.0)*dw*(1.0/(pi))*dt1du*cos(n1*theta1 + n2*theta2) - - # Loop over the radial indices to sum basis contributions - for np=1:basis.nradial - @inbounds restab[np] += pref1*basis.tabUl[np] - end - - # update velocities at end of step 1 - k11 = dw*dt1du - k21 = dw*dt2du - - # RK4 step 2 - # Update the time by half a timestep - u += 0.5*dw - - # Current location of the radius, r=r(u) - rval = OrbitalElements.ru(u,a,e) - - # current value of Theta - gval = OrbitalElements.ThetaRpRaIsochrone(rp,ra,u,bc=bc,Ω₀=params.Orbitalparams.Ω₀) - - # Current value of dtheta1/du and dtheta2/du - dt1du, dt2du = Ω1*gval, (Ω2 - Lval/(rval^(2)))*gval - - # recompute the basis functions for the changed radius value - AstroBasis.tabUl!(basis,params.lharmonic,rval) - - # Common prefactor for all the increments - # Depends on the updated (theta1+0.5*k11,theta2+0.5*k21) - # the factor (1.0/3.0) comes from RK4 - pref2 = (1.0/3.0)*dw*(1.0/(pi))*dt1du*cos(n1*(theta1+0.5*k11) + n2*(theta2+0.5*k21)) - - # Loop over the radial indices to sum basis contributions - #for np=1:basis.nradial - # tabWMat[np,kuval,kvval] += pref2*basis.tabUl[np] - #end - - # update velocities at end of step 2 - k12 = dw*dt1du - k22 = dw*dt2du - - # RK4 step 3 - # The time, u, is not updated for this step - # For this step, no need to re-compute the basis elements, as r has not been updated - - # Common prefactor for all the increments - # depends on the updated (theta1+0.5*k12,theta2+0.5*k22) - # the factor (1.0/3.0) comes from RK4 - pref3 = (1.0/3.0)*dw*(1.0/(pi))*dt1du*cos(n1*(theta1+0.5*k12) + n2*(theta2+0.5*k22)) - - - # Loop over the radial indices to sum basis contributions - for np=1:basis.nradial - #tabWMat[np,kuval,kvval] += pref3*basis.tabUl[np] - @inbounds restab[np] += (pref2+pref3)*basis.tabUl[np] - end - - k13 = k12 # Does not need to be updated - k23 = k22 # Does not need to be updated - - # RK4 step 4 - # update the time by half a timestep: we are now at the next u value - u += 0.5*dw - - # current location of the radius, r=r(u) - rval = OrbitalElements.ru(u,a,e) - - # current value of Theta - gval = OrbitalElements.ThetaRpRaIsochrone(rp,ra,u,bc=bc,Ω₀=params.Orbitalparams.Ω₀) - - # Current value of dtheta1/du and dtheta2/du, always well-posed - dt1du, dt2du = Ω1*gval, (Ω2 - Lval/(rval^(2)))*gval - - # updated basis elements for new rval - AstroBasis.tabUl!(basis,params.lharmonic,rval) - - # Common prefactor for all the increments: - # Depends on the updated (theta1+k13,theta2+k23) - # the factor (1.0/6.0) comes from RK4 - pref4 = (1.0/6.0)*dw*(1.0/(pi))*dt1du*cos(n1*(theta1+k13) + n2*(theta2+k23)) - - # Loop over the radial indices to sum basis contributions - for np=1:basis.nradial - @inbounds restab[np] += pref4*basis.tabUl[np] - end - - # current velocities for theta1,theta2 - k14 = dw*dt1du - k24 = dw*dt2du - - # update the positions using RK4-like sum - theta1 += (k11 + 2.0*k12 + 2.0*k13 + k14)/(6.0) - theta2 += (k21 + 2.0*k22 + 2.0*k23 + k24)/(6.0) - - end # RK4 integration - - return nothing -end - - -""" -with basisFT struct -""" -function WBasisFTIsochrone(a::Float64,e::Float64, - Ω1::Float64,Ω2::Float64, - n1::Int64,n2::Int64, - bc::Float64,M::Float64,G::Float64, - basisFT::FourierTransformedBasis, - params::LinearParameters) - - # Basis FT - WBasisFTIsochrone(a,e,Ω1,Ω2,n1,n2,bc,M,G,basisFT.basis,basisFT.UFT,params) -end - -""" -without Ω1, Ω2 -""" -function WBasisFTIsochrone(a::Float64,e::Float64, - n1::Int64,n2::Int64, - bc::Float64,M::Float64,G::Float64, - basisFT::FourierTransformedBasis, - params::LinearParameters) - - # Frequencies - Ω1, Ω2 = IsochroneOmega12FromAE(a,e,bc,M,G) - - # Basis FT - WBasisFTIsochrone(a,e,Ω1,Ω2,n1,n2,bc,M,G,basisFT.basis,basisFT.UFT,params) -end - - -""" - MakeWmatUVIsochrone(n1::Int64, n2::Int64, tabu::Vector{Float64}, basisFT::FourierTransformedBasis, bc::Float64, M::Float64, G::Float64, params::LinearParameters) - -Construct the matrix `Wdata` based on the given parameters for an isochrone potential. - -# Arguments -- `n1::Int64`: Integer representing some parameter. -- `n2::Int64`: Integer representing some parameter. -- `tabu::Vector{Float64}`: Vector of `Float64` representing values of `u`. -- `basisFT::FourierTransformedBasis`: Type representing some basis Fourier transform. -- `bc::Float64`: Float representing some parameter. -- `M::Float64`: Float representing mass. -- `G::Float64`: Float representing gravitational constant. -- `params::LinearParameters`: Type representing some linear parameters. - -# Returns -- `Wdata`: A structure containing computed values. - -# Description -`MakeWmatUVIsochrone` constructs the matrix `Wdata` based on the provided parameters. It iterates over each combination of `u` and `v` values, computes various orbital elements such as eccentricity, semi-major axis, frequencies, and more. It then computes the Jacobian of the transformation from `(α, β)` to `(E, L)` and saves it in `Wdata`. Finally, it computes `W(u, v)` for each basis element using a Runge-Kutta 4th order scheme and stores the result in `Wdata.tabW`. - -""" -function MakeWmatUVIsochrone(n1::Int64,n2::Int64, - tabu::Vector{Float64}, - basisFT::FourierTransformedBasis, - bc::Float64,M::Float64,G::Float64, - params::LinearParameters) - - @assert length(tabu) == params.Ku "LinearResponse.WMat.MakeWmatUV: tabu length is not Ku." - - Orbitalparams = params.Orbitalparams - Ω₀ = Orbitalparams.Ω₀ - - # allocate the results matrices - Wdata = WMatdataCreateIsochrone(n1,n2,params) - ωmin, ωmax = Wdata.ωmin, Wdata.ωmax - - # start the loop - for kuval = 1:params.Ku - - # get the current u value - uval = tabu[kuval] - - (params.VERBOSE > 2) && println("LinearResponse.WMat.MakeWMat: on step $kuval of $Ku: u=$uval.") - - # get the corresponding v boundary values - vmin,vmax = OrbitalElements.FindVminVmaxIsochrone(n1,n2,uval) - - # saving them - Wdata.tabvminmax[1,kuval], Wdata.tabvminmax[2,kuval] = vmin, vmax - - # determine the step size in v - δvp = 1.0/params.Kv - - - for kvval = 1:params.Kv - - # get the current v value - vp = δvp*(kvval-0.5) - vval = vFromvp(vp,vmin,vmax,params.VMAPN) - - #### - # (u,v) -> (a,e) - #### - # (u,v) -> (α,β) - α,β = OrbitalElements.αβFromUV(uval,vval,n1,n2,ωmin,ωmax) - # (α,β) -> (Ω1,Ω2) - Ω₁, Ω₂ = OrbitalElements.FrequenciesFromαβ(α,β,Ω₀) - # (Ω1,Ω2) -> (a,e) - #a,e = OrbitalElements.ComputeAEFromFrequencies(ψ,dψ,d2ψ,Ω₁,Ω₂,Orbitalparams) - a,e = OrbitalElements.IsochroneAEFromOmega1Omega2(Ω₁,Ω₂,bc,M,G) - - # get (rp,ra) - rp,ra = OrbitalElements.RpRaFromAE(a,e) - - # need angular momentum - Lval = OrbitalElements.isochroneLfromrpra(rp,ra,bc,M,G) - - (params.VERBOSE > 2) && print("v=$kvval,o1=$Ω1,o2=$Ω2;") - - # save (u,v) values for later - Wdata.tabUV[1,kvval,kuval], Wdata.tabUV[2,kvval,kuval] = uval, vval - # save (Ω1,Ω2) values for later - Wdata.tabΩ1Ω2[1,kvval,kuval], Wdata.tabΩ1Ω2[2,kvval,kuval] = Ω₁, Ω₂ - # save (a,e) values for later - Wdata.tabAE[1,kvval,kuval], Wdata.tabAE[2,kvval,kuval] = a, e - # save (E,L) values for later - Wdata.tabEL[1,kvval,kuval], Wdata.tabEL[2,kvval,kuval] = OrbitalElements.isochroneELfromαβ(α,β,bc,M,G) - - # compute the Jacobian of the (α,β) ↦ (E,L) mapping here. a little more expensive, but savings in the long run - Wdata.tabJ[kvval,kuval] = OrbitalElements.IsochroneJacELtoαβ(α,β,bc,M,G) - - # Compute W(u,v) for every basis element using RK4 scheme - WBasisFTIsochrone(a,e,Ω₁,Ω₂,n1,n2,bc,M,G,basisFT,params) - - for np = 1:basisFT.basis.nradial - Wdata.tabW[np,kvval,kuval] = basisFT.UFT[np] - end - end - end - - return Wdata -end - - -######################################################################## -# -# Store FT on all (u,v) points for all resonances -# -######################################################################## - -""" - RunWmat(ψ,dψ,d2ψ,FHT,basis[,params]) - -@TO DESCRIBE -""" -function RunWmatIsochrone(FHT::FiniteHilbertTransform.AbstractFHT, - bc::Float64,M::Float64,G::Float64, - basis::AstroBasis.AbstractAstroBasis, - params::LinearParameters) - - # check the directories + basis and FHT values against the Parameters - CheckDirectories(params.wmatdir) - CheckBasisCompatibility(basis,params) - CheckFHTCompatibility(FHT,params) - - # FT bases prep. - basisFT = BasisFTcreate(basis) - basesFT = [deepcopy(basisFT) for k=1:Threads.nthreads()] - - # print the length of the list of resonance vectors - (params.VERBOSE > 0) && println("LinearResponse.WMat.RunWmatIsochrone: Number of resonances to compute: $(params.nbResVec)") - - Threads.@threads for i = 1:params.nbResVec - k = Threads.threadid() - n1,n2 = params.tabResVec[1,i],params.tabResVec[2,i] - - (params.VERBOSE > 0) && println("LinearResponse.WMat.RunWmatIsochrone: Computing W for the ($n1,$n2) resonance.") - - # Output file name - outputfilename = WMatFilename(n1,n2,params) - - # Check if the file already exist / has enough basis elements / overwritting imposed - # false if no computation needed, then continue - CheckFileNradial(outputfilename,params,"LinearResponse.WMat.RunWMat: ($n1,$n2) resonance") || continue - - # compute the W matrices in UV space: timing optional - if (params.VERBOSE > 1) && (k == 1) - @time Wdata = MakeWmatUVIsochrone(n1,n2,FHT.tabu, - basesFT[k],bc,M,G,params) - else - Wdata = MakeWmatUVIsochrone(n1,n2,FHT.tabu, - basesFT[k],bc,M,G,params) - end - - # now save: we are saving not only W(u,v), but also a(u,v) and e(u,v). - # could consider saving other quantities as well to check mappings. - h5open(outputfilename, "w") do file - # Mappings parameters - write(file, "omgmin",Wdata.ωmin) - write(file, "omgmax",Wdata.ωmax) - write(file, "tabvminmax",Wdata.tabvminmax) - # Mappings - write(file, "UVmat",Wdata.tabUV) - write(file, "Omgmat",Wdata.tabΩ1Ω2) - write(file, "AEmat",Wdata.tabAE) - write(file, "ELmat",Wdata.tabEL) - # Jacobians - write(file, "jELABmat",Wdata.tabJ) - # Basis FT - write(file, "wmat",Wdata.tabW) - # Parameters - WriteParameters(file,params) - end - end -end diff --git a/src/LinearResponse.jl b/src/LinearResponse.jl index ec04c29..3773c68 100644 --- a/src/LinearResponse.jl +++ b/src/LinearResponse.jl @@ -41,10 +41,4 @@ include("ModeComputation/Matrices.jl") include("ModeComputation/Mode.jl") include("ModeComputation/FindPole.jl") -# include code to compute isochrone-specific quantities: -# the isochrone mode is a specific case that we have pre-defined -include("Isochrone/WMatIsochrone.jl") - -include("TimeResponse.jl") - end # module diff --git a/src/TimeResponse.jl b/src/TimeResponse.jl deleted file mode 100644 index 5282f9f..0000000 --- a/src/TimeResponse.jl +++ /dev/null @@ -1,207 +0,0 @@ - - - - -""" - inverseIMinusM!(MMat::Array{Array{Complex{Float64},2},1}, inverse::Array{Array{Complex{Float64},2},1}) - -Function filling the argument `inverse` with the inverse of the argument `matrix`. `matrix` is given as a list of square matrices, corresponding to the values of the response matrix at different times, and is considered as a single triangular block-Toeplitz matrix to compute the inverse. Since the inverse will also be a triangular block-Toeplitz matrix, the `inverse` also has the structure of a list of square matrices. - -# Arguments -- `MMat`: list of square matrices, corresponding to the response matrix at different times. -- `inverse`: list of square matrices to fill. - -# Output -None -""" -function inverseIMinusM!(MMat::Array{Array{Complex{Float64},2},1}, - inverse::Array{Array{Complex{Float64},2},1}, - DeltaT::Float64, - GridTimesSize::Int, - NBasisElements::Int) - - IMat = zeros(Complex{Float64},NBasisElements,NBasisElements) - for np=1:NBasisElements - IMat[np,np] = 1.0 + 0.0im - end - inverse[1] = inv(IMat - DeltaT * MMat[1]) - for k=2:GridTimesSize - matrixSum = zeros(Complex{Float64},NBasisElements,NBasisElements) - for i=1:k-1 - matrixSum += - DeltaT * MMat[i + 1] * inverse[k - i] - end - inverse[k] = - inverse[1] * matrixSum - end -end - - -""" - response!(inverse::Array{Array{Complex{Float64},2},1}, perturber::Array{Array{Complex{Float64},1},1},response::Array{Array{Complex{Float64},1},1}) - -Function filling the argument `response` with the product of the argument `inverse` minus the identity matrix with `perturber`. `inverse` is given as a list of square matrices, considered as a triangular block-Toeplitz matrix. - -# Arguments -- `inverse`: list of square matrices, corresponding to the matrix [(I-M)^-1] at different times. -- `perturber`: list of vectors corresponding to the perturber. The outer dimension concerns different time steps, while the inner dimension concerns the order of radial basis functions. -- `response`: list of vectors to fill. - -# Output -None -""" -function response!(inverse::Array{Array{Complex{Float64},2},1}, - perturber::Array{Array{Complex{Float64},1},1}, - response::Array{Array{Complex{Float64},1},1}, - GridTimesSize::Int, - NBasisElements::Int) - - IMat = zeros(Complex{Float64},NBasisElements,NBasisElements) - for np=1:NBasisElements - IMat[np,np] = 1.0 + 0.0im - end - - for k=1:GridTimesSize - matrixSum = (inverse[1] - IMat) * perturber[k] - for i=2:k - matrixSum += inverse[i] * perturber[k - i + 1] - end - response[k] = matrixSum - end -end - -""" - bareResponse!(matrix::Array{Array{Complex{Float64},2},1}, perturber::Array{Array{Complex{Float64},1},1},response::Array{Array{Complex{Float64},1},1}) - -Function filling the argument `response` with the product of the argument `matrix` with `perturber`. `matrix` is given as a list of square matrices, considered as a triangular block-Toeplitz matrix. - -# Arguments -- `matrix`: list of square matrices, corresponding to the matrix [(I-M)^-1] at different times. -- `perturber`: list of vectors corresponding to the perturber. The outer dimension concerns different time steps, while the inner dimension concerns the order of radial basis functions. -- `response`: list of vectors to fill. - -# Output -None -""" -function bareResponse!(matrix::Array{Array{Complex{Float64},2},1}, - perturber::Array{Array{Complex{Float64},1},1}, - response::Array{Array{Complex{Float64},1},1}, - DeltaT::Float64, - GridTimesSize::Int) - for k=1:GridTimesSize - matrixSum = matrix[1] * perturber[k] - for i=2:k - matrixSum += DeltaT * matrix[i] * perturber[k - i + 1] - end - response[k] = matrixSum - end -end - - - -# Base cardinal sine and cosine are defined as -# Base.sinc(x) = sin(πx)/(πx) -# Base.cosc(x) = cos(πx)/x - sin(πx)/(πx^2) -# while we want -# sinc(t) = sin(t)/t -# cosc(t) = - sin(t)/t^2 -mysinc(t::Float64) = Base.sinc(t/π) -mycosc(t::Float64) = Base.cosc(t/π)/π - - -function timeresponse!(res::Array{ComplexF64,1},t::Float64) - - # Initial values of PFT_k(t) - val_0_PFT = 2im*mysinc(t) # Finite Fourier Transform of P_0 - val_1_PFT = - 2*mycosc(t)+0im # Finite Fourier Transform of P_1 - - FiniteHilbertTransform.tabQLeg!(t+0im, val_0_PFT, val_1_PFT, res) # Computing the tabPFTLeg - - return res -end - -function Getτ( - tnodim::Number, - ωmin::Float64, - ωmax::Float64) - - return (ωmax - ωmin)*tnodim/2 -end - - -function tabM!(t::Float64, - tabM::AbstractMatrix{ComplexF64}, - tabaMcoef::Array{Float64,4}, - tabωminωmax::Matrix{Float64}, - fht::FiniteHilbertTransform.AbstractFHT, - params::LinearParameters) - - # get dimensions from the relevant tables - nbResVec, tabResVec = params.nbResVec, params.tabResVec - KuTruncation = params.KuTruncation - nradial = params.nradial - VERBOSE = params.VERBOSE - Ku = fht.Ku - - # initialise the array to 0. - fill!(tabM,0.0 + 0.0*im) - - # Rescale to get dimensionless frequency - Ω₀ = params.Ω₀ - tnodim = Ω₀*t - - # loop over the resonances: no threading here because we parallelise over frequencies - for nres = 1:nbResVec - - # get current resonance numbers (n1,n2) - n1, n2 = tabResVec[1,nres], tabResVec[2,nres] - - # get ωmin and ωmax values - ωmin, ωmax = tabωminωmax[1,nres], tabωminωmax[2,nres] - - # get the rescaled frequency - τ = Getτ(tnodim,ωmin,ωmax) - - # get the integration values - tabD = timeresponse!(fht.tabDLeg,τ) - - # Add the prefactors that are different from the frequency computations - tabD .*= Ω₀ * (ωmax-ωmin)/2 * exp(-im*tnodim*(ωmax+ωmin)/2) - - # loop over the basis indices to consider - for np = 1:nradial - for nq = np:nradial - - res = 0.0 + 0.0*im - - # loop over the Legendre functions to add all contributions - for k=1:Ku - - # hard check for nans - @inbounds val = tabaMcoef[k,nq,np,nres]*tabD[k] - if !isnan(val) - res += val - else - (k==1) && (VERBOSE>1) && println("LinearResponse.Xi.tabM!: NaN found for n=($n1,$n2), npnq=($np,$nq), k=$k") - end - - - end - - # fill the full M matrix: - # as tab_npnq is the upper triangular matrix (with the diagonal), - # we need to duplicate for symmetries - - # fill in the element (np,nq) - @inbounds tabM[np,nq] += res - - # fill in the element (nq,np). @WARNING: added twice for diagonal elements np=nq. - @inbounds tabM[nq,np] += res - end - end # basis index loop - - end # resonance loop - - # contributions were added twice for diagonal elements: reset - for np=1:nradial - @inbounds tabM[np,np] *= 0.5 - end -end \ No newline at end of file diff --git a/src/WMat.jl b/src/WMat.jl index e6c276e..2dfbc34 100644 --- a/src/WMat.jl +++ b/src/WMat.jl @@ -1,179 +1,112 @@ ######################################################################## # -# Data structure +# Fourier Transform over angles for a given orbit # ######################################################################## - -""" -structure to store Fourier Transform values of basis elements -""" -struct FourierTransformedBasis{BT<:AstroBasis.AbstractAstroBasis} - basis::BT - UFT::Array{Float64} -end - -function BasisFTcreate(basis::BT) where {BT<:AstroBasis.AbstractAstroBasis} - - return FourierTransformedBasis(basis,zeros(Float64,basis.nradial)) -end - - -""" -structure to store the W matrix computation results -""" -struct FourierTransformedBasisData - ωmin::Float64 - ωmax::Float64 - Ω₀::Float64 - tabvminmax::Array{Float64,2} - - tabW::Array{Float64,3} - tabUV::Array{Float64,3} - tabΩ1Ω2::Array{Float64,3} - tabAE::Array{Float64,3} - tabEL::Array{Float64,3} - tabJ::Array{Float64,2} -end - """ -@TO DESCRIBE -""" -function WMatdataCreate(model::OrbitalElements.Potential, - n1::Int64,n2::Int64, - params::LinearParameters) - - # compute the frequency scaling factors for this resonance - ωmin, ωmax = OrbitalElements.frequency_extrema(n1,n2,model,params.Orbitalparams) - - # Useful parameters - nradial = params.nradial - Ku, Kv = params.Ku, params.Kv - - return FourierTransformedBasisData(ωmin,ωmax,OrbitalElements.frequency_scale(model),zeros(Float64,2,Ku), - zeros(Float64,nradial,Kv,Ku), - zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku), # Orbital mappings - zeros(Float64,Kv,Ku)) -end - - -######################################################################## -# -# Mapping functions -# -######################################################################## - -""" -@TO DESCRIBE -""" -function vFromvp(vp::Float64,vmin::Float64,vmax::Float64,n::Int64=2)::Float64 - return (vmax-vmin)*(vp^n)+vmin -end - -""" -@TO DESCRIBE -""" -function DvDvp(vp::Float64,vmin::Float64,vmax::Float64,n::Int64=2)::Float64 - return n*(vmax-vmin)*(vp^(n-1)) -end - -""" -@TO DESCRIBE -""" -function vpFromv(v::Float64,vmin::Float64,vmax::Float64,n::Int64=2)::Float64 - return ((v-vmin)/(vmax-vmin))^(1/n) -end - - -######################################################################## -# -# Fourier Transform at one location -# -######################################################################## - + angle_fouriertransform!( + result::Vector{Float64}, + fun::Function, + a::Float64, + e::Float64, + n1::Int64, + n2::Int64, + model::Potential, + params::LinearParameters; + L::Float64=0.0, + Ω1::Float64=0.0, + Ω2::Float64=0.0 + ) + +Perform the angle Fourier transform of a given function `fun` using the specified parameters. +The result is stored in the `result` vector. + +Assuming the function is of the form `fun(r)`, the Fourier transform is given by: + `result = ∫[0,π] fun(r) * cos(n1*θ1 + n2*(θ2-ϕ)) dθ1` +as `r` is an even function of `θ1` alone. + +# Arguments +- `result::Vector{Float64}`: The vector to store the result of the Fourier transform. +- `fun::Function`: The function to be transformed. Has to be a function of radius only. +- `a::Float64`: Semi-major axis. +- `e::Float64`: Eccentricity. +- `n1::Int64`: Fourier mode for the first angle. +- `n2::Int64`: Fourier mode for the second angle. +- `model::Potential`: The potential model. +- `params::LinearParameters`: Additional parameters for the transformation. + +# Optional Arguments +- `L::Float64=0.0`: Angular momentum. If not provided, it will be calculated +from `a`, `e`, `model`, and `params`. +- `Ω1::Float64=0.0`: Radial frequency. If not provided, it will be calculated +from `a`, `e`, `model`, and `params`. +- `Ω2::Float64=0.0`: Azimuthal frequency. If not provided, it will be calculated +from `a`, `e`, `model`, and `params`. """ - Wintegrand(model, w, a, e, L, Ω1, Ω2, basis, params) - -Integrand computation/update for FT of basis elements -""" -function Wintegrand( - model::Potential, - w::Float64, +function angle_fouriertransform!( + result::Vector{Float64}, + fun::F0, a::Float64, e::Float64, - L::Float64, - Ω1::Float64, - Ω2::Float64, - basis::AstroBasis.AbstractAstroBasis, - params::LinearParameters -)::Tuple{Float64,Float64} - # Current location of the radius, r=r(w) - rval = OrbitalElements.radius_from_anomaly(w, a, e, model, params.Orbitalparams) - - # Current value of the radial frequency integrand (almost dθ/dw) - gval = OrbitalElements.Θ(w, a, e, model, params.Orbitalparams) - - # collect the basis elements (in place!) - AstroBasis.tabUl!(basis,params.lharmonic,rval) - - # the velocity for integration (dθ1dw, dθ2dw) - return Ω1*gval, (Ω2 - L/(rval^(2)))*gval -end - -""" - WBasisFT(a,e,Ω1,Ω2,n1,n2,) - -Fourier Transform of basis elements using RK4 scheme -result stored in place -""" -function WBasisFT(model::OrbitalElements.Potential, - a::Float64,e::Float64, - Ω1::Float64,Ω2::Float64, - n1::Int64,n2::Int64, - basis::AstroBasis.AbstractAstroBasis, - restab::Vector{Float64}, - params::LinearParameters) - - @assert length(restab) == basis.nradial "LinearResponse.WBasisFT: FT array not of the same size as the basis" + n1::Int64, + n2::Int64, + model::Potential, + params::LinearParameters; + L::Float64=0.0, + Ω1::Float64=0.0, + Ω2::Float64=0.0 +) where {F0<:Function} + output_length = length(fun(0.0)) + @assert length(result) == output_length "Result array length does not + match the function's output length." + + if L == 0.0 + # need angular momentum + _, L = EL_from_ae(a, e, model, params) + end + if Ω1 == 0.0 || Ω2 == 0.0 + # need frequencies + Ω1, Ω2 = frequencies_from_ae(a, e, model, params) + end # Integration step Kwp = (params.ADAPTIVEKW) ? ceil(Int64,params.Kw/(0.1+(1-e))) : params.Kw + Orbitalparams = params.Orbitalparams # Caution : Reverse integration (lower error at apocenter than pericenter) # -> Result to multiply by -1 - dw = -(2.0)/(Kwp) - - # need angular momentum - _,Lval = OrbitalElements.EL_from_ae(a,e,model,params.Orbitalparams) + dw = -2/Kwp # Initialise the state vectors: w, θ1, (θ2-psi) + # @WARNING: in this function θ2 really stands for (θ2-psi) which is a function + # of the anomaly `w`. # Reverse integration, starting at apocenter w, θ1, θ2 = 1.0, pi, 0.0 # Initialize integrand - dθ1dw, dθ2dw = Wintegrand(model,w,a,e,Lval,Ω1,Ω2,basis,params) + dθ1dw, dθ2dw = angles_gradient(w, a, e, model, Orbitalparams, L=L, Ω1=Ω1, Ω2=Ω2) + integrand = fun(radius_from_anomaly(w, a, e, model, Orbitalparams)) # Initialize container - fill!(restab,0.0) + fill!(result,0.0) # start the integration loop now that we are initialised # at each step, we are performing an RK4-like calculation - for istep=1:Kwp + for step in 1:Kwp #### # RK4 Step 1 #### # compute the first prefactor - pref1 = (1.0/6.0)*dw*(1.0/(pi))*dθ1dw*cos(n1*θ1 + n2*θ2) + pref1 = (1/6) * dw * (1/pi) * dθ1dw * cos(n1*θ1 + n2*θ2) - # Loop over the radial indices to sum basis contributions - for np=1:basis.nradial - @inbounds restab[np] += pref1*basis.tabUl[np] - end + # Add contribution to the result + result .+= pref1 .* integrand # update velocities at end of step 1 - dθ1_1 = dw*dθ1dw - dθ2_1 = dw*dθ2dw + dθ1_1 = dw * dθ1dw + dθ2_1 = dw * dθ2dw #### # RK4 Step 2 @@ -182,13 +115,13 @@ function WBasisFT(model::OrbitalElements.Potential, w += 0.5*dw # Update integrand - - dθ1dw, dθ2dw = Wintegrand(model,w,a,e,Lval,Ω1,Ω2,basis,params) + dθ1dw, dθ2dw = angles_gradient(w, a, e, model, Orbitalparams, L=L, Ω1=Ω1, Ω2=Ω2) + integrand = fun(radius_from_anomaly(w, a, e, model, Orbitalparams)) # common prefactor for all the increments # depends on the updated (θ1+0.5*dθ1_1,θ2+0.5*dθ2_1) - # the factor (1.0/3.0) comes from RK4 - pref2 = (1.0/3.0)*dw*(1.0/(pi))*dθ1dw*cos(n1*(θ1+0.5*dθ1_1) + n2*(θ2+0.5*dθ2_1)) + # the factor 1/3 comes from RK4 + pref2 = (1/3) * dw * (1/pi) * dθ1dw * cos(n1*(θ1+0.5*dθ1_1) + n2*(θ2+0.5*dθ2_1)) # update velocities at end of step 2 dθ1_2 = dw*dθ1dw @@ -201,14 +134,11 @@ function WBasisFT(model::OrbitalElements.Potential, # For this step, no need to re-compute the basis elements, as r has not been updated # Common prefactor for all the increments # Depends on the updated (θ1+0.5*dθ1_2,θ2+0.5*dθ2_2) - # the factor (1.0/3.0) comes from RK4 - pref3 = (1.0/3.0)*dw*(1.0/(pi))*dθ1dw*cos(n1*(θ1+0.5*dθ1_2) + n2*(θ2+0.5*dθ2_2)) + # the factor 1/3 comes from RK4 + pref3 = (1/3) * dw * (1/pi) * dθ1dw * cos(n1*(θ1+0.5*dθ1_2) + n2*(θ2+0.5*dθ2_2)) - # Loop over the radial indices to sum basis contributions - # Contribution of steps 2 and 3 together - for np=1:basis.nradial - @inbounds restab[np] += (pref2+pref3)*basis.tabUl[np] - end + # Add contribution of steps 2 and 3 together + result .+= (pref2 + pref3) .* integrand dθ1_3 = dθ1_2 # Does not need to be updated dθ2_3 = dθ2_2 # Does not need to be updated @@ -219,74 +149,119 @@ function WBasisFT(model::OrbitalElements.Potential, w += 0.5*dw # Updating the time by half a timestep: we are now at the next u value # Update integrand - - dθ1dw, dθ2dw = Wintegrand(model,w,a,e,Lval,Ω1,Ω2,basis,params) + dθ1dw, dθ2dw = angles_gradient(w, a, e, model, Orbitalparams, L=L, Ω1=Ω1, Ω2=Ω2) + integrand = fun(radius_from_anomaly(w, a, e, model, Orbitalparams)) # Common prefactor for all the increments # Depends on the updated (θ1+dθ1_3,θ2+dθ2_3) - # The factor (1.0/6.0) comes from RK4 - pref4 = (1.0/6.0)*dw*(1.0/(pi))*dθ1dw*cos(n1*(θ1+dθ1_3) + n2*(θ2+dθ2_3)) + # The factor (1/6) comes from RK4 + pref4 = (1/6) * dw * (1/pi) * dθ1dw * cos(n1*(θ1+dθ1_3) + n2*(θ2+dθ2_3)) # Loop over the radial indices to sum basis contributions - for np=1:basis.nradial - @inbounds restab[np] += pref4*basis.tabUl[np] - end + result .+= pref4 .* integrand dθ1_4 = dw*dθ1dw # Current velocity for θ1 dθ2_4 = dw*dθ2dw # Current velocity for θ2 # Update the positions using RK4-like sum (Simpson's 1/3 rule) - θ1 += (dθ1_1 + 2.0*dθ1_2 + 2.0*dθ1_3 + dθ1_4)/(6.0) - θ2 += (dθ2_1 + 2.0*dθ2_2 + 2.0*dθ2_3 + dθ2_4)/(6.0) + θ1 += (dθ1_1 + 2*dθ1_2 + 2*dθ1_3 + dθ1_4)/6 + θ2 += (dθ2_1 + 2*dθ2_2 + 2*dθ2_3 + dθ2_4)/6 # clean or check nans? end # RK4 integration - # check the state of θ1,θ2: - #println("(a,e)=",a," ",e," T1=",θ1," T2=",θ2) - # -1 factor (reverse integration) - for np=1:basis.nradial - @inbounds restab[np] *= -1.0 - end + result .*= -1 - return nothing + return result end +function angle_fouriertransform( + fun::F0, + a::Float64, + e::Float64, + n1::Int64, + n2::Int64, + model::Potential, + params::LinearParameters; + L::Float64=0.0, + Ω1::Float64=0.0, + Ω2::Float64=0.0 +) where {F0<:Function} + result = zeros(Float64,length(fun(0.0))) + return angle_fouriertransform!( + result, fun, a, e, n1, n2, model, params, L=L, Ω1=Ω1, Ω2=Ω2 + ) +end + +""" +for a basis element, compute the Fourier transform over the angles. +""" +function angle_fouriertransform( + basis::AstroBasis.AbstractAstroBasis, + a::Float64, + e::Float64, + n1::Int64, + n2::Int64, + model::Potential, + params::LinearParameters; + L::Float64=0.0, + Ω1::Float64=0.0, + Ω2::Float64=0.0 +) + function basisft_integrand(r::Float64) + # collect the basis elements (in place!) + AstroBasis.tabUl!(basis, params.lharmonic, r) + return basis.tabUl + end + return angle_fouriertransform( + basisft_integrand, a, e, n1, n2, model, params, L=L, Ω1=Ω1, Ω2=Ω2 + ) +end +######################################################################## +# +# Data structure +# +######################################################################## """ -with basisFT struct +structure to store the W matrix computation results """ -function WBasisFT(model::OrbitalElements.Potential, - a::Float64,e::Float64, - Ω1::Float64,Ω2::Float64, - n1::Int64,n2::Int64, - basisFT::FourierTransformedBasis, - params::LinearParameters) - - # Basis FT - WBasisFT(model,a,e,Ω1,Ω2,n1,n2,basisFT.basis,basisFT.UFT,params) +struct FourierTransformedBasisData + ωmin::Float64 + ωmax::Float64 + Ω₀::Float64 + tabvminmax::Array{Float64,2} + + tabW::Array{Float64,3} + tabUV::Array{Float64,3} + tabΩ1Ω2::Array{Float64,3} + tabAE::Array{Float64,3} + tabEL::Array{Float64,3} + tabJ::Array{Float64,2} end """ -without Ω1, Ω2 +@TO DESCRIBE """ -function WBasisFT(model::OrbitalElements.Potential, - a::Float64,e::Float64, - n1::Int64,n2::Int64, - basisFT::FourierTransformedBasis, - params::LinearParameters) +function WMatdataCreate(model::OrbitalElements.Potential, + n1::Int64,n2::Int64, + params::LinearParameters) - # Frequencies - Ω1, Ω2 = OrbitalElements.frequencies_from_ae(a,e,model,params.Orbitalparams) + # compute the frequency scaling factors for this resonance + ωmin, ωmax = frequency_extrema(n1,n2,model,params.Orbitalparams) - # Basis FT + # Useful parameters + nradial = params.nradial + Ku, Kv = params.Ku, params.Kv - WBasisFT(model,a,e,Ω1,Ω2,n1,n2,basisFT.basis,basisFT.UFT,params) + return FourierTransformedBasisData(ωmin,ωmax,frequency_scale(model),zeros(Float64,2,Ku), + zeros(Float64,nradial,Kv,Ku), + zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku),zeros(Float64,2,Kv,Ku), # Orbital mappings + zeros(Float64,Kv,Ku)) end - ######################################################################## # # Store FT on all (u,v) points for one resonance number @@ -299,17 +274,16 @@ end function MakeWmatUV(model::OrbitalElements.Potential, n1::Int64,n2::Int64, tabu::Vector{Float64}, - basisFT::FourierTransformedBasis, + basis::AstroBasis.AbstractAstroBasis, params::LinearParameters) @assert length(tabu) == params.Ku "LinearResponse.WMat.MakeWmatUV: tabu length is not Ku." Orbitalparams = params.Orbitalparams - Ω₀ = OrbitalElements.frequency_scale(model) + Ω₀ = frequency_scale(model) # allocate the results matrices Wdata = WMatdataCreate(model,n1,n2,params) - ωmin, ωmax = Wdata.ωmin, Wdata.ωmax # start the loop for kuval = 1:params.Ku @@ -320,8 +294,8 @@ function MakeWmatUV(model::OrbitalElements.Potential, (params.VERBOSE > 2) && println("LinearResponse.WMat.MakeWMat: on step $kuval of $Ku: u=$uval.") # get the corresponding v boundary values - resonance = OrbitalElements.Resonance(n1,n2,model,Orbitalparams) - vmin,vmax = OrbitalElements.v_boundaries(uval,resonance,model,Orbitalparams) + resonance = Resonance(n1,n2,model,Orbitalparams) + vmin,vmax = v_boundaries(uval,resonance,model,Orbitalparams) # saving them Wdata.tabvminmax[1,kuval], Wdata.tabvminmax[2,kuval] = vmin, vmax @@ -334,17 +308,17 @@ function MakeWmatUV(model::OrbitalElements.Potential, # get the current v value vp = δvp*(kvval-0.5) - vval = vFromvp(vp,vmin,vmax,params.VMAPN) + vval = v_from_vp(vp, vmin, vmax, n=params.VMAPN) #### # (u,v) -> (a,e) #### # (u,v) -> (α,β) - α,β = OrbitalElements.αβ_from_uv(uval,vval,resonance) + α,β = αβ_from_uv(uval,vval,resonance) # (α,β) -> (Ω1,Ω2) - Ω₁, Ω₂ = OrbitalElements.frequencies_from_αβ(α,β,Ω₀) + Ω₁, Ω₂ = frequencies_from_αβ(α,β,Ω₀) # (Ω1,Ω2) -> (a,e) - a,e = OrbitalElements.ae_from_frequencies(Ω₁,Ω₂,model,Orbitalparams) + a,e = ae_from_frequencies(Ω₁,Ω₂,model,Orbitalparams) (params.VERBOSE > 2) && print("v=$kvval,o1=$Ω1,o2=$Ω2;") @@ -355,16 +329,19 @@ function MakeWmatUV(model::OrbitalElements.Potential, # save (a,e) values for later Wdata.tabAE[1,kvval,kuval], Wdata.tabAE[2,kvval,kuval] = a, e # save (E,L) values for later - Wdata.tabEL[1,kvval,kuval], Wdata.tabEL[2,kvval,kuval] = OrbitalElements.EL_from_ae(a,e,model,Orbitalparams) + E, L = EL_from_ae(a,e,model,Orbitalparams) + Wdata.tabEL[1,kvval,kuval], Wdata.tabEL[2,kvval,kuval] = E, L # compute the Jacobian of the (α,β) ↦ (E,L) mapping here. a little more expensive, but savings in the long run - Wdata.tabJ[kvval,kuval] = OrbitalElements.ae_to_EL_jacobian(a,e,model,Orbitalparams)/OrbitalElements.ae_to_αβ_jacobian(a,e,model,Orbitalparams) + Wdata.tabJ[kvval,kuval] = ae_to_EL_jacobian(a,e,model,Orbitalparams)/ae_to_αβ_jacobian(a,e,model,Orbitalparams) # Compute W(u,v) for every basis element using RK4 scheme - WBasisFT(model,a,e,Ω₁,Ω₂,n1,n2,basisFT,params) + basisft = angle_fouriertransform( + basis, a, e, n1, n2, model, params, L=L, Ω1=Ω₁, Ω2=Ω₂ + ) - for np = 1:basisFT.basis.nradial - Wdata.tabW[np,kvval,kuval] = basisFT.UFT[np] + for np = 1:basis.nradial + Wdata.tabW[np,kvval,kuval] = basisft[np] end end end @@ -395,8 +372,7 @@ function RunWmat(model::OrbitalElements.Potential, CheckFHTCompatibility(FHT,params) # FT bases prep. - basisFT = BasisFTcreate(basis) - basesFT = [deepcopy(basisFT) for k=1:Threads.nthreads()] + bases = [deepcopy(basis) for k=1:Threads.nthreads()] # print the length of the list of resonance vectors (params.VERBOSE > 0) && println("LinearResponse.WMat.RunWmat: Number of resonances to compute: $(params.nbResVec)") @@ -418,11 +394,11 @@ function RunWmat(model::OrbitalElements.Potential, if (params.VERBOSE > 1) && (k == 1) @time Wdata = MakeWmatUV(model, n1,n2,FHT.tabu, - basesFT[k],params) + bases[k],params) else Wdata = MakeWmatUV(model, n1,n2,FHT.tabu, - basesFT[k],params) + bases[k],params) end # now save: we are saving not only W(u,v), but also a(u,v) and e(u,v).