diffaaable 1.0.1__py3-none-any.whl

diffaaable/__init__.py

@@ -0,0 +1,10 @@
+"""diffaaable - JAX-differentiable AAA algorithm"""
+
+__version__ = "1.0.1"
+
+__all__ = ["aaa", "adaptive_aaa", "vectorial_aaa", "lorentz_aaa"]
+
+from diffaaable.core import aaa
+from diffaaable.adaptive import adaptive_aaa
+from diffaaable.vectorial import vectorial_aaa
+from diffaaable.lorentz import lorentz_aaa

diffaaable/adaptive.py

@@ -0,0 +1,305 @@
+from typing import Union
+import jax.numpy as np
+import numpy.typing as npt
+from diffaaable import aaa
+import jax
+from jax import random
+import matplotlib.pyplot as plt
+from functools import partial
+Domain = tuple[complex, complex]
+
+def top_right(a: npt.NDArray[complex], b: npt.NDArray[complex]):
+  return np.logical_and(a.imag>=b.imag, a.real>=b.real)
+
+def domain_mask(domain: Domain, z_n):
+    larger_min = top_right(z_n, domain[0])
+    smaller_max = top_right(domain[1], z_n)
+    return np.logical_and(larger_min, smaller_max)
+
+@jax.tree_util.Partial
+def next_samples(z_n, prev_z_n, z_k, domain: Domain, radius, randkey, tolerance=1e-9, min_samples=0, max_samples=0):
+  z_n = z_n[domain_mask(domain, z_n)]
+  movement = np.min(np.abs(z_n[:, None]-prev_z_n[None, :]), axis=-1)
+  ranking = np.argsort(-movement)
+  unstable = movement > tolerance
+  if np.sum(unstable) > min_samples:
+    if max_samples != 0 and np.sum(unstable)>max_samples:
+      z_n_unstable = z_n[ranking[:max_samples]]
+    else:
+      z_n_unstable = z_n[unstable]
+  else:
+    z_n_unstable = z_n[ranking[:min_samples]]
+
+  add_z_k = z_n_unstable
+  add_z_k += radius*np.exp(1j*2*np.pi*jax.random.uniform(randkey, add_z_k.shape))
+  return add_z_k
+
+
+def heat(poles, samples, mesh, sigma):
+  #jax.debug.print("poles: {}", poles)
+  f_p = np.nansum(
+    np.exp(-np.abs(poles[:, None, None]-mesh[None, :])**2/sigma**2),
+    axis=0
+  )
+
+  #jax.debug.print("{}", f_p)
+
+  f = f_p / np.nansum(
+    sigma**2/np.abs(mesh[None, :]-samples[:, None, None])**2,
+    axis=0
+  )
+  return f
+
+@partial(jax.jit, static_argnames=["resolution", "batchsize"])
+def _next_samples_heat(
+  poles, prev_poles, samples, domain, radius, randkey, resolution=(101, 101),
+  batchsize=1, stop=0.2
+  ):
+
+  x = np.linspace(domain[0].real, domain[1].real, resolution[0])
+  y = np.linspace(domain[0].imag, domain[1].imag, resolution[1])
+
+  X, Y = np.meshgrid(x,y, indexing="ij")
+  mesh = X +1j*Y
+
+  add_samples = np.empty(0, dtype=complex)
+  for j in range(batchsize):
+    heat_map = heat(poles, np.concat([samples, add_samples]), mesh, sigma=radius)
+    next_i = np.unravel_index(np.nanargmax(heat_map), heat_map.shape)
+
+    next = np.where(heat_map[next_i] < stop, np.nan, mesh[next_i])
+    add_samples = np.append(add_samples, next)
+
+  return add_samples, X, Y, heat(poles, np.concat([samples]), mesh, sigma=radius)
+
+@jax.tree_util.Partial
+def next_samples_heat(
+  poles, prev_poles, samples, domain, radius, randkey, resolution=(101, 101),
+  batchsize=1, stop=0.2, debug=False, debug_known_poles=None
+  ):
+
+  add_samples, X, Y, heat_map = _next_samples_heat(
+    poles, prev_poles, samples, domain, radius, randkey, resolution,
+    batchsize, stop
+  )
+
+  add_samples = add_samples[~np.isnan(add_samples)]
+
+  if debug:
+    ax = plt.gca()
+    plt.figure()
+    plt.title(f"radius={radius}")
+    if resolution[1]>1:
+      plt.pcolormesh(X, Y, heat_map, vmax=1)#, alpha=np.clip(heat_map, 0, 1))
+      plt.colorbar()
+      plt.scatter(samples.real, samples.imag, label="samples", zorder=1)
+      plt.scatter(poles.real, poles.imag, color="C1", marker="x", label="est. pole", zorder=2)
+      plt.scatter(add_samples.real, add_samples.imag, color="C2", label="next samples", zorder=3)
+      if debug_known_poles is not None:
+        plt.scatter(debug_known_poles.real, debug_known_poles.imag, color="C3", marker="+", label="known pole")
+      plt.xlim(domain[0].real, domain[1].real)
+      plt.ylim(domain[0].imag, domain[1].imag)
+      plt.legend(loc="lower right")
+    else:
+      plt.plot(np.squeeze(X), np.squeeze(heat_map))
+      plt.scatter(samples.real, np.zeros(len(samples)))
+      plt.scatter(poles.real, np.zeros(len(poles)), color="C1", marker="x", label="est. pole", zorder=2)
+      plt.xlim(domain[0].real, domain[1].real)
+    plt.savefig(f"{debug}/{len(samples)}.png")
+    plt.close()
+    plt.sca(ax)
+
+  return add_samples
+
+aaa = jax.tree_util.Partial(aaa)
+
+def _adaptive_aaa(z_k_0: npt.NDArray,
+                 f: callable,
+                 evolutions: int = 2,
+                 cutoff: float = None,
+                 tol: float = 1e-9,
+                 mmax: int = 100,
+                 radius: float = None,
+                 domain: tuple[complex, complex] = None,
+                 f_k_0: npt.NDArray = None,
+                 sampling: Union[callable, str] = next_samples,
+                 prev_z_n: npt.NDArray = None,
+                 return_samples: bool = False,
+                 aaa: callable = aaa,
+                 f_dot: callable = None):
+  """
+  Implementation of `adaptive_aaa`
+
+  Parameters
+  ----------
+    see `adaptive_aaa`
+
+  f_dot: callable
+    Tangent of `f`. If provided JVPs of `f` will be collected throughout the
+    iterations. For use in custom_jvp
+  """
+
+  if sampling == "heat":
+    sampling = next_samples_heat
+
+  collect_tangents = f_dot is not None
+  z_k = z_k_0
+  max_dist = np.max(np.abs(z_k_0[:, np.newaxis] - z_k_0[np.newaxis, :]))
+
+  if collect_tangents:
+    f_unpartial = f.func
+    args, _ = jax.tree.flatten(f)
+    args_dot, _ = jax.tree.flatten(f_dot)
+    z_k_dot = np.zeros_like(z_k)
+    f_k, f_k_dot = jax.jvp(
+      f_unpartial, (*args, z_k), (*args_dot, z_k_dot)
+    )
+  else:
+    if f_k_0 is None:
+      f_k = f(z_k)
+    else:
+      f_k = f_k_0
+    f_k_dot = np.zeros_like(f_k)
+
+  if cutoff is None:
+    cutoff = 1e10*np.median(np.abs(f_k))
+
+  if domain is None:
+    center = np.mean(z_k)
+    disp = max_dist*(1+1j)
+    domain = (center-disp, center+disp)
+
+  if radius is None:
+    radius = 1e-3 * max_dist
+
+  if prev_z_n is None:
+    prev_z_n = np.array([np.inf], dtype=complex)
+
+  def mask(z_k, f_k, f_k_dot):
+    m = np.abs(f_k)<cutoff
+    if m.ndim == 2:
+      m = np.all(m, axis=1) #filter out values, that have diverged too strongly
+    return z_k[m], f_k[m], f_k_dot[m]
+
+  key = random.key(0)
+  for i in range(evolutions):
+    z_k, f_k, f_k_dot = mask(z_k, f_k, f_k_dot)
+    z_j, f_j, w_j, z_n = aaa(z_k, f_k, tol, mmax)
+
+    if i==evolutions-1:
+      break
+
+    key, subkey = jax.random.split(key)
+    #print(f"{z_n=}")
+    #print(f"{prev_z_n=}")
+    add_z_k = sampling(z_n, prev_z_n, z_k, domain, radius, subkey)
+    prev_z_n = z_n
+    add_z_k_dot = np.zeros_like(add_z_k)
+
+    if len(add_z_k) == 0:
+      break
+
+    if collect_tangents:
+      f_k_new , f_k_dot_new = jax.jvp(
+        f_unpartial, (*args, add_z_k), (*args_dot, add_z_k_dot)
+      )
+    else:
+      f_k_new = f(add_z_k)
+      f_k_dot_new = np.zeros_like(f_k_new)
+
+    z_k = np.append(z_k, add_z_k)
+    f_k = np.concatenate([f_k, f_k_new])
+    f_k_dot = np.concatenate([f_k_dot, f_k_dot_new])
+
+  z_k, f_k, f_k_dot = mask(z_k, f_k, f_k_dot)
+
+  if collect_tangents:
+    return z_k, f_k, f_k_dot
+  if return_samples:
+    return z_j, f_j, w_j, z_n, z_k, f_k
+  return z_j, f_j, w_j, z_n
+
+@jax.custom_jvp
+def adaptive_aaa(z_k_0: npt.NDArray,
+                 f:callable,
+                 evolutions: int = 2,
+                 cutoff: float = None,
+                 tol: float = 1e-9,
+                 mmax: int = 100,
+                 radius: float = None,
+                 domain: Domain = None,
+                 f_k_0: npt.NDArray = None,
+                 sampling: callable = next_samples,
+                 prev_z_n: npt.NDArray = None,
+                 return_samples: bool = False,
+                 aaa: callable = aaa):
+  """ 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.
+
+  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 nex 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:
+      Allows user to provide f evaluated at z_k_0
+
+
+  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
+  """
+  return _adaptive_aaa(
+    z_k_0=z_k_0, f=f, evolutions=evolutions, cutoff=cutoff, tol=tol, mmax=mmax,
+    radius=radius, domain=domain, f_k_0=f_k_0, sampling=sampling,
+    prev_z_n=prev_z_n, return_samples=return_samples, aaa=aaa
+  )
+
+@adaptive_aaa.defjvp
+def adaptive_aaa_jvp(primals, tangents):
+  z_k_0, f = primals[:2]
+  z_dot, f_dot = tangents[:2]
+
+  if np.any(z_dot):
+    raise NotImplementedError(
+      "Parametrizing the sampling positions z_k is not supported"
+    )
+
+  z_k, f_k, f_k_dot = \
+    _adaptive_aaa(z_k_0, f, *primals[2:], f_dot=f_dot)
+
+  z_k_dot = np.zeros_like(z_k)
+
+  return jax.jvp(aaa, (z_k, f_k), (z_k_dot, f_k_dot))

