jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __getitem__(self, key_idx):
        """Return item indexed from class

        Parameters
        ----------
        key_idx : str/list/tuple/slice/int
            Query for array items

        Returns
        -------
        arr_slice : np.ndarray
            Array of data containing row names and time indexed
            columns. The return is squeezed meaning that all dimensions
            of the output that are length of one are removed
        """
        rows, cols = self._parse_key_idx(key_idx)
        row_list, row_str = self._get_str_rows(rows)
        assert n.all(row_str) or n.all(n.logical_not(row_str)), \
                "Cannot assign/return combination of strings and numbers"
        if n.all(row_str):
            # Return sliced strings
            arr_slice = np.atleast_2d(np.empty_like(self.array[rows, cols], dtype=object))
            for row_num, row in enumerate(row_list):
                str_arr = self._get_strings(self.inv_map[row])
                arr_slice[row_num, :] = str_arr[cols]
        else:
            arr_slice = self.array[rows, cols]

        if isinstance(rows,list) and len(rows) == 1 \
        and self.inv_map[rows[0]] in self.orig_dtypes:
            dtype = self.orig_dtypes[self.inv_map[rows[0]]]
            if np.issubdtype(dtype, np.integer):
                arr_slice = arr_slice.astype(np.float64)
                if np.any(np.isnan(arr_slice)):
                    nan_indexes = np.isnan(arr_slice)
                    arr_slice[~nan_indexes] = arr_slice[~nan_indexes].astype(dtype)
                else:
                    arr_slice = arr_slice.astype(dtype)
            else:
                arr_slice = arr_slice.astype(dtype)

        # remove all dimensions of length one
        arr_slice = np.squeeze(arr_slice)

        return arr_slice
//  [python]
def __getitem__(self, key_idx):
        """Return item indexed from class

        Parameters
        ----------
        key_idx : str/list/tuple/slice/int
            Query for array items

        Returns
        -------
        arr_slice : np.ndarray
            Array of data containing row names and time indexed
            columns. The return is squeezed meaning that all dimensions
            of the output that are length of one are removed
        """
        rows, cols = self._parse_key_idx(key_idx)
        row_list, row_str = self._get_str_rows(rows)
        assert np.all(row_str) or np.all(np.logical_not(row_str)), \
                "Cannot assign/return combination of strings and numbers"
        if np.all(row_str):
            # Return sliced strings
            arr_slice = np.atleast_2d(np.empty_like(self.array[rows, cols], dtype=object))
            for row_num, row in enumerate(row_list):
                str_arr = self._get_strings(self.inv_map[row])
                arr_slice[row_num, :] = str_arr[cols]
        else:
            arr_slice = self.array[rows, cols]

        if isinstance(rows,list) and len(rows) == 1 \
        and self.inv_map[rows[0]] in self.orig_dtypes:
            dtype = self.orig_dtypes[self.inv_map[rows[0]]]
            if np.issubdtype(dtype, np.integer):
                arr_slice = arr_slice.astype(np.float64)
                if np.any(np.isnan(arr_slice)):
                    nan_indexes = np.isnan(arr_slice)
                    arr_slice[~nan_indexes] = arr_slice[~nan_indexes].astype(dtype)
                else:
                    arr_slice = arr_slice.astype(dtype)
            else:
                arr_slice = arr_slice.astype(dtype)

        # remove all dimensions of length one
        arr_slice = np.squeeze(arr_slice)

        return arr_slice
//  [python]
def clipSegmentToBounds(segment, bounds):
    """Clip segment to volume aligned with axes.

    :param numpy.ndarray segment: (p0, p1)
    :param numpy.ndarray bounds: (lower corner, upper corner)
    :return: Either clipped (p0, p1) or None if outside volume
    :rtype: Union[None,List[numpy.ndarray]]
    """
    segment = numpy.asarray(segment)
    bounds = numpy.asarray(bounds)

    p0, p1 = segment
    # Get intersection points of ray with volume boundary planes
    # Line equation: P = offset * delta + p0
    delta = p1 - p0
    deltaNotZero = numpy.array(delta, copy=True)
    deltaNotZero[deltaNotZero == 0] = numpy.nan  # Invalidated to avoid division by zero
    offsets = ((bounds - p0) / deltaNotZero).reshape(-1)
    points = offsets.reshape(-1, 1) * delta + p0

    # Avoid precision errors by using bounds value
    points.shape = 2, 3, 3  # Reshape 1 point per bound value
    for dim in range(3):
        points[:, dim, dim] = bounds[:, dim]
    points.shape = -1, 3  # Set back to 2D array

    # Find intersection points that are included in the volume
    mask = numpy.logical_and(
        numpy.alltrue(bounds[0] <= points, axis=1), numpy.alltrue(points <= bounds[1], axis=1)
    )
    intersections = numpy.unique(offsets[mask])
    if len(intersections) != 2:
        return None

    intersections.sort()
    # Do p1 first as p0 is need to compute it
    if intersections[1] < 1:  # clip p1
        segment[1] = intersections[1] * delta + p0
    if intersections[0] > 0:  # clip p0
        segment[0] = intersections[0] * delta + p0
    return segment
//  [python]
def clipSegmentToBounds(segment, bounds):
    """Clip segment to volume aligned with axes.

    :param numpy.ndarray segment: (p0, p1)
    :param numpy.ndarray bounds: (lower corner, upper corner)
    :return: Either clipped (p0, p1) or None if outside volume
    :rtype: Union[None,List[numpy.ndarray]]
    """
    segment = numpy.asarray(segment)
    bounds = numpy.asarray(bounds)

    p0, p1 = segment
    # Get intersection points of ray with volume boundary planes
    # Line equation: P = offset * delta + p0
    delta = p1 - p0
    deltaNotZero = numpy.array(delta, copy=True)
    deltaNotZero[deltaNotZero == 0] = numpy.nan  # Invalidated to avoid division by zero
    offsets = ((bounds - p0) / deltaNotZero).reshape(-1)
    points = offsets.reshape(-1, 1) * delta + p0

    # Avoid precision errors by using bounds value
    points.shape = 2, 3, 3  # Reshape 1 point per bound value
    for dim in range(3):
        points[:, dim, dim] = bounds[:, dim]
    points.shape = -1, 3  # Set back to 2D array

    # Find intersection points that are included in the volume
    mask = numpy.logical_and(
        numpy.all(bounds[0] <= points, axis=1), numpy.all(points <= bounds[1], axis=1)
    )
    intersections = numpy.unique(offsets[mask])
    if len(intersections) != 2:
        return None

    intersections.sort()
    # Do p1 first as p0 is need to compute it
    if intersections[1] < 1:  # clip p1
        segment[1] = intersections[1] * delta + p0
    if intersections[0] > 0:  # clip p0
        segment[0] = intersections[0] * delta + p0
    return segment
