jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def random_pick_and_place(pick_num=10, pick_scale=0.01):
    """ Random pick a particle up and the drop it for pick_num times"""
    curr_pos = pyflex.get_positions().reshape(-1, 4)
    num_particles = curr_pos.shape[0]
    for i in range(pick_num):
        pick_id = np.random.randint(num_particles)
        pick_dir = np.random.random(3) * 2 - 1
        pick_dir[1] = (pick_dir[1] + 1)
        pick_dir *= pick_scale
        original_inv_mass = curr_pos[pick_id, 3]
        for _ in range(60):
            curr_pos = pyflex.get_positions().reshape(-1, 4)
            curr_pos[pick_id, :3] += pick_dir
            curr_pos[pick_id, 3] = 0
            pyflex.set_positions(curr_pos.flatten())
            pyflex.step()

        # Revert mass
        curr_pos = pyflex.get_positions().reshape(-1, 4)
        curr_pos[pick_id, 3] = original_inv_mass
        pyflex.set_positions(curr_pos.flatten())
        pyflex.step()

        # Wait to stabalize
        for _ in range(100):
            pyflex.step()
            curr_vel = pyflex.get_velocities()
            if np.alltrue(curr_vel < 0.01):
                break
    for _ in range(500):
        pyflex.step()
        curr_vel = pyflex.get_velocities()
        if np.alltrue(curr_vel < 0.01):
            break
//  [python]
def random_pick_and_place(pick_num=10, pick_scale=0.01):
    """ Random pick a particle up and the drop it for pick_num times"""
    curr_pos = pyflex.get_positions().reshape(-1, 4)
    num_particles = curr_pos.shape[0]
    for i in range(pick_num):
        pick_id = np.random.randint(num_particles)
        pick_dir = np.random.random(3) * 2 - 1
        pick_dir[1] = (pick_dir[1] + 1)
        pick_dir *= pick_scale
        original_inv_mass = curr_pos[pick_id, 3]
        for _ in range(60):
            curr_pos = pyflex.get_positions().reshape(-1, 4)
            curr_pos[pick_id, :3] += pick_dir
            curr_pos[pick_id, 3] = 0
            pyflex.set_positions(curr_pos.flatten())
            pyflex.step()

        # Revert mass
        curr_pos = pyflex.get_positions().reshape(-1, 4)
        curr_pos[pick_id, 3] = original_inv_mass
        pyflex.set_positions(curr_pos.flatten())
        pyflex.step()

        # Wait to stabalize
        for _ in range(100):
            pyflex.step()
            curr_vel = pyflex.get_velocities()
            if np.all(curr_vel < 0.01):
                break
    for _ in range(500):
        pyflex.step()
        curr_vel = pyflex.get_velocities()
        if np.all(curr_vel < 0.01):
            break
//  [python]
def test_alltrue(self):
        with pytest.raises(TypeError):
            np.alltrue(self.q)
//  [python]
def test_alltrue(self):
        with pytest.raises(TypeError):
            np.all(self.q)
//  [python]
def testChunking(self):
        # Override testChunking.

        index = self.cls(corpus, self.similarity_matrix)
        query = [dictionary.doc2bow(document) for document in texts[:3]]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.alltrue(sims[i, i] == 1.0))  # Similarity of a document with itself is 1.0.

        # test the same thing but with num_best
        index.num_best = 5
        sims = index[query]
        for i, chunk in enumerate(sims):
            expected = i
            self.assertAlmostEqual(expected, chunk[0][0], places=2)
            expected = 1.0
            self.assertAlmostEqual(expected, chunk[0][1], places=2)
//  [python]
def testChunking(self):
        # Override testChunking.

        index = self.cls(corpus, self.similarity_matrix)
        query = [dictionary.doc2bow(document) for document in texts[:3]]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.all(sims[i, i] == 1.0))  # Similarity of a document with itself is 1.0.

        # test the same thing but with num_best
        index.num_best = 5
        sims = index[query]
        for i, chunk in enumerate(sims):
            expected = i
            self.assertAlmostEqual(expected, chunk[0][0], places=2)
            expected = 1.0
            self.assertAlmostEqual(expected, chunk[0][1], places=2)
//  [python]
def __reader(self):
        # Load if not already loaded
        if self._reader_ is None:
            if not os.path.isfile(self.filename):
                raise FileNotFoundError(
                    f"Could not find filename video filename named {self.filename}"
                )

            # Try and open the file either locally in current directory or with full
            # path
            self._reader_ = cv2.VideoCapture(self.filename)

            # If the user specified None for grayscale bool, figure it out based on the
            # the first frame of data.
            if self._detect_grayscale is True:
                self.grayscale = bool(
                    np.alltrue(self.test_frame[..., 0] == self.test_frame[..., -1])
                )

        # Return cached reader
        return self._reader_
