Custom props (#77)

* fixing the __repr__ and __str__ for the properties.

* Adding set_custom and remove_custom methods.

This will work for StructureBuilder, as they are in the setter_mixin.py

Author: mikibonacci

docs/source/how_to/creation_mutation.md

@@ -108,18 +108,6 @@ Number of sites: 4
 Kinds: {'Fe1', 'O1'}
 ```
 
-**Advantages of kinds-based format:**
-- More compact for structures with many equivalent atoms
-- Explicit grouping of atoms with same properties
-- Matches the internal storage format used by AiiDA
-- Useful when migrating from legacy `orm.StructureData`
-
-**When to use:**
-- Structures with many atoms of the same type and properties
-- Converting from legacy AiiDA format
-- When you already have kind information from calculations
-- Large structures where compactness matters
-
 ### From ASE
 
 Convert ASE Atoms objects to StructureData:

docs/source/how_to/define_custom.md

@@ -27,3 +27,9 @@ structure_dict = {
 structure = StructureData(**structure_dict)
 print(structure.properties.custom["my_special_property"])
 ```
+
+To add a custom properties on an already initialised `StructureBuilder`, you can directly define `structure.properties.custom` or use the `set_custom` method.
+In the latter case, if `structure.properties.custom` is already defined, it will be updated with the new values provided in the `set_custom` method, but other keys will not be removed.
+
+To delete a set of defined custom properties, use the `remove_custom` method, passing as input a list of keys (strings). To remove the entire custom dictionary, you
+can call the same method without specifying any input parameter.

docs/source/tutorials/1_first_structure.md

@@ -84,20 +84,18 @@ Information on the supported properties can be obtained by using the `get_suppor
 
 ```python
 # Get all supported properties
-supported = structure.get_supported_properties()
+supported = StructureData.get_supported_properties()
 print(f"Supported properties: {supported}")
+```
+
+In a similar way, it is possible to see the properties which are defined for a given `StructureData` instance:
 
+```python
 # Get defined properties in this structure
 defined = structure.get_defined_properties()
 print(f"Defined properties: {defined}")
 ```
 
-**Output:**
-```
-Supported properties: {'global': {'custom', 'tot_magnetization', 'sites', 'cell', 'hubbard', 'tot_charge', 'pbc'}, 'site': {'magnetization', 'kind_name', 'symbol', 'mass', 'magmom', 'position', 'weight', 'charge'}}
-Defined properties: {'positions', 'symbols', 'masses', 'sites', 'cell', 'kind_names', 'pbc'}
-```
-
 ## Modifying Structures
 
 For modifications, use `StructureBuilder` (which is the mutable non-AiiDA version of the `StructureData`):

src/aiida_atomistic/data/structure/models.py

@@ -122,6 +122,13 @@ def freeze_sites(cls, v):
             return freeze_nested(v)
         return v
 
+    @field_validator('custom', mode='after')
+    def freeze_custom(cls, v):
+        """Freeze the list of sites if the structure is immutable."""
+        if not cls._mutable and v is not None:
+            return freeze_nested(v)
+        return v
+
     # computed properties
     @computed_field
     def cell_volume(self) -> float:
@@ -314,65 +321,12 @@ def kinds(self) -> list[Kind]:
         return FrozenList(kinds_list)
 
     def __repr__(self) -> str:
-        """Return a concise string representation of the structure."""
-        # Basic info
-        nsites = len(self.sites)
-        formula = self.formula
-
-        # PBC info
-        pbc_dims = sum(self.pbc)
-        if pbc_dims == 3:
-            pbc_str = "3D"
-        elif pbc_dims == 2:
-            pbc_str = "2D"
-        elif pbc_dims == 1:
-            pbc_str = "1D"
-        else:
-            pbc_str = "0D"
-
-        # Cell volume
-        volume = self.cell_volume
-
-        parts = [
-            f"formula: {formula}",
-            f"sites: {nsites}",
-            f"dimensionality: {pbc_str}",
-            f"V={volume:.2f} A^3"
-        ]
-
-        # Add magnetic info if present
-        if self.tot_magnetization is not None:
-            parts.append(f"tot_mag={self.tot_magnetization:.2f}")
-        elif any(s.magnetization is not None or s.magmom is not None for s in self.sites):
-            parts.append("magnetic")
-
-        # Add charge info if present
-        if self.tot_charge is not None:
-            parts.append(f"tot_charge={self.tot_charge:.2f}")
-        elif any(s.charge is not None for s in self.sites):
-            parts.append("charged")
-
-        # Add alloy/vacancy info
-        if self.is_alloy:
-            parts.append("alloy")
-        if self.has_vacancies:
-            parts.append("vacancies")
-
-        # First line with summary
-        repr_str = f" | {', '.join(parts)} |"
-
-        # Add sites info (limit to first 5 sites to avoid too long representations)
-        max_sites_to_show = 5
-        if nsites > 0:
-            repr_str += "\n Sites:"
-            for i, site in enumerate(self.sites):
-                if i >= max_sites_to_show:
-                    repr_str += f"\n  ... (+{nsites - max_sites_to_show} more sites)"
-                    break
-                repr_str += f"\n  {site}"
-            repr_str += "\n"
-
-        return repr_str
+        from pprint import pformat
+        pformatted = pformat(self.model_dump())
+        return f"StructureModel({pformatted})"
+
+    def __str__(self):
+        return self.__repr__()
 
 class MutableStructureModel(StructureBaseModel):
     """
@@ -432,5 +386,5 @@ def freeze_pbc(cls, v):
     def __setattr__(self, key, value):
         # Customizing the exception message when trying to mutate attributes
         if key in self.model_fields:
-            raise ValueError("The AiiDA `StructureData` is immutable. You can create a mutable copy of it using its `get_value` method.")
+            raise ValueError("The AiiDA `StructureData` is immutable. You can create a mutable copy of it using its `to_builder` method.")
         super().__setattr__(key, value)

src/aiida_atomistic/data/structure/setter_mixin.py

@@ -162,7 +162,6 @@ def clear_sites(self,):
         del self.properties.sites
         return
 
-
     def remove_property(self, property_name):
         """Clear the given property."""
 
@@ -257,3 +256,27 @@ def set_kind_names(self, value: list):
         for site, kind_name in zip(self.properties.sites, value):
             site.kind_name = kind_name
         return
+
+    def set_custom(self, value: dict):
+        """Set the custom properties."""
+        if not self.properties.custom:
+            self.properties.custom = {}
+        self.properties.custom.update(value)
+        return
+
+    def remove_custom(self, keys: t.List[str] = None):
+        """Remove the custom properties.
+
+        If keys is None, remove all custom properties.
+        If keys is provided, remove only the specified keys.
+        """
+
+        if not self.properties.custom:
+            return
+        if keys:
+            for key in keys:
+                if key in self.properties.custom:
+                    del self.properties.custom[key]
+        else:
+            self.remove_property('custom')
+        return

src/aiida_atomistic/data/structure/structure.py

@@ -67,13 +67,16 @@ def to_builder(self) -> 'StructureBuilder':
 
     def __repr__(self) -> str:
         """Return a concise string representation of the structure."""
+
+        from aiida_atomistic.data.structure.utils import get_structure_repr
+
         # Build UUID string without calling super().__repr__() to avoid recursion
         if self.is_stored:
             uuid_str = f'<{self.__class__.__name__}: uuid: {self.uuid} (pk: {self.pk})>'
         else:
             uuid_str = f'<{self.__class__.__name__}: uuid: {self.uuid} (unstored)>'
 
-        prop_repr_str = self.properties.__repr__()
+        prop_repr_str = get_structure_repr(self)
         return uuid_str + f'\n {prop_repr_str.replace("ImmutableStructureModel","")}'
 
     def __str__(self) -> str:
@@ -113,7 +116,10 @@ def to_aiida(self) -> 'StructureData':
 
     def __repr__(self) -> str:
         """Return a concise string representation of the structure."""
-        prop_repr_str = self.properties.__repr__()
+
+        from aiida_atomistic.data.structure.utils import get_structure_repr
+
+        prop_repr_str = get_structure_repr(self)
         return super().__repr__() + f'\n {prop_repr_str.replace("MutableStructureModel","")}'
 
     def __str__(self) -> str:

src/aiida_atomistic/data/structure/utils.py

@@ -936,3 +936,65 @@ def build_sites_from_expanded_properties(expanded):
     structure_dict["sites"] = sites
 
     return structure_dict
+
+def get_structure_repr(structure):
+    """Return a concise string representation of the structure."""
+
+    # Basic info
+    nsites = len(structure.properties.sites)
+    formula = structure.properties.formula
+
+    # PBC info
+    pbc_dims = sum(structure.properties.pbc)
+    if pbc_dims == 3:
+        pbc_str = "3D"
+    elif pbc_dims == 2:
+        pbc_str = "2D"
+    elif pbc_dims == 1:
+        pbc_str = "1D"
+    else:
+        pbc_str = "0D"
+
+    # Cell volume
+    volume = structure.properties.cell_volume
+
+    parts = [
+        f"formula: {formula}",
+        f"sites: {nsites}",
+        f"dimensionality: {pbc_str}",
+        f"V={volume:.2f} A^3"
+    ]
+
+    # Add magnetic info if present
+    if structure.properties.tot_magnetization is not None:
+        parts.append(f"tot_mag={structure.properties.tot_magnetization:.2f}")
+    elif any(s.magnetization is not None or s.magmom is not None for s in structure.properties.sites):
+        parts.append("magnetic")
+
+    # Add charge info if present
+    if structure.properties.tot_charge is not None:
+        parts.append(f"tot_charge={structure.properties.tot_charge:.2f}")
+    elif any(s.charge is not None for s in structure.properties.sites):
+        parts.append("charged")
+
+    # Add alloy/vacancy info
+    if structure.properties.is_alloy:
+        parts.append("alloy")
+    if structure.properties.has_vacancies:
+        parts.append("vacancies")
+
+    # First line with summary
+    repr_str = f" | {', '.join(parts)} |"
+
+    # Add sites info (limit to first 5 sites to avoid too long representations)
+    max_sites_to_show = 5
+    if nsites > 0:
+        repr_str += "\n Sites:"
+        for i, site in enumerate(structure.properties.sites):
+            if i >= max_sites_to_show:
+                repr_str += f"\n  ... (+{nsites - max_sites_to_show} more sites)"
+                break
+            repr_str += f"\n  {site}"
+        repr_str += "\n"
+
+    return repr_str

tests/data/test_property_setters.py

@@ -341,3 +341,36 @@ def test_convert_to_immutable_after_remove(self):
         # Verify charges are not present
         assert immutable.properties.charges is None
         assert 'charges' not in immutable.get_defined_properties()
+
+class TestCustomPropertySettersRemovers:
+    """Test setter and remover methods for custom properties in StructureBuilder."""
+
+    def test_set_then_remove_charges(self):
+        """Test setting charges and then removing them."""
+        structure = StructureBuilder(
+            cell=[[3.0, 0, 0], [0, 3.0, 0], [0, 0, 3.0]],
+            pbc=[True, True, True],
+            sites=[
+                {"symbol": "Fe", "position": [0, 0, 0]},
+                {"symbol": "O", "position": [1.5, 1.5, 1.5]},
+            ],
+            custom = {'first_custom_property': 'Hello'}
+        )
+
+
+        # Initially only one custom property
+        assert structure.properties.custom == {'first_custom_property': 'Hello'}
+        ## testing also for StructureData
+        structuredata = structure.to_aiida()
+        assert structuredata.properties.custom == {'first_custom_property': 'Hello'}
+
+        # Set another property
+        structure.set_custom({'second_custom_property': "World"})
+        assert structure.properties.custom == {'first_custom_property': 'Hello', 'second_custom_property': 'World'}
+
+        # Remove custom properties: first only one, then the whole dictionary
+        structure.remove_custom(['first_custom_property'])
+        assert structure.properties.custom == {'second_custom_property': 'World'}
+
+        structure.remove_custom()
+        assert structure.properties.custom is None