# 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/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/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):


from typing import Union
from alphafold.model.geometry import rotation_matrix
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import vector
import jax
import jax.numpy as jnp
# Copyright 2021 DeepMind Technologies Limited
#
# 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
#
#      http://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.
"""Rigid3Array Transformations represented by a Matrix and a Vector."""

from __future__ import annotations


Float = Union[float, jnp.ndarray]

VERSION = '0.1'



# 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/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/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/rigid_matrix_vector.py
from typing import Union
from alphafold.model.geometry import rotation_matrix
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import vector
import jax
import jax.numpy as jnp

# Copyright 2021 DeepMind Technologies Limited
#
# 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
#
#      http://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.
"""Rigid3Array Transformations represented by a Matrix and a Vector."""

from __future__ import annotations


Float = Union[float, jnp.ndarray]

VERSION = '0.1'



@struct_of_array.StructOfArray(same_dtype=True)