jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_1_1(self):
        # Test PermuteRowElements(vector, vector)
        input = dvector()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(5,))
        p_val = rng.permutation(5).astype('int32')
        out_val = permute(input_val, p_val)

        # Should be equivalent to advanced indexing
        out_bis = input_val[p_val]
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_1(self):
        # Test PermuteRowElements(vector, vector)
        input = dvector()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(5,))
        p_val = rng.permutation(5).astype('int32')
        out_val = permute(input_val, p_val)

        # Should be equivalent to advanced indexing
        out_bis = input_val[p_val]
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_1(self):
        # Test broadcasting in PermuteRowElements(matrix, vector)
        input = matrix()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = rng.permutation(5).astype('int32')
        out_val = permute(input_val, p_val)

        # The same permutation should be applied to every row of the input matrix.
        out_bis = np.asarray([r[p_val] for r in input_val])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_1(self):
        # Test broadcasting in PermuteRowElements(matrix, vector)
        input = matrix()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = rng.permutation(5).astype('int32')
        out_val = permute(input_val, p_val)

        # The same permutation should be applied to every row of the input matrix.
        out_bis = np.asarray([r[p_val] for r in input_val])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_2(self):
        # Test PermuteRowElements(matrix, matrix)
        input = matrix()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)],
                           dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the corresponding
        # row of input
        out_bis = np.asarray([i_row[p_row]
                              for i_row, p_row in zip(input_val, p_val)])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_2(self):
        # Test PermuteRowElements(matrix, matrix)
        input = matrix()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)],
                           dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the corresponding
        # row of input
        out_bis = np.asarray([i_row[p_row]
                              for i_row, p_row in zip(input_val, p_val)])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_2(self):
        # Test PermuteRowElements(vector, matrix)
        # Different permutations will be applied to the same input vector
        input = vector()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(5,)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5)
                            for i in range(3)], dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the input vector
        out_bis = np.asarray([input_val[p_row] for p_row in p_val])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_2(self):
        # Test PermuteRowElements(vector, matrix)
        # Different permutations will be applied to the same input vector
        input = vector()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(5,)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5)
                            for i in range(3)], dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the input vector
        out_bis = np.asarray([input_val[p_row] for p_row in p_val])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_3b_2(self):
        # Test permute_row_elements on a more complex broadcasting pattern:
        # input.type.broadcastable = (False, True, False),
        # p.type.broadcastable = (False, False).

        input = TensorType('floatX', (False, True, False))()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(4, 1, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)],
                           dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to each row
        # of the input tensor
        out_bis = np.asarray([[in_mat[0, p_row] for p_row in p_val]
                              for in_mat in input_val])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_3b_2(self):
        # Test permute_row_elements on a more complex broadcasting pattern:
        # input.type.broadcastable = (False, True, False),
        # p.type.broadcastable = (False, False).

        input = TensorType('floatX', (False, True, False))()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.RandomState(utt.fetch_seed())
        input_val = rng.uniform(size=(4, 1, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)],
                           dtype='int32')
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to each row
        # of the input tensor
        out_bis = np.asarray([[in_mat[0, p_row] for p_row in p_val]
                              for in_mat in input_val])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)
        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_fit_int64(self):
        bitwidth = theano.configdefaults.python_int_bitwidth()
        for exponent in xrange(bitwidth):
            val = L(2 ** exponent - 1)
            scalar_ct = constant(val)

            assert scalar_ct.dtype in tensor.int_dtypes, (exponent, val, scalar_ct.dtype)
            assert scalar_ct.value == val

            vector_ct = constant([val, val])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if PY3 and bitwidth == 32:
                assert vector_ct.dtype == 'int32'
            else:
                assert vector_ct.dtype == 'int64'
            assert n.all(vector_ct.value == val)

            matrix_ct = constant([[val, val]])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if PY3 and bitwidth == 32:
                assert matrix_ct.dtype == 'int32'
            else:
                assert matrix_ct.dtype == 'int64'
            assert n.all(matrix_ct.value == val)
//  [python]
def test_fit_int64(self):
        bitwidth = theano.configdefaults.python_int_bitwidth()
        for exponent in xrange(bitwidth):
            val = L(2 ** exponent - 1)
            scalar_ct = constant(val)

            assert scalar_ct.dtype in tensor.int_dtypes, (exponent, val, scalar_ct.dtype)
            assert scalar_ct.value == val

            vector_ct = constant([val, val])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if PY3 and bitwidth == 32:
                assert vector_ct.dtype == 'int32'
            else:
                assert vector_ct.dtype == 'int64'
            assert np.all(vector_ct.value == val)

            matrix_ct = constant([[val, val]])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if PY3 and bitwidth == 32:
                assert matrix_ct.dtype == 'int32'
            else:
                assert matrix_ct.dtype == 'int64'
            assert np.all(matrix_ct.value == val)
//  [python]
def nband(self) -> int:
        try:
            return self._nband
        except AttributeError:
            assert n.all(self.nband_sk == self.nband_sk[0])
            self._nband = self.nband_sk[0, 0]
            return self._nband
