Add AddControlled

quantum/gate/gate.go

@@ -133,6 +133,26 @@ func RZ(theta float64) matrix.Matrix {
 	}
 }
 
+// AddControlled returns a controlled-u gate with control bit.
+// u is a (2**n x 2**n) unitary matrix and returns a (2**n x 2**n) matrix.
+func AddControlled(u matrix.Matrix, c int) matrix.Matrix {
+	d, _ := u.Dimension()
+	n := number.Log2(d)
+	g := I(n)
+
+	for i := 0; i < d; i++ {
+		if (i>>(n-1-c))&1 == 0 {
+			continue
+		}
+
+		for j := 0; j < d; j++ {
+			g[i][j] = u[i][j]
+		}
+	}
+
+	return g
+}
+
 // Controlled returns a controlled-u gate.
 // u is a (2x2) unitary matrix and returns a (2**n x 2**n) matrix.
 func Controlled(u matrix.Matrix, n int, c []int, t int) matrix.Matrix {

quantum/gate/gate_test.go

@@ -3,6 +3,7 @@ package gate_test
 import (
 	"fmt"
 	"math"
+	"math/rand/v2"
 	"strconv"
 	"testing"
 
@@ -167,6 +168,74 @@ func ExampleTensorProduct() {
 	// [(0+0i) (1+0i) (0+0i) (0+0i)]
 }
 
+func TestAddControlled(t *testing.T) {
+	u := gate.U(rand.Float64(), rand.Float64(), rand.Float64())
+
+	cases := []struct {
+		in   matrix.Matrix
+		want matrix.Matrix
+		bit  int
+	}{
+		{
+			in:   gate.TensorProduct(gate.X(), 2, []int{1}),
+			want: gate.ControlledNot(2, []int{0}, 1),
+			bit:  0,
+		},
+		{
+			in:   gate.TensorProduct(gate.X(), 2, []int{0}),
+			want: gate.ControlledNot(2, []int{1}, 0),
+			bit:  1,
+		},
+		{
+			in:   gate.ControlledNot(3, []int{0}, 2),
+			want: gate.ControlledNot(3, []int{0, 1}, 2),
+			bit:  1,
+		},
+		{
+			in:   gate.ControlledNot(3, []int{0}, 2),
+			want: gate.ControlledNot(3, []int{0}, 2),
+			bit:  0,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{0}, 1),
+			want: gate.Controlled(u, 3, []int{0, 2}, 1),
+			bit:  2,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{0}, 2),
+			want: gate.Controlled(u, 3, []int{0, 1}, 2),
+			bit:  1,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{1}, 2),
+			want: gate.Controlled(u, 3, []int{0, 1}, 2),
+			bit:  0,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{1}, 0),
+			want: gate.Controlled(u, 3, []int{1, 2}, 0),
+			bit:  2,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{2}, 0),
+			want: gate.Controlled(u, 3, []int{2, 1}, 0),
+			bit:  1,
+		},
+		{
+			in:   gate.Controlled(u, 3, []int{2}, 1),
+			want: gate.Controlled(u, 3, []int{0, 2}, 1),
+			bit:  0,
+		},
+	}
+
+	for _, c := range cases {
+		got := gate.AddControlled(c.in, c.bit)
+		if !got.Equals(c.want) {
+			t.Fail()
+		}
+	}
+}
+
 func TestControlledModExp2(t *testing.T) {
 	g1 := matrix.Apply(
 		gate.CNOT(7, 3, 5),