feat: separate velocity autocorrelation function

pydiffuser/tracer/component.py

@@ -52,13 +52,13 @@ def get_cosine_moment(
     ) -> float:
         raise NotImplementedError
 
-    @helpers.checktime()
+    @abc.abstractmethod
     def get_velocity_autocorrelation(
         self, lagtime: int, rolling: bool = True, epsilon: int = 1
     ) -> float:
-        """Normalized velocity autocorrelation function"""
+        """_summary_"""
 
-        return self.get_cosine_moment(lagtime, 1, rolling, epsilon)
+        raise NotImplementedError
 
     @helpers.checktime()
     def get_real_time(self, lagtime: int) -> float:

pydiffuser/tracer/ensemble.py

@@ -175,6 +175,21 @@ def get_cosine_moment(
         m = _get_cosine_moment(self.microstate, lagtime, order, rolling, epsilon)
         return float(m)
 
+    @helpers.checktime()
+    def get_velocity_autocorrelation(
+        self, lagtime: int, rolling: bool = True, epsilon: int = 1
+    ) -> float:
+        if self.dimension != 2:
+            raise InvalidDimensionError(
+                f"Unsupported dimension {self.dimension} is encountered"
+            )
+        if self.length - lagtime - epsilon <= 0 or epsilon <= 0:
+            raise InvalidTimeError("Only positive integers are allowed for lagtime")
+        m = _get_velocity_autocorrelation(
+            self.microstate, lagtime, self.dt, rolling, epsilon
+        )
+        return float(m)
+
     def __reduce__(self) -> Tuple[Any, ...]:
         info = super().__reduce__()
         args = self.dt
@@ -245,3 +260,26 @@ def _fix(_cos: Array) -> Array:
     return jnp.mean(
         jnp.mean(jnp.cos(order * jnp.arccos(_fix(cos))), axis=-1)
     )  # TA -> EA
+
+
+@partial(jit, static_argnums=(1, 2, 3, 4))
+def _get_velocity_autocorrelation(
+    x: LongLongPosType,
+    lagtime: int,
+    dt: ConstType,
+    rolling: bool = True,
+    epsilon: int = 1,
+) -> Array:
+    x = jnp.asarray(x)
+    vel = jnp.roll(x, shift=-epsilon, axis=NDAXIS.L) - x
+    vel = vel[:, :-epsilon] / (epsilon * dt)
+    if not rolling:
+        dot = jnp.sum(vel[:, lagtime] * vel[:, 0], axis=-1)
+        return jnp.mean(dot)
+    mul = jnp.roll(vel, shift=-lagtime, axis=NDAXIS.L) * vel
+    dot = (
+        jnp.sum(mul, axis=NDAXIS.D)
+        if not lagtime
+        else jnp.sum(mul[:, :-lagtime], axis=NDAXIS.D)
+    )
+    return jnp.mean(jnp.mean(dot, axis=-1))  # TA -> EA

pydiffuser/tracer/trajectory.py

@@ -98,6 +98,21 @@ def get_cosine_moment(
         m = _get_cosine_moment(self.position, lagtime, order, rolling, epsilon)
         return float(m)
 
+    @helpers.checktime()
+    def get_velocity_autocorrelation(
+        self, lagtime: int, rolling: bool = True, epsilon: int = 1
+    ) -> float:
+        if self.dimension != 2:
+            raise InvalidDimensionError(
+                f"Unsupported dimension {self.dimension} is encountered"
+            )
+        if self.length - lagtime - epsilon <= 0 or epsilon <= 0:
+            raise InvalidTimeError("Only positive integers are allowed for lagtime")
+        m = _get_velocity_autocorrelation(
+            self.position, lagtime, self.dt, rolling, epsilon
+        )
+        return float(m)
+
     def to_npy(self, npy_path: PathType) -> None:
         jnp.save(file=npy_path, arr=self.position)
 
@@ -159,3 +174,26 @@ def _fix(_cos: Array) -> Array:
         else jnp.sum(mul[:-lagtime], axis=NDAXIS.D)
     )
     return jnp.mean(jnp.cos(order * jnp.arccos(_fix(cos))))
+
+
+@partial(jit, static_argnums=(1, 2, 3, 4))
+def _get_velocity_autocorrelation(
+    x: LongPosType,
+    lagtime: int,
+    dt: ConstType,
+    rolling: bool = True,
+    epsilon: int = 1,
+) -> Array:
+    x = jnp.asarray(x)
+    vel = jnp.roll(x, shift=-epsilon, axis=NDAXIS.L) - x
+    vel = vel[:-epsilon] / (epsilon * dt)
+    if not rolling:
+        dot = jnp.sum(vel[lagtime] * vel[0], axis=-1)
+        return dot
+    mul = jnp.roll(vel, shift=-lagtime, axis=NDAXIS.L) * vel
+    dot = (
+        jnp.sum(mul, axis=NDAXIS.D)
+        if not lagtime
+        else jnp.sum(mul[:-lagtime], axis=NDAXIS.D)
+    )
+    return jnp.mean(dot, axis=-1)