diffaaable/core.py

@@ -0,0 +1,224 @@
+from jax import config
+config.update("jax_enable_x64", True) #important -> else aaa fails
+import jax.numpy as jnp
+import numpy.typing as npt
+import jax
+import numpy as np
+from baryrat import aaa as oaaa # ordinary aaa
+import functools
+import scipy.linalg
+
+@functools.wraps(oaaa)
+@jax.custom_jvp
+def aaa(z_k: npt.NDArray, f_k: npt.NDArray, tol: float=1e-13, mmax: int=100):
+  """
+  Wraped aaa to enable JAX based autodiff.
+  """
+  r = oaaa(z_k, f_k, tol=tol, mmax=mmax)
+  z_j = r.nodes
+  f_j = r.values
+  w_j = r.weights
+
+  mask = w_j!=0
+  z_j = z_j[mask]
+  f_j = f_j[mask]
+  w_j = w_j[mask]
+
+  z_n = poles(z_j, w_j)
+
+  z_n = z_n[jnp.argsort(-jnp.abs(z_n))]
+
+  return z_j, f_j, w_j, z_n
+
+delimiter = '    \n'
+
+aaa.__doc__ = f"""\
+  This is a wrapped version of `aaa` as provided by `baryrat`,
+  providing a custom jvp to enable differentiability.
+
+  For detailed information on the usage of `aaa` please refer to
+  the original documentation::
+
+    {delimiter.join(aaa.__doc__.splitlines())}
+
+  .. attention::
+    Returns nodes, values, weights and poles, in contrast to the
+    baryrat implementation that returns the BarycentricRational ready to be
+    evaluated. This is done to facilitate differentiability.
+
+  Parameters
+  ----------
+    z_k : array (M,)
+      the sampling points of the function. Unlike for interpolation
+      algorithms, where a small number of nodes is preferred, since the
+      AAA algorithm chooses its support points adaptively, it is better
+      to provide a finer mesh over the support.
+    f_k : array (M,)
+      the function to be approximated; can be given as a callable function
+      or as an array of function values over `z_k`.
+    tol : float
+      the approximation tolerance
+    mmax : int
+      the maximum number of iterations/degree of the resulting approximant
+
+  Returns
+  -------
+    z_j : array (m,)
+      nodes of the barycentric approximant
+
+    f_j : array (m,)
+      values of the barycentric approximant
+
+    w_j : array (m,)
+      weights of the barycentric approximant
+
+    z_n : array (m-1,)
+      poles of the barycentric approximant (for convenience)
+
+"""
+
+@aaa.defjvp
+def aaa_jvp(primals, tangents):
+  r"""Derivatives according to [1].
+  The implemented matrix expressions are motivated in the appendix of [2]:
+
+  .. topic:: AAA Derivatives
+
+    Here we will briefly elaborate how the derivatives introduced in [1] are
+    implemented as JAX compatible Jacobian Vector products (JVPs) in `diffaaable`.
+
+    Given the tangents $\frac{\partial f_k}{\partial p}$ we will use the chain rule on $r(w_j, f_j, z)$ along its weights $w_j$ and values $f_j$ (the nodes $z_j$ are treated as independent of $p$) (Equation 4 of [1]):
+    \[
+    \frac{\partial f_k}{\partial p} \approx \frac{\partial r_k}{\partial p}= \sum_{j=1}^m\frac{\partial r_k}{\partial f_j}\frac{\partial f_j}{\partial p}+\sum_{j=1}^m\frac{\partial r_k}{\partial w_j}\frac{\partial w_j}{\partial p}
+    \]
+
+    To solve this system of equations for $\frac{\partial w_j}{\partial p}$ we express it in matrix form:
+
+    .. math::
+      \mathbf{A}\mathbf{w}^\prime = \mathbf{b}
+
+    where $\mathbf{w}^\prime$ is the column vector containing $\frac{\partial w_j}{\partial p}$ and $\mathbf{b}$ and $\mathbf{A}$ are defined element wise:
+
+    .. math::
+      \begin{aligned}
+          b_k &= \frac{\partial f_k}{\partial p} - \sum_{j=1}^m\frac{\partial r_k}{\partial f_j}\frac{\partial f_j}{\partial p}\\
+          A_{kj} &= \frac{\partial r_k}{\partial w_j}
+      \end{aligned}
+
+
+    These are augmented with Equation 5 of \cite{betzEfficientRationalApproximation2024} which removes the ambiguity in $\mathbf{w}^\prime$ associated with a shared phase of all weights.
+
+    The expressions for the derivatives of $r_k$ are found in the definition of $r(z)$ (Equation 1 in [2]): \[ \frac{\partial r_k}{\partial f_j}= \frac{1}{d(z_k)} \frac{w_j}{z_k-z_j}\] and
+    \[
+    \frac{\partial r_k}{\partial w_j}= \frac{1}{d(z_k)} \frac{f_j-r_k}{z_k-z_j} \approx \frac{1}{d(z_k)} \frac{f_j-f_k}{z_k-z_j}
+    \]
+
+  """
+  z_k_full, f_k = primals[:2]
+  z_dot, f_dot = tangents[:2]
+
+  primal_out = aaa(z_k_full, f_k)
+  z_j, f_j, w_j, z_n = primal_out
+
+  chosen = np.isin(z_k_full, z_j)
+
+  z_k = z_k_full[~chosen]
+  f_k = f_k[~chosen]
+
+  # z_dot should be zero anyways
+  if np.any(z_dot):
+    raise NotImplementedError("Parametrizing the sampling positions z_k is not supported")
+  z_k_dot = z_dot[~chosen]
+  f_k_dot = f_dot[~chosen] # $\del f_k / \del p$
+
+  ##################################################
+  # We have to track which f_dot corresponds to z_k
+  sort_orig = jnp.argsort(jnp.abs(z_k_full[chosen]))
+  sort_out = jnp.argsort(jnp.argsort(jnp.abs(z_j)))
+
+  z_j_dot = z_dot[chosen][sort_orig][sort_out]
+  f_j_dot = f_dot[chosen][sort_orig][sort_out]
+  ##################################################
+
+  C = 1/(z_k[:, None]-z_j[None, :]) # Cauchy matrix k x j
+
+  d = C @ w_j # denominator in barycentric formula
+  via_f_j = C @ (f_j_dot * w_j) / d # $\sum_j f_j^\prime \frac{\del r}{\del f_j}$
+
+  A = (f_j[None, :] - f_k[:, None])*C/d[:, None]
+  b = f_k_dot - via_f_j
+
+  # make sure system is not underdetermined according to eq. 5 of [1]
+  A = jnp.concatenate([A, np.conj(w_j.reshape(1, -1))])
+  b = jnp.append(b, 0)
+
+  with jax.disable_jit(): #otherwise backwards differentiation led to error
+    w_j_dot, _, _, _ = jnp.linalg.lstsq(A, b)
+
+  denom = z_n.reshape(1, -1)-z_j.reshape(-1, 1)
+  # jax.debug.print("wj: {}", w_j.reshape(-1, 1))
+  # jax.debug.print("denom^2: {}", denom**2)
+  z_n_dot = (
+    jnp.sum(w_j_dot.reshape(-1, 1)/denom,    axis=0)/
+    jnp.sum(w_j.reshape(-1, 1)    /denom**2, axis=0)
+  )
+
+  tangent_out = z_j_dot, f_j_dot, w_j_dot, z_n_dot
+
+  return primal_out, tangent_out
+
+def poles(z_j,w_j):
+  """
+  The poles of a barycentric rational with given nodes and weights.
+  Poles lifted by zeros of the nominator are included.
+  Thus the values $f_j$ do not contribute and don't need to be provided
+  The implementation was modified from `baryrat` to support JAX AD.
+
+  Parameters
+  ----------
+    z_j : array (m,)
+      nodes of the barycentric rational
+    w_j : array (m,)
+      weights of the barycentric rational
+
+  Returns
+  -------
+    z_n : array (m-1,)
+      poles of the barycentric rational (more strictly zeros of the denominator)
+  """
+  f_j = np.ones_like(z_j)
+
+  B = np.eye(len(w_j) + 1)
+  B[0,0] = 0
+  E = np.block([[0, w_j],
+                [f_j[:,None], np.diag(z_j)]])
+  evals = scipy.linalg.eigvals(E, B)
+  return evals[np.isfinite(evals)]
+
+def residues(z_j,f_j,w_j,z_n):
+  '''
+  Residues for given poles via formula for simple poles
+  of quotients of analytic functions.
+  The implementation was modified from `baryrat` to support JAX AD.
+
+  Parameters
+  ----------
+    z_j : array (m,)
+      nodes of the barycentric rational
+    w_j : array (m,)
+      weights of the barycentric rational
+    z_n : array (n,)
+      poles of interest of the barycentric rational (n<=m-1)
+
+  Returns
+  -------
+    r_n : array (n,)
+      residues of poles `z_n`
+  '''
+
+  C_pol = 1.0 / (z_n[:,None] - z_j[None,:])
+  N_pol = C_pol.dot(f_j*w_j)
+  Ddiff_pol = (-C_pol**2).dot(w_j)
+  res = N_pol / Ddiff_pol
+
+  return jnp.nan_to_num(res)

