jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_dm_3__simulation_3(n_qubits, n_gates):
    from hybridq.gate.measure import _Measure
    from scipy.linalg import eigvalsh
    from string import ascii_letters

    # Get RQC
    circuit = _get_rqc_unitary(n_qubits, n_gates)

    # Get random initial state
    initial_state = ''.join(np.random.choice(list('01+-'), size=n_qubits))

    # Get indexes for open qubits
    index_open_qubits = sorted(np.random.choice(n_qubits, size=2,
                                                replace=False))

    # Get final state, including the right projections
    final_state = ''.join('.' if i in index_open_qubits else c
                          for i, c in enumerate(ascii_letters[:n_qubits]))
    final_state += final_state

    # Get state by using state evolution
    psi_1 = simulation.simulate(circuit,
                                initial_state=initial_state,
                                verbose=True)

    # Get the expectation value using tensor contraction
    rho_1 = dm_simulation.simulate(circuit,
                                   initial_state=initial_state,
                                   final_state=final_state,
                                   optimize='tn',
                                   verbose=True)

    # Get matrix
    _rho_1 = np.reshape(rho_1, (2**len(index_open_qubits),) * 2)

    # Density matrix should be hermitian
    assert_allclose(_rho_1, _rho_1.conj().T)

    # Density matrix should be semi-positive definite
    assert (np.all(np.round(eigvalsh(_rho_1), 5) >= 0))

    # Trace qubits using einsum
    _rho = np.einsum(
        ''.join(ascii_letters[i] for i in range(n_qubits)) + ',' +
        ''.join(ascii_letters[i + n_qubits if i in index_open_qubits else i]
                for i in range(n_qubits)), psi_1, psi_1.conj())

    # Check
    assert_allclose(_rho, rho_1)

    # Get probabilities of the projected states
    probs = _Measure(psi_1, axes=index_open_qubits, get_probs_only=True)

    # Check normalization
    assert (np.isclose(np.sum(probs), 1, atol=_atol))

    # Check
    assert_allclose(np.diag(_rho_1), probs)