//  [python]
def nband(self) -> int:
        try:
            return self._nband
        except AttributeError:
            assert np.all(self.nband_sk == self.nband_sk[0])
            self._nband = self.nband_sk[0, 0]
            return self._nband
//  [python]
def __call__(self, x1):
        if numpy.isscalar(x1) or x1.shape==():
            if x1 <= self.x[0]:
                return self._exfuncLow(x1)
            elif x1 >= self.x[-1]:
                return self._exfuncHigh(x1)
            else:
                return self._interpfunc(x1)
        lowmask = x1 <= self.x[0]
        highmask = x1 >= self.x[-1]
        inmask = numpy.logical_not(numpy.logical_or(lowmask,highmask))
        if numpy.all(inmask):
            return self._interpfunc(x1)
        
        y1 = numpy.empty(x1.shape)
        y1[inmask] = self._interpfunc(x1[inmask])
        y1[lowmask] = self._exfuncLow(x1[lowmask])
        y1[highmask] = self._exfuncHigh(x1[highmask])

        return y1
//  [python]
def test_alloc_diag_values(self):
        for x, test_val in self._generator():
            for offset, axis1, axis2 in [(0, 0, 1), (0, 1, 2), (1, 0, 1),
                                         (0, 1, 3), (0, 2, 3), (1, 2, 3),
                                         (-1, 0, 1), (-2, 0, 1), (-1, 1, 2)]:
                # Test AllocDiag values
                if np.maximum(axis1, axis2) > len(test_val.shape):
                    continue
                adiag_op = self.alloc_diag(offset=offset,
                                           axis1=axis1,
                                           axis2=axis2)
                f = theano.function([x], adiag_op(x))
                # AllocDiag and extract the diagonal again
                # to check
                diag_arr = f(test_val)
                rediag = np.diagonal(
                    diag_arr,
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                )
                assert n.all(rediag == test_val)

                # Test infer_shape
                f_shape = theano.function([x], adiag_op(x).shape, mode='FAST_RUN')

                theano.printing.debugprint(f_shape.maker.fgraph.outputs[0])
                output_shape = f_shape(test_val)
                assert not any(isinstance(node.op, self.alloc_diag)
                               for node in f_shape.maker.fgraph.toposort())
                rediag_shape = np.diagonal(
                    np.ones(output_shape),
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                ).shape
                assert n.all(rediag_shape == test_val.shape)

                diag_x = adiag_op(x)
                sum_diag_x = tensor.sum(diag_x)
                grad_x = tensor.grad(sum_diag_x, x)
                grad_diag_x = tensor.grad(sum_diag_x, diag_x)
                f_grad_x = theano.function([x], grad_x, mode=self.mode)
                f_grad_diag_x = theano.function([x], grad_diag_x, mode=self.mode)
                grad_input = f_grad_x(test_val)
                grad_diag_input = f_grad_diag_x(test_val)
                true_grad_input = np.diagonal(
                    grad_diag_input,
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                )

                assert n.all(true_grad_input == grad_input)
//  [python]
def test_alloc_diag_values(self):
        for x, test_val in self._generator():
            for offset, axis1, axis2 in [(0, 0, 1), (0, 1, 2), (1, 0, 1),
                                         (0, 1, 3), (0, 2, 3), (1, 2, 3),
                                         (-1, 0, 1), (-2, 0, 1), (-1, 1, 2)]:
                # Test AllocDiag values
                if np.maximum(axis1, axis2) > len(test_val.shape):
                    continue
                adiag_op = self.alloc_diag(offset=offset,
                                           axis1=axis1,
                                           axis2=axis2)
                f = theano.function([x], adiag_op(x))
                # AllocDiag and extract the diagonal again
                # to check
                diag_arr = f(test_val)
                rediag = np.diagonal(
                    diag_arr,
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                )
                assert np.all(rediag == test_val)

                # Test infer_shape
                f_shape = theano.function([x], adiag_op(x).shape, mode='FAST_RUN')

                theano.printing.debugprint(f_shape.maker.fgraph.outputs[0])
                output_shape = f_shape(test_val)
                assert not any(isinstance(node.op, self.alloc_diag)
                               for node in f_shape.maker.fgraph.toposort())
                rediag_shape = np.diagonal(
                    np.ones(output_shape),
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                ).shape
                assert np.all(rediag_shape == test_val.shape)

                diag_x = adiag_op(x)
                sum_diag_x = tensor.sum(diag_x)
                grad_x = tensor.grad(sum_diag_x, x)
                grad_diag_x = tensor.grad(sum_diag_x, diag_x)
                f_grad_x = theano.function([x], grad_x, mode=self.mode)
                f_grad_diag_x = theano.function([x], grad_diag_x, mode=self.mode)
                grad_input = f_grad_x(test_val)
                grad_diag_input = f_grad_diag_x(test_val)
                true_grad_input = np.diagonal(
                    grad_diag_input,
                    offset=offset,
                    axis1=axis1,
                    axis2=axis2
                )

                assert np.all(true_grad_input == grad_input)
