NewtonRaphsonGradient2ProcedureTest.java

package gdsc.smlm.fitting.nonlinear.gradient;

import java.util.ArrayList;
import java.util.Arrays;

import org.apache.commons.math3.random.RandomDataGenerator;
import org.apache.commons.math3.random.RandomGenerator;
import org.apache.commons.math3.random.Well19937c;
import org.junit.Assert;
import org.junit.Test;

import gdsc.core.utils.DoubleEquality;
import gdsc.smlm.TestSettings;
import gdsc.smlm.function.Gradient2Function;
import gdsc.smlm.function.ValueProcedure;
import gdsc.smlm.function.gaussian.Gaussian2DFunction;
import gdsc.smlm.function.gaussian.GaussianFunctionFactory;
import gdsc.smlm.function.gaussian.HoltzerAstimatismZModel;
import gdsc.smlm.function.gaussian.erf.ErfGaussian2DFunction;
import gdsc.smlm.function.gaussian.erf.SingleAstigmatismErfGaussian2DFunction;
import gdsc.smlm.function.gaussian.erf.SingleFreeCircularErfGaussian2DFunction;

/**
 * Contains speed tests for the methods for calculating the Hessian and gradient vector
 * for use in the LVM algorithm.
 */
public class NewtonRaphsonGradient2ProcedureTest
{
	boolean speedTests = true;
	DoubleEquality eq = new DoubleEquality(6, 1e-16);

	int MAX_ITER = 20000;
	int blockWidth = 10;
	double Background = 0.5;
	double Signal = 100;
	double Angle = Math.PI;
	double Xpos = 5;
	double Ypos = 5;
	double Xwidth = 1.2;
	double Ywidth = 1.2;

	RandomDataGenerator rdg;

	@Test
	public void gradientProcedureFactoryCreatesOptimisedProcedures()
	{
		double[] y = new double[0];
		Assert.assertEquals(
				NewtonRaphsonGradient2ProcedureFactory.createUnrolled(y, new DummyGradientFunction(4)).getClass(),
				NewtonRaphsonGradient2Procedure4.class);
		Assert.assertEquals(
				NewtonRaphsonGradient2ProcedureFactory.createUnrolled(y, new DummyGradientFunction(5)).getClass(),
				NewtonRaphsonGradient2Procedure5.class);
		Assert.assertEquals(
				NewtonRaphsonGradient2ProcedureFactory.createUnrolled(y, new DummyGradientFunction(6)).getClass(),
				NewtonRaphsonGradient2Procedure6.class);
	}

	@Test
	public void gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator()
	{
		gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(4);
		gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(5);
		gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(6);
		gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(11);
		gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(21);
	}

	@Test
	public void gradientProcedureIsNotSlowerThanGradientCalculator()
	{
		// Note: The procedure does not have a lot of work within loops. It is only a single loop
		// so unrolling does not produce performance gains. The JVM can optimise this. 

		gradientProcedureIsNotSlowerThanGradientCalculator(4);
		gradientProcedureIsNotSlowerThanGradientCalculator(5);
		gradientProcedureIsNotSlowerThanGradientCalculator(6);
		gradientProcedureIsNotSlowerThanGradientCalculator(11);
		gradientProcedureIsNotSlowerThanGradientCalculator(21);
	}

