xrd: Re-introduced label tensorization

Author: Somerandomguy10111

tests/t_labels.py

@@ -1,89 +1,66 @@
 from holytools.devtools import Unittest
-from xrdpattern.xrd import PowderExperiment
+from pymatgen.core import Lattice
+
+from xrdpattern.crystal import CrystalStructure, CrystalBasis
+from xrdpattern.xrd import PowderExperiment, XrayInfo
+from xrdpattern.xrd.experiment import ExperimentTensor
 
 
 class TestPowderExperiment(Unittest):
-    def test_empty(self):
+    def test_is_empty(self):
+        
+        # empty is empty
         empty_experiment = PowderExperiment.make_empty()
         self.assertTrue(not empty_experiment.is_nonempty())
+        
+        #nonempty is nonempty
+
 
+    def test_has_label(self):
+        pass
+    
 
-class TestTensorRegions(Unittest):
+class TestTensorization(Unittest):
     def setUp(self):
-        raise ValueError(f'Tensor regions currently broken! Restore tensorizaation + tests')
-        # self.label : PowderExperiment = self.make_example_label()
-        # self.label_tensor : LabelTensor = self.label.to_tensor()
-        # self.crystal_structure : CrystalPhase = self.label.powder.phases
+        self.label : PowderExperiment = self.make_example_label()
+        self.crystal_structure: CrystalStructure = self.label.phases[0]
+        self.experiment_tensor : ExperimentTensor = self.label.to_tensordict()
+
+    def test_lattice_params(self):
+        expected = (*self.crystal_structure.lengths, *self.crystal_structure.angles)
+        actual = self.experiment_tensor.get_lattice_params().tolist()
+        print(f'Tensor lattice params are {actual}')
+        for x,y in zip(expected, actual):
+            self.assertEqual(x, y)
+
+    def test_spacegroups(self):
+        expected = [1.0 if j ==  self.crystal_structure.spacegroup else 0.0 for j in range(1,231)]
+        actual = self.experiment_tensor.get_spg_probabilities().tolist()
+        print(f'Tensor spacegroups probabilities are {actual}; \nSpacegroups tensor length = {len(actual)}')
+        self.assertEqual(actual, expected, f'Expected: {expected}, Actual: {actual}')
+        self.assertEqual(len(actual), 230)
+
+    def test_wavelength(self):
+        expected = self.label.xray_info.primary_wavelength
+        actual = self.experiment_tensor.get_primary_wavelength()
+        print(f'Tensor primary wavelength is {actual}')
+        self.assertAlmostEqual(actual, expected)
 
+        expected_secondary = self.label.xray_info.secondary_wavelength
+        actual_secondary = self.experiment_tensor.get_secondary_wavelength()
+        print(f'Tensor secondary wavelength is {actual_secondary}')
+        self.assertAlmostEqual(actual_secondary, expected_secondary)
 