diffaaable/lorentz.py

@@ -0,0 +1,88 @@
+from jax import config
+config.update("jax_enable_x64", True) #important -> else aaa fails
+import jax.numpy as np
+import jax
+from diffaaable.vectorial import check_inputs
+import jaxopt
+
+def optimal_weights(A, A_hat, stepsize=0.5):
+  # Initial guess (from uncorrected A)
+  _, _, Vh = np.linalg.svd(A)
+  w_j = Vh[-1, :].conj()
+
+  def obj_fun(w):
+     w /= np.linalg.norm(w)
+     err = (A @ w) + (A_hat @ w.conj())
+     return np.linalg.norm(err)
+
+  solver = jaxopt.LBFGS(fun=obj_fun, maxiter=300, tol=1e-13)
+  res = solver.run(w_j)
+  w_j, state = res
+  w_j /= np.linalg.norm(w_j)
+
+  return w_j
+
+def lorentz_aaa(z_k, f_k, tol=1e-9, mmax=100, return_errors=False):
+  """
+  """
+  z_k, f_k, M, V = check_inputs(z_k, f_k)
+
+  J = np.ones(M, dtype=bool)
+  z_j = np.empty(0, dtype=z_k.dtype)
+  f_j = np.empty((0, V), dtype=f_k.dtype)
+  errors = []
+
+  reltol = tol * np.linalg.norm(f_k, np.inf)
+
+  r_k = np.mean(f_k) * np.ones_like(f_k)
+  # approx.
+
+  for m in range(mmax):
+      # find largest residual
+      jj = np.argmax(np.linalg.norm(f_k - r_k, axis=-1)) #Next sample point to include
+      z_j = np.append(z_j, np.array([z_k[jj]]))
+      f_j = np.concatenate([f_j, f_k[jj][None, :]])
+      J = J.at[jj].set(False)
+
+      # Cauchy matrix containing the basis functions as columns
+      C = 1.0 / (z_k[J,None] - z_j[None,:])
+      # Loewner matrix
+      A = (f_k[J,None] - f_j[None,:]) * C[:,:,None]
+      A = np.concatenate(np.moveaxis(A, -1, 0))
+
+      # Lorentz Correction
+      C_hat = 1.0 / (z_k[J,None] + np.conj(z_j)[None,:])
+      # Loewner matrix
+      A_hat = (f_k[J,None] - np.conj(f_j)[None,:]) * C_hat[:,:,None]
+      A_hat = np.concatenate(np.moveaxis(A_hat, -1, 0))
+
+      w_j = optimal_weights(A, A_hat)
+
+      # approximation: numerator / denominator
+      N = C.dot(w_j[:, None] * f_j)
+      N_hat = C_hat.dot(np.conj(w_j[:, None] * f_j))
+
+      D = C.dot(w_j)[:, None]
+      D_hat = C_hat.dot(np.conj(w_j))[:, None]
+
+
+
+      # update approximation
+      r_k = f_k.at[J].set((N + N_hat) / (D+D_hat))
+
+      # check for convergence
+      errors.append(np.linalg.norm(f_k - r_k, np.inf))
+      if errors[-1] <= reltol:
+          break
+
+  if V == 1:
+    f_j = f_j[:, 0]
+
+
+  z_j = np.concatenate([z_j, -np.conj(z_j)])
+  f_j = np.concatenate([f_j, np.conj(f_j)])
+  w_j = np.concatenate([w_j, np.conj(w_j)])
+
+  if return_errors:
+    return z_j, f_j, w_j, errors
+  return z_j, f_j, w_j

