def get_particles(num_particles):
particle_coords = gen_particles(num_particles)
charges = np.array(
[-1 if i % 2 == 0 else 1 for i in range(num_particles)],
dtype="int32",
)
return particle_coords, charges
def run_and_plot(num_particles, grid_resolution, dist="random", func="numpy"):
particle_coords = gen_particles(num_particles, dist)
charges = np.random.choice(np.int32([-1, 1]), (num_particles, ))
pot_grid = potential_mpi(particle_coords, grid_resolution, charges, func)
if MPI.COMM_WORLD.rank == 0:
plot_potential_grid(pot_grid)
plt.show()
p, ext = os.path.splitext(path)
while os.path.exists(path):
path = f"{p}{iters:02}{ext}"
iters += 1
return path
if __name__ == "__main__":
# pyplot often complains about not finding ffmpeg, even if it's on
# the system path. Pointing to it explicitly seems to work
plt.rcParams['animation.ffmpeg_path'] = _find_ffmpeg()
n_frames = 200
grid_resolution = 1080
n_particles = 2
particle_coords = gen_particles(n_particles, "circle")
charges = gen_charges(n_particles)
delta_t = 0.03
initial_velocities = np.zeros((n_particles, 2))
frames = simulate_particles(particle_coords,
grid_resolution,
charges,
initial_velocities,
n_frames,
delta_t,
func="numba")
out_dir = os.path.join(ROOT_DIR, ".animations")
base_anim_name = "particle_anim.mp4"
with open(file_path, "w") as f:
json.dump(dict_object, f, indent=2)
def _plot_benchmarks_recursive(times, labels, plot_dir):
to_plot = []
for i, label in enumerate(labels):
to_plot.append(label)
fig = plot_benchmarks(times, labels=to_plot)
out_name = _build_plot_name(f"sub_bench-{i}")
fig.savefig(os.path.join(plot_dir, out_name))
if __name__ == "__main__":
num_particles = 10
particle_coords = gen_particles(num_particles)
grid_resolutions = [50, 100, 200, 300, 500, 750, 1000, 1500, 3000, 5000]
charges = np.ones(num_particles, dtype=np.int32)
max_time = 3 # seconds
repetitions = 5
functions_dict = {
"python": {
"func": "python"
},
"numpy": {
"func": "numpy"
},
"cython": {
"func": "cython"
},
parser.add_argument(
"--kwargs",
nargs="*",
help=("keyword args to pass to calculate_grid, should be in the form "
"'keyword1=value1 [keyword2=value2 ...]'"),
default={})
parser.add_argument("--plot",
action="store_true",
help="plot the output grid.")
parser.add_argument("--rng_seed",
type=int,
help="seed for the random number generator",
default=None)
cli_args = parser.parse_args()
particle_coords = gen_particles(cli_args.num_particles, cli_args.dist)
charges = gen_charges(cli_args.num_particles)
args = (particle_coords, cli_args.grid_resolution, charges)
kwargs = {k: v for k, v in [x.split("=") for x in cli_args.kwargs]}
grids = []
for func_name in cli_args.func:
kwargs["func"] = func_name
grid, elapsed_time = time_func(calculate_grid, *args, **kwargs)
grids.append(grid)
print(f"Function call took {elapsed_time} seconds using {func_name}")
if cli_args.plot:
fig, _ = plot_potential_grid(*grids)
fig.show()