//  [python]
def isnot_ibz_sampling(self, require_gamma_centered=False) -> str:
        """
        Test whether the k-points in the band structure represent an IBZ with an associated k-mesh
        Return string with error message if the condition is not fullfilled.

        Args:
            require_gamma_centered: True if the k-mesh should be Gamma-centered
        """
        errors = []
        eapp = errors.append

        if not self.kpoints.is_ibz:
            eapp("Expecting an IBZ sampling but got type `%s`" % type(self.kpoints))

        if not self.kpoints.is_mpmesh:
            eapp("Note that homogeneous k-meshes are required for the FS.\nksampling:` %s`" % str(self.kpoints.ksampling))

        if self.kpoints.is_mpmesh and require_gamma_centered:
            mpdivs, shifts = self.kpoints.mpdivs_shifts
            if shifts is not None and not n.all(shifts == 0.0):
                eapp(f"k-mesh should be gamma-centered but shifts: {shifts}")

        return "\n".join(errors)
//  [python]
def isnot_ibz_sampling(self, require_gamma_centered=False) -> str:
        """
        Test whether the k-points in the band structure represent an IBZ with an associated k-mesh
        Return string with error message if the condition is not fullfilled.

        Args:
            require_gamma_centered: True if the k-mesh should be Gamma-centered
        """
        errors = []
        eapp = errors.append

        if not self.kpoints.is_ibz:
            eapp("Expecting an IBZ sampling but got type `%s`" % type(self.kpoints))

        if not self.kpoints.is_mpmesh:
            eapp("Note that homogeneous k-meshes are required for the FS.\nksampling:` %s`" % str(self.kpoints.ksampling))

        if self.kpoints.is_mpmesh and require_gamma_centered:
            mpdivs, shifts = self.kpoints.mpdivs_shifts
            if shifts is not None and not np.all(shifts == 0.0):
                eapp(f"k-mesh should be gamma-centered but shifts: {shifts}")

        return "\n".join(errors)
//  [python]
def test_transpose():
    x1 = tensor.dvector('x1')
    x2 = tensor.dmatrix('x2')
    x3 = tensor.dtensor3('x3')

    x1v = np.arange(24)
    x2v = np.arange(24).reshape(2, 12)
    x3v = np.arange(24).reshape(2, 3, 4)

    f = theano.function([x1, x2, x3], [
        tensor.transpose(x1),
        tensor.transpose(x2),
        tensor.transpose(x3),
        x1.transpose(),
        x2.transpose(),
        x3.transpose(),
        x2.transpose(0, 1),
        x3.transpose((0, 2, 1)),
        tensor.transpose(x2, [0, 1]),
        tensor.transpose(x3, [0, 2, 1]),
        ])

    t1, t2, t3, t1b, t2b, t3b, t2c, t3c, t2d, t3d = f(x1v, x2v, x3v)
    assert t1.shape == np.transpose(x1v).shape
    assert t2.shape == np.transpose(x2v).shape
    assert t3.shape == np.transpose(x3v).shape
    assert np.all(t1 == np.transpose(x1v))
    assert np.all(t2 == np.transpose(x2v))
    assert np.all(t3 == np.transpose(x3v))
    assert np.all(t1b == x1v.transpose())
    assert np.all(t2b == x2v.transpose())
    assert np.all(t3b == x3v.transpose())
    assert t2c.shape == (2, 12)
    assert t3c.shape == (2, 4, 3)
    assert np.all(t2c == x2v.transpose([0, 1]))
    assert np.all(t3c == x3v.transpose([0, 2, 1]))
    assert t2d.shape == (2, 12)
    assert t3d.shape == (2, 4, 3)
    assert np.all(t2d == np.transpose(x2v, [0, 1]))
    assert np.all(t3d == np.transpose(x3v, [0, 2, 1]))

    # Check that we create a name.
    assert tensor.transpose(x1).name == 'x1.T'
    assert tensor.transpose(x2).name == 'x2.T'
    assert tensor.transpose(x3).name == 'x3.T'
    assert tensor.transpose(tensor.dmatrix()).name is None
//  [python]
def test_iid_args_kwargs(bs_setup):
    bs1 = IIDBootstrap(bs_setup.y, seed=0)
    bs2 = IIDBootstrap(y=bs_setup.y, seed=0)
    for a, b in zip(bs1.bootstrap(1), bs2.bootstrap(1)):
        assert n.all(a[0][0] == b[1]["y"])
//  [python]
def test_iid_args_kwargs(bs_setup):
    bs1 = IIDBootstrap(bs_setup.y, seed=0)
    bs2 = IIDBootstrap(y=bs_setup.y, seed=0)
    for a, b in zip(bs1.bootstrap(1), bs2.bootstrap(1)):
        assert np.all(a[0][0] == b[1]["y"])
