Move the solver to Colour and update all tests

10 files changed, 133 insertions, 292 deletions

diff --git a/gsoc_common.py b/gsoc_common.py
index 5d4aea0..b73d81c 100644
--- a/gsoc_common.py
+++ b/gsoc_common.py
@@ -24,137 +24,6 @@ def model_sd(ccp, primed=True):
 	grid = wvlp if primed else wvl.range()
 	return SpectralDistribution(model(grid, ccp), wvl.range(), name="Model")
 
-# The goal is to minimize the color difference between a given distrbution
-# and the one computed from the model above.
-# This function also calculates the first derivatives with respect to c's.
-def error_function(c, target, wvl, cmfs, illuminant, illuminant_XYZ):
-	U = c[0] * wvl**2 + c[1] * wvl + c[2]
-	t1 = np.sqrt(1 + U**2)
-	R = 1 / 2 + U / (2 * t1)
-
-	t2 = 1 / (2 * t1) - U**2 / (2 * t1**3)
-	dR_dc = [wvl**2 * t2, wvl * t2, t2]
-
-	E = illuminant.values * R / 100
-	dE_dc = illuminant.values * dR_dc / 100
-
-	XYZ = np.empty(3)
-	dXYZ_dc = np.empty((3, 3))
-
-	dlambda = cmfs.wavelengths[1] - cmfs.wavelengths[0]
-	for i in range(3):
-		XYZ[i] = np.dot(E, cmfs.values[:, i]) * dlambda
-		for j in range(3):
-			dXYZ_dc[i, j] = np.dot(dE_dc[j], cmfs.values[:, i]) * dlambda
-
-	# FIXME: this isn't the full CIE 1976 lightness function
-	f = (XYZ / illuminant_XYZ)**(1/3)
-
-	# FIXME: this can be vectorized
-	df_dc = np.empty((3, 3))
-	for i in range(3):
-		for j in range(3):
-			df_dc[i, j] = 1 / (3 * illuminant_XYZ[i]**(1/3)
-			                   * XYZ[i]**(2/3)) * dXYZ_dc[i, j]
-
-	Lab = np.array([
-		116 * f[1] - 16,
-		500 * (f[0] - f[1]),
-		200 * (f[1] - f[2])
-	])
-
-	dLab_dc = np.array([
-		116 * df_dc[1],
-		500 * (df_dc[0] - df_dc[1]),
-		200 * (df_dc[1] - df_dc[2])
-	])
-
-	error = np.sqrt(np.sum((Lab - target)**2))
-
-	derror_dc = np.zeros(3)
-	for i in range(3):
-		for j in range(3):
-			derror_dc[i] += dLab_dc[j, i] * (Lab[j] - target[j])
-		derror_dc[i] /= error
-
-	#print("c=[%.3g, %.3g, %.3g], XYZ=[%.3g, %.3g, %.3g], Lab=[%.3g, %.3g, %.3g], error=%g" % (*c, *XYZ, *Lab, error))
-
-	return error, derror_dc
-
-# 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):
-	ill_sd = ill_sd.align(wvl)
-	ill_XYZ = sd_to_XYZ(ill_sd) / 100
-
-	def do_optimize(XYZ, ccp0):
-		Lab = XYZ_to_Lab(XYZ, ill_xy)
-
-		def fun(*args):
-			error, derror_dc = error_function(*args)
-			return error
-
-		def jac(*args):
-			error, derror_dc = error_function(*args)
-			return derror_dc
-
-		cmfs = STANDARD_OBSERVER_CMFS["CIE 1931 2 Degree Standard Observer"].align(wvl)
-
-		opt = optimize.minimize(
-			fun, ccp0, (Lab, wvlp, cmfs, ill_sd, ill_XYZ), jac=jac,
-			method="L-BFGS-B", options={"disp": verbose}
-		)
-		if verbose:
-			print(opt)
-		return opt
-
-	# 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?
-
-	if verbose:
-		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
-
-	good_XYZ = (1/3, 1/3, 1/3)
-	good_ccp = (2.1276356, -1.07293026, -0.29583292) # FIXME: valid only for D65
-
-	divisions = 3
-	while divisions < 30:
-		if verbose:
-			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 > 1e-3:
-				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
-
-	raise Exception("optimization failed for target_XYZ=%s, ccp0=%s" \
-	                % (target_XYZ, ccp0))
-
 # Makes a comparison plot with SDs and swatches
 def plot_comparison(target, matched_sd, label, error, ill_sd, show=True):
 	if type(target) is SpectralDistribution:
