Sweep Visualizations

.gitignore

@@ -138,3 +138,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+viz

elk/__main__.py

@@ -5,6 +5,7 @@
 from simple_parsing import ArgumentParser
 
 from elk.evaluation.evaluate import Eval
+from elk.plotting.command import Plot
 from elk.training.sweep import Sweep
 from elk.training.train import Elicit
 
@@ -13,7 +14,7 @@
 class Command:
  """Some top-level command"""
 
- command: Elicit | Eval | Sweep
+ command: Elicit | Eval | Sweep | Plot
 
  def execute(self):
  return self.command.execute()

elk/files.py

@@ -6,6 +6,10 @@
 from pathlib import Path
 
 
+def sweeps_dir() -> Path:
+ return elk_reporter_dir() / "sweeps"
+
+
 def elk_reporter_dir() -> Path:
  """Return the directory where reporter checkpoints and logs are stored."""
  env_dir = os.environ.get("ELK_DIR", None)

elk/plotting/command.py

@@ -0,0 +1,27 @@
+import os
+from dataclasses import dataclass, field
+
+from elk.files import sweeps_dir
+from elk.plotting.visualize import visualize_sweep
+
+
+@dataclass
+class Plot:
+ sweeps: list[str] = field(default_factory=list)
+
+ def execute(self):
+ sweeps_root_dir = sweeps_dir()
+ sweep = max(sweeps_root_dir.iterdir(), key=os.path.getctime)
+ if self.sweeps:
+ sweep = sweeps_root_dir / self.sweeps[0]
+ if not sweep.exists():
+ print(f"No sweep with name {self.sweeps[0]} found in {sweeps_root_dir}")
+ return
+ if len(self.sweeps) > 1:
+ # TODO support more than one sweep
+ print(
+ f"""{len(self.sweeps)} paths specified.
+ Only one sweep is supported at this time."""
+ )
+ else:
+ visualize_sweep(sweep)

elk/plotting/utils.py

@@ -0,0 +1,78 @@
+import shutil
+from pathlib import Path
+
+from rich.console import Console
+from rich.table import Table
+
+
+def set_subplot_title_font_size(fig, font_size=14, font_family="Arial"):
+ fig.update_annotations(
+ font=dict(
+ family=font_family,
+ size=font_size,
+ )
+ )
+ return fig
+
+
+def set_legend_font_size(fig, font_size=8, font_family="Arial"):
+ fig.update_layout(
+ xaxis=dict(
+ tickfont=dict(
+ family=font_family,
+ size=font_size,
+ )
+ ),
+ yaxis=dict(
+ tickfont=dict(
+ family=font_family,
+ size=font_size,
+ )
+ ),
+ legend=dict(
+ font=dict(
+ family=font_family,
+ size=font_size,
+ )
+ ),
+ )
+ return fig
+
+
+def display_table(pivot_table):
+ console = Console()
+
+ table = Table(show_header=True, header_style="bold magenta", show_lines=True)
+
+ table.add_column("Run Name")
+ for column in pivot_table.columns:
+ table.add_column(str(column))
+
+ for index, row in pivot_table.iterrows():
+ table.add_row(str(index), *[str(value) for value in row])
+
+ console.print(table)
+
+
+def restructure_to_sweep(
+ elk_reporters: Path, data_path, new_name: str
+): # usually /elk-reporters/*
+ for model_repo_path in elk_reporters.iterdir():
+ for model_path in model_repo_path.iterdir():
+ for dataset_path in model_path.iterdir():
+ for run_path in dataset_path.iterdir():
+ new_path = (
+ data_path
+ / new_name
+ / run_path.name
+ / model_repo_path.name
+ / model_path.name
+ / dataset_path.name
+ )
+ if not new_path.exists():
+ new_path.mkdir(parents=True)
+ for file in run_path.iterdir():
+ if file.is_file():
+ shutil.copy(file, new_path / file.name)
+ else:
+ shutil.copytree(file, new_path / file.name)