diffaaable 1.1.0__tar.gz → 1.2.1__tar.gz

{diffaaable-1.1.0 → diffaaable-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: diffaaable
-Version: 1.1.0
+Version: 1.2.1
 Summary: JAX-differentiable AAA algorithm
 Keywords: python
 Author-email: Jan David Fischbach <fischbach@kit.edu>
@@ -10,11 +10,13 @@ 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
+License-File: LICENSE
+Requires-Dist: numpy>2
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: baryrat
 Requires-Dist: jaxopt
+Requires-Dist: pandas>=2.2.3
 Requires-Dist: tbump>=6.11.0 ; extra == "dev"
 Requires-Dist: towncrier ; extra == "dev"
 Requires-Dist: pre-commit ; extra == "dev"
@@ -25,12 +27,12 @@ Requires-Dist: pytest-benchmark ; extra == "dev"
 Requires-Dist: matplotlib ; extra == "dev"
 Requires-Dist: jupytext ; extra == "docs"
 Requires-Dist: matplotlib ; extra == "docs"
-Requires-Dist: jupyter-book ; extra == "docs"
+Requires-Dist: jupyter-book<2 ; extra == "docs"
 Requires-Dist: sphinx_math_dollar ; extra == "docs"
 Provides-Extra: dev
 Provides-Extra: docs
 
-# diffaaable 1.1.0
+# diffaaable 1.2.1
 
 ![](docs/assets/diffaaable.png)
 
@@ -39,8 +41,9 @@ A detailed derivation of the used matrix expressions is provided in the appendix
 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 (called Iterative Sample Refinement (ISR) in the corresponding paper) 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].
+- **Adaptive**: Adaptive refinement strategy (called *Iterative Sample Refinement* (ISR) in the corresponding paper) to minimize the number of function evaluation needed to precisely locate poles within some domain
+- **Vectorial**: AAA algorithm acting on vector valued functions $\mathbf{f}(z)$ as presented in [^3].
+- **Tensor**: Convenience alternative to the vector valued AAA algorithm (`vectorial`) accepting a tensor valued function F_k (so arbitrary dimensionality) instead of the single dimension that `vectorial` requires.
 - **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].
 
@@ -60,7 +63,7 @@ When using this software package for scientific work please cite the associated
 +++
 
 [^1]: https://arxiv.org/pdf/2403.19404
-[^2]: "A framework to compute resonances arising from multiple scattering", https://arxiv.org/abs/2409.05563
+[^2]: "A framework to compute resonances arising from multiple scattering", https://doi.org/10.1002/adts.202400989
 [^3]: https://doi.org/10.1093/imanum/draa098
 [^4]: https://doi.org/10.48550/arXiv.2405.19582
 

{diffaaable-1.1.0 → diffaaable-1.2.1}/README.md

@@ -1,4 +1,4 @@
-# diffaaable 1.1.0
+# diffaaable 1.2.1
 
 ![](docs/assets/diffaaable.png)
 
@@ -7,8 +7,9 @@ A detailed derivation of the used matrix expressions is provided in the appendix
 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 (called Iterative Sample Refinement (ISR) in the corresponding paper) 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].
+- **Adaptive**: Adaptive refinement strategy (called *Iterative Sample Refinement* (ISR) in the corresponding paper) to minimize the number of function evaluation needed to precisely locate poles within some domain
+- **Vectorial**: AAA algorithm acting on vector valued functions $\mathbf{f}(z)$ as presented in [^3].
+- **Tensor**: Convenience alternative to the vector valued AAA algorithm (`vectorial`) accepting a tensor valued function F_k (so arbitrary dimensionality) instead of the single dimension that `vectorial` requires.
 - **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].
 
@@ -28,6 +29,6 @@ When using this software package for scientific work please cite the associated
 +++
 
 [^1]: https://arxiv.org/pdf/2403.19404