//  [python]
def to_pymatgen(self):
        """
        Return a pymatgen band structure object from an Abipy |ElectronBands| object.
        """
        from pymatgen.electronic_structure.bandstructure import BandStructure, BandStructureSymmLine
        assert n.all(self.nband_sk == self.nband_sk[0, 0])

        # eigenvals is a dict of energies for spin up and spin down
        # {Spin.up:[][], Spin.down:[][]}, the first index of the array
        # [][] refers to the band and the second to the index of the
        # kpoint. The kpoints are ordered according to the order of the
        # kpoints array. If the band structure is not spin polarized, we
        # only store one data set under Spin.up
        eigenvals = {PmgSpin.up: self.eigens[0,:,:].T.copy().tolist()}
        if self.nsppol == 2:
            eigenvals[PmgSpin.down] = self.eigens[1,:,:].T.copy().tolist()

        if self.kpoints.is_path:
            labels_dict = {k.name: k.frac_coords for k in self.kpoints if k.name is not None}
            return BandStructureSymmLine(self.kpoints.frac_coords, eigenvals, self.reciprocal_lattice, self.fermie,
                                         labels_dict, coords_are_cartesian=False,
                                         structure=self.structure, projections=None)
        else:
            return BandStructure(self.kpoints.frac_coords, eigenvals, self.reciprocal_lattice, self.fermie,
                                 labels_dict=None, coords_are_cartesian=False,
                                 structure=self.structure, projections=None)
//  [python]
def to_pymatgen(self):
        """
        Return a pymatgen band structure object from an Abipy |ElectronBands| object.
        """
        from pymatgen.electronic_structure.bandstructure import BandStructure, BandStructureSymmLine
        assert np.all(self.nband_sk == self.nband_sk[0, 0])

        # eigenvals is a dict of energies for spin up and spin down
        # {Spin.up:[][], Spin.down:[][]}, the first index of the array
        # [][] refers to the band and the second to the index of the
        # kpoint. The kpoints are ordered according to the order of the
        # kpoints array. If the band structure is not spin polarized, we
        # only store one data set under Spin.up
        eigenvals = {PmgSpin.up: self.eigens[0,:,:].T.copy().tolist()}
        if self.nsppol == 2:
            eigenvals[PmgSpin.down] = self.eigens[1,:,:].T.copy().tolist()

        if self.kpoints.is_path:
            labels_dict = {k.name: k.frac_coords for k in self.kpoints if k.name is not None}
            return BandStructureSymmLine(self.kpoints.frac_coords, eigenvals, self.reciprocal_lattice, self.fermie,
                                         labels_dict, coords_are_cartesian=False,
                                         structure=self.structure, projections=None)
        else:
            return BandStructure(self.kpoints.frac_coords, eigenvals, self.reciprocal_lattice, self.fermie,
                                 labels_dict=None, coords_are_cartesian=False,
                                 structure=self.structure, projections=None)
//  [python]
def numba_isfinite(a):
    return n.all(np.isfinite(a))
//  [python]
def numba_isfinite(a):
    return np.all(np.isfinite(a))
//  [python]
def __init__(self, structure, ibz, has_timrev, eigens, fermie):
        """
        This object reconstructs by symmetry the eigenvalues in the full BZ starting from the IBZ.
        Provides methods to extract and visualize isosurfaces.

        Args:
            structure:
            ibz:
            has_timrev:
            eigens:
            fermie
        """
        self.ibz = ibz
        self._structure = structure
        self.reciprocal_lattice = structure.lattice.reciprocal_lattice
        self.has_timrev = has_timrev
        self.fermie = fermie
        self.eigens = np.atleast_3d(eigens)
        self.nsppol, _, self.nband = self.eigens.shape

        # Sanity check.
        errors = []; eapp = errors.append
        if not self.ibz.is_ibz:
            eapp("Expecting an IBZ sampling but got %s" % type(self.ibz))
        if not self.ibz.is_mpmesh:
            eapp("Monkhorst-Pack meshes are required.\nksampling: %s" % str(self.ibz.ksampling))

        mpdivs, shifts = self.ibz.mpdivs_shifts
        if shifts is not None and not n.all(shifts == 0.0):
            eapp("Gamma-centered k-meshes are required by Xcrysden.")
        if errors:
            raise ValueError("\n".join(errors))

        # Xcrysden requires points in the unit cell (C-order)
        # and the mesh must include the periodic images hence pbc=True.
        self.uc2ibz = map_grid2ibz(self.structure, self.ibz.frac_coords, mpdivs, self.has_timrev, pbc=True)
        self.mpdivs = mpdivs
        self.kdivs = mpdivs + 1
        self.spacing = 1.0 / mpdivs
        self.ucdata_shape = (self.nsppol, self.nband) + tuple(self.kdivs)
        #self.ibzdata_shape = (self.nsppol, self.nband, len(self.ibz))

        # Construct energy bands on unit cell grid: e_{TSk} = e_{k}
        self.ucdata_sbk = self.symmetrize_ibz_scalars(self.eigens)

        self.ucell_scalars = OrderedDict()
        self.ucell_vectors = OrderedDict()
//  [python]
def __init__(self, structure, ibz, has_timrev, eigens, fermie):
        """
        This object reconstructs by symmetry the eigenvalues in the full BZ starting from the IBZ.
        Provides methods to extract and visualize isosurfaces.

        Args:
            structure:
            ibz:
            has_timrev:
            eigens:
            fermie
        """
        self.ibz = ibz
        self._structure = structure
        self.reciprocal_lattice = structure.lattice.reciprocal_lattice
        self.has_timrev = has_timrev
        self.fermie = fermie
        self.eigens = np.atleast_3d(eigens)
        self.nsppol, _, self.nband = self.eigens.shape

        # Sanity check.
        errors = []; eapp = errors.append
        if not self.ibz.is_ibz:
            eapp("Expecting an IBZ sampling but got %s" % type(self.ibz))
        if not self.ibz.is_mpmesh:
            eapp("Monkhorst-Pack meshes are required.\nksampling: %s" % str(self.ibz.ksampling))

        mpdivs, shifts = self.ibz.mpdivs_shifts
        if shifts is not None and not np.all(shifts == 0.0):
            eapp("Gamma-centered k-meshes are required by Xcrysden.")
        if errors:
            raise ValueError("\n".join(errors))

        # Xcrysden requires points in the unit cell (C-order)
        # and the mesh must include the periodic images hence pbc=True.
        self.uc2ibz = map_grid2ibz(self.structure, self.ibz.frac_coords, mpdivs, self.has_timrev, pbc=True)
        self.mpdivs = mpdivs
        self.kdivs = mpdivs + 1
        self.spacing = 1.0 / mpdivs
        self.ucdata_shape = (self.nsppol, self.nband) + tuple(self.kdivs)
        #self.ibzdata_shape = (self.nsppol, self.nband, len(self.ibz))

        # Construct energy bands on unit cell grid: e_{TSk} = e_{k}
        self.ucdata_sbk = self.symmetrize_ibz_scalars(self.eigens)

        self.ucell_scalars = OrderedDict()
        self.ucell_vectors = OrderedDict()
//  [python]
def _not_all(array, axis):
    return ~n.all(array, axis)
//  [python]
def _not_all(array, axis):
    return ~numpy.all(array, axis)
//  [python]
def test_conf_int_parametric(bs_setup):
    def param_func(x, params=None, state=None):
        if state is not None:
            mu = params
            e = state.standard_normal(x.shape)
            return (mu + e).mean(0)
        else:
            return x.mean(0)

    def semi_func(x, params=None):
        if params is not None:
            mu = params
            e = x - mu
            return (mu + e).mean(0)
        else:
            return x.mean(0)

    reps = 100
    bs = IIDBootstrap(bs_setup.x)

    ci = bs.conf_int(func=param_func, reps=reps, sampling="parametric")
    assert len(ci) == 2
    assert np.all(ci[0] < ci[1])
    bs.reset()
    results = np.zeros((reps, 2))
    count = 0
    mu = bs_setup.x.mean(0)
    for pos, _ in bs.bootstrap(100):
        results[count] = param_func(*pos, params=mu, state=bs.generator)
        count += 1
    assert_equal(bs._results, results)

    bs.reset()
    ci = bs.conf_int(func=semi_func, reps=100, sampling="semi")
    assert len(ci) == 2
    assert np.all(ci[0] < ci[1])
    bs.reset()
    results = np.zeros((reps, 2))
    count = 0
    for pos, _ in bs.bootstrap(100):
        results[count] = semi_func(*pos, params=mu)
        count += 1
    assert_allclose(bs._results, results)

    with pytest.raises(ValueError, match="sampling must be one"):
        bs.conf_int(func=semi_func, reps=100, sampling="other")
//  [python]
def __init__(self, limits, breaks, values):
		assert len(breaks) == len(limits), [len(breaks), len(limits)]
		
		gridall = []
		for l, b in zip(limits, breaks):
			assert b == [] or n.all(n.asarray(b) > l[0]), [b, l[0]]
			assert b == [] or n.all(n.asarray(b) < l[1]), [b, l[1]]
			grid = n.array([l[0]] + list(b) + [l[1]])
			gridall.append(grid)
			assert n.all(grid[1:] > grid[:-1]), 'breaks need to be ascending'
		self.grid = gridall
		assert values.shape == tuple([len(b)+2 for b in breaks]), [
			values.shape, [len(b)+2 for b in breaks]]
		#assert len(values) == n.product([len(b)+2 for b in breaks])
		self.values = values
//  [python]
def __init__(self, limits, breaks, values):
		assert len(breaks) == len(limits), [len(breaks), len(limits)]
		
		gridall = []
		for l, b in zip(limits, breaks):
			assert b == [] or numpy.all(numpy.asarray(b) > l[0]), [b, l[0]]
			assert b == [] or numpy.all(numpy.asarray(b) < l[1]), [b, l[1]]
			grid = numpy.array([l[0]] + list(b) + [l[1]])
			gridall.append(grid)
			assert numpy.all(grid[1:] > grid[:-1]), 'breaks need to be ascending'
		self.grid = gridall
		assert values.shape == tuple([len(b)+2 for b in breaks]), [
			values.shape, [len(b)+2 for b in breaks]]
		#assert len(values) == numpy.product([len(b)+2 for b in breaks])
		self.values = values
//  [python]
def _stride_comb_iter(x):
    """
    Generate cartesian product of strides for all axes
    """

    if not isinstance(x, np.ndarray):
        yield x, "nop"
        return

    stride_set = [(1,)] * x.ndim
    stride_set[-1] = (1, 3, -4)
    if x.ndim > 1:
        stride_set[-2] = (1, 3, -4)
    if x.ndim > 2:
        stride_set[-3] = (1, -4)

    for repeats in itertools.product(*tuple(stride_set)):
        new_shape = [abs(a * b) for a, b in zip(x.shape, repeats)]
        slices = tuple([slice(None, None, repeat) for repeat in repeats])

        # new array with different strides, but same data
        xi = np.empty(new_shape, dtype=x.dtype)
        xi.view(np.uint32).fill(0xdeadbeef)
        xi = xi[slices]
        xi[...] = x
        xi = xi.view(x.__class__)
        assert n.all(xi == x)
        yield xi, "stride_" + "_".join(["%+d" % j for j in repeats])

        # generate also zero strides if possible
        if x.ndim >= 1 and x.shape[-1] == 1:
            s = list(x.strides)
            s[-1] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0"
            except NotImplementedError:
                pass
        if x.ndim >= 2 and x.shape[-2] == 1:
            s = list(x.strides)
            s[-2] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0_x"
            except NotImplementedError:
                pass
        if x.ndim >= 2 and x.shape[:-2] == (1, 1):
            s = list(x.strides)
            s[-1] = 0
            s[-2] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0_0"
            except NotImplementedError:
                pass
//  [python]
def _stride_comb_iter(x):
    """
    Generate cartesian product of strides for all axes
    """

    if not isinstance(x, np.ndarray):
        yield x, "nop"
        return

    stride_set = [(1,)] * x.ndim
    stride_set[-1] = (1, 3, -4)
    if x.ndim > 1:
        stride_set[-2] = (1, 3, -4)
    if x.ndim > 2:
        stride_set[-3] = (1, -4)

    for repeats in itertools.product(*tuple(stride_set)):
        new_shape = [abs(a * b) for a, b in zip(x.shape, repeats)]
        slices = tuple([slice(None, None, repeat) for repeat in repeats])

        # new array with different strides, but same data
        xi = np.empty(new_shape, dtype=x.dtype)
        xi.view(np.uint32).fill(0xdeadbeef)
        xi = xi[slices]
        xi[...] = x
        xi = xi.view(x.__class__)
        assert np.all(xi == x)
        yield xi, "stride_" + "_".join(["%+d" % j for j in repeats])

        # generate also zero strides if possible
        if x.ndim >= 1 and x.shape[-1] == 1:
            s = list(x.strides)
            s[-1] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0"
            except NotImplementedError:
                pass
        if x.ndim >= 2 and x.shape[-2] == 1:
            s = list(x.strides)
            s[-2] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0_x"
            except NotImplementedError:
                pass
        if x.ndim >= 2 and x.shape[:-2] == (1, 1):
            s = list(x.strides)
            s[-1] = 0
            s[-2] = 0
            try:
                xi = np.lib.stride_tricks.as_strided(x, strides=s)
                yield xi, "stride_xxx_0_0"
            except NotImplementedError:
                pass
//  [python]
def get_start_index(data):
    if n.all(data['vehicle_speed'] == 0):
        return 0

    start_ind = np.where(data['brake_percentage'] == 40)

    if len(start_ind[0] > 0):
        ind = start_ind[0][0]
        while ind < len(data):
            if data['brake_percentage'][ind] == 40:
                ind += 1
            else:
                break
        return ind
    else:
        ind = 0
        while ind < len(data):
            if abs(data['vehicle_speed'][ind]) < 0.01:
                ind += 1
            else:
                break
        return ind
//  [python]
def get_start_index(data):
    if np.all(data['vehicle_speed'] == 0):
        return 0

    start_ind = np.where(data['brake_percentage'] == 40)

    if len(start_ind[0] > 0):
        ind = start_ind[0][0]
        while ind < len(data):
            if data['brake_percentage'][ind] == 40:
                ind += 1
            else:
                break
        return ind
    else:
        ind = 0
        while ind < len(data):
            if abs(data['vehicle_speed'][ind]) < 0.01:
                ind += 1
            else:
                break
        return ind
