From 533836efe7879308f723ba166c447461c61a2c27 Mon Sep 17 00:00:00 2001
From: itsubaki <1759459+itsubaki@users.noreply.github.com>
Date: Fri, 19 Apr 2024 22:21:38 +0900
Subject: [PATCH] Add Log2 to math/number package

---
 math/number/log.go        | 25 +++++++++++++++++++++++++
 math/number/log_test.go   | 36 ++++++++++++++++++++++++++++++++++++
 math/number/must_test.go  | 26 ++++++++++++++++++++++++++
 math/number/parse_test.go | 18 ------------------
 q.go                      |  7 ++++---
 quantum/density/matrix.go |  5 ++---
 quantum/qubit/qubit.go    |  4 +---
 7 files changed, 94 insertions(+), 27 deletions(-)
 create mode 100644 math/number/log.go
 create mode 100644 math/number/log_test.go
 create mode 100644 math/number/must_test.go

diff --git a/math/number/log.go b/math/number/log.go
new file mode 100644
index 0000000..0b5a323
--- /dev/null
+++ b/math/number/log.go
@@ -0,0 +1,25 @@
+package number
+
+import "fmt"
+
+func Log2(N int) (int, error) {
+	if N == 1 {
+		return 0, nil
+	}
+
+	if N%2 != 0 {
+		return 0, fmt.Errorf("N must be a power of 2. N=%v", N)
+	}
+
+	var n int = 1
+	for {
+		if N/2 == 1 {
+			break
+		}
+
+		N = N / 2
+		n++
+	}
+
+	return n, nil
+}
diff --git a/math/number/log_test.go b/math/number/log_test.go
new file mode 100644
index 0000000..c54344c
--- /dev/null
+++ b/math/number/log_test.go
@@ -0,0 +1,36 @@
+package number_test
+
+import (
+	"testing"
+
+	"github.com/itsubaki/q/math/number"
+)
+
+func TestLog2(t *testing.T) {
+	cases := []struct {
+		n    int
+		want int
+	}{
+		{1, 0},
+		{2, 1},
+		{4, 2},
+		{8, 3},
+		{16, 4},
+		{32, 5},
+		{64, 6},
+		{128, 7},
+		{256, 8},
+		{512, 9},
+		{1024, 10},
+		{2048, 11},
+		{4096, 12},
+		{8192, 13},
+	}
+
+	for _, c := range cases {
+		got := number.Must(number.Log2(c.n))
+		if got != c.want {
+			t.Errorf("got=%v want=%d", got, c.want)
+		}
+	}
+}
diff --git a/math/number/must_test.go b/math/number/must_test.go
new file mode 100644
index 0000000..7be8d6d
--- /dev/null
+++ b/math/number/must_test.go
@@ -0,0 +1,26 @@
+package number_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/itsubaki/q/math/number"
+)
+
+func TestMustPanic(t *testing.T) {
+	defer func() {
+		if rec := recover(); rec != nil {
+			err, ok := rec.(error)
+			if !ok {
+				t.Fail()
+			}
+
+			if err.Error() != "something went wrong" {
+				t.Fail()
+			}
+		}
+	}()
+
+	number.Must(-1, fmt.Errorf("something went wrong"))
+	t.Fail()
+}
diff --git a/math/number/parse_test.go b/math/number/parse_test.go
index 5f3733f..cbd389c 100644
--- a/math/number/parse_test.go
+++ b/math/number/parse_test.go
@@ -56,24 +56,6 @@ func TestParseFloat(t *testing.T) {
 	}
 }
 
-func TestMustPanic(t *testing.T) {
-	defer func() {
-		if rec := recover(); rec != nil {
-			err, ok := rec.(error)
-			if !ok {
-				t.Fail()
-			}
-
-			if err.Error() != "something went wrong" {
-				t.Fail()
-			}
-		}
-	}()
-
-	number.Must(-1, fmt.Errorf("something went wrong"))
-	t.Fail()
-}
-
 func FuzzParseFloat(f *testing.F) {
 	seed := []string{"123", "101", "1.0101", "abc", "a.bc"}
 	for i := range seed {
diff --git a/q.go b/q.go
index 58adac6..7437f6d 100644
--- a/q.go
+++ b/q.go
@@ -83,7 +83,7 @@ func (q *Q) One() Qubit {
 	return q.New(0, 1)
 }
 
-// ZeroWith returns a qubit in the zero state with n qubits.
+// ZeroWith returns n qubits in the zero state.
 func (q *Q) ZeroWith(n int) []Qubit {
 	qb := make([]Qubit, n)
 	for i := 0; i < n; i++ {
@@ -93,7 +93,7 @@ func (q *Q) ZeroWith(n int) []Qubit {
 	return qb
 }
 
-// One returns a qubit in the one state with n qubits.
+// One returns n qubit in the one state.
 func (q *Q) OneWith(n int) []Qubit {
 	qb := make([]Qubit, n)
 	for i := 0; i < n; i++ {
@@ -103,7 +103,8 @@ func (q *Q) OneWith(n int) []Qubit {
 	return qb
 }
 
-// ZeroLog2 returns a qubit in the zero state with log2(N) qubits.
+// ZeroLog2 returns n qubit in the zero state.
+// n is greater than or equal to log2(N).
 func (q *Q) ZeroLog2(N int) []Qubit {
 	n := int(math.Log2(float64(N))) + 1
 	return q.ZeroWith(n)
diff --git a/quantum/density/matrix.go b/quantum/density/matrix.go
index e094f9f..23b76f9 100644
--- a/quantum/density/matrix.go
+++ b/quantum/density/matrix.go
@@ -73,8 +73,7 @@ func (m *Matrix) Dimension() (int, int) {
 // NumberOfBit returns the number of qubits.
 func (m *Matrix) NumberOfBit() int {
 	p, _ := m.Dimension()
-	log := math.Log2(float64(p))
-	return int(log)
+	return number.Must(number.Log2(p))
 }
 
 // Apply applies a unitary matrix to the density matrix.
@@ -166,7 +165,7 @@ func Flip(p float64, m matrix.Matrix) (matrix.Matrix, matrix.Matrix, error) {
 	}
 
 	d, _ := m.Dimension()
-	n := int(math.Log2(float64(d)))
+	n := number.Must(number.Log2(d))
 
 	e0 := gate.I(n).Mul(complex(math.Sqrt(p), 0))
 	e1 := m.Mul(complex(math.Sqrt(1-p), 0))
diff --git a/quantum/qubit/qubit.go b/quantum/qubit/qubit.go
index 08ea76f..4f60ea4 100644
--- a/quantum/qubit/qubit.go
+++ b/quantum/qubit/qubit.go
@@ -48,9 +48,7 @@ func One(n ...int) *Qubit {
 
 // NumberOfBit returns the number of qubits.
 func (q *Qubit) NumberOfBit() int {
-	d := float64(q.Dimension())
-	n := math.Log2(d)
-	return int(n)
+	return number.Must(number.Log2(q.Dimension()))
 }
 
 // IsZero returns true if q is zero qubit.