wukevin commited on
Commit
1eff7fb
1 Parent(s): 54edb18

Add a molecule viewer

Browse files
Files changed (1) hide show
  1. app.py +63 -2
app.py CHANGED
@@ -5,6 +5,66 @@ import torch
5
  from foldingdiff import sampling
6
  from foldingdiff import angles_and_coords as ac
7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
  def sample_at_length(l:int, seed:int):
9
  """
10
  Sample a single structure at the given length
@@ -16,7 +76,7 @@ def sample_at_length(l:int, seed:int):
16
  outdir = os.path.join(os.getcwd(), "output")
17
  os.makedirs(outdir, exist_ok=True)
18
  pdb_file = ac.create_new_chain_nerf(os.path.join(outdir, "generated.pdb"), s)
19
- return s, pdb_file
20
 
21
  interface = gr.Interface(
22
  fn=sample_at_length,
@@ -25,8 +85,9 @@ interface = gr.Interface(
25
  gr.Number(value=42, label="Random seed", show_label=True, precision=0),
26
  ],
27
  outputs=[
28
- gr.Dataframe(label="Generated angles defining structure", max_rows=8),
29
  gr.File(label="Generated structure in PDB format (cartesian coordinates)"),
 
30
  ],
31
  )
32
  interface.launch()
 
5
  from foldingdiff import sampling
6
  from foldingdiff import angles_and_coords as ac
7
 
8
+ def read_mol(molpath: str) -> str:
9
+ with open(molpath, "r") as fp:
10
+ lines = fp.readlines()
11
+ mol = ""
12
+ for l in lines:
13
+ mol += l
14
+ return mol
15
+
16
+ def molecule(input_pdb: str) -> str:
17
+ """Get the string to view the given pdb in 3dmol.js"""
18
+ mol = read_mol(input_pdb)
19
+
20
+ x = (
21
+ """<!DOCTYPE html>
22
+ <html>
23
+ <head>
24
+ <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
25
+ <style>
26
+ body{
27
+ font-family:sans-serif
28
+ }
29
+ .mol-container {
30
+ width: 100%;
31
+ height: 600px;
32
+ position: relative;
33
+ }
34
+ .mol-container select{
35
+ background-image:None;
36
+ }
37
+ </style>
38
+ <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
39
+ </head>
40
+ <body>
41
+ <div id="container" class="mol-container"></div>
42
+
43
+ <script>
44
+ let pdb = `"""
45
+ + mol
46
+ + """`
47
+
48
+ $(document).ready(function () {
49
+ let element = $("#container");
50
+ let config = { backgroundColor: "black" };
51
+ let viewer = $3Dmol.createViewer(element, config);
52
+ viewer.addModel(pdb, "pdb");
53
+ viewer.getModel(0).setStyle({}, { stick: { colorscheme:"whiteCarbon" } });
54
+ viewer.zoomTo();
55
+ viewer.render();
56
+ viewer.zoom(0.8, 2000);
57
+ })
58
+ </script>
59
+ </body></html>"""
60
+ )
61
+
62
+ return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera;
63
+ display-capture; encrypted-media;" sandbox="allow-modals allow-forms
64
+ allow-scripts allow-same-origin allow-popups
65
+ allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
66
+ allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
67
+
68
  def sample_at_length(l:int, seed:int):
69
  """
70
  Sample a single structure at the given length
 
76
  outdir = os.path.join(os.getcwd(), "output")
77
  os.makedirs(outdir, exist_ok=True)
78
  pdb_file = ac.create_new_chain_nerf(os.path.join(outdir, "generated.pdb"), s)
79
+ return molecule(s), pdb_file, s
80
 
81
  interface = gr.Interface(
82
  fn=sample_at_length,
 
85
  gr.Number(value=42, label="Random seed", show_label=True, precision=0),
86
  ],
87
  outputs=[
88
+ gr.HTML(),
89
  gr.File(label="Generated structure in PDB format (cartesian coordinates)"),
90
+ gr.Dataframe(label="Generated angles defining structure", max_rows=8),
91
  ],
92
  )
93
  interface.launch()