From 6f3b6b49d33571ee45669a21e2d1365ec8b2fca8 Mon Sep 17 00:00:00 2001
From: Ashley Milsted <ashmilsted@gmail.com>
Date: Wed, 10 Jul 2024 10:39:05 -0700
Subject: [PATCH] Don't tuplify time-dependent ops if types are homogenous
 (#401)

---
 src/time_dependent_operators.jl  | 52 ++++++++++++++++++++++++--------
 test/test_timeevolution_tdops.jl | 35 +++++++++++++++++++++
 2 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/src/time_dependent_operators.jl b/src/time_dependent_operators.jl
index e94f8133..024a2bb1 100644
--- a/src/time_dependent_operators.jl
+++ b/src/time_dependent_operators.jl
@@ -1,7 +1,21 @@
 # Convert storage of heterogeneous stuff to tuples for maximal compilation
 # and to avoid runtime dispatch.
-_tuplify(o::TimeDependentSum) = TimeDependentSum(Tuple, o)
-_tuplify(o::LazySum) = LazySum(eltype(o.factors), o.factors, (o.operators...,))
+function _tuplify(o::TimeDependentSum)
+    if isconcretetype(eltype(o.coefficients)) && isconcretetype(eltype(o.static_op.operators))
+        # No need to tuplify is types are concrete.
+        # We will save on compile time this way.
+        return o
+    end
+    return TimeDependentSum(Tuple, o)
+end
+function _tuplify(o::LazySum)
+    if isconcretetype(eltype(o.factors)) && isconcretetype(eltype(o.operators))
+        return o
+    end
+    return LazySum(eltype(o.factors), o.factors, (o.operators...,))
+end
+_tuplify(o::AbstractVector{T}) where T = isconcretetype(T) ? o : (o...,)
+_tuplify(o::Tuple) = o
 _tuplify(o::AbstractOperator) = o
 
 """
@@ -23,6 +37,7 @@ function _tdopdagger(o::TimeDependentSum)
     # that requires that the original operator sticks around and is always
     # updated first (though this is checked).
     # Copies and conjugates the coefficients from the original op.
+    # TODO: Make an Adjoint wrapper for TimeDependentSum instead?
     o_ls = QuantumOpticsBase.static_operator(o)
     facs = o_ls.factors
     c1 = (t)->(@assert current_time(o) == t; conj(facs[1]))
@@ -43,13 +58,18 @@ operators.
 """
 function master_h_dynamic_function(H::AbstractTimeDependentOperator, Js)
     Htup = _tuplify(H)
-    Js_tup = ((_tuplify(J) for J in Js)...,)
-
-    Jdags_tup = _tdopdagger.(Js_tup)
-    function _getfunc(Hop, Jops, Jdops)
-        return (@inline _tdop_master_wrapper_1(t, _) = (set_time!(Hop, t), set_time!.(Jops, t), set_time!.(Jdops, t)))
+    Js_tup = _tuplify(map(_tuplify, Js))
+    Jdags_tup = map(_tdopdagger, Js_tup)
+
+    return let Hop = Htup, Jops = Js_tup, Jdops = Jdags_tup
+        function _tdop_master_wrapper_1(t, _)
+            f = Base.Fix2(set_time!, t)
+            foreach(f, Jops)
+            foreach(f, Jdops)
+            set_time!(Hop, t)
+            return Hop, Jops, Jdops
+        end
     end
-    return _getfunc(Htup, Js_tup, Jdags_tup)
 end
 
 """
@@ -64,15 +84,21 @@ where `Hnh` is represents the non-Hermitian Hamiltonian and `Js` are the
 """
 function master_nh_dynamic_function(Hnh::AbstractTimeDependentOperator, Js)
     Hnhtup = _tuplify(Hnh)
-    Js_tup = ((_tuplify(J) for J in Js)...,)
+    Js_tup = _tuplify(map(_tuplify, Js))
 
-    Jdags_tup = _tdopdagger.(Js_tup)
+    Jdags_tup = map(_tdopdagger, Js_tup)
     Htdagup = _tdopdagger(Hnhtup)
 
-    function _getfunc(Hop, Hdop, Jops, Jdops)
-        return (@inline _tdop_master_wrapper_2(t, _) = (set_time!(Hop, t), set_time!(Hdop, t), set_time!.(Jops, t), set_time!.(Jdops, t)))
+    return let Hop = Hnhtup, Hdop = Htdagup, Jops = Js_tup, Jdops = Jdags_tup
+        function _tdop_master_wrapper_2(t, _)
+            f = Base.Fix2(set_time!, t)
+            foreach(f, Jops)
+            foreach(f, Jdops)
+            set_time!(Hop, t)
+            set_time!(Hdop, t)
+            return Hop, Hdop, Jops, Jdops
+        end
     end
-    return _getfunc(Hnhtup, Htdagup, Js_tup, Jdags_tup)
 end
 
 """
diff --git a/test/test_timeevolution_tdops.jl b/test/test_timeevolution_tdops.jl
index 1d45faad..a3884857 100644
--- a/test/test_timeevolution_tdops.jl
+++ b/test/test_timeevolution_tdops.jl
@@ -1,6 +1,13 @@
 using Test
 using QuantumOptics
 
+function test_settime(op)
+    t = current_time(op)
+    set_time!(op, randn())
+    set_time!(op, t)
+    return nothing
+end
+
 @testset "time-dependent operators" begin
 
 b = FockBasis(7)
@@ -9,8 +16,26 @@ a = destroy(b)
 
 H0 = number(b)
 Hd = (a + a')
+
+# function and op types homogeneous
+H = TimeDependentSum(cos=>H0, cos=>Hd)
+Htup = timeevolution._tuplify(H)
+@test Htup === H
+test_settime(Htup)
+@test (@allocated(test_settime)) == 0
+
+# op types not homogeneous
+H = TimeDependentSum(cos=>H0, cos=>dense(Hd), cos=>LazySum(H0), cos=>LazySum(dense(Hd)))
+Htup = timeevolution._tuplify(H)
+@test Htup !== H
+test_settime(Htup)
+@test (@allocated(test_settime)) == 0
+
 H = TimeDependentSum(1.0=>H0, cos=>Hd)
 
+# function types not homogeneous
+@test timeevolution._tuplify(H) !== H
+
 ts = [0.0, 0.4]
 ts_half = 0.5 * ts
 
@@ -54,6 +79,8 @@ ts_out, rhos = timeevolution.master_dynamic(ts, psi0, H, Js)
 ts_out2, rhos2 = timeevolution.master_dynamic(ts, psi0, fman)
 @test rhos[end].data ≈ rhos2[end].data
 
+set_time!(H, 0.0)
+set_time!.(Js, 0.0)
 Hnh = H - 0.5im * sum(J' * J for J in Js)
 
 _getf = (H0, Hd, a) -> (t,_) -> (
@@ -90,4 +117,12 @@ allocs1 = @allocated timeevolution.master_nh_dynamic(ts, psi0, Hnh, Js)
 allocs2 = @allocated timeevolution.master_nh_dynamic(ts_half, psi0, Hnh, Js)
 @test allocs1 == allocs2
 
+Jstup = (Js...,)
+ts_out2, rhos2 = timeevolution.master_nh_dynamic(ts, psi0, Hnh, Jstup)
+@test rhos[end].data ≈ rhos2[end].data
+
+allocs1 = @allocated timeevolution.master_nh_dynamic(ts, psi0, Hnh, Jstup)
+allocs2 = @allocated timeevolution.master_nh_dynamic(ts_half, psi0, Hnh, Jstup)
+@test allocs1 == allocs2
+
 end