diff --git a/pystencils/plot.py b/pystencils/plot.py
index de1b1e32ebdf1ffd036ed1f95bb564485c91b7a0..1555b94fb87e86e697c4a615e1986d8ee25eab8f 100644
--- a/pystencils/plot.py
+++ b/pystencils/plot.py
@@ -3,6 +3,7 @@ This module extends the pyplot module with functions to show scalar and vector f
 simulation coordinate system (y-axis goes up), instead of the "image coordinate system" (y axis goes down) that
 matplotlib normally uses.
 """
+import warnings
 from itertools import cycle
 
 from matplotlib.pyplot import *
diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py
index dcd89004fc03eb8d0564f80906380c018bda8820..fa437a5d0fbe49ed04cabeaaf3e22d335b939819 100644
--- a/pystencils/simp/assignment_collection.py
+++ b/pystencils/simp/assignment_collection.py
@@ -5,8 +5,7 @@ import sympy as sp
 
 from pystencils.assignment import Assignment
 from pystencils.simp.simplifications import (
-    sort_assignments_topologically,
-    transform_lhs_and_rhs, transform_rhs)
+    sort_assignments_topologically, transform_lhs_and_rhs, transform_rhs)
 from pystencils.sympyextensions import count_operations, fast_subs