//  [python]
def __reader(self):
        # Load if not already loaded
        if self._reader_ is None:
            if not os.path.isfile(self.filename):
                raise FileNotFoundError(
                    f"Could not find filename video filename named {self.filename}"
                )

            # Try and open the file either locally in current directory or with full
            # path
            self._reader_ = cv2.VideoCapture(self.filename)

            # If the user specified None for grayscale bool, figure it out based on the
            # the first frame of data.
            if self._detect_grayscale is True:
                self.grayscale = bool(
                    np.all(self.test_frame[..., 0] == self.test_frame[..., -1])
                )

        # Return cached reader
        return self._reader_
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        # Single query
        index = self.cls(corpus, self.similarity_matrix, num_best=num_best)
        query = dictionary.doc2bow(texts[0])
        sims = index[query]
        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.alltrue(sim <= 1.0))
                self.assertTrue(numpy.alltrue(sim >= 0.0))
        else:
            self.assertAlmostEqual(1.0, sims[0])  # Similarity of a document with itself is 1.0.
            self.assertTrue(numpy.alltrue(sims[1:] >= 0.0))
            self.assertTrue(numpy.alltrue(sims[1:] < 1.0))

        # Corpora
        for query in (
                corpus,  # Basic text corpus.
                self.tfidf[corpus]):  # Transformed corpus without slicing support.
            index = self.cls(query, self.similarity_matrix, num_best=num_best)
            sims = index[query]
            if num_best is not None:
                # Sparse array.
                for result in sims:
                    for i, sim in result:
                        self.assertTrue(numpy.alltrue(sim <= 1.0))
                        self.assertTrue(numpy.alltrue(sim >= 0.0))
            else:
                for i, result in enumerate(sims):
                    self.assertAlmostEqual(1.0, result[i])  # Similarity of a document with itself is 1.0.
                    self.assertTrue(numpy.alltrue(result[:i] >= 0.0))
                    self.assertTrue(numpy.alltrue(result[:i] < 1.0))
                    self.assertTrue(numpy.alltrue(result[i + 1:] >= 0.0))
                    self.assertTrue(numpy.alltrue(result[i + 1:] < 1.0))
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        # Single query
        index = self.cls(corpus, self.similarity_matrix, num_best=num_best)
        query = dictionary.doc2bow(texts[0])
        sims = index[query]
        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.all(sim <= 1.0))
                self.assertTrue(numpy.all(sim >= 0.0))
        else:
            self.assertAlmostEqual(1.0, sims[0])  # Similarity of a document with itself is 1.0.
            self.assertTrue(numpy.all(sims[1:] >= 0.0))
            self.assertTrue(numpy.all(sims[1:] < 1.0))

        # Corpora
        for query in (
                corpus,  # Basic text corpus.
                self.tfidf[corpus]):  # Transformed corpus without slicing support.
            index = self.cls(query, self.similarity_matrix, num_best=num_best)
            sims = index[query]
            if num_best is not None:
                # Sparse array.
                for result in sims:
                    for i, sim in result:
                        self.assertTrue(numpy.all(sim <= 1.0))
                        self.assertTrue(numpy.all(sim >= 0.0))
            else:
                for i, result in enumerate(sims):
                    self.assertAlmostEqual(1.0, result[i])  # Similarity of a document with itself is 1.0.
                    self.assertTrue(numpy.all(result[:i] >= 0.0))
                    self.assertTrue(numpy.all(result[:i] < 1.0))
                    self.assertTrue(numpy.all(result[i + 1:] >= 0.0))
                    self.assertTrue(numpy.all(result[i + 1:] < 1.0))
//  [python]
def test_concat_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.concat([a, b, c], axis=0)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=0)))
//  [python]
def test_concat_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.concat([a, b, c], axis=0)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=0)))
//  [python]
def test_concat_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.concat([a, b, c], axis=1)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=1)))
//  [python]
def test_concat_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.concat([a, b, c], axis=1)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=1)))
//  [python]
def test_concat_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.concat([a, b, c], axis=2)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=2)))
//  [python]
def test_concat_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.concat([a, b, c], axis=2)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=2)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-1)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=-1)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-1)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=-1)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-3)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=-3)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-3)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=-3)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-2)
      r = xdl.execute(r)
      self.assertTrue(np.alltrue(r == np.concatenate([a, b, c], axis=-2)))
//  [python]
def test_concat_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.concat([a, b, c], axis=-2)
      r = xdl.execute(r)
      self.assertTrue(np.all(r == np.concatenate([a, b, c], axis=-2)))
//  [python]
def testIter(self):
        # Override testIter.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        for sims in index:
            self.assertTrue(numpy.alltrue(sims >= 0.0))
            self.assertTrue(numpy.alltrue(sims <= 1.0))
//  [python]
def testIter(self):
        # Override testIter.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        for sims in index:
            self.assertTrue(numpy.all(sims >= 0.0))
            self.assertTrue(numpy.all(sims <= 1.0))
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        index.num_best = num_best
        query = texts[0]
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.alltrue(sim > 0.0))  # Note that similarities are bigger than zero, as they are the 1/ 1 + distances.
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 0.0.
            self.assertTrue(numpy.alltrue(sims[1:] > 0.0))
            self.assertTrue(numpy.alltrue(sims[1:] < 1.0))
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        index.num_best = num_best
        query = texts[0]
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.all(sim > 0.0))  # Note that similarities are bigger than zero, as they are the 1/ 1 + distances.
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 0.0.
            self.assertTrue(numpy.all(sims[1:] > 0.0))
            self.assertTrue(numpy.all(sims[1:] < 1.0))
