Performances (aiidateam#86)

* adding key sorting in npz writing for hashing

* improving performances

symbols and kinds names and site_indices now are in the repository.
the node size in the db does not change increasing the number of sites.

Author: mikibonacci

src/aiida_atomistic/data/structure/models.py

@@ -40,7 +40,7 @@ class StructureBaseModel(BaseModel):
     sites: list[Site] = Field(
         default=[],
         description="List of sites in the structure",
-        json_schema_extra={"store_in": "db", "property_type": "global"},
+        json_schema_extra={"property_type": "global"},
     )
 
     # global and more specific properties
@@ -197,7 +197,7 @@ def positions(self) -> np.ndarray:
             return None
         return np.array([site.position for site in self.sites])
 
-    @computed_field(json_schema_extra={"store_in": "db","singular_form": "kind_name"})
+    @computed_field(json_schema_extra={"store_in": "repository","singular_form": "kind_name"})
     @property
     def kind_names(self) -> t.List[str]:
         """
@@ -210,7 +210,7 @@ def kind_names(self) -> t.List[str]:
             return None
         return FrozenList([site.kind_name if site.kind_name is not None else site.symbol for site in self.sites])
 
-    @computed_field(json_schema_extra={"store_in": "db","singular_form": "symbol"})
+    @computed_field(json_schema_extra={"store_in": "repository","singular_form": "symbol"})
     @property
     def symbols(self) -> t.List[str]:
         """
@@ -393,6 +393,12 @@ def min_magnetization(self) -> t.Optional[float]:
     def n_sites(self) -> int:
         """Total number of sites in the structure."""
         return len(self.sites)
+    
+    @computed_field(json_schema_extra={"store_in": "db"})
+    @property
+    def n_kinds(self) -> int:
+        """Total number of sites in the structure."""
+        return len(self.kinds) if self.kinds is not None else 0
 
     def __repr__(self) -> str:
         from pprint import pformat

src/aiida_atomistic/data/structure/structure.py

@@ -83,6 +83,15 @@ def _is_numeric_array(value) -> bool:
             return arr.dtype.kind in {"i", "f", "u", "c"}
         return False
 
+    @staticmethod
+    def _is_string_array(value) -> bool:
+        """Return True if the value is a string array (list of strings or numpy unicode array)."""
+        if isinstance(value, np.ndarray):
+            return value.dtype.kind in {"U", "S"}
+        if isinstance(value, list) and value and all(isinstance(v, str) for v in value):
+            return True
+        return False
+
     @classmethod
     def get_queryable_properties(cls, include_internal: bool = False) -> dict:
         """
@@ -272,7 +281,7 @@ def detect_storage_backend(cls, prop_name: str) -> str:
             store_in = json_schema_extra.get(cls._storage_metadata_key, '').lower()
         else:
             store_in = 'db'  # default to db for unknown properties
-
+        
         return store_in
 
     def _store_properties(self):
@@ -309,6 +318,23 @@ def _store_properties(self):
                 repository_dict[prop_name] = arr
                 if self._store_shape_metadata:
                     database_dict[f'shape|{prop_name}'] = list(arr.shape)
+            elif target == "repository" and self._is_string_array(value):
+                arr = np.asarray(value, dtype=str)
+                repository_dict[prop_name] = arr
+                if self._store_shape_metadata:
+                    database_dict[f'shape|{prop_name}'] = list(arr.shape)
+            elif prop_name == 'site_indices':
+                # site_indices is a ragged list-of-lists (one per kind, variable length because of different number of sites per kind).
+                # Encode as two flat 1D int arrays using CSR format so the npz stays
+                # homogeneous (allow_pickle=False compatible):
+                #   site_indices_flat    : all indices concatenated, shape (total_sites,)
+                #   site_indices_offsets : cumulative start positions, shape (n_kinds + 1,)
+                # e.g. [[0,1],[2],[3,4,5]] → flat=[0,1,2,3,4,5], offsets=[0,2,3,6]
+                flat = np.array([idx for sublist in value for idx in sublist], dtype=np.int64)
+                lengths = np.array([len(sublist) for sublist in value], dtype=np.int64)
+                offsets = np.concatenate([[0], np.cumsum(lengths)]).astype(np.int64)
+                repository_dict['site_indices_flat'] = flat
+                repository_dict['site_indices_offsets'] = offsets
             else:
                 database_dict[prop_name] = value
 
@@ -343,7 +369,23 @@ def _load_properties_from_npz(self) -> dict:
         with self.base.repository.open(self._properties_filename, mode='rb') as handle:
             npz_data = np.load(handle, allow_pickle=False)
             # Convert to regular dict (npz returns NpzFile object)
-            properties = {key: npz_data[key] for key in npz_data.files}
+            # String arrays (dtype 'U' or 'S') are converted back to Python lists
+            properties = {}
+            for key in npz_data.files:
+                arr = npz_data[key]
+                if arr.dtype.kind in {"U", "S"}:
+                    properties[key] = arr.tolist()
+                else:
+                    properties[key] = arr
+
+            # Decode CSR-encoded site_indices back into list-of-lists
+            if 'site_indices_flat' in properties and 'site_indices_offsets' in properties:
+                flat = properties.pop('site_indices_flat')
+                offsets = properties.pop('site_indices_offsets')
+                properties['site_indices'] = [
+                    flat[offsets[i]:offsets[i + 1]].tolist()
+                    for i in range(len(offsets) - 1)
+                ]
 
         # Cache if stored
         if self.is_stored:

tests/data/test_kinds.py

@@ -178,16 +178,16 @@ def test_kinds_compression_storage(self, complex_example_structure_dict_for_kind
         structure = StructureData(**complex_example_structure_dict_for_kinds)
 
         # Check that kind_names is in attributes
-        assert "kind_names" in structure.base.attributes.all
+        assert "n_kinds" in structure.base.attributes.all
 
         # Check that sites maintain their structure (not compressed, site-based model)
-        stored_symbols = structure.base.attributes.get("symbols")
-        assert len(stored_symbols) == 8  # 8 sites (site-based, not kind-compressed)
+        n_sites = structure.base.attributes.get("n_sites")
+        assert n_sites == 8  # 8 sites (site-based, not kind-compressed)
 
         # But kind_names should have been assigned
-        kind_names = structure.base.attributes.get("kind_names")
-        assert kind_names is not None
-        assert len(set(kind_names)) == 4  # 4 unique kinds
+        n_kinds = structure.base.attributes.get("n_kinds")
+        assert n_kinds is not None
+        assert n_kinds == 4  # 4 unique kinds
 
 
 class TestKindsWorkflow:

tests/data/test_structure.py

@@ -644,7 +644,8 @@ def test_store_properties_with_kind_compression(aiida_profile_clean):
     )
 
     # Should have kind_names in attributes
-    assert 'kind_names' in structure.base.attributes.all
+    assert 'n_kinds' in structure.base.attributes.all
+    assert structure.base.attributes.all['n_kinds'] == 1
 
 
 def test_load_properties_from_npz(aiida_profile_clean):