Skip to content

Commit

Permalink
docs corrected
Browse files Browse the repository at this point in the history
  • Loading branch information
pankajpopli committed Jul 12, 2023
1 parent 270c41d commit 86994e6
Show file tree
Hide file tree
Showing 6 changed files with 318 additions and 0 deletions.
1 change: 1 addition & 0 deletions docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ pages = [
"GPU" => "gpu.md",
"Examples" => [
"literated/OneDShallowWaterGeostrophicAdjustment.md",
"literated/XYAdvectionSetupRun.md"
],
"Contributor's guide" => "contributing.md",
"Library" => [
Expand Down
8 changes: 8 additions & 0 deletions docs/src/library/internals.md
Original file line number Diff line number Diff line change
Expand Up @@ -65,3 +65,11 @@ Modules = [FourierFlows.Diffusion]
Public = false
Pages = ["diffusion.jl"]
```

## XYAdvection Module

```@autodocs
Modules = [FourierFlows.XYAdvection]
Public = false
Pages = ["XYAdvection.jl"]
```
8 changes: 8 additions & 0 deletions docs/src/library/public.md
Original file line number Diff line number Diff line change
Expand Up @@ -67,3 +67,11 @@ Modules = [FourierFlows.Diffusion]
Private = false
Pages = ["diffusion.jl"]
```

## XYAdvection Module

```@autodocs
Modules = [FourierFlows.XYAdvection]
Private = false
Pages = ["XYAdvection.jl"]
```
48 changes: 48 additions & 0 deletions examples/XYAdvectionSetupRun.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
using FourierFlows, JLD2
using FourierFlows:XYAdvection as xy
using LinearAlgebra: mul!, ldiv!
using Plots

dev = CPU()
Nx = 80 # grid resolution
Lx = 80 # box size
dt = 0.02 # timestep (s)
nsteps = 100 # total number of time-steps
printFreq = 1 # data print frequency

stepper = "ForwardEuler" # timestepper_Scheme


J = 1.0 # diffusion coefficient
A = 0.0 # advection coefficient
Γ = 1.0 # GL coefficient
α = 1.0 # average order parameter


grid = xy.set_2Dgrid(Nx,Lx; dev,aliased_fraction = 0)
params = xy.set_Params(;J=J, A=A, Γ=Γ, α=α)
vars = xy.set_Vars(grid)
prob = xy.set_Problem(grid,params,vars,dt,stepper)


px = rand(Nx,Nx)
py = rand(Nx,Nx)
xy.set_initialConditions!(prob,px,py)
heatmap(grid.x,grid.y,atan.(py,px)',clims=(-π,π),c=:hsv,axis=true,grid=false,colorbar=true,ratio=1,size=(400,400))


out = xy.run_and_save(prob,nsteps,printFreq)
path=out.path

# plot order parameter time series
orderP = Float64[]
for titr in 0:printFreq:nsteps
px,py = xy.get_data(titr,path,Nx)
mag = ((sum(px)^2 + sum(py)^2)^0.5)/(Nx*Nx)
push!(orderP,mag)
end
plot(orderP,ylims=(0,1),ylabel="m", xlabel="t", size=(400,400))

titr = nsteps
px,py = xy.get_data(titr,path,Nx)
heatmap(grid.x,grid.y,atan.(py,px)',clims=(-π,π),c=:hsv,axis=true,grid=false,colorbar=true,ratio=1,size=(400,400))
1 change: 1 addition & 0 deletions src/FourierFlows.jl
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ include("timesteppers.jl")

# Physics
include("diffusion.jl")
include("XYAdvection.jl")


function __init__()
Expand Down
252 changes: 252 additions & 0 deletions src/XYAdvection.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,252 @@
"""
# Exports
$(EXPORTS)
"""
module XYAdvection

export set_2Dgrid, set_Params, set_Vars, set_Problem, set_initialConditions!,
run_and_save, get_data, updatevars!, stepforward!, saveoutput

using DocStringExtensions

using FourierFlows, JLD2
using LinearAlgebra: mul!, ldiv!

"""
struct Params{T} <: AbstractParams
The parameters for XYAdvection problem:
$(FIELDS)
"""
struct Params{T} <: AbstractParams
"Diffusion coefficient"
J :: T
"Advection coefficient"
A :: T
"LG coefficient"
Γ :: T
"Average order parameter"
α :: T
end

"""
struct Vars2D{Aphys, Atrans} <: AbstractVars
The variables for XYAdvection problem:
$(FIELDS)
"""
struct Vars2D{Aphys, Atrans} <: AbstractVars
"px and py field and its x and y derivative pxx, pxy, pyx, pyy"
px :: Aphys
py :: Aphys

pxx :: Aphys
pxy :: Aphys
pyx :: Aphys
pyy :: Aphys

"px_hat, py_hat; FFT of px, py and its derivative fields"
pxh :: Atrans
pyh :: Atrans
pxxh :: Atrans
pxyh :: Atrans
pyxh :: Atrans
pyyh :: Atrans
end

"""
Vars2D(grid)
Return the variables `vars` for a XYAdvection problem on `grid`.
$(TYPEDFIELDS)
"""
function Vars2D(grid::TwoDGrid{T}) where T
Dev = typeof(grid.device)

@devzeros Dev T (grid.nx, grid.ny) px py pxx pxy pyx pyy
@devzeros Dev Complex{T} (grid.nkr, grid.nl) pxh pyh pxxh pxyh pyxh pyyh

return Vars2D(px, py, pxx, pxy, pyx, pyy, pxh, pyh, pxxh, pxyh, pyxh, pyyh)
end

"""
calcN!(N, sol, t, clock, vars, params, grid)
Calculate the nonlinear term for the XYAdvection equation.
"""
function calcN!(N, sol, t, clock, vars, params, grid)
# multiply p__h with ik to get derivatives
@. vars.pxxh = im * grid.kr .* sol[:,:,1]
@. vars.pxyh = im * grid.l .* sol[:,:,1]

@. vars.pyxh = im * grid.kr .* sol[:,:,2]
@. vars.pyyh = im * grid.l .* sol[:,:,2]

# get ik*p__h in physical space
ldiv!(vars.pxx, grid.rfftplan, vars.pxxh) # destroys vars.pxxh when using fftw
ldiv!(vars.pxy, grid.rfftplan, vars.pxyh) # destroys vars.pxyh when using fftw
ldiv!(vars.pyx, grid.rfftplan, vars.pyxh) # destroys vars.pyxh when using fftw
ldiv!(vars.pyy, grid.rfftplan, vars.pyyh) # destroys vars.pyyh when using fftw

# non-linear term
@. vars.pxx = params.A * ((vars.px * vars.pxx) + (vars.py * vars.pxy)) + params.Γ * (params.α - (vars.px^2 + vars.py^2))*vars.px
@. vars.pyx = params.A * ((vars.px * vars.pyx) + (vars.py * vars.pyy)) + params.Γ * (params.α - (vars.px^2 + vars.py^2))*vars.py

# go to fourier space and define N
mul!(vars.pxxh, grid.rfftplan, vars.pxx)
mul!(vars.pyxh, grid.rfftplan, vars.pyx)
N[:,:, 1] = vars.pxxh
N[:,:, 2] = vars.pyxh

dealias!(N[:,:,1], grid)
dealias!(N[:,:,2], grid)
return nothing
end

"""
Equation(params,grid)
"""
function Equation(params::Params, grid::TwoDGrid)
dev = grid.device

# Linear operator
L = zeros(dev, eltype(grid), (grid.nkr, grid.nl,2))
@. L[:,:,1] = - params.J * grid.kr^2 - params.J * grid.l^2
@. L[:,:,2] = - params.J * grid.kr^2 - params.J * grid.l^2

# full equation
return FourierFlows.Equation(L, calcN!, grid)
end

#################################### Exported Functions #################################
"""
To be called from main
set_2Dgrid(nx::Int64,Lx; dev::Device=CPU(),aliased_fraction = 0,kwargs...)
setup 2D grid given the parameters
"""
function set_2Dgrid(nx::Int64,Lx; dev::Device=CPU(),aliased_fraction = 0,kwargs...)
grid = TwoDGrid(dev; nx, Lx, aliased_fraction,kwargs...)
return grid
end

"""
To be called from main
set_Params(;J::T,A::T,Γ::T,α::T)
setup parameter values for the problem
"""
function set_Params(;J::T,A::T::T::T) where {T<:Float64}
params = Params(J, A, Γ, α)
return params
end

"""
To be called from main
set_Vars(grid::TwoDGrid)
setup variables for the system
"""
function set_Vars(grid::TwoDGrid)
vars = Vars2D(grid)
return vars
end


"""
To be called from main
set_Problem(grid::TwoDGrid,params::Params,vars::Vars2D,dt::Float64=0.02,stepper = "ForwardEuler";stepperkwargs...)
setup the FourierFlows.Problem
"""
function set_Problem(grid::TwoDGrid,params::Params,vars::Vars2D,dt::Float64=0.02,stepper = "ForwardEuler";stepperkwargs...)

equation = Equation(params,grid)

prob = FourierFlows.Problem(equation, stepper, dt, grid, vars, params; stepperkwargs...)
return prob
end

"""
updatevars!(prob)
"""
function updatevars!(prob)
vars, grid, sol = prob.vars, prob.grid, prob.sol

@. vars.pxh = sol[:,:, 1]
@. vars.pyh = sol[:,:, 2]

ldiv!(vars.px, grid.rfftplan, deepcopy(sol[:,:, 1])) # use deepcopy() because irfft destroys its input
ldiv!(vars.py, grid.rfftplan, deepcopy(sol[:,:, 2])) # use deepcopy() because irfft destroys its input
return nothing
end

"""
set_initialConditions!(prob, px, py)
Set the solution `sol` as the transform of `px` and `py` and update `vars`.
"""
function set_initialConditions!(prob, u, v)
vars, grid, sol = prob.vars, prob.grid, prob.sol

cast_type = typeof(vars.px) # determine the type of vars.px

# below, e.g., A(px0) converts px0 to the same type as vars expects
# (useful when px0 is a CPU array but grid.device is GPU)
mul!(vars.pxh, grid.rfftplan, cast_type(u))
mul!(vars.pyh, grid.rfftplan, cast_type(v))

@. sol[:,:,1] = vars.pxh
@. sol[:,:,2] = vars.pyh

updatevars!(prob)

return nothing
end

"""
To be called from main
run_and_save(prob,nsteps::T=10^5,printFreq::T=10^3;filepath::T2 = ".",filename::T2 = "XYAdvection_data.jld2") where {T<:Integer,T2<:String}
Run the problem and save output to file.
"""
function run_and_save(prob,nsteps::T=10^5,printFreq::T=10^3;filepath::T2 = ".",filename::T2 = "XYAdvection_data.jld2") where {T<:Integer,T2<:String}
# assert nsteps/printFreq is an integer
@assert isinteger(nsteps/printFreq) "requires nsteps/printFreq == Integer"

fname = joinpath(filepath, filename)
get_sol(prob) = prob.sol
out = Output(prob, fname, (:sol, get_sol))
saveproblem(out)
for i in 0:Int(nsteps)
if i % Int(printFreq)==0
saveoutput(out)
end
stepforward!(prob)
updatevars!(prob)
end
return out
end

"""
To be called from main
get_data(titr::Integer,filepath::String,nx::Integer)
Read output data from file and convert to physical space for analysis.
"""
function get_data(titr::Integer,filepath::String,nx::Integer)
file = jldopen(filepath)
px = irfft(file[string("snapshots/sol/", titr)][:,:, 1], nx)
py = irfft(file[string("snapshots/sol/", titr)][:,:, 2], nx)
return px,py
end
############################################################################################################
end #end module

0 comments on commit 86994e6

Please sign in to comment.