From 0f68f17994823597b473fed193d7eb3ac19704b4 Mon Sep 17 00:00:00 2001
From: Ashley Milsted <ammilste@amazon.com>
Date: Thu, 4 Jul 2024 11:52:09 -0700
Subject: [PATCH] Don't tuplify if types are homogenous.

---
 src/time_dependent_operators.jl  | 17 +++++++++++++++--
 test/test_timeevolution_tdops.jl | 14 ++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/time_dependent_operators.jl b/src/time_dependent_operators.jl
index e94f8133..6a22d8db 100644
--- a/src/time_dependent_operators.jl
+++ b/src/time_dependent_operators.jl
@@ -1,7 +1,19 @@
 # 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::AbstractOperator) = o
 
 """
@@ -23,6 +35,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]))
diff --git a/test/test_timeevolution_tdops.jl b/test/test_timeevolution_tdops.jl
index 1d45faad..b8efc4ce 100644
--- a/test/test_timeevolution_tdops.jl
+++ b/test/test_timeevolution_tdops.jl
@@ -9,8 +9,20 @@ a = destroy(b)
 
 H0 = number(b)
 Hd = (a + a')
+
+# function and op types homogeneous
+H = TimeDependentSum(cos=>H0, cos=>Hd)
+@test timeevolution._tuplify(H) === H
+
+# op types not homogeneous
+H = TimeDependentSum(cos=>H0, cos=>dense(Hd))
+@test timeevolution._tuplify(H) !== H
+
 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 +66,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,_) -> (