Add new utility functions (#361)

p-zubieta · web-flow · commit 38732971c1ec · 2025-10-06T15:38:58.000Z
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/examples/openmm/metad/Metadynamics-ADP.ipynb b/examples/openmm/metad/Metadynamics-ADP.ipynb
@@ -351,7 +351,7 @@
    "outputs": [],
    "source": [
     "plot_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True)\n",
-    "xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower\n",
+    "xi = compute_mesh(plot_grid)\n",
     "\n",
     "alpha = 1 if method.deltaT is None else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT\n",
     "kT = kB * T.value_in_unit(unit.kelvin)\n",
diff --git a/examples/openmm/metad/Metadynamics-ADP.md b/examples/openmm/metad/Metadynamics-ADP.md
@@ -211,7 +211,7 @@ Next we use the biasing potential to estimate the free energy surface. For well-
 
 ```python id="6W7Xf0ilqAcm"
 plot_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True)
-xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower
+xi = compute_mesh(plot_grid)
 
 alpha = 1 if method.deltaT is None else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT
 kT = kB * T.value_in_unit(unit.kelvin)
diff --git a/examples/openmm/metad/alanine-dipeptide.py b/examples/openmm/metad/alanine-dipeptide.py
@@ -88,7 +88,7 @@ def get_args(argv):
         ("time-steps", "t", int, 5e5, "Number of simulation steps"),
     ]
     parser = argparse.ArgumentParser(description="Example script to run metadynamics")
-    for (name, short, T, val, doc) in available_args:
+    for name, short, T, val, doc in available_args:
         parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
     return parser.parse_args(argv)
 
@@ -126,7 +126,7 @@ def main(argv=[]):
 
     # generate CV values on a grid to evaluate bias potential
     plot_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True)
-    xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower
+    xi = compute_mesh(plot_grid)
 
     # determine bias factor depending on method (for standard = 1
     # and for well-tempered = (T+deltaT)/deltaT)
diff --git a/examples/openmm/metad/nacl/Metadynamics_NaCl.ipynb b/examples/openmm/metad/nacl/Metadynamics_NaCl.ipynb
@@ -358,7 +358,7 @@
    "source": [
     "# generate CV values on a grid to evaluate bias potential\n",
     "plot_grid = pysages.Grid(lower=(0), upper=(4), shape=(500), periodic=False)\n",
-    "xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower\n",
+    "xi = compute_mesh(plot_grid)\n",
     "xi = xi.flatten()\n",
     "\n",
     "alpha = 1 if method.deltaT is None else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT\n",
diff --git a/examples/openmm/metad/nacl/Metadynamics_NaCl.md b/examples/openmm/metad/nacl/Metadynamics_NaCl.md
@@ -221,7 +221,7 @@ Next we use the biasing potential to estimate the free energy surface. For well-
 ```python id="6W7Xf0ilqAcm"
 # generate CV values on a grid to evaluate bias potential
 plot_grid = pysages.Grid(lower=(0), upper=(4), shape=(500), periodic=False)
-xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower
+xi = compute_mesh(plot_grid)
 xi = xi.flatten()
 
 alpha = 1 if method.deltaT is None else (T.value_in_unit(unit.kelvin) + method.deltaT) / method.deltaT
diff --git a/examples/openmm/metad/nacl/nacl.py b/examples/openmm/metad/nacl/nacl.py
@@ -106,7 +106,7 @@ def get_args(argv):
         ("time-steps", "t", int, 5e5, "Number of simulation steps"),
     ]
     parser = argparse.ArgumentParser(description="Example script to run metadynamics")
-    for (name, short, T, val, doc) in available_args:
+    for name, short, T, val, doc in available_args:
         parser.add_argument("--" + name, "-" + short, type=T, default=T(val), help=doc)
     return parser.parse_args(argv)
 
@@ -147,7 +147,7 @@ def main(argv=[]):
 
     # Generate CV values on a grid to evaluate bias potential
     plot_grid = pysages.Grid(lower=(0), upper=(4), shape=(500), periodic=False)
-    xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower
+    xi = compute_mesh(plot_grid)
     xi = xi.flatten()
 
     # Determine bias factor depending on method
diff --git a/examples/openmm/spectral_abf/alanine-dipeptide.py b/examples/openmm/spectral_abf/alanine-dipeptide.py
@@ -104,7 +104,7 @@ def main(argv=[]):
 
     # generate CV values on a grid to evaluate bias potential
     plot_grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(64, 64), periodic=True)
