lattice-sub 1.1.4__py3-none-any.whl → 1.3.0__py3-none-any.whl

{lattice_sub-1.1.4.dist-info → lattice_sub-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lattice-sub
-Version: 1.1.4
+Version: 1.3.0
 Summary: Lattice subtraction for cryo-EM micrographs - removes periodic crystal signals to reveal non-periodic features
 Author-email: George Stephenson <george.stephenson@colorado.edu>, Vignesh Kasinath <vignesh.kasinath@colorado.edu>
 License: MIT
@@ -94,7 +94,17 @@ lattice-sub batch input_folder/ output_folder/ --pixel-size 0.56
 lattice-sub batch input_folder/ output_folder/ --pixel-size 0.56 --vis comparisons/
 ```
 
-This creates side-by-side PNG images showing before/after/difference for each micrograph.
+This creates 4-panel PNG comparison images for each micrograph showing:
+1. **Original** - Input micrograph
+2. **Subtracted** - Lattice-removed result  
+3. **Difference** - What was removed (5x amplified)
+4. **Threshold Curve** - Threshold vs lattice removal efficacy
+
+**Limit the number of visualizations:**
+```bash
+# Generate visualizations for first 10 images only
+lattice-sub batch input_folder/ output_folder/ -p 0.56 --vis comparisons/ -n 10
+```
 
 ---
 
@@ -105,6 +115,8 @@ This creates side-by-side PNG images showing before/after/difference for each mi
 | `-p, --pixel-size` | **Required.** Pixel size in Ångstroms |
 | `-o, --output` | Output file path (default: `sub_<input>`) |
 | `-t, --threshold` | Peak detection sensitivity (default: **auto** - optimized per image) |
+| `--vis DIR` | Generate 4-panel comparison PNGs in DIR |
+| `-n, --num-vis N` | Limit visualizations to first N images |
 | `--cpu` | Force CPU processing (GPU is used by default) |
 | `-q, --quiet` | Hide the banner and progress messages |
 | `-v, --verbose` | Show detailed processing information |
@@ -172,6 +184,76 @@ python -c "import torch; print(torch.cuda.get_device_name(0) if torch.cuda.is_av
 
 ---
 
+## Multi-GPU Support
+
+When processing batches on systems with multiple GPUs, files are automatically distributed across all available GPUs for faster processing. No extra flags needed!
+
+```bash
+# Automatically uses all available GPUs
+lattice-sub batch input_folder/ output_folder/ -p 0.56
+```
+
+**Example with 2 GPUs and 100 images:**
+- GPU 0: processes images 1-50
+- GPU 1: processes images 51-100
+- Single progress bar shows combined progress
+
+This provides near-linear speedup with additional GPUs.
+
+---
+
+## HPC Example (CU Boulder Alpine)
+
+Using [Open OnDemand Core Desktop](https://curc.readthedocs.io/en/latest/open_ondemand/core_desktop.html) with 2× RTX 8000 GPUs:
+
+```bash
+# Create environment
+module load anaconda
+conda create -n lattice_test python=3.11 -y
+conda activate lattice_test
+pip install lattice-sub
+
+# Process 100 micrographs
+lattice-sub batch input/ output/ -p 0.56
+```
+
+**Output:**
+```
+Phase-preserving FFT inpainting for cryo-EM  |  v1.3.0
+
+  Configuration
+  -------------
+    Pixel size:  0.56 A
+    Threshold:   auto
+    Backend:     Auto → GPU (Quadro RTX 8000)
+
+  Batch Processing
+  ----------------
+    Files:    100
+    Output:   /projects/user/output
+    Workers:  1
+
+✓ Using 2 GPUs: GPU 0, GPU 1
+
+  ✓ GPU 0: Quadro RTX 8000
+  ✓ GPU 1: Quadro RTX 8000
+
+  Processing: 100%|█████████████████████████| 100/100 [05:12<00:00,  3.13s/file]
+
+  [OK] Batch complete (312.9s)
+```
+
+**100 images processed in ~5 minutes** with automatic multi-GPU distribution.
+
+For compute-focused workloads, use Alpine's [GPU partitions](https://curc.readthedocs.io/en/latest/clusters/alpine/alpine-hardware.html) (A100, L40, MI100):
+
+```bash
+# Request 3 GPUs for 1 hour
+sinteractive --partition=aa100 --gres=gpu:3 --ntasks=16 --nodes=1 --time=01:00:00 --qos=normal
+```
+
+---
+
 ## Python API
 
 ```python

lattice_sub-1.3.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+lattice_sub-1.3.0.dist-info/licenses/LICENSE,sha256=2kPoH0cbEp0cVEGqMpyF2IQX1npxdtQmWJB__HIRSb0,1101
+lattice_subtraction/__init__.py,sha256=TNaJXvSgCQdvYUYJfS5scn92YjORiGfLot9WadZ8u28,1737
+lattice_subtraction/batch.py,sha256=zJzvUnr8dznvxE8jaPKDLJ7AcJg8Cbfv5nVo0FzZz1I,20891
+lattice_subtraction/cli.py,sha256=W99XQClUMKaaFQxle0W-ILQ6UuYRFXZVJWD4qXpcIj4,24063
+lattice_subtraction/config.py,sha256=uzwKb5Zi3phHUk2ZgoiLsQdwFdN-rTiY8n02U91SObc,8426
+lattice_subtraction/core.py,sha256=VzcecSZHRuBuHUc2jHGv8LalINL75RH0aTpABI708y8,16265
+lattice_subtraction/io.py,sha256=uHku6rJ0jeCph7w-gOIDJx-xpNoF6PZcLfb5TBTOiw0,4594
+lattice_subtraction/masks.py,sha256=HIamrACmbQDkaCV4kXhnjMDSwIig4OtQFLig9A8PMO8,11741
+lattice_subtraction/processing.py,sha256=tmnj5K4Z9HCQhRpJ-iMd9Bj_uTRuvDEWyUenh8MCWEM,8341
+lattice_subtraction/threshold_optimizer.py,sha256=yEsGM_zt6YjgEulEZqtRy113xOFB69aHJIETm2xSS6k,15398
+lattice_subtraction/ui.py,sha256=Sp_a-yNmBRZJxll8h9T_H5-_KsI13zGYmHcbcpVpbR8,9176
+lattice_subtraction/visualization.py,sha256=hWFz49NBBrS7d6ofO0VyJ6-v8Q6hPG1dijbDtecMOQs,11890
+lattice_sub-1.3.0.dist-info/METADATA,sha256=pKwt8TcftbZGm1gvWZGO1n3iQiI4JB3E_ix3InB-4D0,14901
+lattice_sub-1.3.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+lattice_sub-1.3.0.dist-info/entry_points.txt,sha256=o8PzJR8kFnXlKZufoYGBIHpiosM-P4PZeKZXJjtPS6Y,61
+lattice_sub-1.3.0.dist-info/top_level.txt,sha256=BOuW-sm4G-fQtsWPRdeLzWn0WS8sDYVNKIMj5I3JXew,20
+lattice_sub-1.3.0.dist-info/RECORD,,

{lattice_sub-1.1.4.dist-info → lattice_sub-1.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

lattice_subtraction/__init__.py

@@ -19,7 +19,7 @@ Example:
     >>> result.save("output.mrc")
 """
 
-__version__ = "1.1.4"
+__version__ = "1.3.0"
 __author__ = "George Stephenson & Vignesh Kasinath"
 
 from .config import Config

lattice_subtraction/batch.py

@@ -1,12 +1,14 @@
 """
 Batch processing for multiple micrographs.
 
-This module provides parallel processing capabilities for large datasets.
+This module provides parallel processing capabilities for large datasets,
+including automatic multi-GPU support for systems with multiple CUDA devices.
 """
 
 import os
+import multiprocessing as mp
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 from pathlib import Path
 from typing import List, Tuple, Optional, Callable
 import logging
@@ -63,6 +65,91 @@ def _process_single_file(args: tuple) -> Tuple[Path, Optional[str]]:
         return (Path(input_path), str(e))
 
 
+def _gpu_worker(
+    gpu_id: int,
+    file_pairs: List[Tuple[str, str]],
+    config_dict: dict,
+    progress_queue: mp.Queue,
+    error_queue: mp.Queue,
+):
+    """
+    Worker function for multi-GPU processing.
+    
+    Each worker processes its assigned files on a specific GPU and reports
+    progress through a shared queue.
+    
+    Args:
+        gpu_id: CUDA device ID to use
+        file_pairs: List of (input_path, output_path) tuples
+        config_dict: Configuration dictionary
+        progress_queue: Queue to report progress (sends 1 for each completed file)
+        error_queue: Queue to report errors (sends (gpu_id, file_path, error_msg))
+    """
+    import torch
+    
+    # Set this process to use the specific GPU
+    torch.cuda.set_device(gpu_id)
+    
+    # Reconstruct config with the specific device_id and quiet mode
+    config_dict = config_dict.copy()
+    config_dict['device_id'] = gpu_id
+    config_dict['_quiet'] = True  # Suppress messages - main process handles this
+    config = Config(**config_dict)
+    
+    # Create subtractor (messages suppressed via _quiet flag)
+    subtractor = LatticeSubtractor(config)
+    
+    for input_path, output_path in file_pairs:
+        try:
+            result = subtractor.process(input_path)
+            result.save(output_path, pixel_size=config.pixel_ang)
+            progress_queue.put(1)
+        except Exception as e:
+            error_queue.put((gpu_id, input_path, str(e)))
+            return  # Fail-fast: exit on first error
+
+
+def _check_gpu_memory(device_id: int, image_shape: Tuple[int, int]) -> Tuple[bool, str]:
+    """
+    Check if GPU has sufficient memory for processing.
+    
+    Args:
+        device_id: CUDA device ID
+        image_shape: (height, width) of image
+        
+    Returns:
+        (is_ok, message) - True if sufficient memory, False with warning message
+    """
+    try:
+        import torch
+        free_mem, total_mem = torch.cuda.mem_get_info(device_id)
+        
+        # Estimate memory needed: image + FFT (complex) + masks + overhead
+        # Roughly 16x image size for safe margin (complex FFT, intermediate buffers)
+        image_bytes = image_shape[0] * image_shape[1] * 4  # float32
+        estimated_need = image_bytes * 16
+        
+        if free_mem < estimated_need:
+            return False, (
+                f"GPU {device_id}: {free_mem / 1e9:.1f}GB free, "
+                f"need ~{estimated_need / 1e9:.1f}GB"
+            )
+        return True, ""
+    except Exception as e:
+        return True, ""  # If we can't check, proceed anyway
+
+
+def _get_available_gpus() -> List[int]:
+    """Get list of available CUDA GPU device IDs."""
+    try:
+        import torch
+        if torch.cuda.is_available():
+            return list(range(torch.cuda.device_count()))
+        return []
+    except ImportError:
+        return []
+
+
 class BatchProcessor:
     """
     Parallel batch processor for micrograph datasets.
@@ -157,6 +244,9 @@ class BatchProcessor:
         """
         Process a list of input/output file pairs.
         
+        Automatically uses multi-GPU processing when multiple GPUs are available.
+        Files are distributed evenly across GPUs in chunks.
+        
         Args:
             file_pairs: List of (input_path, output_path) tuples
             show_progress: If True, show progress bar
@@ -168,8 +258,7 @@ class BatchProcessor:
         successful = 0
         failed_files = []
         
-        # Check if using GPU - if so, process sequentially to avoid CUDA fork issues
-        # With "auto" backend, check if PyTorch + CUDA is actually available
+        # Check if using GPU - if so, check for multi-GPU capability
         use_gpu = self.config.backend == "pytorch"
         if self.config.backend == "auto":
             try:
@@ -179,10 +268,19 @@ class BatchProcessor:
                 use_gpu = False
         
         if use_gpu:
-            # Sequential processing for GPU (CUDA doesn't support fork multiprocessing)
-            successful, failed_files = self._process_sequential(
-                file_pairs, show_progress
-            )
+            # Check how many GPUs are available
+            available_gpus = _get_available_gpus()
+            
+            if len(available_gpus) > 1 and total > 1:
+                # Multi-GPU processing
+                successful, failed_files = self._process_multi_gpu(
+                    file_pairs, available_gpus, show_progress
+                )
+            else:
+                # Single GPU - sequential processing
+                successful, failed_files = self._process_sequential(
+                    file_pairs, show_progress
+                )
         else:
             # Parallel processing for CPU
             successful, failed_files = self._process_parallel(
@@ -284,6 +382,159 @@ class BatchProcessor:
         
         return successful, failed_files
     
+    def _process_multi_gpu(
+        self,
+        file_pairs: List[Tuple[Path, Path]],
+        gpu_ids: List[int],
+        show_progress: bool = True,
+    ) -> Tuple[int, List[Tuple[Path, str]]]:
+        """
+        Process files in parallel across multiple GPUs.
+        
+        Files are distributed evenly across GPUs in chunks.
+        Uses spawn-based multiprocessing to avoid CUDA fork issues.
+        
+        Args:
+            file_pairs: List of (input_path, output_path) tuples
+            gpu_ids: List of CUDA device IDs to use
+            show_progress: If True, show unified progress bar
+            
+        Returns:
+            (successful_count, failed_files_list)
+        """
+        import time
+        
+        total = len(file_pairs)
+        num_gpus = len(gpu_ids)
+        
+        # Print multi-GPU info with GPU names
+        try:
+            import torch
+            gpu_names = [torch.cuda.get_device_name(i) for i in gpu_ids]
+            print(f"✓ Using {num_gpus} GPUs: {', '.join(f'GPU {i}' for i in gpu_ids)}")
+            print("")
+            for i, name in zip(gpu_ids, gpu_names):
+                print(f"  ✓ GPU {i}: {name}")
+        except Exception:
+            print(f"✓ Using {num_gpus} GPUs")
+        
+        # Check GPU memory on first GPU (assume similar for all)
+        if file_pairs:
+            try:
+                sample_image = read_mrc(file_pairs[0][0])
+                is_ok, msg = _check_gpu_memory(gpu_ids[0], sample_image.shape)
+                if not is_ok:
+                    print(f"⚠ Memory warning: {msg}")
+            except Exception:
+                pass  # Proceed anyway
+        
+        # Distribute files evenly across GPUs (chunked distribution)
+        chunk_size = (total + num_gpus - 1) // num_gpus  # Ceiling division
+        gpu_file_assignments = []
+        
+        for i, gpu_id in enumerate(gpu_ids):
+            start_idx = i * chunk_size
+            end_idx = min(start_idx + chunk_size, total)
+            if start_idx < total:
+                chunk = [(str(inp), str(out)) for inp, out in file_pairs[start_idx:end_idx]]
+                gpu_file_assignments.append((gpu_id, chunk))
+        
+        # Create shared queues for progress and errors
+        # Use 'spawn' context to avoid CUDA fork issues
+        ctx = mp.get_context('spawn')
+        progress_queue = ctx.Queue()
+        error_queue = ctx.Queue()
+        
+        # Create progress bar (after all GPU info printed)
+        if show_progress:
+            print()  # Blank line for visual separation
+            pbar = tqdm(
+                total=total,
+                desc="  Processing",
+                unit="file",
+                ncols=80,
+                leave=True,
+            )
+        else:
+            pbar = None
+        
+        # Start worker processes
+        processes = []
+        for gpu_id, file_chunk in gpu_file_assignments:
+            p = ctx.Process(
+                target=_gpu_worker,
+                args=(gpu_id, file_chunk, self._config_dict, progress_queue, error_queue),
+            )
+            p.start()
+            processes.append(p)
+        
+        # Monitor progress and check for errors
+        successful = 0
+        failed_files = []
+        completed = 0
+        
+        while completed < total:
+            # Check for progress updates (non-blocking with timeout)
+            try:
+                while True:
+                    progress_queue.get(timeout=0.1)
+                    successful += 1
+                    completed += 1
+                    if pbar:
+                        pbar.update(1)
+            except:
+                pass  # Queue empty, continue
+            
+            # Check for errors (non-blocking)
+            try:
+                while True:
+                    gpu_id, file_path, error_msg = error_queue.get_nowait()
+                    failed_files.append((Path(file_path), error_msg))
+                    completed += 1
+                    if pbar:
+                        pbar.update(1)
+                    
+                    # Fail-fast: terminate all workers and report
+                    print(f"\n✗ GPU {gpu_id} failed on {Path(file_path).name}: {error_msg}")
+                    print(f"\nTip: Try a different configuration:")
+                    print(f"  lattice-sub batch <input> <output> -p {self.config.pixel_ang} --cpu -j 8")
+                    
+                    # Terminate all processes
+                    for p in processes:
+                        if p.is_alive():
+                            p.terminate()
+                    
+                    if pbar:
+                        pbar.close()
+                    
+                    return successful, failed_files
+            except:
+                pass  # No errors, continue
+            
+            # Check if all processes have finished
+            all_done = all(not p.is_alive() for p in processes)
+            if all_done:
+                # Drain remaining queue items
+                try:
+                    while True:
+                        progress_queue.get_nowait()
+                        successful += 1
+                        completed += 1
+                        if pbar:
+                            pbar.update(1)
+                except:
+                    pass
+                break
+        
+        # Wait for all processes to finish
+        for p in processes:
+            p.join(timeout=1.0)
+        
+        if pbar:
+            pbar.close()
+        
+        return successful, failed_files
+
     def process_numbered_sequence(
         self,
         input_pattern: str,

lattice_subtraction/cli.py

@@ -153,9 +153,12 @@ def main():
         # Batch process directory (GPU handles parallelism)
         lattice-sub batch input_dir/ output_dir/ --pixel-size 0.56
         
-        # Batch with visualizations
+        # Batch with visualizations (4-panel with threshold curve)
         lattice-sub batch input_dir/ output_dir/ -p 0.56 --vis viz_dir/
         
+        # Batch with limited visualizations (only first 10)
+        lattice-sub batch input_dir/ output_dir/ -p 0.56 --vis viz_dir/ -n 10
+        
         # CPU batch with parallel workers (use -j only with --cpu)
         lattice-sub batch input_dir/ output_dir/ -p 0.56 --cpu -j 8
         
@@ -502,6 +505,12 @@ def process(
     default=None,
     help="Generate comparison visualizations in this directory",
 )
+@click.option(
+    "-n", "--num-vis",
+    type=int,
+    default=None,
+    help="Number of visualizations to generate (default: all)",
+)
 @click.option(
     "-v", "--verbose",
     is_flag=True,
@@ -528,6 +537,7 @@ def batch(
     config: Optional[str],
     recursive: bool,
     vis: Optional[str],
+    num_vis: Optional[int],
     verbose: bool,
     quiet: bool,
     cpu: bool,
@@ -617,7 +627,8 @@ def batch(
     
     # Generate visualizations if requested
     if vis:
-        logger.info(f"Generating visualizations in: {vis}")
+        limit_msg = f" (limit: {num_vis})" if num_vis else ""
+        logger.info(f"Generating visualizations in: {vis}{limit_msg}")
         viz_success, viz_total = generate_visualizations(
             input_dir=input_dir,
             output_dir=output_dir,
@@ -625,6 +636,8 @@ def batch(
             prefix=prefix,
             pattern=pattern,
             show_progress=True,
+            limit=num_vis,
+            config=cfg,
         )
         logger.info(f"Visualizations: {viz_success}/{viz_total} created")
 

lattice_subtraction/config.py

@@ -64,6 +64,13 @@ class Config:
     # Enabled by default when GPU is available
     use_kornia: bool = True
     
+    # GPU device ID for multi-GPU support. None = auto-select (GPU 0 for single-GPU mode)
+    # When using multi-GPU batch processing, this is set automatically per worker
+    device_id: Optional[int] = None
+    
+    # Internal flag to suppress status messages (used by batch workers)
+    _quiet: bool = False
+    
     def __post_init__(self):
         """Validate and set auto-calculated parameters."""
         if self.pixel_ang <= 0:

lattice_subtraction/core.py

@@ -83,21 +83,31 @@ class LatticeSubtractor:
         
         Auto mode tries PyTorch+CUDA first, then PyTorch CPU, then NumPy.
         Prints user-friendly status message about which backend is active.
+        
+        Uses config.device_id if specified for multi-GPU support.
         """
         backend = self.config.backend
         self._gpu_message_shown = getattr(self, '_gpu_message_shown', False)
         
+        # Check if quiet mode (suppress messages for batch workers)
+        quiet = getattr(self.config, '_quiet', False)
+        if quiet:
+            self._gpu_message_shown = True
+        
+        # Get device ID from config (None means auto-select GPU 0)
+        device_id = self.config.device_id if self.config.device_id is not None else 0
+        
         # Auto mode: try GPU first, then CPU
         if backend == "auto":
             try:
                 import torch
                 if torch.cuda.is_available():
-                    self.device = torch.device('cuda')
+                    self.device = torch.device(f'cuda:{device_id}')
                     self.use_gpu = True
                     # Only print once per session (batch processing reuses subtractor)
                     if not self._gpu_message_shown:
-                        gpu_name = torch.cuda.get_device_name(0)
-                        print(f"✓ Using GPU: {gpu_name}")
+                        gpu_name = torch.cuda.get_device_name(device_id)
+                        print(f"✓ Using GPU {device_id}: {gpu_name}")
                         self._gpu_message_shown = True
                 else:
                     self.device = torch.device('cpu')
@@ -116,7 +126,7 @@ class LatticeSubtractor:
             try:
                 import torch
                 if torch.cuda.is_available():
-                    self.device = torch.device('cuda')
+                    self.device = torch.device(f'cuda:{device_id}')
                     self.use_gpu = True
                 else:
                     import warnings

lattice_subtraction/visualization.py

@@ -2,12 +2,12 @@
 Visualization utilities for lattice subtraction results.
 
 This module provides functions to create comparison visualizations
-showing original, processed, and difference images.
+showing original, processed, difference images, and threshold optimization curves.
 """
 
 import logging
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 
 import numpy as np
 
@@ -22,6 +22,56 @@ logging.getLogger('PIL.PngImagePlugin').setLevel(logging.WARNING)
 logger = logging.getLogger(__name__)
 
 
+def compute_threshold_curve(
+    image: np.ndarray,
+    config,
+    n_points: int = 21,
+) -> Tuple[np.ndarray, np.ndarray, float, float]:
+    """
+    Compute quality scores across a range of thresholds.
+    
+    Args:
+        image: Original image array
+        config: Config object with processing parameters
+        n_points: Number of threshold points to evaluate
+        
+    Returns:
+        Tuple of (thresholds, quality_scores, optimal_threshold, optimal_quality)
+    """
+    from .threshold_optimizer import ThresholdOptimizer
+    
+    optimizer = ThresholdOptimizer(config)
+    
+    # Prepare FFT data once
+    subtracted, radial_mask, box_size = optimizer._prepare_fft_data(image)
+    
+    # Evaluate across threshold range
+    thresholds = np.linspace(
+        optimizer.min_threshold, 
+        optimizer.max_threshold, 
+        n_points
+    )
+    
+    # Use GPU batch if available
+    if optimizer.use_gpu:
+        qualities, peak_counts = optimizer._compute_quality_batch_gpu(
+            subtracted, radial_mask, thresholds
+        )
+    else:
+        qualities = []
+        for t in thresholds:
+            q, _ = optimizer._compute_quality(subtracted, radial_mask, t)
+            qualities.append(q)
+        qualities = np.array(qualities)
+    
+    # Find optimal
+    best_idx = np.argmax(qualities)
+    optimal_threshold = thresholds[best_idx]
+    optimal_quality = qualities[best_idx]
+    
+    return thresholds, qualities, optimal_threshold, optimal_quality
+
+
 def create_comparison_figure(
     original: np.ndarray,
     processed: np.ndarray,
@@ -81,19 +131,106 @@ def create_comparison_figure(
     return fig
 
 
+def create_comparison_figure_with_threshold(
+    original: np.ndarray,
+    processed: np.ndarray,
+    thresholds: np.ndarray,
+    quality_scores: np.ndarray,
+    optimal_threshold: float,
+    optimal_quality: float,
+    title: str = "Lattice Subtraction Comparison",
+    figsize: Tuple[int, int] = (24, 6),
+    dpi: int = 150,
+):
+    """
+    Create a 4-panel comparison figure with threshold optimization curve.
+    
+    Layout: [Original] [Subtracted] [Difference] [Threshold vs Quality]
+    
+    Args:
+        original: Original image array
+        processed: Processed (lattice-subtracted) image array
+        thresholds: Array of threshold values tested
+        quality_scores: Array of quality scores for each threshold
+        optimal_threshold: The optimal threshold that was selected
+        optimal_quality: Quality score at optimal threshold
+        title: Figure title
+        figsize: Figure size in inches (width, height)
+        dpi: Resolution for saving
+        
+    Returns:
+        matplotlib Figure object
+    """
+    import matplotlib.pyplot as plt
+    
+    # Compute difference
+    difference = original - processed
+    
+    # Create figure with 4 panels (1 row, 4 columns)
+    fig, axes = plt.subplots(1, 4, figsize=figsize)
+    
+    # Contrast limits from original
+    vmin, vmax = np.percentile(original, [1, 99])
+    
+    # Panel 1: Original
+    axes[0].imshow(original, cmap='gray', vmin=vmin, vmax=vmax)
+    axes[0].set_title(f'Original\n{original.shape}')
+    axes[0].axis('off')
+    
+    # Panel 2: Lattice Subtracted
+    axes[1].imshow(processed, cmap='gray', vmin=vmin, vmax=vmax)
+    axes[1].set_title(f'Lattice Subtracted\n{processed.shape}')
+    axes[1].axis('off')
+    
+    # Panel 3: Difference (removed lattice)
+    diff_std = np.std(difference)
+    axes[2].imshow(
+        difference, 
+        cmap='RdBu_r',
+        vmin=-diff_std * 3, 
+        vmax=diff_std * 3
+    )
+    axes[2].set_title('Difference (Removed Lattice)')
+    axes[2].axis('off')
+    
+    # Panel 4: Threshold vs Quality Score curve
+    axes[3].plot(thresholds, quality_scores, 'b-', linewidth=2, label='Quality Score')
+    axes[3].axvline(x=optimal_threshold, color='r', linestyle='--', linewidth=2, 
+                    label=f'Optimal: {optimal_threshold:.3f}')
+    axes[3].scatter([optimal_threshold], [optimal_quality], color='r', s=100, zorder=5)
+    axes[3].set_xlabel('Threshold', fontsize=11)
+    axes[3].set_ylabel('Lattice Removal Efficacy', fontsize=11)
+    axes[3].set_title(f'Threshold Optimization\nOptimal = {optimal_threshold:.3f}')
+    axes[3].legend(loc='best', fontsize=9)
+    axes[3].grid(True, alpha=0.3)
+    axes[3].set_xlim(thresholds.min(), thresholds.max())
+    
+    # Title
+    plt.suptitle(title, fontsize=14)
+    plt.tight_layout()
+    
+    return fig
+
+
 def save_comparison_visualization(
     original_path: Path,
     processed_path: Path,
     output_path: Path,
+    config = None,
     dpi: int = 150,
 ) -> None:
     """
-    Create and save a comparison visualization for a single image pair.
+    Create and save a 4-panel comparison visualization for a single image pair.
+    
+    Includes threshold optimization curve showing how the optimal threshold
+    was selected based on lattice removal efficacy.
     
     Args:
         original_path: Path to original MRC file
         processed_path: Path to processed MRC file
         output_path: Path for output PNG file
+        config: Config object for threshold computation (optional, will use defaults)
+        dpi: Resolution for output images
     """
     import matplotlib.pyplot as plt
     import mrcfile
@@ -106,11 +243,32 @@ def save_comparison_visualization(
     
     # Create title
     name = original_path.name
-    short_name = name[:60] + "..." if len(name) > 60 else name
-    title = f"Lattice Subtraction Comparison: {short_name}"
+    short_name = name[:50] + "..." if len(name) > 50 else name
+    title = f"Lattice Subtraction: {short_name}"
+    
+    # Try to compute threshold curve if config available
+    try:
+        if config is None:
+            # Create default config for threshold computation
+            from .config import Config
+            config = Config(pixel_ang=0.56)  # Default K3 pixel size
+        
+        # Compute threshold optimization curve
+        thresholds, quality_scores, optimal_threshold, optimal_quality = \
+            compute_threshold_curve(original, config)
+        
+        # Create 4-panel figure with threshold curve
+        fig = create_comparison_figure_with_threshold(
+            original, processed,
+            thresholds, quality_scores,
+            optimal_threshold, optimal_quality,
+            title=title, dpi=dpi
+        )
+    except Exception as e:
+        # Fallback to 3-panel if threshold computation fails
+        logger.debug(f"Could not compute threshold curve: {e}, using 3-panel view")
+        fig = create_comparison_figure(original, processed, title=title, dpi=dpi)
     
-    # Create and save figure
-    fig = create_comparison_figure(original, processed, title=title, dpi=dpi)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     fig.savefig(output_path, dpi=dpi, bbox_inches='tight')
     plt.close(fig)
@@ -124,9 +282,11 @@ def generate_visualizations(
     pattern: str = "*.mrc",
     dpi: int = 150,
     show_progress: bool = True,
+    limit: Optional[int] = None,
+    config = None,
 ) -> Tuple[int, int]:
     """
-    Generate comparison visualizations for all processed images in a directory.
+    Generate comparison visualizations for processed images in a directory.
     
     Args:
         input_dir: Directory containing original MRC files
@@ -136,6 +296,8 @@ def generate_visualizations(
         pattern: Glob pattern for finding processed files
         dpi: Resolution for output images
         show_progress: Show progress bar
+        limit: Maximum number of visualizations to generate (None = all)
+        config: Config object for threshold computation (optional)
         
     Returns:
         Tuple of (successful_count, total_count)
@@ -152,6 +314,12 @@ def generate_visualizations(
         logger.warning(f"No processed files found matching '{prefix}{pattern}' in {output_dir}")
         return 0, 0
     
+    # Apply limit if specified
+    total_available = len(output_files)
+    if limit is not None and limit > 0:
+        output_files = output_files[:limit]
+        logger.info(f"Limiting to {limit} visualizations (of {total_available} available)")
+    
     successful = 0
     total = len(output_files)
     
@@ -184,11 +352,12 @@ def generate_visualizations(
                 successful += 1
                 continue
             
-            # Generate visualization
+            # Generate visualization with 4-panel layout
             save_comparison_visualization(
                 original_path=input_path,
                 processed_path=processed_path,
                 output_path=viz_path,
+                config=config,
                 dpi=dpi,
             )
             successful += 1

lattice_sub-1.1.4.dist-info/RECORD

@@ -1,17 +0,0 @@
-lattice_sub-1.1.4.dist-info/licenses/LICENSE,sha256=2kPoH0cbEp0cVEGqMpyF2IQX1npxdtQmWJB__HIRSb0,1101
-lattice_subtraction/__init__.py,sha256=zLLC1T_bGZMDpZTcywjloIsV1vvIhsBbj_P9h7YFPqw,1737
-lattice_subtraction/batch.py,sha256=sTDWEL5FlEx2HFaJsTZRXyzLQoNCgUqRo900eZ6kq68,12005
-lattice_subtraction/cli.py,sha256=fWZ3ueahQXPQ_ApzS345AdslsMF0-6HIAIvan2IFHkk,23597
-lattice_subtraction/config.py,sha256=dh8EJFzJEEXwwggQ46rBMHsuVOExQWM-kCfonT94_fE,8111
-lattice_subtraction/core.py,sha256=QzE5CLv92XPoyuw8JcMAGIeSEVgfwSkHgG86WlHAjMo,15790
-lattice_subtraction/io.py,sha256=uHku6rJ0jeCph7w-gOIDJx-xpNoF6PZcLfb5TBTOiw0,4594
-lattice_subtraction/masks.py,sha256=HIamrACmbQDkaCV4kXhnjMDSwIig4OtQFLig9A8PMO8,11741
-lattice_subtraction/processing.py,sha256=tmnj5K4Z9HCQhRpJ-iMd9Bj_uTRuvDEWyUenh8MCWEM,8341
-lattice_subtraction/threshold_optimizer.py,sha256=yEsGM_zt6YjgEulEZqtRy113xOFB69aHJIETm2xSS6k,15398
-lattice_subtraction/ui.py,sha256=Sp_a-yNmBRZJxll8h9T_H5-_KsI13zGYmHcbcpVpbR8,9176
-lattice_subtraction/visualization.py,sha256=pMZKcz6Xgs98lLaZbvGjoMIyEYA_MLRracVxpQStC3w,5935
-lattice_sub-1.1.4.dist-info/METADATA,sha256=hPnPdJSHoP4DL_2R-frH5982ftSw4ggVpxiZsFcoNRM,12581
-lattice_sub-1.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lattice_sub-1.1.4.dist-info/entry_points.txt,sha256=o8PzJR8kFnXlKZufoYGBIHpiosM-P4PZeKZXJjtPS6Y,61
-lattice_sub-1.1.4.dist-info/top_level.txt,sha256=BOuW-sm4G-fQtsWPRdeLzWn0WS8sDYVNKIMj5I3JXew,20
-lattice_sub-1.1.4.dist-info/RECORD,,