diff --git a/test_bad_convergence.py b/test_bad_convergence.py
deleted file mode 100644
index a47af5a..0000000
--- a/test_bad_convergence.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import numpy as np
-from colour import *
-from gsoc_common import D65, D65_xy, jakob_hanika, model_sd, plot_comparison
-
-# A set of inputs that make the Nelder-Mead solver work really hard.
-# This might not be interesting in future versions of jakob_hanika.py, once
-# the solver is more reliable. Make sure you've checked out the latest commit
-# that updates *this* file.
-if __name__ == "__main__":
-	XYZ = np.array([0.00263241, 0.00174905, 0.00092602])
-	ccp0 = np.array([5.11263599, -2.65344447, -2.37301856])
-
-	ccp, error = jakob_hanika(XYZ, D65, D65_xy, ccp0)
-	plot_comparison(XYZ, model_sd(ccp), "Target", error, D65)
diff --git a/test_colorchecker.py b/test_colorchecker.py
index fc54d84..612d7fb 100644
--- a/test_colorchecker.py
+++ b/test_colorchecker.py
@@ -1,6 +1,7 @@
 import numpy as np
 from colour import *
-from gsoc_common import D65, D65_xy, jakob_hanika, model_sd, plot_comparison
+from colour.recovery import XYZ_to_sd_Jakob2019
+from gsoc_common import D65, plot_comparison
 
 # This demo goes through SDs in a color checker
 if __name__ == "__main__":
@@ -9,7 +10,5 @@ if __name__ == "__main__":
 
 		print("Color checker: The target is '%s' with X=%g, Y=%g, Z=%g"
 		      % (name, *XYZ))
-		ccp, error = jakob_hanika(XYZ, D65, D65_xy)
-		matched_sd = model_sd(ccp)
-
-		plot_comparison(sd, matched_sd, name, error, D65)
+		recovered_sd, error = XYZ_to_sd_Jakob2019(XYZ, return_error=True)
+		plot_comparison(sd, recovered_sd, name, error, D65)
diff --git a/test_coverage.py b/test_coverage.py
index 90e3c79..ee519f4 100644
--- a/test_coverage.py
+++ b/test_coverage.py
@@ -1,99 +1,75 @@
-import numpy as np, os, multiprocessing
+import numpy as np
+import multiprocessing
 from colour import *
 from colour.models import eotf_inverse_sRGB
-from gsoc_common import D65, D65_xy, jakob_hanika, model_sd, plot_comparison
-
+from colour.recovery import coefficients_Jakob2019
+from gsoc_common import model_sd, D65, D65_xy, plot_comparison
 
+# Solve for a specific RGB color
+def optimize_RGB(linear_RGB, coeffs_0):
+	RGB = eotf_inverse_sRGB(linear_RGB)
+	XYZ = sRGB_to_XYZ(RGB, D65_xy)
 
-# Resolution of the discretization cubes
-# Both should be 64, but this code is a bit too slow at the moment
-CHROMA_STEPS = 8
-LIGHTNESS_STEPS = 8
+	coeffs, error = coefficients_Jakob2019(
+		XYZ,
+		dimensionalise=False,
+		coefficients_0=coeffs_0
+	)
 
+	if error > 1e-6:
+		with open("out/bad.txt", "a") as fd:
+			fd.write("\n")
+			fd.write("linear_RGB=%s\n" % linear_RGB)
+			fd.write("XYZ=%s\n" % XYZ)
+			fd.write("coeffs_0=%s\n" % coeffs_0)
+			fd.write("coeffs=%s, error=%g\n" % (coeffs, error))
 
