diff --git a/.gitignore b/.gitignore index 6e876b03..8c3d916d 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,7 @@ CLAUDE.md .claude/ # Profiling results -traces/ +tests/profiling/results/ # Trajectories results tests/trajectories/results/ @@ -18,7 +18,7 @@ uv.lock .idea/ # Generated images -images/ +tests/plots/results/ # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/tests/paths.py b/tests/paths.py index 0ed56e2d..5300c93c 100644 --- a/tests/paths.py +++ b/tests/paths.py @@ -1,4 +1,7 @@ from pathlib import Path TORCHJD_DIR = Path(__file__).parent.parent -TRACES_DIR = TORCHJD_DIR / "traces" +TESTS_DIR = Path(__file__).parent +PLOTS_RESULTS_DIR = TESTS_DIR / "plots" / "results" +PROFILING_RESULTS_DIR = TESTS_DIR / "profiling" / "results" +TRAJECTORIES_RESULTS_DIR = TESTS_DIR / "trajectories" / "results" diff --git a/tests/plots/static_plotter.py b/tests/plots/static_plotter.py index de5d6836..2fa8bf89 100644 --- a/tests/plots/static_plotter.py +++ b/tests/plots/static_plotter.py @@ -1,6 +1,5 @@ -import os - import torch +from paths import PLOTS_RESULTS_DIR from plotly import graph_objects as go from plots._utils import ( @@ -268,14 +267,14 @@ def main( ) fig.update_yaxes(range=[-0.1, 2.1], showgrid=False, zeroline=False, visible=False) - os.makedirs("images/", exist_ok=True) - fig.write_image(f"images/{filename}.pdf") + PLOTS_RESULTS_DIR.mkdir(exist_ok=True) + fig.write_image(PLOTS_RESULTS_DIR / f"{filename}.pdf") # Alternative: use .svg here and then convert to pdf using rsvg-convert. Install # [rsvg-convert](https://manpages.ubuntu.com/manpages/bionic/man1/rsvg-convert.1.html) and run: # `rsvg-convert -f pdf -o filename.pdf filename.svg` # To do that on all files at ones, run: # ``` - # for file in images/*.svg; do rsvg-convert -f pdf -o "${file%.svg}.pdf" "$file"; done + # for file in tests/plots/results/*.svg; do rsvg-convert -f pdf -o "${file%.svg}.pdf" "$file"; done # ``` diff --git a/tests/profiling/plot_memory_timeline.py b/tests/profiling/plot_memory_timeline.py index 4c2d9439..398dbc1e 100644 --- a/tests/profiling/plot_memory_timeline.py +++ b/tests/profiling/plot_memory_timeline.py @@ -13,7 +13,7 @@ matplotlib.use("Agg") # Use non-GUI backend to avoid tkinter dependency import matplotlib.pyplot as plt import numpy as np -from paths import TRACES_DIR +from paths import PROFILING_RESULTS_DIR @dataclass @@ -63,7 +63,7 @@ def plot_memory_timelines(experiment: str, folders: list[str]) -> None: cpu_timelines = [] cuda_timelines = [] for folder in folders: - path = TRACES_DIR / folder / f"{experiment}.json" + path = PROFILING_RESULTS_DIR / folder / f"{experiment}.json" cpu_timeline, cuda_timeline = extract_memory_timelines(path) cpu_timelines.append(cpu_timeline) cuda_timelines.append(cuda_timeline) @@ -104,7 +104,7 @@ def plot_memory_timelines(experiment: str, folders: list[str]) -> None: fig.tight_layout() - output_dir = Path(TRACES_DIR / "memory_timelines") + output_dir = Path(PROFILING_RESULTS_DIR / "memory_timelines") output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"{experiment}.png" print(f"\nSaving plot to: {output_path}") diff --git a/tests/profiling/run_profiler.py b/tests/profiling/run_profiler.py index 3cb46b12..9ee042b3 100644 --- a/tests/profiling/run_profiler.py +++ b/tests/profiling/run_profiler.py @@ -23,7 +23,7 @@ ) from utils.tensors import make_inputs_and_targets -from tests.paths import TRACES_DIR +from tests.paths import PROFILING_RESULTS_DIR from torchjd.aggregation import UPGrad, UPGradWeighting from torchjd.autogram import Engine @@ -98,7 +98,7 @@ def _save_and_print_trace( batch_size: int, ) -> None: filename = f"{factory}-bs{batch_size}-{DEVICE.type}.json" - output_dir = TRACES_DIR / method_name + output_dir = PROFILING_RESULTS_DIR / method_name output_dir.mkdir(parents=True, exist_ok=True) trace_path = output_dir / filename diff --git a/tests/trajectories/_paths.py b/tests/trajectories/_paths.py index e92f0514..764feded 100644 --- a/tests/trajectories/_paths.py +++ b/tests/trajectories/_paths.py @@ -1,23 +1,23 @@ from pathlib import Path -RESULTS_DIR = Path(__file__).parent / "results" +from tests.paths import TRAJECTORIES_RESULTS_DIR def get_params_dir(objective_key: str) -> Path: - return RESULTS_DIR / objective_key / "X" + return TRAJECTORIES_RESULTS_DIR / objective_key / "X" def get_values_dir(objective_key: str) -> Path: - return RESULTS_DIR / objective_key / "Y" + return TRAJECTORIES_RESULTS_DIR / objective_key / "Y" def get_param_plots_dir(objective_key: str) -> Path: - return RESULTS_DIR / objective_key / "param_plots" + return TRAJECTORIES_RESULTS_DIR / objective_key / "param_plots" def get_value_plots_dir(objective_key: str) -> Path: - return RESULTS_DIR / objective_key / "value_plots" + return TRAJECTORIES_RESULTS_DIR / objective_key / "value_plots" def get_distance_to_pf_plots_dir(objective_key: str) -> Path: - return RESULTS_DIR / objective_key / "distance_to_pf" + return TRAJECTORIES_RESULTS_DIR / objective_key / "distance_to_pf" diff --git a/tests/trajectories/optimize.py b/tests/trajectories/optimize.py index 9bd1b35c..1e78b480 100644 --- a/tests/trajectories/optimize.py +++ b/tests/trajectories/optimize.py @@ -18,6 +18,7 @@ import numpy as np import torch +from tests.paths import TRAJECTORIES_RESULTS_DIR from torchjd.aggregation import Stateful from trajectories._constants import ( AGGREGATORS, @@ -29,7 +30,7 @@ OBJECTIVES, ) from trajectories._optimization import optimize -from trajectories._paths import RESULTS_DIR, get_params_dir, get_values_dir +from trajectories._paths import get_params_dir, get_values_dir warnings.filterwarnings("ignore") @@ -88,7 +89,7 @@ def main() -> None: "learning_rates": learning_rates, "initial_points": initial_points, } - with open(RESULTS_DIR / objective_key / "metadata.json", "w") as f: + with open(TRAJECTORIES_RESULTS_DIR / objective_key / "metadata.json", "w") as f: json.dump(metadata, f) for aggregator_key in aggregator_keys: diff --git a/tests/trajectories/plot_distance_to_pf.py b/tests/trajectories/plot_distance_to_pf.py index de5074f0..c8d0eee2 100644 --- a/tests/trajectories/plot_distance_to_pf.py +++ b/tests/trajectories/plot_distance_to_pf.py @@ -15,9 +15,10 @@ import numpy as np import torch +from tests.paths import TRAJECTORIES_RESULTS_DIR from trajectories._constants import AGGREGATOR_ORDER, AGGREGATORS, LATEX_NAMES, OBJECTIVES from trajectories._pareto_utils import make_2d_pf_distance_fn -from trajectories._paths import RESULTS_DIR, get_distance_to_pf_plots_dir, get_values_dir +from trajectories._paths import get_distance_to_pf_plots_dir, get_values_dir from trajectories._plotters import ( MultiEvolutionPlotter, SquareBoxAspectSetter, @@ -51,7 +52,7 @@ def main() -> None: args = parser.parse_args() objective_key = args.objective - with open(RESULTS_DIR / objective_key / "metadata.json") as f: + with open(TRAJECTORIES_RESULTS_DIR / objective_key / "metadata.json") as f: metadata = json.load(f) values_dir = get_values_dir(objective_key) diff --git a/tests/trajectories/plot_params.py b/tests/trajectories/plot_params.py index d726eaa9..c5884035 100644 --- a/tests/trajectories/plot_params.py +++ b/tests/trajectories/plot_params.py @@ -14,6 +14,7 @@ import matplotlib.pyplot as plt import numpy as np +from tests.paths import TRAJECTORIES_RESULTS_DIR from trajectories._constants import ( AGGREGATOR_ORDER, AGGREGATORS, @@ -24,7 +25,7 @@ from trajectories._objectives import ElementWiseQuadratic, WithSPSMappingMixin from trajectories._optimization import compute_gradient_cosine_similarities from trajectories._pareto_utils import sample_2d_spss -from trajectories._paths import RESULTS_DIR, get_param_plots_dir, get_params_dir +from trajectories._paths import get_param_plots_dir, get_params_dir from trajectories._plotters import ( AxesPlotter, ContentLimAdjuster, @@ -61,7 +62,7 @@ def main() -> None: args = parser.parse_args() objective_key = args.objective - with open(RESULTS_DIR / objective_key / "metadata.json") as f: + with open(TRAJECTORIES_RESULTS_DIR / objective_key / "metadata.json") as f: metadata = json.load(f) params_dir = get_params_dir(objective_key) diff --git a/tests/trajectories/plot_values.py b/tests/trajectories/plot_values.py index 08c33968..145aaf1c 100644 --- a/tests/trajectories/plot_values.py +++ b/tests/trajectories/plot_values.py @@ -14,6 +14,7 @@ import matplotlib.pyplot as plt import numpy as np +from tests.paths import TRAJECTORIES_RESULTS_DIR from trajectories._constants import ( AGGREGATOR_ORDER, AGGREGATORS, @@ -23,7 +24,7 @@ ) from trajectories._objectives import WithSPSMappingMixin from trajectories._pareto_utils import compute_normalized_2d_pf_distances, sample_2d_pf -from trajectories._paths import RESULTS_DIR, get_value_plots_dir, get_values_dir +from trajectories._paths import get_value_plots_dir, get_values_dir from trajectories._plotters import ( ContentLimAdjuster, HeatmapPlotter, @@ -58,7 +59,7 @@ def main() -> None: args = parser.parse_args() objective_key = args.objective - with open(RESULTS_DIR / objective_key / "metadata.json") as f: + with open(TRAJECTORIES_RESULTS_DIR / objective_key / "metadata.json") as f: metadata = json.load(f) values_dir = get_values_dir(objective_key)