Improving performances

Author: mikibonacci

pyproject.toml

@@ -26,7 +26,8 @@ dependencies = [
     "voluptuous",
     "ase",
     "pymatgen>=2022.1.20",
-    "pydantic>=2.6"
+    "pydantic>=2.6",
+    "scipy>=1.13.1",
 ]
 
 [project.urls]

src/aiida_atomistic/data/structure/getter_mixin.py

@@ -13,7 +13,7 @@
     HubbardGetterMixin,
 )
 
-from aiida_atomistic.data.structure.utils import classify_site_kinds, check_kinds_match, efficient_copy
+from aiida_atomistic.data.structure.utils import classify_site_kinds, check_kinds_match
 
 try:
     import ase  # noqa: F401
@@ -397,7 +397,7 @@ def to_dict(self, exclude_kinds=False):
             :return: The structure as a dictionary.
             :rtype: dict
             """
-            dict_repr = efficient_copy(self.properties.model_dump(exclude_unset=True, exclude_none=True, warnings=False, exclude={'kinds'} if exclude_kinds else {}))
+            dict_repr = copy.deepcopy(self.properties.model_dump(exclude_unset=True, exclude_none=True, warnings=False, exclude={'kinds'} if exclude_kinds else {}))
 
             return dict_repr
 

src/aiida_atomistic/data/structure/models.py

@@ -116,9 +116,9 @@ def validate_sites(cls, v):
 
         if v is None:
             return v
-        else:
-            # test if they can be converted to Site
-            sites = [Site.model_validate(site) if not isinstance(site, Site) else site for site in v]
+        # else:
+        #     # test if they can be converted to Site
+        #     sites = [Site.model_validate(site) if not isinstance(site, Site) else site for site in v]
 
         _check_valid_sites(v)
 
@@ -291,20 +291,33 @@ def kinds(self) -> list[Kind]:
             #raise ValueError("Kind names must be defined to access kinds.")
             return None
 
+        # Mapping of kind_name -> site indices
+        kind_to_indices = defaultdict(list)
+        for i, name in enumerate(self.kind_names):
+            kind_to_indices[name].append(i)
+
+        positions_array = self.positions
+
         kinds_list = []
-        kind_name_set = set(self.kind_names)
-        for idx, site in enumerate(self.sites):
+        seen_kinds = set()
+
+        for site in self.sites:
             kind_name = site.kind_name if site.kind_name else site.symbol
-            if kind_name in kind_name_set:
-                site_indices = [i for i, name in enumerate(self.kind_names) if name == kind_name]
-                positions=np.array([self.positions[i] for i in site_indices])
-                kind = Kind(
-                    **site.model_dump(exclude={'position'}),
-                    site_indices=site_indices,
-                    positions=positions,
-                )
-                kinds_list.append(kind)
-                kind_name_set.remove(kind_name)  # Ensure we don't add the same kind multiple
+
+            # Skip if we've already processed this kind
+            if kind_name in seen_kinds:
+                continue
+
+            seen_kinds.add(kind_name)
+            site_indices = kind_to_indices[kind_name]
+            positions = positions_array[site_indices]
+
+            kind = Kind(
+                **site.model_dump(exclude={'position'}),
+                site_indices=site_indices,
+                positions=positions,
+            )
+            kinds_list.append(kind)
 
         return FrozenList(kinds_list)
 

src/aiida_atomistic/data/structure/utils.py

@@ -1,9 +1,11 @@
 import copy
 import functools
 import re
+import numpy as np
 
 import typing as t
-import numpy as np
+
+from scipy.spatial import cKDTree
 
 from aiida.common.constants import elements
 from aiida.common.exceptions import UnsupportedSpeciesError
@@ -98,6 +100,7 @@ def efficient_copy(obj):
 
     Handles both dictionaries and lists, as well as other types.
     """
+
     if obj is None:
         return None
     elif isinstance(obj, dict):
@@ -167,30 +170,27 @@ def _get_valid_pbc(inputpbc):
 
     return the_pbc
 
+def any_close_pairs(points, eps):
+    tree = cKDTree(points)
+    pairs = tree.query_pairs(r=eps)
+    return len(pairs) > 0,pairs
+
 def _check_valid_sites(sites):
     """Check that no two sites have positions that are too close to each other."""
 
     positions = np.array([site['position'] for site in sites])
+
     n_sites = len(positions)
 
     if n_sites <= 1:
         return
 
-    # Calculate pairwise distances using broadcasting (this is much more efficient than loops...)
-    diff = positions[:, np.newaxis, :] - positions[np.newaxis, :, :]  # Shape: (n_sites, n_sites, 3)
-    distances = np.linalg.norm(diff, axis=2)  # Shape: (n_sites, n_sites)
-
-    # Set diagonal to large value to ignore self-comparisons
-    np.fill_diagonal(distances, np.inf)
-
     # Check if any distance is below threshold
     min_distance = 1e-3  # You can adjust this threshold
-    close_pairs = np.where(distances < min_distance)
+    close_pairs = any_close_pairs(positions, eps=min_distance)
 
