@jax.jit
def update(discrete_velocities_prev):
    discrete_velocities_prev = discrete_velocities_prev.at[(- 1), :, :, LEFT_VELOCITIES].set(discrete_velocities_prev[(- 2), :, :, LEFT_VELOCITIES])
    density_prev = get_density(discrete_velocities_prev)
    macroscopic_velocities_prev = get_macroscopic_velocities(discrete_velocities_prev, density_prev)
    macroscopic_velocities_prev = macroscopic_velocities_prev.at[0, 1:(- 1), 1:(- 1), :].set(VELOCITY_PROFILE[0, 1:(- 1), 1:(- 1), :])
    lateral_densities = get_density(jnp.transpose(discrete_velocities_prev[0, :, :, YZ_VELOCITIES], axes=(1, 2, 0)))
    left_densities = get_density(jnp.transpose(discrete_velocities_prev[0, :, :, LEFT_VELOCITIES], axes=(1, 2, 0)))
    density_prev = density_prev.at[0, :, :].set(((lateral_densities + (2 * left_densities)) / (1 - macroscopic_velocities_prev[0, :, :, 0])))
    equilibrium_discrete_velocities = get_equilibrium_discrete_velocities(macroscopic_velocities_prev, density_prev)
    discrete_velocities_prev = discrete_velocities_prev.at[0, :, :, RIGHT_VELOCITIES].set(equilibrium_discrete_velocities[0, :, :, RIGHT_VELOCITIES])
    discrete_velocities_post_collision = (discrete_velocities_prev - (RELAXATION_OMEGA * (discrete_velocities_prev - equilibrium_discrete_velocities)))
    for i in range(N_DISCRETE_VELOCITIES):
        discrete_velocities_post_collision = discrete_velocities_post_collision.at[(obstacle_mask, LATTICE_INDICES[i])].set(discrete_velocities_prev[(obstacle_mask, OPPOSITE_LATTICE_INDICES[i])])
    discrete_velocities_streamed = discrete_velocities_post_collision
    for i in range(N_DISCRETE_VELOCITIES):
        discrete_velocities_streamed = discrete_velocities_streamed.at[:, :, :, i].set(jnp.roll(jnp.roll(jnp.roll(discrete_velocities_post_collision[:, :, :, i], LATTICE_VELOCITIES[(0, i)], axis=0), LATTICE_VELOCITIES[(1, i)], axis=1), LATTICE_VELOCITIES[(2, i)], axis=2))
    return discrete_velocities_streamed
@jax.jit
def update(discrete_velocities_prev):
    discrete_velocities_prev = discrete_velocities_prev.at[(- 1), :, :, LEFT_VELOCITIES].set(discrete_velocities_prev[(- 2), :, :, LEFT_VELOCITIES])
    density_prev = get_density(discrete_velocities_prev)
    macroscopic_velocities_prev = get_macroscopic_velocities(discrete_velocities_prev, density_prev)
    macroscopic_velocities_prev = macroscopic_velocities_prev.at[0, 1:(- 1), 1:(- 1), :].set(VELOCITY_PROFILE[0, 1:(- 1), 1:(- 1), :])
    lateral_densities = get_density(jnp.transpose(discrete_velocities_prev[0, :, :, YZ_VELOCITIES], axes=(1, 2, 0)))
    left_densities = get_density(jnp.transpose(discrete_velocities_prev[0, :, :, LEFT_VELOCITIES], axes=(1, 2, 0)))
    density_prev = density_prev.at[0, :, :].set(((lateral_densities + (2 * left_densities)) / (1 - macroscopic_velocities_prev[0, :, :, 0])))
    equilibrium_discrete_velocities = get_equilibrium_discrete_velocities(macroscopic_velocities_prev, density_prev)
    discrete_velocities_prev = discrete_velocities_prev.at[0, :, :, RIGHT_VELOCITIES].set(equilibrium_discrete_velocities[0, :, :, RIGHT_VELOCITIES])
    discrete_velocities_post_collision = (discrete_velocities_prev - (RELAXATION_OMEGA * (discrete_velocities_prev - equilibrium_discrete_velocities)))
    for i in range(N_DISCRETE_VELOCITIES):
        discrete_velocities_post_collision = discrete_velocities_post_collision.at[(obstacle_mask, LATTICE_INDICES[i])].set(discrete_velocities_prev[(obstacle_mask, OPPOSITE_LATTICE_INDICES[i])])
    discrete_velocities_streamed = discrete_velocities_post_collision
    for i in range(N_DISCRETE_VELOCITIES):
        discrete_velocities_streamed = discrete_velocities_streamed.at[:, :, :, i].set(jnp.roll(jnp.roll(jnp.roll(discrete_velocities_post_collision[:, :, :, i], LATTICE_VELOCITIES[(0, i)], axis=0), LATTICE_VELOCITIES[(1, i)], axis=1), LATTICE_VELOCITIES[(2, i)], axis=2))
    return discrete_velocities_streamed