diff --git a/emmet-core/emmet/core/xrd.py b/emmet-core/emmet/core/xrd.py new file mode 100644 index 0000000000..bff2e76e6b --- /dev/null +++ b/emmet-core/emmet/core/xrd.py @@ -0,0 +1,112 @@ +from enum import Enum +from typing import Dict, List, Optional, Tuple + +import numpy as np +from pydantic import BaseModel, Field, root_validator +from pymatgen.analysis.diffraction.xrd import WAVELENGTHS, XRDCalculator +from pymatgen.core.periodic_table import Element + +from emmet.core.spectrum import SpectrumDoc +from emmet.core.utils import ValueEnum, jsanitize +from emmet.stubs import DiffractionPattern, Structure + + +class Edge(ValueEnum): + K_Alpha = "Ka" + K_Alpha1 = "Ka1" + K_Alpha2 = "Ka2" + K_Beta = "Kb" + K_Beta1 = "Kb1" + K_Beta2 = "Kb2" + + +class XRDDoc(SpectrumDoc): + """ + Document describing a XRD Diffraction Pattern + """ + + spectrum: DiffractionPattern + min_two_theta: float + max_two_theta: float + wavelength: float = Field(..., description="Wavelength for the diffraction source") + target: Element = Field( + None, description="Target element for the diffraction source" + ) + edge: Edge = Field(None, description="Atomic edge for the diffraction source") + + @root_validator(pre=True) + def get_target_and_edge(cls, values: Dict): + print("Validations") + # Only do this if neither target not edge is defined + if "target" not in values and "edge" not in values: + print("Are we even getting here?") + try: + pymatgen_wavelength = next( + k + for k, v in WAVELENGTHS.items() + if np.allclose(values["wavelength"], v) + ) + values["target"] = pymatgen_wavelength[:2] + values["edge"] = pymatgen_wavelength[2:] + + except Exception: + return values + return values + + @classmethod + def from_structure( + cls, + material_id: str, + spectrum_id: str, + structure: Structure, + wavelength: float, + min_two_theta=0, + max_two_theta=180, + symprec=0.1, + **kwargs, + ) -> "XRDDoc": + calc = XRDCalculator(wavelength=wavelength, symprec=symprec) + pattern = calc.get_pattern( + structure, two_theta_range=(min_two_theta, max_two_theta) + ) + + return super().from_structure( + material_id=material_id, + spectrum_id=spectrum_id, + structure=structure, + spectrum=pattern, + wavelength=wavelength, + min_two_theta=min_two_theta, + max_two_theta=max_two_theta, + **kwargs, + ) + + @classmethod + def from_target( + cls, + material_id: str, + structure: Structure, + target: Element, + edge: Edge, + min_two_theta=0, + max_two_theta=180, + symprec=0.1, + **kwargs, + ) -> "XRDDoc": + if f"{target}{edge}" not in WAVELENGTHS: + raise ValueError(f"{target}{edge} not in pymatgen wavelenghts dictionarty") + + wavelength = WAVELENGTHS[f"{target}{edge}"] + spectrum_id = f"{material_id}-{target}{edge}" + + return cls.from_structure( + material_id=material_id, + spectrum_id=spectrum_id, + structure=structure, + wavelength=wavelength, + target=target, + edge=edge, + min_two_theta=min_two_theta, + max_two_theta=max_two_theta, + **kwargs, + )