//  [python]
def _data_in_bounds(self) -> tuple[np.ndarray, np.ndarray]:
        if self.data:
            points = np.array(list(self.data.keys()))
            values = np.array(list(self.data.values()), dtype=float)
            ll, ur = np.reshape(self.bounds, (2, 2)).T
            inds = np.logical_and(ll <= points, points <= ur).all(axis=1)
            return points[inds], values[inds].reshape(-1, self.vdim)
        return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
//  [python]
def _data_in_bounds(self) -> tuple[np.ndarray, np.ndarray]:
        if self.data:
            points = np.array(list(self.data.keys()))
            values = np.array(list(self.data.values()), dtype=float)
            ll, ur = np.reshape(self.bounds, (2, 2)).T
            inds = np.all(np.logical_and(ll <= points, points <= ur), axis=1)
            return points[inds], values[inds].reshape(-1, self.vdim)
        return np.zeros((0, 2)), np.zeros((0, self.vdim), dtype=float)
//  [python]
def setPlane(self, point=None, normal=None):
        """Set plane point and normal and notify.

        :param point: A point on the plane.
        :type point: 3-tuple of float or None.
        :param normal: Normal of the plane.
        :type normal: 3-tuple of float or None.
        """
        planechanged = False

        if point is not None:
            assert len(point) == 3
            point = numpy.array(point, copy=True, dtype=numpy.float32)
            if not numpy.alltrue(numpy.equal(self._point, point)):
                self._point = point
                planechanged = True

        if normal is not None:
            assert len(normal) == 3
            normal = numpy.array(normal, copy=True, dtype=numpy.float32)

            norm = numpy.linalg.norm(normal)
            if norm != 0.0:
                normal /= norm

            if not numpy.alltrue(numpy.equal(self._normal, normal)):
                self._normal = normal
                planechanged = True

        if planechanged:
            _logger.debug(
                "Plane updated:\n\tpoint: %s\n\tnormal: %s",
                str(self._point),
                str(self._normal),
            )
            self.notify()
//  [python]
def setPlane(self, point=None, normal=None):
        """Set plane point and normal and notify.

        :param point: A point on the plane.
        :type point: 3-tuple of float or None.
        :param normal: Normal of the plane.
        :type normal: 3-tuple of float or None.
        """
        planechanged = False

        if point is not None:
            assert len(point) == 3
            point = numpy.array(point, copy=True, dtype=numpy.float32)
            if not numpy.all(numpy.equal(self._point, point)):
                self._point = point
                planechanged = True

        if normal is not None:
            assert len(normal) == 3
            normal = numpy.array(normal, copy=True, dtype=numpy.float32)

            norm = numpy.linalg.norm(normal)
            if norm != 0.0:
                normal /= norm

            if not numpy.all(numpy.equal(self._normal, normal)):
                self._normal = normal
                planechanged = True

        if planechanged:
            _logger.debug(
                "Plane updated:\n\tpoint: %s\n\tnormal: %s",
                str(self._point),
                str(self._normal),
            )
            self.notify()
//  [python]
def __eq__(self, other: 'Coordinate') -> bool:
        if not isinstance(other, self.__class__):
            return NotImplemented
        if len(self.indices) != len(other.indices):
            return False
        if np.all(self.indices == other.indices):
            return True
        return False
//  [python]
def __eq__(self, other: object) -> bool:
        if not Coordinate.__eq__(self, other):
            return False
        srev = self.reverse()
        if not Coordinate.__eq__(srev, other):
            return False
        if n.all(self.kwargs['ncvecs'] == other.kwargs['ncvecs']):
            return True
        if n.all(srev.kwargs['ncvecs'] == other.kwargs['ncvecs']):
            return True
        return False
//  [python]
def __eq__(self, other: object) -> bool:
        if not Coordinate.__eq__(self, other):
            return False
        srev = self.reverse()
        if not Coordinate.__eq__(srev, other):
            return False
        if np.all(self.kwargs['ncvecs'] == other.kwargs['ncvecs']):
            return True
        if np.all(srev.kwargs['ncvecs'] == other.kwargs['ncvecs']):
            return True
        return False
//  [python]
def __eq__(self, other: Coordinate) -> bool:
        if not Coordinate.__eq__(self, other):
            return False
        return n.all(self.refpos == other.refpos)
//  [python]
def __eq__(self, other: Coordinate) -> bool:
        if not Coordinate.__eq__(self, other):
            return False
        return np.all(self.refpos == other.refpos)
