Support per-iteration CTF correction

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:

docs/inference_config_template.yaml

@@ -4,7 +4,9 @@ general:
   # <model_run_directory>/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:

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
 

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,

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()

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,
         )

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,14 +246,15 @@ 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),
         reconstruction_accelerators=devices_reconstruction,
         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,
         )

src/miss_alignment/utils.py

@@ -4,13 +4,35 @@
 import os
 from pathlib import Path
 from shutil import copyfile
+from typing import Any
 
 import torch
 
 # Env var controlling miss-alignment's own log verbosity (DEBUG/INFO/WARNING/...).
 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``.
 

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,13 +66,15 @@ def fake_run_alignment_parallel(
         output_directory,
         setting,
         downsample,
+        apply_ctf,
         **_,
     ):
         calls.append(
             {
                 "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):

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},
+        )