read_abinit_WFK stores k-point mesh correctly

CHANGELOG.md

@@ -7,6 +7,13 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Added
+  - New `RealBasis`/`ReciprocalBasis` constructors for `Electrum.ABINITHeader`.
+  - New constructor for an empty `PlanewaveWavefunction` with known `KPointMesh`.
+
+### Fixed
+  - `read_abinit_WFK` now correctly stores the `KPointMesh` from the wavefunction file header.
+
 ## [0.1.13]: 2023-10-31
 
 The developers would like to wish you a happy Halloween!

src/data/wavefunctions.jl

@@ -139,7 +139,7 @@ struct PlanewaveWavefunction{D,T} <: AbstractArray{T,D}
         basis::LatticeBasis,
         spins::AbstractVector{<:StaticVector{D,<:Real}},
         kpoints::AbstractVector{KPoint{D}},
-        energies::AbstractArray{<:Real,3},
+        energies::AbstractArray{<:Real,3},      # Not specific to 3D, this is from kpt/band/spin
         occupancies::AbstractArray{<:Real,3},
         grange::NTuple{D,<:AbstractUnitRange{<:Integer}},
         data::Array{T,4}
@@ -155,6 +155,37 @@ end
 
 Base.has_offset_axes(::PlanewaveWavefunction) = true
 
+"""
+    PlanewaveWavefunction{D,T}(
+        basis::LatticeBasis,
+        nspin::Integer,
+        nkpt::AbstractVector,
+        nband::Integer,
+        grange::AbstractUnitRange{<:Integer}...
+    )
+
+Constructs an empty `PlanewaveWavefunction` with `nspin` spins, a list of k-points `kptmesh`, 
+`nband` bands, and G-vectors in the ranges given by `grange`.
+"""
+function PlanewaveWavefunction{D,T}(
+    basis::LatticeBasis,
+    nspin::Integer,
+    kptmesh::AbstractVector{<:AbstractVector},
+    nband::Integer,
+    grange::Vararg{AbstractUnitRange{<:Integer},D}
+) where {D,T}
+    nkpt = length(kptmesh)
+    return PlanewaveWavefunction(
+        basis,
+        zeros(SVector{D,Float64}, nspin),
+        kptmesh,
+        zeros(Float64, nband, nkpt, nspin),
+        zeros(Float64, nband, nkpt, nspin),
+        grange,
+        zeros(T, prod(length.(grange)), nband, nkpt, nspin)
+    )
+end
+
 """
     PlanewaveWavefunction{D,T}(
         basis::LatticeBasis,

src/software/abinit.jl

@@ -158,6 +158,9 @@ end
 Base.getindex(h::ABINITHeader, name::Symbol) = getfield(h, name)
 Base.setindex!(h::ABINITHeader, x, name::Symbol) = setfield!(h, name, x)
 
+RealBasis(h::ABINITHeader) = RealBasis(h.rprimd)
+ReciprocalBasis(h::ABINITHeader) = ReciprocalBasis(RealBasis(h))
+
 #=
 """
     Electrum.symrel_to_sg(symrel::AbstractVector{<:AbstractMatrix{<:Integer}}) -> Int
@@ -178,11 +181,8 @@ symrel_to_sg(h::ABINITHeader) = symrel_to_sg(h.symrel)
 
 function Crystal(h::ABINITHeader)
     atomlist = PeriodicAtomList(
-        RealBasis(h.rprimd),
-        FractionalAtomPosition.(
-            Int.(h.znucltypat[h.typat]),
-            h.xred
-        )
+        RealBasis(h),
+        FractionalAtomPosition.(Int.(h.znucltypat[h.typat]), h.xred)
     )
     # Don't include space group data since all atomic positions are generated (use defaults)
     return Crystal(atomlist)
@@ -673,7 +673,7 @@ function read_abinit_DEN(io::IO)
     # Add each dataset to the dictionary
     data = Dict{String, RealDataGrid{3,T}}()
     # Convert the basis
-    basis = RealBasis(header.rprimd)
+    basis = RealBasis(header)
     # Fill the dictionary
     data["density_total"] = RealDataGrid(rho[1], basis)
     if header.nspden == 2
@@ -714,7 +714,7 @@ function read_abinit_POT(io::IO)
     # No conversion will occur here: assume units of Hartree
     rho = read_abinit_datagrids(T, io, header.nspden, header.ngfft)
     data = Dict{String, RealDataGrid{3,T}}()
-    basis = RealBasis{3}(header.rprimd)
+    basis = RealBasis(header)
     if header.nspden == 1
         data["potential_total"] = RealDataGrid(rho[1], basis)
     elseif header.nspden == 2
@@ -751,7 +751,7 @@ function read_abinit_WFK(io::IO; quiet = false)
     # Get the header from the file
     header = read_abinit_header(io)
     # Get the reciprocal lattice
-    rlatt = convert(ReciprocalBasis, RealBasis(header.rprimd))
+    rlatt = ReciprocalBasis(header)
     # Get the minimum and maximum HKL values needed
     # Units for c (2*m_e/ħ^2) are hartree^-1 bohr^-2
     # this should affect only the size of the preallocated arrays
@@ -763,7 +763,13 @@ function read_abinit_WFK(io::IO; quiet = false)
         "Calculated from ecut = ", header.ecut, " Hartree"
     )
     nb = maximum(header.nband)
-    wf = PlanewaveWavefunction{3,Complex{Float64}}(rlatt, header.nsppol, header.nkpt, nb, bounds...)
+    wf = PlanewaveWavefunction{3,Complex{Float64}}(
+        rlatt,
+        header.nsppol,
+        KPointMesh(header),
+        nb,
+        bounds...
+    )
     # Loop over all the spin polarizations
     for sppol in 1:header.nsppol
         # Loop over all the k-points

test/filetypes.jl

@@ -36,6 +36,7 @@ end
     @test size(den) == (24, 24, 36)
     # Check that there's 1 spin, 4 k-points, and 8 bands
     @test size(wfk)[1:3] == (1, 4, 8)
+    @test KPointMesh(header_wfk) == KPointMesh(wfk)
 end
 
 @testset "VASP outputs" begin