+	# A low budget graph to quickly see what's going on
+	log_error = np.log10(error)
+	bars = "|" * max(0, int(5 * (log_error + 9)))
 
-# Solve for a specific RGB color
-def optimize_RGB(linear_RGB, ccp0):
-	RGB = eotf_inverse_sRGB(linear_RGB)
-	target = sRGB_to_XYZ(RGB, D65_xy)
-	ccp, error = jakob_hanika(target, D65, D65_xy, ccp0, verbose=False)
-	if error > 0.1:
-		print("WARNING: convergence failed with L=%s, starting from %s"
-		      % (linear_RGB, ccp0))
-
-	return ccp, error
+	print("%12.5g %12.5g %12.5g %5.3f %s" % (*XYZ, log_error, bars))
+	return coeffs
 
 # Solve for all lightness values of a fully saturated RGB color
-def optimize_chromaticity(argv):
-	linear_RGB, dirname = argv
-
+def optimize_chromaticity(linear_RGB):
 	# alpha's aren't spaced equally (see the article)
 	def smoothstep(x):
 		return x**2 * (3 - 2 * x)
 
-	ii = np.arange(0, LIGHTNESS_STEPS)
-	aa = smoothstep(smoothstep(ii / LIGHTNESS_STEPS))
-	ccps = np.empty((LIGHTNESS_STEPS, 3)) # FIXME: dtype?
-
-	def f(alpha, ccp0):
-		path = "%s/%.3f" % (dirname, alpha)
-		print("%s..." % path)
+	steps = np.arange(0, LIGHTNESS_STEPS)
+	alphas = smoothstep(smoothstep(steps / LIGHTNESS_STEPS))
 
-		ccp, error = optimize_RGB(linear_RGB * alpha, ccp0)
-
-		try:
-			os.makedirs(os.path.dirname(path))
-		except:
-			pass
-		with open(path, "w") as fd:
-			fd.write("%s, %s" % (ccp, error))
-			if error > 0.1:
-				fd.write(" FAILED")
-			fd.write("\n")
+	i_mid = LIGHTNESS_STEPS // 5
+	coeffs_mid = optimize_RGB(linear_RGB * alphas[i_mid], (0, 0, 0))
 
-		return ccp, error
-
-
-	i_mid = LIGHTNESS_STEPS // 2
-	ccps[i_mid, :], _ = f(aa[i_mid], (0, 0, 0))
-
-	ccp0 = ccps[i_mid, :]
+	coeffs_0 = coeffs_mid
 	for i in range(i_mid + 1, LIGHTNESS_STEPS):
-		ccps[i, :], _ = f(aa[i], ccp0)
+		coeffs_0 = optimize_RGB(linear_RGB * alphas[i], coeffs_0)
 
-	ccp0 = ccps[i_mid, :]
+	coeffs_0 = coeffs_mid
 	for i in reversed(range(0, i_mid)):
-		ccps[i, :], _ = f(aa[i], ccp0)
+		coeffs_0 = optimize_RGB(linear_RGB * alphas[i], coeffs_0)
 
 
 
-# This demo discretizes the sRGB space to three 8x8x8 cubes and tries to find
-# the model parameters for all the colors. Coefficients are saved to a file
-# somewhere in the 'out' directory (each color gets its own file).
-# FIXME: This dumpster fire of an output format is to be replaced ASAP.
-# FIXME: This takes a good while (on the order of an hour).
+# This program runs XYZ_to_sd_Jakob2019 converges on a large number of inputs,
+# covering an entire RGB gamut, to see if it'll diverge somewhere.
 if __name__ == "__main__":
-	args = []
+	CHROMA_STEPS = 8
+	LIGHTNESS_STEPS = 64
+
+	with open("out/bad.txt", "w") as fd:
+		fd.write("Going through %dx%dx%d cubes\n"
+		         % (LIGHTNESS_STEPS, CHROMA_STEPS, CHROMA_STEPS))
 
