From 682b260a29912a3513422b2b834c55a21031940e Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 16 Apr 2026 15:54:33 -0700 Subject: [PATCH 1/9] scaffold computed{structure}entry --- emmet-core/emmet/core/entry.py | 132 ++++++++++++++++++++++++++++ emmet-core/emmet/core/types/mson.py | 37 ++++++++ 2 files changed, 169 insertions(+) create mode 100644 emmet-core/emmet/core/entry.py create mode 100644 emmet-core/emmet/core/types/mson.py diff --git a/emmet-core/emmet/core/entry.py b/emmet-core/emmet/core/entry.py new file mode 100644 index 0000000000..1da076dc90 --- /dev/null +++ b/emmet-core/emmet/core/entry.py @@ -0,0 +1,132 @@ +"""Define schemas for pymatgen entry-like objects.""" + +from pydantic import BaseModel, Field, field_validator, field_serializer + +# from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry +from typing import Literal + +from emmet.core.mpid_ext import ThermoID +from emmet.core.types.enums import IgnoreCaseEnum +from emmet.core.types.mson import MSONType +from emmet.core.types.pymatgen_types.composition_adapter import CompositionType +from emmet.core.types.pymatgen_types.structure_adapter import StructureType +from emmet.core.types.typing import DateTimeType, IdentifierType, JsonDictType +from emmet.core.vasp.calc_types.enums import RunType +from emmet.core.vasp.calculation import PotcarSpec + +CORRECTION_NAME = { + "MP GGA(+U)/r2SCAN mixing adjustment", + "MP2020 GGA/GGA+U mixing correction (Co)", + "MP2020 GGA/GGA+U mixing correction (Cr)", + "MP2020 GGA/GGA+U mixing correction (Fe)", + "MP2020 GGA/GGA+U mixing correction (Mn)", + "MP2020 GGA/GGA+U mixing correction (Mo)", + "MP2020 GGA/GGA+U mixing correction (Ni)", + "MP2020 GGA/GGA+U mixing correction (V)", + "MP2020 GGA/GGA+U mixing correction (W)", + "MP2020 anion correction (Br)", + "MP2020 anion correction (Cl)", + "MP2020 anion correction (F)", + "MP2020 anion correction (H)", + "MP2020 anion correction (I)", + "MP2020 anion correction (N)", + "MP2020 anion correction (S)", + "MP2020 anion correction (Sb)", + "MP2020 anion correction (Se)", + "MP2020 anion correction (Si)", + "MP2020 anion correction (Te)", + "MP2020 anion correction (oxide)", + "MP2020 anion correction (ozonide)", + "MP2020 anion correction (peroxide)", + "MP2020 anion correction (superoxide)", +} +"""These are all the valid correction names that exist in our DB.""" + + +class OxideType(IgnoreCaseEnum): + """Define oxide types used in corrections schemes.""" + + HYDROXIDE = "hydroxide" + PEROXIDE = "peroxide" + SUPEROXIDE = "superoxide" + OXIDE = "oxide" + OZONIDE = "ozonide" + NONE = None + + @classmethod + def _missing_(cls, value): + if isinstance(value, str) and value == "None": + # pymatgen uses a str "None" instead of null + return cls.NONE + super(cls)._missing_(value) + + +class EntryParameters(BaseModel): + """Schematize entry parameters.""" + + run_type: RunType | None = None + is_hubbard: bool = False + hubbards: JsonDictType = None + potcar_spec: list[PotcarSpec] | None = None + + +class EntryData(BaseModel): + """Schematize entry run data.""" + + oxide_type: OxideType = OxideType.NONE + aspherical: bool = False + last_updated: DateTimeType + task_id: IdentifierType + material_id: IdentifierType + oxidation_states: JsonDictType + license: Literal["BY-C", "BY-NC"] = "BY-C" + run_type: RunType | None = None + + +class EnergyAdjustment(BaseModel): + """Schematize energy adjustment/correction from pymatgen.""" + + value: float | None = None + adj_per_atom: float | None = None + n_atoms: int | None = None + uncertainty_per_atom: float | None = None + name: Literal[*CORRECTION_NAME] = None + description: str | None = None + klass: MSONType | None = Field(None, validation_alias="cls") + + @property + def correction(self) -> float | None: + """Get the actual value of the correction.""" + if self.value is not None: + return self.value + elif self.adj_per_atom is not None and self.n_atoms is not None: + return self.adj_per_atom * self.n_atoms + return None + + +class Entry(BaseModel): + """Schematize pymatgen ComputedEntry.""" + + composition: CompositionType + + energy: float | None = None + correction: float | None = None + entry_id: ThermoID | None = None + + energy_adjustments: list[EnergyAdjustment] = Field([]) + parameters: EntryParameters | None = None + data: EntryData | None = None + + @field_validator("entry_id", mode="before") + def _deser_thermo_id(cls, v) -> ThermoID | None: + return ThermoID._deserialize(v) if v is not None else None + + @field_serializer("entry_id") + def _ser_thermo_id(self, v: ThermoID) -> str: + return str(v) + + +class StructureEntry(Entry): + """Schematize pymatgen ComputedStructureEntry.""" + + structure: StructureType diff --git a/emmet-core/emmet/core/types/mson.py b/emmet-core/emmet/core/types/mson.py new file mode 100644 index 0000000000..16aef45972 --- /dev/null +++ b/emmet-core/emmet/core/types/mson.py @@ -0,0 +1,37 @@ +"""Define generic monty serdes.""" + +from __future__ import annotations + +from pydantic import BaseModel, Field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +class MSONType(BaseModel): + + module: str = Field(validation_alias="@module") + klass: str = Field(validation_alias="@class") + version: str | None = Field(None, validation_alias="@version") + callable: str | None = Field(None, validation_alias="@callable") + bound: str | None = Field(None, validation_alias="@bound") + + def as_dict(self) -> dict[str, str | None]: + """Return MSON-style dict.""" + dct = { + field.validation_alias: getattr(self, k) + for k, field in self.__class__.model_fields.items() + } + return {k: v for k, v in dct.items() if v is not None} + + @classmethod + def from_dict(cls, dct) -> Any: + """Mimic monty decoding. + + Doesn't need to be a classmethod, but is included here + for duck-typing. + """ + from monty.json import MontyDecoder + + return MontyDecoder().process_decoded(dct) From ef9d3959f139b9ee9635093f9e6c41b80848c1e3 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 17 Apr 2026 16:56:54 -0700 Subject: [PATCH 2/9] base pymatgen object replacements --- emmet-core/emmet/core/atoms/__init__.py | 7 + emmet-core/emmet/core/atoms/base.py | 169 ++++++++++++++++++ emmet-core/emmet/core/atoms/elements.py | 173 ++++++++++++++++++ emmet-core/emmet/core/atoms/periodic.py | 225 ++++++++++++++++++++++++ 4 files changed, 574 insertions(+) create mode 100644 emmet-core/emmet/core/atoms/__init__.py create mode 100644 emmet-core/emmet/core/atoms/base.py create mode 100644 emmet-core/emmet/core/atoms/elements.py create mode 100644 emmet-core/emmet/core/atoms/periodic.py diff --git a/emmet-core/emmet/core/atoms/__init__.py b/emmet-core/emmet/core/atoms/__init__.py new file mode 100644 index 0000000000..2e1ab73ae2 --- /dev/null +++ b/emmet-core/emmet/core/atoms/__init__.py @@ -0,0 +1,7 @@ +"""Pull in important atomistic classes.""" + +from emmet.core.atoms.base import Molecule +from emmet.core.atoms.elements import Element +from emmet.core.atoms.periodic import Material + +__all__ = ["Element", "Material", "Molecule"] diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py new file mode 100644 index 0000000000..7ee27cc1c1 --- /dev/null +++ b/emmet-core/emmet/core/atoms/base.py @@ -0,0 +1,169 @@ +"""Define core atomistic data structures and analysis.""" + +from __future__ import annotations + +from functools import cached_property +from math import gcd +from typing import Any, TYPE_CHECKING + +import numpy as np +from pydantic import BaseModel + +from emmet.core.atoms.elements import Element, ELEMENT_DATA +from emmet.core.math import Matrix3D, Vector3D + +if TYPE_CHECKING: + from typing import Any + from typing_extensions import Self + + from pymatgen.core.sites import Site, PeriodicSite + from pymatgen.core.structure import Molecule as PmgMolecule + + +class Site(BaseModel): + """Schematize a site in a molecule or material.""" + + element: Element + cart_coords: Vector3D + charge: float | None = None + spin: float | None = None + degrees_of_freedom: tuple[bool, bool, bool] | None = None + velocity: Vector3D | None = None + + def __str__(self) -> str: + base_name = f"{self.element.name}" + if self.charge is not None: + charge_sign = "+" if self.charge >= 0 else "-" + if (abs_charge := abs(self.charge)) == 1: + base_name += charge_sign + else: + if (abs_charge - round(abs_charge)) < 1e-6: + abs_charge = round(abs_charge) + base_name += f"{abs_charge}{charge_sign}" + return base_name + + @property + def Z(self) -> int: + return ELEMENT_DATA[self.element].Z + + @classmethod + def from_pmg(cls, site: Site | PeriodicSite) -> Self: + + from pymatgen.core.periodic_table import Species + + if len(site.species.elements) > 1: + raise ValueError("`Site` currently cannot represent a disordered site!") + + charge = None + if isinstance(species := site.species.elements[0], Species): + charge = species.oxi_state + + return cls( + element=species.element.name, + cart_coords=site.coords, + charge=charge, + spin=site.properties.get("magmom"), + degrees_of_freedom=site.properties.get("selective_dynamics"), + velocity=site.properties.get("velocities"), + ) + + def to_pmg(self, cell: Matrix3D | None = None) -> Site | PeriodicSite: + + from pymatgen.core.sites import Site, PeriodicSite + + if cell is not None: + from pymatgen.core.lattice import Lattice + + species = {str(self): 1.0} + properties = { + v: getattr(self, k) + for k, v in { + "spin": "magmom", + "degrees_of_freedom": "selective_dynamics", + "velocity": "velocities", + }.items() + } + for k in list(properties): + if properties[k] is None: + properties.pop(k) + + if cell is None: + return Site( + species, + self.cart_coords, + properties=properties, + ) + + return PeriodicSite( + species, + self.cart_coords, + lattice=Lattice(cell), + coords_are_cartesian=True, + properties=properties, + ) + + +class Molecule(BaseModel): + """Schematize a molecular structure.""" + + sites: list[Site] + + def composition(self, include_charges: bool = True) -> dict[str, int]: + comp = { + str(site) if include_charges else site.element.name: 0 + for site in self.sites + } + for site in self.sites: + comp[str(site)] += 1 + return comp + + def reduced_composition(self, include_charges: bool = True) -> dict[str, int]: + base_comp = self.composition(include_charges=include_charges) + factor = gcd(*base_comp.values()) + return {k: v // factor for k, v in base_comp.items()} + + @cached_property + def mass(self) -> float: + """Mass in atomic mass units.""" + return sum(ELEMENT_DATA[site.element].atomic_mass for site in self.sites) + + def _aggregate_site_properties(self, prop: str, default: Any = None) -> np.ndarray: + return np.array([getattr(site, prop, None) or default for site in self.sites]) + + def _sum_scalar_site_properties(self, prop: str, default: float = 0.0) -> float: + return sum(getattr(site, prop, None) or default for site in self.sites) + + @cached_property + def cart_coords(self) -> np.ndarray[float]: + return self._aggregate_site_properties("cart_coords") + + @cached_property + def charge(self) -> float: + return self._sum_scalar_site_properties("charge") + + @cached_property + def spin(self) -> float: + return self._sum_scalar_site_properties("spin") + + @cached_property + def spins(self) -> np.ndarray[float]: + return self._aggregate_site_properties("spin") + + @cached_property + def degrees_of_freedom(self) -> np.ndarray[bool]: + return self._aggregate_site_properties( + "degrees_of_freedom", default=(True, True, True) + ) + + @classmethod + def from_pmg(cls, molecule: PmgMolecule) -> Self: + return cls( + sites=[Site.from_pmg(site) for site in molecule], + ) + + def to_pmg(self) -> PmgMolecule: + from pymatgen.core.structure import Molecule as PmgMolecule + + return PmgMolecule.from_sites( + [site.to_pmg(cell=self.cell) for site in self.sites] + ) diff --git a/emmet-core/emmet/core/atoms/elements.py b/emmet-core/emmet/core/atoms/elements.py new file mode 100644 index 0000000000..f61174d890 --- /dev/null +++ b/emmet-core/emmet/core/atoms/elements.py @@ -0,0 +1,173 @@ +"""Define elements / isotopes.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import TYPE_CHECKING + +from pydantic import BaseModel, Field + +if TYPE_CHECKING: + from typing import Any + + +class Element(StrEnum): + """Map short chemical symbol to long name.""" + + H = "Hydrogen" + He = "Helium" + Li = "Lithium" + Be = "Beryllium" + B = "Boron" + C = "Carbon" + N = "Nitrogen" + O = "Oxygen" + F = "Fluorine" + Ne = "Neon" + Na = "Sodium" + Mg = "Magnesium" + Al = "Aluminum" + Si = "Silicon" + P = "Phosphorus" + S = "Sulfur" + Cl = "Chlorine" + Ar = "Argon" + K = "Potassium" + Ca = "Calcium" + Sc = "Scandium" + Ti = "Titanium" + V = "Vanadium" + Cr = "Chromium" + Mn = "Manganese" + Fe = "Iron" + Co = "Cobalt" + Ni = "Nickel" + Cu = "Copper" + Zn = "Zinc" + Ga = "Gallium" + Ge = "Germanium" + As = "Arsenic" + Se = "Selenium" + Br = "Bromine" + Kr = "Krypton" + Rb = "Rubidium" + Sr = "Strontium" + Y = "Yttrium" + Zr = "Zirconium" + Nb = "Niobium" + Mo = "Molybdenum" + Tc = "Technetium" + Ru = "Ruthenium" + Rh = "Rhodium" + Pd = "Palladium" + Ag = "Silver" + Cd = "Cadmium" + In = "Indium" + Sn = "Tin" + Sb = "Antimony" + Te = "Tellurium" + I = "Iodine" + Xe = "Xenon" + Cs = "Cesium" + Ba = "Barium" + La = "Lanthanum" + Ce = "Cerium" + Pr = "Praseodymium" + Nd = "Neodymium" + Pm = "Promethium" + Sm = "Samarium" + Eu = "Europium" + Gd = "Gadolinium" + Tb = "Terbium" + Dy = "Dysprosium" + Ho = "Holmium" + Er = "Erbium" + Tm = "Thulium" + Yb = "Ytterbium" + Lu = "Lutetium" + Hf = "Hafnium" + Ta = "Tantalum" + W = "Tungsten" + Re = "Rhenium" + Os = "Osmium" + Ir = "Iridium" + Pt = "Platinum" + Au = "Gold" + Hg = "Mercury" + Tl = "Thallium" + Pb = "Lead" + Bi = "Bismuth" + Po = "Polonium" + At = "Astatine" + Rn = "Radon" + Fr = "Francium" + Ra = "Radium" + Ac = "Actinium" + Th = "Thorium" + Pa = "Protactinium" + U = "Uranium" + Np = "Neptunium" + Pu = "Plutonium" + Am = "Americium" + Cm = "Curium" + Bk = "Berkelium" + Cf = "Californium" + Es = "Einsteinium" + Fm = "Fermium" + Md = "Mendelevium" + No = "Nobelium" + Lr = "Lawrencium" + Rf = "Rutherfordium" + Db = "Dubnium" + Sg = "Seaborgium" + Bh = "Bohrium" + Hs = "Hassium" + Mt = "Meitnerium" + Ds = "Darmstadtium" + Rg = "Roentgenium" + Cn = "Copernicium" + Nh = "Nihonium" + Fl = "Flerovium" + Mc = "Moscovium" + Lv = "Livermorium" + Ts = "Tennessine" + Og = "Oganesson" + + @classmethod + def _missing_(cls, value: Any) -> "Element" | None: + """Permit search for element based on symbol or name.""" + if value in cls: + return cls(value) + elif value in cls.__members__: + return cls[value] + + +class ElementData(BaseModel): + """Data for the elements.""" + + Z: int = Field(description="The number of protons in this element") + atomic_mass: float = Field( + description="The atomic mass (number of protons + neutrons) in atomic mass units (amu)" + ) + + +class ElementDatabase(dict): + + def _load_data(self) -> dict[Element, ElementData]: + """Cache atom data from pymatgen.""" + from pymatgen.core.periodic_table import Element as PmgElement + + # ignore isotopes + type_map = {e: PmgElement(e.name) for e in Element} + self.update( + { + ele: ElementData(Z=pmg_ele.Z, atomic_mass=pmg_ele.atomic_mass) + for ele, pmg_ele in type_map.items() + } + ) + + def __init__(self, **kwargs) -> None: + self._load_data() + + +ELEMENT_DATA = ElementDatabase() diff --git a/emmet-core/emmet/core/atoms/periodic.py b/emmet-core/emmet/core/atoms/periodic.py new file mode 100644 index 0000000000..54e366c834 --- /dev/null +++ b/emmet-core/emmet/core/atoms/periodic.py @@ -0,0 +1,225 @@ +"""Define data structures for periodic materials.""" + +from __future__ import annotations + +from functools import cached_property +import re +from typing import TYPE_CHECKING + +import numpy as np +from scipy.constants import atomic_mass +import spglib + +from emmet.core.atoms.base import Molecule, Site +from emmet.core.atoms.elements import ELEMENT_DATA +from emmet.core.math import Matrix3D +from emmet.core.settings import EmmetSettings + +if TYPE_CHECKING: + from typing import Literal + from typing_extensions import Self + + from pymatgen.core.structure import Structure + +SETTINGS = EmmetSettings() + + +class Cell(np.ndarray): + + def __new__(cls, data, **kwargs): + arr = np.asarray(data, dtype=float) + return arr.view(cls) + + def __array_finalize__(self, obj): + # If shape/dtype are wrong, demote to plain ndarray instead of raising + if self.shape != (3, 3) or self.dtype != np.float64: + # Can't mutate self's type in-place, so we flag it for __array_ufunc__ + self._invalid = True + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + # Delegate to plain ndarray, then re-wrap only if result is still 3x3 float + plain_inputs = [np.asarray(x) for x in inputs] + result = getattr(ufunc, method)(*plain_inputs, **kwargs) + if ( + isinstance(result, np.ndarray) + and result.shape == (3, 3) + and result.dtype == np.float64 + ): + return result.view(Cell) + return result # return plain ndarray if result doesn't qualify + + @cached_property + def volume(self) -> float: + return abs(np.linalg.det(self)) + + @cached_property + def _reciprocal(self) -> Self: + return Cell( + np.array([np.cross(self[(i + 1) % 3], self[(i + 2) % 3]) for i in range(3)]) + / self.volume + ) + + @property + def reciprocal(self) -> Self: + return 2 * np.pi * self._reciprocal + + @cached_property + def _vector_norms(self) -> np.ndarray[float]: + return np.linalg.norm(self, axis=1) + + @cached_property + def _angles(self) -> np.ndarray: + return [ + 180 + / np.pi + * np.arccos( + np.dot(self.matrix[i], self.matrix[(i + 1) % 3]) + / (self._vector_norms[i] * self._vector_norms[(i + 1) % 3]) + ) + for i in range(3) + ] + + @property + def a(self) -> float: + return self._vector_norms[0] + + @property + def b(self) -> float: + return self._vector_norms[1] + + @property + def c(self) -> float: + return self._vector_norms[2] + + @property + def alpha(self) -> float: + return self._angles[1] + + @property + def beta(self) -> float: + return self._angles[2] + + @property + def gamma(self) -> float: + return self._angles[0] + + @staticmethod + def _get_coords( + cell: Cell, coords: np.ndarray, to: Literal["cartesian", "direct"] + ) -> np.ndarray[float]: + if to == "direct": + return np.einsum("ij,ki->kj", cell._reciprocal.T, coords) + elif to == "cartesian": + return np.einsum("ij,ki->kj", cell, coords) + raise ValueError( + f'Unknown transformation {to}. Please select "cartesian" or "direct".' + ) + + def get_coords( + self, coords: np.ndarray, to: Literal["cartesian", "direct"] + ) -> np.ndarray[float]: + return self._get_coords(self, coords, to=to) + + +class Material(Molecule): + """Schematize an ordered material crystal structure.""" + + lattice: Matrix3D + + @cached_property + def cell(self) -> Cell: + return Cell(self.lattice) + + @property + def volume(self) -> float: + return self.cell.volume + + @cached_property + def frac_coords(self) -> np.ndarray[float]: + return self.cell.get_coords(self.cart_coords, to="direct") + + def density_g_cm3(self) -> float: + """Get density of material in g/cm^3.""" + return self.mass * atomic_mass * 1e27 / self.cell.volume + + def __len__( + self, + ) -> int: + return len(self.sites) + + @property + def num_sites(self) -> int: + return len(self) + + @classmethod + def from_pmg(cls, structure: Structure) -> Self: + return cls( + lattice=structure.lattice.matrix, + sites=[Site.from_pmg(site) for site in structure], + ) + + def to_pmg(self) -> Structure: + from pymatgen.core.structure import Structure + + return Structure.from_sites( + [site.to_pmg(cell=self.cell) for site in self.sites] + ) + + @cached_property + def _to_spglib(self) -> tuple[Cell, np.ndarray[float], np.ndarray[int]]: + """Create an spglib-compatible representation of the atoms.""" + return ( + self.cell, + self.frac_coords, + [site.Z for site in self.sites], + ) + + @classmethod + def _from_spglib( + cls, spglib_rep: tuple[Matrix3D, np.ndarray[float], np.ndarray[int]] + ) -> Self: + cell, frac_coords, atomic_numbers = spglib_rep + + zmap = { + z: next(ele for ele, data in ELEMENT_DATA.items() if data.Z == z) + for z in set(atomic_numbers) + } + cart_coords = Cell(cell).get_coords(frac_coords, to="cartesian") + return cls( + lattice=cell, + sites=[ + Site( + element=zmap[z], + cart_coords=cart_coords[idx], + ) + for idx, z in enumerate(atomic_numbers) + ], + ) + + def primitive( + self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL + ) -> Material: + return self._from_spglib( + spglib.find_primitive( + self._to_spglib, + symprec=symprec, + angle_tolerance=angle_tol, + ) + ) + + def conventional( + self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL + ) -> Material: + return self._from_spglib( + spglib.standardize_cell( + self._to_spglib, symprec=symprec, angle_tol=angle_tol + ) + ) + + def get_space_group_info( + self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL + ) -> tuple[str, int]: + sg_info = spglib.get_spacegroup( + self._to_spglib, symprec=symprec, angle_tolerance=angle_tol + ) + return tuple(re.match(r"(.*) \((.*)\)", sg_info).groups()) From a5663b2e36d064fcedac44a715462f78714ba97d Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 20 Apr 2026 14:13:55 -0700 Subject: [PATCH 3/9] basic compound replacement for composition --- emmet-core/emmet/core/atoms/base.py | 123 ++++++++++++++++++++---- emmet-core/emmet/core/atoms/elements.py | 22 +++++ emmet-core/emmet/core/atoms/periodic.py | 11 +-- 3 files changed, 129 insertions(+), 27 deletions(-) diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index 7ee27cc1c1..896bc1858d 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -7,15 +7,16 @@ from typing import Any, TYPE_CHECKING import numpy as np -from pydantic import BaseModel +from pydantic import BaseModel, model_validator -from emmet.core.atoms.elements import Element, ELEMENT_DATA +from emmet.core.atoms.elements import Element, ELEMENT_DATA, parse_species_str from emmet.core.math import Matrix3D, Vector3D if TYPE_CHECKING: from typing import Any from typing_extensions import Self + from pymatgen.core.composition import Composition from pymatgen.core.sites import Site, PeriodicSite from pymatgen.core.structure import Molecule as PmgMolecule @@ -103,29 +104,117 @@ def to_pmg(self, cell: Matrix3D | None = None) -> Site | PeriodicSite: ) -class Molecule(BaseModel): - """Schematize a molecular structure.""" +class Compound(BaseModel): - sites: list[Site] + species: list[str] + coefficients: list[int] - def composition(self, include_charges: bool = True) -> dict[str, int]: - comp = { - str(site) if include_charges else site.element.name: 0 - for site in self.sites + @model_validator(mode="before") + @classmethod + def _reduce(cls, config) -> Self: + + if not all(config.get(k) for k in cls.model_fields) or len( + config["species"] + ) != len(config["coefficients"]): + raise ValueError(f"Invalid input specified to {cls.__name__}.") + + base_config: dict[str, int] = {} + for idx, spec in enumerate(config["species"]): + if spec not in base_config: + base_config[spec] = 0 + base_config[spec] += config["coefficients"][idx] + sorted_species = sorted(base_config.keys()) + return { + "species": sorted_species, + "coefficients": [base_config[spec] for spec in sorted_species], } - for site in self.sites: - comp[str(site)] += 1 - return comp - def reduced_composition(self, include_charges: bool = True) -> dict[str, int]: - base_comp = self.composition(include_charges=include_charges) - factor = gcd(*base_comp.values()) - return {k: v // factor for k, v in base_comp.items()} + def __str__(self) -> str: + return ( + f"{self.__class__.__name__}(" + + ", ".join( + f"{spec}: {self.coefficients[idx]}" + for idx, spec in enumerate(self.species) + ) + + ")" + ) + + def __repr__(self) -> str: + return self.__str__() + + @cached_property + def elements(self) -> list[Element]: + return [parse_species_str(spec)[0] for spec in self.species] + + @classmethod + def from_dict(cls, dct: dict[str, int]): + ordered_species = sorted(dct.keys()) + return cls( + species=ordered_species, + coefficients=[dct[k] for k in ordered_species], + ) + + def to_dict(self) -> dict[str, int]: + return dict( + [(spec, self.coefficients[idx]) for idx, spec in enumerate(self.species)] + ) + + def to_pmg(self) -> Composition: + from pymatgen.core.composition import Composition + + return Composition(self.to_dict()) + + @property + def reduced(self) -> Compound: + factor = gcd(*self.coefficients) + return Compound( + species=self.species, coefficients=[v // factor for v in self.coefficients] + ) @cached_property def mass(self) -> float: """Mass in atomic mass units.""" - return sum(ELEMENT_DATA[site.element].atomic_mass for site in self.sites) + return sum(ELEMENT_DATA[ele].atomic_mass for ele in self.elements) + + def __getitem__(self, species: str) -> Any: + """Return coefficient of species if present, otherwise raise an exception.""" + if species in self.species: + return next( + self.coefficients[idx] + for idx, spec in enumerate(self.species) + if spec == species + ) + raise KeyError(species) + + def get(self, item: str, default: Any = None) -> Any: + """Return a model field `item`, or `default` if it doesn't exist.""" + try: + return self.__getitem__(item) + except KeyError: + return default + + +class Molecule(BaseModel): + """Schematize a molecular structure.""" + + sites: list[Site] + + def __len__(self) -> int: + return len(self.sites) + + @property + def num_sites(self) -> int: + return len(self) + + @property + def composition(self) -> Compound: + return Compound( + species=[str(site) for site in self.sites], coefficients=[1] * len(self) + ) + + @property + def reduced_composition(self) -> Compound: + return self.composition.reduced def _aggregate_site_properties(self, prop: str, default: Any = None) -> np.ndarray: return np.array([getattr(site, prop, None) or default for site in self.sites]) diff --git a/emmet-core/emmet/core/atoms/elements.py b/emmet-core/emmet/core/atoms/elements.py index f61174d890..6a53fcef2a 100644 --- a/emmet-core/emmet/core/atoms/elements.py +++ b/emmet-core/emmet/core/atoms/elements.py @@ -3,6 +3,7 @@ from __future__ import annotations from enum import StrEnum +import re from typing import TYPE_CHECKING from pydantic import BaseModel, Field @@ -171,3 +172,24 @@ def __init__(self, **kwargs) -> None: ELEMENT_DATA = ElementDatabase() + + +def parse_species_str(rep: str) -> tuple[Element, float]: + """Parse an Atom from a string, including possible oxieation state.""" + _parsed = re.match(r"([A-Z][a-z]?)([0-9.0-9]+)?([+-])?", rep) + if not _parsed: + raise ValueError(f"Unknown element symbol {rep}") + parsed = _parsed.groups() + charge_str = parsed[1] + + charge_sign = "+" + if parsed[2] is not None: + charge_sign = parsed[2] + if charge_str is None: + charge_str = "1" + + charge = 0.0 + if charge_str is not None: + charge = float(charge_sign + charge_str) + + return Element[parsed[0]], charge diff --git a/emmet-core/emmet/core/atoms/periodic.py b/emmet-core/emmet/core/atoms/periodic.py index 54e366c834..934f3a6507 100644 --- a/emmet-core/emmet/core/atoms/periodic.py +++ b/emmet-core/emmet/core/atoms/periodic.py @@ -140,16 +140,7 @@ def frac_coords(self) -> np.ndarray[float]: def density_g_cm3(self) -> float: """Get density of material in g/cm^3.""" - return self.mass * atomic_mass * 1e27 / self.cell.volume - - def __len__( - self, - ) -> int: - return len(self.sites) - - @property - def num_sites(self) -> int: - return len(self) + return self.composition.mass * atomic_mass * 1e27 / self.cell.volume @classmethod def from_pmg(cls, structure: Structure) -> Self: From 47455a1448abf600144962f9a9e18fb536233b6b Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Tue, 28 Apr 2026 11:33:06 -0700 Subject: [PATCH 4/9] ensure entries use material class --- emmet-core/emmet/core/atoms/base.py | 6 ++++- emmet-core/emmet/core/entry.py | 42 ++++++++++++++++++++++------- 2 files changed, 38 insertions(+), 10 deletions(-) diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index 896bc1858d..65ede9a327 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -60,7 +60,7 @@ def from_pmg(cls, site: Site | PeriodicSite) -> Self: charge = species.oxi_state return cls( - element=species.element.name, + element=getattr(species, "element", species).name, cart_coords=site.coords, charge=charge, spin=site.properties.get("magmom"), @@ -159,6 +159,10 @@ def to_dict(self) -> dict[str, int]: [(spec, self.coefficients[idx]) for idx, spec in enumerate(self.species)] ) + @classmethod + def from_pmg(cls, comp: Composition) -> Self: + return cls.from_dict(comp.as_dict()) + def to_pmg(self) -> Composition: from pymatgen.core.composition import Composition diff --git a/emmet-core/emmet/core/entry.py b/emmet-core/emmet/core/entry.py index 1da076dc90..66f5fb2827 100644 --- a/emmet-core/emmet/core/entry.py +++ b/emmet-core/emmet/core/entry.py @@ -1,19 +1,25 @@ """Define schemas for pymatgen entry-like objects.""" +from __future__ import annotations + +from importlib import import_module from pydantic import BaseModel, Field, field_validator, field_serializer -# from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry -from typing import Literal +from typing import Literal, TYPE_CHECKING +from emmet.core.atoms.base import Compound +from emmet.core.atoms.periodic import Material from emmet.core.mpid_ext import ThermoID from emmet.core.types.enums import IgnoreCaseEnum from emmet.core.types.mson import MSONType -from emmet.core.types.pymatgen_types.composition_adapter import CompositionType -from emmet.core.types.pymatgen_types.structure_adapter import StructureType from emmet.core.types.typing import DateTimeType, IdentifierType, JsonDictType from emmet.core.vasp.calc_types.enums import RunType from emmet.core.vasp.calculation import PotcarSpec +if TYPE_CHECKING: + from typing_extensions import Self + from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry + CORRECTION_NAME = { "MP GGA(+U)/r2SCAN mixing adjustment", "MP2020 GGA/GGA+U mixing correction (Co)", @@ -76,9 +82,9 @@ class EntryData(BaseModel): oxide_type: OxideType = OxideType.NONE aspherical: bool = False last_updated: DateTimeType - task_id: IdentifierType - material_id: IdentifierType - oxidation_states: JsonDictType + task_id: IdentifierType = None + material_id: IdentifierType = None + oxidation_states: JsonDictType = None license: Literal["BY-C", "BY-NC"] = "BY-C" run_type: RunType | None = None @@ -107,7 +113,7 @@ def correction(self) -> float | None: class Entry(BaseModel): """Schematize pymatgen ComputedEntry.""" - composition: CompositionType + composition: Compound energy: float | None = None correction: float | None = None @@ -125,8 +131,26 @@ def _deser_thermo_id(cls, v) -> ThermoID | None: def _ser_thermo_id(self, v: ThermoID) -> str: return str(v) + def to_pmg(self) -> ComputedEntry | ComputedStructureEntry: + data = self.model_dump() + pmg_cls = "ComputedEntry" + if data.get("structure"): + data["structure"] = Material(data["structure"]).to_pmg() + pmg_cls = "ComputedStructureEntry" + + pmg_entries = import_module("pymatgen.entries.computed_entries") + return getattr(pmg_entries, pmg_cls).from_dict(data) + + @classmethod + def from_pmg(cls, entry: ComputedEntry | ComputedStructureEntry) -> Self: + config = entry.as_dict() + config["composition"] = Compound.from_dict(config["composition"]) + if config.get("structure"): + config["structure"] = Material.from_pmg(entry.structure) + return cls(**config) + class StructureEntry(Entry): """Schematize pymatgen ComputedStructureEntry.""" - structure: StructureType + structure: Material From 0dd3e8249a7d26addd6ea4dd371aad0dbbf73d87 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Tue, 28 Apr 2026 13:30:38 -0700 Subject: [PATCH 5/9] start migrating away from pmg classes --- emmet-core/emmet/core/atoms/base.py | 26 +++++++++++++-------- emmet-core/emmet/core/atoms/periodic.py | 10 ++++++-- emmet-core/emmet/core/chemenv.py | 6 ++--- emmet-core/emmet/core/grain_boundary.py | 11 +++++---- emmet-core/emmet/core/io/__init__.py | 0 emmet-core/emmet/core/io/pymatgen.py | 17 ++++++++++++++ emmet-core/emmet/core/mpcomplete.py | 4 ++-- emmet-core/emmet/core/surface_properties.py | 15 ++++++------ 8 files changed, 60 insertions(+), 29 deletions(-) create mode 100644 emmet-core/emmet/core/io/__init__.py create mode 100644 emmet-core/emmet/core/io/pymatgen.py diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index 65ede9a327..941351c045 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -16,10 +16,7 @@ from typing import Any from typing_extensions import Self - from pymatgen.core.composition import Composition - from pymatgen.core.sites import Site, PeriodicSite - from pymatgen.core.structure import Molecule as PmgMolecule - + from emmet.core.io.pymatgen import Composition, PeriodicSite, PmgMolecule, PmgSite class Site(BaseModel): """Schematize a site in a molecule or material.""" @@ -48,9 +45,9 @@ def Z(self) -> int: return ELEMENT_DATA[self.element].Z @classmethod - def from_pmg(cls, site: Site | PeriodicSite) -> Self: + def from_pmg(cls, site: PmgSite | PeriodicSite) -> Self: - from pymatgen.core.periodic_table import Species + from emmet.core.io.pymatgen import Species if len(site.species.elements) > 1: raise ValueError("`Site` currently cannot represent a disordered site!") @@ -68,12 +65,12 @@ def from_pmg(cls, site: Site | PeriodicSite) -> Self: velocity=site.properties.get("velocities"), ) - def to_pmg(self, cell: Matrix3D | None = None) -> Site | PeriodicSite: + def to_pmg(self, cell: Matrix3D | None = None) -> PmgSite | PeriodicSite: - from pymatgen.core.sites import Site, PeriodicSite + from emmet.core.io.pymatgen import PmgSite, PeriodicSite if cell is not None: - from pymatgen.core.lattice import Lattice + from emmet.core.io.pymatgen import Lattice species = {str(self): 1.0} properties = { @@ -89,7 +86,7 @@ def to_pmg(self, cell: Matrix3D | None = None) -> Site | PeriodicSite: properties.pop(k) if cell is None: - return Site( + return PmgSite( species, self.cart_coords, properties=properties, @@ -260,3 +257,12 @@ def to_pmg(self) -> PmgMolecule: return PmgMolecule.from_sites( [site.to_pmg(cell=self.cell) for site in self.sites] ) + + @classmethod + def from_sites(cls, sites : list[PmgSite | Site]) -> Self: + from py + return cls( + sites = [ + Site.from_pmg(site) if + ] + ) diff --git a/emmet-core/emmet/core/atoms/periodic.py b/emmet-core/emmet/core/atoms/periodic.py index 934f3a6507..a520270dd8 100644 --- a/emmet-core/emmet/core/atoms/periodic.py +++ b/emmet-core/emmet/core/atoms/periodic.py @@ -19,7 +19,7 @@ from typing import Literal from typing_extensions import Self - from pymatgen.core.structure import Structure + from emmet.core.io.pymatgen import Structure SETTINGS = EmmetSettings() @@ -150,7 +150,7 @@ def from_pmg(cls, structure: Structure) -> Self: ) def to_pmg(self) -> Structure: - from pymatgen.core.structure import Structure + from emmet.core.io.pymatgen import Structure return Structure.from_sites( [site.to_pmg(cell=self.cell) for site in self.sites] @@ -214,3 +214,9 @@ def get_space_group_info( self._to_spglib, symprec=symprec, angle_tolerance=angle_tol ) return tuple(re.match(r"(.*) \((.*)\)", sg_info).groups()) + + @classmethod + def from_cif(cls, cif_str: str) -> Self: + from emmet.core.io.pymatgen import Structure + + return Material.from_pmg(Structure.from_str(cif_str, fmt="str")) diff --git a/emmet-core/emmet/core/chemenv.py b/emmet-core/emmet/core/chemenv.py index 6e8ac8c8ba..7d49267282 100644 --- a/emmet-core/emmet/core/chemenv.py +++ b/emmet-core/emmet/core/chemenv.py @@ -16,15 +16,15 @@ LightStructureEnvironments, ) from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer -from pymatgen.core.structure import Molecule, Structure +from pymatgen.core.structure import Structure +from emmet.core.atoms.base import Molecule from emmet.core.material_property import PropertyDoc if TYPE_CHECKING: from emmet.core.types.typing import IdentifierType from emmet.core.types.pymatgen_types.structure_adapter import ( - MoleculeType, StructureType, ) @@ -367,7 +367,7 @@ class ChemEnvDoc(PropertyDoc): description="Method used to compute chemical environments" ) - mol_from_site_environments: list[MoleculeType | None] = Field( + mol_from_site_environments: list[Molecule | None] = Field( description="List of Molecule Objects describing the detected environment." ) diff --git a/emmet-core/emmet/core/grain_boundary.py b/emmet-core/emmet/core/grain_boundary.py index 89ce842ade..827722077a 100644 --- a/emmet-core/emmet/core/grain_boundary.py +++ b/emmet-core/emmet/core/grain_boundary.py @@ -3,11 +3,10 @@ from typing import TYPE_CHECKING from pydantic import BaseModel, Field -from pymatgen.core import Structure +from emmet.core.atoms.periodic import Material from emmet.core.types.enums import ValueEnum from emmet.core.types.pymatgen_types.grain_boundary_adapter import GrainBoundaryType -from emmet.core.types.pymatgen_types.structure_adapter import StructureType from emmet.core.types.typing import DateTimeType, MaterialIdentifierType if TYPE_CHECKING: @@ -79,7 +78,7 @@ class GrainBoundaryDoc(BaseModel): w_sep: float | None = Field(None, description="Work of separation in J/m^2.") - structure: StructureType | None = Field(None, description="Structure.") + structure: Material | None = Field(None, description="Structure.") chemsys: str | None = Field( None, description="Dash-delimited string of elements in the material." @@ -95,12 +94,14 @@ def _migrate_schema(cls, config: dict[str, Any]) -> Self: if isinstance(cif_str := config.pop("cif", None), str) and not config.get( "structure" ): - config["structure"] = Structure.from_str(cif_str, fmt="cif") + from emmet.core.io.pymatgen import cif_to_material + + config["structure"] = cif_to_material(cif_str) return cls(**config) @property def cif(self) -> str | None: """Support accessing legacy CIF from structure field.""" if self.structure: - return self.structure.to(fmt="cif") + return self.structure.to_pmg().to(fmt="cif") return None diff --git a/emmet-core/emmet/core/io/__init__.py b/emmet-core/emmet/core/io/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/emmet-core/emmet/core/io/pymatgen.py b/emmet-core/emmet/core/io/pymatgen.py new file mode 100644 index 0000000000..aacab87a0c --- /dev/null +++ b/emmet-core/emmet/core/io/pymatgen.py @@ -0,0 +1,17 @@ +"""Define interfaces to pymatgen.""" + +from pymatgen.core.composition import Composition +from pymatgen.core.lattice import Lattice +from pymatgen.core.periodic_table import Species +from pymatgen.core.sites import Site as PmgSite, PeriodicSite +from pymatgen.core.structure import Molecule as PmgMolecule, Structure + +__all__ = [ + "Composition", + "Lattice", + "PmgSite", + "PeriodicSite", + "PmgMolecule", + "Species", + "Structure", +] diff --git a/emmet-core/emmet/core/mpcomplete.py b/emmet-core/emmet/core/mpcomplete.py index 3700d03e3d..1cc677f7d2 100644 --- a/emmet-core/emmet/core/mpcomplete.py +++ b/emmet-core/emmet/core/mpcomplete.py @@ -1,8 +1,8 @@ from pydantic import Field from pydantic.main import BaseModel +from emmet.core.atoms.periodic import Material from emmet.core.types.enums import ValueEnum -from emmet.core.types.pymatgen_types.structure_adapter import StructureType class MPCompleteDoc(BaseModel): @@ -10,7 +10,7 @@ class MPCompleteDoc(BaseModel): Defines data for MPComplete structure submissions """ - structure: StructureType | None = Field( + structure: Material | None = Field( None, description="Structure submitted by the user.", ) diff --git a/emmet-core/emmet/core/surface_properties.py b/emmet-core/emmet/core/surface_properties.py index 1cf357d89b..ab796dde5f 100644 --- a/emmet-core/emmet/core/surface_properties.py +++ b/emmet-core/emmet/core/surface_properties.py @@ -5,9 +5,8 @@ from typing import TYPE_CHECKING from pydantic import BaseModel, Field, field_validator -from pymatgen.core import Structure -from emmet.core.types.pymatgen_types.structure_adapter import StructureType +from emmet.core.atoms.periodic import Material from emmet.core.types.typing import MaterialIdentifierType if TYPE_CHECKING: @@ -39,7 +38,7 @@ class SurfaceEntry(BaseModel): description="Whether it is a reconstructed surface.", ) - structure: StructureType | None = Field( + structure: Material | None = Field( None, description="Slab structure.", ) @@ -65,10 +64,12 @@ class SurfaceEntry(BaseModel): ) @field_validator("structure", mode="before") - def get_structure_from_cif(cls, v: Any) -> Structure | None: - """Transform legacy CIF data to pymatgen Structure.""" + def get_structure_from_cif(cls, v: Any) -> Material | None: + """Transform legacy CIF data to emmet.core Material.""" if isinstance(v, str): - return Structure.from_str(v, fmt="cif") + from emmet.core.io.pymatgen import cif_to_material + + return cif_to_material(v) return v @@ -122,7 +123,7 @@ class SurfacePropDoc(BaseModel): description="The Materials Project ID of the material. This comes in the form: mp-******.", ) - structure: StructureType | None = Field( + structure: Material | None = Field( None, description="The conventional crystal structure of the material.", ) From 831f3dde25ca0991453ccca6c3db7e2dc0e04049 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 20 May 2026 14:14:57 -0700 Subject: [PATCH 6/9] sync with id changes --- emmet-core/emmet/core/entry.py | 17 ++++++----------- emmet-core/emmet/core/io/pymatgen.py | 1 + 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/emmet-core/emmet/core/entry.py b/emmet-core/emmet/core/entry.py index 66f5fb2827..5aa362b0d9 100644 --- a/emmet-core/emmet/core/entry.py +++ b/emmet-core/emmet/core/entry.py @@ -3,16 +3,17 @@ from __future__ import annotations from importlib import import_module -from pydantic import BaseModel, Field, field_validator, field_serializer +from pydantic import BaseModel, Field from typing import Literal, TYPE_CHECKING from emmet.core.atoms.base import Compound from emmet.core.atoms.periodic import Material -from emmet.core.mpid_ext import ThermoID from emmet.core.types.enums import IgnoreCaseEnum from emmet.core.types.mson import MSONType +from emmet.core.types.pymatgen_types.computed_entries_adapter import EntryID from emmet.core.types.typing import DateTimeType, IdentifierType, JsonDictType +from emmet.core.utils import type_override from emmet.core.vasp.calc_types.enums import RunType from emmet.core.vasp.calculation import PotcarSpec @@ -67,6 +68,7 @@ def _missing_(cls, value): super(cls)._missing_(value) +@type_override({"hubbards": str}) class EntryParameters(BaseModel): """Schematize entry parameters.""" @@ -76,6 +78,7 @@ class EntryParameters(BaseModel): potcar_spec: list[PotcarSpec] | None = None +@type_override({"oxidation_states": str}) class EntryData(BaseModel): """Schematize entry run data.""" @@ -117,20 +120,12 @@ class Entry(BaseModel): energy: float | None = None correction: float | None = None - entry_id: ThermoID | None = None + entry_id: EntryID | None = None energy_adjustments: list[EnergyAdjustment] = Field([]) parameters: EntryParameters | None = None data: EntryData | None = None - @field_validator("entry_id", mode="before") - def _deser_thermo_id(cls, v) -> ThermoID | None: - return ThermoID._deserialize(v) if v is not None else None - - @field_serializer("entry_id") - def _ser_thermo_id(self, v: ThermoID) -> str: - return str(v) - def to_pmg(self) -> ComputedEntry | ComputedStructureEntry: data = self.model_dump() pmg_cls = "ComputedEntry" diff --git a/emmet-core/emmet/core/io/pymatgen.py b/emmet-core/emmet/core/io/pymatgen.py index 41998cac88..2351367772 100644 --- a/emmet-core/emmet/core/io/pymatgen.py +++ b/emmet-core/emmet/core/io/pymatgen.py @@ -24,6 +24,7 @@ "Species": "core.periodic_table", "DummySpecies": "core.periodic_table", "get_el_sp": "core.periodic_table", + "Site": "core.sites", "PeriodicSite": "core.sites", "IStructure": "core.structure", "Structure": "core.structure", From dcf8dc865371017853aa88906aa536a8bdce5855 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 22 May 2026 11:42:21 -0700 Subject: [PATCH 7/9] mypy --- emmet-core/emmet/core/atoms/base.py | 18 +++++----- emmet-core/emmet/core/atoms/elements.py | 3 +- emmet-core/emmet/core/atoms/periodic.py | 47 ++++++++++++++----------- emmet-core/emmet/core/entry.py | 6 ++-- emmet-core/emmet/core/types/mson.py | 2 +- 5 files changed, 41 insertions(+), 35 deletions(-) diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index 92aa12218d..b425af6c8a 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -123,10 +123,10 @@ def _reduce(cls, config) -> Self: base_config[spec] = 0 base_config[spec] += config["coefficients"][idx] sorted_species = sorted(base_config.keys()) - return { - "species": sorted_species, - "coefficients": [base_config[spec] for spec in sorted_species], - } + return cls( + species=sorted_species, + coefficients=[base_config[spec] for spec in sorted_species], + ) def __str__(self) -> str: return ( @@ -226,7 +226,7 @@ def _sum_scalar_site_properties(self, prop: str, default: float = 0.0) -> float: return sum(getattr(site, prop, None) or default for site in self.sites) @cached_property - def cart_coords(self) -> np.ndarray[float]: + def cart_coords(self) -> np.ndarray: return self._aggregate_site_properties("cart_coords") @cached_property @@ -238,11 +238,11 @@ def spin(self) -> float: return self._sum_scalar_site_properties("spin") @cached_property - def spins(self) -> np.ndarray[float]: + def spins(self) -> np.ndarray: return self._aggregate_site_properties("spin") @cached_property - def degrees_of_freedom(self) -> np.ndarray[bool]: + def degrees_of_freedom(self) -> np.ndarray: return self._aggregate_site_properties( "degrees_of_freedom", default=(True, True, True) ) @@ -256,9 +256,7 @@ def from_pmg(cls, molecule: PmgMolecule) -> Self: def to_pmg(self) -> PmgMolecule: from emmet.core.io.pymatgen import Molecule as PmgMolecule - return PmgMolecule.from_sites( - [site.to_pmg(cell=self.cell) for site in self.sites] - ) + return PmgMolecule.from_sites([site.to_pmg(cell=None) for site in self.sites]) @classmethod def from_sites(cls, sites: list[PmgSite | Site]) -> Self: diff --git a/emmet-core/emmet/core/atoms/elements.py b/emmet-core/emmet/core/atoms/elements.py index 04a87ad597..3bbb04ff00 100644 --- a/emmet-core/emmet/core/atoms/elements.py +++ b/emmet-core/emmet/core/atoms/elements.py @@ -141,6 +141,7 @@ def _missing_(cls, value: Any) -> "Element" | None: return cls(value) elif value in cls.__members__: return cls[value] + return None class ElementData(BaseModel): @@ -154,7 +155,7 @@ class ElementData(BaseModel): class ElementDatabase(dict): - def _load_data(self) -> dict[Element, ElementData]: + def _load_data(self) -> None: """Cache atom data from pymatgen.""" from emmet.core.io.pymatgen import Element as PmgElement diff --git a/emmet-core/emmet/core/atoms/periodic.py b/emmet-core/emmet/core/atoms/periodic.py index a520270dd8..4b9b864d5a 100644 --- a/emmet-core/emmet/core/atoms/periodic.py +++ b/emmet-core/emmet/core/atoms/periodic.py @@ -19,6 +19,8 @@ from typing import Literal from typing_extensions import Self + from spglib import SpgCell + from emmet.core.io.pymatgen import Structure SETTINGS = EmmetSettings() @@ -53,27 +55,27 @@ def volume(self) -> float: return abs(np.linalg.det(self)) @cached_property - def _reciprocal(self) -> Self: + def _reciprocal(self) -> Cell: return Cell( np.array([np.cross(self[(i + 1) % 3], self[(i + 2) % 3]) for i in range(3)]) / self.volume ) @property - def reciprocal(self) -> Self: - return 2 * np.pi * self._reciprocal + def reciprocal(self) -> Cell: + return Cell(2 * np.pi * self._reciprocal) @cached_property - def _vector_norms(self) -> np.ndarray[float]: + def _vector_norms(self) -> np.ndarray: return np.linalg.norm(self, axis=1) @cached_property - def _angles(self) -> np.ndarray: + def _angles(self) -> list[float]: return [ 180 / np.pi * np.arccos( - np.dot(self.matrix[i], self.matrix[(i + 1) % 3]) + np.dot(self[i], self[(i + 1) % 3]) / (self._vector_norms[i] * self._vector_norms[(i + 1) % 3]) ) for i in range(3) @@ -106,7 +108,7 @@ def gamma(self) -> float: @staticmethod def _get_coords( cell: Cell, coords: np.ndarray, to: Literal["cartesian", "direct"] - ) -> np.ndarray[float]: + ) -> np.ndarray: if to == "direct": return np.einsum("ij,ki->kj", cell._reciprocal.T, coords) elif to == "cartesian": @@ -117,7 +119,7 @@ def _get_coords( def get_coords( self, coords: np.ndarray, to: Literal["cartesian", "direct"] - ) -> np.ndarray[float]: + ) -> np.ndarray: return self._get_coords(self, coords, to=to) @@ -135,7 +137,7 @@ def volume(self) -> float: return self.cell.volume @cached_property - def frac_coords(self) -> np.ndarray[float]: + def frac_coords(self) -> np.ndarray: return self.cell.get_coords(self.cart_coords, to="direct") def density_g_cm3(self) -> float: @@ -157,17 +159,17 @@ def to_pmg(self) -> Structure: ) @cached_property - def _to_spglib(self) -> tuple[Cell, np.ndarray[float], np.ndarray[int]]: + def _to_spglib(self) -> SpgCell: """Create an spglib-compatible representation of the atoms.""" - return ( + return ( # type: ignore[return-value] self.cell, self.frac_coords, - [site.Z for site in self.sites], + np.asarray([site.Z for site in self.sites]), ) @classmethod def _from_spglib( - cls, spglib_rep: tuple[Matrix3D, np.ndarray[float], np.ndarray[int]] + cls, spglib_rep: tuple[Matrix3D, np.ndarray, np.ndarray] | SpgCell ) -> Self: cell, frac_coords, atomic_numbers = spglib_rep @@ -175,7 +177,7 @@ def _from_spglib( z: next(ele for ele, data in ELEMENT_DATA.items() if data.Z == z) for z in set(atomic_numbers) } - cart_coords = Cell(cell).get_coords(frac_coords, to="cartesian") + cart_coords = Cell(cell).get_coords(frac_coords, to="cartesian") # type: ignore[arg-type] return cls( lattice=cell, sites=[ @@ -190,8 +192,8 @@ def _from_spglib( def primitive( self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL ) -> Material: - return self._from_spglib( - spglib.find_primitive( + return self._from_spglib( # type: ignore[arg-type] + spglib.find_primitive( # type: ignore[arg-type] self._to_spglib, symprec=symprec, angle_tolerance=angle_tol, @@ -202,18 +204,23 @@ def conventional( self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL ) -> Material: return self._from_spglib( - spglib.standardize_cell( - self._to_spglib, symprec=symprec, angle_tol=angle_tol + spglib.standardize_cell( # type: ignore[arg-type] + self._to_spglib, symprec=symprec, angle_tolerance=angle_tol ) ) def get_space_group_info( self, symprec: float = SETTINGS.SYMPREC, angle_tol: float = SETTINGS.ANGLE_TOL - ) -> tuple[str, int]: + ) -> tuple[str, int] | None: sg_info = spglib.get_spacegroup( self._to_spglib, symprec=symprec, angle_tolerance=angle_tol ) - return tuple(re.match(r"(.*) \((.*)\)", sg_info).groups()) + if ( + isinstance(sg_info, str) + and (matches := re.match(r"(.*) \((.*)\)", sg_info)) is not None + ) and len(groups := matches.groups()) >= 2: + return (groups[0], int(groups[1])) + return None @classmethod def from_cif(cls, cif_str: str) -> Self: diff --git a/emmet-core/emmet/core/entry.py b/emmet-core/emmet/core/entry.py index 5aa362b0d9..2cd2c8a554 100644 --- a/emmet-core/emmet/core/entry.py +++ b/emmet-core/emmet/core/entry.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: from typing_extensions import Self - from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry + from emmet.core.io.pymatgen import ComputedEntry, ComputedStructureEntry CORRECTION_NAME = { "MP GGA(+U)/r2SCAN mixing adjustment", @@ -82,7 +82,7 @@ class EntryParameters(BaseModel): class EntryData(BaseModel): """Schematize entry run data.""" - oxide_type: OxideType = OxideType.NONE + oxide_type: OxideType = OxideType.NONE # type: ignore[assignment] aspherical: bool = False last_updated: DateTimeType task_id: IdentifierType = None @@ -99,7 +99,7 @@ class EnergyAdjustment(BaseModel): adj_per_atom: float | None = None n_atoms: int | None = None uncertainty_per_atom: float | None = None - name: Literal[*CORRECTION_NAME] = None + name: Literal[*CORRECTION_NAME] = None # type: ignore[valid-type] description: str | None = None klass: MSONType | None = Field(None, validation_alias="cls") diff --git a/emmet-core/emmet/core/types/mson.py b/emmet-core/emmet/core/types/mson.py index 16aef45972..17def20751 100644 --- a/emmet-core/emmet/core/types/mson.py +++ b/emmet-core/emmet/core/types/mson.py @@ -23,7 +23,7 @@ def as_dict(self) -> dict[str, str | None]: field.validation_alias: getattr(self, k) for k, field in self.__class__.model_fields.items() } - return {k: v for k, v in dct.items() if v is not None} + return {str(k): v for k, v in dct.items() if v is not None} @classmethod def from_dict(cls, dct) -> Any: From a78f1e37a4ee64486b8911a4dae4ab775739329f Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 22 May 2026 14:09:36 -0700 Subject: [PATCH 8/9] fix composition model validator --- emmet-core/emmet/core/atoms/base.py | 30 ++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index b425af6c8a..bd1295d530 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -110,9 +110,9 @@ class Compound(BaseModel): @model_validator(mode="before") @classmethod - def _reduce(cls, config) -> Self: + def _reduce(cls, config) -> dict[str, list[str | int]]: - if not all(config.get(k) for k in cls.model_fields) or len( + if not all(config.get(k) for k in ("species", "coefficients")) or len( config["species"] ) != len(config["coefficients"]): raise ValueError(f"Invalid input specified to {cls.__name__}.") @@ -123,10 +123,10 @@ def _reduce(cls, config) -> Self: base_config[spec] = 0 base_config[spec] += config["coefficients"][idx] sorted_species = sorted(base_config.keys()) - return cls( - species=sorted_species, - coefficients=[base_config[spec] for spec in sorted_species], - ) + return { + "species": sorted_species, + "coefficients": [base_config[spec] for spec in sorted_species], + } def __str__(self) -> str: return ( @@ -168,12 +168,19 @@ def to_pmg(self) -> Composition: return Composition(self.to_dict()) @property - def reduced(self) -> Compound: + def reduced_composition(self) -> Compound: factor = gcd(*self.coefficients) return Compound( species=self.species, coefficients=[v // factor for v in self.coefficients] ) + @property + def formula(self) -> str: + return " ".join( + f"{ele}{stoich if stoich != 1 else ''}" + for ele, stoich in zip(self.species, self.coefficients) + ) + @cached_property def mass(self) -> float: """Mass in atomic mass units.""" @@ -196,6 +203,15 @@ def get(self, item: str, default: Any = None) -> Any: except KeyError: return default + def keys(self) -> list[str]: + return list(self.species) + + def value(self) -> list[int]: + return list(self.coefficients) + + def items(self) -> list[tuple[str, int]]: + return list(zip(self.species, self.coefficients)) + class Molecule(BaseModel): """Schematize a molecular structure.""" From 5f2daa2fd2fe9708c1f00f03e248880be904d9b7 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 22 May 2026 14:11:06 -0700 Subject: [PATCH 9/9] mypy --- emmet-core/emmet/core/atoms/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/emmet-core/emmet/core/atoms/base.py b/emmet-core/emmet/core/atoms/base.py index bd1295d530..453d8fcccf 100644 --- a/emmet-core/emmet/core/atoms/base.py +++ b/emmet-core/emmet/core/atoms/base.py @@ -110,7 +110,7 @@ class Compound(BaseModel): @model_validator(mode="before") @classmethod - def _reduce(cls, config) -> dict[str, list[str | int]]: + def _reduce(cls, config) -> dict[str, list[str] | list[int]]: if not all(config.get(k) for k in ("species", "coefficients")) or len( config["species"] @@ -233,7 +233,7 @@ def composition(self) -> Compound: @property def reduced_composition(self) -> Compound: - return self.composition.reduced + return self.composition.reduced_composition def _aggregate_site_properties(self, prop: str, default: Any = None) -> np.ndarray: return np.array([getattr(site, prop, None) or default for site in self.sites])