From 0f06fe9bfcb0ab85403a5b2fecad3434ef8c6c60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Tue, 27 Sep 2022 17:50:48 +0200
Subject: [PATCH] Include the output in the `gathered` environment

---
 aeppl/printing.py      |  9 ++++++---
 tests/test_printing.py | 32 ++++++++++++++++----------------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/aeppl/printing.py b/aeppl/printing.py
index 32355f27..7b188102 100644
--- a/aeppl/printing.py
+++ b/aeppl/printing.py
@@ -561,9 +561,12 @@ def __call__(self, *args, latex_env="equation", latex_label: str = None):
             preamble_lines.append(newline_str.join(comma_str.join(z) for z in v_new))
 
         if preamble_lines and latex_out:
-            preamble_body = newline_str.join(preamble_lines)
-            preamble_str = f"\\begin{{gathered}}\n{textwrap.indent(preamble_body, indent_str)}\n\\end{{gathered}}"
-            res = newline_str.join([preamble_str] + body_strs)
+            preamble_body = newline_str.join(preamble_lines + body_strs)
+            preamble_str = (
+                f"\\begin{{gathered}}\n{textwrap.indent(preamble_body, indent_str)}"
+                f"\n\\end{{gathered}}"
+            )
+            res = newline_str.join([preamble_str])
         else:
             res = newline_str.join(preamble_lines + body_strs)
 
diff --git a/tests/test_printing.py b/tests/test_printing.py
index aeba51f6..2c47d52d 100644
--- a/tests/test_printing.py
+++ b/tests/test_printing.py
@@ -112,9 +112,9 @@ def test_tex_print():
         b \in \mathbb{R}, \,\mu \in \mathbb{R}, \,\sigma \in \mathbb{R}
         \\
         a \sim \operatorname{N}\left(\mu, {\sigma}^{2}\right)\,  \in \mathbb{R}
+        \\
+        (b \odot a)
       \end{gathered}
-      \\
-      (b \odot a)
     \end{equation}
     """
     )
@@ -130,9 +130,9 @@ def test_tex_print():
         b \in \mathbb{R}, \,\mu \in \mathbb{R}, \,\sigma \in \mathbb{R}
         \\
         X \sim \operatorname{N}\left(\mu, {\sigma}^{2}\right)\,  \in \mathbb{R}^{2 \times 1}
+        \\
+        (b \odot X)
       \end{gathered}
-      \\
-      (b \odot X)
     \end{equation}
     """
     )
@@ -157,9 +157,9 @@ def test_tex_print():
         a \sim \operatorname{N}\left(\mu_2, {\sigma_2}^{2}\right)\,  \in \mathbb{R}
         \\
         d \sim \operatorname{N}\left((M \odot a), {\sigma}^{2}\right)\,  \in \mathbb{R}^{N^{d}_{0} \times N^{d}_{1}}
+        \\
+        ((M \odot a) \odot ((b \odot d) + c))
       \end{gathered}
-      \\
-      ((M \odot a) \odot ((b \odot d) + c))
     \end{equation}
     """
     )
@@ -170,9 +170,9 @@ def test_tex_print():
     \begin{equation}
       \begin{gathered}
         b \in \mathbb{Z}, \,c \in \mathbb{Z}, \,M \in \mathbb{R}^{N^{M}_{0} \times N^{M}_{1}}
+        \\
+        M\left[b, \,c\right]
       \end{gathered}
-      \\
-      M\left[b, \,c\right]
     \end{equation}
     """
     )
@@ -187,9 +187,9 @@ def test_tex_print():
     \begin{equation}
       \begin{gathered}
         M \in \mathbb{R}^{N^{M}_{0} \times N^{M}_{1}}
+        \\
+        M\left[1\right]
       \end{gathered}
-      \\
-      M\left[1\right]
     \end{equation}
     """
     )
@@ -200,9 +200,9 @@ def test_tex_print():
     \begin{equation}
       \begin{gathered}
         M \in \mathbb{N}^{N^{M}_{0}}
+        \\
+        M\left[2:4:0\right]
       \end{gathered}
-      \\
-      M\left[2:4:0\right]
     \end{equation}
     """
     )
@@ -220,9 +220,9 @@ def test_tex_print():
         S \sim \operatorname{Gamma^{-1}}\left(0.5, 0.5\right)\,  \in \mathbb{R}
         \\
         Y \sim \operatorname{N}\left(T, {\sqrt{S}}^{2}\right)\,  \in \mathbb{R}
+        \\
+        Y
       \end{gathered}
-      \\
-      Y
     \end{equation}
     """
     )
@@ -242,9 +242,9 @@ def test_tex_print_support_dimension():
         U \sim \operatorname{U}\left(1, 1\right)\,  \in \left[0, 1\right]
         \\
         T \sim \operatorname{C^{+}}\left(1, 1\right)\,  \in \mathbb{R}^{+}
+        \\
+        T
       \end{gathered}
-      \\
-      T
     \end{equation}
     """
     )