Coverage for /builds/tms-localization/papers/tmsloc_proto/scripts/08_calc_opt_coil_pos.py: 87%
200 statements
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-12 08:31 +0000
1#!/usr/bin/env python
3"""
4Determine optimal coil position/orientation from localization experiment.
5"""
7import os
8import re
9import h5py
10import pynibs
11import simnibs
12import warnings
13import argparse
14import numpy as np
15from simnibs import opt_struct, mesh_io
16if int(simnibs.__version__[0]) < 4:
17 from simnibs.utils.nnav import write_tms_navigator_im
18else:
19 from simnibs.utils.nnav import localite
20 loc = localite()
21 write_tms_navigator_im = loc.write
23# set up command line argument parser
24parser = argparse.ArgumentParser()
25parser.add_argument('-s', '--fn_subject', help='Path to *.hdf5-file of subject', required=True)
26parser.add_argument('-m', '--mesh_idx', help='Mesh ID', required=True, type=str)
27parser.add_argument('-e', '--exp_idx', help='Experiment ID', required=True, type=str)
28parser.add_argument('-n', '--n_cpu', help='How many cpus to use', type=int, default=4)
29parser.add_argument('-a', '--anisotropy_type', help='Anisotropy "vn" or "scalar"', type=str, default="vn")
30parser.add_argument('-p', '--pardiso', action='store_true', help="Use pardiso solver")
31parser.add_argument('-t', '--target', nargs='+', required=True, help='R2 results file or x/y/z coordinates')
32parser.add_argument('-q', '--qoi', help='Electric field component ("E_mag", "E_norm", "E_tan")', required=False,
33 default='E_mag')
34parser.add_argument('-c', '--fn_coil', help='Path to TMS coil', required=True)
35parser.add_argument('--opt_search_radius', help='Optimization search radius', default=20, type=float)
36parser.add_argument('--opt_search_angle', help='Optimization search angle', default=180, type=float)
37parser.add_argument('--opt_angle_resolution', help='Optimization angle resolution', default=7.5, type=float)
38parser.add_argument('--opt_spatial_resolution', help='Optimization spatial resolution', default=2, type=float)
39parser.add_argument('--opt_smooth', help='Surface smoothing', default=100, type=float)
40parser.add_argument('--opt_distance', help='Skin-coil distance', default=1, type=float)
41parser.description = 'Determine optimal coil position/orientation for cortical target.\n'
42args = parser.parse_args()
44# print simulation info
45# ================================================================================
46print("-" * 64)
47print(f"{parser.description}")
48args_dict = vars(args)
50print('Parameters:')
51for key in args_dict.keys():
52 print(f"{key: >15}: {args_dict[key]}")
53print("-" * 64)
55fn_subject = os.path.abspath(args.fn_subject)
56mesh_idx = args.mesh_idx
57exp_idx = args.exp_idx
58target = args.target
59n_cpu = args.n_cpu
60anisotropy_type = args.anisotropy_type
61pardiso = args.pardiso
62qoi = args.qoi
63fn_coil = os.path.abspath(args.fn_coil)
64opt_search_radius = args.opt_search_radius
65opt_search_angle = args.opt_search_angle
66opt_angle_resolution = args.opt_angle_resolution
67opt_spatial_resolution = args.opt_spatial_resolution
68opt_smooth = args.opt_smooth
69opt_distance = args.opt_distance
71if len(target) > 1:
72 target = np.array(target).astype(float)
73 print(f"Defining user defined target: {target}")
74else:
75 target = target[0]
77# determine target coordinates from r2 map
78if type(target) is str:
79 print(f"Reading target from file: {target}")
80 # read r2 data
81 fn_r2_data = target
82 fn_r2_geo = os.path.splitext(fn_r2_data)[0][:-4] + "geo.hdf5"
84 # find element with maximum goodness-of-fit
85 with h5py.File(fn_r2_data, 'r') as gof_res:
86 best_elm_id = np.argmax(gof_res[f'data/tris/c_{qoi}'][:])
88 # get location of best element on cortex
89 with h5py.File(fn_r2_geo, 'r') as gof_geo:
90 nodes_ids = gof_geo['/mesh/elm/triangle_number_list'][best_elm_id]
91 target = np.mean(gof_geo['/mesh/nodes/node_coord'][:][nodes_ids], axis=0)
93# Subject information
94# ========================================================================================================
95subject_dir = os.path.split(fn_subject)[0]
96subject = pynibs.load_subject(fn_subject)
98# file containing the head model (.msh format)
99fn_mesh_msh = subject.mesh[mesh_idx]["fn_mesh_msh"]
101assert os.path.exists(fn_coil), f"Coil file '{fn_coil}' is missing."
102assert os.path.exists(fn_mesh_msh), f"Mesh file '{fn_mesh_msh}' is missing."
104optim_dir = os.path.join(subject_dir,
105 f"opt",
106 f"exp_{exp_idx}",
107 f"target_{np.array2string(target, formatter={'float_kind': '{0:+0.2f}'.format})}",
108 f"mesh_{mesh_idx}",
109 f"{qoi}")
110if not os.path.exists(optim_dir):
111 os.makedirs(optim_dir)
113print(f'\ttarget: {target}')
114print(f'\theadmesh: {fn_mesh_msh}')
115print(f'\toptim_dir: {optim_dir}')
117fn_conform_nii = os.path.join(subject.mesh[mesh_idx]["mesh_folder"], subject.mesh[mesh_idx]["fn_mri_conform"])
118try:
119 fn_exp_nii = subject.exp[exp_idx]["fn_mri_nii"][0][0]
120except KeyError:
121 print(f"Experiment {exp_idx} not found.")
122 fn_exp_nii = fn_conform_nii
124# Initialize structure
125# ========================================================================================================
126tms_opt = opt_struct.TMSoptimize()
127tms_opt.fnamehead = fn_mesh_msh
128tms_opt.pathfem = optim_dir
129tms_opt.fnamecoil = fn_coil
130tms_opt.target = target
131tms_opt.anisotropy_type = anisotropy_type
133if subject.mesh[mesh_idx]['fn_tensor_vn'] is not None:
134 tms_opt.fname_tensor = os.path.join(subject.mesh[mesh_idx]["mesh_folder"],
135 subject.mesh[mesh_idx]['fn_tensor_vn'])
136else:
137 tms_opt.fname_tensor = None
139if pardiso:
140 tms_opt.solver_options = 'pardiso' # 'pardiso' # faster, lots of RAM
141else:
142 tms_opt.solver_options = None
143tms_opt.open_in_gmsh = False
145tms_opt.search_radius = opt_search_radius
146tms_opt.search_angle = opt_search_angle
147tms_opt.angle_resolution = opt_angle_resolution
148tms_opt.spatial_resolution = opt_spatial_resolution
149tms_opt.smooth = opt_smooth
150tms_opt.distance = opt_distance
152# optimization target area in mm
153tms_opt.target_size = .5
154mesh = mesh_io.read_msh(fn_mesh_msh)
156# check if the target is within GM
157# ========================================================================================================
158bar = mesh.elements_baricenters()[:]
159dist = np.linalg.norm(bar - target, axis=1)
160elm = mesh.elm.elm_number[
161 (dist < tms_opt.target_size) *
162 np.isin(mesh.elm.tag1, [2]) *
163 np.isin(mesh.elm.elm_type, [4])
164 ]
165if not len(elm):
166 warnings.warn(f"\tNo elements found at {target}. Changing target.")
167 sub_elms = bar[np.isin(mesh.elm.tag1, [2]) * np.isin(mesh.elm.elm_type, [4])]
168 new_sub_coords = sub_elms[np.where(
169 np.linalg.norm(sub_elms - target, axis=1)
170 ==
171 np.min(np.linalg.norm(sub_elms - target, axis=1)))[0][0]]
173 tms_opt.target = new_sub_coords
174 print(f'\tcoords: {new_sub_coords}')
176# Run optimization
177# ========================================================================================================
178if not (os.path.exists(f'{optim_dir}{os.sep}opt_coil_pos.xml') and
179 os.path.exists(f'{optim_dir}{os.sep}opt_coil_pos.hdf5')):
180 tms_opt._prepare()
181 pos_matrices = tms_opt._get_coil_positions()
182 tms_opt.pos_ydir = []
183 tms_opt.centre = []
184 pynibs.write_coil_pos_hdf5(fn_hdf=os.path.join(optim_dir, "search_positions.hdf5"),
185 centers=np.vstack([p[0:3, 3] for p in pos_matrices]),
186 m0=np.vstack([p[0:3, 0] for p in pos_matrices]),
187 m1=np.vstack([p[0:3, 1] for p in pos_matrices]),
188 m2=np.vstack([p[0:3, 2] for p in pos_matrices]),
189 datanames=None, data=None, overwrite=True)
191 opt_matsim = np.squeeze(tms_opt.run(allow_multiple_runs=True, cpus=n_cpu))
193 if len(opt_matsim.shape) == 2:
194 opt_matsim = opt_matsim[:, :, np.newaxis]
196 matlocalite = pynibs.nnav2simnibs(fn_exp_nii=fn_exp_nii,
197 fn_conform_nii=fn_conform_nii,
198 m_nnav=opt_matsim,
199 target='nnav',
200 nnav_system='localite')
202 # write instrument marker file
203 write_tms_navigator_im(matlocalite, os.path.join(optim_dir, 'opt_coil_pos.xml'), names='opt', overwrite=True)
204 pynibs.create_stimsite_from_matsimnibs(os.path.join(optim_dir, 'opt_coil_pos.hdf5'), opt_matsim, overwrite=True)
206else:
207 print(f"Optimal coordinates already exist in:")
208 print(f" > {os.path.join(optim_dir, 'opt_coil_pos.xml')}")
209 print(f" > {os.path.join(optim_dir, 'opt_coil_pos.hdf5')}")
210 print(f"Skipping optimization and reading optimal coordinates...")
212 with h5py.File(os.path.join(optim_dir, 'opt_coil_pos.hdf5'), "r") as f:
213 opt_matsim = f["matsimnibs"][:]
215# save all coil positions together with electric field values in .xmdf file for plotting
216fn_search_positions = os.path.join(optim_dir, "search_positions.hdf5")
217fn_e = os.path.join(optim_dir, subject.id + "_TMS_optimize_" + os.path.split(fn_coil)[1].replace('.nii.gz', '_nii') +
218 ".hdf5")
220with h5py.File(fn_search_positions, "r") as f:
221 centers = f["centers"][:]
222 m0 = f["m0"][:]
223 m1 = f["m1"][:]
224 m2 = f["m2"][:]
226os.unlink(fn_search_positions)
228# get e-fields at target for each coil position
229if os.path.exists(fn_e):
230 with h5py.File(fn_e, "r") as f:
231 try:
232 e = f["tms_optimization/E_norm"][:] # SimNIBS 3
233 except KeyError:
234 e = f["tms_optimization/E_magn"][:] # SimNIBS 4
235else:
236 # read e-field from .pos file
237 fn_e = os.path.join(optim_dir, "coil_positions.geo")
238 # line starting with 'SP'(float, float, float) {e_mag};
239 exp = r"SP\([+-]?[0-9]*[.]?[0-9]+, [+-]?[0-9]*[.]?[0-9]+, [+-]?[0-9]*[.]?[0-9]+\){([+-]?[0-9]*[.]?[0-9]+)};"
240 try:
241 with open(fn_e, 'r') as f:
242 e = []
243 for line in f:
244 try:
245 e.append(float(re.findall(exp, line)[0]))
246 except IndexError:
247 pass
248 e = np.array(e)
249 except Exception as exc:
250 print("Cannot read electrical fields in target for all search positions:")
251 print(f"{exc}")
252 e = None
254pynibs.write_coil_pos_hdf5(fn_hdf=fn_search_positions,
255 centers=centers,
256 m0=m0,
257 m1=m1,
258 m2=m2,
259 datanames=["E_mag"],
260 data=e,
261 overwrite=True)
263# run final simulation at optimal coil position
264# ========================================================================================================
266# Setting up SimNIBS SESSION
267# ========================================================================================================
268print("Setting up SimNIBS SESSION")
269S = simnibs.sim_struct.SESSION()
270S.fnamehead = fn_mesh_msh
271S.pathfem = os.path.join(optim_dir, "e_opt")
272S.fields = "vDeE"
273S.open_in_gmsh = False
274S.fname_tensor = None
275S.map_to_surf = True
276S.map_to_fsavg = False
277S.map_to_MNI = False
278S.subpath = os.path.join(subject.mesh[mesh_idx]["mesh_folder"], f"m2m_{subject.id}")
280if not os.path.exists(S.pathfem):
281 os.makedirs(S.pathfem)
283# Define the TMS simulation and setting up conductivities
284# ========================================================================================================
285tms = [S.add_tmslist()]
286tms[0].fnamecoil = fn_coil
287tms[0].cond[0].value = 0.126 # WM
288tms[0].cond[1].value = 0.275 # GM
289tms[0].cond[2].value = 1.654 # CSF
290tms[0].cond[3].value = 0.01 # Skull
291tms[0].cond[4].value = 0.465 # Scalp
292tms[0].anisotropy_type = anisotropy_type
294if subject.mesh[mesh_idx]['fn_tensor_vn'] is not None:
295 tms[0].fn_tensor_nifti = os.path.join(subject.mesh[mesh_idx]["mesh_folder"],
296 subject.mesh[mesh_idx]['fn_tensor_vn'])
297else:
298 tms[0].fn_tensor_nifti = None
300tms[0].excentricity_scale = None
302# Define the optimal coil position
303# ========================================================================================================
304pos = tms[0].add_position()
305pos.matsimnibs = opt_matsim[:, :, 0]
306pos.distance = 0.1
307pos.didt = 1e6 # coil current rate of change in A/s (1e6 A/s = 1 A/us)
308pos.postprocess = S.fields
310# Running simulations
311# ========================================================================================================
312print("Running SimNIBS")
313filenames = simnibs.run_simnibs(S, cpus=1)
315print("Saving results in .hdf5 format.")
316if type(filenames) is not list:
317 filenames = [filenames]
319pynibs.simnibs_results_msh2hdf5(fn_msh=filenames,
320 fn_hdf5=[os.path.join(S.pathfem, "e")],
321 S=S,
322 pos_tms_idx=[0], # indices of different "coils"
323 pos_local_idx=[0], # indices of associated positions
324 subject=subject,
325 mesh_idx=mesh_idx,
326 overwrite=True,
327 mid2roi=True,
328 verbose=True)
329print("=== DONE ===")