Python numpy.rint() Examples

def visualize_2D_trip(self,trip,tw_open,tw_close):
        plt.figure(figsize=(30,30))
        rcParams.update({'font.size': 22})
        # Plot cities
        colors = ['red'] # Depot is first city
        for i in range(len(tw_open)-1):
            colors.append('blue')
        plt.scatter(trip[:,0], trip[:,1], color=colors, s=200)
        # Plot tour
        tour=np.array(list(range(len(trip))) + [0])
        X = trip[tour, 0]
        Y = trip[tour, 1]
        plt.plot(X, Y,"--", markersize=100)
        # Annotate cities with TW
        tw_open = np.rint(tw_open)
        tw_close = np.rint(tw_close)
        time_window = np.concatenate((tw_open,tw_close),axis=1)
        for tw, (x, y) in zip(time_window,(zip(X,Y))):
            plt.annotate(tw,xy=(x, y))  
        plt.xlim(0,60)
        plt.ylim(0,60)
        plt.show()


    # Heatmap of permutations (x=cities; y=steps) 

def get_elevation(self, interpolation='nearest_neighbor'):
        assert interpolation in ('nearest_neighbor', 'bivariate')

        if self._elevation.get(interpolation, None) is None:
            self._log.debug('Getting elevation...')
            # region = llLon, urLon, llLat, urLon
            coords = self.frame.coordinates
            lons = coords[:, 0]
            lats = coords[:, 1]

            region = (lons.min(), lons.max(), lats.min(), lats.max())
            if not srtmgl3.covers(region):
                raise AssertionError(
                    'Region is outside of SRTMGL3 topo dataset')

            tile = srtmgl3.get(region)
            if not tile:
                raise AssertionError('Cannot get SRTMGL3 topo dataset')

            if interpolation == 'nearest_neighbor':
                iy = num.rint((lats - tile.ymin) / tile.dy).astype(num.intp)
                ix = num.rint((lons - tile.xmin) / tile.dx).astype(num.intp)

                elevation = tile.data[(iy, ix)]

            elif interpolation == 'bivariate':
                interp = interpolate.RectBivariateSpline(
                    tile.y(), tile.x(), tile.data)
                elevation = interp(lats, lons, grid=False)

            elevation = elevation.reshape(self.rows, self.cols)
            self._elevation[interpolation] = elevation

        return self._elevation[interpolation] 

