Coverage for python / aubellhop / compute.py: 85%
87 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-24 14:11 +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.check()
238 env = env.copy()
239 if np.size(env['source_depth']) > 1:
240 env['source_depth'] = env['source_depth'][source_depth_ndx]
241 if np.size(env['receiver_depth']) > 1:
242 env['receiver_depth'] = env['receiver_depth'][receiver_depth_ndx]
243 if np.size(env['receiver_range']) > 1:
244 env['receiver_range'] = env['receiver_range'][receiver_range_ndx]
245 output = compute(env, model, BHStrings.eigenrays, debug, fname_base, overwrite)
246 assert isinstance(output, dict), "Single env should return single result"
247 return output['results']
249def 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:
250 """Compute rays from a given transmitter.
252 Parameters
253 ----------
254 env : dict
255 Environment definition
256 source_depth_ndx : int, default=0
257 Transmitter depth index
258 model : str, optional
259 Propagation model to use (None to auto-select)
260 debug : bool, default=False
261 Generate debug information for propagation model
262 fname_base : str, optional
263 Base file name for Bellhop working files, default (None), creates a temporary file
264 overwrite : bool, default=False
265 If True, remove any existing working/output files that share the same `fname_base`.
267 Returns
268 -------
269 pandas.DataFrame
270 Ray paths
272 Examples
273 --------
274 >>> import aubellhop as bh
275 >>> env = bh.Environment()
276 >>> rays = bh.compute_rays(env)
277 >>> bh.plot_rays(rays, width=1000)
278 """
279 env.check()
280 if np.size(env['source_depth']) > 1:
281 env = env.copy()
282 env['source_depth'] = env['source_depth'][source_depth_ndx]
283 output = compute(env, model, BHStrings.rays, debug, fname_base, overwrite)
284 assert isinstance(output, dict), "Single env should return single result"
285 return output['results']
287def 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:
288 """Compute transmission loss from a given transmitter to all receviers.
290 Parameters
291 ----------
292 env : dict
293 Environment definition
294 source_depth_ndx : int, default=0
295 Transmitter depth index
296 mode : str, optional
297 Coherent, incoherent or semicoherent
298 model : str, optional
299 Propagation model to use (None to auto-select)
300 debug : bool, default=False
301 Generate debug information for propagation model
302 fname_base : str, optional
303 Base file name for Bellhop working files, default (None), creates a temporary file
304 overwrite : bool, default=False
305 If True, remove any existing working/output files that share the same `fname_base`.
307 Returns
308 -------
309 numpy.ndarray
310 Complex transmission loss at each receiver depth and range
312 Examples
313 --------
314 >>> import aubellhop as bh
315 >>> env = bh.Environment()
316 >>> tloss = bh.compute_transmission_loss(env, mode="incoherent")
317 >>> bh.plot_transmission_loss(tloss, width=1000)
318 """
319 env = env.copy()
320 task = mode or env.get("interference_mode") or EnvDefaults.interference_mode
321 env['interference_mode'] = task
322 env.check()
323 if np.size(env['source_depth']) > 1:
324 env['source_depth'] = env['source_depth'][source_depth_ndx]
325 output = compute(env, model, task, debug, fname_base, overwrite)
326 assert isinstance(output, dict), "Single env should return single result"
327 return output['results']
329def arrivals_to_impulse_response(arrivals: Any, fs: float, abs_time: bool = False) -> Any:
330 """Convert arrival times and coefficients to an impulse response.
332 Parameters
333 ----------
334 arrivals : pandas.DataFrame
335 Arrivals times (s) and coefficients
336 fs : float
337 Sampling rate (Hz)
338 abs_time : bool, default=False
339 Absolute time (True) or relative time (False)
341 Returns
342 -------
343 numpy.ndarray
344 Impulse response
346 Notes
347 -----
348 If `abs_time` is set to True, the impulse response is placed such that
349 the zero time corresponds to the time of transmission of signal.
351 Examples
352 --------
353 >>> import aubellhop as bh
354 >>> env = bh.Environment()
355 >>> arrivals = bh.compute_arrivals(env)
356 >>> ir = bh.arrivals_to_impulse_response(arrivals, fs=192000)
357 """
358 t0 = 0 if abs_time else min(arrivals.time_of_arrival)
359 irlen = int(np.ceil((max(arrivals.time_of_arrival)-t0)*fs))+1
360 ir = np.zeros(irlen, dtype=np.complex128)
361 for _, row in arrivals.iterrows():
362 ndx = int(np.round((row.time_of_arrival.real-t0)*fs))
363 ir[ndx] = row.arrival_amplitude
364 return ir