fitting correlation maximum with LMA

src/main/java/Correlations2dLMA.java

+import java.util.ArrayList;
+import java.util.HashMap;
+
+import Jama.Matrix;
+
+/**
+ **
+ ** Correlation2dLMA - Fit multi - baseline correaltion pairs to the model
+ **
+ ** Copyright (C) 2018 Elphel, Inc.
+ **
+ ** -----------------------------------------------------------------------------**
+ **
+ **  Correlation2dLMA.java is free software: you can redistribute it and/or modify
+ **  it under the terms of the GNU General Public License as published by
+ **  the Free Software Foundation, either version 3 of the License, or
+ **  (at your option) any later version.
+ **
+ **  This program is distributed in the hope that it will be useful,
+ **  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ **  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ **  GNU General Public License for more details.
+ **
+ **  You should have received a copy of the GNU General Public License
+ **  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ ** -----------------------------------------------------------------------------**
+ **
+ */
+
+
+/*
+ * Fitting parabola for multiple grids
+ * Difference between ortho and diagonals - just point coordinates and extra overall scale (weight account for number averaged)
+ * Each group of compatible is already averaged, so each group has a single individual variable - scale.
+ *
+ * Parabolas (1 for each group) are Ag * (1 - ((x-x0)/Wx) ^ 2 - (y/Wy)^2), where As is a per-group scale
+ * Wy = Wm * scale +Wyd
+ * Wx = Wm * scale +Wyd + Wxy
+ *
+ * Wm is a correlation measurement parameter, it does not depend on x/y and on particular pair, it depends on the LPF, so the
+ * total contribution is proportional to the baseline reduction (scale)
+ *
+ * Wyd is widening caused the image itself - probably noise and other factors of poor correlation contrast. When multiple
+ * orthogonal directions are combined it influences equally all directions (x,y) so Wx includes that term also
+ *
+ * Wxy widens maximum in disparity direction, it is caused by multiple overlapping maximums for different disparities and for
+ * strong enough matches can indicate miz of disparities in the same tile
+ *
+ * Fitting of a single scale groups (1 or 2) has to have Wm constant.
+ *
+ *
+
+ *
+ */
+
+public class Correlations2dLMA {
+	final static int X0_INDEX =  0;
+	final static int WM_INDEX =  1; // may be frozen
+	final static int WYD_INDEX = 2;
+	final static int WXY_INDEX = 3;
+	final static int AG_INDEX =  4; // and above
+	final static String [] PAR_NAMES = {"X0","HALF-WIDTH","WIDTH-EXTRA","WIDTH_X/Y","AMPLITUDE"};
+	double  [] all_pars;
+	boolean [] par_mask;
+	double  [] vector;
+	double  [] scales = {1.0, 2.0, 4.0};
+	ArrayList<Sample> samples = new ArrayList<Sample>();
+	HashMap<Integer,Integer> groups = new HashMap<Integer,Integer>();
+	double [] weights; // normalized so sum is 1.0 for all - samples and extra regularization terms
+	double    pure_weight; // weight of samples only
+	double [] values;
+	// next values are only updated after success
+	double []   last_rms =        null; // {rms, rms_pure}, matching this.vector
+	double []   good_or_bad_rms = null; // just for diagnostics, to read last (failed) rms
+	double []   initial_rms =     null; // {rms, rms_pure}, first-calcualted rms
+	double []   last_ymfx =       null;
+	double [][] last_jt =         null;
+
+
+	public class Sample{
+		double x;      // x coordinate on the common scale (corresponding to the largest baseline), along the disparity axis
+		double y;      // y coordinate (0 - disparity axis)
+		double v;      // correlation value at that point
+		double w;
+		int    si;     // baseline scale index
+		int    gi;     // group index
+		Sample (
+				double x,      // x coordinate on the common scale (corresponding to the largest baseline), along the disparity axis
+				double y,      // y coordinate (0 - disparity axis)
+				double v,      // correlation value at that point
+				double w,
+				int    si,     // baseline scale index
+				int    gi)     // baseline scale index
+			{
+				this.x =  x;
+				this.y =  y;
+				this.v =  v;
+				this.w =  w;
+				this.si = si;
+				this.gi = gi;
+			}
+	}
+
+	public void printParams() {
+		for (int np = 0; np < all_pars.length; np++) {
+			System.out.println(String.format("%2d%1s %22s %f",
+					np,
+					(par_mask[np]?"+":" "),
+					((np < AG_INDEX)?PAR_NAMES[np]:(PAR_NAMES[AG_INDEX]+(np-AG_INDEX))),
+					all_pars[np]));
+		}
+	}
+
+	public void printInputData(){
+		int ns =0;
+		for (Sample s:samples){
+			System.out.println(String.format("%3d: x=%8.4f y=%8.4f v=%9.6f w=%9.7f si=%1d gi=%1d", ns++, s.x, s.y, s.v, s.w, s.si, s.gi ));
+		}
+	}
+	public double [] getRMS() {
+		return last_rms;
+	}
+	public double [] getGoodOrBadRMS() {
+		return good_or_bad_rms;
+	}
+
+
+
+	public Correlations2dLMA (
+			double [] scales // null - use default table
+			) {
+		if (scales != null)	this.scales = scales.clone();
+	}
+
+	public void addSample(
+			double x,      // x coordinate on the common scale (corresponding to the largest baseline), along the disparity axis
+			double y,      // y coordinate (0 - disparity axis)
+			double v,      // correlation value at that point
+			double w,
+			int    si,     // baseline scale index
+			int    gi){     // baseline scale index
+		samples.add(new Sample(x,y,v,w,si,gi));
+	}
+
+// TODO: add auto x0, half-width?
+// should be called ater all samples are entered (to list groups)
+	public void initVector(
+			boolean adjust_wm,
+			boolean adjust_wy,
+			boolean adjust_wxy,
+			boolean adjust_Ag,
+			double  x0,
+			double  half_width,
+			double  cost_wm,     // cost of non-zero this.all_pars[WM_INDEX]
+			double  cost_wxy     // cost of non-zero this.all_pars[WXY_INDEX]
+			) {
+		for (Sample s:samples) if (!groups.containsKey(s.gi)) groups.put(s.gi,groups.size());
+		int num_groups = groups.size();
+		this.all_pars =        new double[AG_INDEX + num_groups];
+		this.all_pars[X0_INDEX] =  x0;
+		this.all_pars[WM_INDEX] =  half_width;
+		this.all_pars[WYD_INDEX] = 0.0;
+		this.all_pars[WXY_INDEX] = 0.0;
+		this.par_mask = new boolean[AG_INDEX + num_groups];
+		this.par_mask[X0_INDEX] = true;
+		this.par_mask[WM_INDEX] = adjust_wm;
+		this.par_mask[WYD_INDEX] = adjust_wy;
+		this.par_mask[WXY_INDEX] = adjust_wxy;
+		for (int i = 0; i <num_groups; i++) {
+			this.par_mask[AG_INDEX + i] = adjust_Ag;
+		}
+		setWeightsValues(half_width, cost_wm, cost_wxy);
+		int imx = 0;
+		int num_samples = samples.size();
+		for (int i = 1; i<num_samples; i++) if (values[i] > values[imx]) imx = i;
+		for (int i = 0; i < num_groups; i++) this.all_pars[AG_INDEX + i] = values[imx];
+		toVector();
+	}
+
+	public void setWeightsValues(
+			double half_width,   // expected width
+			double  cost_wm,     // cost of non-zero this.all_pars[WYD_INDEX]
+			double  cost_wxy) {  // cost of non-zero this.all_pars[WXY_INDEX]
+		int np = samples.size();
+		weights = new double [np+2];
+		values =  new double [np+2];
+		weights[np] =   cost_wm;
+		weights[np+1] = cost_wxy;
+		values[np] =    half_width;
+		values[np+1] =  0.0;
+		double sw = 0;
+		for (int i = 0; i < np; i++) {
+			weights[i] = samples.get(i).w;
+			values[i] =  samples.get(i).v;
+			if (Double.isNaN(values[i]) || Double.isNaN(weights[i])) {
+				weights[i] = 0.0;
+				values[i] = 0.0;
+			}
+			sw += weights[i];
+		}
+		pure_weight = sw;
+		sw += weights[np] + weights[np];
+		if (sw != 0.0) {
+			sw = 1.0/sw;
+			for (int i = 0; i < weights.length; i++) weights[i] *= sw;
+		}
+		pure_weight *= sw;
+	}
+
+	public void toVector() {
+		int np = 0;
+		for (int i = 0; i < par_mask.length; i++) if (par_mask[i]) np++;
+		vector = new double[np];
+		np = 0;
+		for (int i = 0; i < par_mask.length; i++) if (par_mask[i]) vector[np++] = all_pars[i];
+	}
+
+	public void updateFromVector() {
+		int np = 0;
+		for (int i = 0; i < par_mask.length; i++) if (par_mask[i]) all_pars[i] = vector[np++];
+	}
+
+	public double [] fromVector(double [] vector) { // mix fixed and variable parameters
+		double [] ap = all_pars.clone();
+		int np = 0;
+		for (int i = 0; i < par_mask.length; i++) if (par_mask[i]) ap[i] = vector[np++];
+		return ap;
+	}
+
+    public void debugJt(
+    		double      delta,
+    		double []   vector) {
+    	int num_points = this.values.length;
+    	int num_pars = vector.length;
+    	double [] max_diff = new double [num_pars];
+
+//    	delta = 0.001;
+
+    	double [][] jt =       new double [num_pars][num_points];
+    	double [][] jt_delta = new double [num_pars][num_points];
+    	double [] fx = getFxJt( vector,jt);
+    	getFxJt(delta, vector,jt_delta);
+    	System.out.println("Test of jt-jt_delta difference,  delta = "+delta);
+    	System.out.print(String.format("%3s: %10s ", "#", "fx"));
+    	for (int anp = 0; anp< all_pars.length; anp++) if(par_mask[anp]){
+        	System.out.print(String.format("%17s ", ((anp < AG_INDEX)?PAR_NAMES[anp]:(PAR_NAMES[AG_INDEX]+(anp-AG_INDEX)))));
+    	}
+    	System.out.println();
+
+    	for (int i = 0; i < num_points; i++) {
+        	System.out.print(String.format("%3d: %10.7f ", i, fx[i]));
+        	for (int np = 0; np < num_pars; np++) {
+            	System.out.print(String.format("%8.5f %8.5f ", jt_delta[np][i], 1000*(jt[np][i] - jt_delta[np][i])));
+            	double adiff = Math.abs(jt[np][i] - jt_delta[np][i]);
+            	if (adiff > max_diff[np]) {
+            		max_diff[np] = adiff;
+            	}
+        	}
+        	System.out.println();
+    	}
+    	System.out.print(String.format("%15s ", "Maximal diff"));
+    	for (int np = 0; np < num_pars; np++) {
+        	System.out.print(String.format("%8s %8.5f ", "1000×",  1000*max_diff[np]));
+    	}
+    	System.out.println();
+
+    }
+
+
+	public double [] getFxJt(
+			double      delta, // for testing derivatives: calculates as delta-F/delta_x
+			double []   vector,
+			double [][] jt) { // should be either [vector.length][samples.size()] or null - then only fx is calculated
+		double [] fx0=getFxJt(vector,null);
+		for (int np = 0; np < vector.length; np++) {
+			double [] vector1 = vector.clone();
+			vector1[np]+= delta;
+			double [] fxp=getFxJt(vector1,null);
+			vector1 = vector.clone();
+			vector1[np]-= delta;
+			double [] fxm=getFxJt(vector1,null);
+
+			jt[np] = new double [fxp.length];
+			for (int i = 0; i < fxp.length; i++) {
+				jt[np][i] = (fxp[i] - fxm[i])/delta/2;
+			}
+		}
+		return fx0;
+	}
+
+
+
+
+
+	public double [] getFxJt(
+			double []   vector,
+			double [][] jt) { // should be either [vector.length][samples.size()] or null - then only fx is calculated
+		double [] av = fromVector(vector);
+		int num_samples = samples.size();
+		double [] fx= new double [num_samples + 2];
+		int num_groups = groups.size();
+		for (int ns = 0; ns < num_samples; ns++) {
+			Sample s = samples.get(ns);
+			int grp = groups.get(s.gi);
+			double Wy = av[WM_INDEX] * scales[s.si] + av[WYD_INDEX];
+			double Wx = Wy + av[WXY_INDEX];
+			double dx = s.x - av[X0_INDEX];
+			double Ag = av[AG_INDEX+grp];
+			double dxw = dx/Wx;
+			double dyw = s.y/Wy;
+			double d = (1.0 - dxw*dxw - dyw*dyw);
+			fx[ns] = d * Ag;
+			int np = 0;
+			if (jt != null) {
+				if (par_mask[X0_INDEX])  jt[np++][ns] = Ag * 2 * dxw/Wx; // d/dx0
+				double dfdWx = Ag * 2 * dxw * dxw / Wx;
+				double dfdWy = Ag * 2 * dyw * dyw / Wy;
+				if (par_mask[WM_INDEX])  jt[np++][ns] =  scales[s.si] * (dfdWx + dfdWy); // d/dWm
+				if (par_mask[WYD_INDEX]) jt[np++][ns] = (dfdWx + dfdWy); // d/dWyd
+				if (par_mask[WXY_INDEX]) jt[np++][ns] =  dfdWx;          // d/dWxy
+				for (int i = 0; i < num_groups; i++) {
+					if (par_mask[AG_INDEX + i]) jt[np++][ns] =  (i == grp) ? d : 0.0;    // d/dWAg
+				}
+			}
+		}
+		// add 2 extra samples - damping cost_wy, cost_wxy
+		fx[num_samples] =     av[WM_INDEX];
+		fx[num_samples + 1] = av[WXY_INDEX];
+		// and derivatives
+		if (jt != null) {
+			int np = 0;
+			for (int i = 0; i < par_mask.length; i++) if (par_mask[i]) {
+				jt[np][num_samples] =     (i == WM_INDEX)? 1.0 : 0.0;
+				jt[np++][num_samples+1] = (i == WXY_INDEX)? 1.0 : 0.0;
+			}
+		}
+		return fx;
+	}
+
+	public double [][] getWJtJlambda(
+			double      lambda,
+			double [][] jt){
+		int num_pars = jt.length;
+		int nup_points = jt[0].length;
+		double [][] wjtjl = new double [num_pars][num_pars];
+		for (int i = 0; i < num_pars; i++) {
+			for (int j = i; j < num_pars; j++) {
+				double d = 0.0;
+				for (int k = 0; k < nup_points; k++) {
+					d += this.weights[k]*jt[i][k]*jt[j][k];
+				}
+				wjtjl[i][j] = d;
+				if (i == j) {
+					wjtjl[i][j] += d * lambda;
+				} else {
+					wjtjl[j][i] = d;
+				}
+			}
+		}
+		return wjtjl;
+	}
+
+	// returns {rms, rms_pure}
+	public double [] getWYmFxRms(
+			double []   fx) { // will be replaced with y-fx
+		int num_samples = samples.size();
+		int num_points = fx.length; // includes 2 extra for regularization
+		double rms = 0, rms_pure = 0;
+		for (int i = 0; i < num_samples; i++) {
+			double d = (values[i] - fx[i]);
+			fx[i] = this.weights[i] * d;
+			rms += fx[i]*d; // sum of weights
+		}
+		rms_pure = Math.sqrt(rms)/this.pure_weight;
+		for (int i = num_samples; i < num_points; i++) {
+			double d = (values[i] - fx[i]);
+			fx[i] = this.weights[i] * d;
+			rms += fx[i]*d; // sum of weights
+		}
+		rms = Math.sqrt(rms);
+		double [] rslt = {rms, rms_pure};
+		return rslt;
+	}
+
+	public double [] getJtWdiff(
+			double []   wdiff,
+			double [][] jt){
+		int num_pars = jt.length;
+		int nup_points = jt[0].length;
+		double [] wjtymfx = new double [num_pars];
+		for (int i = 0; i < num_pars; i++) {
+			double d = 0;
+			for (int j = 0; j < nup_points; j++) d += wdiff[j] + jt[i][j];
+			wjtymfx[i] = d;
+		}
+		return wjtymfx;
+	}
+
+	public boolean runLma(
+			double lambda,           // 0.1
+			double lambda_scale_good,// 0.5
+			double lambda_scale_bad, // 8.0
+			double lambda_max,       // 100
+			double rms_diff,         // 0.001
+			int    num_iter,         // 20
+			int    debug_level)
+	{
+		boolean [] rslt = {false,false};
+		int iter = 0;
+		for (iter = 0; iter < num_iter; iter++) {
+			rslt =  lmaStep(
+					lambda,
+					rms_diff,
+					debug_level);
+			if (debug_level > 1) {
+				System.out.println("LMA step "+iter+": {"+rslt[0]+","+rslt[1]+"} full RMS="+good_or_bad_rms[0]+
+						" ("+initial_rms[0]+"), pure RMS="+good_or_bad_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
+			}
+			if (rslt[1]) {
+				break;
+			}
+			if (rslt[0]) { // good
+				lambda *= lambda_scale_good;
+			} else {
+				lambda *= lambda_scale_bad;
+				if (lambda > lambda_max) {
+					break;
+				}
+			}
+		}
+		if (rslt[0]) { // better, but num tries exceeded
+			if (iter >= num_iter) {
+				if (debug_level > 0) System.out.println("Step "+iter+": Improved, but number of steps exceeded maximal");
+			} else {
+				if (debug_level > 0) System.out.println("Step "+iter+": LMA: Success");
+			}
+
+		} else { // improved over initial ?
+			if (last_rms[0] < initial_rms[0]) {
+				rslt[0] = true;
+				if (debug_level > 0) System.out.println("Step "+iter+": Failed to converge, but result improved over initial");
+			} else {
+				if (debug_level > 0) System.out.println("Step "+iter+": Failed to converge");
+			}
+		}
+		if (debug_level > 0) {
+			System.out.println("LMA: full RMS="+last_rms[0]+" ("+initial_rms[0]+"), pure RMS="+last_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
+//			System.out.println("LMA: full RMS="+good_or_bad_rms[0]+" ("+initial_rms[0]+"), pure RMS="+good_or_bad_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
+		}
+
+		return rslt[0];
+	}
+
+/*
+	double []   last_rms =     null; // {rms, rms_pure}, matching this.vector
+	double []   initial_rms =  null; // {rms, rms_pure}, first-calcualted rms
+
+ */
+
+	// returns {success, done}
+	public boolean [] lmaStep(
+			double lambda,
+			double rms_diff,
+			int debug_level) {
+		int num_points = this.weights.length; // includes 2 extra for regularization
+		int num_pars = vector.length;
+		boolean [] rslt = {false,false};
+		if (this.last_rms == null) { //first time, need to calculate all (vector is valid)
+			this.last_jt = new double [num_pars][num_points];
+			this.last_ymfx = getFxJt(
+					this.vector, // double []   vector,
+					this.last_jt); // double [][] jt) { // should be either [vector.length][samples.size()] or null - then only fx is calculated
+			this.last_rms = getWYmFxRms(this.last_ymfx); // modifies this.last_ymfx
+			this.initial_rms = this.last_rms.clone();
+			this.good_or_bad_rms = this.last_rms.clone();
+			if (debug_level > 3) {
+				debugJt(
+						0.000001, // double      delta,
+						this.vector); // double []   vector);
+			}
+		}
+		Matrix y_minus_fx_weighted = new Matrix(this.last_ymfx, this.last_ymfx.length);
+
+		Matrix wjtjlambda = new Matrix(getWJtJlambda(
+				lambda, // *10, // temporary
+				this.last_jt)); // double [][] jt)
+		if (debug_level>2) {
+			System.out.println("JtJ + lambda*diag(JtJ");
+			wjtjlambda.print(18, 6);
+		}
+		Matrix jtjl_inv = null;
+		try {
+			jtjl_inv = wjtjlambda.inverse(); // check for errors
+		} catch (RuntimeException e) {
+			rslt[1] = true;
+			if (debug_level > 0) {
+				System.out.println("Singular Matrix");
+			}
+
+			return rslt;
+		}
+/*
+				try {
+					this.SYNC_COMMAND.wait();
+				} catch (InterruptedException e) {
+					// TODO Auto-generated catch block
+					e.printStackTrace();
+				}
+
+ */
+		if (debug_level>2) {
+			System.out.println("(JtJ + lambda*diag(JtJ).inv()");
+			jtjl_inv.print(18, 6);
+		}
+
+		Matrix jty = (new Matrix(this.last_jt)).times(y_minus_fx_weighted);
+		if (debug_level>2) {
+			System.out.println("Jt * (y-fx)");
+			jty.print(18, 6);
+		}
+
+		Matrix mdelta = jtjl_inv.times(jty);
+		if (debug_level>2) {
+			System.out.println("mdelta");
+			mdelta.print(18, 6);
+		}
+
+
+		double []  delta = mdelta.getColumnPackedCopy();
+		double [] new_vector = this.vector.clone();
+		for (int i = 0; i < num_pars; i++) new_vector[i]+= delta[i];
+		// being optimistic, modify jt and last_ymfx in place, restore if failed
+		this.last_ymfx = getFxJt(
+				new_vector, // double []   vector,
+				this.last_jt); // double [][] jt) { // should be either [vector.length][samples.size()] or null - then only fx is calculated
+		double [] rms = getWYmFxRms(this.last_ymfx); // modifies this.last_ymfx
+		this.good_or_bad_rms = rms.clone();
+		if (rms[0] < this.last_rms[0]) { // improved
+			rslt[0] = true;
+			rslt[1] = rms[0] >=(this.last_rms[0] * (1.0 - rms_diff));
+			this.last_rms = rms.clone();
+			this.vector = new_vector.clone();
+			if (debug_level > 2) {
+				System.out.print("New vector: ");
+				for (int np = 0; np < vector.length; np++) {
+					System.out.print(this.vector[np]+" ");
+				}
+				System.out.println();
+			}
+
+		} else { // worsened
+			rslt[0] = false;
+			rslt[1] = false; // do not know, caller will decide
+			// restore state
+			this.last_ymfx = getFxJt( // recalculate fx
+					this.vector, // double []   vector,
+					this.last_jt); // double [][] jt) { // should be either [vector.length][samples.size()] or null - then only fx is calculated
+			this.last_rms = getWYmFxRms(this.last_ymfx); // modifies this.last_ymfx
+			if (debug_level > 2) {
+				debugJt(
+						0.000001, // double      delta,
+						this.vector); // double []   vector);
+			}
+
+
+		}
+		return rslt;
+	}
+
+
+//	public double [] getValues() {
+//		double [] values= new double [samples.size()];
+//		for (int i = 0; i < values.length; i++) values[i] = samples.get(i).v;
+//		return values;
+//	}
+
+	/*
+	 *
+		if (debugLevel>-1) {
+			jtj.print(18, 6);
+		}
+		Matrix jtj_inv = jtj.inverse();
+		Matrix jty = jt.times(y_minus_fx_weighted);
+		Matrix mrslt = jtj_inv.times(jty);
+		double []  drslt = mrslt.getColumnPackedCopy();
+	 *
+	 * Fitting parabola for multiple grids
+	 * Difference between ortho and diagonals - just point coordinates and extra overall scale (weight account for number averaged)
+	 * Each group of compatible is already averaged, so each group has a single individual variable - scale.
+	 *
+	 * Parabolas (1 for each group) are Ag * (1 - ((x-x0)/Wx) ^ 2 - (y/Wy)^2), where As is a per-group scale
+	 * Wy = Wm * scale +Wyd
+	 * Wx = Wm * scale +Wyd + Wxy
+	 *
+	 * Wm is a correlation measurement parameter, it does not depend on x/y and on particular pair, it depends on the LPF, so the
+	 * total contribution is proportional to the baseline reduction (scale)
+	 *
+	 * Wyd is widening caused the image itself - probably noise and other factors of poor correlation contrast. When multiple
+	 * orthogonal directions are combined it influences equally all directions (x,y) so Wx includes that term also
+	 *
+	 * Wxy widens maximum in disparity direction, it is caused by multiple overlapping maximums for different disparities and for
+	 * strong enough matches can indicate miz of disparities in the same tile
+	 *
+	 * Fitting of a single scale groups (1 or 2) has to have Wm constant.
+	 *
+	 *
+
+	 *
+	 */
+
+
+
+
+}