//  [python]
def segmentVolumeIntersect(segment, nbins):
    """Get bin indices intersecting with segment

    It should work with N dimensions.
    Coordinate convention (z, y, x) or (x, y, z) should not matter
    as long as segment and nbins are consistent.

    :param numpy.ndarray segment:
        Segment end points as a 2xN array of coordinates
    :param numpy.ndarray nbins:
        Shape of the volume with same coordinates order as segment
    :return: List of bins indices as a 2D array or None if no bins
    :rtype: Union[None,numpy.ndarray]
    """
    segment = numpy.asarray(segment)
    nbins = numpy.asarray(nbins)

    assert segment.ndim == 2
    assert segment.shape[0] == 2
    assert nbins.ndim == 1
    assert segment.shape[1] == nbins.size

    dim = len(nbins)

    bounds = numpy.array((numpy.zeros_like(nbins), nbins))
    segment = clipSegmentToBounds(segment, bounds)
    if segment is None:
        return None  # Segment outside volume
    p0, p1 = segment

    # Get intersections

    # Get coordinates of bin edges crossing the segment
    clipped = numpy.ceil(numpy.clip(segment, 0, nbins))
    start = numpy.min(clipped, axis=0)
    stop = numpy.max(clipped, axis=0)  # stop is NOT included
    edgesByDim = [numpy.arange(start[i], stop[i]) for i in range(dim)]

    # Line equation: P = t * delta + p0
    delta = p1 - p0

    # Get bin edge/line intersections as sorted points along the line
    # Get corresponding line parameters
    t = []
    if n.all(0 <= p0) and n.all(p0 <= nbins):
        t.append([0.0])  # p0 within volume, add it
    t += [(edgesByDim[i] - p0[i]) / delta[i] for i in range(dim) if delta[i] != 0]
    if n.all(0 <= p1) and n.all(p1 <= nbins):
        t.append([1.0])  # p1 within volume, add it
    t = numpy.concatenate(t)
    t.sort(kind="mergesort")

    # Remove duplicates
    unique = numpy.ones((len(t),), dtype=bool)
    numpy.not_equal(t[1:], t[:-1], out=unique[1:])
    t = t[unique]

    if len(t) < 2:
        return None  # Not enough intersection points

    # bin edges/line intersection points
    points = t.reshape(-1, 1) * delta + p0
    centers = (points[:-1] + points[1:]) / 2.0
    bins = numpy.floor(centers).astype(numpy.int64)
    return bins
//  [python]
def segmentVolumeIntersect(segment, nbins):
    """Get bin indices intersecting with segment

    It should work with N dimensions.
    Coordinate convention (z, y, x) or (x, y, z) should not matter
    as long as segment and nbins are consistent.

    :param numpy.ndarray segment:
        Segment end points as a 2xN array of coordinates
    :param numpy.ndarray nbins:
        Shape of the volume with same coordinates order as segment
    :return: List of bins indices as a 2D array or None if no bins
    :rtype: Union[None,numpy.ndarray]
    """
    segment = numpy.asarray(segment)
    nbins = numpy.asarray(nbins)

    assert segment.ndim == 2
    assert segment.shape[0] == 2
    assert nbins.ndim == 1
    assert segment.shape[1] == nbins.size

    dim = len(nbins)

    bounds = numpy.array((numpy.zeros_like(nbins), nbins))
    segment = clipSegmentToBounds(segment, bounds)
    if segment is None:
        return None  # Segment outside volume
    p0, p1 = segment

    # Get intersections

    # Get coordinates of bin edges crossing the segment
    clipped = numpy.ceil(numpy.clip(segment, 0, nbins))
    start = numpy.min(clipped, axis=0)
    stop = numpy.max(clipped, axis=0)  # stop is NOT included
    edgesByDim = [numpy.arange(start[i], stop[i]) for i in range(dim)]

    # Line equation: P = t * delta + p0
    delta = p1 - p0

    # Get bin edge/line intersections as sorted points along the line
    # Get corresponding line parameters
    t = []
    if numpy.all(0 <= p0) and numpy.all(p0 <= nbins):
        t.append([0.0])  # p0 within volume, add it
    t += [(edgesByDim[i] - p0[i]) / delta[i] for i in range(dim) if delta[i] != 0]
    if numpy.all(0 <= p1) and numpy.all(p1 <= nbins):
        t.append([1.0])  # p1 within volume, add it
    t = numpy.concatenate(t)
    t.sort(kind="mergesort")

    # Remove duplicates
    unique = numpy.ones((len(t),), dtype=bool)
    numpy.not_equal(t[1:], t[:-1], out=unique[1:])
    t = t[unique]

    if len(t) < 2:
        return None  # Not enough intersection points

    # bin edges/line intersection points
    points = t.reshape(-1, 1) * delta + p0
    centers = (points[:-1] + points[1:]) / 2.0
    bins = numpy.floor(centers).astype(numpy.int64)
    return bins
//  [python]
def __add__(self, other: object) -> 'Internal':
        if self.union is None:
            return NotImplemented
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self == other:
            raise NoValidInternalError(
                'Cannot add {} object to itself.'
                .format(self.__class__.__name__)
            )

        for s, o in product([self, self.reverse()], [other, other.reverse()]):
            if (
                n.all(s.indices[1:] == o.indices[:-1])
                and n.all(s.kwargs['ncvecs'][1:] == o.kwargs['ncvecs'][:-1])
            ):
                new_indices = [*s.indices, o.indices[-1]]
                new_ncvecs = [*s.kwargs['ncvecs'], o.kwargs['ncvecs'][-1]]
                return self.union(new_indices, new_ncvecs)
        raise NoValidInternalError(
            '{} indices do not overlap!'.format(self.__class__.__name__)
        )
//  [python]
def __add__(self, other: object) -> 'Internal':
        if self.union is None:
            return NotImplemented
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self == other:
            raise NoValidInternalError(
                'Cannot add {} object to itself.'
                .format(self.__class__.__name__)
            )

        for s, o in product([self, self.reverse()], [other, other.reverse()]):
            if (
                np.all(s.indices[1:] == o.indices[:-1])
                and np.all(s.kwargs['ncvecs'][1:] == o.kwargs['ncvecs'][:-1])
            ):
                new_indices = [*s.indices, o.indices[-1]]
                new_ncvecs = [*s.kwargs['ncvecs'], o.kwargs['ncvecs'][-1]]
                return self.union(new_indices, new_ncvecs)
        raise NoValidInternalError(
            '{} indices do not overlap!'.format(self.__class__.__name__)
        )
//  [python]
def find_all_dihedrals(self) -> None:
        for a1, a2 in combinations(self.internals['angles'], 2):
            try:
                new = a1 + a2
            except NoValidInternalError:
                continue
            # this is a dihedral that has the same exact atom as both
            # the first and last atom.
            if (
                new.indices[0] == new.indices[3]
                and n.all(
                    n.sum(new.kwargs['ncvecs'], axis=0) == n.array((0, 0, 0))
                )
            ):
                continue
            try:
                self.add_dihedral(new)
            except DuplicateInternalError:
                continue
//  [python]
def find_all_dihedrals(self) -> None:
        for a1, a2 in combinations(self.internals['angles'], 2):
            try:
                new = a1 + a2
            except NoValidInternalError:
                continue
            # this is a dihedral that has the same exact atom as both
            # the first and last atom.
            if (
                new.indices[0] == new.indices[3]
                and np.all(
                    np.sum(new.kwargs['ncvecs'], axis=0) == np.array((0, 0, 0))
                )
            ):
                continue
            try:
                self.add_dihedral(new)
            except DuplicateInternalError:
                continue
//  [python]
def _get_neighbors(self, dx: np.ndarray) -> Iterator[np.ndarray]:
        pbc = self.atoms.pbc
        if self.cell is None or not np.all(self.cell == self.atoms.cell):
            self.cell = self.atoms.cell.array.copy()
            rcell, self.op = minkowski_reduce(
                complete_cell(self.cell), pbc=pbc
            )
            self.rcell = Cell(rcell)
        dx_sc = dx @ self.rcell.reciprocal().T
        offset = np.zeros(3, dtype=np.int32)
        for _ in range(2):
            offset += pbc * ((dx_sc - offset) // 1.).astype(np.int32)

        for ts in product(*[np.arange(-1 * p, p + 1) for p in pbc]):
            yield (np.array(ts) - offset) @ self.op
//  [python]
def test_batch_pickle() -> None:
    batch = Batch(obs=Batch(a=0.0, c=torch.Tensor([1.0, 2.0])), np=np.zeros([3, 4]))
    batch_pk = pickle.loads(pickle.dumps(batch))
    assert batch.obs.a == batch_pk.obs.a
    assert torch.all(batch.obs.c == batch_pk.obs.c)
    assert np.all(batch.np == batch_pk.np)
//  [python]
def test1(self):
        original = self.TestInterface()

        instance = interface.InCodeComponentImplementation(original)

        handler = instance.get_handler('METHOD')
        handler.add_method('get_a', (handler.INDEX, handler.INDEX, handler.INDEX,), (units.kg,))
        handler.add_method('set_a', (handler.INDEX, handler.INDEX, handler.INDEX, units.kg,), ())

        self.assertEqual(instance.get_a([1], [2], [3]), [38] | units.kg)

        handler = instance.get_handler('PARTICLES')
        handler.define_grid('grid',)
        handler.add_setter('grid', 'set_a', names=('mass',))
        handler.add_getter('grid', 'get_a', names=('mass',))

        grid = instance.grid

        self.assertEqual(grid[1][2][3].mass, 38 | units.kg)
        self.assertEqual(grid[0:2][1][2][3].mass, 38 | units.kg)
        self.assertEqual(len(grid[1][2].mass), 5)
        self.assertTrue(n.all(grid[1][2].mass == [35, 36, 37, 38, 39] | units.kg))
//  [python]
def test1(self):
        original = self.TestInterface()

        instance = interface.InCodeComponentImplementation(original)

        handler = instance.get_handler('METHOD')
        handler.add_method('get_a', (handler.INDEX, handler.INDEX, handler.INDEX,), (units.kg,))
        handler.add_method('set_a', (handler.INDEX, handler.INDEX, handler.INDEX, units.kg,), ())

        self.assertEqual(instance.get_a([1], [2], [3]), [38] | units.kg)

        handler = instance.get_handler('PARTICLES')
        handler.define_grid('grid',)
        handler.add_setter('grid', 'set_a', names=('mass',))
        handler.add_getter('grid', 'get_a', names=('mass',))

        grid = instance.grid

        self.assertEqual(grid[1][2][3].mass, 38 | units.kg)
        self.assertEqual(grid[0:2][1][2][3].mass, 38 | units.kg)
        self.assertEqual(len(grid[1][2].mass), 5)
        self.assertTrue(numpy.all(grid[1][2].mass == [35, 36, 37, 38, 39] | units.kg))
//  [python]
def test_batch_standard_compatibility() -> None:
    batch = Batch(a=np.array([[1.0, 2.0], [3.0, 4.0]]), b=Batch(), c=np.array([5.0, 6.0]))
    batch_mean = np.mean(batch)
    assert isinstance(batch_mean, Batch)  # type: ignore  # mypy doesn't know but it works, cf. `batch.rst`
    assert sorted(batch_mean.get_keys()) == ["a", "b", "c"]  # type: ignore
    with pytest.raises(TypeError):
        len(batch_mean)
    assert np.all(batch_mean.a == np.mean(batch.a, axis=0))
    assert batch_mean.c == np.mean(batch.c, axis=0)
    with pytest.raises(IndexError):
        Batch()[0]
//  [python]
def translate(self, x=None, y=None, z=None):
        '''
        mesh.translate(dx, dy, dz) copies the given mesh and returns the duplicate mesh after its
          coordinates have been translated by the vector (dx, dy, dz). Any of the dx, dy, and dz
          values may be None to indicate 0 translation.
        mesh.translate([dx, dy, dz]) is equivalent to mesh.translate(dx, dy, dz).
        '''
        if y is None and z is None and x is not None:
            x = np.asarray(x)
            if len(x.shape) > 0 and x.shape[0] == 3:
                (x, y, z) = x
        if x is None: x = 0
        if y is None: y = 0
        if z is None: z = 0
        if np.all((x == 0) & (y == 0) & (z == 0)): return self
        (x0, y0, z0) = self.coordinates
        return self.copy(coordinates=[x0 + x, y0 + y, z0 + z])
//  [python]
def HPO(x, data, index=None):
    # breakpoint()
    if index is not None:
        ret = data["accs"][index]
        return np.array(ret).reshape(-1, 1)
    else:
        X = get_HPO_domain(data)
        try:
            idx = np.asscalar(np.where(np.all(x == X, axis=1))[0][0])
            # idx = np.asscalar(np.where(np.all(np.isclose(x, X), axis=1))[0][0])
        except:
            print('[HPO] except')
            print(x)
            print(np.all(x == X, axis=1))
            print(np.where(np.all(x == X, axis=1)))
            print((np.where(np.all(x == X, axis=1))[0]))
            exit(10)
        ret = data["accs"][idx]
        return np.array(ret).reshape(-1,1)
//  [python]
def neighborhood_face_partition(indexed_neighborhoods, indexed_faces, labels, index):
        '''
        tess.neighborhood_face_partition is a tuple of subsets of the face matrix that each
        provides a set of faces that can be operated on simultaneously in a multi-step vectorized
        neighborhood calculatiion. Each matrix of the tuple has size 3 x n where n varies, and each
        columnn of the matrix has been rolled such that the first row contains a unique set of
        vertices over which neighborhood calculations can be implemented. See also the
        neighborhood_vectorize method.
        '''
        nnei = np.array([len(uu) for uu in indexed_neighborhoods])
        emax = np.max(nnei)
        whs  = tuple([np.where(nnei > k)[0] for k in range(emax)])
        neis = np.array([u + (-1,)*(emax - len(u)) for u in indexed_neighborhoods], dtype=int)
        # Prep the triangle orderings.
        neifaces = []
        for (k,a) in enumerate(whs):
            b = neis[a, k]
            c = neis[a, np.mod(k+1, nnei[a])]
            fs = index[[labels[a],labels[b],labels[c]]]
            ii = ~(fs == None)
            if not n.all(ii): (a,b,c,fs) = [u[ii].astype(int) for u in (a,b,c,fs)]
            abc = indexed_faces[:, fs]
            aii = np.argmax(abc == (a,a,a), axis=0)
            allfs = np.arange(len(a))*3
            abc = abc.T.flatten()
            abc = np.array([abc[(aii+k)%3 + allfs] for k in range(3)])
            neifaces.append(abc)
        for nf in neifaces: nf.setflags(write=False)
        return tuple(neifaces)
//  [python]
def neighborhood_face_partition(indexed_neighborhoods, indexed_faces, labels, index):
        '''
        tess.neighborhood_face_partition is a tuple of subsets of the face matrix that each
        provides a set of faces that can be operated on simultaneously in a multi-step vectorized
        neighborhood calculatiion. Each matrix of the tuple has size 3 x n where n varies, and each
        columnn of the matrix has been rolled such that the first row contains a unique set of
        vertices over which neighborhood calculations can be implemented. See also the
        neighborhood_vectorize method.
        '''
        nnei = np.array([len(uu) for uu in indexed_neighborhoods])
        emax = np.max(nnei)
        whs  = tuple([np.where(nnei > k)[0] for k in range(emax)])
        neis = np.array([u + (-1,)*(emax - len(u)) for u in indexed_neighborhoods], dtype=int)
        # Prep the triangle orderings.
        neifaces = []
        for (k,a) in enumerate(whs):
            b = neis[a, k]
            c = neis[a, np.mod(k+1, nnei[a])]
            fs = index[[labels[a],labels[b],labels[c]]]
            ii = ~(fs == None)
            if not np.all(ii): (a,b,c,fs) = [u[ii].astype(int) for u in (a,b,c,fs)]
            abc = indexed_faces[:, fs]
            aii = np.argmax(abc == (a,a,a), axis=0)
            allfs = np.arange(len(a))*3
            abc = abc.T.flatten()
            abc = np.array([abc[(aii+k)%3 + allfs] for k in range(3)])
            neifaces.append(abc)
        for nf in neifaces: nf.setflags(write=False)
        return tuple(neifaces)
//  [python]
def test_batch_over_batch() -> None:
    batch = Batch(a=[3, 4, 5], b=[4, 5, 6])
    batch2 = Batch({"c": [6, 7, 8], "b": batch})
    batch2.b.b[-1] = 0
    print(batch2)
    for k, v in batch2.items():
        assert n.all(batch2[k] == v)
    assert batch2[-1].b.b == 0
    batch2.cat_(Batch(c=[6, 7, 8], b=batch))
    assert n.allclose(batch2.c, [6, 7, 8, 6, 7, 8])
    assert n.allclose(batch2.b.a, [3, 4, 5, 3, 4, 5])
    assert n.allclose(batch2.b.b, [4, 5, 0, 4, 5, 0])
    batch2.update(batch2.b, six=[6, 6, 6])
    assert n.allclose(batch2.c, [6, 7, 8, 6, 7, 8])
    assert n.allclose(batch2.a, [3, 4, 5, 3, 4, 5])
    assert n.allclose(batch2.b, [4, 5, 0, 4, 5, 0])
    assert n.allclose(batch2.six, [6, 6, 6])
    d = {"a": [3, 4, 5], "b": [4, 5, 6]}
    batch3 = Batch(c=[6, 7, 8], b=d)
    batch3.cat_(Batch(c=[6, 7, 8], b=d))
    assert n.allclose(batch3.c, [6, 7, 8, 6, 7, 8])
    assert n.allclose(batch3.b.a, [3, 4, 5, 3, 4, 5])
    assert n.allclose(batch3.b.b, [4, 5, 6, 4, 5, 6])
    batch4 = Batch(({"a": {"b": np.array([1.0])}},))
    assert batch4.a.b.ndim == 2
    assert batch4.a.b[0, 0] == 1.0
    # advanced slicing
    batch5 = Batch(a=[[1, 2]], b={"c": np.zeros([3, 2, 1])})
    assert batch5.shape == [1, 2]
    with pytest.raises(IndexError):
        batch5[2]
    with pytest.raises(IndexError):
        batch5[:, 3]
    with pytest.raises(IndexError):
        batch5[:, :, -1]
    batch5[:, -1] += np.int_(1)
    assert n.allclose(batch5.a, [1, 3])
    assert n.allclose(batch5.b.c.squeeze(), [[0, 1]] * 3)
    with pytest.raises(ValueError):
        batch5[:, -1] = 1
    batch5[:, 0] = {"a": -1}
    assert n.allclose(batch5.a, [-1, 3])
    assert n.allclose(batch5.b.c.squeeze(), [[0, 1]] * 3)
//  [python]
def test_batch_over_batch() -> None:
    batch = Batch(a=[3, 4, 5], b=[4, 5, 6])
    batch2 = Batch({"c": [6, 7, 8], "b": batch})
    batch2.b.b[-1] = 0
    print(batch2)
    for k, v in batch2.items():
        assert np.all(batch2[k] == v)
    assert batch2[-1].b.b == 0
    batch2.cat_(Batch(c=[6, 7, 8], b=batch))
    assert np.allclose(batch2.c, [6, 7, 8, 6, 7, 8])
    assert np.allclose(batch2.b.a, [3, 4, 5, 3, 4, 5])
    assert np.allclose(batch2.b.b, [4, 5, 0, 4, 5, 0])
    batch2.update(batch2.b, six=[6, 6, 6])
    assert np.allclose(batch2.c, [6, 7, 8, 6, 7, 8])
    assert np.allclose(batch2.a, [3, 4, 5, 3, 4, 5])
    assert np.allclose(batch2.b, [4, 5, 0, 4, 5, 0])
    assert np.allclose(batch2.six, [6, 6, 6])
    d = {"a": [3, 4, 5], "b": [4, 5, 6]}
    batch3 = Batch(c=[6, 7, 8], b=d)
    batch3.cat_(Batch(c=[6, 7, 8], b=d))
    assert np.allclose(batch3.c, [6, 7, 8, 6, 7, 8])
    assert np.allclose(batch3.b.a, [3, 4, 5, 3, 4, 5])
    assert np.allclose(batch3.b.b, [4, 5, 6, 4, 5, 6])
    batch4 = Batch(({"a": {"b": np.array([1.0])}},))
    assert batch4.a.b.ndim == 2
    assert batch4.a.b[0, 0] == 1.0
    # advanced slicing
    batch5 = Batch(a=[[1, 2]], b={"c": np.zeros([3, 2, 1])})
    assert batch5.shape == [1, 2]
    with pytest.raises(IndexError):
        batch5[2]
    with pytest.raises(IndexError):
        batch5[:, 3]
    with pytest.raises(IndexError):
        batch5[:, :, -1]
    batch5[:, -1] += np.int_(1)
    assert np.allclose(batch5.a, [1, 3])
    assert np.allclose(batch5.b.c.squeeze(), [[0, 1]] * 3)
    with pytest.raises(ValueError):
        batch5[:, -1] = 1
    batch5[:, 0] = {"a": -1}
    assert np.allclose(batch5.a, [-1, 3])
    assert np.allclose(batch5.b.c.squeeze(), [[0, 1]] * 3)
//  [python]
def Antigen(x, data, index=None):
    # breakpoint()
    if index is not None:
        ret = data["accs"][index]
        return np.array(ret).reshape(-1, 1)
    else:
        X = get_Antigen_domain(data)
        try:
            idx = np.asscalar(np.where(np.logical_and.reduce([x == X[:,i] for i in range(x.shape[0])]))[0][0])
            # idx = np.asscalar(np.where(np.all(np.isclose(x, X), axis=1))[0][0])
        except:
            print('FAILED to find x in data.')
            print(x)
            print(np.logical_and.reduce([x == X[:,i] for i in range(x.shape[0])]))
            print(np.where(np.logical_and.reduce([x == X[:,i] for i in range(x.shape[0])])))
            print((np.where(np.logical_and.reduce([x == X[:,i] for i in range(x.shape[0])]))[0]))
            exit(10)
        ret = data["accs"][idx]
        return np.array(ret).reshape(-1,1)
//  [python]
def Antigen(x, data, index=None):
    # breakpoint()
    if index is not None:
        ret = data["accs"][index]
        return np.array(ret).reshape(-1, 1)
    else:
        X = get_Antigen_domain(data)
        try:
            idx = np.asscalar(np.where(np.all(x == X, axis=1))[0][0])
            # idx = np.asscalar(np.where(np.all(np.isclose(x, X), axis=1))[0][0])
        except:
            print('FAILED to find x in data.')
            print(x)
            print(np.all(x == X, axis=1))
            print(np.where(np.all(x == X, axis=1)))
            print((np.where(np.all(x == X, axis=1))[0]))
            exit(10)
        ret = data["accs"][idx]
        return np.array(ret).reshape(-1,1)
//  [python]
def _lazyselect(condlist, choicelist, arrays, default=0):
    """
    Mimic `np.select(condlist, choicelist)`.

    Notice, it assumes that all `arrays` are of the same shape or can be
    broadcasted together.

    All functions in `choicelist` must accept array arguments in the order
    given in `arrays` and must return an array of the same shape as broadcasted
    `arrays`.

    Examples
    --------
    >>> x = np.arange(6)
    >>> np.select([x <3, x > 3], [x**2, x**3], default=0)
    array([  0,   1,   4,   0,  64, 125])

    >>> _lazyselect([x < 3, x > 3], [lambda x: x**2, lambda x: x**3], (x,))
    array([   0.,    1.,    4.,   0.,   64.,  125.])

    >>> a = -np.ones_like(x)
    >>> _lazyselect([x < 3, x > 3],
    ...             [lambda x, a: x**2, lambda x, a: a * x**3],
    ...             (x, a), default=np.nan)
    array([   0.,    1.,    4.,   nan,  -64., -125.])

    """
    arrays = np.broadcast_arrays(*arrays)
    tcode = np.mintypecode([a.dtype.char for a in arrays])
    out = _valarray(np.shape(arrays[0]), value=default, typecode=tcode)
    for index in range(len(condlist)):
        func, cond = choicelist[index], condlist[index]
        if np.all(cond is False):
            continue
        cond, _ = np.broadcast_arrays(cond, arrays[0])
        temp = tuple(np.extract(cond, arr) for arr in arrays)
        np.place(out, cond, func(*temp))
    return out
//  [python]
def check(self, function, exact_fill_value=None, masked_result=True, **kwargs):
        result = function(self.ma, **kwargs)
        expected_data = function(self.ma.filled(np.nan), **kwargs)
        expected_mask = np.isnan(expected_data)
        if masked_result:
            assert isinstance(result, Masked)
            assert_array_equal(result.mask, expected_mask)
            assert np.all(result == expected_data)
        else:
            assert not isinstance(result, Masked)
            assert_array_equal(result, expected_data)
            assert not np.any(expected_mask)
        out = np.zeros_like(result)
        result2 = function(self.ma, out=out, **kwargs)
        assert result2 is out
        assert_array_equal(result2, result)
//  [python]
def check_expected(image, expected, mask=None, ignore=False):
    if mask is None and not image.has_crop_mask:
        numpy.testing.assert_array_almost_equal(image.pixel_data, expected)

        assert not image.has_mask
    elif mask is not None and ignore:
        numpy.testing.assert_array_almost_equal(image.pixel_data, expected)

        assert image.has_mask
    elif mask is not None and not ignore:
        assert image.has_mask

        if not image.has_crop_mask:
            assert numpy.alltrue(mask == image.mask)

        numpy.testing.assert_array_almost_equal(
            image.pixel_data, expected * mask
        )
//  [python]
def check_expected(image, expected, mask=None, ignore=False):
    if mask is None and not image.has_crop_mask:
        numpy.testing.assert_array_almost_equal(image.pixel_data, expected)

        assert not image.has_mask
    elif mask is not None and ignore:
        numpy.testing.assert_array_almost_equal(image.pixel_data, expected)

        assert image.has_mask
    elif mask is not None and not ignore:
        assert image.has_mask

        if not image.has_crop_mask:
            assert numpy.all(mask == image.mask)

        numpy.testing.assert_array_almost_equal(
            image.pixel_data, expected * mask
        )
//  [python]
def __eq__(self, other: object) -> bool:
        """Returns ``True`` if matrices are equal."""
        if not isinstance(other, Matrix):
            raise TypeError("Matrix class required.")
        if self.shape != other.shape:
            raise TypeError("Matrices have different shapes.")
        return bool(np.alltrue(self.matrix == other.matrix))
//  [python]
def __eq__(self, other: object) -> bool:
        """Returns ``True`` if matrices are equal."""
        if not isinstance(other, Matrix):
            raise TypeError("Matrix class required.")
        if self.shape != other.shape:
            raise TypeError("Matrices have different shapes.")
        return bool(np.all(self.matrix == other.matrix))
//  [python]
def isclose(self, other: object) -> bool:
        """Returns ``True`` if matrices are close to equal, tolerance value for
        comparison is adjustable by the attribute :attr:`Matrix.abs_tol`.

        """
        if not isinstance(other, Matrix):
            raise TypeError("Matrix class required.")
        if self.shape != other.shape:
            raise TypeError("Matrices have different shapes.")
        return bool(np.all(np.isclose(self.matrix, other.matrix, atol=self.abs_tol)))
//  [python]
def gto_norm(l, expnt):
    r'''Normalized factor for GTO radial part   :math:`g=r^l e^{-\alpha r^2}`

    .. math::

        \frac{1}{\sqrt{\int g^2 r^2 dr}}
        = \sqrt{\frac{2^{2l+3} (l+1)! (2a)^{l+1.5}}{(2l+2)!\sqrt{\pi}}}

    Ref: H. B. Schlegel and M. J. Frisch, Int. J. Quant.  Chem., 54(1995), 83-87.

    Args:
        l (int):
            angular momentum
        expnt :
            exponent :math:`\alpha`

    Returns:
        normalization factor

    Examples:

    >>> print(gto_norm(0, 1))
    2.5264751109842591
    '''
    if n.all(l >= 0):
        #f = 2**(2*l+3) * math.factorial(l+1) * (2*expnt)**(l+1.5) \
        #        / (math.factorial(2*l+2) * math.sqrt(math.pi))
        #return math.sqrt(f)
        return 1/numpy.sqrt(gaussian_int(l*2+2, 2*expnt))
    else:
        raise ValueError('l should be >= 0')
//  [python]
def gto_norm(l, expnt):
    r'''Normalized factor for GTO radial part   :math:`g=r^l e^{-\alpha r^2}`

    .. math::

        \frac{1}{\sqrt{\int g^2 r^2 dr}}
        = \sqrt{\frac{2^{2l+3} (l+1)! (2a)^{l+1.5}}{(2l+2)!\sqrt{\pi}}}

    Ref: H. B. Schlegel and M. J. Frisch, Int. J. Quant.  Chem., 54(1995), 83-87.

    Args:
        l (int):
            angular momentum
        expnt :
            exponent :math:`\alpha`

    Returns:
        normalization factor

    Examples:

    >>> print(gto_norm(0, 1))
    2.5264751109842591
    '''
    if numpy.all(l >= 0):
        #f = 2**(2*l+3) * math.factorial(l+1) * (2*expnt)**(l+1.5) \
        #        / (math.factorial(2*l+2) * math.sqrt(math.pi))
        #return math.sqrt(f)
        return 1/numpy.sqrt(gaussian_int(l*2+2, 2*expnt))
    else:
        raise ValueError('l should be >= 0')
//  [python]
def test_borrow_input(self):
        """
        Tests that the contract for io.In is respected. When borrow=False, it should be
        impossible for outputs to be aliased to the input variables provided by the user,
        either through a view-map or a destroy map. New tests should be added in the future
        when borrow=True is implemented.
        """
        a = T.dmatrix()
        aval = numpy.random.rand(3, 3)

        # when borrow=False, test that a destroy map cannot alias output to input
        f = theano.function([In(a, borrow=False)], Out(a+1, borrow=True))
        assert n.all(f(aval) == aval+1)
        assert not numpy.may_share_memory(aval, f(aval))

        # when borrow=False, test that a viewmap cannot alias output to input
        f = theano.function([In(a, borrow=False)], Out(a[0, :], borrow=True))
        assert n.all(f(aval) == aval[0, :])
        assert not numpy.may_share_memory(aval, f(aval))
//  [python]
def test_borrow_input(self):
        """
        Tests that the contract for io.In is respected. When borrow=False, it should be
        impossible for outputs to be aliased to the input variables provided by the user,
        either through a view-map or a destroy map. New tests should be added in the future
        when borrow=True is implemented.
        """
        a = T.dmatrix()
        aval = numpy.random.rand(3, 3)

        # when borrow=False, test that a destroy map cannot alias output to input
        f = theano.function([In(a, borrow=False)], Out(a+1, borrow=True))
        assert numpy.all(f(aval) == aval+1)
        assert not numpy.may_share_memory(aval, f(aval))

        # when borrow=False, test that a viewmap cannot alias output to input
        f = theano.function([In(a, borrow=False)], Out(a[0, :], borrow=True))
        assert numpy.all(f(aval) == aval[0, :])
        assert not numpy.may_share_memory(aval, f(aval))
//  [python]
def __eq__(self, other):
        result = NotImplemented
        if isinstance(other, PPField):
            result = True
            for attr in self.__slots__:
                attrs = [hasattr(self, attr), hasattr(other, attr)]
                if all(attrs):
                    self_attr = getattr(self, attr)
                    other_attr = getattr(other, attr)
                    if not np.all(self_attr == other_attr):
                        result = False
                        break
                elif any(attrs):
                    result = False
                    break
        return result
//  [python]
def test_borrow_output(self):
        a = T.dmatrix()
        f = function([a], Out(a, borrow=False))
        o = N.ones((3, 3))
        assert o is not f(o)  # function no longer permits aliasing outputs to inputs

        f = function([a], Out(a*4, borrow=False))
        o = N.ones((3, 3))
        four = f(o)
        assert N.all(four == 4)
        f(o+.1)  # should not clobber the memory used to store four
        assert N.all(four == 4)

        f = function([a], Out(a*4, borrow=True), mode=theano.Mode('c|py_nogc', 'fast_run'))
        o = N.ones((3, 3))
        four = f(o)
        assert N.all(four == 4)
        f(o+.1)  # should clobber the memory used to store four
        if theano.config.cxx:
            assert not N.all(four == 4)
        else:
            # The Elemwise.perform method don't reuse memory
            # as some numpy version don't support that correctly.
            assert N.all(four == 4)
//  [python]
def test_borrow_output(self):
        a = T.dmatrix()
        f = function([a], Out(a, borrow=False))
        o = N.ones((3, 3))
        assert o is not f(o)  # function no longer permits aliasing outputs to inputs

        f = function([a], Out(a*4, borrow=False))
        o = N.ones((3, 3))
        four = f(o)
        assert numpy.all(four == 4)
        f(o+.1)  # should not clobber the memory used to store four
        assert numpy.all(four == 4)

        f = function([a], Out(a*4, borrow=True), mode=theano.Mode('c|py_nogc', 'fast_run'))
        o = N.ones((3, 3))
        four = f(o)
        assert numpy.all(four == 4)
        f(o+.1)  # should clobber the memory used to store four
        if theano.config.cxx:
            assert not numpy.all(four == 4)
        else:
            # The Elemwise.perform method don't reuse memory
            # as some numpy version don't support that correctly.
            assert numpy.all(four == 4)
//  [python]
def standard_deviation(self):
        """float: The error estimate on the mean."""
        if numpy.all(self.data == self.data[0]):
            return 0

        return numpy.sqrt(numpy.max(self._block_variance))
//  [python]
def __call__(self, actions, parent, sibling, dangling, finished):

        # Filter out finished sequences to save compute
        unfinished = ~finished
        final_combined_prior = np.zeros((actions.shape[0], self.L), dtype=np.float32)
        unfinished_actions = actions[unfinished]
        unfinished_parent = parent[unfinished]
        unfinished_sibling = sibling[unfinished]
        unfinished_dangling = dangling[unfinished]

        # Sum the individual priors
        zero_prior = np.zeros((unfinished_actions.shape[0], self.L), dtype=np.float32)
        ind_priors = [zero_prior.copy() for _ in range(len(self.priors))]
        for i in range(len(self.priors)):
            ind_priors[i] += self.priors[i](unfinished_actions, unfinished_parent, unfinished_sibling, unfinished_dangling)
        combined_prior = sum(ind_priors) + zero_prior # TBD FIX HACK

        # Count number of constrained tokens per prior
        if self.do_count:
            self.total_tokens += unfinished_actions.shape[0] * self.library.L
            for i in range(len(self.constraint_counts)):
                self.constraint_counts[i] += np.count_nonzero(ind_priors[i] == -np.inf)
            self.total_constraints += np.count_nonzero(combined_prior == -np.inf)

        # Give status report if a prior contains all -inf
        collision_mask = np.all(np.isneginf(combined_prior), axis=1)
        if np.any(collision_mask):
            collisions = collision_mask.nonzero()[0] # Indices of collisions
            msg = []
            msg.append("ERROR in {}:".format(__file__))
            msg.append("Encountered collision(s) in prior. This occurs when a prior contains all -inf values. " +
                       "This typically indicates a logic error in a prior formulation, or a 'collision' when " +
                       "configuring multiple priors that are not always compatible. See the table(s) below " +
                       "for which priors caused each collision. X's indicate a value of -inf.")
            for i, collision in enumerate(collisions):
                msg.append("\n----- Collision {} of {} -----".format(i + 1, len(collisions)))
                msg.append("Traversal: {}".format(self.library.tokenize(unfinished_actions[collision])))
                table = PrettyTable(["Prior"] + self.library.names)
                for j, ind_prior in enumerate(ind_priors):
                    if np.any(np.isneginf(ind_prior[collision])):
                        name = self.priors[j].__class__.__name__
                        row = ["X" if x == -np.inf else "" for x in ind_prior[collision]]
                        table.add_row([name] + row)
                msg.append(repr(table))
            # msg.append("\nApplication will now exit.")
            print("\n".join(msg))
            # os._exit(1) # Bypass tensorflow exception-handling

        final_combined_prior[unfinished] = combined_prior

        return final_combined_prior