diffbio 0.1.0__py3-none-any.whl

diffbio/operators/drug_discovery/maccs_keys.py

@@ -0,0 +1,267 @@
+"""Differentiable MACCS structural keys fingerprint operator.
+
+This module implements a differentiable version of the 166 MACCS
+(Molecular ACCess System) structural keys fingerprint.
+
+MACCS keys are predefined structural patterns (SMARTS) that encode
+the presence/absence of specific molecular substructures. This
+implementation provides a differentiable approximation using
+learned pattern matching networks.
+
+References:
+    - https://rdkit.org/docs/source/rdkit.Chem.MACCSkeys.html
+    - Durant et al. "Reoptimization of MDL Keys for Use in Drug Discovery" JCIM 2002
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+from diffbio.core import soft_ops
+from diffbio.operators.drug_discovery._graph_utils import (
+    attach_fingerprint,
+    ensure_rngs,
+    initialize_graph_encoder,
+    stabilize_operator_id,
+    unpack_graph_inputs,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MACCSKeysConfig(OperatorConfig):
+    """Configuration for MACCS keys fingerprint operator.
+
+    Attributes:
+        n_bits: Number of fingerprint bits (default: 166 for standard MACCS).
+        differentiable: Use learned pattern matching (default: True).
+        temperature: Temperature for soft bit assignment (default: 1.0).
+        hidden_dim: Hidden dimension for pattern networks (default: 64).
+        num_layers: Number of message passing layers (default: 2).
+        in_features: Number of input node features (default: 4).
+    """
+
+    n_bits: int = 166
+    differentiable: bool = True
+    temperature: float = 1.0
+    hidden_dim: int = 64
+    num_layers: int = 2
+    in_features: int = 4
+
+
+class MACCSKeysOperator(OperatorModule):
+    """Differentiable MACCS structural keys fingerprint operator.
+
+    For differentiable=True:
+        Uses message passing and learned pattern detectors to approximate
+        MACCS key detection. Each of the 166 keys is represented by a
+        learned pattern matching network that outputs a soft presence score.
+
+    For differentiable=False:
+        Would use RDKit's exact MACCS implementation (not differentiable).
+
+    The differentiable version enables gradient flow for end-to-end
+    optimization while approximating the structural pattern detection
+    of traditional MACCS keys.
+
+    MACCS keys encode various structural features:
+
+        - Atom types (C, N, O, S, halides, etc.)
+        - Functional groups (carbonyl, hydroxyl, amine, etc.)
+        - Ring systems (aromatic, aliphatic)
+        - Bond patterns and connectivity
+
+    Example:
+        ```python
+        config = MACCSKeysConfig(temperature=1.0)
+        op = MACCSKeysOperator(config, rngs=nnx.Rngs(42))
+        data = {"node_features": nodes, "adjacency": adj}
+        result, _, _ = op.apply(data, {}, None)
+        fingerprint = result["fingerprint"]  # shape: (166,)
+        ```
+
+    References:
+        - Durant et al. "Reoptimization of MDL Keys" JCIM 2002
+    """
+
+    def __init__(
+        self,
+        config: MACCSKeysConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize MACCS keys operator.
+
+        Args:
+            config: MACCS keys configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+
+        rngs = ensure_rngs(rngs)
+
+        if config.differentiable:
+            # Message passing for local structure aggregation
+            rngs = initialize_graph_encoder(
+                self,
+                rngs=rngs,
+                hidden_dim=config.hidden_dim,
+                num_layers=config.num_layers,
+                in_features=config.in_features,
+                attr="encoder",
+            )
+
+            # Pattern detectors: one network per MACCS key
+            # Each outputs a score indicating pattern presence
+            self.pattern_detectors = nnx.Sequential(
+                nnx.Linear(config.hidden_dim, config.hidden_dim, rngs=rngs),
+                nnx.relu,
+                nnx.Linear(config.hidden_dim, config.n_bits, rngs=rngs),
+            )
+        else:
+            # RDKit mode
+            stabilize_operator_id(self)
+            try:
+                from rdkit import Chem
+                from rdkit.Chem import MACCSkeys as RDKitMACCS
+
+                self._Chem = Chem
+                self._MACCSkeys = RDKitMACCS
+            except ImportError as e:
+                raise ImportError(
+                    "MACCSKeysOperator with differentiable=False requires RDKit: pip install rdkit"
+                ) from e
+
+    def _compute_differentiable_fp(
+        self,
+        node_features: jnp.ndarray,
+        adjacency: jnp.ndarray,
+        edge_features: jnp.ndarray | None = None,
+        node_mask: jnp.ndarray | None = None,
+    ) -> jnp.ndarray:
+        """Compute differentiable MACCS-style fingerprint.
+
+        Uses message passing to aggregate local structure information,
+        then applies learned pattern detectors with soft thresholding.
+
+        Args:
+            node_features: (num_nodes, num_features) atom features
+            adjacency: (num_nodes, num_nodes) adjacency matrix
+            edge_features: Optional edge features
+            node_mask: Optional mask for valid nodes
+
+        Returns:
+            Fingerprint vector of shape (n_bits,)
+        """
+        # Message passing to capture local structure
+        node_hidden = self.encoder(node_features, adjacency, edge_features)
+
+        # Apply node mask
+        if node_mask is not None:
+            node_hidden = node_hidden * node_mask[:, None]
+
+        # Pattern detection at each atom
+        # pattern_logits: (num_nodes, n_bits)
+        pattern_logits = self.pattern_detectors(node_hidden)
+
+        # Apply temperature-scaled sigmoid for soft pattern matching
+        soft_patterns = soft_ops.greater(pattern_logits, 0.0, softness=self.config.temperature)
+
+        # Aggregate across atoms using max (OR-like)
+        # If any atom matches a pattern, the bit is set
+        fingerprint = soft_ops.max(soft_patterns, axis=0, softness=self.config.temperature)
+
+        return fingerprint
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,  # noqa: ARG002
+        stats: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Compute MACCS keys fingerprint.
+
+        Args:
+            data: Input data containing:
+                For differentiable=True:
+                    - node_features: (num_nodes, num_features) atom features
+                    - adjacency: (num_nodes, num_nodes) adjacency matrix
+                    - node_mask: (num_nodes,) optional mask for valid nodes
+                For differentiable=False:
+                    - smiles: SMILES string
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with added "fingerprint" key
+                - unchanged state
+                - unchanged metadata
+        """
+        del random_params, stats  # Unused
+
+        if self.config.differentiable:
+            node_features, adjacency, edge_features, node_mask = unpack_graph_inputs(data)
+            fp = self._compute_differentiable_fp(node_features, adjacency, edge_features, node_mask)
+        else:
+            smiles = data["smiles"]
+            fp = self._compute_rdkit_fp(smiles)
+
+        return attach_fingerprint(data, fp), state, metadata
+
+    def _compute_rdkit_fp(self, smiles: str) -> jnp.ndarray:
+        """Compute exact MACCS keys using RDKit.
+
+        Args:
+            smiles: SMILES string
+
+        Returns:
+            Binary fingerprint vector of shape (166,)
+        """
+        import numpy as np
+
+        mol = self._Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return jnp.zeros(self.config.n_bits, dtype=jnp.float32)
+
+        fp = self._MACCSkeys.GenMACCSKeys(mol)
+
+        # Convert to numpy (RDKit returns 167 bits, we use 1-166)
+        arr = np.zeros(167, dtype=np.float32)
+        for i in range(167):
+            arr[i] = fp.GetBit(i)
+
+        # Return bits 1-166 (index 0 is unused in standard MACCS)
+        return jnp.asarray(arr[1:167])
+
+
+def create_maccs_operator(
+    differentiable: bool = True,
+    temperature: float = 1.0,
+    seed: int = 42,
+) -> MACCSKeysOperator:
+    """Create a MACCS keys fingerprint operator.
+
+    Args:
+        differentiable: Use learned pattern matching.
+        temperature: Temperature for soft matching.
+        seed: Random seed.
+
+    Returns:
+        Configured MACCSKeysOperator.
+    """
+    config = MACCSKeysConfig(
+        differentiable=differentiable,
+        temperature=temperature,
+    )
+    return MACCSKeysOperator(config, rngs=nnx.Rngs(seed))

