jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _parse_mm_header(fn, sparse_threshold=0.01):
  with open(fn) as fp:
    line = fp.readline().strip()
    header = line
    while line[0] == '%':
      line = fp.readline()
    shape_info = line
    data_begin = fp.tell()

  if header.find('real'):
    dtype = np.float
  elif header.find('integer'):
    dtype = np.int

  shape, edges = _extract_mm_coordinate(shape_info)
  edges = int(edges)
  if (edges * 1.0) / np.product(shape) < sparse_threshold:
    sparse = True
  else:
    sparse = False

  return (shape, dtype, sparse, data_begin)
//  [python]
def _parse_mm_header(fn, sparse_threshold=0.01):
  with open(fn) as fp:
    line = fp.readline().strip()
    header = line
    while line[0] == '%':
      line = fp.readline()
    shape_info = line
    data_begin = fp.tell()

  if header.find('real'):
    dtype = np.float
  elif header.find('integer'):
    dtype = np.int

  shape, edges = _extract_mm_coordinate(shape_info)
  edges = int(edges)
  if (edges * 1.0) / np.prod(shape) < sparse_threshold:
    sparse = True
  else:
    sparse = False

  return (shape, dtype, sparse, data_begin)
//  [python]
def flatten_data(data):
    """Input is a matrix of d n-dimensional data points. Data points are
    assumed to be indexed by axis zero. E.g. data[i, ...] is the ith data point.
    Returns a 2-dimensional matrix where each data point has been flattened to
    a 1-D vector. Also returns the original shape of the data points."""
    data = numpy.asarray(data)
    data_count = data.shape[0]
    data_point_shape = data.shape[1:]
    if len(data.shape) > 1:
        flat = numpy.reshape(data, (data_count, numpy.product(data_point_shape)))
    else:
        flat = numpy.atleast_2d(data).transpose()
    return flat, data_point_shape
//  [python]
def flatten_data(data):
    """Input is a matrix of d n-dimensional data points. Data points are
    assumed to be indexed by axis zero. E.g. data[i, ...] is the ith data point.
    Returns a 2-dimensional matrix where each data point has been flattened to
    a 1-D vector. Also returns the original shape of the data points."""
    data = numpy.asarray(data)
    data_count = data.shape[0]
    data_point_shape = data.shape[1:]
    if len(data.shape) > 1:
        flat = numpy.reshape(data, (data_count, numpy.prod(data_point_shape)))
    else:
        flat = numpy.atleast_2d(data).transpose()
    return flat, data_point_shape
//  [python]
def deserialized_length(self):
        if self.data_length > 0:
            return self.data_length
        element_size: int = numpy.dtype(self.dtype).itemsize
        num_elements: int = numpy.product(self.shape)
        return element_size * num_elements
//  [python]
def deserialized_length(self):
        if self.data_length > 0:
            return self.data_length
        element_size: int = numpy.dtype(self.dtype).itemsize
        num_elements: int = numpy.prod(self.shape)
        return element_size * num_elements
//  [python]
def L(alpha,beta): 
    prod = numpy.product(scipy.stats.binom.pmf(y,n,invlogit(alpha+beta*x)))
    return prod
//  [python]
def L(alpha,beta): 
    prod = numpy.prod(scipy.stats.binom.pmf(y,n,invlogit(alpha+beta*x)))
    return prod
//  [python]
def reshape(self, *shape):
		if len(shape)==0:
			raise ValueError('no shape given')
		if len(shape)==1 and isinstance(shape[0], tuple):
			shape = shape[0]
		if isinstance(shape, int):
			shape = (shape,)
		if not isinstance(shape, tuple):
			raise TypeError('reshape must be int or tuple')
		if len(shape)==1 and shape[0] == -1:
			shape = (int(numpy.product(self.shape)), )
		elif len(shape)==2:
			if shape[0] == -1 and shape[1] > 0:
				shape = (int(numpy.product(self.shape) / shape[1]), shape[1])
			if shape[1] == -1 and shape[0] > 0:
				shape = (shape[0], int(numpy.product(self.shape) / shape[0]))
		if numpy.product(shape) != numpy.product(self.shape):
			raise ValueError(f'incompatible reshape {shape} for current shape {self.shape}')
		#print("reshape to", shape)
		for k in self._groupnode._v_children.keys():
			#print("reshape",k,shape)
			self._groupnode._v_children[k].shape = shape
		self.shape = shape
//  [python]
def reshape(self, *shape):
		if len(shape)==0:
			raise ValueError('no shape given')
		if len(shape)==1 and isinstance(shape[0], tuple):
			shape = shape[0]
		if isinstance(shape, int):
			shape = (shape,)
		if not isinstance(shape, tuple):
			raise TypeError('reshape must be int or tuple')
		if len(shape)==1 and shape[0] == -1:
			shape = (int(numpy.prod(self.shape)), )
		elif len(shape)==2:
			if shape[0] == -1 and shape[1] > 0:
				shape = (int(numpy.prod(self.shape) / shape[1]), shape[1])
			if shape[1] == -1 and shape[0] > 0:
				shape = (shape[0], int(numpy.prod(self.shape) / shape[0]))
		if numpy.prod(shape) != numpy.prod(self.shape):
			raise ValueError(f'incompatible reshape {shape} for current shape {self.shape}')
		#print("reshape to", shape)
		for k in self._groupnode._v_children.keys():
			#print("reshape",k,shape)
			self._groupnode._v_children[k].shape = shape
		self.shape = shape
