add wrapper to use mol prop for augmented graph

Author: aditya0by0

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -1,11 +1,12 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 from typing import Dict, List, Optional
 
 from rdkit import Chem
 
 from chebai_graph.preprocessing.property_encoder import OneHotEncoder, PropertyEncoder
 
 from . import properties as pr
+from .base import FrozenPropertyAlias
 from .constants import *
 
 
@@ -103,7 +104,8 @@ def _check_modify_atom_prop_value(self, atom: Chem.rdchem.Atom | Dict, prop: str
             return 0
 
 
-class AugNodeValueDefaulter(AugmentedAtomProperty, ABC):
+class AugNodeValueDefaulter(AugmentedAtomProperty, FrozenPropertyAlias, ABC):
+
     def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
         if isinstance(atom, Chem.rdchem.Atom):
             # Delegate to superclass method for atom
@@ -245,7 +247,7 @@ def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
         return self._check_modify_bond_prop_value(bond, EDGE_LEVEL)
 
 
-class AugBondValueDefaulter(AugmentedBondProperty, ABC):
+class AugBondValueDefaulter(AugmentedBondProperty, FrozenPropertyAlias, ABC):
     def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
         if isinstance(bond, Chem.rdchem.Bond):
             # Delegate to superclass method for bond

chebai_graph/preprocessing/properties/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from types import MappingProxyType
 
 import rdkit.Chem as Chem
 
@@ -49,3 +50,59 @@ def get_bond_value(self, bond: Chem.rdchem.Bond):
 
 class MoleculeProperty(MolecularProperty):
     """Global property of a molecule."""
+
+
+class FrozenPropertyAlias(MolecularProperty, ABC):
+    """
+    Wrapper base class for augmented graph properties that want to reuse existing molecular properties.
+
+    This class allows augmented graph property classes to inherit both from this wrapper and a standard
+    molecular property (from `.properties`), enabling reuse of their encoders and index files without
+    modifying them.
+
+    Key Features:
+    - Prevents new tokens from being added to the encoder cache by freezing it.
+    - Automatically aligns the property name (used for encoder/index resolution) with the inherited
+      base property by removing the "Aug" prefix from the class name.
+
+    Usage:
+        The derived class should:
+        - Inherit from `FrozenPropertyAlias` **and** a valid base molecular property class.
+        - Have a name starting with "Aug" (e.g., `AugAtomType`), which will be resolved to `AtomType`.
+
+    Example:
+        ```python
+        class AugAtomType(FrozenPropertyAlias, AtomType):
+            ...
+        ```
+    Note:
+        Subclass name of this class should with prefix "Aug" for above effect to take place.
+
+    This allows `AugAtomType` to reuse the encoder, index files, and logic of `AtomType` while
+    integrating into augmented graph pipelines.
+    """
+
+    def __init__(self, encoder: PropertyEncoder | None = None):
+        super().__init__(encoder)
+        # Lock the encoder's cache to prevent adding new tokens
+        if hasattr(self.encoder, "cache") and isinstance(self.encoder.cache, dict):
+            self.encoder.cache = MappingProxyType(self.encoder.cache)
+
+    @property
+    def name(self):
+        """
+        Unique identifier for this property, with 'Aug' prefix removed if present.
+        This allows the encoder to reuse index files of the corresponding base property.
+        """
+        class_name = self.__class__.__name__
+        return class_name[3:] if class_name.startswith("Aug") else class_name
+
+    def on_finish(self):
+        if (
+            hasattr(self.encoder, "cache")
+            and len(self.encoder.cache) > self.encoder.index_length_start
+        ):
+            raise ValueError(
+                f"{self.__class__.__name__} attempted to add new tokens to a {self.encoder.index_path}"
+            )
+        super().on_finish()