Skills › Scientific Agent Skills › relion › spa_workflow.py
spa_workflow.py
260 lines · skills/relion/scripts/spa_workflow.py at 5f2bc03
1#!/usr/bin/env python3
2"""Validate RELION SPA inputs, run refinement/postprocessing, and inspect half-map FSC."""
3from __future__ import annotations
5import argparse
6import csv
7import json
8import math
9from pathlib import Path
10import re
11import subprocess
14def read_map(path: Path, max_voxels: int | None = None):
15 import mrcfile
16 import numpy as np
17 with mrcfile.mmap(path, permissive=False) as handle:
18 if np.iscomplexobj(handle.data):
19 raise ValueError(f'{path}: expected a real-space density map, not complex MRC data')
20 if max_voxels is not None and handle.data.size > max_voxels:
21 raise ValueError("Python diagnostic limited to 256-cubed maps; use relion_image_handler --fsc for larger maps")
22 data = handle.data.copy()
23 voxel = tuple(float(handle.voxel_size[axis]) for axis in ('x', 'y', 'z'))
24 origin = tuple(float(handle.header.origin[axis]) for axis in ('x', 'y', 'z'))
25 axes = tuple(int(handle.header[axis]) for axis in ('mapc', 'mapr', 'maps'))
26 starts = tuple(int(handle.header[axis]) for axis in ('nxstart', 'nystart', 'nzstart'))
27 angles = tuple(float(handle.header.cellb[axis]) for axis in ('alpha', 'beta', 'gamma'))
28 if data.ndim != 3 or len(set(data.shape)) != 1 or min(data.shape) < 4 or data.shape[0] % 2:
29 raise ValueError(f'{path}: expected an even-sized cubic 3D map')
30 if not np.isfinite(data).all() or not all(math.isfinite(x) for x in voxel + origin):
31 raise ValueError(f'{path}: nonfinite map or header')
32 if min(voxel) <= 0 or not np.allclose(voxel, voxel[0], rtol=1e-5):
33 raise ValueError(f'{path}: positive isotropic voxel size required')
34 if axes != (1, 2, 3):
35 raise ValueError(f'{path}: noncanonical map axes; explicitly reorder before analysis')
36 if starts != (0, 0, 0) or not np.allclose(angles, 90, rtol=0, atol=1e-4):
37 raise ValueError(f'{path}: this runner requires zero MRC start indices and orthogonal cell angles; verify the grid before conversion')
38 return data, voxel[0], origin
41def validate_star(path: Path, project: Path, check_stacks: bool = True) -> dict:
42 import mrcfile
43 import numpy as np
44 import starfile
45 blocks = starfile.read(path, always_dict=True)
46 if not {'optics', 'particles'}.issubset(blocks):
47 raise ValueError('Requires RELION 3.1+ data_optics and data_particles blocks')
48 optics, particles = blocks['optics'], blocks['particles']
49 if not hasattr(optics, 'columns') or not hasattr(particles, 'columns'):
50 raise ValueError('Optics and particles must be STAR loop tables')
51 needed_optics = {'rlnOpticsGroup', 'rlnVoltage', 'rlnSphericalAberration', 'rlnAmplitudeContrast',
52 'rlnImagePixelSize', 'rlnImageSize', 'rlnImageDimensionality'}
53 needed_particles = {'rlnOpticsGroup', 'rlnImageName', 'rlnDefocusU', 'rlnDefocusV', 'rlnDefocusAngle'}
54 if not needed_optics.issubset(optics.columns) or not needed_particles.issubset(particles.columns):
55 raise ValueError('Missing optics/acquisition/CTF columns for a 2D-particle SPA refinement')
56 if len(optics) == 0 or len(particles) == 0 or optics.rlnOpticsGroup.duplicated().any():
57 raise ValueError('Empty tables or duplicate optics group IDs')
58 optional_numeric = {'rlnOriginXAngst', 'rlnOriginYAngst', 'rlnAngleRot', 'rlnAngleTilt',
59 'rlnAnglePsi', 'rlnCoordinateX', 'rlnCoordinateY', 'rlnPhaseShift'}
60 for frame, columns in [(optics, needed_optics),
61 (particles, (needed_particles - {'rlnImageName'}) | (optional_numeric & set(particles.columns)))]:
62 for column in columns:
63 values = np.asarray(frame[column], dtype=float)
64 if not np.isfinite(values).all():
65 raise ValueError(f'Nonfinite values in {column}')
66 if any(optics.rlnOpticsGroup <= 0) or any(optics.rlnOpticsGroup % 1):
67 raise ValueError('Optics group IDs must be positive integers')
68 if not set(particles.rlnOpticsGroup).issubset(set(optics.rlnOpticsGroup)):
69 raise ValueError('Particle references an undefined optics group')
70 if any(optics.rlnImagePixelSize <= 0) or any(optics.rlnVoltage <= 0) or any(optics.rlnSphericalAberration < 0):
71 raise ValueError('Pixel size/voltage must be positive and spherical aberration nonnegative')
72 if any((optics.rlnAmplitudeContrast < 0) | (optics.rlnAmplitudeContrast >= 1)):
73 raise ValueError('Amplitude contrast must be in [0,1)')
74 if any(optics.rlnImageDimensionality != 2) or any(optics.rlnImageSize < 4) or any(optics.rlnImageSize % 2):
75 raise ValueError('This runner requires 2D particles with an even box size >=4')
76 if any(particles.rlnDefocusU <= 0) or any(particles.rlnDefocusV <= 0):
77 raise ValueError('Underfocus values must be positive Angstroms; verify CTF import units/sign')
78 if particles.rlnImageName.duplicated().any():
79 raise ValueError('Duplicate particle image references can leak between half sets')
80 subsets = None
81 if 'rlnRandomSubset' in particles:
82 values = set(particles.rlnRandomSubset)
83 if values != {1, 2}:
84 raise ValueError('Existing rlnRandomSubset must contain both half sets, encoded 1 and 2')
85 subsets = {str(k): int(v) for k, v in particles.rlnRandomSubset.value_counts().items()}
86 groups = optics.set_index('rlnOpticsGroup')
87 stack_cache, canonical_images = {}, set()
88 for row in particles.itertuples(index=False):
89 match = re.fullmatch(r'([0-9]+)@(.+)', str(row.rlnImageName))
90 if not match or int(match[1]) < 1:
91 raise ValueError(f'Expected 1-based index@stack.mrcs: {row.rlnImageName}')
92 index, stack_name = int(match[1]), match[2]
93 stack = Path(stack_name)
94 stack = (stack if stack.is_absolute() else project / stack).resolve()
95 identity = (index, stack)
96 if identity in canonical_images:
97 raise ValueError('Duplicate particle image after path normalization')
98 canonical_images.add(identity)
99 if check_stacks:
100 if stack not in stack_cache:
101 with mrcfile.mmap(stack, permissive=False) as handle:
102 if np.iscomplexobj(handle.data):
103 raise ValueError(f'{stack}: expected real-space particle images')
104 if tuple(int(handle.header[a]) for a in ('mapc', 'mapr', 'maps')) != (1, 2, 3):
105 raise ValueError(f'{stack}: noncanonical particle stack axes')
106 shape = handle.data.shape
107 # MRC may expose a single image as 2D rather than a stack of length one.
108 stack_cache[stack] = ((1, *shape) if len(shape) == 2 else shape)
109 shape = stack_cache[stack]
110 box = int(groups.loc[row.rlnOpticsGroup, 'rlnImageSize'])
111 if len(shape) != 3 or index > shape[0] or shape[1:] != (box, box):
112 raise ValueError(f'Particle stack index/box does not agree with optics: {row.rlnImageName}')
113 warnings = []
114 if any((optics.rlnVoltage < 60) | (optics.rlnVoltage > 400)):
115 warnings.append('Unusual voltage; confirm kV rather than V')
116 if any(particles.rlnDefocusU > 100000) or any(particles.rlnDefocusV > 100000):
117 warnings.append('Defocus exceeds 10 micrometers; inspect CTF fits and units')
118 return {'particles': len(particles), 'optics_groups': len(optics), 'half_sets': subsets,
119 'stack_checks_performed': check_stacks,
120 'pixel_sizes_angstrom': sorted(set(float(x) for x in optics.rlnImagePixelSize)),
121 'box_sizes': sorted(set(int(x) for x in optics.rlnImageSize)), 'warnings': warnings}
124def check_halves(first: Path, second: Path, mask: Path | None = None, max_voxels: int | None = None):
125 import numpy as np
126 a, pixel, origin = read_map(first, max_voxels)
127 b, other_pixel, other_origin = read_map(second, max_voxels)
128 if a.shape != b.shape or not np.isclose(pixel, other_pixel) or not np.allclose(origin, other_origin):
129 raise ValueError('Half maps disagree in grid, pixel size, or origin')
130 if np.array_equal(a, b):
131 raise ValueError('Half maps are identical; provide independent unfiltered half reconstructions')
132 if a.min() == a.max() or b.min() == b.max():
133 raise ValueError('Constant half map has no measurable signal')
134 if mask:
135 m, mp, mo = read_map(mask)
136 if m.shape != a.shape or not np.isclose(pixel, mp) or not np.allclose(origin, mo):
137 raise ValueError('Mask grid differs from half maps')
138 if m.min() < 0 or m.max() > 1 or m.max() == m.min():
139 raise ValueError('Mask must be nonconstant and lie in [0,1]')
140 if not np.any((m > 0) & (m < 1)):
141 raise ValueError('Mask has no soft edge; generate a soft solvent mask')
142 return a, b, pixel
145def fsc(first: Path, second: Path, output: Path) -> dict:
146 import numpy as np
147 a, b, pixel = check_halves(first, second, max_voxels=256**3)
148 n = a.shape[0]
149 fa, fb = np.fft.rfftn(a), np.fft.rfftn(b)
150 z, y, x = np.ogrid[-n//2:n-n//2, -n//2:n-n//2, 0:n//2+1]
151 shells = np.floor(np.sqrt(z*z+y*y+x*x)+0.5).astype(int)
152 # Shift only the fully represented axes; weight interior rFFT x planes twice.
153 fa, fb = np.fft.fftshift(fa, axes=(0,1)), np.fft.fftshift(fb, axes=(0,1))
154 weights = np.ones(fa.shape); weights[:,:,1:-1] = 2
155 cross = np.bincount(shells.ravel(), weights=(weights*(fa*np.conj(fb)).real).ravel())
156 power_a = np.bincount(shells.ravel(), weights=(weights*abs(fa)**2).ravel())
157 power_b = np.bincount(shells.ravel(), weights=(weights*abs(fb)**2).ravel())
158 counts = np.bincount(shells.ravel(), weights=weights.ravel())
159 denominator = np.sqrt(power_a*power_b)
160 values = np.divide(cross, denominator, out=np.full_like(cross, np.nan), where=denominator>0)
161 rows = [{'shell': k, 'frequency_inv_angstrom': k/(n*pixel),
162 'fsc': float(values[k]) if np.isfinite(values[k]) else None,
163 'fourier_samples': int(counts[k])} for k in range(n//2+1)]
164 crossing = None
165 for previous, current in zip(rows[1:], rows[2:]):
166 a_fsc, b_fsc = previous['fsc'], current['fsc']
167 if a_fsc is not None and b_fsc is not None and a_fsc >= 0.143 > b_fsc:
168 frequency = previous['frequency_inv_angstrom'] + (0.143-a_fsc)/(b_fsc-a_fsc)/(n*pixel)
169 crossing = 1/frequency
170 break
171 output.parent.mkdir(parents=True, exist_ok=True)
172 with output.open('w', newline='') as handle:
173 writer = csv.DictWriter(handle, fieldnames=list(rows[0]), delimiter='\t')
174 writer.writeheader(); writer.writerows(rows)
175 return {'unmasked_diagnostic_fsc_0143_angstrom': crossing, 'nyquist_angstrom': 2*pixel,
176 'warning': 'Diagnostic unmasked FSC; use corrected RELION postprocessing FSC for reported global resolution',
177 'shells': len(rows)}
180def run_native_command(command: list[str], output: Path, project: Path, executable: str) -> None:
181 if output.exists() and any(output.iterdir()):
182 raise ValueError('Use an empty output directory; preserve prior jobs for restart/provenance')
183 version = subprocess.run([executable, '--version'], capture_output=True, text=True, check=True)
184 version_text = version.stdout + version.stderr
185 if not re.search(r'(?<![\w.])5\.0\.1(?![\w.])', version_text):
186 raise ValueError('This runner targets RELION5.0.1; verify the installed executable before adapting it')
187 output.mkdir(parents=True, exist_ok=True)
188 (output/'version.txt').write_text(version_text)
189 (output/'command.json').write_text(json.dumps(command, indent=2)+'\n')
190 with (output/'run.log').open('w') as log:
191 subprocess.run(command, cwd=project, stdout=log, stderr=subprocess.STDOUT, check=True)
194def main() -> None:
195 parser = argparse.ArgumentParser(description=__doc__)
196 sub = parser.add_subparsers(dest='action', required=True)
197 star = sub.add_parser('validate-star')
198 star.add_argument('star', type=Path); star.add_argument('--project', type=Path, default=Path.cwd())
199 star.add_argument('--metadata-only', action='store_true', help='Explicitly omit opening particle stacks')
200 comparison = sub.add_parser('fsc')
201 comparison.add_argument('half1', type=Path); comparison.add_argument('half2', type=Path)
202 comparison.add_argument('--output', type=Path, required=True)
203 refine = sub.add_parser('refine')
204 refine.add_argument('--star', type=Path, required=True); refine.add_argument('--reference', type=Path, required=True)
205 refine.add_argument('--diameter', type=float, required=True, help='Particle diameter in Angstroms')
206 refine.add_argument('--symmetry', default='C1'); refine.add_argument('--initial-lowpass', type=float, default=40.0)
207 refine.add_argument('--threads', type=int, default=1); refine.add_argument('--executable', default='relion_refine_mpi')
208 refine.add_argument('--mpirun', default='mpirun'); refine.add_argument('--mpi-ranks', type=int, default=3)
209 refine.add_argument('--seed', type=int, default=1)
210 refine.add_argument('--project', type=Path, default=Path.cwd()); refine.add_argument('--output', type=Path, required=True)
211 post = sub.add_parser('postprocess')
212 post.add_argument('--half1', type=Path, required=True); post.add_argument('--half2', type=Path, required=True)
213 post.add_argument('--mask', type=Path, required=True); post.add_argument('--output', type=Path, required=True)
214 post.add_argument('--executable', default='relion_postprocess')
215 args = parser.parse_args()
216 try:
217 if args.action == 'validate-star':
218 result = validate_star(args.star, args.project, not args.metadata_only)
219 elif args.action == 'fsc':
220 result = fsc(args.half1, args.half2, args.output)
221 elif args.action == 'refine':
222 result = validate_star(args.star, args.project)
223 reference, pixel, _ = read_map(args.reference)
224 if result['box_sizes'] != [reference.shape[0]] or any(abs(p-pixel)>1e-5*pixel for p in result['pixel_sizes_angstrom']):
225 raise ValueError('Reference box/pixel size must match all particle optics groups for this bounded runner')
226 if (not math.isfinite(args.diameter) or not math.isfinite(args.initial_lowpass)
227 or not 0 < args.diameter < reference.shape[0]*pixel or args.initial_lowpass < 2*pixel
228 or args.threads < 1 or args.mpi_ranks < 3 or args.mpi_ranks % 2 == 0
229 or not 0 <= args.seed <= 2147483647):
230 raise ValueError('Invalid diameter, lowpass, threads, seed (0..2147483647) or MPI ranks (use odd ranks >=3)')
231 if not re.fullmatch(r'(?:[CD][1-9][0-9]*|[TO]|I[1-4]?)', args.symmetry):
232 raise ValueError('Use an explicit supported RELION point-group symmetry such as C1, D2, O, I1')
233 command = [args.mpirun, '-np', str(args.mpi_ranks), args.executable, '--o', str(args.output.resolve()/'run'), '--i', str(args.star.resolve()),
234 '--ref', str(args.reference.resolve()), '--auto_refine', '--split_random_halves', '--ctf',
235 '--firstiter_cc', '--ini_high', str(args.initial_lowpass), '--sym', args.symmetry,
236 '--particle_diameter', str(args.diameter), '--flatten_solvent', '--zero_mask',
237 '--oversampling', '1', '--healpix_order', '2', '--auto_local_healpix_order', '4',
238 '--offset_range', '5', '--offset_step', '2', '--j', str(args.threads), '--random_seed', str(args.seed)]
239 run_native_command(command, args.output.resolve(), args.project.resolve(), args.executable)
240 if not all((args.output/f'run_half{k}_class001_unfil.mrc').is_file() for k in (1,2)):
241 raise ValueError('Refinement exited without converged unfiltered half maps; inspect run.log/optimiser')
242 result['refined_particles'] = validate_star(args.output/'run_data.star', args.project)
243 result['runtime_warnings'] = list(dict.fromkeys(line.split('WARNING:', 1)[1].strip()
244 for line in (args.output/'run.log').read_text().splitlines() if 'WARNING:' in line))
245 result['command'] = command
246 else:
247 _, _, pixel = check_halves(args.half1, args.half2, args.mask)
248 command = [args.executable, '--i', str(args.half1.resolve()), '--i2', str(args.half2.resolve()),
249 '--mask', str(args.mask.resolve()), '--angpix', str(pixel), '--o', str(args.output.resolve()/'postprocess')]
250 run_native_command(command, args.output.resolve(), Path.cwd(), args.executable)
251 if not (args.output/'postprocess.star').is_file():
252 raise ValueError('Postprocess did not produce postprocess.star')
253 result = {'command': command, 'pixel_size_angstrom': pixel}
254 except (ValueError, OSError, KeyError, subprocess.CalledProcessError) as exc:
255 parser.exit(1, f'Error: {exc}\n')
256 print(json.dumps(result, indent=2))
259if __name__ == '__main__':
260 main()