Merge pull request #2764 from pybamm-team/issue-2763-latexify

#2763 start improving latexify

Author: valentinsulzer

CHANGELOG.md

@@ -10,6 +10,7 @@
 
 ## Bug fixes
 
+- Improved `model.latexify()` to have a cleaner and more readable output ([#2764](https://github.com/pybamm-team/PyBaMM/pull/2764))
 - Fixed electrolyte conservation in the case of concentration-dependent transference number ([#2758](https://github.com/pybamm-team/PyBaMM/pull/2758))
 - Fixed `plot_voltage_components` so that the sum of overpotentials is now equal to the voltage ([#2740](https://github.com/pybamm-team/PyBaMM/pull/2740))
 

examples/notebooks/models/latexify.ipynb

@@ -380,16 +380,12 @@ def __init__(self, *children):
                 raise ValueError("Cannot concatenate symbols with different bounds")
         super().__init__(*children, name=name)
 
-        if not any(c._raw_print_name is None for c in children):
-            print_name = intersect(
-                children[0]._raw_print_name, children[1]._raw_print_name
-            )
-            for child in children[2:]:
-                print_name = intersect(print_name, child._raw_print_name)
-            if print_name.endswith("_"):
-                print_name = print_name[:-1]
-        else:
-            print_name = None
+        print_name = intersect(children[0]._raw_print_name, children[1]._raw_print_name)
+        for child in children[2:]:
+            print_name = intersect(print_name, child._raw_print_name)
+        if print_name.endswith("_"):
+            print_name = print_name[:-1]
+
         self.print_name = print_name
 
 

examples/notebooks/models/spm_equations.png

@@ -2,7 +2,6 @@
 # Latexify class
 #
 import copy
-import operator
 import re
 import warnings
 
@@ -55,58 +54,6 @@ def __init__(self, model, filename=None, newline=True):
         self.filename = filename
         self.newline = newline
 
-    def _get_concat_displays(self, node):
-        """
-        Returns all the concatenation nodes by doing a depth first search through the
-        entire equation tree with ranges in front of all nodes.
-        """
-        concat_displays = []
-        dfs_nodes = [node]
-        while dfs_nodes:
-            node = dfs_nodes.pop()
-            if (
-                hasattr(node, "concat_latex")
-                and getattr(node, "print_name", None) is not None
-            ):
-                # Combine list of concatenations with list of ranges
-                concat_geo = map(
-                    operator.add,
-                    node.concat_latex,
-                    self._get_concat_geometry_displays(node),
-                )
-
-                # Add cases and split by new line
-                concat_sym = (
-                    r"\begin{cases}" + r" \\ ".join(concat_geo) + r"\end{cases}"
-                )
-                concat_eqn = sympy.Eq(
-                    sympy.Symbol(node.print_name),
-                    sympy.Symbol(concat_sym),
-                    evaluate=False,
-                )
-                concat_displays.append(concat_eqn)
-            dfs_nodes.extend(node.children)
-
-        # Remove duplicates from the list whilst preserving order
-        return list(dict.fromkeys(concat_displays))
-
-    def _get_concat_geometry_displays(self, var):
-        """Returns a list of min/max ranges of all concatenation nodes in latex."""
-        geo = []
-
-        # Loop through all subdomains for concatenations
-        for domain in var.domain:
-            for var_name, rng in self.model.default_geometry[domain].items():
-                if "min" in rng and "max" in rng:
-                    rng_min = get_rng_min_max_name(rng, "min")
-                    rng_max = get_rng_min_max_name(rng, "max")
-
-                    name = sympy.latex(var_name)
-                    geo_latex = f"& {rng_min} < {name} < {rng_max}"
-                    geo.append(geo_latex)
-
-        return geo
-
     def _get_geometry_displays(self, var):
         """
         Returns min range from the first domain and max range from the last domain of
@@ -136,7 +83,7 @@ def _get_geometry_displays(self, var):
 
         return geo
 
-    def _get_bcs_displays(self, lhs_dr, var):
+    def _get_bcs_displays(self, var):
         """
         Returns a list of boundary condition equations with ranges in front of
         the equations.
@@ -146,27 +93,25 @@ def _get_bcs_displays(self, lhs_dr, var):
 
         if bcs:
             # Take range minimum from the first domain
-            for var_name, rng in self.model.default_geometry[var.domain[0]].items():
-                # Trim name (r_n --> r)
-                name = re.findall(r"(.)_*.*", str(var_name))[0]
-                rng_min = get_rng_min_max_name(rng, "min")
-
-                bcs_left = sympy.latex(bcs["left"][0].to_equation())
-                bcs_left_latex = bcs_left + f"\quad {name} = {rng_min}"
-                bcs_eqn = sympy.Eq(lhs_dr, sympy.Symbol(bcs_left_latex), evaluate=False)
-                bcs_eqn_list.append(bcs_eqn)
+            var_name = list(self.model.default_geometry[var.domain[0]].keys())[0]
+            rng_left = list(self.model.default_geometry[var.domain[0]].values())[0]
+            rng_right = list(self.model.default_geometry[var.domain[-1]].values())[0]
 
-            # Take range maximum from the last domain
-            for var_name, rng in self.model.default_geometry[var.domain[-1]].items():
-                # Trim name (r_n --> r)
-                name = re.findall(r"(.)_*.*", str(var_name))[0]
-                rng_max = get_rng_min_max_name(rng, "max")
+            # Trim name (r_n --> r)
+            var_name = re.findall(r"(.)_*.*", str(var_name))[0]
 
-                bcs_right = sympy.latex(bcs["right"][0].to_equation())
-                bcs_right_latex = bcs_right + f"\quad {name} = {rng_max}"
-                bcs_eqn = sympy.Eq(
-                    lhs_dr, sympy.Symbol(bcs_right_latex), evaluate=False
-                )
+            rng_min = get_rng_min_max_name(rng_left, "min")
+            rng_max = get_rng_min_max_name(rng_right, "max")
+
+            for side, rng in [("left", rng_min), ("right", rng_max)]:
+                bc_value, bc_type = bcs[side]
+                bcs_side = sympy.latex(bc_value.to_equation())
+                bcs_side_latex = bcs_side + f"\\quad  \\text{{at }} {var_name} = {rng}"
+                if bc_type == "Dirichlet":
+                    lhs = sympy.Symbol(var.print_name)
+                else:
+                    lhs = sympy.Symbol(r"\nabla " + var.print_name)
+                bcs_eqn = sympy.Eq(lhs, sympy.Symbol(bcs_side_latex), evaluate=False)
                 bcs_eqn_list.append(bcs_eqn)
 
         return bcs_eqn_list
@@ -214,7 +159,14 @@ def _get_param_var(self, node):
 
         return param_list, var_list
 
-    def latexify(self):
+    def latexify(self, output_variables=None):
+        # Voltage is the default output variable if it exists
+        if output_variables is None:
+            if "Voltage [V]" in self.model.variables:
+                output_variables = ["Voltage [V]"]
+            else:
+                output_variables = []
+
         eqn_list = []
         param_list = []
         var_list = []
@@ -234,22 +186,24 @@ def latexify(self):
                 eqn_list.append(sympy.Symbol(r"\\ \textbf{" + str(var) + "}"))
 
                 # Set lhs derivative
-                lhs = sympy.Derivative(var_symbol, "t")
-                lhs_dr = sympy.Derivative(var_symbol, "r")
+                ddt = sympy.Derivative(var_symbol, "t")
 
                 # Override lhs for algebraic
-                if eqn_type == "algebraic":
+                if eqn_type == "rhs":
+                    lhs = ddt
+                else:
                     lhs = 0
 
                 # Override derivative to partial derivative
-                if len(var.domain) != 0 and var.domain != "current collector":
-                    lhs_dr.force_partial = True
-
-                    if not eqn_type == "algebraic":
-                        lhs.force_partial = True
+                if (
+                    len(var.domain) != 0
+                    and var.domain != "current collector"
+                    and eqn_type == "rhs"
+                ):
+                    lhs.force_partial = True
 
                 # Boundary conditions equations
-                bcs = self._get_bcs_displays(lhs_dr, var)
+                bcs = self._get_bcs_displays(var)
 
                 # Add ranges from geometry in rhs
                 geo = self._get_geometry_displays(var)
@@ -263,14 +217,13 @@ def latexify(self):
                 if not eqn_type == "algebraic":
                     init = self.model.initial_conditions.get(var, None)
                     init_eqn = sympy.Eq(var_symbol, init.to_equation(), evaluate=False)
-                    init_eqn = sympy.Symbol(sympy.latex(init_eqn) + r"\quad at\; t=0")
+                    init_eqn = sympy.Symbol(
+                        sympy.latex(init_eqn) + r"\quad \text{at}\; t=0"
+                    )
 
                 # Make equation from lhs and rhs
                 lhs_rhs = sympy.Eq(lhs, rhs, evaluate=False)
 
-                # Get all concatenation nodes
-                concat_displays = self._get_concat_displays(eqn)
-
                 # Set SymPy's init printing to use CustomPrint from sympy_overrides.py
                 sympy.init_printing(
                     use_latex=True,
@@ -282,10 +235,6 @@ def latexify(self):
                 # Add model equations to the list
                 eqn_list.append(lhs_rhs)
 
-                # Add concatenation to the list
-                if concat_displays:
-                    eqn_list.append(concat_displays)
-
                 # Add initial conditions to the list
                 if not eqn_type == "algebraic":
                     eqn_list.extend([init_eqn])
@@ -298,23 +247,21 @@ def latexify(self):
                 param_list.extend(list1)
                 var_list.extend(list2)
 
-        # Add voltage expression to the list
-        if "Voltage [V]" in self.model.variables:
-            voltage = self.model.variables["Voltage [V]"].to_equation()
-            voltage_eqn = sympy.Eq(sympy.Symbol("V"), voltage, evaluate=False)
-            # Add voltage to the list
-            eqn_list.append(sympy.Symbol(r"\\ \textbf{Voltage [V]}"))
-            eqn_list.extend([voltage_eqn])
+        # Add output variables to the list
+        for var_name in output_variables:
+            var = self.model.variables[var_name].to_equation()
+            var_eqn = sympy.Eq(sympy.Symbol("V"), var, evaluate=False)
+            # Add var to the list
+            eqn_list.append(sympy.Symbol(r"\\ \textbf{" + var_name + "}"))
+            eqn_list.extend([var_eqn])
 
         # Remove duplicates from the list whilst preserving order
         param_list = list(dict.fromkeys(param_list))
         var_list = list(dict.fromkeys(var_list))
         # Add Parameters and Variables to the list
         eqn_list.append(sympy.Symbol(r"\\ \textbf{Parameters and Variables}"))
-        # Add parameters to the list
-        eqn_list.extend(param_list)
-        # Add names to the list
         eqn_list.extend(var_list)
+        eqn_list.extend(param_list)
 
         # Split list with new lines
         eqn_new_line = sympy.Symbol(r"\\\\".join(map(custom_print_func, eqn_list)))

pybamm/expression_tree/concatenations.py

@@ -1,14 +1,11 @@
 #
 # Prettify print_name
 #
-import re
-
 PRINT_NAME_OVERRIDES = {
     "current_with_time": "I",
-    "eps_c_e": r"\epsilon{c_e}",
-    "thermodynamic_factor": r"1+\frac{dlnf}{dlnc}",
-    "negative_particle_concentration_scale": r"c_{n}^{max}",
-    "positive_particle_concentration_scale": r"c_{p}^{max}",
+    "current_density_with_time": r"i_{\mathrm{cell}}",
+    "thermodynamic_factor": r"\left(1+\frac{dlnf}{dlnc}\right)",
+    "t_plus": r"t_{\mathrm{+}}",
 }
 
 GREEK_LETTERS = [
@@ -45,58 +42,64 @@ def prettify_print_name(name):
     if name is None or "{" in name or "\\" in name:
         return name
 
-    # Return print_name if name exists in the dictionary
-    if name in PRINT_NAME_OVERRIDES:
-        return PRINT_NAME_OVERRIDES[name]
-
-    # Superscripts with comma separated (U_ref_n --> U_{n}^{ref})
-    sup_re1 = re.search(r"^[\da-zA-Z]+_?((?:init|ref|typ|max|0))_?(.*)", name)
-    if sup_re1:
-        sup_str = (
-            r"{"
-            + sup_re1.group(2).replace("_", "\,")
-            + r"}^{"
-            + sup_re1.group(1)
-            + r"}"
-        )
-        sup_var = sup_re1.group(1) + "_" + sup_re1.group(2)
-        name = name.replace(sup_var, sup_str)
+    # Find subscripts, superscripts, and averaging
+    # Remove them from the name one by one and add them later in processed form
+    subscripts = []
+    superscripts = []
+    average = False
 
-    # Superscripts with comma separated (U_n_ref --> U_{n}^{ref})
-    sup_re2 = re.search(r"^[\da-zA-Z]+_?(.*?)_?((?:init|ref|typ|max|0))", name)
-    if sup_re2:
-        sup_str = (
-            r"{"
-            + sup_re2.group(1).replace("_", "\,")
-            + r"}^{"
-            + sup_re2.group(2)
-            + r"}"
-        )
-        sup_var = sup_re2.group(1) + "_" + sup_re2.group(2)
-        name = name.replace(sup_var, sup_str)
+    processing = True
+    while processing:
+        # Set processing to False. If any of the following conditions are met,
+        # it will be set to True again
+        processing = False
+        for superscript in ["init", "ref", "typ", "max", "0", "surf"]:
+            if f"_{superscript}_" in name or name.endswith(f"_{superscript}"):
+                superscripts.append(superscript)
+                name = name.replace(f"_{superscript}", "")
+                processing = True
+                break
+        # "0" might also appear without a preceding underscore
+        for superscript in ["0"]:
+            if superscript in name:
+                superscripts.append(superscript)
+                name = name.replace(superscript, "")
+                processing = True
+                break
+        for subscript in ["cc", "dl", "R", "e", "s", "n", "p", "amb"]:
+            if f"_{subscript}_" in name or name.endswith(f"_{subscript}"):
+                subscripts.append(subscript)
+                name = name.replace(f"_{subscript}", "")
+                processing = True
+                break
+        for av in ["av", "xav"]:
+            if f"_{av}_" in name or name.endswith(f"_{av}"):
+                average = True
+                name = name.replace(f"_{av}", "")
+                processing = True
+                break
 
-    # Subscripts with comma separated (a_R_p --> a_{R\,p})
-    sub_re = re.search(r"^a_+(\w+)", name)
-    if sub_re:
-        sub_str = r"{" + sub_re.group(1).replace("_", "\,") + r"}"
-        name = name.replace(sub_re.group(1), sub_str)
-
-    # Bar with comma separated (c_s_n_xav --> \bar{c}_{s\,n})
-    bar_re = re.search(r"^([a-zA-Z]+)_*(\w*?)_(?:av|xav)", name)
-    if bar_re:
-        name = (
-            r"\bar{"
-            + bar_re.group(1)
-            + r"}_{"
-            + bar_re.group(2).replace("_", "\,")
-            + r"}"
-        )
+    # Process name
+    # Override print_name if name exists in the dictionary
+    if name in PRINT_NAME_OVERRIDES:
+        name = PRINT_NAME_OVERRIDES[name]
 
-    # Replace eps with epsilon (eps_n --> epsilon_n)
-    name = re.sub(r"(eps)(?![0-9a-zA-Z])", "epsilon", name)
+    # Replace eps with epsilon (e.g. eps_n --> epsilon_n)
+    if name == "eps":
+        name = r"\epsilon"
+    if name == "eps_c":
+        name = r"(\epsilon c)"
 
     # Greek letters (delta --> \delta)
-    greek_re = r"(?<!\\)(" + "|".join(GREEK_LETTERS) + r")(?![0-9a-zA-Z])"
-    name = re.sub(greek_re, r"\\\1", name, flags=re.IGNORECASE)
+    if name.lower() in GREEK_LETTERS:
+        name = "\\" + name
+
+    # Add subscripts and superscripts
+    if average:
+        name = r"\overline{" + name + "}"
+    if subscripts:
+        name += r"_{\mathrm{" + ",".join(subscripts) + "}}"
+    if superscripts:
+        name += r"^{\mathrm{" + ",".join(superscripts) + "}}"
 
     return name