//  [python]
def do_test_with_param(
            self, index_key, ps_params, params):
        """
        Test equivalence between setting
        1. param_name_2 = value with ParameterSpace
        2. pass in a SearchParameters with param_name = value
        """
        ds = datasets.SyntheticDataset(32, 1000, 100, 20)
        index = faiss.index_factory(ds.d, index_key)
        if index_key.startswith("PQ"):
            index.polysemous_training.n_iter = 50000
            index.polysemous_training.n_redo = 1
        index.train(ds.get_train())
        index.add(ds.get_database())

        I0, D0 = index.search(ds.get_queries(), 10)

        Dnew, Inew = index.search(ds.get_queries(), 10, params=params)

        # make sure rhe parameter does indeed change the result...
        self.assertFalse(n.all(Inew == I0))

        for param_name, value in ps_params.items():
            faiss.ParameterSpace().set_index_parameter(
                index, param_name, value)
        Dref, Iref = index.search(ds.get_queries(), 10)

        np.testing.assert_array_equal(Iref, Inew)
        np.testing.assert_array_equal(Dref, Dnew)
//  [python]
def do_test_with_param(
            self, index_key, ps_params, params):
        """
        Test equivalence between setting
        1. param_name_2 = value with ParameterSpace
        2. pass in a SearchParameters with param_name = value
        """
        ds = datasets.SyntheticDataset(32, 1000, 100, 20)
        index = faiss.index_factory(ds.d, index_key)
        if index_key.startswith("PQ"):
            index.polysemous_training.n_iter = 50000
            index.polysemous_training.n_redo = 1
        index.train(ds.get_train())
        index.add(ds.get_database())

        I0, D0 = index.search(ds.get_queries(), 10)

        Dnew, Inew = index.search(ds.get_queries(), 10, params=params)

        # make sure rhe parameter does indeed change the result...
        self.assertFalse(np.all(Inew == I0))

        for param_name, value in ps_params.items():
            faiss.ParameterSpace().set_index_parameter(
                index, param_name, value)
        Dref, Iref = index.search(ds.get_queries(), 10)

        np.testing.assert_array_equal(Iref, Inew)
        np.testing.assert_array_equal(Dref, Dnew)
//  [python]
def combine_restricted_orbitals(self):
    '''Combines unrestricted orbitals with same MO coefficients to restricted orbitals.
    (For example Psi4 creates a molden file with unrestricted orbitals after a ROHF calculation.)
    '''

    if not self.spinpolarized:
      # orbitals are already restricted
      return

    # assert same energy
    ener_a, ener_b = self.eig[self.alpha_index], self.eig[self.beta_index]
    assert n.all(ener_a == ener_b), "MO eigenvalues do not match"

    # assert same MO coefficients
    coeff_a, coeff_b = self.coeffs[self.alpha_index], self.coeffs[self.beta_index]
    assert n.all(coeff_a == coeff_b), "MO coeffs do not match"

    # add occupation and select alpha orbitals only
    occ = self.occ[self.alpha_index] + self.occ[self.beta_index]
    self.data = self[self.alpha_index].data
    self.spinpolarized = False
    self.update()
    self.set_occ(occ)
    self.update()
//  [python]
def combine_restricted_orbitals(self):
    '''Combines unrestricted orbitals with same MO coefficients to restricted orbitals.
    (For example Psi4 creates a molden file with unrestricted orbitals after a ROHF calculation.)
    '''

    if not self.spinpolarized:
      # orbitals are already restricted
      return

    # assert same energy
    ener_a, ener_b = self.eig[self.alpha_index], self.eig[self.beta_index]
    assert numpy.all(ener_a == ener_b), "MO eigenvalues do not match"

    # assert same MO coefficients
    coeff_a, coeff_b = self.coeffs[self.alpha_index], self.coeffs[self.beta_index]
    assert numpy.all(coeff_a == coeff_b), "MO coeffs do not match"

    # add occupation and select alpha orbitals only
    occ = self.occ[self.alpha_index] + self.occ[self.beta_index]
    self.data = self[self.alpha_index].data
    self.spinpolarized = False
    self.update()
    self.set_occ(occ)
    self.update()
