diff --git a/particlecloud/util.py b/particlecloud/util.py new file mode 100644 index 0000000..f590b06 --- /dev/null +++ b/particlecloud/util.py @@ -0,0 +1,47 @@ +import random + +from particlecloud.stuff import Particle + + +PARTICLE_FILE_COLUMN_NAMES = [ + "number", "mass", "vx", "vy", "vz", "px", "py", "pz", +] + +def make_particles(n=10): + particles = [] + for i in range(n): + particles.append( + Particle( + mass=random.gauss(mu=3., sigma=0.5), + velocity=( + random.gauss(mu=0., sigma=1.5), + random.gauss(mu=0., sigma=1.5), + random.gauss(mu=0., sigma=1.5), + ), + position=( + random.gauss(mu=0., sigma=25), + random.gauss(mu=0., sigma=25), + random.gauss(mu=0., sigma=25), + ) + ) + ) + return particles + + +def write_particles_file(filename, n=10): + """Write a csv file with particle data + """ + particles = make_particles(n=n) + with open(filename, "w") as fp: + fp.write(", ".join(PARTICLE_FILE_COLUMN_NAMES) + ",\n") + for i, p in enumerate(particles): + vals = [] + vals.append(p.mass) + vals.extend(p.velocity) + vals.extend(p.position) + line = f"{i}, " + ", ".join(f"{v:7.5f}" for v in vals) + ",\n" + fp.write(line) + + +if __name__ == '__main__': + write_particles_file("sample_data.csv")