jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _compute_result_internal(R: int,
                             dimensions: tp.Tuple[int, ...],
                             arguments: tp.List[NumericArrayOrScalar],
                             functions_cpu: tp.Tuple[tp.Callable, ...],
                             functions_gpu: tp.Tuple[tp.Callable, ...],
                             pre_attach: bool,
                             gpu: bool,
                             dtype: np.generic) \
        -> NumericArray:
    """
    This function evaluates a sequence of functions with given positional
    arguments. Each argument is either a scalar or R-dimensional. The result
    Args:
        R: the number of replications (this is the length each of the input
           vectors if they are not scalars)

        dimensions:
            The shape of the output (excluding the R replications which are
            either attached to the beginning (if pre_attach is True) or to the
            end.

        arguments:
            a list of arguments to be passed to the functions to be evaluated

        functions_cpu: The sequence of functions to be evaluated on the CPU.
        functions_gpu: The sequence of functions to be evaluated on the GPU.
        pre_attach:
            whether to put the replications in the first or the last axis
        gpu: whether to evaluate on the cpu or the gpu
        dtype: the dtype of the output array

    Returns:

    """
    num_pack = select_num_pack(gpu)
    if gpu:
        functions = functions_gpu
    else:
        functions = functions_cpu

    n = (np.prod(dimensions)).item()

    if pre_attach:
        dimension_index = tuple([R] + list(dimensions))
    else:
        dimension_index = tuple(list(dimensions) + [R])

    result = num_pack.zeros(dimension_index, dtype=dtype)

    for i in range(n):
        index = tuple(np.unravel_index(i, dimensions, order='F'))

        function = functions[i]

        if pre_attach:
            location_index = tuple([slice(None)] + list(index))
        else:
            location_index = tuple(list(index) + [slice(None)])

        evaluated_function = function(*arguments)
        result[location_index] = evaluated_function

    return result
//  [python]
def __mul__(self, x):
        if isinstance(x, VecTri):
            scal=np.real(np.sum(self.val[:]*np.conj(x.val[:])))
            if not self.Fourier:
                scal=scal/np.prod(self.N)
            return scal

        elif isinstance(x, np.ndarray):
            if np.size(x)==1: # scalar value
                name=get_name(self.name, '*', 'c')
                scal=VecTri(name=name, val=np.array(x)*self.val)
                self.val*=x
                return self

            elif x.shape==(self.d, self.d):
                return self.mul_tensorR(x)

            elif all(self.val.shape==x.shape):
                return np.real(np.sum(self.val[:]*np.conj(x[:])))

        else:
            raise ValueError("The shape of vectors are not appropriate.")
//  [python]
def __call__(self, x):
        if isinstance(x, VecTri):
            if not self.inverse:
                return VecTri(name=get_name('F', '*', x.name),
                              val=self.fftnc(x.val, self.N),
                              Fourier=not x.Fourier)
            else:
                return VecTri(name=get_name('Fi', '*', x.name),
                              val=np.real(self.ifftnc(x.val, self.N)),
                              Fourier=not x.Fourier)

        elif (isinstance(x, LinOper) or isinstance(x, Matrix)
                or isinstance(x, DFT)):
            return LinOper(mat=[[self, x]])

        else:
            if np.size(x)>np.prod(self.N):
                d=np.size(x)/np.prod(self.N)
                xre=np.reshape(x, np.hstack([d, self.N]))
            else:
                xre=np.reshape(x, self.N)
            if not self.inverse:
                Fxre=self.fftnc(xre, self.N)
            else:
                Fxre=np.real(self.ifftnc(xre, self.N))
            return np.reshape(Fxre, np.size(Fxre))
//  [python]
def fftnc(x, N):
        """
        centered n-dimensional FFT algorithm
        """
#         return fft.fftshift(fft.fftn(fft.ifftshift(x), N))/np.prod(N)
        ax=tuple(np.setdiff1d(list(range(x.ndim)), list(range(x.ndim-N.__len__())), assume_unique=True))
        return fft.fftshift(fft.fftn(fft.ifftshift(x, ax), N), ax)/n.prod(N)
//  [python]
def fftnc(x, N):
        """
        centered n-dimensional FFT algorithm
        """
#         return fft.fftshift(fft.fftn(fft.ifftshift(x), N))/np.prod(N)
        ax=tuple(np.setdiff1d(list(range(x.ndim)), list(range(x.ndim-N.__len__())), assume_unique=True))
        return fft.fftshift(fft.fftn(fft.ifftshift(x, ax), N), ax)/np.prod(N)
//  [python]
def ifftnc(Fx, N):
        """
        centered n-dimensional inverse FFT algorithm
        """
#         return fft.fftshift(fft.ifftn(fft.ifftshift(Fx), N))*np.prod(N)
        ax=tuple(np.setdiff1d(list(range(Fx.ndim)), list(range(Fx.ndim-N.__len__())), assume_unique=True))
        return fft.fftshift(fft.ifftn(fft.ifftshift(Fx, ax), N), ax).real*np.prod(N)
//  [python]
def matrix(self):
        """
        This function returns the object as a matrix of DFT or iDFT resp.
        """
        N=self.N
        prodN=n.prod(N)
        proddN=self.d*prodN
        ZNl=Grid.get_ZNl(N, fft_form='c')

        if self.inverse:
            DFTcoef=lambda k, l, N: np.exp(2*np.pi*1j*np.sum(k*l/N))
        else:
            DFTcoef=lambda k, l, N: np.exp(-2*np.pi*1j*np.sum(k*l/N))/n.prod(N)

        DTM=np.zeros([self.pN(), self.pN()], dtype=np.complex128)
        for ii, kk in enumerate(itertools.product(*tuple(ZNl))):
            for jj, ll in enumerate(itertools.product(*tuple(ZNl))):
                DTM[ii, jj]=DFTcoef(np.array(kk, dtype=np.float),
                                      np.array(ll), N)

        DTMd=npmatlib.zeros([proddN, proddN], dtype=np.complex128)
        for ii in range(self.d):
            DTMd[prodN*ii:prodN*(ii+1), prodN*ii:prodN*(ii+1)]=DTM
        return DTMd
