EllipticalReg

EllipticalReg

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize


class EllipticalReg():
    """   2         2
    (x-x0)  + (y-y0)
    ------    ------  =  1
      a^2       b^2
    """

    def __init__(self, **kw):
        self.vars_ = {'a':1, 'b':1, 'x0':0, 'y0':0}
        self.args_ = {}

        for k,v in kw.items():
            if (k in self.vars_):
                self.args_[k] = v
                self.vars_.pop(k)

    def get_all(self, var):
        for v in ['a', 'b', 'x0', 'y0']:
            if (v in self.vars_):
                yield var[list(self.vars_).index(v)]
            else:
                yield self.args_[v]

    def fun(self, var):
        a, b, x0, y0 = self.get_all(var)
        dx = (self.x_-x0)
        x = np.empty(len(dx))
        y = np.empty(len(dx))

        ri = (self.y_-y0)[dx!=0]/dx[dx!=0]
        x[dx!=0] = x0 + np.where(self.x_[dx!=0]>x0, 1, -1)*abs(a*b)/np.sqrt( b**2 + (ri*a)**2 )
        y[dx!=0] = y0 + ri*(x[dx!=0]-x0)

        x[dx==0] = x0
        y[dx==0] = y0+np.where(self.y_[dx==0] > y0, 1, -1)*b
        return np.sum( ( np.sqrt((x-self.x_)**2+(y-self.y_)**2) )**self.n_ )

    def reg(self, x, y, **kw):
        self.x_ = np.array(x)
        self.y_ = np.array(y)
        self.n_ = kw.pop('n', 1)
        self.vars_.update(kw)
        m = minimize(self.fun, list(self.vars_.values()), options={'maxiter':10})

        return self.get_all(m.x)


def plot_elipse(a, b, x0, y0, theta, *args, **kwargs):
    x = a*np.cos(theta)+x0
    y = b*np.sin(theta)+y0

    plt.plot(x, y, *args, **kwargs)
    return x, y

if __name__ == '__main__':
    a, b, x0, y0 = 10, 20, 5, 10
    theta = np.linspace(0, 2*np.pi, 100)
    np.random.seed(10)
    rand1 = 2*(np.random.rand(len(theta))-0.5)
    rand2 = 2*(np.random.rand(len(theta))-0.5)
    x = a*np.cos(theta) + x0 + rand1
    y = b*np.sin(theta) + y0 + rand2

    plt.plot(x, y, 'o')

    e = EllipticalReg(a=a)

    a, b, x0, y0 = e.reg(x, y, n=2)

    plot_elipse(a, b, x0, y0, theta)

    plt.show()