Numpy行列をLatex用に出力

目標

普通にnumpyのndarrayをprintすると

[[ 9.999e+00 -1.000e-03 -1.000e-03]
 [-1.000e-03  9.999e+00 -1.000e-03]
 [-1.000e-03 -1.000e-03  9.999e+00]]

のようになる。
このndarrayをlatexでコンパイルできる形に出力させたい。
つまり、このndarrayを与えれば、

\begin{pmatrix}
    \ 9.999 & -0.001 & -0.001 \\
    -0.001 & \ 9.999 & -0.001 \\
    -0.001 & -0.001 & \ 9.999 \\
\end{pmatrix}

みたいな出力をする関数を作る。こればあれば

ができる

\documentclass[12pt]{article}
\usepackage{amsmath}
\begin{document}
\begin{align*}
\begin{pmatrix}
    \ 9.999 & -0.001 & -0.001 \\
    -0.001 & \ 9.999 & -0.001 \\
    -0.001 & -0.001 & \ 9.999 \\
\end{pmatrix}
\end{align*}
\end{document}

作り方

まず、行列の行数と列の数は、numpy.array.shapeでtupleとして取得できる。
tupleに(行数、列数)として入っているのでこれをそれぞれ変数に入れる。
5. データ構造 — Python 3.8.5 ドキュメント

import numpy as np
matrix = np.ones((3,3))
 _, max_column = matrix.shape

今回は行数は使わないので、使わないという意味を込めて_と言う変数として書いた。

作った関数

def str_2dmatrix_latex(matrix, indent=4, digit=3, env="pmatrix"):
    r"""
    Convert 2d matrix to str for latex.
    *In your latex preset, you must add
    amsmath

    Parameters
    ----------
    matrix : numpy.ndarray
        2d matrix
    indent : int
        indent width
        if indent is -1, indent become tab
    digit : int
        digit=2 0.11
        digit=3 0.111
    env : str
        env is inserted ... in \begin{...}
        and \end{...}

    Return
    ------
    str
    """
    _, max_column = matrix.shape
    if indent == -1:
        str_indent = "\t"
    else:
        str_indent = " " * indent
    list_str = []
    list_str.append(r"\begin{" + f"{env}" + "}\n")
    for row in matrix:
        for j, elem in enumerate(row):
            str_elem = f"{elem:.{digit}f}"
            # format
            if elem > 0:
                str_elem = r"\ " + str_elem
            # write matrix element
            if j == 0:
                list_str.append(str_indent + str_elem + " & ")
            elif j == max_column - 1:
                list_str.append(str_elem + r" \\" + "\n")
            else:
                list_str.append(str_elem + " & ")
    list_str.append(r"\end{" + f"{env}" + "}\n")
    return "".join(list_str)


def main():
    one_mat = np.ones((3, 3))
    """
     1 1 1
     1 1 1
     1 1 1
    """
    id_mat = np.eye(3)
    """
     1 0 0
     0 1 0
     0 0 1
    """
    mat = 10 * id_mat - 0.001 * one_mat
    print(str_2dmatrix_latex(mat))


if __name__ == "__main__":
    main()

実行結果

\begin{pmatrix}
    \ 9.999 & -0.001 & -0.001 \\
    -0.001 & \ 9.999 & -0.001 \\
    -0.001 & -0.001 & \ 9.999 \\
\end{pmatrix}