//  [python]
def testChunking(self):
        # Override testChunking.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        query = texts[:3]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.alltrue(sims[i, i] == 1.0))  # Similarity of a document with itself is 0.0.

        # test the same thing but with num_best
        index.num_best = 3
        sims = index[query]
        for sims_temp in sims:
            for i, sim in sims_temp:
                self.assertTrue(numpy.alltrue(sim > 0.0))
                self.assertTrue(numpy.alltrue(sim <= 1.0))
//  [python]
def testChunking(self):
        # Override testChunking.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        query = texts[:3]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.all(sims[i, i] == 1.0))  # Similarity of a document with itself is 0.0.

        # test the same thing but with num_best
        index.num_best = 3
        sims = index[query]
        for sims_temp in sims:
            for i, sim in sims_temp:
                self.assertTrue(numpy.all(sim > 0.0))
                self.assertTrue(numpy.all(sim <= 1.0))
//  [python]
def _compatibility_check(contours):
    if not utils.all_same_shape([c.points for c in contours]):
        raise RuntimeError("All contours must have the same number of points in order to align them.")
    if numpy.alltrue([isinstance(c, contour_class.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 RuntimeError("If all contours have landmarks, they must all have the same number of landmarks.")
//  [python]
def _compatibility_check(contours):
    if not utils.all_same_shape([c.points for c in contours]):
        raise RuntimeError("All contours must have the same number of points in order to align them.")
    if numpy.all([isinstance(c, contour_class.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 RuntimeError("If all contours have landmarks, they must all have the same number of landmarks.")
//  [python]
def _should_allow_reverse(contours, allow_reflection):
    # If we are to allow for reflections, we ought to allow for reversing
    # orientations too, because even if the contours start out oriented in the
    # same direction, reflections can change that.
    # Then check if all of the contours are oriented in the same direction.
    # If they're not, we need to allow for reversing their orientation in the
    # alignment process.
    if allow_reflection:
        return True
    orientations = numpy.array([numpy.sign(contour.signed_area()) for contour in contours])
    homogenous_orientations = numpy.alltrue(orientations == -1) or numpy.alltrue(orientations == 1)
    return not homogenous_orientations
//  [python]
def _should_allow_reverse(contours, allow_reflection):
    # If we are to allow for reflections, we ought to allow for reversing
    # orientations too, because even if the contours start out oriented in the
    # same direction, reflections can change that.
    # Then check if all of the contours are oriented in the same direction.
    # If they're not, we need to allow for reversing their orientation in the
    # alignment process.
    if allow_reflection:
        return True
    orientations = numpy.array([numpy.sign(contour.signed_area()) for contour in contours])
    homogenous_orientations = numpy.all(orientations == -1) or numpy.all(orientations == 1)
    return not homogenous_orientations
//  [python]
def check_callable(callables, n_scales):
    r"""
    Checks the callable type per level.

    Parameters
    ----------
    callables : `callable` or `list` of `callables`
        The callable to be used per scale.
    n_scales : `int`
        The number of scales.

    Returns
    -------
    callable_list : `list`
        A `list` of callables.

    Raises
    ------
    ValueError
        callables must be a callable or a list/tuple of callables with the same
        length as the number of scales
    """
    if callable(callables):
        return [callables] * n_scales
    elif len(callables) == 1 and np.alltrue([callable(f) for f in callables]):
        return list(callables) * n_scales
    elif len(callables) == n_scales and np.alltrue([callable(f)
                                                    for f in callables]):
        return list(callables)
    else:
        raise ValueError("callables must be a callable or a list/tuple of "
                         "callables with the same length as the number "
                         "of scales")
//  [python]
def check_callable(callables, n_scales):
    r"""
    Checks the callable type per level.

    Parameters
    ----------
    callables : `callable` or `list` of `callables`
        The callable to be used per scale.
    n_scales : `int`
        The number of scales.

    Returns
    -------
    callable_list : `list`
        A `list` of callables.

    Raises
    ------
    ValueError
        callables must be a callable or a list/tuple of callables with the same
        length as the number of scales
    """
    if callable(callables):
        return [callables] * n_scales
    elif len(callables) == 1 and np.all([callable(f) for f in callables]):
        return list(callables) * n_scales
    elif len(callables) == n_scales and np.all([callable(f)
                                                    for f in callables]):
        return list(callables)
    else:
        raise ValueError("callables must be a callable or a list/tuple of "
                         "callables with the same length as the number "
                         "of scales")
//  [python]
def check_sampling(sampling, n_scales):
    r"""
    Function that checks the value of a `sampling` parameter defined for
    multiple scales. It must be `int` or `ndarray` or `list` of those.

    Parameters
    ----------
    sampling : `int` or `ndarray` or `list` of those
        The value to check.
    n_scales : `int`
        The number of scales.

    Returns
    -------
    sampling : `list` of `int` or `ndarray`
        The list of values per scale.

    Raises
    ------
    ValueError
        A sampling list can only contain 1 element or {n_scales} elements
    ValueError
        sampling can be an integer or ndarray, a integer or ndarray list
        containing 1 or {n_scales} elements or None
    """
    if (isinstance(sampling, (list, tuple)) and
        np.alltrue([isinstance(s, (np.ndarray, np.int)) or sampling is None
                    for s in sampling])):
        if len(sampling) == 1:
            return sampling * n_scales
        elif len(sampling) == n_scales:
            return sampling
        else:
            raise ValueError('A sampling list can only '
                             'contain 1 element or {} '
                             'elements'.format(n_scales))
    elif isinstance(sampling, (np.ndarray, np.int)) or sampling is None:
        return [sampling] * n_scales
    else:
        raise ValueError('sampling can be an integer or ndarray, '
                         'a integer or ndarray list '
                         'containing 1 or {} elements or '
                         'None'.format(n_scales))
//  [python]
def check_sampling(sampling, n_scales):
    r"""
    Function that checks the value of a `sampling` parameter defined for
    multiple scales. It must be `int` or `ndarray` or `list` of those.

    Parameters
    ----------
    sampling : `int` or `ndarray` or `list` of those
        The value to check.
    n_scales : `int`
        The number of scales.

    Returns
    -------
    sampling : `list` of `int` or `ndarray`
        The list of values per scale.

    Raises
    ------
    ValueError
        A sampling list can only contain 1 element or {n_scales} elements
    ValueError
        sampling can be an integer or ndarray, a integer or ndarray list
        containing 1 or {n_scales} elements or None
    """
    if (isinstance(sampling, (list, tuple)) and
        np.all([isinstance(s, (np.ndarray, np.int)) or sampling is None
                    for s in sampling])):
        if len(sampling) == 1:
            return sampling * n_scales
        elif len(sampling) == n_scales:
            return sampling
        else:
            raise ValueError('A sampling list can only '
                             'contain 1 element or {} '
                             'elements'.format(n_scales))
    elif isinstance(sampling, (np.ndarray, np.int)) or sampling is None:
        return [sampling] * n_scales
    else:
        raise ValueError('sampling can be an integer or ndarray, '
                         'a integer or ndarray list '
                         'containing 1 or {} elements or '
                         'None'.format(n_scales))
//  [python]
def test_zeros_0(self):
      a = xdl.zeros(np.array([], np.int64), xdl.DataType.float)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == 0))
      self.assertTrue(a.shape == ())
      self.assertTrue(a.dtype == np.float32)
//  [python]
def test_zeros_0(self):
      a = xdl.zeros(np.array([], np.int64), xdl.DataType.float)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == 0))
      self.assertTrue(a.shape == ())
      self.assertTrue(a.dtype == np.float32)
//  [python]
def test_zeros_1(self):
      a = xdl.zeros([100], xdl.DataType.float)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == 0))
      self.assertTrue(a.shape == (100,))
      self.assertTrue(a.dtype == np.float32)