-	done = 0
-	total = 0
+	args = []
 	for A in np.linspace(0, 1, CHROMA_STEPS):
 		for B in np.linspace(0, 1, CHROMA_STEPS):
 			for RGB in [np.array([1, A, B]),
 			            np.array([A, 1, B]),
 			            np.array([A, B, 1])]:
-
-				dirname = "out/R%.3fG%.3fB%.3f" % (*RGB,)
-				total += 1
-
-				if os.path.exists(dirname) and \
-				   len(os.listdir(dirname)) == LIGHTNESS_STEPS:
-					done += 1
-					continue
-
-				args.append((RGB, dirname))
-
-	print("%d/%d already done" % (done, total))
+				args.append(RGB)
 
 	pool = multiprocessing.Pool()
 	pool.map(optimize_chromaticity, args)
diff --git a/test_diff.py b/test_diff.py
index b634ca6..0ee5f58 100644
--- a/test_diff.py
+++ b/test_diff.py
@@ -1,51 +1,50 @@
 import numpy as np
 from scipy.optimize import minimize
 from colour import *
-from colour.colorimetry import STANDARD_OBSERVER_CMFS, ILLUMINANT_SDS
-from colour.models import eotf_inverse_sRGB, sRGB_to_XYZ
+from colour.recovery import error_function_Jakob2019
 from matplotlib import pyplot as plt
-from gsoc_common import plot_comparison, error_function, model_sd, D65_xy
-
-shape = SpectralShape(360, 830, 1)
-cmfs = STANDARD_OBSERVER_CMFS["CIE 1931 2 Degree Standard Observer"].align(shape)
-
-illuminant = SpectralDistribution(ILLUMINANT_SDS["D65"]).align(shape)
-illuminant_XYZ = sd_to_XYZ(illuminant) / 100
-wvl = np.linspace(0, 1, len(shape.range()))
-
-target = np.array([50, -20, 30]) # Some arbitrary Lab coordinates
-xs = np.linspace(-10, 10, 500)
-h = xs[1] - xs[0]
+from gsoc_common import plot_comparison
 
 # This test checks if derivatives are calculated correctly by comparing them
 # to finite differences.
-for c_index in range(3):
-	errors = np.empty(len(xs))
-	derrors = np.empty(len(xs))
+if __name__ == "__main__":
+	shape = SpectralShape(360, 830, 1)
+	cmfs = STANDARD_OBSERVER_CMFS["CIE 1931 2 Degree Standard Observer"].align(shape)
+
+	illuminant = SpectralDistribution(ILLUMINANT_SDS["D65"]).align(shape)
+	illuminant_XYZ = sd_to_XYZ(illuminant) / 100
+
+	target = np.array([50, -20, 30]) # Some arbitrary Lab colour
+	xs = np.linspace(-10, 10, 500)
+	h = xs[1] - xs[0]
 
-	for i, x in enumerate(xs):
-		c = np.array([1.0, 1, 1])
-		c[c_index] = x
+	# Vary one coefficient at a time
+	for c_index in range(3):
+		errors = np.empty(len(xs))
+		derrors = np.empty(len(xs))
 
-		error, derror_dc = error_function(
-			c, target, wvl, cmfs, illuminant, illuminant_XYZ
-		)
+		for i, x in enumerate(xs):
+			c = np.array([1.0, 1, 1])
+			c[c_index] = x
 
-		errors[i] = error
-		derrors[i] = derror_dc[c_index]
+			error, derror_dc = error_function_Jakob2019(
+				c, target, shape, cmfs, illuminant, illuminant_XYZ
+			)
 
+			errors[i] = error
+			derrors[i] = derror_dc[c_index]
 
-	plt.subplot(2, 3, 1 + c_index)
-	plt.xlabel("c%d" % c_index)
-	plt.ylabel("ΔE")
-	plt.plot(xs, errors)
 
