@@ -105,88 +105,68 @@ def generate_output(plddt_data, name, out_dir, generate_tsv, pdb):
105105 ) as out_file :
106106 out_file .write (html_content )
107107
108+ def align_structures (structures ):
108109
109- def align_structures_old (structures ):
110- parser = PDB .PDBParser (QUIET = True )
111- structures = [
112- parser .get_structure (f"Structure_{ i } , pdb ) for i , pdb in enumerate (structures )
113- ]
114-
115- ref_structure = structures [0 ]
116- ref_atoms = [atom for atom in ref_structure .get_atoms ()]
110+ if not structures :
111+ raise ValueError ("No structures provided for alignment." )
117112
113+ if structures [0 ].endswith (".pdb" ):
114+ parser = PDB .PDBParser (QUIET = True )
115+ elif structures [0 ].endswith (".cif" ):
116+ parser = PDB .MMCIFParser (QUIET = True )
117+ else :
118+ raise ValueError (f"{ structure } )
119+
120+ parsed_structures = [parser .get_structure (f"structure-{ idx } , structure ) for idx , structure in enumerate (structures )]
121+ ref_structure = parsed_structures [0 ]
122+
123+ def get_atom_ids (structure ):
124+ # Note: this is a *set* of atom_ids due to the {} surrounding the comprehension
125+ return {(atom .get_parent ().get_id (), atom .name ) for atom in structure .get_atoms ()}
126+
127+ # TODO: do we want to raise and error if the structures are not identical atomically, or keep the ability to sub-align?
128+ # Update the atoms shared between structures with progressive intersections
129+ common_atoms = get_atom_ids (ref_structure )
130+ print ("commons: " , len (common_atoms ))
131+ for structure in parsed_structures [1 :]:
132+ common_atoms .intersection_update (get_atom_ids (structure ))
133+ print ("commons: " , len (common_atoms ))
134+
135+ if not common_atoms :
136+ raise ValueError ("No common atoms found between structures." )
137+ #print(common_atoms)
138+ def extract_atoms (structure , atom_ids ):
139+ # Note: this comprehension returns an atom *object* for each atom in the structure
140+ return [atom for atom in structure .get_atoms () if (atom .get_parent ().get_id (), atom .name ) in atom_ids ]
141+
142+ ref_atoms = extract_atoms (ref_structure , common_atoms )
143+ # The aligned structures will be the parsed structures aligned to the common atoms of the reference structure
118144 super_imposer = PDB .Superimposer ()
119- aligned_structures = [structures [0 ]] # Include the reference structure in the list
120-
121- for i , structure in enumerate (structures [1 :], start = 1 ):
122- target_atoms = [atom for atom in structure .get_atoms ()]
123-
145+ aligned_structures = []
146+ for idx , structure in enumerate (parsed_structures ):
147+ # The reference structure doesn't need to be aligned so can be skipped
148+ if idx == 0 :
149+ aligned_structures .append (structure )
150+ continue
151+ target_atoms = extract_atoms (structure , common_atoms )
152+ print (len (ref_atoms ), len (target_atoms ), len (common_atoms ))
124153 super_imposer .set_atoms (ref_atoms , target_atoms )
125154 super_imposer .apply (structure .get_atoms ())
126155
127- aligned_structure = f"aligned_structure_{ i }
128156 io = PDB .PDBIO ()
129157 io .set_structure (structure )
130- io .save (aligned_structure )
131- aligned_structures .append (aligned_structure )
158+ io .save (f"aligned_structure_ { idx } .pdb"
159+ aligned_structures .append (f"aligned_structure_ { idx } .pdb"
132160
161+ # Technically, parsed_structures now also points to the same aligned structures, but I've kept for readability
133162 return aligned_structures
134163
135164
136- def align_structures (structures ):
137- parser = PDB .PDBParser (QUIET = True )
138- structures = [
139- parser .get_structure (f"Structure_{ i } , pdb ) for i , pdb in enumerate (structures )
140- ]
141- ref_structure = structures [0 ]
142-
143- common_atoms = set (
144- f"{ atom .get_parent ().get_id ()[1 ]} { atom .name }
145- for atom in ref_structure .get_atoms ()
146- )
147- for i , structure in enumerate (structures [1 :], start = 1 ):
148- common_atoms = common_atoms .intersection (
149- set (
150- f"{ atom .get_parent ().get_id ()[1 ]} { atom .name }
151- for atom in structure .get_atoms ()
152- )
153- )
154-
155- ref_atoms = [
156- atom
157- for atom in ref_structure .get_atoms ()
158- if f"{ atom .get_parent ().get_id ()[1 ]} { atom .name } in common_atoms
159- ]
160- # print(ref_atoms)
161- super_imposer = PDB .Superimposer ()
162- aligned_structures = [structures [0 ]] # Include the reference structure in the list
163-
164- for i , structure in enumerate (structures [1 :], start = 1 ):
165- target_atoms = [
166- atom
167- for atom in structure .get_atoms ()
168- if f"{ atom .get_parent ().get_id ()[1 ]} { atom .name } in common_atoms
169- ]
170-
171- super_imposer .set_atoms (ref_atoms , target_atoms )
172- super_imposer .apply (structure .get_atoms ())
173-
174- aligned_structure = f"aligned_structure_{ i }
175- io = PDB .PDBIO ()
176- io .set_structure (structure )
177- io .save (aligned_structure )
178- aligned_structures .append (aligned_structure )
179-
180- return aligned_structures
181-
182165def pdb_to_lddt (struct_files , generate_tsv ):
183- struct_files_sorted = struct_files
184- struct_files_sorted .sort ()
185-
186166 output_lddt = []
187167 averages = []
188168
189- for struct_file in struct_files_sorted :
169+ for struct_file in struct_files :
190170 plddt_values = []
191171
192172 if struct_file .endswith ('.pdb' ):
@@ -254,21 +234,18 @@ def pdb_to_lddt(struct_files, generate_tsv):
254234
255235print ("generating html report..." )
256236
257- # structures = args.pdb
258- # # structures.sort()
259- # aligned_structures = align_structures(structures)
260-
261237# Preprocess "esmfold" PDB files, to reset residues on additional chains
262238processed_pdbs = [
263- (pdb_file .replace (".pdb" , "_align_residues.pdb" ) if "esmfold" in pdb_file else pdb_file )
264- for pdb_file in args .pdb
239+ pdb_file .replace (".pdb" , "_aligned.pdb" ) for pdb_file in args .pdb
265240]
266241
267- for pdb_file , output_pdb in zip ( args .pdb , processed_pdbs ) :
268- if "esmfold" in pdb_file :
269- reset_residue_numbers (pdb_file , output_pdb )
242+ for pdb_file in args .pdb :
243+ print ( "Reseting" , pdb_file , " into " , pdb_file . replace ( ".pdb" , "_aligned.pdb" ))
244+ reset_residue_numbers (pdb_file , pdb_file . replace ( ".pdb" , "_aligned.pdb" ) )
270245
271246structures = processed_pdbs # Use the final processed list
247+ print ("reference structure:" , processed_pdbs [0 ])
248+ print ("target structures:" , "," .join (processed_pdbs [1 :]))
272249aligned_structures = align_structures (structures )
273250
274251io = PDB .PDBIO ()
@@ -277,38 +254,49 @@ def pdb_to_lddt(struct_files, generate_tsv):
277254io .save (ref_structure_path )
278255aligned_structures [0 ] = ref_structure_path
279256
280- alphafold_template = open (args .html_template , "r" ).read ()
281- alphafold_template = alphafold_template .replace ("*sample_name*" , args .name )
282- alphafold_template = alphafold_template .replace ("*prog_name*" , args .in_type )
257+ comparision_template = open (args .html_template , "r" ).read ()
258+ comparision_template = comparision_template .replace ("*sample_name*" , args .name )
259+ comparision_template = comparision_template .replace ("*prog_name*" , args .in_type )
283260
284261args_pdb_array_js = (
285262 "const MODELS = [" + ",\n " .join ([f'"{ model } for model in structures ]) + "];"
286263)
287- alphafold_template = alphafold_template .replace ("const MODELS = [];" , args_pdb_array_js )
264+ comparision_template = comparision_template .replace ("const MODELS = [];" , args_pdb_array_js )
288265
289266seq_cov_imgs = []
290- for item in args .msa :
291- if item != "NO_FILE" :
292- image_path = item
267+ seq_cov_methods = []
268+ for msa , pdb in zip (args .msa , args .pdb ):
269+ if msa != "NO_FILE" :
270+ image_path = msa
271+ method = pdb .split (".pdb" )[0 ]
272+ seq_cov_methods .append (method )
293273 with open (image_path , "rb" ) as in_file :
294274 encoded_image = base64 .b64encode (in_file .read ()).decode ("utf-8" )
295275 seq_cov_imgs .append (f"data:image/png;base64,{ encoded_image } )
296276
277+ #MSA IMAGES
297278args_msa_array_js = (
298279 f"""const SEQ_COV_IMGS = [{ ", " .join ([f'"{ img } for img in seq_cov_imgs ])}
299280)
300- alphafold_template = alphafold_template .replace (
281+ comparision_template = comparision_template .replace (
301282 "const SEQ_COV_IMGS = [];" , args_msa_array_js
302283)
284+ #MSA IMAGE LABELS
285+ args_msa_method_array_js = (
286+ f"""const SEQ_COV_METHODS = [{ ", " .join ([f'"{ method } for method in seq_cov_methods ])}
287+ )
288+ comparision_template = comparision_template .replace (
289+ "const SEQ_COV_METHODS = [];" , args_msa_method_array_js
290+ )
303291
304292averages_js_array = f"const LDDT_AVERAGES = { lddt_averages }
305- alphafold_template = alphafold_template .replace (
293+ comparision_template = comparision_template .replace (
306294 "const LDDT_AVERAGES = [];" , averages_js_array
307295)
308296
309297i = 0
310298for structure in aligned_structures :
311- alphafold_template = alphafold_template .replace (
299+ comparision_template = comparision_template .replace (
312300 f"*_data_ranked_{ i } , open (structure , "r" ).read ().replace ("\n " , "\\ n" )
313301 )
314302 i += 1
@@ -318,11 +306,11 @@ def pdb_to_lddt(struct_files, generate_tsv):
318306 "r" ,
319307) as in_file :
320308 lddt_html = in_file .read ()
321- alphafold_template = alphafold_template .replace (
309+ comparision_template = comparision_template .replace (
322310 '<div id="lddt_placeholder"></div>' , lddt_html
323311 )
324312
325313with open (
326314 f"{ args .output_dir } { args .name } { args .in_type .lower ()} , "w"
327315) as out_file :
328- out_file .write (alphafold_template )
316+ out_file .write (comparision_template )
0 commit comments