# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for jax_md.space."""

from absl.testing import absltest
from absl.testing import parameterized

from jax.config import config as jax_config
from jax import random
import jax.numpy as jnp

from jax import grad, jit, jacfwd

from jax import test_util as jtu

from jax_md import space, test_util, quantity, energy
from jax_md.util import *

from functools import partial
from unittest import SkipTest

test_util.update_test_tolerance(5e-5, 5e-13)

jax_config.parse_flags_with_absl()
jax_config.enable_omnistaging()
FLAGS = jax_config.FLAGS

PARTICLE_COUNT = 10
STOCHASTIC_SAMPLES = 10
SHIFT_STEPS = 10
SPATIAL_DIMENSION = [2, 3]
BOX_FORMATS = ['scalar', 'vector', 'matrix']

if FLAGS.jax_enable_x64:
  POSITION_DTYPE = [f32, f64]
else:
  POSITION_DTYPE = [f32]


def make_periodic_general_test_system(N, dim, dtype, box_format):
  assert box_format in BOX_FORMATS

  box_size = quantity.box_size_at_number_density(N, 1.0, dim)
  box = dtype(box_size)
  if box_format == 'vector':
    box = jnp.array(jnp.ones(dim) * box_size, dtype)
  elif box_format == 'matrix':
    box = jnp.array(jnp.eye(dim) * box_size, dtype)

  d, s = space.periodic(jnp.diag(box) if box_format == 'matrix' else box)
  d_gf, s_gf = space.periodic_general(box)
  d_g, s_g = space.periodic_general(box, fractional_coordinates=False)

  key = random.PRNGKey(0)

  R_f = random.uniform(key, (N, dim), dtype=dtype)
  R = space.transform(box, R_f)

  E = jit(energy.soft_sphere_pair(d))
  E_gf = jit(energy.soft_sphere_pair(d_gf))
  E_g = jit(energy.soft_sphere_pair(d_g))

  return R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g)


