flax_model/alphafold3/jax/geometry/rigid_matrix_vector.py

"""Rigid3Array Transformations represented by a Matrix and a Vector."""

from typing import Any, Final, Self, TypeAlias

from flax_model.alphafold3.jax.geometry import rotation_matrix
from flax_model.alphafold3.jax.geometry import struct_of_array
from flax_model.alphafold3.jax.geometry import utils
from flax_model.alphafold3.jax.geometry import vector
import jax
import jax.numpy as jnp


Float: TypeAlias = float | jnp.ndarray

VERSION: Final[str] = '0.1'


# Disabling name in pylint, since the relevant variable in math are typically
# referred to as X, Y in mathematical literature.
def _compute_covariance_matrix(
    row_values: vector.Vec3Array,
    col_values: vector.Vec3Array,
    weights: jnp.ndarray,
    epsilon=1e-6,
) -> jnp.ndarray:
  """Compute covariance matrix.

  The quantity computes is
  cov_xy = weighted_avg_i(row_values[i, x] col_values[j, y]).
  Here x and y run over the xyz coordinates.
  This is used to construct frames when aligning points.

  Args:
    row_values: Values used for rows of covariance matrix, shape [..., n_point]
    col_values: Values used for columns of covariance matrix, shape [...,
      n_point]
    weights: weights to weight points by, shape broacastable to [...]
    epsilon: small value to add to denominator to avoid Nan's when all weights
      are 0.

  Returns:
    Covariance Matrix as [..., 3, 3] array.
  """
  weights = jnp.asarray(weights)
  weights = jnp.broadcast_to(weights, row_values.shape)

  out = []

  normalized_weights = weights / (weights.sum(axis=-1, keepdims=True) + epsilon)

  weighted_average = lambda x: jnp.sum(normalized_weights * x, axis=-1)

  out.append(
      jnp.stack(
          (
              weighted_average(row_values.x * col_values.x),
              weighted_average(row_values.x * col_values.y),
              weighted_average(row_values.x * col_values.z),
          ),
          axis=-1,
      )
  )

  out.append(
      jnp.stack(
          (
              weighted_average(row_values.y * col_values.x),
              weighted_average(row_values.y * col_values.y),
              weighted_average(row_values.y * col_values.z),
          ),
          axis=-1,
      )
  )

  out.append(
      jnp.stack(
          (
              weighted_average(row_values.z * col_values.x),
              weighted_average(row_values.z * col_values.y),
              weighted_average(row_values.z * col_values.z),
          ),
          axis=-1,
      )
  )

  return jnp.stack(out, axis=-2)


@struct_of_array.StructOfArray(same_dtype=True)
class Rigid3Array:
  """Rigid Transformation, i.e. element of special euclidean group."""

  rotation: rotation_matrix.Rot3Array
  translation: vector.Vec3Array

  def __matmul__(self, other: Self) -> Self:
    new_rotation = self.rotation @ other.rotation
    new_translation = self.apply_to_point(other.translation)
    return Rigid3Array(new_rotation, new_translation)

  def inverse(self) -> Self:
    """Return Rigid3Array corresponding to inverse transform."""
    inv_rotation = self.rotation.inverse()
    inv_translation = inv_rotation.apply_to_point(-self.translation)
    return Rigid3Array(inv_rotation, inv_translation)

  def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Apply Rigid3Array transform to point."""
    return self.rotation.apply_to_point(point) + self.translation

  def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
    """Apply inverse Rigid3Array transform to point."""
    new_point = point - self.translation
    return self.rotation.apply_inverse_to_point(new_point)

  def compose_rotation(self, other_rotation: rotation_matrix.Rot3Array) -> Self:
    rot = self.rotation @ other_rotation
    trans = jax.tree.map(
        lambda x: jnp.broadcast_to(x, rot.shape), self.translation
    )
    return Rigid3Array(rot, trans)

  @classmethod
  def identity(cls, shape: Any, dtype: jnp.dtype = jnp.float32) -> Self:
    """Return identity Rigid3Array of given shape."""
    return cls(
        rotation_matrix.Rot3Array.identity(shape, dtype=dtype),
        vector.Vec3Array.zeros(shape, dtype=dtype),
    )  # pytype: disable=wrong-arg-count  # trace-all-classes

  def scale_translation(self, factor: Float) -> Self:
    """Scale translation in Rigid3Array by 'factor'."""
    return Rigid3Array(self.rotation, self.translation * factor)

  def to_array(self):
    rot_array = self.rotation.to_array()
    vec_array = self.translation.to_array()
    return jnp.concatenate([rot_array, vec_array[..., None]], axis=-1)

  @classmethod
  def from_array(cls, array):
    rot = rotation_matrix.Rot3Array.from_array(array[..., :3])
    vec = vector.Vec3Array.from_array(array[..., -1])
    return cls(rot, vec)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def from_array4x4(cls, array: jnp.ndarray) -> Self:
    """Construct Rigid3Array from homogeneous 4x4 array."""
    if array.shape[-2:] != (4, 4):
      raise ValueError(f'array.shape({array.shape}) must be [..., 4, 4]')
    rotation = rotation_matrix.Rot3Array(
        *(array[..., 0, 0], array[..., 0, 1], array[..., 0, 2]),
        *(array[..., 1, 0], array[..., 1, 1], array[..., 1, 2]),
        *(array[..., 2, 0], array[..., 2, 1], array[..., 2, 2]),
    )
    translation = vector.Vec3Array(
        array[..., 0, 3], array[..., 1, 3], array[..., 2, 3]
    )
    return cls(rotation, translation)  # pytype: disable=wrong-arg-count  # trace-all-classes

  @classmethod
  def from_point_alignment(
      cls,
      points_to: vector.Vec3Array,
      points_from: vector.Vec3Array,
      weights: Float | None = None,
      epsilon: float = 1e-6,
  ) -> Self:
    """Constructs Rigid3Array by finding transform aligning points.

    This constructs the optimal Rigid Transform taking points_from to the
    arrangement closest to points_to.

    Args:
      points_to: Points to align to.
      points_from: Points to align from.
      weights: weights for points.
      epsilon: epsilon used to regularize covariance matrix.

    Returns:
      Rigid Transform.
    """
    if weights is None:
      weights = 1.0

    def compute_center(value):
      return utils.weighted_mean(value=value, weights=weights, axis=-1)

    points_to_center = jax.tree.map(compute_center, points_to)
    points_from_center = jax.tree.map(compute_center, points_from)
    centered_points_to = points_to - points_to_center[..., None]
    centered_points_from = points_from - points_from_center[..., None]
    cov_mat = _compute_covariance_matrix(
        centered_points_to,
        centered_points_from,
        weights=weights,
        epsilon=epsilon,
    )
    rots = rotation_matrix.Rot3Array.from_svd(
        jnp.reshape(cov_mat, cov_mat.shape[:-2] + (9,))
    )

    translations = points_to_center - rots.apply_to_point(points_from_center)

    return cls(rots, translations)  # pytype: disable=wrong-arg-count  # trace-all-classes

  def __getstate__(self):
    return (VERSION, (self.rotation, self.translation))

  def __setstate__(self, state):
    version, (rot, trans) = state
    del version
    object.__setattr__(self, 'rotation', rot)
    object.__setattr__(self, 'translation', trans)