change bfactor

app.py

@@ -196,6 +196,21 @@ def prepare_input(input, jobname, baseconfig, hard_case):
     
     return yaml.dump(yaml_dict).replace("'", "\""),os.path.join("/tmp/", f"{jobname}.zip")
 
+
+def convert_bfactors(pdb_path):
+    with open(pdb_path, 'r') as f:
+        lines = f.readlines()
+    for i,line in enumerate(lines):
+        # multiple each bfactor by 100
+        if line[0:6] == 'ATOM  ' or line[0:6] == 'HETATM':
+            bfactor = float(line[60:66])
+            bfactor *= 100
+            line = line[:60] + f'{bfactor:6.2f}' + line[66:]
+            lines[i] = line
+    with open(pdb_path.replace(".pdb", "_processed.pdb"), 'w') as f:
+        f.write(''.join(lines))
+
+
 def run_rf2aa(jobname, zip_archive):
     current_dir = os.getcwd()
     try:
@@ -203,20 +218,11 @@ def run_rf2aa(jobname, zip_archive):
             zip_ref.extractall(os.path.join(current_dir))
         os.system(f"python -m rf2aa.run_inference --config-name {jobname}.yaml --config-path {current_dir}/{jobname}")
         # scale pLDDT to 0-100 range in pdb output file
-        parser = PDBParser(QUIET=True)
-        structure = parser.get_structure(jobname, f"{current_dir}/{jobname}/{jobname}.pdb")
-        for model in structure:
-            for chain in model:
-                for residue in chain:
-                    for atom in residue:
-                        atom.bfactor = atom.bfactor * 100
-        io = PDBIO()
-        io.set_structure(structure)
-        io.save(f"{current_dir}/{jobname}/{jobname}.pdb")
+        convert_bfactors(f"{current_dir}/{jobname}/{jobname}.pdb")
 
     except Exception as e:
         raise gr.Error(f"Error running RFAA: {e}")
-    return f"{current_dir}/{jobname}/{jobname}.pdb"
+    return f"{current_dir}/{jobname}/{jobname}_processed.pdb"
     
 
 
@@ -287,7 +293,7 @@ with gr.Blocks() as demo:
     runfiles = gr.File(label="files to run RFAA", visible=False)
     instructions = gr.Markdown(visible=False)
 
-    out = Molecule3D(visible=False)
+      out = Molecule3D(visible=False)
 
     btn.click(predict, inputs=[inp, jobname, dry_run, base_config, hard_case], outputs=[config_file, runfiles, instructions, out])