nonzero_g_indices and nonzero_g_vectors functions

CHANGELOG.md

@@ -13,12 +13,16 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 this type should be defined for all data structures containing energies and occupancies.
   - New functions: `energy`, `occupancy`, `energies`, `occupancies`
   - `EnergiesOccupancies` constructor for `PlanewaveWavefunction`
+  - New `nonzero_g_indices` and `nonzero_g_vectors` functions for `PlanewaveWavefunction`
 
 ### Changed
   - `min_energy`, `max_energy`, `min_occupancy`, `max_occupancy`, and `fermi` all have generic
 definitions based on `AbstractArray{<:EnergyOccupancy}`.
   - `PlanewaveWavefunction` uses these generic definitions via the constructor.
 
+### Fixed
+  - Missing export of `nspin`
+
 ## [0.1.14]: 2023-11-02
 
 ### Added

docs/src/api/data.md

@@ -55,6 +55,8 @@ Electrum.PlanewaveWavefunction
 Electrum.PlanewaveIndex
 Electrum.nspin
 Electrum.nband(::PlanewaveWavefunction)
+Electrum.nonzero_g_indices
+Electrum.nonzero_g_vectors
 ```
 
 ## Band structures

src/Electrum.jl

@@ -142,6 +142,7 @@ export shift, fft, ifft, fftfreq, voxelsize, integrate, partial_derivative, cell
 # Planewave wavefunctions
 include("data/wavefunctions.jl")
 export PlanewaveIndex, PlanewaveWavefunction
+export nspin, nband, nonzero_g_indices, nonzero_g_vectors
 # Band structures
 include("data/bands.jl")
 export BandAtKPoint, BandStructure

src/data/wavefunctions.jl

@@ -318,3 +318,28 @@ EnergiesOccupancies(wf::PlanewaveWavefunction) = EnergyOccupancy.(wf.energies, w
 function EnergiesOccupancies{T}(wf::PlanewaveWavefunction) where T
     return EnergyOccupancy{T}.(wf.energies, wf.occpuancies)
 end
+
+#---Occupied portions of wavefunctions-------------------------------------------------------------#
+"""
+    nonzero_g_indices(wf::PlanewaveWavefunction{D}) -> Vector{CartesianIndex{D}}
+
+Returns a vector of `CartesianIndex` objects corresponding to planewave G-vector indices that are
+not all zero at each band and k-point.
+
+To return the G-vectors as objects which subtype `AbstractVector`, `nonzero_g_vectors(wf)` may be 
+used instead, which is equivalent to calling `SVector.(Tuple.(nonzero_g_indices(wf)))`.
+"""
+function nonzero_g_indices(wf::PlanewaveWavefunction)
+    return FFTBins(wf)[findall(any.(!iszero, eachslice(wf.data, dims = 1)))]
+end
+
+"""
+    nonzero_g_vectors(wf::PlanewaveWavefunction{D}) -> Vector{SVector{D,Int}}
+
+Returns a vector of `SVector{D,Int}` objects corresponding to planewave G-vectors that are not all
+zero at each band and k-point.
+
+This is equivalent to calling `SVector.(Tuple.(nonzero_g_indices(wf)))`, and can be used whenever
+`AbstractVector` inputs are needed instead of a `CartesianIndex`.
+"""
+nonzero_g_vectors(wf::PlanewaveWavefunction) = SVector.(Tuple.(nonzero_g_indices(wf)))

test/wavefunctions.jl

@@ -29,4 +29,8 @@
     @test min_occupancy(wavecar) == min_occupancy(EnergiesOccupancies(wavecar))
     @test max_occupancy(wavecar) == max_occupancy(EnergiesOccupancies(wavecar))
     @test min_energy(wavecar) <= fermi(wavecar) <= max_energy(wavecar)
+    @test all(
+        length(nonzero_g_vectors(wavecar)) >= count(!iszero(wavecar.data[:,b,k,s]))
+        for b in 1:nband(wavecar), k in 1:nkpt(wavecar), s in 1:nspin(wavecar)
+    )
 end