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")] })