def get_indices(pop, state, masked=None):
    """
    Finds the population agents with a given state.
    Flattens the population for easier index finding.
    """
    # use flattened pop
    pop = pop.flatten()
    if masked is None:
        indices = [i for i, agent in enumerate(pop) if agent.cur_state() == state]
    else:
        assert isinstance(masked, bool), 'masked input must be a boolean'
        indices = [i for i, agent in enumerate(pop) if agent.cur_state() == state and agent.masked == masked]
    return indices