refactor: revamp AuxData

subpex/configs/subpex/data.py

@@ -1,7 +1,18 @@
-from collections.abc import MutableSequence
+from typing import Any
 
+from collections.abc import Callable, MutableSequence
+
+import MDAnalysis as mda
 from pydantic import BaseModel, Field, computed_field
 
+from ...fop.compute import get_fop_volume_convenience
+from ...pocket.compute import (
+    get_pocket_jaccard_convenience,
+    get_pocket_rmsd_convenience,
+    get_pocket_rog_convenience,
+)
+from ...protein.compute import get_backbone_rmsd
+
 
 class AuxData(BaseModel):
     label: str
@@ -18,6 +29,27 @@ class AuxData(BaseModel):
     """File extension if saving."""
     write: bool = Field(default=False)
     """Flag to write file containing data."""
+    compute_function: Callable[[Any], float | MutableSequence[float]] = Field(
+        default=None, exclude=True
+    )
+    """Function to compute the descriptor value."""
+
+    def compute_frame(  # type: ignore
+        self,
+        atoms_frame: mda.AtomGroup,
+        subpex_config,
+        atoms_ref: mda.AtomGroup | None = None,
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
+        """Will use `compute_function` to compute and then append the value to
+        `values`.
+        """
+        if self.compute_function is not None:
+            result = self.compute_function(
+                atoms_frame, subpex_config, atoms_ref, *args, **kwargs
+            )
+            self.values.append(result)
 
     @computed_field  # type: ignore[misc]
     @property
@@ -29,14 +61,11 @@ class DataConfig(BaseModel):
 
     aux: MutableSequence[AuxData] = Field(
         default_factory=lambda: [
-            AuxData(label="progress_coord"),
-            AuxData(label="pocket_volume"),
-            AuxData(label="pocket_rog"),
-            AuxData(label="pocket_rmsd"),
-            AuxData(label="pocket_fop"),
-            AuxData(label="pocket_jd"),
-            AuxData(label="backbone_rmsd"),
-            AuxData(label="composite"),
+            AuxData(label="pocket_volume", compute_function=get_fop_volume_convenience),
+            AuxData(label="pocket_rog", compute_function=get_pocket_rog_convenience),
+            AuxData(label="pocket_rmsd", compute_function=get_pocket_rmsd_convenience),
+            AuxData(label="pocket_jd", compute_function=get_pocket_jaccard_convenience),
+            AuxData(label="backbone_rmsd", compute_function=get_backbone_rmsd),
         ]
     )