sample data generator

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")