Vectorize the interpolation code

1 files changed, 25 insertions, 16 deletions

diff --git a/test_coeff.py b/test_coeff.py
index e53c6f1..733dc45 100644
--- a/test_coeff.py
+++ b/test_coeff.py
@@ -3,6 +3,7 @@ from scipy.interpolate import RegularGridInterpolator
 from colour import *
 from colour.plotting import *
 from colour.models import eotf_inverse_sRGB
+from colour.difference import delta_E_CIE1976
 from jakob_hanika import jakob_hanika, model_sd
 from test_colorchecker import plot_comparison
 
@@ -29,21 +30,19 @@ class JakobHanikaInterpolator:
 			coeffs = coeffs.reshape(3, self.res, self.res, self.res, 3)
 
 		t = np.linspace(0, 1, self.res)
-		axes = [self.scale, t, t]
-		self.cubes = [RegularGridInterpolator(axes, coeffs[j, :, :, :, :], bounds_error=False) for j in range(3)]
+		axes = [[0, 1, 2], self.scale, t, t]
+		self.cubes = RegularGridInterpolator(axes, coeffs[:, :, :, :, :], bounds_error=False)
 
 	def __call__(self, RGB):
 		RGB = np.asarray(RGB, dtype=np.float) # FIXME: default dtype
-		vmax = RGB.max()
-		imax = RGB.argmax()
-		cube = self.cubes[imax]
-		chroma = RGB / (vmax + 1e-10)
-		chroma = np.array([
-			vmax,
-			chroma[..., (imax + 2) % 3],
-			chroma[..., (imax + 1) % 3]
-		])
-		return cube(chroma).squeeze()
+		vmax = np.max(RGB, axis=-1)
+		imax = np.argmax(RGB, axis=-1)
+		chroma = RGB / (np.expand_dims(vmax, -1) + 1e-10) # Avoid division by zero
+		vmax = np.max(RGB, axis=-1)
+		v2 = np.take_along_axis(chroma, np.expand_dims((imax + 2) % 3, axis=-1), axis=-1).squeeze(axis=-1)
+		v3 = np.take_along_axis(chroma, np.expand_dims((imax + 1) % 3, axis=-1), axis=-1).squeeze(axis=-1)
+		coords = np.stack([imax, vmax, v2, v3], axis=-1)
+		return self.cubes(coords).squeeze()
 
 
 
@@ -51,9 +50,19 @@ if __name__ == "__main__":
 	jh = JakobHanikaInterpolator()
 	jh.read_file("data/srgb.coeff")
 
-	linear_RGB = np.array([1, 0, 0]) # A random test color
-	RGB = eotf_inverse_sRGB(linear_RGB)
+	RGBs = np.random.random((7, 6, 5, 4, 3))
+	ccs = jh(RGBs)
+
+	RGB = eotf_inverse_sRGB(RGBs[0, 0, 0, 0, :])
 	XYZ = sRGB_to_XYZ(RGB)
-	cc = jh(linear_RGB)
+	Lab = XYZ_to_Lab(XYZ, ill_xy)
+	cc = ccs[0, 0, 0, 0, :]
+
+	matched_sd = model_sd(cc, primed=False)
+	matched_XYZ = sd_to_XYZ(matched_sd, illuminant=ill_sd)
+	matched_Lab = XYZ_to_Lab(matched_XYZ, ill_xy)
+	error = delta_E_CIE1976(Lab, matched_Lab)
+
+	print(matched_sd)
 
-	plot_comparison(XYZ, model_sd(cc, primed=False), "Model", 123, ill_sd)
+	plot_comparison(XYZ, matched_sd, "Model", error, ill_sd, ill_xy)