Sample batch of trajectories and return them
334020
def sample_trajectories(self, rollout_length, batch_info) -> Trajectories:
        
        indexes = self.backend.sample_batch_trajectories(rollout_length)
        transition_tensors = self.backend.get_trajectories(indexes, rollout_length)

        return Trajectories(
            num_steps=rollout_length,
            num_envs=self.backend.num_envs,
            environment_information=None,
            transition_tensors={k: torch.from_numpy(v) for k, v in transition_tensors.items()},
            rollout_tensors={}
        )
480883
def get_trajectories(self, indexes, rollout_length):
        
        # assert indexes.shape[0] > 1, "There must be multiple indexes supplied"
        assert rollout_length > 1, "Rollout length must be greater than 1"

        batch_indexes = (
                indexes.reshape(1, indexes.shape[0]) - np.arange(rollout_length - 1, -1, -1).reshape(rollout_length, 1)
        )

        return self.get_transitions(batch_indexes)
1072559
def _sample_batch():
    
    if _sample_probability == 1.0 or random.random() < _sample_probability:
        return True

    # Pop off all the metrics for the batch
    for database in _measurements:
        _measurements[database] = _measurements[database][_max_batch_size:]
    return False
149534
def _get_batch_name(sample):
    
    batch = dd.get_batch(sample) or dd.get_sample_name(sample)
    if isinstance(batch, (list, tuple)) and len(batch) > 1:
        batch = dd.get_sample_name(sample)
    return batch
887954
def _group_by_batches(samples, check_fn):
    
    batch_groups = collections.defaultdict(list)
    extras = []
    for data in [x[0] for x in samples]:
        if check_fn(data):
            batch_groups[multi.get_batch_for_key(data)].append(data)
        else:
            extras.append([data])
    return batch_groups, extras
331984
def results(self):
        """"""

        # Add all saved samples (and ancillary quantities) to the results.
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            results = [('niter', self.it - 1),
                       ('ncall', np.array(self.saved_nc)),
                       ('eff', self.eff),
                       ('samples', np.array(self.saved_v)),
                       ('samples_id', np.array(self.saved_id)),
                       ('samples_batch', np.array(self.saved_batch,
                                                  dtype='int')),
                       ('samples_it', np.array(self.saved_it)),
                       ('samples_u', np.array(self.saved_u)),
                       ('samples_n', np.array(self.saved_n)),
                       ('logwt', np.array(self.saved_logwt)),
                       ('logl', np.array(self.saved_logl)),
                       ('logvol', np.array(self.saved_logvol)),
                       ('logz', np.array(self.saved_logz)),
                       ('logzerr', np.sqrt(np.array(self.saved_logzvar))),
                       ('information', np.array(self.saved_h)),
                       ('batch_nlive', np.array(self.saved_batch_nlive,
                                                dtype='int')),
                       ('batch_bounds', np.array(self.saved_batch_bounds))]

        # Add any saved bounds (and ancillary quantities) to the results.
        if self.sampler.save_bounds:
            results.append(('bound', copy.deepcopy(self.bound)))
            results.append(('bound_iter',
                            np.array(self.saved_bounditer, dtype='int')))
            results.append(('samples_bound',
                            np.array(self.saved_boundidx, dtype='int')))
            results.append(('scale', np.array(self.saved_scale)))

        return Results(results)
986580
def _combine_qc_samples(samples):
    
    by_bam = collections.defaultdict(list)
    for data in [utils.to_single_data(x) for x in samples]:
        batch = dd.get_batch(data) or dd.get_sample_name(data)
        if not isinstance(batch, (list, tuple)):
            batch = [batch]
        batch = tuple(batch)
        by_bam[(dd.get_align_bam(data) or dd.get_work_bam(data), batch)].append(data)
    out = []
    for data_group in by_bam.values():
        data = data_group[0]
        alg_qc = []
        qc = {}
        metrics = {}
        for d in data_group:
            qc.update(dd.get_summary_qc(d))
            metrics.update(dd.get_summary_metrics(d))
            alg_qc.extend(dd.get_algorithm_qc(d))
        data["config"]["algorithm"]["qc"] = alg_qc
        data["summary"]["qc"] = qc
        data["summary"]["metrics"] = metrics
        out.append([data])
    return out
