diff --git a/maxwellian_equilibrium.py b/maxwellian_equilibrium.py
index e5d958a5b67c56e895ebd71b8eb5e3c4e87c38fa..25a57796ee0b11bb200a7e71bb9ba5805b0ca74b 100644
--- a/maxwellian_equilibrium.py
+++ b/maxwellian_equilibrium.py
@@ -15,7 +15,7 @@ from pystencils.cache import disk_cache
 def get_weights(stencil, c_s_sq):
     q = len(stencil)
 
-    if c_s_sq != sp.Rational(1, 3):
+    if c_s_sq != sp.Rational(1, 3) and c_s_sq != sp.Symbol("c_s") ** 2:
         warnings.warn("Weights of discrete equilibrium are only valid if c_s^2 = 1/3")
 
     def weight_for_direction(direction):
@@ -55,12 +55,13 @@ def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=tuple(sp.sy
     """
     Returns the common discrete LBM equilibrium as a list of sympy expressions
 
-    :param stencil: tuple of directions
-    :param rho: sympy symbol for the density
-    :param u: symbols for macroscopic velocity, only the first 'dim' entries are used
-    :param order: highest order of velocity terms (for hydrodynamics order 2 is sufficient)
-    :param c_s_sq: square of speed of sound
-    :param compressible: compressibility
+    Args:
+        stencil: tuple of directions
+        rho: sympy symbol for the density
+        u: symbols for macroscopic velocity, only the first 'dim' entries are used
+        order: highest order of velocity terms (for hydrodynamics order 2 is sufficient)
+        c_s_sq: square of speed of sound
+        compressible: compressibility
     """
     weights = get_weights(stencil, c_s_sq)
     assert len(stencil) == len(weights)
diff --git a/quadratic_equilibrium_construction.py b/quadratic_equilibrium_construction.py
index 1be1b008b40f46264516a96b1dd0072109226189..d4159476b7440d2c2a37ff8a8944719e63d2dcda 100644
--- a/quadratic_equilibrium_construction.py
+++ b/quadratic_equilibrium_construction.py
@@ -48,7 +48,8 @@ def match_generic_equilibrium_ansatz(stencil, equilibrium, u=sp.symbols("u_:3"))
           >>> from lbmpy.stencils import get_stencil
           >>> from lbmpy.maxwellian_equilibrium import discrete_maxwellian_equilibrium
           >>> stencil = get_stencil("D2Q9")
-          >>> result = match_generic_equilibrium_ansatz(get_stencil('D2Q9'), discrete_maxwellian_equilibrium(stencil) )
+          >>> eq = discrete_maxwellian_equilibrium(stencil)
+          >>> result = match_generic_equilibrium_ansatz(stencil, eq)
           >>> result[sp.Symbol('A_0')]
           4*rho/9
           >>> result[sp.Symbol('B_1')]