You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
101 lines
3.5 KiB
101 lines
3.5 KiB
3 years ago
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from random import uniform
|
||
|
|
import typing
|
||
|
|
import time
|
||
|
|
import turtle as t
|
||
|
|
|
||
|
|
if typing.TYPE_CHECKING:
|
||
|
|
from nptyping import NDArray
|
||
|
|
|
||
|
|
|
||
|
|
HALF_WIDTH, HALF_HEIGHT = 200, 200
|
||
|
|
INITIAL_VELOCITY_COEF = 1000
|
||
|
|
FORCE_LAW = lambda r: 600000 / (r ** 2)
|
||
|
|
DT, N, VELOCITY_CUTOFF, ACCELERATION_CUTOFF = 0.01, 10, 1500, 10000
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Particle:
|
||
|
|
coordinates: NDArray[(2,), float] = field(default_factory=lambda: np.array(
|
||
|
|
[uniform(-HALF_WIDTH, HALF_WIDTH), uniform(-HALF_HEIGHT, HALF_HEIGHT)]
|
||
|
|
))
|
||
|
|
velocity: NDArray[(2,)] = field(
|
||
|
|
default_factory=lambda: INITIAL_VELOCITY_COEF * np.array([uniform(-1, 1), uniform(-1, 1)]))
|
||
|
|
|
||
|
|
def adjust_velocity(self, other_particles: typing.List[Particle], dt: float = DT):
|
||
|
|
a = np.array([0., 0.])
|
||
|
|
for particle in other_particles:
|
||
|
|
if list(particle.coordinates) == list(self.coordinates):
|
||
|
|
continue
|
||
|
|
r = self.coordinates - particle.coordinates
|
||
|
|
a += (-1) * r * FORCE_LAW(np.abs(r))
|
||
|
|
if np.linalg.norm(a) > ACCELERATION_CUTOFF:
|
||
|
|
a *= ACCELERATION_CUTOFF / np.linalg.norm(a)
|
||
|
|
self.velocity += a * dt
|
||
|
|
if np.linalg.norm(self.velocity) > VELOCITY_CUTOFF:
|
||
|
|
self.velocity *= VELOCITY_CUTOFF / np.linalg.norm(self.velocity)
|
||
|
|
if abs(self.coordinates[0]) >= HALF_WIDTH:
|
||
|
|
self.velocity[0] *= -1
|
||
|
|
self.coordinates[0] = np.sign(self.coordinates[0]) * min(abs(self.coordinates[0]), HALF_WIDTH - 2)
|
||
|
|
if abs(self.coordinates[1]) >= HALF_HEIGHT:
|
||
|
|
self.velocity[1] *= -1
|
||
|
|
self.coordinates[1] = np.sign(self.coordinates[1]) * min(abs(self.coordinates[1]), HALF_HEIGHT - 2)
|
||
|
|
|
||
|
|
def move(self, dt: float = DT):
|
||
|
|
self.coordinates += self.velocity * dt
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Gas:
|
||
|
|
particles: typing.List[Particle]
|
||
|
|
dt: float = DT
|
||
|
|
|
||
|
|
def __init__(self, number: int = N) -> Gas:
|
||
|
|
self.particles = [Particle() for _ in range(N)]
|
||
|
|
|
||
|
|
def step(self):
|
||
|
|
for particle in self.particles:
|
||
|
|
particle.adjust_velocity(self.particles, self.dt)
|
||
|
|
for particle in self.particles:
|
||
|
|
particle.move(self.dt)
|
||
|
|
|
||
|
|
def visualize_with_turtles(self, turtles: typing.List[t.Turtle]):
|
||
|
|
for turtle, particle in zip(turtles, self.particles):
|
||
|
|
turtle.goto(*particle.coordinates)
|
||
|
|
|
||
|
|
def run_with_turtles(self, turtles: typing.List[t.Turtle]):
|
||
|
|
while True:
|
||
|
|
last_time: float = time.time()
|
||
|
|
self.visualize_with_turtles(turtles)
|
||
|
|
self.step()
|
||
|
|
time.sleep(max(0, self.dt - (time.time() - last_time)))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def create_turtles_and_run(cls, number: int = N):
|
||
|
|
self = cls()
|
||
|
|
turtles = [t.Turtle(shape='circle') for i in range(number)]
|
||
|
|
for turtle in turtles:
|
||
|
|
turtle.penup()
|
||
|
|
turtle.speed(0)
|
||
|
|
master_turtle = t.Turtle()
|
||
|
|
master_turtle.penup()
|
||
|
|
master_turtle.speed(0)
|
||
|
|
delta = 5
|
||
|
|
master_turtle.goto(-HALF_WIDTH - delta, -HALF_HEIGHT - delta)
|
||
|
|
master_turtle.pendown()
|
||
|
|
master_turtle.goto(HALF_WIDTH + delta, -HALF_HEIGHT - delta)
|
||
|
|
master_turtle.goto(HALF_WIDTH + delta, HALF_HEIGHT + delta)
|
||
|
|
master_turtle.goto(-HALF_WIDTH - delta, HALF_HEIGHT + delta)
|
||
|
|
master_turtle.goto(-HALF_WIDTH - delta, -HALF_HEIGHT - delta)
|
||
|
|
master_turtle.up()
|
||
|
|
master_turtle.goto(1000, 1000)
|
||
|
|
del master_turtle
|
||
|
|
self.run_with_turtles(turtles)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
Gas.create_turtles_and_run()
|