AllRGB Annealing code

#include <stdio.h>
#include <stdlib.h>
#include <strings.h>
#include <math.h>
// Bits per plane. Must be even and <=8, namely either 2, 4, 6, 8
#define CBIT 4

// temperature range
#define BMAX 10.0
#define BMIN (1e-2)
// steps
#define TDIV 1000
#define MAXS 1000

// Penalty for breaking all-rgb restriction
//  for each same color pixel pair.
//  if there are N pixels with the same color,
//   the penalty is oc * N * (N-1)/2 .
//   commenting out OCMAX is equivalent to prohibit such
//   same color occurence, namely OCMAX=infinity
//#define OCMAX 256

/////////////////////////////////////////////////////
// no need to edit below

// Misc. consts
#define VMAX (1<<CBIT)
#define CBMASK (VMAX-1)
#define rgb2idx(r,g,b) ( (r)|((g)<<CBIT)|((b)<<(CBIT*2)))

#define L (1<<(CBIT + CBIT/2))
#define L2 (L*L)
#define LMASK (L2-1)
#define LMASKX (L-1)
#define LMASKY (LMASK ^ LMASKX)

// global var block
typedef unsigned char rgbtype;
// main color data
//unsigned int col[L2];
rgbtype colr[L2],colg[L2], colb[L2];
// save the best snapshot
rgbtype colrbest[L2],colgbest[L2], colbbest[L2];
// energy
long long ebest,  ene;
// penalty for breaking allRGB rule
long long penalty;
// color->pos index
int cidx[L2];

// max possible 1-pixel energy
#ifdef OCMAX
#define E1MAX (VMAX*VMAX*4 + OCMAX*200)
#else
#define E1MAX (VMAX*VMAX*4)
#endif
// Metropolis threshold table
// energy difference of swapping 2 pixels can be as large as 2*E1MAX
static int thtab[E1MAX*2+10];

void init_th(double beta, int imax)
{
    int i;
    thtab[0] = imax/2;  // 50% chance for microcanonical change
    if (beta<0.0)
    {
	// quench case
	for(i=1;i<E1MAX*2+10;i++)
	    thtab[i] = 0;
	return;
    }
    for(i=1;i<E1MAX*2+10;i++)
    {
	thtab[i] = (int)(imax * exp( - i * beta));
    }
}


void saveit()
{
    int i;
    for(i=0;i<L2;i++)
    {
	colrbest[i]=colr[i];
	colgbest[i]=colg[i];
	colbbest[i]=colb[i];
    }
}

#ifdef OCMAX
int oc;  // actual coefficient of penalty
// number of pixel for each color
unsigned int cn[L2];
// pointer string to the same color pixel at each posision
int cnext[L2];
int cprev[L2];

// pointer string management

void addcol(int c, int adr)
{
    int adr_next=cidx[c];
    cidx[c]= adr;
    cnext[adr]= adr_next;
    cprev[adr]=-1;
    if(adr_next != -1) cprev[adr_next]=adr;
    cn[c]++;
}

void removecol(int c, int adr)
{
    int adr_next=cnext[adr];
    int adr_prev=cprev[adr];

    if (adr_next != -1) cprev[adr_next] = adr_prev;
    if (adr_prev != -1) cnext[adr_prev] = adr_next;
    else cidx[c]=adr_next; 
    cn[c]--;
}

void init_cidx()
{
    //assume col is ready
    int i;
    for(i=0;i<L2;i++)
    {
	cn[i]=0; 
	cnext[i]= -1;
	cprev[i]= -1;
	cidx[i] = -1;
    }
    for(i=0;i<L2;i++)
    {
	int r=colr[i];
	int g=colg[i];
	int b=colb[i];
	addcol(rgb2idx(r,g,b),i);
    }
}