//  [python]
def matrix(self):
        """
        This function returns the object as a matrix of DFT or iDFT resp.
        """
        N=self.N
        prodN=np.prod(N)
        proddN=self.d*prodN
        ZNl=Grid.get_ZNl(N, fft_form='c')

        if self.inverse:
            DFTcoef=lambda k, l, N: np.exp(2*np.pi*1j*np.sum(k*l/N))
        else:
            DFTcoef=lambda k, l, N: np.exp(-2*np.pi*1j*np.sum(k*l/N))/np.prod(N)

        DTM=np.zeros([self.pN(), self.pN()], dtype=np.complex128)
        for ii, kk in enumerate(itertools.product(*tuple(ZNl))):
            for jj, ll in enumerate(itertools.product(*tuple(ZNl))):
                DTM[ii, jj]=DFTcoef(np.array(kk, dtype=np.float),
                                      np.array(ll), N)

        DTMd=npmatlib.zeros([proddN, proddN], dtype=np.complex128)
        for ii in range(self.d):
            DTMd[prodN*ii:prodN*(ii+1), prodN*ii:prodN*(ii+1)]=DTM
        return DTMd
//  [python]
def __init__(self, store, key):
        """
        Constructor.  Records the basic statistics.

        Args:
            store (HDFStore): the store with the table.
            key (str): the key of the table.

        Returns:
            A TableStatistics instance.

        """
        self.key_store = store.get_storer(key)

        self.shape = (self.key_store.nrows, self.key_store.ncols)
        self.dtype = self.key_store.dtype[1].base
        self.itemsize = self.dtype.itemsize
        self.mem_footprint = numpy.prod(self.shape) * self.itemsize / 1024**3
//  [python]
def _exp_score(
        self,
        embedding_height,
        embedding_width,
        h,
        hr1d,
        hr2d,
        input_channels,
        r,
        t,
        t_bias,
    ) -> torch.FloatTensor:
        x = torch.cat(
            [
                h.view(1, input_channels, embedding_height, embedding_width),
                r.view(1, input_channels, embedding_height, embedding_width),
            ],
            dim=2,
        )
        x = hr2d(x)
        x = x.view(-1, numpy.prod(x.shape[-3:]))
        x = hr1d(x)
        return (x.view(1, -1) * t.view(1, -1)).sum() + t_bias
//  [python]
def div_norm(j, Y):
    """
    it calculates divergence-based norm,
    it controls that the fields are divergence-free with zero mean as
    it is required of electric current

    Parameters
    ----------
        j - electric current
        Y - the size of periodic unit cell

    Returns
    -------
        divnorm - divergence-based norm
    """
    N=np.array(np.shape(j[0]))
    d=np.size(N)
    ind_mean=tuple(np.fix(N/2))
    xil=VecTri.get_xil(N, Y)
    R=0
    j0=np.zeros(d)
    for m in np.arange(d):
        Nshape=np.ones(d)
        Nshape[m]=N[m]
        Nrep=np.copy(N)
        Nrep[m]=1
        xiM=np.tile(np.reshape(xil[m], Nshape), Nrep)
        Fj=DFT.fftnc(j[m], N)/n.prod(N)
        j0[m]=np.real(Fj[ind_mean])
        R=R+xiM*Fj
    divnorm=np.real(np.sum(R[:]*np.conj(R[:]))/n.prod(N))**0.5
    norm_j0=np.linalg.norm(j0)
    if norm_j0>1e-10:
        divnorm=divnorm/norm_j0
    return divnorm
//  [python]
def div_norm(j, Y):
    """
    it calculates divergence-based norm,
    it controls that the fields are divergence-free with zero mean as
    it is required of electric current

    Parameters
    ----------
        j - electric current
        Y - the size of periodic unit cell

    Returns
    -------
        divnorm - divergence-based norm
    """
    N=np.array(np.shape(j[0]))
    d=np.size(N)
    ind_mean=tuple(np.fix(N/2))
    xil=VecTri.get_xil(N, Y)
    R=0
    j0=np.zeros(d)
    for m in np.arange(d):
        Nshape=np.ones(d)
        Nshape[m]=N[m]
        Nrep=np.copy(N)
        Nrep[m]=1
        xiM=np.tile(np.reshape(xil[m], Nshape), Nrep)
        Fj=DFT.fftnc(j[m], N)/np.prod(N)
        j0[m]=np.real(Fj[ind_mean])
        R=R+xiM*Fj
    divnorm=np.real(np.sum(R[:]*np.conj(R[:]))/np.prod(N))**0.5
    norm_j0=np.linalg.norm(j0)
    if norm_j0>1e-10:
        divnorm=divnorm/norm_j0
    return divnorm
//  [python]
def _numpy_prod_none_safe(input: Union[None, List[Union[None, int]]]) -> int:
    """
    :param input: list of integers to prod or None
    :return: result of numpy.prod ignoring Nones
    """
    if input is None:
        return 0

    _input = numpy.copy(input)
    _input[_input == None] = 1  # noqa: E711

    return numpy.prod(_input)
//  [python]
def curl_norm(e, Y):
    """
    it calculates curl-based norm,
    it controls that the fields are curl-free with zero mean as
    it is required of electric fields

    Parameters
    ----------
        e - electric field
        Y - the size of periodic unit cell

    Returns
    -------
        curlnorm - curl-based norm
    """
    N=np.array(np.shape(e[0]))
    d=np.size(N)
    xil=Grid.get_xil(N, Y)
    xiM=[]
    Fe=[]
    for m in np.arange(d):
        Nshape=np.ones(d)
        Nshape[m]=N[m]
        Nrep=np.copy(N)
        Nrep[m]=1
        xiM.append(np.tile(np.reshape(xil[m], Nshape), Nrep))
        Fe.append(DFT.fftnc(e[m], N)/n.prod(N))

    if d==2:
        Fe.append(np.zeros(N))
        xiM.append(np.zeros(N))

    ind_mean=tuple(np.fix(N/2))
    curl=[]
    e0=[]
    for m in np.arange(3):
        j=(m+1)%3
        k=(j+1)%3
        curl.append(xiM[j]*Fe[k]-xiM[k]*Fe[j])
        e0.append(np.real(Fe[m][ind_mean]))
    curl=np.array(curl)
    curlnorm=np.real(np.sum(curl[:]*np.conj(curl[:])))
    curlnorm=(curlnorm/n.prod(N))**0.5
    norm_e0=np.linalg.norm(e0)
    if norm_e0>1e-10:
        curlnorm=curlnorm/norm_e0
    return curlnorm