-[^2]: "A framework to compute resonances arising from multiple scattering", https://arxiv.org/abs/2409.05563
+[^2]: "A framework to compute resonances arising from multiple scattering", https://doi.org/10.1002/adts.202400989
 [^3]: https://doi.org/10.1093/imanum/draa098
 [^4]: https://doi.org/10.48550/arXiv.2405.19582

diffaaable-1.2.1/diffaaable/__init__.py

@@ -0,0 +1,15 @@
+"""diffaaable - JAX-differentiable AAA algorithm"""
+
+__version__ = "1.2.1"
+
+
+from diffaaable.core import aaa
+from diffaaable.adaptive import adaptive_aaa
+from diffaaable.lorentz import lorentz_aaa
+from diffaaable.selective import selective_subdivision_aaa
+from diffaaable.set_aaa import set_aaa
+from diffaaable.tensor import tensor_aaa
+from diffaaable.vectorial import vectorial_aaa
+from diffaaable.util import poles, residues
+
+__all__ = ["aaa", "adaptive_aaa", "lorentz_aaa", "selective_subdivision_aaa", "set_aaa", "tensor_aaa", "vectorial_aaa", "poles", "residues"]

{diffaaable-1.1.0 → diffaaable-1.2.1}/diffaaable/adaptive.py

@@ -1,3 +1,4 @@
+import functools
 from typing import Union
 import jax.numpy as np
 import numpy.typing as npt
@@ -8,7 +9,7 @@ import matplotlib.pyplot as plt
 from functools import partial
 Domain = tuple[complex, complex]
 
-def top_right(a: npt.NDArray[complex], b: npt.NDArray[complex]):
+def top_right(a: npt.NDArray, b: npt.NDArray):
   return np.logical_and(a.imag>=b.imag, a.real>=b.real)
 
 def domain_mask(domain: Domain, z_n):
