nberliner / pointobject Goto Github PK

import numpy as np
import matplotlib.pylab as plt
from scipy.ndimage import gaussian_filter1d

class Color:
    """
    Helper to assign colors to float or integer values mapped to a given range.
    """
    def __init__(self, scaleMin=None, scaleMax=None):
        self.Nglobal = dict()
        self.cmap = plt.get_cmap('seismic_r')

        self.scaleMin = scaleMin
        self.scaleMax = scaleMax

        if scaleMin == scaleMax and scaleMin is not None:
            print('Warning: trying to set zero scaling range!')
            self.scaleMin = scaleMin
            self.scaleMax = scaleMin * 1.1

    def __call__(self, N):

        if self.scaleMin is None and self.scaleMax is not None:
            c = float(N) / self.scaleMax
        elif self.scaleMin is not None and self.scaleMax is None:
            c = float(N)  - self.scaleMin
        elif self.scaleMin is None and self.scaleMax is None:
            c = float(N)
        else:
            c = (float(N) - self.scaleMin) / ( self.scaleMax - self.scaleMin )

        return self.cmap(c)

def curvature(x, y, sigma=1, absolute=False):
    # Credit goes here: http://stackoverflow.com/a/28310758/1922650
    #first and second derivative
    x1 = gaussian_filter1d(x,  sigma=sigma, order=1, mode='wrap')
    x2 = gaussian_filter1d(x1, sigma=sigma, order=1, mode='wrap')
    y1 = gaussian_filter1d(y,  sigma=sigma, order=1, mode='wrap')
    y2 = gaussian_filter1d(y1, sigma=sigma, order=1, mode='wrap')
    if absolute:
        return np.abs(x1*y2 - y1*x2) / np.power(x1**2 + y1**2, 3./2)
    else:
        return (x1*y2 - y1*x2) / np.power(x1**2 + y1**2, 3./2)


# The user defined parameters
radius   = 5.0
sampling = 30
sigma    = 1

# Generate the test data
alpha = np.linspace(-np.pi/2,np.pi/2, sampling)
X = radius*np.cos(alpha)
Y = radius*np.sin(alpha)

# Calculate the radius from the curvature
C = 1/curvature(X, Y, sigma=sigma)
curvatureMax   = np.max(C)
curvatureMin   = np.min(C)

color  = Color(scaleMin=curvatureMin, scaleMax=curvatureMax)
cColor = [ color(i) for i in C ]

# Plot the test data color coded for curvature
fig = plt.figure(figsize=(14,7), dpi=120)
ax1  = fig.add_subplot(121)
ax1.set_title("Generated data")
ax1.set_xlim( [-radius-1,radius+1] )
ax1.set_ylim( [-radius-1,radius+1] )
ax1.scatter(x=X, y=Y, c=cColor, alpha=1, edgecolor='none', s=10, cmap=plt.cm.seismic_r)

ax2 = fig.add_subplot(122)
ax2.set_title("Calculated curvature")
ax2.plot(C)
ax2.axhline(y=radius, color='red')

/home/berliner/08 GitHub/PointObject/lib/contour.py:202: UserWarning: Warning: see bandwidth parameter was estimated to be optimal at one sample boundary
  warnings.warn("Warning: see bandwidth parameter was estimated to be optimal at one sample boundary")
/home/berliner/08 GitHub/PointObject/lib/contour.py:203: UserWarning: Try shifting the sample window!
  warnings.warn("Try shifting the sample window!")