//  [python]
def test_zeros_1(self):
      a = xdl.zeros([100], xdl.DataType.float)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == 0))
      self.assertTrue(a.shape == (100,))
      self.assertTrue(a.dtype == np.float32)
//  [python]
def diag_indices_from(arr):
    """Return the indices to access the main diagonal of an n-dimensional
    array.

    See diag_indices() for full details.

    Parameters
    ----------
    arr : array, at least 2-d
    """
    if not arr.ndim >= 2:
        raise ValueError("input array must be at least 2-d")
    # For more than d=2, the strided formula is only valid for arrays with
    # all dimensions equal, so we check first.
    if not np.alltrue(np.diff(arr.shape) == 0):
        raise ValueError("All dimensions of input must be of equal length")

    return diag_indices(arr.shape[0], arr.ndim)
//  [python]
def diag_indices_from(arr):
    """Return the indices to access the main diagonal of an n-dimensional
    array.

    See diag_indices() for full details.

    Parameters
    ----------
    arr : array, at least 2-d
    """
    if not arr.ndim >= 2:
        raise ValueError("input array must be at least 2-d")
    # For more than d=2, the strided formula is only valid for arrays with
    # all dimensions equal, so we check first.
    if not np.all(np.diff(arr.shape) == 0):
        raise ValueError("All dimensions of input must be of equal length")

    return diag_indices(arr.shape[0], arr.ndim)
