fix whenever atom have type 3 and 1 but not 2

src/WarrenCowleyParameters/__init__.py

@@ -68,121 +68,25 @@ def _validate_neigh_order(self):
         if not np.all(np.diff(self.nneigh) > 0):
             raise ValueError("'Max atoms in shells' must be strictly increasing.")
 
-        if (self.only_selected) and "Selection" not in data.particles.keys():
-            raise KeyError("No selection defined")
+    def _validate_selection_existence(self):
+        if self.only_selected and "Selection" not in self.data.particles.keys():
+            raise KeyError("No selection defined in the data")
 
-    @staticmethod
-    def get_type_name(data, id):
-        # Get the name of a particle type by its ID
-        ptype = data.particles["Particle Type"].type_by_id(id)
-        name = ptype.name
-        if name:
-            return name
-        return f"Type {id}"
-
-    def get_concentration(self, all_particles_types, selected_indices, selected_neigh_idx):
-        # If we are using only selected particle, get concentration from selected atoms + their neighbors
-        if self.only_selected:
-            unique_neighs = np.unique(selected_neigh_idx)
-            combined_mask = np.isin(
-                np.arange(len(all_particles_types)),
-                np.concatenate([selected_indices, unique_neighs]),
-            )
-
-            particle_types = all_particles_types[combined_mask]
-        else:
-            # if no selection, we use all particles
-            particle_types = all_particles_types
-
-        # Calculate the concentration of unique particle types
-        unique_types, counts = np.unique(particle_types, return_counts=True)
-        return unique_types, counts / len(particle_types)
-
-    @staticmethod
-    def get_central_atom_type_mask(unique_types, particles_types):
-        # Create a mask for central atom types
-        central_atom_type_mask = []
-        for atom_type in unique_types:
-            central_atom_type_mask.append(np.where(particles_types == atom_type))
-        return central_atom_type_mask
-
-    @staticmethod
-    def get_wc_from_neigh_in_shell_types(
-        neigh_in_shell_types, central_atom_type_mask, c, unique_types
-    ):
-        # Calculate Warren-Cowley parameters for atoms in shells
-        ntypes = len(c)
-        neight_in_shell = neigh_in_shell_types.shape[1]
-        wc = np.zeros((ntypes, ntypes))
-
-        for i in range(ntypes):
-            neight_type_aroud_itype = neigh_in_shell_types[central_atom_type_mask[i]]
-            neight_type_aroud_itype_flat = neight_type_aroud_itype.flatten()
-
-            counts = np.bincount(neight_type_aroud_itype_flat, minlength=ntypes + 1)
-
-            pij = counts[unique_types] / (neight_type_aroud_itype.shape[0] * neight_in_shell)
-
-            wc[i, :] = 1 - pij / c
-
-        return wc
-
-    @staticmethod
-    def get_particle_types(data):
-        particle_types = np.array(data.particles.particle_type)
-
-        return particle_types
-
-    @staticmethod
-    def get_nearest_neighbors(data, max_number_of_neigh):
-        # Find nearest neighbors for each particle
-        finder = NearestNeighborFinder(max_number_of_neigh, data)
-        neigh_idx, _ = finder.find_all()
-
-        return neigh_idx
 
-    def create_visualization_tables(self, unique_types, nshells, wc_for_shells, data):
-        labels = []
-        warrenCowley = []
-        idx = list(range(len(unique_types)))
-
-        # Get labels and values for Warren-Cowley parameters
-        for m in range(nshells):
-            labels.append([])
-            warrenCowley.append([])
-            for i, j in itertools.combinations_with_replacement(idx, 2):
-                namei = self.get_type_name(data, unique_types[i])
-                namej = self.get_type_name(data, unique_types[j])
-                labels[-1].append(f"{namei}-{namej}")
-                warrenCowley[-1].append(wc_for_shells[m, i, j])
-
-        # Create DataTable objects for visualization
-        for m in range(nshells):
-            table = DataTable(
-                title=f"Warren-Cowley parameter (shell={m+1})",
-                plot_mode=DataTable.PlotMode.BarChart,
-            )
-            table.x = table.create_property("i-j pair", data=range(len(labels[m])))
-            table.x.types = [ElementType(id=idx, name=l) for idx, l in enumerate(labels[m])]
-            table.y = table.create_property(
-                f"Warren-Cowley parameter (shell={m+1})", data=warrenCowley[m]
-            )
-            data.objects.append(table)
-
-    @staticmethod
-    def check_symmetry(arr):
-        try:
-            assert np.allclose(arr.T, arr)
-        except:
-            print("WARNING: WCs are not exactly symmetric.")
-
-    def modify(self, data: DataCollection, frame: int, **kwargs):
-        self.validateInput(data)
+class WarrenCowleyCalculator:
+    def __init__(self, data, nneigh, only_selected):
+        self.data = data
+        self.nneigh = nneigh
+        self.max_number_of_neigh = max(nneigh)
+        self.only_selected = only_selected
 
