Output the value of the Ulam spiral

ulam.py

import numpy as np


def ulam_polar(x, y, verbose=False):
    """
    Solves for the value at location (x, y) in the Ulam spiral by

     (1) converting to (x, y) polar coordinates (r, theta)
     (2) developing a `Ulam square` for a given ring-size
     (3) finding the intersection of the the Ulam square and the input vector

    Known issues,

     (1) Rounding error from representing a square in polar coordinates


    Figure 1:

                   . (x, y)
             |y   /
             |   / r
         +-------+
         |   | / | rad(theta)
         |   |/__|___ x
         |       |
         +-------+  <--- rad(-pi/4) at every point r can be calculated
                            by the next consecutive odd number squared
         Ulam Square
                               ulam(r, -pi/4) = (2r + 1) ^ 2


    The intersection of r and the Ulam square is the return value.

    The Ulam square is the values of the ulam spiral represented in
    polar-coordinates

    """
    r = max(abs(x), abs(y)) + 1
    theta = np.arctan2(y, x)
    if theta < 0.:
        theta += 2 * np.pi
    if theta == 0.:
        theta = 2 * np.pi

    width = 2 * r - 1
    last = width ** 2  # on the -0.25 pi theta line

    n_ring = 4 * width - 4
    step = (2 * np.pi) / n_ring

    first = last - n_ring + 1

    angle = np.arange(-0.25 * np.pi + step, 1.75 * np.pi + step, step)
    angle = np.array([a if a > 0. else 2 * np.pi + a for a in angle])
    val = np.arange(first, last+1, 1)

    d_theta = np.abs(angle - theta)
    i = np.argmin(d_theta)

    if verbose:
        print('Angle {:3f} found at index {}, Value {}'.format(theta, i, val[i]))
        print(np.vstack((d_theta, angle, val)).T)

    return val[i]