diffbio/operators/drug_discovery/message_passing.py

@@ -0,0 +1,200 @@
+"""Message passing neural network layers for molecular graphs.
+
+This module implements directed message passing neural network (D-MPNN)
+layers following the ChemProp architecture for molecular property prediction.
+"""
+
+import logging
+
+import jax.numpy as jnp
+from flax import nnx
+
+logger = logging.getLogger(__name__)
+
+
+class MessagePassingLayer(nnx.Module):
+    """Directed message passing layer for molecular graphs.
+
+    Implements the D-MPNN message passing scheme where messages are passed
+    along directed edges. Each node aggregates messages from its neighbors
+    and updates its representation.
+
+    Attributes:
+        hidden_dim: Dimension of hidden node representations.
+        in_features: Number of input node features.
+        num_edge_features: Number of edge features (default 4 for bond types).
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        in_features: int = 4,
+        num_edge_features: int = 4,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize message passing layer.
+
+        Args:
+            hidden_dim: Dimension of hidden representations.
+            in_features: Number of input node features (default 4 for tests).
+            num_edge_features: Number of edge/bond features.
+            rngs: Flax NNX random number generators.
+        """
+        self.hidden_dim = hidden_dim
+        self.in_features = in_features
+        self.num_edge_features = num_edge_features
+
+        # Node encoder - eagerly initialized with specified in_features
+        self.node_encoder = nnx.Linear(
+            in_features=in_features,
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+
+        # Edge encoder
+        self.edge_encoder = nnx.Linear(
+            in_features=num_edge_features,
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+
+        # Message transformation
+        self.message_layer = nnx.Linear(
+            in_features=hidden_dim * 3,  # src_node + edge + dst_node
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+
+        # Update function (GRU-like update)
+        self.update_layer = nnx.Linear(
+            in_features=hidden_dim * 2,  # current + aggregated
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        node_features: jnp.ndarray,
+        adjacency: jnp.ndarray,
+        edge_features: jnp.ndarray | None = None,
+    ) -> jnp.ndarray:
+        """Perform one step of message passing.
+
+        Args:
+            node_features: Node features of shape (num_nodes, in_features).
+            adjacency: Adjacency matrix of shape (num_nodes, num_nodes).
+            edge_features: Optional edge features of shape
+                (num_nodes, num_nodes, num_edge_features).
+
+        Returns:
+            Updated node features of shape (num_nodes, hidden_dim).
+        """
+        num_nodes = node_features.shape[0]
+
+        # Encode node features to hidden dimension
+        node_hidden = nnx.relu(self.node_encoder(node_features))
+
+        # Handle edge features
+        if edge_features is not None:
+            edge_hidden = nnx.relu(self.edge_encoder(edge_features))
+        else:
+            # Use zeros if no edge features provided
+            edge_hidden = jnp.zeros((num_nodes, num_nodes, self.hidden_dim), dtype=jnp.float32)
+
+        # Compute messages for all pairs
+        # For each edge (i, j), message = f(node_i, edge_ij, node_j)
+        # Expand dimensions for broadcasting
+        src_nodes = node_hidden[:, None, :]  # (N, 1, H)
+        dst_nodes = node_hidden[None, :, :]  # (1, N, H)
+
+        # Broadcast to (N, N, H)
+        src_expanded = jnp.broadcast_to(src_nodes, (num_nodes, num_nodes, self.hidden_dim))
+        dst_expanded = jnp.broadcast_to(dst_nodes, (num_nodes, num_nodes, self.hidden_dim))
+
+        # Concatenate [src, edge, dst]
+        message_input = jnp.concatenate([src_expanded, edge_hidden, dst_expanded], axis=-1)
+
+        # Compute messages
+        messages = nnx.relu(self.message_layer(message_input))
+
+        # Mask messages by adjacency (only neighbors contribute)
+        masked_messages = messages * adjacency[:, :, None]
+
+        # Aggregate messages (sum over neighbors)
+        aggregated = jnp.sum(masked_messages, axis=1)  # (N, H)
+
+        # Update node representations
+        update_input = jnp.concatenate([node_hidden, aggregated], axis=-1)
+        updated = nnx.relu(self.update_layer(update_input))
+
+        return updated
+
+
+class StackedMessagePassing(nnx.Module):
+    """Stack of message passing layers.
+
+    Applies multiple rounds of message passing to capture higher-order
+    neighborhood information.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_layers: int,
+        in_features: int = 4,
+        num_edge_features: int = 4,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize stacked message passing.
+
+        Args:
+            hidden_dim: Hidden dimension for all layers.
+            num_layers: Number of message passing iterations.
+            in_features: Number of input node features (default 4 for tests).
+            num_edge_features: Number of edge features.
+            rngs: Flax NNX random number generators.
+        """
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.in_features = in_features
+
+        # Build layers with proper input dimensions:
+        # - First layer: in_features -> hidden_dim
+        # - Subsequent layers: hidden_dim -> hidden_dim
+        layers = []
+        for i in range(num_layers):
+            layer_in_features = in_features if i == 0 else hidden_dim
+            layers.append(
+                MessagePassingLayer(
+                    hidden_dim=hidden_dim,
+                    in_features=layer_in_features,
+                    num_edge_features=num_edge_features,
+                    rngs=rngs,
+                )
+            )
+        self.layers = nnx.List(layers)
+
+    def __call__(
+        self,
+        node_features: jnp.ndarray,
+        adjacency: jnp.ndarray,
+        edge_features: jnp.ndarray | None = None,
+    ) -> jnp.ndarray:
+        """Apply multiple rounds of message passing.
+
+        Args:
+            node_features: Initial node features.
+            adjacency: Adjacency matrix.
+            edge_features: Optional edge features.
+
+        Returns:
+            Final node representations.
+        """
+        h = node_features
+
+        for layer in self.layers:
+            h = layer(h, adjacency, edge_features)
+
+        return h

diffbio/operators/drug_discovery/primitives.py

@@ -0,0 +1,242 @@
+"""Primitive functions for molecular graph processing.
+
+This module provides utility functions for converting SMILES strings to
+molecular graphs suitable for differentiable neural network processing.
+RDKit is used for parsing only; all graph operations use JAX arrays.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from rdkit import Chem
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AtomFeatureConfig:
+    """Configuration for atom feature extraction.
+
+    The default configuration produces 34 features:
+
+    - Atom type: 12 dimensions (C, N, O, S, F, Cl, Br, I, P, Si, B, Other)
+    - Degree: 7 dimensions (0-6)
+    - Formal charge: 5 dimensions (-2 to +2)
+    - Hybridization: 4 dimensions (SP, SP2, SP3, SP3D)
+    - Aromaticity: 1 dimension (binary)
+    - Num hydrogens: 5 dimensions (0-4)
+    """
+
+    num_atom_types: int = 12
+    max_degree: int = 6  # Creates max_degree + 1 dimensions
+    charge_range: tuple[int, int] = (-2, 2)  # Creates 5 dimensions
+    num_hybridization_types: int = 4
+    max_num_hydrogens: int = 4  # Creates max_num_hydrogens + 1 dimensions
+
+    @property
+    def total_features(self) -> int:
+        """Calculate total number of atom features."""
+        return (
+            self.num_atom_types
+            + (self.max_degree + 1)
+            + (self.charge_range[1] - self.charge_range[0] + 1)
+            + self.num_hybridization_types
+            + 1  # aromaticity
+            + (self.max_num_hydrogens + 1)
+        )
+
+
+# Default configuration
+DEFAULT_ATOM_CONFIG = AtomFeatureConfig()
+DEFAULT_ATOM_FEATURES = DEFAULT_ATOM_CONFIG.total_features  # 34 features
+
+# Atom type vocabulary for one-hot encoding
+ATOM_TYPES = ["C", "N", "O", "S", "F", "Cl", "Br", "I", "P", "Si", "B", "Other"]
+
+
+def get_atom_features(atom: Any, config: AtomFeatureConfig | None = None) -> jnp.ndarray:
+    """Extract features from an RDKit atom object.
+
+    Features include:
+    - Atom type (one-hot, config.num_atom_types)
+    - Degree (one-hot, 0 to config.max_degree)
+    - Formal charge (one-hot, config.charge_range)
+    - Hybridization (one-hot, config.num_hybridization_types)
+    - Aromaticity (binary)
+    - Number of hydrogens (one-hot, 0 to config.max_num_hydrogens)
+
+    Args:
+        atom: RDKit atom object.
+        config: Feature extraction configuration. Defaults to DEFAULT_ATOM_CONFIG.
+
+    Returns:
+        Feature vector of shape (config.total_features,).
+    """
+    if config is None:
+        config = DEFAULT_ATOM_CONFIG
+
+    features: list[float] = []
+
+    # Atom type
+    symbol = atom.GetSymbol()
+    atom_type_idx = ATOM_TYPES.index(symbol) if symbol in ATOM_TYPES else len(ATOM_TYPES) - 1
+    atom_type_onehot = [0.0] * config.num_atom_types
+    atom_type_onehot[min(atom_type_idx, config.num_atom_types - 1)] = 1.0
+    features.extend(atom_type_onehot)
+
+    # Degree
+    degree = min(atom.GetDegree(), config.max_degree)
+    degree_onehot = [0.0] * (config.max_degree + 1)
+    degree_onehot[degree] = 1.0
+    features.extend(degree_onehot)
+
+    # Formal charge
+    charge = atom.GetFormalCharge()
+    charge_min, charge_max = config.charge_range
+    charge_idx = max(charge_min, min(charge_max, charge)) - charge_min
+    charge_dim = charge_max - charge_min + 1
+    charge_onehot = [0.0] * charge_dim
+    charge_onehot[charge_idx] = 1.0
+    features.extend(charge_onehot)
+
+    # Hybridization
+    hybridization = atom.GetHybridization()
+    hyb_types = [
+        Chem.rdchem.HybridizationType.SP,
+        Chem.rdchem.HybridizationType.SP2,
+        Chem.rdchem.HybridizationType.SP3,
+        Chem.rdchem.HybridizationType.SP3D,
+    ]
+    hyb_onehot = [0.0] * config.num_hybridization_types
+    for i, h in enumerate(hyb_types[: config.num_hybridization_types]):
+        if hybridization == h:
+            hyb_onehot[i] = 1.0
+            break
+    features.extend(hyb_onehot)
+
+    # Aromaticity
+    features.append(1.0 if atom.GetIsAromatic() else 0.0)
+
+    # Number of hydrogens
+    num_h = min(atom.GetTotalNumHs(), config.max_num_hydrogens)
+    h_onehot = [0.0] * (config.max_num_hydrogens + 1)
+    h_onehot[num_h] = 1.0
+    features.extend(h_onehot)
+
+    return jnp.array(features, dtype=jnp.float32)
+
+
+def get_bond_features(bond: Any) -> jnp.ndarray:
+    """Extract features from an RDKit bond object.
+
+    Features include:
+    - Bond type (one-hot: single, double, triple, aromatic)
+
+    Args:
+        bond: RDKit bond object.
+
+    Returns:
+        Feature vector of shape (4,).
+    """
+    bond_type = bond.GetBondType()
+    features = [
+        1 if bond_type == Chem.rdchem.BondType.SINGLE else 0,
+        1 if bond_type == Chem.rdchem.BondType.DOUBLE else 0,
+        1 if bond_type == Chem.rdchem.BondType.TRIPLE else 0,
+        1 if bond_type == Chem.rdchem.BondType.AROMATIC else 0,
+    ]
+    return jnp.array(features, dtype=jnp.float32)
+
+
+def smiles_to_graph(smiles: str) -> dict[str, Any]:
+    """Convert a SMILES string to a molecular graph.
+
+    Args:
+        smiles: SMILES string representing a molecule.
+
+    Returns:
+        Dictionary containing:
+            - node_features: (num_atoms, num_features) atom feature matrix
+            - adjacency: (num_atoms, num_atoms) adjacency matrix
+            - edge_features: (num_atoms, num_atoms, num_edge_features) bond features
+            - num_nodes: number of atoms
+
+    Raises:
+        ValueError: If SMILES string is invalid.
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES string: {smiles}")
+
+    num_atoms = mol.GetNumAtoms()
+
+    # Extract node features
+    node_features = []
+    for atom in mol.GetAtoms():
+        node_features.append(get_atom_features(atom))
+    node_features = jnp.stack(node_features)
+
+    # Build adjacency matrix and edge features
+    adjacency = jnp.zeros((num_atoms, num_atoms), dtype=jnp.float32)
+    edge_features = jnp.zeros((num_atoms, num_atoms, 4), dtype=jnp.float32)
+
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        bond_feat = get_bond_features(bond)
+
+        # Symmetric (undirected graph)
+        adjacency = adjacency.at[i, j].set(1.0)
+        adjacency = adjacency.at[j, i].set(1.0)
+        edge_features = edge_features.at[i, j].set(bond_feat)
+        edge_features = edge_features.at[j, i].set(bond_feat)
+
+    return {
+        "node_features": node_features,
+        "adjacency": adjacency,
+        "edge_features": edge_features,
+        "num_nodes": num_atoms,
+    }
+
+
+def batch_smiles_to_graphs(smiles_list: list[str]) -> dict[str, Any]:
+    """Convert a batch of SMILES strings to padded graph tensors.
+
+    Args:
+        smiles_list: List of SMILES strings.
+
+    Returns:
+        Dictionary containing:
+            - node_features: (batch_size, max_nodes, num_features)
+            - adjacency: (batch_size, max_nodes, max_nodes)
+            - edge_features: (batch_size, max_nodes, max_nodes, num_edge_features)
+            - node_mask: (batch_size, max_nodes) mask for valid nodes
+    """
+    graphs = [smiles_to_graph(s) for s in smiles_list]
+
+    max_nodes = max(g["num_nodes"] for g in graphs)
+    batch_size = len(graphs)
+    num_features = graphs[0]["node_features"].shape[1]
+    num_edge_features = graphs[0]["edge_features"].shape[2]
+
+    # Initialize padded tensors
+    node_features = jnp.zeros((batch_size, max_nodes, num_features))
+    adjacency = jnp.zeros((batch_size, max_nodes, max_nodes))
+    edge_features = jnp.zeros((batch_size, max_nodes, max_nodes, num_edge_features))
+    node_mask = jnp.zeros((batch_size, max_nodes))
+
+    for i, g in enumerate(graphs):
+        n = g["num_nodes"]
+        node_features = node_features.at[i, :n, :].set(g["node_features"])
+        adjacency = adjacency.at[i, :n, :n].set(g["adjacency"])
+        edge_features = edge_features.at[i, :n, :n, :].set(g["edge_features"])
+        node_mask = node_mask.at[i, :n].set(1.0)
+
+    return {
+        "node_features": node_features,
+        "adjacency": adjacency,
+        "edge_features": edge_features,
+        "node_mask": node_mask,
+    }