diff --git a/femwell/laplace.py b/femwell/laplace.py
index 4052fd34..3ab4fce8 100644
--- a/femwell/laplace.py
+++ b/femwell/laplace.py
@@ -2,7 +2,6 @@
 
 from collections import OrderedDict
 
-import matplotlib
 import numpy as np
 from shapely import LineString, box
 from skfem import (
diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index ae1782ce..3403f6be 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -291,6 +291,16 @@ def plot_component(
         colorbar: bool = False,
         ax: Axes = None,
     ):
+        """Plots a component of the electric or magnetic field.
+
+        Args:
+            field ("E", "H"): Field of interest, can be the electric field or the magnetic field.
+            component ("x", "y", "z", "n", "t"): Component of the field to plot.
+            part ("real", "imag", "abs", callable): Function to use to preprocess the field to be plotted. Defaults to "real".
+            boundaries (bool): Superimpose the mesh boundaries on the plot. Defaults to True.
+            colorbar (bool): Adds a colorbar to the plot. Defaults to False.
+            ax (Axes, optional): Axes onto which the plot is drawn. Defaults to None.
+        """
         from mpl_toolkits.axes_grid1 import make_axes_locatable
 
         if part == "real":
@@ -457,6 +467,7 @@ def plot_intensity(
 @dataclass(frozen=True)
 class Modes:
     modes: List
+    """List of modes"""
 
     def __getitem__(self, idx) -> Mode:
         return self.modes[idx]
@@ -489,6 +500,21 @@ def compute_modes(
     n_guess=None,
     solver="scipy",
 ) -> Modes:
+    """Computes the modes of a waveguide.
+
+    Args:
+        basis_epsilon_r (Basis): Basis on which epsilon_r is defined.
+        epsilon_r (NDArray): Relative permittivity of the waveguide.
+        wavelength (float): Wavelength of the light.
+        mu_r (float, optional): Relative permeability of the waveguide. Defaults to 1.
+        num_modes (int, optional): Number of modes to compute. Defaults to 1.
+        order (int, optional): Order of the basis functions. Defaults to 1.
+        metallic_boundaries (bool, optional): If True, the boundaries are considered to be metallic. Defaults to False.
+        radius (float, optional): Radius of the waveguide. Defaults to np.inf.
+        n_guess (float, optional): Initial guess for the effective index. Defaults to None.
+        solver (str, optional): Solver to use. Defaults to "scipy".
+    """
+
     if solver == "scipy":
         solver = solver_eigen_scipy
     elif solver == "slepc":
@@ -710,6 +736,16 @@ def overlap(w):
 
 
 def plot_mode(basis, mode, plot_vectors=False, colorbar=True, title="E", direction="y"):
+    """Plots a mode.
+
+    Args:
+        basis (Basis): Basis of the mode.
+        mode (np.ndarray): Mode to plot.
+        plot_vectors (bool, optional): If True, plot the normal and tangential components. Defaults to False.
+        colorbar (bool, optional): Adds a colorbar to the plot. Defaults to True.
+        title (str, optional): Title of the plot. Defaults to "E".
+
+    """
     from mpl_toolkits.axes_grid1 import make_axes_locatable
 
     (et, et_basis), (ez, ez_basis) = basis.split(mode)
@@ -775,7 +811,7 @@ def interior_residual(w):
 
     # facet jump
     fbasis = [InteriorFacetBasis(basis.mesh, basis.elem, side=i) for i in [0, 1]]
-    w = {"u" + str(i + 1): fbasis[i].interpolate(u) for i in [0, 1]}
+    w = {f"u{str(i + 1)}": fbasis[i].interpolate(u) for i in [0, 1]}
 
     @Functional
     def edge_jump(w):
diff --git a/femwell/mode_solver_1d.py b/femwell/mode_solver_1d.py
index e16b1ce0..a4125034 100644
--- a/femwell/mode_solver_1d.py
+++ b/femwell/mode_solver_1d.py
@@ -35,12 +35,15 @@ def rhs(u, v, w):
     return inner(u, v)
 
 
-A = lhs.assemble(basis, epsilon=basis_epsilon.interpolate(epsilon))
-B = rhs.assemble(basis)
+if __name__ == "__main__":
 
-lams, xs = solve(
-    *condense(A, B, D=basis.get_dofs()), solver=solver_eigen_scipy_sym(sigma=3.55**2, which="LM")
-)
+    A = lhs.assemble(basis, epsilon=basis_epsilon.interpolate(epsilon))
+    B = rhs.assemble(basis)
+
+    lams, xs = solve(
+        *condense(A, B, D=basis.get_dofs()),
+        solver=solver_eigen_scipy_sym(sigma=3.55**2, which="LM")
+    )
 
-print(np.sqrt(lams))
-basis.plot(xs[:, -1]).show()
+    print(np.sqrt(lams))
+    basis.plot(xs[:, -1]).show()
diff --git a/femwell/mode_solver_2d_periodic.py b/femwell/mode_solver_2d_periodic.py
index a9d4edb7..375b1dcb 100644
--- a/femwell/mode_solver_2d_periodic.py
+++ b/femwell/mode_solver_2d_periodic.py
@@ -64,6 +64,16 @@ def I_phi(phi, k_phi, v, k_v, w):
 
 
 def plot_periodic(k, a, basis_phi, phi, num, ax):
+    """Plot the periodic mode.
+
+    Args:
+        k (float): The wavenumber.
+        a (float): The period.
+        basis_phi (Basis): The basis for the periodic function.
+        phi (np.ndarray): The periodic function.
+        num (int): The number of periods to plot.
+        ax (matplotlib.axes.Axes): The axes to plot on.
+    """
     vminmax = np.max(np.abs(basis_phi.interpolate(phi)))
     for i_plot in range(num):
         phases = basis_phi.project(
diff --git a/femwell/thermal.py b/femwell/thermal.py
index e00cfd51..695e2536 100644
--- a/femwell/thermal.py
+++ b/femwell/thermal.py
@@ -34,6 +34,8 @@ def solve_thermal(
         thermal_conductivity: thermal conductivity in W/m‧K.
         specific_conductivity: specific conductivity in S/m.
         current_densities: current densities flowing through the layer in A.
+        fixed_boundaries: fixed boundaries.
+        order: order of the basis.
 
     Returns:
         basis, temperature profile
diff --git a/femwell/thermal_transient.py b/femwell/thermal_transient.py
index 87fd539c..1cade4b6 100644
--- a/femwell/thermal_transient.py
+++ b/femwell/thermal_transient.py
@@ -23,6 +23,20 @@ def solve_thermal_transient(
     dt,
     steps,
 ):
+    """Thermal transient simulation.
+
+    Args:
+        basis0: Basis of the thermal_conductivity.
+        thermal_conductivity_p0: thermal conductivity in W/m‧K.
+        thermal_diffusivity_p0: thermal diffusivity in m^2/s.
+        specific_conductivity: specific conductivity in S/m.
+        current_densities_0: current densities at time 0.
+        current_densities: current densities as a function of time.
+        fixed_boundaries: fixed boundaries.
+        dt: time step.
+        steps: number of steps.
+
+    """
     basis, temperature = solve_thermal(
         basis0,
         thermal_conductivity_p0,
diff --git a/femwell/utils.py b/femwell/utils.py
index 56458c36..758a2337 100644
--- a/femwell/utils.py
+++ b/femwell/utils.py
@@ -24,16 +24,14 @@ def mpc_symmetric(
 ) -> CondensedSystem:
     """Apply a multipoint constraint on the linear system.
 
-    Parameters
-    ----------
-    A
-    b
-        The linear system to constrain.
-    S
-    M
-    T
-    g
-        The constraint is of the form `x[S] = T @ x[M] + g`.
+    Args:
+        A: spmatrix.
+        b: ndarray.
+        S: ndarray.
+        M: ndarray.
+        T
+        g
+            The constraint is of the form `x[S] = T @ x[M] + g`.
 
     """
     if M is None:
@@ -67,9 +65,9 @@ def mpc_symmetric(
 
     if b.ndim == 1:
         y = np.concatenate((b[U] - A[U][:, S] @ g, b[M] - A[M][:, S] @ g))
+    elif np.any(np.nonzero(g)):
+        raise NotImplementedError("Not yet implemented for g != 0")
     else:
-        if np.any(np.nonzero(g)):
-            raise NotImplementedError("Not yet implemented for g != 0")
         y = bmat(
             [
                 [
@@ -101,7 +99,7 @@ def inside_bbox(bbox):
     given bounding box coordinates.
 
     Args:
-        bbox: 4 element list with [xmin, ymin, xmax, ymax]
+        bbox: 4 element list with [xmin, ymin, xmax, ymax].
     """
 
     def sel_fun(x):
diff --git a/femwell/waveguide.py b/femwell/waveguide.py
index 0d7bd0fc..f440d281 100644
--- a/femwell/waveguide.py
+++ b/femwell/waveguide.py
@@ -13,6 +13,16 @@
 
 
 def mesh_waveguide(filename, wsim, hclad, hbox, wcore, hcore):
+    """Create a mesh for a waveguide.
+
+    Args:
+        filename: path to save the mesh.
+        wsim: width of the simulation window.
+        hclad: height of the cladding.
+        hbox: height of the box.
+        wcore: width of the core.
+        hcore: height of the core.
+    """
     polygons = OrderedDict(
         core=Polygon(
             [