Add typing-extensions, revert removal of Array alias

See:
- python/mypy#7866
- https://peps.python.org/pep-0613/

setup.py

@@ -94,6 +94,7 @@ def build_extension(self, ext):
         "pymbar>3.0.4",
         "pyyaml",
         "scipy",
+        "typing-extensions",
     ],
     extras_require={
         "dev": [

timemachine/potentials/jax_utils.py

@@ -2,9 +2,12 @@
 import numpy as np
 from jax import vmap
 from numpy.typing import NDArray
+from typing_extensions import TypeAlias
 
+Array: TypeAlias = NDArray
 
-def get_all_pairs_indices(n: int) -> NDArray:
+
+def get_all_pairs_indices(n: int) -> Array:
     """all indices i, j such that i < j < n"""
     n_interactions = n * (n - 1) / 2
 
@@ -15,7 +18,7 @@ def get_all_pairs_indices(n: int) -> NDArray:
     return pairs
 
 
-def pairs_from_interaction_groups(group_a_indices: NDArray, group_b_indices: NDArray) -> NDArray:
+def pairs_from_interaction_groups(group_a_indices: Array, group_b_indices: Array) -> Array:
     """(a, b) for a in group_a_indices, b in group_b_indices"""
     n_interactions = len(group_a_indices) * len(group_b_indices)
 
@@ -46,7 +49,7 @@ def compute_lifting_parameter(lamb, lambda_plane_idxs, lambda_offset_idxs, cutof
     return w
 
 
-def augment_dim(x3: NDArray, w: NDArray) -> NDArray:
+def augment_dim(x3: Array, w: Array) -> Array:
     """(x,y,z) -> (x,y,z,w)"""
 
     d4 = jnp.expand_dims(w, axis=-1)

timemachine/potentials/nonbonded.py

@@ -5,6 +5,7 @@
 from jax import vmap
 from jax.scipy.special import erfc
 from numpy.typing import NDArray
+from typing_extensions import TypeAlias
 
 from timemachine.potentials import jax_utils
 from timemachine.potentials.jax_utils import (
@@ -15,6 +16,8 @@
     pairs_from_interaction_groups,
 )
 
+Array: TypeAlias = NDArray
+
 
 def switch_fn(dij, cutoff):
     return jnp.power(jnp.cos((jnp.pi * jnp.power(dij, 8)) / (2 * cutoff)), 2)
@@ -377,7 +380,7 @@ def coulomb_prefactor_on_atom(x_i, x_others, q_others, box=None, beta=2.0, cutof
     return prefactor_i
 
 
-def coulomb_prefactors_on_snapshot(x_ligand, x_env, q_env, box=None, beta=2.0, cutoff=np.inf) -> NDArray:
+def coulomb_prefactors_on_snapshot(x_ligand, x_env, q_env, box=None, beta=2.0, cutoff=np.inf) -> Array:
     """Map coulomb_prefactor_on_atom over atoms in x_ligand
 
     Parameters
@@ -431,7 +434,7 @@ def f_snapshot(coords, box):
     return vmap(f_snapshot)(traj, boxes)
 
 
-def coulomb_interaction_group_energy(q_ligand: NDArray, q_prefactors: NDArray) -> float:
+def coulomb_interaction_group_energy(q_ligand: Array, q_prefactors: Array) -> float:
     """Assuming q_prefactors = coulomb_prefactors_on_snapshot(x_ligand, ...),
     cheaply compute the energy of ligand-environment interaction group