Spaces:
Sleeping
Sleeping
igashov
commited on
Commit
•
fffb501
1
Parent(s):
c438a2a
update COM
Browse files
app.py
CHANGED
|
@@ -154,6 +154,7 @@ def generate(input_file, n_steps):
|
|
| 154 |
print('Created dataloader')
|
| 155 |
|
| 156 |
ddpm.edm.T = n_steps
|
|
|
|
| 157 |
|
| 158 |
for data in dataloader:
|
| 159 |
chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
|
|
@@ -161,6 +162,7 @@ def generate(input_file, n_steps):
|
|
| 161 |
x = chain[0][:, :, :ddpm.n_dims]
|
| 162 |
h = chain[0][:, :, ddpm.n_dims:]
|
| 163 |
|
|
|
|
| 164 |
pos_masked = data['positions'] * data['fragment_mask']
|
| 165 |
N = data['fragment_mask'].sum(1, keepdims=True)
|
| 166 |
mean = torch.sum(pos_masked, dim=1, keepdim=True) / N
|
|
|
|
| 154 |
print('Created dataloader')
|
| 155 |
|
| 156 |
ddpm.edm.T = n_steps
|
| 157 |
+
assert ddpm.center_of_mass == 'fragments'
|
| 158 |
|
| 159 |
for data in dataloader:
|
| 160 |
chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
|
|
|
|
| 162 |
x = chain[0][:, :, :ddpm.n_dims]
|
| 163 |
h = chain[0][:, :, ddpm.n_dims:]
|
| 164 |
|
| 165 |
+
# Put the molecule back to the initial orientation
|
| 166 |
pos_masked = data['positions'] * data['fragment_mask']
|
| 167 |
N = data['fragment_mask'].sum(1, keepdims=True)
|
| 168 |
mean = torch.sum(pos_masked, dim=1, keepdim=True) / N
|
output.py
CHANGED
|
@@ -62,6 +62,7 @@ SAMPLES_RENDERING_TEMPLATE = """<!DOCTYPE html>
|
|
| 62 |
$(document).ready(function() {{
|
| 63 |
viewer.addModel(`{fragments}`, "{fragments_fmt}")
|
| 64 |
viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
|
|
|
|
| 65 |
viewer.addModel(`{molecule}`, "{molecule_fmt}")
|
| 66 |
viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
|
| 67 |
viewer.zoomTo();
|
|
|
|
| 62 |
$(document).ready(function() {{
|
| 63 |
viewer.addModel(`{fragments}`, "{fragments_fmt}")
|
| 64 |
viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
|
| 65 |
+
viewer.getModel().hide();
|
| 66 |
viewer.addModel(`{molecule}`, "{molecule_fmt}")
|
| 67 |
viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
|
| 68 |
viewer.zoomTo();
|