chore: senior-eng audit fixes (datasets, sa docs, viz heatmap, public api)

Surgical fixes from a top-to-bottom code audit. No behaviour change for
working code paths; each fix repairs a real corner case the audit found.

datasets._default_data_dir
  Wheel installs from PyPI have no repo above ``site-packages/qqa/``,
  so ``parents[2] / "data"`` resolved to a non-existent path and the
  loaders crashed deep inside ``os.listdir`` with no actionable hint.
  Now: prefer ``QQA_DATA_DIR``, fall back to the source-tree path when
  it exists, then ``Path.cwd() / "data"``. ``_resolve`` raises a
  FileNotFoundError with the remediation when the directory is missing.

sa.simulated_annealing docstring
  Used to claim BinaryInstanceRelaxation was supported, but the
  function rejects batched-instance problems with NotImplementedError.
  Docstring now matches the implementation.

visualization.plot_solution_heatmap
  Square (N, N) one-hot solutions (TSP, QAP, NQueens) used to silently
  collapse to row 0 when the caller did not pass ``problem``. Show the
  full matrix in that fall-through branch instead.

tests/test_public_api
  Listed only the pre-0.4 names. Extend with PSpinGlass /
  RandomFieldIsing / MinimumDominatingSet / SAResult / enable_tf32 /
  polish / simulated_annealing / warmstart / *Instance, and add one
  mechanical test that ``__all__`` ⊆ what is actually importable.

Made-with: Cursor

Author: Yuma Ichikawa

src/qqa/datasets.py

@@ -33,20 +33,49 @@
 
 
 def _default_data_dir() -> Path:
+    """Resolve the on-disk benchmark dataset directory.
+
+    Priority order:
+    1. ``$QQA_DATA_DIR`` if set — explicit user override.
+    2. ``<repo_root>/data`` if this module lives inside the source tree
+       (``src/qqa/datasets.py`` => ``parents[2] == repo_root``).
+    3. ``./data`` next to the current working directory — a sensible
+       fallback for wheel installs from PyPI where no source tree exists.
+
+    Loaders raise a clear ``FileNotFoundError`` when the resolved
+    directory does not contain the requested benchmark, so callers always
+    get an actionable message ("set QQA_DATA_DIR or pass path=") rather
+    than an opaque ``listdir`` error.
+    """
+
     env = os.environ.get("QQA_DATA_DIR")
     if env:
         return Path(env).expanduser().resolve()
-    # src/qqa/datasets.py -> repo_root = parents[2]
-    return _THIS.parents[2] / "data"
+    # When installed in editable / source mode, ``parents[2]`` is the
+    # repository root and ``parents[2] / "data"`` ships ``mis/er-small``.
+    repo_data = _THIS.parents[2] / "data"
+    if repo_data.is_dir():
+        return repo_data
+    # Wheel install fallback — let the user opt in via cwd/data.
+    return Path.cwd() / "data"
 
 
 DATA_DIR: Path = _default_data_dir()
 
 
 def _resolve(path: str | os.PathLike | None, default_subpath: str) -> Path:
     if path is not None:
-        return Path(path).expanduser().resolve()
-    return _default_data_dir() / default_subpath
+        resolved = Path(path).expanduser().resolve()
+    else:
+        resolved = _default_data_dir() / default_subpath
+    if not resolved.is_dir():
+        raise FileNotFoundError(
+            f"Benchmark directory {resolved!s} does not exist. "
+            "Set the QQA_DATA_DIR environment variable to point at the "
+            "QQA4CO repository's ``data/`` directory, or pass an explicit "
+            "``path=`` argument to this loader."
+        )
+    return resolved
 
 
 def _load_pickle(p: Path):

src/qqa/sa.py

@@ -109,9 +109,15 @@ def simulated_annealing(
     Parameters
     ----------
     problem:
-        Any :class:`~qqa.problems.COProblem` exposing ``loss_fn(x)`` and a
-        ``relaxation`` attribute (BinaryRelaxation, BinaryInstanceRelaxation,
-        or SpinRelaxation). CategoricalRelaxation is not yet supported.
+        Any single-instance :class:`~qqa.problems.COProblem` exposing
+        ``loss_fn(x)`` and a ``relaxation`` attribute
+        (:class:`~qqa.relaxation.BinaryRelaxation` or
+        :class:`~qqa.relaxation.SpinRelaxation`).
+        :class:`~qqa.relaxation.CategoricalRelaxation` and batched-instance
+        problems (those exposing ``num_instance``) are rejected at the API
+        boundary with :class:`NotImplementedError`; iterate over instances
+        and call :func:`simulated_annealing` per instance, or use
+        :func:`qqa.anneal` which handles both natively.
     sol_size:
         Number of independent SA chains run in parallel on GPU.
     num_sweeps:

src/qqa/visualization.py

@@ -602,18 +602,23 @@ def plot_solution_heatmap(
         sol = sol.detach().cpu().numpy()
     sol = np.asarray(sol)
     if sol.ndim == 2:
-        # Two valid 2-D shapes reach here:
+        # Three valid 2-D shapes reach here:
         #   * batched-instance problems: ``(num_instance, max_node)`` — keep
         #     the whole matrix so each row is one instance's solution.
         #   * single-instance categorical problems: ``(N, K)`` — collapse to
         #     the chosen category per variable via argmax so the heatmap
         #     shows a single row of class indices.
+        #   * permutation problems (TSP, QAP, NQueens) where ``best_sol`` is
+        #     a square one-hot ``(N, N)`` matrix — show the full matrix so
+        #     the assignment structure is visible. When ``problem`` is not
+        #     provided we still show the full matrix rather than silently
+        #     collapsing to the first row.
         if problem is not None and getattr(problem, "num_instance", None) is not None:
             arr = sol
         elif problem is not None and getattr(problem, "num_category", None) is not None:
             arr = np.argmax(sol, axis=1)[None, :].astype(float)
         else:
-            arr = sol[0][None, :]
+            arr = sol
     elif sol.ndim == 1:
         arr = sol[None, :]
     else:

tests/test_public_api.py

@@ -64,6 +64,18 @@
     "PopulationTracker",
     "TrajectoryTracker",
     "Relaxation",
+    # Post-0.4.0 additions.
+    "MinimumDominatingSet",
+    "PSpinGlass",
+    "RandomFieldIsing",
+    "SAResult",
+    "enable_tf32",
+    "polish",
+    "simulated_annealing",
+    "warmstart",
+    "MaxCliqueInstance",
+    "MaxCutInstance",
+    "MaximumIndependentSetInstance",
 ]
 
 
@@ -73,6 +85,22 @@ def test_top_level_export_exists(name: str) -> None:
     assert name in qqa.__all__, f"qqa.{name} must be listed in qqa.__all__"
 
 
+def test_public_api_set_matches_dunder_all() -> None:
+    """``__all__`` must agree with what is reachable as ``qqa.<name>``.
+
+    Catches the common slip of adding a class to ``__all__`` while
+    forgetting to import it (or vice versa).
+    """
+
+    advertised = set(qqa.__all__) - {"__version__"}
+    actually_present = {n for n in advertised if hasattr(qqa, n)}
+    missing = sorted(advertised - actually_present)
+    assert not missing, (
+        f"{missing!r} is listed in qqa.__all__ but not importable from "
+        "the top level — fix the import in src/qqa/__init__.py."
+    )
+
+
 def test_callback_re_export_is_the_same_object() -> None:
     """The re-exports must be identical objects, not look-alike copies."""
     from qqa.callbacks import (