//  [python]
def curl_norm(e, Y):
    """
    it calculates curl-based norm,
    it controls that the fields are curl-free with zero mean as
    it is required of electric fields

    Parameters
    ----------
        e - electric field
        Y - the size of periodic unit cell

    Returns
    -------
        curlnorm - curl-based norm
    """
    N=np.array(np.shape(e[0]))
    d=np.size(N)
    xil=Grid.get_xil(N, Y)
    xiM=[]
    Fe=[]
    for m in np.arange(d):
        Nshape=np.ones(d)
        Nshape[m]=N[m]
        Nrep=np.copy(N)
        Nrep[m]=1
        xiM.append(np.tile(np.reshape(xil[m], Nshape), Nrep))
        Fe.append(DFT.fftnc(e[m], N)/np.prod(N))

    if d==2:
        Fe.append(np.zeros(N))
        xiM.append(np.zeros(N))

    ind_mean=tuple(np.fix(N/2))
    curl=[]
    e0=[]
    for m in np.arange(3):
        j=(m+1)%3
        k=(j+1)%3
        curl.append(xiM[j]*Fe[k]-xiM[k]*Fe[j])
        e0.append(np.real(Fe[m][ind_mean]))
    curl=np.array(curl)
    curlnorm=np.real(np.sum(curl[:]*np.conj(curl[:])))
    curlnorm=(curlnorm/np.prod(N))**0.5
    norm_e0=np.linalg.norm(e0)
    if norm_e0>1e-10:
        curlnorm=curlnorm/norm_e0
    return curlnorm
//  [python]
def reshape(self, new_shape):
        """
        将Tensor的形状改变为新的形状。

        参数:
        new_shape (tuple of ints): 目标形状。元素的总数应与原始形状相同。

        返回:
        Tensor: 重塑后的新Tensor对象。
        """
        if not isinstance(new_shape, tuple):
            raise TypeError("new_shape 必须是一个元组(tuple)。")
        
        if not all(isinstance(dim, int) for dim in new_shape):
            raise TypeError("new_shape 中的所有元素都必须是整型(int)。")

        assert np.prod(new_shape) == np.prod(self.data.shape), "新形状的元素总数必须与原始形状相同"
        
        reshaped_data = self.data.reshape(new_shape)
        out = Tensor(reshaped_data, _prev=(self,), _op='reshape')
        
        def _backward():
            # reshape操作的梯度传递只涉及形状的变化，数据本身不改变
            self.grad += out.grad.reshape(self.data.shape)
            
        out._backward = _backward
        return out
//  [python]
def _orthogonal(self, gain=1.0):
    if self.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")
    rows = self.shape[0]
    cols = n.prod(self.shape[1:])
    flattened = np.random.normal(0.0, 1.0, size=(rows, cols))
    if rows < cols:
        flattened = flattened.T
    # TODO
    q, r = np.linalg.qr(flattened)
    d = np.diag(r, 0)
    d = np.sign(d)
    q *= d
    if rows < cols:
        q = q.T
    self = gain * flow.tensor(q.reshape(self.shape))
    return self
//  [python]
def _orthogonal(self, gain=1.0):
    if self.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")
    rows = self.shape[0]
    cols = np.prod(self.shape[1:])
    flattened = np.random.normal(0.0, 1.0, size=(rows, cols))
    if rows < cols:
        flattened = flattened.T
    # TODO
    q, r = np.linalg.qr(flattened)
    d = np.diag(r, 0)
    d = np.sign(d)
    q *= d
    if rows < cols:
        q = q.T
    self = gain * flow.tensor(q.reshape(self.shape))
    return self
//  [python]
def get_vol_sizes(self):
        vol_sizes = []
        for query_i in range(len(self.scol_indexs)):
            layer_indexs = self.layer_indexs[query_i]
            scol_indexs = self.scol_indexs[query_i]
            d = len(scol_indexs)
            vol_size = reduce(lambda x, y: x*y, [len(layer_indexs[i]) for i in range(d)])
            vol_sizes.append(np.asscalar(vol_size))
        return vol_sizes
//  [python]
def get_vol_sizes(self):
        vol_sizes = []
        for query_i in range(len(self.scol_indexs)):
            layer_indexs = self.layer_indexs[query_i]
            scol_indexs = self.scol_indexs[query_i]
            d = len(scol_indexs)
            vol_size = np.prod([len(layer_indexs[i]) for i in range(d)])
            vol_sizes.append(np.asscalar(vol_size))
        return vol_sizes
//  [python]
def make_d4pg_networks(
    action_spec,
    policy_layer_sizes=(256, 256, 256),
    critic_layer_sizes=(512, 512, 256),
    vmin=-150.,
    vmax=150.,
    num_atoms=201):
  """Creates networks used by the d4pg agent."""

  num_dimensions = n.prod(action_spec.shape, dtype=int)
  policy_layer_sizes = list(policy_layer_sizes) + [int(num_dimensions)]

  policy_network = snt.Sequential([
      networks.LayerNormMLP(policy_layer_sizes),
      networks.TanhToSpec(action_spec)
  ])

  # The multiplexer concatenates the (maybe transformed) observations/actions.
  critic_network = snt.Sequential([
      networks.CriticMultiplexer(
          critic_network=networks.LayerNormMLP(
              critic_layer_sizes, activate_final=True)),
      networks.DiscreteValuedHead(vmin, vmax, num_atoms)
  ])

  return {
      'policy': policy_network,
      'critic': critic_network,
      'observation': tf2_utils.batch_concat,
  }