//  [python]
def _local_read_sparse_mm(array, ex, fn, data_begin):
  '''
  1. Noted that Matrix Market format doesn't require (row, col) to be sorted.
     If the file is sorted (by either row or col), each worker will return
     only a part of the array. If the file is unsorted, each worker may
     return a very big and sparser sub-array of the original array. In the
     worst case, the sub-array can be as large as the original array but
     sparser.
  2. We can't know how many lines without reading the whole file. So we simply
     decide the region this worker should read based on the file size.
  '''
  data_size = os.path.getsize(fn) - data_begin
  array_size = np.product(array.shape)
  begin = extent.ravelled_pos(ex.ul, array.shape)
  begin = math.ceil(((begin * 1.0) / array_size) * data_size) + data_begin
  end = extent.ravelled_pos([(i - 1) for i in ex.lr], array.shape)
  end = math.floor(((end * 1.0) / array_size) * data_size) + data_begin

  ul = [array.shape[0], array.shape[1]]
  lr = [0, 0]
  rows = []
  cols = []
  data = []
  with open(fn) as fp:
    fp.seek(begin)
    if begin != data_begin:
      fp.seek(begin - 1)
      a = fp.read(1)
      if a != '\n':
        line = fp.readline()

    pos = fp.tell()
    for line in fp:
      if pos > end + 1:  # +1 in case end locates on \n
        break
      pos += len(line)
      (_row, _col), val = _extract_mm_coordinate(line)
      _row -= 1
      _col -= 1
      rows.append(_row)
      cols.append(_col)
      data.append(float(val))
      ul[0] = _row if _row < ul[0] else ul[0]
      ul[1] = _col if _col < ul[1] else ul[1]
      lr[0] = _row if _row > lr[0] else lr[0]
      lr[1] = _col if _col > lr[1] else lr[1]

  # Adjust rows and cols based on the ul of this submatrix.
  for i in xrange(len(rows)):
    rows[i] -= ul[0]
    cols[i] -= ul[1]

  new_ex = extent.create(ul, [lr[0] + 1, lr[1] + 1], array.shape)
  new_array = sp.coo_matrix((data, (rows, cols)), new_ex.shape)
  return new_ex, sparse.convert_sparse_array(new_array)
//  [python]
def _local_read_sparse_mm(array, ex, fn, data_begin):
  '''
  1. Noted that Matrix Market format doesn't require (row, col) to be sorted.
     If the file is sorted (by either row or col), each worker will return
     only a part of the array. If the file is unsorted, each worker may
     return a very big and sparser sub-array of the original array. In the
     worst case, the sub-array can be as large as the original array but
     sparser.
  2. We can't know how many lines without reading the whole file. So we simply
     decide the region this worker should read based on the file size.
  '''
  data_size = os.path.getsize(fn) - data_begin
  array_size = np.prod(array.shape)
  begin = extent.ravelled_pos(ex.ul, array.shape)
  begin = math.ceil(((begin * 1.0) / array_size) * data_size) + data_begin
  end = extent.ravelled_pos([(i - 1) for i in ex.lr], array.shape)
  end = math.floor(((end * 1.0) / array_size) * data_size) + data_begin

  ul = [array.shape[0], array.shape[1]]
  lr = [0, 0]
  rows = []
  cols = []
  data = []
  with open(fn) as fp:
    fp.seek(begin)
    if begin != data_begin:
      fp.seek(begin - 1)
      a = fp.read(1)
      if a != '\n':
        line = fp.readline()

    pos = fp.tell()
    for line in fp:
      if pos > end + 1:  # +1 in case end locates on \n
        break
      pos += len(line)
      (_row, _col), val = _extract_mm_coordinate(line)
      _row -= 1
      _col -= 1
      rows.append(_row)
      cols.append(_col)
      data.append(float(val))
      ul[0] = _row if _row < ul[0] else ul[0]
      ul[1] = _col if _col < ul[1] else ul[1]
      lr[0] = _row if _row > lr[0] else lr[0]
      lr[1] = _col if _col > lr[1] else lr[1]

  # Adjust rows and cols based on the ul of this submatrix.
  for i in xrange(len(rows)):
    rows[i] -= ul[0]
    cols[i] -= ul[1]

  new_ex = extent.create(ul, [lr[0] + 1, lr[1] + 1], array.shape)
  new_array = sp.coo_matrix((data, (rows, cols)), new_ex.shape)
  return new_ex, sparse.convert_sparse_array(new_array)
//  [python]
def __array__(self, dtype=None):
        """ Create a Numpy array containing the whole dataset.  DON'T THINK
        THIS MEANS DATASETS ARE INTERCHANGABLE WITH ARRAYS.  For one thing,
        you have to read the whole dataset everytime this method is called.
        """
        arr = numpy.empty(self.shape, dtype=self.dtype if dtype is None else dtype)

        # Special case for (0,)*-shape datasets
        if numpy.product(self.shape) == 0:
            return arr

        self.read_direct(arr)
        return arr
//  [python]
def __array__(self, dtype=None):
        """ Create a Numpy array containing the whole dataset.  DON'T THINK
        THIS MEANS DATASETS ARE INTERCHANGABLE WITH ARRAYS.  For one thing,
        you have to read the whole dataset everytime this method is called.
        """
        arr = numpy.empty(self.shape, dtype=self.dtype if dtype is None else dtype)

        # Special case for (0,)*-shape datasets
        if numpy.prod(self.shape) == 0:
            return arr

        self.read_direct(arr)
        return arr
//  [python]
def get_surface(self, dest_surf=None):
        camera = self.camera

        im = highgui.cvQueryFrame(camera)
        #convert Ipl image to PIL image
        #print type(im)
        if im:
            xx = opencv.adaptors.Ipl2NumPy(im)
            #print type(xx)
            #print xx.iscontiguous()
            #print dir(xx)
            #print xx.shape
            xxx = numpy.reshape(xx, (numpy.product(xx.shape),))

            if xx.shape[2] != 3:
                raise ValueError("not sure what to do about this size")

            pg_img = pygame.image.frombuffer(xxx, (xx.shape[1], xx.shape[0]), "RGB")

            # if there is a destination surface given, we blit onto that.
            if dest_surf:
                dest_surf.blit(pg_img, (0, 0))
            return dest_surf
