Improving Pastas performance with caching#
This notebook shows how pastas performance can be improved by caching computation results.
import pandas as pd
import pastas as ps
ps.set_log_level("WARNING")
Load some test data for the examples.
head = pd.read_csv("data/heby_head.csv", index_col=0, parse_dates=True).squeeze(
"columns"
)
evap = pd.read_csv("data/heby_evap.csv", index_col=0, parse_dates=True).squeeze(
"columns"
)
prec = pd.read_csv("data/heby_prec.csv", index_col=0, parse_dates=True).squeeze(
"columns"
)
temp = pd.read_csv("data/heby_temp.csv", index_col=0, parse_dates=True).squeeze(
"columns"
)
If the cachetools module is installed, Pastas can cache intermediate results for certain stressmodels. The cache essentially works as a dictionary that checks if the input arguments to a function, e.g. StressModel.simulate() are already stored in the cache. If so, it returns the stored solution, otherwise it computes the solution and adds it to the cache.
Note: The tradeoff with caching is that it can speed up Pastas by skipping some computations, but it uses more internal memory to store those intermediate results.
By default, caching is turned off.
ps.get_use_cache()
False
You can turn it on with ps.set_use_cache(True):
ps.set_use_cache(True)
ps.get_use_cache() # show that value has changed
True
The cache is stored in StressModels under the ._cache attribute. The following StressModel should indicate it contains an LRUCache that is currently empty.
The size of the cache can be set when creating a StressModel using the max_cache_size parameter. The default size is 32, meaning 32 solutions can be stored. For models with many parameters or complex optimization, you may want to increase this value.
Note: if
cachetools is not available the ._cache attribute will be None.
# sm = ps.StressModel(model=None, stress=prec, rfunc=ps.Exponential(), name="test")
# sm._cache
Let’s use an example to show the effect of caching on the performance of ml.solve(). This is an example from example_snow.py, a time series model with a non-linear recharge model.
def build_model():
ml = ps.Model(head)
ps.RechargeModel(
model=ml,
prec=prec,
evap=evap,
recharge=ps.rch.FlexModel(snow=True),
rfunc=ps.Gamma(),
name="rch",
temp=temp,
)
ml.set_parameter("rch_kv", vary=False)
return ml
Here we create a function to solve the model using a two-step method. First we solve for the optimal parameters without a noise model, then using those estimates of the parameters, we add a noise model and solve for the parameters again.
def two_step_solve(ml):
tmin = "1985"
tmax = "2018"
ml.solve(tmin=tmin, tmax=tmax, fit_constant=False, report=False)
ps.ArNoiseModel(model=ml)
ml.set_parameter("rch_ks", vary=False)
ml.solve(tmin=tmin, tmax=tmax, fit_constant=False, initial=False, report=False)
Now let’s test how long it takes to solve the model without using the cache.
ps.set_use_cache(False) # turn off caching
ml = build_model()
t0 = %timeit -r 1 -n 1 -o two_step_solve(ml) # timing the solve
---------------------------------------------------------------------------
KeyboardInterrupt Traceback (most recent call last)
Cell In[8], line 3
1 ps.set_use_cache(False) # turn off caching
2 ml = build_model()
----> 3 t0 = get_ipython().run_line_magic('timeit', '-r 1 -n 1 -o two_step_solve(ml) # timing the solve')
File ~/.asdf/installs/python/3.14.6/lib/python3.14/timeit.py:205, in Timer.repeat(self, repeat, number)
203 r = []
204 for i in range(repeat):
--> 205 t = self.timeit(number)
206 r.append(t)
207 return r
File <magic-timeit>:1, in inner(_it, _timer)
----> 1 'Could not get source, probably due dynamically evaluated source code.'
Cell In[7], line 4, in two_step_solve(ml)
1 def two_step_solve(ml):
2 tmin = "1985"
3 tmax = "2018"
----> 4 ml.solve(tmin=tmin, tmax=tmax, fit_constant=False, report=False)
5 ps.ArNoiseModel(model=ml)
6 ml.set_parameter("rch_ks", vary=False)
7 ml.solve(tmin=tmin, tmax=tmax, fit_constant=False, initial=False, report=False)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/model.py:1017, in Model.solve(self, tmin, tmax, freq, warmup, solver, report, initial, weights, fit_constant, freq_obs, initialize, reset_settings, noise, **kwargs)
1014 LeastSquares(model=self)
1016 # Solve model
-> 1017 solve_success, result = self.solver.solve(weights=weights, **kwargs)
1018 # Update the parameters with the results from the optimization
1019 for column in result.columns:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/solver/least_squares.py:830, in LeastSquares.solve(self, weights, **kwargs)
815 bounds = Bounds(
816 lb=pmin,
817 ub=pmax,
818 keep_feasible=True,
819 )
821 objfunction = partial(
822 self.objfunction,
823 noise=noise,
(...) 827 callback=kwargs.pop("callback", None),
828 )
--> 830 self.result = least_squares(
831 fun=objfunction,
832 x0=initial[vary],
833 jac=self.jac,
834 bounds=bounds,
835 method=self.method,
836 ftol=self.ftol,
837 xtol=self.xtol,
838 gtol=self.gtol,
839 x_scale=self.x_scale,
840 loss=self.loss,
841 f_scale=self.f_scale,
842 max_nfev=self.max_nfev,
843 diff_step=self.diff_step,
844 tr_solver=self.tr_solver,
845 **kwargs,
846 )
848 self.pcov = DataFrame(
849 self.get_covariances(
850 self.result.jac,
(...) 856 columns=parameters.index,
857 )
859 # Prepare return values
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/_lib/_util.py:660, in _workers_wrapper.<locals>.inner(*args, **kwds)
658 with MapWrapper(_workers) as mf:
659 kwargs['workers'] = mf
--> 660 return func(*args, **kwargs)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_lsq/least_squares.py:1019, in least_squares(fun, x0, jac, bounds, method, ftol, xtol, gtol, x_scale, loss, f_scale, diff_step, tr_solver, tr_options, jac_sparsity, max_nfev, verbose, args, kwargs, callback, workers)
1015 result = call_minpack(vector_fun.fun, x0, vector_fun.jac, ftol, xtol, gtol,
1016 max_nfev, x_scale, jac_method=jac)
1018 elif method == 'trf':
-> 1019 result = trf(vector_fun.fun, vector_fun.jac, x0, f0, J0, lb, ub, ftol, xtol,
1020 gtol, max_nfev, x_scale, loss_function, tr_solver,
1021 tr_options.copy(), verbose, callback=callback_wrapped)
1023 elif method == 'dogbox':
1024 if tr_solver == 'lsmr' and 'regularize' in tr_options:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_lsq/trf.py:124, in trf(fun, jac, x0, f0, J0, lb, ub, ftol, xtol, gtol, max_nfev, x_scale, loss_function, tr_solver, tr_options, verbose, callback)
120 return trf_no_bounds(
121 fun, jac, x0, f0, J0, ftol, xtol, gtol, max_nfev, x_scale,
122 loss_function, tr_solver, tr_options, verbose, callback=callback)
123 else:
--> 124 return trf_bounds(
125 fun, jac, x0, f0, J0, lb, ub, ftol, xtol, gtol, max_nfev, x_scale,
126 loss_function, tr_solver, tr_options, verbose, callback=callback)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_lsq/trf.py:376, in trf_bounds(fun, jac, x0, f0, J0, lb, ub, ftol, xtol, gtol, max_nfev, x_scale, loss_function, tr_solver, tr_options, verbose, callback)
372 f_true = f.copy()
374 cost = cost_new
--> 376 J = jac(x)
377 njev += 1
379 if loss_function is not None:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_differentiable_functions.py:750, in VectorFunction.jac(self, x)
748 def jac(self, x):
749 self._update_x(x)
--> 750 self._update_jac()
751 if hasattr(self.J, "astype"):
752 # returns a copy so that downstream can't overwrite the
753 # internal attribute. But one can't copy a LinearOperator
754 return self.J.astype(self.J.dtype)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_differentiable_functions.py:719, in VectorFunction._update_jac(self)
716 else:
717 self._njev += 1
--> 719 self.J = self.jac_wrapped(xp_copy(self.x), f0=self.f)
720 self.J_updated = True
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_differentiable_functions.py:454, in _VectorJacWrapper.__call__(self, x, f0, **kwds)
452 self.njev += 1
453 elif self.jac in FD_METHODS:
--> 454 J, dct = approx_derivative(
455 self.fun,
456 x,
457 f0=f0,
458 **self.finite_diff_options,
459 )
460 self.nfev += dct['nfev']
462 if self.sparse_jacobian:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_numdiff.py:611, in approx_derivative(fun, x0, method, rel_step, abs_step, f0, bounds, sparsity, as_linear_operator, args, kwargs, full_output, workers)
609 with MapWrapper(workers) as mf:
610 if sparsity is None:
--> 611 J, _nfev = _dense_difference(fun_wrapped, x0, f0, h,
612 use_one_sided, method,
613 mf)
614 else:
615 if not issparse(sparsity) and len(sparsity) == 2:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_numdiff.py:739, in _dense_difference(fun, x0, f0, h, use_one_sided, method, workers)
736 u = next(gen)
738 f1 = next(f_evals)
--> 739 f2 = next(f_evals)
740 if one_sided:
741 dx.append(u[i] - x0[i])
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_numdiff.py:912, in _Fun_Wrapper.__call__(self, x)
909 if xp.isdtype(x.dtype, "real floating"):
910 x = xp.astype(x, self.x0.dtype)
--> 912 f = np.atleast_1d(self.fun(x, *self.args, **self.kwargs))
913 if f.ndim > 1:
914 raise RuntimeError("`fun` return value has "
915 "more than 1 dimension.")
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_differentiable_functions.py:424, in _VectorFunWrapper.__call__(self, x)
422 def __call__(self, x):
423 self.nfev += 1
--> 424 return np.atleast_1d(self.fun(x))
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/optimize/_lsq/least_squares.py:263, in _WrapArgsKwargs.__call__(self, x)
262 def __call__(self, x):
--> 263 return self.f(x, *self.args, **self.kwargs)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/solver/least_squares.py:772, in LeastSquares.objfunction(self, p, noise, weights, initial, vary, callback)
770 par = initial
771 par[vary] = p
--> 772 return misfit(
773 ml=self.model, p=par, noise=noise, weights=weights, callback=callback
774 )
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/solver/objective_function.py:45, in misfit(ml, p, noise, weights, callback, returnseparate)
43 rv = ml.noise(p) * ml._noise_weights(p)
44 else:
---> 45 rv = ml.residuals(p)
47 # Apply weights if provided
48 if weights is not None:
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/model.py:611, in Model.residuals(self, p, tmin, tmax, freq, warmup)
606 freq_obs = (
607 freq if self.settings["freq_obs"] is None else self.settings["freq_obs"]
608 )
610 # simulate model
--> 611 sim = self.simulate(
612 p=p, tmin=tmin, tmax=tmax, freq=freq, warmup=warmup, return_warmup=False
613 )
615 # Get the oseries calibration series
616 obs = self.observations(tmin=tmin, tmax=tmax, freq=freq_obs)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/model.py:531, in Model.simulate(self, p, tmin, tmax, freq, warmup, return_warmup)
529 istart = 0 # Track parameters index to pass to stressmodel object
530 for sm in self.stressmodels.values():
--> 531 contrib = sm.simulate(
532 p=p[istart : istart + sm.nparam],
533 tmin=sim_index.min(),
534 tmax=tmax,
535 freq=freq,
536 dt=dt,
537 )
538 sim = sim.add(contrib)
539 istart += sm.nparam
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/stressmodels.py:2092, in RechargeModel.simulate(self, p, tmin, tmax, freq, dt, istress)
2082 def simulate(
2083 self,
2084 p: ArrayLike | None = None,
(...) 2089 istress: int | None = None,
2090 ) -> Series:
2091 """Simulate the stressmodel's contribution."""
-> 2092 return self._simulate(tuple(p), tmin, tmax, freq, dt, istress)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/decorators.py:342, in conditional_cachedmethod.<locals>.decorator.<locals>.wrapper(self, *args, **kwargs)
340 return cached_func.__get__(self, type(self))(*args, **kwargs)
341 else:
--> 342 return func(self, *args, **kwargs)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/pastas/stressmodels.py:2150, in RechargeModel._simulate(self, p, tmin, tmax, freq, dt, istress)
2146 if self.stresses[istress].name is not None:
2147 name = f"{self.name} ({self.stresses[istress].name})"
2149 return Series(
-> 2150 data=fftconvolve(stress, b, "full")[: stress.size],
2151 index=self.prec.series.index,
2152 name=name,
2153 dtype=p.dtype,
2154 )
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/signal/_signaltools.py:712, in fftconvolve(in1, in2, mode, axes)
707 s2 = in2.shape
709 shape = [max((s1[i], s2[i])) if i not in axes else s1[i] + s2[i] - 1
710 for i in range(in1.ndim)]
--> 712 ret = _freq_domain_conv(xp, in1, in2, axes, shape, calc_fast_len=True)
714 return _apply_conv_mode(ret, s1, s2, mode, axes, xp=xp)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/signal/_signaltools.py:542, in _freq_domain_conv(xp, in1, in2, axes, shape, calc_fast_len)
539 sp1 = fft(in1, fshape, axes=axes)
540 sp2 = fft(in2, fshape, axes=axes)
--> 542 ret = ifft(sp1 * sp2, fshape, axes=axes)
544 if calc_fast_len:
545 fslice = tuple([slice(sz) for sz in shape])
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/fft/_backend.py:29, in _ScipyBackend.__ua_function__(method, args, kwargs)
27 if fn is None:
28 return NotImplemented
---> 29 return fn(*args, **kwargs)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/fft/_basic_backend.py:149, in irfftn(x, s, axes, norm, overwrite_x, workers, plan)
147 def irfftn(x, s=None, axes=None, norm=None,
148 overwrite_x=False, workers=None, *, plan=None):
--> 149 return _execute_nD('irfftn', _duccfft.irfftn, x, s=s, axes=axes, norm=norm,
150 overwrite_x=overwrite_x, workers=workers, plan=plan)
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/fft/_basic_backend.py:57, in _execute_nD(func_str, duccfft_func, x, s, axes, norm, overwrite_x, workers, plan)
55 if is_numpy(xp):
56 x = np.asarray(x)
---> 57 return duccfft_func(x, s=s, axes=axes, norm=norm,
58 overwrite_x=overwrite_x, workers=workers, plan=plan)
60 norm = _validate_fft_args(workers, plan, norm)
61 if hasattr(xp, 'fft'):
File ~/checkouts/readthedocs.org/user_builds/pastas/envs/latest/lib/python3.14/site-packages/scipy/fft/_duccfft/basic.py:218, in c2rn(forward, x, s, axes, norm, overwrite_x, workers, plan)
215 tmp, _ = tuple(_fix_shape(tmp, shape, axes))
217 # Note: overwrite_x is not utilized
--> 218 return pfft.c2r(tmp, axes, lastsize, forward, norm, None, workers)
KeyboardInterrupt:
And now with the cache:
ps.set_use_cache(True) # turn on caching
ml = build_model()
t1 = %timeit -r 1 -n 1 -o two_step_solve(ml) # timing the solve again
Note that the cache is now in use:
ml.stressmodels["rch"]._cache.currsize
And some statistics about the performance gain achieved by activating caching:
diff = t0.average - t1.average
speedup = diff / (t0.average) * 100
print(f"Model solve time reduced by {diff:.1f}s ({speedup:.0f}%) by using caching.")
Finally, another nice bonus when using caching is that the optimal solution is already stored, meaning any calls to functions in which the optimal solution is needed (e.g. plotting) will be extra fast.
ml.plots.results();
Important Notes#
Cache Invalidation: The cache does not automatically detect when you modify stress data or response function settings. If you change these after creating a model, you must manually clear the cache using
stressmodel._cache.clear(), or rebuild the stress model.
Cache Performance: The cache uses exact floating-point comparison for parameter values. During optimization, minor variations in parameter values (e.g., 1.0000000001 vs 1.0) create separate cache entries.
In normal workflows where you create a model, solve it, and use the results, cache invalidation is not an issue.
Temporarily enabling or disabling caching#
For operations where you know caching won’t help, you can temporarily disable caching using the context manager:
with ps.temporarily_disable_cache():
# Caching is disabled here
print(f"Caching is temporarily: {'enabled' if ps.get_use_cache() else 'disabled'}")
params = ml.get_parameters()
result = ml.simulate(params)
# Caching is automatically re-enabled after the block
print(f"Cache is now: {'enabled' if ps.get_use_cache() else 'disabled'}")
Or if you wish to temporarily enable caching:
ps.set_use_cache(False) # turn off caching
with ps.temporarily_enable_cache():
# Caching is enable here
print(f"Caching is temporarily: {'enabled' if ps.get_use_cache() else 'disabled'}")
params = ml.get_parameters()
result = ml.simulate(params)
# Caching is automatically re-enabled after the block
print(f"Cache is now: {'enabled' if ps.get_use_cache() else 'disabled'}")