feat(drug-discovery)!: featurize molecules as AttentiveFP does

CHANGELOG.md

@@ -18,6 +18,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Molecular graphs carry the AttentiveFP featurization (Xiong et al. 2019) as DGL-LifeSci's
+  `AttentiveFPAtomFeaturizer` and `AttentiveFPBondFeaturizer` define it: 39 atom features and 10
+  bond features, against 34 and 4 before. The new features are chirality and the CIP label,
+  radical electrons, the elements As, Se, Te and At, SP3D2 hybridization, bond conjugation, ring
+  membership and bond stereo. Formal charge is a signed count rather than a one-hot clipped to
+  -2..2, and a degree above 5 or more than 4 hydrogens leaves its block all zero, as in the
+  reference. Stereoisomers now get distinct graphs: on BBBP as MolNet supplies it, 12 molecules
+  in 6 groups shared one graph, 2 of the groups with conflicting labels, and none do now.
+  `DEFAULT_ATOM_FEATURES` and `DEFAULT_BOND_FEATURES` follow, so every operator built with its
+  defaults takes the new graphs. `AtomFeatureConfig`, `DEFAULT_ATOM_CONFIG` and `ATOM_TYPES` are
+  removed, and `get_atom_features` takes no configuration.
 - `smiles_to_graph` and `batch_smiles_to_graphs` assemble each graph on the host with NumPy and
   convert it to JAX once per array. They built it with eager JAX operations, one array per atom
   and four scatter updates per bond, each copying the whole adjacency or edge-feature array. The

docs/api/operators/drug-discovery.md

@@ -233,12 +233,17 @@ Differentiable operators for molecular property prediction, fingerprint computat
     options:
       show_root_heading: true
 
-### AtomFeatureConfig
+### get_atom_features
 
-::: diffbio.operators.drug_discovery.primitives.AtomFeatureConfig
+::: diffbio.operators.drug_discovery.primitives.get_atom_features
+    options:
+      show_root_heading: true
+
+### get_bond_features
+
+::: diffbio.operators.drug_discovery.primitives.get_bond_features
     options:
       show_root_heading: true
-      members: []
 
 ## Constants
 
@@ -248,9 +253,9 @@ Differentiable operators for molecular property prediction, fingerprint computat
     options:
       show_root_heading: true
 
-### DEFAULT_ATOM_CONFIG
+### DEFAULT_BOND_FEATURES
 
-::: diffbio.operators.drug_discovery.primitives.DEFAULT_ATOM_CONFIG
+::: diffbio.operators.drug_discovery.primitives.DEFAULT_BOND_FEATURES
     options:
       show_root_heading: true
 
@@ -403,9 +408,9 @@ node_features, adjacency, edge_features, masks = batch_smiles_to_graphs(
     smiles_list,
     max_atoms=20,
 )
-# node_features: (3, 20, 34)
+# node_features: (3, 20, 39)
 # adjacency: (3, 20, 20)
-# edge_features: (3, 20, 20, 4)
+# edge_features: (3, 20, 20, 10)
 # masks: (3, 20)
 ```
 
@@ -513,14 +518,30 @@ grads = nnx.grad(loss_fn)(predictor, data)
 
 ## Atom Feature Dimensions
 
+The AttentiveFP featurization (Xiong et al. 2019), as DGL-LifeSci's `AttentiveFPAtomFeaturizer`
+defines it.
+
 | Feature | Dimensions | Description |
 |---------|------------|-------------|
-| Atom type | 10 | C, N, O, S, F, Cl, Br, I, P, other |
-| Degree | 6 | 0-5+ neighbors |
-| Formal charge | 5 | -2 to +2 |
-| Hybridization | 5 | SP, SP2, SP3, SP3D, SP3D2 |
+| Element | 16 | B, C, N, O, F, Si, P, S, Cl, As, Se, Br, Te, I, At, unknown |
+| Degree | 6 | 0-5 neighbors (all zero above 5) |
+| Formal charge | 1 | Signed count |
+| Radical electrons | 1 | Count |
+| Hybridization | 6 | SP, SP2, SP3, SP3D, SP3D2, unknown |
 | Aromaticity | 1 | Binary |
-| Hydrogens | 5 | 0-4+ |
+| Hydrogens | 5 | 0-4 total hydrogens (all zero above 4) |
+| Chiral centre | 1 | Binary |
+| CIP label | 2 | R, S (all zero when unassigned) |
+| **Total** | **39** | `DEFAULT_ATOM_FEATURES` |
+
+## Bond Feature Dimensions
+
+DGL-LifeSci's `AttentiveFPBondFeaturizer`.
+
+| Feature | Dimensions | Description |
+|---------|------------|-------------|
+| Bond type | 4 | Single, double, triple, aromatic |
+| Conjugated | 1 | Binary |
 | In ring | 1 | Binary |
-| Chiral | 1 | Binary |
-| **Total** | **34** | `DEFAULT_ATOM_FEATURES` |
+| Stereo | 4 | None, any, Z, E |
+| **Total** | **10** | `DEFAULT_BOND_FEATURES` |

docs/examples/basic/molecular-fingerprints.md

@@ -52,9 +52,9 @@ SMILES: CC(=O)OC1=CC=CC=C1C(=O)O
 
 Graph representation:
   Number of atoms: 13
-  Node features shape: (13, 34)
+  Node features shape: (13, 39)
   Adjacency matrix shape: (13, 13)
-  Edge features shape: (13, 13, 4)
+  Edge features shape: (13, 13, 10)
   Number of bonds: 13
 ```
 

docs/user-guide/operators/drug-discovery.md

