aplotter.py

#-----------------------------------------------
#aplotter.py - ascii art function plotter
#Copyright (c) 2006, Imri Goldberg
#All rights reserved.
#
#Redistribution and use in source and binary forms,
#with or without modification, are permitted provided
#that the following conditions are met:
#
#    * Redistributions of source code must retain the
#		above copyright notice, this list of conditions
#		and the following disclaimer.
#    * Redistributions in binary form must reproduce the
#		above copyright notice, this list of conditions
#		and the following disclaimer in the documentation
#		and/or other materials provided with the distribution.
#    * Neither the name of the <ORGANIZATION> nor the names of
#		its contributors may be used to endorse or promote products
#		derived from this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
#LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#-----------------------------------------------
#ported to python3 by Helmut Carter in 2022
#-----------------------------------------------
import argparse
import math
import sys


EPSILON = 0.000001

def transposed(mat):
	result = []
	for i in range(len(mat[0])):
		result.append([x[i] for x in mat])
	return result

def y_reversed(mat):    
	result = []
	for i in range(len(mat)):
		result.append(list(reversed(mat[i])))
	return result

def sign(x):
	if 0<x:
		return 1
	if 0 == x:
		return 0
	return -1

class Plotter(object):

	class PlotData(object):
		def __init__(self, x_size, y_size, min_x, max_x, min_y, max_y, x_mod, y_mod):
			self.x_size = x_size
			self.y_size = y_size
			self.min_x = min_x
			self.max_x = max_x
			self.min_y = min_y
			self.max_y = max_y
			self.x_mod = x_mod
			self.y_mod = y_mod

			self.x_step = float(max_x - min_x)/float(self.x_size)
			self.y_step = float(max_y - min_y)/float(self.y_size)
			self.inv_x_step = 1/self.x_step
			self.inv_y_step = 1/self.y_step

			self.ratio = self.y_step / self.x_step
		def __repr__(self):
			s = "size: %s, bl: %s, tr: %s, step: %s" % ((self.x_size, self.y_size), (self.min_x, self.min_y), (self.max_x, self.max_y),
														 (self.x_step, self.y_step))
			return s
	
	def __init__(self, **kwargs):

		self.x_size = kwargs.get("x_size", 80)
		self.y_size = kwargs.get("y_size", 20)

		self.will_draw_axes = kwargs.get("draw_axes", True)

		self.new_line = kwargs.get("newline", "\n")

		self.dot = kwargs.get("dot", "*")

		self.plot_slope = kwargs.get("plot_slope", True)

		self.x_margin = kwargs.get("x_margin", 0.05)
		self.y_margin = kwargs.get("y_margin", 0.1)

		self.will_plot_labels = kwargs.get("plot_labels", True)

	@staticmethod
	def get_symbol_by_slope(slope, default_symbol):
		draw_symbol = default_symbol
		if slope > math.tan(3*math.pi/8):
			draw_symbol = "|"
		elif slope > math.tan(math.pi/8) and slope < math.tan(3*math.pi/8):
			draw_symbol = "/"
		elif abs(slope) < math.tan(math.pi/8):
			draw_symbol = "-"
		elif slope < math.tan(-math.pi/8) and slope > math.tan(-3*math.pi/8):
			draw_symbol = "\\"
		elif slope < math.tan(-3*math.pi/8):
			draw_symbol = "|"
		return draw_symbol    

	@staticmethod
	def write_clipped(output_buffer, x_start, y_coord, text):
		if y_coord < 0 or y_coord >= len(output_buffer[0]):
			return
		for i, c in enumerate(text):
			x_coord = x_start + i
			if 0 <= x_coord < len(output_buffer):
				output_buffer[x_coord][y_coord] = c

	@staticmethod
	def get_adjacent_coord(coord, size):
		if coord > 0:
			return coord - 1
		if coord + 1 < size:
			return coord + 1
		return coord

	def plot_point(self, x, y, output_buffer, plot_data, draw_symbol=None):
		if x < plot_data.min_x or x >= plot_data.max_x or y < plot_data.min_y or y >= plot_data.max_y:
			return

		x_coord = self.get_coord(x, plot_data.min_x, plot_data.x_step)
		y_coord = self.get_coord(y, plot_data.min_y, plot_data.y_step)
		if x_coord >= 0 and x_coord < len(output_buffer) and y_coord >= 0 and y_coord < len(output_buffer[0]):
			if draw_symbol is None:
				draw_symbol = self.dot
			if self.will_draw_axes:
				if y_coord == self.get_coord(0, plot_data.min_y, plot_data.y_step) and draw_symbol == "-":
					draw_symbol = "="
			output_buffer[x_coord][y_coord] = draw_symbol

	def plot_labels(self, output_buffer, plot_data):
		if plot_data.y_size < 1 or plot_data.x_size < 1:
			return

		margin_factor = 1

		do_plot_x_label = True
		do_plot_y_label = True

		x_str = "%+g"
		if plot_data.x_size < 16:
			do_plot_x_label = False
		elif plot_data.x_size < 23:
			x_str = "%+.2g" 

		y_str = "%+g"    
		if plot_data.x_size < 8:
			do_plot_y_label = False
		elif plot_data.x_size < 11:
			y_str = "%+.2g"
			
		act_min_x = (plot_data.min_x + plot_data.x_mod*margin_factor)
		act_max_x = (plot_data.max_x - plot_data.x_mod*margin_factor)
		act_min_y = (plot_data.min_y + plot_data.y_mod*margin_factor)
		act_max_y = (plot_data.max_y - plot_data.y_mod*margin_factor)

		if abs(act_min_x) < 1:
			min_x_str = "%+.2g" % act_min_x
		else:
			min_x_str = x_str % act_min_x

		if abs(act_max_x) < 1:
			max_x_str = "%+.2g" % act_max_x
		else:
			max_x_str = x_str % act_max_x
		
		if abs(act_min_y) < 1:
			min_y_str = "%+.2g" % act_min_y
		else:
			min_y_str = y_str % act_min_y

		if abs(act_max_y) < 1:
			max_y_str = "%+.2g" % act_max_y
		else:
			max_y_str = y_str % act_max_y
						 
		min_x_coord = self.get_coord(act_min_x,plot_data.min_x,plot_data.x_step)
		max_x_coord = self.get_coord(act_max_x,plot_data.min_x,plot_data.x_step)
		min_y_coord = self.get_coord(act_min_y,plot_data.min_y,plot_data.y_step)
		max_y_coord = self.get_coord(act_max_y,plot_data.min_y,plot_data.y_step)
								 

		has_y_axis = plot_data.min_x <= 0 and plot_data.max_x > 0
		has_x_axis = plot_data.min_y <= 0 and plot_data.max_y > 0

		x_zero_coord = self.get_coord(0, plot_data.min_x, plot_data.x_step) if has_y_axis else 0
		y_zero_coord = self.get_coord(0, plot_data.min_y, plot_data.y_step) if has_x_axis else 0

		if has_y_axis:
			output_buffer[x_zero_coord][min_y_coord] = "+"
			output_buffer[x_zero_coord][max_y_coord] = "+"

		if has_x_axis:
			output_buffer[min_x_coord][y_zero_coord] = "+"
			output_buffer[max_x_coord][y_zero_coord] = "+"

		if do_plot_x_label:        
			x_label_y_coord = self.get_adjacent_coord(y_zero_coord, plot_data.y_size)
			if not has_x_axis:
				x_label_y_coord = self.get_adjacent_coord(0, plot_data.y_size)

			self.write_clipped(output_buffer, min_x_coord, x_label_y_coord, min_x_str)
			self.write_clipped(output_buffer, max_x_coord-len(max_x_str)+1, x_label_y_coord, max_x_str)

		if do_plot_y_label:
			self.write_clipped(output_buffer, x_zero_coord, max_y_coord, max_y_str)
			self.write_clipped(output_buffer, x_zero_coord, min_y_coord, min_y_str)
			

		
		
	
	def plot_data(self, xy_seq, output_buffer, plot_data):
		if self.plot_slope:
			xy_seq = list(xy_seq)
			if len(xy_seq) == 1:
				self.plot_point(xy_seq[0][0], xy_seq[0][1], output_buffer, plot_data)
				return
			#sort according to the x coord
			xy_seq.sort(key = lambda c: c[0])
			prev_p = xy_seq[0]
			e_xy_seq = enumerate(xy_seq)
			next(e_xy_seq)
			for i,(x,y) in e_xy_seq:
				draw_symbol = self.dot
				line_drawn = self.plot_line(prev_p, (x,y), output_buffer, plot_data)
				prev_p = (x,y)
				if not line_drawn:
					if i > 0 and i < len(xy_seq)-1:
						px,py = xy_seq[i-1]
						nx,ny = xy_seq[i+1]

						if abs(nx-px) > EPSILON:
							slope = (1.0/plot_data.ratio)*(ny-py)/(nx-px)
							draw_symbol = self.get_symbol_by_slope(slope, draw_symbol)
					self.plot_point(x, y, output_buffer, plot_data, draw_symbol)
		else:
			for x,y in xy_seq:
				self.plot_point(x, y, output_buffer, plot_data)


	def plot_line(self, start, end, output_buffer, plot_data):

		start_coord = self.get_coord(start[0], plot_data.min_x, plot_data.x_step), self.get_coord(start[1], plot_data.min_y, plot_data.y_step)
		end_coord = self.get_coord(end[0], plot_data.min_x, plot_data.x_step), self.get_coord(end[1], plot_data.min_y, plot_data.y_step)

		x0,y0 = start_coord
		x1,y1 = end_coord
		if (x0,y0) == (x1,y1):
			return True    
		
		clipped_line = clip_line(start, end, (plot_data.min_x, plot_data.min_y), (plot_data.max_x, plot_data.max_y))
		if clipped_line != None:
			start,end = clipped_line
		else:
			return False
		start_coord = self.get_coord(start[0], plot_data.min_x, plot_data.x_step), self.get_coord(start[1], plot_data.min_y, plot_data.y_step)
		end_coord = self.get_coord(end[0], plot_data.min_x, plot_data.x_step), self.get_coord(end[1], plot_data.min_y, plot_data.y_step)

		x0,y0 = start_coord
		x1,y1 = end_coord
		if (x0,y0) == (x1,y1):
			return True
		x_zero_coord = self.get_coord(0, plot_data.min_x, plot_data.x_step)
		y_zero_coord = self.get_coord(0, plot_data.min_y, plot_data.y_step)   

		if start[0]-end[0] == 0:
			draw_symbol = "|"
		else:
			slope = (1.0/plot_data.ratio)*(end[1]-start[1])/(end[0]-start[0])
			draw_symbol = self.get_symbol_by_slope(slope, self.dot)
		try:

			delta = x1-x0, y1-y0
			if abs(delta[0])>abs(delta[1]):
				s = sign(delta[0])
				slope = float(delta[1])/delta[0]
				for i in range(0,abs(int(delta[0]))):
					cur_draw_symbol = draw_symbol
					x = i*s
					cur_y = int(y0+slope*x)
					if self.will_draw_axes and cur_y == y_zero_coord and draw_symbol == "-":
						cur_draw_symbol = "="
					output_buffer[x0+x][cur_y] = cur_draw_symbol
				
				
			else:
				s = sign(delta[1])
				slope = float(delta[0])/delta[1]
				for i in range(0,abs(int(delta[1]))):
					y = i*s
					cur_draw_symbol = draw_symbol
					cur_y = y0+y
					if self.will_draw_axes and cur_y == y_zero_coord and draw_symbol == "-":
						cur_draw_symbol = "="
					output_buffer[int(x0+slope*y)][cur_y] = cur_draw_symbol
		except:
			print(start, end)
			print(start_coord, end_coord)
			print(plot_data)
			raise

		return False            
		
		
	def plot_single(self, seq, min_x = None, max_x = None, min_y = None, max_y = None):
		return self.plot_double(range(len(seq)),seq, min_x, max_x, min_y, max_y)
		



	def plot_double(self, x_seq, y_seq, min_x = None, max_x = None, min_y = None, max_y = None):
		x_seq = list(x_seq)
		y_seq = list(y_seq)

		if len(x_seq) != len(y_seq):
			raise ValueError("x_seq and y_seq must have the same length")
		if not x_seq:
			raise ValueError("x_seq and y_seq must not be empty")

		if min_x == None:
			min_x = min(x_seq)
		if max_x == None:
			max_x = max(x_seq)
		if min_y == None:
			min_y = min(y_seq)
		if max_y == None:
			max_y = max(y_seq)

		if max_x == min_x:
			max_x += 1

		if max_y == min_y:
			max_y += 1

		x_mod = (max_x-min_x)*self.x_margin
		y_mod = (max_y-min_y)*self.y_margin
		min_x-=x_mod
		max_x+=x_mod
		min_y-=y_mod
		max_y+=y_mod


		plot_data = self.PlotData(self.x_size, self.y_size, min_x, max_x, min_y, max_y, x_mod, y_mod)

		output_buffer = [[" "]*self.y_size for i in range(self.x_size)]

		if self.will_draw_axes:
			self.draw_axes(output_buffer, plot_data)

		self.plot_data(zip(x_seq, y_seq), output_buffer, plot_data)

		if self.will_plot_labels:
			self.plot_labels(output_buffer, plot_data)

		trans_result = transposed(y_reversed(output_buffer))

		result = self.new_line.join(["".join(row) for row in trans_result])
		return result

	def draw_axes(self, output_buffer, plot_data):
		
		
		draw_x = False
		draw_y = False

		if plot_data.min_x <= 0 and plot_data.max_x > 0:
			draw_y = True
			zero_x = self.get_coord(0, plot_data.min_x, plot_data.x_step)
			for y in range(plot_data.y_size):
				output_buffer[zero_x][y] = "|"
				
		if plot_data.min_y <= 0 and plot_data.max_y > 0:
			draw_x = True
			zero_y = self.get_coord(0, plot_data.min_y, plot_data.y_step)    
			for x in range(plot_data.x_size):
				output_buffer[x][zero_y] = "-"

		if draw_x and draw_y:
			output_buffer[zero_x][zero_y] = "+"
		
		
	@staticmethod
	def get_coord(val, min, step):
		result = int((val - min)/step)
		return result

