def get_equilibrium_discrete_velocities(macroscopic_velocities, density):
    projected_discrete_velocities = jnp.einsum('dQ,NMLd->NMLQ', LATTICE_VELOCITIES, macroscopic_velocities)
    macroscopic_velocity_magnitude = jnp.linalg.norm(macroscopic_velocities, axis=(- 1), ord=2)
    equilibrium_discrete_velocities = ((density[(..., jnp.newaxis)] * LATTICE_WEIGHTS[jnp.newaxis, jnp.newaxis, jnp.newaxis, :]) * (((1 + (3 * projected_discrete_velocities)) + ((9 / 2) * (projected_discrete_velocities ** 2))) - ((3 / 2) * (macroscopic_velocity_magnitude[(..., jnp.newaxis)] ** 2))))
    return equilibrium_discrete_velocities
def get_equilibrium_discrete_velocities(macroscopic_velocities, density):
    projected_discrete_velocities = jnp.einsum('dQ,NMLd->NMLQ', LATTICE_VELOCITIES, macroscopic_velocities)
    macroscopic_velocity_magnitude = jnp.linalg.norm(macroscopic_velocities, axis=(- 1), ord=2)
    equilibrium_discrete_velocities = ((density[(..., jnp.newaxis)] * LATTICE_WEIGHTS[jnp.newaxis, jnp.newaxis, jnp.newaxis, :]) * (((1 + (3 * projected_discrete_velocities)) + ((9 / 2) * (projected_discrete_velocities ** 2))) - ((3 / 2) * (macroscopic_velocity_magnitude[(..., jnp.newaxis)] ** 2))))
    return equilibrium_discrete_velocities