Special array representation

src/latexify/codegen/function_codegen.py

@@ -416,15 +416,76 @@ def generate_matrix_from_array(data: list[list[str]]) -> str:
 
         return generate_matrix_from_array(rows)
 
+    def _generate_zeros(self, node: ast.Call) -> str | None:
+        """Generates LaTeX for numpy.zeros.
+
+        Args:
+            node: ast.Call node containing the appropriate method invocation.
+
+        Returns:
+            Generated LaTeX, or None if the node has unsupported syntax.
+        """
+        name = ast_utils.extract_function_name_or_none(node)
+        assert name == "zeros"
+
+        if len(node.args) != 1:
+            return None
+
+        # All args to np.zeros should be numeric.
+        if isinstance(node.args[0], ast.Tuple):
+            dims = [ast_utils.extract_int_or_none(x) for x in node.args[0].elts]
+            if any(x is None for x in dims):
+                return None
+            if not dims:
+                return "0"
+            if len(dims) == 1:
+                dims = [1, dims[0]]
+
+            dims_latex = r" \times ".join(str(x) for x in dims)
+        else:
+            dim = ast_utils.extract_int_or_none(node.args[0])
+            if not isinstance(dim, int):
+                return None
+            # 1 x N array of zeros
+            dims_latex = rf"1 \times {dim}"
+
+        return rf"\mathbf{{0}}^{{{dims_latex}}}"
+
+    def _generate_identity(self, node: ast.Call) -> str | None:
+        """Generates LaTeX for numpy.identity.
+
+        Args:
+            node: ast.Call node containing the appropriate method invocation.
+
+        Returns:
+            Generated LaTeX, or None if the node has unsupported syntax.
+        """
+        name = ast_utils.extract_function_name_or_none(node)
+        assert name == "identity"
+
+        if len(node.args) != 1:
+            return None
+
+        ndims = ast_utils.extract_int_or_none(node.args[0])
+        if ndims is None:
+            return None
+
+        return rf"\mathbf{{I}}_{{{ndims}}}"
+
     def visit_Call(self, node: ast.Call) -> str:
         """Visit a call node."""
         func_name = ast_utils.extract_function_name_or_none(node)
 
         # Special treatments for some functions.
+        # TODO(odashi): Move these functions to some separate utility.
         if func_name in ("fsum", "sum", "prod"):
             special_latex = self._generate_sum_prod(node)
         elif func_name in ("array", "ndarray"):
             special_latex = self._generate_matrix(node)
+        elif func_name == "zeros":
+            special_latex = self._generate_zeros(node)
+        elif func_name == "identity":
+            special_latex = self._generate_identity(node)
         else:
             special_latex = None
 

src/latexify/codegen/function_codegen_test.py

@@ -959,3 +959,58 @@ def test_numpy_array(code: str, latex: str) -> None:
     tree = ast_utils.parse_expr(code)
     assert isinstance(tree, ast.Call)
     assert function_codegen.FunctionCodegen().visit(tree) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("zeros(0)", r"\mathbf{0}^{1 \times 0}"),
+        ("zeros(1)", r"\mathbf{0}^{1 \times 1}"),
+        ("zeros(2)", r"\mathbf{0}^{1 \times 2}"),
+        ("zeros(())", r"0"),
+        ("zeros((0,))", r"\mathbf{0}^{1 \times 0}"),
+        ("zeros((1,))", r"\mathbf{0}^{1 \times 1}"),
+        ("zeros((2,))", r"\mathbf{0}^{1 \times 2}"),
+        ("zeros((0, 0))", r"\mathbf{0}^{0 \times 0}"),
+        ("zeros((1, 1))", r"\mathbf{0}^{1 \times 1}"),
+        ("zeros((2, 3))", r"\mathbf{0}^{2 \times 3}"),
+        ("zeros((0, 0, 0))", r"\mathbf{0}^{0 \times 0 \times 0}"),
+        ("zeros((1, 1, 1))", r"\mathbf{0}^{1 \times 1 \times 1}"),
+        ("zeros((2, 3, 5))", r"\mathbf{0}^{2 \times 3 \times 5}"),
+        # Unsupported
+        ("zeros()", r"\mathrm{zeros} \mathopen{}\left( \mathclose{}\right)"),
+        ("zeros(x)", r"\mathrm{zeros} \mathopen{}\left( x \mathclose{}\right)"),
+        ("zeros(0, x)", r"\mathrm{zeros} \mathopen{}\left( 0, x \mathclose{}\right)"),
+        (
+            "zeros((x,))",
+            r"\mathrm{zeros} \mathopen{}\left("
+            r" \mathopen{}\left( x \mathclose{}\right)"
+            r" \mathclose{}\right)",
+        ),
+    ],
+)
+def test_zeros(code: str, latex: str) -> None:
+    tree = ast_utils.parse_expr(code)
+    assert isinstance(tree, ast.Call)
+    assert function_codegen.FunctionCodegen().visit(tree) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("identity(0)", r"\mathbf{I}_{0}"),
+        ("identity(1)", r"\mathbf{I}_{1}"),
+        ("identity(2)", r"\mathbf{I}_{2}"),
+        # Unsupported
+        ("identity()", r"\mathrm{identity} \mathopen{}\left( \mathclose{}\right)"),
+        ("identity(x)", r"\mathrm{identity} \mathopen{}\left( x \mathclose{}\right)"),
+        (
+            "identity(0, x)",
+            r"\mathrm{identity} \mathopen{}\left( 0, x \mathclose{}\right)",
+        ),
+    ],
+)
+def test_identity(code: str, latex: str) -> None:
+    tree = ast_utils.parse_expr(code)
+    assert isinstance(tree, ast.Call)
+    assert function_codegen.FunctionCodegen().visit(tree) == latex