this had better be a clean notebook or else im never coding again

[brookehus]

OMG. See commit message.

Must be too tired to rebase correctly. Anyway:

Check this out! I would appreciate a bit more embellishing on the descriptions. I'm ok if you PR directly to this branch for descriptions, but would prefer to discuss changes you want to make to the code.

Note that you probably don't want to "run all". I have two places where you can uncomment to make the dataset and simulation data smaller.

It would be good to switch device to GPU and also test this out.

@Dom1L I know you're busy so no pressure to review! Just adding you b/c you did a lot of the work.

[nec4]

Will take a look!

[nec4]

Unlike the previous example, we not only need coordinates and forces, 
but also _embeddings_. Here we're going to go with the SchNet convention
[2] of embedding the atomic number of the element corresponding to the
bead. Since our beads are the backbone C-N-C-C-N, we embed the properties
(6, 7, 6, 6, 7). We'll create the embeddings ourselves as follows:

I think we should explain a little more about embeddings. For the uninitiated, they might be initially confusing. What about something like this:

In neural networks, embeddings provide a way of encoding inputs into some
property space. That is, a single input example can be mapped to a corresponding
property vector. After training, the embedding weights can be inspected and
processed in order to visualize the inputs in this property space. 

Or some variation thereof. What do you guys think?

[nec4]

We specify all the hyperparameters now, since we'll need some of them
to build our `SchnetFeature`.

We should also mention the switch from Tanh() to ShiftedSoftplus() in order to follow [2].

[nec4]

We use the `FeatureCombiner` to stack our `GeometryFeature` and our
`SchnetFeature`. We just need to tell our `FeatureCombiner` which integer
indices correspond to distance features. The `FeatureCombiner` is ultimately
the feature we will feed into our `CGnet`.

Would it makes sense to explicitly mention the reason why we need the distances is because they serve as the only geometrical feature inputs into the SchnetFeature()?

[nec4]

It would be especially cool if we could visualize the embedding weights, as this can give a visualization of the the inputs in our property space. Here it is nice because we only have two properties so the embedding does not require further dimensionality reduction to visualize. This will also illustrate the use of the embeddings, as all we have so far is just a means of learning the CG potential and running a subsequent simulation. What do you guys think?

[nec4]

To simulate a system with a `SchnetFeature`, we need to provide the
`Simulation` with embeddings.

Simulations with `SchnetFeature`s tend to be _slow_. If you haven't
switched your device to a GPU, we recommend the second set of
parameters for a meaningless result but quickly executable code.

I think we should expand upon this (possibly in the other notebook as well):

If the free energy from simulations is to be used as a heuristic for CG
model success, the simulations must equilibrate. Short simulations, even
if many are used in parallel, tend to strongly bias the CG free energy to the
high resolution reference system, and therefore can give a deceptive and
potentially wrong representation of the equilibrated CG free energy.

What do you guys think?

Also: we should mention how to switch a model to GPU (using model.mount()) after it has been loaded.

[brookehus]

@nec4 I support all of your changes! I would be happy for you to commit them directly.

[Dom1L]

I also support all of @nec4 changes!

We had a discussion about the visualization of the property embedding. I think in the case of only two different species it doesn't make too much sense, because it will just be two points.

What we could try is to differentiate between the carbon atoms and give for example the C-alpha a different one. Then we have 3 points and maybe see that the carbons are closer in embedding space than the nitrogen.

Otherwise, I would already find it good just to show that you can visualize the embedding somehow to extract some meaning

[nec4]

Also: I don't think we are using the zscore stuff here - should we get rid of it?

[brookehus]

@nec4 all your changes look great! yes, we should remove the zscores - i'll just remove that box now. @Dom1L should we add something about how schnet uses batchnorm but we haven't implemented it?

[brookehus]

Has anyone tried to run this on a GPU?

[Dom1L]

@Dom1L should we add something about how schnet uses batchnorm but we haven't implemented it?

Not sure, I would wait with this until we tried to mimic the Schnet standardization ourselves with batchnorm or so. It is also not explicitly mentioned in the publication, so it might just cause more confusion for now.

Has anyone tried to run this on a GPU?

I did, it works up to the simulation tool. If you remember from the Simulation() class rewrite in the master branch, there is the case that if the tensor is on gpu, you need to do .cpu() before the .numpy(), otherwise this happens:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-76cdac63a925> in <module>
     12                  verbose=True, save_potential=True)
     13 
---> 14 traj = sim.simulate()

/shared/dominik/repos/cgnet/cgnet/network/utils.py in simulate(self, overwrite)
    290         for t in range(self.length):
    291             potential, forces = self.model(x_old, self.embeddings)
--> 292             potential = potential.detach().numpy()
    293             forces = forces.detach().numpy()
    294             noise = self.rng.randn(self.n_sims,

TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

The master branch already has the fix for it, so might be cool to just incorporate it somehow. There are also 2 .to(device) calls missing, but I'm just going to add them now and commit it.

[nec4]

Would it be okay to merge master into this branch for the purpose of updating the simulation tools? I can do it locally first to make sure that nothing important breaks.

[brookehus]

Nb. we know this doesn't work and are merging anyway