Update matrix operation

Author: itsubaki

math/matrix/matrix.go

@@ -44,9 +44,9 @@ func Zero(rows, cols int) Matrix {
 }
 
 // Identity returns an identity matrix.
-func Identity(rows, cols int) Matrix {
-	m := Zero(rows, cols)
-	for i := range rows {
+func Identity(size int) Matrix {
+	m := Zero(size, size)
+	for i := range size {
 		m.Set(i, i, 1)
 	}
 
@@ -68,6 +68,16 @@ func (m Matrix) Set(i, j int, v complex128) {
 	m.Data[i*m.Cols+j] = v
 }
 
+// AddAt adds a value of matrix at (i,j).
+func (m Matrix) AddAt(i, j int, v complex128) {
+	m.Data[i*m.Cols+j] += v
+}
+
+// MulAt multiplies a value of matrix at (i,j).
+func (m Matrix) MulAt(i, j int, v complex128) {
+	m.Data[i*m.Cols+j] *= v
+}
+
 // Seq2 returns a sequence of rows.
 func (m Matrix) Seq2() iter.Seq2[int, []complex128] {
 	return func(yield func(int, []complex128) bool) {
@@ -169,7 +179,7 @@ func (m Matrix) IsUnitary(eps ...float64) bool {
 	}
 
 	mmd := m.Apply(m.Dagger())
-	id := Identity(m.Dimension())
+	id := Identity(m.Rows)
 	return mmd.Equals(id, epsilon.E13(eps...))
 }
 
@@ -181,13 +191,11 @@ func (m Matrix) Apply(n Matrix) Matrix {
 
 	out := Zero(a, p)
 	for i := range a {
-		for j := range p {
-			var c complex128
-			for k := range b {
-				c = c + n.At(i, k)*m.At(k, j)
+		for k := range b {
+			nik := n.At(i, k)
+			for j := range p {
+				out.AddAt(i, j, nik*m.At(k, j))
 			}
-
-			out.Set(i, j, c)
 		}
 	}
 
@@ -276,12 +284,12 @@ func (m Matrix) Inverse() Matrix {
 	p, q := m.Dimension()
 	mm := m.Clone()
 
-	out := Identity(p, q)
+	out := Identity(p)
 	for i := range p {
 		c := 1 / mm.At(i, i)
 		for j := range q {
-			mm.Set(i, j, c*mm.At(i, j))
-			out.Set(i, j, c*out.At(i, j))
+			mm.MulAt(i, j, c)
+			out.MulAt(i, j, c)
 		}
 
 		for j := range q {
@@ -291,8 +299,8 @@ func (m Matrix) Inverse() Matrix {
 
 			c := mm.At(j, i)
 			for k := range q {
-				mm.Set(j, k, mm.At(j, k)-c*mm.At(i, k))
-				out.Set(j, k, out.At(j, k)-c*out.At(i, k))
+				mm.AddAt(j, k, -c*mm.At(i, k))
+				out.AddAt(j, k, -c*out.At(i, k))
 			}
 		}
 	}
@@ -306,14 +314,14 @@ func (m Matrix) TensorProduct(n Matrix) Matrix {
 	a, b := n.Dimension()
 	rows, cols := p*a, q*b
 
-	var idx int
 	data := make([]complex128, rows*cols)
 	for i := range p {
-		for k := range a {
-			for j := range q {
+		for j := range q {
+			mij := m.At(i, j)
+			for k := range a {
 				for l := range b {
-					data[idx] = m.At(i, j) * n.At(k, l)
-					idx++
+					row, col := i*a+k, j*b+l
+					data[row*cols+col] = mij * n.At(k, l)
 				}
 			}
 		}

math/vector/vector.go

@@ -66,12 +66,10 @@ func (v Vector) Mul(z complex128) Vector {
 func (v Vector) TensorProduct(w Vector) Vector {
 	p, q := len(v), len(w)
 
-	var idx int
 	out := make(Vector, p*q)
 	for i := range p {
 		for j := range q {
-			out[idx] = v[i] * w[j]
-			idx++
+			out[i*q+j] = v[i] * w[j]
 		}
 	}
 
@@ -92,21 +90,19 @@ func (v Vector) InnerProduct(w Vector) complex128 {
 
 // OuterProduct returns the outer product of v and w.
 func (v Vector) OuterProduct(w Vector) matrix.Matrix {
-	p, q := len(v), len(w)
+	rows, cols := len(v), len(w)
 	dual := w.Dual()
 
-	var idx int
-	data := make([]complex128, p*q)
+	data := make([]complex128, rows*cols)
 	for i := range v {
 		for j := range dual {
-			data[idx] = v[i] * dual[j]
-			idx++
+			data[i*cols+j] = v[i] * dual[j]
 		}
 	}
 
 	return matrix.Matrix{
-		Rows: p,
-		Cols: q,
+		Rows: rows,
+		Cols: cols,
 		Data: data,
 	}
 }

quantum/density/matrix.go

@@ -39,7 +39,7 @@ func New(ensemble []State) *Matrix {
 		op := s.Qubit.OuterProduct(s.Qubit).Mul(complex(s.Probability, 0))
 		for i := range n {
 			for j := range n {
-				m.Set(i, j, m.At(i, j)+op.At(i, j))
+				m.AddAt(i, j, op.At(i, j))
 			}
 		}
 	}
@@ -129,7 +129,7 @@ func (m *Matrix) PartialTrace(index ...Qubit) (*Matrix, error) {
 
 			r := int(number.Must(strconv.ParseInt(kr, 2, 0)))
 			c := int(number.Must(strconv.ParseInt(lr, 2, 0)))
-			out.Set(r, c, out.At(r, c)+m.m.At(i, j))
+			out.AddAt(r, c, m.m.At(i, j))
 
 			// fmt.Printf("[%v][%v] = [%v][%v] + [%v][%v]\n", r, c, r, c, i, j)
 			//
@@ -160,7 +160,9 @@ func (m *Matrix) Depolarizing(p float64) (*Matrix, error) {
 	n := m.NumQubits()
 	i := gate.I(n).Mul(complex(p/2, 0))
 	r := m.m.Mul(complex(1-p, 0))
-	return &Matrix{i.Add(r)}, nil
+	return &Matrix{
+		m: i.Add(r),
+	}, nil
 }
 
 func take(n, i int, index []Qubit) (string, string) {

quantum/gate/gate.go

@@ -139,12 +139,12 @@ func Controlled(u matrix.Matrix, n int, c []int, t int) matrix.Matrix {
 
 	s := (1 << n)
 	g := I(n)
-	for i := 0; i < s; i++ {
+	for i := range s {
 		if (i & mask) != mask {
 			continue
 		}
 
-		for j := 0; j < s; j++ {
+		for j := range s {
 			if (j & mask) != mask {
 				continue
 			}
@@ -215,9 +215,9 @@ func ControlledZ(n int, c []int, t int) matrix.Matrix {
 	}
 
 	g := I(n)
-	for i := 0; i < (1 << n); i++ {
+	for i := range 1 << n {
 		if (i&mask) == mask && (i&(1<<(n-1-t))) != 0 {
-			g.Set(i, i, -1*g.At(i, i))
+			g.MulAt(i, i, -1)
 		}
 	}
 
@@ -237,9 +237,9 @@ func ControlledS(n int, c []int, t int) matrix.Matrix {
 	}
 
 	g := I(n)
-	for i := 0; i < (1 << n); i++ {
+	for i := range 1 << n {
 		if (i&mask) == mask && (i&(1<<(n-1-t))) != 0 {
-			g.Set(i, i, 1i*g.At(i, i))
+			g.MulAt(i, i, 1i)
 		}
 	}
 
@@ -262,9 +262,9 @@ func ControlledR(theta float64, n int, c []int, t int) matrix.Matrix {
 	}
 
 	g := I(n)
-	for i := 0; i < (1 << n); i++ {
+	for i := range 1 << n {
 		if (i&mask) == mask && (i&(1<<(n-1-t))) != 0 {
-			g.Set(i, i, e*g.At(i, i))
+			g.MulAt(i, i, e)
 		}
 	}