# Path: ldif/torch/sif.py
# def ensure_are_tensors(named_tensors):
# def ensure_type(v, t, name):
# def _load_v1_txt(path):
# def _tile_for_symgroups(elements, symmetry_count):
# def reflect_z(samples):
# def _generate_symgroup_samples(samples, element_count, symmetry_count):
# def _generate_symgroup_frames(world2local, symmetry_count):
#   def __init__(self, constants, centers, radii, rotations, symmetry_count):
#   def from_file(cls, path):
#   def from_flat_tensor(cls, tensor, symmetry_count):
#   def to_flat_tensor(self):
#   def rbf_influence(self, samples):
#   def effective_element_count(self):
#   def constants(self):
#   def _tiled_constants(self):
#   def _tiled_centers(self):
#   def _tiled_radii(self):
#   def _tiled_rotations(self):
#   def world2local(self):
#   def eval(self, samples):
# class Sif(object):


import os
import sys
import unittest
import numpy as np
import torch
from ldif.torch import sif
# limitations under the License.


# Unit tests need LDIF added to the path, because the scripts aren't
# in the repository top-level directory.
ldif_root = os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
sys.path.append(ldif_root)



TEST_DATA_DIR  = './data'

class TestSif(unittest.TestCase):

  def _ensure_sifs_equal(self, a, b, ith=0):
    end=ith+1
    if ith == -1:
      self.assertEqual(a.bs, b.bs)
      ith = 0
      end = b.bs
    self.assertTrue(torch.all(a._constants.eq(b._constants[ith:end, ...])))
    self.assertTrue(torch.all(a._centers.eq(b._centers[ith:end, ...])))
    self.assertTrue(torch.all(a._radii.eq(b._radii[ith:end, ...])))
    self.assertTrue(torch.all(a._rotations.eq(b._rotations[ith:end, ...])))
    self.assertEqual(a.symmetry_count, b.symmetry_count)
    self.assertEqual(a.element_count, b.element_count)

  def test_load_sif(self):
    sif_path = os.path.join(TEST_DATA_DIR, 'b831f60f211435df5bbc861d0124304c.txt')
    constants, centers, radii, rotations, symmetry_count, features = (

# Path: ldif/torch/sif.py
# def ensure_are_tensors(named_tensors):
# def ensure_type(v, t, name):
# def _load_v1_txt(path):
# def _tile_for_symgroups(elements, symmetry_count):
# def reflect_z(samples):
# def _generate_symgroup_samples(samples, element_count, symmetry_count):
# def _generate_symgroup_frames(world2local, symmetry_count):
#   def __init__(self, constants, centers, radii, rotations, symmetry_count):
#   def from_file(cls, path):
#   def from_flat_tensor(cls, tensor, symmetry_count):
#   def to_flat_tensor(self):
#   def rbf_influence(self, samples):
#   def effective_element_count(self):
#   def constants(self):
#   def _tiled_constants(self):
#   def _tiled_centers(self):
#   def _tiled_radii(self):
#   def _tiled_rotations(self):
#   def world2local(self):
#   def eval(self, samples):
# class Sif(object):


# Path: ldif/torch/sif_test.py
import os
import sys
import unittest
import numpy as np
import torch
from ldif.torch import sif

# limitations under the License.


# Unit tests need LDIF added to the path, because the scripts aren't
# in the repository top-level directory.
ldif_root = os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
sys.path.append(ldif_root)



TEST_DATA_DIR  = './data'

class TestSif(unittest.TestCase):

  def _ensure_sifs_equal(self, a, b, ith=0):
    end=ith+1
    if ith == -1:
      self.assertEqual(a.bs, b.bs)
      ith = 0
      end = b.bs
    self.assertTrue(torch.all(a._constants.eq(b._constants[ith:end, ...])))
    self.assertTrue(torch.all(a._centers.eq(b._centers[ith:end, ...])))
    self.assertTrue(torch.all(a._radii.eq(b._radii[ith:end, ...])))
    self.assertTrue(torch.all(a._rotations.eq(b._rotations[ith:end, ...])))
    self.assertEqual(a.symmetry_count, b.symmetry_count)
    self.assertEqual(a.element_count, b.element_count)

  def test_load_sif(self):
    sif_path = os.path.join(TEST_DATA_DIR, 'b831f60f211435df5bbc861d0124304c.txt')
    constants, centers, radii, rotations, symmetry_count, features = (

        sif._load_v1_txt(sif_path))