def add_image(self, img):
        if self.print_progress and self.cur_images % self.progress_interval == 0:
            print('%d / %d\r' % (self.cur_images, self.expected_images), end='', flush=True)
            sys.stdout.flush()
        if self.shape is None:
            self.shape = img.shape
            self.resolution_log2 = int(np.log2(self.shape[1]))
            assert self.shape[0] in [1, 3]
            assert self.shape[1] == self.shape[2]
            assert self.shape[1] == 2**self.resolution_log2
            tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
            for lod in range(self.resolution_log2 - 1):
                tfr_file = self.tfr_prefix + '-r%02d.tfrecords' % (self.resolution_log2 - lod)
                self.tfr_writers.append(tf.python_io.TFRecordWriter(tfr_file, tfr_opt))
        assert img.shape == self.shape
        for lod, tfr_writer in enumerate(self.tfr_writers):
            if lod:
                img = img.astype(np.float32)
                img = (img[:, 0::2, 0::2] + img[:, 0::2, 1::2] + img[:, 1::2, 0::2] + img[:, 1::2, 1::2]) * 0.25
            quant = np.rint(img).clip(0, 255).astype(np.uint8)
            ex = tf.train.Example(features=tf.train.Features(feature={
                'shape': tf.train.Feature(int64_list=tf.train.Int64List(value=quant.shape)),
                'data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[quant.tostring()]))}))
            tfr_writer.write(ex.SerializeToString())
        self.cur_images += 1 

def apply_float_operation(problem, fun):

    # save the original bounds of the problem
    _xl, _xu = problem.xl, problem.xu

    # copy the arrays of the problem and cast them to float
    xl, xu = problem.xl.astype(np.float), problem.xu.astype(np.float)

    # modify the bounds to match the new crossover specifications and set the problem
    problem.xl = xl - (0.5 - 1e-16)
    problem.xu = xu + (0.5 - 1e-16)

    # perform the crossover
    off = fun()

    # now round to nearest integer for all offsprings
    off = np.rint(off).astype(np.int)

    # reset the original bounds of the problem and design space values
    problem.xl = _xl
    problem.xu = _xu

    return off 

def get_kconserv(cell, kpts):
    r'''Get the momentum conservation array for a set of k-points.

    Given k-point indices (k, l, m) the array kconserv[k,l,m] returns
    the index n that satifies momentum conservation,

        (k(k) - k(l) + k(m) - k(n)) \dot a = 2n\pi

    This is used for symmetry e.g. integrals of the form
        [\phi*[k](1) \phi[l](1) | \phi*[m](2) \phi[n](2)]
    are zero unless n satisfies the above.
    '''
    nkpts = kpts.shape[0]
    a = cell.lattice_vectors() / (2*np.pi)

    kconserv = np.zeros((nkpts,nkpts,nkpts), dtype=int)
    kvKLM = kpts[:,None,None,:] - kpts[:,None,:] + kpts
    for N, kvN in enumerate(kpts):
        kvKLMN = np.einsum('wx,klmx->wklm', a, kvKLM - kvN)
        # check whether (1/(2pi) k_{KLMN} dot a) is an integer
        kvKLMN_int = np.rint(kvKLMN)
        mask = np.einsum('wklm->klm', abs(kvKLMN - kvKLMN_int)) < 1e-9
        kconserv[mask] = N
    return kconserv 

def __init__(self, **kw):
    """  
      Constructor of affine, equidistant 3d mesh class
      ucell : unit cell vectors (in coordinate space)
      Ecut  : Energy cutoff to parametrize the discretization 
    """
    from scipy.fftpack import next_fast_len
    
    self.ucell = kw['ucell'] if 'ucell' in kw else 30.0*np.eye(3) # Not even unit cells vectors are required by default
    self.Ecut = Ecut = kw['Ecut'] if 'Ecut' in kw else 50.0 # 50.0 Hartree by default
    luc = np.sqrt(np.einsum('ix,ix->i', self.ucell, self.ucell))
    self.shape = nn = np.array([next_fast_len( int(np.rint(l * np.sqrt(Ecut)/2))) for l in luc], dtype=int)
    self.size  = np.prod(self.shape)
    gc = self.ucell/(nn) # This is probable the best for finite systems, for PBC use nn, not (nn-1)
    self.dv = np.abs(np.dot(gc[0], np.cross(gc[1], gc[2] )))
    rr = [np.array([gc[i]*j for j in range(nn[i])]) for i in range(3)]
    self.rr = rr
    self.origin = kw['origin'] if 'origin' in kw else np.zeros(3) 

def _hrv_get_rri(peaks=None, sampling_rate=1000, interpolate=False, **kwargs):

    rri = np.diff(peaks) / sampling_rate * 1000

    if interpolate is False:
        return rri

    else:

        # Minimum sampling rate for interpolation
        if sampling_rate < 10:
            sampling_rate = 10

        # Compute length of interpolated heart period signal at requested sampling rate.
        desired_length = int(np.rint(peaks[-1] / sampling_rate * sampling_rate))

        rri = signal_interpolate(
            peaks[1:],  # Skip first peak since it has no corresponding element in heart_period
            rri,
            x_new=np.arange(desired_length),
            **kwargs
        )
        return rri, sampling_rate 

def read_segments_as_bool_vec(segments_file):
  """ [ bool_vec ] = read_segments_as_bool_vec(segments_file)
   using kaldi 'segments' file for 1 wav, format : '<utt> <rec> <t-beg> <t-end>'
   - t-beg, t-end is in seconds,
   - assumed 100 frames/second,
  """
  segs = np.loadtxt(segments_file, dtype='object,object,f,f', ndmin=1)
  # Sanity checks,
  assert(len(segs) > 0) # empty segmentation is an error,
  assert(len(np.unique([rec[1] for rec in segs ])) == 1) # segments with only 1 wav-file,
  # Convert time to frame-indexes,
  start = np.rint([100 * rec[2] for rec in segs]).astype(int)
  end = np.rint([100 * rec[3] for rec in segs]).astype(int)
  # Taken from 'read_lab_to_bool_vec', htk.py,
  frms = np.repeat(np.r_[np.tile([False,True], len(end)), False],
                   np.r_[np.c_[start - np.r_[0, end[:-1]], end-start].flat, 0])
  assert np.sum(end-start) == np.sum(frms)
  return frms 

def get_C_hat_transpose():
    probs = []
    net.eval()
    for batch_idx, (data, target) in enumerate(train_gold_deterministic_loader):
        # we subtract 10 because we added 10 to gold so we could identify which example is gold in train_phase2
        data, target = torch.autograd.Variable(data.cuda(), volatile=True),\
                       torch.autograd.Variable((target - num_classes).cuda(), volatile=True)

        # forward
        output = net(data)
        pred = F.softmax(output)
        probs.extend(list(pred.data.cpu().numpy()))

    probs = np.array(probs, dtype=np.float32)
    preds = np.argmax(probs, axis=1)
    C_hat = np.zeros([num_classes, num_classes])
    for i in range(len(train_data_gold.train_labels)):
        C_hat[int(np.rint(train_data_gold.train_labels[i] - num_classes)), preds[i]] += 1

    C_hat /= (np.sum(C_hat, axis=1, keepdims=True) + 1e-7)
    C_hat = C_hat * 0.99 + np.full_like(C_hat, 1/num_classes) * 0.01  # smoothing

    return C_hat.T.astype(np.float32) 

def read_segments_as_bool_vec(segments_file):
  """ [ bool_vec ] = read_segments_as_bool_vec(segments_file)
   using kaldi 'segments' file for 1 wav, format : '<utt> <rec> <t-beg> <t-end>'
   - t-beg, t-end is in seconds,
   - assumed 100 frames/second,
  """
  segs = np.loadtxt(segments_file, dtype='object,object,f,f', ndmin=1)
  # Sanity checks,
  assert(len(segs) > 0) # empty segmentation is an error,
  assert(len(np.unique([rec[1] for rec in segs ])) == 1) # segments with only 1 wav-file,
  # Convert time to frame-indexes,
  start = np.rint([100 * rec[2] for rec in segs]).astype(int)
  end = np.rint([100 * rec[3] for rec in segs]).astype(int)
  # Taken from 'read_lab_to_bool_vec', htk.py,
  frms = np.repeat(np.r_[np.tile([False,True], len(end)), False],
                   np.r_[np.c_[start - np.r_[0, end[:-1]], end-start].flat, 0])
  assert np.sum(end-start) == np.sum(frms)
  return frms 

def extract_segments(wavs, segments=None):
    """ This function returns generator of segmented audio as
        (utterance id, numpy.float32 array)
        TODO?: sampling rate is not converted.
    """
    if segments is not None:
        # segments should be sorted by rec-id
        for seg in segments:
            wav = wavs[seg['rec']]
            data, samplerate = load_wav(wav)
            st_sample = np.rint(seg['st'] * samplerate).astype(int)
            et_sample = np.rint(seg['et'] * samplerate).astype(int)
            yield seg['utt'], data[st_sample:et_sample]
    else:
        # segments file not found,
        # wav.scp is used as segmented audio list
        for rec in wavs:
            data, samplerate = load_wav(wavs[rec])
            yield rec, data 

def __init__(self, matrix_size, input_size, low=0, high=1, round_values=False):
        """Init function.

        @param matrix_size It defines the matrix size 
        @param input_size it defines the vector input size.
        @param low boundary for the random initialization 
        @param high boundary for the random initialization
        @param round_values it is possible to initialize the 
        weights to the closest integer value.
        """
        self._matrix_size = matrix_size
        self._input_size = input_size
        self._weights_matrix = np.random.uniform(low=low, high=high, size=(matrix_size, matrix_size, input_size))

        if (round_values == True):
            self._weights_matrix = np.rint(self._weights_matrix) 

def setUp(self):
        M = 4
        N = 10000
        seed= 0
        sampling_frequency = 30000
        X = np.random.RandomState(seed=seed).normal(0, 1, (M, N))
        geom = np.random.RandomState(seed=seed).normal(0, 1, (M, 2))
        self._X = X
        self._geom = geom
        self._sampling_frequency = sampling_frequency
        self.RX = se.NumpyRecordingExtractor(timeseries=X, sampling_frequency=sampling_frequency, geom=geom)
        self.SX = se.NumpySortingExtractor()
        L = 200
        self._train1 = np.rint(np.random.RandomState(seed=seed).uniform(0, N, L)).astype(int)
        self.SX.add_unit(unit_id=1, times=self._train1)
        self.SX.add_unit(unit_id=2, times=np.random.RandomState(seed=seed).uniform(0, N, L))
        self.SX.add_unit(unit_id=3, times=np.random.RandomState(seed=seed).uniform(0, N, L)) 

def __init__(self, file_path: PathType):
        super().__init__(file_path)
        cluster_classes = self._getfield("cluster_class")
        classes = cluster_classes[:, 0]
        spike_times = cluster_classes[:, 1]
        par = self._getfield("par")
        sample_rate = par[0, 0][np.where(np.array(par.dtype.names) == 'sr')[0][0]][0][0]

        self.set_sampling_frequency(sample_rate)
        self._unit_ids = np.unique(classes[classes > 0]).astype('int')

        self._spike_trains = {}
        for uid in self._unit_ids:
            mask = (classes == uid)
            self._spike_trains[uid] = np.rint(spike_times[mask]*(sample_rate/1000))
        self._unsorted_train = np.rint(spike_times[classes == 0] * (sample_rate / 1000)) 

def optimize_model(task, param_name, test_size: float, binary=False) -> None:
    x, y = task.create_train_data()

    def objective(trial):
        train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=test_size)
        param = redshells.factory.get_optuna_param(param_name, trial)
        model = task.create_model()
        model.set_params(**param)
        model.fit(train_x, train_y)
        predictions = model.predict(test_x)

        if binary:
            predictions = np.rint(predictions)

        return 1.0 - sklearn.metrics.accuracy_score(test_y, predictions)

    study = optuna.create_study()
    study.optimize(objective, n_trials=100)
    task.dump(dict(best_params=study.best_params, best_value=study.best_value)) 

def get_k_mesh_by_cell(self, kpoints_per_reciprocal_angstrom, cell=None):
        """
            get k-mesh density according to the box size.

            Args:
                kpoints_per_reciprocal_angstrom: (float) number of k-points per reciprocal angstrom (i.e. per 2*pi*box_length)
                cell: (list/ndarray) 3x3 cell. If not set, the current cell is used.
        """
        if cell is None:
            if self.structure is None:
                raise AssertionError('structure not set')
            cell = self.structure.cell
        latlens = np.linalg.norm(cell, axis=-1)
        kmesh = np.rint( 2 * np.pi / latlens * kpoints_per_reciprocal_angstrom)
        if kmesh.min() <= 0:
            raise AssertionError("kpoint per angstrom too low")
        return [int(k) for k in kmesh] 

def BGR2YCbCr(im):
    mat = np.array([[24.966, 128.553, 65.481],[112, -74.203, -37.797], [-18.214, -93.786, 112]])
    mat = mat.T
    offset = np.array([[[16, 128, 128]]])
    if im.dtype == 'uint8':
        mat = mat/255
        out = np.dot(im,mat) + offset
        out = np.clip(out, 0, 255)
        out = np.rint(out).astype('uint8')
    elif im.dtype == 'float':
        mat = mat/255
        offset = offset/255
        out = np.dot(im, mat) + offset
        out = np.clip(out, 0, 1)
    else:
        assert False
    return out 

def YCbCr2BGR(im):
    mat = np.array([[24.966, 128.553, 65.481],[112, -74.203, -37.797], [-18.214, -93.786, 112]])
    mat = mat.T
    mat = np.linalg.inv(mat)
    offset = np.array([[[16, 128, 128]]])
    if im.dtype == 'uint8':
        mat = mat * 255
        out = np.dot((im - offset),mat)
        out = np.clip(out, 0, 255)
        out = np.rint(out).astype('uint8')
    elif im.dtype == 'float':
        mat = mat * 255
        offset = offset/255
        out = np.dot((im - offset),mat)
        out = np.clip(out, 0, 1)
    else:
        assert False
    return out 

def visualize_voxel_spectral(points, vis_size=128):
  """Function to visualize voxel (spectral)."""
  points = np.rint(points)
  points = np.swapaxes(points, 0, 2)
  fig = p.figure(figsize=(1, 1), dpi=vis_size)
  verts, faces = measure.marching_cubes_classic(points, 0, spacing=(0.1, 0.1, 0.1))
  ax = fig.add_subplot(111, projection='3d')
  ax.plot_trisurf(
      verts[:, 0], verts[:, 1], faces, verts[:, 2], cmap='Spectral_r', lw=0.1)
  ax.set_axis_off()
  fig.tight_layout(pad=0)
  fig.canvas.draw()
  data = np.fromstring(
      fig.canvas.tostring_rgb(), dtype=np.uint8, sep='').reshape(
          vis_size, vis_size, 3)
  p.close('all')
  return data 

def _get_centerline_lonlat(gdir):
    """Quick n dirty solution to write the centerlines as a shapefile"""

    cls = gdir.read_pickle('centerlines')
    olist = []
    for j, cl in enumerate(cls[::-1]):
        mm = 1 if j == 0 else 0
        gs = dict()
        gs['RGIID'] = gdir.rgi_id
        gs['LE_SEGMENT'] = np.rint(np.max(cl.dis_on_line) * gdir.grid.dx)
        gs['MAIN'] = mm
        tra_func = partial(gdir.grid.ij_to_crs, crs=wgs84)
        gs['geometry'] = shp_trafo(tra_func, cl.line)
        olist.append(gs)

    return olist 

def read_svgcoords(svg_file):
    """Get the vertices coordinates out of a SVG file"""

    from xml.dom import minidom
    doc = minidom.parse(svg_file)
    coords = [path.getAttribute('d') for path
              in doc.getElementsByTagName('path')]
    doc.unlink()
    _, _, coords = coords[0].partition('C')
    x = []
    y = []
    for c in coords.split(' '):
        if c == '':
            continue
        c = c.split(',')
        x.append(np.float(c[0]))
        y.append(np.float(c[1]))
    x.append(x[0])
    y.append(y[0])

    return np.rint(np.asarray((x, y)).T).astype(np.int64) 

def convert_to_pil_image(image, drange=[0,1]):
    assert image.ndim == 2 or image.ndim == 3
    if image.ndim == 3:
        if image.shape[0] == 1:
            image = image[0] # grayscale CHW => HW
        else:
            image = image.transpose(1, 2, 0) # CHW -> HWC

    image = adjust_dynamic_range(image, drange, [0,255])
    image = np.rint(image).clip(0, 255).astype(np.uint8)
    format = 'RGB' if image.ndim == 3 else 'L'
    return PIL.Image.fromarray(image, format) 

def format_time(seconds):
    s = int(np.rint(seconds))
    if s < 60:         return '%ds'                % (s)
    elif s < 60*60:    return '%dm %02ds'          % (s // 60, s % 60)
    elif s < 24*60*60: return '%dh %02dm %02ds'    % (s // (60*60), (s // 60) % 60, s % 60)
    else:              return '%dd %02dh %02dm'    % (s // (24*60*60), (s // (60*60)) % 24, (s // 60) % 60)

#----------------------------------------------------------------------------
# Locating results. 

def draw_text_label(img, text, x, y, alignx=0.5, aligny=0.5, color=255, opacity=1.0, glow_opacity=1.0, **kwargs):
    color = np.array(color).flatten().astype(np.float32)
    assert img.ndim == 3 and img.shape[2] == color.size or color.size == 1
    alpha, glow = setup_text_label(text, **kwargs)
    xx, yy = int(np.rint(x - alpha.shape[1] * alignx)), int(np.rint(y - alpha.shape[0] * aligny))
    xb, yb = max(-xx, 0), max(-yy, 0)
    xe, ye = min(alpha.shape[1], img.shape[1] - xx), min(alpha.shape[0], img.shape[0] - yy)
    img = np.array(img)
    slice = img[yy+yb : yy+ye, xx+xb : xx+xe, :]
    slice[:] = slice * (1.0 - (1.0 - (1.0 - alpha[yb:ye, xb:xe]) * (1.0 - glow[yb:ye, xb:xe] * glow_opacity)) * opacity)[:, :, np.newaxis]
    slice[:] = slice + alpha[yb:ye, xb:xe, np.newaxis] * (color * opacity)[np.newaxis, np.newaxis, :]
    return img 

def generate_interpolation_video(run_id, snapshot=None, grid_size=[1,1], image_shrink=1, image_zoom=1, duration_sec=60.0, smoothing_sec=1.0, mp4=None, mp4_fps=30, mp4_codec='libx265', mp4_bitrate='16M', random_seed=1000, minibatch_size=8):
    network_pkl = misc.locate_network_pkl(run_id, snapshot)
    if mp4 is None:
        mp4 = misc.get_id_string_for_network_pkl(network_pkl) + '-lerp.mp4'
    num_frames = int(np.rint(duration_sec * mp4_fps))
    random_state = np.random.RandomState(random_seed)

    print('Loading network from "%s"...' % network_pkl)
    G, D, Gs = misc.load_network_pkl(run_id, snapshot)

    print('Generating latent vectors...')
    shape = [num_frames, np.prod(grid_size)] + Gs.input_shape[1:] # [frame, image, channel, component]
    all_latents = random_state.randn(*shape).astype(np.float32)
    all_latents = scipy.ndimage.gaussian_filter(all_latents, [smoothing_sec * mp4_fps] + [0] * len(Gs.input_shape), mode='wrap')
    all_latents /= np.sqrt(np.mean(np.square(all_latents)))

    # Frame generation func for moviepy.
    def make_frame(t):
        frame_idx = int(np.clip(np.round(t * mp4_fps), 0, num_frames - 1))
        latents = all_latents[frame_idx]
        labels = np.zeros([latents.shape[0], 0], np.float32)
        images = Gs.run(latents, labels, minibatch_size=minibatch_size, num_gpus=config.num_gpus, out_mul=127.5, out_add=127.5, out_shrink=image_shrink, out_dtype=np.uint8)
        grid = misc.create_image_grid(images, grid_size).transpose(1, 2, 0) # HWC
        if image_zoom > 1:
            grid = scipy.ndimage.zoom(grid, [image_zoom, image_zoom, 1], order=0)
        if grid.shape[2] == 1:
            grid = grid.repeat(3, 2) # grayscale => RGB
        return grid

    # Generate video.
    import moviepy.editor # pip install moviepy
    result_subdir = misc.create_result_subdir(config.result_dir, config.desc)
    moviepy.editor.VideoClip(make_frame, duration=duration_sec).write_videofile(os.path.join(result_subdir, mp4), fps=mp4_fps, codec='libx264', bitrate=mp4_bitrate)
    open(os.path.join(result_subdir, '_done.txt'), 'wt').close()

#----------------------------------------------------------------------------
# Generate MP4 video of training progress for a previous training run.
# To run, uncomment the appropriate line in config.py and launch train.py. 

def test_multinomial_generator():
    # This test fails with dtype float16 if the probabilities themselves cannot be
    # well-represented in float16.  When the float16 random picks are assigned to buckets,
    # only certain bucket-probabilities are possible.  Here we map the desired probabilites
    # (e.g. 0.1) to nearby float16 probabilities (e.g. 0.10009766) that are achievable.
    def quantize_probs(probs, dtype):
        if dtype == 'float16':
            # float16 has a 10-bit fraction plus an implicit leading 1, so all probabilities
            # of the form N/2^11 (where N is an integer) are representable.
            num_quanta = 2048.0
            quantized_probs = np.rint(np.array(probs) * num_quanta) / num_quanta
            # Ensure probabilities add to 1
            quantized_probs[0] += 1.0 - quantized_probs.sum()
        else:
            # no need to quantize probs with this data precision
            quantized_probs = np.array(probs)
        return quantized_probs

    ctx = mx.context.current_context()
    probs = [0.1, 0.2, 0.3, 0.05, 0.15, 0.2]
    samples = 1000000
    trials = 5
    buckets = list(range(6))
    for dtype in ['float16', 'float32', 'float64']:
        print("ctx=%s, dtype=%s" %(ctx, dtype))
        quantized_probs = quantize_probs(probs, dtype)
        generator_mx = lambda x: mx.nd.random.multinomial(data=mx.nd.array(quantized_probs, ctx=ctx, dtype=dtype),
                                                          shape=x).asnumpy()
        verify_generator(generator=generator_mx, buckets=buckets, probs=quantized_probs,
                         nsamples=samples, nrepeat=trials)
        generator_mx_same_seed = \
            lambda x: np.concatenate(
                [mx.nd.random.multinomial(data=mx.nd.array(quantized_probs, ctx=ctx, dtype=dtype),
                                                          shape=x // 10).asnumpy()
                 for _ in range(10)])
        verify_generator(generator=generator_mx_same_seed, buckets=buckets, probs=quantized_probs,
                         nsamples=samples, nrepeat=trials) 

def get_kconserv_ee_r2(self, kshift=0):
        r'''Get the momentum conservation array for a set of k-points.

        Given k-point indices (k, l, m) the array kconserv_r2[k,l,m] returns
        the index n that satisfies momentum conservation,

            (k(k) - k(l) + k(m) - k(n) - kshift) \dot a = 2n\pi

        This is used for symmetry of 2p-2h excitation operator vector
        R_{k k_k, m k_m}^{l k_l n k_n} is zero unless n satisfies the above.

        Note that this method is adapted from `kpts_helper.get_kconserv()`.
        '''
        cell = self._cc._scf.cell
        kpts = self.kpts
        nkpts = kpts.shape[0]
        a = cell.lattice_vectors() / (2 * np.pi)

        kconserv_r2 = np.zeros((nkpts, nkpts, nkpts), dtype=int)
        kvKLM = kpts[:, None, None, :] - kpts[:, None, :] + kpts
        # Apply k shift
        kvKLM = kvKLM - kpts[kshift]
        for N, kvN in enumerate(kpts):
            kvKLMN = np.einsum('wx,klmx->wklm', a, kvKLM - kvN)
            # check whether (1/(2pi) k_{KLMN} dot a) is an integer
            kvKLMN_int = np.rint(kvKLMN)
            mask = np.einsum('wklm->klm', abs(kvKLMN - kvKLMN_int)) < 1e-9
            kconserv_r2[mask] = N
        return kconserv_r2 

def get_kconserv3(cell, kpts, kijkab):
    r'''Get the momentum conservation array for a set of k-points.

    This function is similar to get_kconserv, but instead finds the 'kc'
    that satisfies momentum conservation for 5 k-points,

        (ki + kj + kk - ka - kb - kc) dot a = 2n\pi

    where these kpoints are stored in kijkab[ki, kj, kk, ka, kb].
    '''
    a = cell.lattice_vectors() / (2*np.pi)

    kpts_i, kpts_j, kpts_k, kpts_a, kpts_b = \
            [kpts[x].reshape(-1,3) for x in kijkab]
    shape = [np.size(x) for x in kijkab]
    kconserv = np.zeros(shape, dtype=int)

    kv_kab = kpts_k[:,None,None,:] - kpts_a[:,None,:] - kpts_b
    for i, kpti in enumerate(kpts_i):
        for j, kptj in enumerate(kpts_j):
            kv_ijkab = kv_kab + kpti + kptj
            for c, kptc in enumerate(kpts):
                s = np.einsum('kabx,wx->kabw', kv_ijkab - kptc, a)
                s_int = np.rint(s)
                mask = np.einsum('kabw->kab', abs(s - s_int)) < 1e-9
                kconserv[i,j,mask] = c

    new_shape = [shape[i] for i, x in enumerate(kijkab)
                 if not isinstance(x, (int,np.int))]
    kconserv = kconserv.reshape(new_shape)
    return kconserv 

def _mandelbrot_width2height(size=1000, real_range=(-2, 2), imaginary_range=(-2, 2)):
    return int(np.rint((imaginary_range[1] - imaginary_range[0]) / (real_range[1] - real_range[0]) * size)) 

def test_rint_big_int():
    # np.rint bug for large integer values on Windows 32-bit and MKL
    # https://github.com/numpy/numpy/issues/6685
    val = 4607998452777363968
    # This is exactly representable in floating point
    assert_equal(val, int(float(val)))
    # Rint should not change the value
    assert_equal(val, np.rint(val))