-	plt.subplot(2, 3, 4 + c_index)
-	plt.xlabel("c%d" % c_index)
-	plt.ylabel("dΔE/dc%d" % c_index)
+		plt.subplot(2, 3, 1 + c_index)
+		plt.xlabel("c%d" % c_index)
+		plt.ylabel("ΔE")
+		plt.plot(xs, errors)
 
-	plt.plot(xs, derrors, "k-")
-	plt.plot(xs[:-1] + h / 2, np.diff(errors) / h, "r:")
+		plt.subplot(2, 3, 4 + c_index)
+		plt.xlabel("c%d" % c_index)
+		plt.ylabel("dΔE/dc%d" % c_index)
 
+		plt.plot(xs, derrors, "k-")
+		plt.plot(xs[:-1] + h / 2, np.diff(errors) / h, "r:")
 
-plt.show()
+	plt.show()
diff --git a/test_divergence.py b/test_divergence.py
new file mode 100644
index 0000000..fc7956c
--- /dev/null
+++ b/test_divergence.py
@@ -0,0 +1,9 @@
+import numpy as np
+from colour import *
+from colour.recovery import XYZ_to_sd_Jakob2019
+from gsoc_common import D65, plot_comparison
+
+if __name__ == "__main__":
+	XYZ = np.array([0.1788, 0.3576, 0.0596])
+	found_sd, error = XYZ_to_sd_Jakob2019(XYZ, return_error=True)
+	plot_comparison(XYZ, found_sd, "Target", error, D65)
diff --git a/test_error_function.py b/test_error_function.py
new file mode 100644
index 0000000..99e1268
--- /dev/null
+++ b/test_error_function.py
@@ -0,0 +1,29 @@
+import numpy as np
+from colour import *
+from colour.recovery import error_function_Jakob2019
+from gsoc_common import model_sd, D65_xy
+
+if __name__ == "__main__":
+	shape = SpectralShape(360, 830, 1)
+	cmfs = STANDARD_OBSERVER_CMFS["CIE 1931 2 Degree Standard Observer"].align(shape)
+
+	illuminant = SpectralDistribution(ILLUMINANT_SDS["D65"]).align(shape)
+	illuminant_XYZ = sd_to_XYZ(illuminant) / 100
+
+	coefficients = np.array([8.70184886, -13.19804478, 2.12180137])
+	target = np.array([20, 50, 30])
+
+	error, derror_dc, R, XYZ, Lab = error_function_Jakob2019(
+		coefficients, target, shape, cmfs,
+		illuminant, illuminant_XYZ, True
+	)
+
+	sd = model_sd(coefficients)
+	good_XYZ = sd_to_XYZ(sd, illuminant=illuminant)
+	good_Lab = XYZ_to_Lab(good_XYZ / 100, D65_xy)
+
+	print("Good XYZ: %g %g %g" % (*good_XYZ,))
+	print("  EF XYZ: %g %g %g" % (*XYZ,))
+	print("Good Lab: %g %g %g" % (*good_Lab,))
+	print("  EF Lab: %g %g %g" % (*Lab,))
+	print(" EF* Lab: %g %g %g" % (*XYZ_to_Lab(XYZ / 100, D65_xy),))
diff --git a/test_example.py b/test_example.py
index 80b07c7..a4104b1 100644
--- a/test_example.py
+++ b/test_example.py
@@ -1,31 +1,25 @@
 import numpy as np
 from colour import *
-from colour.difference import delta_E_CIE1976
 from colour.models import eotf_inverse_sRGB
-from gsoc_common import D65, D65_xy, jakob_hanika, model_sd, plot_comparison
-
-
-
-# These numbers are taken from Jakob and Hanika's Jupyter notebook.
-RGB_ref = np.array([0.79264853,  0.4, 0.63703843]) # *linear* sRGB
-cc_ref = np.array([ 18.70184886, -13.19804478, 2.12180137])
-
-
+from colour.difference import delta_E_CIE1976
+from colour.recovery import XYZ_to_sd_Jakob2019
+from gsoc_common import model_sd, D65, D65_xy, plot_comparison
 
 if __name__ == "__main__":
