Yingyang classification animation

#!/usr/bin/python3
#
#Yingyang 2D function classification:
#A 4-class unbalanced dataset, for if you're tired of the classical XOR.
#Using a (large) number of randomly-chosen radial basis functions (RBFs) in front of a logistic regression (LR).
#The individual RBFs 'measure' the distance from their centres to a point in the dataset.
#Note there's no separare test set, as this would make the script a bit more complicated and
#the performance is obvious from the images.
#The dataset is unbalanced (less green/yellow dots than red/blue), so it makes sense that the first
#models ignore the less-present classes.
#
#larsupilami73


#imports
import os,sys
from scipy import *
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score


#generator for the yingyang dataset
def yingyang(n=10000):
    while (n>0):
        n -=1
        r = random.uniform(0,0.5)
        a = random.uniform(0,2*pi)
        x = 0.5 + r * cos(a)
        y = 0.5 + r * sin(a)
        if sqrt((x-0.25)**2 + (y-0.5)**2)<0.12:
            c = 0
        elif sqrt((x-0.75)**2 + (y-0.5)**2)<0.12:
            c = 1
        elif sqrt((x-0.25)**2 + (y-0.5)**2)<0.25:
            c = 2
        elif sqrt((x-0.75)**2 + (y-0.5)**2)<0.25:
            c = 3
        elif y > 0.5:
            c = 3
        else:
            c = 2
        yield (x,y,c)


#main
if __name__ == '__main__':

    numframes = 25
    try:
        os.mkdir('frames')
    except:
        pass
    try:
        os.system('rm -f ./frames/*.png')   #remove previous run frames
    except:
        pass

    #generate the 'ground truth' dataset
    L,C = [],[]
    X,Y = [[],[],[],[]],[[],[],[],[]]
    for x,y,c in yingyang():
        L.append([x,y,c])
        C.append(c)
        X[c].append(x)
        Y[c].append(y)

    L = array(L)
    #savetxt('yingyang.txt',L,fmt='%.6f %.6f %i',header='yingyang dataset.\nX   Y   Class') #save it

    #plot the ground truth
    fig,(ax1,ax2) = plt.subplots(1,2,figsize=(8,4))
    colors = ['green','orange','blue','red']
    for c in range(4):
        ax1.scatter(X[c],Y[c],s=5,c=colors[c],label='class %i'%c)
    ax1.legend(loc='upper right',ncol=1,markerscale=2.0,fontsize='small')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    ax1.set_title('Ground truth')
    ax1.set_aspect(1)

    #lets first try a logistic regression on the ground truth
    P = array([[x_,y_] for x_,y_,c_ in L])
    clf = LogisticRegression(random_state=0,solver='lbfgs',multi_class='auto')
    scores = cross_val_score(clf, P, C, cv=5)
    print('Only LR scores: ',scores)
    print("Only LR accuracy: %0.2f (+/- %0.2f)" % (scores.mean(), scores.std() * 2))

    #make a drawing to show what the logistic regression resulted in
    clf.fit(P,C)
    score = clf.score(P,C)
    print('Train score entire dataset: %2.3g'%score)
    prediction = clf.predict(P)
    XP,YP = [[],[],[],[]],[[],[],[],[]]
    for c,p in zip(prediction,P):
        XP[c].append(p[0])
        YP[c].append(p[1])

    for c in range(4):
        ax2.scatter(XP[c],YP[c],s=5,c=colors[c],label=None)#'class %i'%c)
    #draw the errors more clearly
    fe = False
    for ce,p,c in zip(prediction,P,C):
        if ce != c:
            if fe:
             ax2.scatter([p[0]],[p[1]],s=8,linewidth=0.75,c='none',edgecolors= "black",label=None)
            else:
                ax2.scatter([p[0]],[p[1]],s=8,c='none',linewidth=0.75,edgecolors= "black",label='errors')
                fe = True

    ax2.legend(loc='upper right',ncol=2,markerscale=2.0,fontsize='small')
    ax2.set_xlabel('X')
    #ax2.set_ylabel('Y')
    ax2.set_title('LR only. Score: %2.3g'%score)
    ax2.set_aspect(1)

    plt.savefig('./frames/frame%05d.png'%0)
    ax2.cla()


    #Now lets construct new features which are nonlinear functions of X and Y.
    #Here we choose random radial basis functions (RBF).
    #Then check how well the LR does when fed by these new features
    R,S = [],[] #center positions and radii of the RBFs
    for frame in range(1,numframes+1):
        print('frame %i/%i'%(frame,numframes))
        nrbf = frame
        R.append(random.normal(size=2,loc=[0.5,0.5],scale=0.25))  #center positions of the RBFs
        S.append(random.uniform(low=0.1,high=0.3))             #radii of the RBFs
        Z = array([[ exp(-((p_[0]-r_[0])**2+(p_[1]-r_[1])**2)/s_**2) for r_,s_ in zip(R,S) ] for p_ in P])
        clf2 = LogisticRegression(random_state=0,solver='lbfgs',multi_class='multinomial',max_iter=1000)
        clf2.fit(Z,C)
        score = clf2.score(Z,C)
        print('Train score with %i RBFs: %2.3g'%(nrbf,score))
        XP,YP = [[],[],[],[]],[[],[],[],[]]
        prediction = clf2.predict(Z)
        for c,p in zip(prediction,P):
            XP[c].append(p[0])
            YP[c].append(p[1])

        for c in range(4):
            ax2.scatter(XP[c],YP[c],s=5,c=colors[c],label=None)#'class %i'%c) too many labels overloads picture

        #draw the errors more clearly
        fe = False
        for ce,p,c in zip(prediction,P,C):
            if ce != c:
                if fe:
                 ax2.scatter([p[0]],[p[1]],s=8,linewidth=0.75,c='none',edgecolors= "black",label=None)
                else:
                    ax2.scatter([p[0]],[p[1]],s=8,linewidth=0.75,c='none',edgecolors= "black",label='error')
                    fe = True

        ax2.legend(loc='upper right',ncol=2,markerscale=2.0,fontsize='small')
        ax2.set_xlabel('X')
        #ax2.set_label('Y')
        ax2.set_title('%i RBFs --> LR. Score: %2.3g' %(nrbf,score))
        ax2.set_aspect(1)
        #draw the RBFs position and radius
        #for r,s in zip(R,S):
        #    ax2.add_artist(plt.Circle(r,s,linestyle='--',linewidth=0.5,alpha=0.7,fill=False))

        plt.savefig('./frames/frame%05d.png'%frame)
        ax2.cla()


plt.close()

#now assemble all the pictures in a movie
print('converting to gif!')
os.system('convert -delay 40 -loop 0 ./frames/*.png yingyang.gif')
#or to mp4 with .mp4
#print('converting to mp4!')
#os.system("ffmpeg -y -r 40 -i ./frames/frame%05d.png -c:v libx264 -vf fps=20 yingyang.mp4")