one_chain.py

import os
from string import ascii_uppercase

def consolidate_chains(pdb_file_path):
    try:
        with open(pdb_file_path, 'r') as file:
            lines = file.readlines()

        processed_lines = []

        for line in lines:
            if line.startswith("ATOM"):
                # Extract chain ID and residue ID from the line
                chain_id = line[21]
                residue_id = int(line[22:26].strip())

                # Calculate the new residue ID
                if chain_id == 'B':
                    new_residue_id = 200 + residue_id
                else:
                    chain_index = ascii_uppercase.index(chain_id)
                    new_residue_id = (200 * (chain_index - ascii_uppercase.index('A'))) + residue_id

                # Rebuild the line with new residue ID and chain set to 'A'
                new_line = line[:21] + 'A' + f"{new_residue_id:4}" + line[26:]
                processed_lines.append(new_line)
            else:
                # Non-ATOM lines are added as-is
                processed_lines.append(line)

        # Construct the output file path
        file_dir, file_name = os.path.split(pdb_file_path)
        file_base, file_ext = os.path.splitext(file_name)
        output_file_path = os.path.join(file_dir, f"{file_base}_one_chain{file_ext}")

        # Write the processed lines to the new file
        with open(output_file_path, 'w') as file:
            file.writelines(processed_lines)

        return output_file_path
    except Exception as e:
        return str(e)

def consolidate_chains_2(pdb_file_path):
    try:
        with open(pdb_file_path, 'r') as file:
            lines = file.readlines()

        processed_lines = []
        last_residue_id_of_A = None
        residue_id_increment = 0

        for line in lines:
            if line.startswith("ATOM"):
                chain_id = line[21]
                residue_id = int(line[22:26].strip())

                if chain_id == 'A':
                    last_residue_id_of_A = residue_id
                    new_residue_id = residue_id
                else:
                    if last_residue_id_of_A is not None and residue_id_increment == 0:
                        residue_id_increment = last_residue_id_of_A - residue_id + 1
                    new_residue_id = residue_id + residue_id_increment

                new_line = line[:21] + 'A' + f"{new_residue_id:4}" + line[26:]
                processed_lines.append(new_line)
            elif line.startswith("TER"):
                continue

        # Add a 'TER' line at the end to indicate termination of the single chain
        processed_lines.append("TER\n")

        # Construct the output file path
        file_dir, file_name = os.path.split(pdb_file_path)
        file_base, file_ext = os.path.splitext(file_name)
        output_file_path = os.path.join(file_dir, f"{file_base}_one_chain{file_ext}")

        # Write the processed lines to the new file
        with open(output_file_path, 'w') as file:
            file.writelines(processed_lines)

        return output_file_path
    except Exception as e:
        return str(e)
import os

def consolidate_chains_3(pdb_file_path, new_path, increments):
    try:
        with open(pdb_file_path, 'r') as file:
            lines = file.readlines()

        processed_lines = []
        # Initialize a dictionary to store the last residue ID encountered for each chain
        last_residue_id = {}
        # Initialize a dictionary to store the increment for each chain
        chain_increments = {chr(65 + i): inc for i, inc in enumerate(increments)}
        # Flag to track if we are inside a model
        in_model = False
        print(chain_increments)
        for i, line in enumerate(lines):
            if line.find("MODEL") != -1:
                in_model = True
                # Reset last_residue_id for each new model
                last_residue_id = {}
                processed_lines.append(line)
            elif line.startswith("ENDMDL"):
                in_model = False
                processed_lines.append(line)
            elif line.startswith("ATOM") and in_model:
                chain_id = line[21]
                residue_id = int(line[22:26].strip())

                # Calculate the new residue ID by adding the specified increment
                if chain_id in chain_increments:
                    # Use the increment specified for the chain, fallback to 0 if not specified
                    increment = chain_increments.get(chain_id, 0)
                else:
                    increment = 0
#012345678901234567890123456789
#ATOM   1861  CZ  ARG B 217      12
                new_residue_id = residue_id + increment

                # Update the last residue ID for the chain
                last_residue_id[chain_id] = new_residue_id

                # Construct the new line with the updated chain ID ('A') and new residue ID
                new_line = line[:21] + 'A' + f"{new_residue_id:4}" + line[26:]
                processed_lines.append(new_line)
            elif not in_model:
                processed_lines.append(line)

        # Construct the output file path
        file_dir, file_name = os.path.split(pdb_file_path)
        file_base, file_ext = os.path.splitext(file_name)
        output_file_path = os.path.join(new_path, f"{file_base}_one_chain{file_ext}")

        # Write the processed lines to the new file
        with open(output_file_path, 'w') as file:
            file.writelines(processed_lines)

        return output_file_path
    except Exception as e:
        return str(e)

import sys
from paths import *

if len(sys.argv) < 4:
    print("Usage: python one_chain.py <pdb> <min_residue_id> <max_residue_id>")
    sys.exit(1)

pdb = sys.argv[1].upper()
min_residue_id = int(sys.argv[2])
max_residue_id = int(sys.argv[3])

path = f"{PDB_FILES}{pdb}/"
new_path = f"{PDB_FILES}{pdb}_RPF/"
if not os.path.exists(new_path):
    os.makedirs(new_path)

for root, dirs, files in os.walk(path):
    for filename in files:
        if not os.path.exists(new_path + filename):
            fn = os.path.abspath(os.path.join(root, filename))
            print(fn)
            output_file = consolidate_chains_3(fn, new_path, [min_residue_id, max_residue_id])
            print(f"New file saved at: {output_file}")