-	#XYZ = sRGB_to_XYZ(eotf_inverse_sRGB(RGB_ref), D65_xy)
-	XYZ = np.array([1/3, 1/4, 1/3])
-	Lab = XYZ_to_Lab(XYZ, D65_xy)
+	# These numbers are taken from Jakob and Hanika's Jupyter notebook.
+	RGB_ref = np.array([0.79264853,  0.4, 0.63703843]) # *linear* sRGB
+	cc_ref = np.array([ 18.70184886, -13.19804478, 2.12180137])
 
+	XYZ = sRGB_to_XYZ(eotf_inverse_sRGB(RGB_ref), D65_xy)
+	Lab = XYZ_to_Lab(XYZ, D65_xy)
 	print("Target: X=%g, Y=%g, Z=%g, L=%g, a=%g, b=%g" % (*XYZ, *Lab))
-	ccp, error = jakob_hanika(XYZ, D65, D65_xy)
 
 	reference_sd = model_sd(cc_ref)
-	reference_XYZ = sd_to_XYZ(reference_sd, illuminant=D65)
+	reference_XYZ = sd_to_XYZ(reference_sd, illuminant=D65) / 100
 	reference_Lab = XYZ_to_Lab(reference_XYZ, D65_xy)
 
-	found_sd = model_sd(ccp)
-	found_XYZ = sd_to_XYZ(found_sd, illuminant=D65)
+	found_sd, error = XYZ_to_sd_Jakob2019(XYZ, return_error=True)
+	found_XYZ = sd_to_XYZ(found_sd, illuminant=D65) / 100
 	found_Lab = XYZ_to_Lab(found_XYZ, D65_xy)
 
 	print("Our results differ from the reference by ΔE = %g" \
diff --git a/test_integration.py b/test_integration.py
deleted file mode 100644
index d1b564f..0000000
--- a/test_integration.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import numpy as np
-from colour import *
-from colour.models import eotf_inverse_sRGB
-from colour.recovery import Jakob2019Interpolator
-from colour.plotting import plot_multi_sds
-from gsoc_common import D65_xy
-
-
-
-# This is a basic test of the feature/jakob2019 branch
-if __name__ == "__main__":
-	RGB = np.array([0.79264853,  0.4, 0.63703843])
-	XYZ = sRGB_to_XYZ(eotf_inverse_sRGB(RGB), D65_xy)
-	sd = XYZ_to_sd(XYZ, method="Jakob 2019", verbose=True)
-
-	interp = Jakob2019Interpolator()
-	interp.from_file("data/srgb.coeff")
-	sdi = interp.RGB_to_sd(RGB)
-
-	plot_multi_sds([sd, sdi])
diff --git a/test_interpolator.py b/test_interpolator.py
index e4d2919..ab542f3 100644
--- a/test_interpolator.py
+++ b/test_interpolator.py
@@ -3,7 +3,7 @@ from colour import *
 from colour.difference import delta_E_CIE1976
 from colour.models import eotf_inverse_sRGB
 from colour.recovery import Jakob2019Interpolator
-from gsoc_common import D65, D65_xy, jakob_hanika, model_sd, plot_comparison
+from gsoc_common import D65, D65_xy, model_sd, plot_comparison
 
 # This script tests if the interpolator correctly handles multi-dimensional
 # inputs.
@@ -20,7 +20,7 @@ if __name__ == "__main__":
 	cc = ccs[0, 0, 0, 0, :]
 
 	matched_sd = model_sd(cc, primed=False)
-	matched_XYZ = sd_to_XYZ(matched_sd, illuminant=D65)
+	matched_XYZ = sd_to_XYZ(matched_sd, illuminant=D65) / 100
 	matched_Lab = XYZ_to_Lab(matched_XYZ, D65_xy)
 	error = delta_E_CIE1976(Lab, matched_Lab)