//  [python]
def get_surface(self, dest_surf=None):
        camera = self.camera

        im = highgui.cvQueryFrame(camera)
        #convert Ipl image to PIL image
        #print type(im)
        if im:
            xx = opencv.adaptors.Ipl2NumPy(im)
            #print type(xx)
            #print xx.iscontiguous()
            #print dir(xx)
            #print xx.shape
            xxx = numpy.reshape(xx, (numpy.prod(xx.shape),))

            if xx.shape[2] != 3:
                raise ValueError("not sure what to do about this size")

            pg_img = pygame.image.frombuffer(xxx, (xx.shape[1], xx.shape[0]), "RGB")

            # if there is a destination surface given, we blit onto that.
            if dest_surf:
                dest_surf.blit(pg_img, (0, 0))
            return dest_surf
//  [python]
def _is_vertex_on_constraint(vertex, constraint):
    """
    Given a vertex and a constraint, check if the vertex is on the plane defined
    by the constraint. For inequality constraints, the plane we look at is the
    extremal plane that separates the elements that fulfill an inequality
    constraint from the elements that do not fulfill this constraints.
    """
    r = _vertex_to_rational_point(vertex)

    dims = constraint.space.dim(dim_type.set)
    v = []
    for d in range(dims):
        v.append(constraint.get_coefficient_val(dim_type.set, d).get_num_si())

    summ = 0

    import numpy
    for i in range(dims):
        prod = 1
        for j in range(dims):
            if i == j:
                prod *= r[j][0]
            else:
                prod *= r[j][1]
        summ += v[i] * prod

    constant = constraint.get_constant_val().get_num_si()
    summ += numpy.product([x[1] for x in r]) * constant

    return int(summ) == 0
//  [python]
def _is_vertex_on_constraint(vertex, constraint):
    """
    Given a vertex and a constraint, check if the vertex is on the plane defined
    by the constraint. For inequality constraints, the plane we look at is the
    extremal plane that separates the elements that fulfill an inequality
    constraint from the elements that do not fulfill this constraints.
    """
    r = _vertex_to_rational_point(vertex)

    dims = constraint.space.dim(dim_type.set)
    v = []
    for d in range(dims):
        v.append(constraint.get_coefficient_val(dim_type.set, d).get_num_si())

    summ = 0

    import numpy
    for i in range(dims):
        prod = 1
        for j in range(dims):
            if i == j:
                prod *= r[j][0]
            else:
                prod *= r[j][1]
        summ += v[i] * prod

    constant = constraint.get_constant_val().get_num_si()
    summ += numpy.prod([x[1] for x in r]) * constant

    return int(summ) == 0
//  [python]
def _load_mnist_vectors(fn: str) -> numpy.ndarray:
    import gzip
    import struct

    print("parsing vectors in %s..." % fn)
    f = gzip.open(fn)
    type_code_info = {
        0x08: (1, "!B"),
        0x09: (1, "!b"),
        0x0B: (2, "!H"),
        0x0C: (4, "!I"),
        0x0D: (4, "!f"),
        0x0E: (8, "!d"),
    }
    magic, type_code, dim_count = struct.unpack("!hBB", f.read(4))
    assert magic == 0
    assert type_code in type_code_info

    dimensions = [struct.unpack("!I", f.read(4))[0] for i in range(dim_count)]

    entry_count = dimensions[0]
    entry_size = numpy.product(dimensions[1:])

    b, format_string = type_code_info[type_code]
    vectors = []
    for i in range(entry_count):
        vectors.append([struct.unpack(format_string, f.read(b))[0] for j in range(entry_size)])
    return numpy.array(vectors)
//  [python]
def _load_mnist_vectors(fn: str) -> numpy.ndarray:
    import gzip
    import struct

    print("parsing vectors in %s..." % fn)
    f = gzip.open(fn)
    type_code_info = {
        0x08: (1, "!B"),
        0x09: (1, "!b"),
        0x0B: (2, "!H"),
        0x0C: (4, "!I"),
        0x0D: (4, "!f"),
        0x0E: (8, "!d"),
    }
    magic, type_code, dim_count = struct.unpack("!hBB", f.read(4))
    assert magic == 0
    assert type_code in type_code_info

    dimensions = [struct.unpack("!I", f.read(4))[0] for i in range(dim_count)]

    entry_count = dimensions[0]
    entry_size = numpy.prod(dimensions[1:])

    b, format_string = type_code_info[type_code]
    vectors = []
    for i in range(entry_count):
        vectors.append([struct.unpack(format_string, f.read(b))[0] for j in range(entry_size)])
    return numpy.array(vectors)
//  [python]
def image_circle_patch_diameter(x, x_adv):
    """
    Return smallest circular patch diameter that covers the modified pixels
        This value is normalized by the smallest spatial dimension
        So that a value of 1 is the largest circular patch that fits inside the image
            Values greater than 1 are possible for non-square images

    This metric *assumes* that the modification was circular and contained in the image
        If the modification is non-circular, then this may underestimate diameter

    Note: If this assumption was not made, a more complicated algorithm would be needed
        See: "Smallest enclosing disks (balls and ellipsoids)", Emo Welzl, (1991)
            or https://en.wikipedia.org/wiki/Smallest-circle_problem
    """
    if x.ndim == 2:
        x = np.expand_dims(x, axis=-1)
        x_adv = np.expand_dims(x_adv, axis=-1)
    if x.ndim != 3:
        raise ValueError(
            f"Expected image with 2 (HW) or 3 (HWC or CHW) dimensions. x has shape {x.shape}"
        )

    # Assume channel dim is the smallest dimension
    channel_dim = np.argmin(x.shape)
    # mask is True if the pixel is different
    mask = (x != x_adv).any(axis=channel_dim)

    diff = mask.sum()
    if not diff:
        return 0

    # Get patch diameters for height and width, then take the larger of them
    diameters = []
    for axis in 0, 1:
        marginal = mask.any(axis=axis)
        min_index = np.argmax(marginal)
        max_index = len(marginal) - np.argmax(marginal[::-1]) - 1
        diameters.append(max_index - min_index + 1)
    diameter = max(diameters)

    size = np.product(mask.shape)
    max_diameter = min(mask.shape)  # max in-image diameter
    if diameter > max_diameter:
        log.warning("Circular patch is not contained within the image")
    elif diff >= np.pi * (max_diameter / 2) ** 2:
        log.warning(
            f"x and x_adv differ at {diff/size:.0%} of indices. "
            "Assumption of circular patch within image may be violated."
        )

    return diameter / max_diameter
//  [python]
def image_circle_patch_diameter(x, x_adv):
    """
    Return smallest circular patch diameter that covers the modified pixels
        This value is normalized by the smallest spatial dimension
        So that a value of 1 is the largest circular patch that fits inside the image
            Values greater than 1 are possible for non-square images

    This metric *assumes* that the modification was circular and contained in the image
        If the modification is non-circular, then this may underestimate diameter

    Note: If this assumption was not made, a more complicated algorithm would be needed
        See: "Smallest enclosing disks (balls and ellipsoids)", Emo Welzl, (1991)
            or https://en.wikipedia.org/wiki/Smallest-circle_problem
    """
    if x.ndim == 2:
        x = np.expand_dims(x, axis=-1)
        x_adv = np.expand_dims(x_adv, axis=-1)
    if x.ndim != 3:
        raise ValueError(
            f"Expected image with 2 (HW) or 3 (HWC or CHW) dimensions. x has shape {x.shape}"
        )

    # Assume channel dim is the smallest dimension
    channel_dim = np.argmin(x.shape)
    # mask is True if the pixel is different
    mask = (x != x_adv).any(axis=channel_dim)

    diff = mask.sum()
    if not diff:
        return 0

    # Get patch diameters for height and width, then take the larger of them
    diameters = []
    for axis in 0, 1:
        marginal = mask.any(axis=axis)
        min_index = np.argmax(marginal)
        max_index = len(marginal) - np.argmax(marginal[::-1]) - 1
        diameters.append(max_index - min_index + 1)
    diameter = max(diameters)

    size = np.prod(mask.shape)
    max_diameter = min(mask.shape)  # max in-image diameter
    if diameter > max_diameter:
        log.warning("Circular patch is not contained within the image")
    elif diff >= np.pi * (max_diameter / 2) ** 2:
        log.warning(
            f"x and x_adv differ at {diff/size:.0%} of indices. "
            "Assumption of circular patch within image may be violated."
        )

    return diameter / max_diameter
//  [python]
def apply_along_axis(func1d, axis, arr, *args, **kwargs):
    """
    (This docstring should be overwritten)
    """
    arr = array(arr, copy=False, subok=True)
    nd = arr.ndim
    axis = normalize_axis_index(axis, nd)
    ind = [0] * (nd - 1)
    i = np.zeros(nd, 'O')
    indlist = list(range(nd))
    indlist.remove(axis)
    i[axis] = slice(None, None)
    outshape = np.asarray(arr.shape).take(indlist)
    i.put(indlist, ind)
    res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
    #  if res is a number, then we have a smaller output array
    asscalar = np.isscalar(res)
    if not asscalar:
        try:
            len(res)
        except TypeError:
            asscalar = True
    # Note: we shouldn't set the dtype of the output from the first result
    # so we force the type to object, and build a list of dtypes.  We'll
    # just take the largest, to avoid some downcasting
    dtypes = []
    if asscalar:
        dtypes.append(np.asarray(res).dtype)
        outarr = zeros(outshape, object)
        outarr[tuple(ind)] = res
        Ntot = np.product(outshape)
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= outshape[n]) and (n > (1 - nd)):
                ind[n - 1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist, ind)
            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
            outarr[tuple(ind)] = res
            dtypes.append(asarray(res).dtype)
            k += 1
    else:
        res = array(res, copy=False, subok=True)
        j = i.copy()
        j[axis] = ([slice(None, None)] * res.ndim)
        j.put(indlist, ind)
        Ntot = np.product(outshape)
        holdshape = outshape
        outshape = list(arr.shape)
        outshape[axis] = res.shape
        dtypes.append(asarray(res).dtype)
        outshape = flatten_inplace(outshape)
        outarr = zeros(outshape, object)
        outarr[tuple(flatten_inplace(j.tolist()))] = res
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
                ind[n - 1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist, ind)
            j.put(indlist, ind)
            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
            outarr[tuple(flatten_inplace(j.tolist()))] = res
            dtypes.append(asarray(res).dtype)
            k += 1
    max_dtypes = np.dtype(np.asarray(dtypes).max())
    if not hasattr(arr, '_mask'):
        result = np.asarray(outarr, dtype=max_dtypes)
    else:
        result = asarray(outarr, dtype=max_dtypes)
        result.fill_value = ma.default_fill_value(result)
    return result
//  [python]
def apply_along_axis(func1d, axis, arr, *args, **kwargs):
    """
    (This docstring should be overwritten)
    """
    arr = array(arr, copy=False, subok=True)
    nd = arr.ndim
    axis = normalize_axis_index(axis, nd)
    ind = [0] * (nd - 1)
    i = np.zeros(nd, 'O')
    indlist = list(range(nd))
    indlist.remove(axis)
    i[axis] = slice(None, None)
    outshape = np.asarray(arr.shape).take(indlist)
    i.put(indlist, ind)
    res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
    #  if res is a number, then we have a smaller output array
    asscalar = np.isscalar(res)
    if not asscalar:
        try:
            len(res)
        except TypeError:
            asscalar = True
    # Note: we shouldn't set the dtype of the output from the first result
    # so we force the type to object, and build a list of dtypes.  We'll
    # just take the largest, to avoid some downcasting
    dtypes = []
    if asscalar:
        dtypes.append(np.asarray(res).dtype)
        outarr = zeros(outshape, object)
        outarr[tuple(ind)] = res
        Ntot = np.prod(outshape)
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= outshape[n]) and (n > (1 - nd)):
                ind[n - 1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist, ind)
            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
            outarr[tuple(ind)] = res
            dtypes.append(asarray(res).dtype)
            k += 1
    else:
        res = array(res, copy=False, subok=True)
        j = i.copy()
        j[axis] = ([slice(None, None)] * res.ndim)
        j.put(indlist, ind)
        Ntot = np.prod(outshape)
        holdshape = outshape
        outshape = list(arr.shape)
        outshape[axis] = res.shape
        dtypes.append(asarray(res).dtype)
        outshape = flatten_inplace(outshape)
        outarr = zeros(outshape, object)
        outarr[tuple(flatten_inplace(j.tolist()))] = res
        k = 1
        while k < Ntot:
            # increment the index
            ind[-1] += 1
            n = -1
            while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
                ind[n - 1] += 1
                ind[n] = 0
                n -= 1
            i.put(indlist, ind)
            j.put(indlist, ind)
            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
            outarr[tuple(flatten_inplace(j.tolist()))] = res
            dtypes.append(asarray(res).dtype)
            k += 1
    max_dtypes = np.dtype(np.asarray(dtypes).max())
    if not hasattr(arr, '_mask'):
        result = np.asarray(outarr, dtype=max_dtypes)
    else:
        result = asarray(outarr, dtype=max_dtypes)
        result.fill_value = ma.default_fill_value(result)
    return result
//  [python]
def make_new_dset(parent, shape=None, dtype=None, data=None,
                 chunks=None, compression=None, shuffle=None,
                    fletcher32=None, maxshape=None, compression_opts=None,
                  fillvalue=None, scaleoffset=None, track_times=None):
    """ Return a new low-level dataset identifier

    Only creates anonymous datasets.
    """

    # Convert data to a C-contiguous ndarray
    if data is not None:
        from . import base
        data = numpy.asarray(data, order="C", dtype=base.guess_dtype(data))

    # Validate shape
    if shape is None:
        if data is None:
            raise TypeError("Either data or shape must be specified")
        shape = data.shape
    else:
        shape = tuple(shape)
        if data is not None and (numpy.product(shape) != numpy.product(data.shape)):
            raise ValueError("Shape tuple is incompatible with data")

    tmp_shape = maxshape if maxshape is not None else shape
    # Validate chunk shape
    if isinstance(chunks, tuple) and (-numpy.array([ i>=j for i,j in zip(tmp_shape,chunks) if i is not None])).any():
        errmsg = "Chunk shape must not be greater than data shape in any dimension. "\
                 "{} is not compatible with {}".format(chunks, shape)
        raise ValueError(errmsg)

    if isinstance(dtype, Datatype):
        # Named types are used as-is
        tid = dtype.id
        dtype = tid.dtype  # Following code needs this
    else:
        # Validate dtype
        if dtype is None and data is None:
            dtype = numpy.dtype("=f4")
        elif dtype is None and data is not None:
            dtype = data.dtype
        else:
            dtype = numpy.dtype(dtype)
        tid = h5t.py_create(dtype, logical=1)

    # Legacy
    if any((compression, shuffle, fletcher32, maxshape,scaleoffset)) and chunks is False:
        raise ValueError("Chunked format required for given storage options")

    # Legacy
    if compression is True:
        if compression_opts is None:
            compression_opts = 4
        compression = 'gzip'

    # Legacy
    if compression in _LEGACY_GZIP_COMPRESSION_VALS:
        if compression_opts is not None:
            raise TypeError("Conflict in compression options")
        compression_opts = compression
        compression = 'gzip'

    dcpl = filters.generate_dcpl(shape, dtype, chunks, compression, compression_opts,
                  shuffle, fletcher32, maxshape, scaleoffset)

    if fillvalue is not None:
        fillvalue = numpy.array(fillvalue)
        dcpl.set_fill_value(fillvalue)

    if track_times in (True, False):
        dcpl.set_obj_track_times(track_times)
    elif track_times is not None:
        raise TypeError("track_times must be either True or False")

    if maxshape is not None:
        maxshape = tuple(m if m is not None else h5s.UNLIMITED for m in maxshape)
    sid = h5s.create_simple(shape, maxshape)


    dset_id = h5d.create(parent.id, None, tid, sid, dcpl=dcpl)

    if data is not None:
        dset_id.write(h5s.ALL, h5s.ALL, data)

    return dset_id
//  [python]
def make_new_dset(parent, shape=None, dtype=None, data=None,
                 chunks=None, compression=None, shuffle=None,
                    fletcher32=None, maxshape=None, compression_opts=None,
                  fillvalue=None, scaleoffset=None, track_times=None):
    """ Return a new low-level dataset identifier

    Only creates anonymous datasets.
    """

    # Convert data to a C-contiguous ndarray
    if data is not None:
        from . import base
        data = numpy.asarray(data, order="C", dtype=base.guess_dtype(data))

    # Validate shape
    if shape is None:
        if data is None:
            raise TypeError("Either data or shape must be specified")
        shape = data.shape
    else:
        shape = tuple(shape)
        if data is not None and (numpy.prod(shape) != numpy.prod(data.shape)):
            raise ValueError("Shape tuple is incompatible with data")

    tmp_shape = maxshape if maxshape is not None else shape
    # Validate chunk shape
    if isinstance(chunks, tuple) and (-numpy.array([ i>=j for i,j in zip(tmp_shape,chunks) if i is not None])).any():
        errmsg = "Chunk shape must not be greater than data shape in any dimension. "\
                 "{} is not compatible with {}".format(chunks, shape)
        raise ValueError(errmsg)

    if isinstance(dtype, Datatype):
        # Named types are used as-is
        tid = dtype.id
        dtype = tid.dtype  # Following code needs this
    else:
        # Validate dtype
        if dtype is None and data is None:
            dtype = numpy.dtype("=f4")
        elif dtype is None and data is not None:
            dtype = data.dtype
        else:
            dtype = numpy.dtype(dtype)
        tid = h5t.py_create(dtype, logical=1)

    # Legacy
    if any((compression, shuffle, fletcher32, maxshape,scaleoffset)) and chunks is False:
        raise ValueError("Chunked format required for given storage options")

    # Legacy
    if compression is True:
        if compression_opts is None:
            compression_opts = 4
        compression = 'gzip'

    # Legacy
    if compression in _LEGACY_GZIP_COMPRESSION_VALS:
        if compression_opts is not None:
            raise TypeError("Conflict in compression options")
        compression_opts = compression
        compression = 'gzip'

    dcpl = filters.generate_dcpl(shape, dtype, chunks, compression, compression_opts,
                  shuffle, fletcher32, maxshape, scaleoffset)

    if fillvalue is not None:
        fillvalue = numpy.array(fillvalue)
        dcpl.set_fill_value(fillvalue)

    if track_times in (True, False):
        dcpl.set_obj_track_times(track_times)
    elif track_times is not None:
        raise TypeError("track_times must be either True or False")

    if maxshape is not None:
        maxshape = tuple(m if m is not None else h5s.UNLIMITED for m in maxshape)
    sid = h5s.create_simple(shape, maxshape)


    dset_id = h5d.create(parent.id, None, tid, sid, dcpl=dcpl)

    if data is not None:
        dset_id.write(h5s.ALL, h5s.ALL, data)

    return dset_id
