updated tests added some tests

diffsims/generators/diffraction_generator.py

@@ -357,7 +357,8 @@ def calculate_ed_data(
         )
 
     def calculate_profile_data(
-        self, structure, reciprocal_radius=1.0, minimum_intensity=1e-3
+        self, structure, reciprocal_radius=1.0, minimum_intensity=1e-3,
+        debye_waller_factors={}
     ):
         """Calculates a one dimensional diffraction profile for a
         structure.
@@ -372,6 +373,8 @@ def calculate_profile_data(
         minimum_intensity : float
             The minimum intensity required for a diffraction peak to be
             considered real. Deals with numerical precision issues.
+        debye_waller_factors : dict of str:value pairs
+            Maps element names to their temperature-dependent Debye-Waller factors.
 
         Returns
         -------
@@ -399,7 +402,7 @@ def calculate_profile_data(
             np.asarray(g_hkls),
             prefactor=multiplicities,
             scattering_params=self.scattering_params,
-            debye_waller_factors=self.debye_waller_factors,
+            debye_waller_factors=debye_waller_factors,
         )
 
         if is_lattice_hexagonal(latt):

diffsims/tests/test_generators/test_diffraction_generator.py

@@ -23,10 +23,14 @@
 from diffsims.generators.diffraction_generator import (
     DiffractionGenerator,
     AtomicDiffractionGenerator,
+    _z_sphere_precession,
+    _shape_factor_precession,
+    _average_excitation_error_precession,
 )
 import diffpy.structure
-from diffsims.utils.shape_factor_models import linear, binary, sinc
-
+from diffsims.utils.shape_factor_models import (linear, binary, sinc, sin2c,
+                                                atanc, lorentzian, 
+                                                lorentzian_precession)
 
 @pytest.fixture(params=[(300)])
 def diffraction_calculator(request):
@@ -82,10 +86,45 @@ def probe(x, out=None, scale=None):
         return v + 0 * x[2].reshape(1, 1, -1)
 
 
+@pytest.mark.parametrize("parameters, expected",
+                         [([0, 1, 0.001, 0.5], -0.00822681491001731),
+                          ([0, np.array([1, 2, 20]), 0.001, 0.5],
+                            np.array([-0.00822681, -0.01545354, 0.02547058])),
+                          ([180, 1, 0.001, 0.5], 0.00922693)])
+def test_z_sphere_precession(parameters, expected):
+    result = _z_sphere_precession(*parameters)
+    assert np.allclose(result, expected)
+
+
+@pytest.mark.parametrize("model", [linear, atanc, sin2c, lorentzian])
+def test_shape_factor_precession(model):
+    z = np.array([-0.1, 0.1])
+    r = np.array([1, 5])
+    _ = _shape_factor_precession(z, r, 0.001, 0.5, model, 0.1)
+
+
+def test_average_excitation_error_precession():
+    z = np.array([-0.1, 0.1])
+    r = np.array([1, 5])
+    _ = _average_excitation_error_precession(z, r, 0.001, 0.5)
+
+
+@pytest.mark.parametrize("model, expected",
+                         [("linear", linear),
+                          ("lorentzian", lorentzian),
+                          (binary, binary)],)
+def test_diffraction_generator_init(model, expected):
+    generator = DiffractionGenerator(300, shape_factor_model=model)
+    assert generator.shape_factor_model == expected
+
+
 class TestDiffractionCalculator:
     def test_init(self, diffraction_calculator: DiffractionGenerator):
-        assert diffraction_calculator.debye_waller_factors == {}
-        _ = DiffractionGenerator(300, 2)
+        assert diffraction_calculator.scattering_params == "lobato"
+        assert diffraction_calculator.precession_angle == 0
+        assert diffraction_calculator.shape_factor_model == linear
+        assert diffraction_calculator.approximate_precession == True
+        assert diffraction_calculator.minimum_intensity == 1e-20
 
     def test_matching_results(self, diffraction_calculator, local_structure):
         diffraction = diffraction_calculator.calculate_ed_data(
@@ -124,21 +163,20 @@ def test_appropriate_intensities(self, diffraction_calculator, local_structure):
         )
         assert np.all(smaller)
 
-    @pytest.mark.parametrize("model", [None, linear, binary, sinc])
-    def test_shape_factor_strings(self, diffraction_calculator, local_structure, model):
+    def test_shape_factor_strings(self, diffraction_calculator, local_structure):
         _ = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=model
+            local_structure, 2
         )
 
     def test_shape_factor_custom(self, diffraction_calculator, local_structure):
         def local_excite(excitation_error, maximum_excitation_error, t):
             return (np.sin(t) * excitation_error) / maximum_excitation_error
 
         t1 = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=local_excite, t=0.2
+            local_structure, 2, max_excitation_error=0.02
         )
         t2 = diffraction_calculator.calculate_ed_data(
-            local_structure, 2, shape_factor_model=local_excite, t=0.4
+            local_structure, 2, max_excitation_error=0.4
         )
 
         # softly makes sure the two sims are different
@@ -204,6 +242,11 @@ def test_invalid_scattering_params():
     generator = DiffractionGenerator(300, scattering_params=scattering_param)
 
 
+@pytest.mark.xfail(faises=NotImplementedError)
+def test_invalid_shape_model():
+    generator = DiffractionGenerator(300, shape_factor_model="dracula")
+
+
 @pytest.mark.parametrize("shape", [(10, 20), (20, 10)])
 def test_param_check_atomic(shape):
     detector = [np.linspace(-1, 1, s) for s in shape]