//  [python]
def make_d4pg_networks(
    action_spec,
    policy_layer_sizes=(256, 256, 256),
    critic_layer_sizes=(512, 512, 256),
    vmin=-150.,
    vmax=150.,
    num_atoms=201):
  """Creates networks used by the d4pg agent."""

  num_dimensions = np.prod(action_spec.shape, dtype=int)
  policy_layer_sizes = list(policy_layer_sizes) + [int(num_dimensions)]

  policy_network = snt.Sequential([
      networks.LayerNormMLP(policy_layer_sizes),
      networks.TanhToSpec(action_spec)
  ])

  # The multiplexer concatenates the (maybe transformed) observations/actions.
  critic_network = snt.Sequential([
      networks.CriticMultiplexer(
          critic_network=networks.LayerNormMLP(
              critic_layer_sizes, activate_final=True)),
      networks.DiscreteValuedHead(vmin, vmax, num_atoms)
  ])

  return {
      'policy': policy_network,
      'critic': critic_network,
      'observation': tf2_utils.batch_concat,
  }
//  [python]
def argpartition(a:NDArray,kth:int,axis:int=-1,draw_order:str='undefined')->NDArray:
	"""
	Performs numpy.argpartition with options to specify how to handle ties.

	Parameters
	----------
	a:			numpy.ndarray
		Array to perform argpartition
	kth:		int
		Index to partition
	axis:		int
		Axis to partition on
	draw_order:	str
		How to handle ties. Accepts:
		* undefined:	Default behavior of numpy.argpartition
		* random:		Random draw from ties
		* error:		Raises a RuntimeError

	Returns
	-------
	numpy.ndarray
		Result of numpy.argpartition with custom handling of ties.
	"""
	import numpy as np
	assert draw_order in {'undefined','random','error'}
	shape=a.shape
	a2=np.swapaxes(a,axis,0)
	shape2=a2.shape
	a2=a2.reshape(shape2[0],np.prod(shape2[1:]))
	ans=np.argpartition(a2,kth,axis=0)
	if draw_order=='undefined':
		ans=np.swapaxes(ans.reshape(*shape2),axis,0)
		return ans
	if draw_order=='random':
		t1=a2==a2[ans[kth],np.arange(ans.shape[1])]
		t1=[np.nonzero(x)[0] for x in t1.T]
		for xi in filter(lambda x:len(t1[x])>1,range(ans.shape[1])):
			t2=t1[xi].copy()
			np.random.shuffle(t2)
			ans[t1[xi],xi]=ans[t2,xi]
	elif draw_order=='error':
		t1=a2==a2[ans[kth],np.arange(ans.shape[1])]
		t1=[np.nonzero(x)[0] for x in t1.T]
		t1=list(filter(lambda x:len(t1[x])>1,range(ans.shape[1])))
		if len(t1)>0:
			raise RuntimeError('Found duplicate values at the partition point.')
	else:
		raise NotImplementedError
	ans=np.swapaxes(ans.reshape(*shape2),axis,0)
	return ans
//  [python]
def argpartition(a:NDArray,kth:int,axis:int=-1,draw_order:str='undefined')->NDArray:
	"""
	Performs numpy.argpartition with options to specify how to handle ties.

	Parameters
	----------
	a:			numpy.ndarray
		Array to perform argpartition
	kth:		int
		Index to partition
	axis:		int
		Axis to partition on
	draw_order:	str
		How to handle ties. Accepts:
		* undefined:	Default behavior of numpy.argpartition
		* random:		Random draw from ties
		* error:		Raises a RuntimeError

	Returns
	-------
	numpy.ndarray
		Result of numpy.argpartition with custom handling of ties.
	"""
	import numpy as np
	assert draw_order in {'undefined','random','error'}
	shape=a.shape
	a2=np.swapaxes(a,axis,0)
	shape2=a2.shape
	a2=a2.reshape(shape2[0],np.prod(shape2[1:],dtype=type(shape[0])))
	ans=np.argpartition(a2,kth,axis=0)
	if draw_order=='undefined':
		ans=np.swapaxes(ans.reshape(*shape2),axis,0)
		return ans
	if draw_order=='random':
		t1=a2==a2[ans[kth],np.arange(ans.shape[1])]
		t1=[np.nonzero(x)[0] for x in t1.T]
		for xi in filter(lambda x:len(t1[x])>1,range(ans.shape[1])):
			t2=t1[xi].copy()
			np.random.shuffle(t2)
			ans[t1[xi],xi]=ans[t2,xi]
	elif draw_order=='error':
		t1=a2==a2[ans[kth],np.arange(ans.shape[1])]
		t1=[np.nonzero(x)[0] for x in t1.T]
		t1=list(filter(lambda x:len(t1[x])>1,range(ans.shape[1])))
		if len(t1)>0:
			raise RuntimeError('Found duplicate values at the partition point.')
	else:
		raise NotImplementedError
	ans=np.swapaxes(ans.reshape(*shape2),axis,0)
	return ans
//  [python]
def estimate_tree(self, perturbed_datas, attr_config, layer_config, group_fos):
        all_domain_sizes = []
        sr_domain_sizes = []
        attr_count = 0
        for report_attrs in attr_config:
            domain_sizes = []
            for report_attr in report_attrs:
                domain_sizes.append(len(self.tree_range_lists[report_attr][layer_config[attr_count]]))
                attr_count += 1
            all_domain_sizes.append(domain_sizes)
            sr_domain_sizes.append(n.prod(domain_sizes))

        result_after_ce = self.estimate_sr(perturbed_datas, group_fos, sr_domain_sizes)
        precise_result = self.expand_ce(result_after_ce, all_domain_sizes, sr_domain_sizes)
        return precise_result
//  [python]
def estimate_tree(self, perturbed_datas, attr_config, layer_config, group_fos):
        all_domain_sizes = []
        sr_domain_sizes = []
        attr_count = 0
        for report_attrs in attr_config:
            domain_sizes = []
            for report_attr in report_attrs:
                domain_sizes.append(len(self.tree_range_lists[report_attr][layer_config[attr_count]]))
                attr_count += 1
            all_domain_sizes.append(domain_sizes)
            sr_domain_sizes.append(np.prod(domain_sizes))

        result_after_ce = self.estimate_sr(perturbed_datas, group_fos, sr_domain_sizes)
        precise_result = self.expand_ce(result_after_ce, all_domain_sizes, sr_domain_sizes)
        return precise_result
