fabiannagel / schnax Goto Github PK

It appears that the jax-md standard is neighbor, not neighbors. asax uses a keyword argument for it for some reason, so naming matters.

Is there any particular reason we're returning the energy as (1,) shape array instead of scalar? I assume this is easier for testing, but it makes implementing forces awkward. Can I change this without causing too much trouble?

Many tests currently contain lines like:

import tests.test_utils.initialize as init

which cause my test runner (nosetests) to fail with errors like:

ModuleNotFoundError: No module named 'tests'

My intention was for tests to be run by executing the test runner from within the tests/ directory. This is why the assets are addressed by relative paths, and why originally, imports directly from test_utils were used. Can we revert to this approach?

We want to make sure jitted code performs as expected. chex variants are a way to test this.

A schnax/SchNetPack model consists out of two main components:

1. A SchNet representation block (embeddings, distance expansions, interactions)
2. An atomwise output block (simple 2-layer MLP, aggregation)

This structure could be made clearer in the current implementation. Potential benefits:

• Cleaner separation between representation and atomwise output configurations (see #3).
• Usage of non-standard output blocks.
• Prep for merging the SchNet representation to other repos.

Annoyingly, this would probably break weight mapping from pytorch to haiku.

Is it possible to automatically perform the haiku init call so we can return a clean neighbor_fn, energy_fn tuple in line with the usual jax-md approach? I would guess it's complicated because neighbors needs to be already initialised with something for this to work? Can we stub out out neighbors for this case with just a namedtuple?

@hk.transform_with_state is only needed for testing to extract intermediaries. It'd be nice to have that functionality for testing without bloating production code by having to manage useless states.

Fix weird GitHub CI issue so we can use Colab (runs 3.7)

We should have a helper that accepts a path to a schnetkit model file and returns energy_fn, init_fn, apply_fn.

In principle, multiple interaction blocks are supported, but in practice, that is not exposed. We need to:

• Amend the SchNet class to have a configurable number of interaction blocks
• Amend the get_params method from the schnetkit helper to read the number of interactions from the spec dict and Do The Right Thing

Anything else?