Calculate complex square roots with the correct sign for trajectories in the complex plane.

signed_sqrt.py

#!/usr/bin/env python3

"""
Calculate complex square roots with the correct sign for trajectories in the
complex plane.

Taking the square root of values along a continuous trajectory in the complex
plane is not as simple as calling a sqrt function, because the complex square
root is a multi-valued function whose Riemann surface is composed of two
sheets. If the trajectory crosses the branch cut (conventionally placed along
the negative reals), one must transition to the other sheet and flip the sign
for all square roots until the branch cut is crossed again. Using only the
principal value of the square root will lead to cusps and discontinuities at
the crossing points.

Conceptually, one can simply keep track of the number of crossings thus far (m)
and then multiply all square roots by (-1)^m. This is not difficult, but is
tedious. This module provides SignedSqrt, a class whose instances keep track of
where the trajectory has been and adjust the sign accordingly, and signed_sqrt,
a function which applies the adjusted square root to an entire trajectory at
once.
"""

# This module is made available under the terms of the MIT license:

# Copyright (c) 2015 Dmitri Iouchtchenko
#
# 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.

import numpy as N


class SignedSqrt:
    """
    Continuous square roots for a trajectory in the complex plane.

    Find the square root of complex values with the correct sign (on the
    correct sheet of the surface) for a trajectory that is determined
    iteratively.
    """

    def __init__(self):
        self.factor = 1
        self.sign = None

    def __call__(self, v):
        """
        Args:
            v: Next complex number in the trajectory.

        Returns:
            Square root of v with the correct sign.
        """

        s = N.sign(v.imag)

        # NumPy's sign returns 0 for 0, but its sqrt is continuous from above
        # along the negative reals, so we consider the real line to be in the
        # upper half of the complex plane.
        if s == 0:
            s = 1

        if self.sign is None:
            self.sign = s
        elif self.sign != s:
            self.sign = s

            if N.sign(v.real) < 0:
                self.factor *= -1

        return self.factor * N.sqrt(v)


def signed_sqrt(vs):
    """
    Continuous square roots for a trajectory in the complex plane.

    Find the square root of complex values with the correct sign (on the
    correct sheet of the surface) for a trajectory that is known in its
    entirety.

    Args:
        vs: NumPy vector of complex numbers.

    Returns:
        Square roots of the elements of vs with the correct signs.
    """

    if vs.size == 0:
        return vs

    signs_imag = N.sign(vs.imag)
    # NumPy's sign returns 0 for 0, but its sqrt is continuous from above along
    # the negative reals, so we consider the real line to be in the upper half
    # of the complex plane.
    signs_imag[signs_imag == 0] = 1
    signs_imag_rolled = N.roll(signs_imag, 1)
    # Never flip the first sign.
    signs_imag_rolled[0] = signs_imag[0]

    # Make a vector of booleans, where each element is False iff it corresponds
    # to an appropriate sign change.
    transitions = N.logical_or(N.sign(vs.real) >= 0, signs_imag == signs_imag_rolled)

    # Convert the booleans to 1 and -1 and use the cumulative product to
    # "carry" each occurrence of -1 until the next one.
    return N.cumprod(2 * transitions - 1) * N.sqrt(vs)


# Simple tests.
if __name__ == '__main__':
    ts = N.linspace(0, 4 * N.pi, 400)
    trajectories = []

    # Trajectory 0 (circle, CCW).
    trajectories.append(('Circle', ts / N.pi, N.cos(ts) + 1j * N.sin(ts), 0.8))

    # Trajectory 1 (from http://mathworld.wolfram.com/HeartCurve.html, CW).
    xs = 16 * N.power(N.sin(ts), 3)
    ys = 13 * N.cos(ts) - 5 * N.cos(2 * ts) - 2 * N.cos(3 * ts) - N.cos(4 * ts)
    trajectories.append(('Heart', ts / N.pi, xs + 1j * ys, 0.2))

    # Trajectory 2 (square, CCW).
    # Odd number of elements, so that we hit zero.
    segment = N.linspace(-1.0, 1.0, 51)[:-1]
    traj = N.concatenate(((-segment)+1j, -1-(segment*1j), segment-1j, 1+segment*1j))
    traj = N.concatenate((traj, traj))
    trajectories.append(('Square', N.linspace(0, 4, len(traj)), traj, 0.8))

    # Trajectory 3 (infinity, CCW).
    # It appears that each time we cross the origin, we must have a cusp in
    # either the real or imaginary part.
    trajectories.append(('Infinity', ts / N.pi, N.cos(ts) + 1j * N.sin(2 * ts), 0.8))

    for i, (title, ts, trajectory, scale) in enumerate(trajectories):
        # Compute.
        rooted1 = N.empty_like(trajectory)
        ss = SignedSqrt()

        for step in range(len(trajectory)):
            rooted1[step] = ss(trajectory[step])

        rooted2 = signed_sqrt(trajectory)

        # Check.
        assert N.all(rooted1 == rooted2), '{}: {} != {}'.format(i, rooted1, rooted2)

        # Plot.
        import matplotlib.pyplot as plt

        _, axes = plt.subplots(2, sharex=True)

        axes[0].plot(ts, scale * trajectory.real, color='black', linestyle='dotted', label='Trajectory (scaled)')
        axes[0].plot(ts, N.sqrt(trajectory).real, color='red', linestyle='dashed', linewidth=2.0, label='Unsigned')
        axes[0].plot(ts, rooted1.real, color='blue', label='Signed')
        axes[1].plot(ts, scale * trajectory.imag, color='black', linestyle='dotted', label='Trajectory (scaled)')
        axes[1].plot(ts, N.sqrt(trajectory).imag, color='red', linestyle='dashed', linewidth=2.0, label='Unsigned')
        axes[1].plot(ts, rooted1.imag, color='blue', label='Signed')

        axes[0].set_title(title)

        axes[1].set_xlabel(r'$t$')
        axes[0].set_ylabel('Real')
        axes[1].set_ylabel('Imaginary')

        for axis in axes:
            axis.legend(prop={'size': 6})

        plt.savefig('signed_sqrt_test_{}.pdf'.format(i))