// for debug
void dump_cidx()
{
    int i;
    printf("##############\n");
    for(i=0;i<L2;i++)
    {
	printf("C%d:%03d ",i, cn[i]);
	int h=cidx[i];
	while(h!=-1)
	{
	    printf(" %d ",h);
	    h = cnext[h];
	}
	printf("\n");
    }
}
#else
// macro
#define addcol(c, adr) cidx[c]=adr
#define removecol(c, adr) /*c adr */
void init_cidx()
{
    //assume col is ready
    int i;
    for(i=0;i<L2;i++)
    {
	cidx[i] = -1; 
    }
    for(i=0;i<L2;i++)
    {
	int r=colr[i];
	int g=colg[i];
	int b=colb[i];
	int idx =rgb2idx(r,g,b);
	if (cidx[idx] != -1)
	{
	    printf("# init_cidx: colors are not unique\n");
	    exit(1);
	}
	cidx[idx] = i;
    }
}
#endif

// dump color to a file
void dump(int idx, rgbtype *rr, rgbtype *gg, rgbtype *bb)
{
    unsigned char *pbuf = (unsigned char *)malloc(L2*3);
    char fn[128];
    sprintf(fn, "%04d.ppm",idx);
    FILE *fp=fopen(fn,"w");
    if (fp==NULL)
    {
	printf("cannot open %s for dump\n",fn);
	exit(1);
    }
    int i;
    int vv = 255/(VMAX-1);
    printf("#Dump %s, color scale *%d Size %f Mb\n",fn, vv, L2*3e-6);
    for(i=0;i<L2;i++)
    {
	pbuf[i*3+0] = vv*rr[i];
	pbuf[i*3+1] = vv*gg[i];
	pbuf[i*3+2] = vv*bb[i];
    }
    fprintf(fp, "P6\n%d %d\n255\n", L,L);
    fwrite(pbuf,1,L2*3,fp);
    fclose(fp);
    free(pbuf);
}

// resume from an image file
//  or import smaller size image and expand, add lower bit data
void load_ppm(char *fn)
{
    char buf[128];
    FILE *fp=fopen(fn, "r");
    if (fp==NULL)
    {
	printf("# cannot open file %s\n",fn);
	exit(1);
    }
    int ll;
    fgets(buf,99,fp);
    fgets(buf,99,fp);
    sscanf(buf, "%d", &ll);
    fgets(buf,99,fp);
    printf("#Loading %s L=%d\n",fn, ll);

    if (ll==L)
    {
	//resume
	printf("  #Same size data\n");
	unsigned char *pbuf = (unsigned char *)malloc(L2*3);
	fread(pbuf,1,L2*3,fp);
	fclose(fp);
	int vv = 255/(VMAX-1);
	printf("  # read complete, color scale /%d\n", vv);
	int i;
	for(i=0;i<L2;i++)
	{
	    unsigned char *c = &pbuf[i*3];
	    colr[i] = c[0]/vv;
	    colg[i] = c[1]/vv;
	    colb[i] = c[2]/vv;
	}
	free(pbuf);
	return;
    }
    if (ll*ll==L)
    {
	// self-similar
	printf("  # sqrt(L) size data\n");
	unsigned char *pbuf = (unsigned char *)malloc(ll*ll*3);
	fread(pbuf,1,ll*ll*3,fp);
	fclose(fp);

	int v2 = (1<<(CBIT/2));
	int vv = 255/(v2-1);
	printf("  # read complete, color scale /%d\n", vv);

	int x,y;

	//block
	for(x=0;x<ll;x++){
	    for(y=0;y<ll;y++){
		int ad1 = (x+y*ll)*3;
		int r1=(pbuf[ad1 +0]/vv)<<(CBIT/2);
		int g1=(pbuf[ad1 +1]/vv)<<(CBIT/2);
		int b1=(pbuf[ad1 +2]/vv)<<(CBIT/2);

		int xx,yy;
		for(xx=0;xx<ll;xx++){
		    for(yy=0;yy<ll;yy++){
			int x2=xx;
			int y2=yy;
			if (x&1) x2=ll-1-xx;
			if (y&1) y2=ll-1-yy;
			int ad0 = (x2+y2*ll)*3;

			int r=pbuf[ad0 +0]/vv +r1;
			int g=pbuf[ad0 +1]/vv +g1;
			int b=pbuf[ad0 +2]/vv +b1;
			int ad1 = (x*ll+xx)+(y*ll+yy)*L;
			colr[ad1] = r;
			colg[ad1] = g;
			colb[ad1] = b;
		    }}
	    }}
	free(pbuf);
	return;
    }

    if (ll==L/8)
    {
	// expand CBIT-2 image
	printf("  # 1/8 size data\n");
	unsigned char *pbuf = (unsigned char *)malloc(L2*3/64);
	fread(pbuf,1,L2*3/64,fp);
	fclose(fp);
	printf("  # read complete\n");

	int ofs[8][8][3];
	int x,y;
	int vv = 255/(VMAX/4-1);

	// 8x8 optimum solution
	int ofsR[8][8];
	int ofsG[8][8];
	int ofsRy[8]={1,0,0,1,2,3,3,2};
	int ofsGx[8]={0,0,1,2,3,3,2,1};
	int ofsB[8][8]=
	{
	    {2,3,3,3,3,2,2,2},
	    {2,3,3,3,3,1,1,1},
	    {0,1,2,2,2,0,0,0},
	    {0,1,1,1,1,0,0,0},
	    {0,1,1,1,1,0,0,0},
	    {0,2,2,2,1,0,0,0},
	    {1,3,3,3,3,2,1,1},
	    {2,3,3,3,3,2,2,2}
	};

	// fine bit part
	for(x=0;x<8;x++){
	    for(y=0;y<8;y++){
		ofs[x][y][0] = x/2;
		ofs[x][y][1] = y/2;
		ofs[x][y][2] = (x&1) + (y&1)*2;
		ofsR[y][x]=ofsRy[y];
		ofsG[y][x]=ofsGx[y];
	    }}

	// repeat small image
	//block
	for(x=0;x<8;x++){
	    for(y=0;y<8;y++){
		int r0 =ofsR[x][y]*4;
		int g0 =ofsG[x][y]*4;
		int b0 =ofsB[x][y]*4;
		int xx,yy;
		for(xx=0;xx<L/8;x++){
		    for(yy=0;yy<L/8;yy++){
			int x2=xx;
			int y2=yy;
			if (x&1) x2=L/8-1-xx;
			if (y&1) y2=L/8-1-yy;
			int ad0 = (x2+y2*(L/8))*3;

			int r=pbuf[ad0 +0]/vv +r0;
			int g=pbuf[ad0 +1]/vv +g0;
			int b=pbuf[ad0 +2]/vv + b0;

			int adr=(x*ll+xx)+(y*ll+yy)*L;
			colr[adr]=r;
			colg[adr]=g;
			colb[adr]=b;
		    }}
	    }}

	free(pbuf);
	return;
    }
    printf("#Couldn't load %s\n",fn);
    exit(1);
}

// old;left+right+up+down, screw boundary config
// left+right+up+down, periodic bndry
int sjsum(int ad, rgbtype *d)
{
    int x = ad&LMASKX;
    int ad0 = ad&LMASKY;
    int x1 = (x-1)&LMASKX;
    int x2 = (x+1)&LMASKX;

    int res = 
	+d[ad0|x1] //left
	+d[ad0|x2] //right
	+d[(ad+L)&LMASK] //down
	+d[(ad-L)&LMASK];//up

    return res;
}


// note that (s1-s2)^2 = s1^2 + s2^2 - 2 s1 s2
//  and \sum s^2 remains constant
//  therefore energy= const -2 \sum_(all neighbor pixel pair si, sj) si*sj
//                  = const -2 \sum_(all pixel si)  si * (left+right+up+down)  /2 (/2 to cancel double count)
//   actually use energy/2 = \sum_ij -si*sj for MC
int getene1(int ad)
{
    int sjr = sjsum(ad, colr);
    int sjg = sjsum(ad, colg);
    int sjb = sjsum(ad, colb);

    int r= colr[ad];
    int g= colg[ad];
    int b= colb[ad];

    return -r*sjr -g*sjg - b*sjb;
}

// total energy
//  with x2 factor 
long long calc_ene()
{
    int i;
    long long res=0;

    for(i=0;i<L2;i++)
    {
	int r= colr[i];
	int g= colg[i];
	int b= colb[i];
	res += 4 * (r*r + g*g + b*b) + getene1(i);
    }
    return res;
}

#ifdef OCMAX
// calculate penalty
long long calc_penalty()
{
    int i;
    long long res=0;
    for(i=0;i<L2;i++)
    {
	res += (cn[i]-1)*cn[i]/2;
    }
    return res;
}
#endif

// standard Metropolis step, attempt swapping pixel

void sweep_r(int s)
{
    int r1,g1,b1;
    rgbtype *val = colr;

    for(b1=0;b1<VMAX;b1++){
	for(g1=0;g1<VMAX;g1++){
	    int c0 = rgb2idx(0,g1,b1);
	    for(r1=s;r1+1<VMAX;r1+=2){
		int c1 = c0 | r1;

		//int c1 = rgb2idx(r1,g1,b1);
		// red +1
		int c2 = c1 + 1;
		int val0 = r1;

		int ad1 = cidx[c1];
		int ad2 = cidx[c2];
#ifdef OCMAX
		if (ad1==-1 || ad2==-1) continue;
#endif

		// Calc Address
		int x1 = ad1&LMASKX;
		int x2 = ad2&LMASKX;
		int y1 = ad1&LMASKY;
		int y2 = ad2&LMASKY;

		int x1L = (x1-1)&LMASKX;
		int x1R = (x1+1)&LMASKX;
		int x2L = (x2-1)&LMASKX;
		int x2R = (x2+1)&LMASKX;

		int ad1L= y1 | x1L;
		int ad1R=  y1 | x1R;
		int ad1U= (ad1-L)&LMASK;
		int ad1D= (ad1+L)&LMASK;
		int sj1 = val[ad1L] + val[ad1R] + val[ad1U] + val[ad1D];
		int sj2 = 
		    +val[y2 | x2L] 
		    +val[y2 | x2R] 
		    +val[(ad2-L)&LMASK] 
		    +val[(ad2+L)&LMASK];

		// undo double count
		if (ad2 == ad1L || ad2 == ad1R || ad2 == ad1D || ad2 == ad1U)
		{
		    sj1 -= val0+1;
		    sj2 -= val0;
		}

		// energy increase by swapping
		int de = sj2-sj1; 
		//*(r1-r2) + (sj1g-sj2g)*(g1-g2) + (sj1b-sj2b)*(b1-b2);

		bool accp=true;
		if (de>=0)
		{
		    if (rand() > thtab[de]) accp=false;
		}

		if (accp)
		{
		    //swap
		    val[ad1]=val0+1;
		    val[ad2]=val0;
		    ene += de;
		    removecol(c1, ad1);
		    removecol(c2, ad2);
		    addcol(c1,ad2);
		    addcol(c2,ad1);
		}

	    }}}
}

void sweep_g(int s)
{
    int r1,g1,b1;
    rgbtype *val = colg;

    for(b1=0;b1<VMAX;b1++){
	for(g1=s;g1+1<VMAX;g1+=2){
	    int c0 = rgb2idx(0,g1,b1);
	    for(r1=0;r1<VMAX;r1++){
		int c1 = c0 | r1;

		// green +1
		int c2 = c1 + VMAX;
		int val0 = g1;

		int ad1 = cidx[c1];
		int ad2 = cidx[c2];
#ifdef OCMAX
		if (ad1==-1 || ad2==-1) continue;
#endif

		// Calc Address
		int x1 = ad1&LMASKX;
		int x2 = ad2&LMASKX;
		int y1 = ad1&LMASKY;
		int y2 = ad2&LMASKY;

		int x1L = (x1-1)&LMASKX;
		int x1R = (x1+1)&LMASKX;
		int x2L = (x2-1)&LMASKX;
		int x2R = (x2+1)&LMASKX;

		int ad1L= y1 | x1L;
		int ad1R=  y1 | x1R;
		int ad1U= (ad1-L)&LMASK;
		int ad1D= (ad1+L)&LMASK;
		int sj1 = val[ad1L] + val[ad1R] + val[ad1U] + val[ad1D];
		int sj2 = 
		    +val[y2 | x2L] 
		    +val[y2 | x2R] 
		    +val[(ad2-L)&LMASK] 
		    +val[(ad2+L)&LMASK];

		// undo double count
		if (ad2 == ad1L || ad2 == ad1R || ad2 == ad1D || ad2 == ad1U)
		{
		    sj1 -= val0+1;
		    sj2 -= val0;
		}

		// energy increase by swapping
		int de = sj2-sj1; 
		//*(r1-r2) + (sj1g-sj2g)*(g1-g2) + (sj1b-sj2b)*(b1-b2);

		bool accp=true;
		if (de>=0)
		{
		    if (rand() > thtab[de]) accp=false;
		}

		if (accp)
		{
		    //swap
		    val[ad1]=val0+1;
		    val[ad2]=val0;
		    ene += de;
		    removecol(c1, ad1);
		    removecol(c2, ad2);
		    addcol(c1,ad2);
		    addcol(c2,ad1);
		}

	    }}}
}

