From b75a427aad2d625c955ac96a1b1ec75baf637f82 Mon Sep 17 00:00:00 2001 From: Tim Diller Date: Fri, 15 May 2026 09:38:04 -0500 Subject: [PATCH] simplify the function to create particles goal is to leave lots of room for student improvements in class. also, introduce missingness in the data --- particlecloud/util.py | 67 +++++++++++++++++++++---------------------- 1 file changed, 33 insertions(+), 34 deletions(-) diff --git a/particlecloud/util.py b/particlecloud/util.py index f590b06..54987ae 100644 --- a/particlecloud/util.py +++ b/particlecloud/util.py @@ -1,47 +1,46 @@ import random +import sys -from particlecloud.stuff import Particle +import numpy as np -PARTICLE_FILE_COLUMN_NAMES = [ - "number", "mass", "vx", "vy", "vz", "px", "py", "pz", +FILE_COLUMN_NAMES = [ + "n", "m", "v1", "v2", "v3", "p1", "p2", "p3", ] -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), + +def write_data_file(filename, n=20, missing_values=False): + """Write a csv file with data + """ + + with open(filename, "w") as fp: + fp.write(", ".join(FILE_COLUMN_NAMES) + ",\n") + for i in range(n): + vals = [] + vals.append(random.gauss(3., 0.5)) + vals.extend( + ( + random.gauss(0., 25), + random.gauss(0., 25), + random.gauss(0., 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) + vals.extend( + ( + random.gauss(0., 1.5), + random.gauss(0., 1.5), + random.gauss(0., 1.5), + ) + ) + if missing_values: + if random.randint(0,10) < 1: + idx = random.randint(1, len(vals)) + vals[idx] = np.nan 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") + missing_values = sys.argv[-1] == '-m' + + write_data_file("sample_data.csv", n=20, missing_values=True)