def gen_data()

in data/gen_data_collisions.py [0:0]

• 22 lines of code
• 4 McCabe index (conditional complexity)

def gen_data(args, file_path: str, nsteps: int):
    if os.path.exists(file_path):
        print(f"{file_path} already exists! Skipping generation.")
        return
    print(f"Generating {file_path}")

    data = ParticleCollisionData(
        num_particles=args.num_particles,
        num_colors=args.num_colors,
        grid_size=args.grid_size,
        speed=args.speed,
        color_change_prob=args.color_change_prob,
        easy_q=args.easy_q,
    )

    f = open(file_path, "w")
    g = open(file_path + ".labels", "w")
    x, y = data.generate_sample(nsteps)
    xprint = " ".join([str(c) for c in x.tolist()])
    yprint = " ".join([str(c) for c in y.tolist()])
    f.write(xprint)
    g.write(yprint)
    f.close()
    g.close()