Adaptive AAA#
This module provides the Iterative Sample Refinement (ISR) algorithm outlined in ‘A Framework to Compute Resonances Arising from Multiple Scattering’ (https://doi.org/10.1002/adts.202400989).
- diffaaable.adaptive.adaptive_aaa(z_k_0, f, evolutions=2, cutoff=None, tol=1e-09, mmax=100, radius=None, domain=None, f_k_0=None, sampling=None, prev_z_n=None, return_samples=False, aaa=None)[source]#
An 2x adaptive Antoulas–Anderson algorithm for rational approximation of meromorphic functions that are costly to evaluate.
The algorithm iteratively places additional sample points close to estimated positions of poles identified during the past iteration. By this refinement scheme the number of function evaluations can be reduced. A more detailed description of the iterative sample refinement (ISR) algorithm is provided in (https://doi.org/10.1002/adts.202400989).
It is JAX differentiable wrt. the approximated function f, via its other arguments besides z. f should be provided as a jax.tree_util.Partial with only positional arguments pre-filled!
- Parameters:
z_k_0 (np.ndarray) – Array of initial sample points
f (callable) – function to be approximated. When using gradients f should be provided as a jax.tree_util.Partial with only positional arguments pre-filled. Furthermore it should be compatible to jax.jvp.
evolutions (int) – Number of refinement iterations
cutoff (float) – Maximum absolute value a function evaluation should take to be regarded valid. Otherwise the sample point is discarded. Defaults to 1e10 times the median of f(z_k_0)
tol (float) – Tolerance used in AAA (see diffaaable.aaa)
radius (float) – Distance from the assumed poles for next samples
domain (tuple[complex, complex]) – Tuple of min (lower left) and max (upper right) values defining a rectangle in the complex plane. Assumed poles outside of the domain will not receive refinement.
f_k_0 (ndarray[tuple[int, ...], dtype[_ScalarType_co]]) – Allows user to provide f evaluated at z_k_0
sampling (callable) –
strategy to determine the next sample points. The function should accept the following arguments:
- z_n: np.ndarray
estimated poles
- prev_z_n: np.ndarray
previous estimated poles
- samples: np.ndarray
current sample points
- domain: Domain
domian in wich to refine poles
- radius: float
distance from the assumed poles for next samples
- randkey: jax.random.PRNGKey
random key for sampling
prev_z_n (np.ndarray) – the previous poles that will be passed to the first evaluation of the sampling strategy
return_samples (bool) – If True, the function returns the samples used for the AAA approximation and the function evaluations at these points at the 4t and 5th position.
aaa (callable) – The AAA variant to be used. By default diffaaable.aaa is used. If you want to use the tensor AAA, you can pass diffaaable.tensor.tensor_aaa.
mmax (int)
- Returns:
z_j (np.array) – chosen samples
f_j (np.array) – f(z_j)
w_j (np.array) – Weights of Barycentric Approximation
z_n (np.array) – Poles of Barycentric Approximation
z_k_final (np.array) – all sample points used for the AAA approximation. Only returned if return_samples is True.
f_k_final (np.array) – f(z_k_final) Only returned if return_samples is True.