//  [python]
def diag_indices_from(x1):
    """
    Return the indices to access the main diagonal of an n-dimensional array.

    For full documentation refer to :obj:`numpy.diag_indices_from`.

    See also
    --------
    :obj:`diag_indices` : Return the indices to access the main
                          diagonal of an array.

    """

    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
    if x1_desc:
        # original limitation
        if not x1_desc.ndim >= 2:
            pass

        # original limitation
        # For more than d=2, the strided formula is only valid for arrays with
        # all dimensions equal, so we check first.
        elif not numpy.alltrue(
            numpy.diff(x1_desc.shape) == 0
        ):  # TODO: replace alltrue and diff funcs with dpnp own ones
            pass
        else:
            return dpnp_diag_indices(x1_desc.shape[0], x1_desc.ndim)

    return call_origin(numpy.diag_indices_from, x1)
//  [python]
def diag_indices_from(x1):
    """
    Return the indices to access the main diagonal of an n-dimensional array.

    For full documentation refer to :obj:`numpy.diag_indices_from`.

    See also
    --------
    :obj:`diag_indices` : Return the indices to access the main
                          diagonal of an array.

    """

    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
    if x1_desc:
        # original limitation
        if not x1_desc.ndim >= 2:
            pass

        # original limitation
        # For more than d=2, the strided formula is only valid for arrays with
        # all dimensions equal, so we check first.
        elif not numpy.all(
            numpy.diff(x1_desc.shape) == 0
        ):  # TODO: replace alltrue and diff funcs with dpnp own ones
            pass
        else:
            return dpnp_diag_indices(x1_desc.shape[0], x1_desc.ndim)

    return call_origin(numpy.diag_indices_from, x1)
//  [python]
def test_box_actions_out_of_bound(env: gym.Env):
    """Test out of bound actions in Box action_space.

    Environments with Box actions spaces perform clipping inside `step`.
    The expected behaviour is that an action `out-of-bound` has the same effect
    of an action with value exactly at the upper (or lower) bound.

    Args:
        env (gym.Env): the gym environment
    """
    env.reset(seed=42)

    oob_env = gym.make(env.spec.id, disable_env_checker=True)
    oob_env.reset(seed=42)

    assert isinstance(env.action_space, spaces.Box)
    dtype = env.action_space.dtype
    upper_bounds = env.action_space.high
    lower_bounds = env.action_space.low

    for i, (is_upper_bound, is_lower_bound) in enumerate(
        zip(env.action_space.bounded_above, env.action_space.bounded_below)
    ):
        if is_upper_bound:
            obs, _, _, _, _ = env.step(upper_bounds)
            oob_action = upper_bounds.copy()
            oob_action[i] += np.cast[dtype](OOB_VALUE)

            assert oob_action[i] > upper_bounds[i]
            oob_obs, _, _, _, _ = oob_env.step(oob_action)

            assert np.alltrue(obs == oob_obs)

        if is_lower_bound:
            obs, _, _, _, _ = env.step(
                lower_bounds
            )  # `env` is unwrapped, and in new step API
            oob_action = lower_bounds.copy()
            oob_action[i] -= np.cast[dtype](OOB_VALUE)

            assert oob_action[i] < lower_bounds[i]
            oob_obs, _, _, _, _ = oob_env.step(oob_action)

            assert np.alltrue(obs == oob_obs)

    env.close()
//  [python]
def test_box_actions_out_of_bound(env: gym.Env):
    """Test out of bound actions in Box action_space.

    Environments with Box actions spaces perform clipping inside `step`.
    The expected behaviour is that an action `out-of-bound` has the same effect
    of an action with value exactly at the upper (or lower) bound.

    Args:
        env (gym.Env): the gym environment
    """
    env.reset(seed=42)

    oob_env = gym.make(env.spec.id, disable_env_checker=True)
    oob_env.reset(seed=42)

    assert isinstance(env.action_space, spaces.Box)
    dtype = env.action_space.dtype
    upper_bounds = env.action_space.high
    lower_bounds = env.action_space.low

    for i, (is_upper_bound, is_lower_bound) in enumerate(
        zip(env.action_space.bounded_above, env.action_space.bounded_below)
    ):
        if is_upper_bound:
            obs, _, _, _, _ = env.step(upper_bounds)
            oob_action = upper_bounds.copy()
            oob_action[i] += np.cast[dtype](OOB_VALUE)

            assert oob_action[i] > upper_bounds[i]
            oob_obs, _, _, _, _ = oob_env.step(oob_action)

            assert np.all(obs == oob_obs)

        if is_lower_bound:
            obs, _, _, _, _ = env.step(
                lower_bounds
            )  # `env` is unwrapped, and in new step API
            oob_action = lower_bounds.copy()
            oob_action[i] -= np.cast[dtype](OOB_VALUE)

            assert oob_action[i] < lower_bounds[i]
            oob_obs, _, _, _, _ = oob_env.step(oob_action)

            assert np.all(obs == oob_obs)

    env.close()
//  [python]
def test_zeros_int64(self):
      a = xdl.zeros([100], xdl.DataType.int64)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == 0))
      self.assertTrue(a.shape == (100,))
      self.assertTrue(a.dtype == np.int64)