diffaaable/selective.py

@@ -0,0 +1,197 @@
+import os
+import pathlib
+import jax.numpy as np
+import numpy as onp
+from jax.tree_util import Partial
+from diffaaable.core import aaa, residues
+from diffaaable.adaptive import Domain, domain_mask, adaptive_aaa, next_samples_heat
+import matplotlib.pyplot as plt
+
+def reduced_domain(domain, reduction=1-1/12):
+  r = reduction
+  return (
+    domain[0]*r+domain[1]*(1-r),
+    domain[1]*r+domain[0]*(1-r)
+  )
+
+def sample_cross(domain):
+  center = domain_center(domain)
+  dist = 0.5 * (domain[1]-domain[0])
+  return center+np.array([dist.real, -dist.real, 1j*dist.imag, -1j*dist.imag])
+
+def sample_domain(domain: Domain, N: int):
+  sqrt_N = np.round(np.sqrt(N)).astype(int)
+  domain = reduced_domain(domain)
+  z_k_r = np.linspace(domain[0].real, domain[1].real, sqrt_N)
+  z_k_i = np.linspace(domain[0].imag, domain[1].imag, sqrt_N)
+  Z_r, Z_i = np.meshgrid(z_k_r, z_k_i)
+  z_k = (Z_r+1j*Z_i).flatten()
+  return z_k
+
+def sample_rim(domain: Domain, N: int):
+  side_N = N//4
+  z_k_r = np.linspace(domain[0].real, domain[1].real, side_N+2)[1:-1]
+  z_k_i = np.linspace(domain[0].imag, domain[1].imag, side_N+2)[1:-1] * 1j
+  return np.array([
+    1j*domain[0].imag + z_k_r,
+    1j*domain[1].imag + z_k_r,
+    domain[0].real + z_k_i,
+    domain[1].real + z_k_i
+  ]).flatten()
+
+def anti_domain(domain: Domain):
+  return (
+    domain[0].real + 1j*domain[1].imag,
+    domain[1].real + 1j*domain[0].imag
+    )
+
+def domain_center(domain: Domain):
+  return np.mean(np.array(domain))
+
+def subdomains(domain: Domain, divide_horizontal: bool, center: complex=None):
+  if center is None:
+    center = domain_center(domain)
+  left_up =    domain[0].real + 1j*domain[1].imag
+  right_down = domain[1].real + 1j*domain[0].imag
+
+  subs = [
+    (center, domain[1]),
+    anti_domain((left_up, center)),
+    (domain[0], center),
+    anti_domain((center, right_down)),
+  ]
+
+  if divide_horizontal:
+    return [(subs[1][0], subs[0][1]), (subs[2][0], subs[3][1])]
+  return   [(subs[2][0], subs[1][1]), (subs[3][0], subs[0][1])]
+
+
+def cutoff_mask(z_k, f_k, f_k_dot, cutoff):
+  m = np.abs(f_k)<cutoff #filter out values, that have diverged too strongly
+  return z_k[m], f_k[m], f_k_dot[m]
+
+def plot_domain(domain: Domain, size: float=1):
+  left_up =    domain[0].real + 1j*domain[1].imag
+  right_down = domain[1].real + 1j*domain[0].imag
+
+  points = np.array([domain[0], right_down, domain[1], left_up, domain[0]])
+
+  return plt.plot(points.real, points.imag,
+                  lw=size/30, zorder=1)
+
+def all_poles_known(poles, prev, tol):
+  if prev is None or len(prev)!=len(poles):
+    return False
+  #return True
+
+  dist = np.abs(poles[:, None] - prev[None, :])
+  check = np.all(np.any(dist < tol, axis=1))
+  return check
+
+
+def selective_refinement_aaa(f: callable,
+                domain: Domain,
+                N: int = 36,
+                max_poles: int = 400,
+                cutoff: float = None,
+                tol_aaa: float = 1e-9,
+                tol_pol: float = 1e-5,
+                suggestions = None,
+                on_rim: bool = False,
+                Dmax=30,
+                use_adaptive: bool = True,
+                z_k = None, f_k = None,
+                divide_horizontal=True,
+                debug_name = "d",
+                stop = 0.1,
+                batchsize=10
+                ):
+  """
+  TODO: allow access to samples slightly outside of domain
+  """
+
+  print(f"start domain '{debug_name}', {Dmax=}")
+  folder = f"debug_out/{debug_name:0<33}"
+  domain_size = np.abs(domain[1]-domain[0])/2
+  size = domain_size/2 # for plotting
+  #plot_rect = plot_domain(domain, size=30)
+  #color = plot_rect[0].get_color()
+
+  if cutoff is None:
+    cutoff = np.inf
+
+  eval_count = 0
+  if use_adaptive:
+    pathlib.Path(folder).mkdir(parents=True, exist_ok=True)
+    sampling = Partial(next_samples_heat, debug=folder,
+                       stop=stop, resolution=(101, 101), batchsize=batchsize)
+    if z_k is None:
+      z_k = np.empty((0,), dtype=complex)
+      f_k = z_k.copy()
+
+    if len(z_k) < N/4:
+      z_k_new = sample_domain(domain, N)
+      f_k = np.append(f_k, f(z_k_new))
+      z_k = np.append(z_k, z_k_new)
+
+      eval_count += len(z_k_new)
+      print(f"new eval: {eval_count}")
+    eval_count -= len(z_k)
+    z_j, f_j, w_j, z_n, z_k, f_k = adaptive_aaa(
+      z_k, f, f_k_0=f_k, evolutions=N*16, tol=tol_aaa,
+      domain=reduced_domain(domain, 1.07), radius=4*domain_size/(N), #NOTE: actually increased domain :/
+      return_samples=True, sampling=sampling, cutoff=np.inf
+    )
+    # TODO pass down samples in buffer zone
+    eval_count += len(z_k)
+  else:
+    if on_rim:
+      z_k = sample_rim(domain, N)
+    else:
+      z_k = sample_domain(reduced_domain(domain, 1.05), N)
+    f_k = f(z_k)
+    eval_count += len(f_k)
+    try:
+      z_j, f_j, w_j, z_n = aaa(z_k, f_k, tol=tol_aaa)
+    except onp.linalg.LinAlgError as e:
+      z_n = z_j = f_j = w_j = np.empty((0,))
+
+  print(f"domain '{debug_name}' done: {domain} ->  eval: {eval_count}")
+  poles = z_n[domain_mask(domain, z_n)]
+
+  if (Dmax == 0 or
+    (len(poles)<=max_poles and all_poles_known(poles, suggestions, tol_pol))):
+
+    #plt.scatter(poles.real, poles.imag, color = color, marker="x")#, s=size*3, linewidths=size/2)
+    print("I am done here")
+
+    res = residues(z_j, f_j, w_j, poles)
+    return poles, res, eval_count
+  #plt.scatter(poles.real, poles.imag, color = color, marker="+", s=0.2, zorder=3)#, s=size, linewidths=size/6)
+
+  subs = subdomains(domain, divide_horizontal)
+
+  pol = np.empty((0,), dtype=complex)
+  res = pol.copy()
+  for i,sub in enumerate(subs):
+    sug = poles[domain_mask(sub, poles)]
+    sample_mask = domain_mask(sub, z_k)
+
+    known_z_k = z_k[sample_mask]
+    known_f_k = f_k[sample_mask]
+
+    p, r, e = selective_refinement_aaa(
+      f, sub, N, max_poles, cutoff, tol_aaa, tol_pol,
+      use_adaptive=use_adaptive,
+      suggestions=sug, Dmax=Dmax-1, z_k=known_z_k, f_k=known_f_k,
+      divide_horizontal = not divide_horizontal,
+      debug_name=f"{debug_name}{i+1}",
+    )
+    pol = np.append(pol, p)
+    res = np.append(res, r)
+    eval_count += e
+  # if len(pol) > 0:
+  #   plt.xlim(domain[0].real, domain[1].real)
+  #   plt.ylim(domain[0].imag, domain[1].imag)
+  #   plt.savefig(f"debug_out/{debug_name:0<33}.png")
+  return pol, res, eval_count

