<jupyter_start><jupyter_text>Summary Submitting jobs**Note:** These jobs must be submitted from the ./notebooks folder.**Cedar:**```bashNOTEBOOK_PATH=$(realpath 01_process_pdb_core.ipynb) ORIGINAL_ARRAY_TASK_COUNT=300 sbatch --array=274 --time=24:00:00 --nodes=1 --ntasks-per-node=48 --mem=0 --job-name=process-pdb-core --account=rrg-pmkim --output=/scratch/strokach/tmp/log/run-notebook-cpu-%j-%N.log ../scripts/run_notebook_cpu.sh``` To Do- Remove hydrogen atoms on all structures (use reduce trim?).----  Imports<jupyter_code>import concurrent.futures
import concurrent.futures.process
import gzip
import importlib
import io
import logging
import os
import pickle
import shlex
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import traceback
import warnings
from functools import partial
from itertools import islice
from pathlib import Path

from tqdm.notebook import tqdm

import matplotlib.pyplot as plt
import mdtraj
import numpy as np
import pandas as pd
import psutil
import pyarrow as pa
import pyarrow.parquet as pq
import tenacity
import yaml
from kmbio import PDB
from kmtools import structure_tools
%matplotlib inline

pd.set_option("max_columns", 100)
SRC_PATH = Path.cwd().joinpath('..', 'src').resolve(strict=True)

if SRC_PATH.as_posix() not in sys.path:
    sys.path.insert(0, SRC_PATH.as_posix())

import helper
importlib.reload(helper)<jupyter_output><empty_output><jupyter_text>Parameters<jupyter_code>NOTEBOOK_PATH = Path(os.getenv("CI_JOB_NAME", "01_process_pdb_core"))

NOTEBOOK_PATH
OUTPUT_PATH = Path(os.getenv('OUTPUT_DIR', NOTEBOOK_PATH.name)).resolve()
OUTPUT_PATH.mkdir(parents=True, exist_ok=True)

OUTPUT_PATH
if "scinet" in socket.gethostname():
    CPU_COUNT = 40
