Unify assign_site_occupations signature to accept lattice_matrix

Update the abstract SiteCollection.assign_site_occupations() and all
implementations to accept lattice_matrix (numpy array) instead of
Structure. Trajectory.assign_site_occupations() extracts
structure.lattice.matrix at the boundary before delegating.

PolyhedralSiteCollection accepts lattice_matrix for interface
consistency (not currently used for containment checks).

Author: bjmorgan

site_analysis/polyhedral_site_collection.py

@@ -70,17 +70,19 @@ def analyse_structure(self,
         lattice_matrix = structure.lattice.matrix
         for s in self.sites:
             s.notify_structure_changed(all_frac_coords, lattice_matrix)
-        self.assign_site_occupations(atoms, structure)
-                    
-    def assign_site_occupations(self, atoms, structure) -> None:
+        self.assign_site_occupations(atoms, lattice_matrix)
+
+    def assign_site_occupations(self, atoms, lattice_matrix) -> None:
         """Assign atoms to polyhedral sites based on their positions.
-        
+
         This method implements an optimised assignment logic using a priority-based
         site checking approach.
-        
+
         Args:
-            atoms: List of Atom objects to be assigned to sites
-            structure: Pymatgen Structure containing the atom positions
+            atoms: List of Atom objects to be assigned to sites.
+            lattice_matrix: (3, 3) lattice matrix where rows are lattice
+                vectors. Not currently used for polyhedral containment
+                checks, but accepted for interface consistency.
         """
         self.reset_site_occupations()
         for atom in atoms:

site_analysis/site_collection.py

@@ -210,23 +210,17 @@ def __init__(self, sites: Sequence[Site]) -> None:
             self._site_lookup[site.index] = site
 
     @abstractmethod
-    def assign_site_occupations(self, atoms, structure):
-        """Assigns atoms to sites for a specific structure.
-
-        This method should be implemented in the derived subclass
+    def assign_site_occupations(self, atoms, lattice_matrix):
+        """Assign atoms to sites.
 
         Args:
-            atoms (list(Atom)): List of Atom objects to be assigned to sites.
-            struture (pymatgen.Structure): Pymatgen Structure object used to specificy
-                the atomic coordinates.
-
-        Returns:
-            None
-
-        Notes:
-            The atom coordinates should already be consistent with the coordinates
-            in `structure`. Recommended usage is via the ``analyse_structure()`` method.
+            atoms: List of Atom objects to be assigned to sites.
+            lattice_matrix: (3, 3) lattice matrix where rows are lattice
+                vectors.
 
+        Note:
+            The atom coordinates should already be consistent with the
+            structure. Recommended usage is via ``analyse_structure()``.
         """
         raise NotImplementedError('assign_site_occupations should be implemented in'
             ' the derived class')

site_analysis/trajectory.py

@@ -142,13 +142,13 @@ def analyse_structure(self,
     def assign_site_occupations(self,
             structure: Structure) -> None:
         """Assign atoms to sites for a specific structure.
-        
+
         This delegates the assignment to the site collection's assign_site_occupations method.
-        
+
         Args:
             structure: A pymatgen Structure object to be analysed.
         """
-        self.site_collection.assign_site_occupations(self.atoms, structure)
+        self.site_collection.assign_site_occupations(self.atoms, structure.lattice.matrix)
 
     def site_coordination_numbers(self) -> Counter:
         """Return the coordination numbers of all sites.

tests/benchmark_containment.py

@@ -81,7 +81,7 @@ def _eager_analyse_structure(self, atoms, structure):
         a.assign_coords(frac_coords)
     for s in self.sites:
         s.assign_vertex_coords(structure)
-    self.assign_site_occupations(atoms, structure)
+    self.assign_site_occupations(atoms, structure.lattice.matrix)
 
 
 def benchmark_assignment(trajectory, structures, repeats):

tests/test_polyhedral_site_collection.py

@@ -108,8 +108,11 @@ def test_analyse_structure(self):
             # Verify each site was notified of the new structure
             self.assertEqual(mock_notify.call_count, 3)
 
-            # Verify assign_site_occupations was called with atoms and structure
-            mock_assign.assert_called_once_with(self.atoms, self.structure)
+            # Verify assign_site_occupations was called with atoms and lattice_matrix
+            mock_assign.assert_called_once()
+            args = mock_assign.call_args[0]
+            self.assertIs(args[0], self.atoms)
+            np.testing.assert_array_equal(args[1], self.structure.lattice.matrix)
 
     def test_assign_site_occupations_atom_in_site(self):
         """Test assign_site_occupations when atoms are already in sites."""
@@ -126,7 +129,7 @@ def test_assign_site_occupations_atom_in_site(self):
             mock_contains_atom.return_value = True
 
             # Call method
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
             # Verify site occupations were reset
             self.assertEqual(self.site1.contains_atoms, [])
@@ -154,7 +157,7 @@ def test_assign_site_occupations_atom_moved(self):
             mock_contains_atom.return_value = False
 
             # Call method
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
             # Verify site occupations were reset
             self.assertEqual(self.site1.contains_atoms, [])
@@ -184,7 +187,7 @@ def test_assign_site_occupations_atom_not_in_site(self):
             mock_contains_atom.side_effect = lambda atom, **kwargs: atom is self.atom1 and mock_contains_atom.mock_calls[0][1][0] is atom
 
             # Call method
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
             # Verify site occupations were reset
             self.assertEqual(self.site1.contains_atoms, [])
@@ -217,7 +220,7 @@ def test_empty_atoms_list_polyhedral(self):
         )
 
         # Call the method with empty atom list
