Add additional dependencies for mypy in pre-commit

.pre-commit-config.yaml

@@ -27,8 +27,13 @@ repos:
  rev: v1.5.1
  hooks:
  - id: mypy
+ # include dependencies that export types (i.e. have a py.typed file in the root module) so that they can be used
+ # by mypy in pre-commit
  additional_dependencies:
+ - "jax==0.4.14"
  - "numpy==1.23.5"
+ - "scipy==1.10.1"
+ - "matplotlib==3.7.1"
 
  # Exclude custom_ops.py to work around clash with stub file when typechecking
  exclude: '^timemachine/lib/custom_ops.py$'

tests/nonbonded/test_nonbonded_jit.py

@@ -14,8 +14,8 @@ def test_nonbonded_reference_jittable(num_atom_idxs, rng: np.random.Generator):
 
  U_ref = Nonbonded(
  N,
- exclusion_idxs=jnp.zeros((0,)),
- scale_factors=jnp.zeros((0, 2)),
+ exclusion_idxs=np.zeros((0,), dtype=np.int32),
+ scale_factors=np.zeros((0, 2)),
  beta=1.0,
  cutoff=0.1,
  atom_idxs=np.arange(num_atom_idxs) if num_atom_idxs is not None else None,

tests/test_jax_nonbonded.py

@@ -153,7 +153,7 @@ def generate_random_inputs(n_atoms, dim, instance_flags=difficult_instance_flags
 
  min_dist = 0.1
  conf, box = resolve_clashes(conf, box, min_dist=min_dist)
- box = np.array(box)
+ box = jnp.array(box)
 
  cutoff = 1.2
  if instance_flags["randomize_cutoff"]:

timemachine/fe/energy_decomposition.py

@@ -7,6 +7,7 @@
 
 from timemachine.constants import BOLTZ, DEFAULT_TEMP
 from timemachine.lib.custom_ops import Potential
+from timemachine.potentials.types import Params
 
 Frames = TypeVar("Frames")
 Boxes = List[NDArray]
@@ -25,7 +26,7 @@ class EnergyDecomposedState(Generic[Frames]):
 
 def get_batch_u_fns(
  pots: Sequence[Potential],
- params: Sequence[NDArray],
+ params: Sequence[Params],
  temperature: float = DEFAULT_TEMP,
 ) -> List[Batch_u_fn]:
  """Get a list of functions that take in (coords, boxes), return reduced_potentials

timemachine/fe/free_energy.py

@@ -440,6 +440,7 @@ def get_water_sampler_params(initial_state: InitialState) -> NDArray:
  assert isinstance(summed_pot.potentials[ixn_group_idx], NonbondedInteractionGroup)
  water_params = summed_pot.params_init[ixn_group_idx]
  assert water_params.shape[1] == 4
+ assert isinstance(water_params, np.ndarray)
  return water_params
 
 

timemachine/fe/plots.py

@@ -491,7 +491,7 @@ def plot_hrex_replica_state_distribution_convergence(cumulative_replica_state_co
  ax.set_visible(False)
 
  fig.subplots_adjust(right=0.8, hspace=0.2, wspace=0.2)
- cbar_ax = fig.add_axes([0.85, 0.15, 0.02, 0.7])
+ cbar_ax = fig.add_axes((0.85, 0.15, 0.02, 0.7))
  fig.colorbar(p, cax=cbar_ax, label=r"$\log_{10}$(number of iterations)")
 
 

timemachine/fe/reweighting.py

@@ -5,36 +5,38 @@
  "interpret_as_mixture_potential",
 ]
 
-from typing import Any, Callable, Collection
+from typing import Callable, Collection
 
 import numpy as np
+from jax import Array
 from jax import numpy as jnp
 from jax.scipy.special import logsumexp
+from jax.typing import ArrayLike
 
 Samples = Collection
 Params = Collection
-Array = Any # see https://github.com/google/jax/issues/943
 Energies = Array
 
 BatchedReducedPotentialFxn = Callable[[Samples, Params], Energies]
 
 
-def log_mean(log_values: Array) -> float:
+def log_mean(log_values: ArrayLike) -> Array:
  """stable log(mean(values))
 
  log(mean(values))
  = log(sum(values / len(values)))
  = logsumexp(log(values) - log(len(values))
  """
+ log_values = jnp.asarray(log_values)
  return logsumexp(log_values - jnp.log(len(log_values)))
 
 
-def estimate_log_z_ratio(log_importance_weights: Array) -> float:
+def estimate_log_z_ratio(log_importance_weights: ArrayLike) -> Array:
  """stable log(mean(importance_weights))"""
  return log_mean(log_importance_weights)
 
 
-def one_sided_exp(delta_us: Array) -> float:
+def one_sided_exp(delta_us: ArrayLike) -> Array:
  """exponential averaging
 
  References
@@ -44,10 +46,11 @@ def one_sided_exp(delta_us: Array) -> float:
  """
  # delta_us = -log_importance_weights
  # delta_f = -log_z_ratio
+ delta_us = jnp.asarray(delta_us)
  return -estimate_log_z_ratio(-delta_us)
 
 
-def interpret_as_mixture_potential(u_kn: Array, f_k: Array, N_k: Array) -> Array:
+def interpret_as_mixture_potential(u_kn: ArrayLike, f_k: ArrayLike, N_k: ArrayLike) -> Array:
  r"""Interpret samples from multiple states k as if they originate from a single state
  defined as a weighted mixture:
 
@@ -103,6 +106,9 @@ def interpret_as_mixture_potential(u_kn: Array, f_k: Array, N_k: Array) -> Array
  [3] [Elvira+, 2019] Generalized multiple importance sampling
  https://arxiv.org/abs/1511.03095
  """
+ u_kn = jnp.asarray(u_kn)
+ f_k = jnp.asarray(f_k)
+
  # one-liner: mixture_u_n = -logsumexp(f_k - u_kn.T, b=N_k, axis=1)
 
  # expanding steps:
@@ -141,7 +147,7 @@ def construct_endpoint_reweighting_estimator(
  batched_u_1_fxn: BatchedReducedPotentialFxn,
  ref_params: Params,
  ref_delta_f: float,
-) -> Callable[[Params], float]:
+) -> Callable[[Params], Array]:
  """assuming
  * endpoint samples (samples_0, samples_1)
  * precise estimate of free energy difference at initial params
@@ -181,17 +187,17 @@ def construct_endpoint_reweighting_estimator(
  ref_u_0 = batched_u_0_fxn(samples_0, ref_params)
  ref_u_1 = batched_u_1_fxn(samples_1, ref_params)
 
- def endpoint_correction_0(params) -> float:
+ def endpoint_correction_0(params) -> Array:
  """estimate f(ref, 0) -> f(params, 0) by reweighting"""
  delta_us = batched_u_0_fxn(samples_0, params) - ref_u_0
  return one_sided_exp(delta_us)
 
- def endpoint_correction_1(params) -> float:
+ def endpoint_correction_1(params) -> Array:
  """estimate f(ref, 1) -> f(params, 1) by reweighting"""
  delta_us = batched_u_1_fxn(samples_1, params) - ref_u_1
  return one_sided_exp(delta_us)
 
- def estimate_delta_f(params: Params) -> float:
+ def estimate_delta_f(params: Params) -> Array:
  """estimate f(params, 1) - f(params, 0)
 
  using this thermodynamic cycle:
@@ -218,10 +224,10 @@ def estimate_delta_f(params: Params) -> float:
 
 def construct_mixture_reweighting_estimator(
  samples_n: Samples,
- u_ref_n: Array,
+ u_ref_n: ArrayLike,
  batched_u_0_fxn: BatchedReducedPotentialFxn,
  batched_u_1_fxn: BatchedReducedPotentialFxn,
-) -> Callable[[Params], float]:
+) -> Callable[[Params], Array]:
  r"""assuming
  * samples x_n from a distribution p_ref(x) \propto(exp(-u_ref(x))
  that has good overlap with BOTH p_0(params)(x) and p_1(params)(x),
@@ -275,19 +281,20 @@ def construct_mixture_reweighting_estimator(
  [3] Wieder et al. PyTorch implementation of differentiable reweighting in neutromeratio
  https://github.com/choderalab/neutromeratio/blob/2abf29f03e5175a988503b5d6ceeee8ce5bfd4ad/neutromeratio/parameter_gradients.py#L246-L267
  """
+ u_ref_n = jnp.asarray(u_ref_n)
  assert len(samples_n) == len(u_ref_n)
 
  def f_0(params):
  """estimate f(params, 0) - f(ref) by reweighting"""
  u_0_n = batched_u_0_fxn(samples_n, params)
  return one_sided_exp(u_0_n - u_ref_n)
 
- def f_1(params) -> float:
+ def f_1(params) -> Array:
  """estimate f(params, 1) - f(ref) by reweighting"""
  u_1_n = batched_u_1_fxn(samples_n, params)
  return one_sided_exp(u_1_n - u_ref_n)
 
- def estimate_delta_f(params) -> float:
+ def estimate_delta_f(params) -> Array:
  r"""estimate f(params, 1) - f(params, 0)
 
  using this thermodynamic cycle:

timemachine/fe/utils.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from PIL import Image
 from rdkit import Chem
 from rdkit.Chem import AllChem, Draw
 from rdkit.Chem.Draw import rdMolDraw2D
@@ -214,7 +213,7 @@ def plot_atom_mapping_grid(
  core: NDArray,
  num_rotations: int = 5,
  seed: int = 1234,
-) -> Image:
+):
  mol_a_3d = recenter_mol(mol_a)
  mol_b_3d = recenter_mol(mol_b)
 

timemachine/lib/fixed_point.py

@@ -1,13 +1,15 @@
 import jax.numpy as jnp
+from jax import Array
+from jax.typing import ArrayLike
 
 from timemachine.lib import custom_ops
 
 
-def fixed_to_float(v: int | jnp.uint64) -> jnp.float64:
+def fixed_to_float(v: ArrayLike) -> Array:
  """Meant to imitate the logic of timemachine/cpp/src/fixed_point.hpp::FIXED_TO_FLOAT"""
  return jnp.float64(jnp.int64(jnp.uint64(v))) / custom_ops.FIXED_EXPONENT
 
 
-def float_to_fixed(v: jnp.float32 | float) -> jnp.uint64:
+def float_to_fixed(v: ArrayLike) -> Array:
  """Meant to imitate the logic of timemachine/cpp/src/kernels/k_fixed_point.cuh::FLOAT_TO_FIXED"""
  return jnp.uint64(jnp.int64(v * custom_ops.FIXED_EXPONENT))

timemachine/maps/terminal_bonds.py

@@ -204,7 +204,7 @@ def contains_in_support(self, x) -> bool:
  interval = self.intervals[i]
  bond_valid.append((r <= interval.upper) * (r >= interval.lower))
 
- return jnp.array(bond_valid).all()
+ return jnp.array(bond_valid).all().item()
 
  @classmethod
  def from_harmonic_bond_params(cls, bond_idxs, params, temperature=DEFAULT_TEMP, sigma_thresh=DEFAULT_SIGMA_THRESH):

timemachine/optimize/protocol.py

@@ -40,17 +40,18 @@
  Gradient-based optimization of high-dimensional protocols, using a reweighting-based estimate of a
  a T.I.-tailored objective, stddev(du/dlambda).
 """
-from typing import Any, Callable, cast
+
+from typing import Callable, cast
 
 import numpy as np
-from jax import jit
+from jax import Array, jit
 from jax import numpy as jnp
 from jax import vmap
 from jax.scipy.special import logsumexp
+from jax.typing import ArrayLike
 from scipy.optimize import bisect
 
 Float = float
-Array = Any # see https://github.com/google/jax/issues/943
 DistanceFxn = Callable[[Float, Float], Float]
 WorkStddevEstimator = DistanceFxn
 
@@ -99,7 +100,7 @@ def rebalance_initial_protocol(
  return optimized_protocol
 
 
-def log_weights_from_mixture(u_kn: Array, f_k: Array, N_k: Array) -> Array:
+def log_weights_from_mixture(u_kn: ArrayLike, f_k: ArrayLike, N_k: ArrayLike) -> Array:
  r"""Assuming
  * K reduced potential energy functions u_k
  * N_k samples from each state e^{-u_k} / Z_k
@@ -111,6 +112,9 @@ def log_weights_from_mixture(u_kn: Array, f_k: Array, N_k: Array) -> Array:
  interpret the collection of N = \sum_k N_k samples as coming from a
  mixture of states p(x) = (1 / K) \sum_k e^-u_k / Z_k
  """
+ f_k = jnp.asarray(f_k)
+ u_kn = jnp.asarray(u_kn)
+
  log_q_k = f_k - u_kn.T
  N_k = np.array(N_k, dtype=np.float64) # may be ints, or in a list...
  log_weights = logsumexp(log_q_k, b=N_k, axis=1)
@@ -121,7 +125,7 @@ def linear_u_kn_interpolant(lambdas: Array, u_kn: Array) -> Callable:
  """Given a matrix u_kn[k, n] = u(xs[n], lambdas[k]) produce linear interpolated estimates of u(xs[n], lam)
  at arbitrary new values lam"""
 
- def u_interp(u_n: Array, lam: Float) -> Float:
+ def u_interp(u_n: ArrayLike, lam: ArrayLike) -> Array:
  return jnp.nan_to_num(jnp.interp(lam, lambdas, u_n), nan=+jnp.inf, posinf=+jnp.inf)
 
  @jit

timemachine/potentials/nonbonded.py

@@ -1,9 +1,10 @@
-from typing import Any
+from typing import Tuple, cast
 
 import jax.numpy as jnp
 import numpy as np
-from jax import jit, vmap
+from jax import Array, jit, vmap
 from jax.scipy.special import erfc
+from jax.typing import ArrayLike
 from numpy.typing import NDArray
 from scipy.special import binom
 
@@ -17,8 +18,6 @@
  process_traj_in_chunks,
 )
 
-Array = Any
-
 
 def switch_fn(dij, cutoff):
  return jnp.power(jnp.cos((jnp.pi * jnp.power(dij, 8)) / (2 * cutoff)), 2)
@@ -68,7 +67,7 @@ def nonbonded_block_unsummed(
  params_j: NDArray,
  beta: float,
  cutoff: float,
-):
+) -> Array:
  """
  This is a modified version of `nonbonded` that computes a block of
  NxM interactions between two sets of particles x_i and x_j. It is assumed that
@@ -129,7 +128,7 @@ def nonbonded_block_unsummed(
  lj = lennard_jones(dij, sig_ij, eps_ij)
 
  nrgs = jnp.where(dij < cutoff, es + lj, 0)
- return nrgs
+ return cast(Array, nrgs)
 
 
 def nonbonded_block(xi, xj, box, params_i, params_j, beta, cutoff):
@@ -330,7 +329,7 @@ def nonbonded_on_specific_pairs(
  beta: float,
  cutoff: Optional[float] = None,
  rescale_mask=None,
-):
+) -> Tuple[Array, Array]:
  """See `nonbonded` docstring for more details
 
  Notes
@@ -341,7 +340,7 @@ def nonbonded_on_specific_pairs(
  """
 
  if len(pairs) == 0:
- return np.zeros(1), np.zeros(1)
+ return jnp.zeros(1), jnp.zeros(1)
 
  inds_l, inds_r = pairs.T
 
@@ -377,7 +376,10 @@ def apply_cutoff(x):
  rescale_electrostatics = rescale_mask[:, 0]
  electrostatics = jnp.where(rescale_electrostatics != 0, electrostatics * rescale_electrostatics, 0)
 
- return vdW, electrostatics
+ vdW_arr = cast(Array, vdW)
+ electrostatics_arr = cast(Array, electrostatics)
+
+ return vdW_arr, electrostatics_arr
 
 
 def nonbonded_on_precomputed_pairs(
@@ -474,7 +476,7 @@ def validate_coulomb_cutoff(cutoff=1.0, beta=2.0, threshold=1e-2):
 # TODO: avoid repetition between this and lennard-jones
 
 
-def coulomb_prefactor_on_atom(x_i, x_others, q_others, box=None, beta=2.0, cutoff=jnp.inf) -> float:
+def coulomb_prefactor_on_atom(x_i, x_others, q_others, box=None, beta=2.0, cutoff=jnp.inf) -> Array:
  """Precompute part of (sum_i q_i * q_j / d_ij * rxn_field(d_ij)) that does not depend on q_i
 
  Parameters
@@ -557,7 +559,7 @@ def f_snapshot(coords, box):
  return process_traj_in_chunks(f_snapshot, traj, boxes, chunk_size)
 
 
-def coulomb_interaction_group_energy(q_ligand: Array, q_prefactors: Array) -> float:
+def coulomb_interaction_group_energy(q_ligand: ArrayLike, q_prefactors: ArrayLike) -> Array:
  """Assuming q_prefactors = coulomb_prefactors_on_snapshot(x_ligand, ...),
  cheaply compute the energy of ligand-environment interaction group