diff --git a/py_feyntrop.py b/py_feyntrop.py
index 64d9ea8..8d0085c 100644
--- a/py_feyntrop.py
+++ b/py_feyntrop.py
@@ -6,45 +6,57 @@
 import subprocess
 import os
 
-    #=====
-    # Misc 
-    #=====
+# ======
+#  Misc
+# ======
 
 # s[i,j] = (p_i+p_j)^2 is used in replacement rules
-# eps is used as an expansion parameter 
+# eps is used as an expansion parameter
 sp = IndexedBase('sp')
 eps = symbols('eps')
 
-# un-nest a list
+
 def flatten(list):
+    """
+    un-nest a list
+    """
     return [item for sublist in list for item in sublist]
 
-# vertices of an edge list
+
 def vertices(edges):
+    """
+    vertices of an edge list
+    """
     edges = [e[0] for e in edges]
     edges = flatten(edges)
-    return list(set(edges)) 
+    return list(set(edges))
+
 
-# find largest index in momentum_vars to get V_ext = n_particles
-# last line: + 2 = 1 (zero indexing) + 1 (only provide data for V_ext-1 particles in momentum_vars)
 def get_V_ext(momentum_vars):
+    """
+    find largest index in momentum_vars to get V_ext = n_particles
+    last line: + 2 = 1 (zero indexing) + 1 (only provide data for V_ext-1 particles in momentum_vars)
+    """
     indices = [
-        momentum_vars[i][0] . indices for i in range(len(momentum_vars))
+        momentum_vars[i][0].indices for i in range(len(momentum_vars))
     ]
     indices = flatten(indices)
     return max(indices) + 2
 
-    #==================
-    # Formatting result
-    #==================
+# ===================
+#  Formatting result
+# ===================
+
 
-# round result to two significant digits
 def round_res(res_err_list):
+    """
+    round result to two significant digits
+    """
     res = res_err_list[0]
     err = res_err_list[1]
     if err == 0:
         return [str(res), str(err)]
-    nzeros = abs(int( log10(err) )) # number of zeros after decimal pt.
+    nzeros = abs(int(log10(err)))  # number of zeros after decimal pt.
     nround = nzeros + 2
     res = round(res, nround)
     err = round(err, nround)
@@ -52,93 +64,114 @@ def round_res(res_err_list):
     res = f'{res:.{nround}f}'
     return [res, err]
 
-# print each epsilon order with rounded result
+
 def print_res(res):
-    round_re    , round_im    = [] , []
-    lens_re     , lens_im     = [] , []
-    lens_re_err , lens_im_err = [] , []
+    """
+    print each epsilon order with rounded result
+    """
+    round_re, round_im = [], []
+    lens_re, lens_im = [], []
+    lens_re_err, lens_im_err = [], []
     eps_order = len(res)
     for i in range(eps_order):
         #
-        round_re . append(round_res( res[i][0] ))
-        round_im . append(round_res( res[i][1] ))
+        round_re.append(round_res(res[i][0]))
+        round_im.append(round_res(res[i][1]))
         #
-        lens_re     . append(len( round_re[i][0] ))
-        lens_re_err . append(len( round_re[i][1] ))
+        lens_re    .append(len(round_re[i][0]))
+        lens_re_err.append(len(round_re[i][1]))
         #
-        lens_im     . append(len( round_im[i][0] ))
-        lens_im_err . append(len( round_im[i][1] ))
+        lens_im    .append(len(round_im[i][0]))
+        lens_im_err.append(len(round_im[i][1]))
         #
-    max_re = str(max( lens_re ))
-    max_im = str(max( lens_im ))
-    max_re_err = str(max( lens_re_err ))
-    max_im_err = str(max( lens_im_err ))
+    max_re = str(max(lens_re))
+    max_im = str(max(lens_im))
+    max_re_err = str(max(lens_re_err))
+    max_im_err = str(max(lens_im_err))
     print("")
     for i in range(eps_order):
-        re , re_err = round_re[i]
-        im , im_err = round_im[i]
+        re, re_err = round_re[i]
+        im, im_err = round_im[i]
         eps_str = '-- eps^' + str(i) + ': '
         #
         print(f"{eps_str : <5}{'[' : ^1}{re : ^{max_re}}{' +/- ' : ^5}{re_err : ^{max_re_err}}{']' : ^1}{'  +  i * ' : ^9}{'[' : ^1}{im : ^{max_im}}{' +/- ' : ^5}{im_err : ^{max_im_err}}{']' : ^1}")
 
-# prefactor in Symanzik representation
+
 def prefactor(edges, D0, eps_order):
+    """
+    prefactor in Symanzik representation
+    """
     V = len(vertices(edges))
     E = len(edges)
-    nu = [e[1] for e in edges] # propagator powers
+    nu = [e[1] for e in edges]  # propagator powers
     nu_tot = sum(nu)
-    L = E - V + 1 # loops
-    w = nu_tot - L * (D0-2*eps)/2 # superficial degree of divergence
+    L = E - V + 1  # loops
+    w = nu_tot - L * (D0 - 2 * eps) / 2  # superficial degree of divergence
     num = gamma(w)
     den = prod([gamma(nu[i]) for i in range(E)])
-    return num/den
+    return num / den
+
 
-# mulitplying result by prefactor and expanding in epsilon
 def eps_expansion(res, edges, D0):
+    """
+    multiplying result by prefactor and expanding in epsilon
+    """
     eps_order = len(res)
     terms = 0
-    for i in range(eps_order):
-        re = res[i][0][0]
-        im = res[i][1][0]
-        terms += complex(re, im)*eps**i
+    for i, resi in enumerate(res):
+        re = resi[0][0]
+        im = resi[1][0]
+        terms += complex(re, im) * eps**i
     pf = prefactor(edges, D0, eps_order)
     terms = pf * terms
     terms = series(terms, eps, 0, eps_order)
-    return terms . evalf(10)
+    return terms.evalf(10)
+
+# ========================
+#  Prepare kinematic data
+# ========================
 
-    #=======================
-    # Prepare kinematic data
-    #=======================
 
-# convert phase_space_point into a dictionay with it being optional whether 'val' is string or not
 def to_dict(phase_space_point):
+    """
+    convert phase_space_point into a dictionary with it being optional whether 'val' is string or not
+    """
     n = len(phase_space_point)
-    temp = []
+    temp = {}
     for i in range(n):
         var, val = phase_space_point[i][0], phase_space_point[i][1]
-        temp . append( (var, str(val)) )
-    return dict(temp)
+        temp[var] = str(val)
+    return temp
+
 
-# replace all elements from a dictionary for expressions such as 's + t' containing more than one variable
 def replace_all(expr, dic):
-    for i, j in dic . items():
-        expr = expr . replace(i, j)
+    """
+    replace all elements from a dictionary for expressions such as 's + t' containing more than one variable
+    """
+    for i, j in dic.items():
+        expr = expr.replace(i, j)
     return expr
 
-# replace sp[i,j] in replacement_rules with the chosen values in phase_space_point
+
 def replace_sp(replacement_rules, phase_space_point):
+    """
+    replace sp[i,j] in replacement_rules with the chosen values in phase_space_point
+    """
     n = len(replacement_rules)
     phase_space_point = to_dict(phase_space_point)
     temp = []
     for i in range(n):
-        var, val  = replacement_rules[i][0], replacement_rules[i][1]
+        var, val = replacement_rules[i][0], replacement_rules[i][1]
         val = replace_all(str(val), phase_space_point)
         val = eval(val)
-        temp . append( (var, val) )
+        temp.append((var, val))
     return temp
 
-# list of squared masses
+
 def get_m_sqr(edges, phase_space_point):
+    """
+    list of squared masses
+    """
     E = len(edges)
     phase_space_point = to_dict(phase_space_point)
     temp = []
@@ -146,83 +179,93 @@ def get_m_sqr(edges, phase_space_point):
         val = edges[e][2]
         val = replace_all(str(val), phase_space_point)
         val = eval(val)
-        temp . append(val)
+        temp.append(val)
     return temp
 
-# V x V matrix of scalar products
+
 def get_P_uv(edges, phase_space_point, replacement_rules):
