[time_gf] Define basic arithmetic operators for all GF containers

kleinhenz · May 12, 2023 · b4c6108 · b4c6108
1 parent 8ad2def
commit b4c6108
Show file tree

Hide file tree

Showing 7 changed files with 238 additions and 55 deletions.
diff --git a/src/gf/full.jl b/src/gf/full.jl
@@ -82,6 +82,30 @@ function Base.zero(G::T) where T <: FullTimeGF
   zero(G, G.ξ)
 end
 
+function Base.:+(G1::T, G2::T) where T <: FullTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr + G2.gtr, G1.les + G2.les, G1.rm + G2.rm, G1.mat + G2.mat, G1.ξ)
+end
+
+function Base.:-(G1::T, G2::T) where T <: FullTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr - G2.gtr, G1.les - G2.les, G1.rm - G2.rm, G1.mat - G2.mat, G1.ξ)
+end
+
+function Base.:*(G::T, α::Number) where T <: FullTimeGF
+  return T(G.grid, G.gtr * α, G.les * α, G.rm * α, G.mat * α, G.ξ)
+end
+
+function Base.:*(α::Number, G::T) where T <: FullTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: FullTimeGF
+  return T(G.grid, -G.gtr, -G.les, -G.rm, -G.mat, G.ξ)
+end
+
 function TimeDomain(G::FullTimeGF)
   grid = G.grid
   tplus = grid[forward_branch]

diff --git a/src/gf/generic.jl b/src/gf/generic.jl
@@ -49,6 +49,28 @@ function Base.zero(G::T) where T <: GenericTimeGF
   T(G.grid, zero(G.data))
 end
 
+function Base.:+(G1::T, G2::T) where T <: GenericTimeGF
+  @assert G1.grid == G2.grid
+  return T(G1.grid, G1.data + G2.data)
+end
+
+function Base.:-(G1::T, G2::T) where T <: GenericTimeGF
+  @assert G1.grid == G2.grid
+  return T(G1.grid, G1.data - G2.data)
+end
+
+function Base.:*(G::T, α::Number) where T <: GenericTimeGF
+  return T(G.grid, G.data * α)
+end
+
+function Base.:*(α::Number, G::T) where T <: GenericTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: GenericTimeGF
+  return T(G.grid, -G.data)
+end
+
 function GenericTimeGF(dos::AbstractDOS, grid::FullTimeGrid)
   G = TimeInvariantFullTimeGF(dos, grid)
   GenericTimeGF(grid,1,true) do t1, t2

diff --git a/src/gf/imag.jl b/src/gf/imag.jl
@@ -54,6 +54,30 @@ function Base.zero(G::T) where T <: ImaginaryTimeGF
   zero(G, G.ξ)
 end
 
+function Base.:+(G1::T, G2::T) where T <: ImaginaryTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.mat + G2.mat, G1.ξ)
+end
+
+function Base.:-(G1::T, G2::T) where T <: ImaginaryTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.mat - G2.mat, G1.ξ)
+end
+
+function Base.:*(G::T, α::Number) where T <: ImaginaryTimeGF
+  return T(G.grid, G.mat * α, G.ξ)
+end
+
+function Base.:*(α::Number, G::T) where T <: ImaginaryTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: ImaginaryTimeGF
+  return T(G.grid, -G.mat, G.ξ)
+end
+
 function TimeDomain(G::ImaginaryTimeGF)
   grid = G.grid
   tau = grid[imaginary_branch]

diff --git a/src/gf/keldysh.jl b/src/gf/keldysh.jl
@@ -51,6 +51,30 @@ function Base.zero(G::T) where T <: KeldyshTimeGF
   zero(G, G.ξ)
 end
 
+function Base.:+(G1::T, G2::T) where T <: KeldyshTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr + G2.gtr, G1.les + G2.les, G1.ξ)
+end
+
+function Base.:-(G1::T, G2::T) where T <: KeldyshTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr - G2.gtr, G1.les - G2.les, G1.ξ)
+end
+
+function Base.:*(G::T, α::Number) where T <: KeldyshTimeGF
+  return T(G.grid, G.gtr * α, G.les * α, G.ξ)
+end
+
+function Base.:*(α::Number, G::T) where T <: KeldyshTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: KeldyshTimeGF
+  return T(G.grid, -G.gtr, -G.les, G.ξ)
+end
+
 function TimeDomain(G::KeldyshTimeGF)
   grid = G.grid
   tplus = grid[forward_branch]

diff --git a/src/gf/time_invariant_full.jl b/src/gf/time_invariant_full.jl
@@ -82,6 +82,30 @@ function Base.zero(G::T) where T <: TimeInvariantFullTimeGF
   zero(G, G.ξ)
 end
 
