Sparse set_atom_data with by_species and by_index (closes #65)

docs/_static/generate_images.py

@@ -333,7 +333,7 @@ def logo_scene() -> StructureScene:
     )
 
     data = dict(zip(range(1, 1 + n_verts), colour_vals))
-    scene.set_atom_data("gradient", data)
+    scene.set_atom_data("gradient", by_index=data)
 
     scene.view.look_along(look_dir)
     scene.view.perspective = 0.12
@@ -664,12 +664,12 @@ def red_blue(t: float) -> tuple[float, float, float]:
     type_dict: dict[int, object] = {}
     for i in range(n_outer):
         type_dict[i] = ["Fe", "Co", "Ni"][i % 3]
-    multi_scene.set_atom_data("metal", type_dict)
+    multi_scene.set_atom_data("metal", by_index=type_dict)
     # Inner ring: numerical charge.
     charge_dict: dict[int, object] = {}
     for i in range(n_inner):
         charge_dict[n_outer + i] = float(i) / max(n_inner - 1, 1)
-    multi_scene.set_atom_data("charge", charge_dict)
+    multi_scene.set_atom_data("charge", by_index=charge_dict)
     multi_scene.render_mpl(
         OUT / "colour_by_multi.svg",
         colour_by=["metal", "charge"],

docs/changelog.rst

@@ -4,6 +4,21 @@ Changelog
 0.19.0
 ------
 
+- :meth:`~hofmann.StructureScene.set_atom_data` gains ``by_species``
+  and ``by_index`` keyword arguments for sparse per-atom metadata
+  assignment.  ``by_species`` maps species labels to values;
+  ``by_index`` maps atom indices.  Both can be combined in one call,
+  with ``by_index`` taking precedence at overlapping atoms.
+
+- :meth:`~hofmann.StructureScene.set_atom_data` no longer accepts a
+  dict as its positional ``values`` argument.  Use ``by_index=``
+  instead.
+
+- Missing entries in sparse categorical atom data are now filled with
+  ``None`` (object-dtype) instead of empty strings.  Both are treated
+  as missing by the rendering pipeline; ``None`` cannot collide with
+  a real label.
+
 - The ``AtomData`` container is no longer re-exported from
   ``hofmann`` or ``hofmann.model``.  The only supported way to
   obtain an instance is to read the

docs/colouring.rst

@@ -81,14 +81,46 @@ String arrays assign a distinct colour to each unique value.
    :align: center
    :alt: Ring of atoms coloured by categorical site labels
 
-Atoms with ``NaN`` (numeric) or ``""`` (categorical) values fall
+Atoms with ``NaN`` (numeric) or ``None`` (categorical) values fall
 back to their species colour.  This is useful when metadata is only
 available for a subset of atoms:
 
 .. code-block:: python
 
    # Only colour specific atoms by charge; the rest keep species colours.
-   scene.set_atom_data("charge", {0: 1.2, 3: -0.8, 5: 0.4})
+   scene.set_atom_data("charge", by_index={0: 1.2, 3: -0.8, 5: 0.4})
+
+Sparse assignment
+-----------------
+
+Use ``by_species`` or ``by_index`` to assign metadata to a subset of
+atoms without building a full-length array:
+
+.. code-block:: python
+
+   # All Mn atoms get charge 2.0.
+   scene.set_atom_data("charge", by_species={"Mn": 2.0})
+
+   # Specific atoms by index.
+   scene.set_atom_data("charge", by_index={0: 1.2, 3: -0.8})
+
+Both forms can be combined in a single call.  ``by_index`` values
+take precedence where they overlap with ``by_species``:
+
+.. code-block:: python
+
+   # All Mn atoms charge 2.0, except atom 3 (defect site) at 1.9.
+   scene.set_atom_data(
+       "charge",
+       by_species={"Mn": 2.0},
+       by_index={3: 1.9},
+   )
+
+For trajectory data, ``by_species`` accepts 2-D arrays of shape
+``(n_frames, n_species_atoms)`` and ``by_index`` accepts 1-D arrays
+of length ``n_frames``.  Either of these promotes the output to 2-D.
+Scalar and 1-D ``by_species`` values and scalar ``by_index`` values
+broadcast across frames automatically.
 
 Custom colouring functions
 --------------------------
@@ -128,9 +160,9 @@ the inner ring uses a numerical charge gradient:
 .. code-block:: python
 
    # Outer atoms: categorical type.
-   scene.set_atom_data("metal", {0: "Fe", 1: "Co", 2: "Ni"})
+   scene.set_atom_data("metal", by_index={0: "Fe", 1: "Co", 2: "Ni"})
    # Inner atoms: numerical charge.
-   scene.set_atom_data("charge", {12: 0.0, 13: 0.3})
+   scene.set_atom_data("charge", by_index={12: 0.0, 13: 0.3})
    scene.render_mpl(
        "output.svg",
        colour_by=["metal", "charge"],
@@ -160,7 +192,7 @@ polyhedra without any additional configuration.
    # No colour on the spec -- polyhedra inherit from colour_by.
    spec = PolyhedronSpec(centre="M", alpha=0.4)
 
-   scene.set_atom_data("val", {0: 0.0, 1: 0.5, 2: 1.0})
+   scene.set_atom_data("val", by_index={0: 0.0, 1: 0.5, 2: 1.0})
    scene.render_mpl(
        "output.svg",
        colour_by="val", cmap="coolwarm",

src/hofmann/model/colour.py

@@ -186,8 +186,8 @@ def _resolve_atom_colours(
     non-empty for categorical) determines the atom's colour.  This
     allows different colouring rules for different atom subsets::
 
-        scene.set_atom_data("metal_type", {0: "Fe", 2: "Co"})
-        scene.set_atom_data("o_coord", {1: 4, 3: 6})
+        scene.set_atom_data("metal_type", by_index={0: "Fe", 2: "Co"})
+        scene.set_atom_data("o_coord", by_index={1: 4, 3: 6})
         scene.render_mpl(
             colour_by=["metal_type", "o_coord"],
             cmap=["Set1", "Blues"],