@@ -10,7 +10,7 @@ by the associated classical bitstring.

 import numpy as np

-__all__ = ["from_classical"]  # type: ignore

+__all__ = ["from_classical", "is_normalized", "num_qubits", "zero"]  # type: ignore

 def from_classical(bitstring):

@@ -28,7 +28,6 @@ def from_classical(bitstring):

         The state vector in the computational basis.

         Has :math:`2^n` components.

     """

-

     bitstring = "".join(map(str, bitstring))

     n_qubits = len(bitstring)

     try:

@@ -39,3 +38,41 @@ def from_classical(bitstring):

     state = np.zeros(1 << n_qubits, dtype=complex)

     state[index] = 1

     return state

+

+

+def zero(n: int):

+    """Return the zero state on 'n' qubits."""

+    state = np.zeros(1 << n, dtype=complex)

+    state[0] = 1

+    return state

+

+

+def num_qubits(state):

+    """Return the number of qubits in the state.

+

+    Raises ValueError if 'state' does not have a shape that is

+    an integer power of 2.

+    """

+    _check_valid_state(state)

+    return state.shape[0].bit_length() - 1

+

+

+def is_normalized(state: np.ndarray) -> bool:

+    """Return True if and only if 'state' is normalized."""

+    return np.allclose(np.linalg.norm(state), 1)

+

+

+def _check_valid_state(state):

+    if not (

+        # is an array

+        isinstance(state, np.ndarray)

+        # is complex

+        and np.issubdtype(state.dtype, np.complex128)

+        # is square

+        and len(state.shape) == 1

+        # has size 2**n, n > 1

+        and np.log2(state.shape[0]).is_integer()

+        and state.shape[0].bit_length() > 1

+        and is_normalized(state)

+    ):

+        raise ValueError("State is not valid")

@@ -2,6 +2,7 @@ import string

 from functools import partial

 from hypothesis import given

+from hypothesis.extra import numpy as hnp

 import hypothesis.strategies as st

 import numpy as np

 import pytest

@@ -11,7 +12,14 @@ import qsim.state

 MAX_QUBITS = 5

+

+def is_unit(c: complex):

+    return np.isclose(abs(c), 1)

+

+

 n_qubits = st.integers(1, MAX_QUBITS)

+state_dimensions = n_qubits.map(lambda n: 2 ** n)

+state_shapes = state_dimensions.map(lambda x: (x,))

 classical_bitstrings = n_qubits.flatmap(

     lambda n: st.integers(0, 2 ** n).map(partial("{:0{n}b}".format, n=n))

 )

@@ -19,6 +27,37 @@ invalid_bitstrings = n_qubits.flatmap(

     lambda n: st.text(alphabet=string.digits, min_size=n, max_size=n)

 ).filter(lambda s: any(b not in "01" for b in s))

+invalid_state_dimensions = st.integers(0, 2 ** MAX_QUBITS).filter(

+    lambda n: not np.log2(n).is_integer()

+)

+

+

+valid_complex = st.complex_numbers(

+    max_magnitude=1e10, allow_infinity=False, allow_nan=False

+)

+nonunit_complex = valid_complex.filter(lambda c: not is_unit(c))

+

+

+def ket(n_qubits):

+    return normalized_array(1 << n_qubits)

+

+

+def normalized_array(size):

+    return (

+        hnp.arrays(complex, (size,), elements=valid_complex)

+        .filter(lambda v: np.linalg.norm(v) > 0)  # vectors must be normalizable

+        .map(lambda v: v / np.linalg.norm(v))

+    )

+

+

+valid_states = n_qubits.flatmap(ket)

+invalid_shape_states = invalid_state_dimensions.flatmap(normalized_array)

+invalid_norm_states = valid_states.flatmap(

+    lambda x: nonunit_complex.map(lambda c: c * x)

+)

+zero_states = state_shapes.map(lambda s: np.zeros(s, complex))

+invalid_states = st.one_of(invalid_shape_states, invalid_norm_states, zero_states)

+

 @given(classical_bitstrings)

 def test_from_classical(bitstring):

@@ -47,3 +86,31 @@ def test_from_classical_works_on_integer_lists(bitstring):

 def test_from_classical_raises_on_bad_input(bitstring):

     with pytest.raises(ValueError):

         qsim.state.from_classical(bitstring)

+

+

+@given(classical_bitstrings)

+def test_num_qubits(s):

+    state = qsim.state.from_classical(s)

+    assert qsim.state.num_qubits(state) == len(s)

+

+

+@given(invalid_states)

+def test_num_qubits_raises_exception(state):

+    with pytest.raises(ValueError):

+        qsim.state.num_qubits(state)

+

+

+@given(invalid_norm_states)

+def test_is_not_normalized(state):

+    assert not qsim.state.is_normalized(state)

+

+

+@given(valid_states)

+def test_is_normalized(state):

+    assert qsim.state.is_normalized(state)

+

+

+@given(n_qubits)

+def test_zero(n):

+    z = qsim.state.zero(n)

+    assert qsim.state.num_qubits(z) == n