diff --git a/panama/cli/run.py b/panama/cli/run.py index 9a6f92b..aa2ce62 100644 --- a/panama/cli/run.py +++ b/panama/cli/run.py @@ -103,6 +103,20 @@ def convert(self, value: int | str, param: Any, ctx: Any) -> int | dict[int, int is_flag=True, help="Save CORSIKAs std_out to a log file in output directory.", ) +@click.option( + "--first_run_number", + "-fr", + default=0, + help="Sets the first run number, for multiple output files, the number increments from that number.", + type=int, +) +@click.option( + "--first_event_number", + "-fe", + default=1, + help="Sets the first event number in each produced DAT file.", + type=int, +) @click.option("--debug", "-d", default=False, is_flag=True, help="Enable debug output") def run( template: Path, @@ -114,6 +128,8 @@ def run( seed: int, tmp: Path, save_std: bool, + first_run_number: int, + first_event_number: int, debug: bool, ) -> None: """ @@ -149,6 +165,15 @@ def run( ) runner = CorsikaRunner( - primary, jobs, template, Path(output), corsika, p, seed, save_std + primary, + jobs, + template, + Path(output), + corsika, + p, + seed, + save_std, + first_run_number, + first_event_number, ) runner.run() diff --git a/panama/run.py b/panama/run.py index dabe958..0c55015 100644 --- a/panama/run.py +++ b/panama/run.py @@ -195,9 +195,11 @@ def __init__( corsika_tmp_dir: Path, seed: None | int = None, save_std: bool = False, + first_run_number: int = 0, + first_event_number: int = 1, ) -> None: """ - TODO: Good Docstring, Types + TODO: Docstring Parameters ---------- @@ -208,6 +210,8 @@ def __init__( self.corsika_executable = Path(corsika_executable) self.corsika_tmp_dir = Path(corsika_tmp_dir) self.save_std = save_std + self.first_run_number = first_run_number + self.first_event_number = first_event_number # we always need at least n_showers if we want to run n_jobs if not all(n_jobs <= n_showers for n_showers in primary.values()): @@ -301,11 +305,10 @@ def _get_corsika_config( run_idx: int, n_show: int, primary_corsikaid: int, - first_event_idx: int = 1, ) -> dict[str, str]: return { - "run_idx": f"{run_idx}", - "first_event_idx": f"{first_event_idx}", + "run_idx": f"{run_idx+self.first_run_number}", + "first_event_idx": f"{self.first_event_number}", "n_show": f"{n_show}", "dir": str(self.output.absolute()) + "/", "seed_1": f"{randrange(1, 900_000_000)}", diff --git a/tests/test_cli.py b/tests/test_cli.py index 12136bf..8e0af46 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -149,6 +149,53 @@ def test_cli_fast( assert "DEBUG" in caplog.text +def test_cli_first_number( + tmp_path, + caplog, + test_file_path=Path(__file__).parent / "files" / "example_corsika_low_energy.template", + corsika_path=Path(__file__).parent.parent + / CORSIKA_VERSION + / "run" + / CORSIKA_EXECUTABLE, + compare_files=Path(__file__).parent / "files" / "compare" / "DAT*", +): + runner = CliRunner() + result = runner.invoke( + cli, + [ + "--debug", + "run", + f"{test_file_path}", + "--primary", + "2212", # proton + "-n", + "1", + "--corsika", + f"{corsika_path}", + "--output", + f"{tmp_path}", + "--seed", + "137", + "--jobs", + "1", + "--debug", + "-fr", + "42", + "-fe", + "420" + ], + catch_exceptions=False + ) + + assert result.exit_code == 0 + + run_header_2, event_header_2, ps_2 = read_DAT(glob=f"{tmp_path}/DAT*") + + assert event_header_2.shape[0] == 1 + assert "DEBUG" in caplog.text + assert event_header_2.index[0] == (42, 420) + + def test_cli_double_arg( tmp_path, caplog, diff --git a/tests/test_run.py b/tests/test_run.py index bee0ac5..9c4c342 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -12,8 +12,10 @@ CORSIKA_VERSION = "corsika-77500" CORSIKA_EXECUTABLE = "corsika77500Linux_SIBYLL_urqmd" + def test_run_fail( - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", ): try: runner = CorsikaRunner( @@ -34,7 +36,8 @@ def test_run_fail( def test_no_double_run( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -50,8 +53,8 @@ def test_no_double_run( seed=137, save_std=True ) - - cfg = runner._get_corsika_config(0, 5, 13, 1) + + cfg = runner._get_corsika_config(0, 5, 13) runner.job_pool[0].start(cfg) with pytest.raises(RuntimeError, match="it's still running"): @@ -62,7 +65,8 @@ def test_no_double_run( def test_corsika_runner( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -77,9 +81,9 @@ def test_corsika_runner( corsika_tmp_dir=tmp_path, seed=137, ) - + runner.run() - + run_header_2, event_header_2, ps_2 = read_DAT(glob=f"{tmp_path}/DAT*") assert event_header_2.shape[0] == 2 @@ -87,9 +91,39 @@ def test_corsika_runner( assert len(event_header_2["particle_id"].unique()) == 2 +def test_corsika_runner_first_numbers( + tmp_path, + test_file_path=Path(__file__).parent / "files" / + "example_corsika_low_energy.template", + corsika_path=Path(__file__).parent.parent + / CORSIKA_VERSION + / "run" + / CORSIKA_EXECUTABLE, + compare_files=Path(__file__).parent / "files" / "compare" / "DAT*", +): + runner = CorsikaRunner(primary={2212: 1}, + n_jobs=1, + template_path=test_file_path, + output=tmp_path, + corsika_executable=corsika_path, + corsika_tmp_dir=tmp_path, + seed=137, + first_run_number=42, + first_event_number=420 + ) + + runner.run() + + run_header, event_header, ps = read_DAT(glob=f"{tmp_path}/DAT*") + + assert event_header.shape[0] == 1 + assert len(event_header["particle_id"].unique()) == 1 + assert event_header.index[0] == (42, 420) + def test_corsika_job_no_join( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika_low_energy.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika_low_energy.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -104,16 +138,18 @@ def test_corsika_job_no_join( corsika_tmp_dir=tmp_path, seed=137, ) - + runner.run() - + with pytest.raises(RuntimeError, match="already"): runner.job_pool[0].join() + def test_corsika_error( tmp_path, caplog, - test_file_path=Path(__file__).parent / "files" / "example_corsika_faulty.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika_faulty.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -138,7 +174,7 @@ def test_corsika_error( "--seed", "137", "--jobs", - "1", + "1", "--debug", ], catch_exceptions=False @@ -151,7 +187,8 @@ def test_corsika_error( def test_file_output_compare( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -177,12 +214,14 @@ def test_file_output_compare( dates = list(filter(lambda name: "date" in str(name), df_1.columns)) df_1.drop(dates, axis="columns", inplace=True) df_2.drop(dates, axis="columns", inplace=True) - assert df_1.select_dtypes(exclude=['object']).equals(df_2.select_dtypes(exclude=['object'])) + assert df_1.select_dtypes(exclude=['object']).equals( + df_2.select_dtypes(exclude=['object'])) def test_different_primary_type( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -223,10 +262,11 @@ def test_different_primary_type( def test_multi_job( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION - / "run" + / "run" / CORSIKA_EXECUTABLE, compare_files=Path(__file__).parent / "files" / "compare" / "DAT*", ): @@ -269,7 +309,8 @@ def test_multi_job( def test_multi_job_fail( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -296,12 +337,13 @@ def test_multi_job_fail( "--debug", ], catch_exceptions=False - ) + ) + - def test_save_output( tmp_path, - test_file_path=Path(__file__).parent / "files" / "example_corsika_low_energy.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika_low_energy.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -322,7 +364,7 @@ def test_save_output( "--seed", "137", "--jobs", - "2", + "2", "--debug", "--save-std" ], @@ -339,10 +381,12 @@ def test_save_output( assert "END OF RUN" in content + def test_run_multitemp( tmp_path, caplog, - test_file_path=Path(__file__).parent / "files" / "example_corsika_low_energy.template", + test_file_path=Path(__file__).parent / "files" / + "example_corsika_low_energy.template", corsika_path=Path(__file__).parent.parent / CORSIKA_VERSION / "run" @@ -357,7 +401,7 @@ def test_run_multitemp( "run", f"{test_file_path}", "--primary", - "2212", # proton + "2212", # proton "-n", "1", "--corsika", @@ -369,12 +413,12 @@ def test_run_multitemp( "--seed", "137", "--jobs", - "1", + "1", "--debug", ], catch_exceptions=False ) - + with open(f"{tmp_path}/tmp/touch", "w") as file: file.write("I am touched") @@ -386,7 +430,7 @@ def test_run_multitemp( "run", f"{test_file_path}", "--primary", - "2212", # proton + "2212", # proton "-n", "1", "--corsika", @@ -398,7 +442,7 @@ def test_run_multitemp( "--seed", "69", "--jobs", - "1", + "1", "--debug", ], catch_exceptions=False @@ -413,6 +457,6 @@ def test_run_multitemp( assert event_header_1.shape[0] == 1 assert event_header_2.shape[0] == 1 - assert not event_header_1.select_dtypes(exclude=["object"]).equals(event_header_2.select_dtypes(exclude=["object"])) + assert not event_header_1.select_dtypes(exclude=["object"]).equals( + event_header_2.select_dtypes(exclude=["object"])) assert "DEBUG" in caplog.text -