@@ -111,11 +112,14 @@ def next_samples_heat(
 
   return add_samples
 
-aaa = jax.tree_util.Partial(aaa)
+vanilla_aaa = jax.tree_util.Partial(aaa)
 
 def mask(z_k, f_k, f_k_dot, cutoff):
-    m = np.abs(f_k)<cutoff    #filter out values, that have diverged too strongly
-    m = np.logical_and(m, ~np.isnan(f_k))   #filter out nans
+    def all_f(f):
+      # to make sure all f_k (in case of tensor valued functions) behave nice
+      return np.squeeze(np.apply_over_axes(np.all, f, np.arange(f.ndim-1)+1))
+    m = all_f(np.abs(f_k)<cutoff)    #filter out values, that have diverged too strongly
+    m = np.logical_and(m, all_f(~np.isnan(f_k)))   #filter out nans
     m = np.logical_and(m, ~np.isnan(z_k))   #filter out nans
 
     if m.ndim == 2:
@@ -150,6 +154,11 @@ def _adaptive_aaa(z_k_0: npt.NDArray,
     Tangent of `f`. If provided JVPs of `f` will be collected throughout the
     iterations. For use in custom_jvp
   """
+  if sampling is None:
+    sampling = next_samples
+
+  if aaa is None:
+    aaa = vanilla_aaa
 
   if sampling == "heat":
     sampling = next_samples_heat
@@ -235,16 +244,18 @@ def adaptive_aaa(z_k_0: npt.NDArray,
                  radius: float = None,
                  domain: Domain = None,
                  f_k_0: npt.NDArray = None,
-                 sampling: callable = next_samples,
+                 sampling: callable = None,
                  prev_z_n: npt.NDArray = None,
                  return_samples: bool = False,
-                 aaa: callable = aaa):
+                 aaa: callable = None):
   """ 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.
+  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`
@@ -267,13 +278,36 @@ def adaptive_aaa(z_k_0: npt.NDArray,
   tol: float
       Tolerance used in AAA (see `diffaaable.aaa`)
   radius: float
-      Distance from the assumed poles for nex samples
+      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:
       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`.
 
 
   Returns
@@ -286,6 +320,12 @@ def adaptive_aaa(z_k_0: npt.NDArray,
       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.
   """
   return _adaptive_aaa(
     z_k_0=z_k_0, f=f, evolutions=evolutions, cutoff=cutoff, tol=tol, mmax=mmax,

{diffaaable-1.1.0 → diffaaable-1.2.1}/diffaaable/core.py

@@ -5,8 +5,11 @@ import numpy.typing as npt
 import jax
 import numpy as np
 from baryrat import aaa as oaaa # ordinary aaa
+from .tensor import tensor_aaa
 import functools
-import scipy.linalg
+from .util import poles
+
+USE_SETAAA = False
 
 @functools.wraps(oaaa)
 @jax.custom_jvp
@@ -14,17 +17,22 @@ 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)
+  if USE_SETAAA:
+    print("Using set_aaa")
+    z_j, f_j, w_j, z_n = tensor_aaa(z_k, f_k[:, None], tol_aaa=tol, mmax_aaa=mmax)
+    f_j = f_j[:, 0]
+  else:
+    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))]
 
@@ -166,59 +174,3 @@ def aaa_jvp(primals, tangents):
   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-1.2.1/diffaaable/selective.py

@@ -0,0 +1,267 @@
+import jax.numpy as np
+import numpy as onp
+from diffaaable.core import aaa
+from diffaaable.util import residues
+from diffaaable.adaptive import Domain, domain_mask, adaptive_aaa
+import matplotlib.pyplot as plt
+
+def reduced_domain(domain, reduction=1-1/12):
+  """
+  Utility: Rescale the domain. Can be used to shring or enlarge the domain.
+
+  Parameters
+  ----------
+  domain: Domain
+    The domain to rescale.
+  reduction: float
+    The factor by which to rescale the domain. A value of 1 will not change the domain. A value of 0.5 will shrink the domain by half. A value of 2 will enlarge the domain by a factor of 2.
+  """
+
+  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):
+  """
+  Utility: Subdivide the domain into two subdomains.
+
+  Parameters
+  ----------
+  domain: Domain
+    The domain to subdivide.
+  divide_horizontal: bool
+    If True, the domain is divided horizontally. If False, the domain is divided vertically.
+  center: complex
+    The center through wich to divide the domain. If None, the center is calculated as the mean of the domain.
+  """
+  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 plot_domain(domain: Domain, size: float=1):
+  """
+  Utility: Plot the domain as a rectangle.
+
+  Parameters
+  ----------
+  domain: Domain
+    The domain to plot.
+  size: float
+    The relative size to scale the linewidth of the rectangle.
+
+  TODO
+  ----
+  - Add a color argument to the function.
+  """
+  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
+
+  dist = np.abs(poles[:, None] - prev[None, :])
+  check = np.all(np.any(dist < tol, axis=1))
+  return check
+
+
+def selective_subdivision_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,
+                evolutions_adaptive: int = 5,
+                radius_adaptive: float = 1e-4,
+                z_k = None, f_k = None,
+                divide_horizontal=True,
+                debug_plot_domains: bool = False,
+                ):
+  """
+  When the number of poles that need to be located is large it can be beneficial to subdivide the search domain.
+  This function implements a recursive subdivision of the domain, that automatically terminates once the poles are found to a satisfactory degree.
+
+  Parameters
+  ----------
+  f: callable
+    The function that the pole search is conducted on. It should accept batches of complex numbers and return a complex number per input.
+  domain: Domain
+    The pole search is limited to this domain.
+  N: int
+    The initial number of samples. If the number of samples drops below N the algorithm will add N new samples to the domain.
+  max_poles: int
+    The maximum number of poles that are considered valid within one search. If this number is exceeded the domain will be subdivided.
+  cutoff: float
+    A cutoff to avoid numerical instability due to large samples close to poles. See also `diffaaable.adaptive.adaptive_aaa`.
+  tol_aaa: float
+    The tolerance for the AAA algorithm. See also `diffaaable.core.aaa`.
+  tol_pol: float
+    The tolerance for the pole search. This is used to determine if a pole has moved significantly since the last domain subdivision.
+  suggestions: array
+    A list of poles that are already known. This is used internally to recursively call `selective_subdivision_aaa`.
+  on_rim: bool
+    If True, the initial samples are taken on the rim of the domain. Selecting the samples on the domain border is closely
+    related to similar contour integral approaches. It is however generally recomended to sample within the domain (default: False).
+  Dmax: int
+    The maximum number of subdivisions. If this number is exceeded the algorithm will stop and return the current poles.
+    TODO: allow the user to specify that when reaching Dmax the algorithm should return no poles to avoid false poles.
+  use_adaptive: bool
+    If True, the algorithm will use the adaptive AAA algorithm within the subdomains to locate the poles more accurately.
+    The samples collected while searching the parent domain are passed to the respective subdomains to minimize computational cost.
+    Using the adaptive aaa is generally recommended (default: True). If False, the algorithm will use the standard AAA algorithm.
+  evolutions_adaptive: int
+    The number of evolutions for the adaptive AAA algorithm.
+  radius_adaptive: float
+    The radius for the adaptive AAA algorithm. See also `diffaaable.adaptive.adaptive_aaa`.
+  z_k: array
+    The samples that have already been collected. This is used internally to recursively call `selective_subdivision_aaa`. It can also be used by the user to pass samples that are already known.
+  f_k: array
+    The function values of the samples that have already been collected. This is used internally to recursively call `selective_subdivision_aaa`. It can also be used by the user to pass samples that are already known.
+  divide_horizontal: bool
+    If True, the next domain division will be horizontal. Used internally to alternate between horizontal and vertical divisions during recursion.
+  debug_plot_domains: bool
+    If True, the algorithm will plot the domains that are being searched. This is useful for debugging and understanding the algorithm.
+
+
+  TODO
+  ----
+  - allow access to samples slightly outside of domain
+  - divide horizontal/vertical according to the distribution of poles
+  """
+
+  domain_size = np.abs(domain[1]-domain[0])/2
+
+  if debug_plot_domains:
+    print(f"Domain: {domain}")
+    plot_domain(domain, size=domain_size)
+
+  if cutoff is None:
+    cutoff = np.inf
+
+  eval_count = 0
+  if use_adaptive:
+    if z_k is None:
+      z_k = np.empty((0,), dtype=complex)
+      f_k = z_k.copy()
+
+    if len(z_k) < N:
+      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)
+
+    eval_count -= len(z_k)
+
+    # NOTE: reduced domain with a factor larger than 1
+    # actually increases domain size to avoid missing poles right at the border
+    z_j, f_j, w_j, z_n, z_k, f_k = adaptive_aaa(
+      z_k, f, f_k_0=f_k, evolutions=evolutions_adaptive, tol=tol_aaa,
+      domain=reduced_domain(domain, 1.07), radius=domain_size*radius_adaptive,
+      return_samples=True, cutoff=cutoff
+    )
+    # 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,))
+
+  poles = z_n[domain_mask(domain, z_n)]
+
+  if (Dmax == 0 or
+    (len(poles)<=max_poles and all_poles_known(poles, suggestions, tol_pol))):
+
+    res = residues(z_j, f_j, w_j, poles)
+    return poles, res, eval_count
+
+  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_subdivision_aaa(
+      f, sub, N, max_poles, cutoff, tol_aaa, tol_pol,
+      use_adaptive=use_adaptive,
+      evolutions_adaptive=evolutions_adaptive,
+      radius_adaptive=radius_adaptive, on_rim=on_rim,
+      suggestions=sug, Dmax=Dmax-1, z_k=known_z_k, f_k=known_f_k,
+      divide_horizontal = not divide_horizontal,
+      debug_plot_domains=debug_plot_domains
+    )
+    pol = np.append(pol, p)
+    res = np.append(res, r)
+    eval_count += e
+  return pol, res, eval_count

diffaaable-1.2.1/diffaaable/set_aaa/__init__.py

@@ -0,0 +1 @@
+from .set_aaa import set_aaa