add scientific notation

Author: dbatten5

src/arraytex/api.py

@@ -1,4 +1,5 @@
 """Main package API."""
+
 import re
 from typing import Any
 from typing import Optional
@@ -12,40 +13,51 @@ def to_matrix(
     arr: NDArray[Any],
     style: str = "b",
     num_format: Optional[str] = None,
+    scientific_notation: bool = False,
 ) -> str:
     """Convert a numpy.NDArray to LaTeX matrix.
 
     Args:
         arr: the array to be converted
         style: a style formatter string, either "b" for "bmatrix" or "p" for "pmatrix"
         num_format: a number formatter string, e.g. ".2f"
+        scientific_notation: a flag to determine whether 1 x 10^3 should be used,
+            otherwise e-notation is used (1e3)
 
     Returns:
         the string representation of the array
     """
-
-    def num_formatter(x: Union[np.int64, np.float64, np.float32]) -> str:
-        return f"%{num_format}" % x
-
     environment = f"{style}matrix"
 
     formatter = {}
     if num_format:
+
+        def num_formatter(x: Union[np.int64, np.float64, np.float32]) -> str:
+            return f"%{num_format}" % x
+
         formatter.update({"float_kind": num_formatter, "int_kind": num_formatter})
 
     lines = (
-        np.array2string(arr, max_line_width=np.inf, formatter=formatter)  # type: ignore  # noqa
+        np.array2string(
+            arr,
+            max_line_width=np.inf,  # type: ignore  # noqa
+            formatter=formatter,  # type: ignore  # noqa
+            separator=" & ",
+        )
         .replace("[", "")
         .replace("]", "")
+        .replace(" &\n", "\n")
     )
 
     if num_format and "e" in num_format:
         pattern = r"e(-?\d+)"
         replace = r"\\mathrm{e}{\g<1>}"
+        if scientific_notation:
+            replace = r" \\times 10^{\g<1>}"
         lines = re.sub(pattern, replace, lines)
 
     rv = [f"\\begin{{{environment}}}"]
-    rv += ["  " + " & ".join(line.split()) + r" \\" for line in lines.splitlines()]
+    rv += [line.strip() + r" \\" for line in lines.splitlines()]
     rv += [f"\\end{{{environment}}}"]
 
     return "\n".join(rv)

tests/test_api.py

@@ -21,8 +21,8 @@ def test_default(self) -> None:
         assert (
             out
             == r"""\begin{bmatrix}
-  1 & 2 & 3 \\
-  4 & 5 & 6 \\
+1 & 2 & 3 \\
+4 & 5 & 6 \\
 \end{bmatrix}"""
         )
 
@@ -40,8 +40,8 @@ def test_style(self) -> None:
         assert (
             out
             == r"""\begin{pmatrix}
-  1 & 2 & 3 \\
-  4 & 5 & 6 \\
+1 & 2 & 3 \\
+4 & 5 & 6 \\
 \end{pmatrix}"""
         )
 
@@ -59,13 +59,13 @@ def test_decimal_format(self) -> None:
         assert (
             out
             == r"""\begin{bmatrix}
-  1.00 & 2.00 & 3.00 \\
-  4.00 & 5.00 & 6.00 \\
+1.00 & 2.00 & 3.00 \\
+4.00 & 5.00 & 6.00 \\
 \end{bmatrix}"""
         )
 
-    def test_scientific_notation(self) -> None:
-        """Scientific notation can be selected."""
+    def test_e_notation(self) -> None:
+        """Scientific E notation can be selected."""
         mat = np.array(
             [
                 [0.001, 0.002, 0.003],
@@ -78,7 +78,26 @@ def test_scientific_notation(self) -> None:
         assert (
             out
             == r"""\begin{bmatrix}
-  1.00\mathrm{e}{-03} & 2.00\mathrm{e}{-03} & 3.00\mathrm{e}{-03} \\
-  4.00\mathrm{e}{-03} & 5.00\mathrm{e}{-03} & 6.00\mathrm{e}{-03} \\
+1.00\mathrm{e}{-03} & 2.00\mathrm{e}{-03} & 3.00\mathrm{e}{-03} \\
+4.00\mathrm{e}{-03} & 5.00\mathrm{e}{-03} & 6.00\mathrm{e}{-03} \\
+\end{bmatrix}"""
+        )
+
+    def test_scientific_notation(self) -> None:
+        """Scientific  can be selected."""
+        mat = np.array(
+            [
+                [0.001, 0.002, 0.003],
+                [0.004, 0.005, 0.006],
+            ]
+        )
+
+        out = to_matrix(mat, num_format=".2e", scientific_notation=True)
+
+        assert (
+            out
+            == r"""\begin{bmatrix}
+1.00 \times 10^{-03} & 2.00 \times 10^{-03} & 3.00 \times 10^{-03} \\
+4.00 \times 10^{-03} & 5.00 \times 10^{-03} & 6.00 \times 10^{-03} \\
 \end{bmatrix}"""
         )