From 1e0a4a1f604da6efd6c25b921c1eca92e291421e Mon Sep 17 00:00:00 2001
From: David Plankensteiner <david.plankensteiner@gmx.at>
Date: Sat, 18 Jan 2025 20:04:48 +0100
Subject: [PATCH] Add test for many-atom laser example

---
 src/indexing_sums.jl           |  6 +++
 src/printing.jl                |  2 +-
 test/test_indexed_meanfield.jl | 67 ++++++++++++++++++++++++++++++++--
 3 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/src/indexing_sums.jl b/src/indexing_sums.jl
index 3272fd8e..b97e2196 100644
--- a/src/indexing_sums.jl
+++ b/src/indexing_sums.jl
@@ -123,6 +123,12 @@ end
 
 get_indices!(indices, i::Index) = push!(indices, i)
 get_indices!(indices, op::IndexedOperator) = push!(indices, op.ind)
+get_indices!(indices, eqs::AbstractMeanfieldEquations) = get_indices!(indices, eqs.equations)
+function get_indices!(indices, eq::Symbolics.Equation)
+    get_indices!(indices, eq.lhs)
+    get_indices!(indices, eq.rhs)
+    return indices
+end
 
 # find_equality_for_index --  checks for occurrence of expressions such as i == j
 # which allows to e.g. simplify sums
diff --git a/src/printing.jl b/src/printing.jl
index 25ee348e..b42bdf08 100644
--- a/src/printing.jl
+++ b/src/printing.jl
@@ -102,7 +102,7 @@ function SymbolicUtils.show_term(io::IO, p::SymbolicUtils.Symbolic{<:IndexedPara
         show(io, args[i])
         write(io, ", ")
     end
-    if length(args) > 2
+    if length(args) >= 2
         show(io, args[end])
     end
     write(io, "]")
diff --git a/test/test_indexed_meanfield.jl b/test/test_indexed_meanfield.jl
index b9aa639d..3f79a1a1 100644
--- a/test/test_indexed_meanfield.jl
+++ b/test/test_indexed_meanfield.jl
@@ -6,6 +6,64 @@ using Test
 using Random
 const qc = QuantumCumulants
 
+
+@testset "indexed_many_atom_laser" begin
+    using QuantumCumulants
+using OrdinaryDiffEq, ModelingToolkit
+
+
+# Hilbertspace
+hc = FockSpace(:cavity)
+ha = NLevelSpace(:atom,2)
+h = hc ⊗ ha
+
+# operators
+@qnumbers a::Destroy(h)
+σ(α,β,i) = IndexedOperator(Transition(h, :σ, α, β),i)
+
+
+@cnumbers N Δ κ Γ R ν
+g(i) = IndexedVariable(:g, i)
+
+i = Index(h,:i,N,ha)
+j = Index(h,:j,N,ha)
+k = Index(h,:k,N,ha)
+
+# Hamiltonian
+H = -Δ*a'a + Σ(g(i)*( a'*σ(1,2,i) + a*σ(2,1,i) ),i)
+
+# Jump operators with corresponding rates
+J = [a, σ(1,2,i), σ(2,1,i), σ(2,2,i)]
+rates = [κ, Γ, R, ν]
+
+
+# Derive equations
+ops = [a'*a, σ(2,2,j)]
+eqs = meanfield(ops,H,J;rates=rates,order=2)
+
+# custom filter function
+φ(x::Average) = φ(x.arguments[1])
+φ(::Destroy) = -1
+φ(::Create) =1
+φ(x::QTerm) = sum(map(φ, x.args_nc))
+φ(x::Transition) = x.i - x.j
+φ(x::IndexedOperator) = x.op.i - x.op.j
+φ(x::Sum) = φ(x.term)
+phase_invariant(x) = iszero(φ(x))
+
+# Complete equations -- TODO
+# eqs_c = complete(eqs; filter_func=phase_invariant)
+
+# TODO: need a way to create this more easily
+op = σ(2,1,j)*σ(1,2,k)
+s21_j = SymbolicUtils.arguments(SymbolicUtils.arguments(op)[2])[2]
+s12_k = SymbolicUtils.arguments(SymbolicUtils.arguments(op)[2])[3]
+@test QuantumCumulants.was_merged(s21_j, s12_k)
+ops2 = [a'*a, σ(2,2,j), a'*σ(2,2,j), s21_j * s12_k]
+eqs2 = meanfield(ops2,H,J;rates=rates,order=2)
+
+end
+
 # @testset "indexed_meanfield" begin
 
 order = 2
@@ -57,12 +115,13 @@ ops = [a, σ(2,2,k_ind), σ(1,2,k_ind)]
 
 eqs = meanfield(ops,H,J;rates=rates,order=order)
 
-eqs_2 = meanfield(ops,H,J_2;rates=rates_2,order=order)
+tmp = eqs[1].rhs.arguments[2]
+
+# eqs_2 = meanfield(ops,H,J_2;rates=rates_2,order=order)
 
-@test eqs.equations == eqs_2.equations
+# @test eqs.equations == eqs_2.equations
 
-@test isequal([i_ind,j_ind,k_ind],sort(qc.get_indices_equations(eqs)))
-@test isequal([:i,:j,:k],sort(qc.getIndName.(qc.get_indices_equations(eqs))))
+@test isequal(Set([i_ind,j_ind,k_ind]),qc.get_indices(eqs))
 
 @test length(eqs) == 3