Update density matrix

quantum/density/matrix.go

@@ -11,12 +11,6 @@ import (
 	"github.com/itsubaki/q/quantum/qubit"
 )
 
-// State is a quantum state.
-type State struct {
-	Probability float64
-	Qubit       *qubit.Qubit
-}
-
 // Matrix is a density matrix.
 type Matrix struct {
 	m matrix.Matrix
@@ -25,27 +19,15 @@ type Matrix struct {
 // New returns a new density matrix.
 func New(ensemble []State) (*Matrix, error) {
 	m := &Matrix{matrix.New()}
-	if err := m.Add(ensemble); err != nil {
-		return nil, fmt.Errorf("add: %v", err)
-	}
-
-	return m, nil
-}
-
-// Add adds a quantum state to the density matrix.
-func (m *Matrix) Add(ensemble []State) error {
-	for _, s := range ensemble {
-		if s.Probability < 0 || s.Probability > 1 {
-			return fmt.Errorf("p must be 0 <= p =< 1. p=%v", s.Probability)
-		}
 
+	for _, s := range Normalize(ensemble) {
 		n := s.Qubit.Dimension()
 		if len(m.m) < 1 {
 			m.m = matrix.Zero(n, n)
 		}
 
 		if len(m.m) != n {
-			return fmt.Errorf("invalid dimension. m=%d n=%d", len(m.m), n)
+			return nil, fmt.Errorf("invalid dimension. m=%d n=%d", len(m.m), n)
 		}
 
 		op := s.Qubit.OuterProduct(s.Qubit).Mul(complex(s.Probability, 0))
@@ -54,10 +36,9 @@ func (m *Matrix) Add(ensemble []State) error {
 				m.m[i][j] = m.m[i][j] + op[i][j]
 			}
 		}
-
 	}
 
-	return nil
+	return m, nil
 }
 
 // Raw returns the raw matrix.

quantum/density/matrix_test.go

@@ -473,7 +473,10 @@ func TestMatrixNew(t *testing.T) {
 		in     []density.State
 		hasErr bool
 	}{
-		{[]density.State{{1.5, qubit.Zero()}}, true},
+		{[]density.State{
+			{1.5, qubit.New(1, 0)},
+			{1.0, qubit.New(1, 0, 0, 0)},
+		}, true},
 	}
 
 	for _, c := range cases {
@@ -485,25 +488,6 @@ func TestMatrixNew(t *testing.T) {
 	}
 }
 
-func TestMatrixAdd(t *testing.T) {
-	rho, _ := density.New([]density.State{{0.5, qubit.Zero()}})
-
-	cases := []struct {
-		in     []density.State
-		hasErr bool
-	}{
-		{[]density.State{{0.5, qubit.One()}}, false},
-		{[]density.State{{0.5, qubit.One(2)}}, true},
-	}
-
-	for _, c := range cases {
-		if err := rho.Add(c.in); (err != nil) != c.hasErr {
-			t.Errorf("err: %v", err)
-			continue
-		}
-	}
-}
-
 func TestDepolarizing(t *testing.T) {
 	rho, _ := density.New([]density.State{{1.0, qubit.Zero()}})
 

quantum/density/state.go

@@ -0,0 +1,23 @@
+package density
+
+import "github.com/itsubaki/q/quantum/qubit"
+
+// State is a quantum state.
+type State struct {
+	Probability float64
+	Qubit       *qubit.Qubit
+}
+
+// Normalize normalizes the probabilities of an ensemble.
+func Normalize(ensemble []State) []State {
+	var sum float64
+	for _, s := range ensemble {
+		sum += s.Probability
+	}
+
+	for i := range ensemble {
+		ensemble[i].Probability = ensemble[i].Probability / sum
+	}
+
+	return ensemble
+}

quantum/density/state_test.go

@@ -0,0 +1 @@
+package density_test