From bd02892ab1f46c4363f972af2598792898b6c1b0 Mon Sep 17 00:00:00 2001
From: Augustin Bussy <augustin.bussy@cscs.ch>
Date: Wed, 12 Nov 2025 18:12:22 +0100
Subject: [PATCH] Fix occupation masks and add tests

---
 src/occupation.jl  |  5 +---
 test/occupation.jl | 70 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/src/occupation.jl b/src/occupation.jl
index 301dfa7904..bc3fda24dd 100644
--- a/src/occupation.jl
+++ b/src/occupation.jl
@@ -218,10 +218,7 @@ Return ranges of occupied elements based on a given occupation threshold
 """
 function occupied_empty_masks(occupation, occupation_threshold)
     n_occ = map(occupation) do occ
-        n = count(occ_i -> abs(occ_i) > occupation_threshold, occ)
-        # Check that all occupied elements are contiguous, otherwise range is wrong
-        @assert all(occ[1:n] .> occupation_threshold)
-        n
+        something(findlast(x -> abs(x) > occupation_threshold, occ), 0)
     end
     mask_occ  = [1:n_occ[ik] for ik in 1:length(occupation)]
     mask_empty = [(n_occ[ik] + 1):length(occupation[ik]) for ik in 1:length(occupation)]
diff --git a/test/occupation.jl b/test/occupation.jl
index 2eb036c6dc..a27e35819e 100644
--- a/test/occupation.jl
+++ b/test/occupation.jl
@@ -248,3 +248,73 @@ end
         εF < εF_ref && @test n_electrons < n_electrons_ref
     end
 end
+
+@testitem "Occupied/empty orbitals definition" #=
+    =# tags=[:dont_test_mpi] setup=[Occupation, TestCases] begin
+    using DFTK
+    silicon = TestCases.silicon
+
+    temperatures = [0.0, 0.0005]
+    Ecut = 5
+    kgrid = [2, 2, 2]
+    maxiter = 3
+    threshold = 0.0
+
+    for smearing in Occupation.smearing_methods for temperature in temperatures
+        smearing isa DFTK.Smearing.None && continue
+
+        model = model_DFT(silicon.lattice, silicon.atoms, silicon.positions;
+                          functionals=LDA(), temperature, smearing)
+        basis = PlaneWaveBasis(model; Ecut, kgrid)
+        scfres = self_consistent_field(basis; maxiter)
+
+        # Eigenvalues are monotonically increasing (up to epsilon)
+        @test all(all(diff(εk) .≥ -eps(typeof(basis.Ecut))) for εk in scfres.eigenvalues)
+
+        (mask_occ, mask_empty) = DFTK.occupied_empty_masks(scfres.occupation, threshold)
+
+        # All orbitals are either occupied or empty
+        for ik = 1:length(basis.kpoints)
+            @test length(mask_occ[ik]) + length(mask_empty[ik]) == size(scfres.ψ[ik], 2)
+        end
+
+        # All empty orbitals have occupation below threshold
+        for ik = 1:length(basis.kpoints)
+            @test all(abs(scfres.occupation[ik][j]) ≤ threshold for j in mask_empty[ik])
+        end
+
+        # Only MethfesselPaxton smearing can have negative occupations
+        for ik = 1:length(basis.kpoints)
+            if any(scfres.occupation[ik] .< -threshold)
+                @test smearing isa DFTK.Smearing.MethfesselPaxton
+            end
+        end
+
+        # For all other smearings, all occupied orbitals have occupation above threshold
+        if !(smearing isa DFTK.Smearing.MethfesselPaxton)
+            for ik = 1:length(basis.kpoints)
+                @test all(scfres.occupation[ik][j] ≥ threshold for j in mask_occ[ik])
+            end
+        end
+    end
+    end
+
+    # Special test specifically triggering case of occupied oribitals with occupations
+    # below threshold (oscillations of MethfesselPaxton smearing). If occupations are
+    # not carefully calculated, i.e. the last element with abs(occ) > threshold marks the
+    # boundary, some occupied orbitals can be classified as empty.
+    model = model_DFT(silicon.lattice, silicon.atoms, silicon.positions;
+                      functionals=LDA(), temperature=0.8,
+                      smearing=DFTK.Smearing.MethfesselPaxton(3))
+    basis = PlaneWaveBasis(model; Ecut=15, kgrid=[1, 1, 1])
+    scfres = self_consistent_field(basis; tol=1e-6)
+
+    custom_threshold = 0.01
+    (mask_occ, mask_empty) = DFTK.occupied_empty_masks(scfres.occupation, custom_threshold)
+
+    # All empty orbitals have occupation below threshold
+    for ik = 1:length(basis.kpoints)
+        @test all(abs(scfres.occupation[ik][j]) ≤ custom_threshold for j in mask_empty[ik])
+    end
+
+end
\ No newline at end of file