def clip_line(line_pt_1, line_pt_2, rect_bottom_left, rect_top_right):
	ts = [0.0,1.0]
	if line_pt_1[0] == line_pt_2[0]:
		return ((line_pt_1[0], max(min(line_pt_1[1], line_pt_2[1]), rect_bottom_left[1])),
				(line_pt_1[0], min(max(line_pt_1[1], line_pt_2[1]), rect_top_right[1])))
	if line_pt_1[1] == line_pt_2[1]:
		return ((max(min(line_pt_1[0], line_pt_2[0]), rect_bottom_left[0]), line_pt_1[1]),
				(min(max(line_pt_1[0], line_pt_2[0]), rect_top_right[0]), line_pt_1[1]))

	if ((rect_bottom_left[0] <= line_pt_1[0] and line_pt_1[0] < rect_top_right[0]) and
		(rect_bottom_left[1] <= line_pt_1[1] and line_pt_1[1] < rect_top_right[1]) and
		(rect_bottom_left[0] <= line_pt_2[0] and line_pt_2[0] < rect_top_right[0]) and
		(rect_bottom_left[1] <= line_pt_2[1] and line_pt_2[1] < rect_top_right[1])):
		return line_pt_1, line_pt_2

	ts.append( float(rect_bottom_left[0]-line_pt_1[0])/(line_pt_2[0]-line_pt_1[0]) )
	ts.append( float(rect_top_right[0]-line_pt_1[0])/(line_pt_2[0]-line_pt_1[0]) )
	ts.append( float(rect_bottom_left[1]-line_pt_1[1])/(line_pt_2[1]-line_pt_1[1]) )
	ts.append( float(rect_top_right[1]-line_pt_1[1])/(line_pt_2[1]-line_pt_1[1]) )
	
	ts.sort()
	if ts[2] < 0 or ts[2] >= 1 or ts[3] < 0 or ts[2]>= 1:
		return None
	result = [(pt_1 + t*(pt_2-pt_1)) for t in (ts[2],ts[3]) for (pt_1, pt_2) in zip(line_pt_1, line_pt_2)]
	return (result[0],result[1]), (result[2], result[3])
	


