Skip to content

Refactor sequence and coordinates retrieval logging for clarity and consistency #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisGadek1 merged 1 commit into from
Jul 28, 2026
Merged

Refactor sequence and coordinates retrieval logging for clarity and consistency #24

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 31 additions & 17 deletions toolbox/models/manage_dataset/sequences/sequence_and_coordinates_retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from toolbox.models.manage_dataset.compute_batches import ComputeBatches
from toolbox.models.manage_dataset.index.handle_index import read_index, create_index
from toolbox.models.manage_dataset.index.handle_indexes import HandleIndexes
from toolbox.models.manage_dataset.utils import format_time, read_pdbs_from_h5
from toolbox.models.manage_dataset.utils import format_time, groupby_dict_by_values, read_pdbs_from_h5
from toolbox.models.utils.cif2pdb import aa_dict
from toolbox.config import CarbonAtomType
from toolbox.utlis.logging import log_title, logger
Expand Down Expand Up @@ -274,8 +274,15 @@ def retrieve(
"""
structures_dataset = self.structures_dataset
client: Client = structures_dataset._client

log_title("Retrieving sequences and coordinates")

if is_sequences_retrieved and is_coordinates_retrieved:
retrieved_desc = "sequences and coordinates"
elif is_sequences_retrieved:
retrieved_desc = "sequences"
else:
retrieved_desc = "coordinates"

log_title(f"Retrieving {retrieved_desc}")

start = time.time()

Expand All @@ -285,21 +292,29 @@ def retrieve(
structures_dataset.config.data_path
)

# search sequences
search_sequences_index_result = self.handle_indexes.full_handle(
"sequences", protein_index, structures_dataset.overwrite
)
missing_sequences = []
sequences_coords_index = {}
coordinates_index = {}
missing_protein_files = {}

if is_sequences_retrieved:
search_sequences_index_result = self.handle_indexes.full_handle(
"sequences", protein_index, structures_dataset.overwrite
)
missing_sequences = search_sequences_index_result.missing_protein_files.keys()
sequences_coords_index = search_sequences_index_result.present
missing_protein_files.update(search_sequences_index_result.missing_protein_files)

h5_file_to_codes = search_sequences_index_result.grouped_missing_proteins
missing_sequences = search_sequences_index_result.missing_protein_files.keys()
sequences_coords_index = search_sequences_index_result.present
if is_coordinates_retrieved:
search_coordinates_index_result = self.handle_indexes.full_handle(
"coordinates", protein_index, structures_dataset.overwrite
)
coordinates_index = search_coordinates_index_result.present
missing_protein_files.update(search_coordinates_index_result.missing_protein_files)

search_coordinates_index_result = self.handle_indexes.full_handle(
"coordinates", protein_index, structures_dataset.overwrite
)
coordinates_index = search_coordinates_index_result.present
h5_file_to_codes = groupby_dict_by_values(missing_protein_files)

logger.info("Getting sequences and coordinates from stored PDBs")
logger.info(f"Getting {retrieved_desc} from stored PDBs")

# Get protein entries with range information if available
protein_entries = structures_dataset.get_protein_entries_with_ranges()
Expand Down Expand Up @@ -356,8 +371,7 @@ def collect(result):
if sequences_file is not None:
sequences_file.close()

log_line = f"Saving {'sequences' if is_sequences_retrieved else ''} {'and' if is_sequences_retrieved and is_coordinates_retrieved else ''} {'coordinates' if is_coordinates_retrieved else ''}"
logger.info(log_line)
logger.info(f"Saving {retrieved_desc}")

if is_sequences_retrieved:
# Update index with sequence file path (store as string like the original sequence_retriever)
Expand Down
Loading