-   V = len(vertices(edges)) # total number of vertices (external and internal)
-   P_uv_mat = zeros(V)
-   #
-       # 1-pt function: all scalar products are zero
-   #
-   if replacement_rules == []:
-       return P_uv_mat . tolist()
-   #
-       # 2-pt or higher
-   #
-   V_ext = get_V_ext(replacement_rules) # number of external particles
-   n = V_ext - 1 # to simplify for-loops
-   #
-   # diaonal and upper triangular part
-   #
-   for i in range(n):
-       for j in range(i, n):
-           P_uv_mat[i,j] = sp[i,j]
-   #
-   # copy to upper to lower triangular part
-   #
-   for i in range(n):
-       for j in range(i+1, n):
-           P_uv_mat[j,i] = P_uv_mat[i,j]
-   #
-   # last entry of column = (-1) * everything above 
-   #
-   for j in range(n):
+    """
+    # V x V matrix of scalar products
+    """
+    V = len(vertices(edges))  # total number of vertices (external and internal)
+    P_uv_mat = zeros(V)
+    #
+    # 1-pt function: all scalar products are zero
+    #
+    if not replacement_rules:
+        return P_uv_mat.tolist()
+    #
+    # 2-pt or higher
+    #
+    V_ext = get_V_ext(replacement_rules)  # number of external particles
+    n = V_ext - 1  # to simplify for-loops
+    #
+    # diagonal and upper triangular part
+    #
+    for i in range(n):
+        for j in range(i, n):
+            P_uv_mat[i, j] = sp[i, j]
+    #
+    # copy to upper to lower triangular part
+    #
+    for i in range(n):
+        for j in range(i + 1, n):
+            P_uv_mat[j, i] = P_uv_mat[i, j]
+    #
+    # last entry of column = (-1) * everything above
+    #
+    for j in range(n):
         colsum = 0
         for i in range(n):
-            colsum += P_uv_mat[i,j]
-        P_uv_mat[n,j] = (-1) * colsum
-   #
-   # last entry of row = (-1) * everything to the left
-   #
-   for i in range(V_ext):
+            colsum += P_uv_mat[i, j]
+        P_uv_mat[n, j] = (-1) * colsum
+    #
+    # last entry of row = (-1) * everything to the left
+    #
+    for i in range(V_ext):
         rowsum = 0
         for j in range(n):
-            rowsum += P_uv_mat[i,j]
-        P_uv_mat[i,n] = (-1) * rowsum
-   #
-   # subsitute phase space point
-   #
-   rep = replace_sp(replacement_rules, phase_space_point)
-   return P_uv_mat . subs(rep) . tolist()
-
-    #=====================
-    # Tropical integration
-    #=====================
-
-# returns matrix of scalar products and a list of squared masses
+            rowsum += P_uv_mat[i, j]
+        P_uv_mat[i, n] = (-1) * rowsum
+    #
+    # substitute phase space point
+    #
+    rep = replace_sp(replacement_rules, phase_space_point)
+    return P_uv_mat.subs(rep).tolist()
+
+# ======================
+#  Tropical integration
+# ======================
+
+
 def prepare_kinematic_data(edges, replacement_rules, phase_space_point):
+    """
+    returns matrix of scalar products and a list of squared masses
+    """
     P_uv = get_P_uv(edges, phase_space_point, replacement_rules)
     m_sqr = get_m_sqr(edges, phase_space_point)
     return P_uv, m_sqr
 
-# trop_int is the value of the Feynman integral (without prefactor!)
-# Itr is the normalization factor in the tropical measure
+
 def tropical_integration(N, D0, Lambda, eps_order, edges, replacement_rules, phase_space_point):
+    """
+    trop_int is the value of the Feynman integral (without prefactor!)
+
+    Itr is the normalization factor in the tropical measure
+    """
     P_uv, m_sqr = prepare_kinematic_data(edges, replacement_rules, phase_space_point)
     edges = [e[:-1] for e in edges]
 
-    P_uv = [ [ float(e) for e in row ] for row in P_uv ]
+    P_uv = [[float(e) for e in row] for row in P_uv]
     ft_input = {
-        "graph" : edges,
-        "dimension" : D0,
-        "scalarproducts" : P_uv,
-        "masses_sqr" : m_sqr,
-        "num_eps_terms" : eps_order,
-        "lambda" : Lambda,
-        "N" : N,
-        "seed" : 0
+        "graph": edges,
+        "dimension": D0,
+        "scalarproducts": P_uv,
+        "masses_sqr": m_sqr,
+        "num_eps_terms": eps_order,
+        "lambda": Lambda,
+        "N": N,
+        "seed": 0
     }
 
     json_str = json.dumps(ft_input)
