Add plotting library, simple ODE plot
This commit is contained in:
commit
98ed13132b
6 changed files with 416 additions and 0 deletions
288
lib/plot.py
Normal file
288
lib/plot.py
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
#!/usr/bin/env python
|
||||
|
||||
import cairo
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Paper:
|
||||
def __init__(self, width, height, margins=None):
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.margins = {
|
||||
'top': 0,
|
||||
'right': 0,
|
||||
'bottom': 0,
|
||||
'left': 0
|
||||
}
|
||||
if isinstance(margins, int) or isinstance(margins, float):
|
||||
self.margins['top'] = margins
|
||||
self.margins['right'] = margins
|
||||
self.margins['bottom'] = margins
|
||||
self.margins['left'] = margins
|
||||
elif isinstance(margins, list) and (isinstance(margins[0], int)
|
||||
or isinstance(margins[0], float)):
|
||||
if len(margins) == 2:
|
||||
self.margins['top'] = margins[0]
|
||||
self.margins['bottom'] = margins[0]
|
||||
self.margins['left'] = margins[1]
|
||||
self.margins['right'] = margins[1]
|
||||
elif len(margins) == 3:
|
||||
self.margins['top'] = margins[0]
|
||||
self.margins['left'] = margins[1]
|
||||
self.margins['right'] = margins[1]
|
||||
self.margins['bottom'] = margins[2]
|
||||
elif len(margins) == 4:
|
||||
self.margins['top'] = margins[0]
|
||||
self.margins['right'] = margins[1]
|
||||
self.margins['bottom'] = margins[2]
|
||||
self.margins['left'] = margins[3]
|
||||
else:
|
||||
self.margins = margins
|
||||
|
||||
@property
|
||||
def content_width(self):
|
||||
return self.width - self.margins['left'] - self.margins['right']
|
||||
|
||||
@property
|
||||
def content_height(self):
|
||||
return self.height - self.margins['top'] - self.margins['bottom']
|
||||
|
||||
def size(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return np.array([self.content_width, self.content_height])
|
||||
else:
|
||||
return np.array([self.width, self.height])
|
||||
|
||||
def is_within(self, point, honour_margins=True):
|
||||
return ((self.left(honour_margins) < point[0])
|
||||
and (point[0] < self.right(honour_margins))
|
||||
and (self.top(honour_margins) < point[1])
|
||||
and (point[1] < self.bottom(honour_margins)))
|
||||
|
||||
def top_left(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return np.array([self.margins['left'],
|
||||
self.height - self.margins['top']])
|
||||
else:
|
||||
return np.array([0, self.height])
|
||||
|
||||
def bottom_left(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return np.array([self.margins['left'],
|
||||
self.margins['bottom']])
|
||||
else:
|
||||
return np.array([0, 0])
|
||||
|
||||
def top_right(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return np.array([self.width - self.margins['right'],
|
||||
self.height - self.margins['top']])
|
||||
else:
|
||||
return np.array([self.width, self.height])
|
||||
|
||||
def bottom_right(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return np.array([self.width - self.margins['right'],
|
||||
self.margins['bottom']])
|
||||
else:
|
||||
return np.array([self.width, 0])
|
||||
|
||||
def centre(self, honour_margins=True):
|
||||
return (self.bottom_left(honour_margins)
|
||||
+ self.top_right(honour_margins)) / 2
|
||||
|
||||
def top(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return self.margins['top']
|
||||
else:
|
||||
return 0
|
||||
|
||||
def left(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return self.margins['left']
|
||||
else:
|
||||
return 0
|
||||
|
||||
def right(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return self.width - self.margins['right']
|
||||
else:
|
||||
return self.width
|
||||
|
||||
def bottom(self, honour_margins=True):
|
||||
if honour_margins:
|
||||
return self.height - self.margins['bottom']
|
||||
else:
|
||||
return self.height
|
||||
|
||||
|
||||
A4_PORTRAIT = Paper(210, 297, 20)
|
||||
A6_PORTRAIT = Paper(105, 148, 8)
|
||||
|
||||
|
||||
class MultiPlotter:
|
||||
def __init__(self):
|
||||
self.plotters = []
|
||||
|
||||
def move_to(self, point):
|
||||
for plotter in self.plotters:
|
||||
plotter.move_to(point)
|
||||
|
||||
def line_to(self, point):
|
||||
for plotter in self.plotters:
|
||||
plotter.line_to(point)
|
||||
|
||||
def add_layer(self, *args, **kwargs):
|
||||
for plotter in self.plotters:
|
||||
plotter.add_layer(*args, **kwargs)
|
||||
|
||||
def finalise(self):
|
||||
for plotter in self.plotters:
|
||||
plotter.finalise()
|
||||
|
||||
def register_plotter(self, plotter):
|
||||
self.plotters.append(plotter)
|
||||
|
||||
|
||||
class Plotter:
|
||||
|
||||
def __init__(self, paper, line_width=0, colour=[0, 0, 0, 1]):
|
||||
self.paper = paper
|
||||
self.pen_down = False
|
||||
self.position = np.array([0, 0])
|
||||
|
||||
self.layers = []
|
||||
self.add_layer(colour, line_width)
|
||||
|
||||
def add_layer(self, colour=[0, 0, 0, 1], line_width=0.5, name='',
|
||||
switch=True):
|
||||
self.layers.append({
|
||||
'line_width': line_width,
|
||||
'colour': colour,
|
||||
'name': name
|
||||
})
|
||||
if switch:
|
||||
self.switch_layer(len(self.layers) - 1)
|
||||
return len(self.layers) - 1
|
||||
|
||||
def switch_layer(self, layer):
|
||||
self.current_layer = layer
|
||||
|
||||
def move_to(self, point):
|
||||
self.pen_down = False
|
||||
self.position = point
|
||||
|
||||
def line_to(self, point):
|
||||
self.pen_down = True
|
||||
self.position = point
|
||||
|
||||
|
||||
class SVGPlotter(Plotter):
|
||||
|
||||
def __init__(self, file_name, paper, line_width=0.5, colour=[0, 0, 0, 1]):
|
||||
dpi = 72
|
||||
self.file_name = file_name
|
||||
self.surface = cairo.SVGSurface(file_name,
|
||||
paper.width / 25.4 * dpi,
|
||||
paper.height / 25.4 * dpi)
|
||||
self.context = cairo.Context(self.surface)
|
||||
self.context.scale(dpi / 25.4, dpi / 25.4)
|
||||
self.context.set_line_cap(cairo.LINE_CAP_ROUND)
|
||||
|
||||
super().__init__(paper, line_width, colour)
|
||||
|
||||
def move_to(self, point):
|
||||
if self.pen_down:
|
||||
self.context.stroke()
|
||||
self.context.move_to(point[0], point[1])
|
||||
super().move_to(point)
|
||||
|
||||
def line_to(self, point):
|
||||
self.context.line_to(point[0], point[1])
|
||||
super().line_to(point)
|
||||
|
||||
def finalise(self):
|
||||
if self.pen_down:
|
||||
self.context.stroke()
|
||||
self.pen_down = False
|
||||
self.surface.finish()
|
||||
|
||||
def switch_layer(self, layer):
|
||||
if self.pen_down:
|
||||
self.context.stroke()
|
||||
|
||||
super().switch_layer(layer)
|
||||
self.context.set_source_rgba(
|
||||
*self.layers[self.current_layer]['colour'])
|
||||
self.context.set_line_width(
|
||||
self.layers[self.current_layer]['line_width'])
|
||||
|
||||
return len(self.layers) - 1
|
||||
|
||||
|
||||
class HPGLPlotter(Plotter):
|
||||
|
||||
def __init__(self, paper, file_name_pattern):
|
||||
self.file_name_pattern = file_name_pattern
|
||||
self.resolution = (4000 / 81.5, 4000 / 80.5)
|
||||
self.file = None
|
||||
super().__init__(paper)
|
||||
|
||||
def _layer_file_name(self, index):
|
||||
return self.file_name_pattern.format(
|
||||
index=index,
|
||||
name=self.layers[index]['name'],
|
||||
colour=self.layers[index]['colour']
|
||||
)
|
||||
|
||||
def _mm_to_plotter(self, mm):
|
||||
p = (mm[0] * self.resolution[0],
|
||||
(self.paper.height - mm[1]) * self.resolution[1])
|
||||
return f'{p[0]:.0f},{p[1]:.0f}'
|
||||
|
||||
def move_to(self, point):
|
||||
if self.pen_down:
|
||||
self.file.write(';PU')
|
||||
else:
|
||||
self.file.write(',')
|
||||
self.file.write(self._mm_to_plotter(point))
|
||||
return super().move_to(point)
|
||||
|
||||
def line_to(self, point):
|
||||
if not self.pen_down:
|
||||
self.file.write(';PD')
|
||||
else:
|
||||
self.file.write(',')
|
||||
self.file.write(self._mm_to_plotter(point))
|
||||
return super().line_to(point)
|
||||
|
||||
def add_layer(self, colour='black', line_width=0.5, name='', switch=True):
|
||||
new_index = super().add_layer(colour=colour, line_width=line_width,
|
||||
name=name, switch=False)
|
||||
layer_file = open(self._layer_file_name(new_index), 'wt')
|
||||
|
||||
layer_file.write('IN;PU0,0')
|
||||
|
||||
if switch:
|
||||
if self.file:
|
||||
self.file.close()
|
||||
self.file = layer_file
|
||||
self.current_layer = len(self.layers) - 1
|
||||
else:
|
||||
layer_file.close()
|
||||
|
||||
return len(self.layers) - 1
|
||||
|
||||
def switch_layer(self, layer):
|
||||
if self.pen_down:
|
||||
self.file.write(';PU')
|
||||
self.pen_down = False
|
||||
self.file.close()
|
||||
super().switch_layer(layer)
|
||||
self.file = open(self._layer_file_name(self.current_layer), 'at')
|
||||
|
||||
def finalise(self):
|
||||
self.file.close()
|
||||
for i in range(len(self.layers)):
|
||||
with open(self._layer_file_name(i), 'at') as file:
|
||||
file.write(';PU0,0;IN')
|
||||
Loading…
Add table
Add a link
Reference in a new issue