<jupyter_start><jupyter_text>Schmidt decomposition of a vectorThis problem can be solved using Singular Value Decomposition.<jupyter_code>import numpy as np
import cirq

def schmidt_decompose(vec, dim1, dim2, tol=1e-9):
    """Schmidt decomposition of a unit vector.
    
    Returns a, b, k, s.t. vec = Σ_i (k_i * a_i x b_i).
    """
    assert vec.shape == (dim1 * dim2, )
    assert np.allclose(np.linalg.norm(vec), 1), "Not unit vector."
    A = np.reshape(vec, (dim1, dim2))
    a, k, bT = np.linalg.svd(A)
    size = sum(k>=tol)
    a = [a[:, i] for i in range(size)]
    b = [bT[i, :] for i in range(size)]
    k = k[0:size]
    return a, b, k  
def test_schmidt_decompose(vec):
    def is_orthonormal_basis(vecs):
        V = np.array(vecs)
        return np.allclose(V @ V.T.conj(), np.eye(len(vecs)))
        
    vec = np.array(vec)
    dim = len(vec)
    for dim1 in range(1, dim+1):
        if dim % dim1 != 0: continue
        dim2 = dim // dim1
        a, b, k = schmidt_decompose(vec, dim1, dim2)
        
        assert np.allclose(np.linalg.norm(k), 1)
        assert is_orthonormal_basis(a)
        assert is_orthonormal_basis(b)
        
        vec_restored = sum([k[i] * np.kron(a[i], b[i]) for i in range(len(k))])
        assert np.allclose(vec, vec_restored)

test_schmidt_decompose([1.0])
test_schmidt_decompose([0.5, 0.5, 0.5, 0.5])

for i in range(100):
    vec = np.random.random(size=(8,2)) 
    vec = vec[:,0] + 1j * vec[:,1]
    vec /= np.linalg.norm(vec)
    test_schmidt_decompose(vec)
print("OK")<jupyter_output>OK
{'fraction_non_alphanumeric': 0.15106520335700452, 'fraction_numerical': 0.024531956100710135, 'mean_word_length': 2.7530266343825667, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '19972', 'license': 'MIT', 'n_tokens_mistral': 647, 'n_tokens_neox': 633, 'n_words': 149, 'path': 'Schmidt decomposition of a vector.ipynb', 'repo_name': 'fedimser/quant_comp', 'size': '1549'}