Skip to content
Snippets Groups Projects
Commit dc2e9457 authored by Valentin Bruch's avatar Valentin Bruch
Browse files

initial commit: README and first reasonable version of diagrams.py

parents
No related branches found
No related tags found
No related merge requests found
*.pyc
### Generate TikZ diagrams in Python
This python script generates TikZ code for diagrammatic representations of Wick contractions between different vertices.
License: GPL (version 2 or higher)
#!/bin/env python3
'''
Module for generating TikZ code for Wick contraction diagrams.
Valentin Bruch, 2020
'''
# TODO: rectangular contraction lines
import sys
class BaseTeX(object):
'''
Base class for objects representing a TikZ node.
This base class is probably not needed in python style programming.
'''
def __init__(self, label='', **style):
self.label = label
self.style = style
self.identifier = ''.join(chr(ord(i)+17) for i in hex(hash(self))[-8:])
def pprint(self):
print(self.label, end='')
def tikz(self, position=''):
'Return TikZ code for drawing self.'
return ''
def left(self):
'Return TikZ node id of leftmost subnode.'
return self.identifier
def right(self):
'Return TikZ node id of rightmost subnode.'
return self.identifier
class Vertex(BaseTeX):
'''
Vertex of a diagram.
Each vertex consists of n circles to which contraction lines are attached.
'''
def __init__(self,
n,
label='',
radius=0.5,
radius_unit='ex',
fill_color='white',
draw_color='black'
):
super().__init__()
self.n = n # rank of the vertex
self.label = label
self.radius = radius
self.radius_unit = radius_unit
self.fill_color = fill_color
self.draw_color = draw_color
def pprint(self):
print(self.label + '(', end='')
for i in range(self.n):
print(i, end='')
print(')', end='')
def getId(self, i):
return self.identifier + str(i)
def left(self):
'TikZ node id of the left most circle of the vertex.'
return self.identifier + '0'
def right(self):
'TikZ node id of the right most circle of the vertex.'
return self.identifier + str(self.n-1)
def tikz(self, position=''):
'Return TikZ code for drawing this vertex.'
if position:
string = r'\coordinate[%s] (%s);'%(position, self.identifier)
position = ', ' + position
else:
string = r'\coordinate (%s) at (0,0);'%self.identifier
for i in range(self.n):
string += '\n'
string += r'\coordinate[xshift=%g%s] (%s%d) at (%s);'%(
(2*i-self.n+1)*self.radius,
self.radius_unit,
self.identifier,
i,
self.identifier,
)
string += '\n'
string += r'\fill[%s, draw=%s, thick] (%s%d) circle[radius=%g%s];'%(
self.fill_color,
self.draw_color,
self.identifier,
i,
self.radius,
self.radius_unit
)
if self.label:
string += '\n'
string += r'\node[above of=%s] {%s};'%(self.identifier, self.label);
return string
class Contraction:
'''
Contraction line of two vertices.
The contraction lines knows its vertices, the index position at the vertex, and its style.
'''
def __init__(self, v1, idx1, v2, idx2, string='-', style='', **kwargs):
self.v1 = v1
self.v2 = v2
self.idx1 = idx1
self.idx2 = idx2
self.string = string
if style:
self.style = style
else:
if string == '-':
self.style = 'thick'
elif string == '=':
self.style = 'double, thick'
elif string == '~':
self.style = 'wiggly, thick'
elif string == 'z':
self.style = 'zigzag, thick'
else:
self.style = ''
self.shape_style = ', '.join('%s=%s'%(key, value) for key, value in kwargs.items())
def pprint(self):
print('%s%d%s%s%d'%(self.v1.label, self.idx1, self.string, self.v2.label, self.idx2), end=' ')
def tikz(self):
'Return TikZ code for this contractions.'
if self.shape_style:
shape_style = ', ' + self.shape_style
else:
shape_style = ''
return r'\draw[%s] (%s) to[out=270, in=270%s] (%s);'%(
self.style,
self.v1.getId(self.idx1),
shape_style,
self.v2.getId(self.idx2),
)
class Interrupt(BaseTeX):
'''
Node interrupting the diagram contour.
This can be, e.g., a node ρ separating the two branches of the Keldysh contour.
'''
def __init__(self, label=''):
super().__init__()
self.label = label
def tikz(self, position=''):
'Return TikZ code for this interruption.'
if position:
position = ', ' + position
return r'\node[fill=white%s] (%s) {%s};'%(position, self.identifier, self.label)
class BaseLine:
'''
Propagation line between two vertices.
This object only knows its line style.
'''
def __init__(self, string='-', *args, **kwargs):
self.string = string
self.factor = 1
try:
self.style = kwargs['style']
except KeyError:
if string == '-':
self.style = 'thick'
elif string == '=':
self.style = 'double, thick'
elif string == '~':
self.style = 'wiggly, thick'
elif string == 'z':
self.style = 'zigzag, thick'
else:
self.style = ''
def pprint(self):
print(self.string, end='')
class Contour:
'''
Contour consisting of vertices and base lines.
All vertices and base lines are stored in a list of elements.
'''
def __init__(self):
self.elements = []
def pprint(self):
for e in self.elements:
e.pprint()
class Diagram:
r'''
Class for diagrams on one or multiple horizontally aligned contour(s) and
Wick contractions.
Diagram has interfaces for reading strings representing diagrams and
saving diagrams to TikZ code.
Example:
>>> # Single contour loop diagram containing three two point vertices.
>>> Diagram('- g12 - g23 - g31 -').print_tikz()
\begin{tikzpicture}[node distance=1.5ex, baseline=(HIIrrGss.base)]
\coordinate (HIIrrGss) at (0,0);
\coordinate[xshift=-0.5ex] (HIIrrGss0) at (HIIrrGss);
\fill[white, draw=black, thick] (HIIrrGss0) circle[radius=0.5ex];
\coordinate[xshift=0.5ex] (HIIrrGss1) at (HIIrrGss);
\fill[white, draw=black, thick] (HIIrrGss1) circle[radius=0.5ex];
\node[above of=HIIrrGss] {$g$};
\coordinate[right of=HIIrrGss1, xshift=1.5em] (HIIrrGEJ);
\coordinate[xshift=-0.5ex] (HIIrrGEJ0) at (HIIrrGEJ);
\fill[white, draw=black, thick] (HIIrrGEJ0) circle[radius=0.5ex];
\coordinate[xshift=0.5ex] (HIIrrGEJ1) at (HIIrrGEJ);
\fill[white, draw=black, thick] (HIIrrGEJ1) circle[radius=0.5ex];
\node[above of=HIIrrGEJ] {$g$};
\coordinate[right of=HIIrrGEJ1, xshift=1.5em] (HIIrrGwE);
\coordinate[xshift=-0.5ex] (HIIrrGwE0) at (HIIrrGwE);
\fill[white, draw=black, thick] (HIIrrGwE0) circle[radius=0.5ex];
\coordinate[xshift=0.5ex] (HIIrrGwE1) at (HIIrrGwE);
\fill[white, draw=black, thick] (HIIrrGwE1) circle[radius=0.5ex];
\node[above of=HIIrrGwE] {$g$};
\begin{pgfonlayer}{background}
\coordinate[left of=HIIrrGss, xshift=-1.5em] (left);
\draw[thick] (left) -- (HIIrrGss);
\draw[thick] (HIIrrGss) -- (HIIrrGEJ);
\draw[thick] (HIIrrGEJ) -- (HIIrrGwE);
\coordinate[right of=HIIrrGwE, xshift=1.5em] (right);
\draw[thick] (HIIrrGwE) -- (right);
\draw[thick] (HIIrrGss1) to[out=270, in=270, min distance=2.5ex] (HIIrrGEJ0);
\draw[thick] (HIIrrGEJ1) to[out=270, in=270, min distance=2.5ex] (HIIrrGwE0);
\draw[thick] (HIIrrGss0) to[out=270, in=270, min distance=2.5ex] (HIIrrGwE1);
\end{pgfonlayer}
\end{tikzpicture}
>>> # Diagram with customized vertice, contraction, base line, and interruption.
>>> d = Diagram()
>>> d.setVertexStyle('V', n=2, label='$V$')
>>> d.setBaselineStyle('-', style='thick')
>>> d.setContractionStyle('ijklmn', style='thick')
>>> d.setInterruptStyle('r', label=r'$\rho$')
>>> # Note that all styles must be defined before creating the diagram from a string:
>>> d.readString('- Vij - Vjk - r - Vik -')
>>> d.exportTikZ('filename.tex', standalone=True)
'''
def __init__(self, string='', sep=1.5, sep_unit='em'):
# Initialize empty lists
self.contractions = []
self.contours = []
self.interruptions = []
# Properties
self.sep = sep
self.sep_unit = sep_unit
# Styles
self.vertex_styles = dict(
τ=dict(n=1, label=r'$\tau$'),
g=dict(n=2, label='$g$'),
)
self.contraction_styles = {
'123456789':{'style':'thick', 'min distance':'2.5ex'},
'ijklmnop': {'style':'wiggly, thick', 'min distance':'2.5ex'},
'abcdef': {'style':'double, thick', 'min distance':'2.5ex'},
}
self.baseline_styes = {
'-':dict(style='thick'),
'=':dict(style='double, thick'),
'~':dict(style='wiggly, thick')
}
self.interrupt_styles = {
'|':dict(label=''),
'ρ':dict(label=r'$\rho$'),
}
if string:
self.readString(string)
def setGlobalVertexStyle(self, **properties):
'Set style property for all vertices.'
for v in self.vertex_styles.values():
v.update(properties)
def setVertexStyle(self, vertex, **style):
'Update style of a specific vertex type.'
self.vertex_styles[vertex].update(style)
if not 'label' in self.vertex_styles[vertex]:
self.vertex_styles[vertex]['label'] = vertex
def setGlobalContractionStyle(self, **properties):
'Set style property for all contractions.'
for c in self.contraction_styles.values():
c.update(properties)
def setContractionStyle(self, indices, **style):
'Update style of a specific contraction type.'
self.contraction_styles[indices].update(style)
def setGlobalInterruptStyle(self, **properties):
'Set style property for all interruptions.'
for i in self.interrupt_styles.values():
i.update(properties)
def setInterruptStyle(self, char, **style):
'Update style of a specific interruption type.'
self.interrupt_styles[char].update(style)
if not 'label' in self.interrupt_styles[char]:
self.interrupt_styles[char]['label'] = char
def setGlobalBaselineStyle(self, **properties):
'Set style property for all base lines.'
for b in self.baseline_styles.values():
b.update(properties)
def setBaselineStyle(self, line, **style):
'Update style of a specific base line type.'
self.baseline_styles[line].update(style)
def clear(self):
'Clear diagram. This does not change the style settings.'
self.contractions = []
self.contours = []
self.interruptions = []
def readString(self, string, ignore=' ', clear=True):
'''
Read diagram from string. The string must consist of the following elements:
- interruptions: characters representing interruptions as saved in self.interrupt_styles
- baselines: characters representing base lines as saved in self.baseline_styles.
The type of the baseline cannot be changed between two vertices.
- vertices consisting of the following sequence:
1. a character representing a vertex of rank n as saved in self.vertex_styles
2. n index characters as contained in the keys of self.contraction_styles.
Each index is contracted with next (or previous) vertex which uses the same index.
- ignored characters (by default only space)
Note that this function should be called after defining all styles.
Styles will not be updated once the diagram is created.
'''
# Clear old data:
if clear:
self.clear()
# Remove ignored characters
for c in ignore:
string = string.replace(c, '')
# Generate list of all valid indices
indices = ''.join(self.contraction_styles.keys())
# Start by an interrupt (possibly empty).
try:
self.interruptions = [Interrupt(**self.interrupt_styles[string[0]])]
string = string[1:]
except KeyError:
self.interruptions = [Interrupt()]
# Map of open indices: every open index is mapped to it origin (vertex, subindex).
vertex_indices = {}
# Generate the first contour.
self.contours.append(Contour())
# Read the diagram by iterating over its characters.
i = 0
while i < len(string):
c = string[i]
# Check if c represents a vertex.
if c in self.vertex_styles:
# Add a new vertex to the current contour.
self.contours[-1].elements.append(Vertex(**self.vertex_styles[c]))
# Read the indices for this vertex.
for j in range(self.vertex_styles[c]['n']):
s = string[i+1+j]
assert s in indices
# Check whether the index is open. In this case create a contraction.
if s in vertex_indices:
# Find the style for this contraction.
for key, value in self.contraction_styles.items():
if s in key:
style = value
break
# Create the contraction.
self.contractions.append(
Contraction(*vertex_indices.pop(s), self.contours[-1].elements[-1], j, **style)
)
else:
# Add the index to the open indices.
vertex_indices[s] = self.contours[-1].elements[-1], j
i += self.vertex_styles[c]['n']
# Check if c represents a base line.
elif c in self.baseline_styes.keys():
try:
# Try to increase the length of an existing base line.
assert type(self.contours[-1].elements[-1]) == BaseLine
self.contours[-1].elements[-1].factor += 1
except:
# Add a base line object to the current contour.
self.contours[-1].elements.append(BaseLine(**self.baseline_styes[c]))
# Check if c represents an interruption.
elif c in self.interrupt_styles.keys():
# Add an interruption and start a new contour.
self.interruptions.append(Interrupt(**self.interrupt_styles[c]))
self.contours.append(Contour())
else:
# Invalid character in diagram.
print('Invalid character in diagram: %s at position %d'%(c, i), file=sys.stderr)
i += 1
# Warn about uncontracted indices.
if vertex_indices:
print('Warning: diagram contains uncontracted indices:', *vertex_indices.keys(), file=sys.stderr)
def pprint(self):
'''
Produce more or less human readable output (without contractions!) for debugging.
'''
for i, contour in enumerate(self.contours):
if self.interruptions[i].label:
self.interruptions[i].pprint()
contour.pprint()
for i in range(len(self.contours), len(self.interruptions)):
if self.interruptions[i].label:
self.interruptions[i].pprint()
print()
for c in self.contractions:
c.pprint()
print()
def exportTikZ(self, fname, standalone=False, **kwargs):
'''
Write TikZ code to file.
If standalone == True is set, produce a compilable LaTeX standalone document.
'''
with open(fname, 'w') as tikz:
if standalone:
print(r'\documentclass{standalone}', file=tikz)
print(r'\usepackage{tikz}', file=tikz)
print(r'\usetikzlibrary{backgrounds, decorations.pathmorphing}', file=tikz)
print(r'\tikzset{wiggly/.style={decorate, decoration=snake}}', file=tikz)
print(r'\tikzset{zigzag/.style={decorate, decoration=zigzag}}', file=tikz)
print(r'\begin{document}', file=tikz)
self.print_tikz(file=tikz, **kwargs)
if standalone:
print(r'\end{document}', file=tikz)
def print_tikz(self, file=sys.stdout, node_distance='1.5ex'):
r'''
Produce TikZ code and write it to file buffer or stdout.
Required packages for TikZ are (depending on which line styles you use):
\usepackage{tikz}
\usetikzlibrary{backgrounds, decorations.pathmorphing}
\tikzset{wiggly/.style={decorate, decoration=snake}}
\tikzset{zigzag/.style={decorate, decoration=zigzag}}
'''
# List of all TikZ nodes on the diagram axis combined with the number of prior base lines.
nodes = []
# List of all base line objects with the indices of their starting node: (baseline, index).
baselines = {}
# Iterate over contours, interruptions and elements to collect all nodes and baselines.
for i, contour in enumerate(self.contours):
try:
if self.interruptions[i].label:
nodes.append(self.interruptions[i])
except:
pass
for e in contour.elements:
if type(e) == BaseLine:
baselines[len(nodes)-1] = e
else:
nodes.append(e)
# Collect trailing interruptions.
for i in range(len(self.contours), len(self.interruptions)):
try:
if self.interruptions[i].label:
nodes.append(self.interruptions[i])
except:
pass
# Begin creating the TikZ picture.
print(r'\begin{tikzpicture}[node distance=%s, baseline=(%s.base)]'%(node_distance, nodes[0].identifier), file=file)
# Create the first node.
try:
print(nodes[0].tikz(), file=file)
except:
print('Error while printing node tikz:', nodes[0], file=sys.stderr)
# Draw vertices and interruptions (including their labels).
for i, n in enumerate(nodes[1:]):
try:
factor = baselines[i].factor or 1
except:
factor = 1
print(n.tikz(position='right of=%s, xshift=%g%s'%(nodes[i].right(), factor*self.sep, self.sep_unit)), file=file)
# Draw the rest on the background.
print(r'\begin{pgfonlayer}{background}', file=file)
# Draw base lines.
for i, b in baselines.items():
if i < 0:
# Base line left of all nodes.
print(r'\coordinate[left of=%s, xshift=-%s%s] (left);'%(nodes[0].identifier, b.factor*self.sep, self.sep_unit), file=file)
print(r'\draw[%s] (left) -- (%s);'%(b.style, nodes[0].identifier), file=file)
elif i > len(nodes)-2:
# Base line right of all nodes.
print(r'\coordinate[right of=%s, xshift=%s%s] (right);'%(nodes[-1].identifier, b.factor*self.sep, self.sep_unit), file=file)
print(r'\draw[%s] (%s) -- (right);'%(b.style, nodes[-1].identifier), file=file)
else:
# Base line between two nodes.
print(r'\draw[%s] (%s) -- (%s);'%(b.style, nodes[i].identifier, nodes[i+1].identifier), file=file)
# Draw all contractions.
for c in self.contractions:
print(c.tikz(), file=file)
# End the TikZ picture
print(r'\end{pgfonlayer}', file=file)
print(r'\end{tikzpicture}', file=file)
if __name__ == '__main__':
print(__doc__)
print(Diagram.__doc__)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment