features

features#

Feature extraction module for PrxteinMPNN.

This module contains the ProteinFeatures class that extracts and projects features from raw protein coordinates.

class prxteinmpnn.model.features.ProteinFeatures(node_features, edge_features, k_neighbors, *, key)[source]#

Bases: Module

Extracts and projects features from raw protein coordinates.

This module encapsulates k-NN, RBF, positional encodings, and edge projections. Note: W_e projection is NOT here - it’s in the main model (matches ColabDesign).

Parameters:

  • node_features (int)

  • edge_features (int)

  • k_neighbors (int)

  • key (PRNGKeyArray)

k_neighbors: int#
rbf_dim: int#
pos_embed_dim: int#
w_pos: Linear#
w_e: Linear#
norm_edges: LayerNorm#
w_e_proj: Linear#
Contents