import torch


def initial_state(hparams):
    """ Samples a random initial state. """
    B = hparams.batch_size
    L = hparams.lattice_size
    num_particles = round(hparams.concentration * L**2)

    # choose sites that get a particle. Bose-Hubbard allows multiple particles per site.
    if hparams.potential.startswith('bh'):
        indices = torch.multinomial(torch.ones(B, L**2), num_particles, replacement=True)
    else:
        indices = torch.multinomial(torch.ones(B, L**2), num_particles, replacement=False)
    
    state = torch.zeros(B, L**2)
    state = state.scatter_add_(1, indices, torch.ones_like(indices, dtype=torch.float))
    state = state.view(B, L, L)

    return state