//  [python]
def estimate_sr(perturbed_datas, fos, all_domain_sizes):
        n = len(perturbed_datas)
        dimension = len(all_domain_sizes)
        result = OrderedDict()

        for d in range(1, dimension + 1):

            result[d] = OrderedDict()
            attr_sets = itertools.combinations(range(dimension), d)

            for attr_set in attr_sets:
                result[d][attr_set] = OrderedDict()
                domain_size_subset = [range(all_domain_sizes[x]) for x in attr_set]
                all_possible_digits = itertools.product(*domain_size_subset)

                for digits in all_possible_digits:
                    same_count = 0
                    for i in range(n):
                        same = True
                        for j in range(len(attr_set)):
                            if not fos[attr_set[j]].support_sr(perturbed_datas[i][attr_set[j]], digits[j]):
                                same = False
                                break
                        same_count += same

                    minus_count = n * n.prod([fos[x].q for x in attr_set])
                    for d_ in range(1, d):
                        count_d_ = 0
                        chosen_sets_ = itertools.combinations(range(d), d_)
                        for chosen_set_ in chosen_sets_:
                            coeff = n.prod([fos[d].q for d in attr_set])
                            coeff *= n.prod(
                                [(fos[attr_set[x]].p / fos[attr_set[x]].q - 1) for x in chosen_set_])

                            attr_set_ = tuple([attr_set[chosen_set_[k]] for k in range(d_)])
                            digits_ = tuple([digits[chosen_set_[k]] for k in range(d_)])

                            count_d_ += result[d_][attr_set_][digits_] * coeff
                        minus_count += count_d_

                    est = (same_count - minus_count) / n.prod([(fos[x].p - fos[x].q) for x in attr_set])
                    result[d][attr_set][digits] = est

        return result
//  [python]
def estimate_sr(perturbed_datas, fos, all_domain_sizes):
        n = len(perturbed_datas)
        dimension = len(all_domain_sizes)
        result = OrderedDict()

        for d in range(1, dimension + 1):

            result[d] = OrderedDict()
            attr_sets = itertools.combinations(range(dimension), d)

            for attr_set in attr_sets:
                result[d][attr_set] = OrderedDict()
                domain_size_subset = [range(all_domain_sizes[x]) for x in attr_set]
                all_possible_digits = itertools.product(*domain_size_subset)

                for digits in all_possible_digits:
                    same_count = 0
                    for i in range(n):
                        same = True
                        for j in range(len(attr_set)):
                            if not fos[attr_set[j]].support_sr(perturbed_datas[i][attr_set[j]], digits[j]):
                                same = False
                                break
                        same_count += same

                    minus_count = n * np.prod([fos[x].q for x in attr_set])
                    for d_ in range(1, d):
                        count_d_ = 0
                        chosen_sets_ = itertools.combinations(range(d), d_)
                        for chosen_set_ in chosen_sets_:
                            coeff = np.prod([fos[d].q for d in attr_set])
                            coeff *= np.prod(
                                [(fos[attr_set[x]].p / fos[attr_set[x]].q - 1) for x in chosen_set_])

                            attr_set_ = tuple([attr_set[chosen_set_[k]] for k in range(d_)])
                            digits_ = tuple([digits[chosen_set_[k]] for k in range(d_)])

                            count_d_ += result[d_][attr_set_][digits_] * coeff
                        minus_count += count_d_

                    est = (same_count - minus_count) / np.prod([(fos[x].p - fos[x].q) for x in attr_set])
                    result[d][attr_set][digits] = est

        return result
//  [python]
def elsize(self):
        """
        The total element size, i.e. `product(self.elshape)`

        Returns
        -------
        int
        """
        if self.elshape is None: return 0
        return int(n.prod(self.elshape))
//  [python]
def elsize(self):
        """
        The total element size, i.e. `product(self.elshape)`

        Returns
        -------
        int
        """
        if self.elshape is None: return 0
        return int(_np.prod(self.elshape))
//  [python]
def __init__(self, rng, input, filter_shape, image_shape, poolsize=(2, 2)):
        """
        Allocate a LeNetConvPoolLayer with shared variable internal parameters.

        :type rng: numpy.random.RandomState
        :param rng: a random number generator used to initialize weights

        :type input: theano.tensor.dtensor4
        :param input: symbolic image tensor, of shape image_shape

        :type filter_shape: tuple or list of length 4
        :param filter_shape: (number of filters, num input feature maps,
                              filter height, filter width)

        :type image_shape: tuple or list of length 4
        :param image_shape: (batch size, num input feature maps,
                             image height, image width)

        :type poolsize: tuple or list of length 2
        :param poolsize: the downsampling (pooling) factor (#rows, #cols)
        """
    
        assert image_shape[1] == filter_shape[1]
        self.input = input
        
        # there are "num input feature maps * filter height * filter width"
        # inputs to each hidden unit
        fan_in = n.prod(filter_shape[1:])
    
        # each unit in the lower layer receives a gradient from:
        # "num output feature maps * filter height * filter width" / pooling size
        fan_out = (filter_shape[0] * n.prod(filter_shape[2:]) //
                   n.prod(poolsize))
    
        # initialize weights with random weights
        W_bound = numpy.sqrt(6. / (fan_in + fan_out))
        self.W = theano.shared(
            numpy.asarray(
                rng.uniform(low=-W_bound, high=W_bound, size=filter_shape),
                dtype=theano.config.floatX
            ),
            borrow=True
        ) 
    
        # the bias is a 1D tensor -- one bias per output feature map
        b_values = numpy.zeros((filter_shape[0],), dtype=theano.config.floatX)
        self.b = theano.shared(value=b_values, borrow=True)  
    
        # convolve input feature maps with filters
        conv_out = conv2d(
            input=input,
            filters=self.W,
            filter_shape=filter_shape,
            input_shape=image_shape
        )

        # pool each feature map individually, using maxpooling
        pooled_out = pool.pool_2d(
            input=conv_out,
            ds=poolsize,
            ignore_border=True
        )

        # add the bias term. Since the bias is a vector (1D array), we first
        # reshape it to a tensor of shape (1, n_filters, 1, 1). Each bias will
        # thus be broadcasted across mini-batches and feature map
        # width & height
        self.output = T.tanh(pooled_out + self.b.dimshuffle('x', 0, 'x', 'x'))

        # store parameters of this layer
        self.params = [self.W, self.b]

        # keep track of model input
        self.input = input