-    if len(close_pairs[0]) > 0:
-        i, j = close_pairs[0][0], close_pairs[1][0]  # Get first problematic pair
-        raise ValueError(f"Sites {i} and {j} have positions that are too close: "
-                       f"{positions[i]} and {positions[j]} (distance: {distances[i,j]:.6f})")
+    if close_pairs[0]:
+        raise ValueError(f"The following sites have positions that are too close (less than {min_distance}): {close_pairs[1]}.")
     return
 
 
@@ -390,7 +390,7 @@ def group_symbols(_list):
     :param _list: a list of elements representing a chemical formula
     :return: a list of length-2 lists of the form [ multiplicity , element ]
     """
-    the_list = efficient_copy(_list)
+    the_list = copy.deepcopy(_list)
     the_list.reverse()
     grouped_list = [[1, the_list.pop()]]
     while the_list:
@@ -468,7 +468,7 @@ def group_together(_list, group_size, offset):
             ``group_together(['O','Ba','Ti','Ba','Ti'],2,1) =
                 ['O',['Ba','Ti'],['Ba','Ti']]``
         """
-        the_list = efficient_copy(_list)
+        the_list = copy.deepcopy(_list)
         the_list.reverse()
         grouped_list = []
         for _ in range(offset):
@@ -507,14 +507,14 @@ def group_together_symbols(_list, group_size):
         :return the_symbol_list: the new grouped symbol list
         :return has_grouped: True if we grouped something
         """
-        the_symbol_list = efficient_copy(_list)
+        the_symbol_list = copy.deepcopy(_list)
         has_grouped = False
         offset = 0
         while not has_grouped and offset < group_size:
             grouped_list = group_together(the_symbol_list, group_size, offset)
             new_symbol_list = group_symbols(grouped_list)
             if len(new_symbol_list) < len(grouped_list):
-                the_symbol_list = efficient_copy(new_symbol_list)
+                the_symbol_list = copy.deepcopy(new_symbol_list)
                 the_symbol_list = cleanout_symbol_list(the_symbol_list)
                 has_grouped = True
                 # print get_formula_from_symbol_list(the_symbol_list)
@@ -530,7 +530,7 @@ def group_all_together_symbols(_list):
         """
         has_finished = False
         group_size = 2
-        the_symbol_list = efficient_copy(_list)
+        the_symbol_list = copy.deepcopy(_list)
 
         while not has_finished and group_size <= len(_list) // 2:
             # try to group as much as possible by groups of size group_size
@@ -551,7 +551,7 @@ def group_all_together_symbols(_list):
     # successively apply the grouping procedure until the symbol list does not
     # change anymore
     while new_symbol_list != old_symbol_list:
-        old_symbol_list = efficient_copy(new_symbol_list)
+        old_symbol_list = copy.deepcopy(new_symbol_list)
         new_symbol_list = group_all_together_symbols(old_symbol_list)
 
     return get_formula_from_symbol_list(new_symbol_list, separator=separator)
@@ -879,7 +879,7 @@ def check_is_alloy(data):
     :param data: the data to check. The dict of the SiteCore model.
     :return: True if the data is an alloy, False otherwise.
     """
-    new_data = efficient_copy(data)
+    new_data = copy.deepcopy(data)
     if "weight" not in new_data.keys() or new_data.get("weight", None) is None:
         if isinstance(new_data["symbol"], list) or re.search(r'[A-Z][a-z]*[A-Z]', new_data["symbol"]):
             return new_data
@@ -1003,20 +1003,23 @@ def build_sites_from_expanded_properties(expanded):
     # Use all keys except positions if you want to exclude arrays, or specify your own
     site_props = set(expanded.keys()).difference(_GLOBAL_PROPERTIES + _COMPUTED_PROPERTIES + ["sites", "site_indices"])
 
-    n_sites = len(expanded.get("positions",[]))
-    sites = []
-    for i in range(n_sites):
-        site = {}
-        site["position"] = expanded["positions"][i]
-        for prop in site_props:
-            site[_CONVERSION_PLURAL_SINGULAR[prop]] = expanded[prop][i]
-        sites.append(site)
+    # Pre-compute the conversion mapping to avoid repeated dict lookups
+    prop_conversions = [(prop, _CONVERSION_PLURAL_SINGULAR[prop], expanded[prop])
+                        for prop in site_props]
 
-    structure_dict = {}
-    for prop in _GLOBAL_PROPERTIES:
-        if expanded.get(prop, None) is not None:
-            structure_dict[prop] = expanded[prop]
+    positions = expanded.get("positions", [])
+    n_sites = len(positions)
 
+    sites = [
+        {"position": positions[i], **{singular: values[i] for _, singular, values in prop_conversions}}
+        for i in range(n_sites)
+    ]
+
+    structure_dict = {
+        prop: expanded[prop]
+        for prop in _GLOBAL_PROPERTIES
+        if prop in expanded and expanded[prop] is not None
+    }
     structure_dict["sites"] = sites
 
     return structure_dict
@@ -1145,7 +1148,7 @@ def sites_from_kinds(kinds):
         positions += list(kind['positions'])
     num_sites = len(sites_list)
     for i in range(num_sites):
-        sites_list[i] = efficient_copy(kinds[sites_list[i]])
+        sites_list[i] = copy.deepcopy(kinds[sites_list[i]])
         sites_list[i].pop('site_indices', None)
         sites_list[i].pop('positions', None)
         sites_list[i]['position'] = positions[i]