def plot(*args,**flags):
	limit_flags_names = set(["min_x","min_y","max_x","max_y"])
	limit_flags = dict([(n,flags[n]) for n in limit_flags_names & set(flags)])
	settting_flags = dict([(n,flags[n]) for n in set(flags) - limit_flags_names])
	
	if len(args) == 1:
		p = Plotter(**settting_flags)
		print(p.plot_single(args[0],**limit_flags))
	elif len(args) == 2:
		p = Plotter(**settting_flags)
		print(p.plot_double(args[0],args[1],**limit_flags))
	else:
		raise NotImplementedError("can't draw multiple graphs yet")


def parse_number_list(values):
	result = []
	for value in values:
		result.append(float(value))
	return result


def build_cli_parser():
	parser = argparse.ArgumentParser(description="Render numeric data as ASCII plots.")
	parser.add_argument("y", nargs="*", help="Y values for a single-series plot.")
	parser.add_argument("--x", nargs="+", help="Explicit X values for the plot.")
	parser.add_argument("--y-values", nargs="+", dest="y_values", help="Explicit Y values for use with --x.")
	parser.add_argument("--x-size", type=int, default=80, help="Plot width in characters.")
	parser.add_argument("--y-size", type=int, default=20, help="Plot height in characters.")
	parser.add_argument("--dot", default="*", help="Base character for plotted points.")
	parser.add_argument("--no-axes", action="store_true", help="Disable axis rendering.")
	parser.add_argument("--no-labels", action="store_true", help="Disable axis-bound labels.")
	parser.add_argument("--no-slope", action="store_true", help="Disable slope-aware line drawing.")
	parser.add_argument("--x-margin", type=float, default=0.05, help="Horizontal padding ratio.")
	parser.add_argument("--y-margin", type=float, default=0.1, help="Vertical padding ratio.")
	parser.add_argument("--min-x", type=float, dest="min_x", help="Minimum visible X bound.")
	parser.add_argument("--max-x", type=float, dest="max_x", help="Maximum visible X bound.")
	parser.add_argument("--min-y", type=float, dest="min_y", help="Minimum visible Y bound.")
	parser.add_argument("--max-y", type=float, dest="max_y", help="Maximum visible Y bound.")
	return parser


def main(argv=None):
	parser = build_cli_parser()
	args = parser.parse_args(argv)

	plotter = Plotter(
		x_size=args.x_size,
		y_size=args.y_size,
		draw_axes=not args.no_axes,
		plot_labels=not args.no_labels,
		plot_slope=not args.no_slope,
		dot=args.dot,
		x_margin=args.x_margin,
		y_margin=args.y_margin,
	)

	limit_flags = {
		"min_x": args.min_x,
		"max_x": args.max_x,
		"min_y": args.min_y,
		"max_y": args.max_y,
	}
	limit_flags = dict([(name, value) for name, value in limit_flags.items() if value is not None])

	if args.x is not None:
		if args.y:
			parser.error("positional y values cannot be combined with --x; use --y-values with --x")
		if args.y_values is None:
			parser.error("--y-values is required when --x is provided")
		x_values = parse_number_list(args.x)
		y_values = parse_number_list(args.y_values)
		if len(x_values) != len(y_values):
			parser.error("--x and --y-values must contain the same number of values")
		result = plotter.plot_double(x_values, y_values, **limit_flags)
	else:
		if args.y_values is not None:
			parser.error("--y-values requires --x")
		if not args.y:
			parser.error("provide y values, or provide both --x and --y-values")
		y_values = parse_number_list(args.y)
		result = plotter.plot_single(y_values, **limit_flags)

	sys.stdout.write(result)
	sys.stdout.write(plotter.new_line)
	return 0
	
__all__ = ["Plotter","plot","main"]


if __name__ == "__main__":
	raise SystemExit(main())