[math] Add Extrapolation.shortest_distance()
@@ -125,6 +125,21 @@ class Extrapolation:
def native_grid_sample_mode(self) -> Union[str, None]:
return None
+ def shortest_distance(self, start: Tensor, end: Tensor, domain_size: Tensor):
+ """
+ Computes the shortest distance between two points.
+ Both points are assumed to lie within the domain
+
+ Args:
+ start: Start position.
+ end: End position.
+ domain_size: Domain side lengths as vector.
+
+ Returns:
+ Shortest distance from `start` to `end`.
+ """
+ return end - start
+
def __getitem__(self, item):
return self
@@ -508,6 +523,10 @@ class _PeriodicExtrapolation(_CopyExtrapolation):
lower = {dim: -lo for dim, (lo, _) in widths.items()}
return value.shift(lower, value.shape.after_pad(widths), lambda v: self.pad(v, widths), lambda b: ZERO.pad(b, widths))
+ def shortest_distance(self, start: Tensor, end: Tensor, domain_size: Tensor):
+ dx = end - start
+ return (dx + domain_size / 2) % domain_size - domain_size / 2
+
class _SymmetricExtrapolation(_CopyExtrapolation):
"""Mirror with the boundary value occurring twice."""
