diff --git a/desc/integrals/_bounce_utils.py b/desc/integrals/_bounce_utils.py index 9598dbd33c..89fbf7631c 100644 --- a/desc/integrals/_bounce_utils.py +++ b/desc/integrals/_bounce_utils.py @@ -26,7 +26,7 @@ from desc.backend import dct, ifft, jax, jnp from desc.integrals._interp_utils import ( _JF_BUG, - _eps, + _root_eps, chebder, nufft1d2r, nufft2d2r, @@ -105,7 +105,7 @@ def _bounce_points( ) assert intersect.shape[-2:] == (knots.size - 1, B.shape[-1] - 1) - dB_dz = flatten_mat(jnp.sign(poly_val(x=intersect, c=B[..., None, :], der=True))) + dB_dz = flatten_mat(jnp.sign(poly_val(x=intersect, c=B[..., None, :], der=1))) # Only consider intersect if it is within knots that bound that polynomial. mask = flatten_mat(intersect >= 0) z1 = (dB_dz <= 0) & mask @@ -370,10 +370,11 @@ def bounce_points_jvp(num_well, primals, tangents): dB_dz += dB_dt * dt_dz dB_do += dB_dt * dt_do + regularization = _root_eps() dB_dz = jnp.where( - jnp.abs(dB_dz) > _eps, + jnp.abs(dB_dz) > regularization, dB_dz, - dB_dz + jnp.copysign(_eps, dB_dz.real), + dB_dz + jnp.copysign(regularization, dB_dz.real), ) dz = jnp.where(mask, (dp[..., None] - dB_do) / dB_dz, 0.0) @@ -703,7 +704,7 @@ def get_mins(knots, B, num_mins=-1, fill_value=0.0): a_max=jnp.diff(knots), sentinel=0.0, ) - b = flatten_mat((poly_val(x=mins, c=b[..., None, :], der=True) > 0) & (mins > 0)) + b = flatten_mat((poly_val(x=mins, c=b[..., None, :], der=1) > 0) & (mins > 0)) mins = flatten_mat( jnp.stack( [ diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index 4699618eed..a90a59bf1a 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -3,6 +3,7 @@ import warnings from functools import partial +import numpy as np from interpax import interp1d try: @@ -197,7 +198,7 @@ def interp1d_Hermite_vec(xq, x, f, fx, /): return interp1d(xq, x, f, method="cubic", fx=fx) -def poly_val(*, x, c, der=False): +def poly_val(x, *, c, der=0): """Evaluate polynomial ``c`` at the points ``x``. Parameters @@ -208,8 +209,8 @@ def poly_val(*, x, c, der=False): Last axis should store coefficients of a polynomial. For a polynomial given by ∑ᵢⁿ cᵢ xⁱ, where n is ``c.shape[-1]-1``, coefficient cᵢ should be stored at ``c[...,n-i]``. - der : bool - Whether to evaluate the derivative instead. + der : int + Derivative to evaluate. Returns ------- @@ -227,17 +228,27 @@ def poly_val(*, x, c, der=False): """ if c.shape[-1] == 4: - if der: + if der == 0: + return ((c[..., 0] * x + c[..., 1]) * x + c[..., 2]) * x + c[..., 3] + if der == 1: return (3 * c[..., 0] * x + 2 * c[..., 1]) * x + c[..., 2] - return ((c[..., 0] * x + c[..., 1]) * x + c[..., 2]) * x + c[..., 3] + if der == 2: + return 6 * c[..., 0] * x + 2 * c[..., 1] if c.shape[-1] == 3: - if der: + if der == 0: + return (c[..., 0] * x + c[..., 1]) * x + c[..., 2] + if der == 1: return 2 * c[..., 0] * x + c[..., 1] - return (c[..., 0] * x + c[..., 1]) * x + c[..., 2] + if der == 2: + return 2 * c[..., 0] - if der: + assert 0 <= der <= 2 + if der >= 1: c = c[..., :-1] * jnp.arange(c.shape[-1] - 1, 0, -1) + if der >= 2: + c = c[..., :-1] * jnp.arange(c.shape[-1] - 1, 0, -1) + return jnp.sum(c * x[..., None] ** jnp.arange(c.shape[-1] - 1, -1, -1), axis=-1) @@ -257,18 +268,66 @@ def _subtract_last(c, k): ) -def _filter_distinct(r, sentinel, eps): - """Set all but one of matching adjacent elements in ``r`` to ``sentinel``.""" - # eps needs to be low enough that close distinct roots do not get removed. - # Otherwise, algorithms relying on continuity will fail. - mask = jnp.isclose(jnp.diff(r, axis=-1, prepend=sentinel), 0, atol=eps) - return jnp.where(mask, sentinel, r) - - _root_companion = jnp.vectorize( partial(jnp.roots, strip_zeros=False), signature="(m)->(n)" ) -_eps = max(jnp.finfo(jnp.array(1.0).dtype).eps, 2.5e-12) + + +def _root_eps(): + # Safer to make this a callable since output depends on whether + # double precision is enabled before it is called. + return max(jnp.finfo(jnp.array(1.0).dtype).eps, 1e-11) + + +def _distinct_roots(r, c, eps, keep_extrema=True): + """Returns the distinct roots given sorted roots. + + When we return distinct roots we preserve continuity invariants. + For example, if we return an ordering of distinct roots where the + derivative is nonzero there and does not change sign between adjacent + roots, this violates the behavior implied by intermediate value theorem. + """ + # Due to numerics and condition numbers, roots of multipliciy m > 1 + # may not be at the same spot. To preserve above invariant, we discard + # duplicate roots that lie within ε of each other if the derivative at + # those points has the same sign. + + p = jnp.sign(poly_val(r, c=c[..., None, :], der=1)) + + same_sign = p == jnp.roll(p, shift=1, axis=-1) + if keep_extrema: + same_sign &= p != 0 + is_close = jnp.abs(jnp.diff(r, prepend=jnp.nan)) <= eps + + bad_pair_right_member = same_sign & is_close + bad_pair_left__member = jnp.roll(bad_pair_right_member, shift=-1, axis=-1) + r = jnp.where( + bad_pair_left__member | bad_pair_right_member, + jnp.nan, + r, + ) + return r + + +def _correction_step(r, c, k, backward_stable, eps): + # Schröder first kind correction. + p0 = poly_val(r, c=c) - k + p1 = poly_val(r, c=c, der=1) + p2 = poly_val(r, c=c, der=2) + candidate = r - (p0 * p1) / (p1**2 - p0 * p2) + + res_old = jnp.abs(p0) + res_new = jnp.abs(poly_val(candidate, c=c) - k) + r = jnp.where(res_new < res_old, candidate, r) + + if not backward_stable: + r = jnp.where( + (res_old <= eps) | (res_new <= eps), + r, + jnp.nan, + ) + + return r @partial(jax.custom_jvp, nondiff_argnums=(4, 5, 6, 7)) @@ -279,7 +338,7 @@ def polyroot_vec( a_max=None, sort=False, sentinel=jnp.nan, - eps=_eps, + eps=-1.0, distinct=False, ): """Roots of polynomial with given coefficients. @@ -293,20 +352,14 @@ def polyroot_vec( k : jnp.ndarray Shape (..., *c.shape[:-1]). Specify to find solutions to ∑ᵢⁿ cᵢ xⁱ = ``k``. - a_min : jnp.ndarray + a_min, a_max : jnp.ndarray Shape (..., *c.shape[:-1]). - Minimum ``a_min`` and maximum ``a_max`` value to return roots between. - If specified only real roots are returned, otherwise returns all complex roots. - a_max : jnp.ndarray - Shape (..., *c.shape[:-1]). - Minimum ``a_min`` and maximum ``a_max`` value to return roots between. + If given returns roots in the interval [``a_min``, ``a_max``). If specified only real roots are returned, otherwise returns all complex roots. sort : bool Whether to sort the roots. sentinel : float Value with which to pad array in place of filtered elements. - Anything less than ``a_min`` or greater than ``a_max`` plus some floating point - error buffer will work just like nan while avoiding ``nan`` gradient. eps : float Absolute tolerance with which to consider value as zero. distinct : bool @@ -320,44 +373,60 @@ def polyroot_vec( The roots of the polynomial, iterated over the last axis. """ + if eps < 0: + eps = _root_eps() + get_only_real_roots = not (a_min is None and a_max is None) - num_coef = c.shape[-1] - distinct = distinct and num_coef > 2 - func = {2: _root_linear, 3: _root_quadratic, 4: _root_cubic} - - if ( - num_coef in func - and get_only_real_roots - and jnp.isrealobj(c) - and jnp.isrealobj(k) - ): - # Compute from analytic formula to avoid the issue of complex roots with small - # imaginary parts. Also consumes less memory. + degree = c.shape[-1] - 1 + + if degree <= 3 and get_only_real_roots and jnp.isrealobj(c) and jnp.isrealobj(k): + backward_stable = degree < 3 c = jnp.moveaxis(c, -1, 0) - r = func[num_coef](*c[:-1], c[-1] - k, sentinel, eps, distinct) - if num_coef == 2: - r = r[jnp.newaxis] + r = {1: _root_linear, 2: _root_quadratic, 3: _root_cubic}[degree]( + *c[:-1], c[-1] - k + ) r = jnp.moveaxis(r, 0, -1) - - # We already filtered distinct roots for quadratics. - distinct = distinct and num_coef > 3 + c = jnp.moveaxis(c, 0, -1) else: + backward_stable = True r = _root_companion(_subtract_last(c, k)) + assert r.shape[-1] == degree + # If the complex part is too big, then these would not be real roots of + # a nearby perturbed problem, so we set to nan so that they are not + # classified as candidates after the correction step. if get_only_real_roots: - a_min = -jnp.inf if a_min is None else a_min[..., jnp.newaxis] - a_max = +jnp.inf if a_max is None else a_max[..., jnp.newaxis] r = jnp.where( - (jnp.abs(r.imag) <= eps) & (a_min <= r.real) & (r.real <= a_max), + jnp.abs(r.imag) <= eps**0.5, r.real, - sentinel, + jnp.nan, + ) + if degree > 1: + r = _correction_step( + r, + c[..., None, :], + jnp.expand_dims(k, -1), + backward_stable, + eps**0.5, + ) + if get_only_real_roots: + a_min = -jnp.inf if a_min is None else jnp.expand_dims(a_min, -1) + a_max = +jnp.inf if a_max is None else jnp.expand_dims(a_max, -1) + r = jnp.where( + (a_min <= r.real) & (r.real < a_max), + r.real, + jnp.nan, ) + distinct = distinct and (degree > 1) if sort or distinct: - r = jnp.sort(r, axis=-1) + r = jnp.sort(r, stable=False) if distinct: - r = _filter_distinct(r, sentinel, eps) - assert r.shape[-1] == num_coef - 1 + r = _distinct_roots(r, c, eps) + if not np.isnan(sentinel): + r = jnp.where(jnp.isfinite(r), r, sentinel) + + assert r.shape[-1] == degree return r @@ -377,9 +446,11 @@ def _polyroot_vec_jvp(sort, sentinel, eps, distinct, primals, tangents): c, k, a_min, a_max = primals dc, dk, _, _ = tangents + if eps < 0: + eps = _root_eps() r = polyroot_vec(c, k, a_min, a_max, sort, sentinel, eps, distinct) - dc_dr = poly_val(x=r, c=c[..., None, :], der=True) + dc_dr = poly_val(r, c=c[..., None, :], der=1) dc_dr = jnp.where( jnp.abs(dc_dr) > eps, dc_dr, @@ -388,88 +459,81 @@ def _polyroot_vec_jvp(sort, sentinel, eps, distinct, primals, tangents): dr = jnp.where( r == sentinel, 0.0, - (jnp.expand_dims(dk, -1) - poly_val(x=r, c=dc[..., None, :])) / dc_dr, + (jnp.expand_dims(dk, -1) - poly_val(r, c=dc[..., None, :])) / dc_dr, ) return r, dr -def _root_cubic(a, b, c, d, sentinel, eps, distinct): - """Return real cubic root assuming real coefficients.""" - # numerical.recipes/book.html, page 228 - - def irreducible(Q, R, b): - # Three irrational real roots. - theta = jnp.arccos(R / jnp.sqrt(Q**3)) - return ( - -2 - * jnp.sqrt(Q) - * jnp.stack( - [ - jnp.cos(theta / 3), - jnp.cos((theta + 2 * jnp.pi) / 3), - jnp.cos((theta - 2 * jnp.pi) / 3), - ] - ) - - b / 3 +def _irreducible(Q, R, b): + theta = jnp.arccos(R / jnp.sqrt(Q**3)) + return ( + -2 + * jnp.sqrt(Q) + * jnp.stack( + [ + jnp.cos(theta / 3), + jnp.cos((theta + 2 * jnp.pi) / 3), + jnp.cos((theta - 2 * jnp.pi) / 3), + ] ) + - b / 3 + ) - def reducible(Q, R, b): - # One real and two complex roots. - A = -jnp.sign(R) * jnp.cbrt(jnp.abs(R) + jnp.sqrt(jnp.abs(R**2 - Q**3))) - B = Q / A - r1 = (A + B) - b / 3 - return _concat_sentinel(r1[jnp.newaxis], sentinel, num=2) - - def root(b, c, d): - b = b / a - c = c / a - Q = (b**2 - 3 * c) / 9 - R = (2 * b**3 - 9 * b * c) / 54 + d / (2 * a) - return jnp.where( - R**2 < Q**3, - irreducible(jnp.abs(Q), R, b), - reducible(Q, R, b), - ) +def _reducible(Q, R, b): + A = -jnp.sign(R) * jnp.cbrt(jnp.abs(R) + jnp.sqrt(R**2 - Q**3)) + B = jnp.where(A == 0.0, 0.0, Q / A) + x = A + B + y = (0.5j * 3**0.5) * (A - B) + return jnp.stack( + [ + x - b / 3, + # these can yield true real roots for A near B + -0.5 * x - b / 3 + y, + -0.5 * x - b / 3 - y, + ] + ) + + +def _cubic(a, b, c, d): + b = b / a + c = c / a + Q = (b**2 - 3 * c) / 9 + R = (2 * b**3 - 9 * b * c) / 54 + d / (2 * a) + return jnp.where(R**2 < Q**3, _irreducible(Q, R, b), _reducible(Q, R, b)) + + +def _root_cubic(a, b, c, d): + """Return real cubic root assuming real coefficients. + + Uses numerical.recipes/book.html, page 228, which is not backwards stable. + This can generate fake root with O(1) residual, so post-processing is needed. + Advantage is it is much more performant than eigenvalue solve, especially + when d is higher dimensional than a, b, c. + """ return jnp.where( - # Tests catch failure here if eps < 1e-12 for double precision. - jnp.abs(a) <= eps, - _concat_sentinel( - _root_quadratic(b, c, d, sentinel, eps, distinct), - sentinel, - ), - root(b, c, d), + a == 0.0, + _concat_nan(_root_quadratic(b, c, d)), + _cubic(a, b, c, d), ) -def _root_quadratic(a, b, c, sentinel, eps, distinct): +def _root_quadratic(a, b, c): """Return real quadratic root assuming real coefficients.""" # numerical.recipes/book.html, page 227 - - discriminant = b**2 - 4 * a * c - q = -0.5 * (b + jnp.sign(b) * jnp.sqrt(jnp.abs(discriminant))) - r1 = jnp.where( - discriminant < 0, - sentinel, - jnp.where(a == 0, _root_linear(b, c, sentinel, eps), q / a), - ) - r2 = jnp.where( - # more robust to remove repeated roots with discriminant - (discriminant < 0) | (distinct & (discriminant <= eps)), - sentinel, - c / q, - ) - return jnp.stack([r1, r2]) + q = -0.5 * (b + jnp.where(b >= 0.0, 1.0, -1.0) * jnp.sqrt(b**2 - 4 * a * c)) + # second branch generalizes linear root + return jnp.stack([q / a, c / q]) -def _root_linear(a, b, sentinel, eps, distinct=False): +def _root_linear(a, b): """Return real linear root assuming real coefficients.""" - return jnp.where((a == 0) & (jnp.abs(b) <= eps), 0.0, -b / a) + return (-b / a)[None] -def _concat_sentinel(r, sentinel, num=1): - """Concatenate ``sentinel`` ``num`` times to ``r`` on first axis.""" - return jnp.concatenate((r, jnp.broadcast_to(sentinel, (num, *r.shape[1:])))) +def _concat_nan(r, num=1): + """Concatenate nan ``num`` times to ``r`` on first axis.""" + return jnp.concatenate((r, jnp.broadcast_to(jnp.nan, (num,) + r.shape[1:]))) # TODO: replace the inner loop in orthax with this diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py index 911b14012f..6862e750c1 100644 --- a/desc/integrals/bounce_integral.py +++ b/desc/integrals/bounce_integral.py @@ -49,7 +49,7 @@ ) from desc.integrals._interp_utils import ( _JF_BUG, - _eps, + _root_eps, interp1d_Hermite_vec, interp1d_vec, nufft2d2r, @@ -106,11 +106,6 @@ def pitch_quad(min_B, max_B, num_pitch, **kwargs): """ if isinstance(num_pitch, int): - errorif( - num_pitch > 1e5, - msg="Floating point error impedes detection of bounce points " - f"near global extrema. Choose {num_pitch} < 1e5.", - ) simp = kwargs.get("simp", True) num_pitch = simpson2(num_pitch) if simp else uniform(num_pitch) @@ -683,7 +678,7 @@ def points(self, pitch_inv, num_well=None): # size of 10⁻⁸ is reduced from 10³ to 10⁻⁶ for Γ_c # (which has the singular weight 1/v_∥). z1, z2 = self._B.intersect1d( - self._swap_axes(pitch_inv), num_intersect=num_well, eps=_eps + self._swap_axes(pitch_inv), num_intersect=num_well, eps=_root_eps() ) z1 = move(z1) z2 = move(z2) @@ -992,7 +987,7 @@ def interp_to_argmin(self, f, points, *, nufft_eps=-1.0, **kwargs): # such that all bounce points are at ζ >= 0; and therefore, # junk values in B_mins cannot be selected in argmin. mins, B_mins = ( - self._B.extrema1d(1, num_mins, fill_value=0.0, eps=_eps) + self._B.extrema1d(1, num_mins, fill_value=0.0, eps=_root_eps()) if isinstance(self._B, PiecewiseChebyshevSeries) else get_mins(self._c["knots"], self._B, num_mins, fill_value=0.0) ) @@ -1120,7 +1115,7 @@ def plot(self, l, m, pitch_inv=None, **kwargs): B = B[m] B = PiecewiseChebyshevSeries(B, domain) if pitch_inv is not None: - kwargs["z1"], kwargs["z2"] = B.intersect1d(pitch_inv, eps=_eps) + kwargs["z1"], kwargs["z2"] = B.intersect1d(pitch_inv, eps=_root_eps()) kwargs["k"] = pitch_inv return B.plot1d(B.cheb, **kwargs) @@ -1916,7 +1911,8 @@ def guess( """ errorif( - (surf_batch_size > 1) and (pitch_batch_size is not None), + (surf_batch_size is None or surf_batch_size > 1) + and (pitch_batch_size is not None), msg=f"Expected pitch_batch_size to be None, got {pitch_batch_size}.", ) diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 6c945045a6..f9896bdc06 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index 015815fccd..7a7e021e16 100644 --- a/tests/test_compute_everything.py +++ b/tests/test_compute_everything.py @@ -314,8 +314,7 @@ def fft_grid_data(p): fft_names = ["effective ripple", "Gamma_c", "Gamma_c Velasco"] eq = get("W7-X") - # ci and my laptop differ a bunch at rho = 0, so skip that - rho = np.linspace(1e-2, 1, 10) + rho = np.linspace(0, 1, 10) grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) nufft_eps = 1e-10 @@ -367,7 +366,7 @@ def raz_grid_data(p): eq = get("W7-X") num_transit = 2 Y_B = eq.N_grid * 2 - rho = np.linspace(1e-2, 1, 10) + rho = np.linspace(0, 1, 10) alpha = np.array([0]) zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B * eq.NFP) grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") diff --git a/tests/test_integrals.py b/tests/test_integrals.py index 723bc364d5..52828abf16 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -788,7 +788,11 @@ def test_z2_before_extrema(self): ) pitch_inv = B(B.derivative().roots(extrapolate=False))[2] z1, z2 = _bounce_points(pitch_inv, k, B.c.T, sentinel=start) + z = np.stack((z1, z2)) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) + assert not np.any( + (z > 0) & (z < 2) + ), "This triple root intersect should be removed." z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size intersect = B.solve(pitch_inv, extrapolate=False) @@ -874,9 +878,6 @@ def test_get_mins(self): np.testing.assert_allclose(B_mins[idx], B_ext_scipy) -auto_sin = (automorphism_sin, grad_automorphism_sin) - - class TestBounceQuadrature: """Test bounce quadrature.""" @@ -885,8 +886,12 @@ class TestBounceQuadrature: "is_strong, quad, automorphism", [ (True, tanh_sinh(30), None), - (True, leggauss(25), auto_sin), - (False, leggauss_lob(8, interior_only=True), auto_sin), + (True, leggauss(25), (automorphism_sin, grad_automorphism_sin)), + ( + False, + leggauss_lob(8, interior_only=True), + (automorphism_sin, grad_automorphism_sin), + ), (False, chebgauss2(21), None), ], ) diff --git a/tests/test_interp_utils.py b/tests/test_interp_utils.py index fed2f469e6..8dc351bd41 100644 --- a/tests/test_interp_utils.py +++ b/tests/test_interp_utils.py @@ -233,7 +233,7 @@ class TestPolyUtils: """Test polynomial utilities used for local spline interpolation.""" @pytest.mark.unit - def test_polyroot_vec(self): + def test_polyroot_vec_general(self): """Test vectorized computation of cubic polynomial exact roots.""" c = np.arange(-24, 24).reshape(4, 6, -1).transpose(-1, 1, 0) # Ensure broadcasting won't hide error in implementation. @@ -259,39 +259,84 @@ def test_polyroot_vec(self): err_msg=f"Eigenvalue branch of polyroot_vec failed at {i, *idx}.", ) - # Now test analytic formula branch, Ensure it filters distinct roots, - # and ensure zero coefficients don't bust computation due to singularities - # in analytic formulae which are not present in iterative eigenvalue scheme. + @pytest.mark.unit + def test_polyroot_vec_quadratic_zero_linear_term(self): + """Test quadratic roots when the linear coefficient is zero.""" + c = jnp.array([[1.0, 0.0, -1.0], [2.0, 0.0, -8.0]]) + r = polyroot_vec(c, a_min=-jnp.inf, a_max=jnp.inf, sort=True) + np.testing.assert_allclose(r, jnp.array([[-1.0, 1.0], [-2.0, 2.0]])) + + @pytest.mark.unit + @pytest.mark.parametrize("companion", [True, False]) + def test_polyroot_vec_distinct(self, companion): + """Test vectorized computation of cubic polynomial exact roots.""" + a_min = -jnp.inf + a_max = +jnp.inf + + # coefficient of poly x^3, .... x^0. c = np.array( [ - [1, 0, 0, 0], - [0, 1, 0, 0], - [0, 0, 1, 0], - [0, 0, 0, 1], - [1, -1, -8, 12], - [1, -6, 11, -6], - [0, -6, 11, -2], + [1, 0, 0, 0], # triple root (0, 0, 0) + [0, 1, 0, 0], # double root (0, 0) + [0, 0, 1, 0], # single root 0 + [0, 0, 0, 1], # no roots + [1, -1, -8, 12], # (-3, 2, 2) + [1, -6, 11, -6], # (1, 2, 3) + [0, -6, 11, -2], # 2 distinct irrational ] ) - r = polyroot_vec(c, sort=True, distinct=True) + if companion: + c = np.concat([np.zeros((c.shape[0], 1)), c], axis=-1) + assert c.shape == (7, 5) + + # acceptable outputs for multiple roots are those which are + # consistent with continuity invariant. That is if user requests + # distinct roots only, then acceptable output is + # 1) If sign(derivative) == 0, should return only one of the roots. + # 2) If sign(derivative) != 0 and changes across roots, then you can + # return both since even if they were very close and supposed to + # be the same root in infinite precision, they are distinct roots + # in finite precision, so invariants are preserved by returning both.a_max + # 3) If sign(derivative) != 0 and same across roots. This is not possible + # for continuous functions and it means that numerical error has split + # single root into multiple. If the true root is a single root, + # we should return just one of the roots. If the true root has multiplicity + # greater than one, we should not discard the pairs with same sign. + r = polyroot_vec(c, a_min=a_min, a_max=a_max, sort=True, distinct=True) for j in range(c.shape[0]): root = r[j][~np.isnan(r[j])] - unique_root = np.unique(np.roots(c[j])) - assert root.size == unique_root.size - np.testing.assert_allclose( - root, - unique_root, - err_msg=f"Analytic branch of polyroot_vec failed at {j}.", - ) + unique_root = np.unique(np.roots(c[j]).real) + try: + np.testing.assert_allclose( + root, + unique_root, + err_msg=f"Companion={companion} branch failed at {j}.", + ) + assert root.size == unique_root.size + except AssertionError as e: + if j == 0 or j == 1: + assert root.size == 0 or np.allclose(root, 0) + elif j == 4: + if root.size == 3: + # ok to keep both double roots here because they are + # one red and one green since they lie on opposite + # sides of the true double root. + assert np.allclose(root, (-3, 2, 2)) + else: + # also okay to remove both + assert root.size == 1 and np.allclose(root, -3) + else: + assert False, e + c = np.array([0, 1, -1, -8, 12]) - r = polyroot_vec(c, sort=True, distinct=True) + r = polyroot_vec(c, a_min=a_min, a_max=a_max, sort=True, distinct=True) r = r[~np.isnan(r)] unique_r = np.unique(np.roots(c)) assert r.size == unique_r.size np.testing.assert_allclose(r, unique_r) @pytest.mark.unit - @pytest.mark.parametrize("der", [False, True]) + @pytest.mark.parametrize("der", [0, 1]) def test_cubic_val(self, der): """Test vectorized computation of polynomial evaluation."""