import numpy as np
import matplotlib.pyplot as plt
from celluloid import Camera
# DEFINE FIXED VALUES
# Set how many periods you'd like to run your world
PERIOD_MAX = 1000
# Set the boundaries of the box
BOUND_X_MIN = -15
BOUND_X_MAX = 15
BOUND_Y_MIN = -15
BOUND_Y_MAX = 15
# GATE
GATE_X = 0
#GATE_Y_MIN = -2
#GATE_Y_MAX = 2
# SPEED TYPES
SLOW = 0.2
FAST = 1
# Initial positions
INITIAL_PARTICLE_POSITIONS_X = np.array([6, 5, -3, -8])
INITIAL_PARTICLE_POSITIONS_Y = np.array([4, 20, -15, -11])
# Initial speed
INITIAL_PARTICLE_SPEEDS = np.array([SLOW, FAST, SLOW, FAST])
# Initial Velocity direction as unit vectors
INITIAL_PARTICLE_DIRECTIONS_X = np.array([-np.sqrt(2) / 2, 1, np.sqrt(3)/2, 0.6])
INITIAL_PARTICLE_DIRECTIONS_Y = np.array([-np.sqrt(2) / 2, 0, 1/2, 0.8])
# Total number of particles
particle_total_number = len(INITIAL_PARTICLE_POSITIONS_X)
# Create System State Matrix by defining the correct dimensions
system_state_matrix = np.ndarray(shape=(PERIOD_MAX + 1, 5, particle_total_number))
# Initialize the first entry of the state matrix with the initial values
system_state_matrix[0] = \
INITIAL_PARTICLE_POSITIONS_X, INITIAL_PARTICLE_POSITIONS_Y, INITIAL_PARTICLE_SPEEDS, INITIAL_PARTICLE_DIRECTIONS_X, \
INITIAL_PARTICLE_DIRECTIONS_Y
# Now let's make the particles move
# The Loop to update system
for i in np.arange(0, PERIOD_MAX):
# At each loop, system will move from t=i to t=i+1
# first get the "old" values at t=i
# 2 dimensional position array
particle_positions_at_i = system_state_matrix[i][0:2]
# 1 dimensional speed array
particle_speeds_at_i = system_state_matrix[i][2]
# 2 dimensional velocity direction array
particle_directions_at_i = system_state_matrix[i][3:5]
# First calculate as if no boundary
particle_positions_temporary_array = particle_positions_at_i + particle_speeds_at_i * particle_directions_at_i
# Initialize the next period vectors temporarily
particle_positions_at_i_plus_one = particle_positions_temporary_array
# Make a deep copy to prevent wrong value updating
particle_directions_at_i_plus_one = particle_directions_at_i.copy()
# Now check the boundary conditions and make the necessary updates
# Vertical bounds:
positions_x_out_of_bounds_right = (particle_positions_temporary_array[0] > BOUND_X_MAX)
positions_x_out_of_bounds_left = (particle_positions_temporary_array[0] < BOUND_X_MIN)
positions_x_all_out_of_bound_vertical= np.logical_or(positions_x_out_of_bounds_right, positions_x_out_of_bounds_left)
particle_directions_at_i_plus_one[0][positions_x_all_out_of_bound_vertical] *= -1
particle_positions_at_i_plus_one[0][positions_x_out_of_bounds_right] = 2 * BOUND_X_MAX - \
particle_positions_at_i_plus_one[0][
positions_x_out_of_bounds_right]
particle_positions_at_i_plus_one[0][positions_x_out_of_bounds_left] = 2 * BOUND_X_MIN - \
particle_positions_at_i_plus_one[0][
positions_x_out_of_bounds_left]
# Horizontal bounds
positions_y_out_of_bounds_up = (particle_positions_temporary_array[1] > BOUND_Y_MAX)
positions_x_out_of_bounds_down = (particle_positions_temporary_array[1] < BOUND_X_MIN)
particle_directions_at_i_plus_one[1][
np.logical_or(positions_y_out_of_bounds_up, positions_x_out_of_bounds_down)] *= -1
particle_positions_at_i_plus_one[1][positions_y_out_of_bounds_up] = 2 * BOUND_Y_MAX - \
particle_positions_at_i_plus_one[1][
positions_y_out_of_bounds_up]
particle_positions_at_i_plus_one[1][positions_x_out_of_bounds_down] = 2 * BOUND_Y_MIN - \
particle_positions_at_i_plus_one[1][
positions_x_out_of_bounds_down]
# Check Gate conditions
gate_blocking_left_to_right = np.logical_and.reduce((particle_positions_at_i[0] <= GATE_X,
particle_positions_at_i_plus_one[0] > GATE_X,
particle_speeds_at_i == SLOW))
gate_blocking_right_to_left = np.logical_and.reduce((particle_positions_at_i[0] >= GATE_X,
particle_positions_at_i_plus_one[0] < GATE_X,
particle_speeds_at_i == FAST))
gate_bouncers = np.logical_or(gate_blocking_left_to_right, gate_blocking_right_to_left)
particle_positions_at_i_plus_one[0][gate_bouncers] = 2 * GATE_X - particle_positions_at_i_plus_one[0][gate_bouncers]
particle_directions_at_i_plus_one[0][gate_bouncers] *= -1
# Write the finalized data to the state matrix
system_state_matrix[i + 1] = particle_positions_at_i_plus_one[0], particle_positions_at_i_plus_one[
1], particle_speeds_at_i, particle_directions_at_i_plus_one[0], particle_directions_at_i_plus_one[1]
# Collect the data
print(system_state_matrix)
# Plotting
fig = plt.figure()
camera = Camera(fig)
plt.axis([BOUND_X_MIN, BOUND_X_MAX, BOUND_Y_MIN, BOUND_Y_MAX])
for i in np.arange(0, PERIOD_MAX):
X = system_state_matrix[i][0]
Y = system_state_matrix[i][1]
Z = system_state_matrix[i][2]
plt.scatter(X, Y, c=Z,cmap='bwr')
camera.snap()
animation = camera.animate()
animation.save('ozgun_demon.html')