void sweep_b(int s)
{
    int r1,g1,b1;
    rgbtype *val = colb;

    for(b1=s;b1+1<VMAX;b1+=2){
	for(g1=0;g1<VMAX;g1++){
	    int c0 = rgb2idx(0,g1,b1);
	    for(r1=0;r1<VMAX;r1++){
		int c1 = c0 | r1;

		// blue +1
		int c2 = c1 + VMAX*VMAX;
		int val0 = b1;

		int ad1 = cidx[c1];
		int ad2 = cidx[c2];
#ifdef OCMAX
		if (ad1==-1 || ad2==-1) continue;
#endif

		// Calc Address
		int x1 = ad1&LMASKX;
		int x2 = ad2&LMASKX;
		int y1 = ad1&LMASKY;
		int y2 = ad2&LMASKY;

		int x1L = (x1-1)&LMASKX;
		int x1R = (x1+1)&LMASKX;
		int x2L = (x2-1)&LMASKX;
		int x2R = (x2+1)&LMASKX;

		int ad1L= y1 | x1L;
		int ad1R=  y1 | x1R;
		int ad1U= (ad1-L)&LMASK;
		int ad1D= (ad1+L)&LMASK;
		int sj1 = val[ad1L] + val[ad1R] + val[ad1U] + val[ad1D];
		int sj2 = 
		    +val[y2 | x2L] 
		    +val[y2 | x2R] 
		    +val[(ad2-L)&LMASK] 
		    +val[(ad2+L)&LMASK];

		// undo double count
		if (ad2 == ad1L || ad2 == ad1R || ad2 == ad1D || ad2 == ad1U)
		{
		    sj1 -= val0+1;
		    sj2 -= val0;
		}

		// energy increase by swapping
		int de = sj2-sj1; 
		//*(r1-r2) + (sj1g-sj2g)*(g1-g2) + (sj1b-sj2b)*(b1-b2);

		bool accp=true;
		if (de>=0)
		{
		    if (rand() > thtab[de]) accp=false;
		}

		if (accp)
		{
		    //swap
		    val[ad1]=val0+1;
		    val[ad2]=val0;
		    ene += de;
		    removecol(c1, ad1);
		    removecol(c2, ad2);
		    addcol(c1,ad2);
		    addcol(c2,ad1);
		}

	    }}}
}

#ifdef OCMAX
// increase/decrease one of RGB by 1
void domc1(rgbtype *val, int dv, int dc,
	int ad0, int adL, int adR, int adU, int adD)
{
    int v0= val[ad0];
    if(v0+dv==-1 || v0+dv==VMAX)return;

    int sj = val[adL] + val[adR] + val[adU] + val[adD];
    //int de1 = -sj*dv +2*(2*v0*dv + dv*dv);
    int de = dv*(-sj +4*v0) +2;

    int cc1 = rgb2idx(colr[ad0], colg[ad0], colb[ad0]);
    int cc2 = cc1 + dc*dv;

    int dpenalty = -(cn[cc1]-1) + (cn[cc2]);
    int deall= de+dpenalty*oc;

    bool accp=true;
    if (deall>=0)
    {
	if (rand() > thtab[deall]) accp=false;
    }

    if (accp)
    {
	//change col
	val[ad0]=v0+dv;
	ene += de;
	penalty += dpenalty;

	removecol(cc1, ad0);
	addcol(cc2,ad0);

    }
}

