Accounting for FHI-aims species' defaults (#3877)

* Added species file class

* SAdded a list of species' defaults

* Added a test on species_defaults

* Added new species defaults to AimsControlIn input

* AIMS tests changed to account for new species' defaults object

pymatgen/io/aims/inputs.py

@@ -5,8 +5,8 @@
 
 from __future__ import annotations
 
-import gzip
 import os
+import re
 import time
 from copy import deepcopy
 from dataclasses import dataclass, field
@@ -16,8 +16,9 @@
 import numpy as np
 from monty.io import zopen
 from monty.json import MontyDecoder, MSONable
+from monty.os.path import zpath
 
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core import SETTINGS, Element, Lattice, Molecule, Structure
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -544,8 +545,10 @@ def get_content(
                 content += cube.control_block
 
         content += f"{lim}\n\n"
-        species_dir = self._parameters.get("species_dir", os.environ.get("AIMS_SPECIES_DIR"))
-        content += self.get_species_block(structure, species_dir)
+        species_defaults = self._parameters.get("species_dir", "")
+        if not species_defaults:
+            raise KeyError("Species' defaults not specified in the parameters")
+        content += self.get_species_block(structure, species_defaults)
 
         return content
 
@@ -591,33 +594,24 @@ def write_file(
 
             file.write(content)
 
-    def get_species_block(self, structure: Structure | Molecule, species_dir: str | Path) -> str:
+    def get_species_block(self, structure: Structure | Molecule, basis_set: str | dict[str, str]) -> str:
         """Get the basis set information for a structure
 
         Args:
             structure (Molecule or Structure): The structure to get the basis set information for
-            species_dir (str or Pat:): The directory to find the species files in
+            basis_set (str | dict[str, str]):
+                a name of a basis set (`light`, `tight`...) or a mapping from site labels to basis set names.
+                The name of a basis set can either correspond to the subfolder in `defaults_2020` folder
+                or be a full path from the `FHI-aims/species_defaults` directory.
 
         Returns:
             The block to add to the control.in file for the species
 
         Raises:
             ValueError: If a file for the species is not found
         """
-        block = ""
-        species = np.unique(structure.species)
-        for sp in species:
-            filename = f"{species_dir}/{sp.Z:02d}_{sp.symbol}_default"
-            if Path(filename).exists():
-                with open(filename) as sf:
-                    block += "".join(sf.readlines())
-            elif Path(f"{filename}.gz").exists():
-                with gzip.open(f"{filename}.gz", mode="rt") as sf:
-                    block += "".join(sf.readlines())
-            else:
-                raise ValueError(f"Species file for {sp.symbol} not found.")
-
-        return block
+        species_defaults = SpeciesDefaults.from_structure(structure, basis_set)
+        return str(species_defaults)
 
     def as_dict(self) -> dict[str, Any]:
         """Get a dictionary representation of the geometry.in file."""
@@ -640,3 +634,179 @@ def from_dict(cls, dct: dict[str, Any]) -> Self:
         decoded = {key: MontyDecoder().process_decoded(val) for key, val in dct.items() if not key.startswith("@")}
 
         return cls(_parameters=decoded["parameters"])
+
+
+class AimsSpeciesFile:
+    """An FHI-aims single species' defaults file."""
+
+    def __init__(self, data: str, label: str | None = None) -> None:
+        """
+        Args:
+            data (str): A string of the complete species defaults file
+            label (str): A string representing the name of species
+        """
+        self.data = data
+        self.label = label
+        if self.label is None:
+            for line in data.splitlines():
+                if "species" in line:
+                    self.label = line.split()[1]
+
+    @classmethod
+    def from_file(cls, filename: str, label: str | None = None) -> AimsSpeciesFile:
+        """Initialize from file.
+
+        Args:
+            filename (str): The filename of the species' defaults file
+            label (str): A string representing the name of species
+
+        Returns:
+            The AimsSpeciesFile instance
+        """
+        with zopen(filename, mode="rt") as file:
+            return cls(file.read(), label)
+
+    @classmethod
+    def from_element_and_basis_name(cls, element: str, basis: str, *, label: str | None = None) -> AimsSpeciesFile:
+        """Initialize from element and basis names.
+
+        Args:
+            element (str): the element name (not to confuse with the species)
+            basis (str): the directory in which the species' defaults file is located relative to the
+                root `species_defaults` (or `species_defaults/defaults_2020`) directory.`.
+            label (str): A string representing the name of species. If not specified,
+                then equal to element
+
+        Returns:
+            an AimsSpeciesFile instance
+        """
+        # check if element is in the Periodic Table (+ Emptium)
+        if element != "Emptium":
+            if not hasattr(Element, element):
+                raise ValueError(f"{element} is not a valid element name.")
+            el_obj = Element(element)
+            species_file_name = f"{el_obj.Z:02}_{element}_default"
+        else:
+            species_file_name = "00_Emptium_default"
+
+        aims_species_dir = SETTINGS.get("AIMS_SPECIES_DIR")
+        if aims_species_dir is None:
+            raise ValueError(
+                "No AIMS_SPECIES_DIR variable found in the config file. "
+                "Please set the variable in ~/.config/.pmgrc.yaml to the root of `species_defaults` "
+                "folder in FHIaims/ directory."
+            )
+        paths_to_try = [
+            (Path(aims_species_dir) / basis / species_file_name).expanduser().as_posix(),
+            (Path(aims_species_dir) / "defaults_2020" / basis / species_file_name).expanduser().as_posix(),
+        ]
+        for path in paths_to_try:
+            path = zpath(path)
+            if os.path.isfile(path):
+                return cls.from_file(path, label)
+
+        raise RuntimeError(
+            f"Can't find the species' defaults file for {element} in {basis} basis set. Paths tried: {paths_to_try}"
+        )
+
+    def __str__(self):
+        """String representation of the species' defaults file"""
+        return re.sub(r"^ *species +\w+", f"  species        {self.label}", self.data, flags=re.MULTILINE)
+
+    @property
+    def element(self) -> str:
+        match = re.search(r"^ *species +(\w+)", self.data, flags=re.MULTILINE)
+        if match is None:
+            raise ValueError("Can't find element in species' defaults file")
+        return match.group(1)
+
+    def as_dict(self) -> dict[str, Any]:
+        """Dictionary representation of the species' defaults file."""
+        return {"label": self.label, "data": self.data, "@module": type(self).__module__, "@class": type(self).__name__}
+
+    @classmethod
+    def from_dict(cls, dct: dict[str, Any]) -> AimsSpeciesFile:
+        """Deserialization of the AimsSpeciesFile object"""
+        return AimsSpeciesFile(data=dct["data"], label=dct["label"])
+
+
+class SpeciesDefaults(list, MSONable):
+    """A list containing a set of species' defaults objects with
+    methods to read and write them to files
+    """
+
+    def __init__(
+        self,
+        labels: Sequence[str],
+        basis_set: str | dict[str, str],
+        *,
+        elements: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Args:
+            labels (list[str]): a list of labels, for which to build species' defaults
+            basis_set (str | dict[str, str]):
+                a name of a basis set (`light`, `tight`...) or a mapping from site labels to basis set names.
+                The name of a basis set can either correspond to the subfolder in `defaults_2020` folder
+                or be a full path from the `FHI-aims/species_defaults` directory.
+            elements (dict[str, str] | None):
+                a mapping from site labels to elements. If some label is not in this mapping,
+                it coincides with an element.
+        """
+        super().__init__()
+        self.labels = labels
+        self.basis_set = basis_set
+        if elements is None:
+            elements = {}
+        self.elements = {}
+        for label in self.labels:
+            self.elements[label] = elements.get(label, label)
+        self._set_species()
+
+    def _set_species(self) -> None:
+        """Initialize species defaults from the instance data"""
+        del self[:]
+
+        for label in self.labels:
+            el = self.elements[label]
+            if isinstance(self.basis_set, dict):
+                basis_set = self.basis_set.get(label, None)
+                if basis_set is None:
+                    raise ValueError(f"Basis set not found for specie {label} (represented by element {el})")
+            else:
+                basis_set = self.basis_set
+            self.append(AimsSpeciesFile.from_element_and_basis_name(el, basis_set, label=label))
+
+    def __str__(self):
+        """String representation of the species' defaults"""
+        return "".join([str(x) for x in self])
+
+    @classmethod
+    def from_structure(cls, struct: Structure | Molecule, basis_set: str | dict[str, str]):
+        """Initialize species defaults from a structure."""
+        labels = []
+        elements = {}
+        for label, el in sorted(zip(struct.labels, struct.species)):
+            if not isinstance(el, Element):
+                raise TypeError("FHI-aims does not support fractional compositions")
+            if (label is None) or (el is None):
+                raise ValueError("Something is terribly wrong with the structure")
+            if label not in labels:
+                labels.append(label)
+                elements[label] = el.name
+        return SpeciesDefaults(labels, basis_set, elements=elements)
+
+    def to_dict(self):
+        """Dictionary representation of the species' defaults"""
+        return {
+            "labels": self.labels,
+            "elements": self.elements,
+            "basis_set": self.basis_set,
+            "@module": type(self).__module__,
+            "@class": type(self).__name__,
+        }
+
+    @classmethod
+    def from_dict(cls, dct: dict[str, Any]) -> SpeciesDefaults:
+        """Deserialization of the SpeciesDefaults object"""
+        return SpeciesDefaults(dct["labels"], dct["basis_set"], elements=dct["elements"])

tests/io/aims/conftest.py

@@ -4,9 +4,16 @@
 
 import pytest
 
+from pymatgen.core import SETTINGS
+
 module_dir = os.path.dirname(__file__)
 
 
 @pytest.fixture(autouse=True)
 def _set_aims_species_dir_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setenv("AIMS_SPECIES_DIR", f"{module_dir}/species_directory/light")
+    monkeypatch.setenv("AIMS_SPECIES_DIR", f"{module_dir}/species_directory")
+
+
+@pytest.fixture(autouse=True)
+def _set_aims_species_dir_settings(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setitem(SETTINGS, "AIMS_SPECIES_DIR", f"{module_dir}/species_directory")

tests/io/aims/test_aims_inputs.py

@@ -9,12 +9,15 @@
 from monty.json import MontyDecoder, MontyEncoder
 from numpy.testing import assert_allclose
 
+from pymatgen.core import SETTINGS
 from pymatgen.io.aims.inputs import (
     ALLOWED_AIMS_CUBE_TYPES,
     ALLOWED_AIMS_CUBE_TYPES_STATE,
     AimsControlIn,
     AimsCube,
     AimsGeometryIn,
+    AimsSpeciesFile,
+    SpeciesDefaults,
 )
 from pymatgen.util.testing.aims import compare_single_files as compare_files
 
@@ -164,7 +167,7 @@ def test_aims_control_in(tmp_path: Path):
         "compute_forces": True,
         "relax_geometry": ["trm", "1e-3"],
         "batch_size_limit": 200,
-        "species_dir": str(TEST_DIR.parent / "species_directory/light"),
+        "species_dir": "light",
     }
 
     aims_control = AimsControlIn(parameters.copy())
@@ -204,30 +207,27 @@ def test_aims_control_in(tmp_path: Path):
     compare_files(TEST_DIR / "control.in.si", f"{tmp_path}/control.in")
 
 
-def test_aims_control_in_default_species_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("AIMS_SPECIES_DIR", str(TEST_DIR.parent / "species_directory/light"))
+def test_species_file(monkeypatch: pytest.MonkeyPatch):
+    """Tests an AimsSpeciesFile class"""
+    monkeypatch.setitem(SETTINGS, "AIMS_SPECIES_DIR", str(TEST_DIR.parent / "species_directory"))
+    species_file = AimsSpeciesFile.from_element_and_basis_name("Si", "light", label="Si_surface")
+    assert species_file.label == "Si_surface"
+    assert species_file.element == "Si"
 
-    parameters = {
-        "cubes": [
-            AimsCube(type="eigenstate 1", points=[10, 10, 10]),
-            AimsCube(type="total_density", points=[10, 10, 10]),
-        ],
-        "xc": "LDA",
-        "smearing": ["fermi-dirac", 0.01],
-        "vdw_correction_hirshfeld": True,
-        "compute_forces": True,
-        "relax_geometry": ["trm", "1e-3"],
-        "batch_size_limit": 200,
-        "output": ["band 0 0 0 0.5 0 0.5 10 G X", "band 0 0 0 0.5 0.5 0.5 10 G L"],
-        "k_grid": [1, 1, 1],
-    }
-
-    aims_control = AimsControlIn(parameters.copy())
-
-    for key, val in parameters.items():
-        assert aims_control[key] == val
 
+def test_species_defaults(monkeypatch: pytest.MonkeyPatch):
+    """Tests an AimsSpeciesDefaults class"""
+    monkeypatch.setitem(SETTINGS, "AIMS_SPECIES_DIR", str(TEST_DIR.parent / "species_directory"))
     si = AimsGeometryIn.from_file(TEST_DIR / "geometry.in.si.gz").structure
-
-    aims_control.write_file(si, directory=tmp_path, verbose_header=True, overwrite=True)
-    compare_files(TEST_DIR / "control.in.si.no_sd", f"{tmp_path}/control.in")
+    species_defaults = SpeciesDefaults.from_structure(si, "light")
+    assert species_defaults.labels == [
+        "Si",
+    ]
+    assert species_defaults.elements == {"Si": "Si"}
+
+    si.relabel_sites()
+    species_defaults = SpeciesDefaults.from_structure(si, "light")
+    assert species_defaults.labels == ["Si_1", "Si_2"]
+    assert species_defaults.elements == {"Si_1": "Si", "Si_2": "Si"}
+    assert "Si_1" in str(species_defaults)
+    assert "Si_2" in str(species_defaults)

tests/io/aims/test_sets/test_relax_generator.py

@@ -12,29 +12,17 @@
 
 def test_relax_si(tmp_path):
     params = {
-        "species_dir": str(species_dir / "light"),
+        "species_dir": "light",
         "k_grid": [2, 2, 2],
     }
     comp_system(Si, params, "relax-si/", tmp_path, ref_path, RelaxSetGenerator)
 
 
 def test_relax_si_no_kgrid(tmp_path):
-    params = {"species_dir": str(species_dir / "light")}
+    params = {"species_dir": "light"}
     comp_system(Si, params, "relax-no-kgrid-si", tmp_path, ref_path, RelaxSetGenerator)
 
 
-def test_relax_default_species_dir(tmp_path):
-    params = {"k_grid": [2, 2, 2]}
-
-    comp_system(Si, params, "relax-si", tmp_path, ref_path, RelaxSetGenerator)
-
-
 def test_relax_o2(tmp_path):
     params = {"species_dir": str(species_dir / "light")}
     comp_system(O2, params, "relax-o2", tmp_path, ref_path, RelaxSetGenerator)
-
-
-def test_relax_default_species_dir_o2(tmp_path):
-    params = {"k_grid": [2, 2, 2]}
-
-    comp_system(O2, params, "relax-o2", tmp_path, ref_path, RelaxSetGenerator)