	private void gradientProcedureComputesSameLogLikelihoodAsMLEGradientCalculator(int nparams)
	{
		int iter = 10;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		ArrayList<double[]> paramsList = new ArrayList<double[]>(iter);
		ArrayList<double[]> yList = new ArrayList<double[]>(iter);

		createFakeData(nparams, iter, paramsList, yList);
		FakeGradientFunction func = new FakeGradientFunction(blockWidth, nparams);

		MLEGradientCalculator calc = (MLEGradientCalculator) GradientCalculatorFactory.newCalculator(nparams, true);

		String name = String.format("[%d]", nparams);

		for (int i = 0; i < paramsList.size(); i++)
		{
			NewtonRaphsonGradient2Procedure p = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i),
					func);
			double s = p.computeLogLikelihood(paramsList.get(i));
			double s2 = calc.logLikelihood(yList.get(i), paramsList.get(i), func);
			// Virtually the same ...
			Assert.assertEquals(name + " Result: Not same @ " + i, s, s2, 1e-10);
		}
	}

	@Test
	public void gradientProcedureComputesSameWithPrecomputed()
	{
		int iter = 10;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		ErfGaussian2DFunction f1 = (ErfGaussian2DFunction) GaussianFunctionFactory.create2D(1, 10, 10,
				GaussianFunctionFactory.FIT_ERF_FREE_CIRCLE, null);
		ErfGaussian2DFunction f2 = (ErfGaussian2DFunction) GaussianFunctionFactory.create2D(2, 10, 10,
				GaussianFunctionFactory.FIT_ERF_FREE_CIRCLE, null);

		double[] a1 = new double[7];
		double[] a2 = new double[13];

		final double[] x = new double[f1.size()];
		final double[] b = new double[f1.size()];

		for (int i = 0; i < iter; i++)
		{
			a2[Gaussian2DFunction.BACKGROUND] = rdg.nextUniform(0.1, 0.3);
			a2[Gaussian2DFunction.SIGNAL] = rdg.nextUniform(100, 300);
			a2[Gaussian2DFunction.X_POSITION] = rdg.nextUniform(3, 5);
			a2[Gaussian2DFunction.Y_POSITION] = rdg.nextUniform(3, 5);
			a2[Gaussian2DFunction.X_SD] = rdg.nextUniform(1, 1.3);
			a2[Gaussian2DFunction.Y_SD] = rdg.nextUniform(1, 1.3);
			a2[6 + Gaussian2DFunction.SIGNAL] = rdg.nextUniform(100, 300);
			a2[6 + Gaussian2DFunction.X_POSITION] = rdg.nextUniform(5, 7);
			a2[6 + Gaussian2DFunction.Y_POSITION] = rdg.nextUniform(5, 7);
			a2[6 + Gaussian2DFunction.X_SD] = rdg.nextUniform(1, 1.3);
			a2[6 + Gaussian2DFunction.Y_SD] = rdg.nextUniform(1, 1.3);

			// Simulate Poisson data
			f2.initialise0(a2);
			f1.forEach(new ValueProcedure()
			{
				int k = 0;

				public void execute(double value)
				{
					x[k++] = (value > 0) ? rdg.nextPoisson(value) : 0;
				}
			});

			// Precompute peak 2 (no background)
			a1[Gaussian2DFunction.BACKGROUND] = 0;
			for (int j = 1; j < 7; j++)
				a1[j] = a2[6 + j];
			f1.initialise0(a1);
			f1.forEach(new ValueProcedure()
			{
				int k = 0;

				public void execute(double value)
				{
					b[k++] = value;
				}
			});

			// Reset to peak 1
			for (int j = 0; j < 7; j++)
				a1[j] = a2[j];

			// Compute peak 1+2
			NewtonRaphsonGradient2Procedure p12 = NewtonRaphsonGradient2ProcedureFactory.create(x, f2);
			double[] up1 = Arrays.copyOf(p12.computeUpdate(a2), f1.getNumberOfGradients());

			// Compute peak 1+(precomputed 2)
			NewtonRaphsonGradient2Procedure p1b2 = NewtonRaphsonGradient2ProcedureFactory.create(x, b, f1);
			double[] up2 = p1b2.computeUpdate(a1);

			Assert.assertArrayEquals(" Result: Not same @ " + i, p12.u, p1b2.u, 1e-10);
			Assert.assertArrayEquals(" Update: Not same @ " + i, up1, up2, 1e-10);

			double[] v1 = p12.computeValue(a2);
			double[] v2 = p1b2.computeValue(a1);
			Assert.assertArrayEquals(" Value: Not same @ " + i, v1, v2, 1e-10);

			double[] d1 = Arrays.copyOf(p12.computeFirstDerivative(a2), f1.getNumberOfGradients());
			double[] d2 = p1b2.computeFirstDerivative(a1);
			Assert.assertArrayEquals(" 1st derivative: Not same @ " + i, d1, d2, 1e-10);
		}
	}

	private abstract class Timer
	{
		private int loops;
		int min = 5;

		Timer()
		{
		}

		Timer(int min)
		{
			this.min = min;
		}

		long getTime()
		{
			// Run till stable timing
			long t1 = time();
			for (int i = 0; i < 10; i++)
			{
				long t2 = t1;
				t1 = time();
				if (loops >= min && DoubleEquality.relativeError(t1, t2) < 0.02) // 2% difference
					break;
			}
			return t1;
		}

		long time()
		{
			loops++;
			long t = System.nanoTime();
			run();
			t = System.nanoTime() - t;
			//System.out.printf("[%d] Time = %d\n", loops, t);
			return t;
		}

		abstract void run();
	}

	private void gradientProcedureIsNotSlowerThanGradientCalculator(final int nparams)
	{
		org.junit.Assume.assumeTrue(speedTests || TestSettings.RUN_SPEED_TESTS);

		final int iter = 1000;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		final ArrayList<double[]> paramsList = new ArrayList<double[]>(iter);
		final ArrayList<double[]> yList = new ArrayList<double[]>(iter);

		createFakeData(nparams, iter, paramsList, yList);
		final FakeGradientFunction func = new FakeGradientFunction(blockWidth, nparams);

		MLEGradientCalculator calc = (MLEGradientCalculator) GradientCalculatorFactory.newCalculator(nparams, true);

		for (int i = 0; i < paramsList.size(); i++)
			calc.logLikelihood(yList.get(i), paramsList.get(i), func);

		for (int i = 0; i < paramsList.size(); i++)
		{
			NewtonRaphsonGradient2Procedure p = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i),
					func);
			p.computeLogLikelihood(paramsList.get(i));
		}

		// Realistic loops for an optimisation
		final int loops = 15;

		// Run till stable timing
		Timer t1 = new Timer()
		{
			@Override
			void run()
			{
				for (int i = 0, k = 0; i < iter; i++)
				{
					MLEGradientCalculator calc = (MLEGradientCalculator) GradientCalculatorFactory
							.newCalculator(nparams, true);
					for (int j = loops; j-- > 0;)
						calc.logLikelihood(yList.get(i), paramsList.get(k++ % iter), func);
				}
			}
		};
		long time1 = t1.getTime();

		Timer t2 = new Timer(t1.loops)
		{
			@Override
			void run()
			{
				for (int i = 0, k = 0; i < iter; i++)
				{
					NewtonRaphsonGradient2Procedure p = NewtonRaphsonGradient2ProcedureFactory
							.createUnrolled(yList.get(i), func);
					for (int j = loops; j-- > 0;)
						p.computeLogLikelihood(paramsList.get(k++ % iter));
				}
			}
		};
		long time2 = t2.getTime();

		log("GradientCalculator = %d : NewtonRaphsonGradient2Procedure %d = %d : %fx\n", time1, nparams, time2,
				(1.0 * time1) / time2);
		if (TestSettings.ASSERT_SPEED_TESTS)
		{
			// Add contingency
			Assert.assertTrue(time2 < time1 * 1.5);
		}
	}

	@Test
	public void gradientProcedureUnrolledComputesSameAsGradientProcedure()
	{
		gradientProcedureUnrolledComputesSameAsGradientProcedure(4);
		gradientProcedureUnrolledComputesSameAsGradientProcedure(5);
		gradientProcedureUnrolledComputesSameAsGradientProcedure(6);
	}

	private void gradientProcedureUnrolledComputesSameAsGradientProcedure(int nparams)
	{
		int iter = 10;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		ArrayList<double[]> paramsList = new ArrayList<double[]>(iter);
		ArrayList<double[]> yList = new ArrayList<double[]>(iter);

		createFakeData(nparams, iter, paramsList, yList);
		FakeGradientFunction func = new FakeGradientFunction(blockWidth, nparams);

		NewtonRaphsonGradient2Procedure p1, p2;
		String name = String.format("[%d]", nparams);
		for (int i = 0; i < paramsList.size(); i++)
		{
			p1 = new NewtonRaphsonGradient2Procedure(yList.get(i), func);
			p2 = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i), func);
			double[] a = paramsList.get(i);

			double ll1 = p1.computeLogLikelihood(a);
			double ll2 = p2.computeLogLikelihood(a);
			Assert.assertEquals(name + " LL: Not same @ " + i, ll1, ll2, 0);

			p1 = new NewtonRaphsonGradient2Procedure(yList.get(i), func);
			p2 = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i), func);
			p1.computeFirstDerivative(a);
			p2.computeFirstDerivative(a);
			Assert.assertArrayEquals(name + " first derivative value: Not same @ " + i, p1.u, p2.u, 0);
			Assert.assertArrayEquals(name + " first derivative: Not same @ " + i, p1.d1, p2.d1, 0);

			p1 = new NewtonRaphsonGradient2Procedure(yList.get(i), func);
			p2 = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i), func);
			p1.computeUpdate(a);
			p2.computeUpdate(a);
			Assert.assertArrayEquals(name + " update value: Not same @ " + i, p1.u, p2.u, 0);
			Assert.assertArrayEquals(name + " update: Not same d1 @ " + i, p1.d1, p2.d1, 0);
			Assert.assertArrayEquals(name + " update: Not same d2 @ " + i, p1.d2, p2.d2, 0);
			Assert.assertArrayEquals(name + " update: Not same update @ " + i, p1.getUpdate(), p2.getUpdate(), 0);
		}
	}

	@Test
	public void gradientProcedureIsFasterUnrolledThanGradientProcedure()
	{
		gradientProcedureLinearIsFasterThanGradientProcedure(4);
		gradientProcedureLinearIsFasterThanGradientProcedure(5);
		gradientProcedureLinearIsFasterThanGradientProcedure(6);
	}

	private void gradientProcedureLinearIsFasterThanGradientProcedure(final int nparams)
	{
		org.junit.Assume.assumeTrue(speedTests || TestSettings.RUN_SPEED_TESTS);

		final int iter = 100;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		final ArrayList<double[]> paramsList = new ArrayList<double[]>(iter);
		final ArrayList<double[]> yList = new ArrayList<double[]>(iter);

		createData(1, iter, paramsList, yList);

		// Remove the timing of the function call by creating a dummy function
		final Gradient2Function func = new FakeGradientFunction(blockWidth, nparams);

		for (int i = 0; i < paramsList.size(); i++)
		{
			NewtonRaphsonGradient2Procedure p1 = new NewtonRaphsonGradient2Procedure(yList.get(i), func);
			p1.computeUpdate(paramsList.get(i));
			p1.computeUpdate(paramsList.get(i));

			NewtonRaphsonGradient2Procedure p2 = NewtonRaphsonGradient2ProcedureFactory.createUnrolled(yList.get(i),
					func);
			p2.computeUpdate(paramsList.get(i));
			p2.computeUpdate(paramsList.get(i));

			// Check they are the same
			Assert.assertArrayEquals("Update " + i, p1.getUpdate(), p2.getUpdate(), 0);
		}

		// Realistic loops for an optimisation
		final int loops = 15;

		// Run till stable timing
		Timer t1 = new Timer()
		{
			@Override
			void run()
			{
				for (int i = 0, k = 0; i < paramsList.size(); i++)
				{
					NewtonRaphsonGradient2Procedure p1 = new NewtonRaphsonGradient2Procedure(yList.get(i), func);
					for (int j = loops; j-- > 0;)
						p1.computeUpdate(paramsList.get(k++ % iter));
				}
			}
		};
		long time1 = t1.getTime();

		Timer t2 = new Timer(t1.loops)
		{
			@Override
			void run()
			{
				for (int i = 0, k = 0; i < paramsList.size(); i++)
				{
					NewtonRaphsonGradient2Procedure p2 = NewtonRaphsonGradient2ProcedureFactory
							.createUnrolled(yList.get(i), func);
					for (int j = loops; j-- > 0;)
						p2.computeUpdate(paramsList.get(k++ % iter));
				}
			}
		};
		long time2 = t2.getTime();

		log("Standard = %d : Unrolled %d = %d : %fx\n", time1, nparams, time2, (1.0 * time1) / time2);
		Assert.assertTrue(time2 < time1 * 1.5);
	}

	@Test
	public void gradientCalculatorComputesGradient()
	{
		gradientCalculatorComputesGradient(new SingleFreeCircularErfGaussian2DFunction(blockWidth, blockWidth));

		// Use a reasonable z-depth function from the Smith, et al (2010) paper (page 377)
		double gamma = 0.389;
		double d = 0.531;
		double Ax = -0.0708;
		double Bx = -0.073;
		double Ay = 0.164;
		double By = 0.0417;
		HoltzerAstimatismZModel zModel = HoltzerAstimatismZModel.create(gamma, d, Ax, Bx, Ay, By);
		gradientCalculatorComputesGradient(new SingleAstigmatismErfGaussian2DFunction(blockWidth, blockWidth, zModel));
	}

	private void gradientCalculatorComputesGradient(ErfGaussian2DFunction func)
	{
		// Check the first and second derivatives
		int nparams = func.getNumberOfGradients();
		int[] indices = func.gradientIndices();

		int iter = 100;
		rdg = new RandomDataGenerator(new Well19937c(30051977));

		ArrayList<double[]> paramsList = new ArrayList<double[]>(iter);
		ArrayList<double[]> yList = new ArrayList<double[]>(iter);

		createData(1, iter, paramsList, yList, true);

		double delta = 1e-3;
		DoubleEquality eq = new DoubleEquality(3, 1e-3);

		for (int i = 0; i < paramsList.size(); i++)
		{
			double[] y = yList.get(i);
			double[] a = paramsList.get(i);
			double[] a2 = a.clone();
			NewtonRaphsonGradient2Procedure p = NewtonRaphsonGradient2ProcedureFactory.create(y, func);
			//double ll = p.computeLogLikelihood(a);
			p.computeUpdate(a);
			double[] d1 = p.d1.clone();
			double[] d2 = p.d2.clone();
			for (int j = 0; j < nparams; j++)
			{
				int k = indices[j];
				double d = (a[k] == 0) ? 1e-3 : a[k] * delta;
				a2[k] = a[k] + d;
				double llh = p.computeLogLikelihood(a2);
				p.computeFirstDerivative(a2);
				double[] d1h = p.d1.clone();
				a2[k] = a[k] - d;
				double lll = p.computeLogLikelihood(a2);
				p.computeFirstDerivative(a2);
				double[] d1l = p.d1.clone();
				a2[k] = a[k];

				double gradient1 = (llh - lll) / (2 * d);
				double gradient2 = (d1h[j] - d1l[j]) / (2 * d);
				//System.out.printf("[%d,%d] ll - %f  (%s %f+/-%f) d1 %f ?= %f : d2 %f ?= %f\n", i, k, ll, func.getName(k), a[k], d, 
				//		gradient1, d1[j], gradient2, d2[j]);
				Assert.assertTrue("Not same gradient1 @ " + j, eq.almostEqualComplement(gradient1, d1[j]));
				Assert.assertTrue("Not same gradient2 @ " + j, eq.almostEqualComplement(gradient2, d2[j]));
			}
		}
	}

	/**
	 * Create random elliptical Gaussian data an returns the data plus an estimate of the parameters.
	 * Only the chosen parameters are randomised and returned for a maximum of (background, amplitude, angle, xpos,
	 * ypos, xwidth, ywidth }
	 *
	 * @param npeaks
	 *            the npeaks
	 * @param params
	 *            set on output
	 * @param randomiseParams
	 *            Set to true to randomise the params
	 * @return the double[]
	 */
	private double[] doubleCreateGaussianData(int npeaks, double[] params, boolean randomiseParams)
	{
		int n = blockWidth * blockWidth;

		// Generate a 2D Gaussian
		ErfGaussian2DFunction func = (ErfGaussian2DFunction) GaussianFunctionFactory.create2D(npeaks, blockWidth,
				blockWidth, GaussianFunctionFactory.FIT_ERF_FREE_CIRCLE, null);
		params[0] = random(Background);
		for (int i = 0, j = 1; i < npeaks; i++, j += 6)
		{
			params[j] = random(Signal);
			params[j + 2] = random(Xpos);
			params[j + 3] = random(Ypos);
			params[j + 4] = random(Xwidth);
			params[j + 5] = random(Ywidth);
		}

		double[] y = new double[n];
		func.initialise(params);
		for (int i = 0; i < y.length; i++)
		{
			// Add random Poisson noise
			y[i] = rdg.nextPoisson(func.eval(i));
		}

		if (randomiseParams)
		{
			params[0] = random(params[0]);
			for (int i = 0, j = 1; i < npeaks; i++, j += 6)
			{
				params[j] = random(params[j]);
				params[j + 2] = random(params[j + 2]);
				params[j + 3] = random(params[j + 3]);
				params[j + 4] = random(params[j + 4]);
				params[j + 5] = random(params[j + 5]);
			}
		}

		return y;
	}

	private double random(double d)
	{
		return d + rdg.nextUniform(-d * 0.1, d * 0.1);
	}

	protected int[] createData(int npeaks, int iter, ArrayList<double[]> paramsList, ArrayList<double[]> yList)
	{
		return createData(npeaks, iter, paramsList, yList, true);
	}

	protected int[] createData(int npeaks, int iter, ArrayList<double[]> paramsList, ArrayList<double[]> yList,
			boolean randomiseParams)
	{
		int[] x = new int[blockWidth * blockWidth];
		for (int i = 0; i < x.length; i++)
			x[i] = i;
		for (int i = 0; i < iter; i++)
		{
			double[] params = new double[1 + 6 * npeaks];
			double[] y = doubleCreateGaussianData(npeaks, params, randomiseParams);
			paramsList.add(params);
			yList.add(y);
		}
		return x;
	}

	protected int[] createFakeData(int nparams, int iter, ArrayList<double[]> paramsList, ArrayList<double[]> yList)
	{
		int[] x = new int[blockWidth * blockWidth];
		for (int i = 0; i < x.length; i++)
			x[i] = i;
		for (int i = 0; i < iter; i++)
		{
			double[] params = new double[nparams];
			double[] y = createFakeData(params);
			paramsList.add(params);
			yList.add(y);
		}
		return x;
	}

	private double[] createFakeData(double[] params)
	{
		int n = blockWidth * blockWidth;
		RandomGenerator r = rdg.getRandomGenerator();

		for (int i = 0; i < params.length; i++)
		{
			params[i] = r.nextDouble();
		}

		double[] y = new double[n];
		for (int i = 0; i < y.length; i++)
		{
			y[i] = r.nextDouble();
		}

		return y;
	}

	protected ArrayList<double[]> copyList(ArrayList<double[]> paramsList)
	{
		ArrayList<double[]> params2List = new ArrayList<double[]>(paramsList.size());
		for (int i = 0; i < paramsList.size(); i++)
		{
			params2List.add(copydouble(paramsList.get(i)));
		}
		return params2List;
	}

	private double[] copydouble(double[] d)
	{
		double[] d2 = new double[d.length];
		for (int i = 0; i < d.length; i++)
			d2[i] = d[i];
		return d2;
	}

	void log(String format, Object... args)
	{
		System.out.printf(format, args);
	}
}