+function Base.:+(G1::T, G2::T) where T <: TimeInvariantFullTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr + G2.gtr, G1.les + G2.les, G1.rm + G2.rm, G1.mat + G2.mat, G1.ξ)
+end
+
+function Base.:-(G1::T, G2::T) where T <: TimeInvariantFullTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr - G2.gtr, G1.les - G2.les, G1.rm - G2.rm, G1.mat - G2.mat, G1.ξ)
+end
+
+function Base.:*(G::T, α::Number) where T <: TimeInvariantFullTimeGF
+  return T(G.grid, G.gtr * α, G.les * α, G.rm * α, G.mat * α, G.ξ)
+end
+
+function Base.:*(α::Number, G::T) where T <: TimeInvariantFullTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: TimeInvariantFullTimeGF
+  return T(G.grid, -G.gtr, -G.les, -G.rm, -G.mat, G.ξ)
+end
+
 function TimeDomain(G::TimeInvariantFullTimeGF)
   grid = G.grid
   tplus = grid[forward_branch]

diff --git a/src/gf/time_invariant_keldysh.jl b/src/gf/time_invariant_keldysh.jl
@@ -51,6 +51,30 @@ function Base.zero(G::T) where T <: TimeInvariantKeldyshTimeGF
   zero(G, G.ξ)
 end
 
+function Base.:+(G1::T, G2::T) where T <: TimeInvariantKeldyshTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr + G2.gtr, G1.les + G2.les, G1.ξ)
+end
+
+function Base.:-(G1::T, G2::T) where T <: TimeInvariantKeldyshTimeGF
+  @assert G1.grid == G2.grid
+  @assert G1.ξ == G2.ξ
+  return T(G1.grid, G1.gtr - G2.gtr, G1.les - G2.les, G1.ξ)
+end
+
+function Base.:*(G::T, α::Number) where T <: TimeInvariantKeldyshTimeGF
+  return T(G.grid, G.gtr * α, G.les * α, G.ξ)
+end
+
+function Base.:*(α::Number, G::T) where T <: TimeInvariantKeldyshTimeGF
+  return G * α
+end
+
+function Base.:-(G::T) where T <: TimeInvariantKeldyshTimeGF
+  return T(G.grid, -G.gtr, -G.les, G.ξ)
+end
+
 function TimeDomain(G::TimeInvariantKeldyshTimeGF)
   grid = G.grid
   tplus = grid[forward_branch]

diff --git a/test/time_gf.jl b/test/time_gf.jl
@@ -2,88 +2,129 @@ using Keldysh, Test
 
 @testset "time_gf" begin
 
-  let c = FullContour(tmax=2.0, β=5.0), nt = 21, ntau=51
-    ϵ = -1.0
-    f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, c.β)) * exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
+  function test_arithmetics(G::AbstractTimeGF, f::Function, grid::AbstractTimeGrid)
+    tt_range = Iterators.product(grid, grid)
+
+    A = G + G
+    @test all(map(((t1, t2),) -> isapprox(A[t1,t2], 2*f(t1, t2), atol=1e-10), tt_range))
+    B = A - G
+    @test all(map(((t1, t2),) -> isapprox(B[t1,t2], f(t1, t2), atol=1e-10), tt_range))
+    C = 3 * G
+    @test all(map(((t1, t2),) -> isapprox(C[t1,t2], 3*f(t1, t2), atol=1e-10), tt_range))
+    D = G * 3
+    @test all(map(((t1, t2),) -> isapprox(D[t1,t2], 3*f(t1, t2), atol=1e-10), tt_range))
+    E = -G
+    @test all(map(((t1, t2),) -> isapprox(E[t1,t2], -f(t1, t2), atol=1e-10), tt_range))
+  end
+
+  @testset "FullContour" begin
+    let c = FullContour(tmax=2.0, β=5.0), nt = 21, ntau=51
+
+      ϵ = -1.0
+      f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, c.β)) *
+        exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
 
-    grid = FullTimeGrid(c, nt, ntau)
+      grid = FullTimeGrid(c, nt, ntau)
 
-    G1 = GenericTimeGF(f, grid, 1, true)
-    G1_ = similar(G1)
-    G1_ = zero(G1)
+      G1 = GenericTimeGF(f, grid, 1, true)
+      G1_ = similar(G1)
+      G1_ = zero(G1)
+      test_arithmetics(G1, f, grid)
 
-    @test Keldysh.jump(G1) == 1.0im
+      @test Keldysh.jump(G1) == 1.0im
 
-    G2 = FullTimeGF(f, grid, 1, fermionic, true)
-    G2_ = similar(G2)
-    G2_ = zero(G2)
+      G2 = FullTimeGF(f, grid, 1, fermionic, true)
+      G2_ = similar(G2)
+      G2_ = zero(G2)
+      test_arithmetics(G2, f, grid)
 
-    @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10), Iterators.product(grid, grid)))
+      @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10),
+                    Iterators.product(grid, grid)))
 
-    G3 = TimeInvariantFullTimeGF(f, grid, 1, fermionic, true)
-    G3_ = similar(G3)
-    G3_ = zero(G3)
+      G3 = TimeInvariantFullTimeGF(f, grid, 1, fermionic, true)
+      G3_ = similar(G3)
+      G3_ = zero(G3)
+      test_arithmetics(G3, f, grid)
 
