def log_structure_and_angs(args, pred_ang, pred_coords, true_coords, src_seq, commit, log_angs=True, struct_name="train"): """ Logs a 3D structure prediction to wandb. """ if log_angs: log_angle_distributions(args, pred_ang, src_seq) src_seq_cpu = src_seq.cpu().detach().numpy() # Make dir if needed cur_struct_path = os.path.join(args.structure_dir, struct_name) os.makedirs(cur_struct_path, exist_ok=True) # Remove coordinate level padding (each residue has about 13 atoms, # even if some are missing) gold_item_non_batch_pad = (true_coords != VOCAB.pad_id).any(dim=-1) true_coords = true_coords[gold_item_non_batch_pad] true_coords[torch.isnan(true_coords)] = 0 creator = PDB_Creator(pred_coords.detach().numpy(), seq=VOCAB.ints2str(src_seq_cpu)) creator.save_pdb(f"{cur_struct_path}/{wandb.run.step:05}_pred.pdb", title="pred") t_creator = PDB_Creator(true_coords.cpu().detach().numpy(), seq=VOCAB.ints2str(src_seq_cpu)) if not os.path.isfile(f"{cur_struct_path}/true.pdb") or struct_name == "train": t_creator.save_pdb(f"{cur_struct_path}/true.pdb", title="true") wandb.log({f"{struct_name}_mol_true" : wandb.Molecule(f"{cur_struct_path}/true.pdb")}, commit=False) gltf_out_path = os.path.join(args.gltf_dir, f"{wandb.run.step:05}_{struct_name}.gltf") t_creator.save_gltfs(f"{cur_struct_path}/true.pdb", f"{cur_struct_path}/{wandb.run.step:05}_pred.pdb", gltf_out_path=gltf_out_path, make_pse=True, make_png=args.save_pngs, pse_out_path=f"{cur_struct_path}/{wandb.run.step:05}_both.pse") log_items = {struct_name: wandb.Object3D(gltf_out_path), f"{struct_name}_mol": wandb.Molecule(f"{cur_struct_path}/{wandb.run.step:05}_pred.pdb"), f"{struct_name}_mol_comb": wandb.Molecule(f"{cur_struct_path}/{wandb.run.step:05}_both.pdb")} if args.save_pngs: try: log_items[struct_name + "_img"] = wandb.Image(gltf_out_path.replace("gltf", "png")) except FileNotFoundError: # Account for the possibility that a PyMol session may have failed to create successfully pass wandb.log(log_items, commit=commit)
def test_molecule_file(runner, mocked_run): with runner.isolated_filesystem(): with open("test.pdb", "w") as f: f.write("00000") mol = wandb.Molecule(open("test.pdb", "r")) mol.bind_to_run(mocked_run, "rad", "summary") wandb.Molecule.seq_to_json([mol], mocked_run, "rad", "summary") assert os.path.exists(mol._path)
def test_molecule(mocked_run): with open("test.pdb", "w") as f: f.write("00000") mol = wandb.Molecule("test.pdb") mol.bind_to_run(mocked_run, "rad", "summary") wandb.Molecule.seq_to_json([mol], mocked_run, "rad", "summary") assert os.path.exists(mol._path) wandb.finish()
import wandb wandb.init(project="Corona-Virus") wandb.log({"corona": wandb.Molecule(open("test-files/corona.pdb")) }) wandb.log({"5r84": wandb.Molecule(open("test-files/5r84.pdb"), caption="PanDDA analysis group deposition -- Crystal Structure of COVID-19 main protease in complex with Z31792168" ) }) wandb.log({"6lu7": wandb.Molecule(open("test-files/6lu7.pdb"), caption="The crystal structure of COVID-19 main protease in complex with an inhibitor N3") }) wandb.log({"6vw1": wandb.Molecule(open("test-files/6vw1.pdb"), caption="Structure of 2019-nCoV chimeric receptor-binding domain complexed with its receptor human ACE2") }) wandb.log({"corona_many": [wandb.Molecule(open("test-files/corona.pdb"), caption="molecule")] })