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..453d8fcccf --- /dev/null +++ b/emmet-core/emmet/core/atoms/base.py @@ -0,0 +1,286 @@ +"""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, model_validator + +from emmet.core.atoms.elements import Element, ELEMENT_DATA, parse_species_str +from emmet.core.math import Matrix3D, Vector3D +from emmet.core.io.pymatgen import Site as PmgSite + +if TYPE_CHECKING: + from typing import Any + from typing_extensions import Self + + from emmet.core.io.pymatgen import Composition, PeriodicSite, 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: PmgSite | PeriodicSite) -> Self: + + from emmet.core.io.pymatgen 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=getattr(species, "element", species).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) -> PmgSite | PeriodicSite: + + from emmet.core.io.pymatgen import PmgSite, PeriodicSite + + if cell is not None: + from emmet.core.io.pymatgen 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 PmgSite( + species, + self.cart_coords, + properties=properties, + ) + + return PeriodicSite( + species, + self.cart_coords, + lattice=Lattice(cell), + coords_are_cartesian=True, + properties=properties, + ) + + +class Compound(BaseModel): + + species: list[str] + coefficients: list[int] + + @model_validator(mode="before") + @classmethod + 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"] + ) != 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], + } + + 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)] + ) + + @classmethod + def from_pmg(cls, comp: Composition) -> Self: + return cls.from_dict(comp.as_dict()) + + def to_pmg(self) -> Composition: + from emmet.core.io.pymatgen import Composition + + return Composition(self.to_dict()) + + @property + 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.""" + 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 + + 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.""" + + 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_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]) + + 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: + 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: + return self._aggregate_site_properties("spin") + + @cached_property + def degrees_of_freedom(self) -> np.ndarray: + 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 emmet.core.io.pymatgen import Molecule as PmgMolecule + + return PmgMolecule.from_sites([site.to_pmg(cell=None) for site in self.sites]) + + @classmethod + def from_sites(cls, sites: list[PmgSite | Site]) -> Self: + from emmet.core.io.pymatgen import Site as PmgSite + + return cls( + sites=[ + Site.from_pmg(site) if isinstance(site, PmgSite) else site + for site in 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..3bbb04ff00 --- /dev/null +++ b/emmet-core/emmet/core/atoms/elements.py @@ -0,0 +1,196 @@ +"""Define elements / isotopes.""" + +from __future__ import annotations + +from enum import StrEnum +import re +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] + return None + + +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) -> None: + """Cache atom data from pymatgen.""" + from emmet.core.io.pymatgen 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() + + +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 new file mode 100644 index 0000000000..4b9b864d5a --- /dev/null +++ b/emmet-core/emmet/core/atoms/periodic.py @@ -0,0 +1,229 @@ +"""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 spglib import SpgCell + + from emmet.core.io.pymatgen 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) -> 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) -> Cell: + return Cell(2 * np.pi * self._reciprocal) + + @cached_property + def _vector_norms(self) -> np.ndarray: + return np.linalg.norm(self, axis=1) + + @cached_property + def _angles(self) -> list[float]: + return [ + 180 + / np.pi + * np.arccos( + np.dot(self[i], self[(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: + 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: + 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: + 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.composition.mass * atomic_mass * 1e27 / self.cell.volume + + @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 emmet.core.io.pymatgen import Structure + + return Structure.from_sites( + [site.to_pmg(cell=self.cell) for site in self.sites] + ) + + @cached_property + def _to_spglib(self) -> SpgCell: + """Create an spglib-compatible representation of the atoms.""" + return ( # type: ignore[return-value] + self.cell, + self.frac_coords, + np.asarray([site.Z for site in self.sites]), + ) + + @classmethod + def _from_spglib( + cls, spglib_rep: tuple[Matrix3D, np.ndarray, np.ndarray] | SpgCell + ) -> 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") # type: ignore[arg-type] + 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( # type: ignore[arg-type] + spglib.find_primitive( # type: ignore[arg-type] + 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( # 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] | None: + sg_info = spglib.get_spacegroup( + self._to_spglib, symprec=symprec, angle_tolerance=angle_tol + ) + 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: + 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 c2ed71a7c8..bfc8246146 100644 --- a/emmet-core/emmet/core/chemenv.py +++ b/emmet-core/emmet/core/chemenv.py @@ -5,7 +5,6 @@ from pydantic import Field from emmet.core.io.pymatgen import ( SpacegroupAnalyzer, - Molecule, Structure, SimplestChemenvStrategy, AllCoordinationGeometries, @@ -13,13 +12,13 @@ LightStructureEnvironments, ) +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, ) @@ -362,7 +361,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/entry.py b/emmet-core/emmet/core/entry.py new file mode 100644 index 0000000000..2cd2c8a554 --- /dev/null +++ b/emmet-core/emmet/core/entry.py @@ -0,0 +1,151 @@ +"""Define schemas for pymatgen entry-like objects.""" + +from __future__ import annotations + +from importlib import import_module +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.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 + +if TYPE_CHECKING: + from typing_extensions import Self + from emmet.core.io.pymatgen import ComputedEntry, ComputedStructureEntry + +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) + + +@type_override({"hubbards": str}) +class EntryParameters(BaseModel): + """Schematize entry parameters.""" + + run_type: RunType | None = None + is_hubbard: bool = False + hubbards: JsonDictType = None + potcar_spec: list[PotcarSpec] | None = None + + +@type_override({"oxidation_states": str}) +class EntryData(BaseModel): + """Schematize entry run data.""" + + oxide_type: OxideType = OxideType.NONE # type: ignore[assignment] + aspherical: bool = False + last_updated: DateTimeType + 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 + + +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 # type: ignore[valid-type] + 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: Compound + + energy: float | None = None + correction: float | None = None + entry_id: EntryID | None = None + + energy_adjustments: list[EnergyAdjustment] = Field([]) + parameters: EntryParameters | None = None + data: EntryData | None = None + + 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: Material diff --git a/emmet-core/emmet/core/grain_boundary.py b/emmet-core/emmet/core/grain_boundary.py index bb434e192d..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 emmet.core.io.pymatgen 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/pymatgen.py b/emmet-core/emmet/core/io/pymatgen.py index c898406672..08ed2e9e18 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", 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 7b5706a9be..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 emmet.core.io.pymatgen 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.", ) diff --git a/emmet-core/emmet/core/types/mson.py b/emmet-core/emmet/core/types/mson.py new file mode 100644 index 0000000000..17def20751 --- /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 {str(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)