-    @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G3[t1,t2], atol=1e-10), Iterators.product(grid, grid)))
+      @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G3[t1,t2], atol=1e-10),
+                    Iterators.product(grid, grid)))
 
-    # test interpolation on linear function
-    lin_f = (t1, t2) -> (2.0 * t1.val - im * t2.val) + 5.0 * heaviside(t1, t2)
-    for t1 in grid
-      for t2 in grid
-        G1[t1, t2] = lin_f(t1.bpoint, t2.bpoint)
+      # test interpolation on linear function
+      lin_f = (t1, t2) -> (2.0 * t1.val - im * t2.val) + 5.0 * heaviside(t1, t2)
+      for t1 in grid
+        for t2 in grid
+          G1[t1, t2] = lin_f(t1.bpoint, t2.bpoint)
+        end
       end
-    end
 
-    grid_fine = FullTimeGrid(c, 41, 101)
+      grid_fine = FullTimeGrid(c, 41, 101)
 
-    @test all(map(Iterators.product(grid_fine, grid_fine)) do (t1, t2)
-                isapprox(G1(t1.bpoint, t2.bpoint), lin_f(t1.bpoint, t2.bpoint), atol=1e-10)
-              end
-             )
+      @test all(map(Iterators.product(grid_fine, grid_fine)) do (t1, t2)
+                  isapprox(G1(t1.bpoint, t2.bpoint), lin_f(t1.bpoint, t2.bpoint), atol=1e-10)
+                end
+               )
+    end
   end
 
-  let c = KeldyshContour(tmax=2.0), nt = 21
+  @testset "KeldyshContour" begin
+    let c = KeldyshContour(tmax=2.0), nt = 21
 
-    ϵ = -1.0
-    β = 5.0
-    f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, β)) * exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
+      ϵ = -1.0
+      β = 5.0
+      f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, β)) *
+        exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
 
-    grid = KeldyshTimeGrid(c, nt)
+      grid = KeldyshTimeGrid(c, nt)
 
-    G1 = GenericTimeGF(f, grid, 1, true)
-    G1_ = similar(G1)
-    G1_ = zero(G1)
+      G1 = GenericTimeGF(f, grid, 1, true)
+      G1_ = similar(G1)
+      G1_ = zero(G1)
+      test_arithmetics(G1, f, grid)
 
-    @test Keldysh.jump(G1) == 1.0im
+      @test Keldysh.jump(G1) == 1.0im
 
-    G2 = KeldyshTimeGF(f, grid, 1, fermionic, true)
-    G2_ = similar(G2)
-    G2_ = zero(G2)
+      G2 = KeldyshTimeGF(f, grid, 1, fermionic, true)
+      G2_ = similar(G2)
+      G2_ = zero(G2)
+      test_arithmetics(G2, f, grid)
 
-    @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10), Iterators.product(grid, grid)))
+      @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10),
+                    Iterators.product(grid, grid)))
 
-    G3 = TimeInvariantKeldyshTimeGF(f, grid, 1, fermionic, true)
-    G3_ = similar(G3)
-    G3_ = zero(G3)
+      G3 = TimeInvariantKeldyshTimeGF(f, grid, 1, fermionic, true)
+      G3_ = similar(G3)
+      G3_ = zero(G3)
+      test_arithmetics(G3, f, grid)
 
-    @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G3[t1,t2], atol=1e-10), Iterators.product(grid, grid)))
+      @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G3[t1,t2], atol=1e-10),
+                    Iterators.product(grid, grid)))
 
+    end
   end
 
-  let c = ImaginaryContour(β=5.0), ntau=51
-    ϵ = -1.0
-    f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, c.β)) * exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
-    grid = ImaginaryTimeGrid(c, ntau)
+  @testset "ImaginaryContour" begin
+    let c = ImaginaryContour(β=5.0), ntau=51
+
+      ϵ = -1.0
+      f = (t1, t2) -> -1.0im * (heaviside(t1.bpoint, t2.bpoint) - fermi(ϵ, c.β)) *
+        exp(-1.0im * (t1.bpoint.val - t2.bpoint.val) * ϵ)
+
+      grid = ImaginaryTimeGrid(c, ntau)
 
-    G1 = GenericTimeGF(f, grid, 1, true)
-    G1_ = similar(G1)
-    G1_ = zero(G1)
+      G1 = GenericTimeGF(f, grid, 1, true)
+      G1_ = similar(G1)
+      G1_ = zero(G1)
+      test_arithmetics(G1, f, grid)
 
-    G2 = ImaginaryTimeGF(f, grid, 1, fermionic, true)
-    G2_ = similar(G2)
-    G2_ = zero(G2)
-    @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10), Iterators.product(grid, grid)))
+      G2 = ImaginaryTimeGF(f, grid, 1, fermionic, true)
+      G2_ = similar(G2)
+      G2_ = zero(G2)
+      test_arithmetics(G2, f, grid)
 
+      @test all(map(((t1, t2),) -> isapprox(G1[t1,t2], G2[t1,t2], atol=1e-10),
+                    Iterators.product(grid, grid)))
+
+    end
   end
 
 end