be careful with buoyancy name

firedrakeproject · Nov 14, 2024 · 3f843ce · 3f843ce
1 parent cc3adef
commit 3f843ce
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py
@@ -450,7 +450,8 @@ def __init__(self, domain, parameters, equivalent_buoyancy=False,
             # Coriolis and transport term from depth equation
             linearisation_map = lambda t: \
                 (any(t.has_label(time_derivative, pressure_gradient, coriolis))
-                 or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport)))
+                 or (t.get(prognostic) in ['D', self.b_name]
+                     and t.has_label(transport)))
 
         super().__init__(domain, parameters,
                          equivalent_buoyancy=equivalent_buoyancy,