diff --git a/.gitignore b/.gitignore index f760c6c..f3ffd97 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,4 @@ src/pygenray/__about__.py docs/_build/ dist/ .vscode/ -uv.lock \ No newline at end of file +uv.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..ea9c317 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,18 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: "v6.0.0" + hooks: + - id: check-added-large-files + args: ["--maxkb=5120"] + - id: check-merge-conflict + - id: check-yaml + - id: end-of-file-fixer + - id: trailing-whitespace + - id: debug-statements + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.15.11" + hooks: + - id: ruff-check + args: ["--fix", "--show-fixes"] + - id: ruff-format diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 619330b..0fae287 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -23,4 +23,3 @@ python: path: . extra_requirements: - docs - diff --git a/docs/building_notes.md b/docs/building_notes.md index 8f8cee8..e41b4b0 100644 --- a/docs/building_notes.md +++ b/docs/building_notes.md @@ -38,4 +38,4 @@ actually... scipy might be worth it. Because I can have the ray code done in jus - use launched ray and previous bracketing rays and repeat until convergence -**where I left off**: I have built out `shoot_rays` and `shoot_ray`, and created a single method that can be called from within the functions that take an environment as input, and when accessing shared memory. But I haven't tested all of that functionality yet. \ No newline at end of file +**where I left off**: I have built out `shoot_rays` and `shoot_ray`, and created a single method that can be called from within the functions that take an environment as input, and when accessing shared memory. But I haven't tested all of that functionality yet. diff --git a/docs/conf.py b/docs/conf.py index 9113237..611c68b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -6,7 +6,7 @@ import sys # Add the source directory to the path so Sphinx can import the modules -sys.path.insert(0, os.path.abspath('../src')) +sys.path.insert(0, os.path.abspath("../src")) project = "pygenray" copyright = "2025, John Ragland" diff --git a/docs/quick_start.md b/docs/quick_start.md index 8b13789..e69de29 100644 --- a/docs/quick_start.md +++ b/docs/quick_start.md @@ -1 +0,0 @@ - diff --git a/docs/ray_physics.md b/docs/ray_physics.md index df60fb0..b7529e6 100644 --- a/docs/ray_physics.md +++ b/docs/ray_physics.md @@ -4,4 +4,4 @@ - Environment depth and source / receiver depths are specified with positive z, defined by distance from surface of ocean in meters. - The ray state, $y$, defined by $y = \left[\mathrm{T, z, p_z} \right]^T$, uses sign convention where z is negative and increases in z correspond to moving towards the surface. - Subsequently, a positive ray angle corresponds to a ray moving towards the surface. - - A positive ray ID corresponds to a positive launch angle (towards the surface) \ No newline at end of file + - A positive ray ID corresponds to a positive launch angle (towards the surface) diff --git a/pyproject.toml b/pyproject.toml index dafa407..2e2611f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,9 @@ exclude_lines = [ "if TYPE_CHECKING:", ] +[tool.ruff.lint.per-file-ignores] +"src/pygenray/__init__.py" = ["F403"] + [dependency-groups] dev = [ "pytest>=8.3.5", diff --git a/src/pygenray/eigenrays.py b/src/pygenray/eigenrays.py index 8b09f6b..4a6889b 100644 --- a/src/pygenray/eigenrays.py +++ b/src/pygenray/eigenrays.py @@ -1,25 +1,27 @@ """ Tools and methods for calculating eigenrays for specifed receiver depths. """ + import pygenray as pr import numpy as np import multiprocessing as mp from tqdm import tqdm + def find_eigenrays( - rays, - receiver_depths, - source_depth, - source_range, - receiver_range, - num_range_save, - environment, - ztol=1, - max_iter=20, - num_workers=None, - **kwargs, - ): - ''' + rays, + receiver_depths, + source_depth, + source_range, + receiver_range, + num_range_save, + environment, + ztol=1, + max_iter=20, + num_workers=None, + **kwargs, +): + """ Given an initial ray fan, find eigenrays with [regula falsi](https://en.wikipedia.org/wiki/Regula_falsi#The_regula_falsi_(false_position)_method) method of root finding. Parameters @@ -51,7 +53,7 @@ def find_eigenrays( ------- erays : dict dictionary of eigen rays. Key values are values in `receiver_depths`. - ''' + """ erays_dict = {} num_eigenrays = {} num_eigenrays_found = {} @@ -60,7 +62,7 @@ def find_eigenrays( for rd_idx, receiver_depth in enumerate(receiver_depths): ## get initial bracketing rays # get indices before sign changes - depth_sign = np.sign(rays.zs[:,-1] + receiver_depth) + depth_sign = np.sign(rays.zs[:, -1] + receiver_depth) sign_change = np.diff(depth_sign) bracket_idxs_start = np.where(sign_change)[0] @@ -70,11 +72,11 @@ def find_eigenrays( bracket_idxs = np.column_stack([bracket_idxs_start, bracket_idxs_start + 1]) # compute regula falsi launch angles - z1s = rays.zs[bracket_idxs[:,0].astype(int),-1] - z2s = rays.zs[bracket_idxs[:,1].astype(int),-1] + z1s = rays.zs[bracket_idxs[:, 0].astype(int), -1] + z2s = rays.zs[bracket_idxs[:, 1].astype(int), -1] - theta1s = rays.thetas[bracket_idxs[:,0].astype(int)] - theta2s = rays.thetas[bracket_idxs[:,1].astype(int)] + theta1s = rays.thetas[bracket_idxs[:, 0].astype(int)] + theta2s = rays.thetas[bracket_idxs[:, 1].astype(int)] erays_dict[rd_idx] = [] failed_eray_theta_brackets[rd_idx] = [] @@ -90,7 +92,7 @@ def find_eigenrays( for eray_idx in range(num_eigenrays[receiver_depth]): z1_distance = np.abs(z1s[eray_idx] + receiver_depth) z2_distance = np.abs(z2s[eray_idx] + receiver_depth) - + if (z1_distance < ztol) or (z2_distance < ztol): bracket_winner = np.argmin([z1_distance, z2_distance]) @@ -105,7 +107,7 @@ def find_eigenrays( eray_found_idxs.append(eray_idx) else: continue - + # remove found eigenrays from search z1s = np.delete(z1s, eray_found_idxs) z2s = np.delete(z2s, eray_found_idxs) @@ -113,23 +115,47 @@ def find_eigenrays( theta2s = np.delete(theta2s, eray_found_idxs) """ - regula_falsi_thetas = theta1s - (z1s + receiver_depth) * (theta2s - theta1s) / (z2s - z1s) + regula_falsi_thetas = theta1s - (z1s + receiver_depth) * (theta2s - theta1s) / ( + z2s - z1s + ) - if len(regula_falsi_thetas) > 5: # use parallel processing for large number of rays + if ( + len(regula_falsi_thetas) > 5 + ): # use parallel processing for large number of rays # construct argument iterable for parallel processing args_list = [] - for k in range (len(regula_falsi_thetas)): - args = (k, z1s[k], z2s[k], theta1s[k], theta2s[k], regula_falsi_thetas[k], - receiver_depth, source_depth, source_range, receiver_range, - num_range_save, environment, ztol, max_iter, kwargs) + for k in range(len(regula_falsi_thetas)): + args = ( + k, + z1s[k], + z2s[k], + theta1s[k], + theta2s[k], + regula_falsi_thetas[k], + receiver_depth, + source_depth, + source_range, + receiver_range, + num_range_save, + environment, + ztol, + max_iter, + kwargs, + ) args_list.append(args) # map individual eigen ray finding to different workers if num_workers is None: num_workers = mp.cpu_count() - with mp.get_context('spawn').Pool(num_workers) as pool: - results = list(tqdm(pool.imap(_find_single_eigenray, args_list), total=len(args_list), desc="Finding eigenrays")) - + with mp.get_context("spawn").Pool(num_workers) as pool: + results = list( + tqdm( + pool.imap(_find_single_eigenray, args_list), + total=len(args_list), + desc="Finding eigenrays", + ) + ) + # Filter out None results and add successful rays for result in results: if result is not None: @@ -138,10 +164,26 @@ def find_eigenrays( failed_eray_theta_brackets[rd_idx].append((theta1s[k], theta2s[k])) else: # use sequential processing for small number of rays - for k in tqdm(range(len(regula_falsi_thetas)), desc='Finding eigenrays:'): - ray = _find_single_eigenray((k, z1s[k], z2s[k], theta1s[k], theta2s[k], regula_falsi_thetas[k], - receiver_depth, source_depth, source_range, receiver_range, - num_range_save, environment, ztol, max_iter, kwargs)) + for k in tqdm(range(len(regula_falsi_thetas)), desc="Finding eigenrays:"): + ray = _find_single_eigenray( + ( + k, + z1s[k], + z2s[k], + theta1s[k], + theta2s[k], + regula_falsi_thetas[k], + receiver_depth, + source_depth, + source_range, + receiver_range, + num_range_save, + environment, + ztol, + max_iter, + kwargs, + ) + ) if ray is not None: erays_dict[rd_idx].append(ray) else: @@ -150,7 +192,14 @@ def find_eigenrays( num_eigenrays_found[rd_idx] = len(erays_dict[rd_idx]) # Create EigenRays object after processing all receiver depths - erays = pr.EigenRays(receiver_depths, erays_dict, environment, num_eigenrays, num_eigenrays_found, failed_eray_theta_brackets) + erays = pr.EigenRays( + receiver_depths, + erays_dict, + environment, + num_eigenrays, + num_eigenrays_found, + failed_eray_theta_brackets, + ) return erays @@ -158,20 +207,44 @@ def _find_single_eigenray(args): """ Find single Eigen ray given the bracketing ray depths, and launch angles. """ - k, z1, z2, theta1, theta2, regula_falsi_theta, receiver_depth, source_depth, source_range, receiver_range, num_range_save, environment, ztol, max_iter, kwargs = args - + ( + k, + z1, + z2, + theta1, + theta2, + regula_falsi_theta, + receiver_depth, + source_depth, + source_range, + receiver_range, + num_range_save, + environment, + ztol, + max_iter, + kwargs, + ) = args + iter_count = 0 # Regula Falsi root finding loop while True: - - ray = pr.shoot_ray(source_depth, source_range, regula_falsi_theta, receiver_range, num_range_save, environment, **kwargs) + ray = pr.shoot_ray( + source_depth, + source_range, + regula_falsi_theta, + receiver_range, + num_range_save, + environment, + **kwargs, + ) if ray is None: - print(f'Failed to find eigen ray for receiver depth {receiver_depth} [m] and approximate launch angle {regula_falsi_theta} [m] ray θ = 90°') + print( + f"Failed to find eigen ray for receiver depth {receiver_depth} [m] and approximate launch angle {regula_falsi_theta} [m] ray θ = 90°" + ) return None - - if np.abs(ray.z[-1] + receiver_depth) < ztol: + if np.abs(ray.z[-1] + receiver_depth) < ztol: # flip launch angle to match sign convention ray.launch_angle = -ray.launch_angle return ray @@ -185,13 +258,14 @@ def _find_single_eigenray(args): z2 = ray.z[-1] theta2 = regula_falsi_theta - regula_falsi_theta = theta1 - (z1 + receiver_depth) * (theta2 - theta1) / (z2 - z1) + regula_falsi_theta = theta1 - (z1 + receiver_depth) * (theta2 - theta1) / ( + z2 - z1 + ) if iter_count > max_iter: return None - - iter_count += 1 + iter_count += 1 -__all__ = ['find_eigenrays'] \ No newline at end of file +__all__ = ["find_eigenrays"] diff --git a/src/pygenray/environment.py b/src/pygenray/environment.py index dd6c011..0b42dcb 100644 --- a/src/pygenray/environment.py +++ b/src/pygenray/environment.py @@ -10,6 +10,7 @@ from tqdm import tqdm import multiprocessing as mp + class OceanEnvironment2D: """ Ocean Environment Specification (2D). @@ -44,16 +45,16 @@ class OceanEnvironment2D: bottom_angle_call : callable function that returns the bottom slope angle in degrees at a given range. The angle is computed as the arctangent of the gradient of bathymetry with respect to range. The function takes a single argument, range in meters, and returns the angle in degrees. """ - + def __init__( - self, - sound_speed=None, - bathymetry=None, - lat=35, - flat_earth_transform=True, - verbose=False - ): - + self, + sound_speed=None, + bathymetry=None, + lat=35, + flat_earth_transform=True, + verbose=False, + ): + # save attributes self.latitude = lat @@ -64,12 +65,9 @@ def __init__( c_munk = munk_ssp(z) sound_speed = xr.DataArray( - np.array([c_munk]*100), - dims=['range', 'depth'], - coords={ - 'depth': z, - 'range': np.linspace(0, 100e3,100) - } + np.array([c_munk] * 100), + dims=["range", "depth"], + coords={"depth": z, "range": np.linspace(0, 100e3, 100)}, ) else: # check dimension names and coordinates of sound_speed @@ -77,15 +75,19 @@ def __init__( raise TypeError("sound_speed must be an xarray DataArray.") if sound_speed.ndim not in [1, 2]: raise ValueError("sound_speed must be 1D or 2D.") - if 'depth' not in sound_speed.dims: + if "depth" not in sound_speed.dims: raise ValueError("sound_speed must have a 'depth' dimension.") - if sound_speed.ndim == 2 and 'range' not in sound_speed.dims: + if sound_speed.ndim == 2 and "range" not in sound_speed.dims: raise ValueError("2D sound_speed must have a 'range' dimension.") # Check Bathymetry if bathymetry is None: # Default flat bottom at 5000m depth - bathymetry = xr.DataArray(np.linspace(4500,4900,100), dims=['range'], coords={'range': np.linspace(0, 100e3,100)}) + bathymetry = xr.DataArray( + np.linspace(4500, 4900, 100), + dims=["range"], + coords={"range": np.linspace(0, 100e3, 100)}, + ) else: # check dimension names and coordinates of bathymetry @@ -93,12 +95,12 @@ def __init__( raise TypeError("bathymetry must be an xarray DataArray.") if bathymetry.ndim != 1: raise ValueError("bathymetry must be 1D.") - if 'range' not in bathymetry.dims: + if "range" not in bathymetry.dims: raise ValueError("bathymetry must have a 'range' dimension.") # save sound speed and bathymetry self.sound_speed = sound_speed - self.dcdz = sound_speed.differentiate('depth').values + self.dcdz = sound_speed.differentiate("depth").values self.bathymetry = bathymetry # If selected, do flat-earth transformation @@ -110,110 +112,111 @@ def __init__( bottom_angle_vector = np.degrees(np.arctan(bottom_slope)) self.bottom_angle = bottom_angle_vector - self.bottom_angle_interp= scipy.interpolate.interp1d( + self.bottom_angle_interp = scipy.interpolate.interp1d( self.bathymetry.range.values, bottom_angle_vector, - kind='cubic', + kind="cubic", ) - def flat_earth_transform(self,lat): - ''' + def flat_earth_transform(self, lat): + """ Earth flattening transformation, change depths and sound speeds so that spherical shell can be done as an x-z slice (using WGS-84). Single latitude is used. If latitude changes sufficiently over track, consider using flat_earth_rd() ``range dependenant'' version instead. - + Parameters ---------- lat : float latitude in degrees. Positive is north, negative is south. - ''' + """ # transform sound speed cs_fe = [] - for ridx in range(self.sound_speed.sizes['range']): - c_slice = self.sound_speed.isel({'range': ridx}) + for ridx in range(self.sound_speed.sizes["range"]): + c_slice = self.sound_speed.isel({"range": ridx}) depf, cf = eflat(c_slice.depth.values, lat, c_slice.values) - cs_fe.append(xr.DataArray(cf, dims='depth', coords={'depth': depf})) - cs_fe = xr.concat(cs_fe, dim='range').assign_coords({'range': self.sound_speed.range}) - + cs_fe.append(xr.DataArray(cf, dims="depth", coords={"depth": depf})) + cs_fe = xr.concat(cs_fe, dim="range").assign_coords( + {"range": self.sound_speed.range} + ) + # transform bathymetry - bathy_flat,_ = eflat(self.bathymetry.values, lat) - bathy_fe = xr.DataArray(bathy_flat, dims='range', coords={'range': self.bathymetry.range.values}) - + bathy_flat, _ = eflat(self.bathymetry.values, lat) + bathy_fe = xr.DataArray( + bathy_flat, dims="range", coords={"range": self.bathymetry.range.values} + ) + # save flat earth transformed sound speed and bathymetry self.sound_speed_fe = cs_fe self.bathymetry_fe = bathy_fe return - - def flat_earth_transform_rd(self,): - ''' + + def flat_earth_transform_rd( + self, + ): + """ Earth flattening transformation, change depths and sound speeds so that spherical shell can be done as an x-z slice (using WGS-84). earth flattening is computed for each range / lat value independently. - ''' + """ c_fe = flat_earth_c(self.sound_speed, verbose=False) bathy_fe = self.bathymetry.copy(deep=True) # save flat earth transformed sound speed and bathymetry self.sound_speed_fe = c_fe - self.dcdz = c_fe.differentiate('depth') + self.dcdz = c_fe.differentiate("depth") self.bathymetry_fe = bathy_fe return - - def plot(self,**kwargs): - ''' + + def plot(self, **kwargs): + """ plot 2D slice of environment kwargs are passed to matplotlib.pyplot.pcolormesh through xarray - ''' + """ # whether or not to include color bar - if 'add_colorbar' in kwargs: - add_colorbar=kwargs['add_colorbar'] - del kwargs['add_colorbar'] + if "add_colorbar" in kwargs: + add_colorbar = kwargs["add_colorbar"] + del kwargs["add_colorbar"] else: - add_colorbar=True + add_colorbar = True if add_colorbar: - ssp_kwargs = {'cmap': 'viridis', 'cbar_kwargs': {'label': 'sound speed [m/s]'}} + ssp_kwargs = { + "cmap": "viridis", + "cbar_kwargs": {"label": "sound speed [m/s]"}, + } ssp_kwargs.update(kwargs) - self.sound_speed.plot( - x='range', - y='depth', - **ssp_kwargs - ) + self.sound_speed.plot(x="range", y="depth", **ssp_kwargs) else: - self.sound_speed.plot( - x='range', - y='depth', - add_colorbar=False, - **kwargs - ) - + self.sound_speed.plot(x="range", y="depth", add_colorbar=False, **kwargs) + # plot bathymetry plt.fill_between( self.bathymetry.range, self.bathymetry, 50000, - color='#aaaaaa', + color="#aaaaaa", alpha=1, lw=0, ) # Convert axis ticks to km - plt.xlabel('range [m]') - plt.ylabel('depth [m]') - #ax.set_xticklabels([f'{x/1000:.0f}' for x in ax.get_xticks()]) + plt.xlabel("range [m]") + plt.ylabel("depth [m]") + # ax.set_xticklabels([f'{x/1000:.0f}' for x in ax.get_xticks()]) plt.ylim(self.sound_speed.depth.max(), self.sound_speed.depth.min()) return + def munk_ssp(z, sofar_depth=1300, eps=0.00737): - ''' + """ given vector of depth, return munk sound speed profile munk equations from (here)[https://web.archive.org/web/2/https://oalib-acoustics.org/website_resources/AcousticsToolbox/manual/node8.html] @@ -225,18 +228,21 @@ def munk_ssp(z, sofar_depth=1300, eps=0.00737): depth of the SOFAR channel eps : float parameter to munk equation - - ''' - zh = 2*(z - sofar_depth)/sofar_depth - c = 1500*(1 + eps*(zh - 1 + np.exp(-zh))) + """ + + zh = 2 * (z - sofar_depth) / sofar_depth + c = 1500 * (1 + eps * (zh - 1 + np.exp(-zh))) return c -def flat_earth_c(c: xr.DataArray, verbose: bool = False, n_cpus: int=None, chunk_size: int=None): - ''' - given a 2D sound-speed slice with dimensions ['depth','range'] and coordinates ['depth', 'range', 'latitude'], + +def flat_earth_c( + c: xr.DataArray, verbose: bool = False, n_cpus: int = None, chunk_size: int = None +): + """ + given a 2D sound-speed slice with dimensions ['depth','range'] and coordinates ['depth', 'range', 'latitude'], compute flat earth transformed sound speed array. - + Processing is parallelized across available CPUs or SLURM allocated CPUs, with multiple range points processed per worker for improved efficiency. @@ -250,56 +256,59 @@ def flat_earth_c(c: xr.DataArray, verbose: bool = False, n_cpus: int=None, chunk number of cpus used to parallelize range-dependant transformation chunk_size : int number of range points to process per worker. If None, automatically determined. - + Returns ------- cfs_x : xr.DataArray flattened soundspeed profile - ''' - + """ + if n_cpus is None: n_cpus = mp.cpu_count() # Get total number of range points - total_ranges = c.sizes['range'] - + total_ranges = c.sizes["range"] + # Determine chunk size if not provided if chunk_size is None: # Aim for each CPU to handle at least 4 chunks, but no fewer than 5 points per chunk chunk_size = max(5, min(50, total_ranges // (n_cpus * 4))) - + # Create chunks of range indices r_chunks = [] for i in range(0, total_ranges, chunk_size): r_chunks.append(list(range(i, min(i + chunk_size, total_ranges)))) - + if verbose: - print(f"Processing {total_ranges} range points with {n_cpus} CPUs using {len(r_chunks)} chunks " - f"(~{chunk_size} points per chunk)") - + print( + f"Processing {total_ranges} range points with {n_cpus} CPUs using {len(r_chunks)} chunks " + f"(~{chunk_size} points per chunk)" + ) + # Process data sequentially if only one CPU is available or for very small datasets if n_cpus == 1 or total_ranges < 10: results = [] for r_idx in tqdm(range(total_ranges)): - c_slice = c.isel({'range': r_idx}) + c_slice = c.isel({"range": r_idx}) depf, cf = eflat(c_slice.depth, c_slice.lat, c_slice) - cf_da = xr.DataArray(cf, dims='depth', coords={'depth': depf}) + cf_da = xr.DataArray(cf, dims="depth", coords={"depth": depf}) cf_interp = cf_da.interp(depth=c.depth) results.append(cf_interp) else: # Use multiprocessing for parallel computation results = _process_ranges_chunked_parallel(c, r_chunks, n_cpus) - + # Concatenate results along range dimension - cfs_x = xr.concat(results, dim='range') - cfs_x = cfs_x.assign_coords({'range': c.range, 'lat': c.lat}) - + cfs_x = xr.concat(results, dim="range") + cfs_x = cfs_x.assign_coords({"range": c.range, "lat": c.lat}) + return cfs_x + def _process_chunk_worker(args): """ Worker function for processing a chunk of range points - + Parameters ---------- args : tuple @@ -307,27 +316,28 @@ def _process_chunk_worker(args): sound speed profile r_indices : list list of range indices to process - + Returns ------- list list of processed DataArrays for each range point - """ + """ c, r_indices = args results = [] - + for r_idx in r_indices: - c_slice = c.isel({'range': r_idx}) + c_slice = c.isel({"range": r_idx}) depf, cf = eflat(c_slice.depth, c_slice.lat, c_slice) - cf_da = xr.DataArray(cf, dims='depth', coords={'depth': depf}) + cf_da = xr.DataArray(cf, dims="depth", coords={"depth": depf}) cf_interp = cf_da.interp(depth=c.depth) results.append(cf_interp) - + return results + def _process_ranges_chunked_parallel(c, r_chunks, n_cpus): """Process multiple range chunks in parallel - + Parameters ---------- c : xr.DataArray @@ -336,7 +346,7 @@ def _process_ranges_chunked_parallel(c, r_chunks, n_cpus): list of lists of range indices n_cpus : int number of CPUs to use - + Returns ------- list @@ -344,20 +354,21 @@ def _process_ranges_chunked_parallel(c, r_chunks, n_cpus): """ # Create argument tuples for each chunk args_list = [(c, chunk) for chunk in r_chunks] - + # Run in parallel with progress bar all_results = [] - with mp.get_context('spawn').Pool(processes=n_cpus) as pool: + with mp.get_context("spawn").Pool(processes=n_cpus) as pool: for chunk_results in tqdm( - pool.imap(_process_chunk_worker, args_list), + pool.imap(_process_chunk_worker, args_list), total=len(r_chunks), - desc="Processing chunks" + desc="Processing chunks", ): all_results.extend(chunk_results) - + return all_results -def eflat(dep,lat,cs=None): + +def eflat(dep, lat, cs=None): """ function [depf, csf]=eflat( dep, lat, cs); flat earth transformation @@ -369,28 +380,31 @@ def eflat(dep,lat,cs=None): cs = np.zeros_like(dep) # WGS-84 parameters - wgsa=6378137.0 - wgsb=6356752.314 - wgsfact=(wgsb/wgsa)**4 - Re=wgsa - wgsa=wgsa*wgsa - wgsb=wgsb*wgsb + wgsa = 6378137.0 + wgsb = 6356752.314 + wgsfact = (wgsb / wgsa) ** 4 + # Re = wgsa + wgsa = wgsa * wgsa + wgsb = wgsb * wgsb - ll=np.pi*lat/180.0 + ll = np.pi * lat / 180.0 - ree1=wgsa/np.sqrt(wgsa*np.cos(ll)*np.cos(ll)+wgsb*np.sin(ll)*np.sin(ll)) - re=ree1*np.sqrt(np.cos(ll)*np.cos(ll)+wgsfact*np.sin(ll)*np.sin(ll)) + ree1 = wgsa / np.sqrt( + wgsa * np.cos(ll) * np.cos(ll) + wgsb * np.sin(ll) * np.sin(ll) + ) + re = ree1 * np.sqrt(np.cos(ll) * np.cos(ll) + wgsfact * np.sin(ll) * np.sin(ll)) - E=dep/re - depf=dep*(1.0 + E*(0.50 + E/3.0)) - csf=cs*(1.0+E*(1.0+E)) + E = dep / re + depf = dep * (1.0 + E * (0.50 + E / 3.0)) + csf = cs * (1.0 + E * (1.0 + E)) return depf, csf + def eflatinv(depf, lat, csf=None): """ Inverse flat earth transformation - + Parameters: depf : np.array depth in flat earth coordinates @@ -404,37 +418,38 @@ def eflatinv(depf, lat, csf=None): cs : np.array transformed sound speed """ - + # Ensure inputs are column vectors depf = np.reshape(depf, (-1,)) lat = np.reshape(lat, (-1,)) - + # Default parameter handling if csf is None: csf = np.zeros(depf.shape) csf = np.reshape(csf, (-1,)) - + # WGS-84 parameters wgsa = 6378137.0 wgsb = 6356752.314 - wgsfact = (wgsb/wgsa)**4 - Re = wgsa - wgsa = wgsa*wgsa - wgsb = wgsb*wgsb - + wgsfact = (wgsb / wgsa) ** 4 + wgsa = wgsa * wgsa + wgsb = wgsb * wgsb + # Calculate Earth radius at given latitude - ll = np.pi*lat/180.0 - ree1 = wgsa/np.sqrt(wgsa*np.cos(ll)*np.cos(ll) + wgsb*np.sin(ll)*np.sin(ll)) - re = ree1*np.sqrt(np.cos(ll)*np.cos(ll) + wgsfact*np.sin(ll)*np.sin(ll)) - + ll = np.pi * lat / 180.0 + ree1 = wgsa / np.sqrt( + wgsa * np.cos(ll) * np.cos(ll) + wgsb * np.sin(ll) * np.sin(ll) + ) + re = ree1 * np.sqrt(np.cos(ll) * np.cos(ll) + wgsfact * np.sin(ll) * np.sin(ll)) + # Define accuracy for ridder function - zacc = 0.001*np.ones(depf.shape) - + zacc = 0.001 * np.ones(depf.shape) + # Provide a better bracket for the root finder # We know true depth is less than flat depth, so try a range that will work bracket_lower = depf * 0.5 # Try half the flat depth - bracket_upper = depf # Upper bound is the flat depth - + bracket_upper = depf # Upper bound is the flat depth + # Call ridder function to solve for depth try: dep = _ridder(eflat, bracket_lower, bracket_upper, depf, zacc, lat)[0] @@ -446,18 +461,19 @@ def eflatinv(depf, lat, csf=None): except ValueError: # If all else fails, use an approximation # This is a rough approximation, but better than failing - dep = depf / (1.0 + 0.5*(depf/re) + (depf/re)**2/3.0) - + dep = depf / (1.0 + 0.5 * (depf / re) + (depf / re) ** 2 / 3.0) + # Final calculations - E = dep/re - cs = csf/(1.0 + E*(1.0 + E)) - + E = dep / re + cs = csf / (1.0 + E * (1.0 + E)) + return dep, cs + def _ridder(fhdl, xl, xh, xrhs, xacc, *args): """ Solves f(x)=xrhs using Ridder's method - + Parameters: fhdl: function handle xl, xh: lower and upper bounds @@ -465,58 +481,59 @@ def _ridder(fhdl, xl, xh, xrhs, xacc, *args): xacc: desired accuracy *args: additional arguments to pass to fhdl """ - + # Initial function evaluations fl = fhdl(xl, *args) - xrhs fh = fhdl(xh, *args) - xrhs - + # Check that root is bracketed if np.any(fl * fh > 0): - raise ValueError('root must be bracketed') - + raise ValueError("root must be bracketed") + # Initial midpoint x = (xl + xh) / 2 fx = fhdl(x, *args) - xrhs - + # Main iteration loop while True: xm = (xl + xh) / 2 fm = fhdl(xm, *args) - xrhs dnm = np.sqrt(fm * fm - fl * fh) - if np.any(dnm == 0): + if np.any(dnm == 0): return x, fx xnew = xm + (xm - xl) * np.sign(fl - fh) * fm / dnm - - if np.all(abs(xnew - x) <= xacc): + + if np.all(abs(xnew - x) <= xacc): return x, fx - + x = xnew fnew = fhdl(x, *args) - xrhs fx = fnew - if np.all(fnew == 0): + if np.all(fnew == 0): return x, fx - + ind = np.where(fnew * fm < 0) xl = np.copy(xl) # Create copies to avoid in-place modifications of arrays fl = np.copy(fl) xh = np.copy(xh) fh = np.copy(fh) - + xl[ind] = xm[ind] fl[ind] = fm[ind] xh[ind] = xnew[ind] fh[ind] = fnew[ind] - + ind = np.where(fnew * fh < 0) xl[ind] = xnew[ind] fl[ind] = fnew[ind] - + ind = np.where(fnew * fl < 0) xh[ind] = xnew[ind] fh[ind] = fnew[ind] - - if np.all(abs(xh - xl) <= xacc): + + if np.all(abs(xh - xl) <= xacc): return x, fx - + + # Public API -__all__ = ['OceanEnvironment2D'] \ No newline at end of file +__all__ = ["OceanEnvironment2D"] diff --git a/src/pygenray/integration_processes.py b/src/pygenray/integration_processes.py index 30876fd..6fa7b2a 100644 --- a/src/pygenray/integration_processes.py +++ b/src/pygenray/integration_processes.py @@ -18,21 +18,23 @@ .. [Colosi2016] Colosi, J. A. (2016). Sound Propagation through the Stochastic Ocean, Cambridge University Press, 443 pages. """ + import numba import numpy as np + @numba.njit(fastmath=True, cache=True) def derivsrd( - x : float, - y : np.array, - cin : np.array, - cpin : np.array, - rin : np.array, - zin : np.array, - depths: np.array, - depth_ranges : np.array, - ) -> np.array: - ''' + x: float, + y: np.array, + cin: np.array, + cpin: np.array, + rin: np.array, + zin: np.array, + depths: np.array, + depth_ranges: np.array, +) -> np.array: + """ Compute the differential equations for ray propagation. The ray equations are derived from the Hamiltonian formulation for ray theory [Colosi2016a]_, which consist of three coupled ODEs with range as the independant varibale, given by equations :eq:`ray1d`, :eq:`ray2d`, and :eq:`ray3d`. .. math:: @@ -74,15 +76,15 @@ def derivsrd( References ---------- .. [Colosi2016a] Colosi, J. A. (2016). Sound Propagation through the Stochastic Ocean, Cambridge University Press, 443 pages. - ''' - - #unpack ray variables - z=y[1] # current depth - pz=y[2] # current ray parameter + """ - #interpolate sound speed and its derivative at current depth and range - c = bilinear_interp(x,z,rin,zin,cin) - cp = bilinear_interp(x,z,rin,zin,cpin) + # unpack ray variables + z = y[1] # current depth + pz = y[2] # current ray parameter + + # interpolate sound speed and its derivative at current depth and range + c = bilinear_interp(x, z, rin, zin, cin) + cp = bilinear_interp(x, z, rin, zin, cpin) # calculate derivatives # clamp to avoid division by zero when a Runge-Kutta intermediate stage @@ -90,21 +92,18 @@ def derivsrd( arg = 1.0 - (c**2) * (pz**2) if arg <= 0.0: arg = 1e-30 - fact=1/np.sqrt(arg) - dydx = np.array([ - fact/c, - c*pz*fact, - -fact*cp/(c**2) - ]) + fact = 1 / np.sqrt(arg) + dydx = np.array([fact / c, c * pz * fact, -fact * cp / (c**2)]) return dydx + @numba.njit(fastmath=True, cache=True) def bilinear_interp(x, y, x_grid, y_grid, values): """ Perform bilinear interpolation on a 2D grid. - Fast, purely functional bilinear interpolation for scattered points on a + Fast, purely functional bilinear interpolation for scattered points on a regular 2D grid using Numba JIT compilation for performance. Parameters @@ -114,7 +113,7 @@ def bilinear_interp(x, y, x_grid, y_grid, values): y : float The y-coordinate at which to interpolate. x_grid : array_like - 1-D array of x-coordinates of the grid points, must be sorted in + 1-D array of x-coordinates of the grid points, must be sorted in ascending order. y_grid : array_like 1-D array of y-coordinates of the grid points, must be sorted in @@ -152,29 +151,35 @@ def bilinear_interp(x, y, x_grid, y_grid, values): # Find grid indices i = np.searchsorted(x_grid, x) - 1 j = np.searchsorted(y_grid, y) - 1 - + # Clamp to grid bounds i = max(0, min(i, len(x_grid) - 2)) j = max(0, min(j, len(y_grid) - 2)) - + # Bilinear weights - wx = (x - x_grid[i]) / (x_grid[i+1] - x_grid[i]) - wy = (y - y_grid[j]) / (y_grid[j+1] - y_grid[j]) - + wx = (x - x_grid[i]) / (x_grid[i + 1] - x_grid[i]) + wy = (y - y_grid[j]) / (y_grid[j + 1] - y_grid[j]) + # Interpolate v00 = values[i, j] - v10 = values[i+1, j] - v01 = values[i, j+1] - v11 = values[i+1, j+1] - - return (1-wx)*(1-wy)*v00 + wx*(1-wy)*v10 + (1-wx)*wy*v01 + wx*wy*v11 + v10 = values[i + 1, j] + v01 = values[i, j + 1] + v11 = values[i + 1, j + 1] + + return ( + (1 - wx) * (1 - wy) * v00 + + wx * (1 - wy) * v10 + + (1 - wx) * wy * v01 + + wx * wy * v11 + ) + @numba.njit(fastmath=True, cache=True) def linear_interp(x, xin, yin): """ Perform linear interpolation on a 1D grid. - Fast, purely functional linear interpolation for scattered points on a + Fast, purely functional linear interpolation for scattered points on a regular 1D grid using Numba JIT compilation for performance. Parameters @@ -182,7 +187,7 @@ def linear_interp(x, xin, yin): x_interp : float The x-coordinate at which to interpolate. xin : array_like - 1-D array of x-coordinates of the grid points, must be sorted in + 1-D array of x-coordinates of the grid points, must be sorted in ascending order. yin : array_like 1-D array of shape (len(x_grid),) containing the values @@ -211,31 +216,32 @@ def linear_interp(x, xin, yin): >>> result = linear_interp(0.5, x_grid, values) >>> print(result) # Should be 2.5 """ - + # Find grid index i = np.searchsorted(xin, x) - 1 - + # Clamp to grid bounds i = max(0, min(i, len(xin) - 2)) - + # Linear weight - w = (x - xin[i]) / (xin[i+1] - xin[i]) - + w = (x - xin[i]) / (xin[i + 1] - xin[i]) + # Interpolate v0 = yin[i] - v1 = yin[i+1] - - y_interp = (1-w)*v0 + w*v1 + v1 = yin[i + 1] + + y_interp = (1 - w) * v0 + w * v1 return y_interp + @numba.njit(fastmath=True, cache=True) def surface_bounce(x, y, cin, cpin, rin, zin, depths, depth_ranges): """Surface event: only trigger when approaching surface from below""" ray_depth = y[1] # calculate ray angle - ray_theta,c = ray_angle(x,y,cin, rin, zin) + ray_theta, c = ray_angle(x, y, cin, rin, zin) # trigger event when ray crosses surface boundary and is traveling upwards if (ray_depth < 0) and (ray_theta < 0): @@ -243,6 +249,7 @@ def surface_bounce(x, y, cin, cpin, rin, zin, depths, depth_ranges): else: return -1.0 + @numba.njit(fastmath=True, cache=True) def bottom_bounce(x, y, cin, cpin, rin, zin, depths, depth_ranges): """Bottom event: only trigger when approaching bottom from above""" @@ -250,7 +257,7 @@ def bottom_bounce(x, y, cin, cpin, rin, zin, depths, depth_ranges): ray_depth = y[1] # calculate ray angle - ray_theta, c = ray_angle(x,y,cin, rin, zin) + ray_theta, c = ray_angle(x, y, cin, rin, zin) # trigger event when ray crosses boundary and is traveling downwards if (ray_depth > bottom_depth) and (ray_theta > 0): @@ -258,45 +265,42 @@ def bottom_bounce(x, y, cin, cpin, rin, zin, depths, depth_ranges): else: return -1.0 + @numba.njit(fastmath=True, cache=True) def vertical_ray(x, y, cin, cpin, rin, zin, depths, depth_ranges): """Vertical ray event: trigger when ray is vertical (θ = 90 degrees)""" - ray_theta, _ = ray_angle(x,y,cin, rin, zin) - if np.abs(ray_theta) > (90-1e-3): # within 0.001 degree + ray_theta, _ = ray_angle(x, y, cin, rin, zin) + if np.abs(ray_theta) > (90 - 1e-3): # within 0.001 degree return 1.0 else: return -1.0 + @numba.njit(fastmath=True, cache=True) -def ray_bounding_box_event(x,y,cin, cpin, rin, zin, depths, depth_ranges): - ''' +def ray_bounding_box_event(x, y, cin, cpin, rin, zin, depths, depth_ranges): + """ Ray Bounding Box Event - trigger when ray position goes outside of bounding box. Bounding box is defined as the box where sound speed is defined. Returns ------- bbox : bool True if ray is outside bounding box, False otherwise - ''' + """ z = y[1] bbox = (z > zin[-1]) | (z < zin[0]) | (x < rin[0]) | (x > rin[-1]) if bbox: - print('bbox event triggered', z, zin[-1], zin[0], x, rin[0], rin[-1]) + print("bbox event triggered", z, zin[-1], zin[0], x, rin[0], rin[-1]) return bbox + @numba.njit(fastmath=True, cache=True) -def ray_angle( - x : float, - y : np.array, - cin : np.array, - rin : np.array, - zin : np.array -): +def ray_angle(x: float, y: np.array, cin: np.array, rin: np.array, zin: np.array): """ calculate angle of ray for specific ray state - + Parameters ---------- x : float @@ -320,17 +324,16 @@ def ray_angle( c = bilinear_interp(x, y[1], rin, zin, cin) theta = np.degrees(np.arcsin(y[2] * c)) - return theta,c + return theta, c __all__ = [ - 'derivsrd', - 'bottom_bounce', - 'surface_bounce', - 'ray_bounding_box_event', - 'ray_angle', - 'bilinear_interp', - 'linear_interp', - 'vertical_ray', + "derivsrd", + "bottom_bounce", + "surface_bounce", + "ray_bounding_box_event", + "ray_angle", + "bilinear_interp", + "linear_interp", + "vertical_ray", ] - diff --git a/src/pygenray/launch_rays.py b/src/pygenray/launch_rays.py index 2de81e2..8c022fc 100644 --- a/src/pygenray/launch_rays.py +++ b/src/pygenray/launch_rays.py @@ -7,20 +7,21 @@ from tqdm import tqdm import time + def shoot_rays( - source_depth : float, - source_range : float, - launch_angles : np.array, - receiver_range : float, - num_range_save : int, - environment : pr.OceanEnvironment2D, - rtol = 1e-9, - terminate_backwards : bool = True, - n_processes : int = None, - debug : bool = True, - flatearth : bool = True + source_depth: float, + source_range: float, + launch_angles: np.array, + receiver_range: float, + num_range_save: int, + environment: pr.OceanEnvironment2D, + rtol=1e-9, + terminate_backwards: bool = True, + n_processes: int = None, + debug: bool = True, + flatearth: bool = True, ): - ''' + """ Integrate rays for given environment and launch angles. Different launch angle initial conditions are mapped to available CPUS. Parameters @@ -58,27 +59,35 @@ def shoot_rays( number of bottom bounces n_surf : int number of surface bounces - ''' + """ # flip launch angles to match sign convention if type(launch_angles) is list: launch_angles = np.array(launch_angles) launch_angles = -launch_angles - if n_processes == None: + if n_processes is None: n_processes = mp.cpu_count() # set up initial conditions for ray variable ## unpack environment object - cin, cpin, rin, zin, depths, depth_ranges, bottom_angles = _unpack_envi(environment, flatearth=flatearth) - + cin, cpin, rin, zin, depths, depth_ranges, bottom_angles = _unpack_envi( + environment, flatearth=flatearth + ) + # check that coordinates are monotonically increasing if not (np.all(np.diff(rin) >= 0)): - raise Exception('Sound speed range coordinates must be monotonically increasing.') + raise Exception( + "Sound speed range coordinates must be monotonically increasing." + ) if not (np.all(np.diff(zin) >= 0)): - raise Exception('Sound speed depth coordinates must be monotonically increasing.') + raise Exception( + "Sound speed depth coordinates must be monotonically increasing." + ) if not (np.all(np.diff(depth_ranges) >= 0)): - raise Exception('Bathymetry range coordinates must be monotonically increasing.') + raise Exception( + "Bathymetry range coordinates must be monotonically increasing." + ) # Use multiprocessing if number of rays is high enough # TODO set threshold to accurately reflect overhead trade off @@ -96,7 +105,7 @@ def shoot_rays( rtol=rtol, terminate_backwards=terminate_backwards, debug=debug, - flatearth=flatearth + flatearth=flatearth, ) ) # shoot_ray automatically saves launch angle to ray object @@ -106,15 +115,20 @@ def shoot_rays( rays_ls_nonone = [ray for ray in rays_ls if ray is not None] rays = pr.RayFan(rays_ls_nonone) return rays - - else: # Use multiprocessing + + else: # Use multiprocessing # Create Shared Arrays - array_metadata, shms = pr._init_shared_memory(cin, cpin, rin ,zin, depths, depth_ranges, bottom_angles) + array_metadata, shms = pr._init_shared_memory( + cin, cpin, rin, zin, depths, depth_ranges, bottom_angles + ) try: # calculate initial ray parameter c = pr.bilinear_interp(source_range, source_depth, rin, zin, cin) - y0s = [np.array([0, source_depth, np.sin(np.radians(launch_angle))/c]) for launch_angle in launch_angles] + y0s = [ + np.array([0, source_depth, np.sin(np.radians(launch_angle)) / c]) + for launch_angle in launch_angles + ] shoot_ray_part = partial( _shoot_single_ray_process, @@ -124,13 +138,19 @@ def shoot_rays( num_range_save=num_range_save, array_metadata=array_metadata, rtol=rtol, - terminate_backwards=terminate_backwards + terminate_backwards=terminate_backwards, ) - - with mp.get_context('spawn').Pool(n_processes) as pool: - rays_ls = list(tqdm(pool.imap(shoot_ray_part, y0s), total=len(y0s), desc="Computing ray fan")) - ranges = np.linspace(source_range, receiver_range, num_range_save) + with mp.get_context("spawn").Pool(n_processes) as pool: + rays_ls = list( + tqdm( + pool.imap(shoot_ray_part, y0s), + total=len(y0s), + desc="Computing ray fan", + ) + ) + + # ranges = np.linspace(source_range, receiver_range, num_range_save) # unpack results rays_list = [] @@ -144,39 +164,44 @@ def shoot_rays( # _shoot_single_ray_process does not save launch angle in ray object # need to set manually here - rays_list[rays_list_idx].launch_angle = -launch_angles[k] # Use separate counter + rays_list[rays_list_idx].launch_angle = -launch_angles[ + k + ] # Use separate counter # launch angle sign is flipped to match z sign convention rays_list_idx += 1 # Increment counter - + ray_fan = pr.RayFan(rays_list) - + finally: - time.sleep(0.1) # Ensure all processes have finished before cleaning up shared memory + time.sleep( + 0.1 + ) # Ensure all processes have finished before cleaning up shared memory # Always clean up shared memory, even if an error occurs for var in shms: try: shms[var].unlink() shms[var].close() - except: - pass # ignore cleanup errors + except Exception: + pass # ignore cleanup errors return ray_fan + def shoot_ray( - source_depth : float, - source_range : float, - launch_angle : float, - receiver_range : float, - num_range_save : int, - environment : pr.OceanEnvironment2D, - rtol = 1e-9, - terminate_backwards : bool = True, - debug : bool = True, - flatearth : bool = True + source_depth: float, + source_range: float, + launch_angle: float, + receiver_range: float, + num_range_save: int, + environment: pr.OceanEnvironment2D, + rtol=1e-9, + terminate_backwards: bool = True, + debug: bool = True, + flatearth: bool = True, ): """ Integrate rays for given environment and launch angles. Different launch angle initial conditions are mapped to available CPUS. - + Parameters ---------- source_depth : float @@ -201,8 +226,8 @@ def shoot_ray( whether to print debug information, default is False flatearth : bool whether to transform environment to flat earth coordinates. Default is True. - - Returns + + Returns ------- ray : pr.Ray pr.Ray object @@ -211,24 +236,45 @@ def shoot_ray( # flip launch angle to match sign convention launch_angle = -launch_angle - - cin, cpin, rin, zin ,depths, depth_ranges, bottom_angles = _unpack_envi(environment, flatearth=flatearth) + + cin, cpin, rin, zin, depths, depth_ranges, bottom_angles = _unpack_envi( + environment, flatearth=flatearth + ) # check that coordinates are monotonically increasing if not (np.all(np.diff(rin) >= 0)): - raise Exception('Sound speed range coordinates must be monotonically increasing.') + raise Exception( + "Sound speed range coordinates must be monotonically increasing." + ) if not (np.all(np.diff(zin) >= 0)): - raise Exception('Sound speed depth coordinates must be monotonically increasing.') + raise Exception( + "Sound speed depth coordinates must be monotonically increasing." + ) if not (np.all(np.diff(depth_ranges) >= 0)): - raise Exception('Bathymetry range coordinates must be monotonically increasing.') - + raise Exception( + "Bathymetry range coordinates must be monotonically increasing." + ) + # calculate y0 c = pr.bilinear_interp(source_range, source_depth, rin, zin, cin) - y0 = np.array([0, source_depth, np.sin(np.radians(launch_angle))/c]) + y0 = np.array([0, source_depth, np.sin(np.radians(launch_angle)) / c]) # launch ray at angle theta sols, full_ray, n_bottom, n_surface = _shoot_ray_array( - y0, source_depth, source_range, receiver_range, cin, cpin, rin, zin, depths, depth_ranges, bottom_angles, rtol, terminate_backwards,debug + y0, + source_depth, + source_range, + receiver_range, + cin, + cpin, + rin, + zin, + depths, + depth_ranges, + bottom_angles, + rtol, + terminate_backwards, + debug, ) if full_ray is None: @@ -237,29 +283,37 @@ def shoot_ray( # reinterpolate ray to range grid range_save = np.linspace(source_range, receiver_range, num_range_save) full_ray = _interpolate_ray(sols, range_save) - ray = pr.Ray(full_ray[0,:], full_ray[1:,:], n_bottom, n_surface, launch_angle, source_depth) - + ray = pr.Ray( + full_ray[0, :], + full_ray[1:, :], + n_bottom, + n_surface, + launch_angle, + source_depth, + ) + return ray + def _shoot_ray_array( - y0 : np.array, - source_depth : float, - source_range : float, - receiver_range : float, - cin : np.array, - cpin : np.array, - rin : np.array, - zin : np.array, - depths : np.array, - depth_ranges : np.array, - bottom_angles : np.array, - rtol = 1e-9, - terminate_backwards : bool = True, - debug : bool = True, + y0: np.array, + source_depth: float, + source_range: float, + receiver_range: float, + cin: np.array, + cpin: np.array, + rin: np.array, + zin: np.array, + depths: np.array, + depth_ranges: np.array, + bottom_angles: np.array, + rtol=1e-9, + terminate_backwards: bool = True, + debug: bool = True, ): """ Integrate single ray. Integration is terminated at bottom and surface reflections, and reflection angle is calculated and updated. Integration is looped until ray reaches receiver range. If there is an error in the integration, the function returns None, None, None, None. - + Environment specified by numpy arrays that are returned by {mod}`pr._unpack_envi()`. Parameters @@ -314,16 +368,13 @@ def _shoot_ray_array( # create cubic interpolation of bottom slope bottom_angle_interp = scipy.interpolate.interp1d( - depth_ranges, - bottom_angles, - kind='cubic' + depth_ranges, bottom_angles, kind="cubic" ) # set intermediate ray state to initial ray state y_intermediate = y0.copy() while x_intermediate < receiver_range: - sol = _shoot_ray_segment( x_intermediate, y_intermediate, @@ -338,8 +389,8 @@ def _shoot_ray_array( ) if len(sol.t) == 0: - raise Exception('Integration segment failed, no points returned.') - + raise Exception("Integration segment failed, no points returned.") + sols.append(sol) full_ray = np.append(full_ray, np.vstack((sol.t, sol.y)), axis=1) @@ -348,63 +399,66 @@ def _shoot_ray_array( break elif sol.status == -1: if debug: - print(f'Integration failed with message: {sol.message}') + print(f"Integration failed with message: {sol.message}") return None, None, None, None - y_intermediate = sol.y[:,-1] + y_intermediate = sol.y[:, -1] # Bounce Event if len(sol.t_events[0]) > 0 or len(sol.t_events[1]) > 0: # Bounce event occurred, use the event time as the new range if len(sol.t_events[0]) > 0: # Surface event x_intermediate = sol.t_events[0][0] - elif len(sol.t_events[1]) > 0: # Bottom event + elif len(sol.t_events[1]) > 0: # Bottom event x_intermediate = sol.t_events[1][0] - + # Vertical Ray elif len(sol.t_events[2]) > 0: # Vertical ray event if debug: - print(f'ray is vertical at x={sol.t[-1]}, y={sol.y[1,-1]}, terminating integration') + print( + f"ray is vertical at x={sol.t[-1]}, y={sol.y[1, -1]}, terminating integration" + ) return None, None, None, None # calculate ray angle and sound speed at ray state - theta,c = pr.ray_angle(x_intermediate, y_intermediate, cin, rin, zin) - + theta, c = pr.ray_angle(x_intermediate, y_intermediate, cin, rin, zin) + # Surface Bounce - if len(sol.t_events[0])==1: + if len(sol.t_events[0]) == 1: theta_bounce = -theta n_surface += 1 # Bottom Bounce - elif len(sol.t_events[1])==1: + elif len(sol.t_events[1]) == 1: # β: bottom angle in degrees beta = bottom_angle_interp(x_intermediate) - theta_bounce = 2*beta - theta + theta_bounce = 2 * beta - theta n_bottom += 1 # terminate if ray bounces backwards if terminate_backwards and (np.abs(theta_bounce) > 90): if debug: - print(f'ray bounced backwards, terminating integration') - return None,None,None,None - + print("ray bounced backwards, terminating integration") + return None, None, None, None + # update ray angle y_intermediate[2] = np.sin(np.radians(theta_bounce)) / c - + loop_count += 1 - + return sols, full_ray, n_bottom, n_surface + def _shoot_single_ray_process( - y0 : np.array, - source_range : float, - source_depth : float, - receiver_range : float, - num_range_save : int, - array_metadata : dict, - rtol = 1e-9, - terminate_backwards : bool = True, - debug : bool = False + y0: np.array, + source_range: float, + source_depth: float, + receiver_range: float, + num_range_save: int, + array_metadata: dict, + rtol=1e-9, + terminate_backwards: bool = True, + debug: bool = False, ): """ Shoot a single ray, accessing shared memory for environment data. @@ -444,13 +498,13 @@ def _shoot_single_ray_process( # Access shared arrays shared_arrays, existing_shms = pr._unpack_shared_memory(array_metadata) - cin = shared_arrays['cin'] - cpin = shared_arrays['cpin'] - rin = shared_arrays['rin'] - zin = shared_arrays['zin'] - depths = shared_arrays['depths'] - depth_ranges = shared_arrays['depth_ranges'] - bottom_angles = shared_arrays['bottom_angle'] + cin = shared_arrays["cin"] + cpin = shared_arrays["cpin"] + rin = shared_arrays["rin"] + zin = shared_arrays["zin"] + depths = shared_arrays["depths"] + depth_ranges = shared_arrays["depth_ranges"] + bottom_angles = shared_arrays["bottom_angle"] sols, full_ray, n_bottom, n_surface = _shoot_ray_array( y0, @@ -468,7 +522,7 @@ def _shoot_single_ray_process( terminate_backwards, debug, ) - + range_save = np.linspace(source_range, receiver_range, num_range_save) if full_ray is None: @@ -476,11 +530,11 @@ def _shoot_single_ray_process( else: # reinterpolate ray to range grid - full_ray_interpolated = _interpolate_ray(sols, range_save) + full_ray_interpolated = _interpolate_ray(sols, range_save) ray = pr.Ray( - full_ray_interpolated[0,:], - full_ray_interpolated[1:,:], + full_ray_interpolated[0, :], + full_ray_interpolated[1:, :], n_bottom, n_surface, source_depth=source_depth, @@ -488,30 +542,31 @@ def _shoot_single_ray_process( except Exception as e: if debug: - print(f'Error in ray integration: {e}') + print(f"Error in ray integration: {e}") return None finally: # Always close shared memory handles, even with error for var in existing_shms: try: existing_shms[var].close() - except: - pass # ignore cleanup errors + except Exception: + pass # ignore cleanup errors return ray + def _shoot_ray_segment( - x0 : float, - y0 : np.array, - receiver_range : float, - cin : np.array, - cpin : np.array, - rin : np.array, - zin : np.array, - depths : np.array, - depth_ranges : np.array, - rtol = 1e-9, - **kwargs + x0: float, + y0: np.array, + receiver_range: float, + cin: np.array, + cpin: np.array, + rin: np.array, + zin: np.array, + depths: np.array, + depth_ranges: np.array, + rtol=1e-9, + **kwargs, ): """ Given an initial condition vector and initial range, integrate ray @@ -559,8 +614,8 @@ def _shoot_ray_segment( surface_event = pr.surface_bounce surface_event.terminal = True surface_event.direction = 1 - - bottom_event = pr.bottom_bounce + + bottom_event = pr.bottom_bounce bottom_event.terminal = True bottom_event.direction = 1 @@ -575,26 +630,29 @@ def _shoot_ray_segment( sol = scipy.integrate.solve_ivp( pr.derivsrd, - (x0,receiver_range), + (x0, receiver_range), y0, - args = (cin, cpin, rin, zin, depths, depth_ranges), - events = events, - rtol = rtol, + args=(cin, cpin, rin, zin, depths, depth_ranges), + events=events, + rtol=rtol, dense_output=True, - **kwargs + **kwargs, ) return sol + def _unpack_envi(environment, flatearth=True): if flatearth: # chech that environment.sound_speed_fe exists - if not hasattr(environment, 'sound_speed_fe'): - raise Exception('Flat earth transformation has not been applied. Set `flat_earth_transform=True` when creating the OceanEnvironment2D object.') - + if not hasattr(environment, "sound_speed_fe"): + raise Exception( + "Flat earth transformation has not been applied. Set `flat_earth_transform=True` when creating the OceanEnvironment2D object." + ) + cin = np.array(environment.sound_speed_fe.values) - cpin = np.array(environment.sound_speed_fe.differentiate('depth').values) + cpin = np.array(environment.sound_speed_fe.differentiate("depth").values) rin = np.array(environment.sound_speed_fe.range.values) zin = np.array(environment.sound_speed_fe.depth.values) depths = np.array(environment.bathymetry_fe.values) @@ -602,20 +660,17 @@ def _unpack_envi(environment, flatearth=True): bottom_angles = np.array(environment.bottom_angle) else: cin = np.array(environment.sound_speed.values) - cpin = np.array(environment.sound_speed.differentiate('depth').values) + cpin = np.array(environment.sound_speed.differentiate("depth").values) rin = np.array(environment.sound_speed.range.values) zin = np.array(environment.sound_speed.depth.values) depths = np.array(environment.bathymetry.values) depth_ranges = np.array(environment.bathymetry.range.values) bottom_angles = np.array(environment.bottom_angle) - return cin, cpin, rin, zin ,depths, depth_ranges, bottom_angles + return cin, cpin, rin, zin, depths, depth_ranges, bottom_angles -def _interpolate_ray( - sols : list, - range_save : np.array -): +def _interpolate_ray(sols: list, range_save: np.array): """ Given list of {mod}`scipy.integrate._ivp.ivp.OdeResult` solutions, corresponding to integrated ray segments interpolate ray state to specfied range grid using the order of the numerical solver (using `dense_output=True` in `scipy.integrate.solve_ivp`). @@ -633,25 +688,35 @@ def _interpolate_ray( 4D array of ray state at each range_save point first dimension corresponds to [range, time, depth, pz], m is the number of range values to save """ - full_ray = np.ones((3, len(range_save)-1))*np.nan + full_ray = np.ones((3, len(range_save) - 1)) * np.nan for k, sol in enumerate(sols): idx1 = np.argmin(np.abs(range_save - sol.t[0])) idx2 = np.argmin(np.abs(range_save - sol.t[-1])) if idx1 == idx2: - continue # Skip if no range points exist in segment - - full_ray[:, idx1:idx2] = sol.sol(range_save[idx1:idx2]) + continue # Skip if no range points exist in segment + full_ray[:, idx1:idx2] = sol.sol(range_save[idx1:idx2]) # Append final ray state to full_ray - full_ray = np.concatenate((full_ray, np.expand_dims(sols[-1].y[:,-1], axis=1)), axis=1) + full_ray = np.concatenate( + (full_ray, np.expand_dims(sols[-1].y[:, -1], axis=1)), axis=1 + ) # Append range values to full_ray - full_ray_state = np.concatenate((np.expand_dims(range_save, axis=0), full_ray), axis=0) + full_ray_state = np.concatenate( + (np.expand_dims(range_save, axis=0), full_ray), axis=0 + ) return full_ray_state -__all__ = ['_shoot_ray_segment', 'shoot_rays', 'shoot_ray','_shoot_single_ray_process', '_unpack_envi', '_shoot_ray_array'] +__all__ = [ + "_shoot_ray_segment", + "shoot_rays", + "shoot_ray", + "_shoot_single_ray_process", + "_unpack_envi", + "_shoot_ray_array", +] diff --git a/src/pygenray/multi_processing.py b/src/pygenray/multi_processing.py index f7da661..418faae 100644 --- a/src/pygenray/multi_processing.py +++ b/src/pygenray/multi_processing.py @@ -3,16 +3,9 @@ import uuid import os -def _init_shared_memory( - cin, - cpin, - rin, - zin, - depths, - depth_ranges, - bottom_angle -): - ''' + +def _init_shared_memory(cin, cpin, rin, zin, depths, depth_ranges, bottom_angle): + """ Initialize shared memory for multiprocessing Parameters @@ -36,25 +29,19 @@ def _init_shared_memory( Dictionary containing metadata of shared memory arrays shms : dict Dictionary containing shared memory objects for each array - ''' + """ # Create unique id, uses PID and shotened UUID to handle SLURM better - unique_id = f'{os.getpid()}_{uuid.uuid4().hex[:8]}' - - shared_array_name_bases = [ - 'cin','cpin','rin','zin','depths','depth_ranges','bottom_angle' - ] - - shared_array_names = [f'{name}_{unique_id}' for name in shared_array_name_bases] + unique_id = f"{os.getpid()}_{uuid.uuid4().hex[:8]}" shared_arrays_np = { - f'cin_{unique_id}':cin, - f'cpin_{unique_id}':cpin, - f'rin_{unique_id}':rin, - f'zin_{unique_id}':zin, - f'depths_{unique_id}':depths, - f'depth_ranges_{unique_id}':depth_ranges, - f'bottom_angle_{unique_id}':bottom_angle, + f"cin_{unique_id}": cin, + f"cpin_{unique_id}": cpin, + f"rin_{unique_id}": rin, + f"zin_{unique_id}": zin, + f"depths_{unique_id}": depths, + f"depth_ranges_{unique_id}": depth_ranges, + f"bottom_angle_{unique_id}": bottom_angle, } shms = {} @@ -62,16 +49,23 @@ def _init_shared_memory( array_metadata = {} for var in shared_arrays_np: - shms[var] = shared_memory.SharedMemory(create=True, size=shared_arrays_np[var].nbytes, name=var) - shared_arrays[var] = np.ndarray(shared_arrays_np[var].shape, dtype=shared_arrays_np[var].dtype, buffer=shms[var].buf) + shms[var] = shared_memory.SharedMemory( + create=True, size=shared_arrays_np[var].nbytes, name=var + ) + shared_arrays[var] = np.ndarray( + shared_arrays_np[var].shape, + dtype=shared_arrays_np[var].dtype, + buffer=shms[var].buf, + ) shared_arrays[var][:] = shared_arrays_np[var][:] array_metadata[var] = { - 'shape': shared_arrays_np[var].shape, - 'dtype': shared_arrays_np[var].dtype, + "shape": shared_arrays_np[var].shape, + "dtype": shared_arrays_np[var].dtype, } - + return array_metadata, shms + def _unpack_shared_memory(shared_array_metadata): """ Unpack shared memory arrays from metadata @@ -91,26 +85,28 @@ def _unpack_shared_memory(shared_array_metadata): existing_shms = {} shared_arrays = {} - + # Define the mapping from base names to what we expect - base_names = ['cin', 'cpin', 'rin', 'zin', 'depths', 'depth_ranges', 'bottom_angle'] - + base_names = ["cin", "cpin", "rin", "zin", "depths", "depth_ranges", "bottom_angle"] + for var in shared_array_metadata: # Access shared memory with the unique name existing_shms[var] = shared_memory.SharedMemory(name=var) - + # Create numpy array array = np.ndarray( - shared_array_metadata[var]['shape'], - dtype=shared_array_metadata[var]['dtype'], - buffer=existing_shms[var].buf) - + shared_array_metadata[var]["shape"], + dtype=shared_array_metadata[var]["dtype"], + buffer=existing_shms[var].buf, + ) + # Map back to base name for easy access for base_name in base_names: - if var.startswith(f'{base_name}_'): + if var.startswith(f"{base_name}_"): shared_arrays[base_name] = array break - + return shared_arrays, existing_shms -__all__ = ['_init_shared_memory', '_unpack_shared_memory'] \ No newline at end of file + +__all__ = ["_init_shared_memory", "_unpack_shared_memory"] diff --git a/src/pygenray/ray_objects.py b/src/pygenray/ray_objects.py index c0f6e1c..c3ecb51 100644 --- a/src/pygenray/ray_objects.py +++ b/src/pygenray/ray_objects.py @@ -3,18 +3,20 @@ import pygenray as pr from scipy import io + class Ray: """ Single Ray Object - python object that store all parameters associated with a single ray. """ + def __init__( self, - r : float, - y : np.array, - n_bottom : int, - n_surface : int, - launch_angle : float = None, - source_depth : float = None, + r: float, + y: np.array, + n_bottom: int, + n_surface: int, + launch_angle: float = None, + source_depth: float = None, ): """ Parameters @@ -45,9 +47,9 @@ def __init__( """ self.r = r - self.t = y[0,:] - self.z = -y[1,:] # saving with negative z convention - self.p = -y[2,:] # saving with negative z convention + self.t = y[0, :] + self.z = -y[1, :] # saving with negative z convention + self.p = -y[2, :] # saving with negative z convention self.n_bottom = n_bottom self.n_surface = n_surface if launch_angle is not None: @@ -56,16 +58,16 @@ def __init__( self.source_depth = source_depth return - def plot(self,**kwargs): + def plot(self, **kwargs): """ Plot ray in time-depth space """ - plot_kwargs = {'c':'k', 'lw': 1, 'alpha': 0.5} + plot_kwargs = {"c": "k", "lw": 1, "alpha": 0.5} plot_kwargs.update(kwargs) plt.plot(self.r, self.z, **kwargs) - plt.xlabel('time [s]') - plt.ylabel('depth [m]') + plt.xlabel("time [s]") + plt.ylabel("depth [m]") plt.ylim([self.z.min(), self.z.max()]) return @@ -74,10 +76,8 @@ class RayFan: """ RayFan Object - python object that store all parameters associated with a ray fan. """ - def __init__( - self, - Rays : list - ): + + def __init__(self, Rays: list): """ Parameters ---------- @@ -131,15 +131,17 @@ def __init__( self.n_botts = np.array(n_botts) self.n_surfs = np.array(n_surfs) self.source_depths = np.array(source_depths) - + self.compute_rayids() return def compute_rayids(self): - ''' + """ compute the ray IDs for each ray in ray fan. - ''' - ray_ids = np.sum(np.diff(np.sign(self.ps)) != 0, axis=1) * (np.sign(self.thetas)) + """ + ray_ids = np.sum(np.diff(np.sign(self.ps)) != 0, axis=1) * ( + np.sign(self.thetas) + ) b_mask = (self.n_botts == 0) & (self.n_surfs == 0) ray_ids_str = [] @@ -147,13 +149,20 @@ def compute_rayids(self): if b_mask[i]: ray_ids_str.append(str(ray_ids[i])) else: - ray_ids_str.append(f'{ray_ids[i]}b') + ray_ids_str.append(f"{ray_ids[i]}b") self.ray_ids = np.array(ray_ids_str) return - - def plot_time_front(self,include_lines=False, range_idx=-1, add_colorbar=True, ray_id=False,**kwargs): - ''' + + def plot_time_front( + self, + include_lines=False, + range_idx=-1, + add_colorbar=True, + ray_id=False, + **kwargs, + ): + """ plot time front. Key word arguments are passed to plt.scatter. Parameters @@ -166,13 +175,24 @@ def plot_time_front(self,include_lines=False, range_idx=-1, add_colorbar=True, r if True, a colorbar is added to the plot, Default True ray_id : bool if true, than the colors of the time-front arrivals are the ray IDs - ''' + """ if include_lines: - plt.plot(self.ts[:,range_idx], self.zs[:,range_idx], c='#aaaaaa', lw=0.5, zorder=5) - + plt.plot( + self.ts[:, range_idx], + self.zs[:, range_idx], + c="#aaaaaa", + lw=0.5, + zorder=5, + ) - scatter_kwargs = {'c': self.thetas, 'cmap': 'managua', 's': 2, 'lw': 0, 'zorder': 6} + scatter_kwargs = { + "c": self.thetas, + "cmap": "managua", + "s": 2, + "lw": 0, + "zorder": 6, + } scatter_kwargs.update(kwargs) if ray_id: @@ -185,55 +205,60 @@ def plot_time_front(self,include_lines=False, range_idx=-1, add_colorbar=True, r # Map each point's category to its color point_colors = [category_to_color[cat] for cat in self.ray_ids] - scatter_kwargs.update({'c': point_colors}) + scatter_kwargs.update({"c": point_colors}) add_colorbar = False # Add a legend for i, cat in enumerate(unique_categories): plt.scatter([], [], c=[colors[i]], label=cat) - plt.legend(ncols=3, loc='lower left') + plt.legend(ncols=3, loc="lower left") - - plt.scatter(x=self.ts[:,range_idx], y=self.zs[:,range_idx], **scatter_kwargs) + plt.scatter(x=self.ts[:, range_idx], y=self.zs[:, range_idx], **scatter_kwargs) plt.ylim([self.zs.min(), self.zs.max()]) if add_colorbar: - plt.colorbar(label='launch angle [degrees]') - plt.xlabel('time [s]') - plt.ylabel('depth [m]') - plt.title('Time Front') + plt.colorbar(label="launch angle [degrees]") + plt.xlabel("time [s]") + plt.ylabel("depth [m]") + plt.title("Time Front") - def plot_ray_fan(self,**kwargs): - ''' + def plot_ray_fan(self, **kwargs): + """ plot ray fan - ''' + """ # set alpha value - if (10*1/len(self.thetas) > 1) | (10*1/len(self.thetas) < 0): - alpha_val=1 + if (10 * 1 / len(self.thetas) > 1) | (10 * 1 / len(self.thetas) < 0): + alpha_val = 1 else: - alpha_val = 10*1/len(self.thetas) + alpha_val = 10 * 1 / len(self.thetas) - plot_kwargs = {'c':'k', 'lw': 1, 'alpha': alpha_val} + plot_kwargs = {"c": "k", "lw": 1, "alpha": alpha_val} plot_kwargs.update(kwargs) _ = plt.plot(self.rs.T, self.zs.T, **plot_kwargs) - plt.xlabel('range [m]') - plt.ylabel('depth [m]') + plt.xlabel("range [m]") + plt.ylabel("depth [m]") plt.ylim([self.zs.min(), self.zs.max()]) - plt.title('Ray Fan') + plt.title("Ray Fan") return - - def plot_depth_v_angle(self,include_line=False, **kwargs): - scatter_kwargs = {'c': self.thetas, 'cmap': 'managua', 's': 2, 'lw': 0, 'zorder': 6} + def plot_depth_v_angle(self, include_line=False, **kwargs): + + scatter_kwargs = { + "c": self.thetas, + "cmap": "managua", + "s": 2, + "lw": 0, + "zorder": 6, + } scatter_kwargs.update(kwargs) if include_line: - plt.plot(self.thetas, self.zs[:,-1], c='#aaaaaa', lw=0.5, zorder=5) - plt.scatter(x=self.thetas, y=self.zs[:,-1],**kwargs) + plt.plot(self.thetas, self.zs[:, -1], c="#aaaaaa", lw=0.5, zorder=5) + plt.scatter(x=self.thetas, y=self.zs[:, -1], **kwargs) return - + def save_mat(self, filename): """ Save RayFan object to a .mat file. @@ -242,20 +267,22 @@ def save_mat(self, filename): ---------- filename : str Name of the output .mat file to save the RayFan data. - + """ # Create a dictionary to hold the data - data = {'rayfan': { - 'thetas': self.thetas, - 'xs': self.rs, - 'ts': self.ts, - 'zs': self.zs, - 'ps': self.ps, - 'n_botts': self.n_botts, - 'n_surfs': self.n_surfs, - 'source_depths': self.source_depths - }} + data = { + "rayfan": { + "thetas": self.thetas, + "xs": self.rs, + "ts": self.ts, + "zs": self.zs, + "ps": self.ps, + "n_botts": self.n_botts, + "n_surfs": self.n_surfs, + "source_depths": self.source_depths, + } + } # Save the dictionary to a .mat file io.savemat(filename, data) @@ -263,17 +290,17 @@ def save_mat(self, filename): def __add__(self, other): """ Add two RayFan objects by concatenating along the launch angle dimension (M). - + Parameters ---------- other : RayFan Another RayFan object to add to this one - + Returns ------- RayFan New RayFan object with concatenated rays - + Raises ------ TypeError @@ -283,14 +310,14 @@ def __add__(self, other): """ if not isinstance(other, RayFan): raise TypeError("Can only add RayFan objects together") - + # Check if range arrays are compatible if not np.array_equal(self.rs[0], other.rs[0]): raise ValueError("Range arrays (rs) must be equivalent for concatenation") - + # Create combined Ray objects combined_rays = [] - + # Add rays from self for i in range(len(self.thetas)): ray = Ray( @@ -299,10 +326,10 @@ def __add__(self, other): n_bottom=self.n_botts[i], n_surface=self.n_surfs[i], launch_angle=self.thetas[i], - source_depth=self.source_depths[i] + source_depth=self.source_depths[i], ) combined_rays.append(ray) - + # Add rays from other for i in range(len(other.thetas)): ray = Ray( @@ -311,16 +338,16 @@ def __add__(self, other): n_bottom=other.n_botts[i], n_surface=other.n_surfs[i], launch_angle=other.thetas[i], - source_depth=other.source_depths[i] + source_depth=other.source_depths[i], ) combined_rays.append(ray) - + return RayFan(combined_rays) def __len__(self): """ Return the number of rays in the RayFan. - + Returns ------- int @@ -331,7 +358,7 @@ def __len__(self): def __getitem__(self, key): """ Slice the RayFan along the launch angle dimension (M). - + Parameters ---------- key : int, slice, or array-like @@ -339,7 +366,7 @@ def __getitem__(self, key): - int: single ray index (returns Ray object) - slice: slice object (e.g., 0:10:2, returns RayFan object) - array-like: boolean mask or integer indices (returns RayFan object) - + Returns ------- Ray or RayFan @@ -350,11 +377,13 @@ def __getitem__(self, key): # Handle negative indexing if key < 0: key = len(self.thetas) + key - + # Check bounds if key < 0 or key >= len(self.thetas): - raise IndexError(f"Index {key} is out of bounds for RayFan with {len(self.thetas)} rays") - + raise IndexError( + f"Index {key} is out of bounds for RayFan with {len(self.thetas)} rays" + ) + # Return single Ray object return Ray( r=self.rs[key], @@ -362,12 +391,12 @@ def __getitem__(self, key): n_bottom=self.n_botts[key], n_surface=self.n_surfs[key], launch_angle=self.thetas[key], - source_depth=self.source_depths[key] + source_depth=self.source_depths[key], ) - + # Handle slices and array-like indices - return RayFan object selected_rays = [] - + # Handle the slicing to get the indices if isinstance(key, slice): # Use numpy's advanced indexing to handle slices @@ -377,7 +406,7 @@ def __getitem__(self, key): selected_indices = np.asarray(key) if selected_indices.dtype == bool: selected_indices = np.where(selected_indices)[0] - + # Ensure selected_indices is iterable (handle 0-d arrays and scalars) if np.isscalar(selected_indices) or selected_indices.ndim == 0: selected_indices = [int(selected_indices)] @@ -385,7 +414,7 @@ def __getitem__(self, key): selected_indices = selected_indices.tolist() else: raise ValueError("Invalid indexing array shape") - + # Create Ray objects for selected indices for i in selected_indices: ray = Ray( @@ -394,15 +423,15 @@ def __getitem__(self, key): n_bottom=self.n_botts[i], n_surface=self.n_surfs[i], launch_angle=self.thetas[i], - source_depth=self.source_depths[i] + source_depth=self.source_depths[i], ) selected_rays.append(ray) - + return RayFan(selected_rays) class EigenRays: - ''' + """ EigenRays Object - python object that store all parameters associated with eigen rays for given receiver depths Parameters @@ -446,9 +475,17 @@ class EigenRays: Ray ID string with boundary indicator. ray_id_int : int Ray ID integer with no boundary indicator. - ''' + """ - def __init__(self,receiver_depths, eigenray_dict, environment, num_eigenrays, num_eigenrays_found, failed_eray_theta_brackets): + def __init__( + self, + receiver_depths, + eigenray_dict, + environment, + num_eigenrays, + num_eigenrays_found, + failed_eray_theta_brackets, + ): self.receiver_depths = receiver_depths self.rs = {} @@ -481,69 +518,78 @@ def __init__(self,receiver_depths, eigenray_dict, environment, num_eigenrays, nu ray_ids_int = [] # compute receive angle for eray_idx in range(eray_fan.rs.shape[0]): - y_last = np.stack((eray_fan.ts[eray_idx, -1], eray_fan.zs[eray_idx, -1], eray_fan.ps[eray_idx, -1])) + y_last = np.stack( + ( + eray_fan.ts[eray_idx, -1], + eray_fan.zs[eray_idx, -1], + eray_fan.ps[eray_idx, -1], + ) + ) theta, c = pr.ray_angle( eray_fan.rs[eray_idx, -1], - y_last, environment.sound_speed.values, + y_last, + environment.sound_speed.values, environment.sound_speed.range.values, - environment.sound_speed.depth.values + environment.sound_speed.depth.values, ) received_angles_single.append(theta) - ray_id_single = np.sum(np.diff(np.sign(eray_fan.ps[eray_idx,:])) != 0) * (np.sign(eray_fan.thetas[eray_idx])) + ray_id_single = np.sum( + np.diff(np.sign(eray_fan.ps[eray_idx, :])) != 0 + ) * (np.sign(eray_fan.thetas[eray_idx])) if eray_fan.n_botts[eray_idx] == 0 and eray_fan.n_surfs[eray_idx] == 0: - boundary_flag = '' + boundary_flag = "" else: - boundary_flag = 'b' - ray_ids.append(f'{ray_id_single}{boundary_flag}') + boundary_flag = "b" + ray_ids.append(f"{ray_id_single}{boundary_flag}") ray_ids_int.append(int(ray_id_single)) self.received_angles[ridx] = np.array(received_angles_single) self.launch_angles[ridx] = eray_fan.thetas self.ray_id[ridx] = np.array(ray_ids) self.ray_id_int[ridx] = np.array(ray_ids_int) - def plot_angle_time(self,ridxs = None, **kwargs): + def plot_angle_time(self, ridxs=None, **kwargs): if ridxs is None: ridxs = list(self.received_angles.keys()) - + for ridx in ridxs: - plt.scatter(self.ts[ridx][:,-1], self.received_angles[ridx], **kwargs) - - plt.xlabel('time [s]') - plt.ylabel('received angle [deg]') - plt.title('Received Angle vs Time') - - def plot(self, ridxs = [0], **kwargs): - ''' + plt.scatter(self.ts[ridx][:, -1], self.received_angles[ridx], **kwargs) + + plt.xlabel("time [s]") + plt.ylabel("received angle [deg]") + plt.title("Received Angle vs Time") + + def plot(self, ridxs=[0], **kwargs): + """ Plot all eigen rays in time-depth space Parameters ---------- ridxs : list list of receiver depth indices to plot. Default is [0], which plots the first receiver depth. - ''' + """ # if ridx is int, make list of length 1 if isinstance(ridxs, int): ridxs = [ridxs] - ray_kwargs = {'c':'k'} + ray_kwargs = {"c": "k"} ray_kwargs.update(kwargs) for ridx in ridxs: plt.plot(self.rs[ridx].T, self.zs[ridx].T, **ray_kwargs) - plt.xlabel('range [m]') - plt.ylabel('depth [m]') - plt.title('Eigen Rays') + plt.xlabel("range [m]") + plt.ylabel("depth [m]") + plt.title("Eigen Rays") plt.ylim([self.zs[ridx].min(), self.zs[ridx].max()]) def plot_ducted(self, **kwargs): - ''' + """ Plot all eigen rays that don't interact with boundaries - ''' + """ - ray_kwargs = {'c':'k'} + ray_kwargs = {"c": "k"} ray_kwargs.update(kwargs) for ridx in self.ray_id.keys(): @@ -551,9 +597,9 @@ def plot_ducted(self, **kwargs): mask = (self.n_botts[ridx] == 0) & (self.n_surfs[ridx] == 0) plt.plot(self.rs[ridx][mask].T, -self.zs[ridx][mask].T, **ray_kwargs) - plt.xlabel('range [m]') - plt.ylabel('depth [m]') - plt.title('Ducted Eigen Rays') + plt.xlabel("range [m]") + plt.ylabel("depth [m]") + plt.title("Ducted Eigen Rays") def save_mat(self, filename): """ @@ -566,25 +612,28 @@ def save_mat(self, filename): """ data = {} - for ridx,rdepth in enumerate(self.receiver_depths): - data[f'receiver_depth_{ridx}'] = { - 'receiver_depth' : rdepth, - 'xs' : self.rs[ridx], - 'ts' : self.ts[ridx], - 'zs' : self.zs[ridx], - 'ps' : self.ps[ridx], - 'received_angles' : self.received_angles[ridx], - 'launch_angles' : self.launch_angles[ridx], - 'ray_id' : self.ray_id[ridx], - 'ray_id_int' : self.ray_id_int[ridx], - 'n_bottom' : self.n_botts[ridx] if hasattr(self, 'n_botts') else np.nan, - 'n_surface' : self.n_surfs[ridx] if hasattr(self, 'n_surfs') else np.nan, - 'source_depth' : self.source_depths[ridx] if hasattr(self, 'source_depths') else np.nan, - 'num_eigenrays' : self.num_eigenrays, - 'num_eigenrays_found' : self.num_eigenrays_found, + for ridx, rdepth in enumerate(self.receiver_depths): + data[f"receiver_depth_{ridx}"] = { + "receiver_depth": rdepth, + "xs": self.rs[ridx], + "ts": self.ts[ridx], + "zs": self.zs[ridx], + "ps": self.ps[ridx], + "received_angles": self.received_angles[ridx], + "launch_angles": self.launch_angles[ridx], + "ray_id": self.ray_id[ridx], + "ray_id_int": self.ray_id_int[ridx], + "n_bottom": self.n_botts[ridx] if hasattr(self, "n_botts") else np.nan, + "n_surface": self.n_surfs[ridx] if hasattr(self, "n_surfs") else np.nan, + "source_depth": self.source_depths[ridx] + if hasattr(self, "source_depths") + else np.nan, + "num_eigenrays": self.num_eigenrays, + "num_eigenrays_found": self.num_eigenrays_found, } # Save the dictionary to a .mat file - io.savemat(filename, {'eigenrays':data}) + io.savemat(filename, {"eigenrays": data}) + -__all__ = ['Ray','RayFan','EigenRays'] \ No newline at end of file +__all__ = ["Ray", "RayFan", "EigenRays"] diff --git a/tests/conftest.py b/tests/conftest.py index 4c25277..423c9b3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,26 +1,34 @@ """ Shared test fixtures for pygenray tests. """ + import matplotlib -matplotlib.use('Agg') +import numpy as np +import pytest + +from pygenray.ray_objects import Ray, RayFan + +matplotlib.use("Agg") def pytest_addoption(parser): """Register --regenerate-physics CLI flag for physics regression tests.""" parser.addoption( - '--regenerate-physics', action='store_true', default=False, - help='Regenerate physics regression fixture and skip comparison.', + "--regenerate-physics", + action="store_true", + default=False, + help="Regenerate physics regression fixture and skip comparison.", ) -import numpy as np -import pytest -import xarray as xr - -from pygenray.ray_objects import Ray, RayFan - -def _make_ray(launch_angle: float, source_depth: float, n_bottom: int = 0, - n_surface: int = 0, N: int = 10, R: float = 10000.0) -> Ray: +def _make_ray( + launch_angle: float, + source_depth: float, + n_bottom: int = 0, + n_surface: int = 0, + N: int = 10, + R: float = 10000.0, +) -> Ray: """Helper: build a synthetic Ray without running the ODE solver. The y array uses the positive-z convention expected by Ray.__init__: @@ -34,8 +42,14 @@ def _make_ray(launch_angle: float, source_depth: float, n_bottom: int = 0, z_ode = np.linspace(source_depth, source_depth + R * 0.01, N) p_ode = np.ones(N) * np.sin(np.radians(abs(launch_angle) + 1e-3)) / 1500.0 y = np.vstack([t, z_ode, p_ode]) # shape (3, N) - return Ray(r=r, y=y, n_bottom=n_bottom, n_surface=n_surface, - launch_angle=launch_angle, source_depth=source_depth) + return Ray( + r=r, + y=y, + n_bottom=n_bottom, + n_surface=n_surface, + launch_angle=launch_angle, + source_depth=source_depth, + ) @pytest.fixture @@ -48,8 +62,8 @@ def simple_ray(): def simple_rayfan(): """Small RayFan of 3 synthetic Rays — no ray tracing needed.""" rays = [ - _make_ray(launch_angle=-5.0, source_depth=100.0, n_bottom=0), - _make_ray(launch_angle=5.0, source_depth=150.0, n_bottom=1), + _make_ray(launch_angle=-5.0, source_depth=100.0, n_bottom=0), + _make_ray(launch_angle=5.0, source_depth=150.0, n_bottom=1), _make_ray(launch_angle=-10.0, source_depth=200.0, n_bottom=0), ] return RayFan(rays) diff --git a/tests/test_environment.py b/tests/test_environment.py index a639236..6048b0e 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -1,6 +1,7 @@ """ Tests for pygenray.environment: munk_ssp, OceanEnvironment2D, eflat, eflatinv. """ + import numpy as np import pytest import xarray as xr @@ -18,6 +19,7 @@ # munk_ssp # --------------------------------------------------------------------------- + class TestMunkSSP: def test_output_shape_matches_input(self): z = np.arange(0, 5000, 10) @@ -45,54 +47,63 @@ def test_scalar_input(self): # OceanEnvironment2D construction # --------------------------------------------------------------------------- + class TestOceanEnvironment2DConstruction: def test_default_init_attributes_exist(self): env = OceanEnvironment2D() - for attr in ('sound_speed', 'bathymetry', 'dcdz', 'bottom_angle', - 'bottom_angle_interp'): + for attr in ( + "sound_speed", + "bathymetry", + "dcdz", + "bottom_angle", + "bottom_angle_interp", + ): assert hasattr(env, attr), f"Missing attribute: {attr}" def test_default_sound_speed_is_2d(self): env = OceanEnvironment2D() assert env.sound_speed.ndim == 2 - assert set(env.sound_speed.dims) == {'range', 'depth'} + assert set(env.sound_speed.dims) == {"range", "depth"} def test_default_flat_earth_attributes_exist(self): env = OceanEnvironment2D(flat_earth_transform=True) - assert hasattr(env, 'sound_speed_fe') - assert hasattr(env, 'bathymetry_fe') + assert hasattr(env, "sound_speed_fe") + assert hasattr(env, "bathymetry_fe") def test_flat_earth_false_no_fe_attributes(self): env = OceanEnvironment2D(flat_earth_transform=False) - assert not hasattr(env, 'sound_speed_fe') - assert not hasattr(env, 'bathymetry_fe') + assert not hasattr(env, "sound_speed_fe") + assert not hasattr(env, "bathymetry_fe") def test_custom_1d_sound_speed(self): z = np.arange(0.0, 3000.0, 10.0) c_vals = munk_ssp(z) - ssp = xr.DataArray(c_vals, dims=['depth'], coords={'depth': z}) + ssp = xr.DataArray(c_vals, dims=["depth"], coords={"depth": z}) bathy = xr.DataArray( - np.ones(20) * 4000.0, dims=['range'], - coords={'range': np.linspace(0, 50e3, 20)} + np.ones(20) * 4000.0, + dims=["range"], + coords={"range": np.linspace(0, 50e3, 20)}, + ) + env = OceanEnvironment2D( + sound_speed=ssp, bathymetry=bathy, flat_earth_transform=False ) - env = OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, - flat_earth_transform=False) assert env.sound_speed.ndim == 1 - assert 'depth' in env.sound_speed.dims + assert "depth" in env.sound_speed.dims def test_custom_2d_sound_speed(self): z = np.arange(0.0, 3000.0, 50.0) r = np.linspace(0.0, 50e3, 20) c_2d = np.outer(np.ones(len(r)), munk_ssp(z)) - ssp = xr.DataArray(c_2d, dims=['range', 'depth'], - coords={'range': r, 'depth': z}) + ssp = xr.DataArray( + c_2d, dims=["range", "depth"], coords={"range": r, "depth": z} + ) env = OceanEnvironment2D(sound_speed=ssp, flat_earth_transform=False) assert env.sound_speed.ndim == 2 def test_custom_bathymetry_stored(self): bathy_vals = np.ones(20) * 3500.0 r = np.linspace(0.0, 50e3, 20) - bathy = xr.DataArray(bathy_vals, dims=['range'], coords={'range': r}) + bathy = xr.DataArray(bathy_vals, dims=["range"], coords={"range": r}) env = OceanEnvironment2D(bathymetry=bathy, flat_earth_transform=False) np.testing.assert_array_equal(env.bathymetry.values, bathy_vals) @@ -105,24 +116,26 @@ def test_sound_speed_not_dataarray_raises_type_error(self): def test_sound_speed_3d_raises_value_error(self): da = xr.DataArray( np.ones((5, 10, 20)), - dims=['range', 'depth', 'extra'], - coords={'range': np.arange(5), 'depth': np.arange(10), - 'extra': np.arange(20)} + dims=["range", "depth", "extra"], + coords={ + "range": np.arange(5), + "depth": np.arange(10), + "extra": np.arange(20), + }, ) with pytest.raises(ValueError): OceanEnvironment2D(sound_speed=da) def test_sound_speed_missing_depth_dim_raises_value_error(self): - da = xr.DataArray(np.ones(50), dims=['range'], - coords={'range': np.arange(50)}) + da = xr.DataArray(np.ones(50), dims=["range"], coords={"range": np.arange(50)}) with pytest.raises(ValueError): OceanEnvironment2D(sound_speed=da) def test_2d_sound_speed_missing_range_dim_raises_value_error(self): da = xr.DataArray( np.ones((10, 20)), - dims=['depth', 'extra'], - coords={'depth': np.arange(10), 'extra': np.arange(20)} + dims=["depth", "extra"], + coords={"depth": np.arange(10), "extra": np.arange(20)}, ) with pytest.raises(ValueError): OceanEnvironment2D(sound_speed=da) @@ -132,8 +145,7 @@ def test_bathymetry_not_dataarray_raises_type_error(self): OceanEnvironment2D(bathymetry=np.ones(50)) def test_bathymetry_missing_range_dim_raises_value_error(self): - da = xr.DataArray(np.ones(50), dims=['depth'], - coords={'depth': np.arange(50)}) + da = xr.DataArray(np.ones(50), dims=["depth"], coords={"depth": np.arange(50)}) with pytest.raises(ValueError): OceanEnvironment2D(bathymetry=da) @@ -142,6 +154,7 @@ def test_bathymetry_missing_range_dim_raises_value_error(self): # eflat / eflatinv round-trip # --------------------------------------------------------------------------- + class TestEflat: LAT = 35.0 @@ -149,16 +162,21 @@ def test_depth_roundtrip(self): dep = np.array([100.0, 500.0, 1000.0, 2000.0, 4000.0]) depf, _ = eflat(dep, self.LAT) dep_rec, _ = eflatinv(depf, np.array([self.LAT])) - np.testing.assert_allclose(dep_rec, dep, atol=1.0, - err_msg="Depth round-trip outside 1 m tolerance") + np.testing.assert_allclose( + dep_rec, dep, atol=1.0, err_msg="Depth round-trip outside 1 m tolerance" + ) def test_sound_speed_roundtrip(self): dep = np.array([100.0, 500.0, 1000.0, 2000.0]) cs = np.array([1500.0, 1490.0, 1480.0, 1510.0]) depf, csf = eflat(dep, self.LAT, cs) _, cs_rec = eflatinv(depf, np.array([self.LAT]), csf) - np.testing.assert_allclose(cs_rec, cs, rtol=1e-4, - err_msg="Sound speed round-trip outside 0.01% tolerance") + np.testing.assert_allclose( + cs_rec, + cs, + rtol=1e-4, + err_msg="Sound speed round-trip outside 0.01% tolerance", + ) def test_eflat_increases_depth(self): """Flat-earth transformation should increase effective depths.""" @@ -171,10 +189,11 @@ def test_eflat_increases_depth(self): # OceanEnvironment2D.plot smoke test # --------------------------------------------------------------------------- + class TestOceanEnvironment2DPlot: def test_plot_runs_without_error(self): env = OceanEnvironment2D() fig, ax = plt.subplots() plt.sca(ax) env.plot() - plt.close('all') + plt.close("all") diff --git a/tests/test_physics.py b/tests/test_physics.py index 2f38ee8..8dd80d9 100644 --- a/tests/test_physics.py +++ b/tests/test_physics.py @@ -4,7 +4,7 @@ All tests use flatearth=False to avoid flat-earth correction complicating the comparison with analytical solutions. """ -import os + import pathlib import numpy as np @@ -14,40 +14,41 @@ from pygenray.environment import OceanEnvironment2D, munk_ssp from pygenray.launch_rays import shoot_ray, shoot_rays -FIXTURE_DIR = pathlib.Path(__file__).parent / 'fixtures' +FIXTURE_DIR = pathlib.Path(__file__).parent / "fixtures" # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -def _const_c_env(c0=1500.0, z_max=5000.0, r_max=100e3, - bathy_depth=4500.0, nz=200, nr=20): + +def _const_c_env( + c0=1500.0, z_max=5000.0, r_max=100e3, bathy_depth=4500.0, nz=200, nr=20 +): """Build a range-independent constant sound-speed environment.""" z = np.linspace(0.0, z_max, nz) r = np.linspace(0.0, r_max, nr) c_2d = np.full((nr, nz), c0) - ssp = xr.DataArray(c_2d, dims=['range', 'depth'], - coords={'range': r, 'depth': z}) - bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], - coords={'range': r}) - return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, - flat_earth_transform=False) + ssp = xr.DataArray(c_2d, dims=["range", "depth"], coords={"range": r, "depth": z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=["range"], coords={"range": r}) + return OceanEnvironment2D( + sound_speed=ssp, bathymetry=bathy, flat_earth_transform=False + ) -def _linear_gradient_env(c0=1500.0, g=0.05, z_max=5000.0, r_max=100e3, - bathy_depth=4500.0, nz=500, nr=50): +def _linear_gradient_env( + c0=1500.0, g=0.05, z_max=5000.0, r_max=100e3, bathy_depth=4500.0, nz=500, nr=50 +): """Build a range-independent linear-gradient environment: c(z) = c0 + g*z.""" z = np.linspace(0.0, z_max, nz) r = np.linspace(0.0, r_max, nr) c_1d = c0 + g * z c_2d = np.outer(np.ones(nr), c_1d) - ssp = xr.DataArray(c_2d, dims=['range', 'depth'], - coords={'range': r, 'depth': z}) - bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], - coords={'range': r}) - return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, - flat_earth_transform=False) + ssp = xr.DataArray(c_2d, dims=["range", "depth"], coords={"range": r, "depth": z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=["range"], coords={"range": r}) + return OceanEnvironment2D( + sound_speed=ssp, bathymetry=bathy, flat_earth_transform=False + ) def _munk_env(r_max=100e3, nr=50, nz=600, bathy_depth=5000.0): @@ -56,18 +57,18 @@ def _munk_env(r_max=100e3, nr=50, nz=600, bathy_depth=5000.0): r = np.linspace(0.0, r_max, nr) c_1d = munk_ssp(z) c_2d = np.outer(np.ones(nr), c_1d) - ssp = xr.DataArray(c_2d, dims=['range', 'depth'], - coords={'range': r, 'depth': z}) - bathy = xr.DataArray(np.full(nr, bathy_depth), dims=['range'], - coords={'range': r}) - return OceanEnvironment2D(sound_speed=ssp, bathymetry=bathy, - flat_earth_transform=False) + ssp = xr.DataArray(c_2d, dims=["range", "depth"], coords={"range": r, "depth": z}) + bathy = xr.DataArray(np.full(nr, bathy_depth), dims=["range"], coords={"range": r}) + return OceanEnvironment2D( + sound_speed=ssp, bathymetry=bathy, flat_earth_transform=False + ) # --------------------------------------------------------------------------- # A. Snell's invariant in range-independent constant-c medium # --------------------------------------------------------------------------- + class TestSnellInvariant: """ In a constant sound-speed medium (dc/dz = 0), the ray-parameter @@ -79,15 +80,24 @@ class TestSnellInvariant: def test_p_constant_along_ray(self, user_angle): c0 = 1500.0 env = _const_c_env(c0=c0) - ray = shoot_ray(200.0, 0.0, user_angle, 30e3, 60, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + 200.0, + 0.0, + user_angle, + 30e3, + 60, + env, + rtol=1e-9, + flatearth=False, + debug=False, + ) assert ray is not None, "shoot_ray returned None unexpectedly" # |p| should be constant; sign may flip at boundary reflections abs_p = np.abs(ray.p) p_std = np.std(abs_p) p_mean = np.mean(abs_p) assert p_std / p_mean < 1e-5, ( - f"|p| not constant in constant-c medium: std/mean = {p_std/p_mean:.2e}" + f"|p| not constant in constant-c medium: std/mean = {p_std / p_mean:.2e}" ) @@ -95,6 +105,7 @@ def test_p_constant_along_ray(self, user_angle): # B. Constant sound speed — straight-line rays # --------------------------------------------------------------------------- + class TestConstantSSPStraightLine: """ In a homogeneous medium the ray paths are straight lines. @@ -105,14 +116,15 @@ class TestConstantSSPStraightLine: def test_travel_time_analytical(self): c0 = 1500.0 - user_angle = -10.0 # downward + user_angle = -10.0 # downward theta_ode = abs(user_angle) # internal ODE angle (degrees) - z0 = 200.0 # source depth [m] - R = 20e3 # receiver range [m] + z0 = 200.0 # source depth [m] + R = 20e3 # receiver range [m] env = _const_c_env(c0=c0, r_max=R + 1e3) - ray = shoot_ray(z0, 0.0, user_angle, R, 50, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + z0, 0.0, user_angle, R, 50, env, rtol=1e-9, flatearth=False, debug=False + ) assert ray is not None t_analytical = R / (c0 * np.cos(np.radians(theta_ode))) @@ -126,12 +138,13 @@ def test_final_depth_analytical(self): R = 20e3 env = _const_c_env(c0=c0, r_max=R + 1e3) - ray = shoot_ray(z0, 0.0, user_angle, R, 50, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + z0, 0.0, user_angle, R, 50, env, rtol=1e-9, flatearth=False, debug=False + ) assert ray is not None z_ode_end = z0 + R * np.tan(np.radians(theta_ode)) - z_expected = -z_ode_end # stored sign convention + z_expected = -z_ode_end # stored sign convention assert abs(ray.z[-1] - z_expected) / abs(z_expected) < 1e-3 def test_p_constant_in_const_c(self): @@ -142,20 +155,26 @@ def test_p_constant_in_const_c(self): R = 20e3 env = _const_c_env(c0=c0, r_max=R + 1e3) - ray = shoot_ray(z0, 0.0, user_angle, R, 50, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + z0, 0.0, user_angle, R, 50, env, rtol=1e-9, flatearth=False, debug=False + ) assert ray is not None p_expected = -np.sin(np.radians(theta_ode)) / c0 # stored sign - np.testing.assert_allclose(ray.p, p_expected, - rtol=1e-5, atol=0, - err_msg="p not constant in homogeneous medium") + np.testing.assert_allclose( + ray.p, + p_expected, + rtol=1e-5, + atol=0, + err_msg="p not constant in homogeneous medium", + ) # --------------------------------------------------------------------------- # C. Linear gradient — turning depth # --------------------------------------------------------------------------- + class TestLinearGradientTurningDepth: """ For c(z) = c0 + g·z the Hamiltonian is conserved: @@ -170,9 +189,9 @@ class TestLinearGradientTurningDepth: """ C0 = 1500.0 - G = 0.05 # m/s/m + G = 0.05 # m/s/m Z_SRC = 200.0 # source depth [m] - THETA = 20.0 # degrees downward (user passes -20) + THETA = 20.0 # degrees downward (user passes -20) def _z_turn_analytical(self): c_source = self.C0 + self.G * self.Z_SRC @@ -182,12 +201,21 @@ def test_turning_depth_approx(self): z_turn = self._z_turn_analytical() env = _linear_gradient_env(c0=self.C0, g=self.G) - ray = shoot_ray(self.Z_SRC, 0.0, -self.THETA, 80e3, 400, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + self.Z_SRC, + 0.0, + -self.THETA, + 80e3, + 400, + env, + rtol=1e-9, + flatearth=False, + debug=False, + ) assert ray is not None # ray.z uses sign convention: negative = deep - z_turn_numerical = -np.min(ray.z) # max depth reached + z_turn_numerical = -np.min(ray.z) # max depth reached assert abs(z_turn_numerical - z_turn) < 50.0, ( f"Turning depth: expected {z_turn:.1f} m, got {z_turn_numerical:.1f} m" ) @@ -196,19 +224,28 @@ def test_hamiltonian_conserved_linear_gradient(self): """The Hamiltonian H = sqrt(1/c(z)² − p_ode²) is conserved.""" env = _linear_gradient_env(c0=self.C0, g=self.G) - ray = shoot_ray(self.Z_SRC, 0.0, -self.THETA, 80e3, 400, - env, rtol=1e-9, flatearth=False, debug=False) + ray = shoot_ray( + self.Z_SRC, + 0.0, + -self.THETA, + 80e3, + 400, + env, + rtol=1e-9, + flatearth=False, + debug=False, + ) assert ray is not None - z_ode = -ray.z # positive depth (ODE convention) - p_ode = -ray.p # positive for downward ray + z_ode = -ray.z # positive depth (ODE convention) + p_ode = -ray.p # positive for downward ray c_along = self.C0 + self.G * z_ode H = np.sqrt(1.0 / c_along**2 - p_ode**2) H_std = np.std(H) H_mean = np.mean(H) assert H_std / H_mean < 1e-4, ( - f"Hamiltonian not conserved: std/mean = {H_std/H_mean:.2e}" + f"Hamiltonian not conserved: std/mean = {H_std / H_mean:.2e}" ) @@ -216,6 +253,7 @@ def test_hamiltonian_conserved_linear_gradient(self): # D. Hamiltonian conservation in range-independent Munk profile # --------------------------------------------------------------------------- + class TestMunkHamiltonianConservation: """ In any range-independent environment the Hamiltonian @@ -227,10 +265,19 @@ class TestMunkHamiltonianConservation: @pytest.mark.parametrize("user_angle", [-5.0, -10.0, -15.0]) def test_hamiltonian_conserved_munk(self, user_angle): env = _munk_env(r_max=100e3) - z_src = 1000.0 # below SOFAR axis — ray stays in the water column - - ray = shoot_ray(z_src, 0.0, user_angle, 100e3, 200, - env, rtol=1e-9, flatearth=False, debug=False) + z_src = 1000.0 # below SOFAR axis — ray stays in the water column + + ray = shoot_ray( + z_src, + 0.0, + user_angle, + 100e3, + 200, + env, + rtol=1e-9, + flatearth=False, + debug=False, + ) assert ray is not None, "shoot_ray returned None; ray may have exited domain" z_ode = -ray.z @@ -251,7 +298,7 @@ def test_hamiltonian_conserved_munk(self, user_angle): H_std = np.std(H_valid) H_mean = np.mean(H_valid) assert H_std / H_mean < 1e-3, ( - f"Hamiltonian not conserved in Munk profile: std/mean = {H_std/H_mean:.2e}" + f"Hamiltonian not conserved in Munk profile: std/mean = {H_std / H_mean:.2e}" ) @@ -259,17 +306,22 @@ def test_hamiltonian_conserved_munk(self, user_angle): # E. Regression / golden-file tests # --------------------------------------------------------------------------- + def _regenerate_fixture(): """Run shoot_rays and save regression fixture. Call manually to regenerate.""" env = _munk_env(r_max=50e3, nr=30, nz=400) angles = [-8.0, -4.0, 0.0, 4.0, 8.0] - rf = shoot_rays(1300.0, 0.0, angles, 50e3, 50, env, - n_processes=1, debug=False, flatearth=False) + rf = shoot_rays( + 1300.0, 0.0, angles, 50e3, 50, env, n_processes=1, debug=False, flatearth=False + ) FIXTURE_DIR.mkdir(parents=True, exist_ok=True) np.savez( - FIXTURE_DIR / 'munk_regression.npz', - ts=rf.ts, zs=rf.zs, ps=rf.ps, - n_botts=rf.n_botts, n_surfs=rf.n_surfs, + FIXTURE_DIR / "munk_regression.npz", + ts=rf.ts, + zs=rf.zs, + ps=rf.ps, + n_botts=rf.n_botts, + n_surfs=rf.n_surfs, thetas=rf.thetas, ) return rf @@ -284,24 +336,35 @@ class TestMunkRegression: or run pytest with --regenerate-physics flag. """ - FIXTURE = FIXTURE_DIR / 'munk_regression.npz' + FIXTURE = FIXTURE_DIR / "munk_regression.npz" ANGLES = [-8.0, -4.0, 0.0, 4.0, 8.0] def _run_rays(self): env = _munk_env(r_max=50e3, nr=30, nz=400) - return shoot_rays(1300.0, 0.0, self.ANGLES, 50e3, 50, env, - n_processes=1, debug=False, flatearth=False) + return shoot_rays( + 1300.0, + 0.0, + self.ANGLES, + 50e3, + 50, + env, + n_processes=1, + debug=False, + flatearth=False, + ) def test_regression(self, request): - regenerate = request.config.getoption('--regenerate-physics', - default=False) + regenerate = request.config.getoption("--regenerate-physics", default=False) if regenerate or not self.FIXTURE.exists(): rf = self._run_rays() self.FIXTURE.parent.mkdir(parents=True, exist_ok=True) np.savez( self.FIXTURE, - ts=rf.ts, zs=rf.zs, ps=rf.ps, - n_botts=rf.n_botts, n_surfs=rf.n_surfs, + ts=rf.ts, + zs=rf.zs, + ps=rf.ps, + n_botts=rf.n_botts, + n_surfs=rf.n_surfs, thetas=rf.thetas, ) if regenerate: @@ -312,20 +375,22 @@ def test_regression(self, request): ref = np.load(self.FIXTURE) rf = self._run_rays() - np.testing.assert_allclose(rf.ts, ref['ts'], atol=1e-6, - err_msg="Travel times changed") - np.testing.assert_allclose(rf.zs, ref['zs'], atol=0.1, - err_msg="Depths changed") - np.testing.assert_allclose(rf.ps, ref['ps'], atol=0.1, - err_msg="Ray parameters changed") - np.testing.assert_array_equal(rf.n_botts, ref['n_botts']) - np.testing.assert_array_equal(rf.n_surfs, ref['n_surfs']) + np.testing.assert_allclose( + rf.ts, ref["ts"], atol=1e-6, err_msg="Travel times changed" + ) + np.testing.assert_allclose(rf.zs, ref["zs"], atol=0.1, err_msg="Depths changed") + np.testing.assert_allclose( + rf.ps, ref["ps"], atol=0.1, err_msg="Ray parameters changed" + ) + np.testing.assert_array_equal(rf.n_botts, ref["n_botts"]) + np.testing.assert_array_equal(rf.n_surfs, ref["n_surfs"]) # --------------------------------------------------------------------------- # F. Near-vertical ray — regression for ZeroDivisionError in derivsrd # --------------------------------------------------------------------------- + class TestNearVerticalRay: """ Shooting a ray at or near 90° causes pz = sin(θ)/c ≈ 1/c, making @@ -338,8 +403,17 @@ def test_near_vertical_no_crash(self): """A ray launched at 89.9° (user: -89.9) must not raise ZeroDivisionError.""" env = _const_c_env(c0=1500.0, z_max=5000.0, r_max=100e3, bathy_depth=4500.0) try: - ray = shoot_ray(200.0, 0.0, -89.9, 10e3, 50, env, - rtol=1e-6, flatearth=False, debug=False) + shoot_ray( + 200.0, + 0.0, + -89.9, + 10e3, + 50, + env, + rtol=1e-6, + flatearth=False, + debug=False, + ) except ZeroDivisionError: pytest.fail("ZeroDivisionError raised for near-vertical ray (θ ≈ 89.9°)") @@ -347,8 +421,17 @@ def test_exactly_vertical_no_crash(self): """A ray launched at exactly 90° must not raise ZeroDivisionError.""" env = _const_c_env(c0=1500.0, z_max=5000.0, r_max=100e3, bathy_depth=4500.0) try: - ray = shoot_ray(200.0, 0.0, -90.0, 10e3, 50, env, - rtol=1e-6, flatearth=False, debug=False) + shoot_ray( + 200.0, + 0.0, + -90.0, + 10e3, + 50, + env, + rtol=1e-6, + flatearth=False, + debug=False, + ) except ZeroDivisionError: pytest.fail("ZeroDivisionError raised for exactly vertical ray (θ = 90°)") @@ -357,8 +440,16 @@ def test_steep_rays_no_crash(self, angle): """Steep rays near-vertical must not crash regardless of launch angle.""" env = _munk_env(r_max=50e3) try: - shoot_ray(1000.0, 0.0, angle, 10e3, 50, env, - rtol=1e-6, flatearth=False, debug=False) + shoot_ray( + 1000.0, + 0.0, + angle, + 10e3, + 50, + env, + rtol=1e-6, + flatearth=False, + debug=False, + ) except ZeroDivisionError: pytest.fail(f"ZeroDivisionError raised for steep ray at θ = {angle}°") - diff --git a/tests/test_ray_objects.py b/tests/test_ray_objects.py index 509dd68..9225230 100644 --- a/tests/test_ray_objects.py +++ b/tests/test_ray_objects.py @@ -1,6 +1,7 @@ """ Tests for pygenray.ray_objects: Ray, RayFan. """ + import numpy as np import pytest import scipy.io @@ -13,19 +14,27 @@ # Ray # --------------------------------------------------------------------------- + class TestRay: N = 10 R = 10000.0 - def _make_ray(self, launch_angle=-10.0, source_depth=100.0, - n_bottom=0, n_surface=0): + def _make_ray( + self, launch_angle=-10.0, source_depth=100.0, n_bottom=0, n_surface=0 + ): r = np.linspace(0.0, self.R, self.N) t = r / 1500.0 z_ode = np.linspace(source_depth, source_depth + self.R * 0.01, self.N) p_ode = np.ones(self.N) * np.sin(np.radians(abs(launch_angle) + 1e-3)) / 1500.0 y = np.vstack([t, z_ode, p_ode]) - return Ray(r=r, y=y, n_bottom=n_bottom, n_surface=n_surface, - launch_angle=launch_angle, source_depth=source_depth), y + return Ray( + r=r, + y=y, + n_bottom=n_bottom, + n_surface=n_surface, + launch_angle=launch_angle, + source_depth=source_depth, + ), y def test_attribute_shapes(self): ray, _ = self._make_ray() @@ -57,14 +66,14 @@ def test_optional_launch_angle_not_set(self): t = r / 1500.0 y = np.vstack([t, np.ones(self.N) * 100.0, np.ones(self.N) * 0.1]) ray = Ray(r=r, y=y, n_bottom=0, n_surface=0) - assert not hasattr(ray, 'launch_angle') + assert not hasattr(ray, "launch_angle") def test_optional_source_depth_not_set(self): r = np.linspace(0.0, self.R, self.N) t = r / 1500.0 y = np.vstack([t, np.ones(self.N) * 100.0, np.ones(self.N) * 0.1]) ray = Ray(r=r, y=y, n_bottom=0, n_surface=0) - assert not hasattr(ray, 'source_depth') + assert not hasattr(ray, "source_depth") def test_n_bottom_n_surface_stored(self): ray, _ = self._make_ray(n_bottom=3, n_surface=1) @@ -75,13 +84,14 @@ def test_plot_smoke(self): ray, _ = self._make_ray() plt.figure() ray.plot() - plt.close('all') + plt.close("all") # --------------------------------------------------------------------------- # RayFan # --------------------------------------------------------------------------- + class TestRayFan: M = 3 N = 10 @@ -100,8 +110,14 @@ def _make_rays(self, M=None, N=None, R=None): p_ode = np.ones(N) * np.sin(np.radians(abs(theta) + 1e-3)) / 1500.0 y = np.vstack([t, z_ode, p_ode]) rays.append( - Ray(r=r, y=y, n_bottom=i % 2, n_surface=0, - launch_angle=theta, source_depth=100.0 + i * 50) + Ray( + r=r, + y=y, + n_bottom=i % 2, + n_surface=0, + launch_angle=theta, + source_depth=100.0 + i * 50, + ) ) return rays @@ -119,7 +135,7 @@ def test_shapes(self, simple_rayfan): assert rf.source_depths.shape == (self.M,) def test_ray_ids_set_on_construction(self, simple_rayfan): - assert hasattr(simple_rayfan, 'ray_ids') + assert hasattr(simple_rayfan, "ray_ids") assert len(simple_rayfan.ray_ids) == self.M # --- compute_rayids --- @@ -187,8 +203,9 @@ def test_getitem_int_array_returns_rayfan(self, simple_rayfan): result = simple_rayfan[idx] assert isinstance(result, RayFan) assert len(result) == 2 - np.testing.assert_array_equal(result.thetas, - simple_rayfan.thetas[np.array([0, 2])]) + np.testing.assert_array_equal( + result.thetas, simple_rayfan.thetas[np.array([0, 2])] + ) # --- __add__ --- @@ -225,54 +242,61 @@ def test_add_non_rayfan_raises_type_error(self, simple_rayfan): # --- save_mat --- def test_save_mat_creates_file(self, simple_rayfan, tmp_path): - path = str(tmp_path / 'test_rayfan.mat') + path = str(tmp_path / "test_rayfan.mat") simple_rayfan.save_mat(path) - assert (tmp_path / 'test_rayfan.mat').exists() + assert (tmp_path / "test_rayfan.mat").exists() def test_save_mat_loadable(self, simple_rayfan, tmp_path): - path = str(tmp_path / 'test_rayfan.mat') + path = str(tmp_path / "test_rayfan.mat") simple_rayfan.save_mat(path) data = scipy.io.loadmat(path) - assert 'rayfan' in data + assert "rayfan" in data def test_save_mat_contains_required_keys(self, simple_rayfan, tmp_path): - path = str(tmp_path / 'test_rayfan.mat') + path = str(tmp_path / "test_rayfan.mat") simple_rayfan.save_mat(path) data = scipy.io.loadmat(path) - rayfan = data['rayfan'] + rayfan = data["rayfan"] # MATLAB struct stored as structured array; check dtype field names - expected_keys = {'thetas', 'xs', 'ts', 'zs', 'ps', - 'n_botts', 'n_surfs', 'source_depths'} + expected_keys = { + "thetas", + "xs", + "ts", + "zs", + "ps", + "n_botts", + "n_surfs", + "source_depths", + } actual_keys = set(rayfan.dtype.names) assert expected_keys <= actual_keys def test_save_mat_values_match(self, simple_rayfan, tmp_path): - path = str(tmp_path / 'test_rayfan.mat') + path = str(tmp_path / "test_rayfan.mat") simple_rayfan.save_mat(path) data = scipy.io.loadmat(path) - rayfan = data['rayfan'] - saved_thetas = rayfan['thetas'][0, 0].flatten() - np.testing.assert_allclose(saved_thetas, simple_rayfan.thetas, - atol=1e-10) + rayfan = data["rayfan"] + saved_thetas = rayfan["thetas"][0, 0].flatten() + np.testing.assert_allclose(saved_thetas, simple_rayfan.thetas, atol=1e-10) # --- Plotting smoke tests --- def test_plot_ray_fan_smoke(self, simple_rayfan): plt.figure() simple_rayfan.plot_ray_fan() - plt.close('all') + plt.close("all") def test_plot_time_front_smoke(self, simple_rayfan): plt.figure() simple_rayfan.plot_time_front() - plt.close('all') + plt.close("all") def test_plot_time_front_include_lines_smoke(self, simple_rayfan): plt.figure() simple_rayfan.plot_time_front(include_lines=True) - plt.close('all') + plt.close("all") def test_plot_depth_v_angle_smoke(self, simple_rayfan): plt.figure() simple_rayfan.plot_depth_v_angle() - plt.close('all') + plt.close("all")