Coverage for python/aubellhop/compute.py: 86%
92 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-19 06:49 +0000
1"""Computing wrappers for aubellhop.
3These functions make use of the `Models` registry, selecting appropriate `BellhopSimulator` models (or loading explicitly request ones):
5* `compute(env, ...)` — writes the environment to file and then executes an appropriate Bellhop model;
7* `compute_from_file(model, filename)` uses specified model with pre-existing environment file.
9The `compute()` function allows calculation with multiple environments, tasks, and models, and returns the results in a dictionary of metadata and results.
11Simpler once-off wrapper functions are also provided for convenience (`compute_arrivals()` etc.).
13"""
15from __future__ import annotations
17from typing import Any, cast
18import numpy as np
19import pandas as pd
21from .constants import BHStrings, EnvDefaults, FileExt
22from .environment import Environment
23from .models import Models
25def compute_from_file(
26 model: str,
27 fname: str,
28 debug: bool = False
29 ) -> dict[str, Any]:
30 """Compute Bellhop model directly from existing .env file.
32 Parameters
33 ----------
34 model: str
35 Name of model to run that has been defined in the `Models` registry.
36 fname: str
37 Filename of environment file (with or w/o extension).
38 debug : bool
39 Whether to print diagnostics to the console.
41 Returns
42 -------
43 results : dict
44 Dictionary of metadata and results.
46 Notes
47 -----
48 The environment file is parsed simply to read the specified task; the bellhop executable is run on the original file "in place" in the filesystem. A copy of the parsed environment file is stored in the metadata.
49 """
51 ext = FileExt.env
52 if fname.endswith(ext):
53 nchar = len(ext)
54 fname_base = fname[:-nchar]
55 else:
56 fname_base = fname
57 fname = fname + ext
59 model_fn = Models.get(model)
60 env_tmp = Environment.from_file(fname)
61 task = env_tmp['task']
63 return {
64 "name": env_tmp["name"],
65 "model": model,
66 "task": task,
67 "results": model_fn.run(task, fname_base, rm_files=False, debug=debug),
68 "env": env_tmp.copy(),
69 }
72def compute(
73 env: Environment | list[Environment],
74 model: Any | None = None,
75 task: Any | None = None,
76 debug: bool = False,
77 fname_base: str | None = None,
78 overwrite: bool = False,
79 ) -> dict[str, Any] | tuple[list[dict[str, Any]], pd.DataFrame]:
80 """Compute Bellhop task(s) for given model(s) and environment(s).
82 Parameters
83 ----------
84 env : dict or list of dict
85 Environment definition (which includes the task specification)
86 model : str, optional
87 Propagation model to use (None to auto-select)
88 task : str or list of str, optional
89 Optional task or list of tasks ("arrivals", etc.)
90 debug : bool, default=False
91 Generate debug information for propagation model
92 fname_base : str, optional
93 Base file name for Bellhop working files, default (None), creates a temporary file
94 overwrite : bool, default=False
95 If True, remove any existing working/output files that share the same `fname_base`.
97 Returns
98 -------
99 dict
100 Single run result (and associated metadata) if only one computation is performed.
101 tuple of (list of dict, pandas.DataFrame)
102 List of results and an index DataFrame if multiple computations are performed.
104 Notes
105 -----
106 If any of env, model, and/or task are lists then multiple runs are performed
107 with a list of dictionary outputs returned. The ordering is based on loop iteration
108 but might not be deterministic; use the index DataFrame to extract and filter the
109 output logically.
111 Examples
112 --------
113 Single task based on reading a complete `.env` file:
114 >>> import aubellhop as bh
115 >>> env = bh.Environment.from_file("...")
116 >>> output = bh.compute(env)
117 >>> assert output['task'] == "arrivals"
118 >>> bh.plot_arrivals(output['results'])
120 Multiple tasks:
121 >>> import aubellhop as bh
122 >>> env = bh.Environment()
123 >>> output, ind_df = bh.compute(env,task=["arrivals", "eigenrays"])
124 >>> bh.plot_arrivals(output[0]['results'])
125 """
126 if isinstance(env, list):
127 envs = cast(list[Environment], env)
128 else:
129 envs = [env]
130 models_ = model if isinstance(model, list) else [model]
131 tasks = task if isinstance(task, list) else [task]
132 results: list[dict[str, Any]] = []
133 this_env: Environment
134 for this_env in envs:
135 for this_model in models_:
136 for this_task in tasks:
137 if debug:
138 print(f"Using environment: {this_env['name']}")
139 print(f"Using model: {'[None] (default)' if this_model is None else this_model.get('name')}")
140 print(f"Using task: {this_task}")
141 this_env.check()
142 this_task = this_task or this_env.get('task')
143 if this_task is None:
144 raise ValueError("Task must be specified in env or as parameter")
145 model_fn = Models.select(this_env, this_task, this_model, debug)
146 fname_base = model_fn.write_env(this_env, this_task, fname_base, overwrite=overwrite)
147 results.append({
148 "name": this_env["name"],
149 "model": this_model,
150 "task": this_task,
151 "results": model_fn.run(this_task, fname_base, debug=debug),
152 "env": this_env.copy(),
153 })
154 assert len(results) > 0, "No results generated"
155 index_df = pd.DataFrame([
156 {
157 "i": i,
158 "name": r["name"],
159 "model": getattr(r["model"], "name", str(r["model"])) if r["model"] is not None else None,
160 "task": r["task"],
161 }
162 for i, r in enumerate(results)
163 ])
164 index_df.set_index("i", inplace=True)
165 if len(results) > 1:
166 return results, index_df
167 else:
168 return results[0]
171def compute_arrivals(env: Environment, model: Any | None = None, debug: bool = False, fname_base: str | None = None, overwrite: bool = False) -> Any:
172 """Compute arrivals between each transmitter and receiver.
174 Parameters
175 ----------
176 env : dict
177 Environment definition
178 model : str, optional
179 Propagation model to use (None to auto-select)
180 debug : bool, default=False
181 Generate debug information for propagation model
182 fname_base : str, optional
183 Base file name for Bellhop working files, default (None), creates a temporary file
184 overwrite : bool, default=False
185 If True, remove any existing working/output files that share the same `fname_base`.
187 Returns
188 -------
189 pandas.DataFrame
190 Arrival times and coefficients for all transmitter-receiver combinations
192 Examples
193 --------
194 >>> import aubellhop as bh
195 >>> env = bh.Environment()
196 >>> arrivals = bh.compute_arrivals(env)
197 >>> bh.plot_arrivals(arrivals)
198 """
199 output = compute(env, model, BHStrings.arrivals, debug, fname_base, overwrite)
200 assert isinstance(output, dict), "Single env should return single result"
201 return output['results']
203def compute_eigenrays(env: Environment, source_depth_ndx: int = 0, receiver_depth_ndx: int = 0, receiver_range_ndx: int = 0, model: Any | None = None, debug: bool = False, fname_base: str | None = None, overwrite: bool = False) -> Any:
204 """Compute eigenrays between a given transmitter and receiver.
206 Parameters
207 ----------
208 env : dict
209 Environment definition
210 source_depth_ndx : int, default=0
211 Transmitter depth index
212 receiver_depth_ndx : int, default=0
213 Receiver depth index
214 receiver_range_ndx : int, default=0
215 Receiver range index
216 model : str, optional
217 Propagation model to use (None to auto-select)
218 debug : bool, default=False
219 Generate debug information for propagation model
220 fname_base : str, optional
221 Base file name for Bellhop working files, default (None), creates a temporary file
222 overwrite : bool, default=False
223 If True, remove any existing working/output files that share the same `fname_base`.
225 Returns
226 -------
227 pandas.DataFrame
228 Eigenrays paths
230 Examples
231 --------
232 >>> import aubellhop as bh
233 >>> env = bh.Environment()
234 >>> rays = bh.compute_eigenrays(env)
235 >>> bh.plot_rays(rays, width=1000)
236 """
237 env = env.copy()
238 env.check()
239 if np.size(env['source_depth']) > 1:
240 env['source_depth'] = env['source_depth'][source_depth_ndx]
241 env['source_ndepth'] = None # count is recomputed on the next check()
242 if np.size(env['receiver_depth']) > 1:
243 env['receiver_depth'] = env['receiver_depth'][receiver_depth_ndx]
244 env['receiver_ndepth'] = None
245 if np.size(env['receiver_range']) > 1:
246 env['receiver_range'] = env['receiver_range'][receiver_range_ndx]
247 env['receiver_nrange'] = None
248 output = compute(env, model, BHStrings.eigenrays, debug, fname_base, overwrite)
249 assert isinstance(output, dict), "Single env should return single result"
250 return output['results']
252def compute_rays(env: Environment, source_depth_ndx: int = 0, model: Any | None = None, debug: bool = False, fname_base: str | None = None, overwrite: bool = False) -> Any:
253 """Compute rays from a given transmitter.
255 Parameters
256 ----------
257 env : dict
258 Environment definition
259 source_depth_ndx : int, default=0
260 Transmitter depth index
261 model : str, optional
262 Propagation model to use (None to auto-select)
263 debug : bool, default=False
264 Generate debug information for propagation model
265 fname_base : str, optional
266 Base file name for Bellhop working files, default (None), creates a temporary file
267 overwrite : bool, default=False
268 If True, remove any existing working/output files that share the same `fname_base`.
270 Returns
271 -------
272 pandas.DataFrame
273 Ray paths
275 Examples
276 --------
277 >>> import aubellhop as bh
278 >>> env = bh.Environment()
279 >>> rays = bh.compute_rays(env)
280 >>> bh.plot_rays(rays, width=1000)
281 """
282 env = env.copy()
283 env.check()
284 if np.size(env['source_depth']) > 1:
285 env['source_depth'] = env['source_depth'][source_depth_ndx]
286 env['source_ndepth'] = None # count is recomputed on the next check()
287 output = compute(env, model, BHStrings.rays, debug, fname_base, overwrite)
288 assert isinstance(output, dict), "Single env should return single result"
289 return output['results']
291def compute_transmission_loss(env: Environment, source_depth_ndx: int = 0, mode: str | None = None, model: Any | None = None, debug: bool = False, fname_base: str | None = None, overwrite: bool = False) -> Any:
292 """Compute transmission loss from a given transmitter to all receviers.
294 Parameters
295 ----------
296 env : dict
297 Environment definition
298 source_depth_ndx : int, default=0
299 Transmitter depth index
300 mode : str, optional
301 Coherent, incoherent or semicoherent
302 model : str, optional
303 Propagation model to use (None to auto-select)
304 debug : bool, default=False
305 Generate debug information for propagation model
306 fname_base : str, optional
307 Base file name for Bellhop working files, default (None), creates a temporary file
308 overwrite : bool, default=False
309 If True, remove any existing working/output files that share the same `fname_base`.
311 Returns
312 -------
313 numpy.ndarray
314 Complex transmission loss at each receiver depth and range
316 Examples
317 --------
318 >>> import aubellhop as bh
319 >>> env = bh.Environment()
320 >>> tloss = bh.compute_transmission_loss(env, mode="incoherent")
321 >>> bh.plot_transmission_loss(tloss, width=1000)
322 """
323 env = env.copy()
324 task = mode or env.get("interference_mode") or EnvDefaults.interference_mode
325 env['interference_mode'] = task
326 env.check()
327 if np.size(env['source_depth']) > 1:
328 env['source_depth'] = env['source_depth'][source_depth_ndx]
329 env['source_ndepth'] = None # count is recomputed on the next check()
330 output = compute(env, model, task, debug, fname_base, overwrite)
331 assert isinstance(output, dict), "Single env should return single result"
332 return output['results']
334def arrivals_to_impulse_response(arrivals: Any, fs: float, abs_time: bool = False) -> Any:
335 """Convert arrival times and coefficients to an impulse response.
337 Parameters
338 ----------
339 arrivals : pandas.DataFrame
340 Arrivals times (s) and coefficients
341 fs : float
342 Sampling rate (Hz)
343 abs_time : bool, default=False
344 Absolute time (True) or relative time (False)
346 Returns
347 -------
348 numpy.ndarray
349 Impulse response
351 Notes
352 -----
353 If `abs_time` is set to True, the impulse response is placed such that
354 the zero time corresponds to the time of transmission of signal.
356 Examples
357 --------
358 >>> import aubellhop as bh
359 >>> env = bh.Environment()
360 >>> arrivals = bh.compute_arrivals(env)
361 >>> ir = bh.arrivals_to_impulse_response(arrivals, fs=192000)
362 """
363 t0 = 0 if abs_time else min(arrivals.time_of_arrival)
364 irlen = int(np.ceil((max(arrivals.time_of_arrival)-t0)*fs))+1
365 ir = np.zeros(irlen, dtype=np.complex128)
366 for _, row in arrivals.iterrows():
367 ndx = int(np.round((row.time_of_arrival.real-t0)*fs))
368 ir[ndx] = row.arrival_amplitude
369 return ir