-    #
-    # def test_lattice_params(self):
-    #     expected = (*self.crystal_structure.lengths, *self.crystal_structure.angles)
-    #     actual = self.label_tensor.get_lattice_params().tolist()
-    #     print(f'Tensor lattice params are {actual}')
-    #     for x,y in zip(expected, actual):
-    #         self.assertEqual(x, y)
-    #
-    #
-    # def test_atomic_sites(self):
-    #     base = self.crystal_structure.base
-    #     for i, region in enumerate(self.label.atomic_site_regions):
-    #         expected_site = base[i] if i < len(base) else AtomicSite.make_void()
-    #         expected = expected_site.as_list()
-    #         actual = self.label_tensor.get_atomic_site(i).tolist()
-    #         if i == 0:
-    #             print(f'Tensor atomic site is {actual}\n Atomic site length = {len(actual)}')
-    #         # 3 coordinates, 8 scattering params, 1 occupancy
-    #         self.assertEqual(len(actual), 3+8+1)
-    #         for x,y, in zip(expected, actual):
-    #             if x is None:
-    #                 self.assertTrue(is_nan(y))
-    #             else:
-    #                 self.assertEqual(x,y)
-    #
-    # def test_spacegroups(self):
-    #     expected = [1.0 if j ==  self.crystal_structure.spacegroup else 0.0 for j in range(1,231)]
-    #     actual = self.label_tensor.get_spg_probabilities().tolist()
-    #     print(f'Tensor spacegroups are {actual}; \nSpacegroups tensor length = {len(actual)}')
-    #     self.assertEqual(actual, expected, f'Expected: {expected}, Actual: {actual}')
-    #     self.assertEqual(len(actual), 230)
-    #
-    #
-    # def test_artifacts(self):
-    #     expected = [round(num, 2) for num in self.label.artifacts.as_list()]
-    #     actual = [round(num, 2) for num in self.label_tensor.get_artifacts().tolist()]
-    #     print(f'Tensor artifacts are {actual}')
-    #     self.assertEqual(actual, expected)
-    #
-    # def test_total_length(self):
-    #     expected = len(self.label.list_repr)
-    #     # noinspection PyTypeChecker
-    #     actual = len(self.label_tensor)
-    #     print(f'Tensor length is {actual}')
-    #     self.assertEqual(actual, expected)
-    #
-    #
-    # def test_returns_powder_tensors(self):
-    #     tensors = [self.label_tensor.get_atomic_site(0), self.label_tensor.get_lattice_params(), self.label_tensor.get_lattice_params()]
-    #     for t in tensors:
-    #         print(f'Type of {t} is {type(t)}')
-    #         self.assertEqual(t.__class__, LabelTensor)
-    #
-    # @staticmethod
-    # def make_example_label() -> PowderExperiment:
-    #     primitives = Lengths(a=3, b=3, c=3)
-    #     angles = Angles(alpha=90, beta=90, gamma=90)
-    #     base = CrystalBase([AtomicSite.make_void()])
-    #     crystal_structure = CrystalPhase(lengths=primitives, angles=angles, base=base)
-    #     crystal_structure.spacegroup = 120
-    #
-    #     powder = PowderSample(phases=[crystal_structure], crystallite_size=10.0)
-    #     artifacts = XRayInfo(primary_wavelength=1.54, secondary_wavelength=1.54)
-    #     return PowderExperiment(powder, artifacts, is_simulated=True)
-    #
+    @staticmethod
+    def make_example_label() -> PowderExperiment:
+        lattice = Lattice.from_parameters(3,3,3,90,90,90)
+        basis = CrystalBasis.empty()
+        crystal_structure = CrystalStructure(lattice=lattice, basis=basis)
+        crystal_structure.spacegroup = 120
 
-def is_nan(value):
-    return value != value
+        xray_info = XrayInfo.copper_xray()
+        return PowderExperiment(phases=[crystal_structure], xray_info=xray_info, crystallite_size_nm=10, temp_K=300)
 
 if __name__ == "__main__":
-    TestPowderExperiment.execute_all()
+    # TestPowderExperiment.execute_all()
+    TestTensorization.execute_all()

xrdpattern/parsing/master.py

@@ -122,7 +122,7 @@ def get_float(key: str) -> Optional[float]:
 
         experiment = PowderExperiment.make_empty()
         experiment.xray_info = XrayInfo(primary_wavelength=get_float('ALPHA1'), secondary_wavelength=get_float('ALPHA2'))
-        experiment.temp_in_celcius = get_float('TEMP_CELCIUS')
+        experiment.temp_K = get_float('TEMP_CELCIUS')
 
         return experiment
 

xrdpattern/pattern/pattern.py

@@ -148,7 +148,7 @@ def get_info_as_str(self) -> str:
               f'- Crystal: {crystal_data} \n'
               f'- Experiment Parameters:\n'
               f'    - Primary wavelength: {self.powder_experiment.xray_info.primary_wavelength}\n'
-              f'    - Temperature : {self.powder_experiment.temp_in_celcius}\n'
+              f'    - Temperature : {self.powder_experiment.temp_K}\n'
               f'- Origin Metadata:\n'
                 f'    - Contributor: {self.metadata.contributor_name}\n'
                 f'    - Institution: {self.metadata.institution}\n'

xrdpattern/xrd/experiment.py

@@ -5,6 +5,9 @@
 from importlib.metadata import version
 from typing import Optional
 
