16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134 | def main():
"""Command-line entry point for the multi-model splitter.
The script:
- loads a PDB or mmCIF file,
- determines whether the file contains multiple 3D models,
- groups atom records by model number,
- optionally converts each model to PDB-compatible format,
- writes every model to a separate file in the output directory.
This tool is useful for processing NMR ensembles or any structure
that contains multiple frames in a single file.
"""
parser = argparse.ArgumentParser(
description="Split a multi-model PDB or mmCIF file into separate files per model."
)
parser.add_argument("--output", "-o", help="Output directory", required=True)
parser.add_argument(
"--format",
"-f",
help="Output format (possible values: PDB, mmCIF, keep. Default: keep)",
default="keep",
)
parser.add_argument("file", help="Input PDB or mmCIF file to split")
args = parser.parse_args()
# Check if input file exists
if not os.path.exists(args.file):
print(f"Error: Input file not found: {args.file}", file=sys.stderr)
sys.exit(1)
# Read and parse the input file
input_format = "mmCIF"
try:
with open(args.file) as f:
if is_cif(f):
atoms_df = parse_cif_atoms(f)
model_column = "pdbx_PDB_model_num"
else:
atoms_df = parse_pdb_atoms(f)
input_format = "PDB"
model_column = "model"
except Exception as e:
print(f"Error parsing file {args.file}: {e}", file=sys.stderr)
sys.exit(1)
if atoms_df.empty:
print(f"Warning: No atoms found in {args.file}", file=sys.stderr)
sys.exit(0)
# Check if model column exists
if model_column not in atoms_df.columns:
print(
f"Error: Model column '{model_column}' not found in the parsed data from {args.file}.",
file=sys.stderr,
)
print(
"This might indicate an issue with the input file or the parser.",
file=sys.stderr,
)
sys.exit(1)
# Determine output format
output_format = args.format.upper()
if output_format == "KEEP":
output_format = input_format
elif output_format not in ["PDB", "MMCIF"]:
print(
f"Error: Invalid output format '{args.format}'. Choose PDB, mmCIF, or keep.",
file=sys.stderr,
)
sys.exit(1)
# Ensure output directory exists
os.makedirs(args.output, exist_ok=True)
# Group by model number
grouped_by_model = atoms_df.groupby(model_column)
# Get base name for output files
base_name = os.path.splitext(os.path.basename(args.file))[0]
# Write each model to a separate file
for model_num, model_df in grouped_by_model:
# Ensure model_df is a DataFrame copy to avoid SettingWithCopyWarning
model_df = model_df.copy()
# Set the correct format attribute for the writer function
model_df.attrs["format"] = input_format
# Construct output filename
ext = ".pdb" if output_format == "PDB" else ".cif"
output_filename = f"{base_name}_model_{model_num}{ext}"
output_path = os.path.join(args.output, output_filename)
print(f"Writing model {model_num} to {output_path}...")
try:
if output_format == "PDB":
df_to_write = fit_to_pdb(model_df)
write_pdb(df_to_write, output_path)
else: # mmCIF
write_cif(model_df, output_path)
except ValueError as e:
# Handle errors specifically from fit_to_pdb
print(
f"Error fitting model {model_num} from {args.file} to PDB: {e}. Skipping model.",
file=sys.stderr,
)
continue
except Exception as e:
# Handle general writing errors
print(
f"Error writing file {output_path} for model {model_num}: {e}",
file=sys.stderr,
)
print("Splitting complete.")
|