crystal: Re-implemented AtomicSite to use symbol as init value

CrystalStructure/atomic_constants/atomic_constants.py

@@ -16,11 +16,6 @@ def load_constants_json(fname: str) -> dict:
 
 # ---------------------------------------------------------
 
-class UnknownSite:
-    symbol = 'NaN'
-
-class Void:
-    symbol = '⊥'
 
 
 class AtomicConstants:
@@ -40,13 +35,8 @@ def get_covalent(cls, element_symbol: str) -> float:
         return cls._covalent[element_symbol]
 
     @classmethod
-    def get_scattering_params(cls, species: Union[Element,Species]) -> tuple:
-        if isinstance(species, Species):
-            symbol = str(species.element.symbol)
-        else:
-            symbol = species.symbol
-
-        return cls._scattering_params[symbol]
+    def get_scattering_params(cls, species_symbol: str) -> tuple:
+        return cls._scattering_params[species_symbol]
 
     @classmethod
     def print_all(cls):

CrystalStructure/crystal/atomic_site.py

@@ -3,54 +3,53 @@
 import json
 from dataclasses import dataclass
 from typing import Union, Optional
-from pymatgen.core import Species, Element
+from pymatgen.core import Species, Element, DummySpecies
 
 from holytools.abstract import Serializable
-from CrystalStructure.atomic_constants.atomic_constants import Void, UnknownSite, AtomicConstants
+from pymatgen.util.typing import SpeciesLike
+
+from CrystalStructure.atomic_constants.atomic_constants import AtomicConstants
 
 ScatteringParams = tuple[float, float, float, float, float, float, float, float]
 #---------------------------------------------------------
 
+# Conditions:
+# - AtomicSite should be initializable with only a pymatgen species
+# - The make_void and make_placeholder classmethods should continue to work as before
+
+# ->
+
 @dataclass
 class AtomicSite(Serializable):
     """x,y,z are the coordinates of the site given in the basis of the lattice"""
     x: Optional[float]
     y: Optional[float]
     z: Optional[float]
     occupancy : Optional[float]
-    atom_type : Union[Element,Species, Void, UnknownSite]
+    species_symbol : str
     wyckoff_letter : Optional[str] = None
 
     def __post_init__(self):
-        if isinstance(self.atom_type,Element):
-            self.atom_type = Species(self.atom_type.symbol)
-        self.atom_type : Union[Species, Void, UnknownSite]
+        self.atom_type = AtomType(symbol=self.species_symbol)
 
     @classmethod
     def make_void(cls) -> AtomicSite:
-        return cls(x=None, y=None, z=None, occupancy=0.0, atom_type=Void())
+        return cls(x=None, y=None, z=None, occupancy=0.0, species_symbol=AtomType.void_symbol)
 
     @classmethod
     def make_placeholder(cls):
-        return cls(x=None, y=None, z=None, occupancy=None, atom_type=UnknownSite())
+        return cls(x=None, y=None, z=None, occupancy=None, species_symbol=AtomType.placeholder_symbol)
 
     def is_nonstandard(self) -> bool:
-        if isinstance(self.atom_type, Void) or isinstance(self.atom_type, UnknownSite):
+        if self.species_symbol == AtomType.void_symbol or self.species_symbol == AtomType.placeholder_symbol:
             return True
         return False
 
     # ---------------------------------------------------------
     # properties
 
     def get_symbol(self) -> str:
-        if isinstance(self.atom_type, Species):
-            return self.atom_type.element.symbol
-        elif isinstance(self.atom_type, Void):
-            return Void.symbol
-        elif isinstance(self.atom_type, UnknownSite):
-            return UnknownSite.symbol
-        else:
-            raise ValueError(f'Unknown species type: {self.atom_type}')
+        return self.atom_type.get_symbol()
 
     def as_list(self) -> list[float]:
         site_arr = [*self.get_scattering_params(), self.x, self.y, self.z, self.occupancy]
@@ -60,41 +59,55 @@ def as_list(self) -> list[float]:
     # These are *different* paramters from what you may commonly see e.g. here (https://lampz.tugraz.at/~hadley/ss1/crystaldiffraction/atomicformfactors/formfactors.php)
     # since pymatgen uses a different formula to compute the form factor
     def get_scattering_params(self) -> ScatteringParams:
-        if isinstance(self.atom_type, Species):
-            values = AtomicConstants.get_scattering_params(species=self.atom_type)
-        elif isinstance(self.atom_type, Void):
-            values = (0.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)
-        elif isinstance(self.atom_type, UnknownSite):
-            fnan = float('nan')
-            values = (fnan,fnan), (fnan,fnan), (fnan,fnan), (fnan,fnan)
-        else:
-            raise ValueError(f'Unknown species type: {self.atom_type}')
-
-        (a1, b1), (a2, b2), (a3, b3), (a4, b4) = values
-        return a1, b1, a2, b2, a3, b3, a4, b4
+        return self.atom_type.scattering_params
 
     # ---------------------------------------------------------
     # save/load
 
     def to_str(self) -> str:
         the_dict = {'x': self.x, 'y': self.y, 'z': self.z, 'occupancy': self.occupancy,
-                    'species': str(self.atom_type),
+                    'symbol': self.species_symbol,
                     'wyckoff_letter': self.wyckoff_letter}
 
         return json.dumps(the_dict)
 
     @classmethod
     def from_str(cls, s: str):
         the_dict = json.loads(s)
-        species_symbol = the_dict['species']
-        if species_symbol == Void.symbol:
-            species = Void()
-        elif species_symbol == UnknownSite.symbol:
-            species = UnknownSite()
-        else:
-            species = Species.from_str(species_symbol)
 
         return cls(x=the_dict['x'], y=the_dict['y'], z=the_dict['z'],
                    occupancy=the_dict['occupancy'],
-                   atom_type=species,
+                   species_symbol=the_dict['symbol'],
                    wyckoff_letter=the_dict['wyckoff_letter'])
+
+
+class AtomType:
+    void_symbol = '⊥'
+    placeholder_symbol = 'NaN'
+
+    def __init__(self, symbol : str):
+        if symbol == self.void_symbol:
+            self.species_like : SpeciesLike = DummySpecies(symbol=self.void_symbol)
+        if symbol == self.placeholder_symbol:
+            self.species_like : SpeciesLike = DummySpecies(symbol=self.placeholder_symbol)
+        else:
+            self.species_like : SpeciesLike = Species.from_str(species_string=symbol)
+
+    def get_symbol(self):
+        return str(self.species_like)
+
+    @property
+    def scattering_params(self) -> ScatteringParams:
+        if self.get_symbol() == self.void_symbol:
+            values = (0.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)
+        elif self.get_symbol() == self.placeholder_symbol:
+            fnan = float('nan')
+            values = (fnan, fnan), (fnan, fnan), (fnan, fnan), (fnan, fnan)
+        else:
+            # TODO: This casting only currently exists beacuse the scattering param table only has values for (unoxidized) elements, not ions
+            # TODO: Normally would simply be species_symbol=str(self.species_like)
+            species_symbol = self.species_like.element.symbol
+            values = AtomicConstants.get_scattering_params(species_symbol=species_symbol)
+
+        (a1, b1), (a2, b2), (a3, b3), (a4, b4) = values
+        return a1, b1, a2, b2, a3, b3, a4, b4