Implement core projection

Author: vitartas

app/quoll.jl

@@ -43,7 +43,6 @@ global_logger(
 
 @info "Parsing QUOLL input file"
 
-# input_filepath = "/home/chem/phrpwt/Jobs/Data_Pipeline/Quoll_Workspace/Quoll/examples/SiC/input_file.toml"
 input_filepath = abspath(ARGS[1])
 params, input_dirs, output_dirs = Quoll.Parser.parse_inputfile(input_filepath)
 
@@ -52,6 +51,11 @@ params, input_dirs, output_dirs = Quoll.Parser.parse_inputfile(input_filepath)
 ndirs = length(input_dirs)
 my_idirs = Quoll.split_work(ndirs, global_comm, Quoll.DefaultSplit())
 
+# Construct an MPI sub-communicator for MPI tasks working on the same atomic configuration
+global_color = first(my_idirs)
+config_comm = MPI.Comm_split(global_comm, global_color, global_rank)
+config_rank = MPI.Comm_rank(config_comm)
+
 for idir in my_idirs
     input_dir = input_dirs[idir]
 
@@ -61,7 +65,6 @@ for idir in my_idirs
     operatorkinds = find_operatorkinds(input_dir, params)
     operators = load_operators(input_dir, operatorkinds, params.input.format)
 
-    # TODO: in the future could allow for <other canonical formats/shortcut conversions>
     @info "Converting operators into BSparseOperator format"
 
     operators = convert_operators(operators, BSparseOperator)
@@ -71,7 +74,18 @@ for idir in my_idirs
 
         # Assuming all operators in the directory are related to the same atoms
         atoms = Quoll.get_atoms(first(operators))
-        kgrid = get_kgrid(atoms, params)
+        kgrid = get_kgrid(atoms, operatorkinds, params)
+
+        @info "Splitting k-points across MPI tasks"
+        # TODO: add an option for different splits in the input file
+        nkpoints = Quoll.get_nkpoints(kgrid)
+        my_ikpoints = Quoll.split_work(nkpoints, config_comm, Quoll.FHIaimsLAPACKSplit())
+    end
+
+    if !isnothing(params.basis_projection)
+        @info "Projecting the core states"
+
+        operators = perform_core_projection(operators, kgrid, my_ikpoints, config_comm, params)
     end
 
     if !isnothing(params.output)

examples/SiC/input_file.toml

@@ -11,7 +11,8 @@ directory = "SiC_DeepH_converted"
 # Should also allow different types,
 # i.e. Monkhorst-Pack and even possibly different shifts
 [kpoint_grid]
-mesh = [10, 10, 10]
+mesh = [10, 20, 20]
+# mesh = [1, 1, 3]
 density = 8.0
 
 [basis_projection]
@@ -21,7 +22,7 @@ projected_basis = [
     "Si (2, 0, 0) type = atomic",
     "Si (2, 1, *) type = atomic",
 ]
-method = "FC99V"
+method = "LaikovCore"
 
 [postprocessing]
 fermi_level = true

src/Parser.jl

@@ -7,19 +7,24 @@ using Unitful
 using UnitfulAtomic
 using AtomsBase
 using ArgCheck
+using MPI
 
 using ..Quoll
 using ..Quoll:
     AbstractOperator,
     OperatorKind,
+    KGrid,
     allows_symmetry,
     get_avail_operatorkinds,
     get_operatorkinds,
     get_readformat,
     get_writeformat
 import ..Quoll: find_operatorkinds, get_kgrid
+
 using ..Projections
 using ..Projections: AbstractBasisProjection, get_basis_projection
+import ..Projections: perform_core_projection
+
 using ..Postprocessing
 using ..Postprocessing: SmearingFunction, get_smearing
 

src/Projections.jl

@@ -1,10 +1,42 @@
 module Projections
 
+using LinearAlgebra
+using StaticArrays
+using ArgCheck
+using MPI
+
+using ..Quoll
+using ..Quoll:
+    AbstractOperator,
+    KGrid,
+    precompute_phases,
+    get_float,
+    get_kind,
+    get_data,
+    get_metadata,
+    get_sparsity,
+    get_basisset,
+    get_spins,
+    get_dense_subbasis_mask,
+    construct_subbasis_metadata,
+    build_operator,
+    fourier_transform,
+    inv_fourier_transform_data!,
+    synchronise_data!
+
 abstract type AbstractBasisProjection end
 
-function basis_projection end
+function get_basis_projection end
+
+function perform_core_projection end
+
+export perform_core_projection
+include("projections/common.jl")
 
 export FC99V
 include("projections/fc99v.jl")
 
+export LaikovCore
+include("projections/laikov.jl")
+
 end

src/Quoll.jl

@@ -1,9 +1,11 @@
 module Quoll
 
+using Base
 using Reexport
 using ArgCheck
 using LinearAlgebra
 using StaticArrays
+using OffsetArrays
 using Dictionaries
 using AutoHashEquals
 using AxisKeys
@@ -43,7 +45,7 @@ include("sparsity.jl")
 include("shconversion.jl")
 
 export get_kgrid
-include("kgrid.jl")
+include("kpoints.jl")
 
 # Operator formats and their IO routines
 # (Assuming no dependencies between different format types in each file)
@@ -54,9 +56,10 @@ export write_operator, write_operators
 export find_operatorkinds
 include("operators/interface.jl")
 
-export BSparseOperator
+export BSparseOperator, DenseOperator
 include("operators/canonical/abstract.jl")
 include("operators/canonical/bsparse.jl")
+include("operators/canonical/dense.jl")
 
 export DeepHOperator
 include("operators/deeph/deeph.jl")
@@ -68,8 +71,10 @@ include("operators/fhiaims/fhiaims_csc.jl")
 # Operator format conversions
 
 export convert_operator, convert_operators
+export fourier_transform
 include("conversions/interface.jl")
 
+include("conversions/canonical/bsparse.jl")
 include("conversions/fhiaims/fhiaims_csc.jl")
 include("conversions/deeph/deeph.jl")
 

src/basis.jl

@@ -38,10 +38,47 @@ basis_species(basisset::BasisSetMetadata, species::ChemicalSpecies) = basisset.b
 get_unique_species(basisset::BasisSetMetadata) = unique(basisset.atom2species)
 
 # Assuming ms that belong to a given l are next to each other
-function get_angular_momenta(basis::Vector{T}) where T<:BasisMetadata
+function get_angular_momenta(basis::AbstractVector{T}) where T<:BasisMetadata
     return getfield.(filter(orbital -> orbital.l == orbital.m, basis), :l)
 end
 
+function BasisSetMetadata(basisset::BasisSetMetadata, masks::SpeciesDictionary)
+    unique_species_keys = Indices(keys(basisset.basis))
+    subbasis_dict = map(unique_species_keys) do z
+        basis_z = basis_species(basisset, z)
+        basis_z[masks[z]]
+    end
+    return BasisSetMetadata(Base.ImmutableDict(pairs(subbasis_dict)...), basisset.atom2species)
+end
+
+# Returns Dictionary{ChemicalSpecies, BitVector}
+function get_subbasis_masks(basisset::BasisSetMetadata, subbasis::Vector{T};
+    inverted = false) where T<:BasisMetadata
+
+    unique_species_keys = Indices(keys(basisset.basis))
+
+    return map(unique_species_keys) do z
+        basis_z = basis_species(basisset, z)
+        subbasis_z = filter(orbital -> isequal(orbital.z, z), subbasis)
+
+        # Throw a warning if some orbitals in subbasis are not present in basisset
+        # (which could happen if the orbitals were supplied incorrectly in the input file)
+        if !all(subbasis_z .∈ Ref(basis_z))
+            @warn "Some orbitals in the subbasis are not present in the full basis"
+        end
+
+        inverted ? basis_z .∉ Ref(subbasis_z) : basis_z .∈ Ref(subbasis_z)
+    end
+end
+
+# Returns BitVector to be used with dense operators
+function get_dense_subbasis_mask(basisset::BasisSetMetadata, subbasis::Vector{T};
+    inverted = false) where T<:BasisMetadata
+
+    subbasis_masks = get_subbasis_masks(basisset, subbasis, inverted = inverted)
+    return reduce(vcat, (subbasis_masks[z] for z in basisset.atom2species))
+end
+
 get_atom2nbasis(basisset::BasisSetMetadata) = get_atom2nbasis(basisset.basis, basisset.atom2species)
 get_atom2nbasis(basis, atom2species) = [length(basis[z]) for z in atom2species]
 
@@ -85,4 +122,4 @@ function get_atom2offset(basisset::BasisSetMetadata)
     return get_atom2offset(atom2basis)
 end
 
-get_atom2offset(atom2basis) = [first(interval) - 1 for interval in atom2basis]
+get_atom2offset(atom2basis) = [first(interval) - 1 for interval in atom2basis]

src/constants.jl

@@ -1,12 +1,13 @@
 ### TYPE ALIASES ###
 
 const AbstractAnyDict{K, V} = Union{<:AbstractDict{K, V}, <:AbstractDictionary{K, V}}
-
 const SpeciesPairAnyDict{V} = AbstractAnyDict{NTuple{2, ChemicalSpecies}, V}
 const SpeciesAnyDict{V}  = AbstractAnyDict{ChemicalSpecies, V}
-
 const AtomPairAnyDict{V} = AbstractAnyDict{NTuple{2, Int}, V}
 
+const SpeciesDictionary{V} = AbstractDictionary{ChemicalSpecies, V}
+const SpeciesDict{V} = AbstractDict{ChemicalSpecies, V}
+
 const ThreeDimSliceView{T} = SubArray{
     T, 3, Array{T, 3},
     Tuple{

src/conversions/canonical/bsparse.jl

@@ -0,0 +1,109 @@
+### FOURIER TRANSFORMS ###
+
+# TODO: the same function could be used for both DenseOperator and FHIaimsDenseOperator
+# (except for the data I would need to use SH conversion)
+# TODO: modify to enable float conversions
+function fourier_transform(in_operator::BSparseOperator, kpoints, phases, ::Type{DenseOperator}; float = Float64)
+    in_metadata = get_metadata(in_operator)
+    in_sparsity = get_sparsity(in_operator)
+    in_atoms = get_atoms(in_operator)
+    in_basisset = get_basisset(in_operator)
+    in_spins = get_spins(in_operator)
+    in_kind = get_kind(in_operator)
+
+    # Convert metadata
+    out_sparsity = convert_sparsity(in_metadata, RecipDenseSparsity, hermitian = in_sparsity.hermitian)
+    out_metadata_list = [
+        RecipDenseMetadata(in_atoms, out_sparsity, in_basisset, in_spins, SVector{3}(kpoint))
+        for kpoint in kpoints
+    ]
+    
+    # Initialize out_operator with zeros
+    out_operator_list = [
+        build_operator(DenseOperator, in_kind, out_metadata; uninit = false, value = zero(Complex{float}))
+        for out_metadata in out_metadata_list
+    ]
+
+    # Perform Fourier transform
+    for (out_operator, phases_k) in zip(out_operator_list, phases)
+        fourier_transform_data!(out_operator, in_operator, phases_k)
+    end
+    
+    return length(out_operator_list) == 1 ? first(out_operator_list) : out_operator_list
+end
+
+# TODO: write another method but with shifts (or alternatively write the method with shifts
+# but force the compiler to remove shifts for matching SH convention)
+function fourier_transform_data!(out_operator::DenseOperator, in_operator::BSparseOperator, phases_k)
+    in_keydata = get_keydata(in_operator)
+    in_sparsity = get_sparsity(in_operator)
+    out_data = get_data(out_operator)
+    out_basisset = get_basisset(out_operator)
+    out_float = get_float(out_operator)
+
+    ilocal2iglobal = get_ilocal2iglobal(in_sparsity)
+    atom2offset = get_atom2offset(out_basisset)
+
+    for ((iat, jat), in_block) in pairs(in_keydata)
+
+        phases_kij = phases_k[ilocal2iglobal[(iat, jat)]]
+        atom2offset_i = atom2offset[iat]
+        atom2offset_j = atom2offset[jat]
+
+        for jb in axes(in_block, 2)
+            jb_dense = jb + atom2offset_j
+            for ib in axes(in_block, 1)
+                ib_dense = ib + atom2offset_i
+
+                tmp = zero(out_float)
+                @inbounds for iR in axes(in_block, 3)
+                    tmp += in_block[ib, jb, iR] * phases_kij[iR]
+                end
+
+                out_data[ib_dense, jb_dense] = tmp
+            end
+        end
+
+        # @tullio tmp[ib,jb] := in_block[ib,jb,R] * phases_kij[R]
+        # out_data[atom2basis[iat], atom2basis[jat]] = tmp
+    end
+
+    # TODO: I could use Hermitian array instead, probably directly in build stage
+    if in_sparsity.hermitian
+        for jb_dense in 1:size(out_data, 1)
+            for ib_dense in 1:jb_dense
+                out_data[jb_dense, ib_dense] = conj(out_data[ib_dense, jb_dense])
+            end
+        end
+    end
+end
+
+# Only a contribution to out_operator because in_operator contains only a single k-point
+function inv_fourier_transform_data!(out_operator::BSparseOperator, in_operator::DenseOperator, phases_k, weight)
+    out_basisset = get_basisset(out_operator)
+    out_keydata = get_keydata(out_operator)
+    out_sparsity = get_sparsity(out_operator)
+    in_data = get_data(in_operator)
+
+    ilocal2iglobal = get_ilocal2iglobal(out_sparsity)
+    atom2offset = get_atom2offset(out_basisset)
+
+    for ((iat, jat), out_block) in pairs(out_keydata)
+
+        # Assuming phases is a vector here (i.e. for a single k-point)
+        inv_phases_kij = conj(phases_k[ilocal2iglobal[(iat, jat)]])
+        atom2offset_i = atom2offset[iat]
+        atom2offset_j = atom2offset[jat]
+        
+        for iR in axes(out_block, 3)
+            for jb in axes(out_block, 2)
+                jb_dense = jb + atom2offset_j
+                for ib in axes(out_block, 1)
+                    ib_dense = ib + atom2offset_i
+                    out_block[ib, jb, iR] += weight * real(in_data[ib_dense, jb_dense] * inv_phases_kij[iR])
+                end
+            end
+        end
+    end
+
+end