# Path: alphafold/model/geometry/rigid_matrix_vector.py
# VERSION = '0.1'
# class Rigid3Array:
#   def __matmul__(self, other: Rigid3Array) -> Rigid3Array:
#   def inverse(self) -> Rigid3Array:
#   def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def compose_rotation(self, other_rotation):
#   def identity(cls, shape, dtype=jnp.float32) -> Rigid3Array:
#   def scale_translation(self, factor: Float) -> Rigid3Array:
#   def to_array(self):
#   def from_array(cls, array):
#   def from_array4x4(cls, array: jnp.ndarray) -> Rigid3Array:
#   def __getstate__(self):
#   def __setstate__(self, state):
# 
# Path: alphafold/model/geometry/rotation_matrix.py
# COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']
# VERSION = '0.1'
# class Rot3Array:
#   def inverse(self) -> Rot3Array:
#   def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def __matmul__(self, other: Rot3Array) -> Rot3Array:
#   def identity(cls, shape, dtype=jnp.float32) -> Rot3Array:
#   def from_two_vectors(cls, e0: vector.Vec3Array,
#                        e1: vector.Vec3Array) -> Rot3Array:
#   def from_array(cls, array: jnp.ndarray) -> Rot3Array:
#   def to_array(self) -> jnp.ndarray:
#   def from_quaternion(cls,
#                       w: jnp.ndarray,
#                       x: jnp.ndarray,
#                       y: jnp.ndarray,
#                       z: jnp.ndarray,
#                       normalize: bool = True,
#                       epsilon: float = 1e-6) -> Rot3Array:
#   def random_uniform(cls, key, shape, dtype=jnp.float32) -> Rot3Array:
#   def __getstate__(self):
#   def __setstate__(self, state):
# 
# Path: alphafold/model/geometry/vector.py
# VERSION = '0.1'
# class Vec3Array:
#   def __post_init__(self):
#   def __add__(self, other: Vec3Array) -> Vec3Array:
#   def __sub__(self, other: Vec3Array) -> Vec3Array:
#   def __mul__(self, other: Float) -> Vec3Array:
#   def __rmul__(self, other: Float) -> Vec3Array:
#   def __truediv__(self, other: Float) -> Vec3Array:
#   def __neg__(self) -> Vec3Array:
#   def __pos__(self) -> Vec3Array:
#   def cross(self, other: Vec3Array) -> Vec3Array:
#   def dot(self, other: Vec3Array) -> Float:
#   def norm(self, epsilon: float = 1e-6) -> Float:
#   def norm2(self):
#   def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
#   def zeros(cls, shape, dtype=jnp.float32):
#   def to_array(self) -> jnp.ndarray:
#   def from_array(cls, array):
#   def __getstate__(self):
#   def __setstate__(self, state):
# def square_euclidean_distance(vec1: Vec3Array,
#                               vec2: Vec3Array,
#                               epsilon: float = 1e-6) -> Float:
# def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
# def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
# def euclidean_distance(vec1: Vec3Array,
#                        vec2: Vec3Array,
#                        epsilon: float = 1e-6) -> Float:
# def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
#                    d: Vec3Array) -> Float:
# def random_gaussian_vector(shape, key, dtype=jnp.float32):


import dataclasses
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import rigid_matrix_vector
from alphafold.model.geometry import rotation_matrix
from alphafold.model.geometry import vector
  for field in dataclasses.fields(rotation_matrix.Rot3Array):
    field = field.name
    np.testing.assert_array_equal(
        getattr(matrix1, field), getattr(matrix2, field))


