Converts .RGB files into colorscheme strings that can be used in MathGL

cmap_parser.py

""" This searches for an accurate linear approximation of a ".rgb" file gradient.
	Run it like so: ``python3 cmap_parser.py path/to/color_map.rgb``

	See this site for a list of colormaps
	https://www.ncl.ucar.edu/Document/Graphics/color_table_gallery.shtml
"""

import random, sys, re, math

class RGBApproximation:
	""" A subset of colors used to approximate a full gradient """
	def __init__(self, err, stops:list, parent):
		self.error = err
		""" (list) linear interpolation error """
		self.stops = stops
		""" (list) color indices, making up the gradient approximation """
		self.parent = parent
		""" (RGBDefinition) parent definition where we get color data from """
		
	def mathgl(self) -> str:
		""" Returns a string represented the color stops in MathGL format """
		out = []
		size = len(self.parent.colors)-1
		# multiplier to make them char RGB vals
		mult = 255 if self.parent.mode == "float" else 1;
		for stop in self.stops:
			pos = stop/size
			# convert to char and round
			rgb = list(max(0,min(255,round(v*mult))) for v in self.parent.colors[stop])
			out.append("{x%02x%02x%02x,%.4g}" % (rgb[0],rgb[1],rgb[2],pos))
		return "".join(out)
		

class RGBDefinition:
	""" Container for parsed .rgb file data """
	# some regex to parse the file
	re_prop = re.compile(r"^(\w+)\s?=\s?(\w+)$")
	re_color = re.compile(r"^(?:(\d+(?:\.\d*)?)\s+){2,}(\d+(\.\d*)?)$")
	re_ignore = re.compile(r"(\s+|#.*$)")

	def __init__(self, fname:str):
		""" Parse an .rgb file, creating an RGBDefinition.

			:param fname: The file that should be parsed
		"""
		self.colors = []
		""" (list) list of colors, where each entry is a tuple with 3 values (r,g,b) """
		self.props = {}
		""" (dict) properties that we found inside the .rgb file """
		self.mode = None
		""" (str) "int" for 0-255 color range, or "float" for 0-1 """

		max_val = 0
		with open(fname,"r") as f:
			lines = f.readlines()
			for line in lines:
				# remove comments and extra whitespace
				line = self.re_ignore.sub(" ", line).strip()
				if not line:
					continue
				# rgb definition
				if self.re_color.match(line):
					col = [float(v) for v in line.split(" ")]
					max_val = max(max_val, *col)
					self.colors.append(col)
					continue
				# property definition
				match = self.re_prop.match(line)
				if match:
					g = match.groups()
					self.props[g[0]] = g[1]
					continue
				raise RuntimeError("Don't know how to parse this line in .rgb file:\n{}".format(line))
		self.mode = "float" if max_val <= 1 else "int"

		print("Parsed file:")
		print("\t- max RGB value:", max_val)
		print("\t- color mode:", self.mode)
		print("\t- colors:", len(self.colors))
		print("\t- properties:", self.props)

		if len(self.colors) < 2:
			raise RuntimeError("Found less than 2 color stops in the .rgb file; won't continue")

	def linear_approximate(self, max_stops:int=32):
		""" Approximates the self.colors gradient with linear interpolated color stops

			:param max_stops: Maximum number of color stops to try
		"""
		size = len(self.colors)
		if not isinstance(max_stops, int) or max_stops < 2:
			raise TypeError("max_stops must be int >= 2")
		if max_stops > size:
			print("Limiting max_stops to original gradient size", size)
			max_stops = size

		def eval_line(start:int, end:int):
			""" Measures the error (SSE) if you were to use linear interpolation between two
				points in self.colors

				:param start: the index of the starting point in self.colors
				:param end: the index of the ending point in self.colors
			"""
			tot = 0
			for i in range(start+1, end):
				p = (i-start)/float(end-start);
				for ax in range(3):
					pred = (1-p)*self.colors[start][ax] + p*self.colors[end][ax];
					err = pred - self.colors[i][ax];
					tot += err*err;
			return tot

		def eval_gradient(stops:list):
			""" Measures the error (SSE) for a sequence of stops, using linear interpolation
				between each pair (see eval_line)

				:param stops: list of stops, sorted in ascending order; should exclude the first
					and last stops as those are included automatically (e.g. 0 and size-1)
			"""
			tot = 0
			for i in range(len(stops)-1):
				tot += eval_line(stops[i],stops[i+1])
			return tot
						
		# we always include the start and end colors of gradient as stops;
		# the optimizer is really simple, just trying random subsets self.colors until it finds one with low error
		col_idxs = list(range(0,size-1))
		for stops in range(max_stops-2+1):
			# true number of possibilities is (n k)
			combos = math.comb(size-2, stops)
			print("Trying {} stops ({} combinations)".format(stops+2, combos))
			# monte-carlo optimization; could do brute force search for small combination counts
			best = None
			for sample in range(min(10000,combos)):
				# shuffle the colors; fischer-yates partial shuffle
				for i in range(stops):
					rand_i = random.randrange(i,size-1)
					col_idxs[i], col_idxs[rand_i] = col_idxs[rand_i], col_idxs[i]
				# grab the random number color stops
				approx = col_idxs[:stops]
				approx.sort()
				# always include first/last stop
				approx.insert(0,0)
				approx.append(size-1)
				# TODO: maybe trim if you have three equal color stops all in a row
				# check if this is best approximation we've seen
				err = eval_gradient(approx)
				if best is None or best.error > err:
					best = RGBApproximation(err, approx, self)
			print("\t- error:", best.error)
			print("\t- mathgl:", best.mathgl())
		

if len(sys.argv) != 2:
	raise SyntaxError("Pass a single command line argument, that being the .rgb file to examine")

rgb = RGBDefinition(sys.argv[1])
rgb.linear_approximate()