Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove jax-md and replace with glp #13

Merged
merged 1 commit into from
Aug 12, 2023
Merged

Remove jax-md and replace with glp #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions mlff/mdx/atoms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,6 @@

from mlff.utils import Graph, Neighbors, System

from jax_md.space import periodic_general, free


SpatialPartitioning = namedtuple(
"SpatialPartitioning", ("allocate_fn", "update_fn", "cutoff", "skin", "capacity_multiplier")
Expand Down Expand Up @@ -512,7 +510,9 @@ def _update(x: AtomsX, neighbors):


def to_displacement(atoms):
if atoms.get_cell() is not None:
return periodic_general(atoms.get_cell(), fractional_coordinates=False)[0]
else:
return free()[0]
from glp.periodic import make_displacement

displacement = make_displacement(atoms.cell)

# reverse sign convention for backwards compatibility
return lambda Ra, Rb: raw_disp(Rb, Ra)
15 changes: 7 additions & 8 deletions mlff/mdx/calculator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,8 @@

from flax import struct

from jax_md.space import periodic_general, free

from mlff.utils import Graph


logging.basicConfig(level=logging.INFO)

StackNet = Any
Expand Down Expand Up @@ -136,8 +133,10 @@ def atomsx_to_graph(atoms: Any):
return Graph(edges, nodes, neighbors['idx_j'], neighbors['idx_i'], mask)


def to_displacement(atoms: Any):
if atoms.get_cell() is not None:
return periodic_general(atoms.get_cell(), fractional_coordinates=False)[0]
else:
return free()[0]
def to_displacement(atoms):
from glp.periodic import make_displacement

displacement = make_displacement(atoms.cell)

# reverse sign convention for backwards compatibility
return lambda Ra, Rb: raw_disp(Rb, Ra)