//  [python]
def test_lcth_direct_neighbors(self):
        lcthalgo = LowCloudHeightAlgorithm(**self.testinput)
        elev = self.testinput2['elev']
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 1)
        compare = np.array([800., 900., 1000., 700., 600., 400., 300., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper left corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 0)
        compare = np.array([900., 700., 1500.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower left corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 0)
        compare = np.array([700., 1500., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper right corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 2)
        compare = np.array([900., 1500., 600.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower right corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 2)
        compare = np.array([1500., 600., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test middle right cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 2)
        compare = np.array([900., 1000., 1500., 300., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test middle left cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 0)
        compare = np.array([800., 900., 1500., 400., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper middle cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 1)
        compare = np.array([800., 1000., 700., 1500., 600.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower middle cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 1)
        compare = np.array([700., 1500., 600., 400., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
//  [python]
def test_lcth_direct_neighbors(self):
        lcthalgo = LowCloudHeightAlgorithm(**self.testinput)
        elev = self.testinput2['elev']
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 1)
        compare = np.array([800., 900., 1000., 700., 600., 400., 300., 200.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper left corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 0)
        compare = np.array([900., 700., 1500.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower left corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 0)
        compare = np.array([700., 1500., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper right corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 2)
        compare = np.array([900., 1500., 600.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower right corner
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 2)
        compare = np.array([1500., 600., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test middle right cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 2)
        compare = np.array([900., 1000., 1500., 300., 200.])
        self.assertTrue(np.all(compare == neigh))
        # Test middle left cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 1, 0)
        compare = np.array([800., 900., 1500., 400., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper middle cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 0, 1)
        compare = np.array([800., 1000., 700., 1500., 600.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower middle cell
        center, neigh, ids = lcthalgo.get_neighbors(elev, 2, 1)
        compare = np.array([700., 1500., 600., 400., 200.])
        self.assertTrue(np.all(compare == neigh))
//  [python]
def test_lcth_cell_neighbors(self):
        lcthalgo = LowCloudHeightAlgorithm(**self.testinput)
        elev = self.testinput2['elev']
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 1)
        compare = np.array([800., 900., 1000., 700., 600., 400., 300., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper left corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 0)
        compare = np.array([900., 700., 1500.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower left corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 0)
        compare = np.array([700., 1500., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper right corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 2)
        compare = np.array([900., 1500., 600.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower right corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 2)
        compare = np.array([1500., 600., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test middle right cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 2)
        compare = np.array([900., 1000., 1500., 300., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test middle left cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 0)
        compare = np.array([800., 900., 1500., 400., 300.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test upper middle cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 1)
        compare = np.array([800., 1000., 700., 1500., 600.])
        self.assertTrue(np.alltrue(compare == neigh))
        # Test lower middle cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 1)
        compare = np.array([700., 1500., 600., 400., 200.])
        self.assertTrue(np.alltrue(compare == neigh))
//  [python]
def test_lcth_cell_neighbors(self):
        lcthalgo = LowCloudHeightAlgorithm(**self.testinput)
        elev = self.testinput2['elev']
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 1)
        compare = np.array([800., 900., 1000., 700., 600., 400., 300., 200.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper left corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 0)
        compare = np.array([900., 700., 1500.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower left corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 0)
        compare = np.array([700., 1500., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper right corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 2)
        compare = np.array([900., 1500., 600.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower right corner
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 2)
        compare = np.array([1500., 600., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test middle right cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 2)
        compare = np.array([900., 1000., 1500., 300., 200.])
        self.assertTrue(np.all(compare == neigh))
        # Test middle left cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 1, 0)
        compare = np.array([800., 900., 1500., 400., 300.])
        self.assertTrue(np.all(compare == neigh))
        # Test upper middle cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 0, 1)
        compare = np.array([800., 1000., 700., 1500., 600.])
        self.assertTrue(np.all(compare == neigh))
        # Test lower middle cell
        center, neigh, ids = lcthalgo.cell_neighbors(elev, 2, 1)
        compare = np.array([700., 1500., 600., 400., 200.])
        self.assertTrue(np.all(compare == neigh))
//  [python]
def test_forward_pass_single_channel_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([[
            [[1], [2], [2], [1]],
            [[3], [4], [0], [0]],
            [[5], [2], [1], [1]],
            [[3], [4], [0], [3]]
        ]])

        expected_result = np.array([[
            [[4], [2]],
            [[5], [3]],
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (1, 2, 2, 1)
        assert np.alltrue(expected_result == result)
//  [python]
def test_forward_pass_single_channel_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([[
            [[1], [2], [2], [1]],
            [[3], [4], [0], [0]],
            [[5], [2], [1], [1]],
            [[3], [4], [0], [3]]
        ]])

        expected_result = np.array([[
            [[4], [2]],
            [[5], [3]],
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (1, 2, 2, 1)
        assert np.all(expected_result == result)
//  [python]
def test_forward_pass_two_channels_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([[
            [
                [1, 5],
                [2, 2],
                [2, 2],
                [1, 1]
            ],
            [
                [3, 3],
                [4, 4],
                [0, 3],
                [0, 0]
            ],
            [
                [5, 2],
                [2, 2],
                [1, 1],
                [1, 1]
            ],
            [
                [3, 3],
                [4, 4],
                [0, 2],
                [3, 0]
            ]
        ]])

        expected_result = np.array([[
            [
                [4, 5],
                [2, 3]
            ],
            [
                [5, 4],
                [3, 2]
            ]
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (1, 2, 2, 2)
        assert np.alltrue(expected_result == result)
//  [python]
def test_forward_pass_two_channels_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([[
            [
                [1, 5],
                [2, 2],
                [2, 2],
                [1, 1]
            ],
            [
                [3, 3],
                [4, 4],
                [0, 3],
                [0, 0]
            ],
            [
                [5, 2],
                [2, 2],
                [1, 1],
                [1, 1]
            ],
            [
                [3, 3],
                [4, 4],
                [0, 2],
                [3, 0]
            ]
        ]])

        expected_result = np.array([[
            [
                [4, 5],
                [2, 3]
            ],
            [
                [5, 4],
                [3, 2]
            ]
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (1, 2, 2, 2)
        assert np.all(expected_result == result)
//  [python]
def compare(A, B, ids=None):
    """Safe comparison of general objects

    USAGE:
      compare(A,B)

    DESCRIPTION:
      Return 1 if A and B they are identical, 0 otherwise
    """

    # Keep track of unique id's to protect against infinite recursion
    if ids is None: ids = {}

    # Check if T has already been encountered
    iA = id(A) 
    iB = id(B)     
    
    if (iA, iB) in ids:
        # A and B have been compared already
        return ids[(iA, iB)]
    else:
        ids[(iA, iB)] = True
    
    
    # Check if arguments are of same type
    if type(A) != type(B):
        return False
        
  
    # Compare recursively based on argument type
    if isinstance(A, (tuple, list)):
        N = len(A)
        if len(B) != N: 
            identical = False
        else:
            identical = True
            for i in range(N):
                if not compare(A[i], B[i], ids): 
                    identical = False; break
                    
    elif isinstance(A, dict):
        if len(A) != len(B):
            identical = False
        else:
            # Dictionaries are now ordered as of Python 3.7
            # Make dictionary ordering unique 
            #a = list(A.items()); a.sort()    
            #b = list(B.items()); b.sort()
            
            identical = compare(A, B, ids)
            
    elif isinstance(A, num.ndarray):
        # Use element by element comparison
        identical = num.alltrue(A==B)

    #elif type(A) == types.InstanceType:
    elif type(A) is type:
        # Take care of special case where elements are instances            
        # Base comparison on attributes
        identical = compare(A.__dict__, 
                            B.__dict__, 
                            ids)
    else:       
        # Fall back to general code
        try:
            identical = (A == B)
        except:
            import pickle
            # Use pickle to compare data
            # The native pickler must be used
            # since the faster cPickle does not 
            # guarantee a unique translation
            # FIXME (Ole): Try to fall back on the dill pickler
            try:
                identical = (pickle.dumps(A,0) == pickle.dumps(B,0))
            except:
                identical = False

    # Record result of comparison and return            
    ids[(iA, iB)] = identical
    
    return(identical)
//  [python]
def compare(A, B, ids=None):
    """Safe comparison of general objects

    USAGE:
      compare(A,B)

    DESCRIPTION:
      Return 1 if A and B they are identical, 0 otherwise
    """

    # Keep track of unique id's to protect against infinite recursion
    if ids is None: ids = {}

    # Check if T has already been encountered
    iA = id(A) 
    iB = id(B)     
    
    if (iA, iB) in ids:
        # A and B have been compared already
        return ids[(iA, iB)]
    else:
        ids[(iA, iB)] = True
    
    
    # Check if arguments are of same type
    if type(A) != type(B):
        return False
        
  
    # Compare recursively based on argument type
    if isinstance(A, (tuple, list)):
        N = len(A)
        if len(B) != N: 
            identical = False
        else:
            identical = True
            for i in range(N):
                if not compare(A[i], B[i], ids): 
                    identical = False; break
                    
    elif isinstance(A, dict):
        if len(A) != len(B):
            identical = False
        else:
            # Dictionaries are now ordered as of Python 3.7
            # Make dictionary ordering unique 
            #a = list(A.items()); a.sort()    
            #b = list(B.items()); b.sort()
            
            identical = compare(A, B, ids)
            
    elif isinstance(A, num.ndarray):
        # Use element by element comparison
        identical = num.all(A==B)

    #elif type(A) == types.InstanceType:
    elif type(A) is type:
        # Take care of special case where elements are instances            
        # Base comparison on attributes
        identical = compare(A.__dict__, 
                            B.__dict__, 
                            ids)
    else:       
        # Fall back to general code
        try:
            identical = (A == B)
        except:
            import pickle
            # Use pickle to compare data
            # The native pickler must be used
            # since the faster cPickle does not 
            # guarantee a unique translation
            # FIXME (Ole): Try to fall back on the dill pickler
            try:
                identical = (pickle.dumps(A,0) == pickle.dumps(B,0))
            except:
                identical = False

    # Record result of comparison and return            
    ids[(iA, iB)] = identical
    
    return(identical)
//  [python]
def subblocks(a, *ranges):
    """
    This function produces a distributed array from a subset of the blocks in
    the `a`. The result and `a` will have the same number of dimensions. For
    example,
        subblocks(a, [0, 1], [2, 4])
    will produce a DistArray whose objectids are
        [[a.objectids[0, 2], a.objectids[0, 4]],
         [a.objectids[1, 2], a.objectids[1, 4]]]
    We allow the user to pass in an empty list [] to indicate the full range.
    """
    ranges = list(ranges)
    if len(ranges) != a.ndim:
        raise Exception("sub_blocks expects to receive a number of ranges "
                        "equal to a.ndim, but it received {} ranges and "
                        "a.ndim = {}.".format(len(ranges), a.ndim))
    for i in range(len(ranges)):
        # We allow the user to pass in an empty list to indicate the full
        # range.
        if ranges[i] == []:
            ranges[i] = range(a.num_blocks[i])
        if not np.alltrue(ranges[i] == np.sort(ranges[i])):
            raise Exception("Ranges passed to sub_blocks must be sorted, but "
                            "the {}th range is {}.".format(i, ranges[i]))
        if ranges[i][0] < 0:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be at least 0, but the {}th range is {}.".format(
                                i, ranges[i]))
        if ranges[i][-1] >= a.num_blocks[i]:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be less than the relevant number of blocks, but "
                            "the {}th range is {}, and a.num_blocks = {}."
                            .format(i, ranges[i], a.num_blocks))
    last_index = [r[-1] for r in ranges]
    last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
    shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i]
             for i in range(a.ndim)]
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = a.objectids[tuple(
            ranges[i][index[i]] for i in range(a.ndim))]
    return result
//  [python]
def subblocks(a, *ranges):
    """
    This function produces a distributed array from a subset of the blocks in
    the `a`. The result and `a` will have the same number of dimensions. For
    example,
        subblocks(a, [0, 1], [2, 4])
    will produce a DistArray whose objectids are
        [[a.objectids[0, 2], a.objectids[0, 4]],
         [a.objectids[1, 2], a.objectids[1, 4]]]
    We allow the user to pass in an empty list [] to indicate the full range.
    """
    ranges = list(ranges)
    if len(ranges) != a.ndim:
        raise Exception("sub_blocks expects to receive a number of ranges "
                        "equal to a.ndim, but it received {} ranges and "
                        "a.ndim = {}.".format(len(ranges), a.ndim))
    for i in range(len(ranges)):
        # We allow the user to pass in an empty list to indicate the full
        # range.
        if ranges[i] == []:
            ranges[i] = range(a.num_blocks[i])
        if not np.all(ranges[i] == np.sort(ranges[i])):
            raise Exception("Ranges passed to sub_blocks must be sorted, but "
                            "the {}th range is {}.".format(i, ranges[i]))
        if ranges[i][0] < 0:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be at least 0, but the {}th range is {}.".format(
                                i, ranges[i]))
        if ranges[i][-1] >= a.num_blocks[i]:
            raise Exception("Values in the ranges passed to sub_blocks must "
                            "be less than the relevant number of blocks, but "
                            "the {}th range is {}, and a.num_blocks = {}."
                            .format(i, ranges[i], a.num_blocks))
    last_index = [r[-1] for r in ranges]
    last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
    shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i]
             for i in range(a.ndim)]
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.objectids[index] = a.objectids[tuple(
            ranges[i][index[i]] for i in range(a.ndim))]
    return result
//  [python]
def test_forward_pass_single_channel_two_items(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([
            [
                [[1], [2], [2], [1]],
                [[3], [4], [0], [0]],
                [[5], [2], [1], [1]],
                [[3], [4], [0], [3]]
            ],
            [
                [[5], [2], [2], [1]],
                [[3], [4], [3], [0]],
                [[2], [2], [1], [1]],
                [[3], [4], [2], [0]]
            ]
        ])

        expected_result = np.array([
            [
                [[4], [2]],
                [[5], [3]]
            ],
            [
                [[5], [3]],
                [[4], [2]]
            ]
        ])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (2, 2, 2, 1)
        assert np.alltrue(expected_result == result)
//  [python]
def test_forward_pass_single_channel_two_items(self):
        # given
        pool_size = (2, 2)
        stride = 2
        activation = np.array([
            [
                [[1], [2], [2], [1]],
                [[3], [4], [0], [0]],
                [[5], [2], [1], [1]],
                [[3], [4], [0], [3]]
            ],
            [
                [[5], [2], [2], [1]],
                [[3], [4], [3], [0]],
                [[2], [2], [1], [1]],
                [[3], [4], [2], [0]]
            ]
        ])

        expected_result = np.array([
            [
                [[4], [2]],
                [[5], [3]]
            ],
            [
                [[5], [3]],
                [[4], [2]]
            ]
        ])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        result = layer.forward_pass(activation, training=True)

        # then
        assert result.shape == (2, 2, 2, 1)
        assert np.all(expected_result == result)
//  [python]
def test_backward_pass_single_channel_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        forward_activation = np.array([[
            [[1], [2], [2], [1]],
            [[3], [4], [0], [0]],
            [[5], [2], [1], [1]],
            [[3], [4], [0], [3]]
        ]])

        backward_activation = np.array([[
            [[3], [1]],
            [[8], [2]],
        ]])

        expected_backward_result = np.array([[
            [[0], [0], [1], [0]],
            [[0], [3], [0], [0]],
            [[8], [0], [0], [0]],
            [[0], [0], [0], [2]]
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        _ = layer.forward_pass(forward_activation, training=True)
        backward_result = layer.backward_pass(backward_activation)

        # then
        assert np.alltrue(expected_backward_result == backward_result)
//  [python]
def test_backward_pass_single_channel_single_item(self):
        # given
        pool_size = (2, 2)
        stride = 2
        forward_activation = np.array([[
            [[1], [2], [2], [1]],
            [[3], [4], [0], [0]],
            [[5], [2], [1], [1]],
            [[3], [4], [0], [3]]
        ]])

        backward_activation = np.array([[
            [[3], [1]],
            [[8], [2]],
        ]])

        expected_backward_result = np.array([[
            [[0], [0], [1], [0]],
            [[0], [3], [0], [0]],
            [[8], [0], [0], [0]],
            [[0], [0], [0], [2]]
        ]])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        _ = layer.forward_pass(forward_activation, training=True)
        backward_result = layer.backward_pass(backward_activation)

        # then
        assert np.all(expected_backward_result == backward_result)
//  [python]
def transform(self, A):
        """ Apply cut-below value constraint"""
        if self._excl_mat is None:
            self._make_excl_mat(A.shape)

        if self.axis is None:
            if self.copy:
                return A*((A >= self.value) | self._excl_mat)
            else:
                A *= ((A >= self.value) | self._excl_mat)
                return A
        else:
            if self.copy:
                return A*(_np.alltrue(A < self.value, axis=self.axis, keepdims=True) +
                        (A >= self.value) + self._excl_mat)
            else:
                A *= (_np.alltrue(A < self.value, axis=self.axis, keepdims=True) +
                    (A >= self.value) + self._excl_mat)
                return A
//  [python]
def transform(self, A):
        """ Apply cut-below value constraint"""
        if self._excl_mat is None:
            self._make_excl_mat(A.shape)

        if self.axis is None:
            if self.copy:
                return A*((A >= self.value) | self._excl_mat)
            else:
                A *= ((A >= self.value) | self._excl_mat)
                return A
        else:
            if self.copy:
                return A*(_np.all(A < self.value, axis=self.axis, keepdims=True) +
                        (A >= self.value) + self._excl_mat)
            else:
                A *= (_np.all(A < self.value, axis=self.axis, keepdims=True) +
                    (A >= self.value) + self._excl_mat)
                return A
//  [python]
def transform(self, A):
        """ Apply cut-above value constraint"""

        if self._excl_mat is None:
            self._make_excl_mat(A.shape)
        if self.axis is None:
            if self.copy:
                return A*((A <= self.value) | self._excl_mat)
            else:
                A *= ((A <= self.value) | self._excl_mat)
                return A
        else:
            if self.copy:
                return A*(_np.alltrue(A > self.value, axis=self.axis, keepdims=True) +
                        (A <= self.value) + self._excl_mat)
            else:
                A *= (_np.alltrue(A > self.value, axis=self.axis, keepdims=True) +
                    (A <= self.value) + self._excl_mat)
                return A
//  [python]
def transform(self, A):
        """ Apply cut-above value constraint"""

        if self._excl_mat is None:
            self._make_excl_mat(A.shape)
        if self.axis is None:
            if self.copy:
                return A*((A <= self.value) | self._excl_mat)
            else:
                A *= ((A <= self.value) | self._excl_mat)
                return A
        else:
            if self.copy:
                return A*(_np.all(A > self.value, axis=self.axis, keepdims=True) +
                        (A <= self.value) + self._excl_mat)
            else:
                A *= (_np.all(A > self.value, axis=self.axis, keepdims=True) +
                    (A <= self.value) + self._excl_mat)
                return A
//  [python]
def test_backward_pass_single_channel_two_items(self):
        # given
        pool_size = (2, 2)
        stride = 2
        forward_activation = np.array([
            [
                [[1], [2], [2], [1]],
                [[3], [4], [0], [0]],
                [[5], [2], [1], [1]],
                [[3], [4], [0], [3]]
            ],
            [
                [[5], [2], [2], [1]],
                [[3], [4], [3], [0]],
                [[2], [2], [1], [1]],
                [[3], [4], [2], [0]]
            ]
        ])

        backward_activation = np.array([
            [
                [[7], [2]],
                [[4], [3]]
            ],
            [
                [[1], [5]],
                [[2], [2]]
            ]
        ])

        expected_backward_result = np.array([
            [
                [[0], [0], [2], [0]],
                [[0], [7], [0], [0]],
                [[4], [0], [0], [0]],
                [[0], [0], [0], [3]]
            ],
            [
                [[1], [0], [0], [0]],
                [[0], [0], [5], [0]],
                [[0], [0], [0], [0]],
                [[0], [2], [2], [0]]
            ]
        ])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        _ = layer.forward_pass(forward_activation, training=True)
        backward_result = layer.backward_pass(backward_activation)

        # then
        assert np.alltrue(expected_backward_result == backward_result)
//  [python]
def test_backward_pass_single_channel_two_items(self):
        # given
        pool_size = (2, 2)
        stride = 2
        forward_activation = np.array([
            [
                [[1], [2], [2], [1]],
                [[3], [4], [0], [0]],
                [[5], [2], [1], [1]],
                [[3], [4], [0], [3]]
            ],
            [
                [[5], [2], [2], [1]],
                [[3], [4], [3], [0]],
                [[2], [2], [1], [1]],
                [[3], [4], [2], [0]]
            ]
        ])

        backward_activation = np.array([
            [
                [[7], [2]],
                [[4], [3]]
            ],
            [
                [[1], [5]],
                [[2], [2]]
            ]
        ])

        expected_backward_result = np.array([
            [
                [[0], [0], [2], [0]],
                [[0], [7], [0], [0]],
                [[4], [0], [0], [0]],
                [[0], [0], [0], [3]]
            ],
            [
                [[1], [0], [0], [0]],
                [[0], [0], [5], [0]],
                [[0], [0], [0], [0]],
                [[0], [2], [2], [0]]
            ]
        ])

        # when
        layer = MaxPoolLayer(pool_size=pool_size, stride=stride)
        _ = layer.forward_pass(forward_activation, training=True)
        backward_result = layer.backward_pass(backward_activation)

        # then
        assert np.all(expected_backward_result == backward_result)
//  [python]
def test_stack_0(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=0)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=0)))
//  [python]
def test_stack_0(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=0)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=0)))
//  [python]
def test_stack_1(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=1)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=1)))
//  [python]
def test_stack_1(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=1)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=1)))
//  [python]
def test_stack_2(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=2)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=2)))
//  [python]
def test_stack_2(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=2)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=2)))
//  [python]
N = 1000 * 1000
rnd = numpy.random.randint(N, size=N)

for dtype1 in ('S6', 'b1',
               'i1', 'i2', 'i4', 'i8',
               'u1', 'u2', 'u4', 'u8', 'f4', 'f8'):
    for dtype2 in ('u4', 'i8'):
        print("dtype array1, array2-->", dtype1, dtype2)
        a = numpy.array(rnd, dtype1)
        b = numpy.arange(N, dtype=dtype2)
        c = a.copy()

        t1 = time()
        d = c.argsort()
        # c.sort()
        # e=c
        e = c[d]
        f = b[d]
        tref = time() - t1
        print("normal sort time-->", tref)

        t1 = time()
        keysort(a, b)
        tks = time() - t1
        print("keysort time-->", tks, "    %.2fx" % (tref / tks,))
        assert numpy.alltrue(a == e)
        #assert numpy.alltrue(b == d)
        assert numpy.alltrue(f == d)
//  [python]
N = 1000 * 1000
rnd = numpy.random.randint(N, size=N)

for dtype1 in ('S6', 'b1',
               'i1', 'i2', 'i4', 'i8',
               'u1', 'u2', 'u4', 'u8', 'f4', 'f8'):
    for dtype2 in ('u4', 'i8'):
        print("dtype array1, array2-->", dtype1, dtype2)
        a = numpy.array(rnd, dtype1)
        b = numpy.arange(N, dtype=dtype2)
        c = a.copy()

        t1 = time()
        d = c.argsort()
        # c.sort()
        # e=c
        e = c[d]
        f = b[d]
        tref = time() - t1
        print("normal sort time-->", tref)

        t1 = time()
        keysort(a, b)
        tks = time() - t1
        print("keysort time-->", tks, "    %.2fx" % (tref / tks,))
        assert numpy.all(a == e)
        #assert numpy.all(b == d)
        assert numpy.all(f == d)
//  [python]
def test_stack_3(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=3)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=3)))
//  [python]
def test_stack_3(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=3)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=3)))
//  [python]
def test_stack_x1(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-1)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=-1)))
//  [python]
def test_stack_x1(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-1)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=-1)))
//  [python]
def test_stack_x2(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-2)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=-2)))
//  [python]
def test_stack_x2(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-2)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=-2)))
//  [python]
def test_stack_x3(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-3)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=-3)))
//  [python]
def test_stack_x3(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-3)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=-3)))
//  [python]
def test_stack_x4(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-4)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.stack([a, b, c], axis=-4)))
//  [python]
def test_stack_x4(self):
      a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
      b = np.array([[[11, 12], [13, 14]], [[15, 16], [17, 18]]])
      c = np.array([[[21, 22], [23, 24]], [[25, 26], [27, 28]]])
      r = xdl.stack([a, b, c], axis=-4)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.stack([a, b, c], axis=-4)))
