katabatic 0.2.0__py3-none-any.whl

katabatic/__init__.py

@@ -0,0 +1,15 @@
+"""
+Katabatic package initializer.
+Synthetic tabular data generation, pipelines, and evaluation.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("katabatic")
+except PackageNotFoundError:
+    __version__ = "0.0.0.dev0"
+
+from . import models, pipeline, utils
+
+__all__ = ["__version__", "models", "pipeline", "utils"]

katabatic/artifacts/__init__.py

@@ -0,0 +1,27 @@
+from katabatic.artifacts.base import ArtifactStore
+from katabatic.artifacts.dataset_split import (
+    write_dataset_artifact,
+    write_dataset_artifact_presplit,
+)
+from katabatic.artifacts.ids import new_eval_id, new_split_id, new_train_id
+from katabatic.artifacts.local import LocalArtifactStore
+from katabatic.artifacts.refs import (
+    DatasetRef,
+    EvaluationRef,
+    ModelRef,
+    artifact_path_segment,
+)
+
+__all__ = [
+    "ArtifactStore",
+    "LocalArtifactStore",
+    "DatasetRef",
+    "ModelRef",
+    "EvaluationRef",
+    "artifact_path_segment",
+    "write_dataset_artifact",
+    "write_dataset_artifact_presplit",
+    "new_split_id",
+    "new_train_id",
+    "new_eval_id",
+]

katabatic/artifacts/base.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class ArtifactStore(ABC):
+    @abstractmethod
+    def save_json(self, path: str, data: dict) -> None: ...
+
+    @abstractmethod
+    def load_json(self, path: str) -> dict: ...
+
+    @abstractmethod
+    def save_bytes(self, path: str, data: bytes) -> None: ...
+
+    @abstractmethod
+    def open_path(self, path: str) -> Path: ...
+
+    @abstractmethod
+    def exists(self, path: str) -> bool: ...

katabatic/artifacts/dataset_split.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+import pandas as pd
+
+from katabatic.artifacts.base import ArtifactStore
+from katabatic.artifacts.ids import new_split_id
+from katabatic.artifacts.refs import DatasetRef
+from katabatic.utils.split_dataset import compute_train_test_split
+from katabatic.utils.train_test_consistency import sanity_check_train_test
+
+
+def _copy_extra_assets(
+    source_dir: Path | None,
+    extra_dir: Path,
+) -> None:
+    """Copy info.json and TabSyn-style *.npy from source_dir into extra_dir if present."""
+    if source_dir is None or not source_dir.is_dir():
+        return
+    extra_dir.mkdir(parents=True, exist_ok=True)
+    info = source_dir / "info.json"
+    if info.exists():
+        shutil.copy2(info, extra_dir / "info.json")
+    for name in (
+        "X_num_train.npy",
+        "X_cat_train.npy",
+        "y_train.npy",
+        "X_num_test.npy",
+        "X_cat_test.npy",
+        "y_test.npy",
+    ):
+        p = source_dir / name
+        if p.exists():
+            shutil.copy2(p, extra_dir / name)
+
+
+def _persist_frames_to_store(
+    store: ArtifactStore,
+    ref: DatasetRef,
+    df_train: pd.DataFrame,
+    df_test: pd.DataFrame,
+    extra_source_dir: Path | None,
+) -> tuple[pd.Series, pd.Series]:
+    label_name = df_train.columns[-1]
+    X_train, y_train = df_train.iloc[:, :-1], df_train.iloc[:, -1]
+    X_test, y_test = df_test.iloc[:, :-1], df_test.iloc[:, -1]
+    y_train.name = label_name
+    y_test.name = label_name
+
+    train_dir = store.open_path(ref.train_relpath)
+    test_dir = store.open_path(ref.test_relpath)
+    extra_dir = store.open_path(ref.extra_relpath)
+    train_dir.mkdir(parents=True, exist_ok=True)
+    test_dir.mkdir(parents=True, exist_ok=True)
+    extra_dir.mkdir(parents=True, exist_ok=True)
+
+    df_train.to_csv(train_dir / "train_full.csv", index=False)
+    df_test.to_csv(test_dir / "test_full.csv", index=False)
+    X_train.to_csv(train_dir / "x_train.csv", index=False)
+    y_train.to_csv(train_dir / "y_train.csv", index=False, header=True)
+    X_test.to_csv(test_dir / "x_test.csv", index=False)
+    y_test.to_csv(test_dir / "y_test.csv", index=False, header=True)
+
+    shutil.copy2(train_dir / "x_train.csv", extra_dir / "x_train.csv")
+    shutil.copy2(train_dir / "y_train.csv", extra_dir / "y_train.csv")
+    shutil.copy2(test_dir / "x_test.csv", extra_dir / "x_test.csv")
+    shutil.copy2(test_dir / "y_test.csv", extra_dir / "y_test.csv")
+
+    _copy_extra_assets(extra_source_dir, extra_dir)
+    return y_train, y_test
+
+
+def write_dataset_artifact(
+    store: ArtifactStore,
+    *,
+    input_csv: str | Path,
+    dataset_name: str,
+    test_size: float = 0.2,
+    seed: int = 42,
+    dataset_version: str | None = None,
+    extra_source_dir: str | Path | None = None,
+) -> DatasetRef:
+    """
+    Split input_csv into train/ and test/ under datasets/<name>/<version>/,
+    optional extra/ from extra_source_dir (e.g. TabSyn npy + info.json).
+    """
+    input_csv = Path(input_csv)
+    version = dataset_version or new_split_id()
+    ref = DatasetRef(dataset_name=dataset_name, dataset_version=version)
+
+    df = pd.read_csv(input_csv)
+    print(f"Loaded data with shape: {df.shape}")
+
+    df_train, df_test, _, _, _, _ = compute_train_test_split(
+        df, test_size=test_size, seed=seed
+    )
+
+    src = Path(extra_source_dir) if extra_source_dir else input_csv.parent
+    y_train, y_test = _persist_frames_to_store(store, ref, df_train, df_test, src)
+
+    print("Train label distribution:\n", y_train.value_counts(normalize=True))
+    print("Test label distribution:\n", y_test.value_counts(normalize=True))
+    print(f"Saved dataset artifact under {ref.root_relpath}")
+    return ref
+
+
+def write_dataset_artifact_presplit(
+    store: ArtifactStore,
+    *,
+    train_csv: str | Path,
+    test_csv: str | Path,
+    dataset_name: str,
+    dataset_version: str | None = None,
+    extra_source_dir: str | Path | None = None,
+) -> DatasetRef:
+    """
+    Write user-provided train/test CSVs under datasets/<name>/<version>/
+    without performing a random split.
+    """
+    train_csv = Path(train_csv)
+    test_csv = Path(test_csv)
+    version = dataset_version or new_split_id()
+    ref = DatasetRef(dataset_name=dataset_name, dataset_version=version)
+
+    df_train = pd.read_csv(train_csv)
+    df_test = pd.read_csv(test_csv)
+    print(f"Loaded presplit train {df_train.shape}, test {df_test.shape}")
+    sanity_check_train_test(df_train, df_test)
+
+    src = Path(extra_source_dir) if extra_source_dir else train_csv.parent
+    y_train, y_test = _persist_frames_to_store(store, ref, df_train, df_test, src)
+
+    print("Train label distribution:\n", y_train.value_counts(normalize=True))
+    print("Test label distribution:\n", y_test.value_counts(normalize=True))
+    print(f"Saved dataset artifact under {ref.root_relpath}")
+    return ref

katabatic/artifacts/ids.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+
+def _utc_stamp() -> str:
+    return datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
+
+
+def new_split_id() -> str:
+    return f"split-{_utc_stamp()}"
+
+
+def new_train_id() -> str:
+    return f"train-{_utc_stamp()}"
+
+
+def new_eval_id() -> str:
+    return f"eval-{_utc_stamp()}"

katabatic/artifacts/local.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from katabatic.artifacts.base import ArtifactStore
+
+
+class LocalArtifactStore(ArtifactStore):
+    def __init__(self, root: str | Path = "artifacts") -> None:
+        self.root = Path(root)
+
+    def _full(self, path: str) -> Path:
+        return self.root / path
+
+    def save_json(self, path: str, data: dict) -> None:
+        full = self._full(path)
+        full.parent.mkdir(parents=True, exist_ok=True)
+        full.write_text(json.dumps(data, indent=2), encoding="utf-8")
+
+    def load_json(self, path: str) -> dict:
+        return json.loads(self._full(path).read_text(encoding="utf-8"))
+
+    def save_bytes(self, path: str, data: bytes) -> None:
+        full = self._full(path)
+        full.parent.mkdir(parents=True, exist_ok=True)
+        full.write_bytes(data)
+
+    def open_path(self, path: str) -> Path:
+        return self._full(path)
+
+    def exists(self, path: str) -> bool:
+        return self._full(path).exists()

katabatic/artifacts/refs.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+
+def artifact_path_segment(name: str) -> str:
+    """
+    Sanitize a logical name for use inside artifact directory / file segments.
+    Allows lowercase letters, digits, hyphen; coerces other chars to hyphen.
+    """
+    s = name.strip().lower()
+    s = re.sub(r"[^a-z0-9-]+", "-", s)
+    s = re.sub(r"-{2,}", "-", s).strip("-")
+    return s or "unnamed"
+
+
+@dataclass(frozen=True)
+class DatasetRef:
+    dataset_name: str
+    dataset_version: str
+    """Root relative to artifact store: datasets/<name>/<version>"""
+
+    @property
+    def root_relpath(self) -> str:
+        return f"datasets/{self.dataset_name}/{self.dataset_version}"
+
+    @property
+    def train_relpath(self) -> str:
+        return f"{self.root_relpath}/train"
+
+    @property
+    def test_relpath(self) -> str:
+        return f"{self.root_relpath}/test"
+
+    @property
+    def extra_relpath(self) -> str:
+        return f"{self.root_relpath}/extra"
+
+
+@dataclass(frozen=True)
+class ModelRef:
+    model_name: str
+    dataset_name: str
+    dataset_version: str
+    train_run_id: str
+
+    @property
+    def root_relpath(self) -> str:
+        m = artifact_path_segment(self.model_name)
+        d = artifact_path_segment(self.dataset_name)
+        return f"models/{m}_{d}_{self.train_run_id}"
+
+    @property
+    def state_relpath(self) -> str:
+        return f"{self.root_relpath}/state"
+
+    @property
+    def synthetic_relpath(self) -> str:
+        return f"{self.root_relpath}/synthetic"
+
+
+@dataclass(frozen=True)
+class EvaluationRef:
+    evaluation_type: str
+    eval_run_id: str
+    model_name: str
+    dataset_name: str
+    dataset_version: str
+    train_run_id: str
+    test_dataset_version: str
+
+    @property
+    def root_relpath(self) -> str:
+        """Mirrors :attr:`ModelRef.root_relpath` stem under ``evaluations/``."""
+        m = artifact_path_segment(self.model_name)
+        d = artifact_path_segment(self.dataset_name)
+        return f"evaluations/{m}_{d}_{self.train_run_id}"
+
+    @property
+    def metrics_relpath(self) -> str:
+        et = artifact_path_segment(self.evaluation_type)
+        return f"{self.root_relpath}/{et}_metrics.json"
+
+    @property
+    def report_relpath(self) -> str:
+        et = artifact_path_segment(self.evaluation_type)
+        return f"{self.root_relpath}/{et}_report.csv"

katabatic/cli/__init__.py

@@ -0,0 +1,3 @@
+from .commands import init_model
+
+__all__ = ["init_model"]

katabatic/cli/commands/__init__.py

@@ -0,0 +1,3 @@
+from .model_init import init_model
+
+__all__ = ["init_model"]

katabatic/cli/commands/model_init.py

@@ -0,0 +1,201 @@
+import os
+import re
+from pathlib import Path
+
+
+def _find_models_dict_bounds(text: str) -> tuple[int, int, str]:
+    """
+    Return (open_brace_idx, close_brace_idx, indent) for the _models dict.
+    Uses brace counting and skips over strings.
+    """
+    m = re.search(r"_models\s*:\s*Dict\[[^\]]*\]\s*=\s*{", text)
+    if not m:
+        m = re.search(r"_models\s*=\s*{", text)
+    if not m:
+        raise ValueError("Could not find `_models = { ... }` in registry.py")
+
+    open_idx = text.find("{", m.start())
+    if open_idx == -1:
+        raise ValueError("Could not find opening `{` for _models dictionary")
+
+    line_start = text.rfind("\n", 0, open_idx) + 1
+    indent = re.match(r"[ \t]*", text[line_start:open_idx]).group(0)
+
+    depth = 0
+    i = open_idx
+    in_str: str | None = None
+    escaped = False
+    while i < len(text):
+        ch = text[i]
+        if in_str:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == in_str:
+                in_str = None
+        else:
+            if ch in ("'", '"'):
+                in_str = ch
+            elif ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return open_idx, i, indent
+        i += 1
+
+    raise ValueError("Unbalanced braces in _models dictionary")
+
+
+def _serialize_list_str(xs: list[str]) -> str:
+    return "[" + ", ".join(repr(s) for s in xs) + "]"
+
+
+def _find_project_root() -> Path:
+    """Return repo root when run from a checkout, else cwd."""
+    for candidate in [Path.cwd(), *Path.cwd().parents]:
+        if (candidate / "pyproject.toml").is_file() and (
+            candidate / "katabatic"
+        ).is_dir():
+            return candidate
+    return Path.cwd()
+
+
+def init_model(model_name: str, dependencies: list[str] | None = None) -> None:
+    """Initialize a new model structure aligned with the base `Model` API.
+
+    Args:
+        model_name: Name of the model to create (e.g. "Rlig" or "rlig")
+        dependencies: Optional list of Python package dependencies for the model.
+                      These will be embedded in the subclass' get_required_dependencies().
+    """
+    dir_name = re.sub(r"(?<!^)(?=[A-Z])", "_", model_name).lower()
+    if not re.fullmatch(r"[a-z][a-z0-9_]*", dir_name):
+        raise ValueError(
+            "Invalid model name. Use letters, numbers and underscores; start with a letter."
+        )
+    class_name = "".join(word.capitalize() for word in dir_name.split("_"))
+
+    from katabatic.models.registry import ModelRegistry
+
+    models_dir = Path(__file__).parents[2] / "models"
+    new_model_dir = models_dir / dir_name
+
+    if new_model_dir.exists():
+        raise ValueError(f"Model directory {dir_name} already exists")
+
+    if dir_name in ModelRegistry.get_available_models():
+        raise ValueError(f"Model {dir_name} is already registered")
+
+    new_model_dir.mkdir(parents=True)
+
+    # __init__.py
+    (new_model_dir / "__init__.py").write_text(
+        f'''"""Module for {class_name} model."""
+from .models import {class_name}
+
+__all__ = ["{class_name}"]
+'''
+    )
+
+    deps = dependencies or []
+
+    (new_model_dir / "models.py").write_text(
+        f'''"""Implementation of {class_name} model."""
+from typing import Any, Union
+import numpy as np
+import pandas as pd
+
+from katabatic.models.base_model import Model
+from . import utils as _utils
+
+
+class {class_name}(Model):
+    """Implementation of {class_name} model."""
+
+    def __init__(self):
+        super().__init__()
+        # Validate runtime dependencies early to fail fast
+        self.check_dependencies()
+
+    @classmethod
+    def get_required_dependencies(cls) -> list[str]:
+        # Generated by init_model; adjust as your model evolves.
+        return {deps!r}
+
+    def train(self, *args, **kwargs) -> "Model":
+        \"\"\"Train the model on the given data.\"\"\"
+        raise NotImplementedError
+
+    def evaluate(self, *args, **kwargs) -> float:
+        \"\"\"Evaluate the model performance and return a score.\"\"\"
+        raise NotImplementedError
+
+    def sample(self, *args, **kwargs) -> Union[np.ndarray, pd.DataFrame]:
+        \"\"\"Generate synthetic samples.\"\"\"
+        raise NotImplementedError
+'''
+    )
+
+    (new_model_dir / "utils.py").write_text(
+        f'''"""Utility functions for {class_name} model."""
+
+# Add your utility functions here.
+'''
+    )
+
+    registry_path = models_dir / "registry.py"
+    registry_content = registry_path.read_text()
+
+    open_idx, close_idx, indent = _find_models_dict_bounds(registry_content)
+
+    entry_indent = indent + "    "
+    new_entry = (
+        f"{entry_indent}'{dir_name}': {{\n"
+        f"{entry_indent}    'module': 'katabatic.models.{dir_name}.models',\n"
+        f"{entry_indent}    'class': '{class_name}',\n"
+        f"{entry_indent}    'dependencies': {_serialize_list_str(deps)},\n"
+        f"{entry_indent}    'extra': '{dir_name}'\n"
+        f"{entry_indent}}}"
+    )
+
+    # Preserve the indented closing brace line (close_idx points at `}` only).
+    line_start = registry_content.rfind("\n", 0, close_idx) + 1
+    before = registry_content[:line_start].rstrip()
+    after = registry_content[line_start:]
+
+    needs_comma = not before.endswith("{")
+    if needs_comma and not before.endswith(","):
+        before = before + ","
+
+    if not before.endswith("\n"):
+        before = before + "\n"
+
+    new_registry = before + new_entry + ",\n" + after
+
+    registry_path.write_text(new_registry)
+
+    scripts_dir = _find_project_root() / "scripts"
+    scripts_dir.mkdir(parents=True, exist_ok=True)
+    setup_script_path = scripts_dir / f"setup_{dir_name}.sh"
+    setup_script_path.write_text(
+        f"""#!/bin/bash
+# Setup script for {class_name} model dependencies
+
+pip install katabatic[{dir_name}]
+"""
+    )
+    os.chmod(setup_script_path, 0o755)  # nosec B103: executable bit on generated setup script, intended
+
+    print(f"""Successfully:
+1. Created {class_name} model structure in {new_model_dir}
+2. Added model to registry
+3. Created setup script at {setup_script_path}
+
+Next steps:
+1. Add your model's dependencies to pyproject.toml in:
+   - [tool.poetry.dependencies] (as optional = true)
+   - [tool.poetry.extras] {dir_name} = [...]
+2. Implement train/evaluate/sample in {new_model_dir}/models.py
+3. Add any utility functions in {new_model_dir}/utils.py""")

katabatic/cli/commands/pin_notebook_kernel.py

@@ -0,0 +1,101 @@
+"""Pin Jupyter kernelspec for notebooks to the project virtualenv."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+DEFAULT_KERNEL = "katabatic-venv"
+DEFAULT_DISPLAY = "Python (Katabatic .venv)"
+
+
+def find_project_root(start: Path | None = None) -> Path:
+    """Walk up from *start* (default cwd) to find the Katabatic project root."""
+    for candidate in [start or Path.cwd(), *(start or Path.cwd()).parents]:
+        if (candidate / "pyproject.toml").is_file() and (
+            candidate / "katabatic"
+        ).is_dir():
+            return candidate
+    return start or Path.cwd()
+
+
+def install_kernelspec(
+    *,
+    venv_prefix: Path,
+    python_exe: Path,
+    name: str,
+    display: str,
+) -> Path:
+    subprocess.run(
+        [
+            str(python_exe),
+            "-m",
+            "ipykernel",
+            "install",
+            "--prefix",
+            str(venv_prefix),
+            "--name",
+            name,
+            "--display-name",
+            display,
+        ],
+        check=True,
+    )
+    return venv_prefix / "share" / "jupyter" / "kernels" / name
+
+
+def patch_notebook(path: Path, *, name: str, display: str) -> None:
+    nb = json.loads(path.read_text(encoding="utf-8"))
+    meta = nb.setdefault("metadata", {})
+    meta["kernelspec"] = {
+        "display_name": display,
+        "language": "python",
+        "name": name,
+    }
+    path.write_text(json.dumps(nb, indent=1), encoding="utf-8")
+
+
+def run_pin_notebook_kernel(
+    notebooks: list[str | Path],
+    *,
+    repo_root: Path | None = None,
+    kernel_name: str = DEFAULT_KERNEL,
+    display_name: str = DEFAULT_DISPLAY,
+) -> int:
+    root = (repo_root or find_project_root()).resolve()
+    py = root / ".venv" / "bin" / "python"
+    if not py.is_file():
+        print(f"error: expected venv python at {py}", file=sys.stderr)
+        print("  run: poetry install", file=sys.stderr)
+        return 1
+
+    print(f"Installing kernelspec {kernel_name!r} into {root / '.venv'} ...")
+    dest = install_kernelspec(
+        venv_prefix=root / ".venv",
+        python_exe=py,
+        name=kernel_name,
+        display=display_name,
+    )
+    print(f"  -> {dest}")
+
+    patched = 0
+    for nb_path in notebooks:
+        nb = Path(nb_path).resolve()
+        if nb.suffix != ".ipynb":
+            print(f"skip (not .ipynb): {nb}", file=sys.stderr)
+            continue
+        patch_notebook(nb, name=kernel_name, display=display_name)
+        print(f"Patched kernelspec in {nb}")
+        patched += 1
+
+    if patched == 0:
+        print("error: no notebooks were patched", file=sys.stderr)
+        return 1
+
+    print("\nIn Cursor: reload window, then Select Kernel → pick", repr(display_name))
+    print(
+        "If the picker still spins, use Command Palette → Python: Select Interpreter → .venv"
+    )
+    return 0

katabatic/cli/commands/register_dataset.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from katabatic.artifacts.local import LocalArtifactStore
+from katabatic.datasets.compatibility import check_dataset_for_model
+from katabatic.datasets.registry import DatasetRegistry
+
+
+def register_dataset_cli(
+    dataset_name: str,
+    csv_path: str,
+    *,
+    target_column: str | None = None,
+    artifact_root: str | None = None,
+    check_model: str | None = None,
+) -> None:
+    root = Path(artifact_root or "artifacts")
+    store = LocalArtifactStore(root)
+    reg = DatasetRegistry(store)
+    entry = reg.register_if_absent(dataset_name, csv_path, target_column=target_column)
+    print(
+        f"Registered dataset {dataset_name!r} at {store.root / 'registry' / 'datasets.json'}"
+    )
+    print(
+        f"  task={entry['task']!r}, n_rows={entry['n_rows']}, target={entry['target_column']!r}"
+    )
+    if check_model:
+        ok, msg = check_dataset_for_model(entry, check_model)
+        print(
+            f"  compatibility with {check_model!r}: {msg}" + ("" if ok else " (failed)")
+        )
+        if not ok:
+            raise SystemExit(1)