diffaaable/vectorial.py

@@ -0,0 +1,114 @@
+from jax import config
+config.update("jax_enable_x64", True) #important -> else aaa fails
+import jax.numpy as np
+from diffaaable.core import poles
+
+def check_inputs(z_k, f_k):
+  f_k = np.array(f_k)
+  z_k = np.array(z_k)
+
+  if z_k.ndim != 1:
+    raise ValueError("z_k should be 1D but has shape {z_k.shape}")
+  M = z_k.shape[0]
+
+  if f_k.ndim == 1:
+    f_k = f_k[:, np.newaxis]
+
+  if f_k.ndim != 2 or f_k.shape[0]!=M:
+    raise ValueError("f_k should be 1 or 2D and have the same first"
+                     f"dimension as z_k, {f_k.shape=}, {z_k.shape=}")
+  V = f_k.shape[1]
+
+  return z_k, f_k, M, V
+
+def vectorial_aaa(z_k, f_k, tol=1e-13, mmax=100, return_errors=False):
+  """Find a rational approximation to $\mathbf f(z)$ over the points $z_k$ using
+  a modified AAA algorithm, as presented in [^4]. Importantly the weights and
+  thus also the poles are shared between all entries of $\mathbf f(z)$.
+
+  Parameters
+  ----------
+      z_k : array (M,):
+        M sample points
+      f_k : array (M, V):
+        vector valued function values
+      tol : float
+        the approximation tolerance
+      mmax : int
+        the maximum number of iterations/degree of the resulting approximant
+
+  Returns:
+
+
+
+  """
+  z_k, f_k, M, V = check_inputs(z_k, f_k)
+
+  J = np.ones(M, dtype=bool)
+  z_j = np.empty(0, dtype=z_k.dtype)
+  f_j = np.empty((0, V), dtype=f_k.dtype)
+  errors = []
+
+  reltol = tol * np.linalg.norm(f_k, np.inf)
+
+  r_k = np.mean(f_k) * np.ones_like(f_k)
+
+  mmax = min(mmax, len(f_k)//2)
+
+  for m in range(mmax):
+      # find largest residual
+      jj = np.argmax(np.linalg.norm(f_k - r_k, axis=-1)) #Next sample point to include
+      z_j = np.append(z_j, np.array([z_k[jj]]))
+      f_j = np.concatenate([f_j, f_k[jj][None, :]])
+      J = J.at[jj].set(False)
+
+      # Cauchy matrix containing the basis functions as columns
+      C = 1.0 / (z_k[J,None] - z_j[None,:])
+      # Loewner matrix
+      A = (f_k[J,None] - f_j[None,:]) * C[:,:,None]
+
+      # TODO: stack A
+      A = np.concatenate(np.moveaxis(A, -1, 0))
+
+      # compute weights as right singular vector for smallest singular value
+      if return_errors:
+         print("start SVD")
+      _, _, Vh = np.linalg.svd(A)
+      if return_errors:
+         print("finished SVD")
+
+      w_j = Vh[-1, :].conj()
+
+      # approximation: numerator / denominator
+      N = C.dot(w_j[:, None] * f_j) #TODO check it works
+      D = C.dot(w_j)[:, None]
+
+      # update residual
+      r_k = f_k.at[J].set(N / D)
+
+      # check for convergence
+      errors.append(np.linalg.norm(f_k - r_k, np.inf))
+      if return_errors:
+         print(errors[-1])
+      if errors[-1] <= reltol:
+          break
+
+  z_n = poles(z_j, w_j)
+  if return_errors:
+    return z_j, f_j, w_j, z_n, errors
+  return z_j, f_j, w_j, z_n
+
+
+def residues_vec(z_j,f_j,w_j,z_n):
+  '''Vectorial residues for given poles via formula for simple poles
+  of quotients of analytic functions. For a barycentric rational of order `m`
+  the
+
+  '''
+
+  C_pol = 1.0 / (z_n[:,None] - z_j[None,:])
+  N_pol = C_pol.dot((f_j*w_j[None,:]).T)
+  Ddiff_pol = (-C_pol**2).dot(w_j)
+  res = N_pol / Ddiff_pol[:, None]
+
+  return np.nan_to_num(res.T)

diffaaable-1.0.1.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jan David Fischbach
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

diffaaable-1.0.1.dist-info/METADATA

@@ -0,0 +1,64 @@
+Metadata-Version: 2.1
+Name: diffaaable
+Version: 1.0.1
+Summary: JAX-differentiable AAA algorithm
+Keywords: python
+Author-email: Jan David Fischbach <fischbach@kit.edu>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Operating System :: OS Independent
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: baryrat
+Requires-Dist: jaxopt
+Requires-Dist: pytest
+Requires-Dist: pytest-benchmark
+Requires-Dist: pre-commit ; extra == "dev"
+Requires-Dist: pytest ; extra == "dev"
+Requires-Dist: pytest-cov ; extra == "dev"
+Requires-Dist: pytest_regressions ; extra == "dev"
+Requires-Dist: jupytext ; extra == "docs"
+Requires-Dist: matplotlib ; extra == "docs"
+Requires-Dist: jupyter-book ; extra == "docs"
+Requires-Dist: sphinx_math_dollar ; extra == "docs"
+Provides-Extra: dev
+Provides-Extra: docs
+
+# diffaaable 1.0.1
+
+![](docs/assets/diffaaable.png)
+
+`diffaaable` is a JAX differentiable version of the AAA algorithm. The derivatives are implemented as custom Jacobian Vector products in accordance to [^1].
+A detailed derivation of the used matrix expressions is provided in the appendix of [^2].
+Under the hood `diffaaable` uses the AAA implementation of [`baryrat`](https://github.com/c-f-h/baryrat).
+Additionaly the following application specific extensions to the AAA algorithm are included:
+
+- **Adaptive**: Adaptive refinement strategy to minimize the number of function evaluation needed to precisely locate poles within some domain
+- **Vectorial** (also referred to as set-valued): AAA algorithm acting on vector valued functions $\mathbf{f}(z)$ as presented in [^3].
+- **Lorentz**: Variant that enforces symmetric poles around the imaginary axis.
+- **Selective Refinement**: Use a divide and conquer theme to capture many pole simultaneously and accurately, by limiting the number of poles per AAA solve. Suggested in [^4].
+
+## Installation
+to install `diffaaable` run
+`pip install diffaaable`
+
+## Usage
+Please refer to the [quickstart tutorial](./usage.md)
+
+## Contributing
+Feel free to open issues and/or PRs.
+
+## Citation
+When using this software package for scientific work please cite the associated publication [^2].
+
++++
+
+[^1]: https://arxiv.org/pdf/2403.19404
+[^2]: Multiscat Resonances (to be publlished)
+[^3]: https://doi.org/10.1093/imanum/draa098
+[^4]: https://doi.org/10.48550/arXiv.2405.19582
+

diffaaable-1.0.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+diffaaable/__init__.py,sha256=EiOU3Hg-tdBkK2vZAlL0CkXQM38QzuZyWOnTFxSODnI,310
+diffaaable/adaptive.py,sha256=skpKa9jMpirpAWR-fqnYrD9a0WaXrjrTc9kZk1tMbKI,9667
+diffaaable/core.py,sha256=LOcq4k0IiYHaMffIS8Amgdu9lw9CWVH8WPYtFA5KtbY,7168
+diffaaable/lorentz.py,sha256=BVimKKOo-FuTNGhjwr3di54hziq2_xclQKLNe8YGJ-k,2390
+diffaaable/selective.py,sha256=F0pDH9Sf1vhIzgYSLZrqdRPx0oYY6p_dIJyI67JLHxI,6246
+diffaaable/vectorial.py,sha256=JtuowSk_bfgdjrEK0PsVIQGA1EuKuNOciQqqy_Nxv54,3145
+diffaaable-1.0.1.dist-info/LICENSE,sha256=_GGRQSEhqmML9t-dWxiXXZXwPLdzVHT5HdEISjFjSrU,1076
+diffaaable-1.0.1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+diffaaable-1.0.1.dist-info/METADATA,sha256=CdLhlhvE4EzjNWdlAPeP41j6AlrybpHM00lHx3k_9JA,2512
+diffaaable-1.0.1.dist-info/RECORD,,

diffaaable-1.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: flit 3.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any