//  [python]
def test_zeros_int64(self):
      a = xdl.zeros([100], xdl.DataType.int64)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == 0))
      self.assertTrue(a.shape == (100,))
      self.assertTrue(a.dtype == np.int64)
//  [python]
def __init__(self,
                 atoms, encut, potcar='POTCAR', k=[0.0, 0.0, 0.0],
                 ngrid=None,
                 lgam=False, lsoc=False
                 ):
        '''
        input:
            atoms: ase atom object
            encut: float, energy cutoff in eV
            potcar: the PAW POTCAR file of all the elements in atoms
            k: the k-point vector in fractional coordinate
        '''
        self.atoms = atoms
        self.natoms = len(atoms)
        self.encut = encut
        self.kvec = np.asarray(k, dtype=float)

        if isinstance(potcar, str):
            self.pawpp = [pawpotcar(potstr) for potstr in
                          open(potcar).read().split('End of Dataset')[:-1]]
        elif isinstance(potcar, list):
            assert np.alltrue([
                isinstance(pp, pawpotcar) for pp in potcar
            ])
            self.pawpp = potcar
        else:
            raise ValueError('Argument potcar must be a string or list of pawpotcar type!')

        elements, elem_first_idx, elem_cnts = np.unique(atoms.get_chemical_symbols(),
                                                        return_index=True,
                                                        return_counts=True)
        # Sometimes, the order of the elements returned by np.unique may not be
        # consistent with that in POSCAR/POTCAR
        elem_first_idx = np.argsort(elem_first_idx)
        elements = elements[elem_first_idx]
        self.elem_cnts = elem_cnts[elem_first_idx]

        assert len(self.elem_cnts) == len(self.pawpp), \
            "The number of elements in POTCAR and POSCAR does not match!"
        if not np.alltrue([
            self.pawpp[ii].element.split('_')[0] == elements[ii]
            for ii in range(len(elements))
        ]):
            print(
                "\nWARNING:\nThe name of elements in POTCAR and POSCAR does not match!\n\n" +
                "    POTCAR: {}\n".format(' '.join([pp.element for pp in self.pawpp])) +
                "    POSCAR: {}\n".format(' '.join(elements))
            )

        self.elements = list(elements)
        self.element_idx = [self.elements.index(s) for s in
                            atoms.get_chemical_symbols()]

        self.set_grid(ngrid=ngrid)
        self.rphase()
        self.calc_rproj()
//  [python]
def __init__(self,
                 atoms, encut, potcar='POTCAR', k=[0.0, 0.0, 0.0],
                 ngrid=None,
                 lgam=False, lsoc=False
                 ):
        '''
        input:
            atoms: ase atom object
            encut: float, energy cutoff in eV
            potcar: the PAW POTCAR file of all the elements in atoms
            k: the k-point vector in fractional coordinate
        '''
        self.atoms = atoms
        self.natoms = len(atoms)
        self.encut = encut
        self.kvec = np.asarray(k, dtype=float)

        if isinstance(potcar, str):
            self.pawpp = [pawpotcar(potstr) for potstr in
                          open(potcar).read().split('End of Dataset')[:-1]]
        elif isinstance(potcar, list):
            assert np.all([
                isinstance(pp, pawpotcar) for pp in potcar
            ])
            self.pawpp = potcar
        else:
            raise ValueError('Argument potcar must be a string or list of pawpotcar type!')

        elements, elem_first_idx, elem_cnts = np.unique(atoms.get_chemical_symbols(),
                                                        return_index=True,
                                                        return_counts=True)
        # Sometimes, the order of the elements returned by np.unique may not be
        # consistent with that in POSCAR/POTCAR
        elem_first_idx = np.argsort(elem_first_idx)
        elements = elements[elem_first_idx]
        self.elem_cnts = elem_cnts[elem_first_idx]

        assert len(self.elem_cnts) == len(self.pawpp), \
            "The number of elements in POTCAR and POSCAR does not match!"
        if not np.all([
            self.pawpp[ii].element.split('_')[0] == elements[ii]
            for ii in range(len(elements))
        ]):
            print(
                "\nWARNING:\nThe name of elements in POTCAR and POSCAR does not match!\n\n" +
                "    POTCAR: {}\n".format(' '.join([pp.element for pp in self.pawpp])) +
                "    POSCAR: {}\n".format(' '.join(elements))
            )

        self.elements = list(elements)
        self.element_idx = [self.elements.index(s) for s in
                            atoms.get_chemical_symbols()]

        self.set_grid(ngrid=ngrid)
        self.rphase()
        self.calc_rproj()
//  [python]
def test_example(self):
        exp = [2, 3, 5, 7, 13, 17, 23, 37, 43, 47, 53, 67, 73, 83, 97, 113, 137, 167, 173, 197]
        got = left_truncatable_primes(20)
        self.assertEqual(numpy.alltrue(exp == got), True, 'Example')
//  [python]
def test_example(self):
        exp = [2, 3, 5, 7, 13, 17, 23, 37, 43, 47, 53, 67, 73, 83, 97, 113, 137, 167, 173, 197]
        got = left_truncatable_primes(20)
        self.assertEqual(numpy.all(exp == got), True, 'Example')