-    xi = (compute_mesh(plot_grid) + 1) / 2 * plot_grid.size + plot_grid.lower
+    xi = compute_mesh(plot_grid)
 
     # Set min free energy to zero
     A = fes_fn(xi)
diff --git a/pysages/approxfun/__init__.py b/pysages/approxfun/__init__.py
@@ -14,4 +14,5 @@
     collect_exponents,
     compute_mesh,
     scale,
+    unit_mesh,
 )
diff --git a/pysages/approxfun/core.py b/pysages/approxfun/core.py
@@ -43,7 +43,7 @@ class AbstractFit(ABC):
     def __init__(self, grid: Grid):
         ns = collect_exponents(grid)
         self.grid = grid
-        self.mesh = compute_mesh(grid)
+        self.mesh = unit_mesh(grid)
         self.exponents = ns
         self.pinv = pinv(self)
 
@@ -104,7 +104,15 @@ def scale(x, grid: Grid):
     return (x - grid.lower) * 2 / grid.size - 1
 
 
-def compute_mesh(grid):
+def compute_mesh(grid: Grid):
+    """
+    Returns a dense mesh with the same shape and domain as `grid`.
+    """
+    mesh = unit_mesh(grid)
+    return (mesh + 1) / 2 * grid.size + grid.lower
+
+
+def unit_mesh(grid: Grid):
     """
     Returns a dense mesh with the same shape as `grid`, but on the hypercube [-1, 1]ⁿ,
     where `n` is the dimensionality of `grid`. The resulting mesh is Chebyshev-distributed
@@ -224,7 +232,7 @@ def pinv(model: SpectralSobolev1Fit):  # noqa: F811 # pylint: disable=C0116,E010
 def build_fitter(model: SpectralGradientFit):
     """
     Returns a function which takes an approximation `dy` to the gradient
