Building a NumPy-like Library (1)

Purpose

I'm leaving a memo here because the NumPy-like library I was making last year has started working to some extent.

• It should be usable in a way that feels somewhat like NumPy
• Contraction calculations (einsum) should be available
• The core data structure should be multi-dimensional lists (int, float, complex)

These were the main goals, and I matched the behavior by actually using NumPy. I also referred to the NumPy reference when I wanted to check specific interface details.

Repository: https://github.com/derwind/mynumpy

Implementation details

The implementation consists of two main goals.

1. Support for contraction calculations (einsum)
2. Support for basic MPS (Matrix Product State) calculations

v0.1

Implemented a set of methods that are usable in a basic way.

• array
• zeros
• zeros_like
• ndarray
  • __str__
  • __repr__
  • __eq__
  • __ne__
  • __add__
  • __sub__
  • __mul__
  • __matmul__ (with limitations)
  • __truediv__
  • __len__
  • ndim
  • shape
  • size
  • T
  • flatten
  • reshape

And I implemented broadcast (tensor + scalar only).

v0.2

Implemented minimal contraction calculations.

• einsum (2 tensors only)
  • Vector, Matrix -> Vector ("ij,j->i")
  • Matrix, Matrix -> Matrix ("ij,jk->ik")
  • Higher-order tensor, Higher-order tensor -> Higher-order tensor

v0.3

• ones
• ones_like
• einsum (2 tensors only)
  • Trace ("ii->")
  • Inner product ("i,i->")
  • Transpose ("ij->ji")

v0.4

• broadcast improvement

v0.5

Basically finished for now.

• __getitem__ (excluding fancy indexing)
• __setitem__ (excluding fancy indexing)

v0.6

Started expansion for MPS calculations.

• einsum (multiple tensors)
• ndarray
  • dtype
  • real
• Mathematical functions
  • log, log2, log10, logn, sqrt, power, cos, sin, tan, arccos, arcsin, arctan, arctan2, cosh, sinh, tanh, arccosh, arcsinh, arctanh (scalar only)

v0.7

Implemented linear algebra operations for MPS calculations, excluding gates involving complex numbers such as rotation gates.

• eye
• allclose
• ndarray
  • conj
  • copy
  • astype
  • tolist
• linalg
  • norm
  • matrix_rank (real matrices only)
  • svd (real matrices only)

Core Implementation

Traversing Multi-dimensional Lists

The core part of the basic calculations is the logic that recursively traverses multi-dimensional lists to perform processing. Therefore, I repeatedly use the following logic:

Numbers = Union[int, float, complex]


def walk(data: list, func:  Callable[[Numbers], Numbers]):
    if is_number(data[0]):
        for i, d in enumerate(data):
            data[i] = func(d)
        return

    for subdata in data:
        walk(subdata, func)

einsum Implementation

• Painstakingly implemented calculations for two tensors:
  • Prepare a container.
  • Identify the indices to be summed over, retrieve values in order, multiply them, accumulate, and store the result in the container.
• For multiple tensors, the calculation is decomposed into repeated contraction calculations of two tensors at a time as follows:
  • "Aa,aBb,bC->ABC" => "Aa,aBb->ABb" & "ABb,bC->ABC"

linalg Implementation

The core of this part is the implementation of Singular Value Decomposition (SVD). For the SVD implementation of real matrices, I used the one-sided Jacobi method found in Algorithm 4.1 of Jacobi's method is more accurate than QR. Although the paper covers the square matrix case, an explanation including extensions to non-square cases can be found in convexbrain's research notes (Japanese).

The SVD of $G \in \operatorname{Mat}(m,n; \R)$ for $m \geq n$ is obtained as:

$$\begin{align*} U, \Sigma, V = \operatorname{svd}(G) \end{align*}$$

and can be reconstructed as $G = U \Sigma V^T$. Here, $\Sigma$ is a non-negative real diagonal matrix.

In the case where $m < n$, using $G^\prime = G^T$ results in $G^\prime \in \operatorname{Mat}(n,m; \R)$, so for:

$$\begin{align*} U, \Sigma, V = \operatorname{svd}(G^\prime) \end{align*}$$

it can be reconstructed as $G^\prime = U \Sigma V^T$. Therefore, $G = V \Sigma U^T$. In other words, we can simply swap the roles of $U$ and $V$ in the SVD calculation.

Summary

At first, I was quite skeptical about whether I could emulate NumPy, but as I started working on it from the parts I could implement, I found that it was surprisingly doable. While it still feels like it's lacking features, it has enough functionality to perform a certain level of computation.