//  [python]
def copyAxis(self, axis, newname=None, unlimited=0,
                 index=None, extbounds=None):
        """Copy axis description and data from another axis.

        Parameters
        ----------
        axis : axis to copy (cdms2.axis.FileAxis/cdms2.axis.FileVirtualAxis)

        newname : (None/str) new name for axis (default None)

        unlimited : (int/True/False) unlimited dimension (default 0)

        index : (int/None) (default None)

        extbounds : (numpy.ndarray) new bounds to use bounds (default None)

        Returns
        --------
        copy of input axis (cdms2.axis.FileAxis/cdms2.axis.FileVirtualAxis)
        """
        if newname is None:
            newname = axis.id

        if len(newname) > 127:
            msg = "axis name has more than 127 characters, name will be truncated"
            warnings.warn(msg, UserWarning)
            newname = newname[:127] if len(newname) > 127 else newname

        # If the axis already exists and has the same values, return existing
        if newname in self.axes:
            newaxis = self.axes[newname]
            if newaxis.isVirtual():
                if len(axis) != len(newaxis):
                    raise DuplicateAxisError(DuplicateAxis + newname)
            elif unlimited == 0 or (unlimited == 1 and getNetcdfUseParallelFlag() != 0):
                if len(axis) != len(newaxis) or numpy.alltrue(
                        numpy.less(numpy.absolute(newaxis[:] - axis[:]), 1.e-5)) == 0:
                    raise DuplicateAxisError(DuplicateAxis + newname)
            else:
                if index is None:
                    isoverlap, index = isOverlapVector(axis[:], newaxis[:])
                else:
                    isoverlap = 1
                if isoverlap:
                    self._file_.sync()
                    newaxis[index:index + len(axis)] = axis[:]
                    if extbounds is None:
                        axisBounds = axis.getBounds()
                    else:
                        axisBounds = extbounds
                    if axisBounds is not None:
                        newaxis.setBounds(axisBounds)
                else:
                    raise DuplicateAxisError(DuplicateAxis + newname)

        elif axis.isVirtual():
            newaxis = self.createVirtualAxis(newname, len(axis))

        # Else create the new axis and copy its bounds and metadata
        else:
            newaxis = self.createAxis(newname, axis[:], unlimited)
            bounds = axis.getBounds()
            if bounds is not None:
                if hasattr(axis, 'bounds'):
                    boundsid = axis.bounds
                else:
                    boundsid = None
                newaxis.setBounds(bounds, persistent=1, boundsid=boundsid)
            for attname, attval in axis.attributes.items():
                if attname not in ["datatype", "id", "length",
                                   "isvar", "name_in_file", "partition"]:
                    setattr(newaxis, attname, attval)
        return newaxis
//  [python]
def copyAxis(self, axis, newname=None, unlimited=0,
                 index=None, extbounds=None):
        """Copy axis description and data from another axis.

        Parameters
        ----------
        axis : axis to copy (cdms2.axis.FileAxis/cdms2.axis.FileVirtualAxis)

        newname : (None/str) new name for axis (default None)

        unlimited : (int/True/False) unlimited dimension (default 0)

        index : (int/None) (default None)

        extbounds : (numpy.ndarray) new bounds to use bounds (default None)

        Returns
        --------
        copy of input axis (cdms2.axis.FileAxis/cdms2.axis.FileVirtualAxis)
        """
        if newname is None:
            newname = axis.id

        if len(newname) > 127:
            msg = "axis name has more than 127 characters, name will be truncated"
            warnings.warn(msg, UserWarning)
            newname = newname[:127] if len(newname) > 127 else newname

        # If the axis already exists and has the same values, return existing
        if newname in self.axes:
            newaxis = self.axes[newname]
            if newaxis.isVirtual():
                if len(axis) != len(newaxis):
                    raise DuplicateAxisError(DuplicateAxis + newname)
            elif unlimited == 0 or (unlimited == 1 and getNetcdfUseParallelFlag() != 0):
                if len(axis) != len(newaxis) or numpy.all(
                        numpy.less(numpy.absolute(newaxis[:] - axis[:]), 1.e-5)) == 0:
                    raise DuplicateAxisError(DuplicateAxis + newname)
            else:
                if index is None:
                    isoverlap, index = isOverlapVector(axis[:], newaxis[:])
                else:
                    isoverlap = 1
                if isoverlap:
                    self._file_.sync()
                    newaxis[index:index + len(axis)] = axis[:]
                    if extbounds is None:
                        axisBounds = axis.getBounds()
                    else:
                        axisBounds = extbounds
                    if axisBounds is not None:
                        newaxis.setBounds(axisBounds)
                else:
                    raise DuplicateAxisError(DuplicateAxis + newname)

        elif axis.isVirtual():
            newaxis = self.createVirtualAxis(newname, len(axis))

        # Else create the new axis and copy its bounds and metadata
        else:
            newaxis = self.createAxis(newname, axis[:], unlimited)
            bounds = axis.getBounds()
            if bounds is not None:
                if hasattr(axis, 'bounds'):
                    boundsid = axis.bounds
                else:
                    boundsid = None
                newaxis.setBounds(bounds, persistent=1, boundsid=boundsid)
            for attname, attval in axis.attributes.items():
                if attname not in ["datatype", "id", "length",
                                   "isvar", "name_in_file", "partition"]:
                    setattr(newaxis, attname, attval)
        return newaxis
