diff --git a/docs/config_template.yaml b/docs/config_template.yaml index ab48043..fbff0d8 100644 --- a/docs/config_template.yaml +++ b/docs/config_template.yaml @@ -1,7 +1,9 @@ general: # MissAlignment iteratively trains models and realigns the tilt-series training_directory: /data/set/ # path to directory for writing iteration output - apply_ctf: False # leave False; enabling CTF doubles processing time with no benefit to alignment + # Default for iterations that do not set apply_ctf themselves. Enabling CTF + # increases processing time, so it can be reserved for final refinements. + apply_ctf: False iteration_settings: # defines all iterations to run (length determines number of iterations) # alignment modes: "global" (single pass), "anchoring" (iterative), "spline" (coarse-to-fine) - { downsample: 3, alignment: anchoring } @@ -10,8 +12,8 @@ general: - { downsample: 1, alignment: global } - { downsample: 1, alignment: [3, 3] } # local refinement with a 3x3 image warping grid - { downsample: 1, alignment: [3, 3] } - - { downsample: 1, alignment: [3, 3] } - - { downsample: 1, alignment: [3, 3] } + - { downsample: 1, alignment: [3, 3], apply_ctf: True } + - { downsample: 1, alignment: [3, 3], apply_ctf: True } seed: 45132 model_training: diff --git a/docs/inference_config_template.yaml b/docs/inference_config_template.yaml index a852e26..e12a5f2 100644 --- a/docs/inference_config_template.yaml +++ b/docs/inference_config_template.yaml @@ -4,7 +4,9 @@ general: # /iterN/model.ckpt is loaded and applied. data_directory: /data/newset/ # tilt-series to align; aligned output written here model_run_directory: /data/prevrun/ # finished run holding iter1/.. iterN/model.ckpt - apply_ctf: False # leave False; enabling CTF doubles processing time + # Default for iterations that do not set apply_ctf themselves. Enabling CTF + # increases processing time, so it can be reserved for final refinements. + apply_ctf: False iteration_settings: # length must not exceed the models available in the run # alignment modes: "global" (single pass), "anchoring" (iterative), "spline" (coarse-to-fine) - { downsample: 3, alignment: anchoring } @@ -13,8 +15,8 @@ general: - { downsample: 1, alignment: global } - { downsample: 1, alignment: [3, 3] } # local refinement with a 3x3 image warping grid - { downsample: 1, alignment: [3, 3] } - - { downsample: 1, alignment: [3, 3] } - - { downsample: 1, alignment: [3, 3] } + - { downsample: 1, alignment: [3, 3], apply_ctf: True } + - { downsample: 1, alignment: [3, 3], apply_ctf: True } seed: 45132 tilt_series_alignment: diff --git a/docs/usage.md b/docs/usage.md index a17df66..92decfa 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -21,6 +21,15 @@ Place a miss-alignment config file in the `warp_tiltseries/` directory — use [ - **`training_directory`**: set to `/path/to/your/warp/project/warp_tiltseries/` - **`batch_size`** (in the `tilt_series_alignment` section): controls how many patches are reconstructed simultaneously during alignment. A value of 32 works well for 24 GB cards; reduce it for smaller cards or increase it for larger ones to improve throughput. +- **`apply_ctf`**: the value under `general` is the default for every iteration. + An iteration can override it by adding `apply_ctf: True` or + `apply_ctf: False` to its entry in `iteration_settings`. This allows, for + example, inexpensive non-CTF coarse alignment followed by CTF-aware final + refinement. During training, the resolved value is used consistently for + both model-training reconstructions and alignment in that iteration. The + console reports `Training reconstruction CTF` and `Alignment reconstruction + CTF` separately, including the resolved boolean, so both phases are + auditable in run logs. ### 4. Run miss-alignment diff --git a/src/miss_alignment/alignment/parallel.py b/src/miss_alignment/alignment/parallel.py index a216132..d7249a1 100644 --- a/src/miss_alignment/alignment/parallel.py +++ b/src/miss_alignment/alignment/parallel.py @@ -80,6 +80,11 @@ def run_alignment_parallel( dict[str, float] Dictionary mapping tilt-series names to their final loss values. """ + print( + "Alignment reconstruction CTF: " + f"{'enabled' if apply_ctf else 'disabled'} " + f"(apply_ctf={apply_ctf}, downsample={downsample})" + ) jobs = [ { "model_checkpoint_path": model_checkpoint, diff --git a/src/miss_alignment/data/training_datamodule.py b/src/miss_alignment/data/training_datamodule.py index 7a36f76..79e5226 100644 --- a/src/miss_alignment/data/training_datamodule.py +++ b/src/miss_alignment/data/training_datamodule.py @@ -218,6 +218,11 @@ def setup(self, stage: str = "fit"): f"Pool configuration: {self.n_workers} workers -> " f"{self.n_partitions} partitions ({self.partition_size} files each)" ) + print( + "Training reconstruction CTF: " + f"{'enabled' if self.apply_ctf else 'disabled'} " + f"(apply_ctf={self.apply_ctf}, downsample={self.downsample})" + ) # Start worker processes self.stop_event = ctx.Event() diff --git a/src/miss_alignment/infer.py b/src/miss_alignment/infer.py index d884bd7..e3dec66 100644 --- a/src/miss_alignment/infer.py +++ b/src/miss_alignment/infer.py @@ -11,7 +11,11 @@ from lightning.pytorch import seed_everything from ._cli import OPTION_PROMPT_KWARGS, cli -from .utils import configure_logging, sync_start_iteration_xmls +from .utils import ( + configure_logging, + resolve_iteration_apply_ctf, + sync_start_iteration_xmls, +) from .alignment import run_alignment_parallel from .prepare_stacks import prepare_stacks_parallel from .preprocessing import run_cross_correlation_alignment_parallel @@ -137,11 +141,13 @@ def infer_miss_align( for x in range(start_iter, end_iter): iteration_settings = general_config["iteration_settings"][x] alignment_mode = iteration_settings["alignment"] + apply_ctf = resolve_iteration_apply_ctf(iteration_settings, general_config) model_checkpoint = model_paths[x] print(f"\n{'=' * 60}") print(f"Iteration {x + 1}/{end_iter} - Alignment: {alignment_mode}") print(f" model: {model_checkpoint}") + print(f" apply CTF: {apply_ctf}") print(f"{'=' * 60}\n") # get list of all files to process for alignment @@ -156,7 +162,7 @@ def infer_miss_align( patch_size=alignment_config["patch_size"], patch_overlap=alignment_config["patch_overlap"], batch_size=alignment_config["batch_size"], - apply_ctf=general_config["apply_ctf"], + apply_ctf=apply_ctf, downsample=iteration_settings["downsample"], devices_list=devices_alignment, ) diff --git a/src/miss_alignment/train.py b/src/miss_alignment/train.py index 08538c1..2c6ccea 100644 --- a/src/miss_alignment/train.py +++ b/src/miss_alignment/train.py @@ -18,7 +18,12 @@ from lightning.pytorch.strategies import DDPStrategy from ._cli import OPTION_PROMPT_KWARGS, cli -from .utils import configure_logging, parse_device_list, sync_start_iteration_xmls +from .utils import ( + configure_logging, + parse_device_list, + resolve_iteration_apply_ctf, + sync_start_iteration_xmls, +) from .data import MissAlignmentDataModule from .data.shift_generation import create_default_generator from .models import MissAlignment, MAEarlyStopping, MAProgressBar @@ -241,6 +246,7 @@ def _training_worker( model = MissAlignment(**model_params) # Each GPU uses the full batch size (effective batch size scales with GPUs) + apply_ctf = resolve_iteration_apply_ctf(iteration_settings, general_config) with MissAlignmentDataModule( training_directory, create_default_generator(**shift_generation_config), @@ -248,7 +254,7 @@ def _training_worker( n_training_devices=len(devices_training), batch_size=data_module_config["batch_size"], patch_size=data_module_config["patch_size"], - apply_ctf=general_config["apply_ctf"], + apply_ctf=apply_ctf, downsample=iteration_settings["downsample"], steps_per_epoch=data_module_config["steps_per_epoch"], pool_size=pool_size, @@ -430,8 +436,10 @@ def train_miss_align( # ============================================================ iteration_settings = general_config["iteration_settings"][x] alignment_mode = iteration_settings["alignment"] + apply_ctf = resolve_iteration_apply_ctf(iteration_settings, general_config) print(f"\n{'=' * 60}") print(f"Iteration {x + 1}/{end_iter} - Alignment: {alignment_mode}") + print(f" apply CTF: {apply_ctf}") print(f"{'=' * 60}\n") # Run the training phase. Multi-GPU spawns one DDP worker per training @@ -482,7 +490,7 @@ def train_miss_align( patch_size=alignment_config["patch_size"], patch_overlap=alignment_config["patch_overlap"], batch_size=alignment_config["batch_size"], - apply_ctf=general_config["apply_ctf"], + apply_ctf=apply_ctf, downsample=iteration_settings["downsample"], devices_list=devices_alignment, ) diff --git a/src/miss_alignment/utils.py b/src/miss_alignment/utils.py index 4ee3c62..ea54242 100644 --- a/src/miss_alignment/utils.py +++ b/src/miss_alignment/utils.py @@ -4,6 +4,7 @@ import os from pathlib import Path from shutil import copyfile +from typing import Any import torch @@ -11,6 +12,27 @@ LOG_LEVEL_ENV_VAR = "MISS_ALIGNMENT_LOG_LEVEL" +def resolve_iteration_apply_ctf( + iteration_settings: dict[str, Any], + general_config: dict[str, Any], +) -> bool: + """Resolve whether CTF is applied for one macro-iteration. + + An ``apply_ctf`` value on the iteration takes precedence over the global + ``general.apply_ctf`` value. The global value remains the fallback so + existing configuration files keep their original behavior. + """ + apply_ctf = iteration_settings.get( + "apply_ctf", general_config.get("apply_ctf", False) + ) + if not isinstance(apply_ctf, bool): + raise TypeError( + "apply_ctf must be a YAML boolean (True/False), " + f"got {apply_ctf!r} ({type(apply_ctf).__name__})" + ) + return apply_ctf + + def configure_logging() -> None: """Configure miss-alignment logging from ``MISS_ALIGNMENT_LOG_LEVEL``. diff --git a/tests/test_infer.py b/tests/test_infer.py index 17aedcd..b9eae2a 100644 --- a/tests/test_infer.py +++ b/tests/test_infer.py @@ -52,7 +52,7 @@ def test_infer_applies_models_per_iteration(tmp_path, monkeypatch): model_run_dir = tmp_path / "run" iteration_settings = [ {"downsample": 3, "alignment": "anchoring"}, - {"downsample": 1, "alignment": [3, 3]}, + {"downsample": 1, "alignment": [3, 3], "apply_ctf": True}, ] _make_models(model_run_dir, len(iteration_settings)) config_path = _write_config(tmp_path, data_dir, model_run_dir, iteration_settings) @@ -66,6 +66,7 @@ def fake_run_alignment_parallel( output_directory, setting, downsample, + apply_ctf, **_, ): calls.append( @@ -73,6 +74,7 @@ def fake_run_alignment_parallel( "model_checkpoint": model_checkpoint, "setting": setting, "downsample": downsample, + "apply_ctf": apply_ctf, } ) # mimic the real writer: produce one loss json per tilt series @@ -97,9 +99,11 @@ def fake_run_alignment_parallel( assert calls[0]["model_checkpoint"] == str(model_run_dir / "iter1" / "model.ckpt") assert calls[0]["setting"] == "anchoring" assert calls[0]["downsample"] == 3 + assert calls[0]["apply_ctf"] is False assert calls[1]["model_checkpoint"] == str(model_run_dir / "iter2" / "model.ckpt") assert calls[1]["setting"] == [3, 3] assert calls[1]["downsample"] == 1 + assert calls[1]["apply_ctf"] is True # snapshots created per iteration, with provenance recorded for i in range(1, len(iteration_settings) + 1): diff --git a/tests/test_train.py b/tests/test_train.py index 7306bb0..2d49501 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -24,7 +24,11 @@ _resolve_start_checkpoint, _set_ddp_env, ) -from miss_alignment.utils import is_rank_zero, sync_start_iteration_xmls +from miss_alignment.utils import ( + is_rank_zero, + resolve_iteration_apply_ctf, + sync_start_iteration_xmls, +) class _TinyDataset(Dataset): @@ -205,3 +209,36 @@ def test_sync_start_iteration_xmls_resume_missing_raises(tmp_path): """Resuming at an iteration with no snapshot directory is an error.""" with pytest.raises(FileNotFoundError, match="Cannot resume at iteration 3"): sync_start_iteration_xmls(3, tmp_path) + + +def test_iteration_apply_ctf_uses_global_fallback(): + assert ( + resolve_iteration_apply_ctf( + {"downsample": 1, "alignment": "global"}, + {"apply_ctf": True}, + ) + is True + ) + + +@pytest.mark.parametrize("iteration_value", [True, False]) +def test_iteration_apply_ctf_overrides_global(iteration_value): + assert ( + resolve_iteration_apply_ctf( + { + "downsample": 1, + "alignment": [3, 3], + "apply_ctf": iteration_value, + }, + {"apply_ctf": not iteration_value}, + ) + is iteration_value + ) + + +def test_iteration_apply_ctf_rejects_non_boolean(): + with pytest.raises(TypeError, match="YAML boolean"): + resolve_iteration_apply_ctf( + {"apply_ctf": "True"}, + {"apply_ctf": False}, + )