Fixed eltype and iterate of Configuration and CSF

JuliaAtoms · Apr 19, 2024 · 7db49b9 · 7db49b9
1 parent b3ef25a
commit 7db49b9
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 6 deletions.
diff --git a/src/configurations.jl b/src/configurations.jl
@@ -493,12 +493,14 @@ Base.getindex(conf::Configuration{O}, i::Integer) where O =
 Base.getindex(conf::Configuration{O}, i::Union{<:UnitRange{<:Integer},<:AbstractVector{<:Integer}}) where O =
     Configuration(Tuple{O,Int,Symbol}[conf[ii] for ii in i], sorted=conf.sorted)
 
-Base.iterate(conf::Configuration{O}, (el, i)=(length(conf)>0 ? conf[1] : nothing,1)) where O =
-    i > length(conf) ? nothing : (el, (conf[i==length(conf) ? i : i+1],i+1))
+Base.iterate(conf::Configuration, i=1) =
+    i > length(conf) ? nothing : (conf[i],i+1)
 
 Base.length(conf::Configuration) = length(conf.orbitals)
+Base.firstindex(conf::Configuration) = 1
 Base.lastindex(conf::Configuration) = length(conf)
-Base.eltype(conf::Configuration{O}) where O = (O,Int,Symbol)
+Base.eachindex(conf::Configuration) = Base.OneTo(length(conf))
+Base.eltype(conf::Configuration{O}) where O = Tuple{O,Int,Symbol}
 
 function Base.write(io::IO, conf::Configuration{O}) where O
     write(io, 'c')

diff --git a/src/csfs.jl b/src/csfs.jl
@@ -113,8 +113,12 @@ end
 csfs(configs::Vector{Configuration{O}}) where O = vcat(map(csfs, configs)...)
 
 Base.length(csf::CSF) = length(peel(csf.config))
+Base.firstindex(csf::CSF) = 1
+Base.lastindex(csf::CSF) = length(csf)
+Base.eachindex(csf::CSF) = Base.OneTo(length(csf))
 Base.getindex(csf::CSF, i::Integer) =
     (peel(csf.config)[i],csf.subshell_terms[i],csf.terms[i])
+Base.eltype(csf::CSF{<:Any,T,S}) where {T,S} = Tuple{eltype(csf.config),IntermediateTerm{T,S},T}
 
 Base.iterate(csf::CSF, (el, i)=(length(csf)>0 ? csf[1] : nothing,1)) =
     i > length(csf) ? nothing : (el, (csf[i==length(csf) ? i : i+1],i+1))

diff --git a/test/configurations.jl b/test/configurations.jl
@@ -174,9 +174,14 @@
         @test bound(rc"[Kr] 5s2 5p-2 5p2 ks ld") == rc"[Kr] 5s2 5p-2 5p2"
         @test continuum(rc"[Kr] 5s2 5p-2 5p2 ks ld") == rc"ks ld"
 
-        @test c"[Ne]"[1] == (o"1s",2,:closed)
-        @test c"[Ne]"[1:2] == c"1s2c 2s2c"
-        @test c"[Ne]"[end-1:end] == c"2s2c 2p6c"
+        let c = c"[Ne]"
+            @test c[1] == (o"1s",2,:closed)
+            @test c[1:2] == c"1s2c 2s2c"
+            @test c[end-1:end] == c"2s2c 2p6c"
+            @test c[begin+1:end-1] == c"2s2c"
+            @test eachindex(c) === Base.OneTo(3)
+            @test collect(c) == [c[i] for i in eachindex(c)]
+        end
 
         @test c"1s2 2p kp"[2:3] == c"2p kp"
 

diff --git a/test/csfs.jl b/test/csfs.jl
@@ -74,4 +74,16 @@ using .ATSPParser
                 "[Kr]ᶜ 5s²(₀0|0) 5p-(₁1/2|1/2) 5p⁴(₀0|1/2) kd(₁5/2|3)-"
         end
     end
+
+    @testset "Access subshells" begin
+        c = rc"1s 2p"
+
+        for (i,csf) in enumerate(csfs(c))
+            @test length(csf) == 2
+            @test csf[begin] == ((ro"1s", 1, :open), IntermediateTerm(half(1), Seniority(1)), half(1))
+            # It just so happens that for this particular configuration, the accumulated intermediate term is J = 1, 2
+            @test csf[end] == ((ro"2p", 1, :open), IntermediateTerm(half(3), Seniority(1)), i)
+            @test collect(csf) == [csf[i] for i in eachindex(csf)]
+        end
+    end
 end