//  [python]
async def _auto_color(self, ctx, url: str, ranks):
        phrases = ["Calculating colors..."]  # in case I want more
        await ctx.send("{}".format(random.choice(phrases)))
        clusters = 10

        async with self.session.get(url) as r:
            image = await r.content.read()
        image = BytesIO(image)

        im = Image.open(image).convert("RGBA")
        im = im.resize((290, 290))  # resized to reduce time
        ar = numpy.asarray(im)
        shape = ar.shape
        ar = ar.reshape(numpy.product(shape[:2]), shape[2])

        codes, dist = cluster.vq.kmeans(ar.astype(float), clusters)
        vecs, dist = cluster.vq.vq(ar, codes)  # assign codes
        counts, bins = numpy.histogram(vecs, len(codes))  # count occurrences

        # sort counts
        freq_index = []
        index = 0
        for count in counts:
            freq_index.append((index, count))
            index += 1
        sorted_list = sorted(freq_index, key=operator.itemgetter(1), reverse=True)

        colors = []
        for rank in ranks:
            color_index = min(rank, len(codes))
            peak = codes[sorted_list[color_index][0]]  # gets the original index
            peak = peak.astype(int)

            colors.append("".join(format(c, "02x") for c in peak))
        image.close()
        im.close()
        return colors
//  [python]
async def _auto_color(self, ctx, url: str, ranks):
        phrases = ["Calculating colors..."]  # in case I want more
        await ctx.send("{}".format(random.choice(phrases)))
        clusters = 10

        async with self.session.get(url) as r:
            image = await r.content.read()
        image = BytesIO(image)

        im = Image.open(image).convert("RGBA")
        im = im.resize((290, 290))  # resized to reduce time
        ar = numpy.asarray(im)
        shape = ar.shape
        ar = ar.reshape(numpy.prod(shape[:2]), shape[2])

        codes, dist = cluster.vq.kmeans(ar.astype(float), clusters)
        vecs, dist = cluster.vq.vq(ar, codes)  # assign codes
        counts, bins = numpy.histogram(vecs, len(codes))  # count occurrences

        # sort counts
        freq_index = []
        index = 0
        for count in counts:
            freq_index.append((index, count))
            index += 1
        sorted_list = sorted(freq_index, key=operator.itemgetter(1), reverse=True)

        colors = []
        for rank in ranks:
            color_index = min(rank, len(codes))
            peak = codes[sorted_list[color_index][0]]  # gets the original index
            peak = peak.astype(int)

            colors.append("".join(format(c, "02x") for c in peak))
        image.close()
        im.close()
        return colors
//  [python]
def apply_layer_for_single_spatial_location(self, spatial_location, dependencies_values, output_index=0):
        flat_input = tensorflow.reshape(tensorflow.stack(dependencies_values, axis=1),
                                        shape=[-1, self.layer.input_shape[-1] * numpy.product(self.layer.kernel_size)])
        results = tensorflow.matmul(flat_input, self.reshaped_weights)
        if self.layer.use_bias:
            results = tensorflow.nn.bias_add(results, self.layer.bias)
        if self.layer.activation is not None:
            results = self.layer.activation(results)
        return results
//  [python]
def apply_layer_for_single_spatial_location(self, spatial_location, dependencies_values, output_index=0):
        flat_input = tensorflow.reshape(tensorflow.stack(dependencies_values, axis=1),
                                        shape=[-1, self.layer.input_shape[-1] * numpy.prod(self.layer.kernel_size)])
        results = tensorflow.matmul(flat_input, self.reshaped_weights)
        if self.layer.use_bias:
            results = tensorflow.nn.bias_add(results, self.layer.bias)
        if self.layer.activation is not None:
            results = self.layer.activation(results)
        return results
//  [python]
def check_against_numeric_one_step(self, args, index, out_index):
        if isinstance(args, ParameterTuple):
            x = args[index].data.asnumpy()
        else:
            x = args[index]
        x_shape = x.shape
        x_size = np.product(x_shape)
        for row in self.sampling(list(range(x_size))):
            original = x.ravel().view()[row]
            x.ravel().view()[row] += self.delta
            y_pos = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            x.ravel().view()[row] -= self.delta
            y_neg = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            diff = (y_pos - y_neg) / self.scale
            numeric_grad = diff.sum()
            insert_virtual_grad = False
            if numeric_grad == 0 and not insert_virtual_grad:
                self.gout.insert(0, 0)
                insert_virtual_grad = True
                continue
            theoretical_grad = self.gout[index].ravel().view()[row]

            if np.fabs(numeric_grad - theoretical_grad).max() > self.max_error:
                raise Exception(f'Gradients of df{out_index}/darg{index},{row} do not match, '
                                f'expect {numeric_grad}, actual {theoretical_grad}')

            print(f'GradChecker.check_against_numeric.numeric df{out_index}/darg{index}: '
                  f'{numeric_grad}, theoretical: {theoretical_grad}')
