expand.py

import sys
import os
from paths import PDB_FILES
from os.path import exists

def split_pdb_models(pdb_file_path, output_directory):
    """
    This function takes a PDB file with multiple models and splits it into separate files,
    each named according to the recommendation in the REMARK line following the string "File:".
    The new files are saved to the specified output directory.

    Args:
    pdb_file_path (str): Path to the input PDB file.
    output_directory (str): Path to the directory where output files will be saved.

    Returns:
    list of str: List of paths to the newly created PDB files.
    """
    try:
        with open(pdb_file_path, 'r') as file:
            lines = file.readlines()

        model_number = 0
        model_lines = []
        created_files = []
        model_name = ""

        for line in lines:
            if line.startswith("MODEL"):
                model_number += 1
                model_lines = [line]
            elif line.startswith("REMARK") and "File:" in line:
                # Extract the file name from the REMARK line
                model_name = line.split("File:")[1].strip().split()[0]
            elif line.startswith("REMARK"):
                args = line.split(' ')
                for arg in args:
                    arg = arg.strip()
                    if arg[-4:] == '.pdb':
                        model_name = arg[:-4]
                        print(model_name)
            elif line.startswith("ENDMDL"):
                model_lines.append(line)
                if model_name == "":
                    model_name = f"model_{model_number}"  # Fallback model name if not specified
                output_file_path = os.path.join(output_directory, model_name + ".pdb")
                with open(output_file_path, 'w') as model_file:
                    model_file.writelines(model_lines)
                created_files.append(output_file_path)
                model_lines = []  # Reset for the next model
                model_name = ""  # Reset model name for the next model
            elif model_number > 0:
                model_lines.append(line)

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

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python script.py <pdb_file_path> <output_directory>")
        sys.exit(1)
    
    pdb_file_path = sys.argv[1]
    output_directory = sys.argv[2]
    
    if output_directory.find(PDB_FILES) == -1:
        output_directory = PDB_FILES + output_directory
    if pdb_file_path.find(PDB_FILES) == -1 and not(exists(pdb_file_path)):
        pdb_file_path = PDB_FILES + pdb_file_path

    if not os.path.isdir(output_directory):
        os.makedirs(output_directory, exist_ok=True)

    pdb_file_paths = split_pdb_models(pdb_file_path, output_directory)
    for path in pdb_file_paths:
        print(f"Model file created: {path}")