44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
from crystalpol.shared.system.molecule import Molecule
|
|
|
|
from typing import List
|
|
|
|
|
|
class Crystal:
|
|
"""
|
|
This class represents a crystal, it will be organized in a list structure.
|
|
Each element is a unitary cell in the crystal. And this unitary cell will have
|
|
Molecules in it.
|
|
"""
|
|
|
|
def __init__(self, structure: List[List[str]]):
|
|
self.structure = structure
|
|
|
|
self.cells = []
|
|
|
|
def __iter__(self):
|
|
for cell in self.cells:
|
|
yield cell
|
|
|
|
def __len__(self):
|
|
return len(self.cells)
|
|
|
|
def __getitem__(self, index):
|
|
return self.cells[index]
|
|
|
|
def add_cell(self, cell: List[Molecule]) -> None:
|
|
valid = self._is_valid_cell(cell)
|
|
if not valid:
|
|
raise ValueError(
|
|
"This cell does not obey the declared format for this Crystal."
|
|
)
|
|
else:
|
|
self.cells.append(cell)
|
|
|
|
def _is_valid_cell(self, cell: List[Molecule]) -> bool:
|
|
if len(cell) == len(self.structure):
|
|
for i, molecule in enumerate(cell):
|
|
if len(molecule.atoms) == len(self.structure[i]) \
|
|
and all(atom.symbol == self.structure[i][j] for j, atom in enumerate(molecule.atoms)):
|
|
return True
|
|
return False
|