From f81692b3a623c08e2be3dcf30933951edb7db644 Mon Sep 17 00:00:00 2001
From: Austin Zadoks <austin.zadoks@protonmail.com>
Date: Fri, 20 Feb 2026 11:54:55 +0100
Subject: [PATCH] Add support for GIPAW_ORBITALS and GIPAW_VLOCAL

---
 src/file/upf.jl  | 24 ++++++++++++++
 src/file/upf2.jl | 84 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/src/file/upf.jl b/src/file/upf.jl
index 3f295ca..a046632 100644
--- a/src/file/upf.jl
+++ b/src/file/upf.jl
@@ -311,12 +311,36 @@ struct UpfGipawCoreOrbital
     core_orbital::Vector{Float64}
 end
 
+"""
+UPF `<PP_GIPAW/PP_GIPAW_ORBITALS/PP_GIPAW_ORBITAL.[i]>`
+"""
+struct UpfGipawOrbital
+    index::Int
+    label::Union{Nothing,String}
+    "Angular momentum"
+    l::Int
+    cutoff_radius::Union{Nothing,Float64}
+    ultrasoft_cutoff_radius::Union{Nothing,Float64}
+    wf_ae::Vector{Float64}
+    wf_ps::Vector{Float64}
+end
+
+"""
+UPF `<PP_GIPAW/PP_GIPAW_VLOCAL>`
+"""
+struct UpfGipawVlocal
+    vlocal_ae::Vector{Float64}
+    vlocal_ps::Vector{Float64}
+end
+
 """
 UPF `<PP_GIPAW>`
 """
 struct UpfGipaw
     gipaw_data_format::Int
     core_orbitals::Vector{UpfGipawCoreOrbital}
+    orbitals::Union{Nothing,Vector{UpfGipawOrbital}}
+    vlocal::Union{Nothing,UpfGipawVlocal}
 end
 
 """
diff --git a/src/file/upf2.jl b/src/file/upf2.jl
index e5f0165..0952a05 100644
--- a/src/file/upf2.jl
+++ b/src/file/upf2.jl
@@ -667,6 +667,59 @@ function upf2_dump_gipaw_core_orbital(core_orbital::UpfGipawCoreOrbital)::EzXML.
     return node
 end
 
+function upf2_parse_gipaw_orbital(node::EzXML.Node)
+    index = get_attr(Int, node, "index")
+    label = get_attr(String, node, "label")
+    l = Int(get_attr(Float64, node, "l"))
+    cutoff_radius = get_attr(Float64, node, "cutoff_radius")
+    ultrasoft_cutoff_radius = get_attr(Float64, node, "ultrasoft_cutoff_radius")
+    wf_ae_node = findfirst("PP_GIPAW_WFS_AE", node)
+    wf_ae = parse_nodecontent(Float64, wf_ae_node)
+    wf_ps_node = findfirst("PP_GIPAW_WFS_PS", node)
+    wf_ps = parse_nodecontent(Float64, wf_ps_node)
+    return UpfGipawOrbital(index, label, l, cutoff_radius, ultrasoft_cutoff_radius, wf_ae, wf_ps)
+end
+
+function upf2_dump_gipaw_orbital(orbital::UpfGipawOrbital)::EzXML.Node
+    node = ElementNode("PP_GIPAW_ORBITAL.$(orbital.index)")
+    for n in [n for n in fieldnames(UpfGipawOrbital) if n != :wf_ae && n != :wf_ps]
+        value = getfield(orbital, n)
+        if !isnothing(value)
+            set_attr!(node, n, value)
+        end
+    end
+    node_ae = ElementNode("PP_GIPAW_WFS_AE")
+    text_ae = array_to_text(orbital.wf_ae)
+    link!(node_ae, TextNode(text_ae))
+    link!(node, node_ae)
+    node_ps = ElementNode("PP_GIPAW_WFS_PS")
+    text_ps = array_to_text(orbital.wf_ps)
+    link!(node_ps, TextNode(text_ps))
+    link!(node, node_ps)
+    return node
+end
+
+function upf2_parse_gipaw_vlocal(node::EzXML.Node)
+    vlocal_ae_node = findfirst("PP_GIPAW_VLOCAL_AE", node)
+    vlocal_ae = parse_nodecontent(Float64, vlocal_ae_node)
+    vlocal_ps_node = findfirst("PP_GIPAW_VLOCAL_PS", node)
+    vlocal_ps = parse_nodecontent(Float64, vlocal_ps_node)
+    return UpfGipawVlocal(vlocal_ae, vlocal_ps)
+end
+
+function upf2_dump_gipaw_vlocal(vlocal::UpfGipawVlocal)::EzXML.Node
+    node = ElementNode("PP_GIPAW_VLOCAL")
+    node_ae = ElementNode("PP_GIPAW_VLOCAL_AE")
+    text_ae = array_to_text(vlocal.vlocal_ae)
+    link!(node_ae, TextNode(text_ae))
+    link!(node, node_ae)
+    node_ps = ElementNode("PP_GIPAW_VLOCAL_PS")
+    text_ps = array_to_text(vlocal.vlocal_ps)
+    link!(node_ps, TextNode(text_ps))
+    link!(node, node_ps)
+    return node
+end
+
 function upf2_parse_gipaw(node::EzXML.Node)
     gipaw_data_format = get_attr(Int, node, "gipaw_data_format")
     core_orbitals_node = findfirst("PP_GIPAW_CORE_ORBITALS", node)
@@ -674,7 +727,22 @@ function upf2_parse_gipaw(node::EzXML.Node)
                           for n in eachnode(core_orbitals_node)
                           if occursin("PP_GIPAW_CORE_ORBITAL.", nodename(n))]
     core_orbitals = upf2_parse_gipaw_core_orbital.(core_orbital_nodes)
-    return UpfGipaw(gipaw_data_format, core_orbitals)
+    orbitals_node = findfirst("PP_GIPAW_ORBITALS", node)
+    if isnothing(orbitals_node)
+        orbitals = nothing
+    else
+        orbital_nodes = [n
+                         for n in eachnode(orbitals_node)
+                         if occursin("PP_GIPAW_ORBITAL.", nodename(n))]
+        orbitals = upf2_parse_gipaw_orbital.(orbital_nodes)
+    end
+    vlocal_node = findfirst("PP_GIPAW_VLOCAL", node)
+    if isnothing(vlocal_node)
+        vlocal = nothing
+    else
+        vlocal = upf2_parse_gipaw_vlocal(vlocal_node)
+    end
+    return UpfGipaw(gipaw_data_format, core_orbitals, orbitals, vlocal)
 end
 upf2_parse_gipaw(doc::EzXML.Document) = upf2_parse_gipaw(findfirst("PP_GIPAW", root(doc)))
 
@@ -690,6 +758,20 @@ function upf2_dump_gipaw(gipaw::UpfGipaw)::EzXML.Node
     end
     link!(node, core_orbitals_node)
 
+    # PP_GIPAW_ORBITALS
+    if !isnothing(gipaw.orbitals)
+        orbitals_node = ElementNode("PP_GIPAW_ORBITALS")
+        for orbital in gipaw.orbitals
+            link!(orbitals_node, upf2_dump_gipaw_orbital(orbital))
+        end
+        link!(node, orbitals_node)
+    end
+
+    # PP_GIPAW_VLOCAL
+    if !isnothing(gipaw.vlocal)
+        link!(node, upf2_dump_gipaw_vlocal(gipaw.vlocal))
+    end
+
     return node
 end