Added new tests for Dirax indexing

crystals/indexing/__init__.py

@@ -20,7 +20,9 @@ class IndexingMethod(Enum):
     DirAx = auto()
 
 
-def index_reflections(reflections, length_bounds=(2, 20), method=IndexingMethod.DirAx):
+def index_reflections(
+    reflections, initial=None, length_bounds=(2, 20), method=IndexingMethod.DirAx
+):
     """
     Find the lattice associated with a list of reflections.
 
@@ -32,6 +34,9 @@ def index_reflections(reflections, length_bounds=(2, 20), method=IndexingMethod.
         Iterable of reflections with their three-dimensional reciprocal space
         coordinates, or ndarray where each row is a reflections. Coordinates are
         in inverse Angstroms.
+    initial : :class:`Lattice` or :class:`Crystal`, optional
+        Initial guess for a lattice. Depending on the indexing method,
+        this parameter may not be used.
     length_bounds : 2-tuple of floats, optional
         Minimum and maximum lattice vector lengths to consider, in Angstrom.
     method : IndexingMethod, optional
@@ -64,6 +69,8 @@ def index_reflections(reflections, length_bounds=(2, 20), method=IndexingMethod.
     < Lattice object with parameters 2.464Å, 2.464Å, 6.711Å, 90.00°, 90.00°, 120.00° >
     """
     if method == IndexingMethod.DirAx:
-        return dirax(reflections=reflections, length_bounds=length_bounds)
+        return dirax(
+            reflections=reflections, initial=initial, length_bounds=length_bounds
+        )
     else:
         raise ValueError(f"Expected an indexing method, but got {method}.")

crystals/indexing/dirax.py

@@ -13,7 +13,7 @@
 from .common import IndexingError, row_echelon_form
 
 
-def dirax(reflections, length_bounds=(2, 20)):
+def dirax(reflections, initial=None, length_bounds=(2, 20)):
     """
     Find the lattice associated with a list of
     reflections using the DirAx algorithm.
@@ -26,6 +26,8 @@ def dirax(reflections, length_bounds=(2, 20)):
         Iterable of reflections with their three-dimensional reciprocal space
         coordinates, or ndarray where each row is a reflections. Coordinates are
         in inverse Angstroms.
+    initial : :class:`Lattice` or :class:`Crystal`, optional
+        Initial guess for a lattice.
     length_bounds : 2-tuple of floats, optional
         Minimum and maximum lattice vector lengths to consider, in Angstrom.
 
@@ -60,6 +62,13 @@ def dirax(reflections, length_bounds=(2, 20)):
     ns = np.arange(start=1, stop=13, step=1).reshape((-1, 1))
 
     potential_direct_vectors = set()
+
+    # If a "guess" lattice is provided, we include its
+    # lattice vectors as being high-priority.
+    if initial is not None:
+        for v in initial.lattice_vectors:
+            potential_direct_vectors.add(LatVec(nf=len(reflections), vector=v))
+
     points = [np.squeeze(a) for a in np.vsplit(reflections, reflections.shape[0])]
     for a1, a2, a3 in product(points, repeat=3):
         normal = np.cross(a2 - a1, a3 - a1)

crystals/indexing/tests/test_dirax.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 
-from crystals import Crystal, Lattice, index_reflections, IndexingMethod
+from crystals import Crystal, Lattice, IndexingError, index_reflections, IndexingMethod
 import pytest
 
 np.random.seed(2021)
@@ -23,7 +23,19 @@ def test_dirax_indexing_ideal(name, bound):
     assert np.allclose(hkls - np.rint(hkls), 0, atol=0.1)
 
 
-# The structures to test have been chosen so that it doesn't take too long.
+@pytest.mark.parametrize("name", ["Pu-epsilon", "C", "vo2-m1", "BaTiO3_cubic"])
+def test_dirax_indexing_initial_guess(name):
+    """
+    Test that indexing succeeds with an initial guess and very few reflections.
+    """
+    cryst = Crystal.from_database(name)
+    # We restrict the number of reflections to a single (0,0,0); with
+    # the initial guess, indexing should still succeed!
+    lat, _ = index_reflections([(0, 0, 0)], initial=cryst, method=IndexingMethod.DirAx)
+
+    assert np.allclose(lat.lattice_parameters, cryst.lattice_parameters, atol=1)
+
+
 @pytest.mark.parametrize(
     "name,bound", zip(["Pu-epsilon", "C", "vo2-m1", "BaTiO3_cubic"], [2, 3, 2, 2])
 )
@@ -42,3 +54,19 @@ def test_dirax_indexing_alien_reflections(name, bound):
     # The alien reflections will not be indexed correctly, of course
     hkls = hkls[:num_aliens]
     assert np.allclose(hkls - np.rint(hkls), 0, atol=0.1)
+
+
+def test_dirax_indexing_length_bound():
+    """
+    Test that indexing fails as expected if the lattice length bounds are too restrictive.
+    """
+    cryst = Crystal.from_database("C")
+    refls = [cryst.scattering_vector(r) for r in cryst.bounded_reflections(bound=2)]
+    # We restrict the number of reflections to a single (0,0,0); with
+    # the initial guess, indexing should still succeed!
+    with pytest.raises(IndexingError):
+        index_reflections(
+            refls,
+            length_bounds=(0.01, 2),
+            method=IndexingMethod.DirAx,
+        )