Skip to content

Commit

Permalink
Fix density handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@jtravs
jtravs committed Sep 4, 2021
1 parent 67e0932 commit 380cb30
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 68 deletions.
76 changes: 33 additions & 43 deletions src/Nonlinear.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,33 +5,18 @@ import Luna: Maths, Utils
import FFTW
import LinearAlgebra: mul!, ldiv!

"""
scaled_response(resp, scale)
Scale the response function `resp` by the factor `scale`.
"""
function scaled_response(resp, scale)
let resp=resp
function resp!(out, E)
resp(out, E)
out .*= scale
end
end
end


function KerrScalar(out, E, fac)
@. out = fac*E^3
function KerrScalar!(out, E, fac)
@. out += fac*E^3
end

function KerrVector(out, E, fac)
function KerrVector!(out, E, fac)
for i = 1:size(E,1)
Ex = E[i,1]
Ey = E[i,2]
Ex2 = Ex^2
Ey2 = Ey^2
out[i,1] = fac*(Ex2 + Ey2)*Ex
out[i,2] = fac*(Ex2 + Ey2)*Ey
out[i,1] += fac*(Ex2 + Ey2)*Ex
out[i,2] += fac*(Ex2 + Ey2)*Ey
end
end

Expand All @@ -40,9 +25,9 @@ function Kerr_field(γ3)
Kerr = let γ3 = γ3
function Kerr(out, E, ρ)
if size(E,2) == 1
KerrScalar(out, E, ε_0*γ3)
KerrScalar!(out, E, ρ*ε_0*γ3)
else
KerrVector(out, E, ε_0*γ3)
KerrVector!(out, E, ρ*ε_0*γ3)
end
end
end
Expand All @@ -54,23 +39,23 @@ function Kerr_field_nothg(γ3, n)
hilbert = Maths.plan_hilbert(E)
Kerr = let γ3 = γ3, hilbert = hilbert
function Kerr(out, E, ρ)
out .= 3/4*ε_0*γ3.*abs2.(hilbert(E)).*E
out .+= ρ*3/4*ε_0*γ3.*abs2.(hilbert(E)).*E
end
end
end

function KerrScalarEnv(out, E, fac)
@. out = 3/4*fac*abs2(E)*E
function KerrScalarEnv!(out, E, fac)
@. out += 3/4*fac*abs2(E)*E
end

function KerrVectorEnv(out, E, fac)
function KerrVectorEnv!(out, E, fac)
for i = 1:size(E,1)
Ex = E[i,1]
Ey = E[i,2]
Ex2 = abs2(Ex)
Ey2 = abs2(Ey)
out[i,1] = 3/4*fac*((Ex2 + 2/3*Ey2)*Ex + 1/3*conj(Ex)*Ey^2)
out[i,2] = 3/4*fac*((Ey2 + 2/3*Ex2)*Ey + 1/3*conj(Ey)*Ex^2)
out[i,1] += 3/4*fac*((Ex2 + 2/3*Ey2)*Ex + 1/3*conj(Ex)*Ey^2)
out[i,2] += 3/4*fac*((Ey2 + 2/3*Ex2)*Ey + 1/3*conj(Ey)*Ex^2)
end
end

Expand All @@ -79,9 +64,9 @@ function Kerr_env(γ3)
Kerr = let γ3 = γ3
function Kerr(out, E, ρ)
if size(E,2) == 1
KerrScalarEnv(out, E, ε_0*γ3)
KerrScalarEnv!(out, E, ρ*ε_0*γ3)
else
KerrVectorEnv(out, E, ε_0*γ3)
KerrVectorEnv!(out, E, ρ*ε_0*γ3)
end
end
end
Expand All @@ -93,7 +78,7 @@ function Kerr_env_thg(γ3, ω0, t)
C = exp.(2im*ω0.*t)
Kerr = let γ3 = γ3, C = C
function Kerr(out, E, ρ)
@. out = ε_0*γ3/4*(3*abs2(E) + C*E^2)*E
@. out += ρ*ε_0*γ3/4*(3*abs2(E) + C*E^2)*E
end
end
end
Expand All @@ -107,6 +92,7 @@ struct PlasmaCumtrapz{R, EType, tType}
fraction::tType # buffer to hold the ionization fraction
phase::EType # buffer to hold the plasma induced (mostly) phase modulation
J::EType # buffer to hold the plasma current
P::EType # buffer to hold the plasma polarisation
δt::Float64 # the time step
end