//  [python]
def check_against_numeric_one_step(self, args, index, out_index):
        if isinstance(args, ParameterTuple):
            x = args[index].data.asnumpy()
        else:
            x = args[index]
        x_shape = x.shape
        x_size = np.prod(x_shape)
        for row in self.sampling(list(range(x_size))):
            original = x.ravel().view()[row]
            x.ravel().view()[row] += self.delta
            y_pos = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            x.ravel().view()[row] -= self.delta
            y_neg = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            diff = (y_pos - y_neg) / self.scale
            numeric_grad = diff.sum()
            insert_virtual_grad = False
            if numeric_grad == 0 and not insert_virtual_grad:
                self.gout.insert(0, 0)
                insert_virtual_grad = True
                continue
            theoretical_grad = self.gout[index].ravel().view()[row]

            if np.fabs(numeric_grad - theoretical_grad).max() > self.max_error:
                raise Exception(f'Gradients of df{out_index}/darg{index},{row} do not match, '
                                f'expect {numeric_grad}, actual {theoretical_grad}')

            print(f'GradChecker.check_against_numeric.numeric df{out_index}/darg{index}: '
                  f'{numeric_grad}, theoretical: {theoretical_grad}')
//  [python]
def r_box_norm(psi, a):
    '''
    y = (sum_i,j psi_i,j * dx*dy) = (sum_i,j psi_i,j) * A / (nx * ny)
    psi_i,j = psi(x, y)
    A = nx*dx * ny*dy
    (dy, dx) = grid step size
    (ny, nx) = grid size

    Parameters
    ----------
        psi : numpy array, shape=(ny, nx), dtype=complex
            wave function
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        norm : float = y
    '''
    ndim = np.array(psi.shape)
    nfac = np.product(a) / np.product(ndim)
    return np.sum(psi.real**2 + psi.imag**2) * nfac
//  [python]
def r_box_norm(psi, a):
    '''
    y = (sum_i,j psi_i,j * dx*dy) = (sum_i,j psi_i,j) * A / (nx * ny)
    psi_i,j = psi(x, y)
    A = nx*dx * ny*dy
    (dy, dx) = grid step size
    (ny, nx) = grid size

    Parameters
    ----------
        psi : numpy array, shape=(ny, nx), dtype=complex
            wave function
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        norm : float = y
    '''
    ndim = np.array(psi.shape)
    nfac = np.prod(a) / np.prod(ndim)
    return np.sum(psi.real**2 + psi.imag**2) * nfac
//  [python]
def q_box_norm(psi, a):
    '''
    y = (sum_i,j psi_i,j * dqy*dqy) = (sum_i,j psi_i,j) / A
    psi_i,j = psi(qx, qy)
    A = nx*dx * ny*dy = 1 / (dqx*dqy)
    (dqy, dqx) = grid step size in reciprocal space
    (dy, dx) = grid step size in real space
    (ny, nx) = grid size

    Parameters
    ----------
        psi : numpy array, shape=(ny, nx), dtype=complex
            wave function
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        norm : float = y
    '''
    nfac = 1.0 / np.product(a)
    return np.sum(psi.real**2 + psi.imag**2) * nfac
//  [python]
def q_box_norm(psi, a):
    '''
    y = (sum_i,j psi_i,j * dqy*dqy) = (sum_i,j psi_i,j) / A
    psi_i,j = psi(qx, qy)
    A = nx*dx * ny*dy = 1 / (dqx*dqy)
    (dqy, dqx) = grid step size in reciprocal space
    (dy, dx) = grid step size in real space
    (ny, nx) = grid size

    Parameters
    ----------
        psi : numpy array, shape=(ny, nx), dtype=complex
            wave function
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        norm : float = y
    '''
    nfac = 1.0 / np.prod(a)
    return np.sum(psi.real**2 + psi.imag**2) * nfac
//  [python]
def check_against_numeric_jacobian_one_step(self, args, index, out_index):
        if isinstance(args, ParameterTuple):
            x = args[index].data.asnumpy()
        else:
            x = args[index]
        x_shape = x.shape
        x_size = np.product(x_shape)
        dy = self.to_numpy(self.get_sens(out_index))
        dy_size = np.product(dy.shape)
        numeric_jacobian = np.zeros((x_size, dy_size), dtype=self.to_numpy(x).dtype)
        for row in range(x_size):
            original = x.ravel().view()[row]
            x.ravel().view()[row] += self.delta
            y_pos = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            x.ravel().view()[row] -= self.delta
            y_neg = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            diff = (y_pos - y_neg) / self.scale
            numeric_jacobian[row, :] = diff.ravel().view(numeric_jacobian.dtype)

        dy_mask = np.zeros(dy.shape, dtype=dy.dtype)
        theoretical_jacobian = np.zeros((x_size, dy_size), dtype=self.to_numpy(x).dtype)
        for col in range(dy_size):
            col_jacobian = self.compute_theoretical_jacobian(index, out_index, dy_mask, col)
            theoretical_jacobian[:, col] = col_jacobian.ravel().view(theoretical_jacobian.dtype)

        if np.fabs(numeric_jacobian - theoretical_jacobian).max() > self.max_error:
            raise Exception(f'GradChecker.check_against_numeric_jacobian_one_step expect {out_index}/darg{index}: '
                            f'{numeric_jacobian}, actual: {theoretical_jacobian}')

        print(f'GradChecker.check_against_numeric_jacobian_one_step.numeric jacobian of output{out_index}/darg{index}: '
              f'{numeric_jacobian}, theoretical: {theoretical_jacobian}')
