# Path: alphafold/model/geometry/struct_of_array.py
# def get_item(instance, key):
# def get_shape(instance):
# def get_len(instance):
# def get_dtype(instance):
# def replace(instance, **kwargs):
# def post_init(instance):
# def flatten(instance):
# def make_metadata_class(cls):
# def get_fields(cls_or_instance, filterfn, return_values=False):
# def get_array_fields(cls, return_values=False):
# def get_metadata_fields(cls, return_values=False):
#   def __init__(self, same_dtype=True):
#   def __call__(self, cls):
#     def unflatten(aux, data):
# class StructOfArray:
# 
# Path: alphafold/model/geometry/utils.py
# def unstack(value: jnp.ndarray, axis: int = -1) -> List[jnp.ndarray]:
# 
# 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
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils
from alphafold.model.geometry import vector
from __future__ import annotations


COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']

VERSION = '0.1'


@struct_of_array.StructOfArray(same_dtype=True)
class Rot3Array:
  """Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays."""

  xx: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
  xy: jnp.ndarray
  xz: jnp.ndarray
  yx: jnp.ndarray
  yy: jnp.ndarray
  yz: jnp.ndarray
  zx: jnp.ndarray
  zy: jnp.ndarray
  zz: jnp.ndarray

  __array_ufunc__ = None

  def inverse(self) -> Rot3Array:
    """Returns inverse of Rot3Array."""
    return Rot3Array(self.xx, self.yx, self.zx,
                     self.xy, self.yy, self.zy,
                     self.xz, self.yz, self.zz)


# Path: alphafold/model/geometry/struct_of_array.py
# def get_item(instance, key):
# def get_shape(instance):
# def get_len(instance):
# def get_dtype(instance):
# def replace(instance, **kwargs):
# def post_init(instance):
# def flatten(instance):
# def make_metadata_class(cls):
# def get_fields(cls_or_instance, filterfn, return_values=False):
# def get_array_fields(cls, return_values=False):
# def get_metadata_fields(cls, return_values=False):
#   def __init__(self, same_dtype=True):
#   def __call__(self, cls):
#     def unflatten(aux, data):
# class StructOfArray:
# 
# Path: alphafold/model/geometry/utils.py
# def unstack(value: jnp.ndarray, axis: int = -1) -> List[jnp.ndarray]:
# 
# 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/rotation_matrix.py
import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils
from alphafold.model.geometry import vector

from __future__ import annotations


COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']

VERSION = '0.1'


@struct_of_array.StructOfArray(same_dtype=True)
class Rot3Array:
  """Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays."""

  xx: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32})
  xy: jnp.ndarray
  xz: jnp.ndarray
  yx: jnp.ndarray
  yy: jnp.ndarray
  yz: jnp.ndarray
  zx: jnp.ndarray
  zy: jnp.ndarray
  zz: jnp.ndarray

  __array_ufunc__ = None

  def inverse(self) -> Rot3Array:
    """Returns inverse of Rot3Array."""
    return Rot3Array(self.xx, self.yx, self.zx,
                     self.xy, self.yy, self.zy,
                     self.xz, self.yz, self.zz)


  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array: