bunch of type checks (2)

fplanck/forces.py

@@ -1,16 +1,19 @@
-"""
-pre-defined convenience force functions
-"""
+"""Pre-defined convenience force functions."""
+
 import numpy as np
 from scipy.interpolate import RegularGridInterpolator
+import numpy.typing as npt
+
 
+def force_from_data(grid: npt.ArrayLike, data: npt.ArrayLike) -> npt.ArrayLike:
+ """Create a force function from data on a grid
 
-def force_from_data(grid, data):
- """create a force function from data on a grid
+ Args:
+ grid: list of grid arrays along each dimension
+ data: force data (shape [ndim, ...])
 
- Arguments:
- grid list of grid arrays along each dimension
- data force data (shape [ndim, ...])
+ Returns:
+ force array
  """
  grid = np.asarray(grid)
  if grid.ndim == data.ndim == 1:

fplanck/functions.py

@@ -1,12 +1,20 @@
-"""
-pre-defined convenience probability distribution functions
-"""
+"""Pre-defined convenience probability distribution functions."""
+from typing import Callable
+from typing import Any
 import numpy as np
 from fplanck.utility import value_to_vector
+import numpy.typing as npt
 
 
-def delta_function(r0):
- """a discrete equivalent of a dirac-delta function centered at r0"""
+def delta_function(r0: npt.ArrayLike) -> npt.ArrayLike:
+ """A discrete equivalent of a dirac-delta function centered at r0.
+
+ Args:
+ r0: #TODO
+
+ Returns:
+ #TODO
+ """
  r0 = np.atleast_1d(r0)
 
  def pdf(*args):
@@ -21,12 +29,17 @@ def pdf(*args):
  return pdf
 
 
-def gaussian_pdf(center, width):
- """A Gaussian probability distribution function
+def gaussian_pdf(
+ center: npt.ArrayLike | float, width: npt.ArrayLike | float
+) -> npt.ArrayLike:
+ """A Gaussian probability distribution function.
 
- Arguments:
- center center of Gaussian (scalar or vector)
- width width of Gaussian (scalar or vector)
+ Args:
+ center: center of Gaussian (scalar or vector)
+ width: width of Gaussian (scalar or vector)
+
+ Returns:
+ #TODO
  """
 
  center = np.atleast_1d(center)
@@ -44,11 +57,14 @@ def pdf(*args):
  return pdf
 
 
-def uniform_pdf(func=None):
- """A uniform probability distribution function
+def uniform_pdf(func: Callable[..., Any] | None = None) -> npt.ArrayLike:
+ """A uniform probability distribution function.
+
+ Args:
+ func: a boolean function specifying region of uniform probability (default: everywhere)
 
- Arguments:
- func a boolean function specifying region of uniform probability (default: everywhere)
+ Returns:
+ #TODO
  """
 
  def pdf(*args):

fplanck/potentials.py

@@ -1,17 +1,22 @@
 """
 pre-defined convenience potential functions
 """
+from typing import Callable
 import numpy as np
 from fplanck.utility import value_to_vector
 from scipy.interpolate import RegularGridInterpolator
+import numpy.typing as npt
 
 
-def harmonic_potential(center, k):
- """A harmonic potential
+def harmonic_potential(center: npt.ArrayLike | float, k: npt.ArrayLike | float):
+ """Return a harmonic potential.
 
- Arguments:
- center center of harmonic potential (scalar or vector)
- k spring constant of harmonic potential (scalar or vector)
+ Args:
+ center: center of harmonic potential (scalar or vector)
+ k: spring constant of harmonic potential (scalar or vector)
+
+ Returns:
+ #TODO
  """
 
  center = np.atleast_1d(center)
@@ -29,13 +34,18 @@ def potential(*args):
  return potential
 
 
-def gaussian_potential(center, width, amplitude):
- """A Gaussian potential
+def gaussian_potential(
+ center: npt.ArrayLike | float, width: npt.ArrayLike | float, amplitude: float
+):
+ """Return a Gaussian potential.
+
+ Args:
+ center: center of Gaussian (scalar or vector)
+ width: width of Gaussian (scalar or vector)
+ amplitude: amplitude of Gaussian, (negative for repulsive)
 