//  [python]
def __init__(self, rng, input, filter_shape, image_shape, poolsize=(2, 2)):
        """
        Allocate a LeNetConvPoolLayer with shared variable internal parameters.

        :type rng: numpy.random.RandomState
        :param rng: a random number generator used to initialize weights

        :type input: theano.tensor.dtensor4
        :param input: symbolic image tensor, of shape image_shape

        :type filter_shape: tuple or list of length 4
        :param filter_shape: (number of filters, num input feature maps,
                              filter height, filter width)

        :type image_shape: tuple or list of length 4
        :param image_shape: (batch size, num input feature maps,
                             image height, image width)

        :type poolsize: tuple or list of length 2
        :param poolsize: the downsampling (pooling) factor (#rows, #cols)
        """
    
        assert image_shape[1] == filter_shape[1]
        self.input = input
        
        # there are "num input feature maps * filter height * filter width"
        # inputs to each hidden unit
        fan_in = numpy.prod(filter_shape[1:])
    
        # each unit in the lower layer receives a gradient from:
        # "num output feature maps * filter height * filter width" / pooling size
        fan_out = (filter_shape[0] * numpy.prod(filter_shape[2:]) //
                   numpy.prod(poolsize))
    
        # initialize weights with random weights
        W_bound = numpy.sqrt(6. / (fan_in + fan_out))
        self.W = theano.shared(
            numpy.asarray(
                rng.uniform(low=-W_bound, high=W_bound, size=filter_shape),
                dtype=theano.config.floatX
            ),
            borrow=True
        ) 
    
        # the bias is a 1D tensor -- one bias per output feature map
        b_values = numpy.zeros((filter_shape[0],), dtype=theano.config.floatX)
        self.b = theano.shared(value=b_values, borrow=True)  
    
        # convolve input feature maps with filters
        conv_out = conv2d(
            input=input,
            filters=self.W,
            filter_shape=filter_shape,
            input_shape=image_shape
        )

        # pool each feature map individually, using maxpooling
        pooled_out = pool.pool_2d(
            input=conv_out,
            ds=poolsize,
            ignore_border=True
        )

        # add the bias term. Since the bias is a vector (1D array), we first
        # reshape it to a tensor of shape (1, n_filters, 1, 1). Each bias will
        # thus be broadcasted across mini-batches and feature map
        # width & height
        self.output = T.tanh(pooled_out + self.b.dimshuffle('x', 0, 'x', 'x'))

        # store parameters of this layer
        self.params = [self.W, self.b]

        # keep track of model input
        self.input = input
//  [python]
def test_kuccsd_supercell_vs_kpts_high_cost():
    cell = gto.M(
        unit = 'B',
        a = [[ 0.,          3.37013733,  3.37013733],
             [ 3.37013733,  0.,          3.37013733],
             [ 3.37013733,  3.37013733,  0.        ]],
        mesh = [13]*3,
        atom = '''He 0 0 0
                  He 1.68506866 1.68506866 1.68506866''',
        basis = [[0, (1., 1.)], [0, (.5, 1.)]],
        verbose = 0,
    )

    nmp = [3,3,1]

    # treating supercell at gamma point
    supcell = super_cell(cell,nmp)

    gmf  = scf.UHF(supcell,exxdiv=None)
    ehf  = gmf.kernel()
    gcc  = cc.UCCSD(gmf)
    ecc, t1, t2 = gcc.kernel()
    print('UHF energy (supercell) %f \n' % (float(ehf)/n.prod(nmp)))
    print('UCCSD correlation energy (supercell) %f \n' % (float(ecc)/n.prod(nmp)))
    assert abs(ehf / 9 - -4.343308413289) < 1e-7
    assert abs(ecc / 9 - -0.009470753047083676) < 1e-6

    # treating mesh of k points

    kpts  = cell.make_kpts(nmp)
    kpts -= kpts[0]
    kmf   = scf.KUHF(cell,kpts,exxdiv=None)
    ehf   = kmf.kernel()
    kcc   = cc.KUCCSD(kmf)
    ecc, t1, t2 = kcc.kernel()

    print('UHF energy (kpts) %f \n' % ehf)
    print('UCCSD correlation energy (kpts) %f \n' % ecc)
    assert abs(ehf - -4.343308413289) < 1e-7
    assert abs(ecc - -0.009470753047083676) < 1e-6
//  [python]
def test_kuccsd_supercell_vs_kpts_high_cost():
    cell = gto.M(
        unit = 'B',
        a = [[ 0.,          3.37013733,  3.37013733],
             [ 3.37013733,  0.,          3.37013733],
             [ 3.37013733,  3.37013733,  0.        ]],
        mesh = [13]*3,
        atom = '''He 0 0 0
                  He 1.68506866 1.68506866 1.68506866''',
        basis = [[0, (1., 1.)], [0, (.5, 1.)]],
        verbose = 0,
    )

    nmp = [3,3,1]

    # treating supercell at gamma point
    supcell = super_cell(cell,nmp)

    gmf  = scf.UHF(supcell,exxdiv=None)
    ehf  = gmf.kernel()
    gcc  = cc.UCCSD(gmf)
    ecc, t1, t2 = gcc.kernel()
    print('UHF energy (supercell) %f \n' % (float(ehf)/numpy.prod(nmp)))
    print('UCCSD correlation energy (supercell) %f \n' % (float(ecc)/numpy.prod(nmp)))
    assert abs(ehf / 9 - -4.343308413289) < 1e-7
    assert abs(ecc / 9 - -0.009470753047083676) < 1e-6

    # treating mesh of k points

    kpts  = cell.make_kpts(nmp)
    kpts -= kpts[0]
    kmf   = scf.KUHF(cell,kpts,exxdiv=None)
    ehf   = kmf.kernel()
    kcc   = cc.KUCCSD(kmf)
    ecc, t1, t2 = kcc.kernel()

    print('UHF energy (kpts) %f \n' % ehf)
    print('UCCSD correlation energy (kpts) %f \n' % ecc)
    assert abs(ehf - -4.343308413289) < 1e-7
    assert abs(ecc - -0.009470753047083676) < 1e-6