void mc1sweep(rgbtype *val, int dc)
{
    int y, x;

    for(y=0;y<L;y++)
    {
	int ad0=y*L;
	int adU=(ad0-L)&LMASK;
	int adD=(ad0+L)&LMASK;

	// left edge
	x=0;
	int adL = ad0+L-1;
	int adR = ad0+1;
	int dv = ((rand()>>8)&1)*2-1;
	domc1(val, dv, dc, ad0, adL, adR, adU, adD);

	adL=ad0;
	ad0++;
	adR++;
	adU++;
	adD++;

	// mid section
	for(x=1;x<L-1;x++)
	{
	    dv = ((rand()>>8)&1)*2-1;
	    domc1(val, dv, dc, ad0, adL, adR, adU, adD);
	    adL++;
	    ad0++;
	    adR++;
	    adU++;
	    adD++;
	}//x

	//right edge
	adR = y*L;
	dv = ((rand()>>8)&1)*2-1;
	domc1(val, dv, dc, ad0, adL, adR, adU, adD);
    }//y
}
#endif

///////////////////////////////////////
// observation routine

// energy-corr
int nE;
#define DMAX 100000
double edat[DMAX];

void storee()
{
    edat[nE++] = ene * 1.0/L2;
}

void ecor(double *e1, double *cv, double *cor)
{
    double e1s=0.0;
    double e2s=0.0;
    double ecs=0.0;
    int i, s;

    for(i=0;i<nE;i++)
    {
	double e= edat[i];
	e1s+=e;
	e2s+=e*e;
    }
    e1s/=nE;
    e2s/=nE;
    e2s -= e1s*e1s;

    for(s=1;s<nE/10;s++)
    {
	double sum=0;
	for(i=0;i<nE-s;i++)
	    sum += (edat[i]-e1s)*(edat[i+s]-e1s);
	sum /= (nE-s);
	sum /= e2s;

	ecs +=sum;
	if (sum<0.2)break;
    }

    *e1 = e1s;
    *cv = e2s;
    *cor = ecs;
}

int  main(int argc, char **argv)
{

    int i;

    //init
    for(i=0;i<L2;i++)
    {
	int r=i &CBMASK;
	int g=(i>>CBIT) &CBMASK;
	int b=(i>>(CBIT*2)) &CBMASK;
	colr[i] = r;
	colg[i] = g;
	colb[i] = b;
    }

    if(argc==2) load_ppm(argv[1]);

    init_cidx();

    ene = calc_ene()/2;
    ebest=ene;

#ifdef OCMAX
    penalty = calc_penalty();
    printf("#Energy=%lld Penalty=%lld\n",ene, penalty);
    oc = 1;
#else
    penalty = 0;
    printf("#Color Unique check Ok, energy=%lld\n",ene);
#endif
    saveit();


    int t,s;

    FILE *fp=fopen("elog","w");


    //Begin
    for(t=0;t<TDIV;t++)
    {
	double beta = BMIN * exp( log(BMAX/BMIN)*t/TDIV );
	init_th(beta, RAND_MAX);
	nE=0;

	for(s=0;s<MAXS;s++)
	{
	    sweep_r(0);
	    sweep_g(0);
	    sweep_b(0);
	    sweep_r(1);
	    sweep_g(1);
	    sweep_b(1);
#ifdef OCMAX
	    mc1sweep(colr, 1);
	    mc1sweep(colg, VMAX);
	    mc1sweep(colb, VMAX*VMAX);
#endif
	    if(s>MAXS/10)storee();
	    if (penalty==0 && ebest>ene)
	    {
		ebest=ene;
		saveit();
	    }
	}
	double e1,cv,cr;
	ecor(&e1, &cv, &cr);

	fprintf(fp, "%e %f %e %f %lld %lld\n",beta,e1,cv,cr,ene, penalty);
	fflush(fp);
	if(t%30==0)
	{
	    //dump(2, colr,colg,colb);
	    if(penalty==0)dump(1, colrbest,colgbest,colbbest);
	}
    }
    fprintf(fp, "\n\n");


    fclose(fp);
}