Missing import

- how could this not have been caught by the tests?!

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 *

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