fabiannagel / schnax Goto Github PK
View Code? Open in Web Editor NEWAn implementation of SchNet in JAX and JAX-MD.
An implementation of SchNet in JAX and JAX-MD.
@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.
It appears that the jax-md
standard is neighbor
, not neighbors
. asax
uses a keyword argument for it for some reason, so naming matters.
A schnax
/SchNetPack
model consists out of two main components:
SchNet
representation block (embeddings, distance expansions, interactions)This structure could be made clearer in the current implementation. Potential benefits:
SchNet
representation to other repos.Annoyingly, this would probably break weight mapping from pytorch
to haiku
.
Fix weird GitHub CI issue so we can use Colab (runs 3.7)
In principle, multiple interaction blocks are supported, but in practice, that is not exposed. We need to:
SchNet
class to have a configurable number of interaction blocksget_params
method from the schnetkit
helper to read the number of interactions from the spec
dict and Do The Right ThingAnything else?
We should have a helper that accepts a path to a schnetkit
model file and returns energy_fn, init_fn, apply_fn
.
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
?
We want to make sure jitted code performs as expected. chex
variants are a way to test this.
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?
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?
A declarative, efficient, and flexible JavaScript library for building user interfaces.
๐ Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
An Open Source Machine Learning Framework for Everyone
The Web framework for perfectionists with deadlines.
A PHP framework for web artisans
Bring data to life with SVG, Canvas and HTML. ๐๐๐
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
Some thing interesting about web. New door for the world.
A server is a program made to process requests and deliver data to clients.
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
Some thing interesting about visualization, use data art
Some thing interesting about game, make everyone happy.
We are working to build community through open source technology. NB: members must have two-factor auth.
Open source projects and samples from Microsoft.
Google โค๏ธ Open Source for everyone.
Alibaba Open Source for everyone
Data-Driven Documents codes.
China tencent open source team.