Add more print names

lettneno · Aug 1, 2022 · 5782ec0 · 5782ec0
1 parent df9c582
commit 5782ec0
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/pybamm/expression_tree/printing/print_name.py b/pybamm/expression_tree/printing/print_name.py
@@ -53,15 +53,24 @@ def prettify_print_name(name):
     if name in PRINT_NAME_OVERRIDES:
         return PRINT_NAME_OVERRIDES[name]
 
-    # Superscripts with comma separated (U_n_ref --> U_{n}^{ref})
-    sup_re = re.search(r"^[\da-zA-Z]+_?((?:init|ref|typ|max|0))_?((?:n|s|p))", name)
+    # Superscripts with comma separated (U_ref_n --> U_{n}^{ref})
+    sup_re = re.search(r"^[\da-zA-Z]+_?((?:init|ref|typ|max|0))_?(.*)", name)
     if sup_re:
         sup_str = (
             r"{" + sup_re.group(2).replace("_", "\,") + r"}^{" + sup_re.group(1) + r"}"
         )
         sup_var = sup_re.group(1) + "_" + sup_re.group(2)
         name = name.replace(sup_var, sup_str)
 
+    # Superscripts with comma separated (U_n_ref --> U_{n}^{ref})
+    sup_re = re.search(r"^[\da-zA-Z]+_?(.*?)_?((?:init|ref|typ|max|0))", name)
+    if sup_re:
+        sup_str = (
+            r"{" + sup_re.group(1).replace("_", "\,") + r"}^{" + sup_re.group(2) + r"}"
+        )
+        sup_var = sup_re.group(1) + "_" + sup_re.group(2)
+        name = name.replace(sup_var, sup_str)
+
     # Subscripts with comma separated (a_R_p --> a_{R\,p})
     sub_re = re.search(r"^a_+(\w+)", name)
     if sub_re: