jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testRngUniformF32(self):
    lo, hi = 2., 4.
    shape = (2, 3)
    c = self._NewComputation()
    c.RngUniform(
        c.Constant(NumpyArrayF32(lo)),
        c.Constant(NumpyArrayF32(hi)),
        dims=shape)
    result = xla_client.execute_with_python_values(c.Build().Compile())
    # since the result is random, we just check shape, uniqueness, and range
    self.assertEqual(result.shape, shape)
    self.assertLen(np.unique(result), np.prod(shape))
    self.assertTrue(np.all(lo <= result))
    self.assertTrue(np.all(result < hi))
//  [python]
def testConvGeneralDilatedGroupedConvolutionF32(self):
    c = self._NewComputation()
    a = lambda *dims: np.arange(n.prod(dims)).reshape(dims).astype("float32")
    lhs = a(1, 2, 2, 3)
    rhs = a(2, 1, 1, 2) * 10
    strides = [1, 1]
    pads = [(1, 0), (0, 1)]
    lhs_dilation = (2, 1)
    rhs_dilation = (1, 1)
    dimension_numbers = ("NCHW", "OIHW", "NCHW")
    feature_group_count = 2
    c.ConvGeneralDilated(
        c.Constant(lhs), c.Constant(rhs), strides, pads, lhs_dilation,
        rhs_dilation, dimension_numbers, feature_group_count)
    result = np.array([[[
        [0., 0., 0.],
        [10., 20., 0.],
        [0., 0., 0.],
        [40., 50., 0.],
    ], [
        [0., 0., 0.],
        [330., 380., 160.],
        [0., 0., 0.],
        [480., 530., 220.],
    ]]])
    self._ExecuteAndCompareClose(c, expected=result)
//  [python]
def testConvGeneralDilatedGroupedConvolutionF32(self):
    c = self._NewComputation()
    a = lambda *dims: np.arange(np.prod(dims)).reshape(dims).astype("float32")
    lhs = a(1, 2, 2, 3)
    rhs = a(2, 1, 1, 2) * 10
    strides = [1, 1]
    pads = [(1, 0), (0, 1)]
    lhs_dilation = (2, 1)
    rhs_dilation = (1, 1)
    dimension_numbers = ("NCHW", "OIHW", "NCHW")
    feature_group_count = 2
    c.ConvGeneralDilated(
        c.Constant(lhs), c.Constant(rhs), strides, pads, lhs_dilation,
        rhs_dilation, dimension_numbers, feature_group_count)
    result = np.array([[[
        [0., 0., 0.],
        [10., 20., 0.],
        [0., 0., 0.],
        [40., 50., 0.],
    ], [
        [0., 0., 0.],
        [330., 380., 160.],
        [0., 0., 0.],
        [480., 530., 220.],
    ]]])
    self._ExecuteAndCompareClose(c, expected=result)
//  [python]
def getMemoryRequirement(self):
        """
        Get memory requirements of weights
        :return: memory requirement
        """
        return (numpy.prod(self.filter_shape) + self.filter_shape[0]) * 4
//  [python]
def allocate_hexagon_array(
    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
):
    """
    Allocate a hexagon array which could be a 2D array
    on physical memory defined by axis_separators
    """
    if tensor_shape is None:
        assert data is not None, "Must provide either tensor shape or numpy data array"
        tensor_shape = data.shape
    elif data is not None:
        assert (
            tensor_shape == data.shape
        ), "Mismatch between provided tensor shape and numpy data array shape"

    if dtype is None:
        assert data is not None, "Must provide either dtype or numpy data array"
        dtype = data.dtype.name
    elif data is not None:
        assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"

    if axis_separators is None:
        axis_separators = []

    boundaries = [0, *axis_separators, len(tensor_shape)]
    physical_shape = [
        numpy.multiply.reduce(tensor_shape[dim_i:dim_f])
        for dim_i, dim_f in zip(boundaries[:-1], boundaries[1:])
    ]

    arr = tvm.nd.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)

    if data is not None:
        arr.copyfrom(data.reshape(physical_shape))

    return arr._create_view(tensor_shape)
//  [python]
def allocate_hexagon_array(
    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
):
    """
    Allocate a hexagon array which could be a 2D array
    on physical memory defined by axis_separators
    """
    if tensor_shape is None:
        assert data is not None, "Must provide either tensor shape or numpy data array"
        tensor_shape = data.shape
    elif data is not None:
        assert (
            tensor_shape == data.shape
        ), "Mismatch between provided tensor shape and numpy data array shape"

    if dtype is None:
        assert data is not None, "Must provide either dtype or numpy data array"
        dtype = data.dtype.name
    elif data is not None:
        assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"

    if axis_separators is None:
        axis_separators = []

    boundaries = [0, *axis_separators, len(tensor_shape)]
    physical_shape = [
        numpy.prod(tensor_shape[dim_i:dim_f])
        for dim_i, dim_f in zip(boundaries[:-1], boundaries[1:])
    ]

    arr = tvm.nd.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)

    if data is not None:
        arr.copyfrom(data.reshape(physical_shape))

    return arr._create_view(tensor_shape)
//  [python]
def gradient_class(Nd, axis):
    d_indx, dt_indx = indxmap_diff(Nd)
    I = scipy.sparse.eye(n.prod(Nd), n.prod(Nd))
    data = n.ones((n.prod(Nd),))
    
    row_ind = d_indx[axis]
    col_ind= n.arange(0, n.prod(Nd)).astype(n.int) 
    G = scipy.sparse.csr_matrix(( data, 
                                            (row_ind, col_ind)), shape = (n.prod(Nd),n.prod(Nd))
                                 )
    G = G- I

    G = G.tocsr()
    return G
//  [python]
def gradient_class(Nd, axis):
    d_indx, dt_indx = indxmap_diff(Nd)
    I = scipy.sparse.eye(numpy.prod(Nd), numpy.prod(Nd))
    data = numpy.ones((numpy.prod(Nd),))
    
    row_ind = d_indx[axis]
    col_ind= numpy.arange(0, numpy.prod(Nd)).astype(numpy.int) 
    G = scipy.sparse.csr_matrix(( data, 
                                            (row_ind, col_ind)), shape = (numpy.prod(Nd),numpy.prod(Nd))
                                 )
    G = G- I

    G = G.tocsr()
    return G
//  [python]
def random_unique_float(shape, dtype):
    # create a random float array with 10x length
    numel = np.multiply.reduce(shape)
    arr = np.random.uniform(-10.0, 10.0, numel * 10).astype(dtype)
    arr = np.unique(arr)
    assert (
        arr.shape[0] >= numel
    ), "failed to create enough unique values: %d vs %d" % (arr.shape[0], numel)
    arr = arr[:numel]
    np.random.shuffle(arr)
    arr = arr.reshape(shape)
    return arr
//  [python]
def random_unique_float(shape, dtype):
    # create a random float array with 10x length
    numel = np.prod(shape)
    arr = np.random.uniform(-10.0, 10.0, numel * 10).astype(dtype)
    arr = np.unique(arr)
    assert (
        arr.shape[0] >= numel
    ), "failed to create enough unique values: %d vs %d" % (arr.shape[0], numel)
    arr = arr[:numel]
    np.random.shuffle(arr)
    arr = arr.reshape(shape)
    return arr
//  [python]
def gradient_class2(Nd, axis):
    d_indx, dt_indx = indxmap_diff(Nd)
    I = scipy.sparse.eye(n.prod(Nd), n.prod(Nd))
    data = n.ones((n.prod(Nd),))
    
    row_ind = dt_indx[axis]
    col_ind= n.arange(0, n.prod(Nd)).astype(n.int) 
    G = scipy.sparse.csr_matrix(( data, 
                                            (row_ind, col_ind)), shape = (n.prod(Nd),n.prod(Nd))
                                 )
    G = G- I

    G = G.tocsr()
    return G
//  [python]
def gradient_class2(Nd, axis):
    d_indx, dt_indx = indxmap_diff(Nd)
    I = scipy.sparse.eye(numpy.prod(Nd), numpy.prod(Nd))
    data = numpy.ones((numpy.prod(Nd),))
    
    row_ind = dt_indx[axis]
    col_ind= numpy.arange(0, numpy.prod(Nd)).astype(numpy.int) 
    G = scipy.sparse.csr_matrix(( data, 
                                            (row_ind, col_ind)), shape = (numpy.prod(Nd),numpy.prod(Nd))
                                 )
    G = G- I

    G = G.tocsr()
    return G
//  [python]
def _row_sparse_ndarray_from_definition(data, indices, shape=None, ctx=None,
                                        dtype=None, indices_type=None):
    """Create a `RowSparseNDArray` based on data and indices"""
    storage_type = 'row_sparse'
    # context
    ctx = Context.default_ctx if ctx is None else ctx
    # types
    dtype = _prepare_default_dtype(data, dtype)
    indices_type = _STORAGE_AUX_TYPES[storage_type][0] if indices_type is None else indices_type
    # prepare src array and types
    data = _prepare_src_array(data, dtype)
    indices = _prepare_src_array(indices, indices_type)

    # TODO(junwu): Convert data, indptr, and indices to mxnet NDArrays
    # if they are not for now. In the future, we should provide a c-api
    # to accept np.ndarray types to copy from to result.data and aux_data
    if not isinstance(data, NDArray):
        data = _array(data, ctx, dtype)
    if not isinstance(indices, NDArray):
        indices = _array(indices, ctx, indices_type)
    if shape is None:
        num_indices = indices.shape[0]
        if num_indices == 0:
            raise ValueError('invalid shape')
        dim0 = indices[num_indices - 1].asscalar() + 1
        shape = (dim0, ) + data.shape[1:]
    # verify shapes
    if data.ndim != len(shape) or indices.ndim != 1 or reduce(lambda x, y: x * y, shape[1:]) == 0:
        raise ValueError("invalid shape")
    result = RowSparseNDArray(_new_alloc_handle(storage_type, shape, ctx, False, dtype,
                                                [indices_type], [indices.shape]))
    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, data.handle, ctypes.c_int(-1)))
    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, indices.handle, ctypes.c_int(0)))
    return result
//  [python]
def _row_sparse_ndarray_from_definition(data, indices, shape=None, ctx=None,
                                        dtype=None, indices_type=None):
    """Create a `RowSparseNDArray` based on data and indices"""
    storage_type = 'row_sparse'
    # context
    ctx = Context.default_ctx if ctx is None else ctx
    # types
    dtype = _prepare_default_dtype(data, dtype)
    indices_type = _STORAGE_AUX_TYPES[storage_type][0] if indices_type is None else indices_type
    # prepare src array and types
    data = _prepare_src_array(data, dtype)
    indices = _prepare_src_array(indices, indices_type)

    # TODO(junwu): Convert data, indptr, and indices to mxnet NDArrays
    # if they are not for now. In the future, we should provide a c-api
    # to accept np.ndarray types to copy from to result.data and aux_data
    if not isinstance(data, NDArray):
        data = _array(data, ctx, dtype)
    if not isinstance(indices, NDArray):
        indices = _array(indices, ctx, indices_type)
    if shape is None:
        num_indices = indices.shape[0]
        if num_indices == 0:
            raise ValueError('invalid shape')
        dim0 = indices[num_indices - 1].asscalar() + 1
        shape = (dim0, ) + data.shape[1:]
    # verify shapes
    if data.ndim != len(shape) or indices.ndim != 1 or np.prod(shape[1:]) == 0:
        raise ValueError("invalid shape")
    result = RowSparseNDArray(_new_alloc_handle(storage_type, shape, ctx, False, dtype,
                                                [indices_type], [indices.shape]))
    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, data.handle, ctypes.c_int(-1)))
    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, indices.handle, ctypes.c_int(0)))
    return result
//  [python]
def indxmap_diff(Nd):
    """
    Preindixing for rapid image gradient ()
    
    Diff(x) = x.flat[d_indx[0]] - x.flat
    
    Diff_t(x) =  x.flat[dt_indx[0]] - x.flat
                            
    :param Nd: the dimension of the image
    :type Nd: tuple with integers
    :returns d_indx: image gradient
    :returns  dt_indx:  the transpose of the image gradient 
    :rtype: d_indx: lists with numpy ndarray
    :rtype: dt_indx: lists with numpy ndarray
    """    

    ndims = len(Nd)
    Ndprod = 1
    for x in Nd:
        Ndprod = Ndprod * x
    mylist = numpy.arange(0, Ndprod).astype(numpy.int32)
    mylist = numpy.reshape(mylist, Nd)
    d_indx = []
    dt_indx = []
    for pp in range(0, ndims):
        d_indx = d_indx + [ numpy.reshape(   numpy.roll(  mylist, +1 , pp  ), (Ndprod,)  ,order='C').astype(numpy.int32) ,]
        dt_indx = dt_indx + [ numpy.reshape(   numpy.roll(  mylist, -1 , pp  ) , (Ndprod,) ,order='C').astype(numpy.int32) ,]

    return d_indx,  dt_indx
//  [python]
def indxmap_diff(Nd):
    """
    Preindixing for rapid image gradient ()
    
    Diff(x) = x.flat[d_indx[0]] - x.flat
    
    Diff_t(x) =  x.flat[dt_indx[0]] - x.flat
                            
    :param Nd: the dimension of the image
    :type Nd: tuple with integers
    :returns d_indx: image gradient
    :returns  dt_indx:  the transpose of the image gradient 
    :rtype: d_indx: lists with numpy ndarray
    :rtype: dt_indx: lists with numpy ndarray
    """    

    ndims = len(Nd)
    Ndprod = numpy.prod(Nd)
    mylist = numpy.arange(0, Ndprod).astype(numpy.int32)
    mylist = numpy.reshape(mylist, Nd)
    d_indx = []
    dt_indx = []
    for pp in range(0, ndims):
        d_indx = d_indx + [ numpy.reshape(   numpy.roll(  mylist, +1 , pp  ), (Ndprod,)  ,order='C').astype(numpy.int32) ,]
        dt_indx = dt_indx + [ numpy.reshape(   numpy.roll(  mylist, -1 , pp  ) , (Ndprod,) ,order='C').astype(numpy.int32) ,]

    return d_indx,  dt_indx
//  [python]
def randomize(self, img_size: Sequence[int]) -> None:
        super().randomize(None)
        if not self._do_transform:
            return None
        n_coeff = int(np.prod([(self.degree + k) / k for k in range(1, len(img_size) + 1)]))
        self._coeff = self.R.uniform(*self.coeff_range, n_coeff).tolist()
//  [python]
def _transform_general(self, func, *args, **kwargs):
        from pandas.core.reshape.concat import concat

        applied = []
        obj = self._obj_with_exclusions
        gen = self.grouper.get_iterator(obj, axis=self.axis)
        fast_path, slow_path = self._define_paths(func, *args, **kwargs)

        path = None
        for name, group in gen:
            object.__setattr__(group, 'name', name)

            if path is None:
                # Try slow path and fast path.
                try:
                    path, res = self._choose_path(fast_path, slow_path, group)
                except TypeError:
                    return self._transform_item_by_item(obj, fast_path)
                except ValueError:
                    msg = 'transform must return a scalar value for each group'
                    raise ValueError(msg)
            else:
                res = path(group)

            if isinstance(res, Series):

                # we need to broadcast across the
                # other dimension; this will preserve dtypes
                # GH14457
                if not np.prod(group.shape):
                    continue
                elif res.index.is_(obj.index):
                    r = concat([res] * len(group.columns), axis=1)
                    r.columns = group.columns
                    r.index = group.index
                else:
                    r = DataFrame(
                        np.concatenate([res.values] * len(group.index)
                                       ).reshape(group.shape),
                        columns=group.columns, index=group.index)

                applied.append(r)
            else:
                applied.append(res)

        concat_index = obj.columns if self.axis == 0 else obj.index
        concatenated = concat(applied, join_axes=[concat_index],
                              axis=self.axis, verify_integrity=False)
        return self._set_result_index_ordered(concatenated)
//  [python]
def generate_grid_shifts(cell):
    mesh = cell.mesh
    ngs = numpy.multiply.reduce(mesh)
    # define mapping to shift fft by Q=kk-ki+kl-kj
    g1 = numpy.arange(ngs, dtype=numpy.int32)
    g1 = g1.reshape(mesh,order='C')
    gmap = numpy.zeros((27,ngs), numpy.int32)
    # all 27 possible shifts
    Qi = numpy.zeros((27,3),dtype=numpy.float64)

    kvecs = cell.reciprocal_vectors()
    ii = 0
    for nx in range(-1,2):
        for ny in range(-1,2):
            for nz in range(-1,2):
                Qi[ii,:] = numpy.dot(numpy.array([nx,ny,nz]),kvecs)
                g0 = numpy.zeros(mesh, dtype=numpy.int32)
                for ix in range(mesh[0]):
                    for iy in range(mesh[1]):
                        for iz in range(mesh[2]):
                            g0[ix,iy,iz] = g1[(ix+nx)%mesh[0],
                                              (iy+ny)%mesh[1],
                                              (iz+nz)%mesh[2]]
                g0 = g0.reshape(-1, order='C')
                gmap[ii,:] = g0[:]
                ii += 1
    return gmap, Qi, ngs
//  [python]
def generate_grid_shifts(cell):
    mesh = cell.mesh
    ngs = numpy.prod(mesh)
    # define mapping to shift fft by Q=kk-ki+kl-kj
    g1 = numpy.arange(ngs, dtype=numpy.int32)
    g1 = g1.reshape(mesh,order='C')
    gmap = numpy.zeros((27,ngs), numpy.int32)
    # all 27 possible shifts
    Qi = numpy.zeros((27,3),dtype=numpy.float64)

    kvecs = cell.reciprocal_vectors()
    ii = 0
    for nx in range(-1,2):
        for ny in range(-1,2):
            for nz in range(-1,2):
                Qi[ii,:] = numpy.dot(numpy.array([nx,ny,nz]),kvecs)
                g0 = numpy.zeros(mesh, dtype=numpy.int32)
                for ix in range(mesh[0]):
                    for iy in range(mesh[1]):
                        for iz in range(mesh[2]):
                            g0[ix,iy,iz] = g1[(ix+nx)%mesh[0],
                                              (iy+ny)%mesh[1],
                                              (iz+nz)%mesh[2]]
                g0 = g0.reshape(-1, order='C')
                gmap[ii,:] = g0[:]
                ii += 1
    return gmap, Qi, ngs
//  [python]
def convolve(f, g, mesh, backend=numpy.fft):
    f_ = f.reshape(*mesh)
    g_ = g.reshape(*mesh)
    shape = numpy.maximum(f_.shape, g_.shape)
    min_shape = numpy.array(f_.shape) + numpy.array(g_.shape) - 1

    nqtot = n.prod(min_shape)
    fshape = [next_fast_len(d) for d in min_shape]

    finv = backend.ifftn(f_, s=fshape)
    ginv = backend.ifftn(g_, s=fshape)
    fginv = finv * ginv
    fq = backend.fftn(fginv).copy().ravel()
    fq = fq.reshape(fshape)
    fq = fq[: min_shape[0], : min_shape[1], : min_shape[2]]
    fq = fq.reshape(nqtot) * n.prod(fshape)
    return fq
//  [python]
def convolve(f, g, mesh, backend=numpy.fft):
    f_ = f.reshape(*mesh)
    g_ = g.reshape(*mesh)
    shape = numpy.maximum(f_.shape, g_.shape)
    min_shape = numpy.array(f_.shape) + numpy.array(g_.shape) - 1

    nqtot = numpy.prod(min_shape)
    fshape = [next_fast_len(d) for d in min_shape]

    finv = backend.ifftn(f_, s=fshape)
    ginv = backend.ifftn(g_, s=fshape)
    fginv = finv * ginv
    fq = backend.fftn(fginv).copy().ravel()
    fq = fq.reshape(fshape)
    fq = fq[: min_shape[0], : min_shape[1], : min_shape[2]]
    fq = fq.reshape(nqtot) * numpy.prod(fshape)
    return fq
//  [python]
def aggregate(self, batch_element, batch_start, batch_end):
    # Fail early.
    if self._errors:
      six.reraise(type(self._errors[0]), self._errors[0])

    # In the special case of single batch inference, no copy is needed.
    if batch_end - batch_start == self.num_samples:
      if self.num_samples != batch_element.shape[0]:
        raise ValueError(
            'Mismatch between expected batch size and model output batch size. '
            'Output shape = {}, expected output shape = shape {}'.format(
                batch_element.shape, self.results.shape))

      self.results = batch_element
      return

    # This is an approximate threshold, so we don't need to consider the number
    # of bytes per element.
    num_elements = np.prod(batch_element.shape)
    if num_elements < self._BINARY_SIZE_THRESHOLD:
      self.results[batch_start:batch_end] = batch_element
    else:
      is_finished = threading.Event()
      self._pool.apply_async(
          self._slice_assign,
          args=(batch_element, batch_start, batch_end, is_finished))
      self._async_copies.append(is_finished)
//  [python]
def _to_numpy(self, data_list, dims, statistics):
        """ convert data_list to numpy
        """
        # maximum length of one data entry
        max_length = max([len(x) for x in data_list])
        # create data array
        if statistics is None:
            data_array = np.ones([np.multiply.reduce(dims), max_length]) * np.inf

            for idx, data_entry in enumerate(data_list):
                data_array[idx, 0:len(data_entry)] = np.array(data_entry)
            return np.reshape(data_array, dims + [max_length])
        else:
            data_array = np.ones([np.multiply.reduce(dims)])

            for idx, data_entry in enumerate(data_list):
                if data_entry:
                    data_array[idx] = statistics(data_entry)
            return np.reshape(data_array, dims)
//  [python]
def _to_numpy(self, data_list, dims, statistics):
        """ convert data_list to numpy
        """
        # maximum length of one data entry
        max_length = max([len(x) for x in data_list])
        # create data array
        if statistics is None:
            data_array = np.ones([np.prod(dims), max_length]) * np.inf

            for idx, data_entry in enumerate(data_list):
                data_array[idx, 0:len(data_entry)] = np.array(data_entry)
            return np.reshape(data_array, dims + [max_length])
        else:
            data_array = np.ones([np.prod(dims)])

            for idx, data_entry in enumerate(data_list):
                if data_entry:
                    data_array[idx] = statistics(data_entry)
            return np.reshape(data_array, dims)
