Make sure the solver works for all inputs

This is achieved by using feedback. In case of a failure, a color with known coefficients is used to solve for a point somewhere between the target and the color. The process is repeated, getting closer and closer to the target until it finally converges.

3 files changed, 62 insertions, 56 deletions

diff --git a/jakob_hanika.py b/jakob_hanika.py
index f9629d4..af55c94 100644
--- a/jakob_hanika.py
+++ b/jakob_hanika.py
@@ -1,16 +1,17 @@
 import numpy as np, scipy.optimize as optimize
 from colour import *
-from colour.difference import *
+from colour.difference import delta_E_CIE1976
+from colour.colorimetry import *
 
 
 
 # The same wavelength grid is used throughout
-wvl = np.arange(360, 830, 1)
+wvl = np.arange(360, 830, 5)
 wvlp = (wvl - 360) / (830 - 360)
 
 # This is the model of spectral reflectivity described in the article.
 def model(wvlp, ccp):
-	yy = ccp[0] * wvlp ** 2 + ccp[1] * wvlp + ccp[2]
+	yy = ccp[0] * wvlp**2 + ccp[1] * wvlp + ccp[2]
 	return 1 / 2 + yy / (2 * np.sqrt(1 + yy ** 2))
 
 # Create a SpectralDistribution using given coefficients
@@ -30,54 +31,62 @@ def error_function(ccp, target, ill_sd, ill_xy):
 
 
 
-# Map primed (dimensionless) coefficients to normal
-# FIXME: is this even useful for anything?
-def cc_from_primed(ccp):
-	return np.array([
-		ccp[0] / 220900,
-		ccp[1] / 470 - 36/11045 * ccp[0],
-		ccp[2] - 36/47 * ccp[1] + 1296/2209 * ccp[0]
-	])
-
+# Finds the parameters for Jakob and Hanika's model
+def jakob_hanika(target_XYZ, ill_sd, ill_xy, ccp0=(0, 0, 0), verbose=True, try_hard=True):
+	def do_optimize(XYZ, ccp0):
+		Lab = XYZ_to_Lab(XYZ, ill_xy)
+		opt = optimize.minimize(
+			error_function, ccp0, (Lab, ill_sd, ill_xy),
+			method="Nelder-Mead", options={"disp": verbose}
+		)
+		if verbose:
+			print(opt)
+		return opt
 
-# This callback to scipy.optimize.basinhopping tells the solver to stop once
-# the error is small enough. The threshold was chosen arbitrarily, as a small
-# fraction of the JND (about 2.3 with this metric).
-def cb_basinhopping(x, f, accept):
-	return f < 0.1
+	# A special case that's hard to solve numerically
+	if np.allclose(target_XYZ, [0, 0, 0]):
+		return np.array([0, 0, -1e+9]), 0 # FIXME: dtype?
 
-# Finds the parameters for Jakob and Hanika's model
-def jakob_hanika(target, ill_sd, ill_xy, ccp0=(0, 0, 0), verbose=True):
-	# First, a conventional solver is called. For 'yellow green' this
-	# actually fails: gets stuck at a local minimum that's far away
-	# from the global one.
-	# FIXME: better parameters, a better x0, a better method?
-	# FIXME: stop iterating as soon as delta E is negligible
-	opt = optimize.minimize(
-		error_function, ccp0, (target, ill_sd, ill_xy),
-		bounds=[[-300, 300]] * 3,
-		options={"disp": verbose}
-	)
 	if verbose:
-		print(opt)
+		print("Trying the target directly, XYZ=%s" % target_XYZ)
+	opt = do_optimize(target_XYZ, ccp0)
+	if opt.fun < 0.1 or not try_hard:
+		return opt.x, opt.fun
 
-	error = error_function(opt.x, target, ill_sd, ill_xy)
-	if verbose:
-		print("Delta E is %g" % error)
+	good_XYZ = (1/3, 1/3, 1/3)
+	good_ccp = (2.1276356, -1.07293026, -0.29583292) # FIXME: valid only for D65
 
-	# Basin hopping is far more likely to find the actual minimum we're
-	# looking for, but it's extremely slow in comparison.
-	if error > 0.1:
+	divisions = 3
+	while divisions < 10:
 		if verbose:
-			print("Error too large, trying global optimization")
-		opt = optimize.basinhopping(
-			lambda cc: error_function(cc, target, ill_sd, ill_xy),
-			x0, disp=verbose, callback=cb_basinhopping
-		)
-		if verbose:
-			print(opt)
-		error = error_function(opt.x, target, ill_sd, ill_xy)
-		if verbose:
-			print("Global delta E is %g" % error)
+			print("Trying with %d divisions" % divisions)
+
+		keep_divisions = False
+		ref_XYZ = good_XYZ
+		ref_ccp = good_ccp
+
+		ccp0 = ref_ccp
+		for i in range(1, divisions):
+			intermediate_XYZ = ref_XYZ + (target_XYZ - ref_XYZ) * i / (divisions - 1)
+			if verbose:
+				print("Intermediate step %d/%d, XYZ=%s with ccp0=%s" %
+				      (i + 1, divisions, intermediate_XYZ, ccp0))
+			opt = do_optimize(intermediate_XYZ, ccp0)
+			if opt.fun > 0.1:
+				if verbose:
+					print("WARNING: intermediate optimization failed")
+				break
+			else:
+				good_XYZ = intermediate_XYZ
+				good_ccp = opt.x
+				keep_divisions = True
+
+			ccp0 = opt.x
+		else:
+			return opt.x, opt.fun
+
+		if not keep_divisions:
+			divisions += 2
 
-	return opt.x, error
+	raise Exception("optimization failed for target_XYZ=%s, ccp0=%s" \
+	                % (target_XYZ, ccp0))
diff --git a/test_colorchecker.py b/test_colorchecker.py
index 0368058..83325b6 100644
--- a/test_colorchecker.py
+++ b/test_colorchecker.py
@@ -40,10 +40,9 @@ if __name__ == "__main__":
 	# This demo goes through SDs in a color checker
 	for name, sd in COLOURCHECKERS_SDS['ColorChecker N Ohta'].items():
 		XYZ = sd_to_XYZ(sd, illuminant=ill_sd) / 100
-		target = XYZ_to_Lab(XYZ, ill_xy)
 
-		print("Color checker: The target is '%s' with L=%g, a=%g, b=%g" % (name, *target))
-		ccp, error = jakob_hanika(target, ill_sd, ill_xy)
+		print("Color checker: The target is '%s' with X=%g, Y=%g, Z=%g" % (name, *XYZ))
+		ccp, error = jakob_hanika(XYZ, ill_sd, ill_xy)
 		matched_sd = model_sd(ccp)
 
 		plot_comparison(sd, matched_sd, name, error, ill_sd)
diff --git a/test_example.py b/test_example.py
index f497bbf..d0be8dd 100644
--- a/test_example.py
+++ b/test_example.py
@@ -1,8 +1,6 @@
 import numpy as np
 from colour import *
-from colour.plotting import *
-from colour.difference import *
-from matplotlib import pyplot as plt
+from colour.difference import delta_E_CIE1976
 from jakob_hanika import jakob_hanika, model_sd
 from test_colorchecker import plot_comparison
 
@@ -20,10 +18,10 @@ cc_ref = np.array([ 18.70184886, -13.19804478, 2.12180137])
 
 if __name__ == "__main__":
 	XYZ = sRGB_to_XYZ(1.055 * RGB_ref ** (1/2.4) - 0.055, ill_xy)
-	target = XYZ_to_Lab(XYZ, ill_xy)
+	Lab = XYZ_to_Lab(XYZ, ill_xy)
 
-	print("Target: X=%g, Y=%g, Z=%g, L=%g, a=%g, b=%g" % (*XYZ, *target))
-	ccp, error = jakob_hanika(target, ill_sd, ill_xy)
+	print("Target: X=%g, Y=%g, Z=%g, L=%g, a=%g, b=%g" % (*XYZ, *Lab))
+	ccp, error = jakob_hanika(XYZ, ill_sd, ill_xy)
 
 	reference_sd = model_sd(cc_ref)
 	reference_XYZ = sd_to_XYZ(reference_sd, illuminant=ill_sd)