added: XY advection dynamics example

docs/make.jl

@@ -17,6 +17,7 @@ const OUTPUT_DIR   = joinpath(@__DIR__, "src/literated")
 
 examples = [
   "OneDShallowWaterGeostrophicAdjustment.jl",
+  "XYAdvectionSetupRun.jl",
 ]
 
 for example in examples
@@ -51,6 +52,7 @@ pages = [
     "GPU" => "gpu.md",
     "Examples" => [ 
       "literated/OneDShallowWaterGeostrophicAdjustment.md",
+      "literated/XYAdvectionSetupRun.md"
       ],
     "Contributor's guide" => "contributing.md",
     "Library" => [ 

docs/src/library/internals.md

@@ -65,3 +65,11 @@ Modules = [FourierFlows.Diffusion]
 Public  = false
 Pages   = ["diffusion.jl"]
 ```
+
+## XYAdvection Module
+
+```@autodocs
+Modules = [FourierFlows.XYAdvection]
+Public  = false
+Pages   = ["XYAdvection.jl"]
+```

docs/src/library/public.md

@@ -67,3 +67,11 @@ Modules = [FourierFlows.Diffusion]
 Private = false
 Pages   = ["diffusion.jl"]
 ```
+
+## XYAdvection Module
+
+```@autodocs
+Modules = [FourierFlows.XYAdvection]
+Private = false
+Pages   = ["XYAdvection.jl"]
+```

docs/src/references.bib

@@ -47,11 +47,31 @@ @article{Kassam-Trefethen-2005
 }
 
 @techreport{Carpenter-Kennedy-1994,
-  author = {Carpenter, M. H. and Kennedy, C. A.},
-  title = {Fourth-order $2N$-storage Runge--Kutta schemes},
-  institution = {National Aeronautics and Space Administration},
-  year = {1994},
-  number = {NASA TM-109112},
-  address = {NASA Langley Research Center, Hampton, VA},
-  url = {https://ntrs.nasa.gov/api/citations/19940028444/downloads/19940028444.pdf}
+    author = {Carpenter, M. H. and Kennedy, C. A.},
+    title = {Fourth-order $2N$-storage Runge--Kutta schemes},
+    institution = {National Aeronautics and Space Administration},
+    year = {1994},
+    number = {NASA TM-109112},
+    address = {NASA Langley Research Center, Hampton, VA},
+    url = {https://ntrs.nasa.gov/api/citations/19940028444/downloads/19940028444.pdf}
+}
+
+@article{Toner-Tu-1995,
+    title = {Long-Range order in a two-dimensional dynamical $\mathrm{XY}$ model: How birds fly together},
+    author = {Toner, J. and Tu, Y.},
+    journal = {Physical Review Letters},
+    volume = {75},
+    issue = {23},
+    pages = {4326-4329},
+    year = {1995},
+    doi = {10.1103/PhysRevLett.75.4326},
+}
+
+@book{Chaikin-Lubensky-1995,
+    title = {Principles of condensed matter physics},
+    author = {Chaikin, P. M. and Lubensky, T. C.},
+    publisher = {Cambridge University Press},
+    place = {Cambridge},
+    doi = {10.1017/CBO9780511813467},
+    year = {1995}
 }

examples/XYAdvectionSetupRun.jl

@@ -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))

src/FourierFlows.jl

@@ -109,6 +109,7 @@ include("timesteppers.jl")
 
 # Physics
 include("diffusion.jl")
+include("XYAdvection.jl")
 
 
 function __init__()

src/XYAdvection.jl

@@ -0,0 +1,253 @@
+"""
+# 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)
+      close(file)
+      return px,py
+   end
+   ############################################################################################################
+end #end module