def assert_rotation_matrix_close(mat1: rotation_matrix.Rot3Array,
                                 mat2: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(mat1.to_array(), mat2.to_array(), 6)


def assert_array_equal_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  """Check that array and Matrix match."""
  np.testing.assert_array_equal(matrix.xx, array[..., 0, 0])
  np.testing.assert_array_equal(matrix.xy, array[..., 0, 1])
  np.testing.assert_array_equal(matrix.xz, array[..., 0, 2])
  np.testing.assert_array_equal(matrix.yx, array[..., 1, 0])
  np.testing.assert_array_equal(matrix.yy, array[..., 1, 1])
  np.testing.assert_array_equal(matrix.yz, array[..., 1, 2])
  np.testing.assert_array_equal(matrix.zx, array[..., 2, 0])
  np.testing.assert_array_equal(matrix.zy, array[..., 2, 1])
  np.testing.assert_array_equal(matrix.zz, array[..., 2, 2])


def assert_array_close_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(matrix.to_array(), array, 6)



# Path: alphafold/model/geometry/rigid_matrix_vector.py
# VERSION = '0.1'
# class Rigid3Array:
#   def __matmul__(self, other: Rigid3Array) -> Rigid3Array:
#   def inverse(self) -> Rigid3Array:
#   def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def compose_rotation(self, other_rotation):
#   def identity(cls, shape, dtype=jnp.float32) -> Rigid3Array:
#   def scale_translation(self, factor: Float) -> Rigid3Array:
#   def to_array(self):
#   def from_array(cls, array):
#   def from_array4x4(cls, array: jnp.ndarray) -> Rigid3Array:
#   def __getstate__(self):
#   def __setstate__(self, state):
# 
# Path: alphafold/model/geometry/rotation_matrix.py
# COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']
# VERSION = '0.1'
# class Rot3Array:
#   def inverse(self) -> Rot3Array:
#   def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
#   def __matmul__(self, other: Rot3Array) -> Rot3Array:
#   def identity(cls, shape, dtype=jnp.float32) -> Rot3Array:
#   def from_two_vectors(cls, e0: vector.Vec3Array,
#                        e1: vector.Vec3Array) -> Rot3Array:
#   def from_array(cls, array: jnp.ndarray) -> Rot3Array:
#   def to_array(self) -> jnp.ndarray:
#   def from_quaternion(cls,
#                       w: jnp.ndarray,
#                       x: jnp.ndarray,
#                       y: jnp.ndarray,
#                       z: jnp.ndarray,
#                       normalize: bool = True,
#                       epsilon: float = 1e-6) -> Rot3Array:
#   def random_uniform(cls, key, shape, dtype=jnp.float32) -> Rot3Array:
#   def __getstate__(self):
#   def __setstate__(self, state):
# 
# Path: alphafold/model/geometry/vector.py
# VERSION = '0.1'
# class Vec3Array:
#   def __post_init__(self):
#   def __add__(self, other: Vec3Array) -> Vec3Array:
#   def __sub__(self, other: Vec3Array) -> Vec3Array:
#   def __mul__(self, other: Float) -> Vec3Array:
#   def __rmul__(self, other: Float) -> Vec3Array:
#   def __truediv__(self, other: Float) -> Vec3Array:
#   def __neg__(self) -> Vec3Array:
#   def __pos__(self) -> Vec3Array:
#   def cross(self, other: Vec3Array) -> Vec3Array:
#   def dot(self, other: Vec3Array) -> Float:
#   def norm(self, epsilon: float = 1e-6) -> Float:
#   def norm2(self):
#   def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
#   def zeros(cls, shape, dtype=jnp.float32):
#   def to_array(self) -> jnp.ndarray:
#   def from_array(cls, array):
#   def __getstate__(self):
#   def __setstate__(self, state):
# def square_euclidean_distance(vec1: Vec3Array,
#                               vec2: Vec3Array,
#                               epsilon: float = 1e-6) -> Float:
# def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
# def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
# def euclidean_distance(vec1: Vec3Array,
#                        vec2: Vec3Array,
#                        epsilon: float = 1e-6) -> Float:
# def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
#                    d: Vec3Array) -> Float:
# def random_gaussian_vector(shape, key, dtype=jnp.float32):


# Path: alphafold/model/geometry/test_utils.py
import dataclasses
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import rigid_matrix_vector
from alphafold.model.geometry import rotation_matrix
from alphafold.model.geometry import vector

  for field in dataclasses.fields(rotation_matrix.Rot3Array):
    field = field.name
    np.testing.assert_array_equal(
        getattr(matrix1, field), getattr(matrix2, field))


def assert_rotation_matrix_close(mat1: rotation_matrix.Rot3Array,
                                 mat2: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(mat1.to_array(), mat2.to_array(), 6)


def assert_array_equal_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  """Check that array and Matrix match."""
  np.testing.assert_array_equal(matrix.xx, array[..., 0, 0])
  np.testing.assert_array_equal(matrix.xy, array[..., 0, 1])
  np.testing.assert_array_equal(matrix.xz, array[..., 0, 2])
  np.testing.assert_array_equal(matrix.yx, array[..., 1, 0])
  np.testing.assert_array_equal(matrix.yy, array[..., 1, 1])
  np.testing.assert_array_equal(matrix.yz, array[..., 1, 2])
  np.testing.assert_array_equal(matrix.zx, array[..., 2, 0])
  np.testing.assert_array_equal(matrix.zy, array[..., 2, 1])
  np.testing.assert_array_equal(matrix.zz, array[..., 2, 2])


def assert_array_close_to_rotation_matrix(array: jnp.ndarray,
                                          matrix: rotation_matrix.Rot3Array):
  np.testing.assert_array_almost_equal(matrix.to_array(), array, 6)



def assert_vectors_equal(vec1: vector.Vec3Array, vec2: vector.Vec3Array):