freealg 0.4.1__py3-none-any.whl → 0.5.1__py3-none-any.whl

freealg/_decompress.py

@@ -11,24 +11,566 @@
 # =======
 
 import numpy
+import matplotlib.pyplot as plt
+import texplot
 
-# Fallback to previous API
+# Fallback to previous numpy API
 if not hasattr(numpy, 'trapezoid'):
     numpy.trapezoid = numpy.trapz
 
 __all__ = ['decompress', 'reverse_characteristics']
 
 
+# ==========
+# derivative
+# ==========
+
+def _derivative(f, fz, z, h, fd_order=4, vertical=False, return_second=False):
+    """
+    Compute first or first and second derivatives.
+
+    Parameters
+    ----------
+
+    f : function
+        Calling function
+
+    fz : numpy.array
+        Function values precomputed at the array z. This is used for second
+        order derivative where the stencil also contains the middle point.
+        Since this is already computed outside the function call, it is reused
+        here, essentially making second-order derivative free.
+
+    z : numpy.array
+        A 1D array of complex points to evaluate derivative at.
+
+    h : float or numpy.array
+        Stencil size
+
+    fd_order : {2, 4, 6}, default=2
+        Order of central finite differencing scheme.
+
+    vertical : bool, default=False
+        It `True`, it uses vertical stencil (along y axis) instead of
+        horizontal (along x axis).
+
+    return_second : bool, default=False
+        If `True`, returns both first and second derivatives.
+
+    Returns
+    -------
+
+    df1 : numpy.array
+        First derivative
+
+    If ``return_second=True``:
+
+        df2 : numpy.array
+            Second derivative
+
+    Notes
+    -----
+
+    Uses central finite difference.
+
+    If the function is holomorphic, taking the derivative along horizontal or
+    vertical directions should be identical (in theory), however, in practice,
+    they might not be exactly the same. Note especially that ``vertical=True``
+    is not suitable for points close to branch cut where the stencil points
+    fall into two branches.
+    """
+
+    h = numpy.asarray(h, dtype=z.dtype)
+
+    if vertical:
+        h_ = 1j * h
+    else:
+        h_ = h
+
+    # Stencil indices
+    stencil = numpy.arange(-fd_order//2, fd_order//2 + 1).astype(float)
+
+    # Stencil coefficients
+    if fd_order == 2:
+        coeff_df1 = [-0.5, +0.0, +0.5]
+        coeff_df2 = [+1.0, -2.0, +1.0]
+
+    elif fd_order == 4:
+        coeff_df1 = [+1.0 / 12.0, -2.0 / 3.0, +0.0, +2.0 / 3.0, -1.0 / 12.0]
+        coeff_df2 = [-1.0 / 12.0, +16.0 / 12.0, -30.0 / 12.0, +16.0 / 12.0,
+                     -1.0 / 12.0]
+
+    elif fd_order == 6:
+        coeff_df1 = [-1.0 / 60.0, +3.0 / 20.0, -3.0 / 4.0, +0.0, +3.0 / 4.0,
+                     -3.0 / 20.0, +1.0 / 60.0]
+        coeff_df2 = [+1.0 / 90, -3.0 / 20.0, +3.0 / 2.0, -49.0 / 18.0,
+                     +3.0 / 2.0, -3.0 / 20.0, +1.0 / 90.0]
+
+    else:
+        raise NotImplementedError('"fd_order" is not valid.')
+
+    # Function values at stencil points. Precomputed to avoid redundancy when
+    # both first and second derivatives are needed.
+    val = [None] * stencil.size
+    for i in range(stencil.size):
+
+        # The center stencil for first derivative is zero, so we do not compute
+        # the function value, unless second derivative is needed.
+        if coeff_df1[i] == 0.0:
+            # This is already computed outside this function, for free. This
+            # is only used for second derivative where the stencil has non-zero
+            # coefficient in the middle point. This is not used in the first
+            # derivative where the coefficient of the middle point is zero.
+            val[i] = fz
+        else:
+            val[i] = f(z + stencil[i] * h_)
+
+    # First derivative
+    df1 = numpy.zeros_like(z)
+    for i in range(stencil.size):
+
+        # Skip the center of stencil where the coeff is zero.
+        if coeff_df1[i] != +0.0:
+            df1 += coeff_df1[i] * val[i]
+    df1 /= h_
+
+    # Second derivative
+    if return_second:
+        df2 = numpy.zeros_like(z)
+        for i in range(stencil.size):
+            df2 += coeff_df2[i] * val[i]
+        df2 /= h_**2
+
+    if return_second:
+        return df1, df2
+    else:
+        return df1
+
+
+# ===================
+# adaptive derivative
+# ===================
+
+def _adaptive_derivative(f, fz, z, h, err_low=1e-4, err_high=1e-2, factor=2.0,
+                         vertical=False, return_second=False):
+    """
+    Compute first or first and second derivatives using adaptive refinement of
+    stencil size.
+
+    Parameters
+    ----------
+
+    f : function
+        Calling function
+
+    fz : numpy.array
+        Function values precomputed at the array z. This is used for second
+        order derivative where the stencil also contains the middle point.
+        Since this is already computed outside the function call, it is reused
+        here, essentially making second-order derivative free.
+
+    z : numpy.array
+        A 1D array of complex points to evaluate derivative at.
+
+    h : float or numpy.array
+        Stencil size
+
+    err_low : float, default=1e-4
+        Threshold criteria to increase stencil size h. Namely, the relative
+        error of Richardson test below this threshold is considered as
+        being caused by having too small stencil size, and hence stencil size
+        h has to be increased.
+
+    err_high : float, default=1e-2
+        Threshold criteria to decrease stencil size h. Namely, the relative
+        error of Richardson test above this threshold is considered as
+        being caused by having too large stencil size, and hence stencil size
+        h has to be decreased.
+
+    factor : float, default=2.0
+        Factor to increase or decrease stencil size.
+
+    vertical : bool, default=False
+        It `True`, it uses vertical stencil (along y axis) instead of
+        horizontal (along x axis).
+
+    return_second : bool, default=False
+        If `True`, returns both first and second derivatives.
+
+    Returns
+    -------
+
+    h : numpy.array
+        Updated stencil size
+
+    df1 : numpy.array
+        First derivative
+
+    If ``return_second=True``:
+
+        df2 : numpy.array
+            Second derivative
+
+    Notes
+    -----
+
+    Uses central finite difference. This function uses finite difference of
+    order 2 or 4 as follows.
+
+    First, function f is evaluated on 5 points:
+
+        [f(z-2h), f(z-h), f(z), f(z+h), f(z+2h)]
+
+    This is used for either a:
+
+    1. 2-nd order finite difference with the stencil of size h
+    2. 2-nd order finite difference with the stencil of size 2h
+    3. 4-th order finite difference with the stencil of size h
+
+    To determine which of these should be used, the first derivative for the
+    cases 1 and 2 above are compared using the Richardson test:
+
+    1. If the Richardson test determines the stencil is too small, the output
+       derivative is a second order finite difference on stencil of size 2h.
+       Also, h is grown by factor of two for the next iteration.
+
+    2. If the Richardson test determines the stencil is good (not too small or
+       not too large), the output derivative is a fourth order finite
+       difference on stencil of size h. Also, h is not changed for the next
+       iteration.
+
+    3. If the Richardson test determines the stencil is too large, the output
+       derivative is a second order finite difference on stencil of size h.
+       Also, h is shrunken by factor of two for the next iteration.
+
+    If the function is holomorphic, taking the derivative along horizontal or
+    vertical directions should be identical (in theory), however, in practice,
+    they might not be exactly the same. Note especially that ``vertical=True``
+    is not suitable for points close to branch cut where the stencil points
+    fall into two branches.
+    """
+
+    h = h.copy().astype(z.dtype)
+
+    if vertical:
+        h_ = 1j * h
+    else:
+        h_ = h
+
+    # Center index of stencil of order 4
+    st_center = 2
+
+    # Stencil indices of order 2 of size h
+    st_ord2_h = numpy.array([-1, 0, 1]) + st_center
+
+    # Stencil indices of order 2 of size 2h
+    st_ord2_2h = numpy.array([-2, 0, 2]) + st_center
+
+    # Stencil indices of order 4 of size h
+    st_ord4_h = numpy.array([-2, -1, 0, 1, 2]) + st_center
+
+    # Stencil coefficients for first derivative, order 2
+    coeff_df1_ord2 = [-0.5, +0.0, +0.5]
+
+    # Stencil coefficients for second derivative, order 2
+    coeff_df2_ord2 = [+1.0, -2.0, +1.0]
+
+    # Stencil coefficients for first derivative, order 4
+    coeff_df1_ord4 = [+1.0 / 12.0, -2.0 / 3.0, 0.0, +2.0 / 3.0, -1.0 / 12.0]
+
+    # Stencil coefficients for second derivative, order 4
+    coeff_df2_ord4 = [-1.0 / 12.0, +16.0 / 12.0, -30.0 / 12.0, +16.0 / 12.0,
+                      -1.0 / 12.0]
+
+    # Function values at stencil points of order 4. Precomputed to avoid
+    # redundancy when both first and second derivatives are needed.
+    f_val = [None] * st_ord4_h.size
+    for i in range(st_ord4_h.size):
+
+        if coeff_df1_ord4[i] == 0.0:
+            # This is already computed outside this function, for free. This
+            # is only used for second derivative where the stencil has non-zero
+            # coefficient in the middle point. This is not used in the first
+            # derivative where the coefficient of the middle point is zero.
+            f_val[i] = fz
+        else:
+            f_val[i] = f(z + st_ord4_h[i] * h_)
+
+    # First derivative, fd order 2, stencil of size h
+    df1_ord2_h = numpy.zeros_like(z)
+    for i in range(st_ord2_h.size):
+        # Skip the center of stencil where the coeff is zero.
+        if coeff_df1_ord2[i] != 0.0:
+            df1_ord2_h += coeff_df1_ord2[i] * f_val[st_ord2_h[i]]
+    df1_ord2_h /= h_
+
+    # First derivative, fd order 2, stencil of size 2h
+    df1_ord2_2h = numpy.zeros_like(z)
+    for i in range(st_ord2_2h.size):
+        # Skip the center of stencil where the coeff is zero.
+        if coeff_df1_ord2[i] != 0.0:
+            df1_ord2_2h += coeff_df1_ord2[i] * f_val[st_ord2_2h[i]]
+    df1_ord2_2h /= (2.0 * h_)
+
+    # First derivative, fd order 4, stencil of size h
+    df1_ord4_h = numpy.zeros_like(z)
+    for i in range(st_ord4_h.size):
+        # Skip the center of stencil where the coeff is zero.
+        if coeff_df1_ord4[i] != 0.0:
+            df1_ord4_h += coeff_df1_ord4[i] * f_val[st_ord4_h[i]]
+    df1_ord4_h /= h_
+
+    if return_second:
+        # Second derivative, fd order 2, stencil of size h
+        df2_ord2_h = numpy.zeros_like(z)
+        for i in range(st_ord2_h.size):
+            df2_ord2_h += coeff_df2_ord2[i] * f_val[st_ord2_h[i]]
+        df2_ord2_h /= h_**2
+
+        # Second derivative, fd order 2, stencil of size 2h
+        df2_ord2_2h = numpy.zeros_like(z)
+        for i in range(st_ord2_2h.size):
+            df2_ord2_2h += coeff_df2_ord2[i] * f_val[st_ord2_2h[i]]
+        df2_ord2_2h /= (2.0 * h_)**2
+
+        # Second derivative, fd order 4, stencil of size h
+        df2_ord4_h = numpy.zeros_like(z)
+        for i in range(st_ord4_h.size):
+            df2_ord4_h += coeff_df2_ord4[i] * f_val[st_ord4_h[i]]
+        df2_ord4_h /= h_**2
+
+    # Richardson test
+    fd_order = 2
+    p = 2**fd_order - 1.0
+    abs_err = numpy.abs(df1_ord2_2h - df1_ord2_h) / p
+    rel_err = abs_err / numpy.maximum(1.0, numpy.abs(df1_ord2_h))
+
+    # Grow and Shrink limits criteria on relative error
+    grow = rel_err < err_low
+    shrink = rel_err > err_high
+
+    # Output stencil size
+    h[grow] *= factor
+    h[shrink] /= factor
+
+    # Output first derivative
+    df1 = df1_ord4_h
+    df1[grow] = df1_ord2_2h[grow]
+    df1[shrink] = df1_ord2_h[shrink]
+
+    # Output second derivative
+    if return_second:
+        df2 = df2_ord4_h
+        df2[grow] = df2_ord2_2h[grow]
+        df2[shrink] = df2_ord2_h[shrink]
+
+    if return_second:
+        return h, df1, df2
+    else:
+        return h, df1
+
+
+# =============
+# newton method
+# =============
+
+def _newton_method(f, z_init, a, support, enforce_wall=False, tol=1e-4,
+                   step_size=0.1, max_iter=500, adaptive_stencil=True,
+                   halley=False, vertical=False):
+    """
+    Solves :math:``f(z) = a`` for many starting points simultaneously using the
+    damped Newton method.
+
+    Parameters
+    ----------
+
+    f : function
+        Caller function.
+
+    z_init : numpy.array
+        Initial guess of roots
+
+    a : numpy.array
+        Target value of f in the root finding problem f(z) - a = 0.
+
+    support : tuple
+        The support of the density. This is needed to enforce no crossing
+        through branch cut as wall.
+
+    enforce_wall : bool, default=False
+        If `True`, roots are not allowed to cross branch cut. It is recommended
+        to enable this feature when the initial guess z_init is above the real
+        line. When z_init points are below the real line, this feature is
+        effectively not used.
+
+    tol : float, default=1e-4
+        Tolerance of terminating the iteration when |f(z) - a| < tol.
+
+    step_size : float, default=0.1
+        The step size of Newton (or Halley) iterations, between 0 and 1 where
+        1 corresponds to full Newton (or Halley) step. Note that these methods
+        are generally very sensitive to the step size. Values between 0.05 and
+        0.2 are recommended.
+
+    max_ier : int, default=500
+        Maximum number of iterations.
+
+    adaptive_stencil : bool, default=True
+        If `True`, the stencil size in finite differencing is adaptively
+        selected based on Richardson extrapolation error. If `False`, a fixed
+        stencil size of all points are used.
+
+    halley : bool, default=False
+        If `True`, Halley method is used instead of Newton method, but only
+        when the Newton makes very small increments. Halley method is in
+        general not suitable for the most of the trajectory, and can only be
+        useful at the end of trajectory for convergence. It is recommended to
+        turn this off.
+
+    vertical : bool, default=False
+        If `True`, to compute the derivative of holomorphic function, a
+        vertical stencil (instead of horizontal) is used. This is not suitable
+        for points near branch cut where some points of the stencil might fall
+        intro two sides of the branch cut, but it might be suitable for points
+        inside the support.
+
+    Returns
+    -------
+
+    roots: numpy.array
+        Solutions z of f(z) - a = 0.
+
+    residuals: numpy.array
+        The residuals f(z) - a
+
+    iterations : numpy.array
+        Number of iterations used for each root.
+    """
+
+    # Finite difference order
+    fd_order = 4
+
+    epsilon = numpy.sqrt(numpy.finfo(float).eps)
+    h_base = epsilon**(1.0 / (fd_order + 1.0))
+
+    # Global variables
+    root = z_init.copy()
+    mask = numpy.ones(z_init.shape, dtype=bool)
+    f_val = numpy.zeros(z_init.shape, dtype=z_init.dtype)
+    residual = numpy.ones(z_init.shape, dtype=z_init.dtype) * numpy.inf
+    iterations = numpy.zeros(z_init.shape, dtype=int)
+
+    # Initialize stencil size
+    h = h_base * numpy.maximum(1.0, numpy.abs(z_init))
+
+    # Main loop
+    for _ in range(max_iter):
+
+        if not numpy.any(mask):
+            # No more active point left
+            break
+
+        # Update iterations
+        iterations += mask.astype(int)
+
+        # Mask variables using the previous mask (dash m are masked variables)
+        a_m = a[mask]
+        z_m = root[mask]
+        f_m = f(z_m)
+        f_val[mask] = f_m
+        residual_m = f_m - a_m
+        residual[mask] = residual_m
+
+        # Update mask
+        mask = numpy.abs(residual) >= tol
+
+        # Mask variables again using the new mask
+        a_m = a[mask]
+        z_m = root[mask]
+        f_m = f_val[mask]
+        residual_m = residual[mask]
+        h_m = h[mask]
+
+        if adaptive_stencil:
+            # Use adaptive stencil size from previous update
+            h_m = h[mask]
+
+            # Adaptive stencil size, finite difference order is fixed 2 and 4.
+            h_m, df1_m, df2_m = _adaptive_derivative(f, f_m, z_m, h_m,
+                                                     vertical=vertical,
+                                                     return_second=True)
+
+            # Clip too small or too large stencils
+            h_m = numpy.clip(h_m, epsilon, h_base * numpy.abs(z_m))
+
+            # Update global stencil sizes
+            h[mask] = h_m
+
+        else:
+            # Fixed stencil size for all points
+            h_m = h_base * numpy.maximum(1.0, numpy.abs(z_m))
+
+            # Fixed stencil size, but finite difference order can be set.
+            fd_order = 4  # can be 2, 4, 6
+            df1_m, df2_m = _derivative(f, f_m, z_m, h_m, fd_order=fd_order,
+                                       vertical=False, return_second=True)
+
+        # Handling second order zeros
+        df1_m[numpy.abs(df1_m) < 1e-12 * numpy.abs(z_m)] = 1e-12
+
+        # Newton and Halley steps
+        df0_m = residual_m
+        newton_step = df0_m / df1_m
+
+        if halley:
+            halley_step = (df0_m * df1_m) / (df1_m**2 - 0.5 * df0_m * df2_m)
+
+            # Criteria on where to use Halley instead of Newton
+            use_halley = numpy.abs(newton_step) < 1e-2 * numpy.abs(z_m)
+            step = numpy.where(use_halley, halley_step, newton_step)
+        else:
+            step = newton_step
+
+        # Force new point to not cross branch cut. It has to go from upper to
+        # lower half plane only through the support interval.
+        if enforce_wall:
+
+            root_m_old = root[mask]
+            root_m_new = z_m - step_size * step
+            lam_m, lam_p = support
+
+            x0, y0 = root_m_old.real, root_m_old.imag
+            x1, y1 = root_m_new.real, root_m_new.imag
+
+            # Find which points are crossing the branch cut (either downward
+            # or upward)
+            crossed = \
+                (y0 * y1 <= 0.0) & \
+                ((x0 <= lam_m) | (x0 >= lam_p)) & \
+                ((x1 <= lam_m) | (x1 >= lam_p))
+
+            # Remove imaginary component from step
+            step[crossed] = step[crossed].real + 0.0j
+
+        # Update root
+        root[mask] = z_m - step_size * step
+
+    # Residuals at the final points
+    residual = f(root) - a
+
+    return root, residual, iterations
+
+
 # =============
 # secant method
 # =============
 
-def secant_complex(f, z0, z1, a=0+0j, tol=1e-12, max_iter=100,
-                   alpha=0.5, max_bt=1, eps=1e-30, step_factor=5.0,
-                   post_smooth=True, jump_tol=10.0, verbose=False):
+def _secant_complex(f, z0, z1, a=0+0j, tol=1e-12, max_iter=100,
+                    alpha=0.5, max_bt=1, eps=1e-30, step_factor=5.0,
+                    post_smooth=True, jump_tol=10.0, verbose=False):
     """
-    Solves :math:``f(z) = a`` for many starting points simultaneously
-    using the secant method in the complex plane.
+    Solves :math:``f(z) = a`` for many starting points simultaneously using the
+    secant method in the complex plane.
 
     Parameters
     ----------
@@ -164,7 +706,7 @@ def secant_complex(f, z0, z1, a=0+0j, tol=1e-12, max_iter=100,
         diff_right[:-1] = numpy.abs(roots[:-1] - roots[1:])
         jump = numpy.minimum(diff_left, diff_right)
 
-        # ignore unconverged points
+        # ignore non-converged points
         median_jump = numpy.median(jump[~remaining])
         bad = (jump > jump_tol * median_jump) & ~remaining
 
@@ -178,13 +720,12 @@ def secant_complex(f, z0, z1, a=0+0j, tol=1e-12, max_iter=100,
             z_second = z_first + (roots[bad] - z_first) * 1e-2
 
             # re-solve just the outliers in one vector call
-            new_root, new_res, new_iter = secant_complex(
-                f, z_first, z_second, a[bad],
-                tol=tol, max_iter=max_iter,
-                alpha=alpha, max_bt=max_bt,
-                eps=eps, step_factor=step_factor,
-                post_smooth=False,      # avoid recursion
+            new_root, new_res, new_iter = _secant_complex(
+                f, z_first, z_second, a[bad], tol=tol, max_iter=max_iter,
+                alpha=alpha, max_bt=max_bt, eps=eps, step_factor=step_factor,
+                post_smooth=False,  # avoid recursion
             )
+
             roots[bad] = new_root
             residuals[bad] = new_res
             iterations[bad] = iterations[bad] + new_iter
@@ -199,12 +740,65 @@ def secant_complex(f, z0, z1, a=0+0j, tol=1e-12, max_iter=100,
     )
 
 
+# ================
+# plot diagnostics
+# ================
+
+def _plot_diagnostics(freeform, x, roots, residuals, iterations, tolerance,
+                      max_iter):
+    """
+    Plot the results of root-findings, including the residuals, number of
+    iterations, and the imaginary part of the roots.
+    """
+
+    with texplot.theme(use_latex=False):
+        fig, ax = plt.subplots(ncols=3, figsize=(9.5, 3))
+
+        for i in range(ax.size):
+            ax[i].axvline(x=freeform.lam_m, color='silver', linestyle=':')
+            ax[i].axvline(x=freeform.lam_p, color='silver', linestyle=':')
+
+        ax[0].axhline(y=tolerance, color='silver', linestyle='--',
+                      label='tolerance')
+        ax[0].plot(x, numpy.abs(residuals), color='black')
+
+        ax[1].axhline(y=max_iter, color='silver', linestyle='--',
+                      label='max iter')
+        ax[1].plot(x, iterations, color='black')
+
+        ax[2].axhline(y=0, color='silver', linestyle='-', alpha=0.5)
+        ax[2].plot(x, roots.imag, color='black')
+
+        ax[0].set_yscale('log')
+        ax[0].set_xlim([x[0], x[-1]])
+        ax[0].set_title('Residuals')
+        ax[0].set_xlabel(r'$x$')
+        ax[0].set_ylabel(r'$| f(z_0) - z |$')
+        ax[0].legend(fontsize='xx-small')
+
+        ax[1].set_xlim([x[0], x[-1]])
+        ax[1].set_title('Iterations')
+        ax[1].set_xlabel(r'$x$')
+        ax[1].set_ylabel(r'$n$')
+        ax[1].legend(fontsize='xx-small')
+
+        ax[2].set_xlim([x[0], x[-1]])
+        ax[2].set_title('Roots Imaginary Part')
+        ax[2].set_xlabel(r'$x$')
+        ax[2].set_ylabel(r'Im$(z_0)$')
+        ax[2].set_yscale('symlog', linthresh=1e-5)
+
+        plt.tight_layout()
+        plt.show()
+
+
 # ==========
 # decompress
 # ==========
 
-def decompress(freeform, size, x=None, delta=1e-4, max_iter=500,
-               tolerance=1e-8):
+def decompress(freeform, alpha, x, roots_init=None, method='newton',
+               delta=1e-4, max_iter=500, step_size=0.1, tolerance=1e-4,
+               plot_diagnostics=False):
     """
     Free decompression of spectral density.
 
@@ -214,23 +808,37 @@ def decompress(freeform, size, x=None, delta=1e-4, max_iter=500,
     freeform : FreeForm
         The initial freeform object of matrix to be decompressed
 
-    size : int
-        Size of the decompressed matrix.
+    alpha : float
+        Decompression ratio :math:`\\alpha = n / n_s = e^{t}`.
 
-    x : numpy.array, default=None
-        Positions where density to be evaluated at. If `None`, an interval
-        slightly larger than the support interval will be used.
+    x : numpy.array
+        Positions where density to be evaluated at.
+
+    roots_init : numpy.array, default=None
+        Initial guess for roots. If `None`, all root points are allocated at
+        a point below the center of support. If given, this is usually the
+        root that is found in the previous iteration of the called function.
+
+    method : {``'newton'``, ``'secant'``}, default=``'newton'``
+        Root-finding method.
 
     delta: float, default=1e-4
         Size of the perturbation into the upper half plane for Plemelj's
         formula.
 
     max_iter: int, default=500
-        Maximum number of secant method iterations.
+        Maximum number of iterations of the chosen method.
+
+    step_size: float, default=0.1
+        Step size for Newton iterations.
 
-    tolerance: float, default=1e-12
+    tolerance: float, default=1e-4
         Tolerance for the solution obtained by the secant method solver.
 
+    plot_diagnostics : bool, default=False
+        Plots diagnostics including convergence and number of iterations of
+        root finding method.
+
     Returns
     -------
 
@@ -261,58 +869,90 @@ def decompress(freeform, size, x=None, delta=1e-4, max_iter=500,
         >>> from freealg import FreeForm
     """
 
-    alpha = size / freeform.n
-    m = freeform._eval_stieltjes
-    # Lower and upper bound on new support
-    hilb_lb = (1 / m(freeform.lam_m + delta * 1j)).real
-    hilb_ub = (1 / m(freeform.lam_p + delta * 1j)).real
-    lb = freeform.lam_m - (alpha - 1) * hilb_lb
-    ub = freeform.lam_p - (alpha - 1) * hilb_ub
-
-    # Create x if not given
-    on_grid = (x is None)
-    if on_grid:
-        radius = 0.5 * (ub - lb)
-        center = 0.5 * (ub + lb)
-        scale = 1.25
-        x_min = numpy.floor(center - radius * scale)
-        x_max = numpy.ceil(center + radius * scale)
-        x = numpy.linspace(x_min, x_max, 500)
-    else:
-        x = numpy.asarray(x)
-
+    # Locations where Stieltjes is sought to be found at.
     target = x + delta * 1j
-    if numpy.isclose(alpha, 1.0):
-        return freeform.density(x), x, freeform.support
 
-    # Characteristic curve map
-    def _char_z(z):
-        return z + (1 / m(z)) * (1 - alpha)
+    if numpy.isclose(alpha, 1.0):
+        return freeform.density(x), roots_init
 
-    z0 = numpy.full(target.shape, numpy.mean(freeform.support) + 0.1j,
-                    dtype=numpy.complex128)
-    z1 = z0 - 0.2j
+    # Function that returns the second branch of Stieltjes
+    m = freeform._eval_stieltjes
 
-    roots, _, _ = secant_complex(
-        _char_z, z0, z1,
-        a=target,
-        tol=tolerance,
-        max_iter=max_iter
-    )
+    # ------
+    # char z
+    # ------
+
+    def _char_z(z_0):
+        """
+        Characteristic curve map. Returns z from initial z_0.
+        """
+
+        return z_0 + (1.0 / m(z_0)) * (1.0 - alpha)
+
+    # -----
+
+    # Initialize roots below the real axis
+    if roots_init is None:
+        roots_init = numpy.full(x.shape, numpy.mean(freeform.support) - 0.1j,
+                                dtype=numpy.complex128)
+
+    # Finding roots
+    if method == 'newton':
+        support = (freeform.lam_m, freeform.lam_p)
+
+        # Using initial points below the real line.
+        roots, residuals, iterations = \
+            _newton_method(_char_z, roots_init, target, support,
+                           enforce_wall=False, tol=tolerance,
+                           step_size=step_size, max_iter=max_iter,
+                           adaptive_stencil=True, halley=False, vertical=False)
+
+        # Using target points themselves as initial points. Since these points
+        # are above the real line, enforce branchcut as wall
+        # roots_init_1 = target
+        # roots_1, residuals_1, iterations_1 = \
+        #     _newton_method(_char_z, roots_init_1, target, support,
+        #                    enforce_wall=True, tol=tolerance,
+        #                    step_size=step_size, max_iter=max_iter,
+        #                    adaptive_stencil=True, halley=False,
+        #                    vertical=True)
+        #
+        # # If any result from residuals1 is better than residuals, use them
+        # good = numpy.abs(residuals_1) < numpy.abs(residuals)
+        # roots[good] = roots_1[good]
+        # residuals[good] = residuals_1[good]
+        # iterations[good] = iterations_1[good]
+
+    elif method == 'secant':
+        z0 = numpy.full(x.shape, numpy.mean(freeform.support) + 0.1j,
+                        dtype=numpy.complex128)
+        z1 = z0 - 0.2j
+
+        roots, _, _ = _secant_complex(_char_z, z0, z1, a=target, tol=tolerance,
+                                      max_iter=max_iter)
+    else:
+        raise NotImplementedError('"method" is invalid.')
 
     # Plemelj's formula
     z = roots
-    char_s = m(z) / alpha
+    char_s = numpy.squeeze(m(z)) / alpha
     rho = numpy.maximum(0, char_s.imag / numpy.pi)
-    rho[numpy.isnan(rho) | numpy.isinf(rho)] = 0
-    if on_grid:
-        x, rho = x.ravel(), rho.ravel()
-        # dx = x[1] - x[0]
-        # left_idx, right_idx = support_from_density(dx, rho)
-        # x, rho = x[left_idx-1:right_idx+1], rho[left_idx-1:right_idx+1]
-        rho = rho / numpy.trapezoid(rho, x)
 
-    return rho.reshape(*x.shape), x, (lb, ub)
+    # Check any nans are in the roots
+    num_nan = numpy.sum(numpy.isnan(roots))
+    if num_nan > 0:
+        raise RuntimeWarning(f'"nan" roots detected: num: {num_nan}.')
+
+    # dx = x[1] - x[0]
+    # left_idx, right_idx = support_from_density(dx, rho)
+    # x, rho = x[left_idx-1:right_idx+1], rho[left_idx-1:right_idx+1]
+    rho = rho / numpy.trapezoid(rho, x)
+
+    if plot_diagnostics:
+        _plot_diagnostics(freeform, x, roots, residuals, iterations, tolerance,
+                          max_iter)
+
+    return rho, roots
 
 
 # =======================
@@ -334,7 +974,7 @@ def reverse_characteristics(freeform, z_inits, T, iterations=500,
 
     target_z, target_t = numpy.meshgrid(z_inits, t_eval)
 
-    z = numpy.full(target_z.shape, numpy.mean(freeform.support) - .1j,
+    z = numpy.full(target_z.shape, numpy.mean(freeform.support) - 0.1j,
                    dtype=numpy.complex128)
 
     # Broken Newton steps can produce a lot of warnings. Removing them for now.