diff --git a/pdm.lock b/pdm.lock index 51a9440..7c36855 100644 --- a/pdm.lock +++ b/pdm.lock @@ -5,7 +5,7 @@ groups = ["default", "dev"] strategy = ["inherit_metadata"] lock_version = "4.5.0" -content_hash = "sha256:65d11af716f54b9037cfdcc76bf8521d1c0facc51e383b889157a40397cdce04" +content_hash = "sha256:567bf30edee65b496cbfde3af93faa9d6af967a9a06ddb62d84f370d89b7b139" [[metadata.targets]] requires_python = ">=3.11,<3.13" @@ -265,7 +265,7 @@ name = "asttokens" version = "3.0.0" requires_python = ">=3.8" summary = "Annotate AST trees with source code positions" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2"}, {file = "asttokens-3.0.0.tar.gz", hash = "sha256:0dcd8baa8d62b0c1d118b399b2ddba3c4aff271d0d7a9e0d4c1681c79035bbc7"}, @@ -659,7 +659,7 @@ name = "decorator" version = "5.1.1" requires_python = ">=3.5" summary = "Decorators for Humans" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "decorator-5.1.1-py3-none-any.whl", hash = "sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186"}, {file = "decorator-5.1.1.tar.gz", hash = "sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330"}, @@ -694,7 +694,7 @@ name = "executing" version = "2.1.0" requires_python = ">=3.8" summary = "Get the currently executing AST node of a frame, and other information" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "executing-2.1.0-py2.py3-none-any.whl", hash = "sha256:8d63781349375b5ebccc3142f4b30350c0cd9c79f921cde38be2be4637e98eaf"}, {file = "executing-2.1.0.tar.gz", hash = "sha256:8ea27ddd260da8150fa5a708269c4a10e76161e2496ec3e587da9e3c0fe4b9ab"}, @@ -1072,12 +1072,43 @@ files = [ {file = "installer-0.7.0.tar.gz", hash = "sha256:a26d3e3116289bb08216e0d0f7d925fcef0b0194eedfa0c944bcaaa106c4b631"}, ] +[[package]] +name = "ipdb" +version = "0.13.13" +requires_python = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +summary = "IPython-enabled pdb" +groups = ["default"] +dependencies = [ + "decorator; python_version == \"3.5\"", + "decorator; python_version == \"3.6\"", + "decorator; python_version > \"3.6\" and python_version < \"3.11\"", + "decorator; python_version >= \"3.11\"", + "decorator<5.0.0; python_version == \"2.7\"", + "decorator<5.0.0; python_version == \"3.4\"", + "ipython<6.0.0,>=5.1.0; python_version == \"2.7\"", + "ipython<7.0.0,>=6.0.0; python_version == \"3.4\"", + "ipython<7.10.0,>=7.0.0; python_version == \"3.5\"", + "ipython<7.17.0,>=7.16.3; python_version == \"3.6\"", + "ipython>=7.31.1; python_version > \"3.6\" and python_version < \"3.11\"", + "ipython>=7.31.1; python_version >= \"3.11\"", + "pathlib; python_version == \"2.7\"", + "toml>=0.10.2; python_version == \"2.7\"", + "toml>=0.10.2; python_version == \"3.4\"", + "toml>=0.10.2; python_version == \"3.5\"", + "tomli; python_version == \"3.6\"", + "tomli; python_version > \"3.6\" and python_version < \"3.11\"", +] +files = [ + {file = "ipdb-0.13.13-py3-none-any.whl", hash = "sha256:45529994741c4ab6d2388bfa5d7b725c2cf7fe9deffabdb8a6113aa5ed449ed4"}, + {file = "ipdb-0.13.13.tar.gz", hash = "sha256:e3ac6018ef05126d442af680aad863006ec19d02290561ac88b8b1c0b0cfc726"}, +] + [[package]] name = "ipython" version = "8.31.0" requires_python = ">=3.10" summary = "IPython: Productive Interactive Computing" -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "colorama; sys_platform == \"win32\"", "decorator", @@ -1199,7 +1230,7 @@ name = "jedi" version = "0.19.2" requires_python = ">=3.6" summary = "An autocompletion tool for Python that can be used for text editors." -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "parso<0.9.0,>=0.8.4", ] @@ -1476,7 +1507,7 @@ name = "matplotlib-inline" version = "0.1.7" requires_python = ">=3.8" summary = "Inline Matplotlib backend for Jupyter" -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "traitlets", ] @@ -2007,7 +2038,7 @@ name = "parso" version = "0.8.4" requires_python = ">=3.6" summary = "A Python Parser" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "parso-0.8.4-py2.py3-none-any.whl", hash = "sha256:a418670a20291dacd2dddc80c377c5c3791378ee1e8d12bffc35420643d43f18"}, {file = "parso-0.8.4.tar.gz", hash = "sha256:eb3a7b58240fb99099a345571deecc0f9540ea5f4dd2fe14c2a99d6b281ab92d"}, @@ -2076,7 +2107,7 @@ files = [ name = "pexpect" version = "4.9.0" summary = "Pexpect allows easy control of interactive console applications." -groups = ["dev"] +groups = ["default", "dev"] marker = "sys_platform != \"win32\" and sys_platform != \"emscripten\"" dependencies = [ "ptyprocess>=0.5", @@ -2172,7 +2203,7 @@ name = "prompt-toolkit" version = "3.0.48" requires_python = ">=3.7.0" summary = "Library for building powerful interactive command lines in Python" -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "wcwidth", ] @@ -2259,7 +2290,7 @@ files = [ name = "ptyprocess" version = "0.7.0" summary = "Run a subprocess in a pseudo terminal" -groups = ["dev"] +groups = ["default", "dev"] marker = "sys_platform != \"win32\" and sys_platform != \"emscripten\"" files = [ {file = "ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35"}, @@ -2270,7 +2301,7 @@ files = [ name = "pure-eval" version = "0.2.3" summary = "Safely evaluate AST nodes without side effects" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0"}, {file = "pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42"}, @@ -3083,7 +3114,7 @@ files = [ name = "stack-data" version = "0.6.3" summary = "Extract data from python stack frames and tracebacks for informative displays" -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "asttokens>=2.1.0", "executing>=1.2.0", @@ -3190,7 +3221,7 @@ name = "traitlets" version = "5.14.3" requires_python = ">=3.8" summary = "Traitlets Python configuration system" -groups = ["dev"] +groups = ["default", "dev"] files = [ {file = "traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f"}, {file = "traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7"}, @@ -3289,7 +3320,7 @@ files = [ name = "wcwidth" version = "0.2.13" summary = "Measures the displayed width of unicode strings in a terminal" -groups = ["dev"] +groups = ["default", "dev"] dependencies = [ "backports-functools-lru-cache>=1.2.1; python_version < \"3.2\"", ] diff --git a/pyproject.toml b/pyproject.toml index d63f48c..7bb6756 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "sorcha @ git+https://github.com/B612-Asteroid-Institute/sorcha.git@cd5be9a06c6d24e1277cf2345bda9984f7097ede", "adam-core>=0.3.5", "quivr @ git+https://github.com/B612-Asteroid-Institute/quivr@2d8ae0b40bdfb75bcceff0c73d41a52d4bffb5dc", + "ipdb>=0.13.13", ] requires-python = "<3.13,>=3.11" readme = "README.md" diff --git a/src/adam_impact_study/analysis/main.py b/src/adam_impact_study/analysis/main.py index 1968f3f..3d1b295 100644 --- a/src/adam_impact_study/analysis/main.py +++ b/src/adam_impact_study/analysis/main.py @@ -544,3 +544,5 @@ def summarize_impact_study_results( if plot: make_analysis_plots(results, out_dir) + + return results diff --git a/src/adam_impact_study/cli/impact.py b/src/adam_impact_study/cli/impact.py index 9aba95e..4f3cac2 100644 --- a/src/adam_impact_study/cli/impact.py +++ b/src/adam_impact_study/cli/impact.py @@ -11,7 +11,7 @@ import pyarrow.compute as pc from adam_core.time import Timestamp -from adam_impact_study.analysis import plot_ip_over_time +from adam_impact_study.analysis.main import plot_ip_over_time from adam_impact_study.impacts_study import run_impact_study_all from adam_impact_study.types import ImpactorOrbits, RunConfiguration @@ -24,6 +24,7 @@ def run_impact_study( run_config: RunConfiguration, pointing_file: Optional[str] = None, orbit_id_filter: Optional[str] = None, + overwrite: bool = False, ) -> None: """Run impact study on provided orbits.""" # Load orbits directly from parquet @@ -84,6 +85,7 @@ def run_impact_study( monte_carlo_samples=run_config.monte_carlo_samples, max_processes=run_config.max_processes, seed=run_config.seed, + overwrite=overwrite, ) logger.info("Generating plots...") @@ -115,6 +117,12 @@ def main(): "--orbit-id-filter", help="Comma-delimited list of substrings to filter orbit id by", ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite existing results in run directory", + default=False, + ) parser.add_argument("--debug", action="store_true", help="Enable debug logging") args = parser.parse_args() @@ -154,6 +162,7 @@ def main(): run_config=run_config, pointing_file=args.pointing_file, orbit_id_filter=args.orbit_id_filter, + overwrite=args.overwrite, ) diff --git a/src/adam_impact_study/fo_od.py b/src/adam_impact_study/fo_od.py index 218d268..1becccb 100644 --- a/src/adam_impact_study/fo_od.py +++ b/src/adam_impact_study/fo_od.py @@ -79,9 +79,9 @@ def _create_fo_tmp_directory() -> str: Returns: str: The absolute path to the temporary directory populated with necessary FO files """ - base_tmp_dir = os.path.expanduser("~/.cache/adam_impact_study/ftmp") + base_tmp_dir = os.path.expanduser("~/.cache/adam_impact_study/") os.makedirs(base_tmp_dir, mode=0o770, exist_ok=True) - tmp_dir = tempfile.mkdtemp(dir=base_tmp_dir) + tmp_dir = tempfile.mkdtemp(dir=base_tmp_dir, prefix="fo_") os.chmod(tmp_dir, 0o770) tmp_dir = _populate_fo_directory(tmp_dir) return tmp_dir @@ -186,7 +186,9 @@ def run_fo_od( shutil.rmtree(fo_tmp_dir) if result.returncode != 0: - logger.warning(f"Find_Orb failed with return code {result.returncode}") + logger.warning( + f"Find_Orb failed with return code {result.returncode} for {len(observations)} observations in {fo_result_dir}" + ) logger.warning(f"{result.stdout}\n{result.stderr}") return Orbits.empty(), ADESObservations.empty(), "Find_Orb failed" diff --git a/src/adam_impact_study/impacts_study.py b/src/adam_impact_study/impacts_study.py index 5a9f597..8644301 100644 --- a/src/adam_impact_study/impacts_study.py +++ b/src/adam_impact_study/impacts_study.py @@ -45,7 +45,7 @@ def run_impact_study_all( assist_initial_dt: float, assist_adaptive_mode: int, max_processes: Optional[int] = 1, - overwrite: bool = True, + overwrite: bool = False, seed: Optional[int] = 13612, ) -> Tuple[WindowResult, ResultsTiming]: """ @@ -82,13 +82,14 @@ def __init__(self, *args, **kwargs): # If the run directory already exists, throw an exception # unless the user has specified the overwrite flag - if os.path.exists(f"{run_dir}"): - if not overwrite: - raise ValueError( - f"Run directory {run_dir} already exists. Set overwrite=True to overwrite." + if os.path.exists(run_dir): + if overwrite: + logger.warning(f"Overwriting run directory {run_dir}") + shutil.rmtree(f"{run_dir}") + else: + logger.warning( + f"Run directory {run_dir} already exists, attempting to continue previous run..." ) - logger.warning(f"Overwriting run directory {run_dir}") - shutil.rmtree(f"{run_dir}") os.makedirs(f"{run_dir}", exist_ok=True) @@ -140,9 +141,9 @@ def __init__(self, *args, **kwargs): result, timing = ray.get(finished[0]) impact_results = qv.concatenate([impact_results, result]) results_timings = qv.concatenate([results_timings, timing]) + while len(futures) > 0: finished, futures = ray.wait(futures, num_returns=1) - # import pdb; pdb.set_trace() result, timing = ray.get(finished[0]) impact_results = qv.concatenate([impact_results, result]) results_timings = qv.concatenate([results_timings, timing]) @@ -209,21 +210,54 @@ def run_impact_study_for_orbit( paths = get_study_paths(run_dir, orbit_id) # Serialize the ImpactorOrbit to a file for future analysis use - impactor_orbit.to_parquet( - f"{paths['orbit_base_dir']}/impact_orbits_{orbit_id}.parquet" - ) + impactor_orbit_file = f"{paths['orbit_base_dir']}/impactor_orbit.parquet" + if not os.path.exists(impactor_orbit_file): + impactor_orbit.to_parquet(impactor_orbit_file) + else: + impactor_orbit_saved = ImpactorOrbits.from_parquet(impactor_orbit_file) + impactor_orbit_table = impactor_orbit.flattened_table().drop_columns( + ["coordinates.covariance.values"] + ) + impactor_orbit_saved_table = ( + impactor_orbit_saved.flattened_table().drop_columns( + ["coordinates.covariance.values"] + ) + ) + assert impactor_orbit_table.equals( + impactor_orbit_saved_table + ), "ImpactorOrbit does not match saved version" + + timing_file = f"{paths['orbit_base_dir']}/timings.parquet" + if os.path.exists(timing_file): + timings = ResultsTiming.from_parquet(timing_file) + else: + timings = ResultsTiming.from_kwargs( + orbit_id=[orbit_id], + ) - # Run Sorcha to generate synthetic observations - sorcha_start_time = time.perf_counter() - observations = run_sorcha( - impactor_orbit, - pointing_file, - paths["sorcha_dir"], - seed=seed, - ) - sorcha_runtime = time.perf_counter() - sorcha_start_time - # Serialize the observations to a file for future analysis use - observations.to_parquet(f"{paths['sorcha_dir']}/observations_{orbit_id}.parquet") + observations_file = f"{paths['sorcha_dir']}/observations_{orbit_id}.parquet" + sorcha_runtime = None + if not os.path.exists(observations_file): + # Run Sorcha to generate synthetic observations + sorcha_start_time = time.perf_counter() + observations = run_sorcha( + impactor_orbit, + pointing_file, + paths["sorcha_dir"], + seed=seed, + ) + sorcha_runtime = time.perf_counter() - sorcha_start_time + # Serialize the observations to a file for future analysis use + observations.to_parquet( + f"{paths['sorcha_dir']}/observations_{orbit_id}.parquet" + ) + + # Update timings + timings = timings.set_column("sorcha_runtime", pa.array([sorcha_runtime])) + timings.to_parquet(timing_file) + else: + observations = Observations.from_parquet(observations_file) + logger.info(f"Loaded observations from {observations_file}") if len(observations) == 0: return WindowResult.empty(), ResultsTiming.empty() @@ -294,11 +328,13 @@ def run_impact_study_for_orbit( # Sort the results by observation_end for consistency. results = results.sort_by("observation_end") + # Update timings orbit_end_time = time.perf_counter() + total_runtime = orbit_end_time - orbit_start_time timings = ResultsTiming.from_kwargs( - orbit_id=[orbit_id], - total_runtime=[orbit_end_time - orbit_start_time], - sorcha_runtime=[sorcha_runtime], + orbit_id=timings.orbit_id, + total_runtime=[total_runtime], + sorcha_runtime=timings.sorcha_runtime, mean_od_runtime=[pc.mean(results.od_runtime)], total_od_runtime=[pc.sum(results.od_runtime)], mean_ip_runtime=[pc.mean(results.ip_runtime)], @@ -306,7 +342,7 @@ def run_impact_study_for_orbit( mean_window_runtime=[pc.mean(results.window_runtime)], total_window_runtime=[pc.sum(results.window_runtime)], ) - timings.to_parquet(f"{paths['orbit_base_dir']}/timings.parquet") + timings.to_parquet(timing_file) return results, timings @@ -362,6 +398,10 @@ def calculate_window_impact_probability( window = f"{start_night.as_py()}_{end_night.as_py()}" paths = get_study_paths(run_dir, orbit_id, window) + window_out_file = f"{paths['time_dir']}/window_result.parquet" + if os.path.exists(window_out_file): + return WindowResult.from_parquet(window_out_file) + # Get the start and end date of the observations, the number of # observations, and the number of unique nights observations_count = len(observations) @@ -370,6 +410,7 @@ def calculate_window_impact_probability( rejected_observations = ADESObservations.empty() + od_runtime = None try: od_start_time = time.perf_counter() orbit, rejected_observations, error = run_fo_od( @@ -384,7 +425,7 @@ def calculate_window_impact_probability( ) orbit_with_window.to_parquet(f"{paths['time_dir']}/orbit_with_window.parquet") except Exception as e: - return WindowResult.from_kwargs( + window_result = WindowResult.from_kwargs( orbit_id=[orbit_id], object_id=[object_id], window=[window], @@ -398,9 +439,11 @@ def calculate_window_impact_probability( error=[str(e)], od_runtime=[od_runtime], ) + window_result.to_parquet(window_out_file) + return window_result if error is not None: - return WindowResult.from_kwargs( + window_result = WindowResult.from_kwargs( orbit_id=[orbit_id], object_id=[object_id], window=[window], @@ -415,6 +458,8 @@ def calculate_window_impact_probability( error=[error], od_runtime=[od_runtime], ) + window_result.to_parquet(window_out_file) + return window_result days_until_impact_plus_thirty = ( int( @@ -424,6 +469,7 @@ def calculate_window_impact_probability( + 30 ) + ip_runtime = None try: ip_start_time = time.perf_counter() propagator = propagator_class() @@ -460,7 +506,7 @@ def calculate_window_impact_probability( ip_runtime = time.perf_counter() - ip_start_time except Exception as e: - return WindowResult.from_kwargs( + window_result = WindowResult.from_kwargs( orbit_id=[orbit_id], object_id=[object_id], window=[window], @@ -476,6 +522,8 @@ def calculate_window_impact_probability( od_runtime=[od_runtime], ip_runtime=[ip_runtime], ) + window_result.to_parquet(window_out_file) + return window_result window_end_time = time.perf_counter() diff --git a/src/adam_impact_study/sorcha_utils.py b/src/adam_impact_study/sorcha_utils.py index 165cdaa..d082a2a 100644 --- a/src/adam_impact_study/sorcha_utils.py +++ b/src/adam_impact_study/sorcha_utils.py @@ -108,6 +108,9 @@ def write_config_file_timeframe(impact_date: Timestamp, config_file: str) -> str SSP_track_window = 15 SSP_night_start_utc = 16.0 +[LINKING] +drop_unlinked = False + [OUTPUT] output_format = csv output_columns = basic