mardi 21 mai 2019

Build a numpy array from a random distribution until the last column exceeds a threshold

I want to build a 2d numpy array from a random distribution so that each of the values in the last column of each row exceeds a threshold.

Here's the working code I have now. Is there a cleaner way to build numpy arrays with an arbitrary condition?

def new_array(
        num_rows: int,
        dist: Callable[[int], np.ndarray],
        min_hours: int) -> np.ndarray:
    # Get the 40th percentile as a reasonable guess for how many samples we need.
    # Use a lower percentile to increase num_cols and avoid looping in most cases.
    p40_val = np.quantile(dist(20), 0.4)
    # Generate at least 10 columns each time.
    num_cols = max(int(min_hours / p40_val), 10)

    def create_starts() -> np.ndarray:
        return dist(num_rows * num_cols).reshape((num_rows, num_cols)).cumsum(axis=1)

    max_iters = 20
    starts = create_starts()
    for _ in range(max_iters):
        if np.min(starts[:, -1]) >= min_hours:
            # All the last columns exceed min_hours.
            break

        last_col_vals = starts[:, -1].repeat(num_cols).reshape(starts.shape)
        next_starts = create_starts() + last_col_vals
        starts = np.append(starts, next_starts, axis=1)
    else:
        # We didn't break out of the for loop, so we hit the max iterations.
        raise AssertionError('Failed to create enough samples to exceed '
                             'sim duration for all columns')

    # Only keep columns up to the column where each value > min_hours.
    mins_per_col = np.min(starts, axis=0)
    cols_exceeding_sim_duration = np.nonzero(mins_per_col > min_hours)[0]
    cols_to_keep = cols_exceeding_sim_duration[0]
    return np.delete(starts, np.s_[cols_to_keep:], axis=1)


new_array(5, lambda size: np.random.normal(3, size=size), 7)

# Example output
array([[1.47584632, 4.04034105, 7.19592256],
       [3.10804306, 6.46487043, 9.74177227],
       [1.03633165, 2.62430309, 6.92413189],
       [3.46100139, 6.53068143, 7.37990547],
       [2.70152742, 6.09488369, 9.58376664]])




Aucun commentaire:

Enregistrer un commentaire