Gnina parser
Parser for GNINA docked SDF files and dataset builder.
DOCKING_SCHEMA = pa.schema([pa.field('block_start', pa.int64()), pa.field('block_end', pa.int64()), pa.field('protein_id', pa.string()), pa.field('ligand_id', pa.string()), pa.field('UniqueID', pa.int32()), pa.field('conformer_idx', pa.int32()), pa.field('pose_rank', pa.int32()), pa.field('SMILES', pa.string()), pa.field('n_atoms', pa.int32()), pa.field('Vina_affinity', pa.float64()), pa.field('CNNscore', pa.float64()), pa.field('CNNaffinity', pa.float64()), pa.field('CNN_VS', pa.float64()), pa.field('CNNaffinity_variance', pa.float64()), pa.field('Energy', pa.float64()), pa.field('source_file', pa.string())])
¶
SCORE_DIRECTIONS = {'CNNscore': 'descending', 'CNNaffinity': 'descending', 'CNN_VS': 'descending', 'Vina_affinity': 'ascending', 'Energy': 'ascending', 'CNNaffinity_variance': 'ascending'}
¶
DockedPose(row_idx, table, dr)
¶
Lightweight container for a single docked pose.
Initialize a DockedPose by reading the specified row from the table and storing a reference to the GNINA_Results for block access. Args: row_idx: The index of the row in the table to read. table: The PyArrow Table containing the docking data. dr: The GNINA_Results instance to reference for block access.
__slots__ = ['block_start', 'block_end', 'ligand_id', 'protein_id', 'UniqueID', 'conformer_idx', 'pose_rank', 'SMILES', 'Energy', 'Vina_affinity', 'CNNscore', 'CNNaffinity', 'CNN_VS', 'CNNaffinity_variance', '_dr']
¶
block_text
¶
Get the raw SDF block text for this pose using stored byte offsets. Returns: The raw text of the SDF block for this pose.
coords
¶
Get the atomic coordinates and element symbols for this pose. Returns: A tuple containing: - A NumPy array of shape (n_atoms, 3) with the atomic coordinates. - A list of element symbols corresponding to each atom.
mol
¶
Get the RDKit molecule object for this pose. Returns: An RDKit molecule object representing the structure of this pose.
__repr__()
¶
DockingDataset(dataset_dir)
¶
Manages per protein docking dataset and converts to batched Parquet
one parquet per protein layout::
dataset_dir/
├── parquet/{PROTEIN_ID}.parquet
├── blocks/{PROTEIN_ID}/*.sdf (optional, only if copy_blocks=True)
└── proteins/{PROTEIN_ID}.pdb (optional copy of cleaned receptor PDBs, only if copy_proteins=True)
Parameters¶
dataset_dir : str Root of the dataset directory.
schema
¶
Aggregated schema across all per-protein parquets in this dataset.
__repr__()
¶
add_protein(protein_id, sdf_files, num_modes=9, copy_blocks=False, overwrite=False)
¶
Add a single protein's docking results from given SDF files.
Args:
protein_id: The identifier for the protein (e.g., "1abc").
sdf_files: List of paths to the docked SDF files for this protein.
num_modes: Number of poses per conformer (GNINA default 9).
copy_blocks: If True, copy SDF files into blocks/
build_from_docking_tree(docking_root, num_modes=9, copy_blocks=False, copy_proteins=False, protein_glob='*_cleaned.pdb', overwrite=False)
¶
Scan a raw docking tree and build the per protein dataset.
| PARAMETER | DESCRIPTION |
|---|---|
docking_root
|
Root containing the sdf files
|
num_modes
|
Poses per conformer (GNINA default 9).
|
copy_blocks
|
If True, copy docked SDF files into
blocks/
|
copy_proteins
|
If True, copy cleaned PDB files into proteins/.
|
protein_glob
|
Glob pattern for the cleaned receptor PDB inside each protein directory.
|
overwrite
|
If False (default), skip proteins that already have a parquet file.
|
Returns: A dictionary mapping {protein_id: n_poses} for every protein processed.
iter_batches(protein_id=None, batch_size=200, columns=None)
¶
Stream RecordBatches for loops.
If protein_id is given, streams from that single parquet. Otherwise streams across the whole dataset. Args: protein_id: Optional single protein ID to stream from. batch_size: Number of rows per batch when streaming across the whole dataset. columns: Optional list of columns to read. None = all columns. Returns: An iterator of PyArrow RecordBatch objects.
iter_batches_from(path, batch_size=200, columns=None, **filters)
¶
Stream batches from any parquet file or directory. Args: path: Path to a single parquet file or a directory containing batch_.parquet files. batch_size: Number of rows per batch when streaming. columns: Optional list of columns to read. None = all columns. *filters: Equality filters for pruning, e.g. ligand_id="M3A". Returns: An iterator of PyArrow RecordBatch objects matching the filters.
protein_ids()
¶
List all protein IDs that have parquet files.
read_all(columns=None, **filters)
¶
Cross-protein query via pyarrow.dataset (lazy scan).
columns : list[str], optional Column subset. **filters Equality filters pushed down to row groups, e.g. ligand_id="M3A".
read_batched_parquets(directory, columns=None, **filters)
¶
Read from an batch directory with row-group pruning. Args: directory: Path to the batch directory containing batch_.parquet files. columns: Optional list of columns to read. None = all columns. *filters: Equality filters for pruning, e.g. ligand_id="M3A". Returns: A PyArrow Table containing the filtered results from all batch files.
read_block(protein_id, row_idx)
¶
Read a single SDF block from the dataset by protein and row index.
Works whether blocks were copied into the dataset or not. For copied blocks (decompressed .sdf), this is a fast seek. For original .sdf.gz files, this requires full decompression.
| PARAMETER | DESCRIPTION |
|---|---|
protein_id
|
The protein ID.
TYPE:
|
row_idx
|
Row index within that protein's parquet.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
str
|
The raw SDF block text. |
read_protein(protein_id, columns=None)
¶
Read one protein's parquet.
read_proteins(protein_ids, columns=None)
¶
Read a subset of proteins into one table. Args: protein_ids: List of protein IDs to read. columns: Optional list of columns to read. None = all columns. Returns: A PyArrow Table concatenating the specified proteins. If a protein ID is not found, it is skipped with a warning.
stats()
¶
Quick counts without reading data (parquet metadata only).
to_batched_parquets(output_dir, proteins_per_batch=200, row_group_size=50000, overwrite=False)
¶
Groups proteins into batch files, each containing proteins_per_batch proteins sorted by protein_id. Batch files are named by the first and last PDB ID within e.g. bt_0_10GS_4HVP.parquet. Args: output_dir :Destination directory for batch_000.parquet, etc. proteins_per_batch : How many proteins per batch file. row_group_size : Target rows per row group within each batch file. overwrite : If False (default), skip writing a batch file if it already exists. Returns: Number of batch files written.
topn_per_pair(table, rank_by='CNNscore', n=5, add_rank_col=True, group_keys=('protein_id', 'ligand_id'), rank_col='pair_rank')
¶
Return the top-N rows per protein-lifand with a rank column added if requested. Args: table: PyArrow table. rank_by: Score column used for ranking within each group. n: Keep only top n rowa based on the score. group_keys: Columns defining the group. rank_col: Name of the 0-indexed rank column added to the output. Returns: A PyArrow Table, pre-sorted by (group_keys, rank_by), with rank column added if requested.
GNINA_Results(filepath, num_modes=9, protein_id=None, keep_raw=False, allow_truncated=True)
¶
Lazy-indexed parser for a single GNINA docked SDF file.
Initialize the GNINA_Results parser. Args: filepath: Path to the .sdf or .sdf.gz file containing the docking results. num_modes: The number of modes per ligand/protomer. Can be specified as: - int: A single value applied to all groups. - list[int]: A list of known mode counts to match against group block counts. - None : Auto-detect from block count patterns. Default is 9, which is the standard for GNINA docking outputs. protein_id: Optional identifier for the target protein. If not provided, it will be inferred from the file name. keep_raw: Whether to keep the full decompressed text in memory after parsing. Default is False, the raw text is released after the table is built and a seekable file on disk is used for subsequent block retrieval. allow_truncated: Whether to allow files that appear truncated (e.g. missing some expected blocks).
blocks
¶
All raw SDF blocks as a list of strings. This will load the entire raw text into memory if not already loaded.
filepath
¶
The original file path of the SDF or SDF.GZ file that was parsed.
is_raw_loaded
¶
Whether the full decompressed text is currently held in memory.
n_blocks
¶
The total number of SDF blocks (poses) parsed from the file.
n_conformers
¶
Combine ligand_id, UniqueID, and conformer_idx to count distinct conformers per protomer (e.g., if num_modes=9, expect 9 conformers per protomer).
n_ligands
¶
Count of distinct ligand_id values in the table.
n_protomers
¶
Combine ligand_id and UniqueID to count distinct protomers
num_modes
¶
The unique num_modes values found across all groups, sorted.
Returns a sorted list of distinct mode counts. E.g. [9] if uniform, [7, 9] if one group was truncated.
num_modes_per_group
¶
Per-group num_modes mapping.
Returns a dictionary mapping (ligand_id, UniqueID): num_modes for every group in the file.
protein_id
¶
The identifier for the target protein, either provided or inferred from the file name.
schema
¶
The PyArrow Schema for the docking results table.
table
¶
The full PyArrow Table containing all parsed docking data.
__del__()
¶
clean up if in case cleanup() wasn't called explicitly.
__enter__()
¶
Support with GNINA_Results(...) as dr: usage.
__exit__(exc_type, exc_val, exc_tb)
¶
Clean up temp file when exiting the context.
__len__()
¶
override len() to return the number of blocks (poses) in the dataset.
__repr__()
¶
Return a string representation of the GNINA_Results object, including file name and summary stats.
best_direction(score)
¶
Return the sort direction ('ascending' or 'descending') that corresponds to 'best' for a given score column. Args: score: The column name (e.g. 'CNNscore', 'Vina_affinity'). Returns: 'ascending' or 'descending'.
cleanup()
¶
Remove the temporary seekable file if we created one.
filter(**kwargs)
¶
get_best_per_ligand(by='CNNscore', ascending=None, weights=None)
¶
Convenience method to get the single best pose per ligand_id.
Args:
by: Column name or list of column names to rank by.
ascending: Sort direction(s). None = auto-detect from score type.
weights: Optional weights for composite scoring (one per column in by).
Returns:
A PyArrow Table containing the best pose per ligand_id.
get_block_by_offsets(block_start, block_end)
¶
Retrieve a block by byte offsets into the decompressed raw string or from disk if raw has been released. Args: block_start: Start byte offset of the block. block_end: End byte offset of the block. Returns: The raw text of the SDF block.
get_conformer(ligand_id, unique_id, conformer_idx=0)
¶
Get all poses for a specific conformer/pose. Args: ligand_id: The ligand identifier (e.g., "M3A") to filter by. unique_id: The UniqueID to filter by, which distinguishes different protomers of the same ligand. conformer_idx: The conformer index to filter by (default 0). Returns: A PyArrow Table containing the pose info.
get_coords(block_start, block_end)
¶
Parse atomic coordinates and element symbols from a block. Args: block_start: Start byte offset of the block. block_end: End byte offset of the block. Returns: A tuple containing: - A NumPy array of shape (n_atoms, 3) with the atomic coordinates. - A list of element symbols corresponding to each atom.
get_coords_batch(offsets)
¶
Get coordinates for a batch of blocks. Args: offsets: List of (block_start, block_end) tuples. Returns: A list of (coords_array, elements_list) tuples.
get_ligand(ligand_id)
¶
Get all poses for a specific ligand_id (across all protomers and conformers). Args: ligand_id: The ligand identifier (e.g., "M3A") to filter by. Returns: A PyArrow Table containing all poses for the specified ligand_id.
get_mol(block_start, block_end)
¶
Parse a block into an RDKit molecule object. Args: block_start: Start byte offset of the block. block_end: End byte offset of the block. Returns: An RDKit molecule object.
get_protomer(ligand_id, unique_id)
¶
Get all poses for a specific protomer (ligand_id + UniqueID). Args: ligand_id: The ligand identifier (e.g., "M3A") to filter by. unique_id: The UniqueID to filter by, which distinguishes different protomers of the same ligand. Returns: A PyArrow Table containing all poses for the specified protomer.
get_top_poses(by='CNNscore', ascending=None, weights=None, n=1, per='conformer')
¶
Get the top N poses sorted by one or more score columns
Args:
by: Column name or list of column names to sort/rank by.
ascending: Sort direction(s). Can be:
- None (default): auto-detect from SCORE_DIRECTIONS based on the 'by' column.
- bool: applied to all columns in by.
- list[bool]: one per column in by.
weights: Optional list of floats (one per column in by). When
provided, a weighted composite score is computed. Weights are
normalised to sum to 1. Only valid when by is a list.
n: The number of top poses to return per group. Default is 1.
per: The grouping level for selecting top poses. Options are
'conformer', 'protomer', or 'ligand'.
Returns:
A PyArrow Table containing the top N poses per specified group.
poses_per_conformer()
¶
Count the number of poses per unique conformer (ligand_id + UniqueID + conformer_idx). Returns: A PyArrow Array where each element corresponds to the count of poses for a unique conformer. should be 9 for each conformer if num_modes=9.
protomer_counts()
¶
Count the number of poses and conformers per protomer (ligand_id + UniqueID).
read_block_from_file(source_file, block_start, block_end)
¶
Read a single SDF block directly from a file using stored byte offsets. For .sdf files this is an O(1) seek. For .sdf.gz files the entire file must be decompressed first (offsets are into decompressed text). Args: source_file: Path to the .sdf or .sdf.gz file. block_start: Start byte offset of the block in the decompressed text. block_end: End byte offset of the block in the decompressed text. Returns: The raw text of the SDF block.
release_raw()
¶
Manually release the raw text to free memory. Subsequent block access will read from disk.
summary(per='global', output=None)
¶
write a human-readable summary of the docking results Args: per: The level of detail for the summary. Options are: - 'global': Overall summary for the entire dataset. - 'ligand': Summary broken down by ligand_id, showing score ranges and best pose per ligand. - 'protomer': Summary broken down by protomer (ligand_id + UniqueID), showing score ranges and best pose per protomer. output: The output destination for the summary. Options are: - None (default): Print to console using loguru. - str: A file path to write the summary to.
to_csv(output_path, overwrite=False)
¶
Export the docking table to a CSV file. Args: output_path: Path to the output CSV file. overwrite: If False (default), skip if the file already exists.
to_sdf(output_path, table=None, overwrite=False)
¶
Export SDF blocks to a file. Args: output_path: Path to the output .sdf or .sdf.gz file. table: A PyArrow Table with block_start/block_end columns. None = export all blocks. overwrite: If False (default), skip if the file already exists.
TruncatedSDFError(filepath, reason)
¶
as_poses(table, dr)
¶
Convert a PyArrow Table of docking results into a list of DockedPose objects. Args: table: A PyArrow Table containing the docking data, with required columns. dr: The GNINA_Results instance to reference for block access. Returns: A list of DockedPose objects, one for each row in the table.