- Arguments:
- center center of Gaussian (scalar or vector)
- width width of Gaussian (scalar or vector)
- amplitude amplitude of Gaussian, (negative for repulsive)
+ Returns:
+ #TODO
  """
 
  center = np.atleast_1d(center)
@@ -53,12 +63,15 @@ def potential(*args):
  return potential
 
 
-def uniform_potential(func, U0):
- """A uniform potential
+def uniform_potential(func: Callable, U0: float):
+ """Return a uniform potential.
 
- Arguments:
- func a boolean function specifying region of uniform probability (default: everywhere)
- U0 value of the potential
+ Args:
+ func: a boolean function specifying region of uniform probability (default: everywhere)
+ U0: value of the potential
+
+ Returns:
+ #TODO
  """
 
  def potential(*args):
@@ -71,12 +84,15 @@ def potential(*args):
  return potential
 
 
-def potential_from_data(grid, data):
- """create a potential from data on a grid
+def potential_from_data(grid: npt.ArrayLike, data: npt.ArrayLike) -> npt.ArrayLike:
+ """Create a potential from data on a grid.
+
+ Args:
+ grid: list of grid arrays along each dimension
+ data: potential data
 
- Arguments:
- grid list of grid arrays along each dimension
- data potential data
+ Returns:
+ #TODO
  """
  grid = np.asarray(grid)
  if grid.ndim == data.ndim == 1:

fplanck/solver.py

@@ -4,7 +4,6 @@
 import numpy as np
 import numpy.typing as npt
 from scipy import constants
-from scipy.linalg import expm
 from scipy import sparse
 from scipy.sparse.linalg import expm, eigs, expm_multiply
 import fplanck
@@ -208,7 +207,9 @@ def steady_state(self) -> npt.ArrayLike:
 
  return steady
 
- def propagate(self, initial, time: float, normalize: bool = True, dense: bool = False) -> npt.ArrayLike:
+ def propagate(
+ self, initial, time: float, normalize: bool = True, dense: bool = False
+ ) -> npt.ArrayLike:
  """Propagate an initial probability distribution in time.
 
  Args:
@@ -231,7 +232,14 @@ def propagate(self, initial, time: float, normalize: bool = True, dense: bool =
 
  return pf.reshape(self.Ngrid)
 
- def propagate_interval(self, initial, tf: float, Nsteps: int | None = None, dt: float | None = None, normalize: bool = True):
+ def propagate_interval(
+ self,
+ initial,
+ tf: float,
+ Nsteps: int | None = None,
+ dt: float | None = None,
+ normalize: bool = True,
+ ):
  """Propagate an initial probability distribution over a time interval.
 
  Args:

fplanck/utility.py

@@ -30,7 +30,7 @@ def value_to_vector(value: npt.ArrayLike | float, ndim: int, dtype: type = float
  return vec
 
 
-#TODO: this function seems questionable
+# TODO: this function seems questionable
 def slice_idx(i: int, ndim: int, s0) -> tuple:
  """Return a boolean array for a ndim-1 slice along the i'th axis at value s0"""
  idx = [slice(None)] * ndim
@@ -39,7 +39,7 @@ def slice_idx(i: int, ndim: int, s0) -> tuple:
  return tuple(idx)
 
 
-#TODO: better name + what if funcs is empty
+# TODO: better name + what if funcs is empty
 def combine[T](*funcs: Callable[..., T]) -> Callable[..., T]:
  """Combine a collection of functions into a single function (for probability, potential, and force functions)."""
 
@@ -60,7 +60,9 @@ class Boundary(enum.Enum):
  PERIODIC = enum.auto()
 
 
-def vectorize_force(f: Callable[[npt.ArrayLike], npt.ArrayLike]) -> Callable[[npt.ArrayLike], npt.ArrayLike]:
+def vectorize_force(
+ f: Callable[[npt.ArrayLike], npt.ArrayLike]
+) -> Callable[[npt.ArrayLike], npt.ArrayLike]:
  """Decorator to vectorize a force function."""
  ndim = len(getfullargspec(f).args)
  signature = ",".join(["()"] * ndim)