added XY advection dynamics example

examples/XYAdvectionSetupRun.jl

@@ -0,0 +1,90 @@
+# """
+# # XY dynamics with advected order parameter
+
+# This example solves the 2D vectorial field order parameter governed by XY dynamics and advection.
+# for ``px(x, y, t)``, ``py(x,y, t)``
+
+# The equations for the evolution of ``px(x,y,t)``, ``py(x,y,t)`` are:
+# ```math
+# \begin{aligned}
+# \partial_t \vec{p}  & = J~\nabla^2\vec{p} +  A (\vec{p}\cdot\nabla)\vec{p} + \Gamma~(op-|\vec{p}|^2)\vec{p}\\
+# \end{aligned}
+# ```
+
+# Or, in component,
+
+# ```math
+# \begin{aligned}
+# \partial_t px  & = J~\nabla^2 px +  A (px\partial_x~px + py\partial_y~px) + \Gamma~(op-|\mathbf{p}|^2)px\\
+# \partial_t py  & = J~\nabla^2 py +  A (px\partial_x~py + py\partial_y~py) + \Gamma~(op-|\mathbf{p}|^2)py\\
+# \end{aligned}
+# ```
+# """
+
+"""
+calling required libraries
+"""
+using FourierFlows, JLD2
+using LinearAlgebra: mul!, ldiv!
+include("XYAdvection.jl")
+using .XYAdvection
+using Plots
+
+
+"""
+define system and parameters
+"""
+dev = CPU()
+Nx = 80       # grid resolution
+Lx = 80       # box size
+dt = 0.02     # timestep (s)
+nsteps = 10   # total number of time-steps
+printFreq = 1 # data print frequency
+
+J = 1.0 
+A = 0.0 
+Γ = 1.0 
+op = 1.0
+stepper = "ForwardEuler"  # timestepper_Scheme
+
+
+
+"""
+setup grid, parameters and the problem
+"""
+grid   = set_2Dgrid(Nx,Lx; dev,aliased_fraction = 0)
+params = set_Params(;J=J, A=A, Γ=Γ, op=op)
+vars   = set_Vars(grid)
+prob   = set_Problem(grid,params,vars,dt,stepper)
+
+"""
+set initial conditions
+"""
+u = rand(Nx,Nx)
+v = rand(Nx,Nx)
+set_initialConditions!(prob,u,v)
+
+"""
+check initial condition
+"""
+heatmap(grid.x,grid.y,atan.(v,u)',clims=(-π,π),c=:hsv,axis=true,grid=false,colorbar=true,ratio=1,size=(400,400))
+
+
+"""
+run the problem
+"""
+out = run_and_save(prob,nsteps,printFreq)
+path=out.path
+
+
+"""
+read and plot data
+"""
+# plot order parameter time series
+orderP = Float64[]
+for titr in 0:printFreq:nsteps
+   u,v = get_data(titr,path,Nx)
+   mag = ((sum(u)^2 + sum(v)^2)^0.5)/(Nx*Nx)
+   push!(orderP,mag)
+end
+plot(orderP,marker=:circle,ylims=(0,1))

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,254 @@
+"""
+A test-bed module that solves the 1D diffusion equation.
+
+# 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 diffusion problem:
+
+   $(TYPEDFIELDS)
+   """
+   struct Params{T} <: AbstractParams
+      J  :: T  # Diffusion coefficient
+      A  :: T  # Advection coefficient
+      Γ  :: T  # LG coefficient
+      op :: T  # Average order parameter
+   end
+
+   """
+      struct Vars2D{Aphys, Atrans} <: AbstractVars
+
+   The variables for diffusion 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
+
+   """
+   Vars(grid)
+
+   Return the variables `vars` for a constant diffusivity problem on `grid`.
+   """
+   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 1D diffusion 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.op - (vars.px^2 + vars.py^2))*vars.px
+      @. vars.pyx = params.A * ((vars.px * vars.pyx) + (vars.py * vars.pyy)) + params.Γ * (params.op - (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(grid, params)
+
+   Return the equation for a constant diffusivity problem on `grid` with diffusivity
+   found inside `params`.
+   """
+   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(nx::Int64,Lx; dev::Device=CPU(),aliased_fraction = 0,kwargs...)
+
+   setup parameter values for the problem
+   """
+   function set_Params(;J::T,A::T,Γ::T,op::T) where {T<:Float64}
+      params = Params(J, A, Γ, op)
+      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!(vars, grid, sol)
+
+   Update the variables in `vars` on the `grid` with the solution in `sol`.
+   """
+   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, c)
+
+   Set the solution `sol` as the transform of `c` 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)
+      u = irfft(file[string("snapshots/sol/", titr)][:,:, 1], nx)
+      v = irfft(file[string("snapshots/sol/", titr)][:,:, 2], nx)
+      return u,v
+   end
+   ############################################################################################################
+end #end module