|
|
|
|
| """Vec3Array Class.""" |
|
|
| import dataclasses |
| from typing import Final, Self, TypeAlias |
|
|
| from flax_model.alphafold3.jax.geometry import struct_of_array |
| from flax_model.alphafold3.jax.geometry import utils |
| import jax |
| import jax.numpy as jnp |
| import numpy as np |
|
|
|
|
| Float: TypeAlias = float | jnp.ndarray |
|
|
| VERSION: Final[str] = '0.1' |
|
|
|
|
| @struct_of_array.StructOfArray(same_dtype=True) |
| class Vec3Array: |
| """Vec3Array in 3 dimensional Space implemented as struct of arrays. |
| |
| This is done in order to improve performance and precision. |
| On TPU small matrix multiplications are very suboptimal and will waste large |
| compute ressources, furthermore any matrix multiplication on TPU happens in |
| mixed bfloat16/float32 precision, which is often undesirable when handling |
| physical coordinates. |
| |
| In most cases this will also be faster on CPUs/GPUs since it allows for easier |
| use of vector instructions. |
| """ |
|
|
| x: jnp.ndarray = dataclasses.field(metadata={'dtype': jnp.float32}) |
| y: jnp.ndarray |
| z: jnp.ndarray |
|
|
| def __post_init__(self): |
| if hasattr(self.x, 'dtype'): |
| if not self.x.dtype == self.y.dtype == self.z.dtype: |
| raise ValueError( |
| f'Type mismatch: {self.x.dtype}, {self.y.dtype}, {self.z.dtype}' |
| ) |
| if not self.x.shape == self.y.shape == self.z.shape: |
| raise ValueError( |
| f'Shape mismatch: {self.x.shape}, {self.y.shape}, {self.z.shape}' |
| ) |
|
|
| def __add__(self, other: Self) -> Self: |
| return jax.tree.map(lambda x, y: x + y, self, other) |
|
|
| def __sub__(self, other: Self) -> Self: |
| return jax.tree.map(lambda x, y: x - y, self, other) |
|
|
| def __mul__(self, other: Float) -> Self: |
| return jax.tree.map(lambda x: x * other, self) |
|
|
| def __rmul__(self, other: Float) -> Self: |
| return self * other |
|
|
| def __truediv__(self, other: Float) -> Self: |
| return jax.tree.map(lambda x: x / other, self) |
|
|
| def __neg__(self) -> Self: |
| return jax.tree.map(lambda x: -x, self) |
|
|
| def __pos__(self) -> Self: |
| return jax.tree.map(lambda x: x, self) |
|
|
| def cross(self, other: Self) -> Self: |
| """Compute cross product between 'self' and 'other'.""" |
| new_x = self.y * other.z - self.z * other.y |
| new_y = self.z * other.x - self.x * other.z |
| new_z = self.x * other.y - self.y * other.x |
| return Vec3Array(new_x, new_y, new_z) |
|
|
| def dot(self, other: Self) -> Float: |
| """Compute dot product between 'self' and 'other'.""" |
| return self.x * other.x + self.y * other.y + self.z * other.z |
|
|
| def norm(self, epsilon: float = 1e-6) -> Float: |
| """Compute Norm of Vec3Array, clipped to epsilon.""" |
| |
| norm2 = self.dot(self) |
| if epsilon: |
| norm2 = jnp.maximum(norm2, epsilon**2) |
| return jnp.sqrt(norm2) |
|
|
| def norm2(self): |
| return self.dot(self) |
|
|
| def normalized(self, epsilon: float = 1e-6) -> Self: |
| """Return unit vector with optional clipping.""" |
| return self / self.norm(epsilon) |
|
|
| @classmethod |
| def zeros(cls, shape, dtype=jnp.float32): |
| """Return Vec3Array corresponding to zeros of given shape.""" |
| return cls( |
| jnp.zeros(shape, dtype), |
| jnp.zeros(shape, dtype), |
| jnp.zeros(shape, dtype), |
| ) |
|
|
| def to_array(self) -> jnp.ndarray: |
| return jnp.stack([self.x, self.y, self.z], axis=-1) |
|
|
| @classmethod |
| def from_array(cls, array): |
| return cls(*utils.unstack(array)) |
|
|
| def __getstate__(self): |
| return ( |
| VERSION, |
| [np.asarray(self.x), np.asarray(self.y), np.asarray(self.z)], |
| ) |
|
|
| def __setstate__(self, state): |
| version, state = state |
| del version |
| for i, letter in enumerate('xyz'): |
| object.__setattr__(self, letter, state[i]) |
|
|
|
|
| def square_euclidean_distance( |
| vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-6 |
| ) -> Float: |
| """Computes square of euclidean distance between 'vec1' and 'vec2'. |
| |
| Args: |
| vec1: Vec3Array to compute distance to |
| vec2: Vec3Array to compute distance from, should be broadcast compatible |
| with 'vec1' |
| epsilon: distance is clipped from below to be at least epsilon |
| |
| Returns: |
| Array of square euclidean distances; |
| shape will be result of broadcasting 'vec1' and 'vec2' |
| """ |
| difference = vec1 - vec2 |
| distance = difference.dot(difference) |
| if epsilon: |
| distance = jnp.maximum(distance, epsilon) |
| return distance |
|
|
|
|
| def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float: |
| return vector1.dot(vector2) |
|
|
|
|
| def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float: |
| return vector1.cross(vector2) |
|
|
|
|
| def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float: |
| return vector.norm(epsilon) |
|
|
|
|
| def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array: |
| return vector.normalized(epsilon) |
|
|
|
|
| def euclidean_distance( |
| vec1: Vec3Array, vec2: Vec3Array, epsilon: float = 1e-6 |
| ) -> Float: |
| """Computes euclidean distance between 'vec1' and 'vec2'. |
| |
| Args: |
| vec1: Vec3Array to compute euclidean distance to |
| vec2: Vec3Array to compute euclidean distance from, should be broadcast |
| compatible with 'vec1' |
| epsilon: distance is clipped from below to be at least epsilon |
| |
| Returns: |
| Array of euclidean distances; |
| shape will be result of broadcasting 'vec1' and 'vec2' |
| """ |
| distance_sq = square_euclidean_distance(vec1, vec2, epsilon**2) |
| distance = jnp.sqrt(distance_sq) |
| return distance |
|
|
|
|
| def dihedral_angle( |
| a: Vec3Array, b: Vec3Array, c: Vec3Array, d: Vec3Array |
| ) -> Float: |
| """Computes torsion angle for a quadruple of points. |
| |
| For points (a, b, c, d), this is the angle between the planes defined by |
| points (a, b, c) and (b, c, d). It is also known as the dihedral angle. |
| |
| Arguments: |
| a: A Vec3Array of coordinates. |
| b: A Vec3Array of coordinates. |
| c: A Vec3Array of coordinates. |
| d: A Vec3Array of coordinates. |
| |
| Returns: |
| A tensor of angles in radians: [-pi, pi]. |
| """ |
| v1 = a - b |
| v2 = b - c |
| v3 = d - c |
|
|
| c1 = v1.cross(v2) |
| c2 = v3.cross(v2) |
| c3 = c2.cross(c1) |
|
|
| v2_mag = v2.norm() |
| return jnp.arctan2(c3.dot(v2), v2_mag * c1.dot(c2)) |
|
|
|
|
| def random_gaussian_vector(shape, key, dtype=jnp.float32) -> Vec3Array: |
| vec_array = jax.random.normal(key, shape + (3,), dtype) |
| return Vec3Array.from_array(vec_array) |
|
|