@@ -54,9 +54,9 @@ batch remains valid under `validate_dti_dataset()`.
 ```mermaid
 graph TB
     A["SMILES String"] --> B["smiles_to_graph()<br/>RDKit: atom/bond features"]
-    B --> NF["node_features<br/>(n_atoms, 34)"]
+    B --> NF["node_features<br/>(n_atoms, 39)"]
     B --> ADJ["adjacency<br/>(n_atoms, n_atoms)"]
-    B --> EF["edge_features<br/>(n_atoms, n_atoms, 4)"]
+    B --> EF["edge_features<br/>(n_atoms, n_atoms, 10)"]
     NF --> C["StackedMessagePassing<br/>D-MPNN: message → update"]
     ADJ --> C
     EF --> C
@@ -103,7 +103,7 @@ smiles = "CCO"  # Ethanol
 node_features, adjacency, edge_features = smiles_to_graph(smiles)
 
 print(f"Atoms: {node_features.shape[0]}")      # 3 atoms
-print(f"Features per atom: {node_features.shape[1]}")  # 34 features
+print(f"Features per atom: {node_features.shape[1]}")  # 39 features
 print(f"Adjacency: {adjacency.shape}")         # (3, 3)
 ```
 
@@ -117,24 +117,34 @@ node_features, adjacency, edge_features, masks = batch_smiles_to_graphs(
     max_atoms=20,
 )
 
-print(f"Batch shape: {node_features.shape}")  # (3, 20, 34)
+print(f"Batch shape: {node_features.shape}")  # (3, 20, 39)
 print(f"Masks shape: {masks.shape}")          # (3, 20)
 ```
 
 ### Atom Features
 
-The default atom featurizer extracts 34 features per atom:
-
-| Feature | Dimensions | Description |
-|---------|------------|-------------|
-| Atom type | 10 | One-hot: C, N, O, S, F, Cl, Br, I, P, other |
-| Degree | 6 | One-hot: 0-5+ |
-| Formal charge | 5 | One-hot: -2 to +2 |
-| Hybridization | 5 | One-hot: SP, SP2, SP3, SP3D, SP3D2 |
+Molecular graphs carry the AttentiveFP featurization (Xiong et al. 2019), as DGL-LifeSci's
+`AttentiveFPAtomFeaturizer` and `AttentiveFPBondFeaturizer` define it: 39 features per atom and 10
+per bond. Chirality, CIP labels and bond stereo are included, so stereoisomers get distinct graphs.
+
+| Atom feature | Dimensions | Description |
+|--------------|------------|-------------|
+| Element | 16 | One-hot: B, C, N, O, F, Si, P, S, Cl, As, Se, Br, Te, I, At, unknown |
+| Degree | 6 | One-hot: 0-5 (all zero above 5) |
+| Formal charge | 1 | Signed count |
+| Radical electrons | 1 | Count |
+| Hybridization | 6 | One-hot: SP, SP2, SP3, SP3D, SP3D2, unknown |
 | Aromaticity | 1 | Binary |
-| Hydrogens | 5 | One-hot: 0-4+ |
+| Hydrogens | 5 | One-hot: 0-4 total hydrogens (all zero above 4) |
+| Chiral centre | 1 | Binary |
+| CIP label | 2 | One-hot: R, S (all zero when unassigned) |
+
+| Bond feature | Dimensions | Description |
+|--------------|------------|-------------|
+| Bond type | 4 | One-hot: single, double, triple, aromatic |
+| Conjugated | 1 | Binary |
 | In ring | 1 | Binary |
-| Chiral | 1 | Binary |
+| Stereo | 4 | One-hot: none, any, Z, E |
 
 ## MolecularPropertyPredictor
 
@@ -153,7 +163,7 @@ from diffbio.operators.drug_discovery import (
     DEFAULT_ATOM_FEATURES,
 )
 
-# Create predictor for real molecules (34 atom features)
+# Create predictor for real molecules (39 atom features)
 predictor = create_property_predictor(
     hidden_dim=64,
     num_layers=3,
@@ -929,15 +939,14 @@ grads = nnx.grad(end_to_end_loss)(predictor, fingerprint_op, sim_op, data1, data
 | Parameter | Type | Description |
 |-----------|------|-------------|
 | `smiles` | str | Valid SMILES string |
-| `config` | AtomFeatureConfig | Optional feature configuration |
 
 **Output**
 
 | Return | Shape | Description |
 |--------|-------|-------------|
-| `node_features` | (n_atoms, 34) | Atom feature vectors |
+| `node_features` | (n_atoms, 39) | Atom feature vectors |
 | `adjacency` | (n_atoms, n_atoms) | Binary adjacency matrix |
-| `edge_features` | (n_atoms, n_atoms, 4) | Bond type one-hot |
+| `edge_features` | (n_atoms, n_atoms, 10) | Bond feature vectors |
 
 ### MolecularPropertyPredictor
 

src/diffbio/operators/drug_discovery/__init__.py

@@ -58,8 +58,6 @@
     StackedMessagePassing,
 )
 from diffbio.operators.drug_discovery.primitives import (
-    AtomFeatureConfig,
-    DEFAULT_ATOM_CONFIG,
     DEFAULT_ATOM_FEATURES,
     DEFAULT_BOND_FEATURES,
     batch_smiles_to_graphs,
@@ -85,8 +83,6 @@
     "batch_smiles_to_graphs",
     "DEFAULT_ATOM_FEATURES",
     "DEFAULT_BOND_FEATURES",
-    "DEFAULT_ATOM_CONFIG",
-    "AtomFeatureConfig",
     # Message Passing
     "MessagePassingLayer",
     "StackedMessagePassing",