add Hubbard

1. add read and write for hubbard card
2. add some utility functions for parsing pw inputs

Author: LyKex

src/common/const.jl

@@ -6,3 +6,20 @@ This is the default (Physical constants, SI (NIST 2018)) value in QE
 `Modules/constants.f90`.
 """
 const Bohr_QE::Float64 = 0.529177210903
+
+"""
+pw.x input keywords
+"""
+const PW_KEYWORDS = Set([
+    "atomic_species",
+    "atomic_positions",
+    "k_points",
+    "additional_k_points",
+    "cell_parameters",
+    "constraints",
+    "occupations",
+    "atomic_velocities",
+    "atomic_forces",
+    "solvents",
+    "hubbard",
+])

src/namelist.jl

@@ -84,7 +84,6 @@ function read_namelists(io::IO; all_lines::Bool=false)
         end
         return namelists
     end
-
 end
 
 function read_namelists(filename::AbstractString; kwargs...)
@@ -168,12 +167,61 @@ find_card(lines, name)
 """
 function find_card(lines::AbstractVector, name::AbstractString)
     lname = lowercase(name)
-    return findfirst(
-        line -> startswith(lowercase(remove_comment(line)), lname),
-        lines,
-    )
+    return findfirst(line -> startswith(lowercase(remove_comment(line)), lname), lines)
+end
+
+# TODO only considers pw keywords for now
+function is_namelist_keyword(line; keywords=PW_KEYWORDS)
+    line = strip(line)
+    startswith(line, "&") && return true
+
+    sl = split(line)[1]
+    lowercase(sl) in keywords && return true
+    return false
+end
+
+"""
+    $(SIGNATURES)
+
+Find the last non-empty line in current card, which is any of the following:
+1. end of file
+2. the line before next card keyword
+3. empty line (this is assuming no empty lines in the middle of the card)
+
+# Arguments
+- `lines::AbstractVector`: The lines to read from.
+- `icard::Integer`: Beginning of the current card.
+
+# Returns
+- The index of the last line in the card.
+
+# Examples
+```jldoctest; setup = :(using QuantumEspressoIO: end_of_card)
+lines = [
+    "HUBBARD atomic",
+    "U Ni-3d 5.77",
+    "U Ni1-3d 5.77",
+    "ATOMIC_POSITIONS"
+]
+
+end_of_card(lines, 1)
+# output
+3
+"""
+function end_of_card(lines::AbstractVector, istart::Integer)
+    # advance to next line
+    istart += 1
+    while istart <= length(lines)
+        line = strip(lines[istart])
+        isempty(line) && return istart - 1
+        is_namelist_keyword(line) && return istart - 1
+        istart += 1
+    end
+    # end of lines
+    return istart - 1
 end
 
+
 """
     $(SIGNATURES)
 
@@ -189,6 +237,10 @@ Get the option of a card.
 ```jldoctest; setup = :(using QuantumEspressoIO: parse_card_option)
 julia> parse_card_option("POSITIONS angstrom  ! comment")
 "angstrom"
+julia> parse_card_option("POSITIONS (automatic)  ! comment")
+"automatic"
+julia> parse_card_option("POSITIONS {automatic}  ! comment")
+"automatic"
 ```
 ```jldoctest; setup = :(using QuantumEspressoIO: parse_card_option)
 julia> parse_card_option("POSITIONS {angstrom}  ! comment")
@@ -197,6 +249,7 @@ julia> parse_card_option("POSITIONS {angstrom}  ! comment")
 """
 @inline function parse_card_option(line::AbstractString)
     cardline = remove_comment(line)
+    cardline = replace(cardline, "("=>"", ")"=>"", "{"=>"", "}"=>"")
     parts = split(cardline; limit=2)
     if length(parts) < 2
         return nothing