-        collection.assign_site_occupations([], structure)
+        collection.assign_site_occupations([], lattice.matrix)
 
         # Verify that contains_atoms was reset for both sites
         self.assertEqual(site1.contains_atoms, [])
@@ -311,7 +314,7 @@ def test_checks_recent_site_via_priority_heuristic(self):
             patch.object(collection, 'site_by_index') as mock_site_by_index:
             mock_site_by_index.return_value = mock_site
 
-            collection.assign_site_occupations([mock_atom], mock_structure)
+            collection.assign_site_occupations([mock_atom], np.eye(3) * 10.0)
 
             mock_site_by_index.assert_called_with(5)
             mock_update.assert_called_with(mock_site, mock_atom)
@@ -501,7 +504,7 @@ def setUp(self):
     def test_calls_generator_for_each_atom(self):
         """Test that _get_priority_sites is called once per atom."""
         with patch.object(self.collection, '_get_priority_sites', return_value=[]):
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
             self.collection._get_priority_sites.assert_called_once_with(self.atom)
 
     def test_calls_generator_for_multiple_atoms(self):
@@ -511,7 +514,7 @@ def test_calls_generator_for_multiple_atoms(self):
         atoms = atoms_from_structure(self.structure, "Li")
 
         with patch.object(self.collection, '_get_priority_sites', return_value=[]):
-            self.collection.assign_site_occupations(atoms, self.structure)
+            self.collection.assign_site_occupations(atoms, self.lattice.matrix)
             self.assertEqual(self.collection._get_priority_sites.call_count, 2)
 
     def test_checks_sites_in_generator_order(self):
@@ -523,7 +526,7 @@ def test_checks_sites_in_generator_order(self):
         with patch.object(self.collection, '_get_priority_sites') as mock_gen:
             mock_gen.return_value = [self.site2, self.site1]  # site2 first
 
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
             self.assertEqual(call_order, [2])  # Only site2 checked (found there)
 
@@ -535,7 +538,7 @@ def test_stops_checking_when_atom_found(self):
         with patch.object(self.collection, '_get_priority_sites') as mock_gen:
             mock_gen.return_value = [self.site1, self.site2]
 
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
             self.site1.contains_atom.assert_called_once()
             self.site2.contains_atom.assert_not_called()
@@ -546,7 +549,7 @@ def test_calls_update_occupation_when_found(self):
 
         with patch.object(self.collection, '_get_priority_sites', return_value=[self.site1]):
             with patch.object(self.collection, 'update_occupation') as mock_update:
