# 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()

[((984, 1029), 'jax_md.test_util.update_test_tolerance', 'test_util.update_test_tolerance', (['(5e-05)', '(5e-13)'], {}), '(5e-05, 5e-13)\n', (1015, 1029), False, 'from jax_md import space, test_util, quantity, energy\n'), ((1030, 1064), 'jax.config.config.parse_flags_with_absl', 'jax_config.parse_flags_with_absl', ([], {}), '()\n', (1062, 1064), True, 'from jax.config import config as jax_config\n'), ((1065, 1096), 'jax.config.config.enable_omnistaging', 'jax_config.enable_omnistaging', ([], {}), '()\n', (1094, 1096), True, 'from jax.config import config as jax_config\n'), ((1460, 1508), 'jax_md.quantity.box_size_at_number_density', 'quantity.box_size_at_number_density', (['N', '(1.0)', 'dim'], {}), '(N, 1.0, dim)\n', (1495, 1508), False, 'from jax_md import space, test_util, quantity, energy\n'), ((1788, 1815), 'jax_md.space.periodic_general', 'space.periodic_general', (['box'], {}), '(box)\n', (1810, 1815), False, 'from jax_md import space, test_util, quantity, energy\n'), ((1829, 1886), 'jax_md.space.periodic_general', 'space.periodic_general', (['box'], {'fractional_coordinates': '(False)'}), '(box, fractional_coordinates=False)\n', (1851, 1886), False, 'from jax_md import space, test_util, quantity, energy\n'), ((1896, 1913), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (1910, 1913), False, 'from jax import random\n'), ((1923, 1965), 'jax.random.uniform', 'random.uniform', (['key', '(N, dim)'], {'dtype': 'dtype'}), '(key, (N, dim), dtype=dtype)\n', (1937, 1965), False, 'from jax import random\n'), ((1972, 1997), 'jax_md.space.transform', 'space.transform', (['box', 'R_f'], {}), '(box, R_f)\n', (1987, 1997), False, 'from jax_md import space, test_util, quantity, energy\n'), ((20126, 20141), 'absl.testing.absltest.main', 'absltest.main', ([], {}), '()\n', (20139, 20141), False, 'from absl.testing import absltest\n'), ((2009, 2035), 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', (['d'], {}), '(d)\n', (2032, 2035), False, 'from jax_md import space, test_util, quantity, energy\n'), ((2050, 2079), 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', (['d_gf'], {}), '(d_gf)\n', (2073, 2079), False, 'from jax_md import space, test_util, quantity, energy\n'), ((2093, 2121), 'jax_md.energy.soft_sphere_pair', 'energy.soft_sphere_pair', (['d_g'], {}), '(d_g)\n', (2116, 2121), False, 'from jax_md import space, test_util, quantity, energy\n'), ((2621, 2638), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (2635, 2638), False, 'from jax import random\n'), ((3331, 3348), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (3345, 3348), False, 'from jax import random\n'), ((4239, 4256), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (4253, 4256), False, 'from jax import random\n'), ((5077, 5094), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (5091, 5094), False, 'from jax import random\n'), ((5182, 5208), 'jax_md.space.metric', 'space.metric', (['displacement'], {}), '(displacement)\n', (5194, 5208), False, 'from jax_md import space, test_util, quantity, energy\n'), ((5227, 5282), 'jax_md.space.canonicalize_displacement_or_metric', 'space.canonicalize_displacement_or_metric', (['displacement'], {}), '(displacement)\n', (5268, 5282), False, 'from jax_md import space, test_util, quantity, energy\n'), ((5297, 5322), 'jax_md.space.map_product', 'space.map_product', (['metric'], {}), '(metric)\n', (5314, 5322), False, 'from jax_md import space, test_util, quantity, energy\n'), ((5341, 5371), 'jax_md.space.map_product', 'space.map_product', (['test_metric'], {}), '(test_metric)\n', (5358, 5371), False, 'from jax_md import space, test_util, quantity, energy\n'), ((5957, 5974), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (5971, 5974), False, 'from jax import random\n'), ((7864, 7881), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (7878, 7881), False, 'from jax import random\n'), ((8947, 8964), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (8961, 8964), False, 'from jax import random\n'), ((10376, 10393), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (10390, 10393), False, 'from jax import random\n'), ((11761, 11778), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (11775, 11778), False, 'from jax import random\n'), ((11790, 11816), 'jax.numpy.eye', 'jnp.eye', (['spatial_dimension'], {}), '(spatial_dimension)\n', (11797, 11816), True, 'import jax.numpy as jnp\n'), ((13541, 13558), 'jax.random.PRNGKey', 'random.PRNGKey', (['(0)'], {}), '(0)\n', (13555, 13558), False, 'from jax import random\n'), ((13570, 13609), 'jax.numpy.eye', 'jnp.eye', (['spatial_dimension'], {'dtype': 'dtype'}), '(spatial_dimension, dtype=dtype)\n', (13577, 13609), True, 'import jax.numpy as jnp\n'), ((14658, 14929), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (14677, 14929), True, 'from jax import test_util as jtu\n'), ((15318, 15589), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (15337, 15589), True, 'from jax import test_util as jtu\n'), ((16001, 16272), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (16020, 16272), True, 'from jax import test_util as jtu\n'), ((16792, 17063), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (16811, 17063), True, 'from jax import test_util as jtu\n'), ((17509, 17780), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (17528, 17780), True, 'from jax import test_util as jtu\n'), ((18407, 18678), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (18426, 18678), True, 'from jax import test_util as jtu\n'), ((19234, 19505), 'jax.test_util.cases_from_list', 'jtu.cases_from_list', (["({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)"], {}), "({'testcase_name':\n    f'_dim={dim}_dtype={dtype.__name__}_box_format={box_format}',\n    'spatial_dimension': dim, 'dtype': dtype, 'box_format': box_format} for\n    dim in SPATIAL_DIMENSION for dtype in POSITION_DTYPE for box_format in\n    BOX_FORMATS)\n", (19253, 19505), True, 'from jax import test_util as jtu\n'), ((1723, 1736), 'jax.numpy.diag', 'jnp.diag', (['box'], {}), '(box)\n', (1731, 1736), True, 'import jax.numpy as jnp\n'), ((2708, 2728), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (2720, 2728), False, 'from jax import random\n'), ((2740, 2811), 'jax.random.normal', 'random.normal', (['split1', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (2753, 2811), False, 'from jax import random\n'), ((2831, 2905), 'jax.random.normal', 'random.normal', (['split2', '(spatial_dimension, spatial_dimension)'], {'dtype': 'dtype'}), '(split2, (spatial_dimension, spatial_dimension), dtype=dtype)\n', (2844, 2905), False, 'from jax import random\n'), ((3008, 3029), 'jax_md.space.transform', 'space.transform', (['T', 'R'], {}), '(T, R)\n', (3023, 3029), False, 'from jax_md import space, test_util, quantity, energy\n'), ((3418, 3438), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (3430, 3438), False, 'from jax import random\n'), ((3450, 3508), 'jax.random.normal', 'random.normal', (['split1', '(PARTICLE_COUNT, spatial_dimension)'], {}), '(split1, (PARTICLE_COUNT, spatial_dimension))\n', (3463, 3508), False, 'from jax import random\n'), ((3519, 3580), 'jax.random.normal', 'random.normal', (['split2', '(spatial_dimension, spatial_dimension)'], {}), '(split2, (spatial_dimension, spatial_dimension))\n', (3532, 3580), False, 'from jax import random\n'), ((3598, 3619), 'jax_md.space.transform', 'space.transform', (['T', 'R'], {}), '(T, R)\n', (3613, 3619), False, 'from jax_md import space, test_util, quantity, energy\n'), ((4381, 4401), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (4393, 4401), False, 'from jax import random\n'), ((4413, 4484), 'jax.random.normal', 'random.normal', (['split1', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (4426, 4484), False, 'from jax import random\n'), ((4505, 4579), 'jax.random.normal', 'random.normal', (['split2', '(spatial_dimension, spatial_dimension)'], {'dtype': 'dtype'}), '(split2, (spatial_dimension, spatial_dimension), dtype=dtype)\n', (4518, 4579), False, 'from jax import random\n'), ((4603, 4619), 'jax_md.space.inverse', 'space.inverse', (['T'], {}), '(T)\n', (4616, 4619), False, 'from jax_md import space, test_util, quantity, energy\n'), ((5141, 5167), 'jax.numpy.eye', 'jnp.eye', (['spatial_dimension'], {}), '(spatial_dimension)\n', (5148, 5167), True, 'import jax.numpy as jnp\n'), ((5441, 5461), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (5453, 5461), False, 'from jax import random\n'), ((5473, 5544), 'jax.random.normal', 'random.normal', (['split1', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (5486, 5544), False, 'from jax import random\n'), ((6035, 6052), 'jax.random.split', 'random.split', (['key'], {}), '(key)\n', (6047, 6052), False, 'from jax import random\n'), ((6064, 6135), 'jax.random.uniform', 'random.uniform', (['split', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (6078, 6135), False, 'from jax import random\n'), ((6311, 6329), 'jax_md.space.distance', 'space.distance', (['dR'], {}), '(dR)\n', (6325, 6329), False, 'from jax_md import space, test_util, quantity, energy\n'), ((6348, 6394), 'jax.numpy.reshape', 'jnp.reshape', (['dr_direct', '(dr_direct.shape + (1,))'], {}), '(dr_direct, dr_direct.shape + (1,))\n', (6359, 6394), True, 'import jax.numpy as jnp\n'), ((7403, 7439), 'jax.numpy.array', 'jnp.array', (['dR_direct'], {'dtype': 'dR.dtype'}), '(dR_direct, dtype=dR.dtype)\n', (7412, 7439), True, 'import jax.numpy as jnp\n'), ((7951, 7971), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (7963, 7971), False, 'from jax import random\n'), ((7983, 8055), 'jax.random.uniform', 'random.uniform', (['split1', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (7997, 8055), False, 'from jax import random\n'), ((8384, 8406), 'jax.numpy.all', 'jnp.all', (['(R_shift < 1.0)'], {}), '(R_shift < 1.0)\n', (8391, 8406), True, 'import jax.numpy as jnp\n'), ((8420, 8442), 'jax.numpy.all', 'jnp.all', (['(R_shift > 0.0)'], {}), '(R_shift > 0.0)\n', (8427, 8442), True, 'import jax.numpy as jnp\n'), ((9097, 9117), 'jax.random.split', 'random.split', (['key', '(4)'], {}), '(key, 4)\n', (9109, 9117), False, 'from jax import random\n'), ((9267, 9285), 'jax.numpy.diag', 'jnp.diag', (['box_size'], {}), '(box_size)\n', (9275, 9285), True, 'import jax.numpy as jnp\n'), ((9297, 9369), 'jax.random.uniform', 'random.uniform', (['split2', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (9311, 9369), False, 'from jax import random\n'), ((9421, 9492), 'jax.random.normal', 'random.normal', (['split3', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split3, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (9434, 9492), False, 'from jax import random\n'), ((9529, 9553), 'jax_md.space.periodic', 'space.periodic', (['box_size'], {}), '(box_size)\n', (9543, 9553), False, 'from jax_md import space, test_util, quantity, energy\n'), ((9596, 9629), 'jax_md.space.periodic_general', 'space.periodic_general', (['transform'], {}), '(transform)\n', (9618, 9629), False, 'from jax_md import space, test_util, quantity, energy\n'), ((9647, 9673), 'jax_md.space.map_product', 'space.map_product', (['disp_fn'], {}), '(disp_fn)\n', (9664, 9673), False, 'from jax_md import space, test_util, quantity, energy\n'), ((9698, 9732), 'jax_md.space.map_product', 'space.map_product', (['general_disp_fn'], {}), '(general_disp_fn)\n', (9715, 9732), False, 'from jax_md import space, test_util, quantity, energy\n'), ((10526, 10546), 'jax.random.split', 'random.split', (['key', '(4)'], {}), '(key, 4)\n', (10538, 10546), False, 'from jax import random\n'), ((10696, 10714), 'jax.numpy.diag', 'jnp.diag', (['box_size'], {}), '(box_size)\n', (10704, 10714), True, 'import jax.numpy as jnp\n'), ((10726, 10798), 'jax.random.uniform', 'random.uniform', (['split2', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (10740, 10798), False, 'from jax import random\n'), ((10850, 10921), 'jax.random.normal', 'random.normal', (['split3', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split3, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (10863, 10921), False, 'from jax import random\n'), ((10958, 10982), 'jax_md.space.periodic', 'space.periodic', (['box_size'], {}), '(box_size)\n', (10972, 10982), False, 'from jax_md import space, test_util, quantity, energy\n'), ((11025, 11058), 'jax_md.space.periodic_general', 'space.periodic_general', (['transform'], {}), '(transform)\n', (11047, 11058), False, 'from jax_md import space, test_util, quantity, energy\n'), ((11076, 11102), 'jax_md.space.map_product', 'space.map_product', (['disp_fn'], {}), '(disp_fn)\n', (11093, 11102), False, 'from jax_md import space, test_util, quantity, energy\n'), ((11127, 11161), 'jax_md.space.map_product', 'space.map_product', (['general_disp_fn'], {}), '(general_disp_fn)\n', (11144, 11161), False, 'from jax_md import space, test_util, quantity, energy\n'), ((11899, 11919), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (11911, 11919), False, 'from jax import random\n'), ((11961, 11981), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (11973, 11981), False, 'from jax import random\n'), ((12022, 12042), 'jax.random.split', 'random.split', (['key', '(4)'], {}), '(key, 4)\n', (12034, 12042), False, 'from jax import random\n'), ((12216, 12269), 'jax.numpy.array', 'jnp.array', (['(size_0 * (eye + dtransform_0))'], {'dtype': 'dtype'}), '(size_0 * (eye + dtransform_0), dtype=dtype)\n', (12225, 12269), True, 'import jax.numpy as jnp\n'), ((12471, 12524), 'jax.numpy.array', 'jnp.array', (['(size_1 * (eye + dtransform_1))'], {'dtype': 'dtype'}), '(size_1 * (eye + dtransform_1), dtype=dtype)\n', (12480, 12524), True, 'import jax.numpy as jnp\n'), ((12552, 12579), 'jax_md.space.periodic_general', 'space.periodic_general', (['T_0'], {}), '(T_0)\n', (12574, 12579), False, 'from jax_md import space, test_util, quantity, energy\n'), ((12616, 12643), 'jax_md.space.periodic_general', 'space.periodic_general', (['T_1'], {}), '(T_1)\n', (12638, 12643), False, 'from jax_md import space, test_util, quantity, energy\n'), ((12661, 12686), 'functools.partial', 'partial', (['disp_fn'], {'box': 'T_1'}), '(disp_fn, box=T_1)\n', (12668, 12686), False, 'from functools import partial\n'), ((12704, 12730), 'jax_md.space.map_product', 'space.map_product', (['disp_fn'], {}), '(disp_fn)\n', (12721, 12730), False, 'from jax_md import space, test_util, quantity, energy\n'), ((12752, 12783), 'jax_md.space.map_product', 'space.map_product', (['true_disp_fn'], {}), '(true_disp_fn)\n', (12769, 12783), False, 'from jax_md import space, test_util, quantity, energy\n'), ((12795, 12868), 'jax.random.uniform', 'random.uniform', (['split_R', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split_R, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (12809, 12868), False, 'from jax import random\n'), ((12889, 12962), 'jax.random.normal', 'random.normal', (['split_dR', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split_dR, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (12902, 12962), False, 'from jax import random\n'), ((13736, 13756), 'jax.random.split', 'random.split', (['key', '(3)'], {}), '(key, 3)\n', (13748, 13756), False, 'from jax import random\n'), ((13769, 13844), 'jax.random.normal', 'random.normal', (['split_T', '(spatial_dimension, spatial_dimension)'], {'dtype': 'dtype'}), '(split_T, (spatial_dimension, spatial_dimension), dtype=dtype)\n', (13782, 13844), False, 'from jax import random\n'), ((13904, 13977), 'jax.random.uniform', 'random.uniform', (['split_R', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split_R, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (13918, 13977), False, 'from jax import random\n'), ((14051, 14076), 'jax_md.space.periodic_general', 'space.periodic_general', (['T'], {}), '(T)\n', (14073, 14076), False, 'from jax_md import space, test_util, quantity, energy\n'), ((14104, 14144), 'jax_md.space.periodic_general', 'space.periodic_general', (['T'], {'wrapped': '(False)'}), '(T, wrapped=False)\n', (14126, 14144), False, 'from jax_md import space, test_util, quantity, energy\n'), ((14167, 14198), 'jax_md.space.map_product', 'space.map_product', (['displacement'], {}), '(displacement)\n', (14184, 14198), False, 'from jax_md import space, test_util, quantity, energy\n'), ((16685, 16715), 'jax_md.space.transform', 'space.transform', (['box', 'R_gf_new'], {}), '(box, R_gf_new)\n', (16700, 16715), False, 'from jax_md import space, test_util, quantity, energy\n'), ((19679, 19732), 'unittest.SkipTest', 'SkipTest', (['"""Scalar case fails due to JAX Issue #5849."""'], {}), "('Scalar case fails due to JAX Issue #5849.')\n", (19687, 19732), False, 'from unittest import SkipTest\n'), ((1582, 1595), 'jax.numpy.ones', 'jnp.ones', (['dim'], {}), '(dim)\n', (1590, 1595), True, 'import jax.numpy as jnp\n'), ((2948, 2977), 'jax.numpy.einsum', 'jnp.einsum', (['"""ij,kj->ki"""', 'T', 'R'], {}), "('ij,kj->ki', T, R)\n", (2958, 2977), True, 'import jax.numpy as jnp\n'), ((3653, 3668), 'jax.numpy.sum', 'jnp.sum', (['(R ** 2)'], {}), '(R ** 2)\n', (3660, 3668), True, 'import jax.numpy as jnp\n'), ((3763, 3782), 'jax.grad', 'grad', (['energy_direct'], {}), '(energy_direct)\n', (3767, 3782), False, 'from jax import grad, jit, jacfwd\n'), ((3814, 3838), 'jax.grad', 'grad', (['energy_indirect', '(1)'], {}), '(energy_indirect, 1)\n', (3818, 3838), False, 'from jax import grad, jit, jacfwd\n'), ((4659, 4680), 'jax_md.space.transform', 'space.transform', (['T', 'R'], {}), '(T, R)\n', (4674, 4680), False, 'from jax_md import space, test_util, quantity, energy\n'), ((6156, 6202), 'jax_md.space.map_product', 'space.map_product', (['space.pairwise_displacement'], {}), '(space.pairwise_displacement)\n', (6173, 6202), False, 'from jax_md import space, test_util, quantity, energy\n'), ((8097, 8168), 'jax.random.normal', 'random.normal', (['split2', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split2, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (8110, 8168), False, 'from jax import random\n'), ((9182, 9239), 'jax.random.uniform', 'random.uniform', (['split1', '(spatial_dimension,)'], {'dtype': 'dtype'}), '(split1, (spatial_dimension,), dtype=dtype)\n', (9196, 9239), False, 'from jax import random\n'), ((10611, 10668), 'jax.random.uniform', 'random.uniform', (['split1', '(spatial_dimension,)'], {'dtype': 'dtype'}), '(split1, (spatial_dimension,), dtype=dtype)\n', (10625, 10668), False, 'from jax import random\n'), ((12066, 12100), 'jax.random.uniform', 'random.uniform', (['split_T0_scale', '()'], {}), '(split_T0_scale, ())\n', (12080, 12100), False, 'from jax import random\n'), ((12128, 12194), 'jax.random.normal', 'random.normal', (['split_T0_dT', '(spatial_dimension, spatial_dimension)'], {}), '(split_T0_dT, (spatial_dimension, spatial_dimension))\n', (12141, 12194), False, 'from jax import random\n'), ((12293, 12340), 'jax.random.uniform', 'random.uniform', (['split_T1_scale', '()'], {'dtype': 'dtype'}), '(split_T1_scale, (), dtype=dtype)\n', (12307, 12340), False, 'from jax import random\n'), ((12368, 12447), 'jax.random.normal', 'random.normal', (['split_T1_dT', '(spatial_dimension, spatial_dimension)'], {'dtype': 'dtype'}), '(split_T1_dT, (spatial_dimension, spatial_dimension), dtype=dtype)\n', (12381, 12447), False, 'from jax import random\n'), ((13875, 13892), 'jax.numpy.transpose', 'jnp.transpose', (['dT'], {}), '(dT)\n', (13888, 13892), True, 'import jax.numpy as jnp\n'), ((14256, 14273), 'jax.random.split', 'random.split', (['key'], {}), '(key)\n', (14268, 14273), False, 'from jax import random\n'), ((14287, 14357), 'jax.random.normal', 'random.normal', (['split', '(PARTICLE_COUNT, spatial_dimension)'], {'dtype': 'dtype'}), '(split, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)\n', (14300, 14357), False, 'from jax import random\n'), ((15887, 15894), 'jax.grad', 'grad', (['E'], {}), '(E)\n', (15891, 15894), False, 'from jax import grad, jit, jacfwd\n'), ((15899, 15909), 'jax.grad', 'grad', (['E_gf'], {}), '(E_gf)\n', (15903, 15909), False, 'from jax import grad, jit, jacfwd\n'), ((15940, 15947), 'jax.grad', 'grad', (['E'], {}), '(E)\n', (15944, 15947), False, 'from jax import grad, jit, jacfwd\n'), ((15952, 15961), 'jax.grad', 'grad', (['E_g'], {}), '(E_g)\n', (15956, 15961), False, 'from jax import grad, jit, jacfwd\n'), ((16564, 16571), 'jax.grad', 'grad', (['E'], {}), '(E)\n', (16568, 16571), False, 'from jax import grad, jit, jacfwd\n'), ((16601, 16611), 'jax.grad', 'grad', (['E_gf'], {}), '(E_gf)\n', (16605, 16611), False, 'from jax import grad, jit, jacfwd\n'), ((16639, 16648), 'jax.grad', 'grad', (['E_g'], {}), '(E_g)\n', (16643, 16648), False, 'from jax import grad, jit, jacfwd\n'), ((18153, 18163), 'jax.grad', 'grad', (['E_gf'], {}), '(E_gf)\n', (18157, 18163), False, 'from jax import grad, jit, jacfwd\n'), ((18212, 18221), 'jax.grad', 'grad', (['E_g'], {}), '(E_g)\n', (18216, 18221), False, 'from jax import grad, jit, jacfwd\n'), ((18273, 18285), 'jax.jacfwd', 'jacfwd', (['E_gf'], {}), '(E_gf)\n', (18279, 18285), False, 'from jax import grad, jit, jacfwd\n'), ((18334, 18345), 'jax.jacfwd', 'jacfwd', (['E_g'], {}), '(E_g)\n', (18340, 18345), False, 'from jax import grad, jit, jacfwd\n'), ((19048, 19057), 'jax.grad', 'grad', (['E_g'], {}), '(E_g)\n', (19052, 19057), False, 'from jax import grad, jit, jacfwd\n'), ((1666, 1678), 'jax.numpy.eye', 'jnp.eye', (['dim'], {}), '(dim)\n', (1673, 1678), True, 'import jax.numpy as jnp\n'), ((14568, 14592), 'jax.numpy.all', 'jnp.all', (['(unwrapped_R > 0)'], {}), '(unwrapped_R > 0)\n', (14575, 14592), True, 'import jax.numpy as jnp\n'), ((14597, 14621), 'jax.numpy.all', 'jnp.all', (['(unwrapped_R < 1)'], {}), '(unwrapped_R < 1)\n', (14604, 14621), True, 'import jax.numpy as jnp\n'), ((19139, 19149), 'jax.grad', 'grad', (['E_gf'], {}), '(E_gf)\n', (19143, 19149), False, 'from jax import grad, jit, jacfwd\n'), ((3714, 3735), 'jax_md.space.transform', 'space.transform', (['T', 'R'], {}), '(T, R)\n', (3729, 3735), False, 'from jax_md import space, test_util, quantity, energy\n'), ((6582, 6608), 'jax_md.space.distance', 'space.distance', (['dR_shifted'], {}), '(dR_shifted)\n', (6596, 6608), False, 'from jax_md import space, test_util, quantity, energy\n'), ((6634, 6682), 'jax.numpy.reshape', 'jnp.reshape', (['dr_shifted', '(dr_shifted.shape + (1,))'], {}), '(dr_shifted, dr_shifted.shape + (1,))\n', (6645, 6682), True, 'import jax.numpy as jnp\n'), ((6708, 6764), 'jax.numpy.where', 'jnp.where', (['(dr_shifted < dr_direct)', 'dR_shifted', 'dR_direct'], {}), '(dr_shifted < dr_direct, dR_shifted, dR_direct)\n', (6717, 6764), True, 'import jax.numpy as jnp\n'), ((6789, 6845), 'jax.numpy.where', 'jnp.where', (['(dr_shifted < dr_direct)', 'dr_shifted', 'dr_direct'], {}), '(dr_shifted < dr_direct, dr_shifted, dr_direct)\n', (6798, 6845), True, 'import jax.numpy as jnp\n'), ((6523, 6555), 'jax.numpy.array', 'jnp.array', (['[i, j]'], {'dtype': 'R.dtype'}), '([i, j], dtype=R.dtype)\n', (6532, 6555), True, 'import jax.numpy as jnp\n'), ((7076, 7102), 'jax_md.space.distance', 'space.distance', (['dR_shifted'], {}), '(dR_shifted)\n', (7090, 7102), False, 'from jax_md import space, test_util, quantity, energy\n'), ((7130, 7178), 'jax.numpy.reshape', 'jnp.reshape', (['dr_shifted', '(dr_shifted.shape + (1,))'], {}), '(dr_shifted, dr_shifted.shape + (1,))\n', (7141, 7178), True, 'import jax.numpy as jnp\n'), ((7206, 7262), 'jax.numpy.where', 'jnp.where', (['(dr_shifted < dr_direct)', 'dR_shifted', 'dR_direct'], {}), '(dr_shifted < dr_direct, dR_shifted, dR_direct)\n', (7215, 7262), True, 'import jax.numpy as jnp\n'), ((7308, 7364), 'jax.numpy.where', 'jnp.where', (['(dr_shifted < dr_direct)', 'dr_shifted', 'dr_direct'], {}), '(dr_shifted < dr_direct, dr_shifted, dr_direct)\n', (7317, 7364), True, 'import jax.numpy as jnp\n'), ((7012, 7047), 'jax.numpy.array', 'jnp.array', (['[i, j, k]'], {'dtype': 'R.dtype'}), '([i, j, k], dtype=R.dtype)\n', (7021, 7047), True, 'import jax.numpy as jnp\n')]