Merge pull request #342 from xtalax/complex2

Complex Number Support 2

src/interface/solution/MOLMetadata.jl

@@ -9,13 +9,14 @@ Used to unpack the solution.
           MOLFiniteDifference object.
 - `pdesys`: a PDESystem object, used in the discretization.
 """
-struct MOLMetadata{hasTime,Ds,Disc,PDE,M,Strat} <: SciMLBase.AbstractDiscretizationMetadata{hasTime}
+struct MOLMetadata{hasTime,Ds,Disc,PDE,M,C,Strat} <: SciMLBase.AbstractDiscretizationMetadata{hasTime}
     discretespace::Ds
     disc::Disc
     pdesys::PDE
     use_ODAE::Bool
     metadata::M
-    function MOLMetadata(discretespace, disc, pdesys, boundarymap, metadata=nothing)
+    complexmap::C
+    function MOLMetadata(discretespace, disc, pdesys, boundarymap, complexmap, metadata=nothing)
         metaref = Ref{Any}()
         metaref[] = metadata
         if discretespace.time isa Nothing
@@ -35,14 +36,14 @@ struct MOLMetadata{hasTime,Ds,Disc,PDE,M,Strat} <: SciMLBase.AbstractDiscretizat
         end
         return new{hasTime,typeof(discretespace),
             typeof(disc),typeof(pdesys),
-            typeof(metaref),typeof(disc.disc_strategy)}(discretespace,
+            typeof(metaref),typeof(complexmap),typeof(disc.disc_strategy)}(discretespace,
             disc, pdesys, use_ODAE,
-            metaref)
+            metaref, complexmap)
     end
 end
 
-function PDEBase.generate_metadata(s::DiscreteSpace, disc::MOLFiniteDifference, pdesys::PDESystem, boundarymap, metadata=nothing)
-    return MOLMetadata(s, disc, pdesys, boundarymap, metadata)
+function PDEBase.generate_metadata(s::DiscreteSpace, disc::MOLFiniteDifference, pdesys::PDESystem, boundarymap, complexmap, metadata=nothing)
+    return MOLMetadata(s, disc, pdesys, boundarymap, complexmap, metadata)
 end
 
 # function PDEBase.generate_metadata(s::DiscreteSpace, disc::MOLFiniteDifference{G,D}, pdesys::PDESystem, boundarymap, metadata=nothing) where {G<:StaggeredGrid}

src/interface/solution/common.jl

@@ -23,7 +23,13 @@ function (sol::SciMLBase.PDESolution{T,N,S,D})(args...;
             sol.interp[dv](args[is]...)
         end
     end
-    return sol.interp[dv](args...)
+    if iscomplex(sol) && !any(isequal(safe_unwrap(dv)), sol.dvs)
+        symargs = arguments(safe_unwrap(dv))
+        redv, imdv = sol.disc_data_complexmap[dv]
+        return sol.interp[Num(redv(symargs...))](args...) .+ im * sol.interp[Num(imdv(symargs...))](args...)
+    else
+        return sol.interp[dv](args...)
+    end
 end
 
 Base.@propagate_inbounds function Base.getindex(A::SciMLBase.PDESolution{T,N,S,D},
@@ -43,7 +49,11 @@ Base.@propagate_inbounds function Base.getindex(A::SciMLBase.PDESolution{T,N,S,D
     if SciMLBase.issymbollike(sym) && iv !== nothing && isequal(sym, iv)
         A.ivdomain[iiv]
     elseif SciMLBase.issymbollike(sym) && dv !== nothing && isequal(sym, dv)
-        A.u[sym]
+            A.u[sym]
+    elseif iscomplex(A) && istree(safe_unwrap(sym))
+        symargs = arguments(safe_unwrap(sym))
+        redv, imdv = A.disc_data.complexmap[operation(safe_unwrap(sym))]
+        A.u[Num(redv(symargs...))] .+ im * A.u[Num(imdv(symargs...))]
     else
         error("Invalid indexing of solution. $sym not found in solution.")
     end
@@ -65,6 +75,11 @@ Base.@propagate_inbounds function Base.getindex(A::SciMLBase.PDESolution{T,N,S,D
         A.ivdomains[iiv][args...]
     elseif SciMLBase.issymbollike(sym) && dv !== nothing && isequal(sym, dv)
         A.u[sym][args...]
+    end
+    if iscomplex(A) && SciMlBase.issymbollike(sym) && istree(safe_unwrap(sym))
+        symargs = arguments(safe_unwrap(sym))
+        redv, imdv = A.disc_data.complexmap[operation(safe_unwrap(sym))]
+        A.u[Num(redv(symargs...))][args...] .+ im * A.u[Num(imdv(symargs...))][args...]
     else
         error("Invalid indexing of solution")
     end

src/interface/solution/solution_utils.jl

@@ -29,3 +29,5 @@ end
 # end
 
 sym_to_index(sym, syms) = findfirst(isequal(sym), syms)
+
+iscomplex(A::SciMLBase.AbstractPDESolution) = !isnothing(A.disc_data.complexmap)

src/interface/solution/timedep.jl

@@ -61,6 +61,7 @@ function SciMLBase.PDETimeSeriesSolution(sol::SciMLBase.AbstractODESolution{T},
         end |> Dict
         # Build Interpolations
         interp = build_interpolation(umap, dvs, ivs, ivgrid, sol, pdesys, discretespace.vars.replaced_vars)
+        @show metadata.complexmap
 
         return SciMLBase.PDETimeSeriesSolution{T,length(discretespace.ū),typeof(umap),typeof(metadata),
             typeof(sol),typeof(sol.errors),typeof(sol.t),typeof(ivgrid),

src/system_parsing/pde_system_transformation.jl

@@ -254,7 +254,7 @@ function create_aux_variable!(eqs, bcs, boundarymap, v, term)
     append!(bcs, map(bc -> bc.eq, newbcs))
     # Add the new boundary conditions and initial conditions to the boundarymap
     update_boundarymap!(boundarymap, newbcs, newop, v)
-    # update pmap
+    # update pmap 
 end
 
 function generate_bc_rules(bcs, v)

test/pde_systems/MOLtest2.jl

@@ -474,4 +474,38 @@ end
     prob = discretize(sys, discretization) # ERROR HERE
 
     sol = solve(prob, FBDF(), saveat = s_in_y)
+end
+
+
+@testset "Schroedinger Equation" begin
+    @parameters x z
+    @variables A(..)
+    Dx = Differential(x)
+    Dz = Differential(z)
+    Dzz = Differential(z)^2
+
+    xmin = 0.0
+    xmax = 1e-1
+    zmax = 10.0
+    zmin = -zmax
+
+    c0 = 1.0
+    A0(x, z) = c0 * sech(c0 * z / sqrt(2)) * exp(im * c0^2 * x / 2)
+
+    domains = [x ∈ Interval(xmin, xmax), z ∈ Interval(zmin, zmax)]
+
+    eq = [im * Dx(A(x, z)) + Dzz(A(x, z)) ~ -abs2(A(x, z)) * A(x, z)]
+    bcs = [A(xmin, z) ~ A0(xmin, z),
+        A(x, zmin) ~ 0,
+        A(x, zmax) ~ 0]
+
+    @named pdesys = PDESystem(eq, bcs, domains, [x, z], [A(x, z)])
+
+    N = 100
+    dz = 1 / N
+    order = 2
+
+    discretization = MOLFiniteDifference([z => dz], x)
+
+    @time prob = discretize(pdesys, discretization)
 end