Expand All @@ -122,11 +108,12 @@ function PlasmaCumtrapz(t, E, ratefunc, ionpot)
fraction = similar(t)
phase = similar(E)
J = similar(E)
return PlasmaCumtrapz(ratefunc, ionpot, rate, fraction, phase, J, t[2]-t[1])
P = similar(E)
return PlasmaCumtrapz(ratefunc, ionpot, rate, fraction, phase, J, P, t[2]-t[1])
end

"The plasma response for a scalar electric field"
function PlasmaScalar!(out, Plas::PlasmaCumtrapz, E)
function PlasmaScalar!(Plas::PlasmaCumtrapz, E)
Plas.ratefunc(Plas.rate, E)
Maths.cumtrapz!(Plas.fraction, Plas.rate, Plas.δt)
@. Plas.fraction = 1-exp(-Plas.fraction)
Expand All @@ -137,7 +124,7 @@ function PlasmaScalar!(out, Plas::PlasmaCumtrapz, E)
Plas.J[ii] += Plas.ionpot * Plas.rate[ii] * (1-Plas.fraction[ii])/E[ii]
end
end
Maths.cumtrapz!(out, Plas.J, Plas.δt)
Maths.cumtrapz!(Plas.P, Plas.J, Plas.δt)
end

"""
Expand All @@ -149,7 +136,7 @@ for the vector field.
A similar approach was used in: C Tailliez et al 2020 New J. Phys. 22 103038.
"""
function PlasmaVector!(out, Plas::PlasmaCumtrapz, E)
function PlasmaVector!(Plas::PlasmaCumtrapz, E)
Ex = E[:,1]
Ey = E[:,2]
Em = @. hypot.(Ex, Ey)
Expand All @@ -165,19 +152,22 @@ function PlasmaVector!(out, Plas::PlasmaCumtrapz, E)
Plas.J[ii,2] += pre*Ey[ii]
end
end
Maths.cumtrapz!(out, Plas.J, Plas.δt)
Maths.cumtrapz!(Plas.P, Plas.J, Plas.δt)
end

"Handle plasma polarisation routing to `PlasmaVector` or `PlasmaScalar`."
function (Plas::PlasmaCumtrapz)(out, Et, ρ)
if ndims(Et) > 1
if size(Et, 2) == 1 # handle scalar case but within modal simulation
PlasmaScalar!(out, Plas, reshape(Et, size(Et, 1)))
PlasmaScalar!(Plas, reshape(Et, size(Et,1)))
out .+= ρ .* reshape(Plas.P, size(Et))
else
PlasmaVector!(out, Plas, Et) # vector case
PlasmaVector!(Plas, Et) # vector case
out .+= ρ .* Plas.P
end
else
PlasmaScalar!(out, Plas, Et)
PlasmaScalar!(Plas, Et) # straight scalar case
out .+= ρ .* Plas.P
end
end

Expand Down Expand Up @@ -324,16 +314,16 @@ function (R::RamanPolar)(out, Et, ρ)
R.P .= R.FT \ R.Pω

# calculate full polarisation, extracting only the valid
# grid region
# grid region, which is the first length(E) part.
for i = 1:length(E)
R.Pout[i] = E[i]*R.P[i]#(n ÷ 2) + i]
R.Pout[i] = ρ*E[i]*R.P[i]
end

# copy to output in dimensions requested
if ndims(Et) > 1
out .= reshape(R.Pout, size(Et))
out .+= reshape(R.Pout, size(Et))
else
out .= R.Pout
out .+= R.Pout
end
end

Expand Down
39 changes: 14 additions & 25 deletions src/NonlinearRHS.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,28 +101,25 @@ end


"Accumulate responses induced by Et in Pt"
function Et_to_Pt!(Pt, Ptbuf, Et, responses, density::Number)
function Et_to_Pt!(Pt, Et, responses, density::Number)
fill!(Pt, 0)
for resp! in responses
resp!(Ptbuf, Et, density)
Pt .+= density .* Ptbuf
resp!(Pt, Et, density)
end
end

function Et_to_Pt!(Pt, Ptbuf, Et, responses, density::AbstractVector)
function Et_to_Pt!(Pt, Et, responses, density::AbstractVector)
fill!(Pt, 0)
for ii in eachindex(density)
for resp! in responses[ii]
resp!(Ptbuf, Et, density[ii])
Pt .+= density[ii] .* Ptbuf
resp!(Pt, Et, density[ii])
end
end
end