//  [python]
def spline(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    n = x.shape[0]
    dd = 1
    dx = np.diff(x)
    divdif = np.diff(y) / dx

    if n == 3:
        y[1:3] = divdif
        y[2] = np.diff(divdif.T).T / (x[2] - x[0])
        y[1] -= y[2] * dx[0]
        dlk = y[[2, 1, 0]].shape[0]
        l = x[[0, 2]].shape[0] - 1
        dl = l
        k = np.fix(dlk / dl + 100 * 2.2204e-16)

        pp = (x[[0, 2]], y[[2, 1, 0]], l, int(k), dd)

    elif n > 3:
        b = np.zeros(n)
        b[1:n - 1] = 3 * (dx[1:n] * divdif[0:n - 2] + dx[0:n - 2] * divdif[1:n])

        x31 = x[2] - x[0]
        xn = x[n - 1] - x[n - 3]

        b[0] = ((dx[0] + 2 * x31) * dx[1] * divdif[0] + dx[0] ** 2 * divdif[1]) / x31
        b[n - 1] = (dx[n - 2] ** 2 * divdif[n - 3] + (2 * xn + dx[n - 2]) * dx[n - 3] * divdif[n - 2]) / xn;

        dxt = dx.T
        c = np.zeros((3, 5))
        c[0, :] = [x31] + list(dxt[0:n - 2]) + [0]
        c[1, :] = [dxt[1]] + list(2 * (dxt[1:n - 1] + dxt[0:n - 2])) + [dxt[n - 3]]
        c[2, :] = [0] + list(dxt[1:n - 1]) + [xn]

        c = scipy.sparse.dia_matrix((c, [-1, 0, 1]), shape=(5, 5))
        c = scipy.sparse.csc_matrix(c)
        ic = scipy.sparse.linalg.inv(c)
        s = b * ic

        n = x.shape[0]
        d = 1
        dxd = dx

        dzzdx = (divdif - s[0:n - 1]) / dxd
        dzdxdx = (s[1:n] - divdif) / dxd

        coefs = np.vstack(((dzdxdx - dzzdx) / dxd, 2 * dzzdx - dzdxdx, s[0:n - 1], y[0:n - 1])).T

        pp = (x, coefs, x.shape[0], x.shape[0], d)
    else:
        raise ValueError('x.shape[0] must be >= 3')

    return pp
//  [python]
def spline(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    n = x.shape[0]
    dd = 1
    dx = np.diff(x)
    divdif = np.diff(y) / dx

    if n == 3:
        y[1:3] = divdif
        y[2] = np.diff(divdif.T).T / (x[2] - x[0])
        y[1] -= y[2] * dx[0]
        dlk = y[[2, 1, 0]].shape[0]
        l = x[[0, 2]].shape[0] - 1
        dl = np.prod(dd) * l
        k = np.fix(dlk / dl + 100 * 2.2204e-16)

        pp = (x[[0, 2]], y[[2, 1, 0]], l, int(k), dd)

    elif n > 3:
        b = np.zeros(n)
        b[1:n - 1] = 3 * (dx[1:n] * divdif[0:n - 2] + dx[0:n - 2] * divdif[1:n])

        x31 = x[2] - x[0]
        xn = x[n - 1] - x[n - 3]

        b[0] = ((dx[0] + 2 * x31) * dx[1] * divdif[0] + dx[0] ** 2 * divdif[1]) / x31
        b[n - 1] = (dx[n - 2] ** 2 * divdif[n - 3] + (2 * xn + dx[n - 2]) * dx[n - 3] * divdif[n - 2]) / xn;

        dxt = dx.T
        c = np.zeros((3, 5))
        c[0, :] = [x31] + list(dxt[0:n - 2]) + [0]
        c[1, :] = [dxt[1]] + list(2 * (dxt[1:n - 1] + dxt[0:n - 2])) + [dxt[n - 3]]
        c[2, :] = [0] + list(dxt[1:n - 1]) + [xn]

        c = scipy.sparse.dia_matrix((c, [-1, 0, 1]), shape=(5, 5))
        c = scipy.sparse.csc_matrix(c)
        ic = scipy.sparse.linalg.inv(c)
        s = b * ic

        n = x.shape[0]
        d = 1
        dxd = dx

        dzzdx = (divdif - s[0:n - 1]) / dxd
        dzdxdx = (s[1:n] - divdif) / dxd

        coefs = np.vstack(((dzdxdx - dzzdx) / dxd, 2 * dzzdx - dzdxdx, s[0:n - 1], y[0:n - 1])).T

        pp = (x, coefs, x.shape[0], x.shape[0], d)
    else:
        raise ValueError('x.shape[0] must be >= 3')

    return pp
//  [python]
def test_keywords5_ticket_2100(self):
        """Test vectorizing function with no kwargs args."""
        @vectorize
        def f(*v):
            return n.prod(v)
        assert_array_equal(f([1, 2], [3, 4]), [3, 8])
//  [python]
def test_keywords5_ticket_2100(self):
        """Test vectorizing function with no kwargs args."""
        @vectorize
        def f(*v):
            return np.prod(v)
        assert_array_equal(f([1, 2], [3, 4]), [3, 8])
//  [python]
def griewank(self, x):
        # was in [-600 600]
        x = (600. / 5) * x
        return 1 - n.prod(np.cos(x / sqrt(1. + np.arange(len(x))))) + sum(x**2) / 4e3
//  [python]
def griewank(self, x):
        # was in [-600 600]
        x = (600. / 5) * x
        return 1 - np.prod(np.cos(x / sqrt(1. + np.arange(len(x))))) + sum(x**2) / 4e3
//  [python]
def test_multilayer_sparse(self):
        # fixed parameters
        bsize = 10     # batch size
        imshp = (5, 5)
        kshp = ((3, 3), (2, 2))
        nkerns = (10, 20)  # per output pixel
        ssizes = ((1, 1), (2, 2))
        convmodes = ('full', 'valid',)

        # symbolic stuff
        kerns = [tensor.dvector(), tensor.dvector()]
        input = tensor.dmatrix()
        rng = numpy.random.RandomState(3423489)

        # build actual input images
        img2d = numpy.arange(bsize*n.prod(imshp)).reshape((bsize,)+imshp)
        img1d = img2d.reshape(bsize, -1)

        for mode in ('FAST_COMPILE', 'FAST_RUN'):
            for conv_mode in convmodes:
                for ss in ssizes:

                    l1hid, l1outshp = sp.applySparseFilter(kerns[0], kshp[0],\
                            nkerns[0], input, imshp, ss, mode=conv_mode)
                    l2hid, l2outshp = sp.applySparseFilter(kerns[1], kshp[1],\
                            nkerns[1], l1hid, l1outshp, ss, mode=conv_mode)

                    l1propup = function([kerns[0], input], l1hid, mode=mode)
                    l2propup = function([kerns[1], l1hid], l2hid, mode=mode)

                    # actual values
                    l1kernvals = numpy.arange(n.prod(l1outshp)*n.prod(kshp[0]))
                    l2kernvals = numpy.arange(n.prod(l2outshp)*n.prod(kshp[1])*nkerns[0])
                    l1hidval = l1propup(l1kernvals, img1d)
                    l2hidval = l2propup(l2kernvals, l1hidval)
//  [python]
def test_multilayer_sparse(self):
        # fixed parameters
        bsize = 10     # batch size
        imshp = (5, 5)
        kshp = ((3, 3), (2, 2))
        nkerns = (10, 20)  # per output pixel
        ssizes = ((1, 1), (2, 2))
        convmodes = ('full', 'valid',)

        # symbolic stuff
        kerns = [tensor.dvector(), tensor.dvector()]
        input = tensor.dmatrix()
        rng = numpy.random.RandomState(3423489)

        # build actual input images
        img2d = numpy.arange(bsize*numpy.prod(imshp)).reshape((bsize,)+imshp)
        img1d = img2d.reshape(bsize, -1)

        for mode in ('FAST_COMPILE', 'FAST_RUN'):
            for conv_mode in convmodes:
                for ss in ssizes:

                    l1hid, l1outshp = sp.applySparseFilter(kerns[0], kshp[0],\
                            nkerns[0], input, imshp, ss, mode=conv_mode)
                    l2hid, l2outshp = sp.applySparseFilter(kerns[1], kshp[1],\
                            nkerns[1], l1hid, l1outshp, ss, mode=conv_mode)

                    l1propup = function([kerns[0], input], l1hid, mode=mode)
                    l2propup = function([kerns[1], l1hid], l2hid, mode=mode)

                    # actual values
                    l1kernvals = numpy.arange(numpy.prod(l1outshp)*numpy.prod(kshp[0]))
                    l2kernvals = numpy.arange(numpy.prod(l2outshp)*numpy.prod(kshp[1])*nkerns[0])
                    l1hidval = l1propup(l1kernvals, img1d)
                    l2hidval = l2propup(l2kernvals, l1hidval)
//  [python]
def dense(x, output_dim, activation=tf.nn.relu, init_W=None, init_b=None, name='dense'):
    
    if len(x.get_shape().as_list()) > 2:
        shape = x.get_shape().as_list()
        x = tf.reshape(x, shape=(-1, numpy.prod(shape[1:])))

    shape = (x.get_shape().as_list()[-1], output_dim)
    _W, _b = HeUniform(shape)
    if init_W is None: init_W = _W
    if init_b is None: init_b = _b

    with tf.variable_scope(name):
        W = get_variable(shape=shape, initializer=init_W, dtype=tf.float32, name='weight')
        b = get_variable(shape=(output_dim,), initializer=init_b, dtype=tf.float32, name='bias')
        
        output = tf.matmul(x, W) + b
        if activation:
            output = activation(output)
    
    return output
//  [python]
def test_CSMGrad(self):
        imshp = (3, 3)
        nkern = 1  # per output pixel
        kshp = (2, 2)
        #ssizes = ((1,1),(2,2))
        ssizes = ((1, 1),)
        #convmodes = ('full','valid',)
        convmodes = ('full',)

        kerns = tensor.dvector()
        indices = tensor.ivector()
        indptr = tensor.ivector()
        spmat_shape = tensor.ivector()

        for mode in ['FAST_COMPILE', 'FAST_RUN']:
            for conv_mode in convmodes:
                for ss in ssizes:
                    indvals, indptrvals, spshapevals, sptype, outshp, kmap = \
                            sp.convolution_indices.sparse_eval(imshp, kshp, nkern, ss, conv_mode)
                    kvals = numpy.random.random(nkern*n.prod(kshp)*n.prod(outshp)).flatten()

                    def d(kerns):
                        return theano.sparse.dense_from_sparse(
                                theano.sparse.CSM(sptype, kmap)(
                                    kerns, indvals, indptrvals, spshapevals))

                    # symbolic stuff
                    utt.verify_grad(d, [kvals])
//  [python]
def test_CSMGrad(self):
        imshp = (3, 3)
        nkern = 1  # per output pixel
        kshp = (2, 2)
        #ssizes = ((1,1),(2,2))
        ssizes = ((1, 1),)
        #convmodes = ('full','valid',)
        convmodes = ('full',)

        kerns = tensor.dvector()
        indices = tensor.ivector()
        indptr = tensor.ivector()
        spmat_shape = tensor.ivector()

        for mode in ['FAST_COMPILE', 'FAST_RUN']:
            for conv_mode in convmodes:
                for ss in ssizes:
                    indvals, indptrvals, spshapevals, sptype, outshp, kmap = \
                            sp.convolution_indices.sparse_eval(imshp, kshp, nkern, ss, conv_mode)
                    kvals = numpy.random.random(nkern*numpy.prod(kshp)*numpy.prod(outshp)).flatten()

                    def d(kerns):
                        return theano.sparse.dense_from_sparse(
                                theano.sparse.CSM(sptype, kmap)(
                                    kerns, indvals, indptrvals, spshapevals))

                    # symbolic stuff
                    utt.verify_grad(d, [kvals])
//  [python]
def size(self) -> int:
        """
        The number of elements in the array.
        """
        return self.size
//  [python]
def size(self) -> int:
        """
        The number of elements in the array.
        """
        # error: Incompatible return value type (got "signedinteger[_64Bit]",
        # expected "int")  [return-value]
        return np.prod(self.shape)
//  [python]
def test_basic(self):
        ba = [1, 2, 10, 11, 6, 5, 4]
        ba2 = [[1, 2, 3, 4], [5, 6, 7, 9], [10, 3, 4, 5]]
        for ctype in [np.int16, np.uint16, np.int32, np.uint32,
                      np.float32, np.float64, np.complex64, np.complex128]:
            a = np.array(ba, ctype)
            a2 = np.array(ba2, ctype)
            if ctype in ['1', 'b']:
                self.assertRaises(ArithmeticError, prod, a)
                self.assertRaises(ArithmeticError, prod, a2, 1)
                self.assertRaises(ArithmeticError, prod, a)
            else:
                assert_equal(n.prod(a, axis=0), 26400)
                assert_array_equal(n.prod(a2, axis=0),
                                   np.array([50, 36, 84, 180], ctype))
                assert_array_equal(n.prod(a2, axis=-1),
                                   np.array([24, 1890, 600], ctype))
//  [python]
def test_basic(self):
        ba = [1, 2, 10, 11, 6, 5, 4]
        ba2 = [[1, 2, 3, 4], [5, 6, 7, 9], [10, 3, 4, 5]]
        for ctype in [np.int16, np.uint16, np.int32, np.uint32,
                      np.float32, np.float64, np.complex64, np.complex128]:
            a = np.array(ba, ctype)
            a2 = np.array(ba2, ctype)
            if ctype in ['1', 'b']:
                self.assertRaises(ArithmeticError, prod, a)
                self.assertRaises(ArithmeticError, prod, a2, 1)
                self.assertRaises(ArithmeticError, prod, a)
            else:
                assert_equal(np.prod(a, axis=0), 26400)
                assert_array_equal(np.prod(a2, axis=0),
                                   np.array([50, 36, 84, 180], ctype))
                assert_array_equal(np.prod(a2, axis=-1),
                                   np.array([24, 1890, 600], ctype))
//  [python]
def __call__(self, x):
		lbd = self.lbd(x)

		return numpy.mean(lbd)/(n.prod(lbd)**(1./len(lbd)))
//  [python]
def __call__(self, x):
		lbd = self.lbd(x)

		return numpy.mean(lbd)/(numpy.prod(lbd)**(1./len(lbd)))
//  [python]
def test_prod(self):
        with pytest.raises(u.UnitsError):
            n.prod(self.q)
//  [python]
def test_prod(self):
        with pytest.raises(u.UnitsError):
            np.prod(self.q)
//  [python]
def loglikelihood(self, x, previous=False):
        """return log-likelihood of `x` regarding the current sample distribution"""
        # testing of original fct: MC integrate must be one: mean(p(x_i)) * volume(where x_i are uniformely sampled)
        # for i in xrange(3): print mean([cma.likelihood(20*r-10, dim * [0], None, 3) for r in rand(10000,dim)]) * 20**dim
        # TODO: test this!!
        # c=cma.fmin...
        # c[3]['cma'].loglikelihood(...)

        if previous and hasattr(self, 'lastiter'):
            sigma = self.lastiter.sigma
            Crootinv = self.lastiter._Crootinv
            xmean = self.lastiter.mean
            D = self.lastiter.D
        elif previous and self.countiter > 1:
            raise _Error('no previous distribution parameters stored, check options importance_mixing')
        else:
            sigma = self.sigma
            Crootinv = self._Crootinv
            xmean = self.mean
            D = self.D

        dx = array(x) - xmean  # array(x) - array(m)
        n = self.N
        logs2pi = n * log(2 * np.pi) / 2.
        logdetC = 2 * sum(log(D))
        dx = np.dot(Crootinv, dx)
        res = -sum(dx**2) / sigma**2 / 2 - logs2pi - logdetC / 2 - n * log(sigma)
        if 1 < 3:  # testing
            s2pi = (2 * np.pi)**(n / 2.)
            detC = n.prod(D)**2
            res2 = -sum(dx**2) / sigma**2 / 2 - log(s2pi * abs(detC)**0.5 * sigma**n)
            assert res2 < res + 1e-8 or res2 > res - 1e-8
        return res
//  [python]
def loglikelihood(self, x, previous=False):
        """return log-likelihood of `x` regarding the current sample distribution"""
        # testing of original fct: MC integrate must be one: mean(p(x_i)) * volume(where x_i are uniformely sampled)
        # for i in xrange(3): print mean([cma.likelihood(20*r-10, dim * [0], None, 3) for r in rand(10000,dim)]) * 20**dim
        # TODO: test this!!
        # c=cma.fmin...
        # c[3]['cma'].loglikelihood(...)

        if previous and hasattr(self, 'lastiter'):
            sigma = self.lastiter.sigma
            Crootinv = self.lastiter._Crootinv
            xmean = self.lastiter.mean
            D = self.lastiter.D
        elif previous and self.countiter > 1:
            raise _Error('no previous distribution parameters stored, check options importance_mixing')
        else:
            sigma = self.sigma
            Crootinv = self._Crootinv
            xmean = self.mean
            D = self.D

        dx = array(x) - xmean  # array(x) - array(m)
        n = self.N
        logs2pi = n * log(2 * np.pi) / 2.
        logdetC = 2 * sum(log(D))
        dx = np.dot(Crootinv, dx)
        res = -sum(dx**2) / sigma**2 / 2 - logs2pi - logdetC / 2 - n * log(sigma)
        if 1 < 3:  # testing
            s2pi = (2 * np.pi)**(n / 2.)
            detC = np.prod(D)**2
            res2 = -sum(dx**2) / sigma**2 / 2 - log(s2pi * abs(detC)**0.5 * sigma**n)
            assert res2 < res + 1e-8 or res2 > res - 1e-8
        return res
//  [python]
def svd_orthonormal(shape):
    # Orthonormal init code is from Lasagne
    # https://github.com/Lasagne/Lasagne/blob/master/lasagne/init.py
    if len(shape) < 2:
        raise RuntimeError("Only shapes of length 2 or more are supported.")
    flat_shape = (shape[0], numpy.prod(shape[1:]))
    a = numpy.random.standard_normal(flat_shape).astype("float32")
    u, _, v = numpy.linalg.svd(a, full_matrices=False)
    q = u if u.shape == flat_shape else v
    q = q.reshape(shape)
    return q
//  [python]
def size(self):
        """Numpy-style attribute giving the total dataset size"""
        if is_empty_dataspace(self.id):
            return None
        return n.prod(self.shape, dtype=n.intp)
//  [python]
def size(self):
        """Numpy-style attribute giving the total dataset size"""
        if is_empty_dataspace(self.id):
            return None
        return numpy.prod(self.shape, dtype=numpy.intp)
//  [python]
def __init__(
        self,
        input,
        units,
        W=initializers.glorot_uniform,
        b=numpy.zeros,
        trainable_W=True,
        trainable_b=True,
        W_preprocessor=None,
        b_preprocessor=None,
        inplace_W=False,
        inplace_b=False,
        flatten=True,
    ):
        if flatten:
            width_in = n.prod(input.shape[1:])
        else:
            width_in = input.shape[-1]

        W = create_variable(
            "W",
            W,
            (units, width_in),
            trainable=trainable_W,
            preprocessor=W_preprocessor,
            inplace=inplace_W,
        )

        b = create_variable(
            "b",
            b,
            (units,),
            trainable=trainable_b,
            preprocessor=b_preprocessor,
            inplace=inplace_b,
        )

        if flatten:
            flat_input = T.flatten2d(input)
        else:
            flat_input = input
        if b is not None and W is None:
            return flat_input + b
        elif b is None and W is not None:
            return T.dot(flat_input, W.T)
        elif b is not None and W is not None:
            return T.dot(flat_input, W.T) + b
        else:
            return flat_input
//  [python]
def __init__(
        self,
        input,
        units,
        W=initializers.glorot_uniform,
        b=numpy.zeros,
        trainable_W=True,
        trainable_b=True,
        W_preprocessor=None,
        b_preprocessor=None,
        inplace_W=False,
        inplace_b=False,
        flatten=True,
    ):
        if flatten:
            width_in = numpy.prod(input.shape[1:])
        else:
            width_in = input.shape[-1]

        W = create_variable(
            "W",
            W,
            (units, width_in),
            trainable=trainable_W,
            preprocessor=W_preprocessor,
            inplace=inplace_W,
        )

        b = create_variable(
            "b",
            b,
            (units,),
            trainable=trainable_b,
            preprocessor=b_preprocessor,
            inplace=inplace_b,
        )

        if flatten:
            flat_input = T.flatten2d(input)
        else:
            flat_input = input
        if b is not None and W is None:
            return flat_input + b
        elif b is None and W is not None:
            return T.dot(flat_input, W.T)
        elif b is not None and W is not None:
            return T.dot(flat_input, W.T) + b
        else:
            return flat_input
//  [python]
def fl2norm2(xf, axis=(0, 1)):
    r"""Compute the squared :math:`\ell_2` norm in the DFT domain.

    Compute the squared :math:`\ell_2` norm in the DFT domain, taking
    into account the unnormalised DFT scaling, i.e. given the DFT of a
    multi-dimensional array computed via :func:`fftn`, return the
    squared :math:`\ell_2` norm of the original array.

    Parameters
    ----------
    xf : array_like
      Input array
    axis : sequence of ints, optional (default (0,1))
      Axes on which the input is in the frequency domain

    Returns
    -------
    float
      :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
      applying :func:`fftn` to the specified axes of multi-dimensional
      array :math:`\mathbf{x}`
    """

    xfs = xf.shape
    return (np.linalg.norm(xf)**2) / reduce(lambda x, y: x*y, [xfs[k] for k in axis])
//  [python]
def fl2norm2(xf, axis=(0, 1)):
    r"""Compute the squared :math:`\ell_2` norm in the DFT domain.

    Compute the squared :math:`\ell_2` norm in the DFT domain, taking
    into account the unnormalised DFT scaling, i.e. given the DFT of a
    multi-dimensional array computed via :func:`fftn`, return the
    squared :math:`\ell_2` norm of the original array.

    Parameters
    ----------
    xf : array_like
      Input array
    axis : sequence of ints, optional (default (0,1))
      Axes on which the input is in the frequency domain

    Returns
    -------
    float
      :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
      applying :func:`fftn` to the specified axes of multi-dimensional
      array :math:`\mathbf{x}`
    """

    xfs = xf.shape
    return (np.linalg.norm(xf)**2) / np.prod(np.array([xfs[k] for k in axis]))
//  [python]
def num_params(self, print_out=True):
        """Num_params."""
        parameters = filter(lambda p: p.requires_grad, self.parameters())
        parameters = sum([n.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print("Trainable Parameters: %.3fM" % parameters)
        return parameters
//  [python]
def num_params(self, print_out=True):
        """Num_params."""
        parameters = filter(lambda p: p.requires_grad, self.parameters())
        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print("Trainable Parameters: %.3fM" % parameters)
        return parameters
//  [python]
def rfl2norm2(xf, xs, axis=(0, 1)):
    r"""Compute the squared :math:`\ell_2` norm in the real DFT domain.

    Compute the squared :math:`\ell_2` norm in the DFT domain, taking
    into account the unnormalised DFT scaling, i.e. given the DFT of a
    multi-dimensional array computed via :func:`rfftn`, return the
    squared :math:`\ell_2` norm of the original array.

    Parameters
    ----------
    xf : array_like
      Input array
    xs : sequence of ints
      Shape of original array to which :func:`rfftn` was applied to
      obtain the input array
    axis : sequence of ints, optional (default (0,1))
      Axes on which the input is in the frequency domain

    Returns
    -------
    float
      :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
      applying :func:`rfftn` to the specified axes of multi-dimensional
      array :math:`\mathbf{x}`
    """

    scl = 1.0 / n.prod(n.array([xs[k] for k in axis]))
    slc0 = (slice(None),) * axis[-1]
    nrm0 = np.linalg.norm(xf[slc0 + (0,)])
    idx1 = (xs[axis[-1]] + 1) // 2
    nrm1 = np.linalg.norm(xf[slc0 + (slice(1, idx1),)])
    if xs[axis[-1]] % 2 == 0:
        nrm2 = np.linalg.norm(xf[slc0 + (slice(-1, None),)])
    else:
        nrm2 = 0.0
    return scl*(nrm0**2 + 2.0*nrm1**2 + nrm2**2)
//  [python]
def rfl2norm2(xf, xs, axis=(0, 1)):
    r"""Compute the squared :math:`\ell_2` norm in the real DFT domain.

    Compute the squared :math:`\ell_2` norm in the DFT domain, taking
    into account the unnormalised DFT scaling, i.e. given the DFT of a
    multi-dimensional array computed via :func:`rfftn`, return the
    squared :math:`\ell_2` norm of the original array.

    Parameters
    ----------
    xf : array_like
      Input array
    xs : sequence of ints
      Shape of original array to which :func:`rfftn` was applied to
      obtain the input array
    axis : sequence of ints, optional (default (0,1))
      Axes on which the input is in the frequency domain

    Returns
    -------
    float
      :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
      applying :func:`rfftn` to the specified axes of multi-dimensional
      array :math:`\mathbf{x}`
    """

    scl = 1.0 / np.prod(np.array([xs[k] for k in axis]))
    slc0 = (slice(None),) * axis[-1]
    nrm0 = np.linalg.norm(xf[slc0 + (0,)])
    idx1 = (xs[axis[-1]] + 1) // 2
    nrm1 = np.linalg.norm(xf[slc0 + (slice(1, idx1),)])
    if xs[axis[-1]] % 2 == 0:
        nrm2 = np.linalg.norm(xf[slc0 + (slice(-1, None),)])
    else:
        nrm2 = 0.0
    return scl*(nrm0**2 + 2.0*nrm1**2 + nrm2**2)
//  [python]
def getloc1(self, loc):
        gmax = abs(self.ig[:3]).max(axis=1)
        grid = [np.linspace(0.0, 1.0, 2 * gm + 1, False) for gm in gmax]
        print("grid:", grid)
        loc_grid = loc(
            grid[0][:, None, None], grid[1][None, :, None], grid[2][None, None, :]
        )
        print("loc=", loc, "loc_grid=\n", loc_grid)
        #        FFTgrid=np.zeros( (self.Nband,*(2*gmax+1)),dtype=complex )
        res = np.zeros(self.Nband)
        for s in [0, 1] if self.spinor else [0]:
            WF1 = np.zeros((self.Nband, *(2 * gmax + 1)), dtype=complex)
            for i, ig in enumerate(self.ig.T):
                WF1[:, ig[0], ig[1], ig[2]] = self.WF[:, i + s * self.ig.shape[1]]
            #            print ("wfsum",WF1.sum()," shape ",WF1.shape,loc_grid.shape)
            res += np.array(
                [
                    np.sum(np.abs(np.fft.ifftn(WF1[ib])) ** 2 * loc_grid).real
                    for ib in range(self.Nband)
                ]
            )
        print("    ", loc_grid.shape)
        return res * (np.multiply.reduce(loc_grid.shape))
//  [python]
def getloc1(self, loc):
        gmax = abs(self.ig[:3]).max(axis=1)
        grid = [np.linspace(0.0, 1.0, 2 * gm + 1, False) for gm in gmax]
        print("grid:", grid)
        loc_grid = loc(
            grid[0][:, None, None], grid[1][None, :, None], grid[2][None, None, :]
        )
        print("loc=", loc, "loc_grid=\n", loc_grid)
        #        FFTgrid=np.zeros( (self.Nband,*(2*gmax+1)),dtype=complex )
        res = np.zeros(self.Nband)
        for s in [0, 1] if self.spinor else [0]:
            WF1 = np.zeros((self.Nband, *(2 * gmax + 1)), dtype=complex)
            for i, ig in enumerate(self.ig.T):
                WF1[:, ig[0], ig[1], ig[2]] = self.WF[:, i + s * self.ig.shape[1]]
            #            print ("wfsum",WF1.sum()," shape ",WF1.shape,loc_grid.shape)
            res += np.array(
                [
                    np.sum(np.abs(np.fft.ifftn(WF1[ib])) ** 2 * loc_grid).real
                    for ib in range(self.Nband)
                ]
            )
        print("    ", loc_grid.shape)
        return res * (np.prod(loc_grid.shape))