Coverage for / builds / tms-localization / papers / tmsloc_proto / scripts / 08_calc_opt_coil_pos_online.py: 80%
210 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-01 07:48 +0000
1#!/usr/bin/env python
3"""
4Determine optimal coil position/orientation from localization experiment.
5"""
7import os
8import re
9import sys
10import h5py
11import time
12import pynibs
13import pathlib
14import simnibs
15import warnings
16import shutil
17import argparse
18import numpy as np
19from simnibs import opt_struct, mesh_io
20from simnibs.utils.nnav import localite
21from pynibs.neuron.neuron_regression import calc_e_effective
23loc = localite()
24write_tms_navigator_im = loc.write
26# set up command line argument parser
27parser = argparse.ArgumentParser()
28parser.add_argument('-s', '--fn_subject', help='Path to *.hdf5-file of subject', required=True)
29parser.add_argument('-m', '--mesh_idx', help='Mesh ID', required=True, type=str)
30parser.add_argument('-r', '--roi_idx', help='ROI idx',
31 required=False, type=str, default="midlayer_m1s1pmd")
32parser.add_argument('-e', '--exp_idx', help='Exp idx',
33 required=True, type=str)
34parser.add_argument('-a', '--anisotropy_type',
35 help='Anisotropy "vn" or "scalar"', type=str, default="vn")
36parser.add_argument('-t', '--target', nargs='+', required=True,
37 help='R2 results file or x/y/z coordinates')
38parser.add_argument('--target_radius',
39 required=False,
40 help='Radius around target where e-field is averaged.',
41 default=1)
42parser.add_argument('-q', '--qoi',
43 help='Electric field component ("E_mag", "E_norm", "E_tan")', required=False,
44 default='E_mag')
45parser.add_argument('-c', '--fn_coil',
46 help='Path to TMS coil', required=True)
47parser.add_argument('--opt_search_radius',
48 help='Optimization search radius', default=20, type=float)
49parser.add_argument('--opt_search_angle',
50 help='Optimization search angle', default=180, type=float)
51parser.add_argument('--opt_angle_resolution',
52 help='Optimization angle resolution', default=7.5, type=float)
53parser.add_argument('--opt_spatial_resolution',
54 help='Optimization spatial resolution', default=3, type=float)
55parser.add_argument('--opt_smooth', help='Surface smoothing', default=100, type=float)
56parser.add_argument('--opt_distance', help='Skin-coil distance', default=1, type=float)
57parser.add_argument('--layerid',
58 help='Perform mapping based on neuronal meanfield model on the specified layer. ("L23", "L5")',
59 required=False)
60parser.add_argument('--neuronmodel', help='Select neuron model for localization. '
61 '("sensitivity_weighting", "threshold_subtract", "threshold_binary",'
62 ' "IOcurve")',
63 required=False, default=None)
64parser.add_argument('--waveform', help='Waveform of TMS pulse. ("monophasic", "biphasic")',
65 required=False, default='biphasic')
66parser.add_argument('--cosine_model', dest='cosine_model', action='store_true',
67 help='only consider the normal component of E-field')
68parser.description = 'Determine optimal coil position/orientation for cortical target.\n'
69args = parser.parse_args()
70np.set_printoptions(suppress=True)
72scripts_folder = pathlib.Path(__file__).parent.absolute()
74# print simulation info
75# ================================================================================
76print("-" * 64)
77print(f"{parser.description}")
78args_dict = vars(args)
80print('Parameters:')
81for key in args_dict.keys():
82 print(f"{key: >15}: {args_dict[key]}")
83print("-" * 64)
85fn_subject = os.path.abspath(args.fn_subject)
86mesh_idx = args.mesh_idx
87roi_idx = args.roi_idx
88target = args.target
89exp_idx = args.exp_idx
90target_radius = float(args.target_radius)
91anisotropy_type = args.anisotropy_type
92qoi = args.qoi
93fn_coil = os.path.abspath(args.fn_coil)
94opt_search_radius = args.opt_search_radius
95opt_search_angle = args.opt_search_angle
96opt_angle_resolution = args.opt_angle_resolution
97opt_spatial_resolution = args.opt_spatial_resolution
98opt_smooth = args.opt_smooth
99opt_distance = args.opt_distance
100layerid = args.layerid
101neuronmodel = args.neuronmodel
102waveform = args.waveform
103cosine_model = args.cosine_model
105if len(target) > 1:
106 target = np.array(target).astype(float)
107 print(f"Defining user defined target: {target}")
108else:
109 target = target[0]
111# determine target coordinates from r2 map
112if type(target) is str:
113 print(f"Reading target from file: {target}")
114 # read r2 data
115 fn_r2_data = target
116 fn_r2_geo = os.path.splitext(fn_r2_data)[0][:-4] + "geo.hdf5"
118 # find element with maximum goodness-of-fit
119 with h5py.File(fn_r2_data, 'r') as gof_res:
120 best_elm_id = np.argmax(gof_res[f'data/tris/c_{qoi}'][:])
122 # get location of best element on cortex
123 with h5py.File(fn_r2_geo, 'r') as gof_geo:
124 nodes_ids = gof_geo['/mesh/elm/triangle_number_list'][best_elm_id]
125 target = np.mean(gof_geo['/mesh/nodes/node_coord'][:][nodes_ids], axis=0)
127# Subject information
128# ========================================================================================================
129subject_dir = os.path.split(fn_subject)[0]
130subject = pynibs.load_subject(fn_subject)
132# file containing the head model (.msh format)
133fn_mesh_msh = subject.mesh[mesh_idx]["fn_mesh_msh"]
135assert os.path.exists(fn_coil), f"Coil file '{fn_coil}' is missing."
136assert os.path.exists(fn_mesh_msh), f"Mesh file '{fn_coil}' is missing."
137if args_dict["neuronmodel"]:
138 if layerid == 'L23':
139 optim_dir = os.path.join(subject_dir,
140 f"opt",
141 f"target_{np.array2string(target, formatter={'float_kind': '{0:+0.2f}'.format})}"
142 f"_neuron_L23",
143 f"mesh_{mesh_idx}",
144 f"{qoi}")
145 elif layerid == 'L5':
146 optim_dir = os.path.join(subject_dir,
147 f"opt",
148 f"target_{np.array2string(target, formatter={'float_kind': '{0:+0.2f}'.format})}"
149 f"_neuron_L5",
150 f"mesh_{mesh_idx}",
151 f"{qoi}")
153elif args_dict["cosine_model"]:
154 optim_dir = os.path.join(subject_dir,
155 f"opt",
156 f"target_{np.array2string(target, formatter={'float_kind': '{0:+0.2f}'.format})}"
157 f"_norm",
158 f"mesh_{mesh_idx}",
159 f"{qoi}")
161else:
162 optim_dir = os.path.join(subject_dir,
163 f"opt",
164 f"target_{np.array2string(target, formatter={'float_kind': '{0:+0.2f}'.format})}",
165 f"mesh_{mesh_idx}",
166 f"{qoi}")
168if not os.path.exists(optim_dir):
169 os.makedirs(optim_dir)
170fn_e = os.path.join(optim_dir, f"e_scaled.hdf5")
172print(f'\ttarget: {target}')
173print(f'\ttarget_radius: {target_radius}')
174print(f'\theadmesh: {fn_mesh_msh}')
175print(f'\toptim_dir: {optim_dir}')
177fn_conform_nii = os.path.join(subject.mesh[mesh_idx]["mesh_folder"], subject.mesh[mesh_idx]["fn_mri_conform"])
178try:
179 fn_exp_nii = subject.exp[exp_idx]["fn_mri_nii"][0][0]
180except KeyError:
181 print(f"Experiment {exp_idx} not found.")
182 fn_exp_nii = fn_conform_nii
184# determine target region
185# ========================================================================================================
186roi_pynibs = pynibs.load_roi_surface_obj_from_hdf5(subject.mesh[mesh_idx]['fn_mesh_hdf5'])[roi_idx]
188if args_dict["neuronmodel"]:
189 for layer in roi_pynibs.layers:
190 if layer.id == layerid:
191 nodes = layer.surface.nodes.node_coord
192 con = layer.surface.elm.node_number_list[:, :3] - np.min(layer.surface.elm.node_number_list[:, :3])
193 break
194 else:
195 raise ValueError(f"Layer {layerid} was not found in roi {roi_idx}.")
196 tri_center_layer = np.mean(nodes[con], axis=1)
197 target_mask = np.linalg.norm(tri_center_layer - target, axis=1) < target_radius
198 target_mask_idx = np.where(target_mask)[0]
199else:
200 target_mask = np.linalg.norm(roi_pynibs.tri_center_coord_mid - target, axis=1) < target_radius
201 target_mask_idx = np.where(target_mask)[0]
203if not target_mask.any():
204 raise ValueError("TARGET: No ROI triangle found withing target_radius!")
206# Initialize structure
207# ========================================================================================================
208tms_opt = opt_struct.TMSoptimize()
209tms_opt.fnamehead = fn_mesh_msh
210tms_opt.pathfem = optim_dir
211tms_opt.fnamecoil = fn_coil
212tms_opt.target = target
213tms_opt.search_radius = opt_search_radius
214tms_opt.search_angle = opt_search_angle
215tms_opt.angle_resolution = opt_angle_resolution
216tms_opt.spatial_resolution = opt_spatial_resolution
217tms_opt.smooth = opt_smooth
218tms_opt.distance = opt_distance
219tms_opt.target_size = target_radius
220mesh = mesh_io.read_msh(fn_mesh_msh)
222# Run optimization
223# ========================================================================================================
224if not (os.path.exists(f'{optim_dir}{os.sep}opt_coil_pos.xml') and
225 os.path.exists(f'{optim_dir}{os.sep}opt_coil_pos.hdf5')):
226 # determine all possible coil positions
227 tms_opt._prepare()
228 pos_matrices = tms_opt._get_coil_positions()
229 fn_coilpos_hdf5 = os.path.join(optim_dir, "search_positions.hdf5")
230 pynibs.coil.write_coil_pos_hdf5(fn_hdf=os.path.join(optim_dir, "search_positions.hdf5"),
231 centers=np.vstack([p[0:3, 3] for p in pos_matrices]),
232 m0=np.vstack([p[0:3, 0] for p in pos_matrices]),
233 m1=np.vstack([p[0:3, 1] for p in pos_matrices]),
234 m2=np.vstack([p[0:3, 2] for p in pos_matrices]),
235 datanames=None, data=None, overwrite=True)
236 pos_matrices = np.array(pos_matrices)
237 pos_matrices = np.swapaxes(pos_matrices, 0, 2)
239 # run bruteforce e-field simulations
240 cmd = f'{sys.executable} {os.path.join(scripts_folder, "run_simnibs_online.py")} '
241 cmd += f" --folder '{optim_dir}' "
242 cmd += f" --fn_subject '{fn_subject}' "
243 cmd += f" --fn_coilpos '{fn_coilpos_hdf5}' "
244 cmd += f"--anisotropy_type {anisotropy_type} "
245 cmd += f"-m '{mesh_idx}' "
246 cmd += f"--fn_coil '{fn_coil}' "
247 cmd += f"-r {roi_idx} "
248 cmd += f"--hdf5_fn '{fn_e}' "
249 cmd += " --calc_theta_grad" if args_dict["neuronmodel"] else ""
251 print(f"Running {cmd}")
252 os.system(cmd)
254 # load e-fields (and theta, gradient)
255 with h5py.File(fn_e, "r") as f:
256 if args_dict["neuronmodel"]:
257 datanames = "E_eff"
258 e_qoi = f[f"/{layerid}/{qoi}"][:, target_mask_idx]
259 theta = f[f"/{layerid}/E_theta"][:, target_mask_idx]
260 gradient = f[f"/{layerid}/E_gradient"][:, target_mask_idx]
262 # determine effective e-field in case of neuron optimization
263 e_matrix = calc_e_effective(e=e_qoi,
264 layerid=layerid,
265 theta=theta,
266 gradient=gradient,
267 neuronmodel=neuronmodel,
268 mep=None,
269 waveform=waveform)
270 # use the normal component E if cosine_model selected
271 elif args_dict["cosine_model"]:
272 datanames = "abs_E_norm"
273 e_norm = np.abs(f["E_norm"][:, target_mask_idx])
274 e_matrix = e_norm
275 else:
276 datanames = "E_mag"
277 e_qoi = f[qoi][:, target_mask_idx]
278 e_matrix = e_qoi
280 # determine best coil position
281 idx_best_coil_pos = np.argmax(np.mean(e_matrix, axis=1))
282 opt_matsim = pos_matrices[:, :, idx_best_coil_pos]
283 print("Transposing matsimnibs -- why is this needed?")
284 opt_matsim = np.transpose(opt_matsim)
285 # save best coil position
286 with h5py.File(os.path.join(optim_dir, 'opt_coil_pos.hdf5'), "w") as f:
287 f.create_dataset(name="matsimnibs", data=opt_matsim)
289 if len(opt_matsim.shape) == 2:
290 opt_matsim = opt_matsim[:, :, np.newaxis]
292 matlocalite = pynibs.expio.nnav2simnibs(fn_exp_nii=fn_exp_nii,
293 fn_conform_nii=fn_conform_nii,
294 m_nnav=opt_matsim,
295 target='nnav',
296 fsl_cmd='FSL',
297 nnav_system='localite')
299 # write instrument marker file
300 write_tms_navigator_im(matlocalite, os.path.join(optim_dir, 'opt_coil_pos.xml'), names='opt', overwrite=True)
301 pynibs.coil.create_stimsite_from_matsimnibs(os.path.join(optim_dir, 'opt_coil_pos.hdf5'),
302 opt_matsim, overwrite=True)
304else:
305 print(f"Optimal coordinates already exist in:")
306 print(f" > {os.path.join(optim_dir, 'opt_coil_pos.xml')}")
307 print(f" > {os.path.join(optim_dir, 'opt_coil_pos.hdf5')}")
308 print(f"Skipping optimization and reading optimal coordinates...")
310 with h5py.File(os.path.join(optim_dir, 'opt_coil_pos.hdf5'), "r") as f:
311 opt_matsim = f["matsimnibs"][:]
313# save all coil positions together with electric field values in .xmdf file for plotting
314fn_search_positions = os.path.join(optim_dir, "search_positions.hdf5")
315fn_e = os.path.join(optim_dir, 'e_scaled.hdf5')
317with h5py.File(fn_search_positions, "r") as f:
318 centers = f["centers"][:]
319 m0 = f["m0"][:]
320 m1 = f["m1"][:]
321 m2 = f["m2"][:]
322shutil.copy(fn_search_positions, fn_search_positions.replace('.hdf5', f'_{time.time()}.hdf5'))
324# get e-fields at target for each coil position
325with h5py.File(fn_e, "r") as f:
326 if not 'e_matrix' in globals():
327 if args_dict["neuronmodel"]:
328 datanames = "E_eff"
329 e_qoi = f[f"/{layerid}/{qoi}"][:, target_mask_idx]
330 theta = f[f"/{layerid}/E_theta"][:, target_mask_idx]
331 gradient = f[f"/{layerid}/E_gradient"][:, target_mask_idx]
333 # determine effective e-field in case of neuron optimization
334 e_matrix = calc_e_effective(e=e_qoi,
335 layerid=layerid,
336 theta=theta,
337 gradient=gradient,
338 neuronmodel=neuronmodel,
339 mep=None,
340 waveform=waveform)
341 # use the normal component E if cosine_model selected
342 elif args_dict["cosine_model"]:
343 datanames = "abs_E_norm"
344 e_norm = np.abs(f["E_norm"][:, target_mask_idx])
345 e_matrix = e_norm
346 else:
347 datanames = "E_mag"
348 e_qoi = f[qoi][:, target_mask_idx]
349 e_matrix = e_qoi
350 e_target = np.mean(e_matrix, axis=1)[:, np.newaxis]
352pynibs.coil.write_coil_pos_hdf5(fn_hdf=fn_search_positions,
353 centers=centers,
354 m0=m0,
355 m1=m1,
356 m2=m2,
357 datanames=datanames,
358 data=e_target,
359 overwrite=True)
361# run final simulation at optimal coil position
362# ========================================================================================================
364# Setting up SimNIBS SESSION
365# ========================================================================================================
366# print("Setting up SimNIBS SESSION")
367# S = simnibs.sim_struct.SESSION()
368# S.fnamehead = fn_mesh_msh
369# S.pathfem = os.path.join(optim_dir, "e_opt")
370# S.fields = "vDeE"
371# S.open_in_gmsh = False
372# S.fname_tensor = None
373# S.map_to_surf = True
374# S.map_to_fsavg = False
375# S.map_to_MNI = False
376# S.subpath = os.path.join(subject.mesh[mesh_idx]["mesh_folder"], f"m2m_{subject.id}")
377#
378# if not os.path.exists(S.pathfem):
379# os.makedirs(S.pathfem)
380#
381# # Define the TMS simulation and setting up conductivities
382# # ========================================================================================================
383# tms = [S.add_tmslist()]
384# tms[0].fnamecoil = fn_coil
385# tms[0].cond[0].value = 0.126 # WM
386# tms[0].cond[1].value = 0.275 # GM
387# tms[0].cond[2].value = 1.654 # CSF
388# tms[0].cond[3].value = 0.01 # Skull
389# tms[0].cond[4].value = 0.465 # Scalp
390# tms[0].anisotropy_type = anisotropy_type
391#
392# if subject.mesh[mesh_idx]['fn_tensor_vn'] is not None:
393# tms[0].fn_tensor_nifti = os.path.join(subject.mesh[mesh_idx]["mesh_folder"],
394# subject.mesh[mesh_idx]['fn_tensor_vn'])
395# else:
396# tms[0].fn_tensor_nifti = None
397#
398# tms[0].excentricity_scale = None
399#
400# # Define the optimal coil position
401# # ========================================================================================================
402# pos = tms[0].add_position()
403# pos.matsimnibs = opt_matsim[:, :, 0]
404# pos.distance = 0.1
405# pos.didt = 1e6 # coil current rate of change in A/s (1e6 A/s = 1 A/us)
406# pos.postprocess = S.fields
407#
408# # Running simulations
409# # ========================================================================================================
410# print("Running SimNIBS")
411# filenames = simnibs.run_simnibs(S, cpus=1)
412#
413# print("Saving results in .hdf5 format.")
414# if type(filenames) is not list:
415# filenames = [filenames]
416#
417# pynibs.simnibs_results_msh2hdf5(fn_msh=filenames,
418# fn_hdf5=[os.path.join(S.pathfem, "e")],
419# S=S,
420# pos_tms_idx=[0], # indices of different "coils"
421# pos_local_idx=[0], # indices of associated positions
422# subject=subject,
423# mesh_idx=mesh_idx,
424# overwrite=True,
425# mid2roi=True,
426# verbose=True)
427print("=== DONE ===")