-    of a function `f` evaluated over the set `x = compute_mesh(model.grid)`,
+    of a function `f` evaluated over the set `x = unit_mesh(model.grid)`,
     and returns a `fun: Fun` object which approximates `f` over
     the domain [-1, 1]ⁿ.
     """
@@ -253,7 +261,7 @@ def fit(dy):
 def build_fitter(model: SpectralSobolev1Fit):  # noqa: F811 # pylint: disable=C0116,E0102
     """
     Returns a function which takes approximations `y` and dy` to a function
-    `f` and its gradient evaluated over the set `x = compute_mesh(model.grid)`,
+    `f` and its gradient evaluated over the set `x = unit_mesh(model.grid)`,
     and in turn returns a `fun: Fun` object which approximates `f` over
     the domain [-1, 1]ⁿ.
     """
diff --git a/pysages/methods/analysis.py b/pysages/methods/analysis.py
@@ -95,7 +95,7 @@ def _analyze(result: Result, strategy: GradientLearning, topology):
     method = result.method
     states = result.states
     grid = method.grid
-    mesh = inputs = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    mesh = inputs = compute_mesh(grid)
 
     model = MLP(grid.shape.size, 1, topology, transform=partial(_scale, grid=grid))
     loss = GradientsSSE()
diff --git a/pysages/methods/ann.py b/pysages/methods/ann.py
@@ -182,7 +182,7 @@ def build_free_energy_learner(method: ANN):
     model = method.model
 
     # Training data
-    inputs = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    inputs = compute_mesh(grid)
     smoothing_kernel = blackman_kernel(dims, 7)
     padding = "wrap" if grid.is_periodic else "edge"
     conv = partial(convolve, kernel=smoothing_kernel, boundary=padding)
@@ -284,7 +284,7 @@ def analyze(result: Result[ANN]):
     method = result.method
 
     grid = method.grid
-    mesh = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    mesh = compute_mesh(grid)
     model = method.model
     _, layout = unpack(model.parameters)
 
diff --git a/pysages/methods/cff.py b/pysages/methods/cff.py
@@ -260,7 +260,7 @@ def build_free_energy_learner(method: CFF):
 
     dims = grid.shape.size
     shape = (*grid.shape, 1)
-    inputs = f32((compute_mesh(grid) + 1) * grid.size / 2 + grid.lower)
+    inputs = f32(compute_mesh(grid))
     smoothing_kernel = f32(blackman_kernel(dims, 7))
     padding = "wrap" if grid.is_periodic else "edge"
     conv = partial(convolve, kernel=smoothing_kernel, boundary=padding)
@@ -381,7 +381,7 @@ def analyze(result: Result[CFF]):
     states = result.states
 
     grid = method.grid
-    mesh = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    mesh = compute_mesh(grid)
     model = method.model
     _, layout = unpack(model.parameters)
 
diff --git a/pysages/methods/funn.py b/pysages/methods/funn.py
@@ -225,7 +225,7 @@ def build_free_energy_grad_learner(method: FUNN):
     model = method.model
 
     # Training data
-    inputs = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    inputs = compute_mesh(grid)
     smoothing_kernel = blackman_kernel(dims, 7)
     padding = "wrap" if grid.is_periodic else "edge"
     conv = partial(convolve, kernel=smoothing_kernel, boundary=padding)
diff --git a/pysages/methods/metad.py b/pysages/methods/metad.py
@@ -233,7 +233,7 @@ def next_height(pstate):
         update_grids = jit(lambda *args: (None, None))
         should_deposit = jit(lambda pred, _: pred)
     else:
-        grid_mesh = (compute_mesh(grid) + 1) * (grid.size / 2) + grid.lower
+        grid_mesh = compute_mesh(grid)
         get_grid_index = build_indexer(grid)
         # Reshape so the dimensions are compatible
         accum = jit(lambda total, val: total + val.reshape(total.shape))
diff --git a/pysages/methods/sirens.py b/pysages/methods/sirens.py
@@ -285,7 +285,7 @@ def build_free_energy_learner(method: Sirens):
 
     dims = grid.shape.size
     shape = (*grid.shape, 1)
-    inputs = f32((compute_mesh(grid) + 1) * grid.size / 2 + grid.lower)
+    inputs = f32(compute_mesh(grid))
     smoothing_kernel = f32(blackman_kernel(dims, 5))
     padding = "wrap" if grid.is_periodic else "edge"
     conv = partial(convolve, kernel=smoothing_kernel, boundary=padding)
@@ -456,7 +456,7 @@ def analyze(result: Result[Sirens]):
     method = result.method
 
     grid = method.grid
-    mesh = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    mesh = compute_mesh(grid)
     model = method.model
     _, layout = unpack(model.parameters)
 
diff --git a/pysages/methods/spectral_abf.py b/pysages/methods/spectral_abf.py
@@ -304,7 +304,7 @@ def analyze(result: Result[SpectralABF]):
     method = result.method
 
     grid = method.grid
-    mesh = (compute_mesh(grid) + 1) * grid.size / 2 + grid.lower
+    mesh = compute_mesh(grid)
     evaluate = build_evaluator(method.model)
 
     def average_forces(hist, Fsum):
diff --git a/pysages/typing.py b/pysages/typing.py
@@ -36,6 +36,7 @@
 Any = _typing.Any
 Callable = _typing.Callable
 List = _typing.List
+Iterable = _typing.Iterable
 NamedTuple = _typing.NamedTuple
 Optional = _typing.Optional
 Sequence = _typing.Sequence
diff --git a/pysages/utils/__init__.py b/pysages/utils/__init__.py
@@ -27,6 +27,10 @@
     first_or_all,
     gaussian,
     identity,
+    is_file,
+    last,
     linear_solver,
+    parse_array,
+    splitlines,
 )
 from .transformations import quaternion_from_euler, quaternion_matrix
diff --git a/pysages/utils/core.py b/pysages/utils/core.py
@@ -2,12 +2,13 @@
 # See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES
 
 from copy import deepcopy
+from pathlib import Path
 
 import numpy
 from jax import numpy as np
 from plum import Dispatcher
 
-from pysages.typing import JaxArray, Scalar
+from pysages.typing import Iterable, JaxArray, Scalar, Union
 from pysages.utils.compat import solve_pos_def
 
 # PySAGES main dispatcher
@@ -47,6 +48,16 @@ def identity(x):
     return x
 
 
+@dispatch
+def is_file(path: Path):
+    return path.is_file()
+
+
+@dispatch
+def is_file(path: str):
+    return is_file(Path(path))
+
+
 def first_or_all(seq):
     """
     Returns the only element of a sequence `seq` if its length is one,
@@ -55,6 +66,34 @@ def first_or_all(seq):
     return seq[0] if len(seq) == 1 else seq
 
 
+@dispatch
+def last(iterable, default_value=None):
+    elem = default_value
+    for elem in iterable:
+        pass
+    return elem
+
+
+@dispatch
+def last(a: Union[tuple, JaxArray]):
+    return a[-1]
+
+
+def parse_array(s, shape: tuple = (-1, 3), transpose: bool = False):
+    a = np.fromstring(s, sep=" ").reshape(shape)
+    return a.T if transpose else a
+
+
+@dispatch(precedence=1)
+def splitlines(lines: str):
+    return lines.splitlines()
+
+
+@dispatch
+def splitlines(lines: Iterable[str]):
+    return lines
+
+
 def eps(T: type = np.zeros(0).dtype):
     return np.finfo(T).eps
 
diff --git a/tests/test_approxfun.py b/tests/test_approxfun.py
@@ -10,6 +10,7 @@
     build_fitter,
     build_grad_evaluator,
     compute_mesh,
+    unit_mesh,
 )
 from pysages.grids import Chebyshev, Grid
 
@@ -30,7 +31,7 @@ def f(x):
 def test_fourier_approx():
     grid = Grid(lower=(-np.pi,), upper=(np.pi,), shape=(512,), periodic=True)
 
-    x = np.pi * compute_mesh(grid)
+    x = compute_mesh(grid)
 
     # Periodic function and its gradient
     y = vmap(f)(x.flatten()).reshape(x.shape)
@@ -82,7 +83,7 @@ def test_fourier_approx():
 def test_cheb_approx():
     grid = Grid[Chebyshev](lower=(-1.0,), upper=(1.0,), shape=(256,))
 
-    x = compute_mesh(grid)
+    x = unit_mesh(grid)
 
     y = vmap(g)(x.flatten()).reshape(x.shape)
     dy = vmap(grad(g))(x.flatten()).reshape(x.shape)
diff --git a/tests/test_ml.py b/tests/test_ml.py
@@ -7,6 +7,7 @@
 
 from pysages.approxfun import compute_mesh
 from pysages.approxfun import scale as _scale
+from pysages.approxfun import unit_mesh
 from pysages.grids import Chebyshev, Grid
 from pysages.ml.models import MLP, Siren
 from pysages.ml.objectives import L2Regularization, Sobolev1SSE
@@ -40,8 +41,7 @@ def test_siren_sobolev_training():
     grid = Grid(lower=(-np.pi,), upper=(np.pi,), shape=(64,), periodic=True)
     scale = partial(_scale, grid=grid)
 
-    x_scaled = compute_mesh(grid)
-    x = np.pi * x_scaled
+    x = compute_mesh(grid)
 
     # Periodic function and its gradient
     y = vmap(f)(x.flatten()).reshape(x.shape)
@@ -77,7 +77,7 @@ def test_siren_sobolev_training():
 def test_mlp_training():
     grid = Grid[Chebyshev](lower=(-1.0,), upper=(1.0,), shape=(64,))
 
-    x = compute_mesh(grid)
+    x = unit_mesh(grid)
 
     y = vmap(g)(x.flatten()).reshape(x.shape)
 
@@ -103,7 +103,7 @@ def test_mlp_training():
 def test_adam_optimizer():
     grid = Grid[Chebyshev](lower=(-1.0,), upper=(1.0,), shape=(128,))
 
-    x = compute_mesh(grid)
+    x = unit_mesh(grid)
 
     y = vmap(g)(x.flatten()).reshape(x.shape)
 

Original file line number	Diff line number	Diff line change
`@@ -14,4 +14,5 @@`
`14`	`14`	`collect_exponents,`
`15`	`15`	`compute_mesh,`
`16`	`16`	`scale,`
	`17`	`+ unit_mesh,`
`17`	`18`	`)`