diff --git a/src/beignet/func/_molecular_dynamics/_space/__inverse_transform.py b/src/beignet/func/_molecular_dynamics/_space/__inverse_transform.py deleted file mode 100644 index 649bfbe3f6..0000000000 --- a/src/beignet/func/_molecular_dynamics/_space/__inverse_transform.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch -from torch import Tensor - - -def _inverse_transform(transformation: Tensor) -> Tensor: - """ - Calculates the inverse of an affine transformation matrix. - - Parameters - ---------- - transformation : Tensor - The affine transformation matrix to be inverted. - - Returns - ------- - Tensor - The inverse of the given affine transformation matrix. - """ - if transformation.ndim in {0, 1}: - return 1.0 / transformation - - if transformation.ndim == 2: - return torch.linalg.inv(transformation) - - raise ValueError("Unsupported transformation dimensions.") diff --git a/src/beignet/func/_molecular_dynamics/_space/__transform.py b/src/beignet/func/_molecular_dynamics/_space/__transform.py deleted file mode 100644 index ff4a527493..0000000000 --- a/src/beignet/func/_molecular_dynamics/_space/__transform.py +++ /dev/null @@ -1,45 +0,0 @@ -import torch -from torch import Tensor - - -def _transform(transformation: Tensor, position: Tensor) -> Tensor: - """ - Applies an affine transformation to the position vector. - - Parameters - ---------- - position : Tensor - Position, must have the shape `(..., dimension)`. - - transformation : Tensor - The affine transformation matrix, must be a scalar, a vector, or a - matrix with the shape `(dimension, dimension)`. - - Returns - ------- - Tensor - Affine transformed position vector, has the same shape as the - position vector. - """ - if transformation.ndim == 0: - return position * transformation - - indices = [chr(ord("a") + index) for index in range(position.ndim - 1)] - - indices = "".join(indices) - - if transformation.ndim == 1: - return torch.einsum( - f"i,{indices}i->{indices}i", - transformation, - position, - ) - - if transformation.ndim == 2: - return torch.einsum( - f"ij,{indices}j->{indices}i", - transformation, - position, - ) - - raise ValueError("Unsupported transformation dimensions.") diff --git a/src/beignet/func/_molecular_dynamics/_space/_space.py b/src/beignet/func/_molecular_dynamics/_space/_space.py index 8019947b8e..99257253d0 100644 --- a/src/beignet/func/_molecular_dynamics/_space/_space.py +++ b/src/beignet/func/_molecular_dynamics/_space/_space.py @@ -8,9 +8,6 @@ import torch from torch import Tensor -from .__inverse_transform import _inverse_transform -from .__transform import _transform -from ._transform import transform T = TypeVar("T") @@ -43,7 +40,7 @@ def space( affine transformation, $T$, specified in one of three ways: a cube, $L$; an orthorhombic unit cell, $[L_{x}, L_{y}, L_{z}]$; or a triclinic cell, upper triangular matrix. If `parallelepiped` is `False`, must be - the edge lengths. If `None`, the simulation space has free boudnary + the edge lengths. If `None`, the simulation space has free boundary conditions. normalized : bool, default=True @@ -110,7 +107,6 @@ def space( dimensions = torch.tensor([dimensions]) if dimensions is None: - def _displacement_fn( a: Tensor, b: Tensor, @@ -138,7 +134,6 @@ def _shift_fn(a: Tensor, b: Tensor, **_) -> Tensor: inverse_transformation = _inverse_transform(dimensions) if normalized: - def _displacement_fn( a: Tensor, b: Tensor, @@ -176,7 +171,6 @@ def _displacement_fn( return displacement if remapped: - def _u(a: Tensor, b: Tensor) -> Tensor: return torch.remainder(a + b, 1.0) @@ -258,7 +252,6 @@ def _displacement_fn( return displacement if remapped: - def _u(a: Tensor, b: Tensor) -> Tensor: return torch.remainder(a + b, 1.0) @@ -315,11 +308,9 @@ def _displacement_fn( return displacement if remapped: - def _shift_fn(a: Tensor, b: Tensor, **_) -> Tensor: return torch.remainder(a + b, dimensions) else: - def _shift_fn(a: Tensor, b: Tensor, **_) -> Tensor: return a + b diff --git a/src/beignet/func/_molecular_dynamics/_space/_transform.py b/src/beignet/func/_molecular_dynamics/_space/_transform.py deleted file mode 100644 index 9628f57995..0000000000 --- a/src/beignet/func/_molecular_dynamics/_space/_transform.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import Tuple - -from torch import Tensor -from torch.autograd import Function - -from .__transform import _transform - - -def transform(transformation: Tensor, position: Tensor) -> Tensor: - """ - Return affine transformed position. - - Parameters - ---------- - transformation : Tensor - Affine transformation matrix, must have shape - `(dimension, dimension)`. - - position : Tensor - Position, must have shape `(..., dimension)`. - - Returns - ------- - Tensor - Affine transformed position of shape `(..., dimension)`. - """ - - class _Transform(Function): - generate_vmap_rule = True - - @staticmethod - def forward(transformation: Tensor, position: Tensor) -> Tensor: - """ - Return affine transformed position. - - Parameters - ---------- - transformation : Tensor - Affine transformation matrix, must have shape - `(dimension, dimension)`. - - position : Tensor - Position, must have shape `(..., dimension)`. - - Returns - ------- - Tensor - Affine transformed position of shape `(..., dimension)`. - """ - return _transform(transformation, position) - - @staticmethod - def setup_context(ctx, inputs, output): - transformation, position = inputs - - ctx.save_for_backward(transformation, position, output) - - @staticmethod - def jvp( - ctx, - grad_transformation: Tensor, - grad_position: Tensor, - ) -> Tuple[Tensor, Tensor]: - transformation, position, _ = ctx.saved_tensors - - output = _transform(transformation, position) - - grad_output = grad_position + _transform( - grad_transformation, - position, - ) - - return output, grad_output - - @staticmethod - def backward(ctx, grad_output: Tensor) -> Tuple[Tensor, Tensor]: - _, _, output = ctx.saved_tensors - - return output, grad_output - - return _Transform.apply(transformation, position)