+import torch
+from tensordict import TensorDict
+
 from xrdpattern.crystal import CrystalStructure
 from xrdpattern.serialization import JsonDataclass
 from xrdpattern.xrd.xray import XrayInfo
@@ -19,8 +22,8 @@ class PowderExperiment(JsonDataclass):
     phases: list[CrystalStructure]
     xray_info : XrayInfo
     is_simulated : bool = False
-    crystallite_size: Optional[float] = None
-    temp_in_celcius: Optional[float] = None
+    crystallite_size_nm: Optional[float] = None
+    temp_K: Optional[float] = None
 
     def __post_init__(self):
         if len(self.phases) == 1:
@@ -33,12 +36,12 @@ def make_empty(cls, is_simulated: bool = False) -> PowderExperiment:
 
     @classmethod
     def from_multi_phase(cls, phases : list[CrystalStructure]):
-        return cls(phases=phases, crystallite_size=None, xray_info=XrayInfo.mk_empty(), is_simulated=False)
+        return cls(phases=phases, crystallite_size_nm=None, xray_info=XrayInfo.mk_empty(), is_simulated=False)
 
     @classmethod
     def from_single_phase(cls, phase : CrystalStructure, crystallite_size : Optional[float] = None, is_simulated : bool = False):
         artifacts = XrayInfo.mk_empty()
-        return cls(phases=[phase], crystallite_size=crystallite_size, xray_info=artifacts, is_simulated=is_simulated)
+        return cls(phases=[phase], crystallite_size_nm=crystallite_size, xray_info=artifacts, is_simulated=is_simulated)
 
     @classmethod
     def from_cif(cls, cif_content : str) -> PowderExperiment:
@@ -83,8 +86,28 @@ def is_nonempty(self) -> bool:
     def __eq__(self, other : PowderExperiment):
         return self.to_str() == other.to_str()
 
-    # def to_tensor(self, dtype : torch.dtype = torch.get_default_dtype(), device : torch.device = torch.get_default_device()) -> LabelTensor:
-    #     return LabelTensor(tensor)
+    def to_tensordict(self, dtype : torch.dtype = torch.get_default_dtype(),
+                            device : torch.device = torch.get_default_device()) -> ExperimentTensor:
+
+        def to_tensor(data):
+            return torch.tensor(data=data, dtype=dtype, device=device)
+
+        if len(self.phases) == 0:
+            raise ValueError('No phases in the experiment. Cannot convert to TensorDict.')
+
+        spg_list = [self.phases[0].spacegroup == j for j in range(1,NUM_SPACEGROUPS+1)]
+        feature_dict = {
+            'lengths': to_tensor(self.phases[0].lengths),
+            'angles' : to_tensor(self.phases[0].angles),
+            'spg_probabilities' : to_tensor(spg_list),
+            'crystallite_size' : to_tensor(self.crystallite_size_nm) if self.crystallite_size_nm else None,
+            'temperature' : to_tensor(self.temp_K) if self.temp_K else None,
+            'primary_wavelength' : to_tensor(self.xray_info.primary_wavelength) if self.xray_info.primary_wavelength else None,
+            'secondary_wavelength' : to_tensor(self.xray_info.secondary_wavelength) if self.xray_info.secondary_wavelength else None,
+        }
+        td = ExperimentTensor(feature_dict)
+        return td
+
 
 @dataclass
 class Metadata(JsonDataclass):
@@ -104,5 +127,26 @@ def remove_filename(self):
         self.filename = None
 
 
+class ExperimentTensor(TensorDict):
+    def get_lattice_params(self):
+        lengths, angles = self.get('lengths'), self.get('angles')
+        return torch.cat([lengths, angles], dim=0)
+
+    def get_spg_probabilities(self):
+        return self.get('spg_probabilities')
+
+    def get_crystallite_size(self):
+        return self.get('crystallite_size')
+
+    def get_ambient_temperature(self):
+        return self.get('temperature')
+
+    def get_primary_wavelength(self):
+        return self.get('primary_wavelength')
+
+    def get_secondary_wavelength(self):
+        return self.get('secondary_wavelength')
+
+
 def get_library_version(library_name : str):
     return version(library_name)