Support an explicit frozen key on external projectors

src/aiida_wannier90_workflows/utils/pseudo/__init__.py

@@ -220,9 +220,16 @@ def get_frozen_list_ext(
 ) -> list:
     """Get frozen states (a subset of pseudo wavefunctions) in the external_projectors.
 
+    Whether a projector is frozen is decided by, in order of precedence:
+
+    1. an explicit boolean ``frozen`` key on the projector, if present;
+    2. otherwise the ``alpha`` key: ``"UPF"`` (or a missing ``alpha``) marks a
+       pseudo-atomic orbital, which is frozen; a numeric ``alpha`` marks a
+       generated (e.g. hydrogenic) orbital, which is not.
+
     :param structure: [description]
     :param external_projectors: dict of external projectors, where every external projector
-    contains the `label`, `l` and `j`(optional).
+    contains the `label`, `l` and `j`(optional), and optionally `frozen` and/or `alpha`.
     :param pseudo_orbitals: [description]
     :param spin_non_collinear: [description]
     :return: [description]
@@ -240,10 +247,9 @@ def get_frozen_list_ext(
             else:
                 num_orbs = (2 * orb["l"] + 1) * nspin
 
-            alpha = orb.get(
-                "alpha", "UPF"
-            )  # if not defined, it is better to Lowdin all projectors
-            if alpha == "UPF":
+            # if neither key is defined, it is better to Lowdin all projectors
+            frozen = orb.get("frozen", orb.get("alpha", "UPF") == "UPF")
+            if frozen:
                 frozen_list.extend(list(range(num_projs + 1, num_projs + num_orbs + 1)))
             num_projs += num_orbs
 

tests/utils/test_pseudo.py

@@ -0,0 +1,53 @@
+"""Tests for the :mod:`aiida_wannier90_workflows.utils.pseudo` module."""
+
+import pytest
+
+from aiida_wannier90_workflows.utils.pseudo import get_frozen_list_ext
+
+
+@pytest.mark.parametrize(
+    ("orbital", "expect_frozen"),
+    (
+        # No key at all: a pseudo-atomic orbital, frozen.
+        ({"label": "S", "l": 0}, True),
+        # Yuhao-protocol alpha bookkeeping: "UPF" marks the original
+        # pseudo-atomic orbital, a number marks a generated one.
+        ({"label": "S", "l": 0, "alpha": "UPF"}, True),
+        ({"label": "S", "l": 0, "alpha": 1.5}, False),
+        # An explicit ``frozen`` takes precedence over ``alpha``.
+        ({"label": "S", "l": 0, "frozen": False}, False),
+        ({"label": "S", "l": 0, "frozen": False, "alpha": "UPF"}, False),
+        ({"label": "S", "l": 0, "frozen": True, "alpha": 1.5}, True),
+    ),
+)
+def test_get_frozen_list_ext(generate_structure, orbital, expect_frozen):
+    """Test the per-orbital frozen selection of ``get_frozen_list_ext``."""
+    structure = generate_structure("Si")
+
+    frozen_list = get_frozen_list_ext(
+        structure=structure,
+        external_projectors={"Si": [orbital]},
+        spin_non_collinear=False,
+    )
+
+    # Bulk silicon has two sites; an s orbital contributes one projector each.
+    assert frozen_list == ([1, 2] if expect_frozen else [])
+
+
+def test_get_frozen_list_ext_mixed(generate_structure):
+    """Test frozen indexing across a mixed orbital table."""
+    structure = generate_structure("Si")
+
+    frozen_list = get_frozen_list_ext(
+        structure=structure,
+        external_projectors={
+            "Si": [
+                {"label": "S", "l": 0, "alpha": "UPF"},
+                {"label": "P", "l": 1, "frozen": False},
+            ]
+        },
+        spin_non_collinear=False,
+    )
+
+    # Per site: s (1 projector, frozen) then p (3 projectors, not frozen).
+    assert frozen_list == [1, 5]