From 0a5e946f1cfa7e563e72572a855a1d6e10d68ca4 Mon Sep 17 00:00:00 2001 From: Henry Isaacson Date: Fri, 19 Apr 2024 14:48:27 -0400 Subject: [PATCH] space --- src/beignet/func/__init__.py | 0 .../func/_molecular_dynamics/__init__.py | 0 .../_molecular_dynamics/_space/__init__.py | 0 .../func/_molecular_dynamics/_space/_space.py | 328 ++++++++++++++++++ 4 files changed, 328 insertions(+) create mode 100644 src/beignet/func/__init__.py create mode 100644 src/beignet/func/_molecular_dynamics/__init__.py create mode 100644 src/beignet/func/_molecular_dynamics/_space/__init__.py create mode 100644 src/beignet/func/_molecular_dynamics/_space/_space.py diff --git a/src/beignet/func/__init__.py b/src/beignet/func/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/beignet/func/_molecular_dynamics/__init__.py b/src/beignet/func/_molecular_dynamics/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/beignet/func/_molecular_dynamics/_space/__init__.py b/src/beignet/func/_molecular_dynamics/_space/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/beignet/func/_molecular_dynamics/_space/_space.py b/src/beignet/func/_molecular_dynamics/_space/_space.py new file mode 100644 index 0000000000..bd94b27260 --- /dev/null +++ b/src/beignet/func/_molecular_dynamics/_space/_space.py @@ -0,0 +1,328 @@ +from typing import ( + Callable, + Optional, + Tuple, + TypeVar, +) + +import torch +from torch import Tensor + +from .__inverse_transform import _inverse_transform +from .__transform import _transform +from ._transform import transform + +T = TypeVar("T") + + +def space( + dimensions: Optional[Tensor] = None, + *, + normalized: bool = True, + parallelepiped: bool = True, + remapped: bool = True, +) -> Tuple[Callable, Callable]: + r"""Define a simulation space. + + This function is fundamental in constructing simulation spaces derived from + subsets of $\mathbb{R}^{D}$ (where $D = 1$, $2$, or $3$) and is + instrumental in setting up simulation environments with specific + characteristics (e.g., periodic boundary conditions). The function returns + a a displacement function and a shift function to compute particle + interactions and movements in space. + + This function supports deformation of the simulation cell, crucial for + certain types of simulations, such as those involving finite deformations + or the computation of elastic constants. + + Parameters + ---------- + dimensions : Optional[Tensor], default=None + Dimensions of the simulation space. Interpretation varies based on the + value of `parallelepiped`. If `parallelepiped` is `True`, must be an + 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 + conditions. + + normalized : bool, default=True + If `True`, positions are stored in the unit cube. Displacements and + shifts are computed in a normalized simulation space and can be + transformed back to real simulation space using the provided affine + transformation matrix. If `False`, positions are expressed and + computations performed directly in the real simulation space. + + parallelepiped : bool, default=True + If `True`, the simulation space is defined as a ${1, 2, 3}$-dimensional + parallelepiped with periodic boundary conditions. If `False`, the space + is defined on a ${1, 2, 3}$-dimensional hypercube. + + remapped : bool, default=True + If `True`, positions and displacements are remapped to stay in the + bounds of the defined simulation space. A rempapped simulation space is + topologically equivalent to a torus, ensuring that particles exiting + one boundary re-enter from the opposite side. This is particularly + relevant for simulation spaces with periodic boundary conditions. + + Returns + ------- + Tuple[Callable[[Tensor, Tensor], Tensor], Callable[[Tensor, Tensor], Tensor]] + A tuple containing two functions: + + 1. The displacement function, $\overrightarrow{d}$, measures the + difference between two points in the simulation space, factoring in + the geometry and boundary conditions. This function is used to + calculate particle interactions and dynamics. + 2. The shift function, $u$, applies a displacement vector to a point + in the space, effectively moving it. This function is used to + update simulated particle positions. + + Examples + -------- + transformation = torch.tensor([10.0]) + + displacement_fn, shift_fn = space( + transformation, + normalized=False, + ) + + normalized_displacement_fn, normalized_shift_fn = space( + transformation, + normalized=True, + ) + + normalized_position = torch.rand([4, 3]) + + position = transformation * normalized_position + + displacement = torch.randn([4, 3]) + + torch.testing.assert_close( + displacement_fn(position[0], position[1]), + normalized_displacement_fn( + normalized_position[0], + normalized_position[1], + ), + ) + """ + if isinstance(dimensions, (int, float)): + dimensions = torch.tensor([dimensions]) + + if dimensions is None: + + def _displacement_fn( + a: Tensor, + b: Tensor, + *, + perturbation: Optional[Tensor] = None, + **_, + ) -> Tensor: + if len(a.shape) != 1: + raise ValueError + + if a.shape != b.shape: + raise ValueError + + if perturbation is not None: + return _transform(a - b, perturbation) + + return a - b + + def _shift_fn(a: Tensor, b: Tensor, **_) -> Tensor: + return a + b + + return _displacement_fn, _shift_fn + + if parallelepiped: + inverse_transformation = _inverse_transform(dimensions) + + if normalized: + + def _displacement_fn( + a: Tensor, + b: Tensor, + *, + perturbation: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + _transformation = dimensions + + _inverse_transformation = inverse_transformation + + if "transformation" in kwargs: + _transformation = kwargs["transformation"] + + if "updated_transformation" in kwargs: + _transformation = kwargs["updated_transformation"] + + if len(a.shape) != 1: + raise ValueError + + if a.shape != b.shape: + raise ValueError + + displacement = a - b + + displacement = torch.remainder(displacement + 1.0 * 0.5, 1.0) + + displacement = displacement - 1.0 * 0.5 + + displacement = transform(_transformation, displacement) + + if perturbation is not None: + return _transform(displacement, perturbation) + + return displacement + + if remapped: + + def _u(a: Tensor, b: Tensor) -> Tensor: + return torch.remainder(a + b, 1.0) + + def _shift_fn(a: Tensor, b: Tensor, **kwargs) -> Tensor: + _transformation = dimensions + + _inverse_transformation = inverse_transformation + + if "transformation" in kwargs: + _transformation = kwargs["transformation"] + + _inverse_transformation = _inverse_transform( + _transformation + ) + + if "updated_transformation" in kwargs: + _transformation = kwargs["updated_transformation"] + + return _u(a, transform(_inverse_transformation, b)) + + return _displacement_fn, _shift_fn + + def _shift_fn(a: Tensor, b: Tensor, **kwargs) -> Tensor: + _transformation = dimensions + + _inverse_transformation = inverse_transformation + + if "transformation" in kwargs: + _transformation = kwargs["transformation"] + + _inverse_transformation = _inverse_transform( + _transformation, + ) + + if "updated_transformation" in kwargs: + _transformation = kwargs["updated_transformation"] + + return a + transform(_inverse_transformation, b) + + return _displacement_fn, _shift_fn + + def _displacement_fn( + a: Tensor, + b: Tensor, + *, + perturbation: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + _transformation = dimensions + + _inverse_transformation = inverse_transformation + + if "transformation" in kwargs: + _transformation = kwargs["transformation"] + + _inverse_transformation = _inverse_transform(_transformation) + + if "updated_transformation" in kwargs: + _transformation = kwargs["updated_transformation"] + + a = transform(_inverse_transformation, a) + b = transform(_inverse_transformation, b) + + if len(a.shape) != 1: + raise ValueError + + if a.shape != b.shape: + raise ValueError + + displacement = a - b + + displacement = torch.remainder(displacement + 1.0 * 0.5, 1.0) + + displacement = displacement - 1.0 * 0.5 + + displacement = transform(_transformation, displacement) + + if perturbation is not None: + return _transform(displacement, perturbation) + + return displacement + + if remapped: + + def _u(a: Tensor, b: Tensor) -> Tensor: + return torch.remainder(a + b, 1.0) + + def _shift_fn(a: Tensor, b: Tensor, **kwargs) -> Tensor: + _transformation = dimensions + + _inverse_transformation = inverse_transformation + + if "transformation" in kwargs: + _transformation = kwargs["transformation"] + + _inverse_transformation = _inverse_transform( + _transformation, + ) + + if "updated_transformation" in kwargs: + _transformation = kwargs["updated_transformation"] + + return transform( + _transformation, + _u( + transform(_inverse_transformation, a), + transform(_inverse_transformation, b), + ), + ) + + return _displacement_fn, _shift_fn + + def _shift_fn(a: Tensor, b: Tensor, **_) -> Tensor: + return a + b + + return _displacement_fn, _shift_fn + + def _displacement_fn( + a: Tensor, + b: Tensor, + *, + perturbation: Optional[Tensor] = None, + **_, + ) -> Tensor: + if len(a.shape) != 1: + raise ValueError + + if a.shape != b.shape: + raise ValueError + + displacement = torch.remainder(a - b + dimensions * 0.5, dimensions) + + displacement = displacement - dimensions * 0.5 + + if perturbation is not None: + return _transform(displacement, perturbation) + + 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 + + return _displacement_fn, _shift_fn \ No newline at end of file