-                self.collection.assign_site_occupations(self.atoms, self.structure)
+                self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
                 mock_update.assert_called_once_with(self.site1, self.atom)
 
     def test_handles_atom_not_found(self):
@@ -555,7 +558,7 @@ def test_handles_atom_not_found(self):
 
         with patch.object(self.collection, '_get_priority_sites', return_value=[self.site1]):
             with patch.object(self.collection, 'update_occupation') as mock_update:
-                self.collection.assign_site_occupations(self.atoms, self.structure)
+                self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
 
                 mock_update.assert_not_called()
                 self.assertIsNone(self.atom.in_site)
@@ -564,15 +567,15 @@ def test_resets_site_occupations(self):
         """Test that reset_site_occupations is called at start."""
         with patch.object(self.collection, 'reset_site_occupations') as mock_reset:
             with patch.object(self.collection, '_get_priority_sites', return_value=[]):
-                self.collection.assign_site_occupations(self.atoms, self.structure)
+                self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
                 mock_reset.assert_called_once()
 
     def test_resets_atom_in_site(self):
         """Test that atom.in_site is reset to None."""
         self.atom.in_site = 999  # Set to some previous value
 
         with patch.object(self.collection, '_get_priority_sites', return_value=[]):
-            self.collection.assign_site_occupations(self.atoms, self.structure)
+            self.collection.assign_site_occupations(self.atoms, self.lattice.matrix)
             self.assertIsNone(self.atom.in_site)
 
 

tests/test_site_collection.py

@@ -33,10 +33,10 @@ def test_assign_site_occupations_raises_not_implemented_error(self):
                  Mock(spec=Site, index=1)]
         atoms = [Mock(spec=Atom), 
                  Mock(spec=Atom)]
-        structure = Mock(spec=Structure)
+        lattice_matrix = np.eye(3) * 10.0
         site_collection = ConcreteSiteCollection(sites=sites)
         with self.assertRaises(NotImplementedError):
-            site_collection.assign_site_occupations(atoms, structure)
+            site_collection.assign_site_occupations(atoms, lattice_matrix)
 
     def test_analyse_structure_raises_not_implemented_error(self):
         sites = [Mock(spec=Site, index=0),

tests/test_trajectory.py

@@ -307,20 +307,15 @@ def test_site_labels(self):
         self.assertEqual(labels[1], "site2")
 
     def test_assign_site_occupations(self):
-        """Test that assign_site_occupations delegates to site_collection."""
-        # Setup
-        structure = Mock()
-        
-        # Mock the site_collection
+        """Test that assign_site_occupations extracts lattice_matrix and delegates."""
         self.trajectory.site_collection = Mock()
-        
-        # Call the method
-        self.trajectory.assign_site_occupations(structure)
-        
-        # Check delegation
-        self.trajectory.site_collection.assign_site_occupations.assert_called_once_with(
-            self.atoms, structure
-        )
+
+        self.trajectory.assign_site_occupations(self.structure)
+
+        self.trajectory.site_collection.assign_site_occupations.assert_called_once()
+        args = self.trajectory.site_collection.assign_site_occupations.call_args[0]
+        self.assertIs(args[0], self.atoms)
+        np.testing.assert_array_equal(args[1], self.structure.lattice.matrix)
 
     def test_trajectory_from_structures_with_progress(self):
         """Test trajectory_from_structures wraps iterator with tqdm."""

tests/test_voronoi_site_collection.py

@@ -90,6 +90,10 @@ def test_assign_site_occupations_distance_matrix(self):
 				self.atoms, self.structure.lattice.matrix)
 
 			mock_distances.assert_called_once()
+			args = mock_distances.call_args[0]
+			np.testing.assert_array_equal(args[0], np.array([s.frac_coords for s in self.collection.sites]))
+			np.testing.assert_array_equal(args[1], np.array([a.frac_coords for a in self.atoms]))
+			np.testing.assert_array_equal(args[2], self.structure.lattice.matrix)
 			mock_update.assert_any_call(self.site1, self.atom1)
 			mock_update.assert_any_call(self.site2, self.atom2)
 			self.assertEqual(mock_update.call_count, 2)