//  [python]
def testRange(self):
    """Tests against low MS-SSIM score.

    MS-SSIM is a geometric mean of SSIM and CS scores of various scales.
    If any of the value is negative so that the geometric mean is not
    well-defined, then treat the MS-SSIM score as zero.
    """
    with self.cached_session() as sess:
      img1 = self._LoadTestImage(sess, "checkerboard1.png")
      img2 = self._LoadTestImage(sess, "checkerboard3.png")
      images = [img1, img2, np.zeros_like(img1),
                np.full_like(img1, fill_value=255)]

      images = [ops.convert_to_tensor(x, dtype=dtypes.float32) for x in images]
      msssim_ops = [
          image_ops.ssim_multiscale(
              x, y, 1.0, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
          for x, y in itertools.combinations(images, 2)
      ]
      msssim = self.evaluate(msssim_ops)
      msssim = np.squeeze(msssim)

    self.assertTrue(np.alltrue(msssim >= 0.0))
    self.assertTrue(np.alltrue(msssim <= 1.0))
//  [python]
def testRange(self):
    """Tests against low MS-SSIM score.

    MS-SSIM is a geometric mean of SSIM and CS scores of various scales.
    If any of the value is negative so that the geometric mean is not
    well-defined, then treat the MS-SSIM score as zero.
    """
    with self.cached_session() as sess:
      img1 = self._LoadTestImage(sess, "checkerboard1.png")
      img2 = self._LoadTestImage(sess, "checkerboard3.png")
      images = [img1, img2, np.zeros_like(img1),
                np.full_like(img1, fill_value=255)]

      images = [ops.convert_to_tensor(x, dtype=dtypes.float32) for x in images]
      msssim_ops = [
          image_ops.ssim_multiscale(
              x, y, 1.0, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
          for x, y in itertools.combinations(images, 2)
      ]
      msssim = self.evaluate(msssim_ops)
      msssim = np.squeeze(msssim)

    self.assertTrue(np.all(msssim >= 0.0))
    self.assertTrue(np.all(msssim <= 1.0))
//  [python]
def coerce_to_array(
    values, mask=None, copy: bool = False
) -> tuple[np.ndarray, np.ndarray]:
    """
    Coerce the input values array to numpy arrays with a mask.

    Parameters
    ----------
    values : 1D list-like
    mask : bool 1D array, optional
    copy : bool, default False
        if True, copy the input

    Returns
    -------
    tuple of (values, mask)
    """
    if isinstance(values, BooleanArray):
        if mask is not None:
            raise ValueError("cannot pass mask for BooleanArray input")
        values, mask = values._data, values._mask
        if copy:
            values = values.copy()
            mask = mask.copy()
        return values, mask

    mask_values = None
    if isinstance(values, np.ndarray) and values.dtype == np.bool_:
        if copy:
            values = values.copy()
    elif isinstance(values, np.ndarray) and values.dtype.kind in "iufcb":
        mask_values = isna(values)

        values_bool = np.zeros(len(values), dtype=bool)
        values_bool[~mask_values] = values[~mask_values].astype(bool)

        if not np.all(
            values_bool[~mask_values].astype(values.dtype) == values[~mask_values]
        ):
            raise TypeError("Need to pass bool-like values")

        values = values_bool
    else:
        values_object = np.asarray(values, dtype=object)

        inferred_dtype = lib.infer_dtype(values_object, skipna=True)
        integer_like = ("floating", "integer", "mixed-integer-float")
        if inferred_dtype not in ("boolean", "empty") + integer_like:
            raise TypeError("Need to pass bool-like values")

        # mypy does not narrow the type of mask_values to npt.NDArray[np.bool_]
        # within this branch, it assumes it can also be None
        mask_values = cast("npt.NDArray[np.bool_]", isna(values_object))
        values = np.zeros(len(values), dtype=bool)
        values[~mask_values] = values_object[~mask_values].astype(bool)

        # if the values were integer-like, validate it were actually 0/1's
        if (inferred_dtype in integer_like) and not (
            np.all(
                values[~mask_values].astype(float)
                == values_object[~mask_values].astype(float)
            )
        ):
            raise TypeError("Need to pass bool-like values")

    if mask is None and mask_values is None:
        mask = np.zeros(values.shape, dtype=bool)
    elif mask is None:
        mask = mask_values
    else:
        if isinstance(mask, np.ndarray) and mask.dtype == np.bool_:
            if mask_values is not None:
                mask = mask | mask_values
            else:
                if copy:
                    mask = mask.copy()
        else:
            mask = np.array(mask, dtype=bool)
            if mask_values is not None:
                mask = mask | mask_values

    if values.shape != mask.shape:
        raise ValueError("values.shape and mask.shape must match")

    return values, mask
//  [python]
def testInterpGradNan(self, dtype, period, left, right):
    kwds = dict(period=period, left=left, right=right)
    jnp_fun = partial(jnp.interp, **kwds)
    # Probe values of x and xp that are close to zero and close together.
    x = dtype(np.exp(np.linspace(-90, -20, 1000)))
    g = jax.grad(lambda z: jnp.sum(jnp_fun(z, z, jnp.ones_like(z))))(x)
    np.testing.assert_equal(np.all(np.isfinite(g)), True)
//  [python]
def test_module_input_grads():
    a = mx.sym.Variable('a', __layout__='NC')
    b = mx.sym.Variable('b', __layout__='NC')
    c = mx.sym.Variable('c', __layout__='NC')

    c = a + 2 * b + 3 * c
    net = mx.mod.Module(c, data_names=['b', 'c', 'a'], label_names=None,
                        context=[mx.cpu(0), mx.cpu(1)])
    net.bind(data_shapes=[['b', (5, 5)], ['c', (5, 5)], ['a', (5, 5)]],
             label_shapes=None, inputs_need_grad=True)
    net.init_params()

    net.forward(data_batch=mx.io.DataBatch(data=[nd.ones((5, 5)),
                                                 nd.ones((5, 5)),
                                                 nd.ones((5, 5))]))
    net.backward(out_grads=[nd.ones((5, 5))])
    input_grads = net.get_input_grads()
    b_grad = input_grads[0].asnumpy()
    c_grad = input_grads[1].asnumpy()
    a_grad = input_grads[2].asnumpy()
    assert np.all(a_grad == 1), a_grad
    assert np.all(b_grad == 2), b_grad
    assert np.all(c_grad == 3), c_grad