Coverage for /builds/tms-localization/papers/tmsloc_proto/scripts/06_calc_e.py: 72%
188 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-20 07:31 +0000
1#!/usr/bin/env python
3"""
4This script calculates all E-fields for an experiment.
5Additionally, all fields are interpolated to the GM-EM midlayer ROI.
6"""
8import os
9import sys
10import time
11import h5py
12import pynibs
13import pathlib
14import warnings
15import argparse
16import numpy as np
17import pandas as pd
18from scipy.io import loadmat
19from simnibs.optimization import opt_struct
20from simnibs.simulation import save_matlab_sim_struct
22parser = argparse.ArgumentParser(description='Calculate many E-fields and save midlayer E in ROI in .hdf5 file\n')
24# required arguments
25parser.add_argument('-s', '--fn_subject', help='Path to subject obj',
26 required=True, type=str)
27parser.add_argument('-m', '--mesh_idx', help='Mesh idx',
28 required=True, type=str, default="reg")
29parser.add_argument('-e', '--exp_idx', help='Exp idx',
30 required=True, type=str, default="reg_isi_05")
31parser.add_argument('-r', '--roi_idx', help='ROI idx',
32 required=False, type=str, default="midlayer_m1s1pmd")
33parser.add_argument('-n', '--n_cpu', help='How many cpus to use',
34 type=int, default=10)
35parser.add_argument('-a', '--anisotropy_type', help='Anisotropy "vn" or "scalar". '
36 '"scalar" uses isotropic tissue conductivities.',
37 type=str, default="vn")
38parser.add_argument('-p', '--pardiso', dest='pardiso', action='store_true', help="Use pardiso solver.")
39parser.add_argument('--fail_threshold', type=float, default=5,
40 help="Maximum percentage of not-computed fields.")
41parser.add_argument('--recompute', dest='recompute', action='store_true',
42 help="Force (re-)compution of fields if valid e.hdf is found.")
43parser.add_argument('--calc_theta_grad', default=False, action='store_true',
44 help='Have the E-field angle to surface normal (theta) and E-field gradient between'
45 'the gray and white matter calculated. Required for neuronal meanfield based mapping.')
46parser.add_argument('-o', '--online', dest='online', action='store_true', default=False,
47 help="Use fast OnlineFEM.")
48args = parser.parse_args()
50# print simulation info
51# ================================================================================
52print("-" * 64)
53print(f"{parser.description}")
54args_dict = vars(args)
56print('Parameters:')
57for key in args_dict.keys():
58 print(f"{key: >15}: {args_dict[key]}")
59print("-" * 64)
61# set file names and parameters
62# ========================================================================================================
63# file containing the subject information
64fn_subject = os.path.abspath(args.fn_subject)
66# load subject object
67subject = pynibs.load_subject(fn_subject)
69# exp identifier
70exp_idx = args.exp_idx
72# mesh identifier
73mesh_idx = args.mesh_idx
75# ROI index the fields are extracted for goodness-of-fit calculation
76roi_idx = args.roi_idx
78# force recompute?
79recompute = args.recompute
81# output folder
82fn_out = os.path.join(os.path.split(fn_subject)[0], 'results',
83 f"exp_{exp_idx}", "electric_field",
84 f"mesh_{mesh_idx}",
85 f"roi_{roi_idx}")
87# file containing the head model (.msh format)
88fn_mesh_msh = subject.mesh[mesh_idx]["fn_mesh_msh"]
90# file containing the conductivity tensors
91fn_tensor_vn = subject.mesh[mesh_idx]["fn_tensor_vn"]
93# file containing the magnetic vector potential of the TMS coil
94fn_coil = subject.exp[exp_idx]["fn_coil"][0][0]
96subject.exp[exp_idx]["fn_exp_hdf5"] = [os.path.join(subject.subject_folder, "exp", exp_idx, f"mesh_{mesh_idx}",
97 "experiment.hdf5")]
99# file containing the coil positions
100if 'fn_matsimnibs' in subject.exp[exp_idx].keys():
101 fn_matsimnibs = subject.exp[exp_idx]['fn_matsimnibs']
102else:
103 fn_matsimnibs = "matsimnibs.hdf5"
104fn_coilpos_hdf5 = os.path.join(os.path.split(subject.exp[exp_idx]["fn_exp_hdf5"][0])[0], fn_matsimnibs)
106# file containing the FEM options for SimNIBS (will be created)
107fn_fem_options = os.path.join(fn_out, 'FEM_config.mat')
109# anisotropy type ("scalar" or "vn" for volume normalized)
110anisotropy_type = args.anisotropy_type
111if os.name == 'nt': # 'nt' is Windows
112 warnings.warn(f"anisotropy_type={anisotropy_type} needs diffusion tensory, which cannot be created "
113 f"with Windows at the moment. Is this really what you want?")
114 time.sleep(3)
116# number threads (cpu cores) to use for parallelization
117n_cpu = args.n_cpu
119# use pardiso solver
120pardiso = args.pardiso
122# use fast OnlineFEM (calls script run_simnibs_online.py)
123online = args.online
125if online:
126 try:
127 from simnibs.simulation.onlinefem import OnlineFEM
128 except ImportError:
129 raise NotImplementedError("OnlineFEM not implemented in the installed SimNIBS version.")
131# Set up SimNIBS
132# ========================================================================================================
133if anisotropy_type != 'scalar':
134 if not os.path.exists(os.path.join(subject.mesh[mesh_idx]["mesh_folder"], fn_tensor_vn)):
135 print("No conductivity tensor data available ... changing to anisotropy_type='scalar'")
136 anisotropy_type = "scalar"
138if not online:
139 tms_opt = opt_struct.TMSoptimize()
140 tms_opt.fnamehead = fn_mesh_msh
141 tms_opt.pathfem = fn_out
142 tms_opt.fnamecoil = fn_coil
143 tms_opt.target = None
144 tms_opt.open_in_gmsh = False
145 tms_opt.anisotropy_type = anisotropy_type
147 if pardiso:
148 tms_opt.solver_options = 'pardiso' # faster, eats lots of RAM
149 else:
150 tms_opt.solver_options = None
152 if anisotropy_type != 'scalar':
153 tms_opt.fname_tensor = os.path.join(subject.mesh[mesh_idx]["mesh_folder"], fn_tensor_vn)
155 if not os.path.exists(fn_out):
156 os.makedirs(fn_out)
158 # Save FEM options
159 # ========================================================================================================
160 # let's check if an old fn_fem_options is already existing:
161 if os.path.exists(fn_fem_options) and not recompute:
162 print(f"Found old FEM config file {fn_fem_options}. Checking for differences...")
163 try:
164 # check if old config file has same options set
165 tms_opt_old = opt_struct.TMSoptimize()
166 tms_opt_old = tms_opt_old.read_mat_struct(loadmat(fn_fem_options))
167 for k, v in tms_opt_old.__dict__.items():
168 if k not in ['date', 'time_str'] and v != '':
169 assert v == tms_opt.__dict__[k], print(
170 f"Difference found for '{k}': Old: {v}. New: {tms_opt.__dict__[k]}")
171 print(f"No differences found. Using old FEM config.")
172 except AssertionError:
173 recompute = True
174 save_matlab_sim_struct(tms_opt, fn_fem_options)
175 else:
176 save_matlab_sim_struct(tms_opt, fn_fem_options)
178roi = pynibs.load_roi_surface_obj_from_hdf5(subject.mesh[mesh_idx]['fn_mesh_hdf5'])[roi_idx]
179n_tris_roi = roi.n_tris
181with h5py.File(fn_coilpos_hdf5, 'r') as f:
182 n_coils = f['/matsimnibs'].shape[2]
184# Run electric field simulations
185# ========================================================================================================
186fn_e = os.path.join(fn_out, f"e.hdf5")
187field_fail_threshold = args.fail_threshold
188if not recompute and os.path.exists(fn_e):
189 print(f"Old field simulations found: {fn_e}. Checking file.")
190 with h5py.File(fn_e, 'a') as f:
191 try:
192 if 'tmp' in f.keys():
193 n_fields = f['tmp'].shape[0]
194 assert f['tmp'].shape == (n_coils, n_tris_roi, 6), \
195 f"Wrong shape {f['tmp'].shape} found, expected {(n_coils, n_tris_roi, 6)}. Recomputing fields."
196 n_zero_fields = np.sum(np.sum(np.sum(f['tmp'], axis=2), axis=1) == 0)
197 perc_zero_fields = np.round(n_zero_fields / (n_fields / 100), 2)
198 if n_zero_fields:
199 print(f"{n_zero_fields} empty fields (of {n_fields} total fields) found ({perc_zero_fields}%).")
200 if perc_zero_fields >= field_fail_threshold:
201 print(f"Too many empty fields found.")
202 print(f"This can happen if not enough ressources are available for the computation. \n"
203 f"Re-run this script with less CPUs (currently: '--n_cpu={n_cpu}') "
204 f"and the '--recompute' flag.\n"
205 f"Or, increase '--fail_threshold' (currently: {field_fail_threshold}) to at least "
206 f"{perc_zero_fields} to remove empty fields from the data.\nQuitting.")
207 quit()
209 assert np.all(np.sum(np.sum(f['tmp'], axis=2), axis=1) != 0), f"Empty field found, recomputing fields."
210 if 'E' not in f.keys():
211 print(f"Reshaping fields.")
212 E = f['tmp'][:, :, 0:3]
213 E_mag = f['tmp'][:, :, 3]
214 E_norm = f['tmp'][:, :, 4]
215 E_tan = f['tmp'][:, :, 5]
216 f.create_dataset('/E', data=E)
217 f.create_dataset('/E_mag', data=E_mag)
218 f.create_dataset('/E_norm', data=E_norm)
219 f.create_dataset('/E_tan', data=E_tan)
220 if args_dict["calc_theta_grad"]:
221 E_theta = f['tmp'][:, :, 6]
222 E_grad = f['tmp'][:, :, 7]
223 f.create_dataset('/E_theta', data=E_theta)
224 f.create_dataset('/E_grad', data=E_grad)
226 elif 'E' in f.keys():
227 assert f['E'].shape == (n_coils, n_tris_roi, 3), \
228 f"Field component E: Wrong shape {f['E'].shape} found, expected {(n_coils, n_tris_roi, 3)}. " \
229 f"Recomputing fields."
230 for k in ['E_mag', 'E_norm', 'E_tan']:
231 assert f[k].shape == (n_coils, n_tris_roi), \
232 f"Field component {k}: Wrong shape {f[k].shape} found, expected {(n_coils, n_tris_roi)}. " \
233 f"Recomputing fields."
234 print("Test ok, using existing field simulations.")
235 except AssertionError as e:
236 print(e)
237 recompute = True
238else:
239 recompute = True
241scripts_folder = pathlib.Path(__file__).parent.absolute()
243# compute FEMs
244if online:
245 fn_e = os.path.join(fn_out, f"e_scaled.hdf5")
246 if recompute:
247 cmd = f"{sys.executable} {os.path.join(scripts_folder, 'run_simnibs_online.py')} " \
248 f"--folder '{fn_out}' " \
249 f"--fn_subject '{fn_subject}' " \
250 f"--fn_coilpos '{fn_coilpos_hdf5}' " \
251 f"--anisotropy_type {anisotropy_type} " \
252 f"-m '{mesh_idx}' " \
253 f"--fn_coil '{fn_coil}' " \
254 f"-r {roi_idx} " \
255 f"--hdf5_fn '{fn_e}' " \
256 f"--fn_exp '{subject.exp[exp_idx]['fn_exp_hdf5'][0]}'"
257 cmd += (" --calc_theta_grad" if args_dict["calc_theta_grad"] else "")
258 print(f"Running {cmd}")
259 os.system(cmd)
260else:
261 if recompute:
262 cmd = f"{sys.executable} {os.path.join(scripts_folder, 'run_simnibs.py')}" \
263 f" --folder {fn_out} " \
264 f"--fn_subject {fn_subject} " \
265 f"--fn_coilpos {fn_coilpos_hdf5}" \
266 f" --cpus {n_cpu} " \
267 f"--anisotropy_type {anisotropy_type} " \
268 f"-m {mesh_idx} " \
269 f" -r '{roi_idx}' " \
270 f"--hdf5_fn '{fn_e}'"
271 cmd += (" --calc_theta_grad" if args_dict["calc_theta_grad"] else "")
272 print(f"Running {cmd}")
273 os.system(cmd)
275 with h5py.File(fn_e, "a") as f:
276 print(f"Reshaping fields.")
277 E = f['tmp'][:, :, 0:3]
278 E_mag = f['tmp'][:, :, 3]
279 E_norm = f['tmp'][:, :, 4]
280 E_tan = f['tmp'][:, :, 5]
281 f.create_dataset('/E', data=E)
282 f.create_dataset('/E_mag', data=E_mag)
283 f.create_dataset('/E_norm', data=E_norm)
284 f.create_dataset('/E_tan', data=E_tan)
285 if args_dict["calc_theta_grad"]:
286 E_theta = f['tmp'][:, :, 6]
287 E_grad = f['tmp'][:, :, 7]
288 f.create_dataset('/E_theta', data=E_theta)
289 f.create_dataset('/E_grad', data=E_grad)
291 # Scale electric field with coil current and save it under e_scaled.hdf5
292 # ========================================================================================================
293 # read coil current
294 df_stim_data = pd.read_hdf(subject.exp[exp_idx]["fn_exp_hdf5"][0], "stim_data")
295 didt = np.array(df_stim_data["current"])
297 # read and scale electric fields
298 with h5py.File(fn_e, "r") as f:
299 e_scaled = dict()
300 for e_qoi_path in f.keys():
301 if "E" in e_qoi_path:
302 print(e_qoi_path)
304 # The angle of the E-field with respect to the local surface normal as well as the
305 # change of the efield across the local GM crossection should be be scaled as they
306 # are unaffected by a scaling of the underlying E-field.
307 if e_qoi_path not in ["E_gradient", "E_theta"]:
308 e_scaled[e_qoi_path] = f[e_qoi_path][:] * np.expand_dims(didt, axis=tuple(
309 np.arange(f[e_qoi_path][:].ndim - 1) + 1))
310 else:
311 e_scaled[e_qoi_path] = f[e_qoi_path][:]
313 # save scaled electric fields
314 fn_scaled = os.path.join(fn_out, f"e_scaled.hdf5")
315 print(f"Writing scaled fields to {fn_scaled}.")
316 with h5py.File(fn_scaled, "w") as f:
317 for e_qoi_path in e_scaled.keys():
318 f.create_dataset(e_qoi_path, data=e_scaled[e_qoi_path])
320 # check for any zero rows in simulations
321 print(e_scaled.keys())
322 n_zero_simulations = np.sum(np.linalg.norm(e_scaled[list(e_scaled.keys())[0]], axis=1) == 0)
323 if n_zero_simulations > 0:
324 print(f"WARNING: {n_zero_simulations}/{n_coils} failed simulations detected! "
325 f"Reduce n_cpu and avoid memory overflows to avoid this!")
327print("06_calc_e.py: Done.")