function Et_to_Pt!(Pt, Ptbuf, Et, responses, density, idcs)
function Et_to_Pt!(Pt, Et, responses, density, idcs)
for i in idcs
Et_to_Pt!(view(Pt, :, i), Ptbuf, view(Et, :, i), responses, density)
Et_to_Pt!(view(Pt, :, i), view(Et, :, i), responses, density)
end
end

Expand All @@ -135,7 +132,6 @@ mutable struct TransModal{tsT, lT, TT, FTT, rT, gT, dT, ddT, nT}
Erωo::Array{ComplexF64,2}
Er::Array{TT,2}
Pr::Array{TT,2}
Prbuf::Array{TT,2}
Prω::Array{ComplexF64,2}
Prωo::Array{ComplexF64,2}
Prmω::Array{ComplexF64,2}
Expand Down Expand Up @@ -175,12 +171,11 @@ function TransModal(tT, grid, ts::Modes.ToSpace, FT, resp, densityfun, norm!;
Erωo = Array{ComplexF64,2}(undef, length(grid.ωo), ts.npol)
Er = Array{tT,2}(undef, length(grid.to), ts.npol)
Pr = Array{tT,2}(undef, length(grid.to), ts.npol)
Prbuf = similar(Pr)
Prω = Array{ComplexF64,2}(undef, length(grid.ω), ts.npol)
Prωo = Array{ComplexF64,2}(undef, length(grid.ωo), ts.npol)
Prmω = Array{ComplexF64,2}(undef, length(grid.ω), ts.nmodes)
IFT = inv(FT)
TransModal(ts, full, Modes.dimlimits(ts.ms[1]), Emω, Erω, Erωo, Er, Pr, Prbuf, Prω, Prωo, Prmω,
TransModal(ts, full, Modes.dimlimits(ts.ms[1]), Emω, Erω, Erωo, Er, Pr, Prω, Prωo, Prmω,
FT, resp, grid, densityfun, densityfun(0.0), norm!, 0, 0.0, rtol, atol, mfcn)
end

Expand Down Expand Up @@ -233,7 +228,7 @@ function pointcalc!(fval, xs, t::TransModal)
Modes.to_space!(t.Erω, t.Emω, x, t.ts, z=t.z)
to_time!(t.Er, t.Erω, t.Erωo, inv(t.FT))
# get nonlinear pol at r,θ
Et_to_Pt!(t.Pr, t.Prbuf, t.Er, t.resp, t.density)
Et_to_Pt!(t.Pr, t.Er, t.resp, t.density)
@. t.Pr *= t.grid.towin
to_freq!(t.Prω, t.Prωo, t.Pr, t.FT)
@. t.Prω *= t.grid.ωwin
Expand Down Expand Up @@ -272,7 +267,6 @@ end

struct TransModeAvg{TT, FTT, rT, gT, dT, nT, aT}
Pto::Vector{TT}
Ptobuf::Vector{TT}
Eto::Vector{TT}
Eωo::Vector{ComplexF64}
Pωo::Vector{ComplexF64}
Expand All @@ -296,9 +290,8 @@ function TransModeAvg(TT, grid, FT, resp, densityfun, norm!, aeff)
Eωo = zeros(ComplexF64, length(grid.ωo))
Eto = zeros(TT, length(grid.to))
Pto = similar(Eto)
Ptobuf = similar(Eto)
Pωo = similar(Eωo)
TransModeAvg(Pto, Ptobuf, Eto, Eωo, Pωo, FT, resp, grid, densityfun, norm!, aeff)
TransModeAvg(Pto, Eto, Eωo, Pωo, FT, resp, grid, densityfun, norm!, aeff)
end

function TransModeAvg(grid::Grid.RealGrid, FT, resp, densityfun, norm!, aeff)
Expand All @@ -315,7 +308,7 @@ const nlscale = sqrt(PhysData.ε_0*PhysData.c/2)
function (t::TransModeAvg)(nl, Eω, z)
to_time!(t.Eto, Eω, t.Eωo, inv(t.FT))
@. t.Eto /= nlscale*sqrt(t.aeff(z))
Et_to_Pt!(t.Pto, t.Ptobuf, t.Eto, t.resp, t.densityfun(z))
Et_to_Pt!(t.Pto, t.Eto, t.resp, t.densityfun(z))
@. t.Pto *= t.grid.towin
to_freq!(nl, t.Pωo, t.Pto, t.FT)
t.norm!(nl, z)
Expand Down Expand Up @@ -352,7 +345,6 @@ struct TransRadial{TT, HTT, FTT, nT, rT, gT, dT, iT}
grid::gT # time grid
densityfun::dT # callable which returns density
Pto::Array{TT,2} # Buffer array for NL polarisation on oversampled time grid
Ptobuf::Vector{TT} # Buffer to be used in Et_to_Pt!
Eto::Array{TT,2} # Buffer array for field on oversampled time grid
Eωo::Array{ComplexF64,2} # Buffer array for field on oversampled frequency grid
Pωo::Array{ComplexF64,2} # Buffer array for NL polarisation on oversampled frequency grid
Expand All @@ -373,10 +365,9 @@ function TransRadial(TT, grid, HT, FT, responses, densityfun, normfun)
Eωo = zeros(ComplexF64, (length(grid.ωo), HT.N))
Eto = zeros(TT, (length(grid.to), HT.N))
Pto = similar(Eto)
Ptobuf = zeros(TT, length(grid.to))
Pωo = similar(Eωo)
idcs = CartesianIndices(size(Pto)[2:end])
TransRadial(HT, FT, normfun, responses, grid, densityfun, Pto, Ptobuf, Eto, Eωo, Pωo, idcs)
TransRadial(HT, FT, normfun, responses, grid, densityfun, Pto, Eto, Eωo, Pωo, idcs)
end

"""
Expand Down Expand Up @@ -409,7 +400,7 @@ place the result in `nl`
function (t::TransRadial)(nl, Eω, z)
to_time!(t.Eto, Eω, t.Eωo, inv(t.FT)) # transform ω -> t
ldiv!(t.Eto, t.QDHT, t.Eto) # transform k -> r
Et_to_Pt!(t.Pto, t.Ptobuf, t.Eto, t.resp, t.densityfun(z), t.idcs) # add up responses
Et_to_Pt!(t.Pto, t.Eto, t.resp, t.densityfun(z), t.idcs) # add up responses
@. t.Pto *= t.grid.towin # apodisation
mul!(t.Pto, t.QDHT, t.Pto) # transform r -> k
to_freq!(nl, t.Pωo, t.Pto, t.FT) # transform t -> ω
Expand Down Expand Up @@ -474,7 +465,6 @@ mutable struct TransFree{TT, FTT, nT, rT, gT, xygT, dT, iT}
xygrid::xygT
densityfun::dT # callable which returns density
Pto::Array{TT, 3} # buffer for oversampled time-domain NL polarisation
Ptobuf::Array{TT, 1} # buffer for Et_to_Pt!
Eto::Array{TT, 3} # buffer for oversampled time-domain field
Eωo::Array{ComplexF64, 3} # buffer for oversampled frequency-domain field
Pωo::Array{ComplexF64, 3} # buffer for oversampled frequency-domain NL polarisation
Expand All @@ -498,11 +488,10 @@ function TransFree(TT, scale, grid, xygrid, FT, responses, densityfun, normfun)
Eωo = zeros(ComplexF64, (length(grid.ωo), Ny, Nx))
Eto = zeros(TT, (length(grid.to), Ny, Nx))
Pto = similar(Eto)
Ptobuf = zeros(TT, length(grid.to))
Pωo = similar(Eωo)
idcs = CartesianIndices((Ny, Nx))
TransFree(FT, normfun, responses, grid, xygrid, densityfun,
Pto, Ptobuf, Eto, Eωo, Pωo, scale, idcs)
Pto, Eto, Eωo, Pωo, scale, idcs)
end

"""
Expand Down Expand Up @@ -542,7 +531,7 @@ function (t::TransFree)(nl, Eωk, z)
fill!(t.Eωo, 0)
copy_scale!(t.Eωo, Eωk, length(t.grid.ω), t.scale)
ldiv!(t.Eto, t.FT, t.Eωo) # transform (ω, ky, kx) -> (t, y, x)
Et_to_Pt!(t.Pto, t.Ptobuf, t.Eto, t.resp, t.densityfun(z), t.idcs) # add up responses
Et_to_Pt!(t.Pto, t.Eto, t.resp, t.densityfun(z), t.idcs) # add up responses
@. t.Pto *= t.grid.towin # apodisation
mul!(t.Pωo, t.FT, t.Pto) # transform (t, y, x) -> (ω, ky, kx)
copy_scale!(nl, t.Pωo, length(t.grid.ω), 1/t.scale)
Expand Down

0 comments on commit 380cb30

Please sign in to comment.