@@ -237,7 +290,9 @@ println(lines)
 ["  name2 = value2"]
 ```
 """
-function parse_card!(lines::AbstractVector, name::AbstractString, n_lines::Union{Integer,Nothing}=nothing)
+function parse_card!(
+    lines::AbstractVector, name::AbstractString, n_lines::Union{Integer,Nothing}=nothing
+)
     istart = find_card(lines, name)
     # nothing found
     isnothing(istart) && return nothing
@@ -315,7 +370,7 @@ function write_namelist(io::IO, name::AbstractString, params::AbstractDict)
             @printf(io, "%s%s = %s\n", ' '^indent, k, fv)
         end
     end
-    println(io, "/")
+    return println(io, "/")
 end
 
 """
@@ -337,7 +392,7 @@ end
 function write_namelist(io::Union{IO,AbstractString}, namelist::Pair)
     name = first(namelist)
     params = second(namelist)
-    write_namelist(io, name, params)
+    return write_namelist(io, name, params)
 end
 
 """

src/pw.jl

@@ -239,6 +239,100 @@ function read_k_points!(lines::AbstractVector; name::AbstractString="k_points")
     return name => card
 end
 
+"""
+    $(SIGNATURES)
+
+Count how many valid lines from istart to iend in lines inclusive.
+
+# Examples
+```jldoctest; setup = :(using QuantumEspressoIO: num_valid_lines)
+lines = [
+    "HUBBARD atomic",
+    "! a comment line",
+    "! another comment line",
+    "U Ni-3d 5.77",
+    "V Ni-3d O-2p 1 12 3.5",
+    "ALPHA Ni-3d 0.05",
+    "J0 Ni-3d 0.05",
+    "following",
+]
+nlines = valid_hubbard_lines(lines, 1, 8)
+println(nlines)
+# output
+4
+"""
+function valid_hubbard_lines(lines::AbstractVector, istart::Integer, iend::Integer)
+    i = istart
+    num_lines = iend - istart + 1
+    hubbard_types = ["U", "ALPHA", "J0", "J", "B", "E2", "E3", "V"]
+    for line in lines[istart:iend]
+        line = remove_comment(line)
+        # valid lines should start with hubbard symbol and end with value
+        if isempty(line) || !(endswith(line, r"\d+$") && any([startswith(line, t) for t in hubbard_types]))
+            num_lines -= 1
+        end
+    end
+    return num_lines
+end
+
+"""
+    $(SIGNATURES)
+
+Parse the `hubbard` card of `pw.x` input.
+
+# Arguments
+- `lines::AbstractVector`: The lines of the input file.
+- `n_species::Integer`: The number of species in the input file.
+
+# Returns
+- A `Pair` of card name to card content. The card option is stored under
+    the `:option` key in the card content.
+
+# Examples
+```jldoctest; setup = :(using QuantumEspressoIO: read_hubbard!)
+lines = [
+    "HUBBARD atomic",
+    "! a comment line",
+    "! another comment line",
+    "U Ni-3d 5.77",
+    "U Ni1-3d 5.77",
+    "V Ni-3d O-2p 1 12 3.5",
+    "following lines",
+]
+card = read_hubbard!(lines)
+println(card)
+println(lines)
+# output
+:hubbard => OrderedCollections.OrderedDict{Symbol, Any}(:option => "atomic", :types => ["U", "U", "V"], :manifolds => ["Ni-3d", "Ni1-3d", "Ni-3d O-2p 1 12"], :values => [5.77, 5.77, 3.5])
+["following lines"]
+```
+"""
+function read_hubbard!(lines::AbstractVector)
+    name = "hubbard"
+    icard = find_card(lines, name)
+    # no hubbard card found
+    isnothing(icard) && return nothing
+
+    # hubbard projector must be specified
+    # https://gitlab.com/QEF/q-e/-/blob/develop/Doc/Hubbard_input.tex?ref_type=heads
+    # FIXME in what degree should we consider the integrity of the input?
+    option = parse_card_option(lines[icard])
+    if isempty(option)
+        @warn "Hubbard project is not specified!"
+    end
+
+    # number of hubbard lines is not known a priori
+    iend = end_of_card(lines, icard)
+    nline = valid_hubbard_lines(lines, icard, iend)
+    content = parse_card!(lines, name, nline)
+    card = OrderedDict{Symbol,Any}()
+    card[:option] = content[1]
+    card[:types] = map(x -> String(split(x)[1]), content[2])
+    card[:manifolds] = map(x -> join(split(x)[2:(end - 1)], " "), content[2])
+    card[:values] = map(x -> parse_float(split(x)[end]), content[2])
+    return Symbol(name) => card
+end
+
 """
     $(SIGNATURES)
 
@@ -317,6 +411,17 @@ function read_pw_in(io::Union{IO,AbstractString})
     return params
 end
 
+"""
+    $(SIGNATURES)
+
+Read HUBBARD.dat from hp.x calculations.
+"""
+function read_hubbard_dat(io::Union{IO, AbstractString})
+    _, cards = read_namelists(io; all_lines=true)
+    card = read_hubbard!(cards)
+    return card.second
+end
+
 """
     $(SIGNATURES)
 
@@ -445,7 +550,9 @@ function write_k_points(io::IO, card::AbstractDict; name::AbstractString="K_POIN
     option = get(card, "option", "")
     loption = lowercase(option)
 
-    valid_options_list = ["tpiba", "crystal", "tpiba_b", "crystal_b", "tpiba_c", "crystal_c"]
+    valid_options_list = [
+        "tpiba", "crystal", "tpiba_b", "crystal_b", "tpiba_c", "crystal_c"
+    ]
     valid_options_auto = ["automatic"]
     valid_options_gamma = ["gamma"]
     valid_options = vcat(valid_options_list, valid_options_auto, valid_options_gamma)
@@ -472,6 +579,48 @@ function write_k_points(io::IO, card::AbstractDict; name::AbstractString="K_POIN
     end
 end
 
+"""
+    $(SIGNATURES)
+
+Write the `hubbard` card of `pw.x`.
+
+# Examples
+```jldoctest; setup = :(using QuantumEspressoIO: write_hubbard)
+inputs = Dict(
+    :option => "atomic",
+    :types => ["U", "U"],
+    :values => [5.0, 5.0],
+    :manifolds => ["Ni-3d", "Ni1-3d"],
+)
+write_hubbard(stdout, inputs)
+# output
+HUBBARD atomic
+U Ni-3d 5.0
+U Ni1-3d 5.0
+```
+"""
+function write_hubbard(io::IO, card::AbstractDict)
+    option = get(card, :option, "")
+    # FIXME general question: how do we check for valid input?
+    isempty(option) && @error "Hubbard projector is not specified."
+    projector_types = Set(["atomic", "ortho-atomic", "norm-atomic", "wf", "pseudo"])
+    if option in projector_types
+        println(io, "HUBBARD $option")
+    else
+        @error "Hubbard projector $option is not valid."
+    end
+
+    valid_types = Set(["U", "J0", "J", "B", "E2", "E3", "V"])
+    for t in card[:types]
+        if !in(uppercase(t), valid_types)
+            @error "Hubbard type $t is not valid"
+        end
+    end
+    for (t, m, v) in zip(card[:types], card[:manifolds], card[:values])
+        println(io, "$(uppercase(t)) $m $v")
+    end
+end
+
 """
     $(SIGNATURES)