allow saving/loading OccupationNumberFS

RimuQMC · Jan 29, 2025 · fc15fb6 · fc15fb6
1 parent 0d61c80
commit fc15fb6
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 7 deletions.
diff --git a/src/BitStringAddresses/multicomponent.jl b/src/BitStringAddresses/multicomponent.jl
@@ -20,7 +20,7 @@ end
 
 # Slow constructor - not to be used internallly
 function CompositeFS(adds::Vararg{SingleComponentFockAddress})
-    N = sum(num_particles, adds)
+    N = sum(a -> num_particles(typeof(a)), adds)
     M1, M2 = extrema(num_modes, adds)
     if M1 ≠ M2
         throw(ArgumentError("all addresses must have the same number of modes"))

diff --git a/src/RimuIO/RimuIO.jl b/src/RimuIO/RimuIO.jl
@@ -15,7 +15,7 @@ using StaticArrays: StaticArrays, SVector
 
 using Rimu: mpi_size, mpi_rank, mpi_barrier
 using Rimu.BitStringAddresses: BitStringAddresses, BitString, BoseFS,
-    CompositeFS, FermiFS, SortedParticleList,
+    CompositeFS, OccupationNumberFS, FermiFS, SortedParticleList,
     num_modes, num_particles
 using Rimu.DictVectors: PDVec, DVec, target_segment
 using Rimu.Interfaces: Interfaces, localpart, storage

diff --git a/src/RimuIO/arrowtypes.jl b/src/RimuIO/arrowtypes.jl
@@ -91,6 +91,20 @@ function ArrowTypes.JuliaType(::Val{FERMIFS_SPL}, ::Type{NTuple{X,Y}}, meta) whe
     return FermiFS{parse(Int, N),parse(Int, M),BS}
 end
 
+###
+### OccupationNumberFS
+###
+const OCCUPATIONNUMBERFS = Symbol("Rimu.OccupationNumberFS")
+ArrowTypes.arrowname(::Type{<:OccupationNumberFS}) = OCCUPATIONNUMBERFS
+ArrowTypes.ArrowType(::Type{OccupationNumberFS{M,T}}) where {M,T} = NTuple{M,T}
+ArrowTypes.toarrow(fs::OccupationNumberFS) = Tuple(fs.onr)
+function ArrowTypes.JuliaType(::Val{OCCUPATIONNUMBERFS}, ::Type{NTuple{M,T}}, _) where {M,T}
+    return OccupationNumberFS{M,T}
+end
+function ArrowTypes.fromarrow(::Type{T}, storage) where {T<:OccupationNumberFS}
+    return T(SVector(storage))
+end
+
 ###
 ### CompositeFS
 ###

diff --git a/test/RimuIO.jl b/test/RimuIO.jl
@@ -35,14 +35,16 @@ end
         SortedParticleList((1, 0, 1, 0, 0, 2, 3)),
         near_uniform(BoseFS{10, 10}),
         BoseFS(101, 5 => 10),
-        FermiFS((1,1,1,0,0,0)),
+        FermiFS(1,1,1,0,0,0),
         FermiFS2C(near_uniform(FermiFS{50,100}), FermiFS(100, 1 => 1)),
         CompositeFS(near_uniform(BoseFS{8,9}), near_uniform(BoseFS{1,9})),
+        OccupationNumberFS(1,2,3,4,5),
         CompositeFS(
-            BoseFS((1,1,1,1,1)),
-            FermiFS((1,0,0,0,0)),
-            BoseFS((1,1,0,0,0)),
-            FermiFS((1,1,1,0,0)),
+            BoseFS(1,1,1,1,1),
+            FermiFS(1,0,0,0,0),
+            OccupationNumberFS(5,4,3,2,1),
+            BoseFS(1,1,0,0,0),
+            FermiFS(1,1,1,0,0),
         ),
     )
         @testset "$(typeof(addr))" begin