//  [python]
def atom_perturb(self, eps, group_perturbed_datas, report_count, datas, user_start, user_end, report_attrs, report_layers):
        domain_sizes = [0] * len(report_attrs)
        group_datas = datas[user_start:user_end, report_attrs].copy()
        for i in range(len(report_attrs)):
            report_attr = report_attrs[i]
            report_layer = report_layers[i]
            bins = [self.tree_range_lists[report_attr][report_layer][0][0]] + [x[1] for x in self.tree_range_lists[report_attr][report_layer]]
            # domain_sizes[i] = len(bins) - 1
            domain_sizes[i] = len(self.tree_range_lists[report_attrs[i]][report_layers[i]])
            group_datas[:, i] = np.digitize(group_datas[:, i], bins) - 1

        total_size = reduce(lambda x, y: x * y, domain_sizes)
        fo = FOFactory.create_fo('adap', self.args, eps, total_size)

        datas_to_perturb = np.zeros(len(group_datas), dtype=np.int)
        for domain in range(len(domain_sizes)):
            factor = reduce(lambda x, y: x * y, domain_sizes[domain + 1:])
            datas_to_perturb += group_datas[:, domain] * factor
        perturbed_datas = fo.perturb(datas_to_perturb, reduce(lambda x, y: x * y, domain_sizes))

        for i in range(len(perturbed_datas)):
            group_perturbed_datas[i][report_count].append(perturbed_datas[i])
        return fo
//  [python]
def atom_perturb(self, eps, group_perturbed_datas, report_count, datas, user_start, user_end, report_attrs, report_layers):
        domain_sizes = [0] * len(report_attrs)
        group_datas = datas[user_start:user_end, report_attrs].copy()
        for i in range(len(report_attrs)):
            report_attr = report_attrs[i]
            report_layer = report_layers[i]
            bins = [self.tree_range_lists[report_attr][report_layer][0][0]] + [x[1] for x in self.tree_range_lists[report_attr][report_layer]]
            # domain_sizes[i] = len(bins) - 1
            domain_sizes[i] = len(self.tree_range_lists[report_attrs[i]][report_layers[i]])
            group_datas[:, i] = np.digitize(group_datas[:, i], bins) - 1

        total_size = np.prod(domain_sizes)
        fo = FOFactory.create_fo('adap', self.args, eps, total_size)

        datas_to_perturb = np.zeros(len(group_datas), dtype=np.int)
        for domain in range(len(domain_sizes)):
            factor = int(np.prod(domain_sizes[domain + 1:]))
            datas_to_perturb += group_datas[:, domain] * factor
        perturbed_datas = fo.perturb(datas_to_perturb, np.prod(domain_sizes))

        for i in range(len(perturbed_datas)):
            group_perturbed_datas[i][report_count].append(perturbed_datas[i])
        return fo
//  [python]
def size(self):
        """
        The number of elements (or vector-elements) in the basis.
        """
        return int(n.prod([c.size for c in self.component_bases]))
//  [python]
def size(self):
        """
        The number of elements (or vector-elements) in the basis.
        """
        return int(_np.prod([c.size for c in self.component_bases]))
//  [python]
def dim(self):
        """
        The dimension of the vector space this basis fully or partially
        spans.  Equivalently, the length of the `vector_elements` of the
        basis.
        """
        dim = int(n.prod([c.dim for c in self.component_bases]))

        #NOTE: this is actually to restrictive -- what we need is a test/flag for whether the elements of a
        # basis are in their "natrual" representation where it makes sense to take tensor products.  For
        # example, a direct-sum basis may hold elements in a compact way that violate this... but I'm not sure if they
        # do and this needs to be checked.  For now, we could just disable this overly-restrictive assert:
        assert(all([c.is_simple() for c in self.component_bases])), \
            "Components of a tensor product basis must be *simple* (have vector-dimension == size of elements)"
        # because we use the natural representation to take tensor (kronecker) products.
        # Note: this assertion also means dim == product(component_elsizes) == elsize, so basis is *simple*

        return dim
//  [python]
def dim(self):
        """
        The dimension of the vector space this basis fully or partially
        spans.  Equivalently, the length of the `vector_elements` of the
        basis.
        """
        dim = int(_np.prod([c.dim for c in self.component_bases]))

        #NOTE: this is actually to restrictive -- what we need is a test/flag for whether the elements of a
        # basis are in their "natrual" representation where it makes sense to take tensor products.  For
        # example, a direct-sum basis may hold elements in a compact way that violate this... but I'm not sure if they
        # do and this needs to be checked.  For now, we could just disable this overly-restrictive assert:
        assert(all([c.is_simple() for c in self.component_bases])), \
            "Components of a tensor product basis must be *simple* (have vector-dimension == size of elements)"
        # because we use the natural representation to take tensor (kronecker) products.
        # Note: this assertion also means dim == product(component_elsizes) == elsize, so basis is *simple*

        return dim
//  [python]
def init_stepsize(self, weights_shape, params):
        self.step_sizes = numpy.ones(weights_shape) * self.alpha
        self.h = numpy.zeros((n.prod(weights_shape),))
        self.v = numpy.zeros((n.prod(weights_shape),))
        # Autostep should not be used with eligibility traces (in current form)
        #self.lmbda = 0.0
        self.mu = params.setdefault('autostep_mu', 1.0e-2)
        self.tau = params.setdefault('autostep_tau', 1.0e4)
