def get_macroscopic_velocities(discrete_velocities, density):
    return (jnp.einsum('NMLQ,dQ->NMLd', discrete_velocities, LATTICE_VELOCITIES) / density[(..., jnp.newaxis)])
def get_macroscopic_velocities(discrete_velocities, density):
    return (jnp.einsum('NMLQ,dQ->NMLd', discrete_velocities, LATTICE_VELOCITIES) / density[(..., jnp.newaxis)])