//  [python]
def check_against_numeric_jacobian_one_step(self, args, index, out_index):
        if isinstance(args, ParameterTuple):
            x = args[index].data.asnumpy()
        else:
            x = args[index]
        x_shape = x.shape
        x_size = np.prod(x_shape)
        dy = self.to_numpy(self.get_sens(out_index))
        dy_size = np.prod(dy.shape)
        numeric_jacobian = np.zeros((x_size, dy_size), dtype=self.to_numpy(x).dtype)
        for row in range(x_size):
            original = x.ravel().view()[row]
            x.ravel().view()[row] += self.delta
            y_pos = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            x.ravel().view()[row] -= self.delta
            y_neg = self.to_numpy_and_scale(self.get_ith_elem(self.fn(*self.args), out_index))
            x.ravel().view()[row] = original
            diff = (y_pos - y_neg) / self.scale
            numeric_jacobian[row, :] = diff.ravel().view(numeric_jacobian.dtype)

        dy_mask = np.zeros(dy.shape, dtype=dy.dtype)
        theoretical_jacobian = np.zeros((x_size, dy_size), dtype=self.to_numpy(x).dtype)
        for col in range(dy_size):
            col_jacobian = self.compute_theoretical_jacobian(index, out_index, dy_mask, col)
            theoretical_jacobian[:, col] = col_jacobian.ravel().view(theoretical_jacobian.dtype)

        if np.fabs(numeric_jacobian - theoretical_jacobian).max() > self.max_error:
            raise Exception(f'GradChecker.check_against_numeric_jacobian_one_step expect {out_index}/darg{index}: '
                            f'{numeric_jacobian}, actual: {theoretical_jacobian}')

        print(f'GradChecker.check_against_numeric_jacobian_one_step.numeric jacobian of output{out_index}/darg{index}: '
              f'{numeric_jacobian}, theoretical: {theoretical_jacobian}')
//  [python]
def norm_fft2(psi, a):
    '''
    Norm preserving 2-d FFT

    Can handle stacks of 2d Fourier transforms.

    Parameters
    ----------
        psi : numpy array, shape=(, ny, nx), dtype=complex
            wave function in real space
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        numpy array, shape=(ny, nx), dytpe=complex
            wave function in Fourier space
    '''
    nfac = np.product(a) / np.product(np.array(psi.shape[-2:]))
    return np.fft.fft2(psi) * nfac
//  [python]
def norm_fft2(psi, a):
    '''
    Norm preserving 2-d FFT

    Can handle stacks of 2d Fourier transforms.

    Parameters
    ----------
        psi : numpy array, shape=(, ny, nx), dtype=complex
            wave function in real space
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        numpy array, shape=(ny, nx), dytpe=complex
            wave function in Fourier space
    '''
    nfac = np.prod(a) / np.prod(np.array(psi.shape[-2:]))
    return np.fft.fft2(psi) * nfac
//  [python]
def norm_ifft2(psi, a):
    '''
    Norm preserving 2-d inverse FFT

    Can handle stacks of 2d inverse Fourier transforms.

    Parameters
    ----------
        psi : numpy array, shape=(, ny, nx), dtype=complex
            wave function in Fourier space
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        numpy array, shape=(ny, nx), dytpe=complex
            wave function in real space
    '''
    nfac = np.product(np.array(psi.shape[-2:])) / np.product(a)
    return np.fft.ifft2(psi) * nfac
//  [python]
def norm_ifft2(psi, a):
    '''
    Norm preserving 2-d inverse FFT

    Can handle stacks of 2d inverse Fourier transforms.

    Parameters
    ----------
        psi : numpy array, shape=(, ny, nx), dtype=complex
            wave function in Fourier space
        a : numpy array, shape=(2), dtype=float-type
            box size (y,x)

    Returns
    -------
        numpy array, shape=(ny, nx), dytpe=complex
            wave function in real space
    '''
    nfac = np.prod(np.array(psi.shape[-2:])) / np.prod(a)
    return np.fft.ifft2(psi) * nfac
//  [python]
def average(self):
        """Return the average of the data."""
        n = self.grid.comm.reduce(numpy.sum(self))
        if self.grid.comm.rank == 0:
            n /= numpy.product(self.grid.pdims)
        return n
//  [python]
def average(self):
        """Return the average of the data."""
        n = self.grid.comm.reduce(numpy.sum(self))
        if self.grid.comm.rank == 0:
            n /= numpy.prod(self.grid.pdims)
        return n
//  [python]
def norm(self):
        """Return the :math:`L_2` norm of the data."""
        n = self.grid.comm.reduce(numpy.sum(self*self))
        if self.grid.comm.rank == 0:
            n /= numpy.product(self.grid.pdims)
        return n
//  [python]
def norm(self):
        """Return the :math:`L_2` norm of the data."""
        n = self.grid.comm.reduce(numpy.sum(self*self))
        if self.grid.comm.rank == 0:
            n /= numpy.prod(self.grid.pdims)
        return n
//  [python]
def processDataArrays(index, params, data):
    """Returns a list of arrays corresponding to the ``count`` of
    ``index`` into ``datadescs`` of the current params"""
    datadesc = params['datadescs'][index]
    count = datadesc.get('count', 1)
    shape = datadesc['shape']

    # determine 1D array size
    arraysize = numpy.product(shape)
    arrays = numpy.split(data[:count * arraysize], count)

    # reshape every array in the list
    for i, array in enumerate(arrays):
        arrays[i] = array.reshape(shape)
    return arrays
//  [python]
def processDataArrays(index, params, data):
    """Returns a list of arrays corresponding to the ``count`` of
    ``index`` into ``datadescs`` of the current params"""
    datadesc = params['datadescs'][index]
    count = datadesc.get('count', 1)
    shape = datadesc['shape']

    # determine 1D array size
    arraysize = numpy.prod(shape)
    arrays = numpy.split(data[:count * arraysize], count)

    # reshape every array in the list
    for i, array in enumerate(arrays):
        arrays[i] = array.reshape(shape)
    return arrays