# pylint: disable=invalid-name
class SpaceTest(jtu.JaxTestCase):

  # pylint: disable=g-complex-comprehension
  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_transform(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2 = random.split(key, 3)

      R = random.normal(
        split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      T = random.normal(
        split2, (spatial_dimension, spatial_dimension), dtype=dtype)

      R_prime_exact = jnp.array(jnp.einsum('ij,kj->ki', T, R), dtype=dtype)
      R_prime = space.transform(T, R)

      self.assertAllClose(R_prime_exact, R_prime)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}'.format(dim),
          'spatial_dimension': dim
      } for dim in SPATIAL_DIMENSION))
  def test_transform_grad(self, spatial_dimension):
    key = random.PRNGKey(0)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2 = random.split(key, 3)

      R = random.normal(split1, (PARTICLE_COUNT, spatial_dimension))
      T = random.normal(split2, (spatial_dimension, spatial_dimension))

      R_prime = space.transform(T, R)

      energy_direct = lambda R: jnp.sum(R ** 2)
      energy_indirect = lambda T, R: jnp.sum(space.transform(T, R) ** 2)

      grad_direct = grad(energy_direct)(R_prime)
      grad_indirect = grad(energy_indirect, 1)(T, R)

      self.assertAllClose(grad_direct, grad_indirect)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_transform_inverse(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    tol = 1e-13
    if dtype is f32:
      tol = 1e-5

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2 = random.split(key, 3)

      R = random.normal(
        split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      T = random.normal(
        split2, (spatial_dimension, spatial_dimension), dtype=dtype)
      T_inv = space.inverse(T)

      R_test = space.transform(T_inv, space.transform(T, R))

      self.assertAllClose(R, R_test)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_canonicalize_displacement_or_metric(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    displacement, _ = space.periodic_general(jnp.eye(spatial_dimension))
    metric = space.metric(displacement)
    test_metric = space.canonicalize_displacement_or_metric(displacement)

    metric = space.map_product(metric)
    test_metric = space.map_product(test_metric)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2 = random.split(key, 3)

      R = random.normal(
        split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      self.assertAllClose(metric(R, R), test_metric(R, R))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_displacement(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split = random.split(key)

      R = random.uniform(
        split, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      dR = space.map_product(space.pairwise_displacement)(R, R)

      dR_wrapped = space.periodic_displacement(f32(1.0), dR)

      dR_direct = dR
      dr_direct = space.distance(dR)
      dr_direct = jnp.reshape(dr_direct, dr_direct.shape + (1,))

      if spatial_dimension == 2:
        for i in range(-1, 2):
          for j in range(-1, 2):
            dR_shifted = dR + jnp.array([i, j], dtype=R.dtype)

            dr_shifted = space.distance(dR_shifted)
            dr_shifted = jnp.reshape(dr_shifted, dr_shifted.shape + (1,))

            dR_direct = jnp.where(dr_shifted < dr_direct, dR_shifted, dR_direct)
            dr_direct = jnp.where(dr_shifted < dr_direct, dr_shifted, dr_direct)
      elif spatial_dimension == 3:
        for i in range(-1, 2):
          for j in range(-1, 2):
            for k in range(-1, 2):
              dR_shifted = dR + jnp.array([i, j, k], dtype=R.dtype)

              dr_shifted = space.distance(dR_shifted)
              dr_shifted = jnp.reshape(dr_shifted, dr_shifted.shape + (1,))

              dR_direct = jnp.where(
                  dr_shifted < dr_direct, dR_shifted, dR_direct)
              dr_direct = jnp.where(
                  dr_shifted < dr_direct, dr_shifted, dr_direct)

      dR_direct = jnp.array(dR_direct, dtype=dR.dtype)
      assert dR_wrapped.dtype == dtype
      self.assertAllClose(dR_wrapped, dR_direct)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_shift(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2 = random.split(key, 3)

      R = random.uniform(
        split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      dR = jnp.sqrt(f32(0.1)) * random.normal(
          split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      dR = jnp.where(dR > 0.49, f32(0.49), dR)
      dR = jnp.where(dR < -0.49, f32(-0.49), dR)

      R_shift = space.periodic_shift(f32(1.0), R, dR)

      assert R_shift.dtype == R.dtype
      assert jnp.all(R_shift < 1.0)
      assert jnp.all(R_shift > 0.0)

      dR_after = space.periodic_displacement(f32(1.0), R_shift - R)

      assert dR_after.dtype == R.dtype
      self.assertAllClose(dR_after, dR)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_against_periodic_general(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    tol = 1e-13
    if dtype is f32:
      tol = 1e-5

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2, split3 = random.split(key, 4)

      max_box_size = f32(10.0)
      box_size = max_box_size * random.uniform(
        split1, (spatial_dimension,), dtype=dtype)
      transform = jnp.diag(box_size)

      R = random.uniform(
        split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      R_scaled = R * box_size

      dR = random.normal(
        split3, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      disp_fn, shift_fn = space.periodic(box_size)
      general_disp_fn, general_shift_fn = space.periodic_general(transform)

      disp_fn = space.map_product(disp_fn)
      general_disp_fn = space.map_product(general_disp_fn)

      self.assertAllClose(disp_fn(R_scaled, R_scaled), general_disp_fn(R, R))
      assert disp_fn(R_scaled, R_scaled).dtype == dtype
      self.assertAllClose(
          shift_fn(R_scaled, dR), general_shift_fn(R, dR) * box_size)
      assert shift_fn(R_scaled, dR).dtype == dtype

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_against_periodic_general_grad(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    tol = 1e-13
    if dtype is f32:
      tol = 1e-5

    for _ in range(STOCHASTIC_SAMPLES):
      key, split1, split2, split3 = random.split(key, 4)

      max_box_size = f32(10.0)
      box_size = max_box_size * random.uniform(
        split1, (spatial_dimension,), dtype=dtype)
      transform = jnp.diag(box_size)

      R = random.uniform(
        split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      R_scaled = R * box_size

      dR = random.normal(
        split3, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      disp_fn, shift_fn = space.periodic(box_size)
      general_disp_fn, general_shift_fn = space.periodic_general(transform)

      disp_fn = space.map_product(disp_fn)
      general_disp_fn = space.map_product(general_disp_fn)

      grad_fn = grad(lambda R: jnp.sum(disp_fn(R, R) ** 2)) 
      general_grad_fn = grad(lambda R: jnp.sum(general_disp_fn(R, R) ** 2)) 

      self.assertAllClose(grad_fn(R_scaled), general_grad_fn(R))
      assert general_grad_fn(R).dtype == dtype

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype,
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_general_dynamic(self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    eye = jnp.eye(spatial_dimension)

    for _ in range(STOCHASTIC_SAMPLES):
      key, split_T0_scale, split_T0_dT = random.split(key, 3)
      key, split_T1_scale, split_T1_dT = random.split(key, 3)
      key, split_t, split_R, split_dR = random.split(key, 4)

      size_0 = 10.0 * random.uniform(split_T0_scale, ())
      dtransform_0 = 0.5 * random.normal(
        split_T0_dT, (spatial_dimension, spatial_dimension))
      T_0 = jnp.array(size_0 * (eye + dtransform_0), dtype=dtype)

      size_1 = 10.0 * random.uniform(split_T1_scale, (), dtype=dtype)
      dtransform_1 = 0.5 * random.normal(
          split_T1_dT, (spatial_dimension, spatial_dimension), dtype=dtype)
      T_1 = jnp.array(size_1 * (eye + dtransform_1), dtype=dtype)

      disp_fn, shift_fn = space.periodic_general(T_0)
      true_disp_fn, true_shift_fn = space.periodic_general(T_1)

      disp_fn = partial(disp_fn, box=T_1)

      disp_fn = space.map_product(disp_fn)
      true_disp_fn = space.map_product(true_disp_fn)

      R = random.uniform(
        split_R, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      dR = random.normal(
        split_dR, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      self.assertAllClose(
        disp_fn(R, R), jnp.array(true_disp_fn(R, R), dtype=dtype))
      self.assertAllClose(
        shift_fn(R, dR, box=T_1), jnp.array(true_shift_fn(R, dR), dtype=dtype))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': '_dim={}_dtype={}'.format(dim, dtype.__name__),
          'spatial_dimension': dim,
          'dtype': dtype,
      } for dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE))
  def test_periodic_general_wrapped_vs_unwrapped(
      self, spatial_dimension, dtype):
    key = random.PRNGKey(0)

    eye = jnp.eye(spatial_dimension, dtype=dtype)

    tol = 1e-13
    if dtype is f32:
      tol = 2e-5

    for _ in range(STOCHASTIC_SAMPLES):
      key, split_R, split_T = random.split(key, 3)

      dT = random.normal(
        split_T, (spatial_dimension, spatial_dimension), dtype=dtype)
      T = eye + dT + jnp.transpose(dT)

      R = random.uniform(
        split_R, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      R0 = R
      unwrapped_R = R

      displacement, shift = space.periodic_general(T)
      _, unwrapped_shift = space.periodic_general(T, wrapped=False)

      displacement = space.map_product(displacement)

      for _ in range(SHIFT_STEPS):
        key, split = random.split(key)
        dR = random.normal(
          split, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
        R = shift(R, dR)
        unwrapped_R = unwrapped_shift(unwrapped_R, dR)
        self.assertAllClose(
          displacement(R, R0),
          displacement(unwrapped_R, R0))
      assert not (jnp.all(unwrapped_R > 0) and jnp.all(unwrapped_R < 1))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_energy(self, spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)
    self.assertAllClose(E(R), E_gf(R_f))
    self.assertAllClose(E(R), E_g(R))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_force(self, spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)
    self.assertAllClose(grad(E)(R), grad(E_gf)(R_f))
    self.assertAllClose(grad(E)(R), grad(E_g)(R))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_shift(self, spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)

    R_new = s(R, grad(E)(R))
    R_gf_new = s_gf(R_f, grad(E_gf)(R_f))
    R_g_new = s_g(R, grad(E_g)(R))

    self.assertAllClose(R_new, space.transform(box, R_gf_new))
    self.assertAllClose(R_new, R_g_new)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_deform(self, spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)
    deformed_box = box * 0.9
    self.assertAllClose(E_gf(R_f, box=deformed_box),
                        E_g(R, new_box=deformed_box))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_deform_grad(self,
                                        spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)
    deformed_box = box * 0.9
    self.assertAllClose(grad(E_gf)(R_f, box=deformed_box),
                        grad(E_g)(R, new_box=deformed_box))

    self.assertAllClose(jacfwd(E_gf)(R_f, box=deformed_box),
                        jacfwd(E_g)(R, new_box=deformed_box))

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_deform_shift(self,
                                        spatial_dimension, dtype, box_format):
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)
    deformed_box = box * 0.9

    R_new = s_g(R, grad(E_g)(R), new_box=deformed_box)
    R_gf_new = space.transform(deformed_box, s_gf(R_f, grad(E_gf)(R_f)))

    self.assertAllClose(R_new, R_gf_new)

  @parameterized.named_parameters(jtu.cases_from_list(
      {
          'testcase_name': f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',
          'spatial_dimension': dim,
          'dtype': dtype,
          'box_format': box_format
      } for dim in SPATIAL_DIMENSION
    for dtype in POSITION_DTYPE
    for box_format in BOX_FORMATS))
  def test_periodic_general_grad_box(self, spatial_dimension, dtype, box_format):
    if box_format == 'scalar':
      raise SkipTest('Scalar case fails due to JAX Issue #5849.')
    N = 16
    R_f, R, box, (s, E), (s_gf, E_gf), (s_g, E_g) = \
      make_periodic_general_test_system(N, spatial_dimension, dtype, box_format)

    @grad
    def box_energy_g_fn(box):
      return E_g(R, new_box=box)

    @grad
    def box_energy_gf_fn(box):
      return E_gf(R_f, box=box)

    self.assertAllClose(box_energy_g_fn(box), box_energy_gf_fn(box))
if __name__ == '__main__':
  absltest.main()

[(34, 'jax_md.test_util.update_test_tolerance', 'test_util.update_test_tolerance', False, 'from jax_md import space, test_util, quantity, energy\n'), (36, 'jax.config.config.parse_flags_with_absl', 'jax_config.parse_flags_with_absl', True, 'from jax.config import config as jax_config\n'), (37, 'jax.config.config.enable_omnistaging', 'jax_config.enable_omnistaging', True, 'from jax.config import config as jax_config\n'), (55, 'jax_md.quantity.box_size_at_number_density', 'quantity.box_size_at_number_density', False, 'from jax_md import space, test_util, quantity, energy\n'), (63, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (64, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (66, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (68, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (69, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (566, 'absl.testing.absltest.main', 'absltest.main', False, 'from absl.testing import absltest\n'), (71, 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', False, 'from jax_md import space, test_util, quantity, energy\n'), (72, 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', False, 'from jax_md import space, test_util, quantity, energy\n'), (73, 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', False, 'from jax_md import space, test_util, quantity, energy\n'), (89, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (110, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (135, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (162, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (165, 'jax_md.space.metric', 'space.metric', False, 'from jax_md import space, test_util, quantity, energy\n'), (166, 'jax_md.space.canonicalize_displacement_or_metric', 'space.canonicalize_displacement_or_metric', False, 'from jax_md import space, test_util, quantity, energy\n'), (168, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (169, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (186, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (236, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (267, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (307, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (347, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (349, 'jax.numpy.eye', 'jnp.eye', True, 'import jax.numpy as jnp\n'), (392, 'jax.random.PRNGKey', 'random.PRNGKey', False, 'from jax import random\n'), (394, 'jax.numpy.eye', 'jnp.eye', True, 'import jax.numpy as jnp\n'), (428, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (444, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (460, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (481, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (498, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (519, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (540, 'jax.test_util.cases_from_list', 'jtu.cases_from_list', True, 'from jax import test_util as jtu\n'), (62, 'jax.numpy.diag', 'jnp.diag', True, 'import jax.numpy as jnp\n'), (92, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (94, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (96, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (100, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (113, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (115, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (116, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (118, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (142, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (144, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (147, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (149, 'jax_md.space.inverse', 'space.inverse', False, 'from jax_md import space, test_util, quantity, energy\n'), (164, 'jax.numpy.eye', 'jnp.eye', True, 'import jax.numpy as jnp\n'), (172, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (174, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (189, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (191, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (198, 'jax_md.space.distance', 'space.distance', False, 'from jax_md import space, test_util, quantity, energy\n'), (199, 'jax.numpy.reshape', 'jnp.reshape', True, 'import jax.numpy as jnp\n'), (225, 'jax.numpy.array', 'jnp.array', True, 'import jax.numpy as jnp\n'), (239, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (241, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (252, 'jax.numpy.all', 'jnp.all', True, 'import jax.numpy as jnp\n'), (253, 'jax.numpy.all', 'jnp.all', True, 'import jax.numpy as jnp\n'), (274, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (279, 'jax.numpy.diag', 'jnp.diag', True, 'import jax.numpy as jnp\n'), (281, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (285, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (288, 'jax_md.space.periodic', 'space.periodic', False, 'from jax_md import space, test_util, quantity, energy\n'), (289, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (291, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (292, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (314, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (319, 'jax.numpy.diag', 'jnp.diag', True, 'import jax.numpy as jnp\n'), (321, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (325, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (328, 'jax_md.space.periodic', 'space.periodic', False, 'from jax_md import space, test_util, quantity, energy\n'), (329, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (331, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (332, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (352, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (353, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (354, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (359, 'jax.numpy.array', 'jnp.array', True, 'import jax.numpy as jnp\n'), (364, 'jax.numpy.array', 'jnp.array', True, 'import jax.numpy as jnp\n'), (366, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (367, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (369, 'functools.partial', 'partial', False, 'from functools import partial\n'), (371, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (372, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (374, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (376, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (401, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (403, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (407, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (412, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (413, 'jax_md.space.periodic_general', 'space.periodic_general', False, 'from jax_md import space, test_util, quantity, energy\n'), (415, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (478, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (551, 'unittest.SkipTest', 'SkipTest', False, 'from unittest import SkipTest\n'), (58, 'jax.numpy.ones', 'jnp.ones', True, 'import jax.numpy as jnp\n'), (99, 'jax.numpy.einsum', 'jnp.einsum', True, 'import jax.numpy as jnp\n'), (120, 'jax.numpy.sum', 'jnp.sum', True, 'import jax.numpy as jnp\n'), (123, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (124, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (151, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (193, 'jax_md.space.map_product', 'space.map_product', False, 'from jax_md import space, test_util, quantity, energy\n'), (243, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (277, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (317, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (356, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (357, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (361, 'jax.random.uniform', 'random.uniform', False, 'from jax import random\n'), (362, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (405, 'jax.numpy.transpose', 'jnp.transpose', True, 'import jax.numpy as jnp\n'), (418, 'jax.random.split', 'random.split', False, 'from jax import random\n'), (419, 'jax.random.normal', 'random.normal', False, 'from jax import random\n'), (457, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (457, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (458, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (458, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (474, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (475, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (476, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (513, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (514, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (516, 'jax.jacfwd', 'jacfwd', False, 'from jax import grad, jit, jacfwd\n'), (517, 'jax.jacfwd', 'jacfwd', False, 'from jax import grad, jit, jacfwd\n'), (535, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (60, 'jax.numpy.eye', 'jnp.eye', True, 'import jax.numpy as jnp\n'), (426, 'jax.numpy.all', 'jnp.all', True, 'import jax.numpy as jnp\n'), (426, 'jax.numpy.all', 'jnp.all', True, 'import jax.numpy as jnp\n'), (536, 'jax.grad', 'grad', False, 'from jax import grad, jit, jacfwd\n'), (121, 'jax_md.space.transform', 'space.transform', False, 'from jax_md import space, test_util, quantity, energy\n'), (206, 'jax_md.space.distance', 'space.distance', False, 'from jax_md import space, test_util, quantity, energy\n'), (207, 'jax.numpy.reshape', 'jnp.reshape', True, 'import jax.numpy as jnp\n'), (209, 'jax.numpy.where', 'jnp.where', True, 'import jax.numpy as jnp\n'), (210, 'jax.numpy.where', 'jnp.where', True, 'import jax.numpy as jnp\n'), (204, 'jax.numpy.array', 'jnp.array', True, 'import jax.numpy as jnp\n'), (217, 'jax_md.space.distance', 'space.distance', False, 'from jax_md import space, test_util, quantity, energy\n'), (218, 'jax.numpy.reshape', 'jnp.reshape', True, 'import jax.numpy as jnp\n'), (220, 'jax.numpy.where', 'jnp.where', True, 'import jax.numpy as jnp\n'), (222, 'jax.numpy.where', 'jnp.where', True, 'import jax.numpy as jnp\n'), (215, 'jax.numpy.array', 'jnp.array', True, 'import jax.numpy as jnp\n')]