circular_leastsq.py

import numpy as np
import matplotlib.pyplot as plt
import abel
import scipy.interpolate, scipy.ndimage.interpolation
import time

global X,Y,R,T

size = 301

#define a simple grid
x = np.linspace(-size*0.5+0.5,size*0.5-0.5,size)
y = np.linspace(-size*0.5+0.5,size*0.5-0.5,size)
X,Y = np.meshgrid(x,y)

# convert to polar coords, R and theta
R = np.sqrt(X**2+Y**2)
T = np.arctan2(Y,X)

def flower_scaling(angle, factor=0.02):
    # define a simple scaling that will squish a circle into a "flower"
    return 1 + factor*(np.sin(7*angle))

def rescale(IM, rescale_func, prettify=False):
    global X,Y,R,T
    
    size = IM.shape[0]

    if prettify: IM = IM*R**2 # just making the sample image more beautiful...

    R_rescaled = R*rescale_func(T) # rscale the radial coords by our angle-dependent scaling

    # now, convert our squished angular grid back to a squished X, Y grid
    X_rescaled = R_rescaled*np.cos(T)
    Y_rescaled = R_rescaled*np.sin(T)

    IM_rescaled = scipy.ndimage.interpolation.map_coordinates(IM,(Y_rescaled+size*0.5,X_rescaled+size*0.5))
    
    return IM_rescaled
    
Z = abel.tools.analytical.sample_image(n=size, name='dribinski', sigma=2)
Z_rescaled = rescale(Z, flower_scaling, prettify=True)
# Z_rescaled[Z_rescaled<10e5]=0

# now plot the results:
fig, axs = plt.subplots(1,2,figsize=(8,4))

axs[0].imshow(Z_rescaled,          aspect='auto', origin='lower')
axs[0].set_title('Original')


def spline_rescale(IM,spline_values):
    global X,Y,R,T
    
    def spline_fit(requested_angles, spline_values):
        assumed_angles = np.linspace(-np.pi, np.pi, spline_values.shape[0])
        spline = scipy.interpolate.InterpolatedUnivariateSpline(assumed_angles, spline_values)
    
        return spline(requested_angles)
        
    size = IM.shape[0]

    R_rescaled = R*spline_fit(T, spline_values) # rscale the radial coords by our angle-dependent scaling

    # now, convert our squished angular grid back to a squished X, Y grid
    X_rescaled = R_rescaled*np.cos(T)
    Y_rescaled = R_rescaled*np.sin(T)
   
    IM_rescaled = scipy.ndimage.interpolation.map_coordinates(IM,(Y_rescaled+size*0.5,X_rescaled+size*0.5))  
    return IM_rescaled


def error_func(spline_values, IM):
    print spline_values

    IM_rescaled = spline_rescale(IM, spline_values)
    
    polar, r, t = abel.tools.polar.reproject_image_into_polar(np.rot90(IM_rescaled), dt=360/42.)
    print t
    
    t1 = time.time()
    # calculate error
    
    slice_prev = np.zeros_like(polar)
    slice_next = np.zeros_like(polar)
    slice_oppo = np.zeros_like(polar)
    
    
    slice_prev[:,0:-1] = polar[:,1:]
    slice_prev[:,-1]   = polar[:,0]
    
    slice_next[:,1:]   = polar[:,:-1]
    slice_next[:,0]    = polar[:,-1]
    
    size = polar.shape[1]
    slice_oppo[:,:-size/2] = polar[:, size/2:]
    slice_oppo[:,size/2:] = polar[:, :-size/2]
    
    avg = np.mean(polar, axis=0)
    
    error = np.sum(np.abs(polar-slice_next) + np.abs(polar-slice_prev)  + np.abs(polar-slice_oppo) + np.abs(polar - avg[None,:]), axis=0)
    return error


guess_values = np.ones(20)
bound=0.06

global t1
t1 = time.time()
res = scipy.optimize.least_squares(error_func, guess_values, args=([Z_rescaled]), bounds=(guess_values-bound, guess_values+bound))
x = res.x

Z_recovered = spline_rescale(Z_rescaled, x)

axs[1].set_title('Circularized')
axs[1].imshow(Z_recovered, aspect='auto', origin='lower')
plt.tight_layout()

plt.savefig('Circular.png', dpi=200)



plt.show()