-
Notifications
You must be signed in to change notification settings - Fork 41
/
atoms.py
118 lines (97 loc) · 3.43 KB
/
atoms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""Schemas for storing metadata about Atoms objects."""
from __future__ import annotations
from typing import TYPE_CHECKING
from emmet.core.structure import MoleculeMetadata, StructureMetadata
from pymatgen.io.ase import AseAtomsAdaptor
from quacc.atoms.core import (
copy_atoms,
get_charge_attribute,
get_spin_multiplicity_attribute,
)
if TYPE_CHECKING:
from typing import Any
from ase.atoms import Atoms
from quacc.schemas._aliases.atoms import AtomsSchema
def atoms_to_metadata(
atoms: Atoms,
charge_and_multiplicity: tuple[int, int] | None = None,
get_metadata: bool = True,
store_pmg: bool = True,
additional_fields: dict[str, Any] | None = None,
) -> AtomsSchema:
"""
Convert an ASE Atoms object to a dict suitable for storage in MongoDB.
Parameters
----------
atoms
ASE Atoms object to store in {"atoms": atoms}
charge_and_multiplicity
Charge and spin multiplicity of the Atoms object, only used for Molecule
metadata.
get_metadata
Whether to store atoms metadata in the returned dict.
store_pmg
Whether to store the Pymatgen Structure/Molecule object in {"structure":
Structure} or {"molecule": Molecule}, respectively.
additional_fields
Additional fields to add to the document.
Returns
-------
AtomsSchema
Dict of metadata about the Atoms object.
"""
additional_fields = additional_fields or {}
atoms = copy_atoms(atoms)
results = {}
atoms.calc = None
# Set any charge or multiplicity keys
if not atoms.pbc.any():
_set_charge_and_spin(atoms, charge_and_multiplicity=charge_and_multiplicity)
# Strip the dummy atoms, if present
if "X" in atoms.get_chemical_symbols():
del atoms[[atom.index for atom in atoms if atom.symbol == "X"]]
# Get Atoms metadata, if requested. emmet already has built-in tools for
# generating pymatgen Structure/Molecule metadata, so we'll just use that.
if get_metadata:
if atoms.pbc.any():
struct = AseAtomsAdaptor().get_structure(atoms)
metadata = StructureMetadata().from_structure(struct).model_dump()
if store_pmg:
results["structure"] = struct
else:
mol = AseAtomsAdaptor().get_molecule(atoms, charge_spin_check=False)
metadata = MoleculeMetadata().from_molecule(mol).model_dump()
if store_pmg:
results["molecule"] = mol
else:
metadata = {}
# Store Atoms object
results["atoms"] = atoms
return metadata | results | additional_fields
def _set_charge_and_spin(
atoms: Atoms, charge_and_multiplicity: tuple[int, int] | None = None
) -> None:
"""
Set the charge and spin multiplicity of an Atoms object.
Parameters
----------
atoms
Atoms object
charge_and_multiplicity
Charge and spin multiplicity of the Atoms object, only used for Molecule
metadata.
Returns
-------
None
Modifies the Atoms object in place.
"""
if charge_and_multiplicity:
charge = charge_and_multiplicity[0]
spin_multiplicity = charge_and_multiplicity[1]
else:
charge = get_charge_attribute(atoms)
spin_multiplicity = get_spin_multiplicity_attribute(atoms)
if charge is not None:
atoms.charge = charge
if spin_multiplicity is not None:
atoms.spin_multiplicity = spin_multiplicity