features

features#

Feature extraction for protein structures in the PrxteinMPNN model.

prxteinmpnn.model.features.get_edge_chains_neighbors(chain_index, neighbor_indices)[source]#

Compute edge chains for neighbors.

Return type:

Int[Array, 'num_atoms num_neighbors']

Parameters:

  • chain_index (Int[Array, 'num_residues'])

  • neighbor_indices (Int[Array, 'num_atoms num_neighbors'])

prxteinmpnn.model.features.encode_positions(neighbor_offsets, edge_chains_neighbors, model_parameters)[source]#

Encode positions based on neighbor offsets and edge chains.

Return type:

Int[Array, 'num_atoms num_neighbors (2 * MAXIMUM_RELATIVE_FEATURES + 2)']

Parameters:

  • neighbor_offsets (Int[Array, 'num_residues num_neighbors'])

  • edge_chains_neighbors (Int[Array, 'num_atoms num_neighbors'])

  • model_parameters (PyTree[str, 'P'])

prxteinmpnn.model.features.embed_edges(edge_features, model_parameters)[source]#

Embed edge features using model parameters.

Return type:

Float[Array, 'num_atoms num_neighbors num_features']

Parameters:

  • edge_features (Int[Array, 'num_atoms num_neighbors (2 * MAXIMUM_RELATIVE_FEATURES + 2)'])

  • model_parameters (PyTree[str, 'P'])

prxteinmpnn.model.features.extract_features(prng_key, model_parameters, structure_coordinates, mask, residue_index, chain_index, k_neighbors=48, augment_eps=0.0)[source]#

Extract features from protein structure coordinates.

Parameters:

  • structure_coordinates (Float[Array, 'num_residues num_atoms 3']) – Atomic coordinates of the protein structure.

  • mask (Int[Array, 'num_residues num_atoms']) – Mask indicating valid atoms in the structure.

  • residue_index (Int[Array, 'num_residues']) – Residue indices for each atom.

  • chain_index (Int[Array, 'num_residues']) – Chain indices for each atom.

  • model_parameters (PyTree[str, 'P']) – Model parameters for the feature extraction.

  • prng_key (Union[Key[Array, ''], UInt32[Array, '2']]) – JAX random key for stochastic operations.

  • k_neighbors (int) – Maximum number of neighbors to consider for each atom.

  • augment_eps (float) – Standard deviation for Gaussian noise augmentation.

Returns:

Edge features after concatenation and normalization. edge_indices: Indices of neighboring atoms.

Return type:

edge_features

prxteinmpnn.model.features.project_features(model_parameters, edge_features)[source]#

Project edge features using model parameters.

Return type:

Float[Array, 'num_atoms num_neighbors num_features']

Parameters:

  • model_parameters (PyTree[str, 'P'])

  • edge_features (Float[Array, 'num_atoms num_neighbors num_features'])

Contents