977512
def stft(self, samples: np.ndarray):
        
        window = self.window_function(self.window_size)
        hop_size = self.hop_size

        if len(samples) < self.window_size:
            samples = np.append(samples, np.zeros(self.window_size - len(samples)))

        num_frames = max(1, (len(samples) - self.window_size) // hop_size + 1)

        # Get frames as numpy view with stride_tricks to save RAM
        # Same as: frames = [padded_samples[i*hop_size:i*hop_size+self.window_size] for i in range(num_frames)]
        shape = (num_frames, self.window_size)
        strides = (hop_size * samples.strides[-1], samples.strides[-1])
        frames = np.lib.stride_tricks.as_strided(samples, shape=shape, strides=strides)

        result = np.fft.fft(frames * window, self.window_size) / np.atleast_1d(self.window_size)
        return result
724724
def bootstrap_trajectories(trajs, correlation_length):
    
    from scipy.stats import rv_discrete
    # if we have just one trajectory, put it into a one-element list:
    if (isinstance(trajs[0], (int, int, float))):
        trajs = [trajs]
    ntraj = len(trajs)

    # determine correlation length to be used
    lengths = determine_lengths(trajs)
    Ltot = np.sum(lengths)
    Lmax = np.max(lengths)
    if (correlation_length < 1):
        correlation_length = Lmax

        # assign probabilites to select trajectories
    w_trajs = np.zeros((len(trajs)))
    for i in range(ntraj):
        w_trajs[i] = len(trajs[i])
    w_trajs /= np.sum(w_trajs)  # normalize to sum 1.0
    distrib_trajs = rv_discrete(values=(list(range(ntraj)), w_trajs))

    # generate subtrajectories
    Laccum = 0
    subs = []
    while (Laccum < Ltot):
        # pick a random trajectory
        itraj = distrib_trajs.rvs()
        # pick a starting frame
        t0 = random.randint(0, max(1, len(trajs[itraj]) - correlation_length))
        t1 = min(len(trajs[itraj]), t0 + correlation_length)
        # add new subtrajectory
        subs.append(trajs[itraj][t0:t1])
        # increment available states
        Laccum += (t1 - t0)

    # and return
    return subs
928332
def sample(self, rvec0, nsteps, nskip=1):
        """
        rvec = np.zeros(shape=(nsteps + 1, self.dim))
        rvec[0, :] = rvec0[:]
        for t in range(nsteps):
            q = rvec[t, :]
            for s in range(nskip):
                q = self.step(q)
            rvec[t + 1, :] = q[:]
        return rvec
584612
def sample(self, batch_size, batch_idxs=None):
        
        # It is not possible to tell whether the first state in the memory is terminal, because it
        # would require access to the "terminal" flag associated to the previous state. As a result
        # we will never return this first state (only using `self.terminals[0]` to know whether the
        # second state is terminal).
        # In addition we need enough entries to fill the desired window length.
        assert self.nb_entries >= self.window_length + 2, 'not enough entries in the memory'

        if batch_idxs is None:
            # Draw random indexes such that we have enough entries before each index to fill the
            # desired window length.
            batch_idxs = sample_batch_indexes(
                self.window_length, self.nb_entries - 1, size=batch_size)
        batch_idxs = np.array(batch_idxs) + 1
        assert np.min(batch_idxs) >= self.window_length + 1
        assert np.max(batch_idxs) < self.nb_entries
        assert len(batch_idxs) == batch_size

        # Create experiences
        experiences = []
        for idx in batch_idxs:
            terminal0 = self.terminals[idx - 2]
            while terminal0:
                # Skip this transition because the environment was reset here. Select a new, random
                # transition and use this instead. This may cause the batch to contain the same
                # transition twice.
                idx = sample_batch_indexes(self.window_length + 1, self.nb_entries, size=1)[0]
                terminal0 = self.terminals[idx - 2]
            assert self.window_length + 1 <= idx < self.nb_entries

            # This code is slightly complicated by the fact that subsequent observations might be
            # from different episodes. We ensure that an experience never spans multiple episodes.
            # This is probably not that important in practice but it seems cleaner.
            state0 = [self.observations[idx - 1]]
            for offset in range(0, self.window_length - 1):
                current_idx = idx - 2 - offset
                assert current_idx >= 1
                current_terminal = self.terminals[current_idx - 1]
                if current_terminal and not self.ignore_episode_boundaries:
                    # The previously handled observation was terminal, don't add the current one.
                    # Otherwise we would leak into a different episode.
                    break
                state0.insert(0, self.observations[current_idx])
            while len(state0) < self.window_length:
                state0.insert(0, zeroed_observation(state0[0]))
            action = self.actions[idx - 1]
            reward = self.rewards[idx - 1]
            terminal1 = self.terminals[idx - 1]

            # Okay, now we need to create the follow-up state. This is state0 shifted on timestep
            # to the right. Again, we need to be careful to not include an observation from the next
            # episode if the last state is terminal.
            state1 = [np.copy(x) for x in state0[1:]]
            state1.append(self.observations[idx])

            assert len(state0) == self.window_length
            assert len(state1) == len(state0)
            experiences.append(Experience(state0=state0, action=action, reward=reward,
                                          state1=state1, terminal1=terminal1))
        assert len(experiences) == batch_size
        return experiences
570077
def sample(self, batch_info: BatchInfo, model: RlModel, number_of_steps: int) -> Rollout:
        
        # Sample trajectories
        rollout = self.replay_buffer.sample_trajectories(rollout_length=number_of_steps, batch_info=batch_info)

        last_observations = rollout.transition_tensors['observations_next'][-1].to(self.device)
        final_values = model.value(last_observations).cpu()

        # Add 'final_values' to the rollout
        rollout.rollout_tensors['final_values'] = final_values

        return rollout
470664
def sample_given_context(self, c, c_dims):
        '''
        
        '''
        index = self.discrete_progress.sample_given_context(c, c_dims, self.space)
        return self.space.rand_value(index).flatten()[list(set(range(len(self.space.cardinalities))) - set(c_dims))]
746404
def _process_batch(self):
        """"""
        if self.block_id == len(self.tgroups):
            return None, None

        tgroup = self.tgroups[self.block_id]
        psize = self.chunk_tsamples[self.block_id]
        mid = self.mids[self.block_id]
        stilde = self.data.overwhitened_data(tgroup[0].delta_f)
        psd = stilde.psd

        valid_end = int(psize - self.data.trim_padding)
        valid_start = int(valid_end - self.data.blocksize * self.data.sample_rate)

        seg = slice(valid_start, valid_end)

        self.corr[self.block_id].execute(stilde)
        self.ifts[mid].execute()

        self.block_id += 1

        snr = numpy.zeros(len(tgroup), dtype=numpy.complex64)
        time = numpy.zeros(len(tgroup), dtype=numpy.float64)
        templates = numpy.zeros(len(tgroup), dtype=numpy.uint64)
        sigmasq = numpy.zeros(len(tgroup), dtype=numpy.float32)

        time[:] = self.data.start_time

        result = {}
        tkeys = tgroup[0].params.dtype.names
        for key in tkeys:
            result[key] = []

        veto_info = []

        # Find the peaks in our SNR times series from the various templates
        i = 0
        for htilde in tgroup:
            l = htilde.out[seg].abs_arg_max()

            sgm = htilde.sigmasq(psd)
            norm = 4.0 * htilde.delta_f / (sgm ** 0.5)

            l += valid_start
            snrv = numpy.array([htilde.out[l]])

            # If nothing is above threshold we can exit this template
            s = abs(snrv[0]) * norm
            if s < self.snr_threshold:
                continue

            time[i] += float(l - valid_start) / self.data.sample_rate

            # We have an SNR so high that we will drop the entire analysis
            # of this chunk of time!
            if self.snr_abort_threshold is not None and s > self.snr_abort_threshold:
                logging.info("We are seeing some *really* high SNRs, lets"
                             " assume they aren't signals and just give up")
                return False, []

            veto_info.append((snrv, norm, l, htilde, stilde))

            snr[i] = snrv[0] * norm
            sigmasq[i] = sgm
            templates[i] = htilde.id
            if not hasattr(htilde, 'dict_params'):
                htilde.dict_params = {}
                for key in tkeys:
                    htilde.dict_params[key] = htilde.params[key]

            for key in tkeys:
                result[key].append(htilde.dict_params[key])
            i += 1


        result['snr'] = abs(snr[0:i])
        result['coa_phase'] = numpy.angle(snr[0:i])
        result['end_time'] = time[0:i]
        result['template_id'] = templates[0:i]
        result['sigmasq'] = sigmasq[0:i]

        for key in tkeys:
            result[key] = numpy.array(result[key])

        return result, veto_info
979185
def _collate(self, batch):
        

        if isinstance(batch[0], tuple):
            return [self._cuda(torch.Tensor(samples)) for samples in list(zip(*batch))]
        else:
            return self._cuda(torch.Tensor(batch))
1038366
def batch(self):
        """"""
        _batch = []
        yield _batch
        if _batch:
            new_part = super(type(self), type(self)).empty()
            new_part._from_entities(_batch, _batch[0].__dimension__)
            self.concat(new_part, inplace=True)
727189
def next_batch(self):
        
        is_success, results = AtlasRequest(
            url_path=self.atlas_url,
            user_agent=self._user_agent,
            server=self.server,
            verify=self.verify,
        ).get()

        if not is_success:
            raise APIResponseError(results)

        self.total_count = results.get("count")
        self.atlas_url = self.build_next_url(results.get("next"))
        self.current_batch = results.get("results", [])