//  [python]
def test02b_CompareTable(self):
        """Comparing several written and read time values in a Table."""

        # Create test Table with data.
        tbl = self.h5file.create_table('/', 'test', self.MyTimeRow)

        # Size of the test.
        nrows = tbl.nrowsinbuf + 34  # Add some more rows than buffer.
        # Only for home checks; the value above should check better
        # the I/O with multiple buffers.
        # nrows = 10

        row = tbl.row
        for i in range(nrows):
            row['t32col'] = i
            j = i * 2
            row['t64col'] = (j + 0.012, j+1+0.012)
            row.append()

        self._reopen()

        # Check the written data.
        recarr = self.h5file.root.test.read()
        self.h5file.close()

        # Time32 column.
        orig_val = numpy.arange(nrows, dtype=numpy.int32)
        if common.verbose:
            print("Original values:", orig_val)
            print("Retrieved values:", recarr['t32col'][:])
        self.assertTrue(numpy.alltrue(recarr['t32col'][:] == orig_val),
                        "Stored and retrieved values do not match.")

        # Time64 column.
        orig_val = numpy.arange(0, nrows * 2, dtype=numpy.int32) + 0.012
        orig_val.shape = (nrows, 2)
        if common.verbose:
            print("Original values:", orig_val)
            print("Retrieved values:", recarr['t64col'][:])
        self.assertTrue(allequal(recarr['t64col'][:], orig_val, numpy.float64),
                        "Stored and retrieved values do not match.")
//  [python]
def test02b_CompareTable(self):
        """Comparing several written and read time values in a Table."""

        # Create test Table with data.
        tbl = self.h5file.create_table('/', 'test', self.MyTimeRow)

        # Size of the test.
        nrows = tbl.nrowsinbuf + 34  # Add some more rows than buffer.
        # Only for home checks; the value above should check better
        # the I/O with multiple buffers.
        # nrows = 10

        row = tbl.row
        for i in range(nrows):
            row['t32col'] = i
            j = i * 2
            row['t64col'] = (j + 0.012, j+1+0.012)
            row.append()

        self._reopen()

        # Check the written data.
        recarr = self.h5file.root.test.read()
        self.h5file.close()

        # Time32 column.
        orig_val = numpy.arange(nrows, dtype=numpy.int32)
        if common.verbose:
            print("Original values:", orig_val)
            print("Retrieved values:", recarr['t32col'][:])
        self.assertTrue(numpy.all(recarr['t32col'][:] == orig_val),
                        "Stored and retrieved values do not match.")

        # Time64 column.
        orig_val = numpy.arange(0, nrows * 2, dtype=numpy.int32) + 0.012
        orig_val.shape = (nrows, 2)
        if common.verbose:
            print("Original values:", orig_val)
            print("Retrieved values:", recarr['t64col'][:])
        self.assertTrue(allequal(recarr['t64col'][:], orig_val, numpy.float64),
                        "Stored and retrieved values do not match.")
//  [python]
def get_pdos(self, atoms=':', kpts=':', spd=':'):
        '''
        Get site/k-points/spd-orbital projected partial density of states (PDOS)
        '''

        if self._tdos is None:
            self.init_dos()

        pdos = []
        proj = self.get_pw(atoms, kpts, spd)

        atoms, kpts, spd = self.translate_selection(atoms, kpts, spd)

        if len(np.arange(self._nkpts)[kpts]) == self._nkpts:
            if np.alltrue(
                np.sort(np.arange(self._nkpts)[kpts]) == np.arange(self._nkpts)
            ):
                used_all_kpts = True
            else:
                used_all_kpts = False
        else:
            used_all_kpts = False

        for ispin in range(self._nspin):
            pw = proj[ispin]

            if used_all_kpts:
                td = self._tdos[ispin, kpts]
            else:
                # if not all the k-points are used, then probably we should get
                # rid of the k-point weights
                td = self._tdos[ispin, kpts, ...] / \
                    self._kptw[ispin, kpts, np.newaxis, np.newaxis]

            pdos.append(np.sum(pw[..., np.newaxis] * td, axis=(0, 1)))

            # pwht = np.sum(self._aproj[ispin][kpts,:,atoms,spd], axis=(-1, -2))
            # pdos.append(np.sum(pwht[..., np.newaxis] * self._tdos[ispin][kpts,...], axis=(0, 1)))

        # only return one dos if not spin-polarized
        # p = pdos[0] if self._nspin == 1 else pdos
        pdos = np.array(pdos, dtype=float)

        return self._xen, pdos