//  [python]
def init_stepsize(self, weights_shape, params):
        self.step_sizes = numpy.ones(weights_shape) * self.alpha
        self.h = numpy.zeros((numpy.prod(weights_shape),))
        self.v = numpy.zeros((numpy.prod(weights_shape),))
        # Autostep should not be used with eligibility traces (in current form)
        #self.lmbda = 0.0
        self.mu = params.setdefault('autostep_mu', 1.0e-2)
        self.tau = params.setdefault('autostep_tau', 1.0e4)
//  [python]
def init_stepsize(self, weights_shape, params):
        self.step_sizes = numpy.ones(weights_shape) * self.alpha
        self.prev_grad = None
        self.v = numpy.ones((n.prod(weights_shape),))
        self.almeida_gamma = params.setdefault('almeida_gamma', 0.999)
        self.almeida_stepsize = params.setdefault('almeida_stepsize', 0.00001)
//  [python]
def init_stepsize(self, weights_shape, params):
        self.step_sizes = numpy.ones(weights_shape) * self.alpha
        self.prev_grad = None
        self.v = numpy.ones((numpy.prod(weights_shape),))
        self.almeida_gamma = params.setdefault('almeida_gamma', 0.999)
        self.almeida_stepsize = params.setdefault('almeida_stepsize', 0.00001)
//  [python]
def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        log.info('******** %s ********', self.__class__)
        if self.auxcell is None:
            log.info('auxbasis = %s', self.auxbasis)
        else:
            log.info('auxbasis = %s', self.auxcell.basis)
        if self.eta is not None:
            log.info('eta = %s', self.eta)
        if self.mesh is not None:
            log.info('mesh = %s (%d PWs)', self.mesh, n.prod(self.mesh))
        log.info('exp_to_discard = %s', self.exp_to_discard)
        if isinstance(self._cderi, str):
            log.info('_cderi = %s  where DF integrals are loaded (readonly).',
                     self._cderi)
        elif isinstance(self._cderi_to_save, str):
            log.info('_cderi_to_save = %s', self._cderi_to_save)
        else:
            log.info('_cderi_to_save = %s', self._cderi_to_save.name)
        log.info('len(kpts) = %d', len(self.kpts))
        log.debug1('    kpts = %s', self.kpts)
        if self.kpts_band is not None:
            log.info('len(kpts_band) = %d', len(self.kpts_band))
            log.debug1('    kpts_band = %s', self.kpts_band)
        return self
//  [python]
def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        log.info('******** %s ********', self.__class__)
        if self.auxcell is None:
            log.info('auxbasis = %s', self.auxbasis)
        else:
            log.info('auxbasis = %s', self.auxcell.basis)
        if self.eta is not None:
            log.info('eta = %s', self.eta)
        if self.mesh is not None:
            log.info('mesh = %s (%d PWs)', self.mesh, numpy.prod(self.mesh))
        log.info('exp_to_discard = %s', self.exp_to_discard)
        if isinstance(self._cderi, str):
            log.info('_cderi = %s  where DF integrals are loaded (readonly).',
                     self._cderi)
        elif isinstance(self._cderi_to_save, str):
            log.info('_cderi_to_save = %s', self._cderi_to_save)
        else:
            log.info('_cderi_to_save = %s', self._cderi_to_save.name)
        log.info('len(kpts) = %d', len(self.kpts))
        log.debug1('    kpts = %s', self.kpts)
        if self.kpts_band is not None:
            log.info('len(kpts_band) = %d', len(self.kpts_band))
            log.debug1('    kpts_band = %s', self.kpts_band)
        return self
//  [python]
def partition_users(self):
        num_user = self.users.n
        dimension = len(self.users.column_to_perturb)
        # todo: for now, assume all attributes are used
        d = dimension
        # this solution is brute force
        # more elegant solutions needed!
        configs = []
        attr_configs = {}
        report_size = min(d, self.args.ce_size)
        report_per_user = int(math.ceil(d / report_size))
        used_dim_set = itertools.combinations(range(dimension), d)
        for used_dims in used_dim_set:

            perms = itertools.permutations(used_dims)
            for perm in perms:
                # attributes are sorted
                attr_config = []
                for i in range(report_per_user):
                    sr_set = sorted(perm[i * report_size: (i + 1) * report_size])
                    attr_config.append(tuple(sr_set))
                attr_config = tuple(sorted(attr_config))

                if attr_config in attr_configs:
                    continue
                attr_configs[attr_config] = 1
                num_layers = [range(self.tree_heights[x]) for x in perm]

                # fuse the layers into configuration
                if self.args.layer_style == 'rl':
                    layer_configs = itertools.product(*num_layers)
                    cluster_heights = [range(reduce(lambda x, y: x * y, [self.tree_heights[x] for x in attr_config[i]])) for i in
                                       range(len(attr_config))]
                    report_counts = itertools.product(*cluster_heights)
                    for layer_config in layer_configs:
                        report_count = next(report_counts)
                        exist_star = False
                        for count in report_count:
                            if count == 0:
                                exist_star = True
                                break
                        # some report does not contribute info, which is prohibited
                        if exist_star:
                            continue
                        configs.append(tuple([attr_config, layer_config]))
                elif self.args.layer_style == 'bl':
                    configs.append(tuple([attr_config, tuple([self.tree_heights[x] - 1 for x in perm])]))
                elif self.args.layer_style == 'al':
                    configs.append(tuple([attr_config, tuple([self.tree_heights[x] - 1 for x in perm])]))
                    # layer_configs = itertools.product(*num_layers)
                    # next(layer_configs)
                    # mat_layer_configs = []
                    # for layer_config in layer_configs:
                    #     mat_layer_configs.append(layer_config)
                    # configs.append(tuple([attr_config, mat_layer_configs]))

            num_group = len(configs)
            self.group_configs = configs
            for group_id in range(num_group):
                user_start = int(num_user / num_group * group_id)
                user_end = int(num_user / num_group * (group_id + 1))
                self.group_indexs.append((user_start, user_end))