-        all_particles_types = self.get_particle_types(data)
-        ntypes = len(np.unique(all_particles_types))
-        nparticles = data.particles.count
+    def calculate_warren_cowley_parameters(self):
+        # Extract particle types and NN idices
+        particle_types = self._extract_particle_types()
+        neighbor_indices = self._find_neighbor_indices()
 
+        # If using selected particles, compute the WC based on the selected atom and their NN
+        nparticles = self.data.particles.count
         selected_indices = (
             np.where(self.data.particles["Selection"] == 1)[0]
             if self.only_selected
@@ -213,8 +117,13 @@ def modify(self, data: DataCollection, frame: int, **kwargs):
 
         # Calculate Warren-Cowley parameters for each shell
         for m in range(nshells):
-            neigh_idx_in_shell = selected_neigh_idx[:, self.nneigh[m] : self.nneigh[m + 1]]
-            neigh_in_shell_types = all_particles_types[neigh_idx_in_shell]
+            # Get N shell NN
+            neigh_idx_in_shell = selected_neigh_idx[
+                :, self.nneigh[m] : self.nneigh[m + 1]
+            ]
+
+            # Get atomic types of the NN
+            neigh_in_shell_types = particle_types[neigh_idx_in_shell]
 
             wc = self._compute_wc_params(
                 neigh_in_shell_types, central_atom_mask, concentrations, unique_types
@@ -233,9 +142,17 @@ def _find_neighbor_indices(self):
         return neighbor_indices
 
     def _calculate_concentration(self, particle_types):
-        unique_types, counts = np.unique(particle_types, return_counts=True)
+        #  This is to deal with the case whenever type 3 and 1 are assigned but no atoms have type 2
+        ntypes = np.max(particle_types)
+        unique_types = np.arange(1, ntypes + 1)
+        counts = np.bincount(particle_types)[1:]
         return unique_types, counts / len(particle_types)
 
+        # This used to be the old way
+        # unique_types, counts = np.unique(particle_types, return_counts=True)
+
+    # return unique_types, counts / len(particle_types)
+
     def _create_central_atom_type_mask(self, unique_types, particle_types):
         unique_types_array = unique_types[:, np.newaxis]
         return particle_types == unique_types_array
@@ -275,7 +192,11 @@ def __init__(self, data: DataCollection) -> None:
 
     def get_type_name(self, id: int) -> str:
         """Get the name of a particle type by its ID"""
-        particle_type = self.data.particles["Particle Type"].type_by_id(id)
+        try:
+            particle_type = self.data.particles["Particle Type"].type_by_id(id)
+        except:
+            # This is to deal with the case whenever type 3 and 1 are assigned but no atoms have type 2
+            particle_type.name = False
         return particle_type.name or f"Type {id}"
 
     def create_visualization_tables(self, unique_types, nshells, wc_for_shells):

tests/test_modifier.py

@@ -1,3 +1,6 @@
+import warnings
+
+warnings.filterwarnings("ignore", message=".*OVITO.*PyPI")
 import numpy as np
 import pytest
 from ovito.data import DataCollection
@@ -53,7 +56,6 @@ def test_wc_symmetric(import_data: DataCollection):
 
 
 def test_selection(import_data: DataCollection):
-
     data = import_data
     data.apply(ExpressionSelectionModifier(expression="Position.X > 10"))
     data.apply(WarrenCowleyParameters(only_selected=True))