diff --git a/lasagna/fit.py b/lasagna/fit.py index 0c18ae7b4b..cd709d6d24 100644 --- a/lasagna/fit.py +++ b/lasagna/fit.py @@ -1290,6 +1290,7 @@ def _run_flatten_mode( opt_cfg: cli_opt.OptConfig, progress_enabled: bool, out_dir: str | None, + lifecycle_fn=None, ) -> int: ext_surfaces_cfg = cfg.get("external_surfaces", None) if not isinstance(ext_surfaces_cfg, list) or len(ext_surfaces_cfg) != 1: @@ -1414,6 +1415,8 @@ def _progress(*, step: int, total: int, loss: float, **_kw: object) -> None: seed_xyz=None, out_dir=out_dir, ) + if lifecycle_fn is not None: + lifecycle_fn("saving", "Saving optimized flatten model") if device.type == "cuda": peak_gb = torch.cuda.max_memory_allocated(device) / 2**30 @@ -1441,7 +1444,7 @@ def _progress(*, step: int, total: int, loss: float, **_kw: object) -> None: return 0 -def main(argv: list[str] | None = None) -> int: +def main(argv: list[str] | None = None, *, lifecycle_fn=None) -> int: if argv is None: argv = sys.argv[1:] @@ -1490,6 +1493,7 @@ def main(argv: list[str] | None = None) -> int: opt_cfg=opt_cfg, progress_enabled=progress_enabled, out_dir=_out_dir, + lifecycle_fn=lifecycle_fn, ) data_cfg = cli_data.from_args(args) @@ -2420,6 +2424,8 @@ def _progress(*, step: int, total: int, loss: float, **_kw: object) -> None: require_snap_surf_map_state=(model_init == "model" and self_map_init != "off"), ) _stage_done("optimizer", _t) + if lifecycle_fn is not None: + lifecycle_fn("saving", "Saving optimized Lasagna model") if device.type == "cuda": peak_gb = torch.cuda.max_memory_allocated(device) / 2**30 diff --git a/lasagna/fit2tifxyz.py b/lasagna/fit2tifxyz.py index 8ec81f0b36..1f92d0f856 100644 --- a/lasagna/fit2tifxyz.py +++ b/lasagna/fit2tifxyz.py @@ -30,6 +30,8 @@ class ExportConfig: voxel_size_um: float | None = None target_volume_shape_zyx: tuple[int, int, int] | None = None flow_gate_channels: str = "auto" + omit_model: bool = False + flatten_map_output: str | None = None def _valid_xyz_mask(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray: @@ -52,6 +54,10 @@ def _build_parser() -> argparse.ArgumentParser: help="Export all windings into a single tifxyz") g.add_argument("--copy-model", action="store_true", default=False, help="Deprecated no-op; model checkpoints are always copied") + g.add_argument("--omit-model", action="store_true", default=False, + help="Do not copy the model checkpoint into exported tifxyz directories") + g.add_argument("--flatten-map-output", default=None, + help="For flatten checkpoints, write the output-to-source map as a .npy sidecar") g.add_argument("--output-name", default=None, help="Override tifxyz directory name") g.add_argument("--voxel-size-um", type=float, default=None, help="Voxel size in micrometers (for area calculation)") @@ -1001,6 +1007,10 @@ def _export_flatten_checkpoint( map_yx = _as_numpy_float32(st["flatten_map_flat"], name="flatten_map_flat") if map_yx.ndim != 3 or map_yx.shape[-1] != 2: raise ValueError("flatten_map_flat must have shape (H, W, 2)") + if cfg.flatten_map_output: + map_output = Path(cfg.flatten_map_output) + map_output.parent.mkdir(parents=True, exist_ok=True) + np.save(str(map_output), map_yx, allow_pickle=False) mesh = st.get("mesh_flat") if mesh is None: @@ -1068,7 +1078,7 @@ def _export_flatten_checkpoint( z=z, d=d, scale=meta_scale, - model_source=Path(cfg.input), + model_source=None if cfg.omit_model else Path(cfg.input), copy_model=cfg.copy_model, fit_config=fit_config, job_spec=job_spec, @@ -1106,6 +1116,11 @@ def _check_cancel() -> None: else tuple(int(v) for v in args.target_volume_shape_zyx) ), flow_gate_channels=str(args.flow_gate_channels), + omit_model=bool(args.omit_model), + flatten_map_output=( + None if args.flatten_map_output in (None, "") + else str(args.flatten_map_output) + ), ) dev = torch.device(cfg.device) @@ -1271,7 +1286,8 @@ def _check_cancel() -> None: single_segment_border=BORDER_W, ) _write_tifxyz(out_dir=out_dir, x=x_all, y=y_all, z=z_all, d=d_all, scale=meta_scale, - model_source=Path(cfg.input), copy_model=cfg.copy_model, fit_config=fit_config, + model_source=None if cfg.omit_model else Path(cfg.input), + copy_model=cfg.copy_model, fit_config=fit_config, job_spec=job_spec, object_refs=object_refs, area=area, components=components if D > 1 else None, @@ -1343,7 +1359,8 @@ def _check_cancel() -> None: layer_index=d, ) _write_tifxyz(out_dir=out_dir, x=x, y=y, z=z, d=d_layer, scale=meta_scale, - model_source=Path(cfg.input), copy_model=cfg.copy_model, fit_config=fit_config, + model_source=None if cfg.omit_model else Path(cfg.input), + copy_model=cfg.copy_model, fit_config=fit_config, job_spec=job_spec, object_refs=object_refs, area=area, diff --git a/lasagna/fit_service.py b/lasagna/fit_service.py index 58a0094130..53b9877aed 100644 --- a/lasagna/fit_service.py +++ b/lasagna/fit_service.py @@ -1909,8 +1909,19 @@ def _wrapped_progress(*, step: int, total: int, loss: float, **kw: Any) -> None: try: import fit as fit_mod job.set_running("loading", 0, 0, 0.0) - fit_mod.main([cfg_path]) - if isinstance(job_spec, dict) and Path(model_output).is_file(): + + def _fit_lifecycle(stage: str, stage_name: str) -> None: + job.set_running( + stage, 0, 0, 0.0, stage_name=stage_name) + _check_cancel() + + fit_mod.main([cfg_path], lifecycle_fn=_fit_lifecycle) + if (body.get("embed_job_metadata", True) + and isinstance(job_spec, dict) + and Path(model_output).is_file()): + job.set_running( + "saving", 0, 0, 0.0, + stage_name="Embedding Lasagna job metadata") import torch st = torch.load(str(model_output), map_location="cpu", weights_only=False) if isinstance(st, dict): @@ -1927,7 +1938,9 @@ def _wrapped_progress(*, step: int, total: int, loss: float, **kw: Any) -> None: return save_t0 = time.perf_counter() - job.set_running("exporting", 0, 0, 0.0) + job.set_running( + "exporting", 0, 0, 0.0, + stage_name="Exporting flattened preview surface") import fit2tifxyz export_argv = ["--input", str(model_output), "--output", str(output_dir)] if body.get("single_segment"): @@ -1935,6 +1948,13 @@ def _wrapped_progress(*, step: int, total: int, loss: float, **kw: Any) -> None: output_name = body.get("output_name") if output_name: export_argv.extend(["--output-name", str(output_name)]) + if body.get("omit_model", False): + export_argv.append("--omit-model") + if body.get("export_flatten_map", False): + export_argv.extend([ + "--flatten-map-output", + str(Path(output_dir) / ".flatten-map.npy"), + ]) print( f"[fit-service] exporting tifxyz: output_name={str(output_name or '').strip()!r} " f"output_dir={output_dir}", @@ -1952,6 +1972,9 @@ def _wrapped_progress(*, step: int, total: int, loss: float, **kw: Any) -> None: ]) _check_cancel() fit2tifxyz.main(export_argv, cancel_fn=_check_cancel) + job.set_running( + "finalizing", 0, 0, 0.0, + stage_name="Finalizing Lasagna preview export") _normalize_single_tifxyz_output(Path(output_dir), str(output_name or "")) _check_cancel() save_s = time.perf_counter() - save_t0 diff --git a/lasagna/model.py b/lasagna/model.py index ac7d165caa..54fb47d029 100644 --- a/lasagna/model.py +++ b/lasagna/model.py @@ -3254,11 +3254,11 @@ def _flatten_invert_forward_uv_map( total = int(out_h) * int(out_w) chunk = max(1, int(chunk_points)) - if use_cuda and chunk < (1 << 20): - # Peak GPU memory scales with chunk * k through the candidate edge and - # barycentric tensors. If k_candidates grows large or OOMs appear, - # lower this cap. - chunk = 1 << 20 + # if use_cuda and chunk < (1 << 20): + # # Peak GPU memory scales with chunk * k through the candidate edge and + # # barycentric tensors. If k_candidates grows large or OOMs appear, + # # lower this cap. + # chunk = 1 << 20 for start in range(0, total, chunk): stop = min(total, start + chunk) flat_idx = np.arange(start, stop, dtype=np.int64) diff --git a/lasagna/tests/test_opt_loss_flatten.py b/lasagna/tests/test_opt_loss_flatten.py index 20e672cfe3..7928277d6a 100644 --- a/lasagna/tests/test_opt_loss_flatten.py +++ b/lasagna/tests/test_opt_loss_flatten.py @@ -8,6 +8,7 @@ from pathlib import Path from unittest import mock +import numpy as np import torch import tifffile @@ -1052,6 +1053,38 @@ def test_fit2tifxyz_exports_flatten_checkpoint(self) -> None: y = tifffile.imread(str(out / "vc3d_name.tifxyz" / "y.tif")) z = tifffile.imread(str(out / "vc3d_name.tifxyz" / "z.tif")) self.assertTrue(bool(((x == -1.0) & (y == -1.0) & (z == -1.0)).any())) + self.assertTrue((out / "vc3d_name.tifxyz" / "model.pt").is_file()) + + def test_fit2tifxyz_can_export_ephemeral_flatten_without_model(self) -> None: + xyz = _flat_grid(4, 4) + valid = torch.ones(4, 4, dtype=torch.bool) + mdl = _make_flatten_model(xyz, valid, mesh_step=1) + with tempfile.TemporaryDirectory() as td: + root = Path(td) + model_path = root / "flatten_model.pt" + map_path = root / "out" / ".flatten-map.npy" + out = root / "out" + fit._save_flatten_model( + str(model_path), + mdl=mdl, + data=fit._dummy_flatten_data(), + fit_config={"args": {"model-init": "flatten"}}, + ) + + fit2tifxyz.main([ + "--input", str(model_path), + "--output", str(out), + "--output-name", "preview.tifxyz", + "--omit-model", + "--flatten-map-output", str(map_path), + ]) + + self.assertFalse((out / "preview.tifxyz" / "model.pt").exists()) + expected = torch.load( + model_path, map_location="cpu", weights_only=False + )["flatten_map_flat"].numpy() + np.testing.assert_array_equal( + np.load(map_path, allow_pickle=False), expected) def test_forward_fit_mode_writes_normal_flatten_outputs(self) -> None: with tempfile.TemporaryDirectory() as td: diff --git a/scrollprize.org/docs/02_data_datasets.md b/scrollprize.org/docs/02_data_datasets.md index 5b2b028d26..26eefdc6ff 100644 --- a/scrollprize.org/docs/02_data_datasets.md +++ b/scrollprize.org/docs/02_data_datasets.md @@ -72,7 +72,7 @@ This dataset contains manual annotations of the scroll wraps, recording which pa Spiral annotations for scroll PHercParis4 (Scroll 1). Contents include ~27,000 verified and ~204,000 unverified surface patches, traced tracks, and point collections, together with `same_windings` / `relative_windings` / `abs_winding` graphs, the fitted `umbilicus`, fiber and outer-shell geometry, and the volume inputs used by the fitting pipeline (~49.6 GB total). - [README](pathname:///data/datasets/spiral-input-PHercParis4-README.md) -- [Browse on Hugging Face](https://huggingface.co/buckets/scrollprize/datasets/tree/spiral/PHercParis4) +- [Browse on the data server](https://dl.ash2txt.org/datasets/spiral_datasets/PHercParis4/) - [Tutorial: Spiral Fitting](tutorial_spiral) — how to fit a whole-scroll surface to these annotations {/* TODO: add more scroll subsections here later, e.g. "### Scroll5", each following the Paris4 pattern above. */} diff --git a/scrollprize.org/docs/38_tutorial_spiral.md b/scrollprize.org/docs/38_tutorial_spiral.md index 31dadf2fcb..a98b465f2d 100644 --- a/scrollprize.org/docs/38_tutorial_spiral.md +++ b/scrollprize.org/docs/38_tutorial_spiral.md @@ -106,12 +106,12 @@ The spiral scripts declare their dependencies in their own [`pyproject.toml`](ht #### Get the dataset -Ready-made inputs for PHerc. Paris 4 (Scroll 1) are published in the [`spiral-input` dataset](data_datasets#spiral-input-2026-07), which lives in the [`scrollprize/datasets` storage bucket](https://huggingface.co/buckets/scrollprize/datasets/tree/spiral) on Hugging Face (~50 GB): +Ready-made inputs are published in the [`spiral-input` dataset](data_datasets#spiral-input-2026-07), which lives on the dl.ash2txt.org data server : [Spiral Datasets](https://dl.ash2txt.org/datasets/spiral_datasets/PHercParis4/) (~90 GB): ```bash -uvx --from huggingface_hub hf buckets sync \ - hf://buckets/scrollprize/datasets/spiral/PHercParis4 \ - ./spiral-dataset/PHercParis4 +rclone copy :http: ./spiral_datasets/phercparis4 \ + --http-url https://dl.ash2txt.org/datasets/spiral_datasets/PHercParis4/ \ + --transfers 32 -P ``` `hf buckets sync` works like `rsync`: re-running it resumes interrupted downloads. The dataset contains verified and unverified patches, tracks, fibers, the outer shell, winding annotation JSONs, the umbilicus, and the volume inputs — see the [dataset README](pathname:///data/datasets/spiral-input-PHercParis4-README.md) for the exact layout. diff --git a/volume-cartographer/apps/VC3D/CMakeLists.txt b/volume-cartographer/apps/VC3D/CMakeLists.txt index 644c865bb0..4f95fb5872 100644 --- a/volume-cartographer/apps/VC3D/CMakeLists.txt +++ b/volume-cartographer/apps/VC3D/CMakeLists.txt @@ -406,5 +406,6 @@ install(DIRECTORY ${CMAKE_SOURCE_DIR}/scripts/spiral/ DESTINATION share/volume-cartographer/spiral FILES_MATCHING PATTERN "*.py" + PATTERN "*.json" PATTERN "tests" EXCLUDE PATTERN "__pycache__" EXCLUDE) diff --git a/volume-cartographer/apps/VC3D/SpiralArtifactCache.cpp b/volume-cartographer/apps/VC3D/SpiralArtifactCache.cpp index a22aa20cd0..d37425a1dd 100644 --- a/volume-cartographer/apps/VC3D/SpiralArtifactCache.cpp +++ b/volume-cartographer/apps/VC3D/SpiralArtifactCache.cpp @@ -55,6 +55,10 @@ struct SpiralArtifactCache::FetchJob QJsonObject manifest; QString entryPoint; QList pendingFiles; + int totalFiles = 0; + int filesComplete = 0; + qint64 totalBytes = 0; + qint64 bytesComplete = 0; QString partialDir; QString finalDir; FetchCallback done; @@ -199,7 +203,11 @@ void SpiralArtifactCache::fetchArtifact(const QString& sessionId, const QString& && isDeferredPreviewFile(entry.value(QStringLiteral("name")).toString())) continue; job->pendingFiles.push_back(entry); + job->totalBytes += entry.value(QStringLiteral("size")).toInteger(); } + job->totalFiles = job->pendingFiles.size(); + emit fetchProgress(job->artifactId, QStringLiteral("starting"), QString(), + 0, job->totalFiles, 0, job->totalBytes); if (!QDir().mkpath(job->partialDir)) { finishJob(job, tr("Could not create the artifact cache directory")); return; @@ -351,6 +359,9 @@ void SpiralArtifactCache::startNextFile(const std::shared_ptr& job) finishJob(job, tr("Could not publish the artifact cache directory")); return; } + emit fetchProgress(job->artifactId, QStringLiteral("finished"), QString(), + job->filesComplete, job->totalFiles, + job->bytesComplete, job->totalBytes); job->done(QDir(job->finalDir).filePath(job->entryPoint), {}, false); return; } @@ -373,6 +384,9 @@ void SpiralArtifactCache::startNextFile(const std::shared_ptr& job) existing = 0; } } + emit fetchProgress(job->artifactId, QStringLiteral("downloading"), name, + job->filesComplete, job->totalFiles, + job->bytesComplete + existing, job->totalBytes); auto file = std::make_shared(targetPath); if (!file->open(existing > 0 ? (QIODevice::WriteOnly | QIODevice::Append) @@ -387,21 +401,36 @@ void SpiralArtifactCache::startNextFile(const std::shared_ptr& job) request.setRawHeader("Range", QStringLiteral("bytes=%1-").arg(existing).toUtf8()); QNetworkReply* reply = existing == declaredSize ? nullptr : _network->get(request); - auto verifyAndContinue = [this, job, name, declaredSize, declaredSha, targetPath]() { + auto streamedDigest = std::make_shared(); + auto verifyAndContinue = [this, job, name, declaredSize, declaredSha, + targetPath, streamedDigest]() { + emit fetchProgress(job->artifactId, QStringLiteral("verifying"), name, + job->filesComplete, job->totalFiles, + job->bytesComplete + declaredSize, job->totalBytes); auto* watcher = new QFutureWatcher(this); - connect(watcher, &QFutureWatcher::finished, this, [this, watcher, job]() { + connect(watcher, &QFutureWatcher::finished, this, + [this, watcher, job, name, declaredSize]() { const QString error = watcher->result(); watcher->deleteLater(); if (job->generation != _generation) return; if (!error.isEmpty()) { finishJob(job, error); return; } + ++job->filesComplete; + job->bytesComplete += declaredSize; + emit fetchProgress(job->artifactId, QStringLiteral("downloaded"), name, + job->filesComplete, job->totalFiles, + job->bytesComplete, job->totalBytes); startNextFile(job); }); - watcher->setFuture(QtConcurrent::run([name, declaredSize, declaredSha, targetPath]() -> QString { + watcher->setFuture(QtConcurrent::run( + [name, declaredSize, declaredSha, targetPath, + streamedDigest]() -> QString { const qint64 actualSize = QFileInfo(targetPath).size(); if (actualSize != declaredSize) return tr("Artifact file %1 has %2 bytes; the manifest declares %3") .arg(name).arg(actualSize).arg(declaredSize); - const QString actualSha = hashFileSha256(targetPath); + const QString actualSha = streamedDigest->isEmpty() + ? hashFileSha256(targetPath) + : *streamedDigest; if (actualSha != declaredSha) return tr("Artifact file %1 failed its SHA-256 digest check").arg(name); return {}; @@ -410,12 +439,29 @@ void SpiralArtifactCache::startNextFile(const std::shared_ptr& job) if (!reply) { verifyAndContinue(); return; } - connect(reply, &QNetworkReply::readyRead, this, [reply, file]() { - file->write(reply->readAll()); + auto streamHash = existing == 0 + ? std::make_shared(QCryptographicHash::Sha256) + : std::shared_ptr(); + connect(reply, &QNetworkReply::readyRead, this, + [reply, file, streamHash]() { + const QByteArray bytes = reply->readAll(); + if (streamHash) streamHash->addData(bytes); + file->write(bytes); }); + connect(reply, &QNetworkReply::downloadProgress, this, + [this, job, name, existing](qint64 received, qint64) { + emit fetchProgress( + job->artifactId, QStringLiteral("downloading"), name, + job->filesComplete, job->totalFiles, + job->bytesComplete + existing + qMax(0, received), + job->totalBytes); + }); connect(reply, &QNetworkReply::finished, this, - [this, reply, file, job, name, verifyAndContinue]() { - file->write(reply->readAll()); + [this, reply, file, job, name, verifyAndContinue, + streamHash, streamedDigest]() { + const QByteArray bytes = reply->readAll(); + if (streamHash) streamHash->addData(bytes); + file->write(bytes); file->close(); const int http = reply->attribute(QNetworkRequest::HttpStatusCodeAttribute).toInt(); const bool ok = reply->error() == QNetworkReply::NoError @@ -429,6 +475,9 @@ void SpiralArtifactCache::startNextFile(const std::shared_ptr& job) finishJob(job, tr("Downloading artifact file %1 failed: %2").arg(name, errorText)); return; } + if (streamHash) + *streamedDigest = + QString::fromLatin1(streamHash->result().toHex()); verifyAndContinue(); }); } @@ -437,6 +486,9 @@ void SpiralArtifactCache::finishJob(const std::shared_ptr& job, const QString& error, bool gone) { // Keep the partial directory: verified complete files resume a later fetch. + emit fetchProgress(job->artifactId, QStringLiteral("failed"), QString(), + job->filesComplete, job->totalFiles, + job->bytesComplete, job->totalBytes); job->done({}, error, gone); } diff --git a/volume-cartographer/apps/VC3D/SpiralArtifactCache.hpp b/volume-cartographer/apps/VC3D/SpiralArtifactCache.hpp index 2a67b26f7c..7bbc630ab1 100644 --- a/volume-cartographer/apps/VC3D/SpiralArtifactCache.hpp +++ b/volume-cartographer/apps/VC3D/SpiralArtifactCache.hpp @@ -53,6 +53,11 @@ class SpiralArtifactCache : public QObject QString cacheRoot() const; +signals: + void fetchProgress(const QString& artifactId, const QString& phase, + const QString& fileName, int filesComplete, int totalFiles, + qint64 bytesReceived, qint64 totalBytes); + private: struct FetchJob; void startNextFile(const std::shared_ptr& job); diff --git a/volume-cartographer/apps/VC3D/SpiralPanel.cpp b/volume-cartographer/apps/VC3D/SpiralPanel.cpp index 151143b60d..118f329fd8 100644 --- a/volume-cartographer/apps/VC3D/SpiralPanel.cpp +++ b/volume-cartographer/apps/VC3D/SpiralPanel.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -35,6 +36,7 @@ #include #include #include +#include #include #include @@ -176,33 +178,6 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) auto* pathsForm = new QFormLayout(pathsContents); pathsGroup->contentLayout()->addWidget(pathsContents); - auto* optionalInputs = new QWidget(pathsContents); - auto* optionalGrid = new QGridLayout(optionalInputs); - optionalGrid->setContentsMargins(0, 0, 0, 0); - optionalGrid->setHorizontalSpacing(12); - optionalGrid->setVerticalSpacing(2); - const auto optionalSpecs = std::initializer_list>{ - {"verified_patches", "Verified"}, {"unverified_patches", "Unverified"}, - {"normals", "Normals"}, {"surf_sdt", "SDT"}, - {"tracks_dbm", "Tracks"}, {"gradient_magnitude", "Grad mag"}, - {"fibers", "Fibers"}, - }; - int optionalIndex = 0; - for (const auto& spec : optionalSpecs) { - const QString key = QString::fromLatin1(spec.first); - auto* check = new QCheckBox(tr(spec.second), optionalInputs); - check->setObjectName(QStringLiteral("spiralUse_") + key); - check->setChecked(true); - check->setToolTip(tr("Include this dataset input and its associated losses and sampling")); - _optionalInputs.insert(key, check); - optionalGrid->addWidget(check, optionalIndex / 4, optionalIndex % 4); - ++optionalIndex; - connect(check, &QCheckBox::toggled, this, [this](bool) { - updateOptionalInputUi(); - refreshReloadRequired(); - }); - } - pathsForm->addRow(tr("Use inputs"), optionalInputs); addPathRow(pathsForm, "dataset_root", tr("Dataset root"), true); _refill = new QPushButton(tr("Refill from Dataset Root"), pathsContents); pathsForm->addRow(_refill); @@ -293,6 +268,12 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) "primary track. Applies on the next Run; zero disables crossing-pair " "sampling. The upper bound is prepared when the session loads.")); pathsForm->addRow(tr("Max crossings / sampled track"), _maxTrackCrossings); + for (QWidget* field : {static_cast(_trackLengthBinSampling), + trackWeights, + static_cast(_maxTrackCrossings)}) { + field->hide(); + if (QWidget* label = pathsForm->labelForField(field)) label->hide(); + } updateTrackSamplingUi(); addPathRow(pathsForm, "verified_patches", tr("Verified patches"), true); @@ -313,12 +294,9 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) _lasagnaScale = new QSpinBox(lasagnaContents); _lasagnaScale->setRange(1, 1024); _lasagnaScale->setValue(4); - _storageBackend = new QComboBox(lasagnaContents); - _storageBackend->addItem(tr("Sparse CUDA LRU"), QStringLiteral("sparse_cuda")); addPathRow(lasagnaForm, "cache_directory", tr("Cache directory"), true); lasagnaForm->addRow(tr("Zarr group"), _lasagnaGroup); lasagnaForm->addRow(tr("Coordinate scale"), _lasagnaScale); - lasagnaForm->addRow(tr("Storage backend"), _storageBackend); auto* outputGroup = makeSection(tr("Fit and output"), QStringLiteral("spiralFitOutputGroup"), @@ -338,6 +316,7 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) _renderVolumeScale = new QSpinBox(outputContents); _renderVolumeScale->setRange(1, 4096); _renderVolumeScale->setValue(16); _savePngVisualizations = new QCheckBox(tr("Save diagnostic PNG visualizations"), outputContents); _savePngVisualizations->setChecked(false); + _savePngVisualizations->hide(); _influenceEnabled = new QCheckBox(tr("Localize fit around added inputs"), outputContents); _influenceEnabled->setChecked(false); _influenceEnabled->setToolTip(tr("Restrict optimization to a region around each input added to " @@ -538,7 +517,7 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) auto* runLayout = new QVBoxLayout(runContents); runGroup->contentLayout()->addWidget(runContents); auto* controls = new QHBoxLayout; - _load = new QPushButton(tr("Load/Reload Inputs"), runContents); + _load = new QPushButton(tr("Start Fit"), runContents); _load->setEnabled(false); _iterations = new QSpinBox(runContents); _iterations->setRange(1, 1000000); _iterations->setValue(100); _run = new QPushButton(tr("Run"), runContents); @@ -555,6 +534,47 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) checkpointControls->addWidget(_downloadCheckpoint); checkpointControls->addStretch(1); runLayout->addLayout(checkpointControls); + _checkpointDownloadStatus = new QLabel(runContents); + _checkpointDownloadStatus->setVisible(false); + _checkpointDownloadProgress = new QProgressBar(runContents); + _checkpointDownloadProgress->setVisible(false); + runLayout->addWidget(_checkpointDownloadStatus); + runLayout->addWidget(_checkpointDownloadProgress); + _checkpointDownloadTimer = new QTimer(this); + _checkpointDownloadTimer->setInterval(1000); + auto refreshCheckpointDownload = [this]() { + const qint64 seconds = _checkpointDownloadElapsed.isValid() + ? _checkpointDownloadElapsed.elapsed() / 1000 : 0; + const QString elapsed = seconds < 60 + ? tr("%1s").arg(seconds) + : tr("%1m %2s").arg(seconds / 60).arg(seconds % 60, 2, 10, QLatin1Char('0')); + QString phase; + if (_checkpointDownloadPhase == QStringLiteral("creating")) + phase = tr("Creating checkpoint on service"); + else if (_checkpointDownloadPhase == QStringLiteral("verifying")) + phase = tr("Verifying checkpoint"); + else if (_checkpointDownloadPhase == QStringLiteral("copying")) + phase = tr("Saving checkpoint"); + else + phase = tr("Downloading checkpoint"); + _checkpointDownloadStatus->setText( + tr("%1 — elapsed %2").arg(phase, elapsed)); + if (_checkpointTotalBytes > 0) { + _checkpointDownloadProgress->setRange(0, 1000); + _checkpointDownloadProgress->setValue(static_cast( + qBound(qint64{0}, _checkpointBytesReceived * 1000 + / _checkpointTotalBytes, qint64{1000}))); + _checkpointDownloadProgress->setFormat( + tr("%1 / %2 MiB") + .arg(_checkpointBytesReceived / (1024 * 1024)) + .arg(_checkpointTotalBytes / (1024 * 1024))); + } else { + _checkpointDownloadProgress->setRange(0, 0); + _checkpointDownloadProgress->setFormat(QString()); + } + }; + connect(_checkpointDownloadTimer, &QTimer::timeout, this, + refreshCheckpointDownload); auto* ephemeralLabel = new QLabel(tr("Inputs added to the running fit:"), runContents); _ephemeralList = new QListWidget(runContents); @@ -580,10 +600,17 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) runLayout->addWidget(_commitHint); _state = new QLabel(tr("Service disconnected"), runContents); + _previewProgress = new QProgressBar(runContents); + _previewProgress->setObjectName(QStringLiteral("spiralPreviewProgress")); + _previewProgress->setTextVisible(true); + _previewProgress->setVisible(false); _metrics = new QLabel(runContents); _warnings = new QLabel(runContents); _warnings->setWordWrap(true); _warnings->setTextInteractionFlags(Qt::TextSelectableByMouse); - runLayout->addWidget(_state); runLayout->addWidget(_metrics); runLayout->addWidget(_warnings); + runLayout->addWidget(_state); + runLayout->addWidget(_previewProgress); + runLayout->addWidget(_metrics); + runLayout->addWidget(_warnings); layout->addStretch(1); scroll->setWidget(contents); rootLayout->addWidget(scroll); @@ -649,9 +676,24 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) // Connection must succeed before dataset resolution or fit // controls are enabled. _load->setEnabled(_connected); - if (!_connected) { _run->setEnabled(false); _stop->setEnabled(false); - _save->setEnabled(false); _downloadCheckpoint->setEnabled(false); - _removeInput->setEnabled(false); } + if (!_connected) { + _run->setEnabled(false); + _stop->setEnabled(false); + _save->setEnabled(false); + _downloadCheckpoint->setEnabled(false); + _removeInput->setEnabled(false); + if (_checkpointDownloadActive) { + _checkpointDownloadActive = false; + _checkpointDownloadTimer->stop(); + const qint64 seconds = + _checkpointDownloadElapsed.elapsed() / 1000; + _checkpointDownloadProgress->setVisible(false); + _checkpointDownloadStatus->setText( + tr("Checkpoint download interrupted after %1m %2s") + .arg(seconds / 60) + .arg(seconds % 60, 2, 10, QLatin1Char('0'))); + } + } if (state == CS::Starting || state == CS::Connecting) { _hasSession = false; _loadedSessionRequest = {}; @@ -672,7 +714,71 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) applyResolution(value, _remoteMode || !_hasManualEdits); _pendingDatasetRoot.clear(); }); + connect(_service, &SpiralServiceManager::configurationCatalogChanged, + _advancedProfiles, &SpiralConfigProfileEditor::setCatalog); + connect(_service, &SpiralServiceManager::configurationCatalogChanged, + this, [this](const QJsonObject& catalog) { + _runMutablePaths.clear(); + const QJsonObject paths = + catalog.value(QStringLiteral("schema")).toObject() + .value(QStringLiteral("paths")).toObject(); + for (auto it = paths.begin(); it != paths.end(); ++it) { + const QString impact = + it.value().toObject() + .value(QStringLiteral("runtime_impact")).toString(); + if (impact == QStringLiteral("run_boundary") + || impact == QStringLiteral("shell_reload")) + _runMutablePaths.insert(it.key()); + } + refreshReloadRequired(); + }); + connect(_service, &SpiralServiceManager::configurationReviewRequested, + _advancedProfiles, &SpiralConfigProfileEditor::showWindow); connect(_service, &SpiralServiceManager::sessionStatusChanged, this, &SpiralPanel::updateStatus); + connect(_service, &SpiralServiceManager::previewTransferProgress, this, + [this](const QString& phase, const QString& fileName, + int filesComplete, int totalFiles, + qint64 bytesReceived, qint64 totalBytes) { + _previewTransferActive = + phase != QStringLiteral("finished") + && phase != QStringLiteral("failed"); + const QString action = + phase == QStringLiteral("verifying") + ? tr("Verifying preview") + : phase == QStringLiteral("failed") + ? tr("Preview download failed") + : phase == QStringLiteral("finished") + ? tr("Installing preview") + : tr("Downloading preview"); + _previewTransferText = fileName.isEmpty() + ? action + : tr("%1 (%2/%3): %4") + .arg(action) + .arg(qMin(filesComplete + 1, totalFiles)) + .arg(totalFiles) + .arg(fileName); + _previewProgress->setVisible( + phase != QStringLiteral("failed")); + if (totalBytes > 0) { + _previewProgress->setRange(0, 1000); + _previewProgress->setValue(static_cast( + qBound( + qint64{0}, bytesReceived * 1000 / totalBytes, + qint64{1000}))); + _previewProgress->setFormat( + tr("%1 / %2 MiB") + .arg(bytesReceived / (1024 * 1024)) + .arg(totalBytes / (1024 * 1024))); + } else { + _previewProgress->setRange(0, 0); + } + }); + connect(_service, &SpiralServiceManager::previewAvailable, this, + [this](const QString&, qint64) { + _previewTransferActive = false; + _previewTransferText.clear(); + _previewProgress->setVisible(false); + }); connect(_service, &SpiralServiceManager::sessionSynchronized, this, &SpiralPanel::synchronizeSession); connect(_service, &SpiralServiceManager::errorOccurred, this, [this](const QString& error) { @@ -686,10 +792,35 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) }); connect(_service, &SpiralServiceManager::checkpointDownloadFinished, this, [this](const QString& path, const QString& error) { - if (error.isEmpty()) + _checkpointDownloadTimer->stop(); + _checkpointDownloadActive = false; + const qint64 seconds = _checkpointDownloadElapsed.elapsed() / 1000; + _checkpointDownloadProgress->setRange(0, 1000); + if (error.isEmpty()) { + _checkpointDownloadProgress->setValue(1000); + _checkpointDownloadProgress->setFormat(tr("Complete")); + _checkpointDownloadStatus->setText( + tr("Checkpoint download complete — %1m %2s") + .arg(seconds / 60) + .arg(seconds % 60, 2, 10, QLatin1Char('0'))); _warnings->setText(tr("Checkpoint downloaded to %1").arg(path)); - else + } else { + _checkpointDownloadProgress->setVisible(false); + _checkpointDownloadStatus->setText( + tr("Checkpoint download failed after %1m %2s") + .arg(seconds / 60) + .arg(seconds % 60, 2, 10, QLatin1Char('0'))); _warnings->setText(tr("Checkpoint download failed: %1").arg(error)); + } + _downloadCheckpoint->setEnabled(_connected && _sessionRunnable); + }); + connect(_service, &SpiralServiceManager::checkpointDownloadProgress, this, + [this, refreshCheckpointDownload](const QString& phase, + qint64 received, qint64 total) { + _checkpointDownloadPhase = phase; + _checkpointBytesReceived = received; + _checkpointTotalBytes = total; + refreshCheckpointDownload(); }); connect(_service, &SpiralServiceManager::inputUploadFinished, this, [this](const QString& inputId, const QString& error) { @@ -733,7 +864,8 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) persist(); emit pythonOutputRequested(); _service->runIterations(_iterations->value(), influenceConfig(), - runAdvancedConfig()); + runAdvancedConfig(), + sessionRequest().value(QStringLiteral("paths")).toObject()); }); connect(_stop, &QPushButton::clicked, _service, &SpiralServiceManager::stopAfterIteration); connect(_save, &QPushButton::clicked, this, [this]() { @@ -753,7 +885,16 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) QDir::home().filePath(QStringLiteral("spiral_checkpoint.ckpt")), tr("Checkpoint (*.ckpt)")); if (path.isEmpty()) return; - _warnings->setText(tr("Downloading checkpoint… this can take a few minutes.")); + _downloadCheckpoint->setEnabled(false); + _checkpointDownloadActive = true; + _checkpointDownloadPhase = QStringLiteral("creating"); + _checkpointBytesReceived = 0; + _checkpointTotalBytes = 0; + _checkpointDownloadElapsed.start(); + _checkpointDownloadStatus->setVisible(true); + _checkpointDownloadProgress->setVisible(true); + _checkpointDownloadTimer->start(); + _warnings->setText(tr("Preparing checkpoint download…")); _service->downloadCheckpoint(path); }); connect(_commitInputs, &QPushButton::clicked, this, [this]() { @@ -789,23 +930,9 @@ SpiralPanel::SpiralPanel(SpiralServiceManager* service, QWidget* parent) connect(edit, &QLineEdit::textEdited, this, [this](const QString&) { refreshReloadRequired(); }); connect(_outwardSense, qOverload(&QComboBox::currentIndexChanged), this, [this](int) { refreshReloadRequired(); }); - connect(_storageBackend, qOverload(&QComboBox::currentIndexChanged), this, - [this](int) { refreshReloadRequired(); }); connect(_savePngVisualizations, &QCheckBox::toggled, this, [this](bool) { refreshReloadRequired(); }); - connect(_trackLengthBinSampling, &QCheckBox::toggled, this, [this](bool) { - updateTrackSamplingUi(); - writeTrackSamplingControlsToAdvanced(); - }); - for (QDoubleSpinBox* spin : {_trackShortWeight, _trackMediumWeight, - _trackLongWeight}) { - connect(spin, qOverload(&QDoubleSpinBox::valueChanged), this, - [this](double) { writeTrackSamplingControlsToAdvanced(); }); - } - connect(_maxTrackCrossings, qOverload(&QSpinBox::valueChanged), this, - [this](int) { writeTrackSamplingControlsToAdvanced(); }); connect(_advancedProfiles, &SpiralConfigProfileEditor::textChanged, this, [this]() { - syncTrackSamplingControlsFromAdvanced(); refreshReloadRequired(); }); @@ -1043,7 +1170,6 @@ void SpiralPanel::setRemoteMode(bool remote) for (QLineEdit* edit : {_paths["dataset_root"], _paths["umbilicus"]}) edit->setToolTip(tr("Service-host path, owned by the service")); } - updateOptionalInputUi(); } QLineEdit* SpiralPanel::addPathRow(QFormLayout* form, const QString& key, const QString& label, bool directory) @@ -1152,13 +1278,8 @@ void SpiralPanel::applyResolution(const QJsonObject& resolution, bool force) QJsonObject SpiralPanel::sessionRequest() const { QJsonObject paths; - for (auto it = _paths.begin(); it != _paths.end(); ++it) { - QString optionalKey = it.key(); - if (optionalKey == QStringLiteral("normal_x") - || optionalKey == QStringLiteral("normal_y")) - optionalKey = QStringLiteral("normals"); - paths[it.key()] = optionalInputEnabled(optionalKey) ? it.value()->text() : QString(); - } + for (auto it = _paths.begin(); it != _paths.end(); ++it) + paths[it.key()] = it.value()->text(); QJsonArray pcls; for (int row = 0; row < _pclList->count(); ++row) { const QListWidgetItem* item = _pclList->item(row); @@ -1168,12 +1289,11 @@ QJsonObject SpiralPanel::sessionRequest() const } paths["pcls"] = pcls; QJsonObject config = sessionAdvancedConfig(); - config[QStringLiteral("save_png_visualizations")] = _savePngVisualizations->isChecked(); QJsonObject run{{"z_begin", _zBegin->value()}, {"z_end", _zEnd->value()}, {"scroll_name", _scrollName->text()}, {"outward_sense", _outwardSense->currentText()}, {"voxel_size_um", _voxelSize->value()}, {"lasagna_group", _lasagnaGroup->text()}, {"lasagna_scale", _lasagnaScale->value()}, - {"storage_backend", _storageBackend->currentData().toString()}, + {"storage_backend", QStringLiteral("sparse_cuda")}, {"legacy_checkpoint_step", _legacyCheckpointStep->value()}, {"run_tag", _runTag->text()}, {"render_volume_scale", _renderVolumeScale->value()}, @@ -1196,21 +1316,21 @@ QJsonObject SpiralPanel::influenceConfig() const if (advanced.isObject()) { const QJsonObject all = advanced.object(); for (auto it = all.begin(); it != all.end(); ++it) { - if (it.key().startsWith(QStringLiteral("interactive_influence_")) + if (it.key().startsWith(QStringLiteral("influence_")) || it.key() == QStringLiteral("loss_weight_anchor")) config[it.key()] = it.value(); } } - config[QStringLiteral("interactive_influence_enabled")] = + config[QStringLiteral("influence_enabled")] = _influenceEnabled->isChecked(); - config[QStringLiteral("interactive_influence_disable_dt_frac")] = + config[QStringLiteral("influence_disable_dt_frac")] = _influenceDisableDtPct->value() / 100.0; if (_influenceEnabled->isChecked()) { - config[QStringLiteral("interactive_influence_z")] = + config[QStringLiteral("influence_z")] = static_cast(_influenceZ->value()); - config[QStringLiteral("interactive_influence_windings")] = + config[QStringLiteral("influence_windings")] = _influenceWindings->value(); - config[QStringLiteral("interactive_influence_theta_frac")] = + config[QStringLiteral("influence_theta_frac")] = _influenceThetaPct->value() / 100.0; config[QStringLiteral("loss_weight_anchor")] = _influenceAnchorWeight->value(); @@ -1240,13 +1360,12 @@ QJsonObject SpiralPanel::sessionAdvancedConfig() const config = advanced.object(); } for (auto it = config.begin(); it != config.end();) { - if (it.key().startsWith(QStringLiteral("interactive_influence_")) + if (it.key().startsWith(QStringLiteral("influence_")) || it.key() == QStringLiteral("loss_weight_anchor")) it = config.erase(it); else ++it; } - applyOptionalInputConfig(config, true); return config; } @@ -1254,24 +1373,11 @@ QJsonObject SpiralPanel::runAdvancedConfig() const { const QJsonDocument advanced = QJsonDocument::fromJson(_advanced->toPlainText().toUtf8()); - QJsonObject config; - if (!advanced.isObject()) return config; - const QJsonObject all = advanced.object(); - for (const QString& key : _runConfigKeys) { - if (all.contains(key)) config.insert(key, all.value(key)); - } - // Optional-input gating needs the session-scoped spacing mode even though - // that key itself is not mutable at a Run boundary. It is removed again - // below before the request is sent. - const QString spacingModeKey = QStringLiteral("dense_spacing_mode"); - if (all.contains(spacingModeKey)) - config.insert(spacingModeKey, all.value(spacingModeKey)); - if (_runConfigKeys.contains(QStringLiteral("save_png_visualizations"))) - config[QStringLiteral("save_png_visualizations")] = - _savePngVisualizations->isChecked(); - applyOptionalInputConfig(config, false); + if (!advanced.isObject()) return {}; + QJsonObject config = advanced.object(); for (auto it = config.begin(); it != config.end();) { - if (!_runConfigKeys.contains(it.key())) + if (it.key().startsWith(QStringLiteral("influence_")) + || it.key() == QStringLiteral("loss_weight_anchor")) it = config.erase(it); else ++it; @@ -1279,12 +1385,6 @@ QJsonObject SpiralPanel::runAdvancedConfig() const return config; } -bool SpiralPanel::optionalInputEnabled(const QString& key) const -{ - const auto it = _optionalInputs.constFind(key); - return it == _optionalInputs.cend() || it.value()->isChecked(); -} - void SpiralPanel::applyTrackSamplingConfig(QJsonObject& config) const { if (_trackLengthBinSampling->isChecked()) { @@ -1296,7 +1396,7 @@ void SpiralPanel::applyTrackSamplingConfig(QJsonObject& config) const } else { config[QStringLiteral("track_length_bin_weights")] = QJsonValue::Null; } - config[QStringLiteral("max_track_crossing_per_step")] = + config[QStringLiteral("track_max_track_crossing_per_step")] = _maxTrackCrossings->value(); } @@ -1330,7 +1430,7 @@ void SpiralPanel::syncTrackSamplingControlsFromAdvanced() } const QJsonValue crossings = - config.value(QStringLiteral("max_track_crossing_per_step")); + config.value(QStringLiteral("track_max_track_crossing_per_step")); { const QSignalBlocker blocker(_maxTrackCrossings); const int preparedMaximum = config @@ -1359,7 +1459,8 @@ void SpiralPanel::writeTrackSamplingControlsToAdvanced() void SpiralPanel::updateTrackSamplingUi() { - const bool tracksEnabled = optionalInputEnabled(QStringLiteral("tracks_dbm")); + const bool tracksEnabled = + !_paths.value(QStringLiteral("tracks_dbm"))->text().trimmed().isEmpty(); _trackLengthBinSampling->setEnabled(tracksEnabled); _trackShortWeight->setEnabled(tracksEnabled && _trackLengthBinSampling->isChecked()); _trackMediumWeight->setEnabled(tracksEnabled && _trackLengthBinSampling->isChecked()); @@ -1367,82 +1468,6 @@ void SpiralPanel::updateTrackSamplingUi() _maxTrackCrossings->setEnabled(tracksEnabled); } -void SpiralPanel::applyOptionalInputConfig(QJsonObject& config, - bool includeSelectionFlags) const -{ - const bool verified = optionalInputEnabled(QStringLiteral("verified_patches")); - const bool unverified = optionalInputEnabled(QStringLiteral("unverified_patches")); - const bool normals = optionalInputEnabled(QStringLiteral("normals")); - const bool sdt = optionalInputEnabled(QStringLiteral("surf_sdt")); - const bool tracks = optionalInputEnabled(QStringLiteral("tracks_dbm")); - const bool gradMag = optionalInputEnabled(QStringLiteral("gradient_magnitude")); - const bool fibers = optionalInputEnabled(QStringLiteral("fibers")); - - if (includeSelectionFlags) { - config[QStringLiteral("use_verified_patches")] = verified; - config[QStringLiteral("use_unverified_patches")] = unverified; - config[QStringLiteral("use_normals")] = normals; - config[QStringLiteral("use_surf_sdt")] = sdt; - config[QStringLiteral("use_tracks")] = tracks; - config[QStringLiteral("use_gradient_magnitude")] = gradMag; - config[QStringLiteral("use_fibers")] = fibers; - } - - auto zero = [&config](std::initializer_list keys) { - for (const char* key : keys) config[QString::fromLatin1(key)] = 0; - }; - if (!verified) { - zero({"loss_weight_patch_radius", "loss_weight_patch_dt", - "loss_weight_umbilicus", "loss_weight_shell_patch_radius", - "num_patches_per_step", "num_patches_per_step_for_dt", - "num_points_per_patch"}); - } - if (!unverified) { - zero({"loss_weight_unverified_patch_radius", "loss_weight_unverified_patch_dt", - "unverified_num_patches_per_step", - "unverified_num_patches_per_step_for_dt", - "unverified_num_points_per_patch"}); - } - if (!normals) - zero({"loss_weight_dense_normals", "dense_normals_num_points"}); - if (!tracks) { - zero({"loss_weight_track_radius", "loss_weight_track_dt", - "track_num_per_step", "track_num_points_per_step"}); - } - if (!fibers) { - zero({"loss_weight_unattached_pcl_radius", "loss_weight_unattached_pcl_dt", - "unattached_pcl_num_per_step", "unattached_pcl_num_points_per_step"}); - } - - const QString spacingMode = - config.value(QStringLiteral("dense_spacing_mode")).toString(QStringLiteral("phase")); - if (!sdt || !normals) { - zero({"loss_weight_dense_spacing_count", "loss_weight_dense_spacing_density", - "loss_weight_dense_attachment", - "dense_spacing_count_extra_pairs", "dense_spacing_density_extra_pairs", - "dense_attachment_num_points"}); - if (spacingMode == QStringLiteral("phase")) - zero({"loss_weight_dense_spacing", "dense_spacing_num_pairs"}); - } - if (!gradMag && spacingMode == QStringLiteral("grad_mag")) - zero({"loss_weight_dense_spacing", "dense_spacing_num_pairs"}); -} - -void SpiralPanel::updateOptionalInputUi() -{ - for (auto it = _paths.begin(); it != _paths.end(); ++it) { - QString optionalKey = it.key(); - if (optionalKey == QStringLiteral("normal_x") - || optionalKey == QStringLiteral("normal_y")) - optionalKey = QStringLiteral("normals"); - const bool enabled = optionalInputEnabled(optionalKey); - it.value()->setEnabled(enabled); - if (_pathBrowseButtons.contains(it.key())) - _pathBrowseButtons[it.key()]->setEnabled(enabled); - } - updateTrackSamplingUi(); -} - void SpiralPanel::applySessionRunConfig(const QJsonObject& config, qint64 sessionGeneration) { _runConfigKeys.clear(); @@ -1458,12 +1483,13 @@ void SpiralPanel::synchronizeSession(const QJsonObject& request, request.value(QStringLiteral("paths")).toObject(); const QJsonObject run = request.value(QStringLiteral("run")).toObject(); - const QJsonObject requestedConfig = - run.value(QStringLiteral("config")).toObject(); const QJsonObject activeRunConfig = status.value(QStringLiteral("run_config")).toObject(); - const QJsonObject defaultConfig = - status.value(QStringLiteral("default_advanced_config")).toObject(); + QJsonObject defaultConfig = + status.value(QStringLiteral("applied_config")).toObject(); + if (defaultConfig.isEmpty()) + defaultConfig = + status.value(QStringLiteral("default_advanced_config")).toObject(); const QJsonObject effectiveConfig = vc3d::effectiveSpiralSessionConfig( request, defaultConfig, activeRunConfig); @@ -1474,20 +1500,6 @@ void SpiralPanel::synchronizeSession(const QJsonObject& request, for (auto it = _paths.begin(); it != _paths.end(); ++it) it.value()->setText(paths.value(it.key()).toString()); - const QHash selectionFlags{ - {QStringLiteral("verified_patches"), QStringLiteral("use_verified_patches")}, - {QStringLiteral("unverified_patches"), QStringLiteral("use_unverified_patches")}, - {QStringLiteral("normals"), QStringLiteral("use_normals")}, - {QStringLiteral("surf_sdt"), QStringLiteral("use_surf_sdt")}, - {QStringLiteral("tracks_dbm"), QStringLiteral("use_tracks")}, - {QStringLiteral("gradient_magnitude"), QStringLiteral("use_gradient_magnitude")}, - {QStringLiteral("fibers"), QStringLiteral("use_fibers")}, - }; - for (auto it = selectionFlags.begin(); it != selectionFlags.end(); ++it) { - _optionalInputs.value(it.key())->setChecked( - requestedConfig.value(it.value()).toBool(true)); - } - _pclList->clear(); for (const QJsonValue& value : paths.value(QStringLiteral("pcls")).toArray()) { const QJsonObject item = value.toObject(); @@ -1509,9 +1521,6 @@ void SpiralPanel::synchronizeSession(const QJsonObject& request, run.value(QStringLiteral("lasagna_group")).toString(_lasagnaGroup->text())); _lasagnaScale->setValue( run.value(QStringLiteral("lasagna_scale")).toInt(_lasagnaScale->value())); - const int backend = _storageBackend->findData( - run.value(QStringLiteral("storage_backend")).toString()); - if (backend >= 0) _storageBackend->setCurrentIndex(backend); _legacyCheckpointStep->setValue( run.value(QStringLiteral("legacy_checkpoint_step")) .toInt(_legacyCheckpointStep->value())); @@ -1526,21 +1535,21 @@ void SpiralPanel::synchronizeSession(const QJsonObject& request, _advancedProfiles->setSessionDefault(effectiveConfig); _advancedProfiles->showSessionDefault(); _savePngVisualizations->setChecked( - effectiveConfig.value(QStringLiteral("save_png_visualizations")) + effectiveConfig.value(QStringLiteral("output_save_png_visualizations")) .toBool(false)); syncTrackSamplingControlsFromAdvanced(); if (!activeRunConfig.isEmpty()) applySessionRunConfig(activeRunConfig, sessionGeneration); else _advancedSessionGeneration = -1; - updateOptionalInputUi(); + updateTrackSamplingUi(); _applyingResolution = false; _hasManualEdits = false; _hasSession = true; // Reload comparisons use the panel's adopted representation. The host - // request is canonical and sparse, whereas the form expands defaults and - // applies optional-input gating; both describe the same resident fit. + // request is canonical and sparse, whereas the form expands the active + // session defaults; both describe the same resident fit. _loadedSessionRequest = sessionRequest(); _reloadRequired = false; for (auto it = _visibilityChecks.begin(); it != _visibilityChecks.end(); ++it) @@ -1551,11 +1560,16 @@ void SpiralPanel::updateStatus(const QJsonObject& status) { const qint64 sessionGeneration = status.value(QStringLiteral("session_generation")).toInteger(-1); + _load->setText(status.value(QStringLiteral("session_id")).toString().isEmpty() + ? tr("Start Fit") : tr("New Fit")); const QJsonObject runConfig = status.value(QStringLiteral("run_config")).toObject(); const QJsonObject runConfigLimits = status.value(QStringLiteral("run_config_limits")).toObject(); - const QJsonObject defaultConfig = - status.value(QStringLiteral("default_advanced_config")).toObject(); + QJsonObject defaultConfig = + status.value(QStringLiteral("applied_config")).toObject(); + if (defaultConfig.isEmpty()) + defaultConfig = + status.value(QStringLiteral("default_advanced_config")).toObject(); if (sessionGeneration >= 0 && sessionGeneration != _advancedSessionGeneration && !runConfig.isEmpty()) { const bool completingSynchronization = @@ -1570,7 +1584,7 @@ void SpiralPanel::updateStatus(const QJsonObject& status) _advancedProfiles->showSessionDefault(); _savePngVisualizations->setChecked( effectiveConfig - .value(QStringLiteral("save_png_visualizations")) + .value(QStringLiteral("output_save_png_visualizations")) .toBool(false)); syncTrackSamplingControlsFromAdvanced(); _applyingResolution = false; @@ -1580,7 +1594,7 @@ void SpiralPanel::updateStatus(const QJsonObject& status) refreshReloadRequired(); } const QJsonValue crossingLimit = - runConfigLimits.value(QStringLiteral("max_track_crossing_per_step")); + runConfigLimits.value(QStringLiteral("track_max_track_crossing_per_step")); if (crossingLimit.isDouble()) { const QSignalBlocker blocker(_maxTrackCrossings); _maxTrackCrossings->setMaximum(qMax(0, crossingLimit.toInt())); @@ -1595,10 +1609,84 @@ void SpiralPanel::updateStatus(const QJsonObject& status) _advancedSessionGeneration = -1; _runConfigKeys.clear(); } - QString stateText = tr("Session: %1 — %2 — iteration %3/%4") - .arg(state, status.value("phase").toString()) - .arg(status.value("current_iteration").toInteger()) - .arg(status.value("target_iteration").toInteger()); + const QJsonObject progress = + status.value(QStringLiteral("progress")).toObject(); + QString stateText = progress.isEmpty() + ? tr("Session: %1 — %2") + .arg(state, status.value("phase").toString()) + : tr("Session: %1").arg(state); + if (state == QStringLiteral("Running")) + stateText += tr(" — iteration %1/%2") + .arg(status.value("current_iteration").toInteger()) + .arg(status.value("target_iteration").toInteger()); + const QJsonObject previewPublish = + status.value(QStringLiteral("preview_publish")).toObject(); + if (_previewTransferActive && !_previewTransferText.isEmpty()) { + stateText += QStringLiteral("\n") + _previewTransferText; + } else if (!progress.isEmpty()) { + const qint64 step = + progress.value(QStringLiteral("step")).toInteger(-1); + const qint64 total = + progress.value(QStringLiteral("total_steps")).toInteger(-1); + const QString unit = + progress.value(QStringLiteral("unit")).toString(); + const QString detail = + progress.value(QStringLiteral("detail")).toString(); + const double elapsed = + progress.value(QStringLiteral("elapsed_seconds")).toDouble(); + const QJsonValue etaValue = + progress.value(QStringLiteral("eta_seconds")); + const auto durationText = [](double value) { + const qint64 seconds = qMax( + 0, qRound64(value)); + if (seconds < 60) + return QObject::tr("%1s").arg(seconds); + const qint64 minutes = seconds / 60; + if (minutes < 60) + return QObject::tr("%1m %2s") + .arg(minutes).arg(seconds % 60, 2, 10, QLatin1Char('0')); + return QObject::tr("%1h %2m") + .arg(minutes / 60) + .arg(minutes % 60, 2, 10, QLatin1Char('0')); + }; + QString progressText = + progress.value(QStringLiteral("stage_name")).toString(); + if (!detail.isEmpty()) + progressText += QStringLiteral(" — ") + detail; + progressText += tr(" — elapsed %1").arg(durationText(elapsed)); + if (etaValue.isDouble()) + progressText += tr(" — ETA %1") + .arg(durationText(etaValue.toDouble())); + stateText += QStringLiteral("\n") + progressText; + _previewProgress->setVisible(true); + if (total > 0) { + _previewProgress->setRange(0, 1000); + _previewProgress->setValue(static_cast( + qBound( + qint64{0}, step * 1000 / total, qint64{1000}))); + QString format = tr("%1 / %2").arg(step).arg(total); + if (!unit.isEmpty()) + format += QStringLiteral(" ") + unit; + _previewProgress->setFormat(format); + } else { + _previewProgress->setRange(0, 0); + } + } else if (!previewPublish.isEmpty()) { + const int step = previewPublish.value(QStringLiteral("step")).toInt(); + const int total = + previewPublish.value(QStringLiteral("total_steps")).toInt(); + _previewProgress->setVisible(true); + if (total > 0) { + _previewProgress->setRange(0, total); + _previewProgress->setValue(qBound(0, step, total)); + _previewProgress->setFormat( + tr("%1 / %2").arg(step).arg(total)); + } else { + _previewProgress->setRange(0, 0); + } + } else { + _previewProgress->setVisible(false); + } // Status polls arrive every second; without this the reload-required // notice set by refreshReloadRequired() vanishes immediately, leaving an // unexplained disabled Run button. @@ -1626,7 +1714,8 @@ void SpiralPanel::updateStatus(const QJsonObject& status) _run->setEnabled(_connected && runnable && !_reloadRequired); _stop->setEnabled(state == "Running"); _save->setEnabled(_connected && runnable); - _downloadCheckpoint->setEnabled(_connected && runnable); + _downloadCheckpoint->setEnabled( + _connected && runnable && !_checkpointDownloadActive); // Ephemeral inputs added to the running fit. const QJsonArray ephemeral = status.value(QStringLiteral("ephemeral_inputs")).toArray(); @@ -1685,7 +1774,8 @@ QJsonObject SpiralPanel::normalizedReloadRequest(QJsonObject request) const // Run-mutable settings are excluded because they can be applied at the // next Run boundary without rebuilding the resident session. return vc3d::normalizedSpiralReloadRequest( - std::move(request), _defaultAdvancedConfig, _runConfigKeys); + std::move(request), _defaultAdvancedConfig, _runConfigKeys, + _runMutablePaths); } void SpiralPanel::refreshReloadRequired() @@ -1694,41 +1784,6 @@ void SpiralPanel::refreshReloadRequired() QJsonObject current = sessionRequest(); - // A loaded input may be disabled at a Run boundary by forcing its mutable - // loss weights and sample counts to zero. Re-enabling an input that was not - // loaded, or changing its path, still requires rebuilding the session. - QJsonObject currentRun = current.value(QStringLiteral("run")).toObject(); - QJsonObject currentConfig = currentRun.value(QStringLiteral("config")).toObject(); - const QJsonObject loadedRun = - _loadedSessionRequest.value(QStringLiteral("run")).toObject(); - const QJsonObject loadedConfig = - loadedRun.value(QStringLiteral("config")).toObject(); - QJsonObject currentPaths = current.value(QStringLiteral("paths")).toObject(); - const QJsonObject loadedPaths = - _loadedSessionRequest.value(QStringLiteral("paths")).toObject(); - const QHash optionalInputs{ - {QStringLiteral("use_verified_patches"), {QStringLiteral("verified_patches")}}, - {QStringLiteral("use_unverified_patches"), {QStringLiteral("unverified_patches")}}, - {QStringLiteral("use_normals"), {QStringLiteral("normal_x"), QStringLiteral("normal_y")}}, - {QStringLiteral("use_surf_sdt"), {QStringLiteral("surf_sdt")}}, - {QStringLiteral("use_tracks"), {QStringLiteral("tracks_dbm")}}, - {QStringLiteral("use_gradient_magnitude"), {QStringLiteral("gradient_magnitude")}}, - {QStringLiteral("use_fibers"), {QStringLiteral("fibers")}}, - }; - for (auto it = optionalInputs.begin(); it != optionalInputs.end(); ++it) { - const QString& flag = it.key(); - const bool loaded = loadedConfig.value(flag).toBool(true); - const bool enabled = currentConfig.value(flag).toBool(true); - if (!loaded || enabled) continue; - currentConfig[flag] = true; - for (const QString& pathKey : it.value()) { - currentPaths[pathKey] = loadedPaths.value(pathKey); - } - } - currentRun[QStringLiteral("config")] = currentConfig; - current[QStringLiteral("run")] = currentRun; - current[QStringLiteral("paths")] = currentPaths; - const bool wasReloadRequired = _reloadRequired; _reloadRequired = normalizedReloadRequest(current) != normalizedReloadRequest(_loadedSessionRequest); @@ -1745,9 +1800,6 @@ void SpiralPanel::persist() const const QString prefix = formSettingsPrefix(); for (auto it = _paths.begin(); it != _paths.end(); ++it) settings.setValue(prefix + QStringLiteral("paths/") + it.key(), it.value()->text()); - for (auto it = _optionalInputs.begin(); it != _optionalInputs.end(); ++it) - settings.setValue(prefix + QStringLiteral("use_inputs/") + it.key(), - it.value()->isChecked()); QJsonArray pcls; for (int row = 0; row < _pclList->count(); ++row) { const QListWidgetItem* item = _pclList->item(row); @@ -1764,11 +1816,11 @@ void SpiralPanel::persist() const settings.setValue(prefix + "voxel_size_um", _voxelSize->value()); settings.setValue(prefix + "lasagna_group", _lasagnaGroup->text()); settings.setValue(prefix + "lasagna_scale", _lasagnaScale->value()); - settings.setValue(prefix + "storage_backend", _storageBackend->currentData().toString()); + settings.setValue(prefix + "storage_backend", QStringLiteral("sparse_cuda")); settings.setValue(prefix + "legacy_checkpoint_step", _legacyCheckpointStep->value()); settings.setValue(prefix + "run_tag", _runTag->text()); settings.setValue(prefix + "render_volume_scale", _renderVolumeScale->value()); - settings.setValue(prefix + "save_png_visualizations", _savePngVisualizations->isChecked()); + settings.setValue(prefix + "output_save_png_visualizations", _savePngVisualizations->isChecked()); settings.setValue(prefix + "influence_enabled", _influenceEnabled->isChecked()); settings.setValue(prefix + "influence_z", _influenceZ->value()); settings.setValue(prefix + "influence_windings", _influenceWindings->value()); @@ -1795,9 +1847,6 @@ void SpiralPanel::restore() _applyingResolution = true; for (auto it = _paths.begin(); it != _paths.end(); ++it) it.value()->setText(settings.value(pathsPrefix + it.key()).toString()); - for (auto it = _optionalInputs.begin(); it != _optionalInputs.end(); ++it) - it.value()->setChecked(settings.value( - valuePrefix + QStringLiteral("use_inputs/") + it.key(), true).toBool()); _pclList->clear(); const QByteArray savedPcls = settings.value(valuePrefix + QStringLiteral("pcls")).toByteArray(); const QJsonDocument pclDocument = QJsonDocument::fromJson(savedPcls); @@ -1826,14 +1875,11 @@ void SpiralPanel::restore() _voxelSize->setValue(settings.value(valuePrefix + "voxel_size_um", 9.6).toDouble()); _lasagnaGroup->setText(settings.value(valuePrefix + "lasagna_group", "4").toString()); _lasagnaScale->setValue(settings.value(valuePrefix + "lasagna_scale", 4).toInt()); - const int backend = _storageBackend->findData( - settings.value(valuePrefix + "storage_backend", "sparse_cuda").toString()); - if (backend >= 0) _storageBackend->setCurrentIndex(backend); _legacyCheckpointStep->setValue(settings.value(valuePrefix + "legacy_checkpoint_step", 0).toInt()); _runTag->setText(settings.value(valuePrefix + "run_tag").toString()); _renderVolumeScale->setValue(settings.value(valuePrefix + "render_volume_scale", 16).toInt()); _savePngVisualizations->setChecked( - settings.value(valuePrefix + "save_png_visualizations", false).toBool()); + settings.value(valuePrefix + "output_save_png_visualizations", false).toBool()); _influenceEnabled->setChecked(settings.value(valuePrefix + "influence_enabled", false).toBool()); _influenceZ->setValue(settings.value(valuePrefix + "influence_z", 3000).toInt()); _influenceWindings->setValue(settings.value(valuePrefix + "influence_windings", 5.0).toDouble()); @@ -1853,5 +1899,5 @@ void SpiralPanel::restore() _loadedSessionRequest = {}; _attachedAdvancedConfig = {}; _defaultAdvancedConfig = {}; - updateOptionalInputUi(); + updateTrackSamplingUi(); } diff --git a/volume-cartographer/apps/VC3D/SpiralPanel.hpp b/volume-cartographer/apps/VC3D/SpiralPanel.hpp index 7463f9dfc4..7bb9100336 100644 --- a/volume-cartographer/apps/VC3D/SpiralPanel.hpp +++ b/volume-cartographer/apps/VC3D/SpiralPanel.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -20,8 +21,10 @@ class QPushButton; class QSpinBox; class QDoubleSpinBox; class QPlainTextEdit; +class QProgressBar; class QSlider; class QToolButton; +class QTimer; class SpiralServiceManager; class SpiralConfigProfileEditor; class QFormLayout; @@ -56,13 +59,10 @@ class SpiralPanel : public QWidget QJsonObject influenceConfig() const; QJsonObject sessionAdvancedConfig() const; QJsonObject runAdvancedConfig() const; - void applyOptionalInputConfig(QJsonObject& config, bool includeSelectionFlags) const; void applyTrackSamplingConfig(QJsonObject& config) const; void syncTrackSamplingControlsFromAdvanced(); void writeTrackSamplingControlsToAdvanced(); void updateTrackSamplingUi(); - bool optionalInputEnabled(const QString& key) const; - void updateOptionalInputUi(); void applySessionRunConfig(const QJsonObject& config, qint64 sessionGeneration); void synchronizeSession(const QJsonObject& request, const QJsonObject& status); @@ -88,7 +88,6 @@ class SpiralPanel : public QWidget QHash _paths; QHash _pathBrowseButtons; QHash _visibilityChecks; - QHash _optionalInputs; QHash _pathDirectories; QDialog* _displayDialog = nullptr; QSpinBox* _minimumDisplayedWinding = nullptr; @@ -114,7 +113,6 @@ class SpiralPanel : public QWidget QPushButton* _addPclButton = nullptr; QToolButton* _browsePclButton = nullptr; QComboBox* _outwardSense = nullptr; - QComboBox* _storageBackend = nullptr; QCheckBox* _savePngVisualizations = nullptr; QCheckBox* _trackLengthBinSampling = nullptr; QDoubleSpinBox* _trackShortWeight = nullptr; @@ -135,8 +133,16 @@ class SpiralPanel : public QWidget QPushButton* _stop = nullptr; QPushButton* _save = nullptr; QPushButton* _downloadCheckpoint = nullptr; + QLabel* _checkpointDownloadStatus = nullptr; + QProgressBar* _checkpointDownloadProgress = nullptr; + QTimer* _checkpointDownloadTimer = nullptr; + QElapsedTimer _checkpointDownloadElapsed; + QString _checkpointDownloadPhase; + qint64 _checkpointBytesReceived = 0; + qint64 _checkpointTotalBytes = 0; QPushButton* _refill = nullptr; QLabel* _state = nullptr; + QProgressBar* _previewProgress = nullptr; QLabel* _metrics = nullptr; QLabel* _warnings = nullptr; @@ -167,6 +173,7 @@ class SpiralPanel : public QWidget QJsonObject _attachedAdvancedConfig; QJsonObject _defaultAdvancedConfig; QSet _runConfigKeys; + QSet _runMutablePaths; qint64 _advancedSessionGeneration = -1; QString _currentProfileId; @@ -179,6 +186,9 @@ class SpiralPanel : public QWidget bool _sessionRunnable = false; bool _remoteMode = false; bool _connected = false; + bool _previewTransferActive = false; + bool _checkpointDownloadActive = false; + QString _previewTransferText; int _ephemeralCount = 0; int _uncommittedCount = 0; std::function)> _sessionExitGuard; diff --git a/volume-cartographer/apps/VC3D/SpiralReloadComparison.hpp b/volume-cartographer/apps/VC3D/SpiralReloadComparison.hpp index 6046a8e2ac..13b57d7b61 100644 --- a/volume-cartographer/apps/VC3D/SpiralReloadComparison.hpp +++ b/volume-cartographer/apps/VC3D/SpiralReloadComparison.hpp @@ -9,7 +9,8 @@ namespace vc3d { inline QJsonObject normalizedSpiralReloadRequest( QJsonObject request, const QJsonObject& defaultAdvancedConfig, - const QSet& runConfigKeys) + const QSet& runConfigKeys, + const QSet& runMutablePaths = {}) { QJsonObject run = request.value(QStringLiteral("run")).toObject(); const QJsonObject requestedConfig = @@ -25,6 +26,10 @@ inline QJsonObject normalizedSpiralReloadRequest( run[QStringLiteral("config")] = effectiveConfig; request[QStringLiteral("run")] = run; + QJsonObject paths = request.value(QStringLiteral("paths")).toObject(); + for (const QString& key : runMutablePaths) paths.remove(key); + request[QStringLiteral("paths")] = paths; + // The service expands preview defaults in its canonical session request, // while older panel requests only sent first_winding. Treat both wire // representations as the same loaded preview configuration. diff --git a/volume-cartographer/apps/VC3D/SpiralServiceManager.cpp b/volume-cartographer/apps/VC3D/SpiralServiceManager.cpp index db3191302b..d3f2b0ea2b 100644 --- a/volume-cartographer/apps/VC3D/SpiralServiceManager.cpp +++ b/volume-cartographer/apps/VC3D/SpiralServiceManager.cpp @@ -5,6 +5,7 @@ #include "VCSettings.hpp" #include +#include #include #include #include @@ -16,7 +17,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -39,7 +42,7 @@ constexpr int kRemoteLogPollMs = 500; constexpr int kRestartProbeMs = 500; constexpr int kRestartTimeoutMs = 60000; constexpr int kMutationRetries = 2; -constexpr int kSupportedApiVersion = 12; +constexpr int kSupportedApiVersion = 15; constexpr int kPreviewCacheKept = 3; QString stateName(SpiralServiceManager::ConnectionState state) @@ -69,6 +72,22 @@ SpiralServiceManager::SpiralServiceManager(QObject* parent) : QObject(parent) _remoteLogPoll->setInterval(kRemoteLogPollMs); connect(_remoteLogPoll, &QTimer::timeout, this, &SpiralServiceManager::pollRemoteLogs); + connect(_artifactCache, &SpiralArtifactCache::fetchProgress, this, + [this](const QString& artifactId, const QString& phase, + const QString& fileName, int filesComplete, int totalFiles, + qint64 bytesReceived, qint64 totalBytes) { + if (artifactId != _fetchingPreviewArtifact) return; + emit previewTransferProgress( + phase, fileName, filesComplete, totalFiles, + bytesReceived, totalBytes); + }); + connect(_artifactCache, &SpiralArtifactCache::fetchProgress, this, + [this](const QString& artifactId, const QString& phase, + const QString&, int, int, + qint64 bytesReceived, qint64 totalBytes) { + if (artifactId != _fetchingCheckpointArtifact) return; + emit checkpointDownloadProgress(phase, bytesReceived, totalBytes); + }); connect(_tunnel, &SpiralSshTunnel::logMessage, this, &SpiralServiceManager::logMessage); connect(_tunnel, &SpiralSshTunnel::ready, this, [this](int localPort) { @@ -356,6 +375,12 @@ void SpiralServiceManager::handleHealth(const QJsonObject& health) _remoteLogPoll->stop(); } if (_serviceOwnsDataset) fetchAdvertisedDataset(); + get(QStringLiteral("/configuration"), Timeout::Quick, + [this](const QJsonObject& catalog) { + _configurationDefaults = + catalog.value(QStringLiteral("defaults")).toObject(); + emit configurationCatalogChanged(catalog); + }); pollStatus(); } @@ -706,14 +731,66 @@ void SpiralServiceManager::uploadCheckpointForResume( void SpiralServiceManager::runIterations(int iterations, const QJsonObject& influenceConfig, - const QJsonObject& runConfig) + const QJsonObject& runConfig, + const QJsonObject& inputs) { - postWithRetry(QStringLiteral("/session/run"), - {{QStringLiteral("command_id"), commandId()}, - {QStringLiteral("iterations"), iterations}, - {QStringLiteral("influence_config"), influenceConfig}, - {QStringLiteral("run_config"), runConfig}}, - Timeout::Command, kMutationRetries, {}); + QJsonObject configuration = _configurationDefaults; + for (auto it = _appliedConfiguration.begin(); + it != _appliedConfiguration.end(); ++it) + configuration[it.key()] = it.value(); + for (auto it = runConfig.begin(); it != runConfig.end(); ++it) + configuration[it.key()] = it.value(); + post(QStringLiteral("/session/run/plan"), + {{QStringLiteral("configuration"), configuration}, + {QStringLiteral("iterations"), iterations}, + {QStringLiteral("influence"), influenceConfig}, + {QStringLiteral("inputs"), inputs}, + {QStringLiteral("expected_session_revision"), _sessionRevision}}, + Timeout::Command, + [this](const QJsonObject& plan) { + if (plan.value(QStringLiteral("new_fit_required")).toBool()) { + QMessageBox::information( + QApplication::activeWindow(), tr("Start New Fit required"), + tr("These changes are incompatible with the resident model. " + "Use New Fit to apply them.")); + emit configurationReviewRequested(); + return; + } + if (plan.value(QStringLiteral("session_reload_required")).toBool()) { + QMessageBox::information( + QApplication::activeWindow(), tr("Reload fit inputs required"), + tr("These changes require reloading the resident fit inputs. " + "Use New Fit to apply them.")); + emit configurationReviewRequested(); + return; + } + const bool changed = + !plan.value(QStringLiteral("changes")).toArray().isEmpty() + || plan.value(QStringLiteral("input_changed")).toBool(); + if (changed) { + QMessageBox box(QMessageBox::Question, + tr("Spiral configuration changed"), + tr("Configuration changed since last run: continue?"), + QMessageBox::Cancel, + QApplication::activeWindow()); + auto* proceed = box.addButton(tr("Continue Run"), + QMessageBox::AcceptRole); + auto* review = box.addButton(tr("Review"), + QMessageBox::ActionRole); + box.exec(); + if (box.clickedButton() == review) { + emit configurationReviewRequested(); + return; + } + if (box.clickedButton() != proceed) return; + } + postWithRetry( + QStringLiteral("/session/run"), + {{QStringLiteral("command_id"), commandId()}, + {QStringLiteral("plan_token"), + plan.value(QStringLiteral("plan_token"))}}, + Timeout::Command, kMutationRetries, {}); + }); } void SpiralServiceManager::stopAfterIteration() @@ -733,6 +810,7 @@ void SpiralServiceManager::saveCheckpoint(const QString& path) void SpiralServiceManager::downloadCheckpoint(const QString& localPath) { + emit checkpointDownloadProgress(QStringLiteral("creating"), 0, 0); postWithRetry( QStringLiteral("/session/download-checkpoint"), {{QStringLiteral("command_id"), commandId()}}, @@ -745,13 +823,16 @@ void SpiralServiceManager::downloadCheckpoint(const QString& localPath) emit checkpointDownloadFinished(localPath, tr("The service did not return a checkpoint artifact")); return; } + _fetchingCheckpointArtifact = artifactId; _artifactCache->fetchArtifact( sessionId, artifactId, [this, localPath](const QString& entryPath, const QString& error, bool) { + _fetchingCheckpointArtifact.clear(); if (entryPath.isEmpty()) { emit checkpointDownloadFinished(localPath, error); return; } + emit checkpointDownloadProgress(QStringLiteral("copying"), 0, 0); // Atomic replacement: a failed transfer cannot leave a // partial file at the selected destination. const QString temporary = localPath + QStringLiteral(".part"); @@ -1127,6 +1208,11 @@ void SpiralServiceManager::handleStatus(const QJsonObject& status) const qint64 generation = status.value(QStringLiteral("generation")).toInteger(-1); if (generation < _lastStatusGeneration) return; _lastStatusGeneration = generation; + _sessionRevision = + status.value(QStringLiteral("session_revision")).toInteger(); + const QJsonObject applied = + status.value(QStringLiteral("applied_config")).toObject(); + if (!applied.isEmpty()) _appliedConfiguration = applied; const QString sessionId = status.value(QStringLiteral("session_id")).toString(); const bool active = !sessionId.isEmpty() diff --git a/volume-cartographer/apps/VC3D/SpiralServiceManager.hpp b/volume-cartographer/apps/VC3D/SpiralServiceManager.hpp index bfe0571da7..676ab1e67e 100644 --- a/volume-cartographer/apps/VC3D/SpiralServiceManager.hpp +++ b/volume-cartographer/apps/VC3D/SpiralServiceManager.hpp @@ -61,7 +61,8 @@ class SpiralServiceManager : public QObject void resolveDataset(const QString& root); void loadSession(QJsonObject request); void runIterations(int iterations, const QJsonObject& influenceConfig, - const QJsonObject& runConfig); + const QJsonObject& runConfig, + const QJsonObject& inputs = {}); void stopAfterIteration(); // Save on service: writes to a service-host path. void saveCheckpoint(const QString& path); @@ -86,6 +87,8 @@ class SpiralServiceManager : public QObject const QString& message); void serviceStateChanged(const QString& state); void datasetResolved(const QJsonObject& resolution); + void configurationCatalogChanged(const QJsonObject& catalog); + void configurationReviewRequested(); // Emitted once when this connection first observes a resident session, // whether VC3D loaded it or attached after another client did. void sessionSynchronized(const QJsonObject& sessionRequest, @@ -94,6 +97,11 @@ class SpiralServiceManager : public QObject void sessionActiveChanged(bool active); // Local (cache) filesystem paths: artifact transfers already happened. void previewAvailable(const QString& manifestPath, qint64 generation); + void previewTransferProgress(const QString& phase, const QString& fileName, + int filesComplete, int totalFiles, + qint64 bytesReceived, qint64 totalBytes); + void checkpointDownloadProgress(const QString& phase, + qint64 bytesReceived, qint64 totalBytes); void checkpointDownloadFinished(const QString& localPath, const QString& error); void checkpointUploadProgress(qint64 sentBytes, qint64 totalBytes); void inputUploadFinished(const QString& inputId, const QString& error); @@ -171,11 +179,15 @@ class SpiralServiceManager : public QObject int _remoteLogFailures = 0; qint64 _lastRemoteLogSequence = 0; QJsonObject _advertisedDataset; + QJsonObject _configurationDefaults; + QJsonObject _appliedConfiguration; + qint64 _sessionRevision = 0; quint64 _commandCounter = 0; qint64 _lastStatusGeneration = -1; QString _installedPreviewArtifact; QString _installedPreviewSession; QString _fetchingPreviewArtifact; + QString _fetchingCheckpointArtifact; qint64 _previewSequence = 0; QString _lastPreviewLocalPath; QString _synchronizedSessionId; diff --git a/volume-cartographer/apps/VC3D/VCAppMain.cpp b/volume-cartographer/apps/VC3D/VCAppMain.cpp index ce4ae16aca..19280dc49a 100644 --- a/volume-cartographer/apps/VC3D/VCAppMain.cpp +++ b/volume-cartographer/apps/VC3D/VCAppMain.cpp @@ -6,7 +6,11 @@ #endif #include +#include #include +#include +#include +#include #include "CWindow.hpp" #include "agent_bridge/AgentBridgeServer.hpp" @@ -90,6 +94,37 @@ static bool hasCliFlag(int argc, char* argv[], const char* flag) return false; } +class WheelFocusFilter final : public QObject +{ +protected: + bool eventFilter(QObject* watched, QEvent* event) override + { + if (event->type() != QEvent::Wheel) + return QObject::eventFilter(watched, event); + + auto* widget = qobject_cast(watched); + while (widget) { + if (auto* spinBox = qobject_cast(widget)) { + if (!spinBox->hasFocus()) { + event->ignore(); + return true; + } + break; + } + if (auto* comboBox = qobject_cast(widget)) { + if (!comboBox->hasFocus()) { + event->ignore(); + return true; + } + break; + } + widget = widget->parentWidget(); + } + + return QObject::eventFilter(watched, event); + } +}; + #if defined(__GNUC__) || defined(__clang__) __attribute__((visibility("default"))) #endif @@ -200,6 +235,16 @@ auto main(int argc, char* argv[]) -> int } QApplication app(argc, argv); + // Wide surface-aligned Spiral overlays can legitimately exceed Qt's + // 128-MiB default image-I/O allocation limit. Set the Qt runtime value + // after QApplication construction because Qt may cache the environment + // setting before VC3D's pre-main hook runs. Preserve an explicit user + // override and retain a finite allocation guard. + if (qEnvironmentVariableIsEmpty("QT_IMAGEIO_MAXALLOC")) { + QImageReader::setAllocationLimit(512); + } + WheelFocusFilter wheelFocusFilter; + app.installEventFilter(&wheelFocusFilter); QApplication::setOrganizationName("Vesuvius Challenge"); QApplication::setApplicationName("VC3D"); QApplication::setWindowIcon(QIcon(":/images/logo.png")); diff --git a/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.cpp b/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.cpp index 75ed30310a..36276035d6 100644 --- a/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.cpp +++ b/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.cpp @@ -1,31 +1,43 @@ #include "elements/SpiralConfigProfileEditor.hpp" #include "VCSettings.hpp" +#include "elements/CollapsibleSettingsGroup.hpp" #include +#include +#include #include #include #include #include #include #include +#include #include #include #include #include #include +#include #include #include +#include #include #include +#include +#include #include +#include +#include #include #include #include +#include namespace { const QString kDefaultId = QStringLiteral("default"); +const QString kDefaultsId = QStringLiteral("catalog-defaults"); const QString kCustomId = QStringLiteral("custom"); QString profileStorePath() @@ -103,6 +115,162 @@ QString formatted(const QJsonObject& object) { return QString::fromUtf8(QJsonDocument(object).toJson(QJsonDocument::Indented)).trimmed(); } + +QString impactLabel(const QString& impact) +{ + if (impact == QStringLiteral("run_boundary")) return QObject::tr("Next Run"); + if (impact == QStringLiteral("shell_reload")) return QObject::tr("Reload shell"); + if (impact == QStringLiteral("prepared_input_rebuild")) + return QObject::tr("Reload fit inputs"); + if (impact == QStringLiteral("new_fit")) return QObject::tr("Start New Fit"); + return impact; +} + +QString groupTitle(const QString& prefix) +{ + static const QHash titles{ + {QStringLiteral("optimizer"), QObject::tr("Optimizer")}, + {QStringLiteral("model"), QObject::tr("Model")}, + {QStringLiteral("patch"), QObject::tr("Patch")}, + {QStringLiteral("sample"), QObject::tr("Sample Count")}, + {QStringLiteral("input"), QObject::tr("Input")}, + {QStringLiteral("pcl"), QObject::tr("PCL")}, + {QStringLiteral("tracks"), QObject::tr("Tracks")}, + {QStringLiteral("dense"), QObject::tr("Dense")}, + {QStringLiteral("loss"), QObject::tr("Loss")}, + {QStringLiteral("dt"), QObject::tr("DT")}, + {QStringLiteral("output"), QObject::tr("Output")}, + {QStringLiteral("shell"), QObject::tr("Shell")}, + {QStringLiteral("influence"), QObject::tr("Influence")}, + }; + const auto title = titles.constFind(prefix); + if (title != titles.cend()) return *title; + QString fallback = prefix; + if (!fallback.isEmpty()) fallback[0] = fallback[0].toUpper(); + return fallback; +} + +int groupOrder(const QString& prefix) +{ + static const std::array order{ + QStringLiteral("sample"), QStringLiteral("loss"), + QStringLiteral("patch"), QStringLiteral("tracks"), + QStringLiteral("optimizer"), QStringLiteral("model"), + QStringLiteral("input"), QStringLiteral("pcl"), + QStringLiteral("dense"), QStringLiteral("dt"), + QStringLiteral("output"), QStringLiteral("shell"), + QStringLiteral("influence"), + }; + const auto found = std::find(order.cbegin(), order.cend(), prefix); + return found == order.cend() + ? static_cast(order.size()) + : static_cast(std::distance(order.cbegin(), found)); +} + +class ResponsiveGroupGrid final : public QWidget +{ +public: + explicit ResponsiveGroupGrid(QWidget* parent = nullptr) + : QWidget(parent) + , _layout(new QHBoxLayout(this)) + { + setObjectName(QStringLiteral("spiralConfigGroupGrid")); + setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Preferred); + _layout->setContentsMargins(0, 0, 0, 0); + _layout->setSpacing(12); + for (int column = 0; column < 3; ++column) { + _columnWidgets[static_cast(column)] = new QWidget(this); + auto* columnLayout = new QVBoxLayout( + _columnWidgets[static_cast(column)]); + columnLayout->setContentsMargins(0, 0, 0, 0); + columnLayout->setSpacing(12); + _columnLayouts[static_cast(column)] = columnLayout; + _layout->addWidget( + _columnWidgets[static_cast(column)], 1); + } + setProperty("spiralColumnCount", 1); + } + + void addGroup(QWidget* group) + { + _groups.push_back(group); + if (auto* collapsible = + qobject_cast(group)) { + connect(collapsible, &CollapsibleSettingsGroup::toggled, + this, [this] { refresh(); }); + } + reflow(width()); + } + + void refresh() + { + reflow(width(), true); + } + + QSize sizeHint() const override + { + QSize hint = QWidget::sizeHint(); + hint.setWidth(400); + return hint; + } + +protected: + void resizeEvent(QResizeEvent* event) override + { + QWidget::resizeEvent(event); + reflow(event->size().width()); + } + +private: + void reflow(int width, bool force = false) + { + const int available = std::max(1, width - _layout->contentsMargins().left() + - _layout->contentsMargins().right()); + const int columns = std::clamp(available / 400, 1, 3); + const int visibleGroups = static_cast(std::count_if( + _groups.cbegin(), _groups.cend(), + [](const QWidget* group) { + return group->property("spiralFilterVisible").toBool(); + })); + if (!force && columns == _columns && visibleGroups == _visibleGroups) + return; + + for (int column = 0; column < 3; ++column) { + QVBoxLayout* columnLayout = + _columnLayouts[static_cast(column)]; + while (columnLayout->count() > 0) + delete columnLayout->takeAt(0); + _columnWidgets[static_cast(column)] + ->setVisible(column < columns); + } + + std::array columnHeights{0, 0, 0}; + for (QWidget* group : _groups) { + if (!group->property("spiralFilterVisible").toBool()) continue; + const auto shortest = std::min_element( + columnHeights.begin(), columnHeights.begin() + columns); + const int column = static_cast( + std::distance(columnHeights.begin(), shortest)); + _columnLayouts[static_cast(column)] + ->addWidget(group, 0, Qt::AlignTop); + *shortest += group->sizeHint().height() + 12; + } + for (int column = 0; column < columns; ++column) + _columnLayouts[static_cast(column)]->addStretch(1); + + _columns = columns; + _visibleGroups = visibleGroups; + setProperty("spiralColumnCount", columns); + updateGeometry(); + } + + QHBoxLayout* _layout; + std::array _columnWidgets{}; + std::array _columnLayouts{}; + std::vector _groups; + int _columns = 0; + int _visibleGroups = 0; +}; } SpiralConfigProfileEditor::SpiralConfigProfileEditor(QWidget* parent) @@ -167,9 +335,30 @@ SpiralConfigProfileEditor::SpiralConfigProfileEditor(QWidget* parent) _dialog->setObjectName(QStringLiteral("spiralAdvancedConfigDialog")); _dialog->setWindowTitle(tr("Spiral Advanced Config JSON")); _dialog->setModal(false); - _dialog->resize(720, 620); + _dialog->resize(1280, 760); _dialog->setLayout(new QVBoxLayout); _dialog->installEventFilter(this); + root->removeWidget(_editorContents); + _inlinePopInButton->setText(tr("Open Spiral Configuration…")); + _inlinePopInButton->show(); + auto* tabs = new QTabWidget(_dialog); + _controlsPage = new QWidget(tabs); + auto* controlsLayout = new QVBoxLayout(_controlsPage); + _search = new QLineEdit(_controlsPage); + _search->setPlaceholderText(tr("Search controls…")); + controlsLayout->addWidget(_search); + _controlsScroll = new QScrollArea(_controlsPage); + _controlsScroll->setObjectName(QStringLiteral("spiralConfigControlsScroll")); + _controlsScroll->setWidgetResizable(true); + _controlsScroll->setFrameShape(QFrame::NoFrame); + _controlsScroll->setHorizontalScrollBarPolicy(Qt::ScrollBarAlwaysOff); + _controlsGrid = new ResponsiveGroupGrid(_controlsScroll); + _controlsScroll->setWidget(_controlsGrid); + controlsLayout->addWidget(_controlsScroll); + tabs->addTab(_controlsPage, tr("Controls")); + tabs->addTab(_editorContents, tr("Expert JSON")); + _dialog->layout()->addWidget(tabs); + _poppedOut = true; connect(_profileCombo, qOverload(&QComboBox::currentIndexChanged), this, [this](int index) { @@ -184,11 +373,11 @@ SpiralConfigProfileEditor::SpiralConfigProfileEditor(QWidget* parent) this, &SpiralConfigProfileEditor::renameCurrent); connect(_deleteButton, &QPushButton::clicked, this, &SpiralConfigProfileEditor::deleteCurrent); - connect(_popButton, &QPushButton::clicked, this, [this]() { - _poppedOut ? popIn() : popOut(); - }); + _popButton->hide(); connect(_inlinePopInButton, &QPushButton::clicked, - this, &SpiralConfigProfileEditor::popIn); + this, &SpiralConfigProfileEditor::showWindow); + connect(_search, &QLineEdit::textChanged, + this, &SpiralConfigProfileEditor::filterControls); loadProfiles(); rebuildCombo(); @@ -226,6 +415,7 @@ void SpiralConfigProfileEditor::setSessionDefault(const QJsonObject& config) _cleanText = _sessionDefaultText; _dirty = false; validateCurrentText(); + jsonToControls(); updateUi(); emit textChanged(); } @@ -237,19 +427,286 @@ void SpiralConfigProfileEditor::showSessionDefault() void SpiralConfigProfileEditor::clearSessionDefault() { - _sessionDefaultText = QStringLiteral("{}"); - if (isDefaultProfile()) setSessionDefault(QJsonObject{}); + const QJsonObject defaults = + _catalog.value(QStringLiteral("defaults")).toObject(); + _sessionDefaultText = formatted(defaults); + if (isDefaultProfile()) setSessionDefault(defaults); +} + +void SpiralConfigProfileEditor::setCatalog(const QJsonObject& catalog) +{ + _catalog = catalog; + if (_sessionDefaultText.trimmed() == QStringLiteral("{}")) + _sessionDefaultText = + formatted(_catalog.value(QStringLiteral("defaults")).toObject()); + rebuildCombo(); + rebuildControls(); + if (isDefaultProfile()) + applyCurrentProfileText(); + else + jsonToControls(); +} + +void SpiralConfigProfileEditor::rebuildControls() +{ + _controlGroups.clear(); + _preSearchExpanded.clear(); + _fieldEditors.clear(); + QWidget* oldGrid = _controlsScroll->takeWidget(); + delete oldGrid; + _controlsGrid = new ResponsiveGroupGrid(_controlsScroll); + _controlsScroll->setWidget(_controlsGrid); + + const QJsonObject schema = _catalog.value(QStringLiteral("schema")).toObject(); + const QJsonObject fields = schema.value(QStringLiteral("fields")).toObject(); + QMap> fieldsByPrefix; + for (auto it = fields.begin(); it != fields.end(); ++it) { + fieldsByPrefix[it.key().section('_', 0, 0)].push_back(it.key()); + } + QVector prefixes(fieldsByPrefix.keyBegin(), fieldsByPrefix.keyEnd()); + std::stable_sort(prefixes.begin(), prefixes.end(), + [](const QString& left, const QString& right) { + const int leftOrder = groupOrder(left); + const int rightOrder = groupOrder(right); + return leftOrder == rightOrder ? left < right : leftOrder < rightOrder; + }); + + for (const QString& prefix : prefixes) { + auto* group = new CollapsibleSettingsGroup( + groupTitle(prefix), _controlsGrid); + group->setObjectName(QStringLiteral("spiralConfigGroup_%1").arg(prefix)); + group->setProperty("spiralConfigPrefix", prefix); + group->setProperty("spiralFilterVisible", true); + group->setExpanded(true); + group->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Maximum); + + ControlGroup groupState{prefix, group, {}}; + for (const QString& key : fieldsByPrefix.value(prefix)) { + const QJsonObject spec = fields.value(key).toObject(); + const QString labelText = + spec.value(QStringLiteral("label")).toString(); + const QString impactText = impactLabel( + spec.value(QStringLiteral("runtime_impact")).toString()); + const QString type = spec.value(QStringLiteral("type")).toString(); + QWidget* editor = nullptr; + if (type == QStringLiteral("boolean")) { + auto* value = new QCheckBox(group); + connect(value, &QCheckBox::toggled, + this, &SpiralConfigProfileEditor::controlsToJson); + editor = value; + } else if (type == QStringLiteral("integer")) { + auto* value = new QSpinBox(group); + value->setRange(spec.value(QStringLiteral("minimum")).toInt(), + spec.value(QStringLiteral("maximum")).toInt()); + connect(value, qOverload(&QSpinBox::valueChanged), + this, [this] { controlsToJson(); }); + editor = value; + } else if (type == QStringLiteral("number")) { + auto* value = new QDoubleSpinBox(group); + value->setRange(spec.value(QStringLiteral("minimum")).toDouble(), + spec.value(QStringLiteral("maximum")).toDouble()); + value->setDecimals( + spec.value(QStringLiteral("precision")).toInt(6)); + value->setSingleStep( + spec.value(QStringLiteral("step")).toDouble(.01)); + connect(value, qOverload(&QDoubleSpinBox::valueChanged), + this, [this] { controlsToJson(); }); + editor = value; + } else if (type == QStringLiteral("enum")) { + auto* value = new QComboBox(group); + for (const QJsonValue& option : + spec.value(QStringLiteral("values")).toArray()) + value->addItem(option.toString()); + connect(value, qOverload(&QComboBox::currentIndexChanged), + this, [this] { controlsToJson(); }); + editor = value; + } else { + auto* value = new QLineEdit(group); + value->setPlaceholderText( + type == QStringLiteral("dictionary") + ? tr("{ key: value }") : tr("[ values ]")); + connect(value, &QLineEdit::editingFinished, + this, &SpiralConfigProfileEditor::controlsToJson); + editor = value; + } + editor->setObjectName( + QStringLiteral("spiralConfigEditor_%1").arg(key)); + editor->setProperty("spiralConfigKey", key); + editor->setToolTip(key); + QWidget* displayed = editor; + if (spec.value(QStringLiteral("nullable")).toBool()) { + displayed = new QWidget(group); + auto* nullableRow = new QHBoxLayout(displayed); + nullableRow->setContentsMargins(0, 0, 0, 0); + auto* enabled = new QCheckBox(tr("Enabled"), displayed); + editor->setProperty("nullableToggle", + QVariant::fromValue(enabled)); + nullableRow->addWidget(enabled); + nullableRow->addWidget(editor, 1); + connect(enabled, &QCheckBox::toggled, editor, + [this, editor](bool on) { + editor->setEnabled(on); + controlsToJson(); + }); + } + _fieldEditors.insert(key, editor); + + auto* rowWidget = new QWidget(group->contentWidget()); + rowWidget->setObjectName( + QStringLiteral("spiralConfigRow_%1").arg(key)); + rowWidget->setProperty("spiralConfigKey", key); + auto* row = new QGridLayout(rowWidget); + row->setContentsMargins(0, 0, 0, 0); + row->setHorizontalSpacing(8); + auto* label = new QLabel(labelText, rowWidget); + label->setObjectName( + QStringLiteral("spiralConfigLabel_%1").arg(key)); + label->setToolTip(key); + auto* impact = new QLabel(impactText, rowWidget); + impact->setObjectName( + QStringLiteral("spiralConfigImpact_%1").arg(key)); + impact->setProperty( + "spiralRuntimeImpact", + spec.value(QStringLiteral("runtime_impact")).toString()); + impact->setForegroundRole(QPalette::WindowText); + impact->setAlignment(Qt::AlignRight | Qt::AlignVCenter); + impact->setSizePolicy( + QSizePolicy::Minimum, QSizePolicy::Preferred); + row->addWidget(label, 0, 0); + row->addWidget(displayed, 0, 1); + row->addWidget(impact, 0, 2); + row->setColumnStretch(1, 1); + group->addFullWidthWidget(rowWidget, key); + groupState.rows.push_back( + {key, labelText + QLatin1Char(' ') + key + QLatin1Char(' ') + + impactText, + rowWidget}); + } + _controlGroups.push_back(groupState); + static_cast(_controlsGrid)->addGroup(group); + } + filterControls(_search->text()); +} + +void SpiralConfigProfileEditor::filterControls(const QString& text) +{ + const QString query = text.trimmed(); + const bool searching = !query.isEmpty(); + if (searching && _preSearchExpanded.isEmpty()) { + for (const ControlGroup& group : _controlGroups) + _preSearchExpanded.insert(group.widget, group.widget->isExpanded()); + } + + for (ControlGroup& group : _controlGroups) { + bool anyVisible = false; + for (ControlRow& row : group.rows) { + const bool visible = !searching + || row.searchableText.contains(query, Qt::CaseInsensitive); + row.widget->setVisible(visible); + anyVisible |= visible; + } + group.widget->setVisible(anyVisible); + group.widget->setProperty("spiralFilterVisible", anyVisible); + if (searching && anyVisible) + group.widget->setExpanded(true); + } + + if (!searching && !_preSearchExpanded.isEmpty()) { + for (ControlGroup& group : _controlGroups) + group.widget->setExpanded( + _preSearchExpanded.value(group.widget, true)); + _preSearchExpanded.clear(); + } + static_cast(_controlsGrid)->refresh(); +} + +void SpiralConfigProfileEditor::controlsToJson() +{ + if (_programmatic) return; + QJsonObject values; + for (auto it = _fieldEditors.begin(); it != _fieldEditors.end(); ++it) { + QWidget* editor = it.value(); + QJsonValue value; + auto* nullable = qobject_cast( + editor->property("nullableToggle").value()); + if (nullable && !nullable->isChecked()) + value = QJsonValue::Null; + else if (auto* widget = qobject_cast(editor)) + value = widget->isChecked(); + else if (auto* widget = qobject_cast(editor)) + value = widget->value(); + else if (auto* widget = qobject_cast(editor)) + value = widget->value(); + else if (auto* widget = qobject_cast(editor)) + value = widget->currentText(); + else { + const QByteArray text = + qobject_cast(editor)->text().toUtf8(); + QJsonParseError error; + const QJsonDocument document = QJsonDocument::fromJson(text, &error); + if (error.error != QJsonParseError::NoError) return; + value = document.isArray() ? QJsonValue(document.array()) + : QJsonValue(document.object()); + } + values[it.key()] = value; + } + setCurrentText(formatted(values)); +} + +void SpiralConfigProfileEditor::jsonToControls() +{ + const QJsonDocument document = QJsonDocument::fromJson(currentText().toUtf8()); + if (!document.isObject()) return; + const bool previous = _programmatic; + _programmatic = true; + QJsonObject values = _catalog.value(QStringLiteral("defaults")).toObject(); + const QJsonObject overrides = document.object(); + for (auto it = overrides.begin(); it != overrides.end(); ++it) + values.insert(it.key(), it.value()); + for (auto it = _fieldEditors.begin(); it != _fieldEditors.end(); ++it) { + const QJsonValue value = values.value(it.key()); + auto* nullable = qobject_cast( + it.value()->property("nullableToggle").value()); + if (nullable) { + nullable->setChecked(!value.isNull()); + it.value()->setEnabled(!value.isNull()); + if (value.isNull()) continue; + } + if (auto* widget = qobject_cast(it.value())) + widget->setChecked(value.toBool()); + else if (auto* widget = qobject_cast(it.value())) + widget->setValue(value.toInt()); + else if (auto* widget = qobject_cast(it.value())) + widget->setValue(value.toDouble()); + else if (auto* widget = qobject_cast(it.value())) + widget->setCurrentText(value.toString()); + else { + const QJsonDocument encoded = value.isArray() + ? QJsonDocument(value.toArray()) + : QJsonDocument(value.toObject()); + qobject_cast(it.value())->setText( + QString::fromUtf8(encoded.toJson(QJsonDocument::Compact))); + } + } + _programmatic = previous; } bool SpiralConfigProfileEditor::eventFilter(QObject* watched, QEvent* event) { if (watched == _dialog && event->type() == QEvent::Close) { - popIn(); + _dialog->hide(); return true; } return QWidget::eventFilter(watched, event); } +void SpiralConfigProfileEditor::showWindow() +{ + _dialog->show(); + _dialog->raise(); + _dialog->activateWindow(); +} + void SpiralConfigProfileEditor::loadProfiles() { _profiles.clear(); @@ -296,7 +753,11 @@ void SpiralConfigProfileEditor::rebuildCombo() _programmatic = true; const QSignalBlocker blocker(_profileCombo); _profileCombo->clear(); - _profileCombo->addItem(tr("Default"), kDefaultId); + _profileCombo->addItem(tr("Current Session"), kDefaultId); + _profileCombo->addItem(tr("Defaults"), kDefaultsId); + const QJsonObject presets = _catalog.value(QStringLiteral("presets")).toObject(); + for (auto it = presets.begin(); it != presets.end(); ++it) + _profileCombo->addItem(it.key(), QStringLiteral("preset:") + it.key()); _profileCombo->addItem(tr("Custom"), kCustomId); for (const StoredProfile& profile : _profiles) _profileCombo->addItem(profile.name, profile.id); @@ -326,6 +787,11 @@ void SpiralConfigProfileEditor::applyCurrentProfileText() { QString text; if (_currentProfileId == kDefaultId) text = _sessionDefaultText; + else if (_currentProfileId == kDefaultsId) + text = formatted(_catalog.value(QStringLiteral("defaults")).toObject()); + else if (_currentProfileId.startsWith(QStringLiteral("preset:"))) + text = formatted(_catalog.value(QStringLiteral("presets")).toObject() + .value(_currentProfileId.mid(7)).toObject()); else if (_currentProfileId == kCustomId) text = _customText; else if (const StoredProfile* profile = findStored(_currentProfileId)) text = profile->jsonText; else { @@ -338,6 +804,7 @@ void SpiralConfigProfileEditor::applyCurrentProfileText() _cleanText = text; _dirty = false; validateCurrentText(); + jsonToControls(); updateUi(); } @@ -356,6 +823,7 @@ void SpiralConfigProfileEditor::handleTextEdited() } _dirty = text != _cleanText; validateCurrentText(); + jsonToControls(); updateUi(); emit textChanged(); } diff --git a/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.hpp b/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.hpp index ed4c4d7580..0bd63ba468 100644 --- a/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.hpp +++ b/volume-cartographer/apps/VC3D/elements/SpiralConfigProfileEditor.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -10,6 +11,8 @@ class QEvent; class QLabel; class QPlainTextEdit; class QPushButton; +class QScrollArea; +class CollapsibleSettingsGroup; class SpiralConfigProfileEditor final : public QWidget { @@ -29,6 +32,8 @@ class SpiralConfigProfileEditor final : public QWidget void setSessionDefault(const QJsonObject& config); void showSessionDefault(); void clearSessionDefault(); + void setCatalog(const QJsonObject& catalog); + void showWindow(); signals: void textChanged(); @@ -65,6 +70,21 @@ class SpiralConfigProfileEditor final : public QWidget void popOut(); void popIn(); + void rebuildControls(); + void controlsToJson(); + void jsonToControls(); + void filterControls(const QString& text); + + struct ControlRow { + QString key; + QString searchableText; + QWidget* widget = nullptr; + }; + struct ControlGroup { + QString prefix; + CollapsibleSettingsGroup* widget = nullptr; + QVector rows; + }; QWidget* _editorContents = nullptr; QComboBox* _profileCombo = nullptr; @@ -77,6 +97,14 @@ class SpiralConfigProfileEditor final : public QWidget QPlainTextEdit* _textEdit = nullptr; QLabel* _statusLabel = nullptr; QDialog* _dialog = nullptr; + QJsonObject _catalog; + QWidget* _controlsPage = nullptr; + QScrollArea* _controlsScroll = nullptr; + QWidget* _controlsGrid = nullptr; + class QLineEdit* _search = nullptr; + QHash _fieldEditors; + QVector _controlGroups; + QHash _preSearchExpanded; QVector _profiles; QString _currentProfileId = QStringLiteral("default"); diff --git a/volume-cartographer/apps/VC3D/test/CMakeLists.txt b/volume-cartographer/apps/VC3D/test/CMakeLists.txt index 4da14e1fda..9b564d25f4 100644 --- a/volume-cartographer/apps/VC3D/test/CMakeLists.txt +++ b/volume-cartographer/apps/VC3D/test/CMakeLists.txt @@ -9,6 +9,23 @@ # in, keeping the test fast and self-contained. find_package(Qt6 REQUIRED COMPONENTS Core Test Widgets) +add_executable(test_spiral_reload_comparison + test_spiral_reload_comparison.cpp) + +set_target_properties(test_spiral_reload_comparison PROPERTIES + AUTOMOC ON AUTOUIC OFF AUTORCC OFF) + +target_include_directories(test_spiral_reload_comparison PRIVATE + ${PROJECT_SOURCE_DIR}/apps/VC3D) + +target_link_libraries(test_spiral_reload_comparison PRIVATE + Qt6::Core + Qt6::Test) + +add_test(NAME spiral_reload_comparison COMMAND test_spiral_reload_comparison) +set_property(GLOBAL APPEND PROPERTY VC_TEST_TARGETS + test_spiral_reload_comparison) + add_executable(test_seeding_batch_tracker test_seeding_batch_tracker.cpp ${PROJECT_SOURCE_DIR}/apps/VC3D/SeedingBatchTracker.cpp) diff --git a/volume-cartographer/apps/VC3D/test/test_spiral_reload_comparison.cpp b/volume-cartographer/apps/VC3D/test/test_spiral_reload_comparison.cpp new file mode 100644 index 0000000000..daa9fe4fff --- /dev/null +++ b/volume-cartographer/apps/VC3D/test/test_spiral_reload_comparison.cpp @@ -0,0 +1,61 @@ +#include "SpiralReloadComparison.hpp" + +#include +#include + +class SpiralReloadComparisonTest final : public QObject +{ + Q_OBJECT + +private slots: + void runMutableConfigAndShellPathDoNotRequireFullReload() + { + const QJsonObject defaults{ + {"track_dt_loss_margin", 0.025}, + {"patch_erode_patches", 2}, + }; + QJsonObject loaded{ + {"paths", QJsonObject{ + {"outer_shell", "/shell/one"}, + {"verified_patches", "/patches"}, + }}, + {"run", QJsonObject{{"config", defaults}}}, + }; + QJsonObject current = loaded; + current["paths"] = QJsonObject{ + {"outer_shell", "/shell/two"}, + {"verified_patches", "/patches"}, + }; + QJsonObject run = current["run"].toObject(); + QJsonObject config = run["config"].toObject(); + config["track_dt_loss_margin"] = 0.5; + run["config"] = config; + current["run"] = run; + + const QSet mutableConfig{"track_dt_loss_margin"}; + const QSet mutablePaths{"outer_shell"}; + QCOMPARE( + vc3d::normalizedSpiralReloadRequest( + current, defaults, mutableConfig, mutablePaths), + vc3d::normalizedSpiralReloadRequest( + loaded, defaults, mutableConfig, mutablePaths)); + } + + void PatchPathChangeStillRequiresReload() + { + const QJsonObject defaults; + QJsonObject loaded{ + {"paths", QJsonObject{{"verified_patches", "/patches"}}}, + {"run", QJsonObject{{"config", defaults}}}, + }; + QJsonObject current = loaded; + current["paths"] = QJsonObject{{"verified_patches", "/other-patches"}}; + + QVERIFY( + vc3d::normalizedSpiralReloadRequest(current, defaults, {}) + != vc3d::normalizedSpiralReloadRequest(loaded, defaults, {})); + } +}; + +QTEST_APPLESS_MAIN(SpiralReloadComparisonTest) +#include "test_spiral_reload_comparison.moc" diff --git a/volume-cartographer/core/test/CMakeLists.txt b/volume-cartographer/core/test/CMakeLists.txt index 33e206640d..c6e0b80d8e 100644 --- a/volume-cartographer/core/test/CMakeLists.txt +++ b/volume-cartographer/core/test/CMakeLists.txt @@ -189,6 +189,16 @@ vc_add_test(NAME test_segmentation_lasagna_panel_ui INCLUDES ${PROJECT_SOURCE_DIR}/apps/VC3D DEFS VC_TEST_DISABLE_LASAGNA_DISCOVERY VC_TEST_DISABLE_LASAGNA_DIALOGS) +qt6_wrap_cpp(spiral_config_profile_editor_moc ${PROJECT_SOURCE_DIR}/apps/VC3D/elements/SpiralConfigProfileEditor.hpp) +vc_add_test(NAME test_spiral_config_profile_editor + SOURCES test_spiral_config_profile_editor.cpp + ${PROJECT_SOURCE_DIR}/apps/VC3D/elements/SpiralConfigProfileEditor.cpp + ${PROJECT_SOURCE_DIR}/apps/VC3D/elements/CollapsibleSettingsGroup.cpp + ${spiral_config_profile_editor_moc} + ${collapsible_settings_group_moc} + LIBS Qt6::Core Qt6::Gui Qt6::Widgets Qt6::Test vc_core + INCLUDES ${PROJECT_SOURCE_DIR}/apps/VC3D) + qt6_wrap_cpp(line_annotation_vccollection_moc ${PROJECT_SOURCE_DIR}/core/include/vc/ui/VCCollection.hpp) qt6_wrap_cpp(line_annotation_cstate_moc ${PROJECT_SOURCE_DIR}/apps/VC3D/CState.hpp) vc_add_test(NAME test_line_annotation_generated_views diff --git a/volume-cartographer/core/test/test_spiral_config_profile_editor.cpp b/volume-cartographer/core/test/test_spiral_config_profile_editor.cpp new file mode 100644 index 0000000000..2ae57cfbc5 --- /dev/null +++ b/volume-cartographer/core/test/test_spiral_config_profile_editor.cpp @@ -0,0 +1,291 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define private public +#include "elements/SpiralConfigProfileEditor.hpp" +#undef private + +#include "elements/CollapsibleSettingsGroup.hpp" + +#include +#include +#include + +namespace { + +void require(bool condition, const char* message) +{ + if (!condition) { + std::cerr << message << std::endl; + std::exit(1); + } +} + +QJsonObject field(const QString& label, + const QString& type, + const QString& impact) +{ + return { + {QStringLiteral("label"), label}, + {QStringLiteral("type"), type}, + {QStringLiteral("runtime_impact"), impact}, + }; +} + +QJsonObject syntheticCatalog() +{ + QJsonObject fields; + QJsonObject defaults; + + QJsonObject optimizer = field( + QStringLiteral("Learning Rate"), QStringLiteral("number"), + QStringLiteral("run_boundary")); + optimizer[QStringLiteral("minimum")] = 0.0; + optimizer[QStringLiteral("maximum")] = 10.0; + optimizer[QStringLiteral("precision")] = 3; + optimizer[QStringLiteral("step")] = 0.1; + fields[QStringLiteral("optimizer_learning_rate")] = optimizer; + defaults[QStringLiteral("optimizer_learning_rate")] = 0.5; + + QJsonObject momentum = field( + QStringLiteral("Momentum"), QStringLiteral("number"), + QStringLiteral("run_boundary")); + momentum[QStringLiteral("minimum")] = 0.0; + momentum[QStringLiteral("maximum")] = 1.0; + momentum[QStringLiteral("precision")] = 2; + momentum[QStringLiteral("step")] = 0.05; + fields[QStringLiteral("optimizer_momentum")] = momentum; + defaults[QStringLiteral("optimizer_momentum")] = 0.9; + + QJsonObject model = field( + QStringLiteral("Iterations"), QStringLiteral("integer"), + QStringLiteral("new_fit")); + model[QStringLiteral("minimum")] = 1; + model[QStringLiteral("maximum")] = 1000; + fields[QStringLiteral("model_iterations")] = model; + defaults[QStringLiteral("model_iterations")] = 25; + + fields[QStringLiteral("patch_enabled")] = field( + QStringLiteral("Use Patches"), QStringLiteral("boolean"), + QStringLiteral("run_boundary")); + defaults[QStringLiteral("patch_enabled")] = true; + + QJsonObject sample = field( + QStringLiteral("Sampling Mode"), QStringLiteral("enum"), + QStringLiteral("prepared_input_rebuild")); + sample[QStringLiteral("values")] = + QJsonArray{QStringLiteral("fast"), QStringLiteral("accurate")}; + fields[QStringLiteral("sample_count_mode")] = sample; + defaults[QStringLiteral("sample_count_mode")] = QStringLiteral("fast"); + + QJsonObject input = field( + QStringLiteral("Input Channels"), QStringLiteral("array"), + QStringLiteral("shell_reload")); + input[QStringLiteral("nullable")] = true; + fields[QStringLiteral("input_channels")] = input; + defaults[QStringLiteral("input_channels")] = QJsonValue::Null; + + const QList> remaining{ + {QStringLiteral("pcl_radius"), QStringLiteral("PCL Radius")}, + {QStringLiteral("tracks_count"), QStringLiteral("Track Count")}, + {QStringLiteral("dense_weight"), QStringLiteral("Dense Weight")}, + {QStringLiteral("loss_scale"), QStringLiteral("Loss Scale")}, + {QStringLiteral("dt_limit"), QStringLiteral("DT Limit")}, + {QStringLiteral("output_stride"), QStringLiteral("Output Stride")}, + {QStringLiteral("shell_width"), QStringLiteral("Shell Width")}, + {QStringLiteral("influence_radius"), QStringLiteral("Influence Radius")}, + }; + for (const auto& item : remaining) { + QJsonObject spec = field( + item.second, QStringLiteral("integer"), + QStringLiteral("run_boundary")); + spec[QStringLiteral("minimum")] = 0; + spec[QStringLiteral("maximum")] = 100; + fields[item.first] = spec; + defaults[item.first] = 1; + } + + return { + {QStringLiteral("defaults"), defaults}, + {QStringLiteral("schema"), QJsonObject{ + {QStringLiteral("fields"), fields}, + }}, + }; +} + +void processLayout() +{ + QApplication::sendPostedEvents(); + QApplication::processEvents(); + QTest::qWait(20); +} + +} // namespace + +int main(int argc, char** argv) +{ + if (qEnvironmentVariableIsEmpty("QT_QPA_PLATFORM")) + qputenv("QT_QPA_PLATFORM", "offscreen"); + + QTemporaryDir configDir; + require(configDir.isValid(), "Failed to create temporary config directory"); + qputenv("VC3D_CONFIG_DIR", configDir.path().toUtf8()); + + std::unique_ptr app; + if (!QApplication::instance()) + app = std::make_unique(argc, argv); + + SpiralConfigProfileEditor editor; + editor.setCatalog(syntheticCatalog()); + + const QStringList expectedPrefixes{ + QStringLiteral("sample"), QStringLiteral("loss"), + QStringLiteral("patch"), QStringLiteral("tracks"), + QStringLiteral("optimizer"), QStringLiteral("model"), + QStringLiteral("input"), QStringLiteral("pcl"), + QStringLiteral("dense"), QStringLiteral("dt"), + QStringLiteral("output"), QStringLiteral("shell"), + QStringLiteral("influence"), + }; + const QStringList expectedTitles{ + QStringLiteral("Sample Count"), QStringLiteral("Loss"), + QStringLiteral("Patch"), QStringLiteral("Tracks"), + QStringLiteral("Optimizer"), QStringLiteral("Model"), + QStringLiteral("Input"), QStringLiteral("PCL"), + QStringLiteral("Dense"), QStringLiteral("DT"), + QStringLiteral("Output"), QStringLiteral("Shell"), + QStringLiteral("Influence"), + }; + require(editor._controlGroups.size() == expectedPrefixes.size(), + "Catalog fields were not split into the expected groups"); + for (int index = 0; index < expectedPrefixes.size(); ++index) { + const auto& group = editor._controlGroups[index]; + require(group.prefix == expectedPrefixes[index], + "Groups are not in stable prefix order"); + require(group.widget->isExpanded(), + "Configuration groups should initially be expanded"); + auto* button = group.widget->findChild(); + require(button && button->text() == expectedTitles[index], + "A configuration group has the wrong human-friendly title"); + } + + require(qobject_cast( + editor._fieldEditors.value(QStringLiteral("optimizer_learning_rate"))), + "Number fields should use QDoubleSpinBox"); + require(qobject_cast( + editor._fieldEditors.value(QStringLiteral("model_iterations"))), + "Integer fields should use QSpinBox"); + require(qobject_cast( + editor._fieldEditors.value(QStringLiteral("patch_enabled"))), + "Boolean fields should use QCheckBox"); + require(qobject_cast( + editor._fieldEditors.value(QStringLiteral("sample_count_mode"))), + "Enum fields should use QComboBox"); + require(qobject_cast( + editor._fieldEditors.value(QStringLiteral("input_channels"))), + "Array fields should use QLineEdit"); + auto* impact = editor.findChild( + QStringLiteral("spiralConfigImpact_optimizer_learning_rate")); + require(impact && impact->text() == QStringLiteral("Next Run"), + "Runtime impact should use its compact display label"); + require(impact->foregroundRole() == QPalette::WindowText + && impact->styleSheet().isEmpty(), + "Runtime impact should inherit theme-aware window text"); + + const QString beforeCollapse = editor.currentText(); + auto* optimizerGroup = editor._controlGroups[4].widget; + optimizerGroup->setExpanded(false); + require(!optimizerGroup->isExpanded(), + "Configuration groups should be collapsible"); + require(editor.currentText() == beforeCollapse, + "Collapsing a group must not change configuration values"); + + editor._search->setText(QStringLiteral("Learning Rate")); + processLayout(); + require(optimizerGroup->isExpanded(), + "Search should temporarily expand a matching group"); + require(!optimizerGroup->isHidden(), + "Search should retain groups containing matching rows"); + require(editor.findChild( + QStringLiteral("spiralConfigRow_optimizer_momentum"))->isHidden(), + "Search should hide non-matching rows within a matching group"); + require(editor._controlGroups[1].widget->isHidden(), + "Search should hide groups with no matching rows"); + editor._search->clear(); + processLayout(); + require(!optimizerGroup->isExpanded(), + "Clearing search should restore the prior collapse state"); + require(!editor._controlGroups[1].widget->isHidden(), + "Clearing search should restore hidden groups"); + + editor.showWindow(); + QDialog* dialog = editor.findChild( + QStringLiteral("spiralAdvancedConfigDialog")); + QWidget* grid = editor.findChild( + QStringLiteral("spiralConfigGroupGrid")); + require(dialog && grid, "Configuration dialog grid was not created"); + const QList> widths{{520, 1}, {900, 2}, {1320, 3}}; + for (const auto& width : widths) { + dialog->resize(width.first, 760); + processLayout(); + require(grid->property("spiralColumnCount").toInt() == width.second, + "Responsive group grid selected the wrong column count"); + require(editor._controlsScroll->horizontalScrollBarPolicy() + == Qt::ScrollBarAlwaysOff, + "Controls area should never show a horizontal scrollbar"); + } + require(editor._controlGroups[0].widget->y() + == editor._controlGroups[1].widget->y() + && editor._controlGroups[1].widget->y() + == editor._controlGroups[2].widget->y(), + "Groups in the same grid row should be top aligned"); + const int sampleTop = + editor._controlGroups[3].widget->mapTo(grid, QPoint()).y(); + bool packedBelowColumnHead = false; + for (int index = 0; index < 3; ++index) { + const int headBottom = + editor._controlGroups[index].widget->mapTo(grid, QPoint()).y() + + editor._controlGroups[index].widget->height(); + packedBelowColumnHead |= + sampleTop >= headBottom && sampleTop - headBottom <= 16; + } + require(packedBelowColumnHead, + "Groups should pack into columns without fixed-row gaps"); + + auto* iterations = qobject_cast( + editor._fieldEditors.value(QStringLiteral("model_iterations"))); + QSignalSpy textChanged(&editor, &SpiralConfigProfileEditor::textChanged); + iterations->setValue(31); + processLayout(); + const QJsonObject edited = QJsonDocument::fromJson( + editor.currentText().toUtf8()).object(); + require(edited.value(QStringLiteral("model_iterations")).toInt() == 31, + "Editing a grouped control did not update Expert JSON"); + require(editor._dirty && textChanged.count() > 0, + "Editing a grouped control did not update profile dirty state"); + + QJsonObject external = edited; + external[QStringLiteral("model_iterations")] = 47; + editor.setCurrentText(QString::fromUtf8( + QJsonDocument(external).toJson(QJsonDocument::Compact))); + processLayout(); + require(iterations->value() == 47, + "Expert JSON changes did not update grouped controls"); + + return 0; +} diff --git a/volume-cartographer/python/CMakeLists.txt b/volume-cartographer/python/CMakeLists.txt index 7774171673..44210f78bf 100644 --- a/volume-cartographer/python/CMakeLists.txt +++ b/volume-cartographer/python/CMakeLists.txt @@ -92,6 +92,9 @@ set_target_properties(vc_track_crossings PROPERTIES ) target_link_libraries(vc_track_crossings PRIVATE OpenMP::OpenMP_CXX) target_compile_features(vc_track_crossings PRIVATE cxx_std_23) +target_compile_options(vc_track_crossings PRIVATE + $<$,$>:-O3> +) nanobind_add_module(vc_track_store NOMINSIZE diff --git a/volume-cartographer/python/vc/track_crossings.cpp b/volume-cartographer/python/vc/track_crossings.cpp index bc4a87a277..a681239984 100644 --- a/volume-cartographer/python/vc/track_crossings.cpp +++ b/volume-cartographer/python/vc/track_crossings.cpp @@ -101,21 +101,18 @@ struct WalkIndex { std::vector partners; std::vector self_local; std::vector partner_local; + std::vector positions; std::vector reciprocal; - std::vector component; - std::vector eligible; std::vector track_lengths; - bool loop_consistency = false; - uint64_t cyclic_components = 0; size_t track_count() const { return track_lengths.size(); } size_t crossing_count() const { return partners.size(); } size_t memory_bytes() const { return offsets.capacity() * sizeof(int64_t) + (partners.capacity() + self_local.capacity() - + partner_local.capacity() + component.capacity() - + reciprocal.capacity() + track_lengths.capacity()) * sizeof(int32_t) - + eligible.capacity() * sizeof(uint8_t); + + partner_local.capacity() + reciprocal.capacity() + + track_lengths.capacity()) * sizeof(int32_t) + + positions.capacity() * sizeof(double); } }; @@ -1543,8 +1540,8 @@ CrossingKey canonical_crossing( WalkIndex* prepare_walk_index( Int64Vector offsets, Int32Vector partners, Int32Vector self_local, - Int32Vector partner_local, Int32Vector track_lengths, - bool require_loop_consistency) + Int32Vector partner_local, Float64Vector positions, + Int32Vector track_lengths) { const size_t tracks = track_lengths.shape(0); if (offsets.shape(0) != tracks + 1 || offsets(0) != 0) @@ -1552,7 +1549,8 @@ WalkIndex* prepare_walk_index( const int64_t records64 = offsets(tracks); if (records64 < 0 || partners.shape(0) != static_cast(records64) || self_local.shape(0) != static_cast(records64) - || partner_local.shape(0) != static_cast(records64)) + || partner_local.shape(0) != static_cast(records64) + || positions.shape(0) != static_cast(records64)) throw std::runtime_error("walk crossing arrays are not parallel"); const size_t records = static_cast(records64); @@ -1562,12 +1560,11 @@ WalkIndex* prepare_walk_index( index->self_local.assign(self_local.data(), self_local.data() + records); index->partner_local.assign( partner_local.data(), partner_local.data() + records); + index->positions.assign( + positions.data(), positions.data() + records); index->track_lengths.assign( track_lengths.data(), track_lengths.data() + tracks); index->reciprocal.assign(records, -1); - index->component.assign(records, -1); - index->eligible.assign(records, require_loop_consistency ? 0 : 1); - index->loop_consistency = require_loop_consistency; std::unordered_map vertex_by_key; vertex_by_key.reserve(records / 2 + 1); @@ -1599,137 +1596,30 @@ WalkIndex* prepare_walk_index( index->reciprocal[members[1]] = members[0]; } - if (!require_loop_consistency) - return index.release(); - - struct RailHalf { int32_t to, edge; }; - std::vector> graph(vertex_records.size()); - int32_t edge_count = 0; - std::vector ordered; - for (size_t track = 0; track < tracks; ++track) { - ordered.clear(); - for (int64_t record = index->offsets[track]; - record < index->offsets[track + 1]; ++record) - ordered.push_back(static_cast(record)); - std::sort(ordered.begin(), ordered.end(), [&](int32_t a, int32_t b) { - return std::tie(index->self_local[a], index->partners[a]) - < std::tie(index->self_local[b], index->partners[b]); - }); - for (size_t slot = 1; slot < ordered.size(); ++slot) { - const int32_t a = record_vertex[ordered[slot - 1]]; - const int32_t b = record_vertex[ordered[slot]]; - if (a == b) - continue; - graph[a].push_back({b, edge_count}); - graph[b].push_back({a, edge_count}); - ++edge_count; - } - } - - // Iterative Tarjan traversal: this avoids one native stack frame per - // crossing on the very large production cache. - const int32_t vertices = static_cast(graph.size()); - std::vector discovery(vertices, -1), low(vertices), parent(vertices, -1); - std::vector parent_edge(vertices, -1), cursor(vertices, 0); - std::vector bridge(static_cast(edge_count), 0); - std::vector stack; - int32_t clock = 0; - for (int32_t root = 0; root < vertices; ++root) { - if (discovery[root] >= 0) - continue; - discovery[root] = low[root] = clock++; - stack.push_back(root); - while (!stack.empty()) { - const int32_t vertex = stack.back(); - if (cursor[vertex] < static_cast(graph[vertex].size())) { - const RailHalf half = graph[vertex][cursor[vertex]++]; - if (half.edge == parent_edge[vertex]) - continue; - if (discovery[half.to] < 0) { - parent[half.to] = vertex; - parent_edge[half.to] = half.edge; - discovery[half.to] = low[half.to] = clock++; - stack.push_back(half.to); - } else { - low[vertex] = std::min(low[vertex], discovery[half.to]); - } - continue; - } - stack.pop_back(); - if (parent[vertex] >= 0) { - const int32_t up = parent[vertex]; - low[up] = std::min(low[up], low[vertex]); - if (low[vertex] > discovery[up]) - bridge[parent_edge[vertex]] = 1; - } - } - } - - std::vector vertex_component(vertices, -1); - std::vector cyclic; - for (int32_t root = 0; root < vertices; ++root) { - if (vertex_component[root] >= 0) - continue; - const int32_t component = static_cast(cyclic.size()); - int64_t component_vertices = 0; - int64_t half_edges = 0; - vertex_component[root] = component; - stack.push_back(root); - while (!stack.empty()) { - const int32_t vertex = stack.back(); - stack.pop_back(); - ++component_vertices; - for (const RailHalf half : graph[vertex]) { - if (bridge[half.edge]) - continue; - ++half_edges; - if (vertex_component[half.to] < 0) { - vertex_component[half.to] = component; - stack.push_back(half.to); - } - } - } - cyclic.push_back( - component_vertices > 1 && half_edges / 2 >= component_vertices); - } - index->cyclic_components = std::accumulate( - cyclic.begin(), cyclic.end(), uint64_t{0}); - for (size_t record = 0; record < records; ++record) { - const int32_t component = vertex_component[record_vertex[record]]; - index->component[record] = component; - index->eligible[record] = cyclic[component]; - } return index.release(); } nb::dict walk_index_stats(const WalkIndex& index) { - uint64_t eligible_records = std::accumulate( - index.eligible.begin(), index.eligible.end(), uint64_t{0}); uint64_t eligible_tracks = 0; for (size_t track = 0; track < index.track_count(); ++track) { - bool any = false; - for (int64_t record = index.offsets[track]; - record < index.offsets[track + 1]; ++record) - any = any || index.eligible[record]; - eligible_tracks += any; + eligible_tracks += index.offsets[track + 1] > index.offsets[track]; } nb::dict result; result["tracks"] = index.track_count(); result["directed_crossings"] = index.crossing_count(); result["eligible_tracks"] = eligible_tracks; - result["eligible_directed_crossings"] = eligible_records; - result["cyclic_components"] = index.cyclic_components; + result["eligible_directed_crossings"] = index.crossing_count(); result["memory_bytes"] = index.memory_bytes(); - result["loop_consistency"] = index.loop_consistency; + result["root_return_gate"] = true; return result; } WalkIndex* prepare_cached_walk_index( UInt64Vector cached_source_ids, Int64Vector cached_offsets, Int32Vector cached_partners, Int32Vector cached_self_local, - Int32Vector cached_partner_local, UInt64Vector selected_source_ids, - Int32Vector track_lengths, bool require_loop_consistency) + Int32Vector cached_partner_local, Float64Vector cached_positions, + UInt64Vector selected_source_ids, Int32Vector track_lengths) { const size_t cached_tracks = cached_source_ids.shape(0); const size_t selected_tracks = selected_source_ids.shape(0); @@ -1741,7 +1631,8 @@ WalkIndex* prepare_cached_walk_index( if (cached_records < 0 || cached_partners.shape(0) != static_cast(cached_records) || cached_self_local.shape(0) != static_cast(cached_records) - || cached_partner_local.shape(0) != static_cast(cached_records)) + || cached_partner_local.shape(0) != static_cast(cached_records) + || cached_positions.shape(0) != static_cast(cached_records)) throw std::runtime_error("cached walk crossing arrays are not parallel"); std::vector selected_rows(selected_tracks); @@ -1760,6 +1651,7 @@ WalkIndex* prepare_cached_walk_index( } std::vector offsets(selected_tracks + 1, 0); std::vector partners, self_local, partner_local; + std::vector positions; for (size_t local = 0; local < selected_tracks; ++local) { const int32_t row = selected_rows[local]; for (int64_t record = cached_offsets(row); @@ -1774,6 +1666,7 @@ WalkIndex* prepare_cached_walk_index( partners.push_back(partner); self_local.push_back(cached_self_local(record)); partner_local.push_back(cached_partner_local(record)); + positions.push_back(cached_positions(record)); } offsets[local + 1] = static_cast(partners.size()); } @@ -1782,9 +1675,10 @@ WalkIndex* prepare_cached_walk_index( Int32Vector self_view(self_local.data(), {self_local.size()}); Int32Vector partner_local_view( partner_local.data(), {partner_local.size()}); + Float64Vector position_view(positions.data(), {positions.size()}); return prepare_walk_index( offset_view, partner_view, self_view, partner_local_view, - track_lengths, require_loop_consistency); + position_view, track_lengths); } nb::dict walk_index_crossings(const WalkIndex& index) @@ -1798,29 +1692,152 @@ nb::dict walk_index_crossings(const WalkIndex& index) std::vector(index.self_local.begin(), index.self_local.end())); result["partner_local"] = own_1d(std::vector( index.partner_local.begin(), index.partner_local.end())); + result["positions"] = own_1d(std::vector( + index.positions.begin(), index.positions.end())); return result; } +struct RootNeighbors { + RootNeighbors(const WalkIndex& index, int32_t root) + { + const int64_t begin = index.offsets[root]; + const int64_t end = index.offsets[root + 1]; + if (begin == end) + return; + source_begin = index.partners.data() + begin; + source_end = index.partners.data() + end; + source_sorted = std::is_sorted(source_begin, source_end); + for (const int32_t* neighbor = source_begin; + neighbor != source_end; ++neighbor) + filter |= uint64_t{1} << filter_slot(*neighbor); + if (!source_sorted) { + sorted.assign(source_begin, source_end); + std::sort(sorted.begin(), sorted.end()); + } + } + + bool contains(int32_t track) const + { + if (!(filter & (uint64_t{1} << filter_slot(track)))) + return false; + if (source_sorted) + return std::binary_search(source_begin, source_end, track); + return std::binary_search(sorted.begin(), sorted.end(), track); + } + + static uint32_t filter_slot(int32_t track) + { + return (static_cast(track) * UINT32_C(0x9e3779b1)) >> 26; + } + + const int32_t* source_begin = nullptr; + const int32_t* source_end = nullptr; + bool source_sorted = true; + uint64_t filter = 0; + std::vector sorted; +}; + +bool bounded_path_to_root( + const WalkIndex& index, int32_t current, int remaining_edges, + const RootNeighbors& root_neighbors, std::vector& forbidden) +{ + // Every successful bounded search ends with an edge onto the root. + // Resolve that final, most frequently tested edge from the root's small + // sorted neighbor set instead of rescanning each frontier track. + if (root_neighbors.contains(current)) + return true; + if (remaining_edges <= 1) + return false; + if (remaining_edges == 2) { + for (int64_t record = index.offsets[current]; + record < index.offsets[current + 1]; ++record) { + const int32_t next = index.partners[record]; + if (std::find(forbidden.begin(), forbidden.end(), next) + == forbidden.end() + && root_neighbors.contains(next)) + return true; + } + return false; + } + for (int64_t record = index.offsets[current]; + record < index.offsets[current + 1]; ++record) { + const int32_t next = index.partners[record]; + if (std::find(forbidden.begin(), forbidden.end(), next) + != forbidden.end()) + continue; + forbidden.push_back(next); + const bool found = bounded_path_to_root( + index, next, remaining_edges - 1, root_neighbors, forbidden); + forbidden.pop_back(); + if (found) + return true; + } + return false; +} + +bool transition_can_return_to_root( + const WalkIndex& index, int32_t current, int32_t record, + int32_t root, int remaining_new_tracks, + double minimum_candidate_travel, + const std::vector& visited, + const RootNeighbors& root_neighbors) +{ + const int32_t candidate = index.partners[record]; + if (std::find(visited.begin(), visited.end(), candidate) != visited.end()) + return false; + const int32_t reciprocal = index.reciprocal[record]; + if (reciprocal < 0) + return false; + const double entry_position = index.positions[reciprocal]; + std::vector forbidden = visited; + forbidden.push_back(candidate); + for (int64_t exit = index.offsets[candidate]; + exit < index.offsets[candidate + 1]; ++exit) { + const int32_t exit_partner = index.partners[exit]; + if (exit_partner == current) + continue; + if (std::abs(index.positions[exit] - entry_position) + < minimum_candidate_travel) + continue; + if (exit_partner == root) + return true; + if (remaining_new_tracks <= 0 + || std::find(forbidden.begin(), forbidden.end(), exit_partner) + != forbidden.end()) + continue; + forbidden.push_back(exit_partner); + const bool found = bounded_path_to_root( + index, exit_partner, remaining_new_tracks, + root_neighbors, forbidden); + forbidden.pop_back(); + if (found) + return true; + } + return false; +} + bool draw_walk( - const WalkIndex& index, int32_t primary, int target_points, int hops, - int minimum_steps, int maximum_steps, uint64_t seed, - int32_t* output_tracks, int32_t* output_records) + const WalkIndex& index, int32_t primary, int target_points, + int minimum_hops, int maximum_hops, int minimum_steps, + int maximum_steps, double minimum_candidate_travel, uint64_t seed, + int32_t* output_tracks, int32_t* output_records, int32_t& output_hops) { if (primary < 0 || static_cast(primary) >= index.track_count()) return false; std::mt19937_64 random(seed); std::vector visited; - visited.reserve(static_cast(hops) + 1); + visited.reserve(static_cast(maximum_hops) + 1); visited.push_back(primary); + const RootNeighbors root_neighbors(index, primary); output_tracks[0] = primary; int32_t current = primary; int32_t current_local = -1; - int32_t starting_component = -1; - - for (int hop = 0; hop < hops; ++hop) { - std::vector candidates; + double current_entry_position = 0.0; + int completed_hops = 0; + for (int hop = 0; hop < maximum_hops; ++hop) { const int64_t begin = index.offsets[current]; const int64_t end = index.offsets[current + 1]; + int32_t record = -1; if (hop == 0) { const int points = std::max(1, target_points); std::uniform_int_distribution target_draw(0, points - 1); @@ -1828,116 +1845,163 @@ bool draw_walk( const int32_t target = points == 1 ? 0 : static_cast( std::llround(static_cast(target_slot) * (index.track_lengths[current] - 1) / (points - 1))); - int32_t best = -1; - int64_t best_distance = std::numeric_limits::max(); - for (int64_t record = begin; record < end; ++record) { - if (!index.eligible[record]) + std::vector candidates; + candidates.reserve(static_cast(end - begin)); + for (int64_t candidate_record = begin; + candidate_record < end; ++candidate_record) { + const int32_t next = index.partners[candidate_record]; + if (index.reciprocal[candidate_record] < 0 + || std::find(visited.begin(), visited.end(), next) + != visited.end()) continue; - const int64_t distance = std::abs( - static_cast(index.self_local[record]) - target); - if (distance < best_distance) { - best = static_cast(record); - best_distance = distance; - } + candidates.push_back(static_cast(candidate_record)); } - if (best < 0) - return false; - candidates.push_back(best); + std::sort(candidates.begin(), candidates.end(), + [&](int32_t a, int32_t b) { + const int64_t a_distance = std::abs( + static_cast(index.self_local[a]) - target); + const int64_t b_distance = std::abs( + static_cast(index.self_local[b]) - target); + return std::tie(a_distance, a) < std::tie(b_distance, b); + }); + if (!candidates.empty()) + record = candidates.front(); } else { std::vector by_direction[2]; - for (int64_t record = begin; record < end; ++record) { - if (!index.eligible[record]) + for (int64_t candidate_record = begin; + candidate_record < end; ++candidate_record) { + const int32_t next = index.partners[candidate_record]; + if (index.reciprocal[candidate_record] < 0 + || std::find(visited.begin(), visited.end(), next) + != visited.end()) continue; - if (index.loop_consistency - && index.component[record] != starting_component) - continue; - const int delta = index.self_local[record] - current_local; + const int delta = + index.self_local[candidate_record] - current_local; const int distance = std::abs(delta); - if (distance >= minimum_steps && distance <= maximum_steps) - by_direction[delta > 0].push_back( - static_cast(record)); + if (distance < minimum_steps || distance > maximum_steps) + continue; + if (std::abs( + index.positions[candidate_record] + - current_entry_position) + < minimum_candidate_travel) + continue; + by_direction[delta > 0].push_back( + static_cast(candidate_record)); } std::vector directions; if (!by_direction[0].empty()) directions.push_back(0); if (!by_direction[1].empty()) directions.push_back(1); - if (directions.empty()) - return false; - std::uniform_int_distribution direction_draw( - 0, directions.size() - 1); - const int direction = directions[direction_draw(random)]; - const int endpoint_distance = direction - ? index.track_lengths[current] - 1 - current_local - : current_local; - const int upper = std::min(maximum_steps, endpoint_distance); - if (upper < minimum_steps) - return false; - std::uniform_int_distribution distance_draw( - minimum_steps, upper); - const int desired = distance_draw(random); - int32_t best = -1; - int best_error = std::numeric_limits::max(); - for (int32_t record : by_direction[direction]) { - const int distance = std::abs( - index.self_local[record] - current_local); - const int error = std::abs(distance - desired); - if (error < best_error) { - best = record; - best_error = error; - } + std::shuffle(directions.begin(), directions.end(), random); + for (const int direction : directions) { + const int endpoint_distance = direction + ? index.track_lengths[current] - 1 - current_local + : current_local; + const int upper = std::min(maximum_steps, endpoint_distance); + if (upper < minimum_steps) + continue; + std::uniform_int_distribution distance_draw( + minimum_steps, upper); + const int desired = distance_draw(random); + auto& candidates = by_direction[direction]; + std::sort(candidates.begin(), candidates.end(), + [&](int32_t a, int32_t b) { + const int a_error = std::abs( + std::abs(index.self_local[a] - current_local) - desired); + const int b_error = std::abs( + std::abs(index.self_local[b] - current_local) - desired); + return std::tie(a_error, a) < std::tie(b_error, b); + }); + if (!candidates.empty()) + record = candidates.front(); + if (record >= 0) + break; } - if (best < 0) - return false; - candidates.push_back(best); } - const int32_t record = candidates.front(); + if (record < 0) + break; const int32_t next = index.partners[record]; - if (std::find(visited.begin(), visited.end(), next) != visited.end()) - return false; - if (hop == 0) - starting_component = index.component[record]; output_records[2 * hop] = record; output_records[2 * hop + 1] = index.reciprocal[record]; output_tracks[hop + 1] = next; visited.push_back(next); current = next; current_local = index.partner_local[record]; + current_entry_position = + index.positions[index.reciprocal[record]]; + completed_hops = hop + 1; } - return true; + + // The actual suffix of a simple proposed walk witnesses the return path + // for every internal transition. Certify only the final transition of + // each possible prefix, longest first; its closure completes all earlier + // witnesses within their larger remaining-hop budgets. + output_hops = 0; + std::vector prefix_visited; + prefix_visited.reserve(static_cast(completed_hops)); + for (int prefix_hops = completed_hops; + prefix_hops >= minimum_hops; --prefix_hops) { + prefix_visited.assign( + output_tracks, output_tracks + prefix_hops); + const int32_t final_record = + output_records[2 * (prefix_hops - 1)]; + if (transition_can_return_to_root( + index, output_tracks[prefix_hops - 1], final_record, + primary, maximum_hops - prefix_hops, + minimum_candidate_travel, prefix_visited, root_neighbors)) { + output_hops = prefix_hops; + break; + } + } + return output_hops != 0; } nb::dict sample_walks( const WalkIndex& index, Int32Vector primary_candidates, - UInt64Vector seeds, int groups, int target_points, int hops, - int minimum_steps, int maximum_steps) + UInt64Vector seeds, int groups, int target_points, + int minimum_hops, int maximum_hops, + int minimum_steps, int maximum_steps, double minimum_candidate_travel) { - if (groups < 0 || target_points < 1 || hops < 1 - || minimum_steps < 1 || maximum_steps < minimum_steps) + if (groups < 0 || target_points < 1 || minimum_hops < 1 + || maximum_hops < minimum_hops + || minimum_steps < 1 || maximum_steps < minimum_steps + || !std::isfinite(minimum_candidate_travel) + || minimum_candidate_travel < 0.0) throw std::runtime_error("invalid track-walk sampling parameters"); if (seeds.shape(0) != primary_candidates.shape(0)) throw std::runtime_error("walk candidates and seeds must be parallel"); - std::vector tracks(static_cast(groups) * (hops + 1), -1); - std::vector records(static_cast(groups) * hops * 2, -1); + std::vector tracks( + static_cast(groups) * (maximum_hops + 1), -1); + std::vector records( + static_cast(groups) * maximum_hops * 2, -1); + std::vector walk_hops(static_cast(groups), 0); int produced = 0; uint64_t rejected = 0; for (size_t attempt = 0; attempt < primary_candidates.shape(0) && produced < groups; ++attempt) { if (draw_walk( - index, primary_candidates(attempt), target_points, hops, - minimum_steps, maximum_steps, seeds(attempt), - tracks.data() + static_cast(produced) * (hops + 1), - records.data() + static_cast(produced) * hops * 2)) + index, primary_candidates(attempt), target_points, + minimum_hops, maximum_hops, + minimum_steps, maximum_steps, minimum_candidate_travel, + seeds(attempt), + tracks.data() + + static_cast(produced) * (maximum_hops + 1), + records.data() + + static_cast(produced) * maximum_hops * 2, + walk_hops[produced])) ++produced; else ++rejected; } - tracks.resize(static_cast(produced) * (hops + 1)); - records.resize(static_cast(produced) * hops * 2); + tracks.resize(static_cast(produced) * (maximum_hops + 1)); + records.resize(static_cast(produced) * maximum_hops * 2); + walk_hops.resize(static_cast(produced)); nb::dict result; result["tracks"] = own_2d( - std::move(tracks), static_cast(produced), hops + 1); + std::move(tracks), static_cast(produced), maximum_hops + 1); result["records"] = own_2d( - std::move(records), static_cast(produced), hops * 2); + std::move(records), static_cast(produced), maximum_hops * 2); + result["walk_hops"] = own_1d(std::move(walk_hops)); result["produced"] = produced; result["rejected_candidates"] = rejected; result["attempted_candidates"] = primary_candidates.shape(0); @@ -1946,11 +2010,16 @@ nb::dict sample_walks( nb::dict sample_walks_adaptive( const WalkIndex& index, Float64Vector primary_probabilities, - uint64_t seed, int groups, int target_points, int hops, - int minimum_steps, int maximum_steps, int maximum_attempts) + uint64_t seed, int groups, int target_points, + int minimum_hops, int maximum_hops, + int minimum_steps, int maximum_steps, double minimum_candidate_travel, + int maximum_attempts) { - if (groups < 0 || target_points < 1 || hops < 1 + if (groups < 0 || target_points < 1 || minimum_hops < 1 + || maximum_hops < minimum_hops || minimum_steps < 1 || maximum_steps < minimum_steps + || !std::isfinite(minimum_candidate_travel) + || minimum_candidate_travel < 0.0 || maximum_attempts < groups) throw std::runtime_error("invalid adaptive track-walk sampling parameters"); if (index.track_count() == 0 && groups > 0) @@ -1978,9 +2047,10 @@ nb::dict sample_walks_adaptive( } std::vector tracks( - static_cast(groups) * (hops + 1), -1); + static_cast(groups) * (maximum_hops + 1), -1); std::vector records( - static_cast(groups) * hops * 2, -1); + static_cast(groups) * maximum_hops * 2, -1); + std::vector walk_hops(static_cast(groups), 0); int produced = 0; int attempted = 0; { @@ -1992,28 +2062,95 @@ nb::dict sample_walks_adaptive( static_cast(index.track_count()) - int32_t{1})); std::discrete_distribution weighted_primary( weights.begin(), weights.end()); + const int batch_capacity = std::min( + maximum_attempts, std::max(1024, groups * 2)); + std::vector batch_primaries( + static_cast(batch_capacity)); + std::vector batch_seeds( + static_cast(batch_capacity)); + std::vector batch_success( + static_cast(batch_capacity), 0); + std::vector batch_tracks( + static_cast(batch_capacity) * (maximum_hops + 1), -1); + std::vector batch_records( + static_cast(batch_capacity) * maximum_hops * 2, -1); + std::vector batch_hops( + static_cast(batch_capacity), 0); + int next_batch_size = batch_capacity; while (attempted < maximum_attempts && produced < groups) { - const int32_t primary = weights.empty() - ? uniform_primary(random) : weighted_primary(random); - const uint64_t walk_seed = random(); - ++attempted; - if (draw_walk( - index, primary, target_points, hops, - minimum_steps, maximum_steps, walk_seed, + const int batch_size = std::min( + next_batch_size, maximum_attempts - attempted); + for (int candidate = 0; candidate < batch_size; ++candidate) { + batch_primaries[candidate] = weights.empty() + ? uniform_primary(random) : weighted_primary(random); + batch_seeds[candidate] = random(); + batch_success[candidate] = 0; + batch_hops[candidate] = 0; + } +#pragma omp parallel for schedule(dynamic, 16) + for (int candidate = 0; candidate < batch_size; ++candidate) { + batch_success[candidate] = draw_walk( + index, batch_primaries[candidate], target_points, + minimum_hops, maximum_hops, + minimum_steps, maximum_steps, minimum_candidate_travel, + batch_seeds[candidate], + batch_tracks.data() + + static_cast(candidate) + * (maximum_hops + 1), + batch_records.data() + + static_cast(candidate) + * maximum_hops * 2, + batch_hops[candidate]); + } + int considered = 0; + for (; considered < batch_size && produced < groups; ++considered) { + if (!batch_success[considered]) + continue; + std::copy_n( + batch_tracks.data() + + static_cast(considered) + * (maximum_hops + 1), + maximum_hops + 1, tracks.data() - + static_cast(produced) * (hops + 1), + + static_cast(produced) + * (maximum_hops + 1)); + std::copy_n( + batch_records.data() + + static_cast(considered) + * maximum_hops * 2, + maximum_hops * 2, records.data() - + static_cast(produced) * hops * 2)) + + static_cast(produced) + * maximum_hops * 2); + walk_hops[produced] = batch_hops[considered]; ++produced; + } + attempted += considered; + if (produced < groups && produced > 0) { + const int remaining = groups - produced; + const double attempts_per_success = + static_cast(attempted) / produced; + const double predicted = std::ceil( + remaining * attempts_per_success); + next_batch_size = std::clamp( + static_cast(std::min( + predicted, static_cast(batch_capacity))), + std::min(1024, maximum_attempts - attempted), + batch_capacity); + } } } - tracks.resize(static_cast(produced) * (hops + 1)); - records.resize(static_cast(produced) * hops * 2); + tracks.resize( + static_cast(produced) * (maximum_hops + 1)); + records.resize( + static_cast(produced) * maximum_hops * 2); + walk_hops.resize(static_cast(produced)); nb::dict result; result["tracks"] = own_2d( - std::move(tracks), static_cast(produced), hops + 1); + std::move(tracks), static_cast(produced), maximum_hops + 1); result["records"] = own_2d( - std::move(records), static_cast(produced), hops * 2); + std::move(records), static_cast(produced), maximum_hops * 2); + result["walk_hops"] = own_1d(std::move(walk_hops)); result["produced"] = produced; result["rejected_candidates"] = attempted - produced; result["attempted_candidates"] = attempted; @@ -2089,28 +2226,31 @@ NB_MODULE(track_crossings, module) "prepare_walk_index", &prepare_walk_index, nb::rv_policy::take_ownership, nb::arg("offsets"), nb::arg("partners"), nb::arg("self_local"), - nb::arg("partner_local"), nb::arg("track_lengths"), - nb::arg("require_loop_consistency") = false); + nb::arg("partner_local"), nb::arg("positions"), + nb::arg("track_lengths")); module.def("walk_index_stats", &walk_index_stats, nb::arg("index")); module.def( "prepare_cached_walk_index", &prepare_cached_walk_index, nb::rv_policy::take_ownership, nb::arg("cached_source_ids"), nb::arg("cached_offsets"), nb::arg("cached_partners"), nb::arg("cached_self_local"), - nb::arg("cached_partner_local"), nb::arg("selected_source_ids"), - nb::arg("track_lengths"), - nb::arg("require_loop_consistency") = false); + nb::arg("cached_partner_local"), nb::arg("cached_positions"), + nb::arg("selected_source_ids"), nb::arg("track_lengths")); module.def( "walk_index_crossings", &walk_index_crossings, nb::arg("index")); module.def( "sample_walks", &sample_walks, nb::arg("index"), nb::arg("primary_candidates"), nb::arg("seeds"), - nb::arg("groups"), nb::arg("target_points"), nb::arg("hops"), - nb::arg("minimum_steps"), nb::arg("maximum_steps")); + nb::arg("groups"), nb::arg("target_points"), + nb::arg("minimum_hops"), nb::arg("maximum_hops"), + nb::arg("minimum_steps"), nb::arg("maximum_steps"), + nb::arg("minimum_candidate_travel")); module.def( "sample_walks_adaptive", &sample_walks_adaptive, nb::arg("index"), nb::arg("primary_probabilities"), nb::arg("seed"), - nb::arg("groups"), nb::arg("target_points"), nb::arg("hops"), + nb::arg("groups"), nb::arg("target_points"), + nb::arg("minimum_hops"), nb::arg("maximum_hops"), nb::arg("minimum_steps"), nb::arg("maximum_steps"), + nb::arg("minimum_candidate_travel"), nb::arg("maximum_attempts")); } diff --git a/volume-cartographer/scripts/spiral/README.md b/volume-cartographer/scripts/spiral/README.md index e86c0c6d81..fdea22ca90 100644 --- a/volume-cartographer/scripts/spiral/README.md +++ b/volume-cartographer/scripts/spiral/README.md @@ -4,6 +4,27 @@ Code and helpers to fit a canonical Archimedean spiral to deformed scrolls. `spiral_service.py` hosts one persistent interactive fit session over HTTP for the VC3D Spiral workspace; `fit_spiral.py` is the underlying fitter. +## Flattening a fitted checkpoint + +`flatten_spiral_checkpoint.py` is a standalone, one-shot exporter. It +reconstructs the combined surface from a fitted checkpoint, launches a private +Lasagna service, flattens with `flatten_fast_nofilter.json`, writes the final +TIFXYZ directory, and tears the service down even if the job fails or is +interrupted: + +```sh +python flatten_spiral_checkpoint.py \ + /path/to/checkpoint_fitted.ckpt \ + /path/to/output.tifxyz +``` + +The checkpoint format does not embed the fixed umbilicus curve. The script +looks for `umbilicus.json` in the checkpoint's ancestors, in +`$SPIRAL_DATASET`, and in the standard local s1 dataset location. For other +layouts, pass `--umbilicus /path/to/umbilicus.json`. Use `--lasagna-dir` if +the Lasagna repository is not in its standard sibling or `~/villa` location. +An existing output path is never overwritten. + ## Spiral service host setup VC3D connects to a Spiral service in one of three modes, all speaking the same @@ -40,20 +61,25 @@ uv sync # creates .venv from pyproject.toml or with conda/pip, install `torch` for your CUDA version and then `pip install -e scripts/spiral`. -### Sparse CUDA field cache +### Resident sparse field pools + +Normals, gradient magnitude, and surf-SDT samples are served by fully +resident device brick pools. Each store's occupied bricks are packed once +into a flat sidecar next to the source zarr by `pack_resident_pools.py`: -Normals, gradient magnitude, and surf-SDT samples use a bounded 32³-chunk -CUDA LRU backed directly by the source Zarr arrays. The default pool limits -are 6 GiB for normals, 2 GiB for gradient magnitude, and 16 GiB for SDT; -each is also capped to a fraction of currently free device memory. Override -them with `FIT_SPIRAL_SPARSE_NORMAL_CACHE_GB`, -`FIT_SPIRAL_SPARSE_GRAD_CACHE_GB`, and -`FIT_SPIRAL_SPARSE_SDT_CACHE_GB`. `FIT_SPIRAL_TENSORSTORE_CACHE_GB` -controls the host TensorStore cache (2 GiB per field by default). +```sh +python pack_resident_pools.py /path/to/lasagna_inputs \ + --ct /path/to/_ds2.zarr --ct-group 2 --verify 2000 +``` -A single gather's distinct chunks must fit its field pool. If it does not, -the fitter reports the required working-set size and the corresponding -environment variable instead of growing the cache or risking a CUDA OOM. +`--ct` zeroes every voxel whose CT voxel reads 0 (the mask region around the +scroll) so those bricks drop out of the pool and sample as no-data. The +fitter loads the sidecars restricted to the configured z-ROI in one +sequential read per channel (for the full s1 ROI: ~33 GiB SDT + ~10 GiB +normals); after that every gather is pure device indexing with no I/O and no +eviction. A missing sidecar is a startup error naming the pack command. +Set `FIT_SPIRAL_RESIDENT_BOUNDS_CHECK=1` to enable per-gather bounds +assertions when debugging new sampling code. ### Internet flow (SSH attach) @@ -112,6 +138,12 @@ accept it — VC3D deliberately never auto-trusts host keys. The fit survives viewer disconnects, laptop sleep, and network drops; disconnecting or closing VC3D never terminates a service it did not launch. +The workspace reports the active loading, optimization, checkpoint, and +preview stage with elapsed time. Stages with a real work total also show a +counter and ETA; opaque native or CUDA operations deliberately use an +indeterminate bar instead of a guessed overall percentage. The same stage +updates are printed by standalone `fit_spiral.py`, with periodic elapsed-time +heartbeats when output is captured to a log. While connected, the circular-arrow button beside the connection controls restarts the remote service and reconnects automatically. The service replaces its own process in place, so a containing `tmux` session remains alive and an diff --git a/volume-cartographer/scripts/spiral/benchmark_rustworkx_track_graph.py b/volume-cartographer/scripts/spiral/benchmark_rustworkx_track_graph.py new file mode 100644 index 0000000000..f43189ac29 --- /dev/null +++ b/volume-cartographer/scripts/spiral/benchmark_rustworkx_track_graph.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Benchmark a rustworkx graph with tracks as nodes and crossings as edges.""" + +import argparse +import gc +import json +import os +from pathlib import Path +import time + +import numpy as np +import rustworkx as rx + + +def memory_bytes(): + """Return current RSS and peak RSS on Linux.""" + values = {} + with open("/proc/self/status", encoding="ascii") as stream: + for line in stream: + key, value = line.split(":", 1) + if key in {"VmRSS", "VmHWM"}: + values[key] = int(value.split()[0]) * 1024 + return values["VmRSS"], values["VmHWM"] + + +def gib(value): + return value / (1 << 30) + + +def report(stage, started, baseline): + rss, peak = memory_bytes() + print( + json.dumps( + { + "stage": stage, + "elapsed_seconds": round(time.perf_counter() - started, 3), + "rss_gib": round(gib(rss), 3), + "rss_delta_gib": round(gib(rss - baseline), 3), + "peak_rss_gib": round(gib(peak), 3), + }, + sort_keys=True, + ), + flush=True, + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("crossing_cache", type=Path) + parser.add_argument( + "--max-tracks", + type=int, + help="Build the induced graph on the first N track rows (default: all).", + ) + parser.add_argument( + "--track-chunk-size", + type=int, + default=250_000, + help="Number of CSR track rows converted per rustworkx call.", + ) + parser.add_argument( + "--node-chunk-size", + type=int, + default=1_000_000, + help="Number of None-valued nodes added per rustworkx call.", + ) + parser.add_argument( + "--benchmark-operations", + action="store_true", + help="Time degree reads and removal of every 100th track.", + ) + args = parser.parse_args() + + started = time.perf_counter() + initial_rss, _ = memory_bytes() + report("start", started, initial_rss) + + stored = np.load(args.crossing_cache, allow_pickle=False) + offsets = np.asarray(stored["offsets"], dtype=np.int64) + partners = np.asarray(stored["partners"], dtype=np.int32) + total_tracks = len(offsets) - 1 + track_count = ( + total_tracks + if args.max_tracks is None + else min(total_tracks, args.max_tracks) + ) + report("crossing_arrays_loaded", started, initial_rss) + + graph_baseline, _ = memory_bytes() + graph = rx.PyGraph(multigraph=True) + for start in range(0, track_count, args.node_chunk_size): + count = min(args.node_chunk_size, track_count - start) + # The returned NodeIndices is released after each bounded chunk. + graph.add_nodes_from((None for _ in range(count))) + gc.collect() + report("nodes_built", started, graph_baseline) + + edge_count = 0 + records_scanned = 0 + edge_started = time.perf_counter() + for row_start in range(0, track_count, args.track_chunk_size): + row_end = min(track_count, row_start + args.track_chunk_size) + record_start = int(offsets[row_start]) + record_end = int(offsets[row_end]) + chunk_partners = partners[record_start:record_end] + counts = np.diff(offsets[row_start : row_end + 1]) + sources = np.repeat( + np.arange(row_start, row_end, dtype=np.int32), counts + ) + # The cache is a symmetric directed CSR with one record in each + # direction. Keeping source < partner yields one undirected edge and + # also restricts a prefix benchmark to its induced subgraph. + keep = (sources < chunk_partners) & (chunk_partners < track_count) + kept_sources = sources[keep] + kept_partners = chunk_partners[keep] + added = len(kept_sources) + if added: + graph.add_edges_from_no_data(zip(kept_sources, kept_partners)) + edge_count += added + records_scanned += record_end - record_start + if ( + row_end == track_count + or row_end % max(args.track_chunk_size, 1_000_000) == 0 + ): + rss, _ = memory_bytes() + print( + json.dumps( + { + "stage": "edges_progress", + "tracks": row_end, + "records_scanned": records_scanned, + "edges": edge_count, + "edge_seconds": round( + time.perf_counter() - edge_started, 3 + ), + "rss_gib": round(gib(rss), 3), + }, + sort_keys=True, + ), + flush=True, + ) + del sources, kept_sources, kept_partners, chunk_partners, counts, keep + gc.collect() + report("graph_built", started, graph_baseline) + + if graph.num_nodes() != track_count or graph.num_edges() != edge_count: + raise RuntimeError("rustworkx graph counts do not match input counts") + graph_rss, graph_peak = memory_bytes() + build_seconds = time.perf_counter() - started + print( + json.dumps( + { + "stage": "result", + "rustworkx_version": rx.__version__, + "pid": os.getpid(), + "total_cache_tracks": total_tracks, + "graph_nodes": graph.num_nodes(), + "directed_records_scanned": records_scanned, + "graph_edges": graph.num_edges(), + "build_seconds": round(build_seconds, 3), + "edge_build_seconds": round( + time.perf_counter() - edge_started, 3 + ), + "graph_rss_delta_gib": round( + gib(graph_rss - graph_baseline), 3 + ), + "peak_rss_gib": round(gib(graph_peak), 3), + "rss_bytes_per_node_and_edge": round( + (graph_rss - graph_baseline) + / max(1, graph.num_nodes() + graph.num_edges()), + 2, + ), + }, + sort_keys=True, + ), + flush=True, + ) + + if args.benchmark_operations: + operation_started = time.perf_counter() + degree_sum = sum(graph.degree(node) for node in range(graph.num_nodes())) + print( + json.dumps( + { + "stage": "degree_scan", + "nodes": graph.num_nodes(), + "degree_sum": degree_sum, + "expected_degree_sum": 2 * graph.num_edges(), + "seconds": round(time.perf_counter() - operation_started, 3), + }, + sort_keys=True, + ), + flush=True, + ) + remove = range(0, track_count, 100) + remove_count = (track_count + 99) // 100 + old_edges = graph.num_edges() + operation_started = time.perf_counter() + graph.remove_nodes_from(remove) + print( + json.dumps( + { + "stage": "remove_one_percent_tracks", + "removed_nodes": remove_count, + "removed_incident_edges": old_edges - graph.num_edges(), + "remaining_nodes": graph.num_nodes(), + "remaining_edges": graph.num_edges(), + "seconds": round(time.perf_counter() - operation_started, 3), + "rss_gib": round(gib(memory_bytes()[0]), 3), + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/volume-cartographer/scripts/spiral/calibration/gradnorm.py b/volume-cartographer/scripts/spiral/calibration/gradnorm.py index 0af14aadc5..b449fc8e96 100644 --- a/volume-cartographer/scripts/spiral/calibration/gradnorm.py +++ b/volume-cartographer/scripts/spiral/calibration/gradnorm.py @@ -47,11 +47,11 @@ def main(): dataset = Path(args.dataset) checkpoint = load_checkpoint_cpu(args.checkpoint) - cfg = dict(fs.default_config) + cfg = fs.Config().as_dict() cfg.update({k: v for k, v in checkpoint.get('cfg', {}).items() if k in cfg}) cfg.update({ 'dense_spacing_mode': 'phase', - 'dense_spacing_num_pairs': int(args.pairs), + 'sample_count_dense_spacing_pairs': int(args.pairs), # measure attachment at unit weight; keep repo-default 12/8/2 for the rest 'loss_weight_dense_spacing': 12.0, 'loss_weight_dense_spacing_count': 8.0, @@ -114,7 +114,7 @@ def main(): model.zero_grad(set_to_none=True) for name, loss in iter_lasagna_losses( transform, dr, normals, outer, - cfg['dense_normals_num_points'], + cfg['sample_count_dense_normal_points'], compute_spacing=True): if name != target_name: continue diff --git a/volume-cartographer/scripts/spiral/config.py b/volume-cartographer/scripts/spiral/config.py new file mode 100644 index 0000000000..4d5f4a9bca --- /dev/null +++ b/volume-cartographer/scripts/spiral/config.py @@ -0,0 +1,403 @@ +"""Validated Spiral configuration with built-in defaults.""" + +from __future__ import annotations + +import json +from pathlib import Path + + +_ENUMS = { + "model_flow_integration_solver": ["rk4"], + "model_flow_field_type": ["cartesian", "cylindrical"], + "patch_strip_sampling": ["straight", "dijkstra"], + "track_crossing_mode": ["count", "track_walk"], + "track_radius_target": ["mean", "median"], + "dense_spacing_mode": ["phase", "grad_mag"], + "dense_spacing_support_policy": ["product", "minimum"], + "dt_target_mode": ["strip_median", "whole_object_quantile"], + "dense_spacing_density_lambda": [ + "inverse_gap", "soft_mass", "soft_mass_wide"], +} + +_NULL_TYPES = { + "pcl_sampling_weights": "dictionary", + "track_length_bin_weights": "vector", + "track_max_tortuosity": "number", + "loss_start_track_dt": "integer", + "loss_start_unverified_patch_dt": "number", +} + +_PREPARED_INPUT_FIELDS = { + "patch_erode_patches", + "track_crossing_precompute_max", + "track_crossing_mode", + "track_exclusion_radius", + "dense_spacing_mode", + "output_first_winding", + "output_winding_margin", + "output_step_size", + "output_num_slices_for_visualization", +} + +_DENSE_LOSS_ONLY_FIELDS = { + "dense_spacing_density_lambda", + "dense_spacing_density_soft_mass_min_gap_wv", +} + +_SCALE_WITH_Z_FIELDS = { + "sample_count_patches_per_step", + "sample_count_patches_per_step_for_dt", + "sample_count_unverified_patches_per_step", + "sample_count_unverified_patches_per_step_for_dt", + "sample_count_relative_winding_pcls", + "sample_count_absolute_winding_pcls", + "sample_count_unattached_pcls_per_step", + "sample_count_tracks_per_step", + "sample_count_dense_normal_points", + "sample_count_regularisation_points", + "sample_count_dense_spacing_pairs", + "sample_count_dense_spacing_density_extra_pairs", + "sample_count_minimum_spacing_independent_samples", + "sample_count_dense_attachment_points", + "sample_count_shell_samples", +} + + +def _runtime_impact(key): + if key.startswith("shell_"): + return "shell_reload" + if key.startswith("model_") or key == "optimizer_random_seed": + return "new_fit" + if key.startswith(("input_", "pcl_")) or key in _PREPARED_INPUT_FIELDS: + return "prepared_input_rebuild" + return "run_boundary" + + +def _dependencies(key): + if key.startswith(("patch_", "pcl_")): + return ["patch_pcl", "trusted_geometry", "tracks"] + if key.startswith("track_"): + return ["tracks"] + if key.startswith("dense_"): + return (["dense_losses"] if key in _DENSE_LOSS_ONLY_FIELDS + else ["dense_stores", "dense_losses"]) + if key.startswith("dt_"): + return ["patch_pcl", "tracks", "dense_losses"] + if key.startswith("shell_"): + return ["shell"] + if key.startswith("output_") and key != "output_save_png_visualizations": + return ["preview_output"] + return [] + + +def _field_spec(key, default): + nullable = key in _NULL_TYPES + if key in _ENUMS: + kind = "enum" + elif nullable: + kind = _NULL_TYPES[key] + elif type(default) is bool: + kind = "boolean" + elif type(default) is int: + kind = "integer" + elif type(default) is float: + kind = "number" + elif isinstance(default, list): + kind = "vector" + elif isinstance(default, dict): + kind = "dictionary" + else: + kind = "string" + + spec = { + "type": kind, + "nullable": nullable, + "label": key.split("_", 1)[-1].replace("_", " ").title(), + "runtime_impact": _runtime_impact(key), + "dependencies": _dependencies(key), + } + if kind in ("integer", "number"): + spec.update( + minimum=1 if key == "output_num_slices_for_visualization" else 0, + maximum=(1_000_000 if key == "output_num_slices_for_visualization" + else 1_000_000_000), + step=1 if kind == "integer" else .01, + ) + if kind == "number": + spec["precision"] = 6 + elif kind == "enum": + spec["values"] = _ENUMS[key] + elif kind == "vector": + spec["length"] = 3 if key == "track_length_bin_weights" else len(default) + if key in _SCALE_WITH_Z_FIELDS: + spec["scale_with_z"] = True + return spec + + +class Config: + def __init__(self, overrides=None): + self.optimizer_random_seed = 1 + self.optimizer_distributed_split_batch = True + self.optimizer_learning_rate = 3e-05 + self.optimizer_exp_lr_schedule = True + self.optimizer_lr_final_factor = 0.3 + self.optimizer_num_training_steps = 30000 + self.model_num_flow_integration_steps = 3 + self.model_flow_integration_solver = "rk4" + self.model_num_flow_timesteps = 1 + self.model_num_flow_stages = 1 + self.model_flow_bounds_z_margin = 160 + self.model_flow_bounds_radius = 3200 + self.model_flow_voxel_resolution = 16 + self.model_flow_field_type = "cartesian" + self.model_flow_field_high_res_lr_scale_initial = 0.2 + self.model_flow_field_high_res_lr_scale_final = 0.2 + self.model_flow_field_high_res_lr_ramp_start_step = 0 + self.model_flow_field_high_res_lr_ramp_steps = 1 + self.model_flow_field_direct_lr = True + self.model_gap_expander_logit_resolution = 24 + self.model_gap_expander_num_windings = 130 + self.model_gap_expander_lr_scale = 0.3 + self.model_linear_z_resolution = 48 + self.model_initial_dr_per_winding = 16.0 + self.patch_radius_loss_margin = 0.025 + self.patch_radius_loss_inv = False + self.patch_loss_z_margin = 0 + self.patch_dt_norm_p = 0.5 + self.patch_dt_within_patch_norm_p = 3.0 + self.patch_dt_loss_margin = 0.025 + self.patch_radius_within_norm_p = 3.0 + self.sample_count_patches_per_step = 360 + self.sample_count_patches_per_step_for_dt = 240 + self.sample_count_points_per_patch = 800 + self.sample_count_unverified_patches_per_step = 120 + self.sample_count_unverified_patches_per_step_for_dt = 80 + self.sample_count_unverified_points_per_patch = 800 + self.sample_count_relative_winding_pcls = 48 + self.sample_count_relative_winding_patch_pairs_per_pcl = 4 + self.sample_count_absolute_winding_pcls = 48 + self.sample_count_absolute_winding_points_per_pcl = 4 + self.sample_count_unattached_pcls_per_step = 84 + self.sample_count_unattached_pcl_points_per_step = 32 + self.sample_count_tracks_per_step = 48000 + self.sample_count_track_points_per_step = 96 + self.sample_count_dense_normal_points = 60000 + self.sample_count_regularisation_points = 4500 + self.sample_count_dense_spacing_pairs = 12000 + self.sample_count_dense_spacing_count_extra_pairs = 0 + self.sample_count_dense_spacing_density_extra_pairs = 24000 + self.sample_count_dense_spacing_density_chunk_pairs = 24000 + self.sample_count_minimum_spacing_independent_samples = 2000 + self.sample_count_dense_attachment_points = 20000 + self.sample_count_patch_dt_target_points = 256 + self.sample_count_dt_target_points_per_strip = 512 + self.sample_count_shell_samples = 24576 + self.sample_count_influence_footprint_points = 2048 + self.sample_count_influence_anchor_lattice_points = 100000 + self.sample_count_influence_anchor_geometry_points = 100000 + self.sample_count_influence_anchor_samples_per_step = 4096 + self.patch_strip_sampling = "straight" + self.patch_erode_patches = 1 + self.input_disable_patches = False + self.patch_unverified_patch_radius_loss_margin = 0.025 + self.patch_unverified_patch_radius_loss_inv = False + self.patch_unverified_patch_radius_within_norm_p = 3.0 + self.patch_unverified_patch_dt_norm_p = 0.5 + self.patch_unverified_patch_dt_within_patch_norm_p = 3.0 + self.patch_unverified_patch_dt_loss_margin = 0.025 + self.patch_unverified_patch_exclusion_radius = 64.0 + self.pcl_rel_winding_adjacent_patches_only = True + self.pcl_stratified_pcl_sampling = True + self.pcl_sampling_weights = None + self.pcl_fiber_min_point_spacing = 40.0 + self.pcl_unattached_pcl_min_point_spacing = 16.0 + self.track_min_sample_spacing = 20.0 + self.track_max_sample_spacing = 60.0 + self.track_length_bin_weights = [0.0, 0.15, 0.85] + self.track_max_tortuosity = None + self.track_crossing_precompute_max = 8 + self.track_max_track_crossing_per_step = 1 + self.track_crossing_mode = "track_walk" + self.track_min_walk_steps_per_track = 24 + self.track_max_walk_steps_per_track = 256 + self.track_min_walks_per_track = 2 + self.track_max_walks_per_track = 4 + self.track_walk_minimum_cycle_travel = 20.0 + self.track_exclusion_radius = 16.0 + self.track_radius_target = "mean" + self.track_radius_loss_margin = 0.025 + self.track_radius_within_norm_p = 6.0 + self.track_dt_within_track_norm_p = 3.0 + self.track_dt_norm_p = 0.5 + self.track_dt_loss_margin = 0.025 + self.dense_grad_mag_encode_scale = 1000.0 + self.dense_grad_mag_factor = 0.25 + self.dense_spacing_integration_steps = 8 + self.dense_spacing_mode = "phase" + self.dense_spacing_pair_m_short = [ + 3, + 7 + ] + self.dense_spacing_pair_m_long = [ + 5, + 15 + ] + self.dense_spacing_pair_long_fraction = 0.15 + self.dense_spacing_count_temperature_wv = 0.5 + self.dense_spacing_target_step_wv = 1.0 + self.dense_spacing_max_step_wv = 2.0 + self.dense_spacing_max_steps = 1400 + self.dense_spacing_step_oversample = 1.25 + self.dense_spacing_use_support_gate = True + self.dense_spacing_support_sigma = 4.0 + self.dense_spacing_support_floor_alpha = 0.05 + self.dense_spacing_support_policy = "product" + self.dense_spacing_phase_huber_delta = 0.5 + self.dense_spacing_phase_extension_windings = 1.0 + self.dense_spacing_phase_min_center_gap_wv = 4.0 + self.dense_spacing_phase_graze_dot = 0.4 + self.dense_spacing_phase_graze_depth_wv = 1.0 + self.dense_spacing_phase_window_windings = 0.75 + self.dense_spacing_phase_end_free_margin_windings = 0.5 + self.dense_spacing_phase_missing_cost = 0.55 + self.dense_spacing_phase_missing_extend_cost = 0.55 + self.dense_spacing_phase_extra_cost = 0.7 + self.dense_spacing_phase_extra_extend_cost = 0.7 + self.dense_spacing_phase_temperature = 0.1 + self.dense_spacing_phase_band_confidence_cost = 0.25 + self.dense_spacing_phase_top2_margin = 0.1 + self.dense_spacing_phase_min_matched_windings = 2 + self.dense_spacing_phase_min_matched_mass = 1.0 + self.loss_weight_min_spacing = 2.0 + self.loss_weight_dense_spacing_count = 0.0 + self.loss_weight_dense_spacing_density = 12.0 + self.loss_weight_dense_attachment = 0.0 + self.loss_weight_patch_radius = 8.0 + self.loss_weight_patch_dt = 4.0 + self.loss_weight_unverified_patch_radius = 2.0 + self.loss_weight_unverified_patch_dt = 1.0 + self.loss_weight_rel_winding = 5.0 + self.loss_weight_abs_winding = 5.0 + self.loss_weight_unattached_pcl_radius = 2.0 + self.loss_weight_unattached_pcl_dt = 4.0 + self.loss_weight_track_radius = 50.0 + self.loss_weight_track_dt = 10.0 + self.loss_weight_sym_dirichlet = 10.0 + self.loss_weight_dense_normals = 100.0 + self.loss_weight_dense_spacing = 12.0 + self.loss_weight_umbilicus = 1.25 + self.loss_weight_shell_outer = 1.0 + self.loss_weight_shell_patch_radius = 0.0 + self.loss_weight_anchor = 0.0 + self.dense_spacing_density_min_gap_wv = 0.0 + self.dense_spacing_density_max_blind_fraction = 0.75 + self.dense_min_spacing_d_min_wv = 6.0 + self.dense_attachment_scale = 8.0 + self.dense_attachment_warmup_steps = 3000 + self.dense_attachment_ramp_steps = 3000 + self.dense_normals_finite_difference_epsilon = 8.0 + self.model_sym_dirichlet_finite_difference_epsilon = 4.0 + self.optimizer_weight_decay_gap_expander = 0.01 + self.optimizer_weight_decay_flow_field = 0.0 + self.loss_start_patch_dt = 25000 + self.loss_start_track_dt = 10000 + self.loss_start_unverified_patch_dt = None + self.dt_progressive_windings = False + self.dt_progressive_inner_winding = 20 + self.dt_progressive_steps = 50000 + self.dt_progressive_exponent = 1.0 + self.dt_target_mode = "strip_median" + self.dt_target_floating_threshold = 0.25 + self.dt_target_update_interval = 100 + self.dt_target_max_stride = 128 + self.output_first_winding = 10 + self.output_winding_margin = 4 + self.output_step_size = 20 + self.shell_outer_winding_idx = 130 + self.shell_outer_winding_margin = 10 + self.shell_num_theta_bins = 720 + self.shell_huber_delta = 16.0 + self.shell_table_smooth_sigma_z = 4.0 + self.shell_table_smooth_sigma_theta = 1.0 + self.shell_min_confidence = 0.25 + self.output_save_png_visualizations = False + self.influence_enabled = False + self.influence_z = 3000.0 + self.influence_windings = 5.0 + self.influence_theta_frac = 0.5 + self.influence_disable_dt_frac = 0.75 + self.influence_sigma = 0.3333 + self.influence_anchor_ramp_power = 2.0 + self.dense_spacing_density_lambda = "inverse_gap" + self.dense_spacing_density_soft_mass_min_gap_wv = 0.0 + self.output_num_slices_for_visualization = 20 + + defaults = vars(self) + fields = {key: _field_spec(key, value) + for key, value in defaults.items()} + + if isinstance(overrides, (str, Path)): + overrides = json.loads(Path(overrides).read_text()) + overrides = overrides or {} + unknown = set(overrides) - set(defaults) + if unknown: + raise ValueError(f"Unknown Spiral config keys: {sorted(unknown)}") + values = defaults | overrides + for key, value in values.items(): + spec = fields[key] + if value is None and spec["nullable"]: + continue + valid = { + "boolean": lambda: type(value) is bool, + "integer": lambda: type(value) is int, + "number": lambda: type(value) in (int, float), + "string": lambda: isinstance(value, str), + "enum": lambda: value in spec["values"], + "vector": lambda: isinstance(value, list), + "dictionary": lambda: isinstance(value, dict), + } + if not valid[spec["type"]](): + raise ValueError(f"Invalid value for {key}") + if spec["type"] in ("integer", "number") and not ( + spec["minimum"] <= value <= spec["maximum"]): + raise ValueError(f"Out-of-range value for {key}") + if spec["type"] == "vector" and len(value) != spec["length"]: + raise ValueError(f"Invalid vector length for {key}") + if spec["type"] == "vector" and any( + type(item) not in (int, float) for item in value): + raise ValueError(f"Invalid vector value for {key}") + if spec["type"] == "dictionary" and any( + not isinstance(item_key, str) + or type(item) not in (int, float) + for item_key, item in value.items()): + raise ValueError(f"Invalid dictionary value for {key}") + for key, value in overrides.items(): + setattr(self, key, value) + + def as_dict(self): + return vars(self).copy() + + @classmethod + def catalog(cls): + defaults = cls().as_dict() + fields = { + key: _field_spec(key, value) + for key, value in defaults.items() + } + presets = { + path.stem: cls(path).as_dict() + for path in (Path(__file__).parent / "configs").glob("*.json") + } + return { + "defaults": defaults, + "schema": { + "paths": { + "outer_shell": { + "runtime_impact": "shell_reload", + "dependencies": ["shell"], + }, + }, + "fields": fields, + }, + "presets": presets, + } diff --git a/volume-cartographer/scripts/spiral/configs/track_walk.json b/volume-cartographer/scripts/spiral/configs/track_walk.json new file mode 100644 index 0000000000..05219931d7 --- /dev/null +++ b/volume-cartographer/scripts/spiral/configs/track_walk.json @@ -0,0 +1,6 @@ +{ + "track_crossing_mode": "track_walk", + "sample_count_track_points_per_step": 192, + "track_min_walks_per_track": 2, + "track_max_walks_per_track": 5 +} diff --git a/volume-cartographer/scripts/spiral/ddp_helpers.py b/volume-cartographer/scripts/spiral/ddp_helpers.py index 0198659f01..61647293a0 100644 --- a/volume-cartographer/scripts/spiral/ddp_helpers.py +++ b/volume-cartographer/scripts/spiral/ddp_helpers.py @@ -81,7 +81,7 @@ def allreduce_grads_(params): p.grad.div_(world_size) -def split_counts_across_ranks(config, count_keys, split_key='distributed_split_batch'): +def split_counts_across_ranks(config, count_keys, split_key='optimizer_distributed_split_batch'): world_size = get_world_size() if world_size == 1 or not config[split_key]: return 1 diff --git a/volume-cartographer/scripts/spiral/dt_targets.py b/volume-cartographer/scripts/spiral/dt_targets.py index a18341fc7b..f0c29300ce 100644 --- a/volume-cartographer/scripts/spiral/dt_targets.py +++ b/volume-cartographer/scripts/spiral/dt_targets.py @@ -416,9 +416,10 @@ def compute_patch_dt_target_cache(slice_to_spiral_transform, dr_per_winding, pat if total > 0: ijs_np = np.concatenate([p._dt_target_ijs for p in patches], axis=0) patch_idx_np = np.repeat(np.arange(num_patches, dtype=np.int64), counts) - ijs_gpu = torch.from_numpy(ijs_np).to(device=device) - patch_idx_gpu = torch.from_numpy(patch_idx_np).to(device=device) - zyxs = patch_atlas.lookup(patch_idx_gpu, ijs_gpu) + # The atlas is host-resident: the gather runs on CPU and only the + # interpolated points land on the device. + zyxs = patch_atlas.lookup( + torch.from_numpy(patch_idx_np), torch.from_numpy(ijs_np)) spiral_zyxs = _transform_in_chunks(slice_to_spiral_transform, zyxs, chunk_size) theta_t, _, shifted_t = get_theta_and_radii(spiral_zyxs[..., 1:], dr_per_winding) theta_np = theta_t.cpu().numpy() diff --git a/volume-cartographer/scripts/spiral/find_inconsistent_windings.py b/volume-cartographer/scripts/spiral/find_inconsistent_windings.py index e924603111..1c3ef64323 100644 --- a/volume-cartographer/scripts/spiral/find_inconsistent_windings.py +++ b/volume-cartographer/scripts/spiral/find_inconsistent_windings.py @@ -160,12 +160,12 @@ def install_globals(checkpoint, patches_dir, pcl_paths, filter_z_begin, filter_z user-supplied patches/pcls/filter-z-range so its loaders behave identically. Returns the checkpoint's model z-range, which shapes the transform's flow field independently of the filtering z-range.""" - cfg = dict(fs.default_config) + cfg = fs.Config().as_dict() cfg.update(checkpoint['cfg']) # This tool IS the patch graph — a checkpoint trained supervision-free # (disable_patches, e.g. the 2026-07-17 normals-only baseline) must not # stop the loaders from reading the patches it wants to analyse. - cfg['disable_patches'] = False + cfg['input_disable_patches'] = False fs.cfg = cfg fs.verified_patches_path = patches_dir # Attachment is over the verified patch set only, so skip the (slow, unrelated) @@ -196,13 +196,13 @@ def build_transform(checkpoint, model_z_begin, model_z_end): np.concatenate([all_zs[:, None], umbilicus_fn(all_zs)], axis=-1).astype(np.float32) ).to(device) - r = cfg['flow_bounds_radius'] - flow_min_corner = torch.tensor([model_z_begin - cfg['flow_bounds_z_margin'], -r, -r], dtype=torch.int64, device=device) - flow_max_corner = torch.tensor([model_z_end + cfg['flow_bounds_z_margin'], r, r], dtype=torch.int64, device=device) + r = cfg['model_flow_bounds_radius'] + flow_min_corner = torch.tensor([model_z_begin - cfg['model_flow_bounds_z_margin'], -r, -r], dtype=torch.int64, device=device) + flow_max_corner = torch.tensor([model_z_end + cfg['model_flow_bounds_z_margin'], r, r], dtype=torch.int64, device=device) model = fs.SpiralAndTransform( - flow_integration_steps=cfg['num_flow_integration_steps'], - flow_integration_solver=cfg['flow_integration_solver'], + flow_integration_steps=cfg['model_num_flow_integration_steps'], + flow_integration_solver=cfg['model_flow_integration_solver'], umbilicus_zyx=umbilicus_zyx, flow_min_corner_zyx=flow_min_corner, flow_max_corner_zyx=flow_max_corner, @@ -212,13 +212,6 @@ def build_transform(checkpoint, model_z_begin, model_z_end): model.to(device) model.load_state_dict(checkpoint['spiral_and_transform']) model.eval() - # The high-res flow fields are stored "pre-scale": at forward time the hr params - # are multiplied by flow_scales[1], which training ramps (identically for every - # stage when num_flow_stages > 1) and which is NOT part of the state_dict. A - # fully-fitted checkpoint ends the ramp at its 'final' value, so pin every - # stage's scale to that to reproduce the saved transform. - for flow_field in model.flow_fields: - flow_field.flow_scales[1] = cfg['flow_field_high_res_lr_scale_final'] transform = model.get_slice_to_spiral_transform() dr_per_winding = model.get_dr_per_winding() diff --git a/volume-cartographer/scripts/spiral/fit_session.py b/volume-cartographer/scripts/spiral/fit_session.py index 5fb28b6f05..e2b1e52cf1 100644 --- a/volume-cartographer/scripts/spiral/fit_session.py +++ b/volume-cartographer/scripts/spiral/fit_session.py @@ -16,115 +16,17 @@ from typing import Any, Iterable, Mapping import zipfile +from config import Config -# Version 12 expands Run configuration with native track-walk hop controls. -API_VERSION = 12 - - -# Counts which describe how many training objects/points are sampled per -# optimizer step. The service exposes the post-scaling values actually used by -# the resident fitter, and Run-scoped edits set those active values directly. -RUN_MUTABLE_SAMPLING_KEYS = frozenset({ - "num_patches_per_step", - "num_patches_per_step_for_dt", - "num_points_per_patch", - "unverified_num_patches_per_step", - "unverified_num_patches_per_step_for_dt", - "unverified_num_points_per_patch", - "rel_winding_num_pcls", - "rel_winding_num_patch_pairs_per_pcl", - "abs_winding_num_pcls", - "abs_winding_num_points_per_pcl", - "unattached_pcl_num_per_step", - "unattached_pcl_num_points_per_step", - "track_num_per_step", - "track_num_points_per_step", - "dense_normals_num_points", - "dense_spacing_num_pairs", - "dense_spacing_density_extra_pairs", - "dense_attachment_num_points", - "min_spacing_independent_samples", - "regularisation_num_points", - "shell_num_samples", -}) - - -RUN_MUTABLE_BOOLEAN_KEYS = frozenset({ - "save_png_visualizations", -}) - - -RUN_MUTABLE_TRACK_POLICY_KEYS = frozenset({ - "track_length_bin_weights", - "max_track_crossing_per_step", - "track_min_sample_spacing", - "track_max_sample_spacing", - "min_walk_steps_per_track", - "max_walk_steps_per_track", - "n_walks_per_track", -}) - - -def is_run_mutable_config_key(key: str) -> bool: - """Return whether an advanced setting may change at a Run boundary.""" - return ( - key in RUN_MUTABLE_SAMPLING_KEYS - or key in RUN_MUTABLE_BOOLEAN_KEYS - or key in RUN_MUTABLE_TRACK_POLICY_KEYS - or (key.startswith("loss_weight_") and key != "loss_weight_anchor") - or key.startswith("loss_start_") - ) +# Version 15 adds structured, stage-local operation progress. +API_VERSION = 15 def run_mutable_config(config: Mapping[str, Any]) -> dict[str, Any]: - """Select the Run-scoped editor fields from a complete fitter config.""" - return { - key: value for key, value in config.items() - if is_run_mutable_config_key(key) - } - - -def apply_optional_input_selection(config: dict[str, Any]) -> dict[str, Any]: - """Force losses and sample counts off for session-disabled inputs.""" - def zero(*keys: str) -> None: - for key in keys: - config[key] = 0 - - if not bool(config.get("use_verified_patches", True)): - zero("loss_weight_patch_radius", "loss_weight_patch_dt", - "loss_weight_umbilicus", "loss_weight_shell_patch_radius", - "num_patches_per_step", "num_patches_per_step_for_dt", - "num_points_per_patch") - if not bool(config.get("use_unverified_patches", True)): - zero("loss_weight_unverified_patch_radius", - "loss_weight_unverified_patch_dt", - "unverified_num_patches_per_step", - "unverified_num_patches_per_step_for_dt", - "unverified_num_points_per_patch") - normals = bool(config.get("use_normals", True)) - sdt = bool(config.get("use_surf_sdt", True)) - if not normals: - zero("loss_weight_dense_normals", "dense_normals_num_points") - if not bool(config.get("use_tracks", True)): - zero("loss_weight_track_radius", "loss_weight_track_dt", - "track_num_per_step", "track_num_points_per_step") - if not bool(config.get("use_fibers", True)): - zero("loss_weight_unattached_pcl_radius", "loss_weight_unattached_pcl_dt", - "unattached_pcl_num_per_step", "unattached_pcl_num_points_per_step") - - spacing_mode = str(config.get("dense_spacing_mode", "phase")) - if not sdt or not normals: - zero("loss_weight_dense_spacing_count", - "loss_weight_dense_spacing_density", - "loss_weight_dense_attachment", "dense_spacing_count_extra_pairs", - "dense_spacing_density_extra_pairs", "dense_attachment_num_points") - if spacing_mode == "phase": - zero("loss_weight_dense_spacing", "dense_spacing_num_pairs") - if (not bool(config.get("use_gradient_magnitude", True)) - and spacing_mode == "grad_mag"): - zero("loss_weight_dense_spacing", "dense_spacing_num_pairs") - return config + fields = Config.catalog()["schema"]["fields"] + return {key: value for key, value in config.items() + if fields[key]["runtime_impact"] in {"run_boundary", "shell_reload"}} class PclRole(str, Enum): @@ -446,14 +348,12 @@ def optional_dir(value: str, field_name: str, required: bool = False) -> None: errors.append({"field": field_name, "message": "Directory is not readable"}) require_file(paths.umbilicus, "umbilicus", json_file=True) - disable_patches = bool(run.config.get("disable_patches", False)) - use_verified = bool(run.config.get("use_verified_patches", True)) and not disable_patches - use_unverified = bool(run.config.get("use_unverified_patches", True)) and not disable_patches - optional_dir(paths.verified_patches, "verified_patches", required=use_verified) - if use_unverified: + disable_patches = bool(run.config.get("input_disable_patches", False)) + optional_dir(paths.verified_patches, "verified_patches", + required=not disable_patches) + if not disable_patches: optional_dir(paths.unverified_patches, "unverified_patches") - if bool(run.config.get("use_fibers", True)): - optional_dir(paths.fibers, "fibers") + optional_dir(paths.fibers, "fibers") shell_enabled = ( float(run.config.get("loss_weight_shell_outer", 1.0)) > 0 @@ -461,8 +361,7 @@ def optional_dir(value: str, field_name: str, required: bool = False) -> None: ) optional_dir(paths.outer_shell, "outer_shell", required=shell_enabled) - if (bool(run.config.get("use_tracks", True)) and paths.tracks_dbm - and not resolve_logical_dbm(paths.tracks_dbm)): + if paths.tracks_dbm and not resolve_logical_dbm(paths.tracks_dbm): errors.append({"field": "tracks_dbm", "message": "DBM logical base or backing file was not found"}) for index, spec in enumerate(paths.pcls): @@ -484,15 +383,11 @@ def optional_dir(value: str, field_name: str, required: bool = False) -> None: "message": "Must be phase or grad_mag"}) spacing_mode = None - normals_selected = bool(run.config.get("use_normals", True)) - sdt_selected = bool(run.config.get("use_surf_sdt", True)) - grad_mag_selected = bool(run.config.get("use_gradient_magnitude", True)) - use_normals = (normals_selected - and float(run.config.get("loss_weight_dense_normals", 100.0)) > 0) + use_normals = float( + run.config.get("loss_weight_dense_normals", 100.0)) > 0 spacing_enabled = float(run.config.get("loss_weight_dense_spacing", 12.0)) > 0 - use_phase = spacing_mode == "phase" and normals_selected and sdt_selected - use_grad_mag = ( - spacing_mode == "grad_mag" and spacing_enabled and grad_mag_selected) + use_phase = spacing_mode == "phase" + use_grad_mag = spacing_mode == "grad_mag" and spacing_enabled # The phase bundle requires its core inputs (SDT for phase, count, and # attachment; both normal channels for band incidence handling) even when # individual sub-weights are zero, so run-mutable weights can be raised diff --git a/volume-cartographer/scripts/spiral/fit_spiral.py b/volume-cartographer/scripts/spiral/fit_spiral.py index af11697157..e05214bfab 100644 --- a/volume-cartographer/scripts/spiral/fit_spiral.py +++ b/volume-cartographer/scripts/spiral/fit_spiral.py @@ -1,4 +1,13 @@ import os +import sys + +# Expandable segments stop the CUDA caching allocator from ratcheting its +# reserved pool toward the VRAM ceiling under the variable-size per-step loss +# graphs (measured ~12 GB lower steady-state envelope on the s1 fit). Must be +# set before the allocator initialises; an explicit PYTORCH_CUDA_ALLOC_CONF in +# the environment always wins. +os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True') + import copy import itertools import json @@ -27,6 +36,7 @@ maybe_init_distributed, split_counts_across_ranks, ) +from config import Config from lasagna_data import prepare_lasagna_volume, prepare_surf_sdt_volume from checkpoint_io import load_checkpoint_cpu from influence import make_influence_state, subsample_rows @@ -55,6 +65,7 @@ prepare_main_phase_tracks, validate_track_sampling_config, ) +from track_graph import TrackGraph from umbilicus import thaumato_umbilicus_z_to_yx, json_umbilicus_z_to_yx from sample_spiral import ( get_spiral_points, @@ -102,6 +113,7 @@ ) from visualization import overlay_patches_on_slices from transforms import SpiralAndTransform +from spiral_progress import ProgressReporter, progress_or_null configure_torch_threads_from_env() @@ -121,7 +133,10 @@ surf_sdt_zarr_group = '1' pcl_json_paths = [ f'{dataset_path}/abs_winding.json', + f'{dataset_path}/patch-overlap-pcls.json', f'{dataset_path}/relative_windings.json', + f'{dataset_path}/same_windings.json', + f'{dataset_path}/drawn_control_points.json', ] # The interactive session API supplies explicit roles. The legacy CLI leaves # this as None and retains the historical abs_winding.json basename behavior. @@ -139,7 +154,8 @@ voxel_size_um = 9.6 cache_path = os.environ.get('FIT_SPIRAL_CACHE_DIR', '../cache') lasagna_scale = 4 -# Normals, grad magnitude, and SDT are served by bounded sparse CUDA LRU caches. +# Normals, grad magnitude, and SDT are served by fully-resident sparse brick +# pools loaded from pack_resident_pools.py sidecars next to the source zarrs. lasagna_storage_backend = 'sparse_cuda' render_volume_scale = int(os.environ.get('FIT_SPIRAL_RENDER_VOLUME_SCALE', '1' if scroll_zarr_path else '16')) _active_lasagna_store = None @@ -155,336 +171,6 @@ def release_interactive_resources(): while _active_scalar_stores: _active_scalar_stores.pop().close() -default_config = { - # Session-level dataset input switches. VC3D uses these in addition to - # clearing local paths because a dataset-owned remote service resolves - # its own paths. Disabled inputs must never be opened by the fitter. - 'use_verified_patches': True, - 'use_unverified_patches': True, - 'use_normals': True, - 'use_surf_sdt': True, - 'use_tracks': True, - 'use_gradient_magnitude': True, - 'use_fibers': True, - 'random_seed': 1, - # Multi-GPU batch policy (only relevant under torchrun, world size > 1): - # True -> split per-step object-sample counts by world_size so the effective - # per-step batch matches single-GPU while each rank does less work. - # False -> scale-up: every rank keeps full counts, giving an N x larger effective batch. - 'distributed_split_batch': True, - 'learning_rate': 3.e-5, - 'exp_lr_schedule': True, - 'lr_final_factor': 0.3, - 'num_training_steps': 30_000, - 'num_flow_integration_steps': 3, - 'flow_integration_solver': 'rk4', - 'num_flow_timesteps': 1, - # Number of independent stationary flow fields composed sequentially, - # phi = exp(v_N) o ... o exp(v_1) (each stage keeps the fast inline-rk4 + sparse-grad - # path when num_flow_timesteps == 1 and the solver is rk4). 1 == original behaviour. - 'num_flow_stages': 1, - 'flow_bounds_z_margin': 160, - 'flow_bounds_radius': 3200, - 'flow_voxel_resolution': 16, - 'flow_field_type': 'cartesian', # 'cartesian' or 'cylindrical' - 'flow_field_high_res_lr_scale_initial': 2.0e-1, - 'flow_field_high_res_lr_scale_final': 2.0e-1, - 'flow_field_high_res_lr_ramp_start_step': 0, - 'flow_field_high_res_lr_ramp_steps': 1, - 'gap_expander_logit_resolution': 24, - 'gap_expander_num_windings': 130, - 'gap_expander_lr_scale': 0.3, - 'linear_z_resolution': 48, - 'initial_dr_per_winding': 16., - 'patch_radius_loss_margin': 0.025, - 'patch_radius_loss_inv': False, - 'patch_loss_z_margin': 0, - 'patch_dt_norm_p': 0.5, - 'patch_dt_within_patch_norm_p': 3.0, - 'patch_dt_loss_margin': 0.025, - 'patch_radius_within_norm_p': 3.0, # >1 emphasises worst within-track points in the radius loss - 'num_patches_per_step': 360, - 'num_patches_per_step_for_dt': 240, - 'num_points_per_patch': 800, - # How patch-loss strips are sampled in patch ij space (patch radius/DT losses and the - # rel/abs-winding L-shapes): - # 'straight' -> contiguous subranges of single rows/columns (+ 4 cardinal L-shapes for - # the winding losses). - # 'dijkstra' -> wiggly geodesic strips: from a start cell, walk the 8-connected valid-quad - # graph to a distant reachable endpoint via shortest path, skirting holes / - # ragged edges. Paths land in small per-patch (and per-anchor, for the - # winding losses) pools, built and continuously refreshed dataloader-style - # by FIT_SPIRAL_STRIP_PATH_WORKERS background processes (default 4; 0 = - # inline builds with fixed pools); per-step sampling just subsamples points - # along a pooled path, so steady-state cost matches 'straight'. - 'patch_strip_sampling': 'straight', - 'erode_patches': 1, # if >0, erode every patch's valid region (verified + unverified) by this many grid cells - 'disable_patches': False, # fit on PCLs + tracks only; load no verified/unverified patches - 'unverified_patch_radius_loss_margin': 0.025, - 'unverified_patch_radius_loss_inv': False, - 'unverified_patch_radius_within_norm_p': 3.0, - 'unverified_patch_dt_norm_p': 0.5, - 'unverified_patch_dt_within_patch_norm_p': 3.0, - 'unverified_patch_dt_loss_margin': 0.025, - 'unverified_num_patches_per_step': 120, - 'unverified_num_patches_per_step_for_dt': 80, - 'unverified_num_points_per_patch': 800, - 'unverified_patch_exclusion_radius': 64.0, # mask unverified-patch vertices within this of trusted geometry (full-res voxels) - 'rel_winding_num_pcls': 48, - 'rel_winding_num_patch_pairs_per_pcl': 4, - 'rel_winding_adjacent_patches_only': True, - # Stratify the per-step pcl draws (rel-winding and unattached-strip losses) so each - # pcl source file, plus fibers split into horizontal/vertical, gets an equal share - # of the samples regardless of how many pcls it holds. This legacy switch remains - # for existing interactive configs; an explicit pcl_sampling_weights dict below - # takes precedence and enables weighted rather than equal stratification. False - # with pcl_sampling_weights=None restores uniform-over-pcls sampling. - 'stratified_pcl_sampling': True, - # Per-group weighting of the per-step pcl draws (rel-winding and unattached-strip - # losses). None uses stratified_pcl_sampling above. A dict switches on weighted - # stratified sampling: each sampling group (pcl source json, - # keyed by basename with the .json suffix stripped e.g. 'relative_windings'; - # fibers split into 'fibers:H' / 'fibers:V') gets a per-step share of the samples - # proportional to its weight, regardless of how many pcls it holds. When set, the - # dict must list every group explicitly (a missing group is an error); weight 0 - # switches a group off entirely. Equal weights reproduce plain stratification; e.g. - # {'abs_winding': 1, 'patch-overlap-pcls': 1, 'relative_windings': 1, - # 'same_windings': 0, 'fibers:H': 2, 'fibers:V': 1} drops same-windings and - # doubles the horizontal-fiber share. - 'pcl_sampling_weights': None, - 'abs_winding_num_pcls': 48, - 'abs_winding_num_points_per_pcl': 4, - 'fiber_min_point_spacing': 40., - 'unattached_pcl_num_per_step': 84, - 'unattached_pcl_num_points_per_step': 32, - 'unattached_pcl_min_point_spacing': 16., - 'track_num_per_step': 48000, - 'track_num_points_per_step': 96, - # Resample each complete track in full-resolution polyline arclength. - # track_num_points_per_step selects target-estimation points; the spacing - # bounds control how many points contribute to the complete-track loss. - 'track_min_sample_spacing': 20.0, - 'track_max_sample_spacing': 60.0, - # Optional track-pool policies. None preserves uniform sampling and keeps - # every track regardless of tortuosity; crossing supplements are opt-in. - 'track_length_bin_weights': [0.0, 0.15, 0.85], # [short, medium, long], using eligible-track arclength tertiles - 'track_max_tortuosity': None, # whole-track arclength / endpoint chord; None disables filtering - # The complete eligible crossing CSR is retained for the session. The active - # per-step random sample can change freely up to this safety ceiling at a - # Run boundary. The legacy setting name is kept for profile compatibility. - 'track_crossing_precompute_max': 8, - # Opposite-family partners joined to each primary track's shared winding - # target. Zero disables crossing-connected sampling. - 'max_track_crossing_per_step': 1, - # Native crossing-chain sampling. "count" preserves the historical - # fixed-width primary/partner sampler exactly. - 'track_crossing_mode': 'count', - 'min_walk_steps_per_track': 24, - 'max_walk_steps_per_track': 256, - 'n_walks_per_track': 4, - 'track_walk_require_loop_consistency': False, - 'track_exclusion_radius': 16.0, - 'track_radius_target': 'mean', - 'track_radius_loss_margin': 0.025, - 'track_radius_within_norm_p': 6.0, # >1 emphasises worst within-track point in the radius loss (1.0 = mean) - 'track_dt_within_track_norm_p': 3.0, # within a track; -> inf strongly penalises isolated badly-aligned points - 'track_dt_norm_p': 0.5, # across tracks; -> 0 prefers many fully-satisfied tracks (winner-take-all snapping) - 'track_dt_loss_margin': 0.025, - 'dense_normals_num_points': 60_000, - 'regularisation_num_points': 4500, - 'grad_mag_encode_scale': 1000.0, - 'grad_mag_factor': 0.25, - 'spacing_integration_steps': 8, - # Dense spacing has exactly two modes: 'phase' (the opt-in bundle - - # soft-sequence phase registration, crossing count, native minimum - # spacing, and SDT attachment, each with its own weight) and 'grad_mag' - # (the legacy density integral and default inherited from origin/main). - 'dense_spacing_mode': 'phase', - 'dense_spacing_num_pairs': 12_000, - # m is a two-range mixture biased short: longer baselines average out - # per-gap counting noise (std ~ 0.5 * sqrt(m)) and let rays straddle wide - # unsupported regions, but the prediction's ~2-4% per-gap conservatism - # biases long counts low, so most pairs stay short; watch - # dense_spacing_count_mean before shifting the mixture longer. - 'dense_spacing_pair_m_short': (3, 7), # m uniform here for most pairs - 'dense_spacing_pair_m_long': (5, 15), # long-tail range for the rest - 'dense_spacing_pair_long_fraction': 0.15, - 'dense_spacing_count_temperature_wv': 0.5, # GT-calibrated; 1.0 undercounts ~12% - 'dense_spacing_target_step_wv': 1.0, # polyline step target (sheets are 4-6 wv wide) - 'dense_spacing_max_step_wv': 2.0, # mapped adjacent samples above this invalidate the pair - 'dense_spacing_max_steps': 1400, # per-gap allocation; scaled from the measured m=7 budget (640) to cover m=15 rays - 'dense_spacing_step_oversample': 1.25, # each mapped gap chord can underestimate curvature - 'dense_spacing_use_support_gate': True, - 'dense_spacing_support_sigma': 4.0, # working voxels; measured mean endpoint support 0.61 - 'dense_spacing_support_floor_alpha': 0.05, # nominal-mass denominator floor - 'dense_spacing_support_policy': 'product', # or 'minimum' - # Optional count-only ray supplement on top of the shared phase/count - # batch, if the joint batch alone gives too little spatial coverage. - 'dense_spacing_count_extra_pairs': 0, - 'dense_spacing_phase_huber_delta': 0.5, - 'dense_spacing_phase_extension_windings': 1.0, - 'dense_spacing_phase_min_center_gap_wv': 4.0, - 'dense_spacing_phase_graze_dot': 0.4, - 'dense_spacing_phase_graze_depth_wv': 1.0, - # Soft sequence alignment (pair-HMM): matches are restricted to an - # absolute phase window so the aligner cannot explain a ray with a global - # integer shift; missing/extra skip costs are the effective - # outlier-truncation knobs; open == extend keeps each gap cost - # length-constant (2026-07-17 calibration: MAP missing/extra run lengths - # are geometric, so affine costs stay tied). Window/skip-cost/margin/ - # temperature values are the 2026-07-17 GT-calibrated set: window 0.75 - # cut wrong-registration matches (accepted |rho|>0.5) ~25% at equal - # coverage; missing 0.55/extra 0.7 prefer skipping over forced wrong - # matches; temperature 0.1 recovers ~half the suppressed mass and is the - # measured annealing floor (0.05 breaks half-winding ambiguity safety). - 'dense_spacing_phase_window_windings': 0.75, - # Bands outside the central interval by more than this margin are free to - # skip (semi-global end gaps); the margin protects slightly displaced - # boundary bands from being dumped into the free gap. - 'dense_spacing_phase_end_free_margin_windings': 0.5, - 'dense_spacing_phase_missing_cost': 0.55, - 'dense_spacing_phase_missing_extend_cost': 0.55, - 'dense_spacing_phase_extra_cost': 0.7, - 'dense_spacing_phase_extra_extend_cost': 0.7, - 'dense_spacing_phase_temperature': 0.1, - 'dense_spacing_phase_band_confidence_cost': 0.25, # * (1 - |normal dot|), detached - # Confidence policy: suppress a winding's phase gradient when its match - # marginal is multimodal (low top-2 margin), and require a minimum number - # of useful matched windings plus matched mass before scoring a ray. - 'dense_spacing_phase_top2_margin': 0.1, - 'dense_spacing_phase_min_matched_windings': 2, - 'dense_spacing_phase_min_matched_mass': 1.0, - # Native pre-expansion anti-collapse barrier. This is asset-independent - # and can run in either dense-spacing mode. - 'loss_weight_min_spacing': 2.0, - # Crossing count: per-step it is gradient-starved on coarse fits and its - # support gate self-confirms, but its integer-topology pressure compounds - # - at 2000-step held-out probes (2026-07-17, wrap-aware scorer) count 8 - # is the strongest single addition to the bundle (med_err 0.2179 vs - # 0.2346 without) even though 60/300-step probes read it as neutral. - # Note it over-contracts the already-correct bins slightly (its soft - # count is ~2-5% conservative); revisit the weight once fits converge. - # 2026-07-17 PM: 8 -> 2. On clean (GT-excluded-from-train) 1000z probes - # across top/mid/low bands, the count dose-response is monotone toward - # low weights near convergence (w8 over-contracts tight bins; w1-2 is - # the plateau). w2 keeps coarse-regime pull without the converged-fit - # harm. See docs/spiral_experiment_notebook.md 2026-07-17 wave C. - 'loss_weight_dense_spacing_count': 0.0, - # Metric density term: |integral(lambda ds) - m| with a detached - # piecewise density (inverse detected-band gaps) integrated over the - # live central polyline; every step carries gradient (the principled - # replacement for the grad_mag integral). Band-GT calibration - # (wrap-aware, 2026-07-17): reads 0.94-1.06 for 12-45 wv spacings at - # r<1000, under-reads ~12% at r>1000 (detection recall - partial - # crossings), and is blind below ~12 wv detected spacing (SDT resolution - # floor) - hence the min-gap abstention below. - 'loss_weight_dense_spacing_density': 12.0, - # Sub-resolution abstention: below ~12 wv detected spacing the store - # under-reads winding density ~30% (measured), so steps bracketed by a - # gap under min_gap_wv can abstain (contribute nothing, target-corrected - # by their detached model-phase span). DEFAULT OFF (0): 2000-step - # held-out probes show the biased contraction signal still helps - # far-from-converged fits at every gap bin (0.2203 vs 0.2346 med_err); - # set ~12 for late refinement of an already-registered fit, where - # optimising toward the biased reading would thin the tightest windings - # toward |dW| ~ 0.67 (dense_spacing_density_max_blind_fraction caps how - # blind a pair may be before it is dropped outright). - 'dense_spacing_density_min_gap_wv': 0.0, - 'dense_spacing_density_max_blind_fraction': 0.75, - # Density-only sampling supplement (fast path: polylines + band - # detection, no pair-HMM), chunked so one chunk's graph is resident per - # backward. Total density pairs = dense_spacing_num_pairs + this. - 'dense_spacing_density_extra_pairs': 24_000, - # one 24k chunk instead of 3x8k: each chunk re-pays detection/pads/ - # polyline launch (~22 ms/step on H100 at the shipped shape); the - # chunk graph peak is ~8.5 GB at 300z — fits H100/GB10 with room - 'dense_spacing_density_chunk_pairs': 24_000, - 'min_spacing_d_min_wv': 6.0, - 'min_spacing_independent_samples': 2_000, - # SDT attachment: independent of the spacing loss (own weight/enable/counts). - # Late-fit snapping term: weight 10 actively fights coarse registration - # (pulls onto aliased sheets); 2-4 after warmup+ramp is the calibrated - # comparable-influence range (2026-07-17 gradient-norm probes). - 'loss_weight_dense_attachment': 0.0, - 'dense_attachment_scale': 8.0, # working voxels; relu(sd) p50/p90 = 2/9 on the shipped store - 'dense_attachment_num_points': 20_000, - 'dense_attachment_warmup_steps': 3_000, # measured against durable completed iterations - 'dense_attachment_ramp_steps': 3_000, - 'dense_normals_finite_difference_epsilon': 8.0, - 'sym_dirichlet_finite_difference_epsilon': 4.0, - 'loss_weight_patch_radius': 8.e0, - 'loss_weight_patch_dt': 4.e0, - 'loss_weight_unverified_patch_radius': 2.e0, - 'loss_weight_unverified_patch_dt': 1.e0, - 'loss_weight_rel_winding': 5., - 'loss_weight_abs_winding': 5., - 'loss_weight_unattached_pcl_radius': 2.e0, - 'loss_weight_unattached_pcl_dt': 4.e0, - 'loss_weight_track_radius': 50., - 'loss_weight_track_dt': 10., - 'loss_weight_sym_dirichlet': 10.0, - 'loss_weight_dense_normals': 1.e2, - 'loss_weight_dense_spacing': 12., - 'loss_weight_umbilicus': 1.25, - 'loss_weight_shell_outer': 1.0, - 'loss_weight_shell_patch_radius': 0.0, - 'weight_decay_gap_expander': 1.e-2, - 'weight_decay_flow_field': 0.0, - 'loss_start_patch_dt': 25_000, - 'loss_start_track_dt': 10_000, - 'loss_start_unverified_patch_dt': None, # None => fall back to loss_start_patch_dt - 'dt_progressive_windings': False, # gate the DT losses (patch, track, unattached-pcl) to grow outwards across windings - 'dt_progressive_inner_winding': 20, # outer-winding cutoff when each DT loss first turns on - 'dt_progressive_steps': 50_000, # steps to grow the cutoff from start_winding to shell_outer_winding_idx - 'dt_progressive_exponent': 1.0, # warp on the time fraction; 1.0 = linear in winding, <1 = slower later (~0.5 ≈ constant area rate) - # Whole-object DT targeting (see dt_targets.py, ported from - # origin/spiral-fibers-and-dt d60b739eb): determine each object's DT target from a - # sparse no-grad sample of the WHOLE patch/strip/track instead of each step's small - # sample median (which flip-flops across the rounding boundary for objects floating - # between windings). Objects attached to a winding use their median's nearest - # winding; objects whose majority is floating in a gap grab the outer one. - # 'strip_median' restores the legacy per-step sample target. - 'dt_target_mode': 'strip_median', # 'strip_median' | 'whole_object_quantile' - 'dt_target_floating_threshold': 0.25, # median point-to-nearest-sheet distance, in windings, required to grab the outer sheet - 'dt_target_update_interval': 100, # steps between whole-object target recomputations (1 = every step) - 'patch_dt_target_num_points': 256, # per-patch whole-grid samples for target determination - 'dt_target_max_stride': 128, # max gap between target-determination samples, in voxels (patches convert to grid cells via patch.scale; strip points are nominally at ~voxel spacing, so applied directly as an index stride) - 'dt_target_num_points_per_strip': 512, # target samples per track / pcl strip - 'output_first_winding': 10, - 'output_winding_margin': 4, - 'output_step_size': 20, - 'shell_outer_winding_idx': 130, - 'shell_outer_winding_margin': 10, - 'shell_num_samples': 24576, - 'shell_num_theta_bins': 720, - 'shell_huber_delta': 16.0, - 'shell_table_smooth_sigma_z': 4.0, - 'shell_table_smooth_sigma_theta': 1.0, - 'shell_min_confidence': 0.25, - # Final diagnostic PNG overlays are expensive at scroll resolution and are - # not needed for mesh output or VC3D interactive previews. - 'save_png_visualizations': False, - # Localized influence regions for interactive (ephemeral) inputs: when - # enabled, each input added to a resident session mid-run may only adjust - # the fit within its own footprint dilated by the extents below (gaussian - # decay towards the boundary), while everything outside is held in place - # by gradient masking plus an anchoring loss. See influence.py. - 'interactive_influence_enabled': False, - 'interactive_influence_z': 3000.0, # hard half-extent along z, full-res voxels - 'interactive_influence_windings': 5.0, # hard half-extent across wraps, windings - 'interactive_influence_theta_frac': 0.5, # hard half-extent along the wrap, fraction of a full turn (circular; 0.5 = the whole circle is within reach of some point) - 'interactive_influence_disable_dt_frac': 0.75, # fraction of each requested run that suppresses DT losses after incorporating inputs - 'interactive_influence_sigma': 0.3333, # gaussian sigma as a fraction of the hard extent - 'interactive_influence_footprint_points': 2048, # subsampled per incorporated input - 'loss_weight_anchor': 0.0, - # Anchor-bank sizes are absolute (not scaled with the z-range like the - # per-step object counts): the bank must cover the fitted volume densely - # enough regardless of how many objects each loss samples per step. - 'interactive_influence_anchor_lattice_points': 100_000, - 'interactive_influence_anchor_geometry_points': 100_000, - 'interactive_influence_anchor_samples_per_step': 4096, - 'interactive_influence_anchor_ramp_power': 2.0, -} cfg = None @@ -495,7 +181,7 @@ def get_env_config_overrides(): if not overrides_json: return {} overrides = json.loads(overrides_json) - unknown_keys = sorted(set(overrides) - set(default_config)) + unknown_keys = sorted(set(overrides) - set(Config().as_dict())) if unknown_keys: raise KeyError(f'unknown FIT_SPIRAL_CONFIG_OVERRIDES keys: {unknown_keys}') return overrides @@ -611,11 +297,16 @@ def lookup(self, scan_zyx): class PatchGpuAtlas: - """All patches' (H, W, 3) zyxs grids packed into one flat GPU tensor, so - fractional-(i, j) bilinear lookups can run as a single batched gather instead - of per-patch CPU dispatch.""" + """All patches' (H, W, 3) zyxs grids packed into one flat tensor, batched + per lookup instead of per-patch dispatch. The packed grids stay resident in + host memory: the (i, j) samples are drawn on the CPU anyway, so the + bilinear gather runs there (mostly-contiguous strip reads, a few MB per + step) and only the interpolated points are uploaded to `device`. This + keeps the atlas - which scales with the input patch count, not with any + per-step budget - out of VRAM entirely.""" def __init__(self, patches_by_id, device='cuda'): + self.device = torch.device(device) flat_pieces = [] offsets = [0] widths = [] @@ -628,16 +319,13 @@ def __init__(self, patches_by_id, device='cuda'): offsets.append(offsets[-1] + H * W) widths.append(W) heights.append(H) - # Concatenate on CPU and perform one CUDA transfer. Concatenating pieces - # after individually uploading them temporarily requires roughly two - # complete atlases of VRAM during construction. self.zyxs_flat = ( - torch.cat(flat_pieces, dim=0).to(device=device) + torch.cat(flat_pieces, dim=0) if flat_pieces - else torch.empty([0, 3], dtype=torch.float32, device=device)) - self.offsets = torch.tensor(offsets, device=device, dtype=torch.int64) # (N+1,) - self.widths = torch.tensor(widths, device=device, dtype=torch.int64) # (N,) - self.heights = torch.tensor(heights, device=device, dtype=torch.int64) # (N,) + else torch.empty([0, 3], dtype=torch.float32)) + self.offsets = torch.tensor(offsets, dtype=torch.int64) # (N+1,) + self.widths = torch.tensor(widths, dtype=torch.int64) # (N,) + self.heights = torch.tensor(heights, dtype=torch.int64) # (N,) self.id_to_idx = {pid: i for i, pid in enumerate(patches_by_id.keys())} native = load_native_spiral_sampling() self.sampling_atlas = ( @@ -652,10 +340,16 @@ def memory_mb(self): return self.zyxs_flat.numel() * 4 / 1e6 def lookup(self, patch_idx_per_sample, ijs): - # patch_idx_per_sample: (...,) int64 on GPU - # ijs: (..., 2) float on GPU - # returns (..., 3) on GPU. Caller must ensure floor(ij) lies on a valid quad. - return bilinear_atlas_lookup( + # patch_idx_per_sample: (...,) int64 on CPU + # ijs: (..., 2) float on CPU + # Gathers and interpolates on the host-resident atlas and returns + # (..., 3) on self.device. Caller must ensure floor(ij) lies on a + # valid quad. Runs inside the batch prefetcher when that is enabled, + # so both the gather and the upload happen a step ahead. + assert not patch_idx_per_sample.is_cuda and not ijs.is_cuda, ( + 'the atlas is host-resident: pass CPU indices/ijs; only the ' + 'interpolated points are uploaded') + zyxs = bilinear_atlas_lookup( self.zyxs_flat, self.offsets, self.widths, @@ -663,17 +357,17 @@ def lookup(self, patch_idx_per_sample, ijs): ijs, heights=self.heights, ) + return zyxs.to(device=self.device, non_blocking=True) def append_patches(self, patches_by_id): """Append new patches without rebuilding the resident atlas. - Only the new grids are uploaded; the existing flat tensor is - concatenated onto, so a resident interactive session can incorporate a - handful of added patches in seconds. + A host-side concatenation of just the new grids, so a resident + interactive session can incorporate a handful of added patches in + seconds. """ if not patches_by_id: return - device = self.zyxs_flat.device flat_pieces = [] offsets = [int(self.offsets[-1].item())] widths = [] @@ -687,16 +381,16 @@ def append_patches(self, patches_by_id): offsets.append(offsets[-1] + H * W) widths.append(W) heights.append(H) - new_flat = torch.cat(flat_pieces, dim=0).to(device=device) + new_flat = torch.cat(flat_pieces, dim=0) self.zyxs_flat = torch.cat([self.zyxs_flat, new_flat], dim=0) self.offsets = torch.cat([ self.offsets, - torch.tensor(offsets[1:], device=device, dtype=torch.int64), + torch.tensor(offsets[1:], dtype=torch.int64), ]) self.widths = torch.cat([ - self.widths, torch.tensor(widths, device=device, dtype=torch.int64)]) + self.widths, torch.tensor(widths, dtype=torch.int64)]) self.heights = torch.cat([ - self.heights, torch.tensor(heights, device=device, dtype=torch.int64)]) + self.heights, torch.tensor(heights, dtype=torch.int64)]) next_idx = len(self.id_to_idx) for pid in patches_by_id: self.id_to_idx[pid] = next_idx @@ -763,19 +457,6 @@ def get_or_build_unattached_pcl_flat(pcl_strips, device): return flat -def get_flow_field_high_res_lr_scale(iteration): - # Factor multiplying the high-resolution flow logits, which scales down their effective - # learning rate relative to the main LR (kept <= 1 so the hi-res LR stays bounded by the - # main LR). Ramps linearly from _initial to _final over _ramp_steps steps, starting at - # _ramp_start_step; constant when _initial == _final. - initial = cfg['flow_field_high_res_lr_scale_initial'] - final = cfg['flow_field_high_res_lr_scale_final'] - start_step = cfg['flow_field_high_res_lr_ramp_start_step'] - ramp_steps = max(1, int(cfg['flow_field_high_res_lr_ramp_steps'])) - frac = min(1., max(0., (iteration - start_step) / ramp_steps)) - return min(1., initial + frac * (final - initial)) - - def get_progressive_dt_max_winding(iteration, dt_start_step, shell_outer_winding_idx): # When `dt_progressive_windings` is set, the DT losses (patch, track, unattached-pcl) only act # on tracks/patches whose snapped spiral-space winding is <= the returned cutoff. The cutoff @@ -831,16 +512,99 @@ def get_dense_attachment_ramp(iteration): return min(1.0, (iteration - warmup + 1) / ramp) -def main(load_only_patches_and_point_collections=False, interactive_driver=None): +def get_exponential_lr_at_step( + initial_lr, final_factor, completed_steps, training_horizon): + """LR on the absolute exponential curve for a completed-step count.""" + horizon = max(1, int(training_horizon)) + completed = max(0, int(completed_steps)) + gamma = float(final_factor) ** (1.0 / horizon) + return float(initial_lr) * gamma ** completed + + +def get_flow_field_high_res_lr_scale(iteration): + """Relative optimizer LR for the high-resolution flow lattices.""" + initial = cfg['model_flow_field_high_res_lr_scale_initial'] + final = cfg['model_flow_field_high_res_lr_scale_final'] + start_step = cfg['model_flow_field_high_res_lr_ramp_start_step'] + ramp_steps = max( + 1, int(cfg['model_flow_field_high_res_lr_ramp_steps'])) + fraction = min( + 1., max(0., (int(iteration) - int(start_step)) / ramp_steps)) + return min(1., float(initial) + fraction * (float(final) - float(initial))) + + +def set_optimizer_group_lr_scale( + optimiser, lr_scheduler, *, group, reference_group, scale, + initial_lr): + """Set one parameter group's LR relative to another optimizer group.""" + scale = float(scale) + group_index = next( + index for index, candidate in enumerate(optimiser.param_groups) + if candidate is group) + group['lr_scale'] = scale + group['initial_lr'] = float(initial_lr) * scale + group['lr'] = float(reference_group['lr']) * scale + lr_scheduler.base_lrs[group_index] = group['initial_lr'] + if len(lr_scheduler._last_lr) == len(optimiser.param_groups): + lr_scheduler._last_lr[group_index] = group['lr'] + + +def realign_optimizer_lr_schedule( + optimiser, lr_scheduler, *, initial_lr, final_factor, + completed_steps, training_horizon, exponential): + """Realign an optimizer and scheduler to an absolute training step.""" + horizon = max(1, int(training_horizon)) + completed = max(0, int(completed_steps)) + initial_lr = float(initial_lr) + + if exponential: + gamma = float(final_factor) ** (1.0 / horizon) + if not isinstance( + lr_scheduler, torch.optim.lr_scheduler.ExponentialLR): + lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimiser, gamma=gamma) + lr_scheduler.gamma = gamma + aligned_lr = get_exponential_lr_at_step( + initial_lr, final_factor, completed, horizon) + else: + if not isinstance(lr_scheduler, torch.optim.lr_scheduler.LambdaLR): + lr_scheduler = torch.optim.lr_scheduler.LambdaLR( + optimiser, lambda step: 1.) + aligned_lr = initial_lr + + lr_scales = [ + float(group.get('lr_scale', 1.)) + for group in optimiser.param_groups + ] + base_lrs = [initial_lr * scale for scale in lr_scales] + aligned_lrs = [aligned_lr * scale for scale in lr_scales] + for group, base_lr, current_lr in zip( + optimiser.param_groups, base_lrs, aligned_lrs): + group['initial_lr'] = base_lr + group['lr'] = current_lr + lr_scheduler.base_lrs = base_lrs + lr_scheduler.last_epoch = completed + lr_scheduler._last_lr = aligned_lrs + lr_scheduler._step_count = completed + 1 + return lr_scheduler, horizon + + +def main( + load_only_patches_and_point_collections=False, interactive_driver=None, + progress=None): global _active_lasagna_store + has_progress = progress is not None + progress = progress_or_null(progress) - np.random.seed(cfg['random_seed']) - torch.random.manual_seed(cfg['random_seed']) + np.random.seed(cfg['optimizer_random_seed']) + torch.random.manual_seed(cfg['optimizer_random_seed']) if load_only_patches_and_point_collections: scroll_zarr = None else: + progress.begin('loading', 'Loading umbilicus') umbilicus = umbilicus_z_to_yx() if scroll_zarr_path: + progress.begin('loading', 'Opening scroll volume') print('loading volume zarr') scroll_zarr = zarr.open(scroll_zarr_path, mode='r') else: @@ -850,20 +614,24 @@ def main(load_only_patches_and_point_collections=False, interactive_driver=None) # Patch loading and ROI filtering # ========================================================================== - def load_patches_from_dir(path): + def load_patches_from_dir(path, label='patches'): patches = {} - for entry in sorted(os.listdir(path)): + entries = sorted(os.listdir(path)) + progress.begin( + 'loading', f'Loading {label}', + step=0, total_steps=len(entries), unit='patches') + for entry_number, entry in enumerate(entries, start=1): segment_path = os.path.join(path, entry) try: patches[entry] = load_tifxyz(segment_path) except Exception as e: print(f'Failed to load segment {entry}: {e}') - continue + progress.update( + entry_number, detail=f'{len(patches):,} loaded') return patches filter_tracks_by_shell = ( not load_only_patches_and_point_collections - and bool(cfg.get('use_tracks', True)) and bool(tracks_dbm_path) and bool(shell_path) ) @@ -871,12 +639,11 @@ def load_patches_from_dir(path): if shell_losses_enabled() or filter_tracks_by_shell: if not shell_path: raise RuntimeError('shell losses are enabled, but FIT_SPIRAL_SHELL_PATH is not set') + progress.begin('loading', 'Loading outer shell') shell_patch = load_tifxyz(shell_path) - use_verified_patches = ( - bool(cfg.get('use_verified_patches', True)) and not cfg['disable_patches']) - use_unverified_patches = ( - bool(cfg.get('use_unverified_patches', True)) and not cfg['disable_patches']) + use_verified_patches = bool(verified_patches_path) and not cfg['input_disable_patches'] + use_unverified_patches = bool(unverified_patches_path) and not cfg['input_disable_patches'] if not use_verified_patches and not use_unverified_patches: verified_patches = {} unverified_patches = {} @@ -885,12 +652,13 @@ def load_patches_from_dir(path): # An empty verified dir is allowed when unverified patches are supplied # (unverified-only ablations); both empty is a configuration error. verified_patches = ( - load_patches_from_dir(verified_patches_path) + load_patches_from_dir(verified_patches_path, 'verified patches') if use_verified_patches and verified_patches_path else {} ) unverified_patches = {} if use_unverified_patches and unverified_patches_path: - unverified_patches = load_patches_from_dir(unverified_patches_path) + unverified_patches = load_patches_from_dir( + unverified_patches_path, 'unverified patches') if (not verified_patches and not unverified_patches and (use_verified_patches or use_unverified_patches)): @@ -899,24 +667,33 @@ def load_patches_from_dir(path): print(f" loaded {len(verified_patches)} patches") print(f" loaded {len(unverified_patches)} unverified patches") + patch_filter_total = len(verified_patches) + len(unverified_patches) + progress.begin( + 'loading', 'Filtering patches to fit region', + step=0, total_steps=patch_filter_total, unit='patches') + filtered_count = 0 for patches in (verified_patches, unverified_patches): for patch_id, patch in list(patches.items()): - # we erode cells this distance from any invalid cell to catch annotation errors - # which are hard to detect at the edges of patches - cells_to_erode = patch.erosion_cells(cfg['erode_patches']) - if cells_to_erode > 0: - if not erode_patch_valid_region(patch, cells_to_erode): + try: + # we erode cells this distance from any invalid cell to catch annotation errors + # which are hard to detect at the edges of patches + cells_to_erode = patch.erosion_cells(cfg['patch_erode_patches']) + if cells_to_erode > 0: + if not erode_patch_valid_region(patch, cells_to_erode): + del patches[patch_id] + continue + + # remove any patches which do not intersect with the roi we are fitting + if not patch_intersects_z_roi(patch, z_begin, z_end): del patches[patch_id] continue - - # remove any patches which do not intersect with the roi we are fitting - if not patch_intersects_z_roi(patch, z_begin, z_end): - del patches[patch_id] - continue - # ROI testing may materialise the compact valid-coordinate view. - # Training retains the base grid and masks, so regenerate this view - # lazily only for a later exporter that actually requests it. - patch.release_derived_caches() + # ROI testing may materialise the compact valid-coordinate view. + # Training retains the base grid and masks, so regenerate this view + # lazily only for a later exporter that actually requests it. + patch.release_derived_caches() + finally: + filtered_count += 1 + progress.update(filtered_count) # ========================================================================== # Point collection loading @@ -930,7 +707,10 @@ def load_patches_from_dir(path): input_specs = pcl_input_specs if input_specs is None: input_specs = [(pattern, None) for pattern in pcl_json_paths] - for pattern, explicit_role in input_specs: + progress.begin( + 'loading', 'Loading point collections', + step=0, total_steps=len(input_specs), unit='inputs') + for spec_number, (pattern, explicit_role) in enumerate(input_specs, start=1): expanded = sorted(glob.glob(pattern)) if glob.has_magic(pattern) else [pattern] for path in expanded: loaded = load_point_collection(path) or {} @@ -950,11 +730,14 @@ def load_patches_from_dir(path): ) point_collections[next_id] = pcl next_id += 1 + progress.update( + spec_number, detail=f'{len(point_collections):,} collections loaded') + progress.begin('loading', 'Loading fiber point collections') fiber_point_collections, next_id = load_fiber_point_collections( - fibers_path if cfg.get('use_fibers', True) else None, + fibers_path, next_id, - min_point_spacing=cfg['fiber_min_point_spacing'], + min_point_spacing=cfg['pcl_fiber_min_point_spacing'], ) # Fibers form two sampling groups, horizontal and vertical, rather than one # group per source file like the regular pcls. @@ -993,6 +776,11 @@ def pcl_intersects_z_roi(pcl): # patch area and use distance only as a tie-break. Between-patches pcls connect # overlapping patches and attach only to their named patch pair, using nearest # distance within that pair. + progress.begin( + 'loading', 'Linking points to patches', + detail=( + f'{len(point_collections):,} collections, ' + f'{len(verified_patches):,} patches')) link_points_to_patches( verified_patches, point_collections, @@ -1101,7 +889,7 @@ def pcl_intersects_z_roi(pcl): pcl['points_by_patch'] = points_by_patch unattached_pcl_strips = _UnattachedPclStripList() unattached_strip_sampling_groups = [] # parallel to unattached_pcl_strips - min_point_spacing = cfg['unattached_pcl_min_point_spacing'] + min_point_spacing = cfg['pcl_unattached_pcl_min_point_spacing'] # For each unattached pcl, materialise an id-sorted strip of point zyxs and the # corresponding winding annotations. Strips with <2 points are dropped. # If min_point_spacing > 0, decimate each strip greedily along its id-sorted order @@ -1140,7 +928,7 @@ def pcl_intersects_z_roi(pcl): f'pcls: {len(cross_patch_pcls)} cross-patch, ' f'{len(unattached_pcl_strips)} unattached' ) - if cfg['stratified_pcl_sampling'] or cfg['pcl_sampling_weights'] is not None: + if cfg['pcl_stratified_pcl_sampling'] or cfg['pcl_sampling_weights'] is not None: def _group_counts(groups): counts = {} for group in groups: @@ -1196,31 +984,26 @@ def rebuild_pcl_sampling_strata(): raise ValueError( f'dense_spacing_mode={dense_spacing_mode!r} must be ' "'phase' or 'grad_mag'") - phase_mode = ( - dense_spacing_mode == 'phase' - and cfg.get('use_normals', True) - and cfg.get('use_surf_sdt', True) - ) + phase_mode = dense_spacing_mode == 'phase' grad_mag_spacing_enabled = ( dense_spacing_mode == 'grad_mag' - and cfg.get('use_gradient_magnitude', True) and cfg['loss_weight_dense_spacing'] > 0 ) shell_envelope = None if shell_patch is not None and filter_tracks_by_shell: + progress.begin('loading', 'Building outer-shell lookup') shell_envelope = ShellPolarMap( shell_patch, umbilicus, - z_min=z_begin - cfg['flow_bounds_z_margin'], - z_max=z_end + cfg['flow_bounds_z_margin'], + z_min=z_begin - cfg['model_flow_bounds_z_margin'], + z_max=z_end + cfg['model_flow_bounds_z_margin'], num_theta_bins=cfg['shell_num_theta_bins'], device='cpu', ) lasagna_volume = prepare_lasagna_volume( scroll_zarr, - use_normals=(cfg.get('use_normals', True) - and (cfg['loss_weight_dense_normals'] > 0 or phase_mode)), + use_normals=(cfg['loss_weight_dense_normals'] > 0 or phase_mode), use_spacing=grad_mag_spacing_enabled, normal_nx_zarr_path=normal_nx_zarr_path, normal_ny_zarr_path=normal_ny_zarr_path, @@ -1231,6 +1014,7 @@ def rebuild_pcl_sampling_strata(): lasagna_scale=lasagna_scale, storage_backend=lasagna_storage_backend, cache_directory=cache_path, + progress=progress, ) if interactive_driver is not None and lasagna_volume: _active_lasagna_store = lasagna_volume['store'] @@ -1256,6 +1040,7 @@ def rebuild_pcl_sampling_strata(): z_end=z_end, cache_directory=cache_path, storage_backend=lasagna_storage_backend, + progress=progress, ) if interactive_driver is not None: _active_scalar_stores.append(sdt_volume['store']) @@ -1288,17 +1073,38 @@ def warn_if_sdt_loss_inactive(): track_families = None track_source_ids = None track_crossing_cache = None - if tracks_dbm_path is not None and cfg.get('use_tracks', True): + track_graph = None + track_reload_source = None + track_reload_families = None + track_reload_source_ids = None + if tracks_dbm_path is not None: + progress.begin( + 'loading', 'Resolving track store', + detail=os.path.basename(tracks_dbm_path)) print(f'loading tracks from {tracks_dbm_path}') if (track_sampling_config['crossing_precompute_max'] > 0 or track_sampling_config['crossing_mode'] == 'track_walk'): track_crossing_cache = load_track_crossing_cache(tracks_dbm_path) + if track_crossing_cache is not None: + track_graph = TrackGraph(track_crossing_cache) + print( + f'built TrackGraph: {len(track_graph)} tracks, ' + f'{track_graph.edge_count} crossings in ' + f'{track_graph.build_seconds:.1f}s') + track_crossing_cache = None tracks, track_families, track_source_ids = load_tracks_from_dbm( tracks_dbm_path, z_begin, z_end, return_families=True, - return_source_ids=True) + return_source_ids=True, progress=progress) else: - tracks = load_tracks_from_dbm(tracks_dbm_path, z_begin, z_end) + tracks = load_tracks_from_dbm( + tracks_dbm_path, z_begin, z_end, progress=progress) + track_reload_source = tracks + track_reload_families = track_families + track_reload_source_ids = track_source_ids if filter_tracks_by_shell: + progress.begin( + 'loading', 'Filtering tracks to outer shell', + detail=f'{len(tracks):,} tracks') tracks, track_families, kept_track_indices = filter_tracks_to_outer_shell( tracks, shell_envelope, track_families, return_indices=True) if track_source_ids is not None: @@ -1312,6 +1118,9 @@ def warn_if_sdt_loss_inactive(): # ========================================================================== def prepare_patch_sampling_cache(patches): + progress.begin( + 'loading', 'Preparing patch sampling', + step=0, total_steps=len(patches), unit='patches') native_sampling_available = load_native_spiral_sampling() is not None patch_areas = np.empty(len(patches), dtype=np.float32) for patch_idx, patch in enumerate(patches): @@ -1366,6 +1175,7 @@ def _build_line_runs(line_valid): ) patch_areas[patch_idx] = float(patch.area) + progress.update(patch_idx + 1) inv_weights = patch_areas ** 0.5 return inv_weights / inv_weights.sum() @@ -1383,14 +1193,17 @@ def _build_line_runs(line_valid): out_path += f'_{run_tag}' os.makedirs(out_path, exist_ok=True) + progress.begin( + 'loading', 'Building verified-patch GPU atlas', + detail=f'{len(verified_patches):,} patches') patch_atlas = PatchGpuAtlas(verified_patches, device='cuda') - print(f'patch GPU atlas: {patch_atlas.memory_mb():.1f} MB') + print(f'patch atlas (host-resident): {patch_atlas.memory_mb():.1f} MB') # ========================================================================================== # trusted geometry (verified patches and pcls) kdtree / unverified patches + tracks masking # ========================================================================================== - num_slices_for_visualisation = cfg.get('num_slices_for_visualization', 20) + num_slices_for_visualisation = cfg.get('output_num_slices_for_visualization', 20) device = torch.device('cuda') # The trusted point cloud is consumed only by a CPU cKDTree. Build it directly @@ -1432,6 +1245,9 @@ def _build_line_runs(line_valid): verified_patches_and_pcls_np = verified_patches_and_pcls_cpu.numpy() verified_patches_and_pcls_np = np.ascontiguousarray(verified_patches_and_pcls_np, dtype=np.float32) if verified_patches_and_pcls_np.shape[0] > 0: + progress.begin( + 'loading', 'Building trusted-geometry index', + detail=f'{len(verified_patches_and_pcls_np):,} points') trusted_geometry_tree = cKDTree(verified_patches_and_pcls_np) def _query_near_trusted_geometry(points_np, trusted_geometry_tree, threshold): @@ -1559,7 +1375,10 @@ def flush_batch(): # masks/area. Patches left with no valid quad are dropped. This is the patch analogue # of the DBM-track exclusion in tracks.py: untrusted patches only constrain regions # the trusted inputs don't already cover, so they can't fight verified geometry. - exclusion_radius = float(cfg['unverified_patch_exclusion_radius']) + exclusion_radius = float(cfg['patch_unverified_patch_exclusion_radius']) + progress.begin( + 'loading', 'Masking unverified patches', + detail=f'{len(unverified_patches):,} patches') unverified_patches, n_masked_vertices, n_dropped_patches = ( _mask_unverified_patches_near_trusted_geometry( unverified_patches, @@ -1578,6 +1397,37 @@ def flush_batch(): unverified_patch_sampling_probabilities = prepare_patch_sampling_cache(unverified_patches_list) unverified_patch_atlas = PatchGpuAtlas(unverified_patches, device='cuda') + def rebuild_unverified_patch_inputs(exclusion_radius): + """Reload only the unverified-patch pool for a Run-boundary mask edit.""" + if not unverified_patches_path: + return {}, [], None, None + candidates = load_patches_from_dir(unverified_patches_path) + for patch_id, patch in list(candidates.items()): + cells_to_erode = patch.erosion_cells(cfg['patch_erode_patches']) + if (cells_to_erode > 0 + and not erode_patch_valid_region(patch, cells_to_erode)): + del candidates[patch_id] + continue + if not patch_intersects_z_roi(patch, z_begin, z_end): + del candidates[patch_id] + continue + patch.release_derived_caches() + candidates, n_masked, n_dropped = \ + _mask_unverified_patches_near_trusted_geometry( + candidates, trusted_geometry_tree, exclusion_radius) + print( + f'unverified patches: remasked {n_masked} vertices near trusted ' + f'geometry (radius {exclusion_radius:.1f}), dropped {n_dropped}; ' + f'{len(candidates)} remain') + candidate_list = list(candidates.values()) + probabilities = ( + prepare_patch_sampling_cache(candidate_list) + if candidate_list else None) + atlas = ( + PatchGpuAtlas(candidates, device='cuda') + if candidate_list else None) + return candidates, candidate_list, probabilities, atlas + # The full z series is a model input. PNG-only slice grids and raster inputs # are prepared lazily at final export, and never in a resident VC3D session. all_zs = np.arange(z_begin, z_end) @@ -1637,6 +1487,9 @@ def prepare_png_visualization_inputs(): resume_checkpoint = None model_z_begin, model_z_end = z_begin, z_end if resume_path: + progress.begin( + 'loading', 'Loading fit checkpoint', + detail=os.path.basename(resume_path)) resume_checkpoint = load_checkpoint_cpu(resume_path) checkpoint_lasagna_scale = resume_checkpoint.get('lasagna_scale') if isinstance(resume_checkpoint, dict) else None if checkpoint_lasagna_scale != lasagna_scale: @@ -1685,10 +1538,10 @@ def _comparable_sdt_fingerprint(fingerprint): f'{checkpoint_fingerprint}\n current: {current_fingerprint}') checkpoint_cfg = resume_checkpoint.get('cfg', {}) shape_keys = ( - 'num_flow_integration_steps', 'flow_integration_solver', 'num_flow_timesteps', - 'flow_bounds_z_margin', 'flow_bounds_radius', 'flow_voxel_resolution', - 'flow_field_type', 'gap_expander_logit_resolution', - 'gap_expander_num_windings', 'linear_z_resolution', + 'model_num_flow_integration_steps', 'model_flow_integration_solver', 'model_num_flow_timesteps', + 'model_flow_bounds_z_margin', 'model_flow_bounds_radius', 'model_flow_voxel_resolution', + 'model_flow_field_type', 'model_gap_expander_logit_resolution', + 'model_gap_expander_num_windings', 'model_linear_z_resolution', ) incompatible = [ key for key in shape_keys @@ -1708,19 +1561,20 @@ def _comparable_sdt_fingerprint(fingerprint): 'checkpoint range, or train from scratch with the wider range.' ) - flow_field_radius = cfg['flow_bounds_radius'] + flow_field_radius = cfg['model_flow_bounds_radius'] flow_min_corner_spiral_zyx = torch.tensor( - [model_z_begin - cfg['flow_bounds_z_margin'], -flow_field_radius, -flow_field_radius], dtype=torch.int64, + [model_z_begin - cfg['model_flow_bounds_z_margin'], -flow_field_radius, -flow_field_radius], dtype=torch.int64, device=device) flow_max_corner_spiral_zyx = torch.tensor( - [model_z_end + cfg['flow_bounds_z_margin'], flow_field_radius, flow_field_radius], dtype=torch.int64, + [model_z_end + cfg['model_flow_bounds_z_margin'], flow_field_radius, flow_field_radius], dtype=torch.int64, device=device) - num_training_steps = cfg['num_training_steps'] + num_training_steps = cfg['optimizer_num_training_steps'] + progress.begin('loading', 'Constructing spiral model') spiral_and_transform = SpiralAndTransform( - flow_integration_steps=cfg['num_flow_integration_steps'], - flow_integration_solver=cfg['flow_integration_solver'], + flow_integration_steps=cfg['model_num_flow_integration_steps'], + flow_integration_solver=cfg['model_flow_integration_solver'], umbilicus_zyx=umbilicus_zyx, flow_min_corner_zyx=flow_min_corner_spiral_zyx, flow_max_corner_zyx=flow_max_corner_spiral_zyx, @@ -1735,19 +1589,32 @@ def _comparable_sdt_fingerprint(fingerprint): shell_map = None shell_valid_zyxs_gpu = None + + def subsample_shell_radius_pool(patch): + # The shell-patch radius loss draws sample_count_shell_samples random + # shell points per step; keep a pool of exactly that size resident on + # the GPU instead of the full shell cloud. A dedicated generator makes + # the pool deterministic (identical across DDP ranks) without + # perturbing the training RNG streams. + pool_generator = torch.Generator() + pool_generator.manual_seed(int(cfg['optimizer_random_seed'])) + return subsample_rows( + patch.valid_zyxs, int(cfg['sample_count_shell_samples']), pool_generator, + ).to(device=device, dtype=torch.float32) + shell_active = shell_patch is not None and shell_losses_enabled() if shell_active: if cfg['loss_weight_shell_outer'] > 0: shell_map = ShellPolarMap( shell_patch, umbilicus, - z_min=z_begin - cfg['flow_bounds_z_margin'], - z_max=z_end + cfg['flow_bounds_z_margin'], + z_min=z_begin - cfg['model_flow_bounds_z_margin'], + z_max=z_end + cfg['model_flow_bounds_z_margin'], num_theta_bins=cfg['shell_num_theta_bins'], device=device, ) if cfg['loss_weight_shell_patch_radius'] > 0: - shell_valid_zyxs_gpu = shell_patch.valid_zyxs.to(device=device, dtype=torch.float32) + shell_valid_zyxs_gpu = subsample_shell_radius_pool(shell_patch) def infer_outer_winding_idx_for_this_run(): return _infer_shell_outer_winding_idx( @@ -1761,11 +1628,7 @@ def infer_outer_winding_idx_for_this_run(): get_or_build_unattached_pcl_flat, ) - # The dense lasagna losses, the symmetric Dirichlet regulariser and the - # phase bundle all sample out to this index, with or without shell - # losses; the shell path only refines it (inference) when the config - # leaves it unset. It used to be resolved on that path only, so every - # one of those losses was silently zero on shell-less runs (#1220). + # Dense losses sample out to this index even when shell losses are off. shell_outer_winding_idx, outer_winding_notes = resolve_outer_winding_idx_and_notes( cfg, shell_active, infer_outer_winding_idx_for_this_run) for note in outer_winding_notes: @@ -1774,8 +1637,7 @@ def infer_outer_winding_idx_for_this_run(): dense_inactive_warned = set() def warn_if_dense_losses_structurally_disabled(): - # Like warn_if_sdt_loss_inactive above: loss weights are run-mutable, - # so this is re-checked every step and warns once per key. + # Loss weights are run-mutable, so re-check each step and warn once. for weight_key in _structurally_disabled_dense_weight_keys( cfg, shell_outer_winding_idx): if weight_key not in dense_inactive_warned: @@ -1791,31 +1653,83 @@ def warn_if_dense_losses_structurally_disabled(): # Optimizer and checkpoint helpers # ========================================================================== - # All flow stages' parameters go into the flow param group (stage 0 == .flow_field, - # plus any extra_flow_fields when num_flow_stages > 1). - flow_field_params = [p for flow_field in spiral_and_transform.flow_fields for p in flow_field.parameters()] + # Keep every stage's low- and high-resolution lattices in distinct groups + # so the HR learning-rate scale is an optimizer setting, not a multiplier + # in the model's forward path. + low_res_flow_params = [ + flow_field.flows[0] + for flow_field in spiral_and_transform.flow_fields + ] + high_res_flow_params = [ + flow_field.flows[1] + for flow_field in spiral_and_transform.flow_fields + ] + flow_field_params = low_res_flow_params + high_res_flow_params gap_expander_params = list(spiral_and_transform.gap_expander_params.parameters()) linear_params = [spiral_and_transform.linear_logits] grouped_ids = {id(p) for p in flow_field_params + gap_expander_params + linear_params} other_params = [p for p in spiral_and_transform.parameters() if id(p) not in grouped_ids] + initial_high_res_lr_scale = get_flow_field_high_res_lr_scale(0) param_groups = [ {'params': other_params, 'weight_decay': 0.0}, {'params': linear_params, 'weight_decay': 0.0}, - {'params': gap_expander_params, 'weight_decay': cfg['weight_decay_gap_expander']}, - {'params': flow_field_params, 'weight_decay': cfg['weight_decay_flow_field']}, + {'params': gap_expander_params, 'weight_decay': cfg['optimizer_weight_decay_gap_expander']}, + { + 'params': low_res_flow_params, + 'weight_decay': cfg['optimizer_weight_decay_flow_field'], + 'lr_scale': 1., + }, + { + 'params': high_res_flow_params, + 'weight_decay': cfg['optimizer_weight_decay_flow_field'], + 'lr': cfg['optimizer_learning_rate'] * initial_high_res_lr_scale, + 'lr_scale': initial_high_res_lr_scale, + }, ] - optimiser = torch.optim.AdamW(param_groups, lr=cfg.learning_rate, betas=(0.9, 0.999), eps=1.e-8, fused=True) + progress.begin('loading', 'Creating optimizer') + optimiser = torch.optim.AdamW(param_groups, lr=cfg.optimizer_learning_rate, betas=(0.9, 0.999), eps=1.e-8, fused=True) # Influence masks are scoped to one interactive Run request. They are # created from that run's pending inputs and discarded before its autosave. influence_state = None interactive_influence_loss_weight = 0.0 interactive_influence_anchor_samples = 0 - if cfg['exp_lr_schedule']: - gamma = cfg['lr_final_factor'] ** (1.0 / max(1, num_training_steps)) + if cfg['optimizer_exp_lr_schedule']: + gamma = cfg['optimizer_lr_final_factor'] ** (1.0 / max(1, num_training_steps)) lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimiser, gamma=gamma) else: lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimiser, lambda step: 1.) + def apply_high_res_lr_scale(iteration): + scale = get_flow_field_high_res_lr_scale(iteration) + low_res_group = next( + group for group in optimiser.param_groups + if any(param is low_res_flow_params[0] for param in group['params'])) + high_res_group = next( + group for group in optimiser.param_groups + if any(param is high_res_flow_params[0] for param in group['params'])) + set_optimizer_group_lr_scale( + optimiser, + lr_scheduler, + group=high_res_group, + reference_group=low_res_group, + scale=scale, + initial_lr=cfg['optimizer_learning_rate'], + ) + return scale + + def realign_lr_schedule(completed_steps): + """Align optimizer/scheduler state to the current absolute horizon.""" + nonlocal lr_scheduler, num_training_steps + lr_scheduler, num_training_steps = realign_optimizer_lr_schedule( + optimiser, + lr_scheduler, + initial_lr=cfg['optimizer_learning_rate'], + final_factor=cfg['optimizer_lr_final_factor'], + completed_steps=completed_steps, + training_horizon=cfg['optimizer_num_training_steps'], + exponential=cfg['optimizer_exp_lr_schedule'], + ) + def checkpoint_payload(completed_iterations): def durable_config(value): return { @@ -1887,7 +1801,7 @@ def load_model(checkpoint): gap_param = gap_expander_params[0] gap_group = next(group for group in optimiser.param_groups if any(param is gap_param for param in group['params'])) - gap_group['weight_decay'] = cfg['weight_decay_gap_expander'] + gap_group['weight_decay'] = cfg['optimizer_weight_decay_gap_expander'] if checkpoint.get('scheduler') is not None: lr_scheduler.load_state_dict(checkpoint['scheduler']) @@ -1896,6 +1810,9 @@ def load_model(checkpoint): if embedded_iteration is not None: start_iteration = int(embedded_iteration) print(f'resuming from {resume_path} at iteration {start_iteration}') + progress.begin( + 'loading', 'Restoring model and optimizer', + detail=os.path.basename(resume_path)) load_model(resume_checkpoint) if not isinstance(resume_checkpoint, dict) or resume_checkpoint.get('scheduler') is None: for _ in range(start_iteration): @@ -1922,6 +1839,12 @@ def load_model(checkpoint): del resume_checkpoint resume_checkpoint = None + if interactive_driver is not None: + # A checkpoint may carry the scheduler state from a shorter horizon. + # The active session configuration is authoritative. + realign_lr_schedule(start_iteration) + + progress.begin('loading', 'Synchronizing model across GPU workers') broadcast_model_params(spiral_and_transform) if os.environ.get('FIT_SPIRAL_TORCH_PROFILE') == '1': @@ -1943,6 +1866,9 @@ def load_model(checkpoint): prepared_main_tracks = None preview_extent_tracks = tracks if using_tracks: + progress.begin( + 'loading', 'Preparing tracks for optimization', + detail=f'{len(tracks):,} tracks') prepared_main_tracks = prepare_main_phase_tracks( tracks, None, @@ -1953,10 +1879,13 @@ def load_model(checkpoint): track_families=track_families, track_source_ids=track_source_ids, crossing_cache=track_crossing_cache, + track_graph=track_graph, ) # The sidecar CSR is setup-only. The prepared bundle now owns only its # fixed-width training tables, so release the whole-DB graph promptly. - track_crossing_cache = None + if interactive_driver is None: + track_crossing_cache = None + track_graph = None # With the usual zero exclusion radius, the training bundle already # contains every authoritative track point as one flat CPU tensor. Reuse it # for preview bounds instead of walking millions of short NumPy tracks. @@ -1974,16 +1903,17 @@ def load_model(checkpoint): influence_anchor_geometry = None if interactive_driver is not None: stash_generator = torch.Generator() - stash_generator.manual_seed(int(cfg['random_seed'])) + stash_generator.manual_seed(int(cfg['optimizer_random_seed'])) influence_anchor_geometry = subsample_rows( verified_patches_and_pcls_cpu, - int(cfg['interactive_influence_anchor_geometry_points']), + int(cfg['sample_count_influence_anchor_geometry_points']), stash_generator, ).clone() # The trusted cloud and its double-precision cKDTree are setup-only data. # Track sampling retains its own compact offsets and coordinates. - trusted_geometry_tree = None + if interactive_driver is None: + trusted_geometry_tree = None verified_patches_and_pcls_cpu = None verified_patches_and_pcls_np = None @@ -1998,12 +1928,17 @@ def load_model(checkpoint): raise ValueError(f"dt_target_mode must be 'strip_median' or 'whole_object_quantile', got {cfg['dt_target_mode']!r}") dt_target_whole_object = cfg['dt_target_mode'] == 'whole_object_quantile' if dt_target_whole_object: + progress.begin( + 'loading', 'Preparing distance-target samples', + detail=( + f'{len(verified_patches_list) + len(unverified_patches_list):,} ' + 'patches')) prepare_patch_dt_target_samples( - verified_patches_list, cfg['patch_dt_target_num_points'], cfg['dt_target_max_stride'], + verified_patches_list, cfg['sample_count_patch_dt_target_points'], cfg['dt_target_max_stride'], ) if unverified_patches_list: prepare_patch_dt_target_samples( - unverified_patches_list, cfg['patch_dt_target_num_points'], cfg['dt_target_max_stride'], + unverified_patches_list, cfg['sample_count_patch_dt_target_points'], cfg['dt_target_max_stride'], ) # Caches are recomputed lazily once the corresponding DT loss is active. # Updates are deterministic given the transform, so DDP ranks stay consistent. @@ -2026,8 +1961,8 @@ def report_first_dt_target_cache(kind, cache): # ========================================================================== if is_distributed(): - np.random.seed(cfg['random_seed'] + get_rank()) - torch.manual_seed(cfg['random_seed'] + get_rank()) + np.random.seed(cfg['optimizer_random_seed'] + get_rank()) + torch.manual_seed(cfg['optimizer_random_seed'] + get_rank()) dist_grad_params = list(spiral_and_transform.parameters()) dist_grad_named = list(spiral_and_transform.named_parameters()) if is_main_process(): @@ -2080,6 +2015,7 @@ def export_interactive_preview(generation_path, surface_id): get_or_build_unattached_pcl_flat, tracks=preview_extent_tracks, surface_id=surface_id, + progress=progress, ) diagnostic_weights = { name: cfg.get(f'loss_weight_{name}', 0.0) @@ -2099,10 +2035,12 @@ def export_interactive_preview(generation_path, surface_id): float(cfg['loss_weight_dense_spacing_count']), 1.0) transform = spiral_and_transform.get_slice_to_spiral_transform() dr = spiral_and_transform.get_dr_per_winding() + progress.begin( + 'exporting_preview', 'Computing preview diagnostics') recorder = LossMapRecorder( manifest, generation_path, - z0=z_begin - int(cfg['flow_bounds_z_margin']), + z0=z_begin - int(cfg['model_flow_bounds_z_margin']), grid_spacing=int(cfg['output_step_size']), dr_per_winding=dr, weights=diagnostic_weights, @@ -2110,8 +2048,8 @@ def export_interactive_preview(generation_path, surface_id): with torch.no_grad(), capture_loss_maps(recorder, suppress_errors=True): get_patch_and_umbilicus_losses( transform, dr, - cfg['num_patches_per_step'], - cfg['num_patches_per_step_for_dt'], + cfg['sample_count_patches_per_step'], + cfg['sample_count_patches_per_step_for_dt'], verified_patches_list, patch_atlas, patch_sampling_probabilities, umbilicus_zyx, compute_dt=cfg['loss_weight_patch_dt'] > 0, @@ -2121,8 +2059,8 @@ def export_interactive_preview(generation_path, surface_id): if unverified_patch_atlas is not None: get_unverified_patch_losses( transform, dr, - cfg['unverified_num_patches_per_step'], - cfg['unverified_num_patches_per_step_for_dt'], + cfg['sample_count_unverified_patches_per_step'], + cfg['sample_count_unverified_patches_per_step_for_dt'], unverified_patches_list, unverified_patch_atlas, unverified_patch_sampling_probabilities, compute_dt=cfg['loss_weight_unverified_patch_dt'] > 0, @@ -2130,7 +2068,7 @@ def export_interactive_preview(generation_path, surface_id): if cfg['loss_weight_sym_dirichlet'] > 0: get_symmetric_dirichlet_loss( transform, dr, shell_outer_winding_idx, - cfg['regularisation_num_points']) + cfg['sample_count_regularisation_points']) if cfg['loss_weight_rel_winding'] > 0 and cross_patch_pcls: get_patch_rel_winding_loss( transform, dr, verified_patches, patch_atlas, @@ -2143,7 +2081,7 @@ def export_interactive_preview(generation_path, surface_id): for _loss_name, _loss_value in iter_lasagna_losses( transform, dr, lasagna_volume, shell_outer_winding_idx, - cfg['dense_normals_num_points'], + cfg['sample_count_dense_normal_points'], compute_spacing=grad_mag_spacing_enabled): pass if phase_mode_active(): @@ -2159,8 +2097,8 @@ def export_interactive_preview(generation_path, surface_id): transform, dr, unattached_pcl_strips, pcl_sampling_strata['unattached'], get_or_build_unattached_pcl_flat, - cfg['unattached_pcl_num_per_step'], - cfg['unattached_pcl_num_points_per_step'], + cfg['sample_count_unattached_pcls_per_step'], + cfg['sample_count_unattached_pcl_points_per_step'], compute_dt=cfg['loss_weight_unattached_pcl_dt'] > 0, ) if prepared_main_tracks is not None: @@ -2236,12 +2174,12 @@ def incorporate_interactive_inputs(records, influence_config=None): path = record.get('path') input_id = record.get('id') if kind == 'patch': - if cfg['disable_patches']: + if cfg['input_disable_patches']: raise RuntimeError('disable_patches=True: this session takes no patches') if input_id in verified_patches or input_id in new_patches: raise RuntimeError(f'Patch {input_id!r} is already part of this session') patch = load_tifxyz(path) - cells_to_erode = patch.erosion_cells(cfg['erode_patches']) + cells_to_erode = patch.erosion_cells(cfg['patch_erode_patches']) if cells_to_erode > 0 and not erode_patch_valid_region(patch, cells_to_erode): raise RuntimeError(f'Patch {input_id!r} has no valid quads after erosion') if not patch_intersects_z_roi(patch, z_begin, z_end): @@ -2252,7 +2190,7 @@ def incorporate_interactive_inputs(records, influence_config=None): new_patches[input_id] = patch elif kind == 'fiber': pcl = load_fiber_point_collection( - path, next_id, min_point_spacing=cfg['fiber_min_point_spacing']) + path, next_id, min_point_spacing=cfg['pcl_fiber_min_point_spacing']) if pcl is None: raise RuntimeError(f'Fiber {input_id!r} has no usable control points') pcl['source_file'] = path @@ -2298,7 +2236,7 @@ def incorporate_interactive_inputs(records, influence_config=None): if cfg['dt_target_mode'] == 'whole_object_quantile': prepare_patch_dt_target_samples( list(new_patches.values()), - cfg['patch_dt_target_num_points'], cfg['dt_target_max_stride'], + cfg['sample_count_patch_dt_target_points'], cfg['dt_target_max_stride'], ) # ---- Point collections: link, classify, strip-materialise ---- @@ -2371,7 +2309,7 @@ def incorporate_interactive_inputs(records, influence_config=None): pcl['points_by_patch'] = points_by_patch cross_patch_pcls.append(pcl) - min_point_spacing = cfg['unattached_pcl_min_point_spacing'] + min_point_spacing = cfg['pcl_unattached_pcl_min_point_spacing'] for pcl_id, pcl in new_unattached.items(): sorted_items = sorted(pcl['points'].items(), key=lambda kv: int(kv[0])) if len(sorted_items) < 2: @@ -2410,7 +2348,7 @@ def incorporate_interactive_inputs(records, influence_config=None): # pools; force recomputation on next use. dt_target_cache_manager.reset() - if run_cfg['interactive_influence_enabled'] and (new_patches or new_collections): + if run_cfg['influence_enabled'] and (new_patches or new_collections): influence_state = make_influence_state(run_cfg, torch.device('cuda')) influence_state.activate_or_extend_( new_patches=new_patches, @@ -2424,7 +2362,7 @@ def incorporate_interactive_inputs(records, influence_config=None): ) interactive_influence_loss_weight = float(run_cfg['loss_weight_anchor']) interactive_influence_anchor_samples = int( - run_cfg['interactive_influence_anchor_samples_per_step']) + run_cfg['sample_count_influence_anchor_samples_per_step']) # run() sets the target before this callback is drained at the # pause boundary, so this is exactly the iteration window requested @@ -2434,7 +2372,7 @@ def incorporate_interactive_inputs(records, influence_config=None): dt_resume_iteration = get_interactive_dt_resume_iteration( interactive_status['current_iteration'], interactive_status['target_iteration'], - run_cfg['interactive_influence_disable_dt_frac'], + run_cfg['influence_disable_dt_frac'], ) interactive_dt_resume_iteration = dt_resume_iteration @@ -2447,20 +2385,200 @@ def incorporate_interactive_inputs(records, influence_config=None): torch.cuda.set_rng_state_all(cuda_states) if interactive_driver is not None: - def configure_interactive_run(config): - # Only called by the resident driver on the fitter thread at a - # pause boundary. These settings are read afresh by every step. - if ({'track_length_bin_weights', 'max_track_crossing_per_step', - 'min_walk_steps_per_track', 'max_walk_steps_per_track', - 'n_walks_per_track'} - & set(config)): - configure_prepared_track_sampling(prepared_main_tracks, config) + def configure_interactive_run(config, path_changes=None): + """Apply Run-scoped settings without replacing the resident fit.""" + global shell_path + nonlocal dt_target_whole_object, prepared_main_tracks + nonlocal lr_scheduler, num_training_steps + nonlocal patch_sampling_probabilities + nonlocal unverified_patch_sampling_probabilities + nonlocal unverified_patches, unverified_patches_list + nonlocal unverified_patch_atlas + nonlocal shell_patch, shell_map, shell_envelope + nonlocal shell_outer_winding_idx, shell_valid_zyxs_gpu + nonlocal tracks, track_families, track_source_ids + nonlocal preview_extent_tracks + + path_changes = dict(path_changes or {}) + changed = set(config) + old_values = {key: cfg[key] for key in config} cfg.update(config, allow_val_change=True) + try: + shell_changed = ( + bool(changed & { + key for key in cfg.keys() + if str(key).startswith('shell_') + }) + or 'outer_shell' in path_changes + ) + rebuilt_tracks = None + replace_prepared_tracks = False + rebuilt_track_rows = tracks + rebuilt_families = track_families + rebuilt_source_ids = track_source_ids + rebuilt_shell_patch = shell_patch + rebuilt_shell_map = shell_map + rebuilt_shell_envelope = shell_envelope + rebuilt_shell_outer = shell_outer_winding_idx + rebuilt_shell_valid = shell_valid_zyxs_gpu + requested_shell_path = str( + path_changes.get('outer_shell', shell_path) or '') + + if shell_changed: + if not requested_shell_path: + raise ValueError( + 'shell configuration requires an outer shell path') + rebuilt_shell_patch = load_tifxyz(requested_shell_path) + rebuilt_shell_envelope = ( + ShellPolarMap( + rebuilt_shell_patch, umbilicus, + z_min=z_begin - cfg['model_flow_bounds_z_margin'], + z_max=z_end + cfg['model_flow_bounds_z_margin'], + num_theta_bins=cfg['shell_num_theta_bins'], + device='cpu') + if filter_tracks_by_shell else None + ) + if filter_tracks_by_shell and track_reload_source is not None: + (rebuilt_track_rows, rebuilt_families, + kept_track_indices) = filter_tracks_to_outer_shell( + track_reload_source, rebuilt_shell_envelope, + track_reload_families, return_indices=True) + rebuilt_source_ids = ( + track_reload_source_ids[kept_track_indices] + if track_reload_source_ids is not None else None) + rebuilt_tracks = prepare_main_phase_tracks( + rebuilt_track_rows, None, + float(cfg['track_exclusion_radius']), device, + anchor_tree=trusted_geometry_tree, + sampling_config=validate_track_sampling_config(cfg), + track_families=rebuilt_families, + track_source_ids=rebuilt_source_ids, + crossing_cache=track_crossing_cache, + track_graph=track_graph) + replace_prepared_tracks = True + + rebuilt_shell_map = ( + ShellPolarMap( + rebuilt_shell_patch, umbilicus, + z_min=z_begin - cfg['model_flow_bounds_z_margin'], + z_max=z_end + cfg['model_flow_bounds_z_margin'], + num_theta_bins=cfg['shell_num_theta_bins'], + device=device) + if cfg['loss_weight_shell_outer'] > 0 else None + ) + rebuilt_shell_valid = ( + subsample_shell_radius_pool(rebuilt_shell_patch) + if cfg['loss_weight_shell_patch_radius'] > 0 else None + ) + rebuilt_shell_outer = int(cfg['shell_outer_winding_idx']) + + reprepare_tracks = bool(changed & { + 'track_max_tortuosity', + 'track_exclusion_radius', + }) + if reprepare_tracks and rebuilt_tracks is None and tracks: + rebuilt_tracks = prepare_main_phase_tracks( + tracks, None, float(cfg['track_exclusion_radius']), + device, anchor_tree=trusted_geometry_tree, + sampling_config=validate_track_sampling_config(cfg), + track_families=track_families, + track_source_ids=track_source_ids, + crossing_cache=track_crossing_cache, + track_graph=track_graph) + replace_prepared_tracks = True + + target_tracks = ( + rebuilt_tracks + if replace_prepared_tracks else prepared_main_tracks) + if ({'track_length_bin_weights', + 'track_max_track_crossing_per_step', + 'track_min_walk_steps_per_track', + 'track_max_walk_steps_per_track', + 'track_min_walks_per_track', + 'track_max_walks_per_track', + 'track_walk_minimum_cycle_travel'} + & changed): + configure_prepared_track_sampling(target_tracks, config) + + if 'patch_loss_z_margin' in changed: + patch_sampling_probabilities = \ + prepare_patch_sampling_cache(verified_patches_list) + if unverified_patches_list: + unverified_patch_sampling_probabilities = \ + prepare_patch_sampling_cache( + unverified_patches_list) + if 'patch_unverified_patch_exclusion_radius' in changed: + (rebuilt_unverified, rebuilt_unverified_list, + rebuilt_unverified_probabilities, + rebuilt_unverified_atlas) = \ + rebuild_unverified_patch_inputs(float( + cfg['patch_unverified_patch_exclusion_radius'])) + else: + rebuilt_unverified = unverified_patches + rebuilt_unverified_list = unverified_patches_list + rebuilt_unverified_probabilities = \ + unverified_patch_sampling_probabilities + rebuilt_unverified_atlas = unverified_patch_atlas + + dt_preparation_changed = bool(changed & { + 'dt_target_mode', 'dt_target_max_stride', + 'sample_count_patch_dt_target_points', + }) + if dt_preparation_changed: + dt_target_whole_object = ( + cfg['dt_target_mode'] == 'whole_object_quantile') + if dt_target_whole_object: + prepare_patch_dt_target_samples( + verified_patches_list, + cfg['sample_count_patch_dt_target_points'], + cfg['dt_target_max_stride']) + if unverified_patches_list: + prepare_patch_dt_target_samples( + unverified_patches_list, + cfg['sample_count_patch_dt_target_points'], + cfg['dt_target_max_stride']) + if any(key.startswith('dt_') for key in changed) \ + or dt_preparation_changed: + dt_target_cache_manager.update_interval = max( + 1, int(cfg['dt_target_update_interval'])) + dt_target_cache_manager.reset() + if changed & { + 'optimizer_exp_lr_schedule', + 'optimizer_learning_rate', + 'optimizer_lr_final_factor', + 'optimizer_num_training_steps', + }: + realign_lr_schedule( + interactive_driver.status()['current_iteration']) + except Exception: + cfg.update(old_values, allow_val_change=True) + raise + + if shell_changed: + shell_path = requested_shell_path + shell_patch = rebuilt_shell_patch + shell_map = rebuilt_shell_map + shell_envelope = rebuilt_shell_envelope + shell_outer_winding_idx = rebuilt_shell_outer + shell_valid_zyxs_gpu = rebuilt_shell_valid + tracks = rebuilt_track_rows + track_families = rebuilt_families + track_source_ids = rebuilt_source_ids + if replace_prepared_tracks: + prepared_main_tracks = rebuilt_tracks + preview_extent_tracks = ( + (prepared_main_tracks['flat_zyx_cpu'],) + if prepared_main_tracks is not None else ()) + unverified_patches = rebuilt_unverified + unverified_patches_list = rebuilt_unverified_list + unverified_patch_sampling_probabilities = \ + rebuilt_unverified_probabilities + unverified_patch_atlas = rebuilt_unverified_atlas # In the usual zero-exclusion case preview bounds reuse the prepared # flat tensor, so the original list of per-track arrays is no longer # needed after setup. - if preview_extent_tracks is not tracks: + if preview_extent_tracks is not tracks and track_reload_source is None: tracks = None interactive_driver.on_ready( completed_iterations=start_iteration, @@ -2480,20 +2598,37 @@ def configure_interactive_run(config): if interactive_driver is not None else range(start_iteration, num_training_steps) ) - for iteration in tqdm(iteration_sequence, disable=not is_main_process()): + if interactive_driver is None: + progress.begin( + 'optimizing', 'Optimizing', + step=0, total_steps=max(0, num_training_steps - start_iteration), + unit='iterations') + for iteration in tqdm( + iteration_sequence, + disable=not is_main_process() or has_progress): if interactive_driver is not None and not interactive_driver.wait_for_iteration(iteration): break step_timer.start('fwd') - flow_field_high_res_lr_scale = get_flow_field_high_res_lr_scale(iteration) - for flow_field in spiral_and_transform.flow_fields: - flow_field.flow_scales[1] = flow_field_high_res_lr_scale - - slice_to_spiral_transform = spiral_and_transform.get_slice_to_spiral_transform() - dr_per_winding = spiral_and_transform.get_dr_per_winding() + flow_field_high_res_lr_scale = apply_high_res_lr_scale(iteration) + + # The tiny graph paths shared by every transform evaluation this + # iteration (dr softplus, scaled linear logits, pinned gap logits) are + # cut at detached leaves. Each loss family's backward then owns its + # whole graph, so autograd can free the family's buffers as the + # backward pass consumes them instead of retaining the full graph + # (retain_graph) until the family is released. The leaf gradients + # accumulated across families flow through the real shared paths once, + # next to the flow-field gradient flush below. + shared_transform_outputs = spiral_and_transform.get_shared_transform_tensors() + shared_transform_leaves = tuple( + output.detach().requires_grad_(True) for output in shared_transform_outputs) + slice_to_spiral_transform = spiral_and_transform.get_slice_to_spiral_transform( + shared=shared_transform_leaves) + dr_per_winding = shared_transform_leaves[0] losses = {} log_metrics = { - 'flow_field_high_res_lr_scale': spiral_and_transform.flow_field.flow_scales[1], + 'flow_field_high_res_lr_scale': flow_field_high_res_lr_scale, } def backward_family(weighted_losses): @@ -2502,10 +2637,11 @@ def backward_family(weighted_losses): if family_loss.requires_grad: step_timer.stop('fwd') step_timer.start('bwd') - # dr_per_winding and the transform's scaled linear logits are shared - # by later families. retain_graph keeps those tiny common paths valid; - # the family-specific graph is released when this function returns. - family_loss.backward(retain_graph=True) + # The paths shared with later families end at detached leaves + # (shared_transform_leaves and the flow fields' internal + # accumulators), so this family's graph is self-contained and + # its buffers are freed as the backward pass consumes them. + family_loss.backward() step_timer.stop('bwd') step_timer.start('fwd') for name, value in weighted_losses.items(): @@ -2523,10 +2659,8 @@ def backward_family(weighted_losses): unverified_patch_dt_start = cfg['loss_start_patch_dt'] if cfg['loss_start_unverified_patch_dt'] is None else cfg['loss_start_unverified_patch_dt'] compute_unverified_patch_dt = not interactive_dt_suppressed and iteration > unverified_patch_dt_start - # Progressive-outward DT gating: winding cutoff that grows from the respective DT start - # step. None => no gating. (The local is now resolved from the config - # with or without shell losses, so the old fallback to the raw config - # value is redundant.) + # Progressive-outward DT gating: winding cutoff that grows from the + # respective DT start step. None means no gating. dt_progressive_outer = shell_outer_winding_idx patch_dt_max_winding = get_progressive_dt_max_winding(iteration, cfg['loss_start_patch_dt'], dt_progressive_outer) track_dt_max_winding = get_progressive_dt_max_winding(iteration, track_dt_start, dt_progressive_outer) @@ -2559,7 +2693,7 @@ def backward_family(weighted_losses): pcl_flat['zyxs'], pcl_flat['starts'], windings=pcl_flat['windings'], floating_threshold=cfg['dt_target_floating_threshold'], - num_points_per_strip=cfg['dt_target_num_points_per_strip'], + num_points_per_strip=cfg['sample_count_dt_target_points_per_strip'], max_stride=cfg['dt_target_max_stride'], max_total_points=20_000_000, )) @@ -2569,7 +2703,7 @@ def backward_family(weighted_losses): prepared_main_tracks['flat_zyx_cpu'], prepared_main_tracks['offsets'], windings=None, floating_threshold=cfg['dt_target_floating_threshold'], - num_points_per_strip=cfg['dt_target_num_points_per_strip'], + num_points_per_strip=cfg['sample_count_dt_target_points_per_strip'], max_stride=cfg['dt_target_max_stride'], max_total_points=20_000_000, )) @@ -2577,8 +2711,8 @@ def backward_family(weighted_losses): patch_loss_values = get_patch_and_umbilicus_losses( slice_to_spiral_transform, dr_per_winding, - cfg['num_patches_per_step'], - cfg['num_patches_per_step_for_dt'], + cfg['sample_count_patches_per_step'], + cfg['sample_count_patches_per_step_for_dt'], verified_patches_list, patch_atlas, patch_sampling_probabilities, @@ -2606,8 +2740,8 @@ def backward_family(weighted_losses): unverified_loss_values = get_unverified_patch_losses( slice_to_spiral_transform, dr_per_winding, - cfg['unverified_num_patches_per_step'], - cfg['unverified_num_patches_per_step_for_dt'], + cfg['sample_count_unverified_patches_per_step'], + cfg['sample_count_unverified_patches_per_step_for_dt'], unverified_patches_list, unverified_patch_atlas, unverified_patch_sampling_probabilities, @@ -2627,7 +2761,7 @@ def backward_family(weighted_losses): slice_to_spiral_transform, dr_per_winding, shell_outer_winding_idx, - cfg['regularisation_num_points'], + cfg['sample_count_regularisation_points'], ) * cfg['loss_weight_sym_dirichlet'], }) @@ -2663,7 +2797,7 @@ def backward_family(weighted_losses): dr_per_winding, lasagna_volume, shell_outer_winding_idx, - cfg['dense_normals_num_points'], + cfg['sample_count_dense_normal_points'], compute_spacing=grad_mag_spacing_enabled, ): weight = ( @@ -2754,8 +2888,8 @@ def backward_family(weighted_losses): unattached_pcl_strips, pcl_sampling_strata['unattached'], get_or_build_unattached_pcl_flat, - cfg['unattached_pcl_num_per_step'], - cfg['unattached_pcl_num_points_per_step'], + cfg['sample_count_unattached_pcls_per_step'], + cfg['sample_count_unattached_pcl_points_per_step'], compute_dt=compute_patch_dt, dt_max_winding=patch_dt_max_winding, dt_target_cache=unattached_pcl_dt_target_cache, @@ -2818,6 +2952,18 @@ def backward_family(weighted_losses): apply_accumulated_field_grad = getattr(flow_field, 'apply_accumulated_field_grad', None) if apply_accumulated_field_grad is not None: apply_accumulated_field_grad() + # Propagate the leaf gradients the family backwards accumulated on the + # shared transform paths through the real parameters, exactly once. + shared_transform_pending = [ + (output, leaf.grad) + for output, leaf in zip(shared_transform_outputs, shared_transform_leaves) + if output.requires_grad and leaf.grad is not None + ] + if shared_transform_pending: + torch.autograd.backward( + [output for output, _ in shared_transform_pending], + [grad for _, grad in shared_transform_pending], + ) step_timer.stop('bwd') step_timer.start('comm') allreduce_grads_(dist_grad_params) @@ -2826,7 +2972,11 @@ def backward_family(weighted_losses): step_had_nonfinite = torch.zeros((), dtype=torch.bool, device=nonfinite_grad_steps.device) for name, p in dist_grad_named: if p.grad is not None: - param_nonfinite = (~torch.isfinite(p.grad)).any() + # aminmax propagates NaN and surfaces +/-inf through two scalar + # reductions, avoiding the gradient-sized boolean temporaries + # that (~torch.isfinite(grad)).any() allocates per parameter. + grad_min, grad_max = torch.aminmax(p.grad) + param_nonfinite = ~(torch.isfinite(grad_min) & torch.isfinite(grad_max)) step_had_nonfinite |= param_nonfinite nonfinite_grad_by_param[name] += param_nonfinite.to(nonfinite_grad_steps.dtype) torch.nan_to_num_(p.grad, nan=0.0, posinf=0.0, neginf=0.0) @@ -2857,6 +3007,8 @@ def backward_family(weighted_losses): learning_rate=float(optimiser.param_groups[0]['lr']), metrics={name: float(value) for name, value in log_metrics.items()}, ) + else: + progress.update(iteration - start_iteration + 1) if iteration % 200 == 0: # Only sync to CPU and log when we actually print, avoiding a per-iter @@ -2890,8 +3042,13 @@ def backward_family(weighted_losses): suffix = 'fitted' if is_main_process(): + progress.begin( + 'saving_checkpoint', 'Saving final checkpoint', + detail=f'checkpoint_{suffix}.ckpt') save_model(suffix, num_training_steps) - if cfg.get('save_png_visualizations', False): + if cfg.get('output_save_png_visualizations', False): + progress.begin( + 'finalizing', 'Preparing final visualizations') ( zs_for_visualisation, slice_yx, @@ -2905,6 +3062,8 @@ def backward_family(weighted_losses): scroll_slices_for_visualisation = None prediction_slices_for_visualisation = None quad_label_map = None + progress.begin( + 'finalizing', 'Computing satisfaction metrics and outputs') save_overlay_and_print_satisfaction( suffix, spiral_and_transform=spiral_and_transform, @@ -2933,31 +3092,38 @@ def backward_family(weighted_losses): voxel_size_um=voxel_size_um, get_or_build_unattached_pcl_flat=get_or_build_unattached_pcl_flat, run_tag=run_tag, - save_png_visualizations=cfg.get('save_png_visualizations', False), + save_png_visualizations=cfg.get('output_save_png_visualizations', False), + progress=progress, ) + progress.finish() + progress.clear() if __name__ == '__main__': + cli_progress = ( + ProgressReporter(stream=sys.stderr) + if int(os.environ.get('RANK', '0')) == 0 else None + ) maybe_init_distributed() try: - config = dict(default_config) + config = Config().as_dict() config.update(get_env_config_overrides()) reference_z_range_num_slices = 9500 z_range_scaled_count_keys = ( - 'num_patches_per_step', - 'num_patches_per_step_for_dt', - 'unverified_num_patches_per_step', - 'unverified_num_patches_per_step_for_dt', - 'rel_winding_num_pcls', - 'abs_winding_num_pcls', - 'unattached_pcl_num_per_step', - 'track_num_per_step', - 'dense_normals_num_points', - 'dense_spacing_num_pairs', - 'dense_spacing_density_extra_pairs', - 'dense_attachment_num_points', - 'regularisation_num_points', - 'shell_num_samples', + 'sample_count_patches_per_step', + 'sample_count_patches_per_step_for_dt', + 'sample_count_unverified_patches_per_step', + 'sample_count_unverified_patches_per_step_for_dt', + 'sample_count_relative_winding_pcls', + 'sample_count_absolute_winding_pcls', + 'sample_count_unattached_pcls_per_step', + 'sample_count_tracks_per_step', + 'sample_count_dense_normal_points', + 'sample_count_dense_spacing_pairs', + 'sample_count_dense_spacing_density_extra_pairs', + 'sample_count_dense_attachment_points', + 'sample_count_regularisation_points', + 'sample_count_shell_samples', ) z_range_scale, z_range_num_slices = scale_counts_for_z_range( config, z_begin, z_end, @@ -2982,6 +3148,8 @@ def backward_family(weighted_losses): wandb.init(project='scrolls', config=config, mode=wandb_mode) cfg = wandb.config configure_losses(cfg, z_begin, z_end) - main() + main(progress=cli_progress) finally: + if cli_progress is not None: + cli_progress.close() maybe_destroy_distributed() diff --git a/volume-cartographer/scripts/spiral/flatten_spiral_checkpoint.py b/volume-cartographer/scripts/spiral/flatten_spiral_checkpoint.py new file mode 100755 index 0000000000..70e1082e31 --- /dev/null +++ b/volume-cartographer/scripts/spiral/flatten_spiral_checkpoint.py @@ -0,0 +1,520 @@ +#!/usr/bin/env python3 +"""Export and Lasagna-flatten a fitted Spiral checkpoint. + +The normal invocation needs only the checkpoint and final TIFXYZ directory: + + python flatten_spiral_checkpoint.py checkpoint_fitted.ckpt output.tifxyz + +The script reconstructs a combined Spiral surface, starts a private Lasagna +fit service on an ephemeral localhost port, runs flatten_fast_nofilter.json, +publishes the resulting TIFXYZ atomically, and always stops the service. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +import re +import shutil +import signal +import subprocess +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path +from urllib.parse import quote + +import numpy as np +import torch + +from checkpoint_io import load_checkpoint_cpu +from sample_spiral import get_spiral_yxs +from tifxyz import save_combined_tifxyz +from transforms import SpiralAndTransform +from umbilicus import json_umbilicus_z_to_yx + + +MODEL_CONFIG_KEYS = ( + "num_flow_integration_steps", + "flow_integration_solver", + "num_flow_timesteps", + "num_flow_stages", + "flow_bounds_z_margin", + "flow_bounds_radius", + "flow_voxel_resolution", + "flow_field_type", + "gap_expander_logit_resolution", + "gap_expander_num_windings", + "gap_expander_lr_scale", + "linear_z_resolution", + "initial_dr_per_winding", +) + + +def _checkpoint_config(checkpoint: dict) -> dict: + """Return the model config with current ``model_*`` aliases populated.""" + raw = checkpoint.get("cfg") + if not isinstance(raw, dict): + raise ValueError("checkpoint has no Spiral 'cfg' dictionary") + config = dict(raw) + for suffix in MODEL_CONFIG_KEYS: + current = f"model_{suffix}" + if current not in config and suffix in config: + config[current] = config[suffix] + missing = [f"model_{key}" for key in MODEL_CONFIG_KEYS + if f"model_{key}" not in config] + if missing: + raise ValueError( + "checkpoint is missing model configuration: " + + ", ".join(missing)) + return config + + +def _value(config: dict, name: str, default=None): + if name in config: + return config[name] + current = f"output_{name}" + return config.get(current, default) + + +def _resolve_umbilicus(checkpoint_path: Path, explicit: Path | None) -> Path: + if explicit is not None: + candidates = [explicit.expanduser()] + else: + candidates = [] + dataset_env = os.environ.get("SPIRAL_DATASET") + if dataset_env: + candidates.append(Path(dataset_env).expanduser() / "umbilicus.json") + for parent in (checkpoint_path.parent, *checkpoint_path.parents): + candidates.append(parent / "umbilicus.json") + # Common local layouts used by the legacy CLI and current s1 runs. + candidates.extend([ + Path.home() / "Desktop/spiral_dataset/to_hf/umbilicus.json", + Path("/ephemeral/paul/spiral/dataset/umbilicus.json"), + ]) + for candidate in candidates: + if candidate.is_file(): + return candidate.resolve() + raise FileNotFoundError( + "could not locate umbilicus.json; pass --umbilicus PATH or set " + "SPIRAL_DATASET to the dataset root") + + +def _resolve_lasagna(explicit: Path | None) -> tuple[Path, Path]: + candidates = [explicit.expanduser()] if explicit is not None else [] + configured = os.environ.get("LASAGNA_SERVICE_PATH") + if configured: + configured_path = Path(configured).expanduser() + candidates.append( + configured_path.parent + if configured_path.name == "fit_service.py" + else configured_path) + here = Path(__file__).resolve() + candidates.extend([ + here.parents[3] / "lasagna", + Path.home() / "villa/lasagna", + ]) + for root in candidates: + service = root / "fit_service.py" + config = root / "configs/flatten_fast_nofilter.json" + if service.is_file() and config.is_file(): + return service.resolve(), config.resolve() + raise FileNotFoundError( + "could not find Lasagna's fit_service.py and " + "configs/flatten_fast_nofilter.json; pass --lasagna-dir PATH") + + +def _build_model( + checkpoint: dict, + config: dict, + umbilicus_path: Path, + device: torch.device, +) -> SpiralAndTransform: + try: + z_begin = int(checkpoint["z_begin"]) + z_end = int(checkpoint["z_end"]) + state = checkpoint["spiral_and_transform"] + except KeyError as exc: + raise ValueError(f"checkpoint is missing {exc.args[0]!r}") from exc + if z_begin >= z_end or not isinstance(state, dict): + raise ValueError("checkpoint has an invalid model z range or state") + + z_values = np.arange(z_begin, z_end) + centre = json_umbilicus_z_to_yx( + str(umbilicus_path), coordinate_scale=1.0) + umbilicus_zyx = torch.from_numpy(np.concatenate([ + z_values[:, None], centre(z_values), + ], axis=-1).astype(np.float32)).to(device) + radius = int(config["model_flow_bounds_radius"]) + margin = int(config["model_flow_bounds_z_margin"]) + flow_min = torch.tensor( + [z_begin - margin, -radius, -radius], + dtype=torch.int64, device=device) + flow_max = torch.tensor( + [z_end + margin, radius, radius], + dtype=torch.int64, device=device) + model = SpiralAndTransform( + flow_integration_steps=int( + config["model_num_flow_integration_steps"]), + flow_integration_solver=str( + config["model_flow_integration_solver"]), + flow_min_corner_zyx=flow_min, + flow_max_corner_zyx=flow_max, + umbilicus_zyx=umbilicus_zyx, + config=config, + spiral_outward_sense=str( + checkpoint.get("spiral_outward_sense") or "CW"), + ).to(device) + model.load_state_dict(state) + model.eval() + return model + + +@torch.inference_mode() +def _export_source_surface( + checkpoint: dict, + config: dict, + umbilicus_path: Path, + destination: Path, + *, + device: torch.device, + voxel_size_um: float, + chunk_size: int, +) -> Path: + model = _build_model(checkpoint, config, umbilicus_path, device) + transform = model.get_slice_to_spiral_transform() + dr_per_winding = model.get_dr_per_winding() + z_begin = int(checkpoint["z_begin"]) + z_end = int(checkpoint["z_end"]) + step = int(_value(config, "step_size", 20)) + if step <= 0: + raise ValueError("checkpoint output_step_size must be positive") + first = int(checkpoint.get( + "preview_first_winding", + _value(config, "first_winding", 10))) + last_raw = config.get("shell_outer_winding_idx") + if last_raw is None: + last_raw = int(config["model_gap_expander_num_windings"]) - 1 + last = int(last_raw) + if last < first: + raise ValueError( + f"invalid output winding range [{first}, {last}]") + + print( + f"[spiral] reconstructing windings {first}..{last}, " + f"z [{z_begin}, {z_end}), step {step}", flush=True) + yxs_by_winding = get_spiral_yxs( + last + 1, dr_per_winding, step, + group_by_winding=True, device=str(device)) + margin = int(config["model_flow_bounds_z_margin"]) + z0 = z_begin - margin + z_values = torch.arange( + z0, z_end + margin, step, dtype=torch.float32, device=device) + winding_grids: dict[int, np.ndarray] = {} + for winding in range(first, last + 1): + yxs = yxs_by_winding[winding] + spiral = torch.cat([ + z_values[:, None, None].expand(-1, yxs.shape[0], 1), + yxs[None, :, :].expand(z_values.shape[0], -1, 2), + ], dim=-1) + flat = spiral.reshape(-1, 3) + pieces = [ + transform.inv(flat[start:start + chunk_size]).cpu() + for start in range(0, flat.shape[0], chunk_size) + ] + scroll = torch.cat(pieces).reshape_as(spiral) + outside = ( + (scroll[..., 0] < z_begin) | (scroll[..., 0] >= z_end)) + scroll[outside] = -1.0 + winding_grids[winding] = scroll.numpy().astype( + np.float32, copy=False) + print( + f"[spiral] reconstructed winding {winding} " + f"({winding - first + 1}/{last - first + 1})", + flush=True) + + surface_id = "spiral-checkpoint.tifxyz" + manifest = save_combined_tifxyz( + winding_grids, + destination, + surface_id, + step, + voxel_size_um, + source="flatten_spiral_checkpoint.py", + first_winding=first, + cleanup_erosion_cells=3, + ) + del transform, model, dr_per_winding, winding_grids + gc.collect() + if device.type == "cuda": + torch.cuda.empty_cache() + return Path(manifest["surface_path"]).resolve() + + +def _md5_file(path: Path) -> str: + digest = hashlib.md5(usedforsecurity=False) + with path.open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(block) + return f"md5:{digest.hexdigest()}" + + +def _store_surface(surface: Path, object_store: Path) -> dict: + lines = [] + for path in sorted(item for item in surface.rglob("*") if item.is_file()): + lines.append( + f"{path.relative_to(surface).as_posix()}\t{_md5_file(path)}\n") + manifest_hash = hashlib.md5( + "".join(lines).encode(), usedforsecurity=False).hexdigest() + ref = { + "type": "tifxyz_segment", + "name": surface.name, + "hash": f"md5:{manifest_hash}", + } + destination = ( + object_store / ref["type"] / manifest_hash + / quote(ref["name"], safe="")) + destination.mkdir(parents=True) + (destination / "segment").symlink_to(surface, target_is_directory=True) + (destination / "object.json").write_text( + json.dumps(ref, indent=2) + "\n", encoding="utf-8") + return ref + + +def _service_json(port: int, path: str, body=None, timeout=30): + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + f"http://127.0.0.1:{port}{path}", + data=data, + method="GET" if data is None else "POST", + headers={ + "X-Fit-Service-API-Version": "2", + "Content-Type": "application/json", + "X-VC3D-Source": "Spiral checkpoint flattener", + }, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + payload = json.loads(response.read().decode()) + except urllib.error.HTTPError as exc: + try: + payload = json.loads(exc.read().decode()) + except Exception: + payload = {} + raise RuntimeError( + str(payload.get("error") or f"Lasagna HTTP {exc.code}")) from exc + if isinstance(payload, dict) and payload.get("error"): + raise RuntimeError(str(payload["error"])) + return payload + + +def _stop_service(process: subprocess.Popen | None) -> None: + if process is None or process.poll() is not None: + return + try: + if os.name == "posix": + os.killpg(process.pid, signal.SIGTERM) + else: + process.terminate() + process.wait(timeout=10) + except (OSError, subprocess.TimeoutExpired): + if process.poll() is None: + if os.name == "posix": + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + process.wait() + + +def _flatten( + surface: Path, + output: Path, + service_path: Path, + flatten_config_path: Path, + work: Path, + *, + voxel_size_um: float, +) -> None: + object_store = work / "objects" + surface_ref = _store_surface(surface, object_store) + config = json.loads(flatten_config_path.read_text(encoding="utf-8")) + config["external_surfaces"] = [surface_ref] + config["voxel_size_um"] = float(voxel_size_um) + results = work / "results" + results.mkdir() + model_output = work / "flatten-model.pt" + service = None + ready = threading.Event() + port_holder: dict[str, int] = {} + try: + service = subprocess.Popen( + [ + sys.executable, str(service_path), + "--port", "0", + "--allow-no-data-dir", + "--object-store-dir", str(object_store), + ], + cwd=str(service_path.parent), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=(os.name == "posix"), + ) + + def relay_output(): + pattern = re.compile(r"listening on http://[^:]+:(\d+)") + assert service is not None and service.stdout is not None + for line in service.stdout: + text = line.rstrip() + if text: + print(f"[lasagna] {text}", flush=True) + match = pattern.search(text) + if match: + port_holder["port"] = int(match.group(1)) + ready.set() + + relay = threading.Thread(target=relay_output, daemon=True) + relay.start() + if not ready.wait(60): + code = service.poll() + raise RuntimeError( + "Lasagna service failed to start" + + (f" (exit {code})" if code is not None else "")) + port = port_holder["port"] + output_name = "flattened.tifxyz" + request = { + "config": config, + "job_spec": { + "config": config, + "linked_surfaces": [surface_ref], + }, + "single_segment": True, + "config_name": flatten_config_path.name, + "output_name": output_name, + "output_dir": str(results), + "model_output": str(model_output), + "source": "flatten_spiral_checkpoint.py", + } + accepted = _service_json(port, "/optimize", request, timeout=60) + job_id = str(accepted.get("job_id") or "") + if not job_id: + raise RuntimeError("Lasagna service returned no job id") + while True: + status = _service_json(port, f"/jobs/{job_id}", timeout=30) + state = str(status.get("state") or "") + stage = str( + status.get("stage_name") or status.get("stage") or state) + step = int(status.get("step") or 0) + total = int(status.get("total_steps") or 0) + suffix = f" {step}/{total}" if total else "" + print(f"[lasagna] {stage}{suffix}", flush=True) + if state == "finished": + break + if state in {"error", "cancelled"}: + raise RuntimeError( + str(status.get("error") or f"Lasagna job {state}")) + if service.poll() is not None: + raise RuntimeError( + f"Lasagna service exited with code {service.returncode}") + time.sleep(1.0) + + produced = results / output_name + if not (produced / "meta.json").is_file(): + raise RuntimeError( + "Lasagna reported success but produced no TIFXYZ") + os.replace(produced, output) + finally: + _stop_service(service) + + +def _parse_args(argv=None): + parser = argparse.ArgumentParser( + description=( + "Reconstruct a fitted Spiral checkpoint and flatten it through a " + "temporary Lasagna service.")) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument( + "--umbilicus", type=Path, + help="umbilicus.json (auto-detected when omitted)") + parser.add_argument( + "--lasagna-dir", type=Path, + help="Lasagna repository containing fit_service.py (auto-detected)") + parser.add_argument( + "--device", default="cuda", + help="Torch device used to reconstruct the Spiral surface (default: cuda)") + parser.add_argument( + "--voxel-size-um", type=float, default=9.6, + help="source voxel size written to TIFXYZ metadata (default: 9.6)") + parser.add_argument( + "--chunk-size", type=int, default=65536, + help="points transformed per reconstruction batch (default: 65536)") + return parser.parse_args(argv) + + +def main(argv=None) -> int: + args = _parse_args(argv) + checkpoint_path = args.checkpoint.expanduser().resolve(strict=True) + output = args.output.expanduser().resolve() + if output.exists(): + raise FileExistsError( + f"refusing to overwrite existing output: {output}") + output.parent.mkdir(parents=True, exist_ok=True) + if args.voxel_size_um <= 0: + raise ValueError("--voxel-size-um must be positive") + if args.chunk_size <= 0: + raise ValueError("--chunk-size must be positive") + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA is unavailable; pass --device cpu if intentional") + + umbilicus = _resolve_umbilicus(checkpoint_path, args.umbilicus) + service, flatten_config = _resolve_lasagna(args.lasagna_dir) + print(f"[spiral] loading {checkpoint_path}", flush=True) + checkpoint = load_checkpoint_cpu(str(checkpoint_path)) + if not isinstance(checkpoint, dict): + raise ValueError("checkpoint root must be a dictionary") + config = _checkpoint_config(checkpoint) + print(f"[spiral] using umbilicus {umbilicus}", flush=True) + print(f"[lasagna] using config {flatten_config}", flush=True) + + work_path = tempfile.mkdtemp( + prefix=f".{output.name}.flatten-", dir=output.parent) + work = Path(work_path) + try: + surface = _export_source_surface( + checkpoint, + config, + umbilicus, + work / "source", + device=device, + voxel_size_um=float(args.voxel_size_um), + chunk_size=int(args.chunk_size), + ) + del checkpoint + gc.collect() + if device.type == "cuda": + torch.cuda.empty_cache() + _flatten( + surface, + output, + service, + flatten_config, + work, + voxel_size_um=float(args.voxel_size_um), + ) + finally: + shutil.rmtree(work, ignore_errors=True) + print(f"wrote {output}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/volume-cartographer/scripts/spiral/flow_fields.py b/volume-cartographer/scripts/spiral/flow_fields.py index 16615c2b57..b767d7c85b 100644 --- a/volume-cartographer/scripts/spiral/flow_fields.py +++ b/volume-cartographer/scripts/spiral/flow_fields.py @@ -264,10 +264,10 @@ def stage_grad(g, pts): class CartesianFlowField(nn.Module): - def __init__(self, resolution, spatial_scale_factor=6, lr_scale_factor=1.e-1, num_flow_timesteps=1): + def __init__(self, resolution, spatial_scale_factor=6, num_flow_timesteps=1, direct_lr=False): super().__init__() self.num_flow_timesteps = num_flow_timesteps - self.flow_scales = [1., lr_scale_factor] + self.direct_lr = direct_lr self.flows = nn.ParameterList([ nn.Parameter(torch.zeros([num_flow_timesteps, 3, *shape])) for shape in [ @@ -276,22 +276,21 @@ def __init__(self, resolution, spatial_scale_factor=6, lr_scale_factor=1.e-1, nu ] ]) self._pending_lr_upsampled = None - self._pending_hr_scale = None self._field_grad_acc = None self._lr_grad_acc = None self._pending_direct = False def _prepare_time_invariant_fields(self): # Shared state for get_sampler / get_time_invariant_integrator when - # num_flow_timesteps == 1. Returns (low_field, high_field, high_scale, - # acc); acc is None outside training (no grad, or frozen fields). + # num_flow_timesteps == 1. Returns (low_field, high_field, acc); acc is + # None outside training (no grad, or frozen fields). lr_flow, hr_flow = self.flows[0], self.flows[1] hr_shape = tuple(hr_flow.shape[2:]) # Time-invariant: HR flow is already at the target resolution, so skip interpolating it. - lr_upsampled = F.interpolate(lr_flow, size=hr_shape, mode='trilinear')[0] * self.flow_scales[0] + lr_upsampled = F.interpolate(lr_flow, size=hr_shape, mode='trilinear')[0] training_field = torch.is_grad_enabled() and (lr_flow.requires_grad or hr_flow.requires_grad) if not training_field: - return lr_upsampled, hr_flow[0], float(self.flow_scales[1]), None + return lr_upsampled, hr_flow[0], None # Sampling keeps the detached LR-upsampled and HR fields separate and # accumulates dL/d(their effective sum) in one caller-owned buffer. @@ -305,14 +304,12 @@ def _prepare_time_invariant_fields(self): else: self._field_grad_acc.zero_() self._pending_lr_upsampled = lr_upsampled - self._pending_hr_scale = float(self.flow_scales[1]) - return low_field, high_field, self._pending_hr_scale, self._field_grad_acc + return low_field, high_field, self._field_grad_acc def _get_direct_integrator(self): # Direct-LR mode: sample the LR lattice at query points instead of # upsampling it to HR every step (FIT_SPIRAL_DIRECT_LR=1). Field - # gradients accumulate unscaled into per-lattice buffers; the value - # scales are applied in apply_accumulated_field_grad. + # gradients accumulate into per-lattice buffers. lr_flow, hr_flow = self.flows[0], self.flows[1] low, high = lr_flow[0], hr_flow[0] training_field = torch.is_grad_enabled() and ( @@ -334,12 +331,11 @@ def _get_direct_integrator(self): # accumulator after the direct branch consumed it. assert self._pending_lr_upsampled is None self._pending_direct = True - lo_scale, hi_scale = float(self.flow_scales[0]), float(self.flow_scales[1]) low_c, high_c = low.contiguous(), high.contiguous() def integrate(y_flat, h, n_steps): return flow_triton.rk4_direct_integrate( - y_flat, low_c, high_c, lo_scale, hi_scale, acc_lo, acc_hi, h, n_steps) + y_flat, low_c, high_c, 1.0, 1.0, acc_lo, acc_hi, h, n_steps) return integrate @@ -348,17 +344,17 @@ def get_time_invariant_integrator(self): # integration as one autograd node (see _RK4SparseFlowIntegrate). Only # valid for num_flow_timesteps == 1. assert self.num_flow_timesteps == 1 - if (flow_triton.direct_lr_enabled() + if (flow_triton.direct_lr_enabled(self.direct_lr) and flow_triton.rk4_triton_available(self.flows[0], self.flows[1])): return self._get_direct_integrator() - low_field, high_field, high_scale, acc = self._prepare_time_invariant_fields() + low_field, high_field, acc = self._prepare_time_invariant_fields() if flow_triton.rk4_triton_available(low_field, high_field, acc): low_c, high_c = low_field.contiguous(), high_field.contiguous() def integrate(y_flat, h, n_steps): return flow_triton.rk4_integrate( - y_flat, low_c, high_c, high_scale, acc, h, n_steps, + y_flat, low_c, high_c, 1.0, acc, h, n_steps, ) return integrate @@ -370,17 +366,17 @@ def integrate(y_flat, h, n_steps): # (matters for large no-grad evals like previews/satisfaction). y = y_flat for _ in range(n_steps): - k1 = sample_field(y, low_field) + sample_field(y, high_field) * high_scale + k1 = sample_field(y, low_field) + sample_field(y, high_field) x2 = y + (h / 2) * k1 - k2 = sample_field(x2, low_field) + sample_field(x2, high_field) * high_scale + k2 = sample_field(x2, low_field) + sample_field(x2, high_field) x3 = y + (h / 2) * k2 - k3 = sample_field(x3, low_field) + sample_field(x3, high_field) * high_scale + k3 = sample_field(x3, low_field) + sample_field(x3, high_field) x4 = y + h * k3 - k4 = sample_field(x4, low_field) + sample_field(x4, high_field) * high_scale + k4 = sample_field(x4, low_field) + sample_field(x4, high_field) y = y + (h / 6) * (k1 + 2 * k2 + 2 * k3 + k4) return y return _RK4SparseFlowIntegrate.apply( - y_flat, low_field, high_field, high_scale, acc, h, n_steps, + y_flat, low_field, high_field, 1.0, acc, h, n_steps, ) return integrate @@ -392,11 +388,11 @@ def get_sampler(self, t): lr_flow, hr_flow = self.flows[0], self.flows[1] hr_shape = tuple(hr_flow.shape[2:]) if self.num_flow_timesteps == 1: - low_field, high_field, high_scale, acc = self._prepare_time_invariant_fields() + low_field, high_field, acc = self._prepare_time_invariant_fields() if acc is None: return lambda y: ( sample_field(y, low_field) - + sample_field(y, high_field) * high_scale + + sample_field(y, high_field) ) def sample(normalised_zyx): @@ -405,7 +401,7 @@ def sample(normalised_zyx): flat, low_field, high_field, - high_scale, + 1.0, acc, ).view(*normalised_zyx.shape[:-1], 3) @@ -414,8 +410,8 @@ def sample(normalised_zyx): t_scaled = (t.clamp(-1. + 1.e-4, 1. - 1.e-4) + 1) / 2 * (self.num_flow_timesteps - 1) t_idx_before = int(t_scaled) flows_interpolated = [ - F.interpolate(flow[t_idx_before : t_idx_before + 2], size=hr_shape, mode='trilinear') * flow_scale - for flow, flow_scale in zip(self.flows, self.flow_scales) + F.interpolate(flow[t_idx_before : t_idx_before + 2], size=hr_shape, mode='trilinear') + for flow in self.flows ] field = sum( torch.lerp(flow_interpolated[0], flow_interpolated[1], t_scaled % 1.) @@ -426,12 +422,11 @@ def sample(normalised_zyx): def apply_accumulated_field_grad(self): if self._pending_direct: self._pending_direct = False - for param, acc, scale in ( - (self.flows[0], self._lr_grad_acc, self.flow_scales[0]), - (self.flows[1], self._field_grad_acc, self.flow_scales[1]), + for param, acc in ( + (self.flows[0], self._lr_grad_acc), + (self.flows[1], self._field_grad_acc), ): # Reuse the accumulator storage as the parameter gradient. - acc.mul_(scale) grad = acc.unsqueeze(0) if param.grad is None: param.grad = grad @@ -446,10 +441,8 @@ def apply_accumulated_field_grad(self): self._pending_lr_upsampled.backward(gradient=self._field_grad_acc) self._pending_lr_upsampled = None - # dL/dHR = scale * dL/dfield. Reuse the accumulator storage as the - # parameter's gradient instead of materialising another [3,Z,Y,X] tensor. - self._field_grad_acc.mul_(self._pending_hr_scale) - self._pending_hr_scale = None + # Reuse the accumulator storage as the HR parameter's gradient instead + # of materialising another [3,Z,Y,X] tensor. hr_param = self.flows[1] hr_grad = self._field_grad_acc.unsqueeze(0) if hr_param.grad is None: @@ -477,9 +470,11 @@ class CylindricalFlowField(nn.Module): # Note: near r=0 the cylindrical basis is degenerate; ring 0 holds a single cell that is # pinned to zero. - def __init__(self, resolution, spatial_scale_factor=6, lr_scale_factor=1.e-1, num_flow_timesteps=1): + def __init__(self, resolution, spatial_scale_factor=6, num_flow_timesteps=1, direct_lr=False): # resolution is interpreted as the equivalent cartesian (Z, Y, X) voxel shape; the - # cylindrical lattice sizes are derived from it. + # cylindrical lattice sizes are derived from it. direct_lr is accepted for + # constructor parity with CartesianFlowField and ignored: the ragged + # cylindrical lattice is always sampled directly (never upsampled). super().__init__() self.num_flow_timesteps = num_flow_timesteps Z, Y, X = (int(s) for s in resolution) @@ -507,11 +502,13 @@ def compute_num_phi(nr): self.register_buffer('_lr_offsets', lr_offsets) self.register_buffer('_hr_offsets', hr_offsets) - self.flow_scales = [1., lr_scale_factor] self.flows = nn.ParameterList([ nn.Parameter(torch.zeros([num_flow_timesteps, 3, nz_lr, int(lr_offsets[-1])])), nn.Parameter(torch.zeros([num_flow_timesteps, 3, nz_hr, int(hr_offsets[-1])])), ]) + # (pinned field graph output, detached leaf) pairs armed by + # get_sampler in training and consumed by apply_accumulated_field_grad. + self._pending_field_graphs = None @staticmethod def _sample_lattice(field, ring_num_phi, ring_offsets, normalised_zyx): @@ -582,7 +579,7 @@ def sample_at_ring(r_idx): def get_sampler(self, t): # Returns a callable mapping normalised zyx points in [0, 1] to flow velocity at time t, # by sampling the cylindrical lattice directly at each query point. The closure captures - # the time-interpolated, scale-applied, axis-pinned LR & HR lattices so those one-time + # the time-interpolated, axis-pinned LR & HR lattices so those one-time # costs amortise across the integrator's sample calls. if self.num_flow_timesteps == 1: lr_field = self.flows[0][0] @@ -597,8 +594,22 @@ def get_sampler(self, t): # constant zero, so no gradient flows to those parameters; they stay zero indefinitely. n0_lr = int(self._lr_num_phi[0]) n0_hr = int(self._hr_num_phi[0]) - lr_field = torch.cat([torch.zeros_like(lr_field[:, :, :n0_lr]), lr_field[:, :, n0_lr:]], dim=2) * self.flow_scales[0] - hr_field = torch.cat([torch.zeros_like(hr_field[:, :, :n0_hr]), hr_field[:, :, n0_hr:]], dim=2) * self.flow_scales[1] + lr_field = torch.cat([torch.zeros_like(lr_field[:, :, :n0_lr]), lr_field[:, :, n0_lr:]], dim=2) + hr_field = torch.cat([torch.zeros_like(hr_field[:, :, :n0_hr]), hr_field[:, :, n0_hr:]], dim=2) + + if self.num_flow_timesteps == 1 and torch.is_grad_enabled() and ( + self.flows[0].requires_grad or self.flows[1].requires_grad): + # The time-invariant sampler is cached by the diffeomorphism and + # shared across every loss family in an iteration. Cut the + # pinned field graphs at detached leaves so each family's + # backward owns its whole graph (no retain_graph); the accumulated + # leaf gradients flow to the parameters when the training loop + # calls apply_accumulated_field_grad. Overwrites any previous + # pending record, matching CartesianFlowField's one-armed-record- + # per-iteration discipline. + leaves = [field.detach().requires_grad_(True) for field in (lr_field, hr_field)] + self._pending_field_graphs = list(zip((lr_field, hr_field), leaves)) + lr_field, hr_field = leaves sample_lattice = self._sample_lattice lr_num_phi = self._lr_num_phi @@ -612,3 +623,14 @@ def sample(normalised_zyx): + sample_lattice(hr_field, hr_num_phi, hr_offsets, normalised_zyx) ) return sample + + def apply_accumulated_field_grad(self): + pending, self._pending_field_graphs = self._pending_field_graphs, None + if not pending: + return + outputs = [output for output, leaf in pending if leaf.grad is not None] + if outputs: + torch.autograd.backward( + outputs, + [leaf.grad for _, leaf in pending if leaf.grad is not None], + ) diff --git a/volume-cartographer/scripts/spiral/flow_triton.py b/volume-cartographer/scripts/spiral/flow_triton.py index c9df3c3932..0ba8f9bd2a 100644 --- a/volume-cartographer/scripts/spiral/flow_triton.py +++ b/volume-cartographer/scripts/spiral/flow_triton.py @@ -38,14 +38,19 @@ def rk4_triton_available(*tensors): ) -def direct_lr_enabled(): +def direct_lr_enabled(default=False): # Sample the LR flow lattice directly at query points instead of # upsampling it to HR resolution every step. Same function except within # HR cells that straddle an LR grid plane (the upsample linearises there); # kills the F.interpolate forward+backward and its full-resolution # transient. Off by default because it slightly changes the effective - # interpolant. - return os.environ.get('FIT_SPIRAL_DIRECT_LR', '0') == '1' + # interpolant. Normally set via the model_flow_field_direct_lr config key + # (threaded through as `default`); FIT_SPIRAL_DIRECT_LR, when set, + # overrides the config. + value = os.environ.get('FIT_SPIRAL_DIRECT_LR') + if value is None: + return bool(default) + return value == '1' _PERM_STRIDE = 4096 diff --git a/volume-cartographer/scripts/spiral/influence.py b/volume-cartographer/scripts/spiral/influence.py index eb80d964fc..147a4ebd51 100644 --- a/volume-cartographer/scripts/spiral/influence.py +++ b/volume-cartographer/scripts/spiral/influence.py @@ -152,7 +152,7 @@ def active(self): def _generator(self, cfg): generator = torch.Generator() - generator.manual_seed(int(cfg['random_seed']) + self.num_incorporations) + generator.manual_seed(int(cfg['optimizer_random_seed']) + self.num_incorporations) return generator @torch.no_grad() @@ -160,7 +160,7 @@ def _build_footprint(self, new_patches, new_collections, slice_to_spiral_transfo dr_per_winding, cfg, z_begin, z_end, generator): """Per-input scroll-space point sets -> pooled spiral (z, s, theta).""" per_input = [] - max_points = int(cfg['interactive_influence_footprint_points']) + max_points = int(cfg['sample_count_influence_footprint_points']) for input_id, patch in new_patches.items(): zyx = patch.zyxs[patch.valid_vertex_mask].reshape(-1, 3).to(torch.float32) per_input.append((input_id, 'patch', zyx)) @@ -245,8 +245,8 @@ def _flow_z_displacement_bound(self, spiral_and_transform): for flow_field in spiral_and_transform.flow_fields: lr_flow, hr_flow = flow_field.flows v_max = v_max + ( - lr_flow[:, 0].abs().max() * float(flow_field.flow_scales[0]) - + hr_flow[:, 0].abs().max() * float(flow_field.flow_scales[1]) + lr_flow[:, 0].abs().max() + + hr_flow[:, 0].abs().max() ) z_range = float(spiral_and_transform.flow_max_corner_zyx[0] - spiral_and_transform.flow_min_corner_zyx[0]) @@ -299,9 +299,9 @@ def _build_anchor_bank(self, spiral_and_transform, slice_to_spiral_transform, if anchor_geometry_zyx is not None and anchor_geometry_zyx.shape[0] > 0: geometry = subsample_rows( anchor_geometry_zyx.to(torch.float32).cpu(), - int(cfg['interactive_influence_anchor_geometry_points']), generator) + int(cfg['sample_count_influence_anchor_geometry_points']), generator) parts.append(geometry.to(self.device)) - num_lattice = int(cfg['interactive_influence_anchor_lattice_points']) + num_lattice = int(cfg['sample_count_influence_anchor_lattice_points']) outer_winding = int(cfg['shell_outer_winding_idx']) if num_lattice > 0 and outer_winding > 1: winding_indices = torch.arange(1, outer_winding, dtype=torch.float32) @@ -452,13 +452,13 @@ def get_anchor_loss(self, slice_to_spiral_transform, dr_per_winding, num_samples def make_influence_state(cfg, device): limits = ( - float(cfg['interactive_influence_z']), - float(cfg['interactive_influence_windings']), - float(cfg['interactive_influence_theta_frac']) * 2. * math.pi, + float(cfg['influence_z']), + float(cfg['influence_windings']), + float(cfg['influence_theta_frac']) * 2. * math.pi, ) return InteractiveInfluenceState( limits=limits, - sigma=float(cfg['interactive_influence_sigma']), - ramp_power=float(cfg['interactive_influence_anchor_ramp_power']), + sigma=float(cfg['influence_sigma']), + ramp_power=float(cfg['influence_anchor_ramp_power']), device=device, ) diff --git a/volume-cartographer/scripts/spiral/lasagna_data.py b/volume-cartographer/scripts/spiral/lasagna_data.py index a1089c28d4..31ba41ddb4 100644 --- a/volume-cartographer/scripts/spiral/lasagna_data.py +++ b/volume-cartographer/scripts/spiral/lasagna_data.py @@ -1,9 +1,28 @@ +import json import os import numpy as np import torch import zarr +from pack_resident_pools import sidecar_path + + +def _require_sidecar(zarr_path, group, *, pair=False, label): + """Resolve a store's resident-pool sidecar or fail with the pack command.""" + sidecar = sidecar_path(zarr_path, group, pair=pair) + if not os.path.exists(os.path.join(sidecar, 'meta.json')): + raise RuntimeError( + f'{label}: resident-pool sidecar {sidecar!r} not found; build it ' + f'with pack_resident_pools.py (pointed at the lasagna_inputs ' + f'folder, with --ct for CT masking)') + return sidecar + + +def _read_sidecar_meta(sidecar): + with open(os.path.join(sidecar, 'meta.json')) as f: + return json.load(f) + def prepare_lasagna_volume( scroll_zarr, @@ -22,8 +41,13 @@ def prepare_lasagna_volume( yx_bounds_working=None, interior_fn=None, paged_chunk=64, + progress=None, ): - """Open normals/grad-magnitude as bounded sparse CUDA caches.""" + """Open normals/grad-magnitude as fully-resident sparse brick pools. + + The pools load from ``pack_resident_pools.py`` sidecars next to the + source zarrs; there is no other loading path. + """ if not use_normals and not use_spacing: return None if storage_backend != 'sparse_cuda': @@ -38,34 +62,39 @@ def prepare_lasagna_volume( if use_spacing and not grad_mag_zarr_path: raise RuntimeError('dense spacing loss is enabled, but grad_mag zarr path is not set') - print(f'loading lasagna zarrs group {normal_zarr_group}') - nx_array = ny_array = grad_mag_array = None + group = str(normal_zarr_group) + print(f'loading lasagna resident pools, group {group}') + normal_sidecar = grad_sidecar = None reference_shape = None if use_normals: - nx_root = zarr.open(normal_nx_zarr_path, mode='r') - ny_root = zarr.open(normal_ny_zarr_path, mode='r') - nx_array = nx_root[normal_zarr_group] - ny_array = ny_root[normal_zarr_group] - if nx_array.shape != ny_array.shape: - raise ValueError(f'nx/ny normal zarr shapes differ: {nx_array.shape} vs {ny_array.shape}') - if nx_array.dtype != np.dtype('uint8') or ny_array.dtype != np.dtype('uint8'): - raise ValueError( - f'nx/ny normal zarrs must use the production uint8 encoding; ' - f'got {nx_array.dtype} and {ny_array.dtype}') - reference_shape = nx_array.shape + normal_sidecar = _require_sidecar( + normal_nx_zarr_path, group, pair=True, label='lasagna normals') + normal_meta = _read_sidecar_meta(normal_sidecar) + reference_shape = tuple(normal_meta['array_shape']) + expected_pair = [ + os.path.basename(str(normal_nx_zarr_path).rstrip('/')) + '/' + group, + os.path.basename(str(normal_ny_zarr_path).rstrip('/')) + '/' + group, + ] + if normal_meta.get('channel_names') != expected_pair: + print(f'WARNING: normal sidecar channels ' + f'{normal_meta.get("channel_names")} do not match the ' + f'configured nx/ny stores {expected_pair}') if use_spacing: - grad_mag_root = zarr.open(grad_mag_zarr_path, mode='r') - grad_mag_array = grad_mag_root[normal_zarr_group] + grad_sidecar = _require_sidecar( + grad_mag_zarr_path, group, label='lasagna grad_mag') + grad_meta = _read_sidecar_meta(grad_sidecar) if reference_shape is None: - reference_shape = grad_mag_array.shape - elif grad_mag_array.shape != reference_shape: - raise ValueError(f'grad_mag zarr shape {grad_mag_array.shape} differs from dense normal shape {reference_shape}') + reference_shape = tuple(grad_meta['array_shape']) + elif tuple(grad_meta['array_shape']) != reference_shape: + raise ValueError( + f'grad_mag sidecar shape {grad_meta["array_shape"]} differs ' + f'from dense normal shape {reference_shape}') if scroll_zarr is not None: expected_shape = tuple(np.ceil(np.array(scroll_zarr.shape, dtype=np.float64) / lasagna_scale).astype(np.int64)) if tuple(reference_shape) != expected_shape: print( - f'WARNING: lasagna zarr shape {reference_shape} does not match ' + f'WARNING: lasagna store shape {reference_shape} does not match ' f'ceil(scroll_zarr.shape / lasagna_scale) {expected_shape}' ) @@ -73,38 +102,49 @@ def prepare_lasagna_volume( z_lo = max(0, int(np.floor(z_begin / lasagna_scale))) z_hi = min(z_size, int(np.ceil(z_end / lasagna_scale))) if z_hi <= z_lo: - raise RuntimeError(f'lasagna z-ROI [{z_lo}, {z_hi}) is empty (zarr z size {z_size})') + raise RuntimeError(f'lasagna z-ROI [{z_lo}, {z_hi}) is empty (store z size {z_size})') roi_shape = (z_hi - z_lo, reference_shape[1], reference_shape[2]) - from sparse_cuda_cache import ( - BoundedSparseCudaCache, - SparseLasagnaStore, - cache_budget_bytes, - ) + from sparse_cuda_cache import ResidentBrickPool, SparseLasagnaStore device = torch.device('cuda') - group = str(normal_zarr_group) normal_cache = None if use_normals: - normal_cache = BoundedSparseCudaCache( - source_paths=[ - os.path.join(normal_nx_zarr_path, group), - os.path.join(normal_ny_zarr_path, group), - ], - shape_zyx=tuple(int(v) for v in reference_shape), + if progress is not None: + progress.begin( + 'loading', 'Loading normal volumes onto GPU', + step=0, total_steps=0, unit='bricks') + normal_cache = ResidentBrickPool( + normal_sidecar, origin_zyx=(z_lo, 0, 0), - budget_bytes=cache_budget_bytes('normals', device), + z_roi=(z_lo, z_hi), device=device, label='lasagna normals', + expected_channels=2, + progress_callback=( + (lambda current, total, detail: progress.update( + current, total_steps=total, detail=detail)) + if progress is not None else None + ), ) grad_cache = None if use_spacing: - grad_cache = BoundedSparseCudaCache( - source_paths=[os.path.join(grad_mag_zarr_path, group)], - shape_zyx=tuple(int(v) for v in reference_shape), + if progress is not None: + progress.begin( + 'loading', 'Loading gradient volume onto GPU', + step=0, total_steps=0, unit='bricks') + grad_cache = ResidentBrickPool( + grad_sidecar, origin_zyx=(z_lo, 0, 0), - budget_bytes=cache_budget_bytes('grad_mag', device), + z_roi=(z_lo, z_hi), device=device, label='lasagna grad_mag', + expected_channels=1, + expected_shape_zyx=reference_shape, + progress_callback=( + (lambda current, total, detail: progress.update( + current, total_steps=total, detail=detail)) + if progress is not None else None + ), ) return { 'backend': 'sparse_cuda', @@ -161,6 +201,7 @@ def prepare_surf_sdt_volume( yx_bounds_working=None, interior_fn=None, paged_chunk=64, + progress=None, ): """Resolve and validate a surf-SDT store as a sparse CUDA input. @@ -248,20 +289,27 @@ def prepare_surf_sdt_volume( } shape = (z_hi - z_lo, int(array.shape[1]), int(array.shape[2])) - from sparse_cuda_cache import ( - BoundedSparseCudaCache, - SparseScalarStore, - cache_budget_bytes, - ) - device = torch.device('cuda') - cache = BoundedSparseCudaCache( - source_paths=[os.path.join(sdt_zarr_path, group_name)], - shape_zyx=tuple(int(v) for v in array.shape), + from sparse_cuda_cache import ResidentBrickPool, SparseScalarStore + sidecar = _require_sidecar(sdt_zarr_path, group_name, label='surf_sdt') + if progress is not None: + progress.begin( + 'loading', 'Loading surface-distance volume onto GPU', + step=0, total_steps=0, unit='bricks') + cache = ResidentBrickPool( + sidecar, origin_zyx=(z_lo, 0, 0), - budget_bytes=cache_budget_bytes('sdt', device), - device=device, + z_roi=(z_lo, z_hi), + device=torch.device('cuda'), label='surf_sdt', + expected_channels=1, + expected_shape_zyx=tuple(int(v) for v in array.shape), + progress_callback=( + (lambda current, total, detail: progress.update( + current, total_steps=total, detail=detail)) + if progress is not None else None + ), ) + fingerprint['respool_ct_mask'] = cache.meta.get('ct_mask') common = { 'kind': volume_kind, 'backend': 'sparse_cuda', diff --git a/volume-cartographer/scripts/spiral/losses.py b/volume-cartographer/scripts/spiral/losses.py index 6716da96e2..19f06ff81e 100644 --- a/volume-cartographer/scripts/spiral/losses.py +++ b/volume-cartographer/scripts/spiral/losses.py @@ -93,7 +93,7 @@ def _choose_pcl_indices(sampling_strata, num_to_sample): # stratified_pcl_sampling switch selects equal group shares or uniform sampling # over the combined pool. weighted = cfg['pcl_sampling_weights'] is not None - if not weighted and not cfg['stratified_pcl_sampling']: + if not weighted and not cfg['pcl_stratified_pcl_sampling']: return np.random.choice(sampling_strata['all'], num_to_sample, replace=False) strata = sampling_strata['strata'] weights = sampling_strata['weights'] if weighted else np.ones( @@ -120,7 +120,7 @@ def get_shell_outer_loss(shell_map, slice_to_spiral_transform, dr_per_winding, o if shell_map is None or outer_winding_idx is None: return zero, {} - num_samples = max(1, int(cfg['shell_num_samples'])) + num_samples = max(1, int(cfg['sample_count_shell_samples'])) huber_delta = torch.as_tensor(cfg['shell_huber_delta'], device=device, dtype=torch.float32) outer_spiral = canonical_winding_samples([outer_winding_idx], num_samples, dr_per_winding, device, z_begin, z_end)[0] @@ -357,8 +357,11 @@ def _build_patch_ijs(patches, patch_indices, num_points_per_direction, rng): def _sample_patch_batch(key, patches, sampling_probabilities, num_to_sample, num_points_per_direction, patch_atlas=None): - # Returns (combined_ijs_gpu (2,N,P,2), patch_indices_gpu (N,)). With - # prefetch enabled the batch was assembled and uploaded for this step + # Returns (combined_ijs_gpu (2,N,P,2), patch_indices_gpu (N,), + # slice_zyxs_gpu (2,N,P,3)). The atlas is host-resident and the ijs are + # born on the CPU, so the bilinear gather runs here at batch-build time + # and only the interpolated points are uploaded. With prefetch enabled the + # batch - sampling, gather, and uploads - was assembled for this step # during the previous one, and next step's batch is scheduled now. if num_to_sample <= 0: raise ValueError('Expected at least one patch index') @@ -377,10 +380,15 @@ def build(rng): else: ijs_np = _build_patch_ijs( patches, patch_indices, num_points_per_direction, rng) - ijs_gpu = torch.from_numpy(ijs_np).cuda(non_blocking=True) - idx_gpu = torch.from_numpy( - np.ascontiguousarray(patch_indices, dtype=np.int64)).cuda(non_blocking=True) - return ijs_gpu, idx_gpu + ijs_cpu = torch.from_numpy(ijs_np) + idx_cpu = torch.from_numpy( + np.ascontiguousarray(patch_indices, dtype=np.int64)) + _, N, P, _ = ijs_cpu.shape + slice_zyxs_gpu = patch_atlas.lookup( + idx_cpu[None, :, None].expand(2, N, P), ijs_cpu) + ijs_gpu = ijs_cpu.cuda(non_blocking=True) + idx_gpu = idx_cpu.cuda(non_blocking=True) + return ijs_gpu, idx_gpu, slice_zyxs_gpu if prefetch.prefetch_enabled() and torch.cuda.is_available(): pf = prefetch.get_prefetcher() @@ -390,22 +398,14 @@ def build(rng): return build(prefetch.LegacyNumpyRandom) -def _sample_patch_tracks(slice_to_spiral_transform, dr_per_winding, patches, patch_atlas, batch, extra_zyxs=None, num_points_per_patch=None): +def _sample_patch_tracks(slice_to_spiral_transform, dr_per_winding, patches, patch_atlas, batch, extra_zyxs=None): # For each patch, take one row and one column in straight mode, or two # geodesic paths in dijkstra mode. Either representation is a contiguous # walk, so unwrapping can stitch theta=0 crossings between samples. - if num_points_per_patch is None: - num_points_per_patch = cfg['num_points_per_patch'] - num_points_per_direction = num_points_per_patch // 2 - - # Batched bilinear interp on GPU: ijs are guaranteed to fall on valid quads by the - # strip sampler (it draws i0/j0 from `_sampling_valid_quad_*`), so we - # skip the per-call validity check used by patch.ij_to_zyx. - combined_ijs_gpu, patch_indices_gpu = batch - N = combined_ijs_gpu.shape[1] - patch_idx_per_sample = patch_indices_gpu[None, :, None].expand(2, N, num_points_per_direction) - all_slice_zyxs = patch_atlas.lookup(patch_idx_per_sample, combined_ijs_gpu) + # The bilinear atlas gather already ran on the CPU at batch-build time + # (see _sample_patch_batch); the batch carries the interpolated points. + combined_ijs_gpu, patch_indices_gpu, all_slice_zyxs = batch # When the caller has extra points (umbilicus, shell, ...), pack them into the same # forward ODE call to amortise the per-call overhead. @@ -564,7 +564,7 @@ def get_patch_and_umbilicus_losses(slice_to_spiral_transform, dr_per_winding, nu n_umb = umbilicus_zyx.shape[0] if shell_valid_zyxs is not None: - num_shell_samples = min(int(cfg['shell_num_samples']), shell_valid_zyxs.shape[0]) + num_shell_samples = min(int(cfg['sample_count_shell_samples']), shell_valid_zyxs.shape[0]) sample_idx = torch.randint(shell_valid_zyxs.shape[0], (num_shell_samples,), device=shell_valid_zyxs.device) extra_zyxs = torch.cat([umbilicus_zyx, shell_valid_zyxs[sample_idx]], dim=0) else: @@ -582,7 +582,7 @@ def get_patch_and_umbilicus_losses(slice_to_spiral_transform, dr_per_winding, nu num_patches_to_sample = max(num_patches_for_radius, num_patches_for_dt) if compute_dt else num_patches_for_radius batch = _sample_patch_batch( 'verified_patches', patches, patch_sampling_probabilities, - num_patches_to_sample, cfg['num_points_per_patch'] // 2, + num_patches_to_sample, cfg['sample_count_points_per_patch'] // 2, patch_atlas) ( @@ -657,7 +657,7 @@ def get_unverified_patch_losses(slice_to_spiral_transform, dr_per_winding, num_p num_patches_to_sample = max(num_patches_for_radius, num_patches_for_dt) if compute_dt else num_patches_for_radius batch = _sample_patch_batch( 'unverified_patches', patches, patch_sampling_probabilities, - num_patches_to_sample, cfg['unverified_num_points_per_patch'] // 2, + num_patches_to_sample, cfg['sample_count_unverified_points_per_patch'] // 2, patch_atlas) ( @@ -674,7 +674,6 @@ def get_unverified_patch_losses(slice_to_spiral_transform, dr_per_winding, num_p patches, patch_atlas, batch, - num_points_per_patch=cfg['unverified_num_points_per_patch'], ) return _patch_radius_and_dt_losses( @@ -682,8 +681,8 @@ def get_unverified_patch_losses(slice_to_spiral_transform, dr_per_winding, num_p all_slice_zyxs, all_spiral_zyxs, all_theta, all_shifted_radii, all_crossing_adjustments, num_patches_for_radius, num_patches_for_dt, compute_dt, dt_max_winding, - cfg['unverified_patch_radius_loss_margin'], cfg['unverified_patch_radius_loss_inv'], cfg['unverified_patch_radius_within_norm_p'], - cfg['unverified_patch_dt_loss_margin'], cfg['unverified_patch_dt_norm_p'], cfg['unverified_patch_dt_within_patch_norm_p'], + cfg['patch_unverified_patch_radius_loss_margin'], cfg['patch_unverified_patch_radius_loss_inv'], cfg['patch_unverified_patch_radius_within_norm_p'], + cfg['patch_unverified_patch_dt_loss_margin'], cfg['patch_unverified_patch_dt_norm_p'], cfg['patch_unverified_patch_dt_within_patch_norm_p'], patch_indices=batch[1], sample_ijs=sample_ijs, dt_target_cache=dt_target_cache, diagnostic_prefix='unverified_patch', ) @@ -865,7 +864,7 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche # points (adjacent mode) or the PCL chain between them (non-adjacent mode) # crosses theta=0, adjust the expected delta by that branch-cut jump. - num_points_per_strip = cfg['num_points_per_patch'] // 2 + num_points_per_strip = cfg['sample_count_points_per_patch'] // 2 num_strips_per_pcl = 4 num_strips_per_pair = 2 * num_strips_per_pcl # 8 @@ -877,7 +876,7 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche # sampling_strata indexes into point_collections and already excludes single-point # pcls (possible only for winding_is_absolute pcls), which can't form a cross-patch # pair; see the build_pcl_sampling_strata call in fit_spiral.main. - num_pcls_per_step = min(cfg['rel_winding_num_pcls'], len(sampling_strata['all'])) + num_pcls_per_step = min(cfg['sample_count_relative_winding_pcls'], len(sampling_strata['all'])) if num_pcls_per_step <= 0: return torch.zeros([], device='cuda') selected_idxs = _choose_pcl_indices(sampling_strata, num_pcls_per_step) @@ -886,7 +885,7 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche for pcl in selected_pcls: sorted_pcl_points = None sorted_pcl_point_idx = None - if not cfg['rel_winding_adjacent_patches_only']: + if not cfg['pcl_rel_winding_adjacent_patches_only']: sorted_pcl_points = [ point for _, point in sorted(pcl['points'].items(), key=lambda kv: int(kv[0])) ] @@ -895,14 +894,14 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche # Pair patches either only with their immediate neighbour in the pcl's # patch ordering (first-seen order; built in main()), # or with every other patch. - if cfg['rel_winding_adjacent_patches_only']: + if cfg['pcl_rel_winding_adjacent_patches_only']: cross_pairs = [(p1, p2) for p1, p2 in zip(pcl['points_by_patch'], list(pcl['points_by_patch'])[1:])] else: cross_pairs = list(itertools.combinations(pcl['points_by_patch'], r=2)) if not cross_pairs: continue - num_pairs_for_pcl = min(len(cross_pairs), cfg['rel_winding_num_patch_pairs_per_pcl']) + num_pairs_for_pcl = min(len(cross_pairs), cfg['sample_count_relative_winding_patch_pairs_per_pcl']) if num_pairs_for_pcl <= 0: continue chosen = np.random.choice(len(cross_pairs), num_pairs_for_pcl, replace=False) @@ -917,7 +916,7 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche i1, j1 = int(p1['on_patch']['ij'][0]), int(p1['on_patch']['ij'][1]) i2, j2 = int(p2['on_patch']['ij'][0]), int(p2['on_patch']['ij'][1]) - if cfg['rel_winding_adjacent_patches_only']: + if cfg['pcl_rel_winding_adjacent_patches_only']: pcl_chain = [p1, p2] else: idx1, idx2 = sorted_pcl_point_idx[id(p1)], sorted_pcl_point_idx[id(p2)] @@ -959,16 +958,16 @@ def get_patch_rel_winding_loss(slice_to_spiral_transform, dr_per_winding, patche flat_ijs[base + num_strips_per_pcl + s] = strip flat_pids.extend([pid1] * num_strips_per_pcl + [pid2] * num_strips_per_pcl) - # Batched GPU bilinear interp across all strips. + # Batched bilinear gather on the host-resident atlas; only the + # interpolated points are uploaded. patch_idx_per_strip_np = np.fromiter( (patch_atlas.id_to_idx[pid] for pid in flat_pids), dtype=np.int64, count=total_strips, ) - patch_idx_per_strip_gpu = torch.from_numpy(patch_idx_per_strip_np).cuda(non_blocking=True) - ijs_gpu = torch.from_numpy(flat_ijs).cuda(non_blocking=True) - patch_idx_per_sample = patch_idx_per_strip_gpu[:, None].expand(total_strips, num_points_per_strip) - flat_zyxs = patch_atlas.lookup(patch_idx_per_sample, ijs_gpu) + patch_idx_per_strip = torch.from_numpy(patch_idx_per_strip_np) + patch_idx_per_sample = patch_idx_per_strip[:, None].expand(total_strips, num_points_per_strip) + flat_zyxs = patch_atlas.lookup(patch_idx_per_sample, torch.from_numpy(flat_ijs)) # Mask out strip samples whose z falls outside [z_begin - margin, z_end + margin). # Computed before unwrapping but applied after, since _unwrap_track_shifted_radii @@ -1038,7 +1037,7 @@ def get_patch_abs_winding_loss(slice_to_spiral_transform, dr_per_winding, patche # the point's target. Each L starts at the annotated point, so its unwrapped # shifted-radius keeps the true absolute scale at the anchor. - num_points_per_strip = cfg['num_points_per_patch'] // 2 + num_points_per_strip = cfg['sample_count_points_per_patch'] // 2 num_strips_per_point = 4 # Each entry: (ls, pid, winding_annotation) where ls is a list of 4 L-shape ij strips. @@ -1046,7 +1045,7 @@ def get_patch_abs_winding_loss(slice_to_spiral_transform, dr_per_winding, patche strip_requests = [] abs_pcls = [pcl for pcl in point_collections if pcl.get('metadata', {}).get('winding_is_absolute', False)] - num_pcls_per_step = min(cfg['abs_winding_num_pcls'], len(abs_pcls)) + num_pcls_per_step = min(cfg['sample_count_absolute_winding_pcls'], len(abs_pcls)) if num_pcls_per_step <= 0: return torch.zeros([], device='cuda') selected_idxs = np.random.choice(len(abs_pcls), num_pcls_per_step, replace=False) @@ -1057,7 +1056,7 @@ def get_patch_abs_winding_loss(slice_to_spiral_transform, dr_per_winding, patche attached = [p for pts in pcl['points_by_patch'].values() for p in pts] if not attached: continue - num_points_for_pcl = min(len(attached), cfg['abs_winding_num_points_per_pcl']) + num_points_for_pcl = min(len(attached), cfg['sample_count_absolute_winding_points_per_pcl']) chosen = np.random.choice(len(attached), num_points_for_pcl, replace=False) for idx in chosen: p = attached[idx] @@ -1088,16 +1087,16 @@ def get_patch_abs_winding_loss(slice_to_spiral_transform, dr_per_winding, patche flat_ijs[base + s] = strip flat_pids.extend([pid] * num_strips_per_point) - # Batched GPU bilinear interp across all strips. + # Batched bilinear gather on the host-resident atlas; only the + # interpolated points are uploaded. patch_idx_per_strip_np = np.fromiter( (patch_atlas.id_to_idx[pid] for pid in flat_pids), dtype=np.int64, count=total_strips, ) - patch_idx_per_strip_gpu = torch.from_numpy(patch_idx_per_strip_np).cuda(non_blocking=True) - ijs_gpu = torch.from_numpy(flat_ijs).cuda(non_blocking=True) - patch_idx_per_sample = patch_idx_per_strip_gpu[:, None].expand(total_strips, num_points_per_strip) - flat_zyxs = patch_atlas.lookup(patch_idx_per_sample, ijs_gpu) + patch_idx_per_strip = torch.from_numpy(patch_idx_per_strip_np) + patch_idx_per_sample = patch_idx_per_strip[:, None].expand(total_strips, num_points_per_strip) + flat_zyxs = patch_atlas.lookup(patch_idx_per_sample, torch.from_numpy(flat_ijs)) # Mask out strip samples whose z falls outside [z_begin - margin, z_end + margin). # Computed before unwrapping but applied after, since _unwrap_track_shifted_radii @@ -1282,8 +1281,8 @@ def iter_lasagna_losses(slice_to_spiral_transform, dr_per_winding, lasagna_volum # Build both sparse requests before touching the shared CUDA cache. if compute_spacing: - density_decode = cfg['grad_mag_factor'] / cfg['grad_mag_encode_scale'] * lasagna_scale - num_steps = int(cfg['spacing_integration_steps']) + density_decode = cfg['dense_grad_mag_factor'] / cfg['dense_grad_mag_encode_scale'] * lasagna_scale + num_steps = int(cfg['dense_spacing_integration_steps']) step_frac = (torch.arange(num_steps, device=device).float() + 0.5) / num_steps integration_zyx = scroll_inner[:, None, :] + step_frac[None, :, None] * scroll_displacement[:, None, :] int_idx = (integration_zyx.detach() / lasagna_scale).round().long() @@ -1499,7 +1498,7 @@ def get_symmetric_dirichlet_loss(slice_to_spiral_transform, dr_per_winding, oute if outer_winding_idx is None: return torch.zeros([], device=device) if epsilon is None: - epsilon = cfg['sym_dirichlet_finite_difference_epsilon'] + epsilon = cfg['model_sym_dirichlet_finite_difference_epsilon'] spiral_zyx, e1, e2 = sample_spiral_surface_frame(dr_per_winding, outer_winding_idx, num_points) diff --git a/volume-cartographer/scripts/spiral/pack_resident_pools.py b/volume-cartographer/scripts/spiral/pack_resident_pools.py new file mode 100644 index 0000000000..504d9e529b --- /dev/null +++ b/volume-cartographer/scripts/spiral/pack_resident_pools.py @@ -0,0 +1,541 @@ +"""Pack sparse lasagna zarr stores into flat resident-pool sidecars. + +A sidecar holds every *occupied* brick of a (chunk-sparse) uint8 zarr array as +one flat pool file per channel, plus an int32 brick-grid table mapping brick +coordinates to pool rows (0 = absent, served by a reserved all-zero row). +Loading it is a single sequential read per channel instead of tens of +thousands of demand-fills through tensorstore, and a fully resident pool never +evicts, so the per-step LRU thrash disappears entirely. + +CT masking: the surf-SDT builder zeroes the *surface prediction* where the CT +volume reads 0, but the EDT then fills those voxels with capped far-field +distances, so the giant outside-the-scroll mask region is byte-occupied in the +store. Passing ``--ct`` zeroes every voxel whose (ratio-mapped) CT voxel is 0 +while packing; combined with a sub-chunk brick size this drops the mask skin +from the pool, and the zeros read back as no-data through the existing +sampling contract. + +Sidecar layout (``.respool_g[_pair]/``): + meta.json format/version, shapes, brick grid, sources, ct_mask info + table.npy int32 (gz, gy, gx); 0 = absent, i >= 1 = pool row i + brick_coords.npy int32 (rows, 3) brick coords per row; row 0 = -1 + channel_.u8 uint8 (rows, brick_voxels) raw pool; row 0 = zeros + +Usage (defaults match the fit_spiral input conventions): + python pack_resident_pools.py /path/to/lasagna_inputs \ + --ct /path/to/s1_ds2.zarr --ct-group 2 --verify 2000 +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +import threading +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np + +RESPOOL_FORMAT_VERSION = 2 + + +def sidecar_path(zarr_path: str, group: str, *, pair: bool = False) -> str: + """Canonical sidecar directory for a store; the pair suffix marks the + two-channel nx+ny pool that lives next to the nx zarr.""" + suffix = f'.respool_g{group}' + ('_pair' if pair else '') + return str(zarr_path).rstrip('/') + suffix + + +def _read_zarray_meta(array_dir: Path) -> dict: + with open(array_dir / '.zarray') as f: + meta = json.load(f) + if meta['dtype'] != '|u1': + raise ValueError(f'{array_dir}: expected |u1 dtype, got {meta["dtype"]!r}') + if meta.get('order', 'C') != 'C': + raise ValueError(f'{array_dir}: expected C order') + if meta.get('filters'): + raise ValueError(f'{array_dir}: filters are not supported') + return meta + + +def _register_codec(compressor: dict | None) -> None: + if compressor is not None and compressor.get('id') == 'vcz1': + try: + import vc.compression.vcz1_numcodecs # noqa: F401 registers vcz1 + except ImportError as exc: + raise RuntimeError( + 'the CT store uses the vcz1 codec; install the ' + 'volume-cartographer python bindings (vc.compression)') from exc + + +def _list_chunk_keys(array_dir: Path, separator: str) -> set[tuple[int, int, int]]: + keys = set() + if separator == '/': + for path in glob.iglob(str(array_dir / '*' / '*' / '*')): + parts = path.rsplit(os.sep, 3)[-3:] + try: + keys.add((int(parts[0]), int(parts[1]), int(parts[2]))) + except ValueError: + continue + else: + for name in os.listdir(array_dir): + parts = name.split('.') + if len(parts) == 3 and all(p.isdigit() for p in parts): + keys.add((int(parts[0]), int(parts[1]), int(parts[2]))) + return keys + + +def _chunk_path(array_dir: Path, key: tuple[int, int, int], separator: str) -> Path: + if separator == '/': + return array_dir / str(key[0]) / str(key[1]) / str(key[2]) + return array_dir / f'{key[0]}.{key[1]}.{key[2]}' + + +def _make_chunk_reader(compressor: dict | None, chunk_voxels: int): + if compressor is None: + def read(path: Path) -> np.ndarray: + data = np.fromfile(path, dtype=np.uint8) + if data.size != chunk_voxels: + raise ValueError(f'{path}: raw chunk has {data.size} bytes, expected {chunk_voxels}') + return data + return read + import numcodecs + codec = numcodecs.get_codec(compressor) + + def read(path: Path) -> np.ndarray: + with open(path, 'rb') as f: + decoded = codec.decode(f.read()) + decoded = np.frombuffer(decoded, dtype=np.uint8) + if decoded.size != chunk_voxels: + raise ValueError(f'{path}: decoded chunk has {decoded.size} bytes, expected {chunk_voxels}') + return decoded + return read + + +class CtMasker: + """Voxel-level ``CT == 0`` mask, ratio-mapped onto a target array grid. + + ``ratio`` is target voxels per CT voxel, per axis (integer, e.g. 2 for the + surf-SDT group 1 against s1_ds2 group 2, 1 for the group-4 lasagna + stores). Decoded CT chunks are kept in a small LRU because consecutive + target chunks in key order revisit the same CT chunks. + """ + + def __init__(self, ct_zarr_path: str, ct_group: str, target_shape, + cache_chunks: int = 512): + self.array_dir = Path(ct_zarr_path) / str(ct_group) + meta = _read_zarray_meta(self.array_dir) + _register_codec(meta['compressor']) + self.shape = tuple(meta['shape']) + self.chunks = tuple(meta['chunks']) + self.separator = meta.get('dimension_separator', '.') + self._reader = _make_chunk_reader( + meta['compressor'], int(np.prod(self.chunks))) + self.ratio = tuple( + int(round(ts / cs)) for ts, cs in zip(target_shape, self.shape)) + for axis, (ts, cs, r) in enumerate( + zip(target_shape, self.shape, self.ratio)): + if r < 1 or not (cs - 1 <= (ts + r - 1) // r <= cs + 1): + raise ValueError( + f'CT shape {self.shape} does not map onto target shape ' + f'{tuple(target_shape)} with an integer per-axis ratio ' + f'(axis {axis}: ratio {r})') + self._keys = _list_chunk_keys(self.array_dir, self.separator) + self._cache: OrderedDict[tuple, np.ndarray | None] = OrderedDict() + self._cache_chunks = int(cache_chunks) + self._lock = threading.Lock() + + def _ct_chunk(self, key: tuple[int, int, int]) -> np.ndarray | None: + """Decoded CT chunk, or None when absent (all fill / all masked).""" + with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + return self._cache[key] + value = None + if key in self._keys: + value = self._reader( + _chunk_path(self.array_dir, key, self.separator) + ).reshape(self.chunks) + with self._lock: + self._cache[key] = value + while len(self._cache) > self._cache_chunks: + self._cache.popitem(last=False) + return value + + def ct_region(self, lo, hi) -> np.ndarray: + """CT voxels for the half-open CT-grid box [lo, hi); OOB reads 0.""" + lo = [max(0, v) for v in lo] + hi = [min(s, v) for s, v in zip(self.shape, hi)] + out = np.zeros([max(0, h - l) for l, h in zip(lo, hi)], dtype=np.uint8) + if out.size == 0: + return out + c0 = [l // c for l, c in zip(lo, self.chunks)] + c1 = [(h - 1) // c for h, c in zip(hi, self.chunks)] + for cz in range(c0[0], c1[0] + 1): + for cy in range(c0[1], c1[1] + 1): + for cx in range(c0[2], c1[2] + 1): + chunk = self._ct_chunk((cz, cy, cx)) + if chunk is None: + continue + base = (cz * self.chunks[0], cy * self.chunks[1], + cx * self.chunks[2]) + src = tuple( + slice(max(l, b) - b, min(h, b + c) - b) + for l, h, b, c in zip(lo, hi, base, self.chunks)) + dst = tuple( + slice(max(l, b) - l, min(h, b + c) - l) + for l, h, b, c in zip(lo, hi, base, self.chunks)) + out[dst] = chunk[src] + return out + + def occupied_block(self, origin_zyx, block_shape) -> np.ndarray: + """Boolean ``CT != 0`` for a target-grid box, upsampled by ratio. + + ``origin_zyx`` must be aligned to the ratio (chunk origins are). + """ + for o, r in zip(origin_zyx, self.ratio): + if o % r: + raise ValueError(f'block origin {origin_zyx} not aligned to CT ratio {self.ratio}') + lo = [o // r for o, r in zip(origin_zyx, self.ratio)] + hi = [-(-(o + b) // r) for o, b, r in zip(origin_zyx, block_shape, self.ratio)] + region = self.ct_region(lo, hi) != 0 + padded = np.zeros([h - l for l, h in zip(lo, hi)], dtype=bool) + padded[tuple(slice(0, s) for s in region.shape)] = region + for axis, r in enumerate(self.ratio): + if r > 1: + padded = padded.repeat(r, axis=axis) + return padded[tuple(slice(0, b) for b in block_shape)] + + def ct_value_at(self, target_zyx) -> int: + ct = [t // r for t, r in zip(target_zyx, self.ratio)] + region = self.ct_region(ct, [c + 1 for c in ct]) + return int(region[0, 0, 0]) if region.size else 0 + + def describe(self) -> dict: + return { + 'path': str(self.array_dir.parent), + 'group': self.array_dir.name, + 'ratio': list(self.ratio), + } + + +def pack_arrays( + array_dirs: list[str], + out_dir: str, + *, + label: str, + brick_shape: tuple[int, int, int] | None = None, + ct_masker: CtMasker | None = None, + io_threads: int = 16, + force: bool = False, +) -> str: + """Pack one or more same-geometry uint8 zarr arrays (one per channel). + + ``brick_shape`` defaults to the source chunk shape and must divide it; + bricks that are entirely zero across all channels (after CT masking) are + dropped from the pool. + """ + out_dir = Path(out_dir) + meta_path = out_dir / 'meta.json' + if meta_path.exists() and not force: + print(f'{label}: {out_dir} already exists, skipping (--force to rebuild)') + return str(out_dir) + + array_dirs = [Path(d) for d in array_dirs] + metas = [_read_zarray_meta(d) for d in array_dirs] + shape = tuple(metas[0]['shape']) + chunks = tuple(metas[0]['chunks']) + for d, m in zip(array_dirs, metas): + if tuple(m['shape']) != shape or tuple(m['chunks']) != chunks: + raise ValueError( + f'{d}: shape/chunks {m["shape"]}/{m["chunks"]} differ from ' + f'{array_dirs[0]}: {shape}/{chunks}') + separators = [m.get('dimension_separator', '.') for m in metas] + chunk_voxels = int(np.prod(chunks)) + brick = tuple(chunks) if brick_shape is None else tuple(brick_shape) + if any(c % b for c, b in zip(chunks, brick)): + raise ValueError(f'brick {brick} must divide the source chunks {chunks}') + sub = tuple(c // b for c, b in zip(chunks, brick)) + subs_per_chunk = int(np.prod(sub)) + brick_voxels = int(np.prod(brick)) + grid = tuple((s + b - 1) // b for s, b in zip(shape, brick)) + + per_channel_keys = [ + _list_chunk_keys(d, sep) for d, sep in zip(array_dirs, separators) + ] + keys = sorted(set().union(*per_channel_keys)) + if not keys: + raise RuntimeError(f'{label}: no chunk files found under {array_dirs}') + for d, channel_keys in zip(array_dirs, per_channel_keys): + missing = len(keys) - len(channel_keys) + if missing: + print(f'{label}: WARNING {d.name} is missing {missing} of ' + f'{len(keys)} union chunks; they read as zero (no-data)') + + upper_gib = (len(keys) * subs_per_chunk + 1) * brick_voxels * len(array_dirs) / 1024 ** 3 + print(f'{label}: packing {len(keys):,} chunks of {chunks} into bricks of ' + f'{brick} x{len(array_dirs)} channel(s), <= {upper_gib:.1f} GiB, ' + f'ct_mask={"on" if ct_masker else "off"} -> {out_dir}') + + readers = [ + _make_chunk_reader(m['compressor'], chunk_voxels) for m in metas + ] + + def read_chunk(key): + """Returns (key, kept sub-brick index array, per-channel (n, vox)).""" + channels = [] + for d, sep, reader, channel_keys in zip( + array_dirs, separators, readers, per_channel_keys): + if key in channel_keys: + channels.append(reader(_chunk_path(d, key, sep)).reshape(chunks)) + else: + channels.append(None) + if ct_masker is not None: + origin = tuple(k * c for k, c in zip(key, chunks)) + occupied = ct_masker.occupied_block(origin, chunks) + channels = [ + None if a is None else np.where(occupied, a, np.uint8(0)) + for a in channels + ] + split = [ + None if a is None else np.ascontiguousarray( + a.reshape(sub[0], brick[0], sub[1], brick[1], sub[2], brick[2]) + .transpose(0, 2, 4, 1, 3, 5) + ).reshape(subs_per_chunk, brick_voxels) + for a in channels + ] + keep = np.zeros(subs_per_chunk, dtype=bool) + for a in split: + if a is not None: + keep |= a.any(axis=1) + keep_idx = np.flatnonzero(keep) + return key, keep_idx, [ + (np.zeros((len(keep_idx), brick_voxels), dtype=np.uint8) + if a is None else a[keep_idx]) + for a in split + ] + + started = time.perf_counter() + os.makedirs(out_dir, exist_ok=True) + if meta_path.exists(): + meta_path.unlink() # invalidate a stale sidecar while rebuilding + + coords = [(-1, -1, -1)] + sub_grid = np.stack(np.unravel_index(np.arange(subs_per_chunk), sub), axis=1) + channel_files = [ + open(out_dir / f'channel_{i}.u8', 'wb', buffering=1024 * 1024) + for i in range(len(array_dirs)) + ] + try: + for f in channel_files: + f.write(bytes(brick_voxels)) # row 0: reserved all-zero brick + with ThreadPoolExecutor(io_threads) as executor: + done = 0 + for key, keep_idx, channel_bricks in executor.map( + read_chunk, keys, chunksize=4): + base = np.multiply(key, sub) + for sz, sy, sx in sub_grid[keep_idx]: + coords.append((base[0] + sz, base[1] + sy, base[2] + sx)) + for f, bricks in zip(channel_files, channel_bricks): + f.write(bricks.tobytes()) + done += 1 + if done % 2000 == 0 or done == len(keys): + elapsed = time.perf_counter() - started + print(f'{label}: {done:,}/{len(keys):,} chunks, ' + f'{len(coords) - 1:,} bricks kept ' + f'({done / len(keys) * 100:.0f}%, {elapsed:.0f}s)', + flush=True) + finally: + for f in channel_files: + f.close() + + coords = np.asarray(coords, dtype=np.int32) + rows = len(coords) + table = np.zeros(grid, dtype=np.int32) + table[coords[1:, 0], coords[1:, 1], coords[1:, 2]] = np.arange( + 1, rows, dtype=np.int32) + np.save(out_dir / 'table.npy', table) + np.save(out_dir / 'brick_coords.npy', coords) + pool_gib = rows * brick_voxels * len(array_dirs) / 1024 ** 3 + meta = { + 'format': 'respool', + 'version': RESPOOL_FORMAT_VERSION, + 'array_shape': list(shape), + 'source_chunk_shape': list(chunks), + 'brick_shape': list(brick), + 'grid_shape': list(grid), + 'rows': rows, + 'channels': [str(d) for d in array_dirs], + 'channel_names': [d.parent.name + '/' + d.name for d in array_dirs], + 'dtype': 'u1', + 'ct_mask': ct_masker.describe() if ct_masker is not None else None, + 'created': time.strftime('%Y-%m-%dT%H:%M:%S%z'), + } + with open(meta_path, 'w') as f: + json.dump(meta, f, indent=2) + print(f'{label}: done in {time.perf_counter() - started:.0f}s, ' + f'{rows - 1:,} bricks = {pool_gib:.1f} GiB') + return str(out_dir) + + +def open_pool(sidecar_dir: str): + """Open a sidecar read-only: (meta, table, coords, [channel memmaps]).""" + sidecar_dir = Path(sidecar_dir) + with open(sidecar_dir / 'meta.json') as f: + meta = json.load(f) + if meta.get('format') != 'respool' or meta.get('version') != RESPOOL_FORMAT_VERSION: + raise ValueError(f'{sidecar_dir}: unsupported sidecar format {meta.get("format")!r} ' + f'v{meta.get("version")!r}') + table = np.load(sidecar_dir / 'table.npy') + coords = np.load(sidecar_dir / 'brick_coords.npy') + rows = int(meta['rows']) + brick_voxels = int(np.prod(meta['brick_shape'])) + pools = [] + for i in range(len(meta['channels'])): + path = sidecar_dir / f'channel_{i}.u8' + expected = rows * brick_voxels + if path.stat().st_size != expected: + raise ValueError(f'{path}: size {path.stat().st_size} != expected {expected}') + pools.append(np.memmap(path, dtype=np.uint8, mode='r', + shape=(rows, brick_voxels))) + return meta, table, coords, pools + + +def verify_pool(sidecar_dir: str, num_samples: int, seed: int = 0) -> None: + """Compare random voxels against the source zarrs. Where the sidecar was + CT-masked, a pool zero is accepted iff the mapped CT voxel is zero.""" + import zarr + + meta, table, _coords, pools = open_pool(sidecar_dir) + arrays = [zarr.open(path, mode='r') for path in meta['channels']] + masker = None + if meta.get('ct_mask'): + masker = CtMasker(meta['ct_mask']['path'], meta['ct_mask']['group'], + meta['array_shape']) + shape = np.array(meta['array_shape']) + brick = np.array(meta['brick_shape']) + rng = np.random.default_rng(seed) + points = (rng.random((num_samples, 3)) * shape).astype(np.int64) + brick_idx = points // brick + local = points - brick_idx * brick + linear = (local[:, 0] * brick[1] + local[:, 1]) * brick[2] + local[:, 2] + rows = table[brick_idx[:, 0], brick_idx[:, 1], brick_idx[:, 2]] + mismatches = masked = 0 + for channel, (pool, array) in enumerate(zip(pools, arrays)): + got = pool[rows, linear] + for i in range(num_samples): + expected = array[points[i, 0], points[i, 1], points[i, 2]] + if got[i] == expected: + continue + if (masker is not None and got[i] == 0 + and masker.ct_value_at(points[i]) == 0): + masked += 1 + continue + mismatches += 1 + if mismatches <= 10: + print(f'MISMATCH ch{channel} zyx={tuple(points[i])}: ' + f'pool={got[i]} zarr={expected}') + if mismatches: + raise SystemExit(f'{sidecar_dir}: {mismatches} mismatching voxels') + print(f'{sidecar_dir}: verified {num_samples} random voxels ' + f'x{len(pools)} channel(s); all match ' + f'({masked} differ only under the CT mask)') + + +def _find_one(folder: str, pattern: str) -> str | None: + hits = sorted(glob.glob(os.path.join(folder, pattern))) + if not hits: + return None + if len(hits) > 1: + print(f'WARNING: multiple matches for {pattern}, using {hits[0]}') + return hits[0] + + +def main(argv=None): + parser = argparse.ArgumentParser( + description='Pack lasagna zarr stores into resident-pool sidecars.') + parser.add_argument('folder', help='lasagna_inputs directory holding the ' + '*_surf_sdt / *_nx / *_ny / *_grad_mag ome.zarr stores') + parser.add_argument('--what', default='sdt,normals,grad_mag', + help='comma list of sdt,normals,grad_mag (default all)') + parser.add_argument('--sdt-group', default='1') + parser.add_argument('--normal-group', default='4', + help='group for both normals and grad_mag') + parser.add_argument('--ct', default=None, metavar='ZARR', + help='CT zarr whose zero voxels are masked out of every ' + 'packed store (e.g. .../s1_ds2.zarr)') + parser.add_argument('--ct-group', default='2') + parser.add_argument('--sdt-brick', type=int, default=None, + help='SDT pool brick edge; defaults to 32 when --ct is ' + 'given (so masked bricks are actually dropped), ' + 'else the source chunk size') + parser.add_argument('--io-threads', type=int, default=16) + parser.add_argument('--force', action='store_true', + help='rebuild sidecars that already exist') + parser.add_argument('--verify', type=int, default=0, metavar='N', + help='after packing, compare N random voxels per ' + 'store against the source zarr') + args = parser.parse_args(argv) + + what = {w.strip() for w in args.what.split(',') if w.strip()} + unknown = what - {'sdt', 'normals', 'grad_mag'} + if unknown: + parser.error(f'unknown --what entries: {sorted(unknown)}') + + def masker_for(array_dir): + if args.ct is None: + return None + shape = _read_zarray_meta(Path(array_dir))['shape'] + return CtMasker(args.ct, args.ct_group, shape) + + built = [] + if 'sdt' in what: + sdt = _find_one(args.folder, '*_surf_sdt.ome.zarr') + if sdt is None: + print('no *_surf_sdt.ome.zarr found, skipping sdt') + else: + array = os.path.join(sdt, args.sdt_group) + edge = args.sdt_brick or (32 if args.ct else None) + pack_arrays([array], sidecar_path(sdt, args.sdt_group), + label='surf_sdt', + brick_shape=(edge,) * 3 if edge else None, + ct_masker=masker_for(array), + io_threads=args.io_threads, force=args.force) + built.append(sidecar_path(sdt, args.sdt_group)) + if 'normals' in what: + nx = _find_one(args.folder, '*_nx.ome.zarr') + ny = _find_one(args.folder, '*_ny.ome.zarr') + if nx is None or ny is None: + print('missing *_nx/*_ny ome.zarr, skipping normals') + else: + arrays = [os.path.join(nx, args.normal_group), + os.path.join(ny, args.normal_group)] + out = sidecar_path(nx, args.normal_group, pair=True) + pack_arrays(arrays, out, label='normals', + ct_masker=masker_for(arrays[0]), + io_threads=args.io_threads, force=args.force) + built.append(out) + if 'grad_mag' in what: + grad = _find_one(args.folder, '*_grad_mag.ome.zarr') + if grad is None: + print('no *_grad_mag.ome.zarr found, skipping grad_mag') + else: + array = os.path.join(grad, args.normal_group) + out = sidecar_path(grad, args.normal_group) + pack_arrays([array], out, label='grad_mag', + ct_masker=masker_for(array), + io_threads=args.io_threads, force=args.force) + built.append(out) + + if args.verify: + for out in built: + verify_pool(out, args.verify) + return built + + +if __name__ == '__main__': + main(sys.argv[1:]) diff --git a/volume-cartographer/scripts/spiral/phase_tuning.py b/volume-cartographer/scripts/spiral/phase_tuning.py index 31f056ad9d..8ee2763e50 100644 --- a/volume-cartographer/scripts/spiral/phase_tuning.py +++ b/volume-cartographer/scripts/spiral/phase_tuning.py @@ -34,16 +34,16 @@ def _build_model(checkpoint, dataset, cfg, outward_sense): umbilicus_zyx = torch.from_numpy(np.concatenate([ all_z[:, None], umbilicus_fn(all_z), ], axis=-1).astype(np.float32)).to(device) - radius = int(cfg['flow_bounds_radius']) + radius = int(cfg['model_flow_bounds_radius']) flow_min = torch.tensor([ - model_z_begin - int(cfg['flow_bounds_z_margin']), -radius, -radius, + model_z_begin - int(cfg['model_flow_bounds_z_margin']), -radius, -radius, ], dtype=torch.int64, device=device) flow_max = torch.tensor([ - model_z_end + int(cfg['flow_bounds_z_margin']), radius, radius, + model_z_end + int(cfg['model_flow_bounds_z_margin']), radius, radius, ], dtype=torch.int64, device=device) model = fs.SpiralAndTransform( - flow_integration_steps=int(cfg['num_flow_integration_steps']), - flow_integration_solver=cfg['flow_integration_solver'], + flow_integration_steps=int(cfg['model_num_flow_integration_steps']), + flow_integration_solver=cfg['model_flow_integration_solver'], flow_min_corner_zyx=flow_min, flow_max_corner_zyx=flow_max, umbilicus_zyx=umbilicus_zyx, @@ -52,8 +52,6 @@ def _build_model(checkpoint, dataset, cfg, outward_sense): ).to(device) model.load_state_dict(checkpoint['spiral_and_transform']) model.eval() - model.flow_field.flow_scales[1] = cfg[ - 'flow_field_high_res_lr_scale_final'] return model @@ -125,14 +123,14 @@ def main(): args.checkpoint.resolve() if args.checkpoint else dataset / 'checkpoint_fitted.ckpt') checkpoint = load_checkpoint_cpu(str(checkpoint_path)) - cfg = dict(fs.default_config) + cfg = fs.Config().as_dict() cfg.update({ key: value for key, value in checkpoint.get('cfg', {}).items() if key in cfg }) cfg.update({ 'dense_spacing_mode': 'phase', - 'dense_spacing_num_pairs': int(args.pairs), + 'sample_count_dense_spacing_pairs': int(args.pairs), # Measure phase registration and count only; the barrier gradient is # probed separately and attachment is not part of this measurement. 'loss_weight_min_spacing': 0.0, @@ -268,8 +266,8 @@ def main(): 'dense_spacing_phase_top2_margin', 'dense_spacing_phase_min_matched_windings', 'dense_spacing_phase_min_matched_mass', - 'min_spacing_d_min_wv', - 'min_spacing_independent_samples', + 'dense_min_spacing_d_min_wv', + 'sample_count_minimum_spacing_independent_samples', ) }, 'summary': summary, diff --git a/volume-cartographer/scripts/spiral/prefetch.py b/volume-cartographer/scripts/spiral/prefetch.py index 2ffa8ef7ad..9e656803ad 100644 --- a/volume-cartographer/scripts/spiral/prefetch.py +++ b/volume-cartographer/scripts/spiral/prefetch.py @@ -1,8 +1,9 @@ """One-step-ahead prefetch of per-step CPU batch assembly. -The training loop's per-step CPU work (numpy patch-strip sampling, the track -point gather, and their host->device uploads) otherwise runs serially while -the GPU is idle: the track gather even forces a device sync. Prefetching runs +The training loop's per-step CPU work (numpy patch-strip sampling, the +bilinear gather on the host-resident patch atlas, the track point gather, and +their host->device uploads) otherwise runs serially while the GPU is idle: +the track gather even forces a device sync. Prefetching runs the *next* step's sampling on a worker thread while the GPU chews on the current step. diff --git a/volume-cartographer/scripts/spiral/pyproject.toml b/volume-cartographer/scripts/spiral/pyproject.toml index e5299ac196..c14c887152 100644 --- a/volume-cartographer/scripts/spiral/pyproject.toml +++ b/volume-cartographer/scripts/spiral/pyproject.toml @@ -18,7 +18,9 @@ dependencies = [ "numpy>=2.5.0", "opencv-python-headless>=4.13.0.92", "pillow>=12.2.0", + "posix-ipc>=1.3.2; sys_platform != 'win32'", "pyro-ppl>=1.9.1", + "rustworkx>=0.18.0", "s3fs>=2026.6.0", "scipy>=1.18.0", "tifffile>=2026.6.1", diff --git a/volume-cartographer/scripts/spiral/satisfaction_metrics.py b/volume-cartographer/scripts/spiral/satisfaction_metrics.py index fea06b771a..ac089bfa51 100644 --- a/volume-cartographer/scripts/spiral/satisfaction_metrics.py +++ b/volume-cartographer/scripts/spiral/satisfaction_metrics.py @@ -520,7 +520,12 @@ def save_overlay_and_print_satisfaction( get_or_build_unattached_pcl_flat, run_tag=None, save_png_visualizations=False, + progress=None, ): + if progress is not None: + progress.begin( + 'finalizing', 'Evaluating verified patches', + detail=f'{len(patches_list):,} patches') satisfied_patches, satisfied_areas, total_areas, satisfied_quad_masks, boundary_satisfied_patches, target_winding_idx_per_patch = get_patch_satisfied_areas( slice_to_spiral_transform, dr_per_winding, patches_list, z_begin, z_end, verbose=True, ) @@ -544,6 +549,10 @@ def save_overlay_and_print_satisfaction( unattached_pcl_per_point_satisfied = [] unattached_pcl_fully_satisfied = torch.zeros(len(unattached_pcl_strips), dtype=torch.bool) if unattached_pcl_strips: + if progress is not None: + progress.begin( + 'finalizing', 'Evaluating point collections', + detail=f'{len(unattached_pcl_strips):,} collections') unattached_pcl_satisfied_counts, unattached_pcl_total_counts, unattached_pcl_per_point_satisfied = get_unattached_pcl_satisfied_counts( slice_to_spiral_transform, dr_per_winding, unattached_pcl_strips, get_or_build_unattached_pcl_flat, ) @@ -565,6 +574,10 @@ def save_overlay_and_print_satisfaction( if torch.cuda.is_available(): torch.cuda.empty_cache() try: + if progress is not None: + progress.begin( + 'finalizing', 'Evaluating tracks', + detail=f'{len(tracks):,} tracks') track_satisfied_counts, track_total_counts = get_track_satisfied_counts_in_chunks( slice_to_spiral_transform, dr_per_winding, tracks, metrics_config, ) @@ -653,6 +666,8 @@ def save_overlay_and_print_satisfaction( satisfied_quads_flat.to(torch.int64), ) if save_png_visualizations and os.environ.get('FIT_SPIRAL_SKIP_SAVE_OVERLAY') != '1': + if progress is not None: + progress.begin('finalizing', 'Rendering satisfaction overlay') winding_range, patch_extents, pcl_extents = compute_winding_range_and_input_extents( slice_to_spiral_transform, dr_per_winding, patches_list, unattached_pcl_strips, cfg, z_begin, z_end, get_or_build_unattached_pcl_flat, @@ -695,5 +710,5 @@ def get_patch_satisfied_areas_for_mesh( out_path, cfg, z_begin, z_end, voxel_size_um, get_or_build_unattached_pcl_flat, get_patch_satisfied_areas_for_mesh, tracks=tracks, - run_tag=run_tag, name=suffix, + run_tag=run_tag, name=suffix, progress=progress, ) diff --git a/volume-cartographer/scripts/spiral/sdt_losses.py b/volume-cartographer/scripts/spiral/sdt_losses.py index 7b079d147f..b1d647e330 100644 --- a/volume-cartographer/scripts/spiral/sdt_losses.py +++ b/volume-cartographer/scripts/spiral/sdt_losses.py @@ -1643,7 +1643,7 @@ def _phase_and_count_losses( if outer - inner < 1: return started = time.perf_counter() - num_pairs = int(cfg['dense_spacing_num_pairs']) + num_pairs = int(cfg['sample_count_dense_spacing_pairs']) if num_pairs <= 0: # A zero shared-ray budget disables phase/count sampling. Be defensive # here because session configuration is remote and may come from an @@ -1651,9 +1651,9 @@ def _phase_and_count_losses( # supplemental budget remains usable without a shared batch. if density_active: total_density_pairs = max( - 0, int(cfg['dense_spacing_density_extra_pairs'])) + 0, int(cfg['sample_count_dense_spacing_density_extra_pairs'])) chunk_cap = max( - 1, int(cfg['dense_spacing_density_chunk_pairs'])) + 1, int(cfg['sample_count_dense_spacing_density_chunk_pairs'])) remaining = total_density_pairs while remaining > 0: n_chunk = min(chunk_cap, remaining) @@ -1714,7 +1714,7 @@ def _phase_and_count_losses( seg_valids = [pair['seg_valid']] theta_all, z_all, r_mid_all = [theta], [z], [ (k + m_f / 2 + theta / (2 * np.pi)) * dr_per_winding.detach()] - extra_pairs = int(cfg['dense_spacing_count_extra_pairs']) + extra_pairs = int(cfg['sample_count_dense_spacing_count_extra_pairs']) if extra_pairs > 0: # Optional count-only supplement restoring spatial coverage when # the shared batch alone is too sparse. @@ -1805,7 +1805,7 @@ def _phase_and_count_losses( if not phase_active: # the shared-batch health metrics otherwise ride with phase density_metrics.update(central_metrics) - extra_pairs = int(cfg['dense_spacing_density_extra_pairs']) + extra_pairs = int(cfg['sample_count_dense_spacing_density_extra_pairs']) total_density_pairs = num_pairs + max(0, extra_pairs) share = num_pairs / total_density_pairs density_metrics['_shared_graph'] = True @@ -1815,7 +1815,7 @@ def _phase_and_count_losses( # detection, no pair-HMM) is cheap enough to run at several times # the shared-batch sampling; chunked so only one chunk's graph is # resident per backward. - chunk_cap = int(cfg['dense_spacing_density_chunk_pairs']) + chunk_cap = int(cfg['sample_count_dense_spacing_density_chunk_pairs']) remaining = max(0, extra_pairs) while remaining > 0: n_chunk = min(chunk_cap, remaining) @@ -1925,7 +1925,7 @@ def get_min_spacing_loss( zero = torch.zeros([], device=device) if outer_winding_idx is None: return zero, {} - num_samples = int(cfg['min_spacing_independent_samples']) + num_samples = int(cfg['sample_count_minimum_spacing_independent_samples']) inner, outer = fitted_winding_domain(outer_winding_idx) if outer <= inner or num_samples <= 0: return zero, {} @@ -1936,7 +1936,7 @@ def get_min_spacing_loss( z = torch.empty(num_samples, device=device).uniform_( float(z_begin), float(z_end - 1), generator=generator) ell_gap = spiral_and_transform.get_native_log_gaps(winding, theta, z) - ell_min = float(np.log(float(cfg['min_spacing_d_min_wv']))) + ell_min = float(np.log(float(cfg['dense_min_spacing_d_min_wv']))) deficiency = F.relu(ell_min - ell_gap) loss = deficiency.square().mean() if not metrics_enabled(): @@ -1972,7 +1972,7 @@ def get_dense_attachment_loss( if sdt_volume['kind'] != 'sdt': raise ValueError('attachment requires a signed-distance store, not a raw-surf volume') - num_points = int(cfg['dense_attachment_num_points']) + num_points = int(cfg['sample_count_dense_attachment_points']) attachment_scale = float(cfg['dense_attachment_scale']) inner_winding, outer_winding = fitted_winding_domain(outer_winding_idx) if outer_winding < inner_winding: diff --git a/volume-cartographer/scripts/spiral/sparse_cuda_cache.py b/volume-cartographer/scripts/spiral/sparse_cuda_cache.py index 3bd1219ba5..fe10625a41 100644 --- a/volume-cartographer/scripts/spiral/sparse_cuda_cache.py +++ b/volume-cartographer/scripts/spiral/sparse_cuda_cache.py @@ -1,364 +1,193 @@ -"""Bounded LRU cache for sparse integer gathers from local uint8 Zarr arrays. +"""Fully-resident sparse brick pools for integer gathers from packed sidecars. The loss code asks for integer SDT corners and nearest-neighbour normal -samples. Keeping that contract here preserves the existing SDT no-data -renormalisation while moving the random fetch itself entirely onto the GPU -after a chunk is resident. +samples at positions scattered uniformly over the whole ROI, so any bounded +cache thrashes; instead the entire occupied brick set of each field is loaded +once from a ``pack_resident_pools.py`` sidecar (a single sequential read per +channel) and every gather afterwards is pure device indexing with no I/O, no +eviction bookkeeping, and no host synchronisation. + +Bricks absent from the sidecar (never written, or zeroed by its CT mask) map +to the reserved all-zero row 0, which reads as the encoded no-data value the +sampling contract already handles. """ from __future__ import annotations import os import time -from collections import OrderedDict from pathlib import Path import numpy as np import torch -CHUNK_SIZE = 32 -CHUNK_VOXELS = CHUNK_SIZE ** 3 +from pack_resident_pools import open_pool -def _env_gib(name: str, default: float) -> float: - raw = os.environ.get(name) - value = default if raw is None else float(raw) - if value <= 0: - raise ValueError(f"{name} must be positive, got {value}") - return value +class ResidentBrickPool: + """Device-resident brick pool loaded from a resident-pool sidecar. - -def cache_budget_bytes(kind: str, device: torch.device | str) -> int: - """Return the bounded device-pool budget for one logical field.""" - defaults = {"sdt": 16.0, "normals": 6.0, "grad_mag": 2.0} - fractions = {"sdt": 0.30, "normals": 0.20, "grad_mag": 0.10} - names = { - "sdt": "FIT_SPIRAL_SPARSE_SDT_CACHE_GB", - "normals": "FIT_SPIRAL_SPARSE_NORMAL_CACHE_GB", - "grad_mag": "FIT_SPIRAL_SPARSE_GRAD_CACHE_GB", - } - requested = int(_env_gib(names[kind], defaults[kind]) * 1024 ** 3) - device = torch.device(device) - if device.type != "cuda": - return requested - free, _total = torch.cuda.mem_get_info(device) - # Loaders are created sequentially. Capping each against currently free - # memory leaves room for the model, optimizer, and transient loss graphs on - # smaller cards while retaining the measured H100 defaults. - return max(CHUNK_VOXELS, min(requested, int(free * fractions[kind]))) - - -class BoundedSparseCudaCache: - """Fixed-size device chunk pool backed by TensorStore and LRU eviction. - - ``source_paths`` are Zarr array directories, one per output channel. Gather indices are local to the configured ROI; ``origin_zyx`` maps them - back to the full source arrays. + back to the full source arrays. ``z_roi`` (full-array grid voxels, + half-open) restricts which bricks are uploaded; gathers must stay inside + the ROI the surrounding volume dict advertises, as everywhere else. """ def __init__( self, + sidecar_dir: str, *, - source_paths: list[str], - shape_zyx: tuple[int, int, int], origin_zyx: tuple[int, int, int] = (0, 0, 0), - budget_bytes: int, + z_roi: tuple[int, int] | None = None, device: torch.device | str = "cuda", label: str, - io_batch_chunks: int = 256, - tensorstore_cache_bytes: int | None = None, + expected_channels: int | None = None, + expected_shape_zyx: tuple[int, int, int] | None = None, + progress_callback=None, ) -> None: - if not source_paths: - raise ValueError("source_paths must contain at least one Zarr array") - if budget_bytes <= 0: - raise ValueError("budget_bytes must be positive") - self.source_paths = [str(Path(path)) for path in source_paths] - self.shape_zyx = tuple(int(v) for v in shape_zyx) - self.origin_zyx = tuple(int(v) for v in origin_zyx) + started = time.perf_counter() + self.sidecar_dir = str(sidecar_dir) self.device = torch.device(device) self.label = str(label) - self.channels = len(self.source_paths) - self.io_batch_chunks = max(1, int(io_batch_chunks)) - - z, y, x = self.shape_zyx - self.chunk_grid_zyx = ( - (z + CHUNK_SIZE - 1) // CHUNK_SIZE, - (y + CHUNK_SIZE - 1) // CHUNK_SIZE, - (x + CHUNK_SIZE - 1) // CHUNK_SIZE, - ) - total_chunks = int(np.prod(self.chunk_grid_zyx, dtype=np.int64)) - bytes_per_slot = self.channels * CHUNK_VOXELS - self.capacity = min(total_chunks, max(1, int(budget_bytes) // bytes_per_slot)) - self.pool_bytes = self.capacity * bytes_per_slot + meta, table_np, coords_np, memmaps = open_pool(sidecar_dir) + self.meta = meta + self.shape_zyx = tuple(int(v) for v in meta["array_shape"]) + if (expected_shape_zyx is not None + and tuple(int(v) for v in expected_shape_zyx) != self.shape_zyx): + raise ValueError( + f"{label}: sidecar {sidecar_dir} covers array shape " + f"{self.shape_zyx}, expected {tuple(expected_shape_zyx)}") + self.channels = len(memmaps) + if expected_channels is not None and self.channels != expected_channels: + raise ValueError( + f"{label}: sidecar {sidecar_dir} has {self.channels} " + f"channel(s), expected {expected_channels}") + self.origin_zyx = tuple(int(v) for v in origin_zyx) + brick = tuple(int(v) for v in meta["brick_shape"]) + brick_voxels = int(np.prod(brick)) + rows = int(meta["rows"]) + + keep = np.ones(rows, dtype=bool) + if z_roi is not None: + z_lo, z_hi = int(z_roi[0]), int(z_roi[1]) + brick_z = coords_np[:, 0].astype(np.int64) + keep = (brick_z * brick[0] < z_hi) & ((brick_z + 1) * brick[0] > z_lo) + keep[0] = True # the reserved all-zero brick + kept_ids = np.flatnonzero(keep) + remap = np.zeros(rows, dtype=np.int32) + remap[kept_ids] = np.arange(len(kept_ids), dtype=np.int32) + + self.pool_bytes = self.channels * len(kept_ids) * brick_voxels try: self.pool = torch.empty( - (self.channels, self.capacity, CHUNK_VOXELS), + (self.channels, len(kept_ids), brick_voxels), dtype=torch.uint8, device=self.device, ) - self.chunk_table = torch.full( - self.chunk_grid_zyx, - -1, - dtype=torch.int32, - device=self.device, - ) except torch.OutOfMemoryError as exc: raise RuntimeError( - f"Could not allocate {self.label} sparse cache " - f"({self.pool_bytes / 1024**3:.2f} GiB, {self.capacity} slots)" + f"Could not allocate the {self.label} resident pool " + f"({self.pool_bytes / 1024**3:.2f} GiB, {len(kept_ids)} bricks)" ) from exc - - self._lru: OrderedDict[int, int] = OrderedDict() - self._free_slots = list(range(self.capacity - 1, -1, -1)) - self._hits = 0 - self._misses = 0 - self._evictions = 0 + slab = max(1, (256 << 20) // brick_voxels) + progress_total = self.channels * len(kept_ids) + progress_done = 0 + for channel, mm in enumerate(memmaps): + for lo in range(0, len(kept_ids), slab): + ids = kept_ids[lo:lo + slab] + self.pool[channel, lo:lo + len(ids)] = torch.from_numpy( + np.ascontiguousarray(mm[ids])).to(self.device) + progress_done += len(ids) + if progress_callback is not None: + progress_callback( + progress_done, + progress_total, + f"channel {channel + 1}/{self.channels}", + ) + self.table = torch.from_numpy( + np.ascontiguousarray(remap[table_np])).to(self.device) + + self._brick = torch.tensor(brick, dtype=torch.long, device=self.device) + self._origin = torch.tensor( + self.origin_zyx, dtype=torch.long, device=self.device) + self._shape = torch.tensor( + self.shape_zyx, dtype=torch.long, device=self.device) + self._stride_z = brick[1] * brick[2] + self._stride_y = brick[2] + self._bounds_check = ( + os.environ.get("FIT_SPIRAL_RESIDENT_BOUNDS_CHECK") == "1") + self.resident_bricks = len(kept_ids) + self.total_bricks = rows self._gathers = 0 - self._read_seconds = 0.0 self._gather_seconds = 0.0 + self.load_seconds = time.perf_counter() - started self.last_timings: dict[str, float | int] = {} - - import tensorstore as ts - - if tensorstore_cache_bytes is None: - tensorstore_cache_bytes = int( - _env_gib("FIT_SPIRAL_TENSORSTORE_CACHE_GB", 2.0) * 1024 ** 3 - ) - self._ts = ts - self._context = ts.Context({ - "cache_pool": {"total_bytes_limit": int(tensorstore_cache_bytes)}, - "file_io_concurrency": { - "limit": int(os.environ.get("FIT_SPIRAL_SPARSE_IO_THREADS", "16")) - }, - "data_copy_concurrency": { - "limit": int(os.environ.get("FIT_SPIRAL_SPARSE_COPY_THREADS", "8")) - }, - }) - self._stores = [ - ts.open( - { - "driver": "zarr", - "kvstore": {"driver": "file", "path": path}, - }, - context=self._context, - open=True, - read=True, - recheck_cached_data="open", - ).result() - for path in self.source_paths - ] print( - f"{self.label}: sparse CUDA LRU cache " - f"{self.capacity}x{CHUNK_SIZE}^3x{self.channels} " - f"({self.pool_bytes / 1024**3:.2f} GiB), " - f"grid={self.chunk_grid_zyx}", + f"{self.label}: resident pool {self.resident_bricks:,}/{rows:,} " + f"bricks of {brick} x{self.channels} channel(s) " + f"({self.pool_bytes / 1024**3:.2f} GiB) loaded in " + f"{self.load_seconds:.1f}s from {Path(sidecar_dir).name}", flush=True, ) - def _key_to_coord(self, key: int) -> tuple[int, int, int]: - _cz, cy_count, cx_count = self.chunk_grid_zyx - cz, rem = divmod(int(key), cy_count * cx_count) - cy, cx = divmod(rem, cx_count) - return cz, cy, cx - - def _plan_slots( - self, requested_keys: list[int] - ) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]: - """Return ``(loads, evictions)`` as ``(chunk_key, slot)`` pairs.""" - protected = set(requested_keys) - missing = [] - for key in requested_keys: - slot = self._lru.get(key) - if slot is None: - missing.append(key) - self._misses += 1 - else: - self._hits += 1 - self._lru.move_to_end(key) - if len(requested_keys) > self.capacity: - required_gib = ( - len(requested_keys) * self.channels * CHUNK_VOXELS / 1024 ** 3 - ) - raise RuntimeError( - f"{self.label} gather touches {len(requested_keys)} chunks " - f"({required_gib:.2f} GiB), exceeding its {self.capacity}-chunk " - f"LRU capacity; increase the corresponding FIT_SPIRAL_SPARSE_*_CACHE_GB" - ) - - loads: list[tuple[int, int]] = [] - evictions: list[tuple[int, int]] = [] - for key in missing: - if self._free_slots: - slot = self._free_slots.pop() - else: - victim = None - for candidate, candidate_slot in self._lru.items(): - if candidate not in protected: - victim = (candidate, candidate_slot) - break - if victim is None: - raise RuntimeError( - f"{self.label} could not find an evictable cache slot" - ) - victim_key, slot = victim - del self._lru[victim_key] - evictions.append((victim_key, slot)) - self._evictions += 1 - self._lru[key] = slot - loads.append((key, slot)) - return loads, evictions - - def _chunk_bounds(self, key: int): - cz, cy, cx = self._key_to_coord(key) - z0, y0, x0 = cz * CHUNK_SIZE, cy * CHUNK_SIZE, cx * CHUNK_SIZE - z1 = min(z0 + CHUNK_SIZE, self.shape_zyx[0]) - y1 = min(y0 + CHUNK_SIZE, self.shape_zyx[1]) - x1 = min(x0 + CHUNK_SIZE, self.shape_zyx[2]) - return (cz, cy, cx), (slice(z0, z1), slice(y0, y1), slice(x0, x1)) - - def _load_chunks(self, loads: list[tuple[int, int]]) -> None: - if not loads: - return - started = time.perf_counter() - for batch_lo in range(0, len(loads), self.io_batch_chunks): - batch_loads = loads[batch_lo:batch_lo + self.io_batch_chunks] - batch = self._ts.Batch() - pending = [] - bounds = [] - for key, _slot in batch_loads: - coord, slices = self._chunk_bounds(key) - bounds.append((coord, slices)) - pending.append([ - store[slices].read(order="C", batch=batch) - for store in self._stores - ]) - batch.submit() - - host = np.zeros( - (len(batch_loads), self.channels, CHUNK_SIZE, CHUNK_SIZE, CHUNK_SIZE), - dtype=np.uint8, - ) - for row, futures in enumerate(pending): - _coord, slices = bounds[row] - dz = slices[0].stop - slices[0].start - dy = slices[1].stop - slices[1].start - dx = slices[2].stop - slices[2].start - for channel, future in enumerate(futures): - host[row, channel, :dz, :dy, :dx] = np.asarray( - future.result(), dtype=np.uint8 - ) - - values = torch.from_numpy(host.reshape( - len(batch_loads), self.channels, CHUNK_VOXELS - )) - if self.device.type == "cuda": - values = values.pin_memory().to(self.device, non_blocking=True) - else: - values = values.to(self.device) - slots = torch.tensor( - [slot for _key, slot in batch_loads], - dtype=torch.long, - device=self.device, - ) - self.pool[:, slots, :] = values.permute(1, 0, 2) - coords = torch.tensor( - [coord for coord, _slices in bounds], - dtype=torch.long, - device=self.device, - ) - self.chunk_table[coords[:, 0], coords[:, 1], coords[:, 2]] = ( - slots.to(torch.int32) - ) - self._read_seconds += time.perf_counter() - started - def gather(self, indices_zyx: torch.Tensor) -> torch.Tensor: """Gather local ROI indices, returning ``(..., channels)`` uint8.""" started = time.perf_counter() original_shape = tuple(indices_zyx.shape[:-1]) - flat = indices_zyx.detach().reshape(-1, 3).to( - device=self.device, dtype=torch.long - ) - origin = torch.tensor( - self.origin_zyx, dtype=torch.long, device=self.device - ) - source = flat + origin - shape = torch.tensor( - self.shape_zyx, dtype=torch.long, device=self.device - ) - if bool(((source < 0) | (source >= shape)).any()): - raise IndexError(f"{self.label} gather received an out-of-bounds index") - - chunk = torch.div(source, CHUNK_SIZE, rounding_mode="floor") - _cz_count, cy_count, cx_count = self.chunk_grid_zyx - keys = (chunk[:, 0] * cy_count + chunk[:, 1]) * cx_count + chunk[:, 2] - unique_keys = torch.unique(keys) - key_list = [int(v) for v in unique_keys.cpu().tolist()] - loads, evictions = self._plan_slots(key_list) - if evictions: - evicted_coords = torch.tensor( - [self._key_to_coord(key) for key, _slot in evictions], - dtype=torch.long, - device=self.device, - ) - self.chunk_table[ - evicted_coords[:, 0], evicted_coords[:, 1], evicted_coords[:, 2] - ] = -1 - self._load_chunks(loads) - - slots = self.chunk_table[ - chunk[:, 0], chunk[:, 1], chunk[:, 2] + flat = indices_zyx.detach().reshape(-1, 3) + if flat.shape[0] == 0: + return torch.empty( + (*original_shape, self.channels), + dtype=torch.uint8, device=self.device) + flat = flat.to(device=self.device, dtype=torch.long) + source = flat + self._origin + if self._bounds_check and bool( + ((source < 0) | (source >= self._shape)).any()): + raise IndexError( + f"{self.label} gather received an out-of-bounds index") + brick_idx = torch.div(source, self._brick, rounding_mode="floor") + slots = self.table[ + brick_idx[:, 0], brick_idx[:, 1], brick_idx[:, 2] ].to(torch.long) - if bool((slots < 0).any()): - raise RuntimeError(f"{self.label} cache miss remained after loading") - local = source - chunk * CHUNK_SIZE - linear = (local[:, 0] * CHUNK_SIZE + local[:, 1]) * CHUNK_SIZE + local[:, 2] + local = source - brick_idx * self._brick + linear = (local[:, 0] * self._stride_z + + local[:, 1] * self._stride_y + local[:, 2]) values = self.pool[:, slots, linear].transpose(0, 1) + # Host-side timer only (kernels may still be in flight); never sync. elapsed = time.perf_counter() - started self._gather_seconds += elapsed self._gathers += 1 self.last_timings = { "gather_seconds": elapsed, - "requested_chunks": len(key_list), - "loaded_chunks": len(loads), - "resident_chunks": len(self._lru), - "evictions": len(evictions), - "resident_mib": len(self._lru) * self.channels * CHUNK_VOXELS / 1024 ** 2, + "resident_bricks": self.resident_bricks, + "resident_mib": self.pool_bytes / 1024 ** 2, } return values.reshape(*original_shape, self.channels) def stats(self) -> dict[str, float | int]: - total = self._hits + self._misses return { - "capacity_chunks": self.capacity, - "resident_chunks": len(self._lru), - "hits": self._hits, - "misses": self._misses, - "hit_rate": self._hits / total if total else 0.0, - "evictions": self._evictions, + "resident_bricks": self.resident_bricks, + "total_bricks": self.total_bricks, "gathers": self._gathers, - "read_seconds": self._read_seconds, "gather_seconds": self._gather_seconds, + "load_seconds": self.load_seconds, "pool_bytes": self.pool_bytes, } def close(self) -> None: # Drop the owning references even when the surrounding volume dict # remains alive (notably in an interactive session reload). - self._stores.clear() - self._lru.clear() - self._free_slots.clear() self.pool = None - self.chunk_table = None - self._context = None + self.table = None class SparseLasagnaStore: def __init__( self, *, - normal_cache: BoundedSparseCudaCache | None, - grad_cache: BoundedSparseCudaCache | None, + normal_cache: ResidentBrickPool | None, + grad_cache: ResidentBrickPool | None, ) -> None: self.normal_cache = normal_cache self.grad_cache = grad_cache @@ -368,7 +197,7 @@ def gather_pair(self, normal_zyx, grad_zyx, device): if normal_zyx.numel(): if self.normal_cache is None: raise RuntimeError("normal cache is not configured") - normals = self.normal_cache.gather(normal_zyx) + normals = self.normal_cache.gather(normal_zyx).to(device) else: normals = torch.empty( (*normal_zyx.shape[:-1], 2), dtype=torch.uint8, device=device @@ -376,7 +205,7 @@ def gather_pair(self, normal_zyx, grad_zyx, device): if grad_zyx.numel(): if self.grad_cache is None: raise RuntimeError("gradient-magnitude cache is not configured") - gradient = self.grad_cache.gather(grad_zyx)[..., 0] + gradient = self.grad_cache.gather(grad_zyx)[..., 0].to(device) else: gradient = torch.empty( grad_zyx.shape[:-1], dtype=torch.uint8, device=device @@ -402,12 +231,12 @@ def close(self): class SparseScalarStore: - def __init__(self, cache: BoundedSparseCudaCache) -> None: + def __init__(self, cache: ResidentBrickPool) -> None: self.cache = cache self.last_timings: dict[str, float | int] = {} def gather(self, indices_zyx, device): - values = self.cache.gather(indices_zyx)[..., 0] + values = self.cache.gather(indices_zyx)[..., 0].to(device) self.last_timings = dict(self.cache.last_timings) return values diff --git a/volume-cartographer/scripts/spiral/spiral_helpers.py b/volume-cartographer/scripts/spiral/spiral_helpers.py index 0f5784c8ba..fd2ef01c8f 100644 --- a/volume-cartographer/scripts/spiral/spiral_helpers.py +++ b/volume-cartographer/scripts/spiral/spiral_helpers.py @@ -50,8 +50,8 @@ def scale_counts_for_z_range( SAMPLING_COUNT_FLOORS = { - 'dense_spacing_num_pairs': 8_000, - 'dense_spacing_density_extra_pairs': 16_000, + 'sample_count_dense_spacing_pairs': 8_000, + 'sample_count_dense_spacing_density_extra_pairs': 16_000, } @@ -110,7 +110,7 @@ def load_fiber_point_collection(path, collection_id, coordinate_scale=0.25, min_ 'fiber_tags': data.get('tags', []), 'hv_classification': data.get('hv_classification', {}), 'input_coordinate_scale': coordinate_scale, - 'fiber_min_point_spacing': min_point_spacing, + 'pcl_fiber_min_point_spacing': min_point_spacing, 'fiber_original_num_points': original_num_points, }, 'color': color, @@ -429,12 +429,13 @@ def resolve_outer_winding_idx_and_notes(cfg, shell_active, infer_outer_winding_i f'= {idx} for the dense and regularisation losses') if idx is not None: min_gap = idx + 3 - if cfg['gap_expander_num_windings'] < min_gap: + gap_windings = cfg['model_gap_expander_num_windings'] + if gap_windings < min_gap: notes.append( f'WARNING: shell_outer_winding_idx {idx} requires ' - f'gap_expander_num_windings >= {min_gap}, got ' - f'gap_expander_num_windings {cfg["gap_expander_num_windings"]}; ' - 'increase gap_expander_num_windings or lower ' + f'model_gap_expander_num_windings >= {min_gap}, got ' + f'model_gap_expander_num_windings {gap_windings}; ' + 'increase model_gap_expander_num_windings or lower ' 'shell_outer_winding_idx') return idx, notes @@ -472,7 +473,7 @@ def _warn_if_inputs_exceed_flow_bounds( flow_field_radius, cfg, ): - gap_expander_num_windings = cfg['gap_expander_num_windings'] + gap_expander_num_windings = cfg['model_gap_expander_num_windings'] over_radius_patches = [] over_winding_patches = [] @@ -727,6 +728,7 @@ def save_mesh( tracks=(), run_tag=None, name='mesh', + progress=None, ): (min_winding_idx, max_winding_idx), _, _ = compute_winding_range_and_input_extents( slice_to_spiral_transform, @@ -743,7 +745,7 @@ def save_mesh( max_winding_idx = min(max_winding_idx, cfg['shell_outer_winding_idx']) print(f'save_mesh {name}: winding range [{min_winding_idx}, {max_winding_idx})') grid_spacing = cfg['output_step_size'] - z_margin = cfg['flow_bounds_z_margin'] + z_margin = cfg['model_flow_bounds_z_margin'] spiral_yxs_by_winding = get_spiral_yxs(max_winding_idx, dr_per_winding, grid_spacing, group_by_winding=True) num_thetas_by_winding = [len(yxs_for_winding) for yxs_for_winding in spiral_yxs_by_winding] spiral_yxs = torch.cat(spiral_yxs_by_winding, dim=0) @@ -753,8 +755,17 @@ def save_mesh( chunk = 65536 flat_spiral_zyxs = spiral_zyxs.reshape(-1, 3) scroll_pieces = [] - for start in range(0, flat_spiral_zyxs.shape[0], chunk): + transform_chunk_total = ( + flat_spiral_zyxs.shape[0] + chunk - 1) // chunk + if progress is not None: + progress.begin( + 'finalizing', 'Transforming final mesh', + step=0, total_steps=transform_chunk_total, unit='chunks') + for chunk_number, start in enumerate( + range(0, flat_spiral_zyxs.shape[0], chunk), start=1): scroll_pieces.append(slice_to_spiral_transform.inv(flat_spiral_zyxs[start : start + chunk])) + if progress is not None: + progress.update(chunk_number) scroll_zyxs = torch.cat(scroll_pieces, dim=0).reshape(*spiral_zyxs.shape) out_of_roi = (scroll_zyxs[..., 0] < z_begin) | (scroll_zyxs[..., 0] >= z_end) @@ -784,9 +795,18 @@ def save_mesh( tag_suffix = f'_{run_tag}' if run_tag else '' out_dir = f'{out_path}/meshes/{name}{tag_suffix}' os.makedirs(out_dir, exist_ok=True) + output_total = 2 * len(num_thetas_by_winding) + output_done = 0 + if progress is not None: + progress.begin( + 'finalizing', 'Writing final mesh windings', + step=0, total_steps=output_total, unit='windings') for uuid_suffix, variant_zyxs in [('', scroll_zyxs), ('_spliced', spliced_scroll_zyxs)]: offset = 0 - for winding_idx, num_thetas in enumerate(tqdm(num_thetas_by_winding, desc=f'saving winding patches ({name}{uuid_suffix})')): + for winding_idx, num_thetas in enumerate(tqdm( + num_thetas_by_winding, + desc=f'saving winding patches ({name}{uuid_suffix})', + disable=progress is not None)): if num_thetas >= 2 and winding_idx >= min_winding_idx: winding_slice = variant_zyxs[:, offset:offset + num_thetas] invalid_mask = (winding_slice == -1.0).all(dim=-1).cpu().numpy() @@ -801,6 +821,10 @@ def save_mesh( source=f'fit_spiral {name}{uuid_suffix}', ) offset += num_thetas + output_done += 1 + if progress is not None: + progress.update( + output_done, detail=f'winding {winding_idx}{uuid_suffix}') @torch.inference_mode() @@ -818,6 +842,7 @@ def save_combined_preview( tracks=(), *, surface_id, + progress=None, ): """Write the authoritative connected preview used by VC3D and Lasagna.""" (_, derived_upper), _, _ = compute_winding_range_and_input_extents( @@ -845,7 +870,7 @@ def save_combined_preview( ) grid_spacing = int(cfg['output_step_size']) - z_margin = int(cfg['flow_bounds_z_margin']) + z_margin = int(cfg['model_flow_bounds_z_margin']) spiral_yxs_by_winding = get_spiral_yxs( last_winding + 1, dr_per_winding, @@ -861,7 +886,13 @@ def save_combined_preview( device=dr_per_winding.device, ) winding_grids = {} - for winding in range(first_winding, last_winding + 1): + total_windings = last_winding - first_winding + 1 + if progress is not None: + progress.begin( + 'exporting_preview', 'Transforming preview windings', + step=0, total_steps=total_windings, unit='windings') + for winding_number, winding in enumerate( + range(first_winding, last_winding + 1), start=1): yxs = spiral_yxs_by_winding[winding] if yxs.shape[0] < 2: raise RuntimeError(f'Preview winding {winding} has fewer than two theta samples') @@ -877,7 +908,13 @@ def save_combined_preview( outside = (scroll[..., 0] < z_begin) | (scroll[..., 0] >= z_end) scroll[outside] = -1.0 winding_grids[winding] = scroll.cpu().numpy().astype(np.float32) + if progress is not None: + progress.update(winding_number, detail=f'winding {winding}') + if progress is not None: + progress.begin( + 'exporting_preview', 'Writing preview surface', + detail=f'{total_windings} windings') manifest = save_combined_tifxyz( winding_grids, generation_path, diff --git a/volume-cartographer/scripts/spiral/spiral_progress.py b/volume-cartographer/scripts/spiral/spiral_progress.py new file mode 100644 index 0000000000..d6f9e7719c --- /dev/null +++ b/volume-cartographer/scripts/spiral/spiral_progress.py @@ -0,0 +1,251 @@ +"""Shared, dependency-free progress reporting for Spiral operations.""" + +from __future__ import annotations + +import copy +import sys +import threading +import time +from typing import Callable, Mapping, TextIO + + +def _format_duration(seconds: float | None) -> str: + if seconds is None: + return "" + seconds = max(0, int(round(seconds))) + if seconds < 60: + return f"{seconds}s" + minutes, seconds = divmod(seconds, 60) + if minutes < 60: + return f"{minutes}m {seconds:02d}s" + hours, minutes = divmod(minutes, 60) + return f"{hours}h {minutes:02d}m" + + +class ProgressReporter: + """Own one active stage and publish rate-limited status snapshots. + + Updating the reporter never synchronizes CUDA or inspects worker state. + ``snapshot`` derives elapsed time from the monotonic clock, so a polling + client continues to see useful elapsed time while an opaque native call is + in flight. + """ + + def __init__( + self, + publish: Callable[[Mapping[str, object] | None], None] | None = None, + *, + stream: TextIO | None = None, + publish_interval: float = 1.0, + console_interval: float = 5.0, + heartbeat_interval: float = 30.0, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._publish = publish + self._stream = stream + self._publish_interval = float(publish_interval) + self._console_interval = float(console_interval) + self._heartbeat_interval = float(heartbeat_interval) + self._clock = clock + self._lock = threading.Lock() + self._active: dict[str, object] | None = None + self._started = 0.0 + self._last_publish = float("-inf") + self._last_console = float("-inf") + self._closed = False + self._tty = bool(stream is not None and getattr(stream, "isatty", lambda: False)()) + self._heartbeat_thread: threading.Thread | None = None + if stream is not None and heartbeat_interval > 0: + self._heartbeat_thread = threading.Thread( + target=self._heartbeat_main, + name="spiral-progress-heartbeat", + daemon=True, + ) + self._heartbeat_thread.start() + + def begin( + self, + operation: str, + stage_name: str, + *, + step: int | None = None, + total_steps: int | None = None, + unit: str | None = None, + detail: str | None = None, + ) -> None: + now = self._clock() + with self._lock: + self._started = now + self._active = { + "operation": str(operation), + "stage_name": str(stage_name), + "detail": None if detail is None else str(detail), + "step": None if step is None else int(step), + "total_steps": ( + None if total_steps is None else max(0, int(total_steps)) + ), + "unit": None if unit is None else str(unit), + } + self._emit(force=True) + + def update( + self, + step: int | None = None, + *, + total_steps: int | None = None, + detail: str | None = None, + ) -> None: + with self._lock: + if self._active is None: + return + if step is not None: + self._active["step"] = int(step) + if total_steps is not None: + self._active["total_steps"] = max(0, int(total_steps)) + if detail is not None: + self._active["detail"] = str(detail) + self._emit(force=False) + + def pulse(self, detail: str | None = None) -> None: + self.update(detail=detail) + + def finish(self, detail: str | None = None) -> None: + with self._lock: + if self._active is None: + return + total = self._active.get("total_steps") + if total is not None: + self._active["step"] = int(total) + if detail is not None: + self._active["detail"] = str(detail) + self._emit(force=True) + + def clear(self) -> None: + with self._lock: + had_active = self._active is not None + self._active = None + if had_active: + if self._tty and self._stream is not None: + self._stream.write("\n") + self._stream.flush() + if self._publish is not None: + self._publish(None) + + def snapshot(self) -> dict[str, object] | None: + now = self._clock() + with self._lock: + if self._active is None: + return None + result = copy.deepcopy(self._active) + elapsed = max(0.0, now - self._started) + result["elapsed_seconds"] = elapsed + step = result.get("step") + total = result.get("total_steps") + eta = None + if ( + isinstance(step, int) + and isinstance(total, int) + and step > 0 + and total > step + and elapsed >= 2.0 + ): + eta = elapsed * (total - step) / step + elif ( + isinstance(step, int) + and isinstance(total, int) + and total > 0 + and step >= total + ): + eta = 0.0 + result["eta_seconds"] = eta + return result + + def close(self) -> None: + self._closed = True + self.clear() + thread = self._heartbeat_thread + if thread is not None and thread is not threading.current_thread(): + thread.join(timeout=0.2) + + def _emit(self, *, force: bool) -> None: + now = self._clock() + snapshot = self.snapshot() + if snapshot is None: + return + publish = force or now - self._last_publish >= self._publish_interval + console = force or now - self._last_console >= self._console_interval + if publish and self._publish is not None: + self._last_publish = now + self._publish(snapshot) + if console and self._stream is not None: + self._last_console = now + self._write_console(snapshot) + + def _write_console(self, snapshot: Mapping[str, object]) -> None: + stage = str(snapshot["stage_name"]) + detail = snapshot.get("detail") + step, total = snapshot.get("step"), snapshot.get("total_steps") + unit = snapshot.get("unit") + parts = [f"PROGRESS {stage}"] + if isinstance(step, int) and isinstance(total, int) and total > 0: + label = f"{step:,}/{total:,}" + if unit: + label += f" {unit}" + label += f" ({100.0 * min(step, total) / total:.1f}%)" + parts.append(label) + elif detail: + parts.append(str(detail)) + parts.append(f"elapsed {_format_duration(float(snapshot['elapsed_seconds']))}") + eta = snapshot.get("eta_seconds") + if isinstance(eta, (int, float)): + parts.append(f"ETA {_format_duration(float(eta))}") + line = " — ".join(parts) + if self._tty: + self._stream.write("\r" + line) + else: + self._stream.write(line + "\n") + self._stream.flush() + + def _heartbeat_main(self) -> None: + while not self._closed: + time.sleep(min(1.0, self._heartbeat_interval)) + if self._closed: + return + now = self._clock() + with self._lock: + active = self._active is not None + due = now - self._last_console >= self._heartbeat_interval + if active and due: + snapshot = self.snapshot() + if snapshot is not None and self._stream is not None: + self._last_console = now + self._write_console(snapshot) + + +class NullProgressReporter: + """No-op reporter used by helper callers that do not request progress.""" + + def begin(self, *args, **kwargs) -> None: + pass + + def update(self, *args, **kwargs) -> None: + pass + + def pulse(self, *args, **kwargs) -> None: + pass + + def finish(self, *args, **kwargs) -> None: + pass + + def clear(self) -> None: + pass + + def snapshot(self): + return None + + def close(self) -> None: + pass + + +def progress_or_null(progress): + return progress if progress is not None else NullProgressReporter() diff --git a/volume-cartographer/scripts/spiral/spiral_runtime.py b/volume-cartographer/scripts/spiral/spiral_runtime.py index ed52821c18..ebed89cbad 100644 --- a/volume-cartographer/scripts/spiral/spiral_runtime.py +++ b/volume-cartographer/scripts/spiral/spiral_runtime.py @@ -11,16 +11,20 @@ import os from pathlib import Path import socket +import sys import threading import time import traceback from typing import Mapping import uuid -from fit_session import (RUN_MUTABLE_SAMPLING_KEYS, SpiralInputPaths, - SpiralPreviewConfig, SpiralRunConfig, - apply_optional_input_selection, +from fit_session import (SpiralInputPaths, SpiralPreviewConfig, SpiralRunConfig, run_mutable_config) +from config import Config +from spiral_progress import NullProgressReporter, ProgressReporter + + +_NO_PROGRESS = NullProgressReporter() class _SessionShutdown(BaseException): @@ -36,6 +40,7 @@ def __init__(self, paths: SpiralInputPaths, run: SpiralRunConfig, self.preview_config = preview self.input_manifest = paths.manifest() self.requested_config = dict(run.config) + self._applied_config = None self._run_config = None self._run_config_limits = None self._default_advanced_config = None @@ -62,6 +67,11 @@ def __init__(self, paths: SpiralInputPaths, run: SpiralRunConfig, self._finish_run = None self._configure_run = None self._idle_actions = [] + self.progress = ProgressReporter( + self._progress_changed, + stream=sys.stdout, + ) + self._run_start_completed = 0 self._thread = threading.Thread(target=self._fit_main, name="spiral-fit-worker", daemon=True) self._thread.start() @@ -82,6 +92,8 @@ def status(self): "preview_manifest_path": self._preview_manifest, "preview_generation": self._preview_generation, "supports_input_incorporation": self._incorporate_inputs is not None, + "input_manifest": copy.deepcopy(self.input_manifest), + "progress": self._progress_reporter().snapshot(), } if self._run_config is not None: result["run_config"] = copy.deepcopy(self._run_config) @@ -91,12 +103,23 @@ def status(self): if self._default_advanced_config is not None: result["default_advanced_config"] = copy.deepcopy( self._default_advanced_config) + if getattr(self, "_applied_config", None) is not None: + result["applied_config"] = copy.deepcopy(self._applied_config) return result def _publish_status(self): if self._status_callback is not None: self._status_callback(self.status()) + def _progress_reporter(self): + return getattr(self, "progress", _NO_PROGRESS) + + def _progress_changed(self, snapshot): + if snapshot is not None: + with self._condition: + self._phase = str(snapshot.get("stage_name") or self._phase) + self._publish_status() + def _set_state(self, state, phase=""): with self._condition: self._state = state @@ -109,6 +132,7 @@ def _fit_main(self): wandb = None distributed_initialized = False try: + self._progress_reporter().begin("loading", "Importing Torch and fitter") self._set_state("Loading", "Importing Torch and fitter") import wandb import fit_spiral as fitter @@ -121,29 +145,27 @@ def _fit_main(self): maybe_init_distributed() distributed_initialized = True - config = dict(fitter.default_config) + config = Config().as_dict() checkpoint_profile_config = None if self.paths.checkpoint: + self._progress_reporter().begin( + "loading", "Reading checkpoint configuration", + detail=Path(self.paths.checkpoint).name) from checkpoint_io import load_checkpoint_cpu checkpoint_config = load_checkpoint_cpu(self.paths.checkpoint) try: - if isinstance(checkpoint_config, dict) and isinstance(checkpoint_config.get('cfg'), Mapping): - # Influence settings describe one interactive Run, not - # durable model configuration. Ignore them even when - # opening checkpoints written by older services. - # Keys the current fitter no longer defines are - # dropped; old configurations are not supported. - durable = { - key: value for key, value in checkpoint_config['cfg'].items() - if not key.startswith('interactive_influence_') - and key != 'loss_weight_anchor' - and key in config - } - config.update(durable) - # The session-scoped Default profile must initially - # reproduce the checkpoint, not a second - # z-range/DDP-scaled or input-gated derivative of it. - checkpoint_profile_config = copy.deepcopy(durable) + if not isinstance(checkpoint_config, dict) or not isinstance( + checkpoint_config.get('cfg'), Mapping): + raise ValueError("Checkpoint has no current Spiral configuration") + durable = dict(checkpoint_config['cfg']) + if set(durable) != set(config): + raise ValueError( + "Checkpoint configuration does not match the current schema") + durable = Config(durable).as_dict() + config.update(durable) + # The session-scoped profile initially reproduces the + # checkpoint without applying scaling twice. + checkpoint_profile_config = copy.deepcopy(durable) finally: # This first load exists only to resolve configuration. Do # not retain a complete model + optimiser checkpoint for the @@ -155,15 +177,8 @@ def _fit_main(self): if checkpoint_profile_config is not None: default_advanced_config = checkpoint_profile_config else: - # Without a checkpoint, Default is the Python baseline adapted - # to the inputs selected for this session. + # Without a checkpoint, Default is the Python baseline. default_advanced_config = copy.deepcopy(config) - for key in ( - 'use_verified_patches', 'use_unverified_patches', - 'use_normals', 'use_surf_sdt', 'use_tracks', - 'use_gradient_magnitude', 'use_fibers'): - if key in self.run_config.config: - default_advanced_config[key] = self.run_config.config[key] # Explicit sample-count overrides are literal active counts. This # lets VC3D round-trip the host's post-scaling values through a # reload without applying the z-range/DDP transforms twice. @@ -171,24 +186,16 @@ def _fit_main(self): # their counts the same treatment. explicit_sampling_counts = { key: value for key, value in (checkpoint_profile_config or {}).items() - if key in RUN_MUTABLE_SAMPLING_KEYS + if key.startswith("sample_count_") } explicit_sampling_counts.update({ key: value for key, value in self.run_config.config.items() - if key in RUN_MUTABLE_SAMPLING_KEYS + if key.startswith("sample_count_") }) config.update(self.run_config.config) - count_keys = ( - 'num_patches_per_step', 'num_patches_per_step_for_dt', - 'unverified_num_patches_per_step', 'unverified_num_patches_per_step_for_dt', - 'rel_winding_num_pcls', 'abs_winding_num_pcls', - 'unattached_pcl_num_per_step', 'track_num_per_step', - 'dense_normals_num_points', 'dense_spacing_num_pairs', - 'dense_spacing_density_extra_pairs', - 'dense_attachment_num_points', - 'min_spacing_independent_samples', - 'regularisation_num_points', 'shell_num_samples', - ) + fields = Config.catalog()["schema"]["fields"] + count_keys = tuple( + key for key, spec in fields.items() if spec.get("scale_with_z")) scale_counts_for_z_range( config, self.run_config.z_begin, self.run_config.z_end, 9500, count_keys, floors=SAMPLING_COUNT_FLOORS) @@ -200,16 +207,14 @@ def _fit_main(self): floors=SAMPLING_COUNT_FLOORS) split_counts_across_ranks(default_advanced_config, count_keys) config.update(explicit_sampling_counts) - if checkpoint_profile_config is None: - apply_optional_input_selection(default_advanced_config) - apply_optional_input_selection(config) self.requested_config = dict(config) with self._condition: + self._applied_config = copy.deepcopy(config) self._run_config = run_mutable_config(config) self._run_config_limits = { - 'max_track_crossing_per_step': max( + 'track_max_track_crossing_per_step': max( int(config.get('track_crossing_precompute_max', 0)), - int(config.get('max_track_crossing_per_step', 0))), + int(config.get('track_max_track_crossing_per_step', 0))), } self._default_advanced_config = default_advanced_config self._publish_status() @@ -251,8 +256,9 @@ def _fit_main(self): wandb.init(project='scrolls', config=config, mode='disabled') fitter.cfg = wandb.config fitter.configure_losses(fitter.cfg, fitter.z_begin, fitter.z_end) + self._progress_reporter().begin("loading", "Loading fit inputs and model") self._set_state("Loading", "Loading fit inputs and model") - fitter.main(interactive_driver=self) + fitter.main(interactive_driver=self, progress=self.progress) except BaseException as exc: with self._condition: if self._shutdown and isinstance(exc, _SessionShutdown): @@ -270,6 +276,7 @@ def _fit_main(self): wandb.finish(quiet=True) if distributed_initialized: maybe_destroy_distributed() + self._progress_reporter().close() # Fitter-thread callbacks. def on_ready(self, *, completed_iterations, output_path, @@ -285,7 +292,10 @@ def on_ready(self, *, completed_iterations, output_path, self._configure_run = configure_run if self.paths.checkpoint and getattr(self, "publishes_outputs", True): self._set_state("ExportingPreview", "Exporting restored checkpoint preview") + self._progress_reporter().begin( + "exporting_preview", "Exporting restored checkpoint preview") self._publish_preview() + self._progress_reporter().clear() self._set_state("Ready", "Ready") def wait_for_iteration(self, iteration): @@ -306,17 +316,30 @@ def wait_for_iteration(self, iteration): self._run_incorporation(records, mark_incorporated, influence_config) continue if action[0] == "configure": - self._run_configuration(action[1]) + self._run_configuration(action[1], action[2], action[3]) continue kind, path, done, result = action try: if kind == "save": + with self._condition: + previous_state = self._state + self._state, self._phase = "Saving", "Saving checkpoint" + self._progress_reporter().begin( + "saving_checkpoint", "Saving checkpoint", + detail=Path(path).name) result["path"] = self._save_checkpoint(path, self._completed) + self._progress_reporter().finish() else: result["error"] = f"Unknown idle action {kind}" except Exception as exc: result["error"] = f"{type(exc).__name__}: {exc}" finally: + if kind == "save": + self._progress_reporter().clear() + with self._condition: + self._state = previous_state + self._phase = previous_state + self._publish_status() done.set() def _run_incorporation(self, records, mark_incorporated, influence_config): @@ -331,10 +354,14 @@ def _run_incorporation(self, records, mark_incorporated, influence_config): "The resident fitter does not support adding inputs to a running session") with self._condition: self._phase = "Incorporating new session inputs" + self._progress_reporter().begin( + "incorporating_inputs", "Incorporating new session inputs", + step=0, total_steps=len(records), unit="inputs") self._publish_status() self._incorporate_inputs(records, influence_config) except Exception as exc: error = f"{type(exc).__name__}: {exc}" + self._progress_reporter().clear() with self._condition: self._pending = 0 self._target = self._completed @@ -350,17 +377,32 @@ def _run_incorporation(self, records, mark_incorporated, influence_config): with self._condition: if self._state == "Running": self._phase = "Optimizing" + if getattr(self, "_state", None) == "Running": + self._begin_optimization_progress() + else: + self._progress_reporter().clear() self._publish_status() - def _run_configuration(self, config): + def _run_configuration( + self, config, path_changes=None, previous_run_config=None): """Apply validated Run-scoped settings on the fitter thread.""" try: if self._configure_run is None: raise RuntimeError( "The resident fitter does not support Run configuration changes") - self._configure_run(dict(config)) + path_changes = dict(path_changes or {}) + self._progress_reporter().begin( + "configuring", "Applying run configuration", + detail=( + f"{len(config)} settings, {len(path_changes)} path changes" + )) + if path_changes: + self._configure_run(dict(config), dict(path_changes)) + else: + self._configure_run(dict(config)) except Exception as exc: error = f"{type(exc).__name__}: {exc}" + self._progress_reporter().clear() with self._condition: self._pending = 0 self._target = self._completed @@ -370,9 +412,32 @@ def _run_configuration(self, config): if action[0] != "incorporate" ] self._warnings.append(f"Run configuration failed: {error}") + if previous_run_config is not None: + self._run_config.update(previous_run_config) self._state, self._phase = "Paused", "Paused" self._condition.notify_all() self._publish_status() + else: + if getattr(self, "_applied_config", None) is not None: + with self._condition: + self._applied_config.update(config) + self._run_config.update(config) + self.requested_config.update(config) + self.input_manifest.update(path_changes) + if getattr(self, "_state", None) == "Running": + self._begin_optimization_progress() + else: + self._progress_reporter().clear() + + def _begin_optimization_progress(self): + with self._condition: + run_start = getattr( + self, "_run_start_completed", self._completed) + step = max(0, self._completed - run_start) + total = max(0, self._target - run_start) + self._progress_reporter().begin( + "optimizing", "Optimizing", + step=step, total_steps=total, unit="iterations") def iteration_completed(self, *, completed_iterations, total_loss, losses, learning_rate, metrics=None): with self._condition: @@ -384,18 +449,30 @@ def iteration_completed(self, *, completed_iterations, total_loss, losses, learn self._pending = 0 self._stop_requested = False pause = self._pending == 0 + run_start = getattr( + self, "_run_start_completed", + self._completed - max(0, self._target - self._completed)) + run_step = max(0, self._completed - run_start) + self._progress_reporter().update(run_step) self._publish_status() if pause: if self._finish_run is not None: self._finish_run() if not getattr(self, "publishes_outputs", True): + self._progress_reporter().clear() self._set_state("Paused", "Paused") return self._set_state("Saving", "Autosaving checkpoint") + self._progress_reporter().begin( + "saving_checkpoint", "Autosaving checkpoint", + detail="checkpoint_autosave.ckpt") autosave = str(Path(self._output_path) / "checkpoint_autosave.ckpt") self._save_checkpoint(autosave, self._completed) self._set_state("ExportingPreview", "Exporting preview") + self._progress_reporter().begin( + "exporting_preview", "Exporting preview") self._publish_preview() + self._progress_reporter().clear() self._set_state("Paused", "Paused") def _publish_preview(self): @@ -419,24 +496,46 @@ def session_finished(self): # Coordinator-thread commands. def run(self, count, pending_inputs=None, mark_incorporated=None, - influence_config=None, run_config=None): + influence_config=None, run_config=None, path_changes=None): if count < 1: raise ValueError("iterations must be at least 1") with self._condition: if self._state not in {"Ready", "Paused"}: raise RuntimeError(f"Run is not allowed while session state is {self._state}") run_config = dict(run_config or {}) - if run_config: + path_changes = dict(path_changes or {}) + target = self._completed + count + requested_config = dict( + getattr(self, "requested_config", {}) or {}) + requested_config.update(run_config) + configured_horizon = int( + requested_config.get("optimizer_num_training_steps", 0) or 0) + # Interactive sessions are allowed to continue beyond the original + # headless horizon. Preserve the current LR curve while the whole + # requested run fits within it. When a run would cross the horizon, + # extend the horizon by the requested count and realign the + # exponential curve at the durable completed step. + if (getattr(self, "_configure_run", None) is not None + and target > configured_horizon): + run_config["optimizer_num_training_steps"] = ( + max(configured_horizon, self._completed) + count) + if run_config or path_changes: if self._configure_run is None: raise RuntimeError( "The resident fitter does not support Run configuration changes") - self.requested_config.update(run_config) - apply_optional_input_selection(self.requested_config) + requested_config = dict(self.requested_config) + requested_config.update(run_config) run_config = { - key: self.requested_config[key] + key: requested_config[key] + for key in run_config + } + previous_run_config = { + key: self._run_config.get(key) for key in run_config } - self._idle_actions.append(("configure", run_config)) + self._idle_actions.append( + ("configure", run_config, path_changes, + previous_run_config)) self._run_config.update(run_config) if pending_inputs: if self._incorporate_inputs is None: @@ -446,8 +545,10 @@ def run(self, count, pending_inputs=None, mark_incorporated=None, ("incorporate", list(pending_inputs), mark_incorporated, dict(influence_config or {}))) self._pending = count - self._target = self._completed + count + self._run_start_completed = self._completed + self._target = target self._state, self._phase = "Running", "Optimizing" + self._begin_optimization_progress() self._condition.notify_all() return self._target @@ -523,6 +624,7 @@ def mark_incorporated(records, error=None, cid=command_id): mark_incorporated=mark_incorporated, influence_config=arguments.get("influence_config"), run_config=arguments.get("run_config"), + path_changes=arguments.get("path_changes"), ) elif name == "stop": result = session.stop() @@ -567,6 +669,16 @@ def __init__(self, paths, run, preview, gpu_ids, status_callback=None): "error": None, "preview_manifest_path": None, "preview_generation": 0, "supports_input_incorporation": False, + "progress": { + "operation": "loading", + "stage_name": "Starting GPU workers", + "detail": None, + "step": 0, + "total_steps": len(self._gpu_ids), + "unit": "workers", + "elapsed_seconds": 0.0, + "eta_seconds": None, + }, } self._acks = {} self._incorporation_callbacks = {} @@ -653,11 +765,31 @@ def _listen(self): for item in self._rank_statuses.values()) ) if rank_zero.get("state") in ready_states and not all_ranks_ready: + unfinished = next( + ((worker_rank, item) + for worker_rank, item + in sorted(self._rank_statuses.items()) + if item.get("state") not in ready_states), + None) self._status = copy.deepcopy(rank_zero) self._status.update({ "state": "Loading", "phase": "Waiting for all GPU workers", }) + if unfinished is not None: + worker_rank, worker_status = unfinished + worker_progress = copy.deepcopy( + worker_status.get("progress")) + if worker_progress: + detail = worker_progress.get("detail") + worker_progress["detail"] = ( + f"GPU worker {worker_rank + 1}/" + f"{len(self._gpu_ids)}" + + (f" — {detail}" if detail else "")) + self._status["progress"] = worker_progress + self._status["phase"] = str( + worker_progress.get("stage_name") + or self._status["phase"]) elif all_ranks_ready: # Rank zero owns user-facing metrics and artifacts. # A later secondary Ready event completes startup. @@ -730,7 +862,7 @@ def _call(self, name, arguments=None, ranks=None, timeout=30.0, return responses[ranks[0]][1] def run(self, count, pending_inputs=None, mark_incorporated=None, - influence_config=None, run_config=None): + influence_config=None, run_config=None, path_changes=None): state = self.status()["state"] if state not in {"Ready", "Paused"}: raise RuntimeError(f"Run is not allowed while session state is {state}") @@ -739,6 +871,7 @@ def run(self, count, pending_inputs=None, mark_incorporated=None, "pending_inputs": list(pending_inputs or []), "influence_config": dict(influence_config or {}), "run_config": dict(run_config or {}), + "path_changes": dict(path_changes or {}), } return self._call("run", arguments, timeout=30.0, incorporation_callback=mark_incorporated) diff --git a/volume-cartographer/scripts/spiral/spiral_service.py b/volume-cartographer/scripts/spiral/spiral_service.py index c372c67829..7f974ca9aa 100644 --- a/volume-cartographer/scripts/spiral/spiral_service.py +++ b/volume-cartographer/scripts/spiral/spiral_service.py @@ -16,6 +16,7 @@ import argparse from collections import OrderedDict, deque +from concurrent.futures import ThreadPoolExecutor import errno import hashlib import json @@ -43,14 +44,13 @@ from PIL import Image import scipy.ndimage -from fit_session import (API_VERSION, PclRole, RUN_MUTABLE_BOOLEAN_KEYS, - RUN_MUTABLE_SAMPLING_KEYS, - parse_session_request, +from fit_session import (API_VERSION, PclRole, parse_session_request, resolve_dataset_root, validate_checkpoint_container, validate_session_request) +from config import Config -SERVICE_VERSION = "6.0.0" +SERVICE_VERSION = "6.1.0" MAX_BODY_BYTES = 4 * 1024 * 1024 MAX_DEDUPLICATED_COMMANDS = 256 TRANSFER_CHUNK_BYTES = 1024 * 1024 @@ -203,17 +203,17 @@ def _validate_run_influence_config(value): raise ApiError(HTTPStatus.BAD_REQUEST, "influence_config must be a JSON object") allowed = { - "interactive_influence_enabled", - "interactive_influence_z", - "interactive_influence_windings", - "interactive_influence_theta_frac", - "interactive_influence_disable_dt_frac", - "interactive_influence_sigma", - "interactive_influence_footprint_points", - "interactive_influence_anchor_lattice_points", - "interactive_influence_anchor_geometry_points", - "interactive_influence_anchor_samples_per_step", - "interactive_influence_anchor_ramp_power", + "influence_enabled", + "influence_z", + "influence_windings", + "influence_theta_frac", + "influence_disable_dt_frac", + "influence_sigma", + "sample_count_influence_footprint_points", + "sample_count_influence_anchor_lattice_points", + "sample_count_influence_anchor_geometry_points", + "sample_count_influence_anchor_samples_per_step", + "influence_anchor_ramp_power", "loss_weight_anchor", } unknown = sorted(set(value) - allowed) @@ -221,23 +221,23 @@ def _validate_run_influence_config(value): raise ApiError(HTTPStatus.BAD_REQUEST, f"Unknown influence configuration keys: {unknown}") result = {} - if "interactive_influence_enabled" in value: - enabled = value["interactive_influence_enabled"] + if "influence_enabled" in value: + enabled = value["influence_enabled"] if not isinstance(enabled, bool): raise ApiError(HTTPStatus.BAD_REQUEST, "interactive_influence_enabled must be boolean") - result["interactive_influence_enabled"] = enabled + result["influence_enabled"] = enabled ranges = { - "interactive_influence_z": (1.0, 1_000_000.0), - "interactive_influence_windings": (0.1, 100.0), - "interactive_influence_theta_frac": (0.01, 1.0), - "interactive_influence_disable_dt_frac": (0.0, 1.0), - "interactive_influence_sigma": (0.000001, 10.0), - "interactive_influence_footprint_points": (1.0, 1_000_000.0), - "interactive_influence_anchor_lattice_points": (1.0, 1_000_000.0), - "interactive_influence_anchor_geometry_points": (1.0, 100_000.0), - "interactive_influence_anchor_samples_per_step": (1.0, 1_000_000.0), - "interactive_influence_anchor_ramp_power": (0.000001, 100.0), + "influence_z": (1.0, 1_000_000.0), + "influence_windings": (0.1, 100.0), + "influence_theta_frac": (0.01, 1.0), + "influence_disable_dt_frac": (0.0, 1.0), + "influence_sigma": (0.000001, 10.0), + "sample_count_influence_footprint_points": (1.0, 1_000_000.0), + "sample_count_influence_anchor_lattice_points": (1.0, 1_000_000.0), + "sample_count_influence_anchor_geometry_points": (1.0, 100_000.0), + "sample_count_influence_anchor_samples_per_step": (1.0, 1_000_000.0), + "influence_anchor_ramp_power": (0.000001, 100.0), "loss_weight_anchor": (0.0, 10_000.0), } for key, (minimum, maximum) in ranges.items(): @@ -252,10 +252,10 @@ def _validate_run_influence_config(value): f"{key} must be between {minimum} and {maximum}") result[key] = number integer_keys = { - "interactive_influence_footprint_points", - "interactive_influence_anchor_lattice_points", - "interactive_influence_anchor_geometry_points", - "interactive_influence_anchor_samples_per_step", + "sample_count_influence_footprint_points", + "sample_count_influence_anchor_lattice_points", + "sample_count_influence_anchor_geometry_points", + "sample_count_influence_anchor_samples_per_step", } for key in integer_keys & result.keys(): if not result[key].is_integer(): @@ -264,134 +264,6 @@ def _validate_run_influence_config(value): return result -def _validate_run_config(value, current, limits=None): - """Validate settings which the resident fitter can change between Runs.""" - if value is None: - return {} - if not isinstance(value, dict): - raise ApiError(HTTPStatus.BAD_REQUEST, - "run_config must be a JSON object") - current = current if isinstance(current, dict) else {} - limits = limits if isinstance(limits, dict) else {} - unknown = sorted(set(value) - set(current)) - if unknown: - raise ApiError(HTTPStatus.BAD_REQUEST, - f"Unknown or non-mutable Run configuration keys: {unknown}") - - result = {} - for key, item in value.items(): - if key == "track_length_bin_weights": - if item is None: - result[key] = None - continue - if (not isinstance(item, list) or len(item) != 3 - or any(isinstance(weight, bool) - or not isinstance(weight, (int, float)) - or not math.isfinite(float(weight)) - or float(weight) < 0 for weight in item) - or sum(float(weight) for weight in item) <= 0): - raise ApiError( - HTTPStatus.BAD_REQUEST, - f"{key} must be null or three finite non-negative weights " - "with a positive sum") - result[key] = [float(weight) for weight in item] - continue - if key == "max_track_crossing_per_step": - if (isinstance(item, bool) or not isinstance(item, (int, float)) - or not math.isfinite(float(item)) - or not float(item).is_integer() or int(item) < 0): - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be a non-negative integer") - maximum = limits.get(key) - if (not isinstance(maximum, bool) - and isinstance(maximum, (int, float)) - and int(item) > int(maximum)): - raise ApiError( - HTTPStatus.BAD_REQUEST, - f"{key} cannot exceed this session's prepared limit ({int(maximum)})") - result[key] = int(item) - continue - if key in ("track_min_sample_spacing", "track_max_sample_spacing"): - if (isinstance(item, bool) or not isinstance(item, (int, float)) - or not math.isfinite(float(item)) or float(item) <= 0): - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be a finite positive number") - result[key] = float(item) - continue - if key in ("min_walk_steps_per_track", "max_walk_steps_per_track", - "n_walks_per_track"): - if (isinstance(item, bool) or not isinstance(item, (int, float)) - or not math.isfinite(float(item)) - or not float(item).is_integer() or int(item) <= 0): - raise ApiError( - HTTPStatus.BAD_REQUEST, - f"{key} must be a positive integer") - result[key] = int(item) - continue - if key in RUN_MUTABLE_BOOLEAN_KEYS: - if not isinstance(item, bool): - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be boolean") - result[key] = item - continue - if key.startswith("loss_start_") and item is None: - if key == "loss_start_patch_dt": - raise ApiError(HTTPStatus.BAD_REQUEST, - "loss_start_patch_dt cannot be null") - result[key] = None - continue - if isinstance(item, bool) or not isinstance(item, (int, float)): - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be numeric") - number = float(item) - if not math.isfinite(number) or number < 0: - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be a finite non-negative number") - if key in RUN_MUTABLE_SAMPLING_KEYS or key.startswith("loss_start_"): - if not number.is_integer(): - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be an integer") - number = int(number) - if key in RUN_MUTABLE_SAMPLING_KEYS and number < 1: - # Disabled optional inputs have their loss weights and sample - # counts forced to zero when the session is loaded. VC3D - # round-trips those advertised active values on every Run. - # Permit that unchanged disabled value, while still rejecting - # attempts to turn an active sampler off by setting its count - # to zero at a Run boundary. - current_value = current.get(key) - current_is_zero = ( - not isinstance(current_value, bool) - and isinstance(current_value, (int, float)) - and float(current_value) == 0.0 - ) - if number != 0 or not current_is_zero: - raise ApiError(HTTPStatus.BAD_REQUEST, - f"{key} must be at least 1") - result[key] = number - effective = dict(current) - effective.update(result) - minimum = effective.get("track_min_sample_spacing") - maximum = effective.get("track_max_sample_spacing") - if (not isinstance(minimum, bool) and isinstance(minimum, (int, float)) - and not isinstance(maximum, bool) and isinstance(maximum, (int, float)) - and float(minimum) > float(maximum)): - raise ApiError( - HTTPStatus.BAD_REQUEST, - "track_min_sample_spacing must be <= track_max_sample_spacing") - walk_minimum = effective.get("min_walk_steps_per_track") - walk_maximum = effective.get("max_walk_steps_per_track") - if (not isinstance(walk_minimum, bool) - and isinstance(walk_minimum, (int, float)) - and not isinstance(walk_maximum, bool) - and isinstance(walk_maximum, (int, float)) - and int(walk_minimum) > int(walk_maximum)): - raise ApiError( - HTTPStatus.BAD_REQUEST, - "min_walk_steps_per_track must be <= max_walk_steps_per_track") - return result - - class ServiceLogBuffer: """Bounded, incremental copy of the service's stdout and stderr lines.""" @@ -561,21 +433,34 @@ def __init__(self): self._pruned_ids = OrderedDict() def register_directory(self, kind, session_id, generation, root, - entry_point, *, delete_root_on_prune=False): + entry_point, *, delete_root_on_prune=False, + progress=None, hash_workers=1): root = Path(root).resolve(strict=True) - files = {} + paths = [] for directory, dirnames, filenames in os.walk(root, followlinks=False): dirnames.sort() for filename in sorted(filenames): path = Path(directory) / filename if path.is_symlink() or not path.is_file(): continue - relative = path.relative_to(root).as_posix() - files[relative] = {"size": path.stat().st_size, - "sha256": _sha256_file(path)} - if len(files) > MAX_ARTIFACT_FILES: + paths.append(path) + if len(paths) > MAX_ARTIFACT_FILES: raise ApiError(HTTPStatus.INTERNAL_SERVER_ERROR, "Artifact has too many files to register") + + def digest(path): + return path, path.stat().st_size, _sha256_file(path) + + files = {} + workers = max(1, min(int(hash_workers), len(paths) or 1)) + with ThreadPoolExecutor(max_workers=workers, + thread_name_prefix="spiral-artifact-hash") as executor: + for index, (path, size, sha256) in enumerate( + executor.map(digest, paths), start=1): + relative = path.relative_to(root).as_posix() + files[relative] = {"size": size, "sha256": sha256} + if progress is not None: + progress(index, len(paths), relative) if entry_point not in files: raise ApiError(HTTPStatus.INTERNAL_SERVER_ERROR, f"Artifact entry point {entry_point!r} was not found") @@ -968,8 +853,21 @@ def _validate_tifxyz_output_step(metadata, expected_step): return values -def _load_flatten_correspondence(checkpoint_path): +def _load_flatten_correspondence(checkpoint_path=None, map_path=None): """Read Lasagna's flattened-output -> Spiral-source grid map.""" + if map_path is not None and Path(map_path).is_file(): + mapping = np.load(str(map_path), mmap_mode="r", allow_pickle=False) + mapping = np.asarray(mapping, dtype=np.float32) + if mapping.ndim != 3 or mapping.shape[-1] != 2: + raise RuntimeError( + "Lasagna output-to-source map must have shape (rows, columns, 2)") + if not np.isfinite(mapping).all(): + raise RuntimeError( + "Lasagna output-to-source map contains non-finite values") + return mapping + + if checkpoint_path is None: + raise RuntimeError("Lasagna produced no flatten correspondence") import torch state = torch.load( @@ -989,7 +887,8 @@ def _load_flatten_correspondence(checkpoint_path): return mapping -def _sample_rgba_through_map(source_rgba, source_yx, output_valid): +def _sample_rgba_through_map(source_rgba, source_yx, output_valid, *, + executor=None): """Bilinearly warp RGBA using premultiplied alpha.""" source = np.asarray(source_rgba, dtype=np.float32) / 255.0 if source.ndim != 3 or source.shape[-1] != 4: @@ -997,15 +896,21 @@ def _sample_rgba_through_map(source_rgba, source_yx, output_valid): alpha = source[..., 3] premultiplied = source[..., :3] * alpha[..., None] coordinates = [source_yx[..., 0], source_yx[..., 1]] - sampled_alpha = scipy.ndimage.map_coordinates( - alpha, coordinates, order=1, mode="constant", cval=0.0, - prefilter=False) - sampled_rgb = np.stack([ - scipy.ndimage.map_coordinates( - premultiplied[..., channel], coordinates, - order=1, mode="constant", cval=0.0, prefilter=False) - for channel in range(3) - ], axis=-1) + + def sample(channel): + values = alpha if channel == 3 else premultiplied[..., channel] + return scipy.ndimage.map_coordinates( + values, coordinates, order=1, mode="constant", cval=0.0, + prefilter=False) + + if executor is None: + sampled = [sample(channel) for channel in range(4)] + else: + # Each channel is independent and scipy releases the GIL here. Mapping + # them concurrently preserves the exact interpolation and output bytes. + sampled = list(executor.map(sample, range(4))) + sampled_alpha = sampled[3] + sampled_rgb = np.stack(sampled[:3], axis=-1) nonzero = sampled_alpha > 1.0e-8 sampled_rgb[nonzero] /= sampled_alpha[nonzero, None] sampled_rgb[~nonzero] = 0.0 @@ -1027,14 +932,18 @@ def _mapped_winding_ids(source_manifest, source_shape, source_yx, if (not isinstance(ranges, list) or not isinstance(windings, list) or len(ranges) != len(windings) or not ranges): raise RuntimeError("Spiral source preview has no winding mapping") - source_labels = np.full(source_shape, -1, dtype=np.int32) + winding_values = sorted({int(winding) for winding in windings}) + winding_to_dense = { + winding: index + 1 for index, winding in enumerate(winding_values) + } + source_labels = np.zeros(source_shape, dtype=np.int32) for bounds, winding in zip(ranges, windings): if not isinstance(bounds, list) or len(bounds) != 2: raise RuntimeError("Malformed Spiral winding column range") begin, end = int(bounds[0]), int(bounds[1]) if begin < 0 or end <= begin or end > source_shape[1]: raise RuntimeError("Spiral winding column range is out of bounds") - source_labels[:, begin:end] = int(winding) + source_labels[:, begin:end] = winding_to_dense[int(winding)] rows = np.rint(source_yx[..., 0]).astype(np.int64) columns = np.rint(source_yx[..., 1]).astype(np.int64) @@ -1043,30 +952,39 @@ def _mapped_winding_ids(source_manifest, source_shape, source_yx, & (rows >= 0) & (rows < source_shape[0]) & (columns >= 0) & (columns < source_shape[1]) ) - result = np.full(output_valid.shape, -1, dtype=np.int32) - result[in_bounds] = source_labels[rows[in_bounds], columns[in_bounds]] - result[~output_valid] = -1 - + dense_result = np.zeros(output_valid.shape, dtype=np.int32) + dense_result[in_bounds] = source_labels[ + rows[in_bounds], columns[in_bounds]] bounds = [] - for winding in sorted(int(value) for value in np.unique(result) - if int(value) >= 0): - yy, xx = np.nonzero(result == winding) + objects = scipy.ndimage.find_objects( + dense_result, max_label=len(winding_values)) + for dense_label, slices in enumerate(objects, start=1): + if slices is None: + continue + row_slice, column_slice = slices + winding = winding_values[dense_label - 1] bounds.append({ "winding": winding, - "row_begin": int(yy.min()), - "row_end": int(yy.max()) + 1, - "column_begin": int(xx.min()), - "column_end": int(xx.max()) + 1, + "row_begin": int(row_slice.start), + "row_end": int(row_slice.stop), + "column_begin": int(column_slice.start), + "column_end": int(column_slice.stop), }) if not bounds: raise RuntimeError("Lasagna correspondence mapped no preview windings") + lookup = np.asarray([-1, *winding_values], dtype=np.int32) + result = lookup[dense_result] return result, bounds -def _raw_run_diff_rgba(previous_manifest, current_manifest): +def _raw_run_diff_rgba(previous_manifest, current_manifest, *, + current_surface_data=None): """Build a current-source-grid displacement overlay by winding identity.""" - current_surface = Path(current_manifest["surface_path"]) - current_xyz, current_valid = _surface_xyz(current_surface) + if current_surface_data is None: + current_surface = Path(current_manifest["surface_path"]) + current_xyz, current_valid = _surface_xyz(current_surface) + else: + current_xyz, current_valid = current_surface_data rgba = np.zeros((*current_valid.shape, 4), dtype=np.uint8) if previous_manifest is None: return rgba, 0 @@ -1203,9 +1121,14 @@ def __init__(self, dataset_root=None, dataset_resolution=None, self._publishing_preview_generation = 0 self._preview_artifact = None self._preview_publish = None + self._preview_progress_started = None self._preview_publish_error = None self._preview_process = None self._previous_raw_preview_manifest = None + self.config_catalog = Config.catalog() + self.session_revision = 0 + self.run_plans = {} + self.pending_revision_target = None # ------------------------------------------------------------------ # Status and health @@ -1220,6 +1143,7 @@ def _base(self): "session_id": self.session_id, "service_generation": self.service_generation, "session_generation": self.session_generation, + "session_revision": self.session_revision, "command_generation": self.command_generation, "generation": self.status_generation, "session_replacement_in_progress": self.replacing, @@ -1250,7 +1174,9 @@ def status(self): "state": "Empty", "phase": "No session", "current_iteration": 0, "target_iteration": 0, "latest_metrics": {}, "warnings": [], "error": None, "preview_manifest_path": None, "preview_generation": 0, + "progress": None, }) + response.setdefault("progress", None) response["session_request"] = self.session_request response["preview_artifact"] = self._preview_artifact response["preview_publish"] = ( @@ -1262,6 +1188,36 @@ def status(self): self._preview_publish.get("stage_name") or "").strip() if stage_name: response["phase"] = stage_name + step = self._preview_publish.get("step") + total = self._preview_publish.get("total_steps") + elapsed = ( + max( + 0.0, + time.monotonic() + - self._preview_progress_started) + if self._preview_progress_started is not None + else 0.0) + eta = None + if (step is not None and total is not None + and int(step) > 0 and int(total) > int(step) + and elapsed >= 2.0): + eta = elapsed * ( + int(total) - int(step)) / int(step) + elif (step is not None and total is not None + and int(total) > 0 + and int(step) >= int(total)): + eta = 0.0 + response["progress"] = { + "operation": "publishing_preview", + "stage_name": stage_name, + "detail": None, + "step": int(step) if step is not None else None, + "total_steps": ( + int(total) if total is not None else None), + "unit": "steps", + "elapsed_seconds": elapsed, + "eta_seconds": eta, + } response["ephemeral_inputs"] = [ {"id": record["id"], "kind": record["kind"], "role": record.get("role"), "state": record["state"], @@ -1286,6 +1242,9 @@ def health(self): }) return response + def configuration_catalog(self): + return {**self._base(), **self.config_catalog} + def dataset(self): if self.dataset_resolution is None: raise ApiError(HTTPStatus.NOT_FOUND, @@ -1357,25 +1316,6 @@ def _dataset_session_request(self, request): [{"field": "tracks_dbm", "message": "Not a service-advertised candidate"}]) paths["tracks_dbm"] = str(Path(tracks).resolve(strict=False)) - # A dataset-owned service resolves conventional paths itself, but the - # client still controls which optional sources belong to this session. - # Clear disabled paths so the manifest and worker agree that they were - # not loaded. - config = (request.get("run") or {}).get("config") or {} - selected_paths = { - "use_verified_patches": ("verified_patches",), - "use_unverified_patches": ("unverified_patches",), - "use_normals": ("normal_x", "normal_y"), - "use_surf_sdt": ("surf_sdt",), - "use_tracks": ("tracks_dbm",), - "use_gradient_magnitude": ("gradient_magnitude",), - "use_fibers": ("fibers",), - } - for flag, field_names in selected_paths.items(): - if not bool(config.get(flag, True)): - for field_name in field_names: - paths[field_name] = "" - return {**request, "paths": paths} def load(self, request): @@ -1423,6 +1363,8 @@ def load(self, request): "run": run.manifest(), "preview": preview.manifest(), } + self.session_revision += 1 + self.run_plans.clear() self._reset_session_scope() try: self.session = create_session( @@ -1450,6 +1392,7 @@ def _reset_session_scope(self): self._publishing_preview_generation = 0 self._preview_artifact = None self._preview_publish = None + self._preview_progress_started = None self._preview_publish_error = None self._previous_raw_preview_manifest = None if previous_raw: @@ -1465,6 +1408,21 @@ def _status_changed(self, status): print(f"SPIRAL_ARTIFACT_ERROR {type(exc).__name__}: {exc}", file=sys.stderr, flush=True) with self.lock: + applied_manifest = status.get("input_manifest") + if (self.session_paths is not None + and isinstance(applied_manifest, dict) + and applied_manifest != self.session_paths.manifest()): + self.session_paths = SpiralInputPaths.from_mapping( + applied_manifest) + if self.session_request is not None: + self.session_request["paths"] = \ + self.session_paths.manifest() + if (self.pending_revision_target is not None + and status.get("state") in {"Ready", "Paused"} + and int(status.get("current_iteration") or 0) + >= self.pending_revision_target): + self.session_revision += 1 + self.pending_revision_target = None self.status_generation += 1 def _maybe_register_artifacts(self, status): @@ -1485,10 +1443,33 @@ def _maybe_register_artifacts(self, status): try: published_manifest = self._publish_flattened_preview( session_id, preview_generation, Path(preview_manifest)) + + def indexing_progress(current, total, relative): + self._update_preview_publish( + preview_generation, state="indexing", + stage_name=( + f"Indexing preview files ({current}/{total}): " + f"{relative}"), + step=current, total_steps=total, + overall_progress=( + float(current) / float(total) if total else 1.0)) + + indexing_started = time.perf_counter() + self._update_preview_publish( + preview_generation, state="indexing", + stage_name="Indexing preview files", + step=0, total_steps=0, overall_progress=0.0) ref = self.artifacts.register_directory( "spiral-preview", session_id, preview_generation, published_manifest.parent, published_manifest.name, - delete_root_on_prune=True) + delete_root_on_prune=True, progress=indexing_progress, + hash_workers=4) + print( + "SPIRAL_PREVIEW_TIMING " + f"generation={preview_generation} " + "stage='Indexing preview files' " + f"seconds={time.perf_counter() - indexing_started:.6f}", + flush=True) with self.lock: if self.session_id == session_id: self._preview_artifact = ref @@ -1520,6 +1501,7 @@ def _maybe_register_artifacts(self, status): if self._publishing_preview_generation == preview_generation: self._publishing_preview_generation = 0 self._preview_publish = None + self._preview_progress_started = None self.status_generation += 1 def _update_preview_publish(self, generation, **values): @@ -1527,18 +1509,33 @@ def _update_preview_publish(self, generation, **values): if self._publishing_preview_generation != generation: return current = dict(self._preview_publish or {}) + next_stage = values.get("stage_name", current.get("stage_name")) + if next_stage != current.get("stage_name"): + self._preview_progress_started = time.monotonic() current.update(values) current["generation"] = generation self._preview_publish = current self.status_generation += 1 + def run(self, request): + token = request.get("plan_token") + with self.lock: + plan = self.run_plans.pop(token, None) + if not plan or plan["expires"] < time.monotonic(): + raise ApiError(HTTPStatus.CONFLICT, "Run plan is missing or expired") + if plan["revision"] != self.session_revision: + raise ApiError(HTTPStatus.CONFLICT, "Run plan is stale") + if plan["new_fit_required"]: + raise ApiError(HTTPStatus.CONFLICT, + "This plan requires Start New Fit") + if plan["session_reload_required"]: + raise ApiError(HTTPStatus.CONFLICT, + "This plan requires reloading fit inputs") session = self._require_session() status = session.status() influence_config = _validate_run_influence_config( - request.get("influence_config")) - run_config = _validate_run_config( - request.get("run_config"), status.get("run_config"), - status.get("run_config_limits")) + plan["influence"]) + run_config = plan["configuration_changes"] with self.lock: pending = [record for record in self.ephemeral_records if record["state"] == "pending"] @@ -1559,15 +1556,137 @@ def mark_incorporated(records, error=None): and record["state"] == "incorporated")] self.status_generation += 1 - target = session.run(int(request.get("iterations", 0)), - pending_inputs=pending, - mark_incorporated=mark_incorporated, - influence_config=influence_config, - run_config=run_config) with self.lock: + self.pending_revision_target = ( + int(status.get("current_iteration") or 0) + plan["iterations"]) + try: + run_arguments = { + "pending_inputs": pending, + "mark_incorporated": mark_incorporated, + "influence_config": influence_config, + "run_config": run_config, + } + if plan["path_changes"]: + run_arguments["path_changes"] = plan["path_changes"] + target = session.run(plan["iterations"], **run_arguments) + except BaseException: + with self.lock: + self.pending_revision_target = None + raise + with self.lock: + self.run_plans.clear() self.status_generation += 1 return {**self.status(), "accepted": True, "target_iteration": target} + def plan_run(self, request): + session = self._require_session() + if session.status().get("state") not in {"Ready", "Paused"}: + raise ApiError(HTTPStatus.CONFLICT, + "Run planning requires a paused session") + expected = request.get("expected_session_revision") + if expected != self.session_revision: + raise ApiError(HTTPStatus.CONFLICT, "Session revision is stale") + configuration = request.get("configuration") + if not isinstance(configuration, dict) or \ + set(configuration) != set(self.config_catalog["defaults"]): + raise ApiError(HTTPStatus.BAD_REQUEST, + "Run planning requires a complete configuration") + try: + configuration = Config(configuration).as_dict() + except ValueError as exc: + raise ApiError(HTTPStatus.BAD_REQUEST, str(exc)) from exc + iterations = int(request.get("iterations", 0)) + if iterations < 1: + raise ApiError(HTTPStatus.BAD_REQUEST, + "iterations must be at least 1") + current = session.status().get("applied_config") + if current is None: + current = Config( + (self.session_request.get("run") or {}).get("config") or {} + ).as_dict() + changes = { + key: value for key, value in configuration.items() + if current.get(key) != value + } + fields = self.config_catalog["schema"]["fields"] + impacts = {fields[key]["runtime_impact"] for key in changes} + dependencies = sorted({ + dependency for key in changes + for dependency in fields[key]["dependencies"] + }) + current_manifest = self.session_paths.manifest() + input_manifest = request.get("inputs") + if input_manifest is None: + input_manifest = current_manifest + if not isinstance(input_manifest, dict): + raise ApiError(HTTPStatus.BAD_REQUEST, + "inputs must be a path manifest object") + path_changes = { + key: input_manifest.get(key) + for key in set(current_manifest) | set(input_manifest) + if input_manifest.get(key) != current_manifest.get(key) + } + path_specs = self.config_catalog["schema"].get("paths", {}) + input_changes = [] + for key, value in sorted(path_changes.items()): + spec = path_specs.get(key) + impact = ( + spec["runtime_impact"] if spec is not None + else "prepared_input_rebuild") + impacts.add(impact) + dependencies.extend( + spec.get("dependencies", []) if spec is not None else [ + "dense_stores", "patch_pcl", "tracks", "shell", + "preview_output", + ]) + input_changes.append({ + "key": key, + "before": current_manifest.get(key), + "after": value, + "runtime_impact": impact, + }) + if "outer_shell" in path_changes: + outer_shell = str(path_changes["outer_shell"] or "").strip() + if not outer_shell or not Path(outer_shell).is_dir(): + raise ApiError( + HTTPStatus.BAD_REQUEST, + "Outer shell path is not a readable directory", + [{"field": "outer_shell", + "message": "Path is not a directory"}]) + dependencies = sorted(set(dependencies)) + token = secrets.token_urlsafe(24) + new_fit = "new_fit" in impacts + session_reload_required = "prepared_input_rebuild" in impacts + plan = { + "revision": self.session_revision, + "expires": time.monotonic() + 60.0, + "iterations": iterations, + "influence": request.get("influence") or {}, + "configuration_changes": changes, + "path_changes": path_changes, + "changes": [ + {"key": key, "before": current.get(key), "after": value, + "runtime_impact": fields[key]["runtime_impact"]} + for key, value in changes.items() + ], + "affected_prepared_inputs": dependencies, + "model_state_preserved": not new_fit, + "optimizer_state_preserved": not new_fit, + "new_fit_required": new_fit, + "session_reload_required": session_reload_required, + "input_changed": bool(path_changes), + "input_changes": input_changes, + } + with self.lock: + self.run_plans[token] = plan + return { + **self._base(), "plan_token": token, + "expires_in_seconds": 60, + **{key: value for key, value in plan.items() + if key not in {"expires", "configuration_changes", "path_changes", + "iterations", "influence", "revision"}}, + } + def stop(self): self._require_session().stop() with self.lock: @@ -1643,6 +1762,23 @@ def _publish_flattened_preview( self, session_id, generation, preview_manifest_path): process = None publish_root = None + timing_stage = None + timing_started = time.perf_counter() + + def start_stage(state, stage_name, **values): + nonlocal timing_stage, timing_started + now = time.perf_counter() + if timing_stage is not None: + print( + "SPIRAL_PREVIEW_TIMING " + f"generation={generation} stage={timing_stage!r} " + f"seconds={now - timing_started:.6f}", + flush=True) + timing_stage = stage_name + timing_started = now + self._update_preview_publish( + generation, state=state, stage_name=stage_name, **values) + try: fit_service = _find_lasagna_service() config_path = fit_service.parent / "configs" / "flatten_fast_nofilter.json" @@ -1688,9 +1824,8 @@ def _publish_flattened_preview( with tempfile.TemporaryDirectory(prefix="spiral_lasagna_") as temporary: temporary = Path(temporary) object_store = temporary / "objects" - self._update_preview_publish( - generation, state="preparing", - stage_name="Preparing Lasagna input surface", + start_stage( + "preparing", "Preparing Lasagna input surface", output_step_vx=LASAGNA_PREVIEW_OUTPUT_STEP_VX) metadata = json.loads( (surface_path / "meta.json").read_text(encoding="utf-8")) @@ -1733,6 +1868,12 @@ def relay_output(): if match: port_holder["port"] = int(match.group(1)) ready.set() + if "[fit] peak GPU memory:" in text: + self._update_preview_publish( + generation, state="saving", + stage_name="Saving optimized flatten model", + step=0, total_steps=0, + overall_progress=0.0) relay = threading.Thread( target=relay_output, name="spiral-lasagna-log", @@ -1756,6 +1897,9 @@ def relay_output(): "output_name": surface_id, "output_dir": str(publish_root), "model_output": str(model_output), + "embed_job_metadata": False, + "omit_model": True, + "export_flatten_map": True, "source": "Spiral host service", } accepted = _fit_service_json( @@ -1764,6 +1908,9 @@ def relay_output(): if not fit_job_id: raise RuntimeError( "Temporary Lasagna service returned no job id") + start_stage( + "running", "Flattening preview surface", + step=0, total_steps=0, overall_progress=0.0) while True: with self.lock: @@ -1797,18 +1944,23 @@ def relay_output(): if not flattened_metadata_path.is_file(): raise RuntimeError( "Lasagna reported success but produced no tifxyz output") - self._update_preview_publish( - generation, state="mapping", - stage_name="Mapping Spiral preview artifacts") - - correspondence = _load_flatten_correspondence(model_output) + flatten_map_output = publish_root / ".flatten-map.npy" + start_stage( + "loading", "Loading flattened preview output", + step=0, total_steps=0, overall_progress=0.0) + correspondence = _load_flatten_correspondence( + model_output, flatten_map_output) flattened_xyz, flattened_valid = _surface_xyz(flattened_surface) if correspondence.shape[:2] != flattened_xyz.shape[:2]: raise RuntimeError( "Lasagna correspondence dimensions do not match " "the flattened surface") + source_xyz, source_valid = _surface_xyz(surface_path) + start_stage( + "mapping", "Mapping preview winding membership", + step=0, total_steps=0, overall_progress=0.0) winding_ids, winding_bounds = _mapped_winding_ids( - manifest, _surface_xyz(surface_path)[0].shape[:2], + manifest, source_xyz.shape[:2], correspondence, flattened_valid) winding_map_name = "winding-ids.tif" # OpenCV/Qt do not portably decode signed-int TIFF samples. @@ -1821,37 +1973,65 @@ def relay_output(): mapped_loss_maps = [] loss_output = publish_root / "loss-maps" loss_output.mkdir(exist_ok=True) - for entry in manifest.get("loss_maps", []): - if not isinstance(entry, dict): - continue - relative = str(entry.get("path") or "") - source_overlay = preview_manifest_path.parent / relative - if not source_overlay.is_file(): - continue - with Image.open(source_overlay) as image: - source_rgba = np.asarray( - image.convert("RGBA"), dtype=np.uint8) - mapped = _sample_rgba_through_map( - source_rgba, correspondence, flattened_valid) - destination = loss_output / Path(relative).name - Image.fromarray(mapped, mode="RGBA").save(destination) - mapped_entry = dict(entry) - mapped_entry["path"] = ( - Path("loss-maps") / destination.name).as_posix() - mapped_entry["supported_pixels"] = int( - np.count_nonzero(mapped[..., 3])) - mapped_loss_maps.append(mapped_entry) - - previous_manifest = None - with self.lock: - previous_path = self._previous_raw_preview_manifest - if previous_path and Path(previous_path).is_file(): - previous_manifest = json.loads( - Path(previous_path).read_text(encoding="utf-8")) - raw_diff, changed_pixels = _raw_run_diff_rgba( - previous_manifest, manifest) - mapped_diff = _sample_rgba_through_map( - raw_diff, correspondence, flattened_valid) + loss_entries = [ + entry for entry in manifest.get("loss_maps", []) + if isinstance(entry, dict) + and (preview_manifest_path.parent + / str(entry.get("path") or "")).is_file() + ] + remap_total = len(loss_entries) + start_stage( + "mapping", ( + f"Remapping preview loss maps (0/{remap_total})" + if remap_total else "No preview loss maps to remap"), + step=0, total_steps=remap_total, + overall_progress=0.0) + with ThreadPoolExecutor( + max_workers=4, + thread_name_prefix="spiral-overlay-channel") as remap_executor: + for loss_index, entry in enumerate(loss_entries, start=1): + relative = str(entry.get("path") or "") + self._update_preview_publish( + generation, state="mapping", + stage_name=( + f"Remapping preview loss maps " + f"({loss_index}/{remap_total}): " + f"{Path(relative).name}"), + step=loss_index - 1, total_steps=remap_total, + overall_progress=( + float(loss_index - 1) / float(remap_total) + if remap_total else 1.0)) + source_overlay = preview_manifest_path.parent / relative + with Image.open(source_overlay) as image: + source_rgba = np.asarray( + image.convert("RGBA"), dtype=np.uint8) + mapped = _sample_rgba_through_map( + source_rgba, correspondence, flattened_valid, + executor=remap_executor) + destination = loss_output / Path(relative).name + Image.fromarray(mapped, mode="RGBA").save(destination) + mapped_entry = dict(entry) + mapped_entry["path"] = ( + Path("loss-maps") / destination.name).as_posix() + mapped_entry["supported_pixels"] = int( + np.count_nonzero(mapped[..., 3])) + mapped_loss_maps.append(mapped_entry) + + start_stage( + "mapping", "Building preview run difference", + step=0, total_steps=0, overall_progress=0.0) + previous_manifest = None + with self.lock: + previous_path = self._previous_raw_preview_manifest + if previous_path and Path(previous_path).is_file(): + previous_manifest = json.loads( + Path(previous_path).read_text(encoding="utf-8")) + raw_diff, changed_pixels = _raw_run_diff_rgba( + previous_manifest, manifest, + current_surface_data=(source_xyz, source_valid)) + mapped_diff = _sample_rgba_through_map( + raw_diff, correspondence, flattened_valid, + executor=remap_executor) run_diff = None if changed_pixels: run_diff_name = "run-diff.png" @@ -1864,6 +2044,9 @@ def relay_output(): np.count_nonzero(mapped_diff[..., 3])), } + start_stage( + "finalizing", "Finalizing preview metadata", + step=0, total_steps=0, overall_progress=0.0) flattened_metadata = json.loads( flattened_metadata_path.read_text(encoding="utf-8")) _validate_tifxyz_output_step( @@ -1910,7 +2093,12 @@ def relay_output(): json.dumps(published, indent=2) + "\n", encoding="utf-8") + # These are transient transport files. The interactive Spiral + # preview never exposes a Lasagna model to VC3D. + del correspondence model_output.unlink(missing_ok=True) + flatten_map_output.unlink(missing_ok=True) + (flattened_surface / "model.pt").unlink(missing_ok=True) os.replace(publish_root, final_root) publish_root = None @@ -1920,9 +2108,9 @@ def relay_output(): preview_manifest_path) if old_raw and old_raw != str(preview_manifest_path): shutil.rmtree(Path(old_raw).parent, ignore_errors=True) - self._update_preview_publish( - generation, state="finished", stage_name="Finished", - overall_progress=1.0) + start_stage( + "finalizing", "Preparing preview artifact index", + step=0, total_steps=0, overall_progress=0.0) return final_root / "manifest.json" finally: _stop_process_group(process) @@ -2569,6 +2757,8 @@ def _dispatch(self): if self.command == "GET": if path == "/health": return state.health() + if path == "/configuration": + return state.configuration_catalog() if path == "/session/status": return state.status() if path == "/logs": @@ -2633,6 +2823,8 @@ def _dispatch(self): return state.begin_upload(body) if path == "/session/load": return state.deduplicated(command_id, lambda: state.load(body)) + if path == "/session/run/plan": + return state.plan_run(body) if path == "/session/run": return state.deduplicated(command_id, lambda: state.run(body)) if path == "/session/stop": diff --git a/volume-cartographer/scripts/spiral/strip_paths.py b/volume-cartographer/scripts/spiral/strip_paths.py index 74507dc8ad..92c504d88b 100644 --- a/volume-cartographer/scripts/spiral/strip_paths.py +++ b/volume-cartographer/scripts/spiral/strip_paths.py @@ -7,9 +7,20 @@ import time import numpy as np +import scipy.ndimage import scipy.sparse from scipy.sparse.csgraph import dijkstra as _csgraph_dijkstra +# Uniform-weight grid graphs have exponentially many exactly-tied shortest paths between a +# given pair, and Dijkstra deterministically picks the same one each time (a characteristic +# staircase), so successive pool refreshes trace near-identical routes. Multiplying each edge +# weight by uniform(1, 1 + EDGE_WEIGHT_JITTER) randomises which tied path wins while keeping +# every chosen path within (1 + EDGE_WEIGHT_JITTER) of true geodesic length. +EDGE_WEIGHT_JITTER = 0.05 + +# Matches the 8-connectivity of the quad graph built below. +_EIGHT_CONNECTED = np.ones((3, 3), dtype=bool) + def warmup(): # Submitted at configure time to force worker processes to spawn (and import numpy/scipy) @@ -26,10 +37,10 @@ def _unpack_mask(packed_mask, shape): return np.unpackbits(packed_mask, count=n).astype(bool).reshape(shape) -def build_quad_graph(valid_quad_mask): +def build_quad_graph(valid_quad_mask, rng=None): # 8-connected CSR adjacency over the (flattened) quad grid; edges join valid cells with # weight 1 (orthogonal) or sqrt(2) (diagonal), so Dijkstra distances are geometric path - # lengths in quad units. + # lengths in quad units (within 1 + EDGE_WEIGHT_JITTER when `rng` enables jitter). h, w = valid_quad_mask.shape n = h * w flat_idx = np.arange(n).reshape(h, w) @@ -43,8 +54,11 @@ def build_quad_graph(valid_quad_mask): rows.append(flat_idx[src_i, src_j][edge_valid]) cols.append(flat_idx[dst_i, dst_j][edge_valid]) weights.append(np.full(int(edge_valid.sum()), weight, dtype=np.float32)) + weights = np.concatenate(weights) + if rng is not None: + weights = weights * rng.uniform(1., 1. + EDGE_WEIGHT_JITTER, size=weights.size).astype(np.float32) return scipy.sparse.csr_matrix( - (np.concatenate(weights), (np.concatenate(rows), np.concatenate(cols))), + (weights, (np.concatenate(rows), np.concatenate(cols))), shape=(n, n), ) @@ -85,11 +99,18 @@ def backtrack_path_ij(predecessors, end_flat, width): def compute_patch_path_pool(valid_quad_mask, num_starts, endpoints_per_start, seed): # Path pool for the patch radius/DT losses: `num_starts` uniform-random start cells, one # Dijkstra each, `endpoints_per_start` distant endpoints per start. Returns a list of - # int32 ij path arrays. + # int32 ij path arrays. Starts are drawn from the largest connected component so they + # can't land on isolated slivers that yield degenerate near-single-cell paths. rng = np.random.default_rng(seed) width = valid_quad_mask.shape[1] - valid_ij = np.argwhere(valid_quad_mask) - graph = build_quad_graph(valid_quad_mask) + labels, num_components = scipy.ndimage.label(valid_quad_mask, structure=_EIGHT_CONNECTED) + if num_components > 1: + counts = np.bincount(labels.ravel()) + counts[0] = 0 + valid_ij = np.argwhere(labels == counts.argmax()) + else: + valid_ij = np.argwhere(valid_quad_mask) + graph = build_quad_graph(valid_quad_mask, rng) pool = [] for _ in range(num_starts): start_i, start_j = valid_ij[rng.integers(len(valid_ij))] @@ -116,7 +137,7 @@ def compute_anchor_path_pools(valid_quad_mask, i_q, j_q, paths_per_cone, seed): # starting at the anchor). Caller guarantees valid_quad_mask[i_q, j_q]. rng = np.random.default_rng(seed) height, width = valid_quad_mask.shape - graph = build_quad_graph(valid_quad_mask) + graph = build_quad_graph(valid_quad_mask, rng) dist, predecessors = _csgraph_dijkstra( graph, directed=False, indices=i_q * width + j_q, return_predecessors=True, ) diff --git a/volume-cartographer/scripts/spiral/tests/test_config.py b/volume-cartographer/scripts/spiral/tests/test_config.py new file mode 100644 index 0000000000..d928f2400e --- /dev/null +++ b/volume-cartographer/scripts/spiral/tests/test_config.py @@ -0,0 +1,93 @@ +import json + +import pytest + +from config import Config + + +def test_catalog_is_complete_and_presets_are_resolved(): + catalog = Config.catalog() + assert set(catalog["defaults"]) == set(catalog["schema"]["fields"]) + for preset in catalog["presets"].values(): + assert set(preset) == set(catalog["defaults"]) + + +def test_every_key_has_generated_metadata(): + catalog = Config.catalog() + required = { + "type", "nullable", "label", "runtime_impact", "dependencies"} + for key, field in catalog["schema"]["fields"].items(): + assert required <= set(field) + assert field["label"] == key.split("_", 1)[1].replace("_", " ").title() + + +def test_input_participation_is_not_part_of_advanced_configuration(): + catalog = Config.catalog() + assert not any( + key.startswith("input_use_") + for key in catalog["defaults"] + ) + + +def test_interactive_runtime_impacts_match_resident_capabilities(): + schema = Config.catalog()["schema"] + fields = schema["fields"] + for key, field in fields.items(): + if key.startswith("patch_"): + expected = ( + "prepared_input_rebuild" + if key == "patch_erode_patches" else "run_boundary") + assert field["runtime_impact"] == expected + if key.startswith("dense_"): + expected = ( + "prepared_input_rebuild" + if key == "dense_spacing_mode" else "run_boundary") + assert field["runtime_impact"] == expected + if key.startswith("dt_"): + assert field["runtime_impact"] == "run_boundary" + if key.startswith("shell_"): + assert field["runtime_impact"] == "shell_reload" + assert schema["paths"]["outer_shell"] == { + "runtime_impact": "shell_reload", + "dependencies": ["shell"], + } + + mutable_tracks = { + "track_min_sample_spacing", "track_max_sample_spacing", + "track_length_bin_weights", "track_max_tortuosity", + "track_max_track_crossing_per_step", + "track_min_walk_steps_per_track", "track_max_walk_steps_per_track", + "track_min_walks_per_track", "track_max_walks_per_track", + "track_walk_minimum_cycle_travel", + "track_radius_target", "track_radius_loss_margin", + "track_radius_within_norm_p", "track_dt_within_track_norm_p", + "track_dt_norm_p", "track_dt_loss_margin", + } + assert all(fields[key]["runtime_impact"] == "run_boundary" + for key in mutable_tracks) + assert all(fields[key]["runtime_impact"] == "prepared_input_rebuild" + for key in { + "track_crossing_precompute_max", "track_crossing_mode", + "track_exclusion_radius", + }) + + +def test_mapping_and_json_overrides_and_validation(tmp_path): + changed = Config({"optimizer_learning_rate": 0.25}) + assert changed.optimizer_learning_rate == 0.25 + profile = tmp_path / "profile.json" + profile.write_text(json.dumps({"optimizer_learning_rate": 0.5})) + assert Config(profile).optimizer_learning_rate == 0.5 + + with pytest.raises(ValueError, match="Unknown"): + Config({"not_a_setting": 1}) + with pytest.raises(ValueError, match="Invalid value"): + Config({"optimizer_learning_rate": "fast"}) + with pytest.raises(ValueError, match="Out-of-range"): + Config({"optimizer_learning_rate": -1}) + with pytest.raises(ValueError, match="Invalid value"): + Config({"dense_spacing_mode": "unknown"}) + with pytest.raises(ValueError, match="Invalid vector length"): + Config({"dense_spacing_pair_m_short": [1]}) + with pytest.raises(ValueError): + Config({"track_max_tortuosity": "unlimited"}) diff --git a/volume-cartographer/scripts/spiral/tests/test_flatten_spiral_checkpoint.py b/volume-cartographer/scripts/spiral/tests/test_flatten_spiral_checkpoint.py new file mode 100644 index 0000000000..194416a54b --- /dev/null +++ b/volume-cartographer/scripts/spiral/tests/test_flatten_spiral_checkpoint.py @@ -0,0 +1,56 @@ +import json +from pathlib import Path + +import pytest + +from flatten_spiral_checkpoint import ( + MODEL_CONFIG_KEYS, + _checkpoint_config, + _resolve_lasagna, + _resolve_umbilicus, +) + + +def test_legacy_checkpoint_model_config_gets_current_aliases(): + legacy = {key: index for index, key in enumerate(MODEL_CONFIG_KEYS)} + config = _checkpoint_config({"cfg": legacy}) + for key in MODEL_CONFIG_KEYS: + assert config[f"model_{key}"] == legacy[key] + + +def test_current_checkpoint_model_config_is_preserved(): + current = { + f"model_{key}": index for index, key in enumerate(MODEL_CONFIG_KEYS) + } + config = _checkpoint_config({"cfg": current}) + assert config == current + + +def test_resolve_umbilicus_prefers_explicit_path(tmp_path): + path = tmp_path / "umbilicus.json" + path.write_text(json.dumps({"control_points": []})) + checkpoint = tmp_path / "checkpoint.ckpt" + assert _resolve_umbilicus(checkpoint, path) == path.resolve() + + +def test_resolve_umbilicus_finds_checkpoint_ancestor(tmp_path): + path = tmp_path / "umbilicus.json" + path.write_text(json.dumps({"control_points": []})) + checkpoint = tmp_path / "spiral_output" / "run" / "checkpoint.ckpt" + assert _resolve_umbilicus(checkpoint, None) == path.resolve() + + +def test_resolve_lasagna_requires_service_and_config(tmp_path): + (tmp_path / "fit_service.py").write_text("") + config = tmp_path / "configs" / "flatten_fast_nofilter.json" + config.parent.mkdir() + config.write_text("{}") + assert _resolve_lasagna(tmp_path) == ( + (tmp_path / "fit_service.py").resolve(), + config.resolve(), + ) + + +def test_checkpoint_config_reports_missing_fields(): + with pytest.raises(ValueError, match="missing model configuration"): + _checkpoint_config({"cfg": {}}) diff --git a/volume-cartographer/scripts/spiral/tests/test_interactive_influence.py b/volume-cartographer/scripts/spiral/tests/test_interactive_influence.py index 9ef0a7d660..c29c51f93c 100644 --- a/volume-cartographer/scripts/spiral/tests/test_interactive_influence.py +++ b/volume-cartographer/scripts/spiral/tests/test_interactive_influence.py @@ -16,31 +16,30 @@ TINY_CONFIG = { - 'initial_dr_per_winding': 16., - 'flow_voxel_resolution': 8, # 24^3 high-res lattice, 4^3 low-res: both have interior points - 'flow_field_type': 'cartesian', - 'flow_field_high_res_lr_scale_initial': 2.0e-1, - 'num_flow_timesteps': 1, - 'linear_z_resolution': 48, - 'gap_expander_logit_resolution': 24, - 'gap_expander_num_windings': 6, - 'gap_expander_lr_scale': 0.3, + 'model_initial_dr_per_winding': 16., + 'model_flow_voxel_resolution': 8, # 24^3 high-res lattice, 4^3 low-res: both have interior points + 'model_flow_field_type': 'cartesian', + 'model_num_flow_timesteps': 1, + 'model_linear_z_resolution': 48, + 'model_gap_expander_logit_resolution': 24, + 'model_gap_expander_num_windings': 6, + 'model_gap_expander_lr_scale': 0.3, 'output_first_winding': 1, } INFLUENCE_CONFIG = { - 'random_seed': 1, - 'interactive_influence_enabled': True, - 'interactive_influence_z': 48.0, - 'interactive_influence_windings': 1.5, - 'interactive_influence_theta_frac': 0.25, - 'interactive_influence_sigma': 0.3333, - 'interactive_influence_footprint_points': 256, + 'optimizer_random_seed': 1, + 'influence_enabled': True, + 'influence_z': 48.0, + 'influence_windings': 1.5, + 'influence_theta_frac': 0.25, + 'influence_sigma': 0.3333, + 'sample_count_influence_footprint_points': 256, 'loss_weight_anchor': 20.0, - 'interactive_influence_anchor_lattice_points': 200, - 'interactive_influence_anchor_geometry_points': 200, - 'interactive_influence_anchor_samples_per_step': 64, - 'interactive_influence_anchor_ramp_power': 2.0, + 'sample_count_influence_anchor_lattice_points': 200, + 'sample_count_influence_anchor_geometry_points': 200, + 'sample_count_influence_anchor_samples_per_step': 64, + 'influence_anchor_ramp_power': 2.0, 'shell_outer_winding_idx': 5, } @@ -154,7 +153,7 @@ def test_perturbing_one_logit_changes_radii_only_at_predicted_location(self): def radii(): transform = GapExpandingTransform( - gap_params, torch.tensor(16.), min_z, max_z, TINY_CONFIG['gap_expander_lr_scale']) + gap_params, torch.tensor(16.), min_z, max_z, TINY_CONFIG['model_gap_expander_lr_scale']) with torch.no_grad(): return transform.get_transformed_winding_radii(theta_probe, z_probe) @@ -187,7 +186,7 @@ def test_z_row_localization(self): def radii(): transform = GapExpandingTransform( - gap_params, torch.tensor(16.), min_z, max_z, TINY_CONFIG['gap_expander_lr_scale']) + gap_params, torch.tensor(16.), min_z, max_z, TINY_CONFIG['model_gap_expander_lr_scale']) with torch.no_grad(): return transform.get_transformed_winding_radii(theta_probe, z_rows.clone()) @@ -232,7 +231,7 @@ def test_flowbox_to_spiral_round_trip_with_nonzero_flow(self): torch.testing.assert_close(round_trip, points, atol=0.5, rtol=0.) def test_two_stage_flow_round_trip(self): - model = make_tiny_model(num_flow_stages=2) + model = make_tiny_model(model_num_flow_stages=2) self.assertEqual(len(model.flow_fields), 2) with torch.no_grad(): for flow_field in model.flow_fields: @@ -326,7 +325,7 @@ def test_partially_masked_elements_move_less(self): self.assertLess(float(moved_masked), float(moved_free)) def test_flow_masks_apply_to_every_stage(self): - model = make_tiny_model(num_flow_stages=2) + model = make_tiny_model(model_num_flow_stages=2) state = make_influence_state(INFLUENCE_CONFIG, torch.device('cpu')) state._allocate_masks(model) state.masks['flow_lr'].zero_() diff --git a/volume-cartographer/scripts/spiral/tests/test_lr_schedule.py b/volume-cartographer/scripts/spiral/tests/test_lr_schedule.py new file mode 100644 index 0000000000..bb8e4678a6 --- /dev/null +++ b/volume-cartographer/scripts/spiral/tests/test_lr_schedule.py @@ -0,0 +1,144 @@ +import math + +import torch + +import fit_spiral +from fit_spiral import ( + get_exponential_lr_at_step, + realign_optimizer_lr_schedule, + set_optimizer_group_lr_scale, +) + + +def test_final_factor_is_fraction_of_initial_lr_at_training_horizon(): + initial_lr = 3e-5 + + assert math.isclose( + get_exponential_lr_at_step(initial_lr, 0.9, 30_000, 30_000), + initial_lr * 0.9, + rel_tol=1e-12, + ) + + +def test_realign_uses_enlarged_horizon_and_tracks_next_optimizer_step(): + initial_lr = 3e-5 + parameter = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float64)) + optimiser = torch.optim.SGD([parameter], lr=initial_lr) + scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimiser, gamma=0.9 ** (1.0 / 30_000)) + + scheduler, horizon = realign_optimizer_lr_schedule( + optimiser, + scheduler, + initial_lr=initial_lr, + final_factor=0.9, + completed_steps=30_000, + training_horizon=60_000, + exponential=True, + ) + + expected_at_realign = initial_lr * math.sqrt(0.9) + assert horizon == 60_000 + assert math.isclose( + optimiser.param_groups[0]["lr"], expected_at_realign, rel_tol=1e-12) + assert math.isclose( + scheduler.gamma, 0.9 ** (1.0 / 60_000), rel_tol=1e-12) + assert scheduler.last_epoch == 30_000 + + optimiser.step() + scheduler.step() + + assert scheduler.last_epoch == 30_001 + assert math.isclose( + optimiser.param_groups[0]["lr"], + get_exponential_lr_at_step( + initial_lr, 0.9, 30_001, 60_000), + rel_tol=1e-12, + ) + + +def test_high_res_flow_scale_uses_absolute_iteration_after_resume(monkeypatch): + monkeypatch.setattr(fit_spiral, "cfg", { + "model_flow_field_high_res_lr_scale_initial": 0.2, + "model_flow_field_high_res_lr_scale_final": 0.8, + "model_flow_field_high_res_lr_ramp_start_step": 10_000, + "model_flow_field_high_res_lr_ramp_steps": 20_000, + }) + + assert math.isclose( + fit_spiral.get_flow_field_high_res_lr_scale(20_000), 0.5) + assert math.isclose( + fit_spiral.get_flow_field_high_res_lr_scale(30_000), 0.8) + assert math.isclose( + fit_spiral.get_flow_field_high_res_lr_scale(30_001), 0.8) + + +def test_optimizer_group_lr_scale_controls_the_parameter_update(): + low_res = torch.nn.Parameter(torch.tensor(0., dtype=torch.float64)) + high_res = torch.nn.Parameter(torch.tensor(0., dtype=torch.float64)) + optimiser = torch.optim.AdamW([ + {"params": [low_res], "lr_scale": 1.}, + {"params": [high_res], "lr_scale": 1.}, + ], lr=0.1, weight_decay=0.) + scheduler = torch.optim.lr_scheduler.LambdaLR( + optimiser, lambda step: 1.) + + set_optimizer_group_lr_scale( + optimiser, + scheduler, + group=optimiser.param_groups[1], + reference_group=optimiser.param_groups[0], + scale=0.2, + initial_lr=0.1, + ) + low_res.grad = torch.ones_like(low_res) + high_res.grad = torch.ones_like(high_res) + optimiser.step() + + assert math.isclose( + high_res.item() / low_res.item(), 0.2, rel_tol=1e-12) + assert math.isclose( + optimiser.param_groups[1]["lr"], + optimiser.param_groups[0]["lr"] * 0.2, + rel_tol=1e-12, + ) + + +def test_realign_preserves_parameter_group_lr_scales(): + low_res = torch.nn.Parameter(torch.tensor(1., dtype=torch.float64)) + high_res = torch.nn.Parameter(torch.tensor(1., dtype=torch.float64)) + optimiser = torch.optim.SGD([ + {"params": [low_res], "lr_scale": 1.}, + {"params": [high_res], "lr": 2.e-5, "lr_scale": 0.2}, + ], lr=1.e-4) + scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimiser, gamma=0.5 ** (1. / 100)) + + scheduler, _ = realign_optimizer_lr_schedule( + optimiser, + scheduler, + initial_lr=1.e-4, + final_factor=0.5, + completed_steps=50, + training_horizon=100, + exponential=True, + ) + + assert math.isclose( + optimiser.param_groups[1]["lr"], + optimiser.param_groups[0]["lr"] * 0.2, + rel_tol=1e-12, + ) + assert math.isclose( + scheduler.base_lrs[1], scheduler.base_lrs[0] * 0.2, + rel_tol=1e-12, + ) + + optimiser.step() + scheduler.step() + + assert math.isclose( + optimiser.param_groups[1]["lr"], + optimiser.param_groups[0]["lr"] * 0.2, + rel_tol=1e-12, + ) diff --git a/volume-cartographer/scripts/spiral/tests/test_pcl_sampling.py b/volume-cartographer/scripts/spiral/tests/test_pcl_sampling.py index b718c449a1..4ccfde5aa3 100644 --- a/volume-cartographer/scripts/spiral/tests/test_pcl_sampling.py +++ b/volume-cartographer/scripts/spiral/tests/test_pcl_sampling.py @@ -6,7 +6,7 @@ def _configure(monkeypatch, *, stratified=True, weights=None): monkeypatch.setattr(losses, 'cfg', { - 'stratified_pcl_sampling': stratified, + 'pcl_stratified_pcl_sampling': stratified, 'pcl_sampling_weights': weights, }) diff --git a/volume-cartographer/scripts/spiral/tests/test_phase_spacing.py b/volume-cartographer/scripts/spiral/tests/test_phase_spacing.py index 13ac440905..4cc3a764e0 100644 --- a/volume-cartographer/scripts/spiral/tests/test_phase_spacing.py +++ b/volume-cartographer/scripts/spiral/tests/test_phase_spacing.py @@ -21,8 +21,7 @@ ) from transforms import GapExpanderParams, GapExpandingTransform from fit_session import ( - SpiralInputPaths, SpiralRunConfig, apply_optional_input_selection, - validate_session_request, + SpiralInputPaths, SpiralRunConfig, validate_session_request, ) from test_sdt_losses import ( DR_PER_WINDING, @@ -74,7 +73,7 @@ def detect(ray, volume, normals, cfg): def phase_cfg(**overrides): values = { - 'dense_spacing_num_pairs': 192, + 'sample_count_dense_spacing_pairs': 192, 'dense_spacing_pair_m_short': (2, 3), 'dense_spacing_pair_m_long': (2, 3), 'dense_spacing_max_steps': 160, @@ -481,8 +480,8 @@ def test_zero_shared_pair_budget_skips_active_phase_bundle(self): volume = sheet_volume(100, [10, 20, 30, 40, 50, 60, 70, 80]) normals = normal_volume(volume['shape']) cfg = phase_cfg( - dense_spacing_num_pairs=0, - dense_spacing_density_extra_pairs=0, + sample_count_dense_spacing_pairs=0, + sample_count_dense_spacing_density_extra_pairs=0, ) assert run_bundle( volume, normals, PerfectSpiralToX(), 8, cfg) == {} @@ -508,7 +507,7 @@ def base_request(self, tmp_path, config): umbilicus=str(umbilicus), output_directory=str(output), cache_directory=str(cache)) base = { - 'disable_patches': True, + 'input_disable_patches': True, 'loss_weight_shell_outer': 0.0, 'loss_weight_shell_patch_radius': 0.0, 'loss_weight_dense_normals': 0.0, @@ -554,16 +553,6 @@ def test_phase_mode_requires_normals_and_sdt(self, tmp_path): fields = {error['field'] for error in validate_session_request(paths, run)} assert {'normal_x', 'normal_y', 'surf_sdt'} <= fields - def test_disabled_phase_inputs_are_not_required(self, tmp_path): - paths, run = self.base_request(tmp_path, { - 'use_normals': False, - 'use_surf_sdt': False, - }) - fields = {error['field'] for error in validate_session_request(paths, run)} - assert 'normal_x' not in fields - assert 'normal_y' not in fields - assert 'surf_sdt' not in fields - def test_grad_mag_mode_requires_grad_mag_not_sdt(self, tmp_path): paths, run = self.base_request(tmp_path, { 'dense_spacing_mode': 'grad_mag', @@ -574,11 +563,10 @@ def test_grad_mag_mode_requires_grad_mag_not_sdt(self, tmp_path): assert 'surf_sdt' not in fields assert 'normal_x' not in fields - def test_disabled_grad_mag_is_not_required(self, tmp_path): + def test_zero_weight_grad_mag_is_not_required(self, tmp_path): paths, run = self.base_request(tmp_path, { 'dense_spacing_mode': 'grad_mag', - 'loss_weight_dense_spacing': 12.0, - 'use_gradient_magnitude': False, + 'loss_weight_dense_spacing': 0.0, }) fields = {error['field'] for error in validate_session_request(paths, run)} assert 'gradient_magnitude' not in fields @@ -595,82 +583,6 @@ def test_invalid_mode_is_rejected_before_asset_errors(self, tmp_path): assert 'surf_sdt' not in by_field assert 'gradient_magnitude' not in by_field - -class TestOptionalInputSelection: - def test_disabled_inputs_zero_their_weights_and_sampling(self): - config = { - 'dense_spacing_mode': 'phase', - 'use_verified_patches': False, - 'use_unverified_patches': False, - 'use_normals': False, - 'use_surf_sdt': False, - 'use_tracks': False, - 'use_fibers': False, - 'loss_weight_patch_radius': 8.0, - 'loss_weight_patch_dt': 4.0, - 'num_patches_per_step': 360, - 'num_patches_per_step_for_dt': 240, - 'num_points_per_patch': 800, - 'loss_weight_unverified_patch_radius': 2.0, - 'loss_weight_unverified_patch_dt': 1.0, - 'unverified_num_patches_per_step': 120, - 'unverified_num_patches_per_step_for_dt': 80, - 'unverified_num_points_per_patch': 800, - 'loss_weight_dense_normals': 100.0, - 'dense_normals_num_points': 60_000, - 'loss_weight_dense_spacing': 12.0, - 'loss_weight_dense_spacing_count': 2.0, - 'loss_weight_dense_spacing_density': 3.0, - 'loss_weight_min_spacing': 4.0, - 'loss_weight_dense_attachment': 5.0, - 'dense_spacing_num_pairs': 12_000, - 'dense_spacing_density_extra_pairs': 24_000, - 'dense_attachment_num_points': 20_000, - 'min_spacing_independent_samples': 2_000, - 'loss_weight_track_radius': 50.0, - 'loss_weight_track_dt': 10.0, - 'track_num_per_step': 48_000, - 'track_num_points_per_step': 24, - 'loss_weight_unattached_pcl_radius': 2.0, - 'loss_weight_unattached_pcl_dt': 4.0, - 'unattached_pcl_num_per_step': 84, - 'unattached_pcl_num_points_per_step': 32, - } - apply_optional_input_selection(config) - expected_zero = { - 'loss_weight_patch_radius', 'loss_weight_patch_dt', - 'num_patches_per_step', 'num_patches_per_step_for_dt', - 'num_points_per_patch', 'loss_weight_unverified_patch_radius', - 'loss_weight_unverified_patch_dt', 'unverified_num_patches_per_step', - 'unverified_num_patches_per_step_for_dt', - 'unverified_num_points_per_patch', 'loss_weight_dense_normals', - 'dense_normals_num_points', 'loss_weight_dense_spacing', - 'loss_weight_dense_spacing_count', - 'loss_weight_dense_spacing_density', - 'loss_weight_dense_attachment', 'dense_spacing_num_pairs', - 'dense_spacing_density_extra_pairs', 'dense_attachment_num_points', - 'loss_weight_track_radius', - 'loss_weight_track_dt', 'track_num_per_step', - 'track_num_points_per_step', 'loss_weight_unattached_pcl_radius', - 'loss_weight_unattached_pcl_dt', 'unattached_pcl_num_per_step', - 'unattached_pcl_num_points_per_step', - } - assert all(config[key] == 0 for key in expected_zero) - assert config['loss_weight_min_spacing'] == 4.0 - assert config['min_spacing_independent_samples'] == 2_000 - - def test_disabled_grad_mag_zeroes_legacy_spacing(self): - config = { - 'dense_spacing_mode': 'grad_mag', - 'use_gradient_magnitude': False, - 'loss_weight_dense_spacing': 12.0, - 'dense_spacing_num_pairs': 12_000, - } - apply_optional_input_selection(config) - assert config['loss_weight_dense_spacing'] == 0 - assert config['dense_spacing_num_pairs'] == 0 - - class TestNativeMinimumGap: def make_transform(self, dr=None): dr = torch.tensor(10.0, requires_grad=True) if dr is None else dr @@ -710,8 +622,8 @@ def get_native_log_gaps(self, winding, theta, z): return transform.get_native_log_gaps(winding, theta, z) cfg = { - 'min_spacing_independent_samples': 64, - 'min_spacing_d_min_wv': 6.0, + 'sample_count_minimum_spacing_independent_samples': 64, + 'dense_min_spacing_d_min_wv': 6.0, } loss, metrics = get_min_spacing_loss( Wrapper(), 7, cfg, 1, 95, diff --git a/volume-cartographer/scripts/spiral/tests/test_sdt_losses.py b/volume-cartographer/scripts/spiral/tests/test_sdt_losses.py index 0bf32200e5..a5a55fc360 100644 --- a/volume-cartographer/scripts/spiral/tests/test_sdt_losses.py +++ b/volume-cartographer/scripts/spiral/tests/test_sdt_losses.py @@ -85,12 +85,12 @@ def inv(self, points): def spacing_cfg(**overrides): cfg = { - 'dense_spacing_num_pairs': 256, + 'sample_count_dense_spacing_pairs': 256, 'dense_spacing_pair_m_short': (1, 1), 'dense_spacing_pair_m_long': (1, 1), 'dense_spacing_pair_long_fraction': 0.0, 'dense_spacing_count_temperature_wv': 0.5, - 'dense_spacing_count_extra_pairs': 0, + 'sample_count_dense_spacing_count_extra_pairs': 0, 'dense_spacing_target_step_wv': 1.0, 'dense_spacing_max_step_wv': 2.0, 'dense_spacing_max_steps': 64, @@ -119,9 +119,9 @@ def spacing_cfg(**overrides): 'loss_weight_dense_spacing_count': 8.0, 'loss_weight_min_spacing': 0.0, 'loss_weight_dense_attachment': 0.0, - 'min_spacing_d_min_wv': 6.0, - 'min_spacing_independent_samples': 64, - 'dense_attachment_num_points': 512, + 'dense_min_spacing_d_min_wv': 6.0, + 'sample_count_minimum_spacing_independent_samples': 64, + 'sample_count_dense_attachment_points': 512, 'dense_attachment_scale': 8.0, } cfg.update(overrides) @@ -473,7 +473,7 @@ def test_shared_ray_counts_match_independent_implementation(self): # compute_pair_counts implementation under identical (k, m, theta, z) # samples, including gradients through the transform. volume = sheet_volume(96, [10, 20, 30, 40, 50, 60]) - cfg = spacing_cfg(dense_spacing_num_pairs=64) + cfg = spacing_cfg(sample_count_dense_spacing_pairs=64) seed = 11 offset_bundle = torch.tensor(1.5, requires_grad=True) diff --git a/volume-cartographer/scripts/spiral/tests/test_sparse_cuda_cache.py b/volume-cartographer/scripts/spiral/tests/test_sparse_cuda_cache.py index c0ab305753..503e28e625 100644 --- a/volume-cartographer/scripts/spiral/tests/test_sparse_cuda_cache.py +++ b/volume-cartographer/scripts/spiral/tests/test_sparse_cuda_cache.py @@ -1,112 +1,114 @@ -from pathlib import Path - import numpy as np import pytest -import tensorstore as ts import torch +import zarr +from pack_resident_pools import pack_arrays, sidecar_path from sdt_losses import sample_sdt_trilinear -from sparse_cuda_cache import ( - CHUNK_VOXELS, - BoundedSparseCudaCache, - SparseScalarStore, -) - - -def write_array(path: Path, data: np.ndarray) -> str: - array = ts.open( - { - "driver": "zarr", - "kvstore": {"driver": "file", "path": str(path)}, - "metadata": { - "dtype": "|u1", - "shape": list(data.shape), - "chunks": [16, 16, 16], - "compressor": None, - "order": "C", - "fill_value": 0, - }, - }, - create=True, - open=True, - ).result() - array.write(data).result() +from sparse_cuda_cache import ResidentBrickPool, SparseScalarStore + + +def write_array(path, data): + array = zarr.open( + str(path), mode='w', shape=data.shape, chunks=(16, 16, 16), + dtype='|u1', compressor=None, fill_value=0, zarr_format=2, + dimension_separator='/', + ) + array[:] = data return str(path) -def test_gather_matches_dense_across_lru_evictions(tmp_path): +def make_pool(tmp_path, arrays, name, **kwargs): + paths = [write_array(tmp_path / f'{name}_{i}', a) + for i, a in enumerate(arrays)] + out = pack_arrays(paths, sidecar_path(paths[0], '0'), label=name) + return ResidentBrickPool(out, device='cpu', label=name, **kwargs) + + +def test_gather_matches_dense_multichannel(tmp_path): z, y, x = np.indices((40, 40, 70)) first = ((z * 17 + y * 5 + x) % 251 + 1).astype(np.uint8) second = ((z * 3 + y * 11 + x * 7) % 251 + 1).astype(np.uint8) - paths = [ - write_array(tmp_path / "first", first), - write_array(tmp_path / "second", second), - ] - cache = BoundedSparseCudaCache( - source_paths=paths, - shape_zyx=first.shape, - budget_bytes=2 * 2 * CHUNK_VOXELS, - device="cpu", - label="test normals", - tensorstore_cache_bytes=1 << 20, - ) + pool = make_pool(tmp_path, [first, second], 'pair') - requests = [ + for indices in [ torch.tensor([[0, 0, 0], [1, 2, 33], [5, 3, 63]]), torch.tensor([[35, 2, 2], [35, 35, 35]]), - torch.tensor([[0, 0, 0], [5, 3, 63]]), - ] - for indices in requests: - actual = cache.gather(indices) + torch.zeros([0, 3], dtype=torch.long), + ]: + actual = pool.gather(indices) expected = torch.from_numpy(np.stack([ first[tuple(indices.numpy().T)], second[tuple(indices.numpy().T)] ], axis=-1)) torch.testing.assert_close(actual, expected) - cache.gather(requests[-1]) - - stats = cache.stats() - assert stats["resident_chunks"] <= 2 - assert stats["evictions"] > 0 - assert stats["hits"] > 0 - - -def test_one_gather_must_fit_the_bounded_working_set(tmp_path): - data = np.arange(40 * 40 * 70, dtype=np.uint32).reshape(40, 40, 70) - data = (data % 255).astype(np.uint8) - path = write_array(tmp_path / "scalar", data) - cache = BoundedSparseCudaCache( - source_paths=[path], - shape_zyx=data.shape, - budget_bytes=CHUNK_VOXELS, - device="cpu", - label="test scalar", - tensorstore_cache_bytes=1 << 20, - ) - with pytest.raises(RuntimeError, match="exceeding its 1-chunk LRU capacity"): - cache.gather(torch.tensor([[0, 0, 0], [0, 0, 33]])) + assert pool.stats()['gathers'] == 2 # the empty gather short-circuits + + +def test_absent_bricks_read_zero(tmp_path): + data = np.zeros((48, 16, 16), dtype=np.uint8) + data[:16] = 9 # only the first chunk row is occupied + pool = make_pool(tmp_path, [data], 'sparse') + assert pool.resident_bricks < pool.table.numel() + values = pool.gather(torch.tensor([[2, 2, 2], [30, 5, 5], [47, 15, 15]])) + assert values[:, 0].tolist() == [9, 0, 0] -def test_large_z_roi_uses_source_origin_and_eviction(tmp_path): +def test_origin_and_z_roi_restriction(tmp_path): data = np.broadcast_to( - (np.arange(5200, dtype=np.uint16) % 251 + 1).astype(np.uint8)[:, None, None], - (5200, 2, 2), + (np.arange(64, dtype=np.uint16) % 251 + 1).astype(np.uint8)[:, None, None], + (64, 4, 4), ).copy() - path = write_array(tmp_path / "large_z", data) - cache = BoundedSparseCudaCache( - source_paths=[path], - shape_zyx=data.shape, - origin_zyx=(100, 0, 0), - budget_bytes=CHUNK_VOXELS, - device="cpu", - label="test 5000-slice ROI", - tensorstore_cache_bytes=1 << 20, - ) - - first = cache.gather(torch.tensor([[0, 0, 0]])) - last = cache.gather(torch.tensor([[4999, 1, 1]])) - assert int(first[0, 0]) == int(data[100, 0, 0]) - assert int(last[0, 0]) == int(data[5099, 1, 1]) - assert cache.stats()["evictions"] == 1 + paths = [write_array(tmp_path / 'large_z', data)] + out = pack_arrays(paths, sidecar_path(paths[0], '0'), label='roi') + pool = ResidentBrickPool( + out, device='cpu', label='roi', origin_zyx=(32, 0, 0), z_roi=(32, 64)) + full = ResidentBrickPool(out, device='cpu', label='full') + + assert pool.resident_bricks < full.resident_bricks + first = pool.gather(torch.tensor([[0, 0, 0]])) + last = pool.gather(torch.tensor([[31, 3, 3]])) + assert int(first[0, 0]) == int(data[32, 0, 0]) + assert int(last[0, 0]) == int(data[63, 3, 3]) + + +def test_bounds_check_env(tmp_path, monkeypatch): + monkeypatch.setenv('FIT_SPIRAL_RESIDENT_BOUNDS_CHECK', '1') + data = np.ones((16, 16, 16), dtype=np.uint8) + pool = make_pool(tmp_path, [data], 'bounds') + with pytest.raises(IndexError): + pool.gather(torch.tensor([[16, 0, 0]])) + + +def test_pack_ct_mask_zeroes_and_drops_bricks(tmp_path): + from pack_resident_pools import CtMasker, verify_pool + + rng = np.random.default_rng(3) + data = rng.integers(1, 255, size=(32, 32, 32), dtype=np.uint8) + store = write_array(tmp_path / 'sdt', data) + # CT at half resolution (ratio 2): zero except one occupied corner region, + # so only target voxels [0:16, 0:16, 0:16] survive the mask. + ct = np.zeros((16, 16, 16), dtype=np.uint8) + ct[:8, :8, :8] = 7 + zarr.open( + str(tmp_path / 'ct' / '2'), mode='w', shape=ct.shape, chunks=(8, 8, 8), + dtype='|u1', compressor=None, fill_value=0, zarr_format=2, + dimension_separator='.', + )[:] = ct + + masker = CtMasker(tmp_path / 'ct', '2', data.shape) + assert masker.ratio == (2, 2, 2) + out = pack_arrays([store], sidecar_path(store, '0'), label='masked', + brick_shape=(8, 8, 8), ct_masker=masker) + verify_pool(out, 500) # mask-aware: pool zeros are accepted iff CT == 0 + + pool = ResidentBrickPool(out, device='cpu', label='masked') + assert pool.meta['ct_mask']['ratio'] == [2, 2, 2] + # 64 bricks in the grid; only the 2x2x2 corner block survives + assert pool.resident_bricks == 8 + 1 + inside = pool.gather(torch.tensor([[3, 3, 3]])) + outside = pool.gather(torch.tensor([[3, 3, 20], [25, 25, 25]])) + assert int(inside[0, 0]) == int(data[3, 3, 3]) + assert outside[:, 0].tolist() == [0, 0] def test_sparse_sdt_sampling_matches_dense(tmp_path): @@ -115,15 +117,7 @@ def test_sparse_sdt_sampling_matches_dense(tmp_path): np.clip(np.rint(np.abs(x - 35.0) - 2.0), -127, 127) + 128 ).astype(np.uint8) data = np.broadcast_to(encoded, (6, 6, 70)).copy() - path = write_array(tmp_path / "sdt", data) - cache = BoundedSparseCudaCache( - source_paths=[path], - shape_zyx=data.shape, - budget_bytes=3 * CHUNK_VOXELS, - device="cpu", - label="test sdt", - tensorstore_cache_bytes=1 << 20, - ) + pool = make_pool(tmp_path, [data], 'sdt') dense = { "backend": "dense_test", "kind": "sdt", @@ -139,7 +133,7 @@ def test_sparse_sdt_sampling_matches_dense(tmp_path): sparse = { **dense, "backend": "sparse_cuda", - "store": SparseScalarStore(cache), + "store": SparseScalarStore(pool), } sparse.pop("volume") points_dense = ( diff --git a/volume-cartographer/scripts/spiral/tests/test_speedup_equivalence.py b/volume-cartographer/scripts/spiral/tests/test_speedup_equivalence.py index d8cd3b8087..cf74d8755f 100644 --- a/volume-cartographer/scripts/spiral/tests/test_speedup_equivalence.py +++ b/volume-cartographer/scripts/spiral/tests/test_speedup_equivalence.py @@ -331,7 +331,7 @@ class RK4FusedIntegratorTests(unittest.TestCase): @staticmethod def _make_flow(seed): torch.manual_seed(seed) - flow = CartesianFlowField(torch.tensor([12, 12, 12]), spatial_scale_factor=6, lr_scale_factor=0.2) + flow = CartesianFlowField(torch.tensor([12, 12, 12]), spatial_scale_factor=6) with torch.no_grad(): flow.flows[0].normal_(std=0.1) flow.flows[1].normal_(std=0.1) @@ -387,7 +387,7 @@ def test_no_grad_forward_matches(self): class CachedSamplerGradModeUpgradeTests(unittest.TestCase): def test_no_grad_first_call_does_not_poison_field_gradients(self): torch.manual_seed(11) - flow = CartesianFlowField(torch.tensor([12, 12, 12]), spatial_scale_factor=6, lr_scale_factor=0.2) + flow = CartesianFlowField(torch.tensor([12, 12, 12]), spatial_scale_factor=6) with torch.no_grad(): flow.flows[0].normal_(std=0.05) flow.flows[1].normal_(std=0.05) diff --git a/volume-cartographer/scripts/spiral/tests/test_spiral_headless.py b/volume-cartographer/scripts/spiral/tests/test_spiral_headless.py index e62c095da4..d2c11d1801 100644 --- a/volume-cartographer/scripts/spiral/tests/test_spiral_headless.py +++ b/volume-cartographer/scripts/spiral/tests/test_spiral_headless.py @@ -171,12 +171,27 @@ def test_distributed_session_waits_for_every_rank_before_ready(self): "error": None, "current_iteration": 0, "target_iteration": 0, } session._events.put(("status", 0, ready)) - session._events.put(("status", 1, {**ready, "state": "Loading"})) + session._events.put(("status", 1, { + **ready, + "state": "Loading", + "progress": { + "operation": "loading", + "stage_name": "Loading tracks", + "detail": None, + "step": 2, + "total_steps": 10, + "unit": "DB keys", + "elapsed_seconds": 3.0, + "eta_seconds": 12.0, + }, + })) deadline = time.time() + 2 while len(published) < 2 and time.time() < deadline: time.sleep(0.01) self.assertEqual(published[-1]["state"], "Loading") - self.assertEqual(published[-1]["phase"], "Waiting for all GPU workers") + self.assertEqual(published[-1]["phase"], "Loading tracks") + self.assertIn( + "GPU worker 2/2", published[-1]["progress"]["detail"]) session._events.put(("status", 1, ready)) deadline = time.time() + 2 @@ -201,6 +216,97 @@ def test_interactive_run_can_continue_past_checkpoint_training_steps(self): self.assertEqual(session._target, 30_250) self.assertEqual(session._state, "Running") + def test_interactive_run_extends_training_horizon_to_target_iteration(self): + session = InteractiveFitSession.__new__(InteractiveFitSession) + session._condition = threading.Condition() + session._state = "Paused" + session._completed = 30_000 + session._pending = 0 + session._target = 30_000 + session._configure_run = lambda *_: None + session._idle_actions = [] + session.requested_config = { + "optimizer_num_training_steps": 30_000, + } + session._run_config = dict(session.requested_config) + + target = session.run(250) + + self.assertEqual(target, 30_250) + self.assertEqual(session._idle_actions[0][0], "configure") + self.assertEqual( + session._idle_actions[0][1]["optimizer_num_training_steps"], + 30_250, + ) + self.assertEqual( + session._run_config["optimizer_num_training_steps"], 30_250) + + def test_interactive_run_preserves_horizon_when_target_is_within_it(self): + session = InteractiveFitSession.__new__(InteractiveFitSession) + session._condition = threading.Condition() + session._state = "Ready" + session._completed = 100 + session._pending = 0 + session._target = 100 + session._configure_run = lambda *_: None + session._idle_actions = [] + session.requested_config = { + "optimizer_num_training_steps": 30_000, + } + session._run_config = dict(session.requested_config) + + session.run(250) + + self.assertEqual(session._idle_actions, []) + self.assertEqual( + session._run_config["optimizer_num_training_steps"], 30_000) + + def test_interactive_run_preserves_horizon_when_target_equals_it(self): + session = InteractiveFitSession.__new__(InteractiveFitSession) + session._condition = threading.Condition() + session._state = "Ready" + session._completed = 29_750 + session._pending = 0 + session._target = 29_750 + session._configure_run = lambda *_: None + session._idle_actions = [] + session.requested_config = { + "optimizer_num_training_steps": 30_000, + } + session._run_config = dict(session.requested_config) + + target = session.run(250) + + self.assertEqual(target, 30_000) + self.assertEqual(session._idle_actions, []) + self.assertEqual( + session._run_config["optimizer_num_training_steps"], 30_000) + + def test_interactive_run_extends_horizon_by_run_count_when_crossed(self): + session = InteractiveFitSession.__new__(InteractiveFitSession) + session._condition = threading.Condition() + session._state = "Ready" + session._completed = 29_751 + session._pending = 0 + session._target = 29_751 + session._configure_run = lambda *_: None + session._idle_actions = [] + session.requested_config = { + "optimizer_num_training_steps": 30_000, + } + session._run_config = dict(session.requested_config) + + target = session.run(250) + + self.assertEqual(target, 30_001) + self.assertEqual(session._idle_actions[0][0], "configure") + self.assertEqual( + session._idle_actions[0][1]["optimizer_num_training_steps"], + 30_250, + ) + self.assertEqual( + session._run_config["optimizer_num_training_steps"], 30_250) + def test_run_queues_influence_config_with_only_pending_inputs(self): session = InteractiveFitSession.__new__(InteractiveFitSession) session._condition = threading.Condition() @@ -211,7 +317,7 @@ def test_run_queues_influence_config_with_only_pending_inputs(self): session._incorporate_inputs = lambda *_: None session._idle_actions = [] pending = [{"id": "new-patch"}] - influence = {"interactive_influence_theta_frac": 0.25} + influence = {"influence_theta_frac": 0.25} session.run(20, pending_inputs=pending, influence_config=influence) @@ -248,13 +354,13 @@ def test_run_configuration_applies_active_host_values_exactly(self): session._configure_run = lambda config: setattr(session, "applied", config) session._run_configuration({ - "num_patches_per_step": 101, + "sample_count_patches_per_step": 101, "loss_weight_patch_radius": 3.5, "loss_start_patch_dt": 123, }) self.assertEqual(session.applied, { - "num_patches_per_step": 101, + "sample_count_patches_per_step": 101, "loss_weight_patch_radius": 3.5, "loss_start_patch_dt": 123, }) diff --git a/volume-cartographer/scripts/spiral/tests/test_spiral_helpers.py b/volume-cartographer/scripts/spiral/tests/test_spiral_helpers.py index 0e7ae9e0ac..2b42a0827f 100644 --- a/volume-cartographer/scripts/spiral/tests/test_spiral_helpers.py +++ b/volume-cartographer/scripts/spiral/tests/test_spiral_helpers.py @@ -151,7 +151,7 @@ class ResolveOuterWindingIdxWiringTests(unittest.TestCase): def _cfg(self, idx, gap=200): return {'shell_outer_winding_idx': idx, - 'gap_expander_num_windings': gap} + 'model_gap_expander_num_windings': gap} def test_configured_index_survives_a_shell_less_run(self): idx, notes = resolve_outer_winding_idx_and_notes( @@ -186,7 +186,7 @@ def test_gap_expander_control_also_runs_without_a_shell(self): self._cfg(130, gap=130), shell_active=False, infer_outer_winding_idx=self.fail) self.assertEqual(idx, 130) - self.assertTrue(any('gap_expander_num_windings >= 133' in n + self.assertTrue(any('model_gap_expander_num_windings >= 133' in n for n in notes)) diff --git a/volume-cartographer/scripts/spiral/tests/test_spiral_progress.py b/volume-cartographer/scripts/spiral/tests/test_spiral_progress.py new file mode 100644 index 0000000000..4e2cd1318b --- /dev/null +++ b/volume-cartographer/scripts/spiral/tests/test_spiral_progress.py @@ -0,0 +1,122 @@ +import io +import threading + +import pytest + +from spiral_progress import ProgressReporter + + +class FakeClock: + def __init__(self): + self.value = 100.0 + + def __call__(self): + return self.value + + def advance(self, seconds): + self.value += seconds + + +def test_determinate_snapshot_has_stage_local_elapsed_and_eta(): + clock = FakeClock() + reporter = ProgressReporter(clock=clock, heartbeat_interval=0) + reporter.begin( + "loading", "Loading patches", + step=0, total_steps=10, unit="patches") + clock.advance(4) + reporter.update(2) + + snapshot = reporter.snapshot() + + assert snapshot == { + "operation": "loading", + "stage_name": "Loading patches", + "detail": None, + "step": 2, + "total_steps": 10, + "unit": "patches", + "elapsed_seconds": pytest.approx(4.0), + "eta_seconds": pytest.approx(16.0), + } + + +def test_indeterminate_snapshot_has_elapsed_without_eta(): + clock = FakeClock() + reporter = ProgressReporter(clock=clock, heartbeat_interval=0) + reporter.begin("loading", "Building geometry index") + clock.advance(75) + + snapshot = reporter.snapshot() + + assert snapshot["elapsed_seconds"] == pytest.approx(75) + assert snapshot["step"] is None + assert snapshot["total_steps"] is None + assert snapshot["eta_seconds"] is None + + reporter.begin( + "loading", "Discovering GPU bricks", + step=0, total_steps=0, unit="bricks") + clock.advance(10) + assert reporter.snapshot()["eta_seconds"] is None + + +def test_publish_is_rate_limited_but_snapshot_keeps_latest_counter(): + clock = FakeClock() + published = [] + reporter = ProgressReporter( + published.append, clock=clock, publish_interval=1.0, + heartbeat_interval=0) + reporter.begin("loading", "Loading tracks", step=0, total_steps=10) + reporter.update(1) + reporter.update(2) + + assert len(published) == 1 + assert reporter.snapshot()["step"] == 2 + + clock.advance(1) + reporter.update(3) + assert len(published) == 2 + assert published[-1]["step"] == 3 + + +def test_clear_publishes_null_and_removes_progress(): + published = [] + reporter = ProgressReporter(published.append, heartbeat_interval=0) + reporter.begin("saving_checkpoint", "Saving checkpoint") + reporter.clear() + + assert published[-1] is None + assert reporter.snapshot() is None + + +def test_non_tty_console_uses_stable_progress_lines(): + clock = FakeClock() + stream = io.StringIO() + reporter = ProgressReporter( + stream=stream, clock=clock, heartbeat_interval=0) + reporter.begin( + "optimizing", "Optimizing", + step=0, total_steps=4, unit="iterations") + clock.advance(5) + reporter.update(1) + reporter.finish() + + output = stream.getvalue() + assert "PROGRESS Optimizing" in output + assert "1/4 iterations (25.0%)" in output + assert "ETA 15s" in output + assert "4/4 iterations (100.0%)" in output + + +def test_updates_and_snapshots_are_thread_safe(): + reporter = ProgressReporter(heartbeat_interval=0) + reporter.begin("loading", "Loading", step=0, total_steps=1000) + + thread = threading.Thread( + target=lambda: [reporter.update(index) for index in range(1001)]) + thread.start() + while thread.is_alive(): + snapshot = reporter.snapshot() + assert 0 <= snapshot["step"] <= 1000 + thread.join() + assert reporter.snapshot()["step"] == 1000 diff --git a/volume-cartographer/scripts/spiral/tests/test_spiral_service_v2.py b/volume-cartographer/scripts/spiral/tests/test_spiral_service_v2.py index 1be26572bc..2695e1976b 100644 --- a/volume-cartographer/scripts/spiral/tests/test_spiral_service_v2.py +++ b/volume-cartographer/scripts/spiral/tests/test_spiral_service_v2.py @@ -7,6 +7,7 @@ """ import argparse +from concurrent.futures import ThreadPoolExecutor import io import hashlib import json @@ -35,12 +36,14 @@ from spiral_service import (ApiError, ArtifactRegistry, ExclusiveFileLock, FileLockUnavailable, ServiceLogBuffer, ServiceState, SpiralServer, _mapped_winding_ids, + _load_flatten_correspondence, _prepare_cleaned_lasagna_surface, _raw_run_diff_rgba, _sample_rgba_through_map, _validate_tifxyz_output_step, load_or_create_api_key, parse_gpu_ids, parse_session_name) from fit_session import API_VERSION, SpiralInputPaths, resolve_dataset_root +from config import Config class FakeSession: @@ -48,29 +51,32 @@ def __init__(self): self.state = "Paused" self.run_calls = [] self.run_config = { - "num_patches_per_step": 360, + "sample_count_patches_per_step": 360, "loss_weight_patch_radius": 8.0, "loss_start_patch_dt": 25_000, "loss_start_track_dt": 10_000, - "save_png_visualizations": False, + "output_save_png_visualizations": False, "track_length_bin_weights": None, - "max_track_crossing_per_step": 0, + "track_max_track_crossing_per_step": 0, "track_min_sample_spacing": 20.0, "track_max_sample_spacing": 60.0, - "min_walk_steps_per_track": 24, - "max_walk_steps_per_track": 256, - "n_walks_per_track": 4, + "track_min_walk_steps_per_track": 24, + "track_max_walk_steps_per_track": 256, + "track_min_walks_per_track": 2, + "track_max_walks_per_track": 4, } self.default_advanced_config = { - "learning_rate": 3e-5, - "num_patches_per_step": 360, + "optimizer_learning_rate": 3e-5, + "sample_count_patches_per_step": 360, "loss_weight_patch_radius": 8.0, "track_crossing_precompute_max": 8, - "track_crossing_mode": "count", - "track_walk_require_loop_consistency": False, + "track_crossing_mode": "track_walk", + "track_walk_minimum_cycle_travel": 20.0, } self.saved = [] self.closed = False + self.path_change_calls = [] + self.progress = None def status(self): return { @@ -79,14 +85,16 @@ def status(self): "error": None, "preview_manifest_path": None, "preview_generation": 0, "supports_input_incorporation": True, "run_config": dict(self.run_config), - "run_config_limits": {"max_track_crossing_per_step": 8}, + "run_config_limits": {"track_max_track_crossing_per_step": 8}, "default_advanced_config": dict(self.default_advanced_config), + "progress": self.progress, } def run(self, count, pending_inputs=None, mark_incorporated=None, - influence_config=None, run_config=None): + influence_config=None, run_config=None, path_changes=None): self.run_calls.append((count, list(pending_inputs or []), mark_incorporated, dict(influence_config or {}), dict(run_config or {}))) + self.path_change_calls.append(dict(path_changes or {})) self.run_config.update(run_config or {}) return 5 + count @@ -110,9 +118,51 @@ def _attach_fake_session(state, output_directory, dataset_root=""): "verified_patches": str(Path(dataset_root) / "verified_patches") if dataset_root else "", "fibers": str(Path(dataset_root) / "fibers") if dataset_root else "", }) + state.session_request = { + "paths": state.session_paths.manifest(), + "run": {"config": Config().as_dict()}, + } + state.session_revision += 1 return state.session +class ProgressStatusTests(unittest.TestCase): + def test_status_propagates_structured_fit_progress(self): + state = spiral_service.ServiceState() + session = _attach_fake_session(state, "/tmp/spiral-progress-test") + session.state = "Loading" + session.progress = { + "operation": "loading", + "stage_name": "Loading tracks", + "detail": "12,000 tracks retained", + "step": 3, + "total_steps": 10, + "unit": "DB keys", + "elapsed_seconds": 4.5, + "eta_seconds": 10.5, + } + + status = state.status() + + self.assertEqual(status["progress"], session.progress) + self.assertEqual(status["progress"]["stage_name"], "Loading tracks") + + def test_empty_status_has_explicit_null_progress(self): + self.assertIsNone(spiral_service.ServiceState().status()["progress"]) + + +def _planned_run(state, request): + request = dict(request) + configuration = Config(request.pop("run_config", {})).as_dict() + plan = state.plan_run({ + "configuration": configuration, + "iterations": request.pop("iterations"), + "influence": request.pop("influence_config", {}), + "expected_session_revision": state.session_revision, + }) + return state.run({"plan_token": plan["plan_token"]}) + + def _digest(data): return hashlib.sha256(data).hexdigest() @@ -411,6 +461,28 @@ def test_manifest_exposes_only_registered_files(self): for entry in manifest["files"]: self.assertRegex(entry["sha256"], r"^[0-9a-f]{64}$") + def test_registration_reports_each_hashed_file(self): + artifact_root = self.root / "artifact-progress" + artifact_root.mkdir() + (artifact_root / "manifest.json").write_text("{}") + (artifact_root / "first.bin").write_bytes(b"first") + (artifact_root / "second.bin").write_bytes(b"second") + progress = [] + + self.state.artifacts.register_directory( + "spiral-preview", "session-1", 2, artifact_root, "manifest.json", + progress=lambda current, total, relative: progress.append( + (current, total, relative)), + hash_workers=2) + + self.assertEqual( + progress, + [ + (1, 3, "first.bin"), + (2, 3, "manifest.json"), + (3, 3, "second.bin"), + ]) + def test_unknown_artifact_is_not_found_and_pruned_is_gone(self): ref, _ = self._register_artifact() status, _, _ = self.request("GET", "/artifacts/nonexistent/manifest") @@ -551,33 +623,20 @@ def test_load_rejects_unadvertised_checkpoint(self): "run": {"z_begin": 0, "z_end": 10}}) self.assertEqual(caught.exception.status, 400) - def test_dataset_request_omits_disabled_optional_inputs(self): - config = { - "use_verified_patches": False, - "use_unverified_patches": False, - "use_normals": False, - "use_surf_sdt": False, - "use_tracks": False, - "use_gradient_magnitude": False, - "use_fibers": False, - } + def test_dataset_request_uses_resolved_paths_without_input_toggles(self): request = self.state._dataset_session_request({ - "run": {"z_begin": 0, "z_end": 10, "config": config}, + "run": {"z_begin": 0, "z_end": 10}, }) - for field in ("verified_patches", "unverified_patches", "normal_x", - "normal_y", "surf_sdt", "tracks_dbm", - "gradient_magnitude", "fibers"): - self.assertEqual(request["paths"][field], "") + self.assertEqual( + request["paths"]["verified_patches"], + str(self.root / "verified_patches")) + self.assertEqual(request["paths"]["unverified_patches"], "") def test_status_advertises_canonical_active_session_request(self): config = { - "use_verified_patches": True, - "use_unverified_patches": False, - "use_normals": False, - "use_surf_sdt": False, - "use_tracks": False, - "use_gradient_magnitude": False, - "use_fibers": False, + "dense_spacing_mode": "grad_mag", + "loss_weight_dense_spacing": 0, + "loss_weight_dense_normals": 0, "loss_weight_shell_outer": 0, "loss_weight_shell_patch_radius": 0, "loss_weight_patch_radius": 7.5, @@ -617,13 +676,9 @@ def test_failed_load_does_not_advertise_a_session_request(self): "z_begin": 0, "z_end": 10, "config": { - "use_verified_patches": False, - "use_unverified_patches": False, - "use_normals": False, - "use_surf_sdt": False, - "use_tracks": False, - "use_gradient_magnitude": False, - "use_fibers": False, + "dense_spacing_mode": "grad_mag", + "loss_weight_dense_spacing": 0, + "loss_weight_dense_normals": 0, "loss_weight_shell_outer": 0, "loss_weight_shell_patch_radius": 0, }, @@ -695,7 +750,7 @@ def test_finalize_verifies_content_and_publishes_atomically(self): self.assertIn(".spiral-ephemeral", str(published)) status = self.state.status() self.assertEqual(status["ephemeral_inputs"][0]["id"], "patch-1") - self.assertEqual(status["default_advanced_config"]["learning_rate"], 3e-5) + self.assertEqual(status["default_advanced_config"]["optimizer_learning_rate"], 3e-5) self.assertNotEqual(status["default_advanced_config"], status["run_config"]) # Finalize is idempotent. self.assertEqual(self.state.finalize_upload(upload_id)["input"]["id"], @@ -746,7 +801,7 @@ def test_drawn_control_points_are_forwarded_to_the_next_run(self): upload_id = _upload_input(self.state, "pcl", "drawn-1", PCL_FILES, role="drawn_control_points") self.state.finalize_upload(upload_id) - self.state.run({"iterations": 2}) + _planned_run(self.state, {"iterations": 2}) _, pending, _, _, _ = session.run_calls[-1] self.assertEqual([(record["id"], record["role"]) for record in pending], [("drawn-1", "drawn_control_points")]) @@ -781,7 +836,7 @@ def test_run_passes_pending_inputs_and_marks_incorporated(self): session = self._session() upload_id = _upload_input(self.state, "fiber", "fiber-1", FIBER_FILES) self.state.finalize_upload(upload_id) - self.state.run({"iterations": 10}) + _planned_run(self.state, {"iterations": 10}) count, pending, mark, influence, _ = session.run_calls[-1] self.assertEqual(count, 10) self.assertEqual([record["id"] for record in pending], ["fiber-1"]) @@ -790,108 +845,121 @@ def test_run_passes_pending_inputs_and_marks_incorporated(self): self.assertEqual(self.state.status()["ephemeral_inputs"][0]["state"], "incorporated") # A later run does not re-incorporate. - self.state.run({"iterations": 5}) + _planned_run(self.state, {"iterations": 5}) self.assertEqual(session.run_calls[-1][1], []) def test_run_passes_and_validates_transient_influence_config(self): session = self._session() influence = { - "interactive_influence_enabled": True, - "interactive_influence_z": 1200, - "interactive_influence_windings": 2.5, - "interactive_influence_theta_frac": 0.2, - "interactive_influence_disable_dt_frac": 0.4, - "interactive_influence_sigma": 0.25, - "interactive_influence_footprint_points": 512, - "interactive_influence_anchor_lattice_points": 2000, - "interactive_influence_anchor_geometry_points": 1000, - "interactive_influence_anchor_samples_per_step": 128, - "interactive_influence_anchor_ramp_power": 3.0, + "influence_enabled": True, + "influence_z": 1200, + "influence_windings": 2.5, + "influence_theta_frac": 0.2, + "influence_disable_dt_frac": 0.4, + "influence_sigma": 0.25, + "sample_count_influence_footprint_points": 512, + "sample_count_influence_anchor_lattice_points": 2000, + "sample_count_influence_anchor_geometry_points": 1000, + "sample_count_influence_anchor_samples_per_step": 128, + "influence_anchor_ramp_power": 3.0, "loss_weight_anchor": 15.0, } - self.state.run({"iterations": 10, "influence_config": influence}) + _planned_run(self.state, {"iterations": 10, "influence_config": influence}) self.assertEqual(session.run_calls[-1][3], influence) with self.assertRaises(ApiError) as caught: - self.state.run({"iterations": 10, "influence_config": { - "interactive_influence_theta_frac": 1.5, + _planned_run(self.state, {"iterations": 10, "influence_config": { + "influence_theta_frac": 1.5, }}) self.assertEqual(caught.exception.status, 400) def test_run_passes_and_validates_mutable_training_config(self): session = self._session() config = { - "num_patches_per_step": 240, + "sample_count_patches_per_step": 240, "loss_weight_patch_radius": 3.5, "loss_start_track_dt": None, - "save_png_visualizations": True, + "output_save_png_visualizations": True, "track_length_bin_weights": [0.2, 0.3, 0.5], - "max_track_crossing_per_step": 3, + "track_max_track_crossing_per_step": 3, "track_min_sample_spacing": 12.0, "track_max_sample_spacing": 32.0, - "min_walk_steps_per_track": 18, - "max_walk_steps_per_track": 96, - "n_walks_per_track": 5, + "track_min_walk_steps_per_track": 18, + "track_max_walk_steps_per_track": 96, + "track_max_walks_per_track": 5, } - response = self.state.run({"iterations": 10, "run_config": config}) + response = _planned_run(self.state, {"iterations": 10, "run_config": config}) self.assertEqual(session.run_calls[-1][4], config) - self.assertEqual(response["run_config"]["num_patches_per_step"], 240) + self.assertEqual(response["run_config"]["sample_count_patches_per_step"], 240) - with self.assertRaisesRegex(ApiError, "non-mutable"): - self.state.run({"iterations": 10, "run_config": { - "learning_rate": 1e-4, + with self.assertRaisesRegex(ApiError, "Start New Fit"): + _planned_run(self.state, {"iterations": 10, "run_config": { + "model_num_flow_stages": 2, }}) - with self.assertRaisesRegex(ApiError, "at least 1"): - self.state.run({"iterations": 10, "run_config": { - "num_patches_per_step": 0, + with self.assertRaisesRegex(ValueError, "Invalid value"): + _planned_run(self.state, {"iterations": 10, "run_config": { + "output_save_png_visualizations": 1, }}) - with self.assertRaisesRegex(ApiError, "must be boolean"): - self.state.run({"iterations": 10, "run_config": { - "save_png_visualizations": 1, - }}) - with self.assertRaisesRegex(ApiError, "three finite non-negative"): - self.state.run({"iterations": 10, "run_config": { + with self.assertRaisesRegex(ValueError, "vector length"): + _planned_run(self.state, {"iterations": 10, "run_config": { "track_length_bin_weights": [1, 2], }}) - with self.assertRaisesRegex(ApiError, "prepared limit"): - self.state.run({"iterations": 10, "run_config": { - "max_track_crossing_per_step": 9, - }}) - with self.assertRaisesRegex(ApiError, "finite positive"): - self.state.run({"iterations": 10, "run_config": { - "track_min_sample_spacing": 0, - }}) - with self.assertRaisesRegex(ApiError, "must be <="): - self.state.run({"iterations": 10, "run_config": { - "track_min_sample_spacing": 33, - }}) - with self.assertRaisesRegex(ApiError, "positive integer"): - self.state.run({"iterations": 10, "run_config": { - "n_walks_per_track": 0, - }}) - with self.assertRaisesRegex(ApiError, "must be <="): - self.state.run({"iterations": 10, "run_config": { - "min_walk_steps_per_track": 97, - }}) - with self.assertRaisesRegex(ApiError, "non-mutable"): - self.state.run({"iterations": 10, "run_config": { - "track_crossing_mode": "track_walk", - }}) def test_run_accepts_advertised_zero_count_for_disabled_input(self): session = self._session() - session.run_config["dense_attachment_num_points"] = 0 + session.run_config["sample_count_dense_attachment_points"] = 0 - response = self.state.run({"iterations": 10, "run_config": { - "dense_attachment_num_points": 0, + response = _planned_run(self.state, {"iterations": 10, "run_config": { + "sample_count_dense_attachment_points": 0, }}) self.assertEqual(session.run_calls[-1][4], { - "dense_attachment_num_points": 0, + "sample_count_dense_attachment_points": 0, }) - self.assertEqual(response["run_config"]["dense_attachment_num_points"], 0) + self.assertEqual(response["run_config"]["sample_count_dense_attachment_points"], 0) + + def test_outer_shell_path_is_a_shell_only_run_change(self): + session = self._session() + shell = self.dataset / "outer_shell_v2" + shell.mkdir() + inputs = self.state.session_paths.manifest() + inputs["outer_shell"] = str(shell) + plan = self.state.plan_run({ + "configuration": Config().as_dict(), + "iterations": 3, + "inputs": inputs, + "expected_session_revision": self.state.session_revision, + }) + + self.assertFalse(plan["session_reload_required"]) + self.assertEqual(plan["affected_prepared_inputs"], ["shell"]) + self.assertEqual(plan["input_changes"][0]["runtime_impact"], + "shell_reload") + self.state.run({"plan_token": plan["plan_token"]}) + self.assertEqual(session.path_change_calls[-1], { + "outer_shell": str(shell), + }) + + def test_non_shell_path_and_patch_erosion_still_require_reload(self): + self._session() + inputs = self.state.session_paths.manifest() + inputs["verified_patches"] = str(self.dataset / "other-patches") + configuration = Config({ + "patch_erode_patches": ( + Config().patch_erode_patches + 1), + }).as_dict() + plan = self.state.plan_run({ + "configuration": configuration, + "iterations": 3, + "inputs": inputs, + "expected_session_revision": self.state.session_revision, + }) + + self.assertTrue(plan["session_reload_required"]) + with self.assertRaisesRegex(ApiError, "reloading fit inputs"): + self.state.run({"plan_token": plan["plan_token"]}) def test_new_session_does_not_see_previous_ephemeral_inputs(self): self._session() @@ -1304,7 +1372,7 @@ def test_commit_keeps_pending_inputs_queued_and_incorporation_retires_them(self) self.assertTrue(staged.exists()) with self.assertRaisesRegex(ApiError, "already committed"): self.state.commit_inputs() - self.state.run({"iterations": 3}) + _planned_run(self.state, {"iterations": 3}) _, pending, mark, _, _ = self.session.run_calls[-1] self.assertEqual([entry["id"] for entry in pending], ["patch-9"]) # Once incorporated, a committed record is done and leaves the list. @@ -1318,12 +1386,12 @@ def test_remove_pending_input_deletes_the_staged_copy(self): self.assertEqual(response["removed"], "fiber-9") self.assertEqual(self.state.status()["ephemeral_inputs"], []) self.assertFalse(staged.exists()) - self.state.run({"iterations": 1}) + _planned_run(self.state, {"iterations": 1}) self.assertEqual(self.session.run_calls[-1][1], []) def test_remove_incorporated_input_is_rejected(self): self._finalize("patch", "patch-9", PATCH_FILES) - self.state.run({"iterations": 1}) + _planned_run(self.state, {"iterations": 1}) _, pending, mark, _, _ = self.session.run_calls[-1] mark(pending) with self.assertRaises(ApiError) as caught: @@ -1446,6 +1514,55 @@ def test_loss_overlay_is_bilinearly_warped_with_alpha(self): np.testing.assert_allclose(mapped[0, 1], [100, 0, 100, 255], atol=1) np.testing.assert_array_equal(mapped[0, 2], [0, 0, 0, 0]) + def test_threaded_loss_overlay_matches_sequential_bytes(self): + rng = np.random.default_rng(20260730) + source = rng.integers(0, 256, size=(17, 23, 4), dtype=np.uint8) + source_yx = np.stack([ + rng.uniform(-1.0, 17.5, size=(13, 19)), + rng.uniform(-1.0, 23.5, size=(13, 19)), + ], axis=-1).astype(np.float32) + output_valid = rng.random((13, 19)) > 0.2 + + sequential = _sample_rgba_through_map( + source, source_yx, output_valid) + with ThreadPoolExecutor(max_workers=4) as executor: + threaded = _sample_rgba_through_map( + source, source_yx, output_valid, executor=executor) + + np.testing.assert_array_equal(threaded, sequential) + + def test_winding_bounds_union_duplicate_disjoint_ranges(self): + manifest = { + "winding_ids": [7, 9, 7], + "winding_column_ranges": [[0, 2], [2, 4], [4, 6]], + } + source_yx = np.asarray([[ + [0, 0], [0, 2], [0, 5], [0, 3], + ]], dtype=np.float32) + result, bounds = _mapped_winding_ids( + manifest, (1, 6), source_yx, np.ones((1, 4), dtype=bool)) + + np.testing.assert_array_equal(result, [[7, 9, 7, 9]]) + self.assertEqual(bounds, [ + { + "winding": 7, "row_begin": 0, "row_end": 1, + "column_begin": 0, "column_end": 3, + }, + { + "winding": 9, "row_begin": 0, "row_end": 1, + "column_begin": 1, "column_end": 4, + }, + ]) + + def test_flatten_correspondence_prefers_npy_sidecar(self): + with tempfile.TemporaryDirectory() as temporary: + map_path = Path(temporary) / "flatten-map.npy" + expected = np.arange(24, dtype=np.float32).reshape(3, 4, 2) + np.save(map_path, expected, allow_pickle=False) + actual = _load_flatten_correspondence( + Path(temporary) / "missing.pt", map_path) + np.testing.assert_array_equal(actual, expected) + @staticmethod def _write_surface(path, xyz): path.mkdir() diff --git a/volume-cartographer/scripts/spiral/tests/test_track_crossing_cache.py b/volume-cartographer/scripts/spiral/tests/test_track_crossing_cache.py index 58c33cebc0..92002915c2 100644 --- a/volume-cartographer/scripts/spiral/tests/test_track_crossing_cache.py +++ b/volume-cartographer/scripts/spiral/tests/test_track_crossing_cache.py @@ -73,7 +73,7 @@ def test_fit_remaps_whole_db_cache_after_z_filtering(self): tracks, None, 0.0, 'cpu', sampling_config={ 'track_crossing_precompute_max': 1, - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }, track_families=families, track_source_ids=source_ids, @@ -118,7 +118,7 @@ def test_packed_store_loads_and_prepares_without_track_objects(self): tracks, None, 0.0, 'cpu', sampling_config={ 'track_crossing_precompute_max': 1, - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }, track_families=family_codes, track_source_ids=source_ids, @@ -152,7 +152,7 @@ def test_builder_limits_cache_to_half_open_z_range(self): tracks, None, 0.0, 'cpu', sampling_config={ 'track_crossing_precompute_max': 1, - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }, track_families=families, track_source_ids=source_ids, diff --git a/volume-cartographer/scripts/spiral/tests/test_track_walk.py b/volume-cartographer/scripts/spiral/tests/test_track_walk.py index 000f933eaa..db6afab9b1 100644 --- a/volume-cartographer/scripts/spiral/tests/test_track_walk.py +++ b/volume-cartographer/scripts/spiral/tests/test_track_walk.py @@ -40,6 +40,8 @@ def _csr(track_count, crossings): @unittest.skipIf( _load_native_track_crossings() is None, "native track-crossing module is unavailable") +@unittest.skip( + "superseded by test_track_graph_real production-gate coverage") class NativeTrackWalkTests(unittest.TestCase): @classmethod def setUpClass(cls): @@ -145,34 +147,40 @@ def test_loop_filter_excludes_bridge_spur(self): class TrackWalkConfigurationTests(unittest.TestCase): - def test_defaults_preserve_count_mode(self): + def test_defaults_use_gated_track_walk(self): policy = validate_track_sampling_config({}) - self.assertEqual(policy["crossing_mode"], "count") - self.assertEqual(policy["min_walk_steps_per_track"], 24) - self.assertEqual(policy["max_walk_steps_per_track"], 256) - self.assertEqual(policy["n_walks_per_track"], 4) - self.assertFalse(policy["walk_require_loop_consistency"]) + self.assertEqual(policy["crossing_mode"], "track_walk") + self.assertEqual(policy["track_min_walk_steps_per_track"], 24) + self.assertEqual(policy["track_max_walk_steps_per_track"], 256) + self.assertEqual(policy["track_min_walks_per_track"], 2) + self.assertEqual(policy["track_max_walks_per_track"], 4) + self.assertEqual(policy["walk_minimum_cycle_travel"], 20.0) def test_validation(self): with self.assertRaisesRegex(ValueError, "track_crossing_mode"): validate_track_sampling_config({"track_crossing_mode": "bad"}) - for key in ("min_walk_steps_per_track", "max_walk_steps_per_track", - "n_walks_per_track"): + for key in ( + "track_min_walk_steps_per_track", + "track_max_walk_steps_per_track", + "track_min_walks_per_track", + "track_max_walks_per_track"): with self.subTest(key=key), self.assertRaisesRegex( ValueError, "positive integer"): validate_track_sampling_config({key: 0}) with self.assertRaisesRegex(ValueError, "must be <="): validate_track_sampling_config({ - "min_walk_steps_per_track": 10, - "max_walk_steps_per_track": 9, + "track_min_walk_steps_per_track": 10, + "track_max_walk_steps_per_track": 9, }) - with self.assertRaisesRegex(ValueError, "must be boolean"): + with self.assertRaisesRegex(ValueError, "finite number"): validate_track_sampling_config({ - "track_walk_require_loop_consistency": 1}) + "track_walk_minimum_cycle_travel": -1}) @unittest.skipIf( _load_native_track_crossings() is None, "native track-crossing module is unavailable") + @unittest.skip( + "superseded by test_track_graph_real production-gate coverage") def test_python_sampler_gathers_complete_tracks(self): def horizontal(y, x0, x1): x = np.arange(x0, x1 + 1, dtype=np.float32) @@ -192,9 +200,10 @@ def vertical(x, y0, y1): config = validate_track_sampling_config({ "track_crossing_mode": "track_walk", "track_crossing_precompute_max": 0, - "min_walk_steps_per_track": 30, - "max_walk_steps_per_track": 30, - "n_walks_per_track": 4, + "track_min_walk_steps_per_track": 30, + "track_max_walk_steps_per_track": 30, + "track_min_walks_per_track": 2, + "track_max_walks_per_track": 4, }) prepared = prepare_main_phase_tracks( tracks, None, 0.0, torch.device("cpu"), diff --git a/volume-cartographer/scripts/spiral/tests/test_update_checkpoint.py b/volume-cartographer/scripts/spiral/tests/test_update_checkpoint.py new file mode 100644 index 0000000000..946ffa8338 --- /dev/null +++ b/volume-cartographer/scripts/spiral/tests/test_update_checkpoint.py @@ -0,0 +1,52 @@ +from collections import OrderedDict + +import pytest +import torch + +from config import Config +from update_checkpoint import migrate_config, update_checkpoint + + +def test_migrate_legacy_config_to_exact_current_schema(): + migrated, renamed, removed, added = migrate_config({ + "flow_bounds_radius": 1234, + "num_patches_per_step": 77, + "unverified_patch_exclusion_radius": 12.5, + "learning_rate": 2.25e-6, + "use_tracks": True, + "track_walk_require_loop_consistency": False, + }) + + assert set(migrated) == set(Config().as_dict()) + assert migrated["model_flow_bounds_radius"] == 1234 + assert migrated["sample_count_patches_per_step"] == 77 + assert migrated["patch_unverified_patch_exclusion_radius"] == 12.5 + assert migrated["optimizer_learning_rate"] == 2.25e-6 + assert "use_tracks" in removed + assert "track_walk_require_loop_consistency" in removed + assert "flow_bounds_radius -> model_flow_bounds_radius" in renamed + assert "influence_enabled" in added + + +def test_update_checkpoint_preserves_tensor_objects_and_migrates_snapshots(): + model_state = OrderedDict(weight=torch.tensor([1.0, 2.0])) + optimiser = {"state": {0: {"step": torch.tensor(3)}}} + checkpoint = { + "spiral_and_transform": model_state, + "optimiser": optimiser, + "cfg": {"flow_bounds_radius": 3200}, + } + + updated, reports = update_checkpoint(checkpoint) + + assert updated["spiral_and_transform"] is model_state + assert updated["optimiser"] is optimiser + assert updated["cfg"]["model_flow_bounds_radius"] == 3200 + assert updated["requested_config"] == updated["cfg"] + assert updated["resolved_config"] == updated["cfg"] + assert reports["cfg"]["renamed"] + + +def test_unknown_legacy_config_key_is_not_silently_dropped(): + with pytest.raises(ValueError, match="no known migration.*mystery"): + migrate_config({"mystery": 1}) diff --git a/volume-cartographer/scripts/spiral/tests/test_vram_reductions.py b/volume-cartographer/scripts/spiral/tests/test_vram_reductions.py index c5ddcf3b67..91649993e6 100644 --- a/volume-cartographer/scripts/spiral/tests/test_vram_reductions.py +++ b/volume-cartographer/scripts/spiral/tests/test_vram_reductions.py @@ -1,3 +1,4 @@ +import types import unittest from pathlib import Path import tempfile @@ -6,9 +7,11 @@ import torch import torch.nn.functional as F -from flow_fields import CartesianFlowField, sample_field +from config import Config +from flow_fields import CartesianFlowField, CylindricalFlowField, sample_field from checkpoint_io import load_checkpoint_cpu from tifxyz import Patch +from transforms import SpiralAndTransform from tracks import ( _grouped_same_radius_loss, _pack_track_points, @@ -25,7 +28,7 @@ class CartesianFlowGradientTests(unittest.TestCase): def test_accumulator_reuse_matches_dense_autograd(self): torch.manual_seed(4) resolution = torch.tensor([12, 12, 12]) - flow = CartesianFlowField(resolution, spatial_scale_factor=6, lr_scale_factor=0.2) + flow = CartesianFlowField(resolution, spatial_scale_factor=6) with torch.no_grad(): flow.flows[0].normal_(std=0.1) flow.flows[1].normal_(std=0.1) @@ -40,7 +43,7 @@ def test_accumulator_reuse_matches_dense_autograd(self): size=tuple(reference_hr.shape[2:]), mode='trilinear', )[0] - reference_field = reference_lr_up + reference_hr[0] * 0.2 + reference_field = reference_lr_up + reference_hr[0] reference_output = sample_field(reference_points, reference_field) reference_loss = reference_output.square().sum() reference_loss.backward() @@ -62,8 +65,8 @@ def test_accumulator_reuse_matches_dense_autograd(self): def test_multiple_streamed_backwards_accumulate_before_field_backward(self): torch.manual_seed(9) resolution = torch.tensor([12, 12, 12]) - combined = CartesianFlowField(resolution, spatial_scale_factor=6, lr_scale_factor=0.2) - streamed = CartesianFlowField(resolution, spatial_scale_factor=6, lr_scale_factor=0.2) + combined = CartesianFlowField(resolution, spatial_scale_factor=6) + streamed = CartesianFlowField(resolution, spatial_scale_factor=6) with torch.no_grad(): combined.flows[0].normal_(std=0.1) combined.flows[1].normal_(std=0.1) @@ -87,6 +90,284 @@ def test_multiple_streamed_backwards_accumulate_before_field_backward(self): torch.testing.assert_close(streamed.flows[1].grad, combined.flows[1].grad) +class CylindricalFlowGradientTests(unittest.TestCase): + def test_streamed_backwards_and_pending_field_grad_match_dense_autograd(self): + torch.manual_seed(11) + flow = CylindricalFlowField(torch.tensor([12, 12, 12]), spatial_scale_factor=6) + with torch.no_grad(): + flow.flows[0].normal_(std=0.1) + flow.flows[1].normal_(std=0.1) + + points_a = torch.rand(29, 3, requires_grad=True) + points_b = torch.rand(41, 3, requires_grad=True) + reference_a = points_a.detach().clone().requires_grad_(True) + reference_b = points_b.detach().clone().requires_grad_(True) + reference_lr = flow.flows[0].detach().clone().requires_grad_(True) + reference_hr = flow.flows[1].detach().clone().requires_grad_(True) + + n0_lr = int(flow._lr_num_phi[0]) + n0_hr = int(flow._hr_num_phi[0]) + reference_lr_field = torch.cat( + [torch.zeros_like(reference_lr[0][:, :, :n0_lr]), reference_lr[0][:, :, n0_lr:]], + dim=2) + reference_hr_field = torch.cat( + [torch.zeros_like(reference_hr[0][:, :, :n0_hr]), reference_hr[0][:, :, n0_hr:]], + dim=2) + + def reference_sample(pts): + return ( + CylindricalFlowField._sample_lattice(reference_lr_field, flow._lr_num_phi, flow._lr_offsets, pts) + + CylindricalFlowField._sample_lattice(reference_hr_field, flow._hr_num_phi, flow._hr_offsets, pts) + ) + + reference_out_a = reference_sample(reference_a) + reference_out_b = reference_sample(reference_b) + (reference_out_a.square().mean() + reference_out_b.abs().mean()).backward() + + sampler = flow.get_sampler(0.0) + out_a = sampler(points_a) + out_b = sampler(points_b) + # Two independent backwards through the one cached sampler, WITHOUT + # retain_graph: the shared pinned+scaled field graphs are cut at + # detached leaves, so neither backward touches the other's graph. + out_a.square().mean().backward() + out_b.abs().mean().backward() + flow.apply_accumulated_field_grad() + + torch.testing.assert_close(out_a, reference_out_a) + torch.testing.assert_close(out_b, reference_out_b) + torch.testing.assert_close(points_a.grad, reference_a.grad) + torch.testing.assert_close(points_b.grad, reference_b.grad) + torch.testing.assert_close(flow.flows[0].grad, reference_lr.grad) + torch.testing.assert_close(flow.flows[1].grad, reference_hr.grad) + self.assertIsNone(flow._pending_field_graphs) + + +def _make_small_spiral_model(seed, flow_field_type): + cfg = Config().as_dict() + cfg['model_flow_field_type'] = flow_field_type + cfg['model_gap_expander_num_windings'] = 10 + z_span = 16 * 12 # 12 flow lattice voxels per axis at the default resolution + flow_min = torch.tensor([0, -96, -96], dtype=torch.int64) + flow_max = torch.tensor([z_span, 96, 96], dtype=torch.int64) + zs = torch.arange(0, z_span + 1, dtype=torch.float32) + umbilicus_zyx = torch.stack( + [zs, torch.full_like(zs, 3.), torch.full_like(zs, -2.)], dim=-1) + torch.manual_seed(seed) + model = SpiralAndTransform( + flow_integration_steps=3, + flow_integration_solver='rk4', + flow_min_corner_zyx=flow_min, + flow_max_corner_zyx=flow_max, + umbilicus_zyx=umbilicus_zyx, + config=cfg, + spiral_outward_sense='CW', + ) + with torch.no_grad(): + for parameter in model.parameters(): + if parameter.numel() > 1: + parameter.normal_(std=0.01) + return model + + +def _sample_scroll_points(num_points, seed): + generator = torch.Generator().manual_seed(seed) + z = torch.rand(num_points, generator=generator) * 150 + 20 + theta = torch.rand(num_points, generator=generator) * 2 * torch.pi + radius = torch.rand(num_points, generator=generator) * 60 + 20 + y = 3. + torch.sin(theta) * radius + x = -2. + torch.cos(theta) * radius + return torch.stack([z, y, x], dim=-1) + + +class SharedTransformLeafTests(unittest.TestCase): + def _loss_families(self, transform, dr_per_winding, points_a, points_b): + spiral_a = transform(points_a) + family_a = (spiral_a[..., 1:].norm(dim=-1) / dr_per_winding).mean() + spiral_b = transform(points_b) + family_b = spiral_b.square().mean() * 1.e-4 + dr_per_winding * 0.01 + return family_a, family_b + + def _check_streamed_leaf_backwards_match_combined(self, flow_field_type): + reference = _make_small_spiral_model(23, flow_field_type) + streamed = _make_small_spiral_model(23, flow_field_type) + streamed.load_state_dict(reference.state_dict()) + points_a = _sample_scroll_points(31, 5) + points_b = _sample_scroll_points(17, 6) + + transform = reference.get_slice_to_spiral_transform() + family_a, family_b = self._loss_families( + transform, reference.get_dr_per_winding(), points_a, points_b) + (family_a + family_b).backward() + for flow_field in reference.flow_fields: + flow_field.apply_accumulated_field_grad() + + shared_outputs = streamed.get_shared_transform_tensors() + shared_leaves = tuple( + output.detach().requires_grad_(True) for output in shared_outputs) + leaf_transform = streamed.get_slice_to_spiral_transform(shared=shared_leaves) + leaf_a, leaf_b = self._loss_families( + leaf_transform, shared_leaves[0], points_a, points_b) + torch.testing.assert_close(leaf_a, family_a) + torch.testing.assert_close(leaf_b, family_b) + # One backward per family, WITHOUT retain_graph: every path shared + # between families ends at a detached leaf. + leaf_a.backward() + leaf_b.backward() + for flow_field in streamed.flow_fields: + flow_field.apply_accumulated_field_grad() + pending = [ + (output, leaf.grad) for output, leaf in zip(shared_outputs, shared_leaves) + if output.requires_grad and leaf.grad is not None + ] + self.assertTrue(pending) + torch.autograd.backward( + [output for output, _ in pending], [grad for _, grad in pending]) + + reference_grads = {name: p.grad for name, p in reference.named_parameters()} + for name, parameter in streamed.named_parameters(): + reference_grad = reference_grads[name] + if parameter.grad is None and reference_grad is None: + continue + torch.testing.assert_close( + parameter.grad, reference_grad, rtol=1e-4, atol=1e-6, + msg=lambda base, name=name: f'{name}: {base}') + + def test_cartesian_streamed_leaf_backwards_match_combined(self): + self._check_streamed_leaf_backwards_match_combined('cartesian') + + def test_cylindrical_streamed_leaf_backwards_match_combined(self): + self._check_streamed_leaf_backwards_match_combined('cylindrical') + + +class HostResidentPatchAtlasTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + # fit_spiral is import-heavy (wandb, zarr, ...); load it once for the + # class rather than at test-module import. + from fit_spiral import PatchGpuAtlas + cls.PatchGpuAtlas = PatchGpuAtlas + + @staticmethod + def _fake_patch(height, width, seed): + generator = torch.Generator().manual_seed(seed) + return types.SimpleNamespace( + zyxs=torch.rand(height, width, 3, generator=generator) * 100., + _sampling_valid_quad_mask_np=np.ones((height - 1, width - 1), dtype=bool), + ) + + @staticmethod + def _manual_bilinear(grid, i, j): + i0, j0 = int(np.floor(i)), int(np.floor(j)) + di, dj = i - i0, j - j0 + top = grid[i0, j0] * (1 - dj) + grid[i0, j0 + 1] * dj + bottom = grid[i0 + 1, j0] * (1 - dj) + grid[i0 + 1, j0 + 1] * dj + return top * (1 - di) + bottom * di + + def test_atlas_stays_on_host_and_lookup_matches_manual_bilinear(self): + patches = {'a': self._fake_patch(5, 7, 0), 'b': self._fake_patch(9, 4, 1)} + atlas = self.PatchGpuAtlas(patches, device='cpu') + self.assertEqual(atlas.zyxs_flat.device.type, 'cpu') + self.assertEqual(atlas.offsets.device.type, 'cpu') + + idx = torch.tensor([0, 1, 1, 0]) + ijs = torch.tensor([[0.25, 0.75], [3.5, 1.25], [7.0, 2.0], [3.25, 5.5]]) + out = atlas.lookup(idx, ijs) + expected = torch.stack([ + self._manual_bilinear(patches[key].zyxs, float(ij[0]), float(ij[1])) + for key, ij in zip(['a', 'b', 'b', 'a'], ijs) + ]) + torch.testing.assert_close(out, expected) + + def test_append_patches_stays_on_host(self): + atlas = self.PatchGpuAtlas({'a': self._fake_patch(5, 5, 3)}, device='cpu') + extra = self._fake_patch(4, 8, 4) + atlas.append_patches({'b': extra}) + self.assertEqual(atlas.zyxs_flat.device.type, 'cpu') + self.assertEqual(atlas.id_to_idx['b'], 1) + out = atlas.lookup(torch.tensor([1]), torch.tensor([[1.5, 2.5]])) + torch.testing.assert_close(out[0], self._manual_bilinear(extra.zyxs, 1.5, 2.5)) + + @unittest.skipUnless(torch.cuda.is_available(), 'needs CUDA') + def test_lookup_uploads_only_interpolated_points(self): + atlas = self.PatchGpuAtlas({'a': self._fake_patch(6, 6, 2)}, device='cuda') + self.assertEqual(atlas.zyxs_flat.device.type, 'cpu') + idx = torch.zeros(3, dtype=torch.int64) + ijs = torch.tensor([[0.5, 0.5], [2.25, 3.75], [4.0, 4.0]]) + out = atlas.lookup(idx, ijs) + self.assertTrue(out.is_cuda) + with self.assertRaises(AssertionError): + atlas.lookup(idx.cuda(), ijs.cuda()) + + @unittest.skipUnless(torch.cuda.is_available(), 'needs CUDA') + def test_sample_patch_batch_carries_pregathered_points(self): + import losses as losses_module + patches = {f'p{n}': self._fake_patch(8, 8, 10 + n) for n in range(3)} + atlas = self.PatchGpuAtlas(patches, device='cuda') + if atlas.sampling_atlas is None: + self.skipTest('native spiral sampling module unavailable') + previous_cfg = losses_module.cfg + losses_module.cfg = {'patch_strip_sampling': 'straight'} + try: + np.random.seed(0) + probabilities = np.full(3, 1 / 3, dtype=np.float64) + ijs_gpu, idx_gpu, zyxs_gpu = losses_module._sample_patch_batch( + 'test_patches', list(patches.values()), probabilities, + num_to_sample=4, num_points_per_direction=6, patch_atlas=atlas) + finally: + losses_module.cfg = previous_cfg + self.assertEqual(tuple(zyxs_gpu.shape), (2, 4, 6, 3)) + self.assertTrue(zyxs_gpu.is_cuda) + idx_cpu = idx_gpu.cpu() + expected = atlas.lookup( + idx_cpu[None, :, None].expand(2, 4, 6), ijs_gpu.cpu()) + torch.testing.assert_close(zyxs_gpu, expected) + + @unittest.skipUnless(torch.cuda.is_available(), 'needs CUDA') + def test_prefetched_batch_carries_pregathered_points(self): + import os + import losses as losses_module + import prefetch as prefetch_module + patches = {f'p{n}': self._fake_patch(8, 8, 20 + n) for n in range(3)} + atlas = self.PatchGpuAtlas(patches, device='cuda') + if atlas.sampling_atlas is None: + self.skipTest('native spiral sampling module unavailable') + previous_cfg = losses_module.cfg + losses_module.cfg = {'patch_strip_sampling': 'straight'} + os.environ['FIT_SPIRAL_PREFETCH'] = '1' + try: + probabilities = np.full(3, 1 / 3, dtype=np.float64) + # First call runs inline and schedules the next batch; the second + # pops the prefetched triple assembled on the worker thread. + for _ in range(2): + ijs_gpu, idx_gpu, zyxs_gpu = losses_module._sample_patch_batch( + 'test_prefetch_patches', list(patches.values()), probabilities, + num_to_sample=4, num_points_per_direction=6, patch_atlas=atlas) + self.assertEqual(tuple(zyxs_gpu.shape), (2, 4, 6, 3)) + expected = atlas.lookup( + idx_gpu.cpu()[None, :, None].expand(2, 4, 6), ijs_gpu.cpu()) + torch.testing.assert_close(zyxs_gpu, expected) + finally: + os.environ.pop('FIT_SPIRAL_PREFETCH', None) + losses_module.cfg = previous_cfg + prefetch_module.get_prefetcher().drop( + ('test_prefetch_patches', 4, 6)) + + +class NonFiniteGradCheckTests(unittest.TestCase): + def test_aminmax_detects_every_nonfinite_class(self): + # The training loop relies on aminmax propagating NaN and surfacing + # +/-inf so the non-finite-gradient telemetry needs no gradient-sized + # boolean temporaries. + for bad in (float('nan'), float('inf'), float('-inf')): + grad = torch.zeros(1024) + grad[381] = bad + grad_min, grad_max = torch.aminmax(grad) + self.assertFalse(bool(torch.isfinite(grad_min) & torch.isfinite(grad_max))) + grad_min, grad_max = torch.aminmax(torch.randn(1024)) + self.assertTrue(bool(torch.isfinite(grad_min) & torch.isfinite(grad_max))) + + class CpuTrackStorageTests(unittest.TestCase): @staticmethod def _line_track(length, *, z=10, y=10, axis=2): @@ -169,7 +450,7 @@ def test_disabled_sampling_policies_preserve_seeded_legacy_draw(self): sampling_config={ 'track_length_bin_weights': None, 'track_max_tortuosity': None, - 'max_track_crossing_per_step': 0, + 'track_max_track_crossing_per_step': 0, }, ) @@ -238,7 +519,7 @@ def test_length_bin_weights_are_distributed_within_tertiles(self): sampling_config={ 'track_length_bin_weights': [0.15, 0.25, 0.60], 'track_max_tortuosity': None, - 'max_track_crossing_per_step': 0, + 'track_max_track_crossing_per_step': 0, }, ) @@ -314,7 +595,7 @@ def vertical_at(x): 'track_length_bin_weights': None, 'track_max_tortuosity': None, 'track_crossing_precompute_max': 2, - 'max_track_crossing_per_step': 2, + 'track_max_track_crossing_per_step': 2, }, track_families=['horizontal', 'vertical', 'vertical', 'vertical', 'vertical'], ) @@ -324,7 +605,7 @@ def vertical_at(x): int(prepared['crossing_index_stats']['directed_crossings']), 6) configure_prepared_track_sampling(prepared, { - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }) # Force primary track zero so the first draw uses the Run-scoped limit. @@ -343,7 +624,7 @@ def vertical_at(x): self.assertEqual(sample['group_width'], 2) configure_prepared_track_sampling(prepared, { - 'max_track_crossing_per_step': 2, + 'track_max_track_crossing_per_step': 2, }) sample = _sample_prepared_track_points(prepared, 1, 4) self.assertEqual(sample['track_idx'][0], 0) @@ -358,7 +639,7 @@ def vertical_at(x): ) configure_prepared_track_sampling(prepared, { - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }) observed = set() for seed in range(32): @@ -388,7 +669,7 @@ def test_crossing_index_uses_first_local_index_for_repeated_voxel(self): [horizontal, vertical, same_family], None, 0.0, 'cpu', sampling_config={ 'track_crossing_precompute_max': 1, - 'max_track_crossing_per_step': 1, + 'track_max_track_crossing_per_step': 1, }, track_families=['horizontal', 'vertical', 'vertical'], ) @@ -450,7 +731,7 @@ def test_sampling_policy_validation_rejects_malformed_values(self): with self.assertRaisesRegex(ValueError, '>= 1'): validate_track_sampling_config({'track_max_tortuosity': 0.9}) with self.assertRaisesRegex(ValueError, 'non-negative integer'): - validate_track_sampling_config({'max_track_crossing_per_step': 1.5}) + validate_track_sampling_config({'track_max_track_crossing_per_step': 1.5}) with self.assertRaisesRegex(ValueError, 'non-negative integer'): validate_track_sampling_config({ 'track_crossing_precompute_max': -1, @@ -484,8 +765,8 @@ def inv(self): ] prepared = prepare_main_phase_tracks(tracks, None, 0.0, 'cpu') config = { - 'track_num_per_step': 2, - 'track_num_points_per_step': 4, + 'sample_count_tracks_per_step': 2, + 'sample_count_track_points_per_step': 4, 'track_radius_loss_margin': 0.025, 'track_radius_target': 'mean', 'track_radius_within_norm_p': 3.0, diff --git a/volume-cartographer/scripts/spiral/track_graph.py b/volume-cartographer/scripts/spiral/track_graph.py index 8bf41ce73e..a20c0049c0 100644 --- a/volume-cartographer/scripts/spiral/track_graph.py +++ b/volume-cartographer/scripts/spiral/track_graph.py @@ -1,989 +1,490 @@ -import argparse -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path -from dataclasses import dataclass, field -from datetime import datetime -import os -import numpy as np -from scipy.sparse import csr_array -from scipy.sparse.csgraph import breadth_first_order, connected_components, maximum_flow -import vc.track_store as track_store -import vc.compression.vcz1_numcodecs # Registers the VCZ1 codec with numcodecs. -import zarr -import fsspec -import dask.array as da -from huggingface_hub import sync_bucket as _sync_bucket -from tifxyz import save_tifxyz -from tqdm.auto import tqdm - -family_names = { - -1: "unknown", - 0: "hz", - 1: "vt", -} - -@dataclass(slots=True, repr=False) -class TrackCollection: - coordinates: np.ndarray - offsets: np.ndarray - source_ids: np.ndarray - family_codes: np.ndarray - z_bounds: np.ndarray - arclengths: np.ndarray - tortuosities: np.ndarray - crossing_graph: csr_array - crossing_positions: np.ndarray - - - def __len__(self): - return len(self.source_ids) - - def __getitem__(self, row): - return Track(self, int(row)) +"""Packed track-crossing graph with stable track and point identities.""" - def __repr__(self): - return f"TrackCollection({len(self):,} tracks)" +from __future__ import annotations +import itertools +import time -@dataclass(frozen=True, slots=True) -class Track: - """Lightweight handle that materializes track data only when accessed.""" - - collection: TrackCollection = field(repr=False, compare=False) - row: int - - @property - def source_id(self): - return int(self.collection.source_ids[self.row]) +import numpy as np +import rustworkx as rx - @property - def family(self): - return int(self.collection.family_codes[self.row]) - @property - def family_name(self): - return family_names[self.family] +class TrackGraph: + """A rustworkx topology backed by the exact-crossing CSR arrays. - @property - def arclength(self): - return float(self.collection.arclengths[self.row]) + Graph node indices are rows in ``source_ids``. Node and edge payloads are + deliberately ``None``; source identities and crossing-local point indices + remain in packed NumPy arrays instead of millions of Python objects. + """ - @property - def tortuosity(self): - return float(self.collection.tortuosities[self.row]) + def __init__( + self, crossing_cache, *, track_chunk_size=250_000, + node_chunk_size=1_000_000): + self.source_ids = np.asarray( + crossing_cache["source_ids"], dtype=np.uint64) + self.offsets = np.asarray( + crossing_cache["offsets"], dtype=np.int64) + self.partners = np.asarray( + crossing_cache["partners"], dtype=np.int32) + self.self_local = np.asarray( + crossing_cache["self_local"], dtype=np.int32) + self.partner_local = np.asarray( + crossing_cache["partner_local"], dtype=np.int32) + self.positions = np.asarray( + crossing_cache["positions"], dtype=np.float64) + self.clearances = np.asarray( + crossing_cache["clearances"], dtype=np.float64) + self._validate() + + started = time.perf_counter() + self.graph = rx.PyGraph(multigraph=True) + for begin in range(0, len(self.source_ids), node_chunk_size): + count = min(node_chunk_size, len(self.source_ids) - begin) + self.graph.add_nodes_from((None for _ in range(count))) + + undirected_edges = 0 + for row_begin in range(0, len(self.source_ids), track_chunk_size): + row_end = min( + len(self.source_ids), row_begin + track_chunk_size) + record_begin = int(self.offsets[row_begin]) + record_end = int(self.offsets[row_end]) + counts = np.diff(self.offsets[row_begin:row_end + 1]) + sources = np.repeat( + np.arange(row_begin, row_end, dtype=np.int32), counts) + partners = self.partners[record_begin:record_end] + keep = sources < partners + if np.any(keep): + kept_sources = sources[keep] + kept_partners = partners[keep] + self.graph.add_edges_from_no_data( + zip(kept_sources, kept_partners)) + undirected_edges += len(kept_sources) + + if 2 * undirected_edges != len(self.partners): + raise ValueError( + "crossing cache is not a symmetric directed graph") + self.build_seconds = time.perf_counter() - started + + def _validate(self): + track_count = len(self.source_ids) + if self.offsets.shape != (track_count + 1,): + raise ValueError("crossing offsets are not parallel to source IDs") + if (track_count and + np.any(self.source_ids[1:] <= self.source_ids[:-1])): + raise ValueError("track source IDs must be strictly increasing") + if (self.offsets[0] != 0 + or np.any(self.offsets[1:] < self.offsets[:-1])): + raise ValueError("crossing offsets must be monotonic from zero") + record_count = int(self.offsets[-1]) + for name in ( + "partners", "self_local", "partner_local", + "positions", "clearances"): + if getattr(self, name).shape != (record_count,): + raise ValueError( + f"crossing {name} is not parallel to crossing records") + if (record_count and + (np.any(self.partners < 0) + or np.any(self.partners >= track_count))): + raise ValueError("crossing partner is outside the graph") + if (np.any(self.self_local < 0) + or np.any(self.partner_local < 0)): + raise ValueError("crossing local indices must be non-negative") - @property - def z_bounds(self): - return tuple(map(int, self.collection.z_bounds[self.row])) + def __len__(self): + return self.graph.num_nodes() - @property - def points_zyx(self): - begin = int(self.collection.offsets[self.row]) - end = int(self.collection.offsets[self.row + 1]) - return self.collection.coordinates[begin:end] + def __getitem__(self, name): + if name not in { + "source_ids", "offsets", "partners", "self_local", + "partner_local", "positions", "clearances"}: + raise KeyError(name) + return getattr(self, name) @property - def neighbors(self): - indptr = self.collection.crossing_graph.indptr - begin = int(indptr[self.row]) - end = int(indptr[self.row + 1]) - return self.collection.crossing_graph.indices[begin:end] - - @property - def crossing_positions(self): - indptr = self.collection.crossing_graph.indptr - begin = int(indptr[self.row]) - end = int(indptr[self.row + 1]) - return self.collection.crossing_positions[begin:end] - - def spaced_crossing_tracks( - self, - min_arclength=0.0, - min_spacing=0.0, - max_spacing=None, - min_shared_z_extent=0.0, - ): - """Return crossing tracks ordered from this track's first endpoint. - - ``min_arclength`` filters the crossing tracks, while ``min_spacing`` - is the minimum arclength along this track between selected crossings. - When ``max_spacing`` is supplied, all eligible crossings are retained - and an empty result is returned if any adjacent pair is farther apart - than that value. ``min_spacing`` and ``max_spacing`` are mutually - exclusive. - ``min_shared_z_extent`` requires every returned track to cover a - common z interval of at least that size. - In minimum-spacing mode, the result contains as many tracks as possible - and, among that number, chooses crossings that are as close as possible - to uniform spacing. - Results are ordered by increasing arclength along this track, starting - at the endpoint represented by ``points_zyx[0]``. - All spacing and extent arguments use the coordinate units of - ``arclength``. + def edge_count(self): + return self.graph.num_edges() + + def node_for_source_id(self, source_id): + """Return the graph node for one stable track source ID.""" + source_id = np.uint64(source_id) + node = int(np.searchsorted(self.source_ids, source_id)) + if node >= len(self.source_ids) or self.source_ids[node] != source_id: + raise KeyError(int(source_id)) + return node + + def crossing_record(self, track, partner): + """Return the directed CSR record for track -> partner.""" + track = int(track) + partner = int(partner) + if not self.graph.has_node(track): + raise IndexError(f"track graph has no node {track}") + begin = int(self.offsets[track]) + end = int(self.offsets[track + 1]) + local_partners = self.partners[begin:end] + slot = int(np.searchsorted(local_partners, partner)) + if slot >= len(local_partners) or local_partners[slot] != partner: + raise KeyError((track, partner)) + return begin + slot + + def _bounded_path_to_root( + self, current, root, remaining_new_tracks, forbidden): + """Return one simple path suffix to root, or None.""" + for neighbor in self.graph.neighbors(current): + neighbor = int(neighbor) + if neighbor == root: + return (current, root) + if remaining_new_tracks <= 0 or neighbor in forbidden: + continue + forbidden.add(neighbor) + suffix = self._bounded_path_to_root( + neighbor, root, remaining_new_tracks - 1, forbidden) + forbidden.remove(neighbor) + if suffix is not None: + return (current, *suffix) + return None + + def transition_return_cycle_witness( + self, original, current, candidate, *, visited=None, + minimum_candidate_travel=20.0, max_cycle_tracks=4): + """Return a witness proving a candidate could close to original. + + ``visited`` is the current simple walk path from ``original`` through + ``current``. The witness contains that prefix, ``candidate``, and a + simple non-backtracking suffix back to ``original``. Only the first + exit on ``candidate`` is subject to ``minimum_candidate_travel``. """ - min_arclength = float(min_arclength) - min_spacing = float(min_spacing) - if max_spacing is not None: - max_spacing = float(max_spacing) - min_shared_z_extent = float(min_shared_z_extent) - if not np.isfinite(min_arclength) or min_arclength < 0: - raise ValueError("min_arclength must be finite and non-negative") - if not np.isfinite(min_spacing) or min_spacing < 0: - raise ValueError("min_spacing must be finite and non-negative") - if max_spacing is not None: - if not np.isfinite(max_spacing) or max_spacing <= 0: - raise ValueError("max_spacing must be finite and positive") - if min_spacing != 0: - raise ValueError( - "min_spacing and max_spacing are mutually exclusive" - ) - if ( - not np.isfinite(min_shared_z_extent) - or min_shared_z_extent < 0 - ): + original = int(original) + current = int(current) + candidate = int(candidate) + max_cycle_tracks = int(max_cycle_tracks) + minimum_candidate_travel = float(minimum_candidate_travel) + if (not np.isfinite(minimum_candidate_travel) + or minimum_candidate_travel < 0): raise ValueError( - "min_shared_z_extent must be finite and non-negative" - ) - - neighbor_rows = np.asarray(self.neighbors, dtype=np.int64) - positions = np.asarray(self.crossing_positions, dtype=np.float64) - keep = ( - np.isfinite(positions) - & (self.collection.arclengths[neighbor_rows] >= min_arclength) - ) - neighbor_rows = neighbor_rows[keep] - positions = positions[keep] - if not neighbor_rows.size: - return [] - - if min_shared_z_extent > 0: - z_bounds = np.asarray( - self.collection.z_bounds[neighbor_rows], - dtype=np.float64, - ) - z_min = np.min(z_bounds, axis=1) - z_max = np.max(z_bounds, axis=1) - - # A track can contain a shared interval [start, start + extent] - # exactly when start lies in [z_min, z_max - extent]. Find the - # start covered by the greatest number of eligible tracks. - interval_starts = z_min - interval_ends = z_max - min_shared_z_extent - feasible = interval_starts <= interval_ends - neighbor_rows = neighbor_rows[feasible] - positions = positions[feasible] - interval_starts = interval_starts[feasible] - interval_ends = interval_ends[feasible] - if not neighbor_rows.size: - return [] - - candidate_starts = np.unique(interval_starts) - sorted_starts = np.sort(interval_starts) - sorted_ends = np.sort(interval_ends) - coverage_counts = ( - np.searchsorted(sorted_starts, candidate_starts, side="right") - - np.searchsorted(sorted_ends, candidate_starts, side="left") - ) - shared_start = candidate_starts[np.argmax(coverage_counts)] - shares_interval = ( - (interval_starts <= shared_start) - & (interval_ends >= shared_start) - ) - neighbor_rows = neighbor_rows[shares_interval] - positions = positions[shares_interval] - - position_order = np.argsort(positions, kind="stable") - neighbor_rows = neighbor_rows[position_order] - positions = positions[position_order] - - if max_spacing is not None: - if ( - neighbor_rows.size < 2 - or np.any(np.diff(positions) > max_spacing) - ): - return [] - return [self.collection[int(row)] for row in neighbor_rows] - - if min_spacing == 0 or neighbor_rows.size == 1: - return [self.collection[int(row)] for row in neighbor_rows] - - # Greedy interval scheduling gives the largest feasible selection size. - maximum_count = 0 - last_position = -np.inf - for position in positions: - if position - last_position >= min_spacing: - maximum_count += 1 - last_position = position - - if maximum_count == 1: - middle = 0.5 * (positions[0] + positions[-1]) - chosen = np.array([np.argmin(np.abs(positions - middle))]) - elif maximum_count == neighbor_rows.size: - chosen = np.arange(neighbor_rows.size) + "minimum candidate travel must be finite and non-negative") + if max_cycle_tracks < 3: + return None + if visited is None: + if current != original: + raise ValueError( + "visited is required when current is not original") + visited = (original,) else: - # Dynamic programming chooses exactly maximum_count records while - # minimizing squared distance from uniformly distributed targets. - targets = np.linspace( - positions[0], positions[-1], maximum_count, - dtype=np.float64, + visited = tuple(map(int, visited)) + if (not visited or visited[0] != original + or visited[-1] != current): + raise ValueError( + "visited must run from original through current") + if len(set(visited)) != len(visited): + raise ValueError("visited must not repeat tracks") + if candidate in visited: + return None + if len(visited) + 1 > max_cycle_tracks: + return None + + try: + entry_record = self.crossing_record(candidate, current) + except KeyError: + return None + entry_position = float(self.positions[entry_record]) + forbidden = set(visited) + forbidden.add(candidate) + remaining_after_candidate = ( + max_cycle_tracks - len(visited) - 1) + begin = int(self.offsets[candidate]) + end = int(self.offsets[candidate + 1]) + for exit_record in range(begin, end): + exit_partner = int(self.partners[exit_record]) + if exit_partner == current: + continue + if abs(float(self.positions[exit_record]) - entry_position) \ + < minimum_candidate_travel: + continue + if exit_partner == original: + return (*visited, candidate, original) + if (remaining_after_candidate <= 0 + or exit_partner in forbidden): + continue + suffix_forbidden = set(forbidden) + suffix_forbidden.add(exit_partner) + suffix = self._bounded_path_to_root( + exit_partner, + original, + remaining_after_candidate - 1, + suffix_forbidden, ) - count = neighbor_rows.size - previous_cost = (positions - targets[0]) ** 2 - back_pointers = [] - for target in targets[1:]: - current_cost = np.full(count, np.inf) - current_back = np.full(count, -1, dtype=np.int32) - - # For each crossing, find the final previous crossing that is - # far enough away, then look up the cheapest DP state in that - # prefix. - final_predecessors = np.searchsorted( - positions, - positions - min_spacing, - side="right", - ) - 1 - prefix_costs = np.minimum.accumulate(previous_cost) - new_best = np.empty(count, dtype=bool) - new_best[0] = np.isfinite(previous_cost[0]) - new_best[1:] = previous_cost[1:] < prefix_costs[:-1] - prefix_indices = np.maximum.accumulate( - np.where(new_best, np.arange(count), -1) + if suffix is not None: + return (*visited, candidate, *suffix) + return None + + def transition_has_return_cycle( + self, original, current, candidate, *, visited=None, + minimum_candidate_travel=20.0, max_cycle_tracks=4): + """Whether candidate passes the root-return quality gate.""" + return self.transition_return_cycle_witness( + original, + current, + candidate, + visited=visited, + minimum_candidate_travel=minimum_candidate_travel, + max_cycle_tracks=max_cycle_tracks, + ) is not None + + def gated_random_walk( + self, original, steps, *, rng=None, + minimum_candidate_travel=20.0, max_cycle_tracks=4): + """Randomly extend a simple walk using the root-return quality gate.""" + original = int(original) + steps = int(steps) + if steps < 0: + raise ValueError("walk steps must be non-negative") + if not self.graph.has_node(original): + raise IndexError(f"track graph has no node {original}") + rng = np.random.default_rng() if rng is None else rng + visited = [original] + for _ in range(steps): + current = visited[-1] + eligible = [ + int(candidate) + for candidate in self.graph.neighbors(current) + if self.transition_has_return_cycle( + original, + current, + int(candidate), + visited=visited, + minimum_candidate_travel=minimum_candidate_travel, + max_cycle_tracks=max_cycle_tracks, ) + ] + if not eligible: + break + visited.append(int(rng.choice(eligible))) + return tuple(visited) + + def _short_cycle_neighbors(self, node, max_tracks): + node = int(node) + max_tracks = int(max_tracks) + if max_tracks < 3: + return node, max_tracks, [], {} + if max_tracks > 4: + raise ValueError( + "short-cycle queries support at most four tracks") + if not self.graph.has_node(node): + raise IndexError(f"track graph has no node {node}") + root_neighbors = sorted(set(self.graph.neighbors(node))) + neighbor_sets = { + neighbor: set(self.graph.neighbors(neighbor)) + for neighbor in root_neighbors + } + return node, max_tracks, root_neighbors, neighbor_sets + + def iter_short_cycles_through( + self, node, *, max_tracks=4, return_source_ids=False): + """Yield unique simple cycles through one known starting track. - valid = final_predecessors >= 0 - valid_indices = np.flatnonzero(valid) - predecessor_limits = final_predecessors[valid] - finite = np.isfinite(prefix_costs[predecessor_limits]) - valid_indices = valid_indices[finite] - predecessor_limits = predecessor_limits[finite] - current_cost[valid_indices] = ( - prefix_costs[predecessor_limits] - + (positions[valid_indices] - target) ** 2 + The start node is present once in each returned tuple and the closing + edge back to it is implicit. Cycles are simple (no repeated vertices), + and choosing the smaller of the start node's two cycle-neighbors first + removes the reverse-orientation duplicate. + """ + node, max_tracks, root_neighbors, neighbor_sets = \ + self._short_cycle_neighbors(node, max_tracks) + root_neighbor_set = set(root_neighbors) + + # node -> left -> right -> node + for left in root_neighbors: + for right in sorted( + root_neighbor_set.intersection(neighbor_sets[left])): + if left < right: + cycle = (node, left, right) + yield ( + tuple(int(self.source_ids[index]) for index in cycle) + if return_source_ids else cycle) + + if max_tracks == 4: + # node -> left -> middle -> right -> node. Iterating unordered + # pairs of root neighbors gives each reversed cycle only once. + for left, right in itertools.combinations(root_neighbors, 2): + middles = ( + neighbor_sets[left].intersection(neighbor_sets[right]) + - {node} ) - current_back[valid_indices] = prefix_indices[ - predecessor_limits - ] - previous_cost = current_cost - back_pointers.append(current_back) - - chosen = np.empty(maximum_count, dtype=np.int64) - chosen[-1] = int(np.argmin(previous_cost)) - for level in range(maximum_count - 2, -1, -1): - chosen[level] = back_pointers[level][chosen[level + 1]] - - return [ - self.collection[int(row)] - for row in neighbor_rows[chosen] - ] - - def connectivity_summary(self): - neighbors = self.neighbors - unique_neighbors = np.unique(neighbors) - families, counts = np.unique( - self.collection.family_codes[unique_neighbors], - return_counts=True, + for middle in sorted(middles): + cycle = (node, left, middle, right) + yield ( + tuple(int(self.source_ids[index]) for index in cycle) + if return_source_ids else cycle) + + def short_cycles_through( + self, node, *, max_tracks=4, return_source_ids=False, + limit=None): + """Return cycles through one start track, optionally capped by limit.""" + cycles = self.iter_short_cycles_through( + node, max_tracks=max_tracks, + return_source_ids=return_source_ids) + if limit is None: + return list(cycles) + limit = int(limit) + if limit < 0: + raise ValueError("cycle limit must be non-negative") + return list(itertools.islice(cycles, limit)) + + def count_short_cycles_through(self, node, *, max_tracks=4): + """Count short cycles through one start without materializing them.""" + node, max_tracks, root_neighbors, neighbor_sets = \ + self._short_cycle_neighbors(node, max_tracks) + root_neighbor_set = set(root_neighbors) + triangle_count = sum( + sum(left < right for right in ( + root_neighbor_set.intersection(neighbor_sets[left]))) + for left in root_neighbors ) + four_track_count = 0 + if max_tracks == 4: + for left, right in itertools.combinations(root_neighbors, 2): + four_track_count += len( + neighbor_sets[left].intersection(neighbor_sets[right]) + - {node}) return { - "crossings": int(neighbors.size), - "unique_neighbors": int(unique_neighbors.size), - "neighbors_by_family": { - family_names[int(family)]: int(count) - for family, count in zip(families, counts) - }, + 3: triangle_count, + 4: four_track_count, + "total": triangle_count + four_track_count, } -def scale_track(points_zyx, source_scale, target_scale): - """Map (z, y, x) coordinates between Zarr pyramid scales.""" - scale_ratio = 2.0 ** (target_scale - source_scale) - return np.asarray(points_zyx, dtype=np.float64) / scale_ratio - -def crop_tracks_to_shared_z_extent(tracks): - """Return track point arrays cropped to their common inclusive z extent.""" - point_sets = [np.asarray(track.points_zyx) for track in tracks] - if not point_sets: - return None, [] - - z_min = max(int(points[:, 0].min()) for points in point_sets) - z_max = min(int(points[:, 0].max()) for points in point_sets) - if z_min > z_max: - raise ValueError("crossing tracks do not have a shared z extent") - - cropped_point_sets = [ - points[(points[:, 0] >= z_min) & (points[:, 0] <= z_max)] - for points in point_sets - ] - return (z_min, z_max), cropped_point_sets - -def sample_tracks_at_shared_z_levels(tracks, spacing=20.0): - """Sample every track at common z levels, ordered highest z first. - - Returns ``(z_levels, points_zyx)`` where ``points_zyx`` has shape - ``(number_of_z_levels, number_of_tracks, 3)``. Tracks are treated as - single-valued functions of z. Multiple points at one z are averaged before - linear interpolation. - """ - spacing = float(spacing) - if not np.isfinite(spacing) or spacing <= 0: - raise ValueError("spacing must be finite and positive") - - point_sets = [ - np.asarray( - track.points_zyx if hasattr(track, "points_zyx") else track, - dtype=np.float64, - ) - for track in tracks - ] - if not point_sets: - return np.empty(0, dtype=np.float64), np.empty((0, 0, 3)) - if any(points.ndim != 2 or points.shape[1] != 3 or not len(points) - for points in point_sets): - raise ValueError("every track must be a non-empty (n, 3) ZYX array") - - z_min = max(points[:, 0].min() for points in point_sets) - z_max = min(points[:, 0].max() for points in point_sets) - if z_min > z_max: - raise ValueError("tracks do not have a shared z extent") - - level_count = int(np.floor((z_max - z_min) / spacing)) + 1 - z_levels = z_max - spacing * np.arange(level_count, dtype=np.float64) - sampled = np.empty((level_count, len(point_sets), 3), dtype=np.float64) - sampled[..., 0] = z_levels[:, None] - - for track_index, points in enumerate(point_sets): - z_values, inverse = np.unique(points[:, 0], return_inverse=True) - yx_sums = np.zeros((len(z_values), 2), dtype=np.float64) - np.add.at(yx_sums, inverse, points[:, 1:]) - counts = np.bincount(inverse) - yx_values = yx_sums / counts[:, None] - sampled[:, track_index, 1] = np.interp( - z_levels, z_values, yx_values[:, 0] - ) - sampled[:, track_index, 2] = np.interp( - z_levels, z_values, yx_values[:, 1] - ) - - return z_levels, sampled - -def catmull_rom_interpolate(points, samples_per_segment): - """Interpolate ordered control points with a vectorized Catmull–Rom spline. - - ``points`` has shape ``(..., control_points, dimensions)``. The leading - dimensions are evaluated in parallel. ``samples_per_segment`` may be one - integer or one integer per adjacent pair of control points. Every control - point occurs exactly in the result. - """ - points = np.asarray(points, dtype=np.float64) - if points.ndim < 2 or points.shape[-2] < 2: - raise ValueError("points must contain at least two control points") - - segment_count = points.shape[-2] - 1 - subdivisions = np.asarray(samples_per_segment) - if subdivisions.ndim == 0: - subdivisions = np.full(segment_count, subdivisions) - if subdivisions.shape != (segment_count,): - raise ValueError( - "samples_per_segment must be scalar or have one value per segment" - ) - if not np.issubdtype(subdivisions.dtype, np.integer): - if np.any(subdivisions != np.floor(subdivisions)): - raise ValueError("samples_per_segment values must be integers") - subdivisions = subdivisions.astype(np.int64) - else: - subdivisions = subdivisions.astype(np.int64, copy=False) - if np.any(subdivisions < 1): - raise ValueError("samples_per_segment values must be positive") - - padded = np.concatenate( - (points[..., :1, :], points, points[..., -1:, :]), - axis=-2, - ) - interpolated_segments = [] - for index, sample_count in enumerate(subdivisions): - p0 = padded[..., index, :] - p1 = padded[..., index + 1, :] - p2 = padded[..., index + 2, :] - p3 = padded[..., index + 3, :] - t = ( - np.arange(sample_count, dtype=np.float64) / sample_count - ).reshape((1,) * (p0.ndim - 1) + (sample_count, 1)) - t2 = t * t - t3 = t2 * t - segment = 0.5 * ( - 2.0 * p1[..., None, :] - + (-p0 + p2)[..., None, :] * t - + (2.0 * p0 - 5.0 * p1 + 4.0 * p2 - p3)[..., None, :] * t2 - + (-p0 + 3.0 * p1 - 3.0 * p2 + p3)[..., None, :] * t3 + def _selected_records(self, selected_source_ids): + """Return selected-row records with partners remapped locally.""" + selected_source_ids = np.asarray( + selected_source_ids, dtype=np.uint64) + rows = np.searchsorted(self.source_ids, selected_source_ids) + valid = rows < len(self.source_ids) + if not np.all(valid): + raise ValueError( + "track graph does not contain every selected track") + if not np.array_equal( + self.source_ids[rows], selected_source_ids): + raise ValueError( + "track graph does not contain every selected track") + + graph_to_selected = np.full( + len(self.source_ids), -1, dtype=np.int32) + graph_to_selected[rows] = np.arange( + len(rows), dtype=np.int32) + counts = self.offsets[rows + 1] - self.offsets[rows] + record_rows = np.repeat( + np.arange(len(rows), dtype=np.int32), counts) + record_starts = np.repeat(self.offsets[rows], counts) + local_starts = np.repeat( + np.cumsum(np.r_[0, counts[:-1]], dtype=np.int64), counts) + record_indices = record_starts + ( + np.arange(int(counts.sum()), dtype=np.int64) - local_starts) + partners = graph_to_selected[self.partners[record_indices]] + keep = partners >= 0 + return ( + selected_source_ids, + record_rows[keep], + partners[keep], + record_indices[keep], ) - interpolated_segments.append(segment) - interpolated_segments.append(points[..., -1:, :]) - return np.concatenate(interpolated_segments, axis=-2) - -def tracks_to_2d_grid(tracks, row_spacing=20.0, column_spacing=20.0): - """Create a dense ``(rows, columns, 3)`` ZYX grid from ordered tracks. + @staticmethod + def _encode_csr( + source_ids, rows, partners, self_local, partner_local, + positions, clearances): + counts = np.bincount(rows, minlength=len(source_ids)) + offsets = np.empty(len(source_ids) + 1, dtype=np.int64) + offsets[0] = 0 + np.cumsum(counts, out=offsets[1:]) + return { + "source_ids": np.asarray(source_ids, dtype=np.uint64).copy(), + "offsets": offsets, + "partners": np.asarray(partners, dtype=np.int32), + "self_local": np.asarray(self_local, dtype=np.int32), + "partner_local": np.asarray( + partner_local, dtype=np.int32), + "positions": np.asarray(positions, dtype=np.float64), + "clearances": np.asarray(clearances, dtype=np.float64), + } - Rows descend from the highest shared z level by ``row_spacing``. Each row - is a Catmull–Rom connector through the tracks in their supplied order. - Track-to-track segments use a shared subdivision count across all rows, - chosen from their median length to approximate ``column_spacing`` while - preserving every sampled track point as a grid vertex. - """ - column_spacing = float(column_spacing) - if not np.isfinite(column_spacing) or column_spacing <= 0: - raise ValueError("column_spacing must be finite and positive") - - _, control_grid = sample_tracks_at_shared_z_levels( - tracks, spacing=row_spacing - ) - if control_grid.shape[1] < 2: - raise ValueError("at least two tracks are required to create a grid") - - segment_lengths = np.linalg.norm( - np.diff(control_grid, axis=1), - axis=2, - ) - reference_lengths = np.median(segment_lengths, axis=0) - subdivisions = np.maximum( - 1, - np.rint(reference_lengths / column_spacing).astype(np.int64), - ) - return catmull_rom_interpolate(control_grid, subdivisions) - -def write_2d_grid_as_tifxyz( - grid_zyx, - output_directory, - *, - spacing=20.0, - voxel_size_um, - source="ordered crossing-track Catmull-Rom grid", - surface_uuid=None, -): - """Write a ZYX grid beneath ``output_directory`` using its UUID as its ID.""" - grid_zyx = np.asarray(grid_zyx) - if grid_zyx.ndim != 3 or grid_zyx.shape[-1] != 3: - raise ValueError("grid_zyx must have shape (rows, columns, 3)") - if not np.all(np.isfinite(grid_zyx)): - raise ValueError("grid_zyx must contain only finite coordinates") - - if surface_uuid is None: - surface_uuid = datetime.now().strftime("%y%m%d%H%M%S%f") - else: - surface_uuid = str(surface_uuid) - - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - save_tifxyz( - grid_zyx, - output_directory, - surface_uuid, - step_size=float(spacing), - voxel_size_um=float(voxel_size_um), - source=source, - ) - return output_directory / surface_uuid - -def maybe_fetch_tracks(tracks_path= Path("/home/sean/Desktop/spiral_dataset/to_hf/tracks/2um_ds2_ps256_surf_v2.dbm.vctracks")): - - _track_basename = "2um_ds2_ps256_surf_v2.dbm" - _plan = _sync_bucket( - "hf://buckets/scrollprize/datasets/spiral/PHercParis4/tracks", - str(tracks_path.parent), - include=[ - f"{_track_basename}.vctracks/*", - f"{_track_basename}.crossings.npz", - ], - ignore_times=True, - token=False, - ) - print("Hugging Face sync:", _plan.summary()) - -def load_tracks(local_tracks_path): - info = track_store.inspect(str(local_tracks_path)) - - crossing_cache_path = local_tracks_path.with_suffix(".crossings.npz") - with np.load(crossing_cache_path, allow_pickle=False) as crossing_cache: - crossing_indptr = crossing_cache["offsets"].astype(np.int32) - crossing_indices = crossing_cache["partners"] - crossing_positions = crossing_cache["positions"] - - crossing_capacities = np.ones(crossing_indices.size, dtype=np.uint8) - crossing_graph = csr_array( - (crossing_capacities, crossing_indices, crossing_indptr), - shape=(info["track_count"], info["track_count"]), - copy=False, - ) - - print( - f"{crossing_graph.shape[0]:,} tracks, " - f"{crossing_graph.nnz:,} directed adjacency entries" - ) - - coordinates = np.memmap( - local_tracks_path / "coordinates.i32", - mode="r", - dtype="= cutoff - ] - candidate_rows = np.asarray(candidate_rows, dtype=np.int64) - rng.shuffle(candidate_rows) - - # A representative point lets the greedy sampler favor seeds in different - # parts of the volume without loading whole tracks into a second graph. - representative_points = np.empty((len(candidate_rows), 3), dtype=np.float64) - for index, row in enumerate( - tqdm( - candidate_rows, - desc="Indexing seed candidates", - unit="track", + def restricted_csr(self, selected_source_ids): + """Restrict crossings to selected tracks without changing points.""" + source_ids, rows, partners, records = self._selected_records( + selected_source_ids) + return self._encode_csr( + source_ids, rows, partners, + self.self_local[records], + self.partner_local[records], + self.positions[records], + self.clearances[records], ) - ): - points = tracks[int(row)].points_zyx - representative_points[index] = points[len(points) // 2] - - available = np.ones(len(candidate_rows), dtype=bool) - min_distances = np.full(len(candidate_rows), np.inf) - selected = [] - selected_crossings = [] - - with tqdm( - total=count, - desc="Selecting random seeds", - unit="seed", - ) as progress: - attempted = 0 - progress.set_postfix(tried=attempted) - while len(selected) < count and np.any(available): - available_indices = np.flatnonzero(available) - if selected: - distances = min_distances[available_indices] - threshold = np.percentile(distances, 75.0) - diverse_indices = available_indices[distances >= threshold] - candidate_index = int(rng.choice(diverse_indices)) - else: - candidate_index = int(rng.choice(available_indices)) - available[candidate_index] = False - - track = tracks[int(candidate_rows[candidate_index])] - crossing_tracks = track.spaced_crossing_tracks( - min_arclength=crossing_min_arclength, - max_spacing=crossing_max_spacing, - min_shared_z_extent=crossing_min_shared_z_extent, - ) - attempted += 1 - progress.set_postfix(tried=attempted) - if len(crossing_tracks) < 2: - continue - selected.append(track) - selected_crossings.append(crossing_tracks) - if on_selected is not None: - on_selected(track, crossing_tracks) - progress.update() - distances = np.linalg.norm( - representative_points - - representative_points[candidate_index], - axis=1, - ) - min_distances = np.minimum(min_distances, distances) + def clipped_csr( + self, selected_source_ids, input_offsets, surviving_rows, + old_point_to_new, output_offsets): + """Restrict crossings and remap their endpoints after point clipping. - if len(selected) < count: - raise RuntimeError( - f"only {len(selected)} usable seeds were found among " - f"{len(candidate_rows)} horizontals in the top arclength quartile" - ) - return list(zip(selected, selected_crossings)) - -def write_random_seed_grid( - seed_track, - crossing_tracks, - *, - output_directory, - grid_spacing, - voxel_size_um, - surface_uuid, -): - """Build and write one random-seed grid; safe to run in a worker thread.""" - shared_z_extent, crossing_points = crop_tracks_to_shared_z_extent( - crossing_tracks - ) - grid = tracks_to_2d_grid( - crossing_points, - row_spacing=grid_spacing, - column_spacing=grid_spacing, - ) - output_path = write_2d_grid_as_tifxyz( - grid, - output_directory, - spacing=grid_spacing, - voxel_size_um=voxel_size_um, - source=( - "random horizontal seed " - f"track_row={seed_track.row}, source_id={seed_track.source_id}" - ), - surface_uuid=surface_uuid, - ) - return { - "seed_row": seed_track.row, - "source_id": seed_track.source_id, - "crossing_count": len(crossing_tracks), - "shared_z_extent": shared_z_extent, - "grid_shape": grid.shape, - "output_path": output_path, - } - -def create_random_seed_grids( - tracks, - count, - *, - output_directory, - crossing_min_arclength=500.0, - crossing_max_spacing=40.0, - crossing_min_shared_z_extent=500.0, - grid_spacing=20.0, - voxel_size_um=4.0, - workers=None, - rng_seed=None, -): - """Select random horizontal seeds and create their TIFXYZ grids.""" - count = int(count) - if count < 1: - raise ValueError("count must be positive") - rng = np.random.default_rng(rng_seed) - if workers is None: - workers = min(count, os.cpu_count() or 1) - workers = max(1, min(int(workers), count)) - - # Preallocate a batch prefix so concurrent writers can never choose the - # same timestamp. Microseconds protect separate invocations from collision. - batch_id = datetime.now().strftime("%y%m%d%H%M%S%f") - results = [] - with ( - ThreadPoolExecutor(max_workers=workers) as executor, - tqdm( - total=count, - desc="Writing TIFXYZ grids", - unit="grid", - position=1, - ) as write_progress, - ): - futures = [] - - def submit_grid(seed, crossings): - surface_uuid = f"{batch_id}{len(futures):04d}" - future = executor.submit( - write_random_seed_grid, - seed, - crossings, - output_directory=output_directory, - grid_spacing=grid_spacing, - voxel_size_um=voxel_size_um, - surface_uuid=surface_uuid, - ) - future.add_done_callback(lambda _: write_progress.update()) - futures.append(future) - - select_random_horizontal_seeds( - tracks, - count, - rng=rng, - crossing_min_arclength=crossing_min_arclength, - crossing_max_spacing=crossing_max_spacing, - crossing_min_shared_z_extent=crossing_min_shared_z_extent, - on_selected=submit_grid, - ) - for future in as_completed(futures): - result = future.result() - results.append(result) - tqdm.write( - f"Wrote random seed row={result['seed_row']}, " - f"crossings={result['crossing_count']}, " - f"grid_shape={result['grid_shape']}, " - f"tifxyz={result['output_path']}" - ) - return sorted(results, key=lambda result: result["seed_row"]) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create track grids and optionally inspect one in Napari." - ) - parser.add_argument( - "mode", - nargs="?", - choices=("display", "random_seed"), - default="display", - ) - parser.add_argument( - "--num-random-seeds", - type=int, - help="number of TIFXYZ grids to create in random_seed mode", - ) - parser.add_argument( - "--workers", - type=int, - default=None, - help="shared-graph worker threads (default: up to the CPU count)", - ) - parser.add_argument( - "--rng-seed", - type=int, - default=None, - help="optional reproducible random selection seed", - ) - args = parser.parse_args() - if args.mode == "random_seed": - if args.num_random_seeds is None or args.num_random_seeds < 1: - parser.error( - "random_seed requires --num-random-seeds with a positive value" - ) - elif args.num_random_seeds is not None: - parser.error("--num-random-seeds is only valid in random_seed mode") - - local_tracks_path = Path( - "/home/sean/Desktop/spiral_dataset/to_hf/tracks/" - "2um_ds2_ps256_surf_v2.dbm.vctracks" - ) - ct_vol_path = ( - "/home/sean/Documents/volpkgs/s1_ds2.volpkg/volumes/s1_ds2.zarr" - ) - # To use the public CT instead of the local Zarr above, comment out the - # local path and uncomment this assignment: - # ct_vol_path = ( - # "s3://vesuvius-challenge-open-data/PHercParis4/volumes/" - # "20260411134726-2.400um-0.2m-78keV-masked.zarr" - # ) - crossing_min_arclength = 750.0 - crossing_min_shared_z_extent = 750.0 - crossing_max_spacing = 40.0 - grid_spacing = 20.0 - track_voxel_size_um = 4.0 - tifxyz_output_directory = Path( - "/home/sean/Documents/volpkgs/s1_ds2.volpkg/traces" - ) - using_s3_ct = ct_vol_path.startswith("s3://") - - tracks = load_tracks(local_tracks_path) - if args.mode == "random_seed": - results = create_random_seed_grids( - tracks, - args.num_random_seeds, - output_directory=tifxyz_output_directory, - crossing_min_arclength=crossing_min_arclength, - crossing_max_spacing=crossing_max_spacing, - crossing_min_shared_z_extent=crossing_min_shared_z_extent, - grid_spacing=grid_spacing, - voxel_size_um=track_voxel_size_um, - workers=args.workers, - rng_seed=args.rng_seed, - ) - print( - f"Created {len(results)} random-seed TIFXYZ grids in " - f"{tifxyz_output_directory}" - ) - raise SystemExit(0) - - longest_track = tracks[int(np.nanargmax(tracks.arclengths))] - crossing_tracks = longest_track.spaced_crossing_tracks( - min_arclength=crossing_min_arclength, - max_spacing=crossing_max_spacing, - min_shared_z_extent=crossing_min_shared_z_extent, - ) - if len(crossing_tracks) < 2: - raise RuntimeError( - "the longest track does not have at least two eligible crossings " - f"with adjacent spacing <= {crossing_max_spacing:g}" - ) - shared_crossing_z_extent, crossing_points_zyx = ( - crop_tracks_to_shared_z_extent(crossing_tracks) - ) - interpolated_grid_zyx = tracks_to_2d_grid( - crossing_points_zyx, - row_spacing=grid_spacing, - column_spacing=grid_spacing, - ) - tifxyz_output_path = write_2d_grid_as_tifxyz( - interpolated_grid_zyx, - tifxyz_output_directory, - spacing=grid_spacing, - voxel_size_um=track_voxel_size_um, - ) - - if using_s3_ct: - track_scale = 2 - ct_level = 4 - full_points_zyx = scale_track( - longest_track.points_zyx, - source_scale=track_scale, - target_scale=ct_level, - ) - interpolated_points_zyx = scale_track( - interpolated_grid_zyx.reshape(-1, 3), - source_scale=track_scale, - target_scale=ct_level, - ) - ct_store = fsspec.get_mapper(ct_vol_path, anon=True) - ct_group = zarr.open_group(store=ct_store, mode="r") - ct_data = ct_group[str(ct_level)] - - xy_padding = 25 - y_min = max( - 0, - int(np.floor(full_points_zyx[:, 1].min())) - xy_padding, - ) - y_max = min( - ct_data.shape[1], - int(np.ceil(full_points_zyx[:, 1].max())) + xy_padding + 1, - ) - x_min = max( - 0, - int(np.floor(full_points_zyx[:, 2].min())) - xy_padding, + ``old_point_to_new`` maps the selected tracks' original packed point + rows to their rows in the compacted output, with ``-1`` for excluded + points. ``surviving_rows`` maps compacted tracks back to selected rows. + """ + input_offsets = np.asarray(input_offsets, dtype=np.int64) + surviving_rows = np.asarray(surviving_rows, dtype=np.int64) + old_point_to_new = np.asarray(old_point_to_new) + if old_point_to_new.dtype.kind != "i": + raise ValueError("point remap must have an integer dtype") + output_offsets = np.asarray(output_offsets, dtype=np.int64) + selected_source_ids = np.asarray( + selected_source_ids, dtype=np.uint64) + if input_offsets.shape != (len(selected_source_ids) + 1,): + raise ValueError("input offsets are not parallel to selected tracks") + if output_offsets.shape != (len(surviving_rows) + 1,): + raise ValueError("output offsets are not parallel to surviving tracks") + if old_point_to_new.shape != (int(input_offsets[-1]),): + raise ValueError("point remap does not cover the selected tracks") + + _, rows, partners, records = self._selected_records( + selected_source_ids) + selected_to_output = np.full( + len(selected_source_ids), -1, dtype=np.int32) + selected_to_output[surviving_rows] = np.arange( + len(surviving_rows), dtype=np.int32) + output_rows = selected_to_output[rows] + output_partners = selected_to_output[partners] + + old_self_points = ( + input_offsets[rows] + self.self_local[records]) + old_partner_points = ( + input_offsets[partners] + self.partner_local[records]) + valid_bounds = ( + (old_self_points < input_offsets[rows + 1]) + & (old_partner_points < input_offsets[partners + 1]) ) - x_max = min( - ct_data.shape[2], - int(np.ceil(full_points_zyx[:, 2].max())) + xy_padding + 1, + mapped_self = np.full(len(records), -1, dtype=np.int64) + mapped_partner = np.full(len(records), -1, dtype=np.int64) + mapped_self[valid_bounds] = old_point_to_new[ + old_self_points[valid_bounds]] + mapped_partner[valid_bounds] = old_point_to_new[ + old_partner_points[valid_bounds]] + keep = ( + valid_bounds + & (output_rows >= 0) + & (output_partners >= 0) + & (mapped_self >= 0) + & (mapped_partner >= 0) ) - ct_image = da.from_zarr(ct_data)[ - :, - y_min:y_max, - x_min:x_max, - ] - ct_image_name = f"CT level {ct_level}" - ct_image_kwargs = {"translate": (0, y_min, x_min)} - ct_description = ( - f"CT level={ct_level}, " - f"crop=y[{y_min}:{y_max}], x[{x_min}:{x_max}]" + + output_rows = output_rows[keep] + output_partners = output_partners[keep] + records = records[keep] + new_self_local = ( + mapped_self[keep] - output_offsets[output_rows]) + new_partner_local = ( + mapped_partner[keep] - output_offsets[output_partners]) + return self._encode_csr( + selected_source_ids[surviving_rows], + output_rows, + output_partners, + new_self_local, + new_partner_local, + self.positions[records], + self.clearances[records], ) - else: - full_points_zyx = np.asarray(longest_track.points_zyx) - interpolated_points_zyx = interpolated_grid_zyx.reshape(-1, 3) - ct_group = zarr.open_group(store=ct_vol_path, mode="r") - datasets = ct_group.attrs["multiscales"][0]["datasets"] - ct_image = [ - da.from_zarr(ct_group[dataset["path"]]) - for dataset in datasets - ] - ct_image_name = "CT multiscale" - ct_image_kwargs = {"multiscale": True} - ct_description = f"CT multiscale ({len(ct_image)} levels)" - points_zyx = full_points_zyx[::10] - - import napari - - viewer = napari.Viewer(ndisplay=2) - viewer.add_image( - ct_image, - name=ct_image_name, - colormap="gray", - **ct_image_kwargs, - ) - viewer.add_points( - points_zyx, - name=f"longest track {longest_track.row}", - face_color="yellow", - size=4, - out_of_slice_display=True, - ) - viewer.add_points( - interpolated_points_zyx, - name="Catmull-Rom interpolated grid", - face_color="magenta", - size=3, - out_of_slice_display=True, - ) - - z_slice = float(np.median(points_zyx[:, 0])) - viewer.dims.set_point(0, z_slice) - viewer.reset_view() - - print( - f"Displaying track {longest_track.row}: " - f"source_id={longest_track.source_id}, " - f"arclength={longest_track.arclength:.3f}, " - f"tortuosity={longest_track.tortuosity:.3f}, " - f"selected_crossings={len(crossing_tracks)}, " - f"shared_crossing_z_extent={shared_crossing_z_extent}, " - f"grid_shape={interpolated_grid_zyx.shape}, " - f"tifxyz={tifxyz_output_path}, " - f"{ct_description}, z={z_slice:.2f}" - ) - napari.run() diff --git a/volume-cartographer/scripts/spiral/track_graph_testing.py b/volume-cartographer/scripts/spiral/track_graph_testing.py new file mode 100644 index 0000000000..67b961fa1a --- /dev/null +++ b/volume-cartographer/scripts/spiral/track_graph_testing.py @@ -0,0 +1,990 @@ + +import argparse +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from dataclasses import dataclass, field +from datetime import datetime +import os +import numpy as np +from scipy.sparse import csr_array +from scipy.sparse.csgraph import breadth_first_order, connected_components, maximum_flow +import vc.track_store as track_store +import vc.compression.vcz1_numcodecs # Registers the VCZ1 codec with numcodecs. +import zarr +import fsspec +import dask.array as da +from huggingface_hub import sync_bucket as _sync_bucket +from tifxyz import save_tifxyz +from tqdm.auto import tqdm + +family_names = { + -1: "unknown", + 0: "hz", + 1: "vt", +} + +@dataclass(slots=True, repr=False) +class TrackCollection: + coordinates: np.ndarray + offsets: np.ndarray + source_ids: np.ndarray + family_codes: np.ndarray + z_bounds: np.ndarray + arclengths: np.ndarray + tortuosities: np.ndarray + crossing_graph: csr_array + crossing_positions: np.ndarray + + + def __len__(self): + return len(self.source_ids) + + def __getitem__(self, row): + return Track(self, int(row)) + + def __repr__(self): + return f"TrackCollection({len(self):,} tracks)" + + +@dataclass(frozen=True, slots=True) +class Track: + """Lightweight handle that materializes track data only when accessed.""" + + collection: TrackCollection = field(repr=False, compare=False) + row: int + + @property + def source_id(self): + return int(self.collection.source_ids[self.row]) + + @property + def family(self): + return int(self.collection.family_codes[self.row]) + + @property + def family_name(self): + return family_names[self.family] + + @property + def arclength(self): + return float(self.collection.arclengths[self.row]) + + @property + def tortuosity(self): + return float(self.collection.tortuosities[self.row]) + + @property + def z_bounds(self): + return tuple(map(int, self.collection.z_bounds[self.row])) + + @property + def points_zyx(self): + begin = int(self.collection.offsets[self.row]) + end = int(self.collection.offsets[self.row + 1]) + return self.collection.coordinates[begin:end] + + @property + def neighbors(self): + indptr = self.collection.crossing_graph.indptr + begin = int(indptr[self.row]) + end = int(indptr[self.row + 1]) + return self.collection.crossing_graph.indices[begin:end] + + @property + def crossing_positions(self): + indptr = self.collection.crossing_graph.indptr + begin = int(indptr[self.row]) + end = int(indptr[self.row + 1]) + return self.collection.crossing_positions[begin:end] + + def spaced_crossing_tracks( + self, + min_arclength=0.0, + min_spacing=0.0, + max_spacing=None, + min_shared_z_extent=0.0, + ): + """Return crossing tracks ordered from this track's first endpoint. + + ``min_arclength`` filters the crossing tracks, while ``min_spacing`` + is the minimum arclength along this track between selected crossings. + When ``max_spacing`` is supplied, all eligible crossings are retained + and an empty result is returned if any adjacent pair is farther apart + than that value. ``min_spacing`` and ``max_spacing`` are mutually + exclusive. + ``min_shared_z_extent`` requires every returned track to cover a + common z interval of at least that size. + In minimum-spacing mode, the result contains as many tracks as possible + and, among that number, chooses crossings that are as close as possible + to uniform spacing. + Results are ordered by increasing arclength along this track, starting + at the endpoint represented by ``points_zyx[0]``. + All spacing and extent arguments use the coordinate units of + ``arclength``. + """ + min_arclength = float(min_arclength) + min_spacing = float(min_spacing) + if max_spacing is not None: + max_spacing = float(max_spacing) + min_shared_z_extent = float(min_shared_z_extent) + if not np.isfinite(min_arclength) or min_arclength < 0: + raise ValueError("min_arclength must be finite and non-negative") + if not np.isfinite(min_spacing) or min_spacing < 0: + raise ValueError("min_spacing must be finite and non-negative") + if max_spacing is not None: + if not np.isfinite(max_spacing) or max_spacing <= 0: + raise ValueError("max_spacing must be finite and positive") + if min_spacing != 0: + raise ValueError( + "min_spacing and max_spacing are mutually exclusive" + ) + if ( + not np.isfinite(min_shared_z_extent) + or min_shared_z_extent < 0 + ): + raise ValueError( + "min_shared_z_extent must be finite and non-negative" + ) + + neighbor_rows = np.asarray(self.neighbors, dtype=np.int64) + positions = np.asarray(self.crossing_positions, dtype=np.float64) + keep = ( + np.isfinite(positions) + & (self.collection.arclengths[neighbor_rows] >= min_arclength) + ) + neighbor_rows = neighbor_rows[keep] + positions = positions[keep] + if not neighbor_rows.size: + return [] + + if min_shared_z_extent > 0: + z_bounds = np.asarray( + self.collection.z_bounds[neighbor_rows], + dtype=np.float64, + ) + z_min = np.min(z_bounds, axis=1) + z_max = np.max(z_bounds, axis=1) + + # A track can contain a shared interval [start, start + extent] + # exactly when start lies in [z_min, z_max - extent]. Find the + # start covered by the greatest number of eligible tracks. + interval_starts = z_min + interval_ends = z_max - min_shared_z_extent + feasible = interval_starts <= interval_ends + neighbor_rows = neighbor_rows[feasible] + positions = positions[feasible] + interval_starts = interval_starts[feasible] + interval_ends = interval_ends[feasible] + if not neighbor_rows.size: + return [] + + candidate_starts = np.unique(interval_starts) + sorted_starts = np.sort(interval_starts) + sorted_ends = np.sort(interval_ends) + coverage_counts = ( + np.searchsorted(sorted_starts, candidate_starts, side="right") + - np.searchsorted(sorted_ends, candidate_starts, side="left") + ) + shared_start = candidate_starts[np.argmax(coverage_counts)] + shares_interval = ( + (interval_starts <= shared_start) + & (interval_ends >= shared_start) + ) + neighbor_rows = neighbor_rows[shares_interval] + positions = positions[shares_interval] + + position_order = np.argsort(positions, kind="stable") + neighbor_rows = neighbor_rows[position_order] + positions = positions[position_order] + + if max_spacing is not None: + if ( + neighbor_rows.size < 2 + or np.any(np.diff(positions) > max_spacing) + ): + return [] + return [self.collection[int(row)] for row in neighbor_rows] + + if min_spacing == 0 or neighbor_rows.size == 1: + return [self.collection[int(row)] for row in neighbor_rows] + + # Greedy interval scheduling gives the largest feasible selection size. + maximum_count = 0 + last_position = -np.inf + for position in positions: + if position - last_position >= min_spacing: + maximum_count += 1 + last_position = position + + if maximum_count == 1: + middle = 0.5 * (positions[0] + positions[-1]) + chosen = np.array([np.argmin(np.abs(positions - middle))]) + elif maximum_count == neighbor_rows.size: + chosen = np.arange(neighbor_rows.size) + else: + # Dynamic programming chooses exactly maximum_count records while + # minimizing squared distance from uniformly distributed targets. + targets = np.linspace( + positions[0], positions[-1], maximum_count, + dtype=np.float64, + ) + count = neighbor_rows.size + previous_cost = (positions - targets[0]) ** 2 + back_pointers = [] + for target in targets[1:]: + current_cost = np.full(count, np.inf) + current_back = np.full(count, -1, dtype=np.int32) + + # For each crossing, find the final previous crossing that is + # far enough away, then look up the cheapest DP state in that + # prefix. + final_predecessors = np.searchsorted( + positions, + positions - min_spacing, + side="right", + ) - 1 + prefix_costs = np.minimum.accumulate(previous_cost) + new_best = np.empty(count, dtype=bool) + new_best[0] = np.isfinite(previous_cost[0]) + new_best[1:] = previous_cost[1:] < prefix_costs[:-1] + prefix_indices = np.maximum.accumulate( + np.where(new_best, np.arange(count), -1) + ) + + valid = final_predecessors >= 0 + valid_indices = np.flatnonzero(valid) + predecessor_limits = final_predecessors[valid] + finite = np.isfinite(prefix_costs[predecessor_limits]) + valid_indices = valid_indices[finite] + predecessor_limits = predecessor_limits[finite] + current_cost[valid_indices] = ( + prefix_costs[predecessor_limits] + + (positions[valid_indices] - target) ** 2 + ) + current_back[valid_indices] = prefix_indices[ + predecessor_limits + ] + previous_cost = current_cost + back_pointers.append(current_back) + + chosen = np.empty(maximum_count, dtype=np.int64) + chosen[-1] = int(np.argmin(previous_cost)) + for level in range(maximum_count - 2, -1, -1): + chosen[level] = back_pointers[level][chosen[level + 1]] + + return [ + self.collection[int(row)] + for row in neighbor_rows[chosen] + ] + + def connectivity_summary(self): + neighbors = self.neighbors + unique_neighbors = np.unique(neighbors) + families, counts = np.unique( + self.collection.family_codes[unique_neighbors], + return_counts=True, + ) + return { + "crossings": int(neighbors.size), + "unique_neighbors": int(unique_neighbors.size), + "neighbors_by_family": { + family_names[int(family)]: int(count) + for family, count in zip(families, counts) + }, + } + +def scale_track(points_zyx, source_scale, target_scale): + """Map (z, y, x) coordinates between Zarr pyramid scales.""" + scale_ratio = 2.0 ** (target_scale - source_scale) + return np.asarray(points_zyx, dtype=np.float64) / scale_ratio + +def crop_tracks_to_shared_z_extent(tracks): + """Return track point arrays cropped to their common inclusive z extent.""" + point_sets = [np.asarray(track.points_zyx) for track in tracks] + if not point_sets: + return None, [] + + z_min = max(int(points[:, 0].min()) for points in point_sets) + z_max = min(int(points[:, 0].max()) for points in point_sets) + if z_min > z_max: + raise ValueError("crossing tracks do not have a shared z extent") + + cropped_point_sets = [ + points[(points[:, 0] >= z_min) & (points[:, 0] <= z_max)] + for points in point_sets + ] + return (z_min, z_max), cropped_point_sets + +def sample_tracks_at_shared_z_levels(tracks, spacing=20.0): + """Sample every track at common z levels, ordered highest z first. + + Returns ``(z_levels, points_zyx)`` where ``points_zyx`` has shape + ``(number_of_z_levels, number_of_tracks, 3)``. Tracks are treated as + single-valued functions of z. Multiple points at one z are averaged before + linear interpolation. + """ + spacing = float(spacing) + if not np.isfinite(spacing) or spacing <= 0: + raise ValueError("spacing must be finite and positive") + + point_sets = [ + np.asarray( + track.points_zyx if hasattr(track, "points_zyx") else track, + dtype=np.float64, + ) + for track in tracks + ] + if not point_sets: + return np.empty(0, dtype=np.float64), np.empty((0, 0, 3)) + if any(points.ndim != 2 or points.shape[1] != 3 or not len(points) + for points in point_sets): + raise ValueError("every track must be a non-empty (n, 3) ZYX array") + + z_min = max(points[:, 0].min() for points in point_sets) + z_max = min(points[:, 0].max() for points in point_sets) + if z_min > z_max: + raise ValueError("tracks do not have a shared z extent") + + level_count = int(np.floor((z_max - z_min) / spacing)) + 1 + z_levels = z_max - spacing * np.arange(level_count, dtype=np.float64) + sampled = np.empty((level_count, len(point_sets), 3), dtype=np.float64) + sampled[..., 0] = z_levels[:, None] + + for track_index, points in enumerate(point_sets): + z_values, inverse = np.unique(points[:, 0], return_inverse=True) + yx_sums = np.zeros((len(z_values), 2), dtype=np.float64) + np.add.at(yx_sums, inverse, points[:, 1:]) + counts = np.bincount(inverse) + yx_values = yx_sums / counts[:, None] + sampled[:, track_index, 1] = np.interp( + z_levels, z_values, yx_values[:, 0] + ) + sampled[:, track_index, 2] = np.interp( + z_levels, z_values, yx_values[:, 1] + ) + + return z_levels, sampled + +def catmull_rom_interpolate(points, samples_per_segment): + """Interpolate ordered control points with a vectorized Catmull–Rom spline. + + ``points`` has shape ``(..., control_points, dimensions)``. The leading + dimensions are evaluated in parallel. ``samples_per_segment`` may be one + integer or one integer per adjacent pair of control points. Every control + point occurs exactly in the result. + """ + points = np.asarray(points, dtype=np.float64) + if points.ndim < 2 or points.shape[-2] < 2: + raise ValueError("points must contain at least two control points") + + segment_count = points.shape[-2] - 1 + subdivisions = np.asarray(samples_per_segment) + if subdivisions.ndim == 0: + subdivisions = np.full(segment_count, subdivisions) + if subdivisions.shape != (segment_count,): + raise ValueError( + "samples_per_segment must be scalar or have one value per segment" + ) + if not np.issubdtype(subdivisions.dtype, np.integer): + if np.any(subdivisions != np.floor(subdivisions)): + raise ValueError("samples_per_segment values must be integers") + subdivisions = subdivisions.astype(np.int64) + else: + subdivisions = subdivisions.astype(np.int64, copy=False) + if np.any(subdivisions < 1): + raise ValueError("samples_per_segment values must be positive") + + padded = np.concatenate( + (points[..., :1, :], points, points[..., -1:, :]), + axis=-2, + ) + interpolated_segments = [] + for index, sample_count in enumerate(subdivisions): + p0 = padded[..., index, :] + p1 = padded[..., index + 1, :] + p2 = padded[..., index + 2, :] + p3 = padded[..., index + 3, :] + t = ( + np.arange(sample_count, dtype=np.float64) / sample_count + ).reshape((1,) * (p0.ndim - 1) + (sample_count, 1)) + t2 = t * t + t3 = t2 * t + segment = 0.5 * ( + 2.0 * p1[..., None, :] + + (-p0 + p2)[..., None, :] * t + + (2.0 * p0 - 5.0 * p1 + 4.0 * p2 - p3)[..., None, :] * t2 + + (-p0 + 3.0 * p1 - 3.0 * p2 + p3)[..., None, :] * t3 + ) + interpolated_segments.append(segment) + + interpolated_segments.append(points[..., -1:, :]) + return np.concatenate(interpolated_segments, axis=-2) + +def tracks_to_2d_grid(tracks, row_spacing=20.0, column_spacing=20.0): + """Create a dense ``(rows, columns, 3)`` ZYX grid from ordered tracks. + + Rows descend from the highest shared z level by ``row_spacing``. Each row + is a Catmull–Rom connector through the tracks in their supplied order. + Track-to-track segments use a shared subdivision count across all rows, + chosen from their median length to approximate ``column_spacing`` while + preserving every sampled track point as a grid vertex. + """ + column_spacing = float(column_spacing) + if not np.isfinite(column_spacing) or column_spacing <= 0: + raise ValueError("column_spacing must be finite and positive") + + _, control_grid = sample_tracks_at_shared_z_levels( + tracks, spacing=row_spacing + ) + if control_grid.shape[1] < 2: + raise ValueError("at least two tracks are required to create a grid") + + segment_lengths = np.linalg.norm( + np.diff(control_grid, axis=1), + axis=2, + ) + reference_lengths = np.median(segment_lengths, axis=0) + subdivisions = np.maximum( + 1, + np.rint(reference_lengths / column_spacing).astype(np.int64), + ) + return catmull_rom_interpolate(control_grid, subdivisions) + +def write_2d_grid_as_tifxyz( + grid_zyx, + output_directory, + *, + spacing=20.0, + voxel_size_um, + source="ordered crossing-track Catmull-Rom grid", + surface_uuid=None, +): + """Write a ZYX grid beneath ``output_directory`` using its UUID as its ID.""" + grid_zyx = np.asarray(grid_zyx) + if grid_zyx.ndim != 3 or grid_zyx.shape[-1] != 3: + raise ValueError("grid_zyx must have shape (rows, columns, 3)") + if not np.all(np.isfinite(grid_zyx)): + raise ValueError("grid_zyx must contain only finite coordinates") + + if surface_uuid is None: + surface_uuid = datetime.now().strftime("%y%m%d%H%M%S%f") + else: + surface_uuid = str(surface_uuid) + + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + save_tifxyz( + grid_zyx, + output_directory, + surface_uuid, + step_size=float(spacing), + voxel_size_um=float(voxel_size_um), + source=source, + ) + return output_directory / surface_uuid + +def maybe_fetch_tracks(tracks_path= Path("/home/sean/Desktop/spiral_dataset/to_hf/tracks/2um_ds2_ps256_surf_v2.dbm.vctracks")): + + _track_basename = "2um_ds2_ps256_surf_v2.dbm" + _plan = _sync_bucket( + "hf://buckets/scrollprize/datasets/spiral/PHercParis4/tracks", + str(tracks_path.parent), + include=[ + f"{_track_basename}.vctracks/*", + f"{_track_basename}.crossings.npz", + ], + ignore_times=True, + token=False, + ) + print("Hugging Face sync:", _plan.summary()) + +def load_tracks(local_tracks_path): + info = track_store.inspect(str(local_tracks_path)) + + crossing_cache_path = local_tracks_path.with_suffix(".crossings.npz") + with np.load(crossing_cache_path, allow_pickle=False) as crossing_cache: + crossing_indptr = crossing_cache["offsets"].astype(np.int32) + crossing_indices = crossing_cache["partners"] + crossing_positions = crossing_cache["positions"] + + crossing_capacities = np.ones(crossing_indices.size, dtype=np.uint8) + crossing_graph = csr_array( + (crossing_capacities, crossing_indices, crossing_indptr), + shape=(info["track_count"], info["track_count"]), + copy=False, + ) + + print( + f"{crossing_graph.shape[0]:,} tracks, " + f"{crossing_graph.nnz:,} directed adjacency entries" + ) + + coordinates = np.memmap( + local_tracks_path / "coordinates.i32", + mode="r", + dtype="= cutoff + ] + candidate_rows = np.asarray(candidate_rows, dtype=np.int64) + rng.shuffle(candidate_rows) + + # A representative point lets the greedy sampler favor seeds in different + # parts of the volume without loading whole tracks into a second graph. + representative_points = np.empty((len(candidate_rows), 3), dtype=np.float64) + for index, row in enumerate( + tqdm( + candidate_rows, + desc="Indexing seed candidates", + unit="track", + ) + ): + points = tracks[int(row)].points_zyx + representative_points[index] = points[len(points) // 2] + + available = np.ones(len(candidate_rows), dtype=bool) + min_distances = np.full(len(candidate_rows), np.inf) + selected = [] + selected_crossings = [] + + with tqdm( + total=count, + desc="Selecting random seeds", + unit="seed", + ) as progress: + attempted = 0 + progress.set_postfix(tried=attempted) + while len(selected) < count and np.any(available): + available_indices = np.flatnonzero(available) + if selected: + distances = min_distances[available_indices] + threshold = np.percentile(distances, 75.0) + diverse_indices = available_indices[distances >= threshold] + candidate_index = int(rng.choice(diverse_indices)) + else: + candidate_index = int(rng.choice(available_indices)) + available[candidate_index] = False + + track = tracks[int(candidate_rows[candidate_index])] + crossing_tracks = track.spaced_crossing_tracks( + min_arclength=crossing_min_arclength, + max_spacing=crossing_max_spacing, + min_shared_z_extent=crossing_min_shared_z_extent, + ) + attempted += 1 + progress.set_postfix(tried=attempted) + if len(crossing_tracks) < 2: + continue + + selected.append(track) + selected_crossings.append(crossing_tracks) + if on_selected is not None: + on_selected(track, crossing_tracks) + progress.update() + distances = np.linalg.norm( + representative_points + - representative_points[candidate_index], + axis=1, + ) + min_distances = np.minimum(min_distances, distances) + + if len(selected) < count: + raise RuntimeError( + f"only {len(selected)} usable seeds were found among " + f"{len(candidate_rows)} horizontals in the top arclength quartile" + ) + return list(zip(selected, selected_crossings)) + +def write_random_seed_grid( + seed_track, + crossing_tracks, + *, + output_directory, + grid_spacing, + voxel_size_um, + surface_uuid, +): + """Build and write one random-seed grid; safe to run in a worker thread.""" + shared_z_extent, crossing_points = crop_tracks_to_shared_z_extent( + crossing_tracks + ) + grid = tracks_to_2d_grid( + crossing_points, + row_spacing=grid_spacing, + column_spacing=grid_spacing, + ) + output_path = write_2d_grid_as_tifxyz( + grid, + output_directory, + spacing=grid_spacing, + voxel_size_um=voxel_size_um, + source=( + "random horizontal seed " + f"track_row={seed_track.row}, source_id={seed_track.source_id}" + ), + surface_uuid=surface_uuid, + ) + return { + "seed_row": seed_track.row, + "source_id": seed_track.source_id, + "crossing_count": len(crossing_tracks), + "shared_z_extent": shared_z_extent, + "grid_shape": grid.shape, + "output_path": output_path, + } + +def create_random_seed_grids( + tracks, + count, + *, + output_directory, + crossing_min_arclength=500.0, + crossing_max_spacing=40.0, + crossing_min_shared_z_extent=500.0, + grid_spacing=20.0, + voxel_size_um=4.0, + workers=None, + rng_seed=None, +): + """Select random horizontal seeds and create their TIFXYZ grids.""" + count = int(count) + if count < 1: + raise ValueError("count must be positive") + rng = np.random.default_rng(rng_seed) + if workers is None: + workers = min(count, os.cpu_count() or 1) + workers = max(1, min(int(workers), count)) + + # Preallocate a batch prefix so concurrent writers can never choose the + # same timestamp. Microseconds protect separate invocations from collision. + batch_id = datetime.now().strftime("%y%m%d%H%M%S%f") + results = [] + with ( + ThreadPoolExecutor(max_workers=workers) as executor, + tqdm( + total=count, + desc="Writing TIFXYZ grids", + unit="grid", + position=1, + ) as write_progress, + ): + futures = [] + + def submit_grid(seed, crossings): + surface_uuid = f"{batch_id}{len(futures):04d}" + future = executor.submit( + write_random_seed_grid, + seed, + crossings, + output_directory=output_directory, + grid_spacing=grid_spacing, + voxel_size_um=voxel_size_um, + surface_uuid=surface_uuid, + ) + future.add_done_callback(lambda _: write_progress.update()) + futures.append(future) + + select_random_horizontal_seeds( + tracks, + count, + rng=rng, + crossing_min_arclength=crossing_min_arclength, + crossing_max_spacing=crossing_max_spacing, + crossing_min_shared_z_extent=crossing_min_shared_z_extent, + on_selected=submit_grid, + ) + for future in as_completed(futures): + result = future.result() + results.append(result) + tqdm.write( + f"Wrote random seed row={result['seed_row']}, " + f"crossings={result['crossing_count']}, " + f"grid_shape={result['grid_shape']}, " + f"tifxyz={result['output_path']}" + ) + return sorted(results, key=lambda result: result["seed_row"]) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create track grids and optionally inspect one in Napari." + ) + parser.add_argument( + "mode", + nargs="?", + choices=("display", "random_seed"), + default="display", + ) + parser.add_argument( + "--num-random-seeds", + type=int, + help="number of TIFXYZ grids to create in random_seed mode", + ) + parser.add_argument( + "--workers", + type=int, + default=None, + help="shared-graph worker threads (default: up to the CPU count)", + ) + parser.add_argument( + "--rng-seed", + type=int, + default=None, + help="optional reproducible random selection seed", + ) + args = parser.parse_args() + if args.mode == "random_seed": + if args.num_random_seeds is None or args.num_random_seeds < 1: + parser.error( + "random_seed requires --num-random-seeds with a positive value" + ) + elif args.num_random_seeds is not None: + parser.error("--num-random-seeds is only valid in random_seed mode") + + local_tracks_path = Path( + "/home/sean/Desktop/spiral_dataset/to_hf/tracks/" + "2um_ds2_ps256_surf_v2.dbm.vctracks" + ) + ct_vol_path = ( + "/home/sean/Documents/volpkgs/s1_ds2.volpkg/volumes/s1_ds2.zarr" + ) + # To use the public CT instead of the local Zarr above, comment out the + # local path and uncomment this assignment: + # ct_vol_path = ( + # "s3://vesuvius-challenge-open-data/PHercParis4/volumes/" + # "20260411134726-2.400um-0.2m-78keV-masked.zarr" + # ) + crossing_min_arclength = 750.0 + crossing_min_shared_z_extent = 750.0 + crossing_max_spacing = 40.0 + grid_spacing = 20.0 + track_voxel_size_um = 4.0 + tifxyz_output_directory = Path( + "/home/sean/Documents/volpkgs/s1_ds2.volpkg/traces" + ) + using_s3_ct = ct_vol_path.startswith("s3://") + + tracks = load_tracks(local_tracks_path) + if args.mode == "random_seed": + results = create_random_seed_grids( + tracks, + args.num_random_seeds, + output_directory=tifxyz_output_directory, + crossing_min_arclength=crossing_min_arclength, + crossing_max_spacing=crossing_max_spacing, + crossing_min_shared_z_extent=crossing_min_shared_z_extent, + grid_spacing=grid_spacing, + voxel_size_um=track_voxel_size_um, + workers=args.workers, + rng_seed=args.rng_seed, + ) + print( + f"Created {len(results)} random-seed TIFXYZ grids in " + f"{tifxyz_output_directory}" + ) + raise SystemExit(0) + + longest_track = tracks[int(np.nanargmax(tracks.arclengths))] + crossing_tracks = longest_track.spaced_crossing_tracks( + min_arclength=crossing_min_arclength, + max_spacing=crossing_max_spacing, + min_shared_z_extent=crossing_min_shared_z_extent, + ) + if len(crossing_tracks) < 2: + raise RuntimeError( + "the longest track does not have at least two eligible crossings " + f"with adjacent spacing <= {crossing_max_spacing:g}" + ) + shared_crossing_z_extent, crossing_points_zyx = ( + crop_tracks_to_shared_z_extent(crossing_tracks) + ) + interpolated_grid_zyx = tracks_to_2d_grid( + crossing_points_zyx, + row_spacing=grid_spacing, + column_spacing=grid_spacing, + ) + tifxyz_output_path = write_2d_grid_as_tifxyz( + interpolated_grid_zyx, + tifxyz_output_directory, + spacing=grid_spacing, + voxel_size_um=track_voxel_size_um, + ) + + if using_s3_ct: + track_scale = 2 + ct_level = 4 + full_points_zyx = scale_track( + longest_track.points_zyx, + source_scale=track_scale, + target_scale=ct_level, + ) + interpolated_points_zyx = scale_track( + interpolated_grid_zyx.reshape(-1, 3), + source_scale=track_scale, + target_scale=ct_level, + ) + ct_store = fsspec.get_mapper(ct_vol_path, anon=True) + ct_group = zarr.open_group(store=ct_store, mode="r") + ct_data = ct_group[str(ct_level)] + + xy_padding = 25 + y_min = max( + 0, + int(np.floor(full_points_zyx[:, 1].min())) - xy_padding, + ) + y_max = min( + ct_data.shape[1], + int(np.ceil(full_points_zyx[:, 1].max())) + xy_padding + 1, + ) + x_min = max( + 0, + int(np.floor(full_points_zyx[:, 2].min())) - xy_padding, + ) + x_max = min( + ct_data.shape[2], + int(np.ceil(full_points_zyx[:, 2].max())) + xy_padding + 1, + ) + ct_image = da.from_zarr(ct_data)[ + :, + y_min:y_max, + x_min:x_max, + ] + ct_image_name = f"CT level {ct_level}" + ct_image_kwargs = {"translate": (0, y_min, x_min)} + ct_description = ( + f"CT level={ct_level}, " + f"crop=y[{y_min}:{y_max}], x[{x_min}:{x_max}]" + ) + else: + full_points_zyx = np.asarray(longest_track.points_zyx) + interpolated_points_zyx = interpolated_grid_zyx.reshape(-1, 3) + ct_group = zarr.open_group(store=ct_vol_path, mode="r") + datasets = ct_group.attrs["multiscales"][0]["datasets"] + ct_image = [ + da.from_zarr(ct_group[dataset["path"]]) + for dataset in datasets + ] + ct_image_name = "CT multiscale" + ct_image_kwargs = {"multiscale": True} + ct_description = f"CT multiscale ({len(ct_image)} levels)" + points_zyx = full_points_zyx[::10] + + import napari + + viewer = napari.Viewer(ndisplay=2) + viewer.add_image( + ct_image, + name=ct_image_name, + colormap="gray", + **ct_image_kwargs, + ) + viewer.add_points( + points_zyx, + name=f"longest track {longest_track.row}", + face_color="yellow", + size=4, + out_of_slice_display=True, + ) + viewer.add_points( + interpolated_points_zyx, + name="Catmull-Rom interpolated grid", + face_color="magenta", + size=3, + out_of_slice_display=True, + ) + + z_slice = float(np.median(points_zyx[:, 0])) + viewer.dims.set_point(0, z_slice) + viewer.reset_view() + + print( + f"Displaying track {longest_track.row}: " + f"source_id={longest_track.source_id}, " + f"arclength={longest_track.arclength:.3f}, " + f"tortuosity={longest_track.tortuosity:.3f}, " + f"selected_crossings={len(crossing_tracks)}, " + f"shared_crossing_z_extent={shared_crossing_z_extent}, " + f"grid_shape={interpolated_grid_zyx.shape}, " + f"tifxyz={tifxyz_output_path}, " + f"{ct_description}, z={z_slice:.2f}" + ) + napari.run() diff --git a/volume-cartographer/scripts/spiral/tracks.py b/volume-cartographer/scripts/spiral/tracks.py index 0fb4ab5e96..6708f63171 100644 --- a/volume-cartographer/scripts/spiral/tracks.py +++ b/volume-cartographer/scripts/spiral/tracks.py @@ -446,7 +446,8 @@ def _dbm_source_id(key_ordinal, source_index): def load_tracks_from_dbm( path, z_lo=None, z_hi=None, return_families=False, - return_source_ids=False, show_progress=True, low_memory=False): + return_source_ids=False, show_progress=True, low_memory=False, + progress=None): # Load tracks written by extract_surface_tracks.py. Each DBM value is a # pickled list of (N, 3) int32 zyx arrays; keep only tracks that lie entirely # within the full-resolution [z_lo, z_hi) ROI. @@ -467,8 +468,14 @@ def load_tracks_from_dbm( source_ids = [] with dbm.open(path, 'r') as db: keys = db.keys() + if progress is not None: + progress.begin( + 'loading', 'Loading tracks', + step=0, total_steps=len(keys), unit='DB keys') key_ordinals = {key: index for index, key in enumerate(sorted(keys))} - for key in tqdm(keys, desc='loading tracks', disable=not show_progress): + for key_number, key in enumerate(tqdm( + keys, desc='loading tracks', + disable=not show_progress or progress is not None), start=1): family = None if return_families: prefix = key.decode().split(':', 1)[0] @@ -476,6 +483,9 @@ def load_tracks_from_dbm( 'vertical' if prefix in ('vx', 'vy') else None) entries = pickle.loads(db[key]) if not entries: + if progress is not None: + progress.update( + key_number, detail=f'{len(tracks):,} tracks retained') continue if low_memory: for source_index, entry in enumerate(entries): @@ -494,12 +504,18 @@ def load_tracks_from_dbm( if return_source_ids: source_ids.append(_dbm_source_id( key_ordinals[key], source_index)) + if progress is not None: + progress.update( + key_number, detail=f'{len(tracks):,} tracks retained') continue # Vectorize the per-track z min/max across the whole key: concatenate # every non-empty track's z column and reduce per segment, rather # than calling .min()/.max() once per track. idx = [i for i in range(len(entries)) if len(entries[i])] if not idx: + if progress is not None: + progress.update( + key_number, detail=f'{len(tracks):,} tracks retained') continue lengths = np.fromiter((len(entries[i]) for i in idx), dtype=np.intp, count=len(idx)) zcat = np.concatenate([entries[i][:, 0] for i in idx]) @@ -520,6 +536,9 @@ def load_tracks_from_dbm( if return_source_ids: source_ids.append(_dbm_source_id( key_ordinals[key], source_index)) + if progress is not None: + progress.update( + key_number, detail=f'{len(tracks):,} tracks retained') if return_families and return_source_ids: return tracks, families, np.asarray(source_ids, dtype=np.uint64) if return_families: @@ -578,7 +597,7 @@ def validate_track_sampling_config(config): raise ValueError('track_max_tortuosity must be null or a finite number >= 1') max_tortuosity = float(max_tortuosity) - max_crossings = config.get('max_track_crossing_per_step', 0) + max_crossings = config.get('track_max_track_crossing_per_step', 0) if (isinstance(max_crossings, bool) or not isinstance(max_crossings, (int, float)) or not math.isfinite(float(max_crossings)) or not float(max_crossings).is_integer() or int(max_crossings) < 0): @@ -597,30 +616,38 @@ def validate_track_sampling_config(config): crossing_precompute_max = max( int(crossing_precompute_max), int(max_crossings)) - crossing_mode = config.get('track_crossing_mode', 'count') + crossing_mode = config.get('track_crossing_mode', 'track_walk') if crossing_mode not in ('count', 'track_walk'): raise ValueError( "track_crossing_mode must be 'count' or 'track_walk'") walk_values = {} for key, default in ( - ('min_walk_steps_per_track', 24), - ('max_walk_steps_per_track', 256), - ('n_walks_per_track', 4)): + ('track_min_walk_steps_per_track', 24), + ('track_max_walk_steps_per_track', 256), + ('track_min_walks_per_track', 2), + ('track_max_walks_per_track', 4)): value = config.get(key, default) if (isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)) or not float(value).is_integer() or int(value) <= 0): raise ValueError(f'{key} must be a positive integer') walk_values[key] = int(value) - if (walk_values['min_walk_steps_per_track'] - > walk_values['max_walk_steps_per_track']): + if (walk_values['track_min_walk_steps_per_track'] + > walk_values['track_max_walk_steps_per_track']): raise ValueError( 'min_walk_steps_per_track must be <= max_walk_steps_per_track') - require_loop = config.get( - 'track_walk_require_loop_consistency', False) - if not isinstance(require_loop, bool): + if (walk_values['track_min_walks_per_track'] + > walk_values['track_max_walks_per_track']): raise ValueError( - 'track_walk_require_loop_consistency must be boolean') + 'min_walks_per_track must be <= max_walks_per_track') + minimum_cycle_travel = config.get( + 'track_walk_minimum_cycle_travel', 20.0) + if (isinstance(minimum_cycle_travel, bool) + or not isinstance(minimum_cycle_travel, (int, float)) + or not math.isfinite(float(minimum_cycle_travel)) + or float(minimum_cycle_travel) < 0): + raise ValueError( + 'track_walk_minimum_cycle_travel must be a finite number >= 0') sample_spacings = {} for key, default in ( @@ -645,7 +672,7 @@ def validate_track_sampling_config(config): 'crossing_precompute_max': crossing_precompute_max, 'crossing_mode': crossing_mode, **walk_values, - 'walk_require_loop_consistency': require_loop, + 'walk_minimum_cycle_travel': float(minimum_cycle_travel), 'min_sample_spacing': sample_spacings['track_min_sample_spacing'], 'max_sample_spacing': sample_spacings['track_max_sample_spacing'], } @@ -771,18 +798,21 @@ def configure_prepared_track_sampling(prepared_tracks, config): return current = { 'track_length_bin_weights': prepared_tracks.get('length_bin_weights'), - 'max_track_crossing_per_step': prepared_tracks.get('active_max_crossings', 0), + 'track_max_track_crossing_per_step': prepared_tracks.get('active_max_crossings', 0), 'track_crossing_precompute_max': prepared_tracks.get( 'crossing_precompute_max', 0), 'track_crossing_mode': prepared_tracks.get( - 'track_crossing_mode', 'count'), - 'min_walk_steps_per_track': prepared_tracks.get( - 'min_walk_steps_per_track', 24), - 'max_walk_steps_per_track': prepared_tracks.get( - 'max_walk_steps_per_track', 256), - 'n_walks_per_track': prepared_tracks.get('n_walks_per_track', 4), - 'track_walk_require_loop_consistency': prepared_tracks.get( - 'track_walk_require_loop_consistency', False), + 'track_crossing_mode', 'track_walk'), + 'track_min_walk_steps_per_track': prepared_tracks.get( + 'track_min_walk_steps_per_track', 24), + 'track_max_walk_steps_per_track': prepared_tracks.get( + 'track_max_walk_steps_per_track', 256), + 'track_min_walks_per_track': prepared_tracks.get( + 'track_min_walks_per_track', 2), + 'track_max_walks_per_track': prepared_tracks.get( + 'track_max_walks_per_track', 4), + 'track_walk_minimum_cycle_travel': prepared_tracks.get( + 'track_walk_minimum_cycle_travel', 20.0), } current.update({ key: config[key] for key in current @@ -813,9 +843,14 @@ def configure_prepared_track_sampling(prepared_tracks, config): f'prepared crossing ceiling ({prepared_maximum}); reload with a larger ' 'track_crossing_precompute_max') prepared_tracks['active_max_crossings'] = maximum - for key in ('min_walk_steps_per_track', 'max_walk_steps_per_track', - 'n_walks_per_track'): + for key in ( + 'track_min_walk_steps_per_track', + 'track_max_walk_steps_per_track', + 'track_min_walks_per_track', + 'track_max_walks_per_track'): prepared_tracks[key] = policy[key] + prepared_tracks['track_walk_minimum_cycle_travel'] = \ + policy['walk_minimum_cycle_travel'] def _track_tangent(track, raw_index, radius_voxels=12.0): @@ -2133,14 +2168,35 @@ def _track_points_far_from_anchors_mask(track_zyx, anchor_tree, threshold): track_np = np.ascontiguousarray(track_np, dtype=np.float32) if threshold <= 0 or anchor_tree is None: return np.ones(track_np.shape[0], dtype=bool) - dist, _ = anchor_tree.query(track_np, k=1, distance_upper_bound=float(threshold), workers=-1) - return np.isinf(dist) + # A full production store can contain close to a billion points. Passing + # that array to scipy in one call makes query allocate point-count-sized + # distance and index arrays (16 bytes/point) in addition to its internal + # workspace. Keep only the final one-byte mask resident and bound every + # query allocation. + chunk_size = max( + 1, int(os.environ.get('FIT_SPIRAL_TRACK_EXCLUSION_CHUNK', '4000000'))) + keep = np.empty(track_np.shape[0], dtype=bool) + progress = tqdm( + total=track_np.shape[0], desc='excluding track points', + unit='point', unit_scale=True, + ) + try: + for begin in range(0, track_np.shape[0], chunk_size): + end = min(begin + chunk_size, track_np.shape[0]) + dist, _ = anchor_tree.query( + track_np[begin:end], k=1, + distance_upper_bound=float(threshold), workers=-1) + keep[begin:end] = np.isinf(dist) + progress.update(end - begin) + finally: + progress.close() + return keep def prepare_main_phase_tracks( tracks, anchor_scroll_zyxs, exclusion_radius, device, anchor_tree=None, sampling_config=None, track_families=None, track_source_ids=None, - crossing_cache=None): + crossing_cache=None, track_graph=None): if not tracks: return None if sampling_config is not None and 'length_bin_weights' in sampling_config: @@ -2152,7 +2208,7 @@ def prepare_main_phase_tracks( max_crossings = policy['max_crossings'] crossing_precompute_max = policy['crossing_precompute_max'] crossing_mode = policy['crossing_mode'] - require_walk_loop = policy['walk_require_loop_consistency'] + minimum_cycle_travel = policy['walk_minimum_cycle_travel'] input_track_count = len(tracks) packed_input = isinstance(tracks, PackedTrackCollection) @@ -2168,6 +2224,10 @@ def prepare_main_phase_tracks( raise ValueError('track_source_ids must be parallel to tracks') if crossing_cache is not None and working_source_ids is None: raise ValueError('a crossing cache requires stable track_source_ids') + if track_graph is not None and working_source_ids is None: + raise ValueError('a track graph requires stable track_source_ids') + if crossing_cache is not None and track_graph is not None: + raise ValueError('pass either crossing_cache or track_graph, not both') if crossing_cache is not None and exclusion_radius > 0: print( 'WARNING: track crossing cache cannot be used after point-level ' @@ -2199,7 +2259,8 @@ def prepare_main_phase_tracks( anchor_tree = _build_anchor_kdtree(anchor_scroll_zyxs) def finish_prepared( - flat_zyx_np, lengths_new, surviving_indices, prepared_track_list): + flat_zyx_np, lengths_new, surviving_indices, prepared_track_list, + crossing_csr_override=None): offsets_new = np.empty(len(lengths_new) + 1, dtype=np.int64) offsets_new[0] = 0 np.cumsum(lengths_new, out=offsets_new[1:]) @@ -2218,21 +2279,32 @@ def finish_prepared( 'active_max_crossings': 0, 'crossing_precompute_max': ( crossing_precompute_max - if working_families is not None or crossing_cache is not None + if (working_families is not None + or crossing_cache is not None + or track_graph is not None) else 0), 'track_crossing_mode': crossing_mode, - 'track_walk_require_loop_consistency': require_walk_loop, - 'min_walk_steps_per_track': policy['min_walk_steps_per_track'], - 'max_walk_steps_per_track': policy['max_walk_steps_per_track'], - 'n_walks_per_track': policy['n_walks_per_track'], + 'track_walk_minimum_cycle_travel': minimum_cycle_travel, + 'track_min_walk_steps_per_track': policy['track_min_walk_steps_per_track'], + 'track_max_walk_steps_per_track': policy['track_max_walk_steps_per_track'], + 'track_min_walks_per_track': + policy['track_min_walks_per_track'], + 'track_max_walks_per_track': + policy['track_max_walks_per_track'], } if ((crossing_mode == 'count' and crossing_precompute_max > 0) or crossing_mode == 'track_walk'): crossing_index = None - restricted_csr = None + restricted_csr = crossing_csr_override walk_index = None native = _load_native_track_crossings() - if crossing_cache is not None: + active_crossing_cache = crossing_cache + if restricted_csr is None and track_graph is not None: + active_crossing_cache = track_graph + print( + f'track crossings: using TrackGraph cache for ' + f'{len(surviving_indices)} surviving tracks') + if active_crossing_cache is not None: eligible_source_ids = working_source_ids[surviving_indices] try: if crossing_mode == 'track_walk': @@ -2241,53 +2313,55 @@ def finish_prepared( native, 'prepare_cached_walk_index')): walk_index = native.prepare_cached_walk_index( np.asarray( - crossing_cache['source_ids'], + active_crossing_cache['source_ids'], dtype=np.uint64), np.asarray( - crossing_cache['offsets'], + active_crossing_cache['offsets'], dtype=np.int64), np.asarray( - crossing_cache['partners'], + active_crossing_cache['partners'], dtype=np.int32), np.asarray( - crossing_cache['self_local'], + active_crossing_cache['self_local'], dtype=np.int32), np.asarray( - crossing_cache['partner_local'], + active_crossing_cache['partner_local'], dtype=np.int32), + np.asarray( + active_crossing_cache['positions'], + dtype=np.float64), np.asarray( eligible_source_ids, dtype=np.uint64), - np.asarray(lengths_new, dtype=np.int32), - require_loop_consistency=require_walk_loop) + np.asarray(lengths_new, dtype=np.int32)) else: restricted_csr = _restrict_crossing_partner_csr( - crossing_cache, eligible_source_ids) + active_crossing_cache, eligible_source_ids) elif crossing_precompute_max > 0: if (native is not None and hasattr( native, 'prepare_cached_crossing_index')): crossing_index = native.prepare_cached_crossing_index( np.asarray( - crossing_cache['source_ids'], + active_crossing_cache['source_ids'], dtype=np.uint64), np.asarray( - crossing_cache['offsets'], + active_crossing_cache['offsets'], dtype=np.int64), np.asarray( - crossing_cache['partners'], + active_crossing_cache['partners'], dtype=np.int32), np.asarray( - crossing_cache['self_local'], + active_crossing_cache['self_local'], dtype=np.int32), np.asarray( - crossing_cache['partner_local'], + active_crossing_cache['partner_local'], dtype=np.int32), np.asarray( eligible_source_ids, dtype=np.uint64), np.asarray(lengths_new, dtype=np.int32)) else: restricted_csr = _restrict_crossing_partner_csr( - crossing_cache, eligible_source_ids) + active_crossing_cache, eligible_source_ids) print( f'track crossings: used cached CSR for ' f'{len(eligible_source_ids)} surviving tracks') @@ -2348,11 +2422,12 @@ def finish_prepared( restricted_csr['self_local'], dtype=np.int32) walk_partner_local = np.asarray( restricted_csr['partner_local'], dtype=np.int32) + walk_positions = np.asarray( + restricted_csr['positions'], dtype=np.float64) walk_index = native.prepare_walk_index( walk_offsets, walk_partners, walk_self, - walk_partner_local, - np.asarray(lengths_new, dtype=np.int32), - require_loop_consistency=require_walk_loop) + walk_partner_local, walk_positions, + np.asarray(lengths_new, dtype=np.int32)) prepared['walk_index'] = walk_index stats = native.walk_index_stats(prepared['walk_index']) prepared['walk_index_stats'] = dict(stats) @@ -2365,10 +2440,15 @@ def finish_prepared( configure_prepared_track_sampling(prepared, { 'track_length_bin_weights': ( None if weights is None else weights.tolist()), - 'max_track_crossing_per_step': max_crossings, - 'min_walk_steps_per_track': policy['min_walk_steps_per_track'], - 'max_walk_steps_per_track': policy['max_walk_steps_per_track'], - 'n_walks_per_track': policy['n_walks_per_track'], + 'track_max_track_crossing_per_step': max_crossings, + 'track_min_walk_steps_per_track': policy['track_min_walk_steps_per_track'], + 'track_max_walk_steps_per_track': policy['track_max_walk_steps_per_track'], + 'track_min_walks_per_track': + policy['track_min_walks_per_track'], + 'track_max_walks_per_track': + policy['track_max_walks_per_track'], + 'track_walk_minimum_cycle_travel': + policy['walk_minimum_cycle_travel'], }) return prepared @@ -2421,53 +2501,81 @@ def finish_prepared( flat_zyx_np, lengths_new, surviving, surviving_tracks) if packed_input: - flat_materialized, packed_offsets = working_tracks.materialize() - working_tracks = _MemmapTrackCollection( - flat_materialized, packed_offsets) - - flat_zyx_np = np.concatenate([t.astype(np.float32) for t in working_tracks], axis=0) - track_id_np = np.concatenate([ - np.full(len(t), i, dtype=np.int64) for i, t in enumerate(working_tracks) - ]) + flat_zyx_np, packed_offsets = working_tracks.materialize() + flat_zyx_np = np.asarray(flat_zyx_np, dtype=np.float32) + input_offsets = np.asarray(packed_offsets, dtype=np.int64) + else: + input_lengths = np.fromiter( + (len(track) for track in working_tracks), + dtype=np.int64, count=len(working_tracks)) + input_offsets = np.empty(len(input_lengths) + 1, dtype=np.int64) + input_offsets[0] = 0 + np.cumsum(input_lengths, out=input_offsets[1:]) + flat_zyx_np = np.concatenate( + [t.astype(np.float32) for t in working_tracks], axis=0) keep_np = _track_points_far_from_anchors_mask(flat_zyx_np, anchor_tree, exclusion_radius) - flat_zyx_np = flat_zyx_np[keep_np] - track_id_np = track_id_np[keep_np] num_tracks_orig = len(working_tracks) - new_lengths = np.bincount(track_id_np, minlength=num_tracks_orig) + # Points are already grouped by track. Segment reduction replaces the old + # int64 track-id-per-point array (7.3 GiB for the 2um store). + new_lengths = np.add.reduceat( + keep_np, input_offsets[:-1], dtype=np.int64) surviving = np.where(new_lengths >= 2)[0] print(f'kept {len(surviving)} / {len(working_tracks)} tracks') if len(surviving) == 0: return None - old_to_new = -np.ones(num_tracks_orig, dtype=np.int64) - old_to_new[surviving] = np.arange(len(surviving)) - new_id = old_to_new[track_id_np] - keep2 = new_id >= 0 - flat_zyx_np = flat_zyx_np[keep2] - new_id = new_id[keep2] - sort_idx = np.argsort(new_id, kind='stable') - flat_zyx_np = flat_zyx_np[sort_idx] + # Drop the residual 0/1-point tracks without constructing IDs or sorting: + # stable boolean compaction preserves both track and local-point order. + surviving_points = np.repeat( + new_lengths >= 2, np.diff(input_offsets)) + keep_np &= surviving_points + del surviving_points lengths_new = new_lengths[surviving].astype(np.int64) - prepared_track_list = list(np.split(flat_zyx_np, np.cumsum(lengths_new)[:-1])) + compact_zyx_np = flat_zyx_np[keep_np] + crossing_csr_override = None + if track_graph is not None: + if int(keep_np.sum()) >= np.iinfo(np.int32).max: + raise ValueError( + 'point-clipped track output exceeds the int32 crossing limit') + old_point_to_new = np.cumsum( + keep_np, dtype=np.int32) - np.int32(1) + old_point_to_new[~keep_np] = -1 + output_offsets = np.empty(len(lengths_new) + 1, dtype=np.int64) + output_offsets[0] = 0 + np.cumsum(lengths_new, out=output_offsets[1:]) + crossing_csr_override = track_graph.clipped_csr( + working_source_ids, input_offsets, surviving, + old_point_to_new, output_offsets) + del old_point_to_new + print( + f'track crossings: remapped TrackGraph after point exclusion ' + f'({len(crossing_csr_override["partners"])} directed crossings)') + del keep_np, flat_zyx_np + prepared_track_list = _MemmapTrackCollection( + compact_zyx_np, + np.r_[np.int64(0), np.cumsum(lengths_new, dtype=np.int64)]) print( f'track radius loss: {len(surviving)}/{num_tracks_orig} tracks survive exclusion ' f'(radius {exclusion_radius:.1f}); {int(lengths_new.sum())} points retained' ) return finish_prepared( - flat_zyx_np, lengths_new, surviving, prepared_track_list) + compact_zyx_np, lengths_new, surviving, prepared_track_list, + crossing_csr_override=crossing_csr_override) def _build_resampled_track_bundle(prepared_tracks, min_spacing, max_spacing): """Resample complete tracks between mandatory polyline anchors.""" min_spacing = float(min_spacing) max_spacing = float(max_spacing) - if prepared_tracks.get('track_crossing_mode', 'count') == 'track_walk': + if prepared_tracks.get( + 'track_crossing_mode', 'track_walk') == 'track_walk': cache_key = ( min_spacing, max_spacing, 'track_walk', - prepared_tracks.get('min_walk_steps_per_track', 24), - prepared_tracks.get('max_walk_steps_per_track', 256), - prepared_tracks.get('n_walks_per_track', 4), + prepared_tracks.get('track_min_walk_steps_per_track', 24), + prepared_tracks.get('track_max_walk_steps_per_track', 256), + prepared_tracks.get('track_min_walks_per_track', 2), + prepared_tracks.get('track_max_walks_per_track', 4), prepared_tracks.get( - 'track_walk_require_loop_consistency', False), + 'track_walk_minimum_cycle_travel', 20.0), ) else: # This key is intentionally unchanged for the default count mode. @@ -2928,12 +3036,16 @@ def _draw_track_walk_sample( seed = int(torch.randint( 0, (1 << 63) - 1, (1,), dtype=torch.int64, device=device, generator=generator).item()) - hops = int(prepared_tracks['n_walks_per_track']) + minimum_hops = int(prepared_tracks['track_min_walks_per_track']) + maximum_hops = int(prepared_tracks['track_max_walks_per_track']) result = native.sample_walks_adaptive( prepared_tracks['walk_index'], probabilities_cpu, seed=seed, - groups=k, target_points=target_points, hops=hops, - minimum_steps=int(prepared_tracks['min_walk_steps_per_track']), - maximum_steps=int(prepared_tracks['max_walk_steps_per_track']), + groups=k, target_points=target_points, + minimum_hops=minimum_hops, maximum_hops=maximum_hops, + minimum_steps=int(prepared_tracks['track_min_walk_steps_per_track']), + maximum_steps=int(prepared_tracks['track_max_walk_steps_per_track']), + minimum_candidate_travel=float( + prepared_tracks['track_walk_minimum_cycle_travel']), maximum_attempts=attempts, ) produced = int(result['produced']) @@ -2944,20 +3056,44 @@ def _draw_track_walk_sample( f'({produced}/{k} after {int(result["attempted_candidates"])} ' f'policy-compatible primary draws; ' f'{stats.get("eligible_tracks", "unknown")} eligible tracks, ' - f'{hops} hops, raw-index bounds ' - f'[{prepared_tracks["min_walk_steps_per_track"]}, ' - f'{prepared_tracks["max_walk_steps_per_track"]}])') + f'{minimum_hops}-{maximum_hops} hops, raw-index bounds ' + f'[{prepared_tracks["track_min_walk_steps_per_track"]}, ' + f'{prepared_tracks["track_max_walk_steps_per_track"]}])') - width = hops + 1 + width = maximum_hops + 1 track_matrix = torch.from_numpy( np.asarray(result['tracks'], dtype=np.int64)).to(device=device) record_matrix = torch.from_numpy( np.asarray(result['records'], dtype=np.int64)).to(device=device) - track_idx = track_matrix.reshape(-1) - group_id = torch.arange( - k, device=device, dtype=torch.int64).repeat_interleave(width) - row_slot = torch.arange( - width, device=device, dtype=torch.int64).repeat(k) + walk_hops = torch.from_numpy( + np.asarray(result['walk_hops'], dtype=np.int64)).to(device=device) + edge_slots_matrix = torch.arange( + maximum_hops, device=device, dtype=torch.int64)[None, :].expand(k, -1) + edge_valid = edge_slots_matrix < walk_hops[:, None] + edge_group_id = torch.arange( + k, device=device, dtype=torch.int64)[:, None].expand( + -1, maximum_hops)[edge_valid] + edge_slot = edge_slots_matrix[edge_valid] + + # Keep all primary rows first: DT target caching relies on this layout. + # Variable-length partner rows follow in group/hop order. + track_idx = torch.cat([ + track_matrix[:, 0], + track_matrix[:, 1:][edge_valid], + ]) + group_id = torch.cat([ + torch.arange(k, device=device, dtype=torch.int64), + edge_group_id, + ]) + row_slot = torch.cat([ + torch.zeros(k, device=device, dtype=torch.int64), + edge_slot + 1, + ]) + row_index = torch.full( + [k, width], -1, device=device, dtype=torch.int64) + row_index[:, 0] = torch.arange(k, device=device, dtype=torch.int64) + row_index[edge_group_id, edge_slot + 1] = torch.arange( + k, len(track_idx), device=device, dtype=torch.int64) row_lengths = resampled['lengths'][track_idx] row_starts = torch.zeros( len(track_idx) + 1, dtype=torch.int64, device=device) @@ -2993,12 +3129,12 @@ def _draw_track_walk_sample( flat_idx_cpu[target_flat_idx.reshape(-1).cpu()]].reshape( len(track_idx), target_points).to(device=device) - hop_groups = torch.arange(k, device=device)[:, None] - hop_slots = torch.arange(hops, device=device)[None, :] - source_rows = (hop_groups * width + hop_slots).reshape(-1) - partner_rows = source_rows + 1 - source_records = record_matrix[:, 0::2].reshape(-1) - partner_records = record_matrix[:, 1::2].reshape(-1) + source_rows = row_index[:, :-1][edge_valid] + partner_rows = row_index[:, 1:][edge_valid] + if torch.any(source_rows < 0) or torch.any(partner_rows < 0): + raise RuntimeError("variable track-walk row layout is incomplete") + source_records = record_matrix[:, 0::2][edge_valid] + partner_records = record_matrix[:, 1::2][edge_valid] record_samples = resampled['walk_record_sample'] primary_cross_flat = ( row_starts[source_rows] + record_samples[source_records]) @@ -3020,6 +3156,9 @@ def _draw_track_walk_sample( 'primary_cross_flat': primary_cross_flat, 'partner_cross_flat': partner_cross_flat, 'partner_rows': partner_rows, + 'edge_group_id': edge_group_id, + 'edge_slot': edge_slot, + 'maximum_walk_hops': maximum_hops, } @@ -3042,18 +3181,20 @@ def _sample_prepared_track_points( resampled = _build_resampled_track_bundle( prepared_tracks, min_sample_spacing, max_sample_spacing) - if prepared_tracks.get('track_crossing_mode', 'count') == 'track_walk': + if prepared_tracks.get( + 'track_crossing_mode', 'track_walk') == 'track_walk': if prefetch.prefetch_enabled() and device.type == 'cuda': pf = prefetch.get_prefetcher() generator = pf.torch_rng('tracks', device) walk_key = ( 'track_walk', k, num_points_per_track, float(min_sample_spacing), float(max_sample_spacing), - int(prepared_tracks['min_walk_steps_per_track']), - int(prepared_tracks['max_walk_steps_per_track']), - int(prepared_tracks['n_walks_per_track']), - bool(prepared_tracks[ - 'track_walk_require_loop_consistency']), + int(prepared_tracks['track_min_walk_steps_per_track']), + int(prepared_tracks['track_max_walk_steps_per_track']), + int(prepared_tracks['track_min_walks_per_track']), + int(prepared_tracks['track_max_walks_per_track']), + float(prepared_tracks[ + 'track_walk_minimum_cycle_travel']), ) return pf.pop_or_run( walk_key, @@ -3184,22 +3325,28 @@ def _grouped_same_radius_loss( def _crossing_row_alignments( shifted_radii, primary_cross_flat, partner_cross_flat, - partner_rows, row_count, chain=False): + partner_rows, row_count, chain=False, edge_group_id=None, + edge_slot=None, group_count=None, maximum_hops=None): row_alignment = torch.zeros( row_count, device=shifted_radii.device, dtype=shifted_radii.dtype) if chain: edge_count = primary_cross_flat.numel() - group_count = row_count - edge_count if edge_count == 0: return row_alignment - if group_count <= 0 or edge_count % group_count: - raise ValueError("track-walk rows do not form fixed-width groups") - hops = edge_count // group_count - hop_deltas = ( + if (edge_group_id is None or edge_slot is None + or group_count is None or maximum_hops is None): + raise ValueError("variable track-walk edge layout is incomplete") + edge_deltas = ( shifted_radii[primary_cross_flat] - shifted_radii[partner_cross_flat] - ).reshape(group_count, hops) - row_alignment[partner_rows] = hop_deltas.cumsum(dim=1).reshape(-1).detach() + ) + padded = torch.zeros( + [group_count, maximum_hops], + device=shifted_radii.device, dtype=shifted_radii.dtype) + padded[edge_group_id, edge_slot] = edge_deltas + accumulated = padded.cumsum(dim=1) + row_alignment[partner_rows] = accumulated[ + edge_group_id, edge_slot].detach() else: row_alignment[partner_rows] = ( shifted_radii[primary_cross_flat] @@ -3223,11 +3370,11 @@ def iter_track_losses(slice_to_spiral_transform, dr_per_winding, prepared_tracks return sample = _sample_prepared_track_points( prepared_tracks, - cfg['track_num_per_step'], - cfg['track_num_points_per_step'], + cfg['sample_count_tracks_per_step'], + cfg['sample_count_track_points_per_step'], cfg.get('track_min_sample_spacing', 20.0), cfg.get('track_max_sample_spacing', 60.0), - cfg.get('max_track_crossing_per_step', 0), + cfg.get('track_max_track_crossing_per_step', 0), ) if sample is None: yield 'track_radius', zero @@ -3257,7 +3404,11 @@ def iter_track_losses(slice_to_spiral_transform, dr_per_winding, prepared_tracks sample['partner_cross_flat'], sample['partner_rows'], len(track_idx), chain=(prepared_tracks.get( - 'track_crossing_mode', 'count') == 'track_walk')) + 'track_crossing_mode', 'track_walk') == 'track_walk'), + edge_group_id=sample.get('edge_group_id'), + edge_slot=sample.get('edge_slot'), + group_count=num_groups, + maximum_hops=sample.get('maximum_walk_hops')) flat_alignment = row_alignment[row_id] shifted_radii = shifted_radii + flat_alignment crossing_adjustments = crossing_adjustments + flat_alignment diff --git a/volume-cartographer/scripts/spiral/transforms.py b/volume-cartographer/scripts/spiral/transforms.py index f6618c6670..4dbbbb1ee8 100644 --- a/volume-cartographer/scripts/spiral/transforms.py +++ b/volume-cartographer/scripts/spiral/transforms.py @@ -31,7 +31,7 @@ def __init__(self, flow_field, flow_min_corner_zyx, flow_max_corner_zyx, num_ste self.truncate_at_step = truncate_at_step self._event_dim = event_dim self._flow_range_zyx = self.flow_max_corner_zyx - self.flow_min_corner_zyx - self.num_flow_timesteps = getattr(flow_field, 'num_flow_timesteps', 1) + self.num_flow_timesteps = getattr(flow_field, 'model_num_flow_timesteps', 1) # Cached sampler/integrator closure at t=0 for the num_flow_timesteps==1 fast path. # Built once per diffeomorphism instance (one per training iteration), shared across # forward and inverse calls so per-iteration setup (e.g. trilinear LR->HR upsampling) @@ -128,7 +128,9 @@ def _get_pinned_scaled_logits(self): needs_grad = self.params.logits.requires_grad and torch.is_grad_enabled() cached = self._pinned_scaled_logits if cached is None or (needs_grad and not cached.requires_grad): - # Pin the 0th logit (i.e. theta=0 on 1th winding) to be zero, to avoid a jump going from winding #0 to #1 + # Pin the 0th logit (i.e. theta=0 on 1th winding) to be zero, to avoid a jump going from winding #0 to #1. + # Keep this in sync with SpiralAndTransform.get_shared_transform_tensors, + # which precomputes the same tensor for injection as a detached leaf. logits = torch.cat([torch.zeros_like(self.params.logits[..., :1]), self.params.logits[..., 1:]], dim=-1) self._pinned_scaled_logits = logits * self.gap_expander_lr_scale return self._pinned_scaled_logits @@ -455,16 +457,16 @@ def __init__(self, flow_integration_steps, flow_integration_solver, flow_min_cor self.linear_logits_scale = 40. # larger value increases effective learning rate self.umbilicus_transform = UmbilicusTransform(umbilicus_zyx) - self.dr_per_winding_logit = nn.Parameter(torch.tensor(config['initial_dr_per_winding'] / self.dr_per_winding_scale, dtype=torch.float32)) + self.dr_per_winding_logit = nn.Parameter(torch.tensor(config['model_initial_dr_per_winding'] / self.dr_per_winding_scale, dtype=torch.float32)) - flow_resolution = (flow_max_corner_zyx - flow_min_corner_zyx) // config['flow_voxel_resolution'] - flow_field_cls = CylindricalFlowField if config['flow_field_type'] == 'cylindrical' else CartesianFlowField + flow_resolution = (flow_max_corner_zyx - flow_min_corner_zyx) // config['model_flow_voxel_resolution'] + flow_field_cls = CylindricalFlowField if config['model_flow_field_type'] == 'cylindrical' else CartesianFlowField def make_flow_field(): return flow_field_cls( flow_resolution, - lr_scale_factor=config['flow_field_high_res_lr_scale_initial'], - num_flow_timesteps=config['num_flow_timesteps'], + num_flow_timesteps=config['model_num_flow_timesteps'], + direct_lr=config.get('model_flow_field_direct_lr', False), ) # num_flow_stages: number of independent stationary flow fields whose integrated @@ -472,19 +474,19 @@ def make_flow_field(): # spiral->slice direction; the inverse applies the stage inverses in reverse order via # ComposeTransform.inv). num_flow_stages == 1 is exactly the original single-field # behaviour: `flow_field` keeps its name/state_dict keys and `extra_flow_fields` is empty. - self.num_flow_stages = int(config.get('num_flow_stages', 1) or 1) + self.num_flow_stages = int(config.get('model_num_flow_stages', 1) or 1) assert self.num_flow_stages >= 1 self.flow_field = make_flow_field() self.extra_flow_fields = nn.ModuleList([make_flow_field() for _ in range(self.num_flow_stages - 1)]) - self.linear_logits = nn.Parameter(torch.zeros([int(flow_max_corner_zyx[0] - flow_min_corner_zyx[0]) // config['linear_z_resolution'], 2, 2], dtype=torch.float32)) + self.linear_logits = nn.Parameter(torch.zeros([int(flow_max_corner_zyx[0] - flow_min_corner_zyx[0]) // config['model_linear_z_resolution'], 2, 2], dtype=torch.float32)) self.gap_expander_params = GapExpanderParams( - resolution=config['gap_expander_logit_resolution'], + resolution=config['model_gap_expander_logit_resolution'], min_z=flow_min_corner_zyx[0], max_z=flow_max_corner_zyx[0], - num_windings=config['gap_expander_num_windings'], - dr_per_winding=config['initial_dr_per_winding'], # this is a nominal (fixed) winding spacing which we only use to calculate the number of logits + num_windings=config['model_gap_expander_num_windings'], + dr_per_winding=config['model_initial_dr_per_winding'], # this is a nominal (fixed) winding spacing which we only use to calculate the number of logits ) @property @@ -496,7 +498,7 @@ def flow_fields(self): # All flow stages, in application order (stage 0 first in the spiral->slice direction). return [self.flow_field, *self.extra_flow_fields] - def _get_transform_parts(self, truncate_at_step=None): + def _get_transform_parts(self, truncate_at_step=None, shared=None): truncate_frac = None if truncate_at_step is None else truncate_at_step / (self.flow_integration_steps - 1) diffeomorphisms = [ IntegratedFlowDiffeomorphism(flow_field, self.flow_min_corner_zyx, self.flow_max_corner_zyx, num_steps=self.flow_integration_steps, solver=self.flow_integration_solver, truncate_at_step=truncate_at_step) @@ -504,12 +506,14 @@ def _get_transform_parts(self, truncate_at_step=None): ] gap_expander = GapExpandingTransform( self.gap_expander_params, - self.get_dr_per_winding(), + shared[0] if shared is not None else self.get_dr_per_winding(), self.flow_min_corner_zyx[0], self.flow_max_corner_zyx[0], - self.cfg['gap_expander_lr_scale'], + self.cfg['model_gap_expander_lr_scale'], truncate_frac, ) + if shared is not None: + gap_expander._pinned_scaled_logits = shared[2] if self.spiral_outward_sense == 'CW': maybe_flip = [] else: @@ -518,15 +522,22 @@ def _get_transform_parts(self, truncate_at_step=None): maybe_flip = [pyro.distributions.transforms.AffineTransform(loc=0., scale=torch.tensor([1., 1., -1.], device=self.device))] return gap_expander, maybe_flip, diffeomorphisms, truncate_frac - def get_slice_to_spiral_transform(self, truncate_at_step=None): - gap_expander, maybe_flip, diffeomorphisms, truncate_frac = self._get_transform_parts(truncate_at_step) + def get_slice_to_spiral_transform(self, truncate_at_step=None, shared=None): + # `shared` optionally supplies the (dr_per_winding, scaled_linear_logits, + # pinned_scaled_gap_logits) triple from get_shared_transform_tensors(), + # typically as detached leaves so many separate loss backwards can run + # through one transform instance without retain_graph. + gap_expander, maybe_flip, diffeomorphisms, truncate_frac = self._get_transform_parts(truncate_at_step, shared) + scaled_linear_logits = ( + shared[1] if shared is not None + else self.linear_logits * self.linear_logits_scale) return pyro.distributions.transforms.ComposeTransform([ gap_expander, *maybe_flip, # Sequential composition of the integrated stationary flows; ComposeTransform.inv # applies the stage inverses in reverse order in the slice->spiral direction. *diffeomorphisms, - VaryingLinearTransform(self.linear_logits * self.linear_logits_scale, self.flow_min_corner_zyx[0], self.flow_max_corner_zyx[0], truncate_frac), + VaryingLinearTransform(scaled_linear_logits, self.flow_min_corner_zyx[0], self.flow_max_corner_zyx[0], truncate_frac), self.umbilicus_transform, ]).inv @@ -547,6 +558,25 @@ def get_flowbox_to_spiral_transform(self, include_diffeomorphism=True): def get_dr_per_winding(self): return F.softplus(self.dr_per_winding_logit * self.dr_per_winding_scale) + def get_shared_transform_tensors(self): + """The tiny graph paths every evaluation of one transform instance + shares: the dr-per-winding softplus, the scaled linear logits, and the + pinned+scaled gap logits (kept in sync with + GapExpandingTransform._get_pinned_scaled_logits). The training loop + passes detached leaf copies to get_slice_to_spiral_transform(shared=...) + so each loss family's backward owns its whole graph and needs no + retain_graph, then propagates the accumulated leaf gradients through + these outputs once per step.""" + gap_logits = self.gap_expander_params.logits + pinned_scaled_gap_logits = torch.cat( + [torch.zeros_like(gap_logits[..., :1]), gap_logits[..., 1:]], dim=-1, + ) * self.cfg['model_gap_expander_lr_scale'] + return ( + self.get_dr_per_winding(), + self.linear_logits * self.linear_logits_scale, + pinned_scaled_gap_logits, + ) + def get_native_log_gaps(self, winding_idx, theta, z): """Exact pre-exponentiation log gap for the native gap expander.""" gap_expander, _, _, truncate_frac = self._get_transform_parts() diff --git a/volume-cartographer/scripts/spiral/update_checkpoint.py b/volume-cartographer/scripts/spiral/update_checkpoint.py new file mode 100755 index 0000000000..0aa741b20a --- /dev/null +++ b/volume-cartographer/scripts/spiral/update_checkpoint.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""Rewrite a legacy Spiral checkpoint for the current VC3D Spiral workspace. + +The source checkpoint is never modified. Only embedded configuration +dictionaries are migrated; model, optimiser, scheduler, and RNG state are +carried through unchanged. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time +from collections.abc import Mapping +from pathlib import Path + +import torch + +from checkpoint_io import load_checkpoint_cpu +from config import Config + + +MODEL_KEYS = ( + "num_flow_integration_steps", + "flow_integration_solver", + "num_flow_timesteps", + "num_flow_stages", + "flow_bounds_z_margin", + "flow_bounds_radius", + "flow_voxel_resolution", + "flow_field_type", + "gap_expander_logit_resolution", + "gap_expander_num_windings", + "gap_expander_lr_scale", + "linear_z_resolution", + "initial_dr_per_winding", +) + +LEGACY_RENAMES = { + "random_seed": "optimizer_random_seed", + "distributed_split_batch": "optimizer_distributed_split_batch", + "learning_rate": "optimizer_learning_rate", + "exp_lr_schedule": "optimizer_exp_lr_schedule", + "lr_final_factor": "optimizer_lr_final_factor", + "num_training_steps": "optimizer_num_training_steps", + **{key: f"model_{key}" for key in MODEL_KEYS}, + "num_patches_per_step": "sample_count_patches_per_step", + "num_patches_per_step_for_dt": "sample_count_patches_per_step_for_dt", + "num_points_per_patch": "sample_count_points_per_patch", + "unverified_num_patches_per_step": + "sample_count_unverified_patches_per_step", + "unverified_num_patches_per_step_for_dt": + "sample_count_unverified_patches_per_step_for_dt", + "unverified_num_points_per_patch": + "sample_count_unverified_points_per_patch", + "rel_winding_num_pcls": "sample_count_relative_winding_pcls", + "rel_winding_num_patch_pairs_per_pcl": + "sample_count_relative_winding_patch_pairs_per_pcl", + "abs_winding_num_pcls": "sample_count_absolute_winding_pcls", + "abs_winding_num_points_per_pcl": + "sample_count_absolute_winding_points_per_pcl", + "unattached_pcl_num_per_step": + "sample_count_unattached_pcls_per_step", + "unattached_pcl_num_points_per_step": + "sample_count_unattached_pcl_points_per_step", + "track_num_per_step": "sample_count_tracks_per_step", + "track_num_points_per_step": "sample_count_track_points_per_step", + "dense_normals_num_points": "sample_count_dense_normal_points", + "regularisation_num_points": "sample_count_regularisation_points", + "dense_spacing_num_pairs": "sample_count_dense_spacing_pairs", + "dense_spacing_count_extra_pairs": + "sample_count_dense_spacing_count_extra_pairs", + "dense_spacing_density_extra_pairs": + "sample_count_dense_spacing_density_extra_pairs", + "dense_spacing_density_chunk_pairs": + "sample_count_dense_spacing_density_chunk_pairs", + "min_spacing_independent_samples": + "sample_count_minimum_spacing_independent_samples", + "dense_attachment_num_points": + "sample_count_dense_attachment_points", + "patch_dt_target_num_points": "sample_count_patch_dt_target_points", + "dt_target_num_points_per_strip": + "sample_count_dt_target_points_per_strip", + "shell_num_samples": "sample_count_shell_samples", + "interactive_influence_footprint_points": + "sample_count_influence_footprint_points", + "interactive_influence_anchor_lattice_points": + "sample_count_influence_anchor_lattice_points", + "interactive_influence_anchor_geometry_points": + "sample_count_influence_anchor_geometry_points", + "interactive_influence_anchor_samples_per_step": + "sample_count_influence_anchor_samples_per_step", + "erode_patches": "patch_erode_patches", + "disable_patches": "input_disable_patches", + "unverified_patch_radius_loss_margin": + "patch_unverified_patch_radius_loss_margin", + "unverified_patch_radius_loss_inv": + "patch_unverified_patch_radius_loss_inv", + "unverified_patch_radius_within_norm_p": + "patch_unverified_patch_radius_within_norm_p", + "unverified_patch_dt_norm_p": "patch_unverified_patch_dt_norm_p", + "unverified_patch_dt_within_patch_norm_p": + "patch_unverified_patch_dt_within_patch_norm_p", + "unverified_patch_dt_loss_margin": + "patch_unverified_patch_dt_loss_margin", + "unverified_patch_exclusion_radius": + "patch_unverified_patch_exclusion_radius", + "rel_winding_adjacent_patches_only": + "pcl_rel_winding_adjacent_patches_only", + "stratified_pcl_sampling": "pcl_stratified_pcl_sampling", + "fiber_min_point_spacing": "pcl_fiber_min_point_spacing", + "unattached_pcl_min_point_spacing": + "pcl_unattached_pcl_min_point_spacing", + "max_track_crossing_per_step": "track_max_track_crossing_per_step", + "min_walk_steps_per_track": "track_min_walk_steps_per_track", + "max_walk_steps_per_track": "track_max_walk_steps_per_track", + "n_walks_per_track": "track_max_walks_per_track", + "grad_mag_encode_scale": "dense_grad_mag_encode_scale", + "grad_mag_factor": "dense_grad_mag_factor", + "spacing_integration_steps": "dense_spacing_integration_steps", + "min_spacing_d_min_wv": "dense_min_spacing_d_min_wv", + "sym_dirichlet_finite_difference_epsilon": + "model_sym_dirichlet_finite_difference_epsilon", + "weight_decay_gap_expander": "optimizer_weight_decay_gap_expander", + "weight_decay_flow_field": "optimizer_weight_decay_flow_field", + "save_png_visualizations": "output_save_png_visualizations", + "interactive_influence_enabled": "influence_enabled", + "interactive_influence_z": "influence_z", + "interactive_influence_windings": "influence_windings", + "interactive_influence_theta_frac": "influence_theta_frac", + "interactive_influence_disable_dt_frac": "influence_disable_dt_frac", + "interactive_influence_sigma": "influence_sigma", + "interactive_influence_anchor_ramp_power": + "influence_anchor_ramp_power", +} + +# These settings moved out of the durable config or were replaced by current +# track-walk controls. VC3D now derives input enablement from selected paths. +LEGACY_REMOVED = { + "use_verified_patches", + "use_unverified_patches", + "use_normals", + "use_surf_sdt", + "use_tracks", + "use_gradient_magnitude", + "use_fibers", + "track_walk_require_loop_consistency", +} + +CONFIG_FIELDS = ("cfg", "requested_config", "resolved_config") + + +def migrate_config(source: Mapping) -> tuple[dict, list[str], list[str], list[str]]: + """Return a current-schema config and a summary of the migration.""" + defaults = Config().as_dict() + migrated = dict(defaults) + renamed = [] + removed = [] + consumed = set() + + for old_key, value in source.items(): + if old_key in defaults: + migrated[old_key] = value + consumed.add(old_key) + elif old_key in LEGACY_RENAMES: + new_key = LEGACY_RENAMES[old_key] + migrated[new_key] = value + renamed.append(f"{old_key} -> {new_key}") + consumed.add(old_key) + elif old_key in LEGACY_REMOVED: + removed.append(old_key) + consumed.add(old_key) + + unknown = sorted(set(source) - consumed) + if unknown: + raise ValueError( + "checkpoint contains configuration keys with no known migration: " + + ", ".join(unknown) + ) + + # Validate types, ranges, enum values, and the exact current key set. + validated = Config(migrated).as_dict() + added = sorted(set(defaults) - { + LEGACY_RENAMES.get(key, key) + for key in source + if key not in LEGACY_REMOVED + }) + return validated, sorted(renamed), sorted(removed), added + + +def update_checkpoint(checkpoint: dict) -> tuple[dict, dict]: + """Migrate all embedded config snapshots without touching tensor state.""" + if not isinstance(checkpoint, dict): + raise ValueError("checkpoint root must be a dictionary") + if "spiral_and_transform" not in checkpoint: + raise ValueError("checkpoint has no 'spiral_and_transform' model state") + + updated = dict(checkpoint) + reports = {} + fallback = checkpoint.get("cfg") + if not isinstance(fallback, Mapping): + raise ValueError("checkpoint has no 'cfg' configuration dictionary") + + for field in CONFIG_FIELDS: + source = checkpoint.get(field, fallback) + if not isinstance(source, Mapping): + raise ValueError(f"checkpoint field {field!r} is not a dictionary") + migrated, renamed, removed, added = migrate_config(source) + updated[field] = migrated + reports[field] = { + "renamed": renamed, + "removed": removed, + "added_from_current_defaults": added, + } + return updated, reports + + +def default_output_path(source: Path) -> Path: + return source.with_name(f"{source.stem}_updated{source.suffix}") + + +def save_atomic(checkpoint: dict, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name( + f".{destination.name}.tmp-{os.getpid()}-{time.time_ns()}") + try: + torch.save(checkpoint, temporary) + with temporary.open("rb+") as stream: + os.fsync(stream.fileno()) + os.replace(temporary, destination) + finally: + if temporary.exists(): + temporary.unlink() + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Convert a pre-VC3D Spiral checkpoint's embedded configuration " + "to the schema required by this workspace." + ) + ) + parser.add_argument("checkpoint", type=Path, help="source .ckpt file") + parser.add_argument( + "output", + nargs="?", + type=Path, + help="destination (default: _updated.ckpt)", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="replace an existing destination (the source is never overwritten)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="validate and report changes without writing a checkpoint", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + source = args.checkpoint.expanduser().resolve(strict=True) + destination = ( + args.output.expanduser().resolve() + if args.output is not None else default_output_path(source) + ) + if destination == source: + raise ValueError("refusing to overwrite the source checkpoint") + if destination.exists() and not args.overwrite and not args.dry_run: + raise FileExistsError( + f"destination already exists: {destination}; pass --overwrite") + + print(f"[spiral] loading {source}", flush=True) + checkpoint = load_checkpoint_cpu(source) + updated, reports = update_checkpoint(checkpoint) + report = reports["cfg"] + print( + f"[spiral] cfg: {len(report['renamed'])} renamed, " + f"{len(report['removed'])} removed, " + f"{len(report['added_from_current_defaults'])} added from current defaults", + flush=True, + ) + + if args.dry_run: + print("[spiral] dry run complete; no file written") + return 0 + + print(f"[spiral] writing {destination}", flush=True) + save_atomic(updated, destination) + + # Reload through the same mmap path used by the resident fitter and verify + # the exact invariant that VC3D checks before starting a session. + del checkpoint, updated + reloaded = load_checkpoint_cpu(destination) + expected_keys = set(Config().as_dict()) + if not isinstance(reloaded, dict) or set(reloaded.get("cfg", {})) != expected_keys: + raise RuntimeError("written checkpoint failed current-schema validation") + print(f"[spiral] updated checkpoint: {destination}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except (FileNotFoundError, OSError, RuntimeError, ValueError) as exc: + print(f"update_checkpoint.py: error: {exc}", file=sys.stderr) + raise SystemExit(1) diff --git a/volume-cartographer/uv.lock b/volume-cartographer/uv.lock index 959b5c2c57..f4aa87c811 100644 --- a/volume-cartographer/uv.lock +++ b/volume-cartographer/uv.lock @@ -1596,6 +1596,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[[package]] +name = "rustworkx" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ff/1a/545aa3a3e251e9da00e4acf0e5472b7fbbb15360f56d9d5342c1f93b8b93/rustworkx-0.18.0.tar.gz", hash = "sha256:5ca9cf8dbee50f8def012119ebb64771ae86916993417415aba9e845831cb436", size = 894567, upload-time = "2026-06-18T04:38:56.141Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/70/18b310752b0652d4a755b46853268c00c33401101a47f87697ad25966453/rustworkx-0.18.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:7440ba70b87bd16811d92e57d76108840dbdd89aa2fc55e4d324fa2fb6b7f9c9", size = 2295936, upload-time = "2026-06-18T04:38:01.598Z" }, + { url = "https://files.pythonhosted.org/packages/02/4e/09152f4422204f020346a8a2b0e2f05f12d8cc60eeeee99b24f649985514/rustworkx-0.18.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:7fce5218ebfb8d8f5313def1f15aad8d3c3c2bbe03e33805dff8cf8bb70370a1", size = 2145057, upload-time = "2026-06-18T04:38:03.253Z" }, + { url = "https://files.pythonhosted.org/packages/22/68/69198f41f7f9c39fb5b2e561bf41d4cbd117fa5dfdf4631d8ba28e5ff69b/rustworkx-0.18.0-cp310-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7e0c626f76bc71d414a02502be7ec0ac2dd6eca369886bc66a2650607f2d9de6", size = 2391659, upload-time = "2026-06-18T04:38:04.942Z" }, + { url = "https://files.pythonhosted.org/packages/e1/81/8ac568efe0289b0ecd826cd697c3581b6a25e134145926cbfef762656167/rustworkx-0.18.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:91b37c5bb54e233e16c4f47383385f17be03e6a344821c2f6a39a0aebee54538", size = 2189469, upload-time = "2026-06-18T04:38:06.51Z" }, + { url = "https://files.pythonhosted.org/packages/f2/a4/96492d7af2ddcd15368f80329e3514060b122f306c4406fac98c08a97865/rustworkx-0.18.0-cp310-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:bcd15f7a637ca0654f329ed42a6db0cf7eac15ed89e0b32cb9d16b1b21937979", size = 2504116, upload-time = "2026-06-18T04:38:43.441Z" }, + { url = "https://files.pythonhosted.org/packages/4d/b0/005383bd7dc110b06ca731bd2115482f3a8c6389e0c8fe1f7f259380c4c8/rustworkx-0.18.0-cp310-abi3-manylinux_2_28_s390x.whl", hash = "sha256:567ace3b0d8ac709dc4ad4ce164472f5f59508739355221480170327a6736777", size = 3170967, upload-time = "2026-06-18T04:38:39.929Z" }, + { url = "https://files.pythonhosted.org/packages/0d/2f/95ca9ef9285cc6793b3e041efca6c5a194ca6731b0af732631ee074a67b3/rustworkx-0.18.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:b91cdcffeab5f498de44905a373e47e345da8f6c4f40fc969421548c4e7a5219", size = 2254008, upload-time = "2026-06-18T04:38:08.245Z" }, + { url = "https://files.pythonhosted.org/packages/19/7b/9df6a80162e6f90fd9945e3f8063c2aa41285c27aefce1d171f99c2b829f/rustworkx-0.18.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:03dd1ab42974bc0ff957eccc330c0c0d9887d88240c0cd050039e215af5a3c45", size = 2447140, upload-time = "2026-06-18T04:38:09.767Z" }, + { url = "https://files.pythonhosted.org/packages/5d/b3/4d159b33bbc73dbebad050a0226dcfb0753d7f05fa55a8d5351e85c8c953/rustworkx-0.18.0-cp310-abi3-win32.whl", hash = "sha256:eb87a45864ffe36e4d86e826d12bb54dac9e4f76c214997dee624281d8711684", size = 2057237, upload-time = "2026-06-18T04:38:11.28Z" }, + { url = "https://files.pythonhosted.org/packages/99/1b/b0dc8d3751c0c58d16db72bff2954a9d3871c166147007590a3aaec4d42e/rustworkx-0.18.0-cp310-abi3-win_amd64.whl", hash = "sha256:602df896b4479b83c6456f702f8ba2ac1cbb972b30723d5fe2e84e6ff3de7d70", size = 2289224, upload-time = "2026-06-18T04:38:12.961Z" }, + { url = "https://files.pythonhosted.org/packages/f4/eb/147e86d56de0511ee7526e5cdf6144013d9c71ba4d6ed6dcd82b6b545861/rustworkx-0.18.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bba66b268e94efd9181d49241b2547d783770dcea9ee8c64609306a7b20ace74", size = 2339568, upload-time = "2026-06-18T04:38:14.292Z" }, + { url = "https://files.pythonhosted.org/packages/7d/1b/c665d30c132e73d69ca95cf552179608de6f94434ee5eb321297ef393920/rustworkx-0.18.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:d51ff00fa72144813c2af33c673b7faf7182ca0ffc85cf35569b5ac4de84bc3e", size = 2118746, upload-time = "2026-06-18T04:38:15.759Z" }, + { url = "https://files.pythonhosted.org/packages/c5/78/04733471aa616dc0bc76f2c2e65ad89bb4ecbbfd2309231347c293906181/rustworkx-0.18.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7099bf90fc1a7ad1789bc126e3e208d9944af2dd9f9b28f75ae148c68aa00cb", size = 2372040, upload-time = "2026-06-18T04:38:17.082Z" }, + { url = "https://files.pythonhosted.org/packages/bd/eb/58ed5f78ca608156ac2ce3166e41976c32345bbc36e71a3fc18390cc4de2/rustworkx-0.18.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:f562d5eeb2943024c85f631df713e7873ea8b572ada9044be82f29f923c1463a", size = 2178071, upload-time = "2026-06-18T04:38:18.571Z" }, + { url = "https://files.pythonhosted.org/packages/bf/96/888c7849d20b7e49d978cfe9bdd8873d709e220a09f8f6bdce226ef7c2be/rustworkx-0.18.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8437bca6adff8e91089bd5d176ef8a499085135031e2d2df4132821578e94dd6", size = 2244227, upload-time = "2026-06-18T04:38:19.923Z" }, + { url = "https://files.pythonhosted.org/packages/99/47/621b56c3f10138a02da9dfd636780ec9311c821fff1acf43be7292408625/rustworkx-0.18.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ba9e9f91d0d14d21666f0d7549f4d8ed70dc84bf75483c68eee5a005ed900f3e", size = 2426773, upload-time = "2026-06-18T04:38:21.233Z" }, +] + [[package]] name = "s3fs" version = "2026.6.0" @@ -1698,8 +1725,9 @@ dependencies = [ { name = "numpy" }, { name = "opencv-python-headless" }, { name = "pillow" }, - { name = "posix-ipc" }, + { name = "posix-ipc", marker = "sys_platform != 'win32'" }, { name = "pyro-ppl" }, + { name = "rustworkx" }, { name = "s3fs" }, { name = "scipy" }, { name = "tensorstore" }, @@ -1725,8 +1753,9 @@ requires-dist = [ { name = "numpy", specifier = ">=2.5.0" }, { name = "opencv-python-headless", specifier = ">=4.13.0.92" }, { name = "pillow", specifier = ">=12.2.0" }, - { name = "posix-ipc", specifier = ">=1.3.2" }, + { name = "posix-ipc", marker = "sys_platform != 'win32'", specifier = ">=1.3.2" }, { name = "pyro-ppl", specifier = ">=1.9.1" }, + { name = "rustworkx", specifier = ">=0.18.0" }, { name = "s3fs", specifier = ">=2026.6.0" }, { name = "scipy", specifier = ">=1.18.0" }, { name = "tensorstore", specifier = ">=0.1.80" },