else:
    CPU_COUNT = max(1, len(os.sched_getaffinity(0)) // 2)

CPU_COUNT
TASK_ID = os.getenv("SLURM_ARRAY_TASK_ID")
TASK_COUNT = os.getenv("ORIGINAL_ARRAY_TASK_COUNT") or os.getenv("SLURM_ARRAY_TASK_COUNT")

TASK_ID = int(TASK_ID) if TASK_ID is not None else None
TASK_COUNT = int(TASK_COUNT) if TASK_COUNT is not None else None

TASK_ID, TASK_COUNT
DEBUG = TASK_ID is None

if DEBUG:
    TASK_ID = 87
    TASK_COUNT = 300
else:
    assert TASK_ID is not None
    assert TASK_COUNT is not None

TASK_ID, TASK_COUNT
ADJACENCY_NET_DATA_PATH = (
    Path(os.getenv("DATAPKG_OUTPUT_DIR")).joinpath("adjacency-net-v2", "v0.3").resolve(strict=True)
)

ADJACENCY_NET_DATA_PATH
PDB_DATA_PATH = (
    Path(os.environ["DATAPKG_OUTPUT_DIR"]).joinpath("pdb-ffindex", "2020-01-16", "arrow").resolve(strict=True)
)

PDB_DATA_PATH<jupyter_output><empty_output><jupyter_text>Load data<jupyter_code>with PDB_DATA_PATH.joinpath("pdb-list.pickle").open("rb") as fin:
    pdb_list = pickle.load(fin)
    
pdb_data_reader = pa.RecordBatchFileReader(PDB_DATA_PATH.joinpath("pdb-mmcif.arrow"))

assert len(pdb_list) == pdb_data_reader.num_record_batches
pdb_list[:3]
chunk_size = int(np.ceil(len(pdb_list) / TASK_COUNT))
task_idx = TASK_ID - 1
pdb_chunk = pdb_list[task_idx * chunk_size : (task_idx + 1) * chunk_size]
pdb_chunk_idxs = list(range(task_idx * chunk_size, (task_idx + 1) * chunk_size))[:len(pdb_chunk)]
assert all(pdb_chunk[i] == pdb_list[j] for i, j in enumerate(pdb_chunk_idxs))

chunk_size, task_idx, len(pdb_chunk), pdb_chunk[:3]
# if DEBUG:
#     pdb_chunk = pdb_chunk[:10]
#     pdb_chunk_idxs = pdb_chunk_idxs[:10]<jupyter_output><empty_output><jupyter_text>Process structures<jupyter_code>def structure_from_chain(structure_ref, model_ref, chain):
    model = PDB.Model(model_ref.id, model_ref.serial_num)
    model.add(chain)
    structure = PDB.Structure(structure_ref.id)
    structure.add(model)
    assert len(list(structure.chains)) == 1
    return structure
def process_pdb_core(pdb_id, pdb_idx):
    # Load data
    pdb_data_reader = pa.RecordBatchFileReader(PDB_DATA_PATH.joinpath("pdb-mmcif.arrow"))
    pdb_data = pdb_data_reader.get_record_batch(pdb_idx).to_pydict()
    assert pdb_data["pdb_id"][0] == pdb_id

    # Create structure from data
    buf = io.StringIO()
    buf.write(gzip.decompress(pdb_data["mmcif_data"][0]).decode())
    use_auth_id = False
    try:
        buf.seek(0)
        bioassembly_id = True
        structure = PDB.MMCIFParser(use_auth_id=use_auth_id).get_structure(buf, bioassembly_id=bioassembly_id)
    except PDB.BioassemblyError as e:
        print(f"Encountered error when parsing pdb {pdb_idx} ('{pdb_id}'): {e!s}.")
        buf.seek(0)
        bioassembly_id = False
        structure = PDB.MMCIFParser(use_auth_id=use_auth_id).get_structure(buf, bioassembly_id=bioassembly_id)

    results = []
    _seen = set()
    for model_idx, model in enumerate(structure):
        for chain_idx, chain in enumerate(model):
            aa_sequence = structure_tools.get_chain_sequence(chain, if_unknown="replace")
            if aa_sequence in _seen:
                continue
            _seen.add(aa_sequence)
            if len(aa_sequence.strip()) < 5:
                continue

            schain = structure_from_chain(structure, model, chain.copy())

            with tempfile.NamedTemporaryFile(suffix=".pdb") as pdb_file:
                PDB.save(schain, pdb_file.name)
                traj = mdtraj.load(pdb_file.name)
            assert aa_sequence == traj.top.to_fasta()[0]

            residue_df = helper.construct_residue_df(traj)
            helper.validate_residue_df(residue_df)

            residue_pairs_df = helper.construct_residue_pairs_df(traj)
            helper.validate_residue_pairs_df(residue_pairs_df)

            result = {
                "pdb_id": [pdb_id],
                "pdb_idx": [pdb_idx],
                "use_auth_id": [use_auth_id],
                "bioassembly_id": [bioassembly_id],
                "model_idx": [model_idx],
                "model_id": [model.id],
                "chain_idx": [chain_idx],
                "chain_id": [chain.id],
                "aa_sequence": [aa_sequence],
                **helper.residue_df_to_row(residue_df),
                **helper.residue_pairs_df_to_row(residue_pairs_df),
            }
            result = helper.downcast_and_compress(result)
            results.append(result)

    return results
def worker(pdb_id, pdb_idx):
    try:
        results = process_pdb_core(pdb_id, pdb_idx)
        return results, []
    except Exception as error:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        traceback_string = "\n".join(traceback.format_exception(exc_type, exc_value, exc_traceback))
        failure = {
            "pdb_id": [pdb_id],
            "pdb_idx": [pdb_idx],
            "error_type": [str(type(error))],
            "error_message": [str(error)],
            "error_traceback": [traceback_string],
        }
        return [], [failure]
# result = process_pdb_core("1ju3", 11704)
# result
# with concurrent.futures.ProcessPoolExecutor() as pool:
#     futures = pool.map(worker, pdb_chunk[:10], pdb_chunk_idxs[:10])
#     for result in tqdm(futures, total=100):
#         break
# result
output_dir = ADJACENCY_NET_DATA_PATH.joinpath("pdb-core")
output_dir.mkdir(exist_ok=True)
output_file = output_dir.joinpath(f"pdb-core-{TASK_ID}-{TASK_COUNT}.arrow")

output_dir_failed = output_dir.joinpath("failed")
output_dir_failed.mkdir(exist_ok=True)
output_file_failed = output_dir_failed.joinpath(f"pdb-core-{TASK_ID}-{TASK_COUNT}-failed.arrow")

output_file, output_file_failed
logging.getLogger("kmtools.structure_tools.fixes").setLevel(logging.CRITICAL)
writer = None
writer_failed = None
num_pdbs_processed = 0
while num_pdbs_processed < len(pdb_chunk):
    try:
        with concurrent.futures.ProcessPoolExecutor(CPU_COUNT) as pool:
            futures = pool.map(worker, pdb_chunk[num_pdbs_processed:], pdb_chunk_idxs[num_pdbs_processed:], chunksize=1)
            for (results, results_failed) in tqdm(futures, total=len(pdb_chunk) - num_pdbs_processed):
                num_pdbs_processed += 1
                for result in results:
                    if writer is None:
                        batch = pa.RecordBatch.from_arrays(list(result.values()), list(result.keys()))
                        writer = pa.RecordBatchFileWriter(output_file, batch.schema)
                    batch = pa.RecordBatch.from_arrays(list(result.values()), list(result.keys()))
                    writer.write_batch(batch)
                for result_failed in results_failed:
                    if writer_failed is None:
                        batch = pa.RecordBatch.from_arrays(list(result_failed.values()), list(result_failed.keys()))
                        writer_failed = pa.RecordBatchFileWriter(output_file_failed, batch.schema)
                    batch = pa.RecordBatch.from_arrays(list(result_failed.values()), list(result_failed.keys()))
                    writer_failed.write_batch(batch)
    except concurrent.futures.BrokenExecutor as e:
        print(f"ProcessPoolExecutor crashed with an error ('{type(e)!s}'): '{e!s}'.")
if writer is not None:
    writer.close()
if writer_failed is not None:
    writer_failed.close()
if output_file.is_file():
    reader = pa.RecordBatchFileReader(output_file)
    print(f"Number of successful chains: {reader.num_record_batches}.")
if output_file_failed.is_file():
    reader_failed = pa.RecordBatchFileReader(output_file_failed)
    print(f"Number of failed PDBs: {reader_failed.num_record_batches}.")<jupyter_output><empty_output><jupyter_text>Write a `*.SUCCESS` file<jupyter_code>with output_file.with_suffix(".SUCCESS").open("wt"):
    pass<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11299138112255623, 'fraction_numerical': 0.008198444397729661, 'mean_word_length': 3.6964461994077, 'pattern_counts': {'":': 14, '<': 23, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '169846', 'license': 'MIT', 'n_tokens_mistral': 3388, 'n_tokens_neox': 3250, 'n_words': 576, 'path': 'notebooks2/01_process_pdb_core.ipynb', 'repo_name': 'ostrokach/protein-adjacency-net', 'size': '9514'}