@@ -253,4 +296,3 @@ def tropical_integration(N, D0, Lambda, eps_order, edges, replacement_rules, pha
     print("Prefactor: " + str(prefactor(edges, D0, eps_order)) + ".")
     print_res(trop_res)
     return trop_res, Itr
-
diff --git a/simple_example_2L_3pt.py b/simple_example_2L_3pt.py
index 397c7d6..4868fdc 100644
--- a/simple_example_2L_3pt.py
+++ b/simple_example_2L_3pt.py
@@ -1,34 +1,34 @@
 # Authors: Michael Borinsky & Henrik Munch 2023
 #
 # Email: michael.borinsky@eth-its.ethz.ch
-# 
-# This example script accompanies the publication: 
 #
-# M. Borinsky, H. J. Munch, F. Tellander: 'Tropical Feynman integration in 
-# the Minkowski regime', Computer Physics Communications, 292 (2023), 108874 
+# This example script accompanies the publication:
+#
+# M. Borinsky, H. J. Munch, F. Tellander: 'Tropical Feynman integration in
+# the Minkowski regime', Computer Physics Communications, 292 (2023), 108874
 # arXiv:2302.08955
 
-# Running it in the terminal with 
-# 
+# Running it in the terminal with
+#
 # > python simple_example_2L_3pt.py
 #
 # integrates a 2-loop 3-point Feynman graph.
 #
-# The comments in this file should completely explain the python interface of 
+# The comments in this file should completely explain the python interface of
 # feyntrop.
 
-# The first step is to import the `feyntrop` module. If there is an error 
-# here, make sure the files `feyntrop` and `py_feyntrop.py` are in the working 
+# The first step is to import the `feyntrop` module. If there is an error
+# here, make sure the files `feyntrop` and `py_feyntrop.py` are in the working
 # directory.
 
-from py_feyntrop import *
+from py_feyntrop import sp, tropical_integration, eps_expansion, prepare_kinematic_data
 
-# A Feynman graph is given as a list of edges (propagators). Each edge or 
+# A Feynman graph is given as a list of edges (propagators). Each edge or
 # propagator is put in as a tuple in the format:
 #
 # ((v,u), nu, 'msqr')
-# 
-# where 
+#
+# where
 #
 # * u,v are the vertices that the edge/propagator connects
 # * nu is the edge weight
@@ -38,60 +38,60 @@
 #
 # IMPORTANT:
 #
-# In the Python interface, external vertices have to have smaller numbers 
-# than internal ones. I.e. external vertices *have to be* labeled 0,1,2,... 
-# and internal ones *have to be labeled* labeled n,n+1,... where n is the 
+# In the Python interface, external vertices have to have smaller numbers
+# than internal ones. I.e. external vertices *have to be* labeled 0,1,2,...
+# and internal ones *have to be labeled* labeled n,n+1,... where n is the
 # number of external vertices.
 #
 #
 # The edge weights must be numbers > 0.
 #
 #
-# The squared masses can either be floating point numbers (in quotes) or 
+# The squared masses can either be floating point numbers (in quotes) or
 # strings that represent variables that are later replaced by numbers.
 
-# Here, we integrate the 2-loop 3-point graph from arXiv:2302.08955 (see 
+# Here, we integrate the 2-loop 3-point graph from arXiv:2302.08955 (see
 # the picture on page 21 in the arXiv version). The edge/propagator list is:
 
-edges = [((0,1), 1, 'mm'), ((1,3), 1, 'mm'), ((3,2), 1, 'mm'), 
-         ((2,0), 1, 'mm'), ((0,3), 1, 'mm')]
+edges = [((0, 1), 1, 'mm'), ((1, 3), 1, 'mm'), ((3, 2), 1, 'mm'),
+         ((2, 0), 1, 'mm'), ((0, 3), 1, 'mm')]
 
-# All edge weights are 1. The external vertices are 0, 1 and 2. All 
+# All edge weights are 1. The external vertices are 0, 1 and 2. All
 # propagators have a uniform squared mass m^2 represented by the symbol 'mm'.
 
-# The next step is to fix kinematics. We must fix the three external momenta 
-# p_0, p_1, p_1 in-coming into vertices 0, 1 and 2. This is done by fixing all 
-# the scalar products p_u*p_v for u,v in {0,1}, as the remaining momentum 
+# The next step is to fix kinematics. We must fix the three external momenta
+# p_0, p_1, p_1 in-coming into vertices 0, 1 and 2. This is done by fixing all
+# the scalar products p_u*p_v for u,v in {0,1}, as the remaining momentum
 # coming into vertex 2 is fixed by momentum conservation.
 #
-# feyntrop uses the convention that all momenta are in-coming and that 
+# feyntrop uses the convention that all momenta are in-coming and that
 # Euclidean kinematics give rise to a negative semi-definite matrix sp[u,v].
 # (see the script examples/1L_3pt_conformal.py for a Euclidean example)
 #
-# In the running example, we want to set p_0^2 = p_1^2 = 0, and the remaining 
-# free paramter fixes is the scalar product p_0 * p_1. Due to momentum 
-# conservation this is equivalent to fixing p_2^2, because 
+# In the running example, we want to set p_0^2 = p_1^2 = 0, and the remaining
+# free paramter fixes is the scalar product p_0 * p_1. Due to momentum
+# conservation this is equivalent to fixing p_2^2, because
 # p_2^2 = (-p_0 - p_1)^2 = 2 \cdot p_0 \cdot p_1.
 #
 # In code, we write
 
-replacement_rules = [(sp[0,0], '0'), (sp[1,1], '0'), (sp[0,1], 'pp2/2')]
+replacement_rules = [(sp[0, 0], '0'), (sp[1, 1], '0'), (sp[0, 1], 'pp2/2')]
 
-# sp[u,v] stands for the scalar product p_u * p_v. Implicitly, we define a new 
+# sp[u,v] stands for the scalar product p_u * p_v. Implicitly, we define a new
 # variable here, 'pp2', which represents p_2^2, which is equal to 2 p_1*p_2.
 
-# The two symbolic variables in `edges` and `replacement_rules`, namely 'mm' 
+# The two symbolic variables in `edges` and `replacement_rules`, namely 'mm'
 # and 'pp2', must still get explicit numerical values.
 #
 # Here, we choose $m^2 = 0.2$ and $p_2^2 = 1$:
 
 phase_space_point = [('mm', 0.2), ('pp2', 1)]
 
-# We can print the resulting $V \times V$ dimensional scalar product matrix 
-# $\mathcal{P}^{u,v} = p_u \cdot p_v$ (where $V = |V_\text{ext}| + 
+# We can print the resulting $V \times V$ dimensional scalar product matrix
+# $\mathcal{P}^{u,v} = p_u \cdot p_v$ (where $V = |V_\text{ext}| +
 # |V_\text{int}|$) and the list of squared edge masses as follows:
 
-P_uv, m_sqr_list = prepare_kinematic_data(edges, replacement_rules, 
+P_uv, m_sqr_list = prepare_kinematic_data(edges, replacement_rules,
                                           phase_space_point)
 
 print("P_uv-matrix:", P_uv)
@@ -101,8 +101,8 @@
 # OPTIONAL
 #
 # At Euclidean kinematic points, the matrix `P_uv` is negative semi-definite.
-# At Minkowski kinematic points, it is generally indefinite. If you want to 
-# integrate in the Euclidean regime, you can check the signature of `P_uv` 
+# At Minkowski kinematic points, it is generally indefinite. If you want to
+# integrate in the Euclidean regime, you can check the signature of `P_uv`
 # e.g. using the commands
 #
 # import numpy as np
@@ -121,31 +121,31 @@
 D0 = 2
 eps_order = 5
 Lambda = 7.6
-N = int(1e7) # 1e7 is Python notation for 10^7
+N = int(1e7)  # 1e7 is Python notation for 10^7
 
 # Now, we are ready for the tropical integration:
 
-print("="*80)
+print("=" * 80)
 print("Starting tropical integration")
 
-trop_res, Itr = tropical_integration(N, D0, Lambda, eps_order, edges, 
+trop_res, Itr = tropical_integration(N, D0, Lambda, eps_order, edges,
                                      replacement_rules, phase_space_point)
 
 print("Tropical integration finished")
-print("="*80)
+print("=" * 80)
 
 # The output variable
 #
 # * `trop_res` contains the value of the Feynman integral *without* prefactor.
 # * `Itr` is the normalization factor in the tropical measure.
 #
-# The command also prints statistics or outputs errors if something goes 
+# The command also prints statistics or outputs errors if something goes
 # wrong. E.g., if the input graph has a divergence.
 
-# The following command still translates the output into the epsilon-expansion 
-# *with* the prefactor $\Gamma(\omega)/\Gamma(\nu_1) \cdots \Gamma(\nu_E) = 
-# \Gamma(2\epsilon+3)$ included, where \omega is the superficial degree of 
-# divergence, and $E$ is the number of edges. The output in this convention is 
+# The following command still translates the output into the epsilon-expansion
+# *with* the prefactor $\Gamma(\omega)/\Gamma(\nu_1) \cdots \Gamma(\nu_E) =
+# \Gamma(2\epsilon+3)$ included, where \omega is the superficial degree of
+# divergence, and $E$ is the number of edges. The output in this convention is
 # equal to the usual momentum space representation.
 
 expansion = eps_expansion(trop_res, edges, D0)