//  [python]
def test_move_to_numpy() -> None:
    in_dict = {"k1": torch.zeros(10), "k2": torch.ones(4)}
    for k in in_dict:
        assert isinstance(in_dict[k], torch.Tensor)

    out_dict = move_to_numpy(in_dict)
    assert np.alltrue(list(in_dict.keys()) == list(out_dict.keys()))
    for k in out_dict:
        assert isinstance(out_dict[k], np.ndarray)
//  [python]
def test_move_to_numpy() -> None:
    in_dict = {"k1": torch.zeros(10), "k2": torch.ones(4)}
    for k in in_dict:
        assert isinstance(in_dict[k], torch.Tensor)

    out_dict = move_to_numpy(in_dict)
    assert np.all(list(in_dict.keys()) == list(out_dict.keys()))
    for k in out_dict:
        assert isinstance(out_dict[k], np.ndarray)
//  [python]
def test_flatten_backward_pass_with_tree_axis(self):
        # given
        activation = np.random.rand(100, 28, 28)

        # when
        flatten_layer = FlattenLayer()
        forward_result = flatten_layer.forward_pass(activation, training=True)
        backward_result = flatten_layer.backward_pass(forward_result)

        # then
        assert np.alltrue(activation == backward_result)
//  [python]
def test_flatten_backward_pass_with_tree_axis(self):
        # given
        activation = np.random.rand(100, 28, 28)

        # when
        flatten_layer = FlattenLayer()
        forward_result = flatten_layer.forward_pass(activation, training=True)
        backward_result = flatten_layer.backward_pass(forward_result)

        # then
        assert np.all(activation == backward_result)
//  [python]
def test_move_to_device_trivial() -> None:
    in_dict = {"k1": torch.zeros(10), "k2": torch.ones(4)}
    for k in in_dict:
        assert isinstance(in_dict[k], torch.Tensor)
        assert in_dict[k].device == torch.device("cpu")

    out_dict = move_to_device(in_dict, torch.device("cpu"))
    assert np.alltrue(list(in_dict.keys()) == list(out_dict.keys()))
    for k in out_dict:
        assert isinstance(out_dict[k], torch.Tensor)
        assert out_dict[k].device == torch.device("cpu")
//  [python]
def test_move_to_device_trivial() -> None:
    in_dict = {"k1": torch.zeros(10), "k2": torch.ones(4)}
    for k in in_dict:
        assert isinstance(in_dict[k], torch.Tensor)
        assert in_dict[k].device == torch.device("cpu")

    out_dict = move_to_device(in_dict, torch.device("cpu"))
    assert np.all(list(in_dict.keys()) == list(out_dict.keys()))
    for k in out_dict:
        assert isinstance(out_dict[k], torch.Tensor)
        assert out_dict[k].device == torch.device("cpu")
//  [python]
def test_flatten_backward_pass_with_four_axis(self):
        # given
        activation = np.random.rand(100, 28, 28, 3)

        # when
        flatten_layer = FlattenLayer()
        forward_result = flatten_layer.forward_pass(activation, training=True)
        backward_result = flatten_layer.backward_pass(forward_result)

        # then
        assert np.alltrue(activation == backward_result)
//  [python]
def test_flatten_backward_pass_with_four_axis(self):
        # given
        activation = np.random.rand(100, 28, 28, 3)

        # when
        flatten_layer = FlattenLayer()
        forward_result = flatten_layer.forward_pass(activation, training=True)
        backward_result = flatten_layer.backward_pass(forward_result)

        # then
        assert np.all(activation == backward_result)
//  [python]
def calculate_mean_contour(contours):
    """Calculate the average of a set of contours, while retaining units and
    scaling information, if possible. If all contours have associated landmarks,
    then the average will be such a contour as well."""
    all_points = [c.points for c in contours]
    if not utils.all_same_shape(all_points):
        raise ContourError("Cannot calculate mean of contours with different numbers of points.")
    mean_points = numpy.mean(all_points, axis=0)
    units = [c.units for c in contours]
    if not numpy.alltrue([u == units[0] for u in units]):
        raise ContourError("All contours must have the same units in order calculate their mean.")
    units = contours[0].units
    scales = [utils.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
    if numpy.alltrue([numpy.allclose(scales[0], s) for s in scales[1:]]):
        transform = utils.make_homogenous_transform(transform=scales[0])
    else:
        transform = numpy.eye(3)
    if numpy.alltrue([isinstance(c, ContourAndLandmarks) for c in contours]):
        # if they're all landmark'd contours
        all_landmarks = [c.landmarks for c in contours]
        if not utils.all_same_shape(all_landmarks):
            raise ContourError("Cannot calculate mean of contours with different numbers of landmarks.")
        mean_landmarks = numpy.mean(all_landmarks, axis=0)
        mean_weights = numpy.mean([c.weights for c in contours], axis=0)
        return ContourAndLandmarks(
            points=mean_points,
            units=units,
            landmarks=mean_landmarks,
            weights=mean_weights,
            to_world_transform=transform,
        )
    else:
        return Contour(points=mean_points, units=units, to_world_transform=transform)
//  [python]
def calculate_mean_contour(contours):
    """Calculate the average of a set of contours, while retaining units and
    scaling information, if possible. If all contours have associated landmarks,
    then the average will be such a contour as well."""
    all_points = [c.points for c in contours]
    if not utils.all_same_shape(all_points):
        raise ContourError("Cannot calculate mean of contours with different numbers of points.")
    mean_points = numpy.mean(all_points, axis=0)
    units = [c.units for c in contours]
    if not numpy.all([u == units[0] for u in units]):
        raise ContourError("All contours must have the same units in order calculate their mean.")
    units = contours[0].units
    scales = [utils.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
    if numpy.all([numpy.allclose(scales[0], s) for s in scales[1:]]):
        transform = utils.make_homogenous_transform(transform=scales[0])
    else:
        transform = numpy.eye(3)
    if numpy.all([isinstance(c, ContourAndLandmarks) for c in contours]):
        # if they're all landmark'd contours
        all_landmarks = [c.landmarks for c in contours]
        if not utils.all_same_shape(all_landmarks):
            raise ContourError("Cannot calculate mean of contours with different numbers of landmarks.")
        mean_landmarks = numpy.mean(all_landmarks, axis=0)
        mean_weights = numpy.mean([c.weights for c in contours], axis=0)
        return ContourAndLandmarks(
            points=mean_points,
            units=units,
            landmarks=mean_landmarks,
            weights=mean_weights,
            to_world_transform=transform,
        )
    else:
        return Contour(points=mean_points, units=units, to_world_transform=transform)
//  [python]
def test_get_qm9_smiles():
    # test default behavior
    pp = AtomicNumberPreprocessor()
    dataset, smiles = qm9.get_qm9(preprocessor=pp, return_smiles=True)

    # --- Test dataset is correctly obtained ---
    index = numpy.random.choice(len(dataset), None)
    atoms, label = dataset[index]

    assert atoms.ndim == 1  # (atom, )
    assert atoms.dtype == numpy.int32
    # (atom from, atom to) or (edge_type, atom from, atom to)
    assert label.ndim == 1
    assert label.shape[0] == QM9_NUM_LABEL
    assert label.dtype == numpy.float32

    # --- Test number of dataset ---
    assert len(dataset) == QM9_NUM_DATASET
    assert len(smiles) == QM9_NUM_DATASET

    # --- Test order of dataset ---
    atoms0, labels0 = dataset[0]
    assert smiles[0] == 'C'
    assert numpy.alltrue(atoms0 == numpy.array([6], dtype=numpy.int32))

    atoms7777, labels7777 = dataset[7777]
    assert smiles[7777] == 'CC1=NCCC(C)O1'
    assert numpy.alltrue(
        atoms7777 == numpy.array([6, 6, 7, 6, 6, 6, 6, 8], dtype=numpy.int32))

    atoms133884, labels133884 = dataset[133884]
    assert smiles[133884] == 'C1N2C3C4C5OC13C2C54'
    assert numpy.alltrue(
        atoms133884 == numpy.array([6, 7, 6, 6, 6, 8, 6, 6, 6],
                                   dtype=numpy.int32))