diff --git a/panama/read.py b/panama/read.py index 220c011..89f828f 100644 --- a/panama/read.py +++ b/panama/read.py @@ -128,6 +128,16 @@ def read_DAT( if files is not None and glob is not None: raise TypeError("`file` and `glob` can't both be not None") + if not additional_columns: + if drop_non_particles: + raise ValueError( + "drop_non_particles requires additional_columns to be calculated." + ) + if mother_columns: + raise ValueError( + "mother_columns require additional columns to be calculated" + ) + if glob is not None: basepath = Path(glob).parent files = list(basepath.glob(Path(glob).name)) @@ -409,10 +419,11 @@ def read_DAT( ) # mother is muon (and in early generation) ) - if drop_mothers: - df_particles.drop( - index=df_particles.query("is_mother == True").index.values, inplace=True - ) + if drop_mothers: + df_particles.drop( + index=df_particles.query("particle_description < 0").index.values, + inplace=True, + ) # Numba version... # df["mother_run_idx"], df["mother_event_idx"], df["mother_particle_idx"] = mother_idx_numba(df.loc[:, "is_mother"].values, df.loc[:, "run_number"].values, df.loc[:, "event_number"].values, df.loc[:, "particle_number"].values) diff --git a/tests/test_read.py b/tests/test_read.py index 233489c..7b8f9a2 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -20,21 +20,48 @@ def test_noparse(test_file_path=Path(__file__).parent / "files" / "DAT000000"): assert df_np.equals(df) -def test_read_corsia_file(test_file_path=Path(__file__).parent / "files" / "DAT000000"): - - df_run, df_event, df = panama.read_DAT(test_file_path, drop_non_particles=False) - - with CorsikaParticleFile(test_file_path, parse_blocks=True) as cf: +def check_eq(file, df_run, df_event, particles, skip_mother=False): + with CorsikaParticleFile(file, parse_blocks=True) as cf: num = 0 for idx, event in enumerate(cf): assert df_event.iloc[idx]["total_energy"] == event.header["total_energy"] for particle in event.particles: - if particle["particle_description"] < 0: + if particle["particle_description"] < 0 and skip_mother: continue - assert df.iloc[num]["px"] == particle["px"] + assert particles.iloc[num]["px"] == particle["px"] num += 1 +def test_noadd(test_file_path=Path(__file__).parent / "files" / "DAT000000"): + + try: + df_run, df_event, particles = panama.read_DAT( + test_file_path, drop_non_particles=True, additional_columns=False + ) + + check_eq(test_file_path, df_run, df_event, particles) + except ValueError as e: + assert "requires" in str(e) + + df_run, df_event, particles = panama.read_DAT( + test_file_path, drop_non_particles=False, additional_columns=True + ) + + check_eq(test_file_path, df_run, df_event, particles, skip_mother=True) + + +def test_read_corsia_file(test_file_path=Path(__file__).parent / "files" / "DAT000000"): + + df_run, df_event, df = panama.read_DAT(test_file_path, drop_non_particles=False) + + check_eq(test_file_path, df_run, df_event, df, skip_mother=True) + try: + check_eq(test_file_path, df_run, df_event, df, skip_mother=False) + assert False + except AssertionError: + pass + + def test_cli(tmp_path, test_file_path=Path(__file__).parent / "files" / "DAT000000"): runner = CliRunner() @@ -50,17 +77,10 @@ def test_cli(tmp_path, test_file_path=Path(__file__).parent / "files" / "DAT0000 assert result.exit_code == 0 particles = pd.read_hdf(tmp_path / "output.hdf5", "particles") - event_headers = pd.read_hdf(tmp_path / "output.hdf5", "event_header") + event_header = pd.read_hdf(tmp_path / "output.hdf5", "event_header") + run_header = pd.read_hdf(tmp_path / "output.hdf5", "run_header") - with CorsikaParticleFile(test_file_path, parse_blocks=True) as cf: - num = 0 - for idx, event in enumerate(cf): - assert ( - event_headers.iloc[idx]["total_energy"] == event.header["total_energy"] - ) - for particle in event.particles: - assert particles.iloc[num]["px"] == particle["px"] - num += 1 + check_eq(test_file_path, run_header, event_header, particles) def test_spectral_index(test_file_path=Path(__file__).parent / "files" / "DAT*"):