Fix some problems with the PT-DWBA model

src/echosms/ptdwbamodel.py

@@ -17,7 +17,7 @@ def __init__(self):
         self.shapes = ['unrestricted voxel-based']
         self.max_ka = 20
 
-    def calculate_ts_single(self, volume, theta, phi, f, voxel_size, target_rho, target_c):
+    def calculate_ts_single(self, volume, theta, phi, f, voxel_size, rho, c, **kwargs):
         """Phase-tracking distorted-wave Born approximation scattering model.
 
         Implements the phase-tracking distorted-wave Born approximation
@@ -52,11 +52,11 @@ def calculate_ts_single(self, volume, theta, phi, f, voxel_size, target_rho, tar
             The size of the voxels in `volume` [m], ordered (_x_, _y_, _z_).
             This code assumes that the voxels are cubes so _y_ and _z_ are currently irrevelant.
 
-        target_rho : iterable[float]
+        rho : iterable[float]
             Densities of each material. Must have at least the same number of entries as unique
             integers in `volume` [kg/m³].
 
-        target_c : iterable[float]
+        c : iterable[float]
             Sound speed of each material. Must have at least the same number of entries as unique
             integers in `volume` [m/s].
 
@@ -82,8 +82,8 @@ def calculate_ts_single(self, volume, theta, phi, f, voxel_size, target_rho, tar
         125(1), 73-88. <https://doi.org/10.1121/1.3021298>
         """
         # Make sure things are numpy arrays
-        target_rho = np.array(target_rho)
-        target_c = np.array(target_c)
+        rho = np.array(rho)
+        c = np.array(c)
         voxel_size = np.array(voxel_size)
 
         # volume of the voxels [m^3]
@@ -116,17 +116,17 @@ def calculate_ts_single(self, volume, theta, phi, f, voxel_size, target_rho, tar
             raise ValueError('The integers in volume must include all values in the series '
                              '(0, 1, 2, ..., n), where n is the largest integer in volume.')
 
-        if not len(target_rho) >= len(categories):
+        if not len(rho) >= len(categories):
             raise ValueError('The target_rho variable must contain at least as many values as '
                              'unique integers in the volume variable.')
 
-        if not len(target_c) >= len(categories):
+        if not len(c) >= len(categories):
             raise ValueError('The target_c variable must contain at least as many values '
                              'as unique integers in the volume variable.')
 
         # density and sound speed ratios for all object materials
-        g = target_rho[1:] / target_rho[0]
-        h = target_c[1:] / target_c[0]
+        g = rho[1:] / rho[0]
+        h = c[1:] / c[0]
 
         # Do the pitch and roll rotations
         v = ndimage.rotate(volume, -theta, axes=(0, 2), order=0)
@@ -135,7 +135,7 @@ def calculate_ts_single(self, volume, theta, phi, f, voxel_size, target_rho, tar
         categories = np.unique(v)  # or just take the max?
 
         # wavenumbers in the various media
-        k = 2.0*np.pi * f / target_c
+        k = 2.0*np.pi * f / c
 
         # DWBA coefficients
         # amplitudes in media 1,2,...,n

src/echosms/utils.py

@@ -154,7 +154,8 @@ def spherical_jnpp(n: int, z: float) -> float:
     return 1./z**2 * ((n**2-n-z**2)*spherical_jn(n, z) + 2.*z*spherical_jn(n+1, z))
 
 
-def prolate_swf(m: int, lnum: int, c: float, xi: float, eta: Iterable[float], norm=False):
+def prolate_swf(m: int, lnum: int, c: float, xi: float, eta: Iterable[float],
+                norm=False) -> tuple:
     """Calculate prolate spheroidal wave function values.
 
     Parameters
@@ -179,6 +180,14 @@ def prolate_swf(m: int, lnum: int, c: float, xi: float, eta: Iterable[float], no
         same norm as the corresponding associated Legendre function (that is,
         the Meixner and Schafke normalization scheme). This norm becomes very
         large as `m` becomes large.
+    angular : bool
+        Whether to return values for the angular functions.
+    radial : bool
+        Whether to return values for the radial functions.
+    angular_derivative : bool
+        Whether to return values for the derivative of the angular functions.
+    radial_derivative : bool
+        Whether to return values for the derivative of the radial functions.
 
     Returns
     -------
@@ -215,7 +224,7 @@ def prolate_swf(m: int, lnum: int, c: float, xi: float, eta: Iterable[float], no
     This method encapsulates the prolate spheroidal wave function code for non complex
     arguments (van Buren & Boisvert, 2002, and van Buren & Boisvert, 2024), available on
     [github](https://github.com/MathieuandSpheroidalWaveFunctions). This code is in Fortran90
-    and code was interfaced to Python using `numpy.f2py` then wrapped with the current method to
+    and was interfaced to Python using `numpy.f2py` then wrapped with the current method to
     provide a convenient interface for use in the [`PSMSModel`][psmsmodel] class.
 
     References

src/example_code.py

@@ -19,6 +19,27 @@
 bmf = bm.freq_dataset
 bmt = bm.angle_dataset
 
+
+def plot_compare(f1, ts1, label1, f2, ts2, label2, title):
+    """Plot together two ts results TS."""
+    jech_index = np.mean(np.abs(np.array(ts1) - np.array(ts2)))
+    # Plot the mss model and benchmark results
+    fig, axs = plt.subplots(2, 1, sharex=True)
+    axs[0].plot(f1/1e3, ts1, label=label1)
+    axs[0].plot(f2/1e3, ts2, label=label2)
+    axs[0].set_ylabel('TS re 1 m$^2$ [dB]')
+    axs[0].legend(frameon=False, fontsize=6)
+
+    # Plot difference between the two ts datasets
+    axs[1].plot(f1/1e3, np.array(ts1)-np.array(ts2), color='black')
+    axs[1].set_xlabel('Frequency [kHz]')
+    axs[1].set_ylabel(r'$\Delta$ TS [dB]')
+    axs[1].annotate(f'{jech_index:.2f} dB', (0.05, 0.80), xycoords='axes fraction',
+                    backgroundcolor=[.8, .8, .8])
+    plt.suptitle(title)
+    plt.show()
+
+
 # %% ###############################################################################################
 # Run the benchmark models and compare to the frequency-varying benchmark results.
 
@@ -68,23 +89,7 @@
     # The finite benchmark TS values
     bm_ts = bmf[bm_name][~np.isnan(bmf[bm_name])]
 
-    jech_index = np.mean(np.abs(ts - bm_ts))
-
-    # Plot the mss model and benchmark results
-    fig, axs = plt.subplots(2, 1, sharex=True)
-    axs[0].plot(m['f']/1e3, ts, label='echoSMs')
-    axs[0].plot(m['f']/1e3, bm_ts, label='Benchmark')
-    axs[0].set_ylabel('TS re 1 m$^2$ [dB]')
-    axs[0].legend(frameon=False, fontsize=6)
-
-    # Plot difference between benchmark values and newly calculated mss model values
-    axs[1].plot(m['f']*1e-3, ts-bm_ts, color='black')
-    axs[1].set_xlabel('Frequency [kHz]')
-    axs[1].set_ylabel(r'$\Delta$ TS [dB]')
-    axs[1].annotate(f'{jech_index:.2f} dB', (0.05, 0.80), xycoords='axes fraction',
-                    backgroundcolor=[.8, .8, .8])
-    plt.suptitle(name)
-    plt.show()
+    plot_compare(m['f'], ts, s['benchmark_model'], m['f'], bm_ts, 'Benchmark', name)
 
 # %% ###############################################################################################
 # Run the benchmark models and compare to the angle-varying benchmark results.
@@ -223,39 +228,46 @@
 # %% ###############################################################################################
 # Example of PT-DWBA model
 
-# A simple, minimal example
-p = {
-     'f': 38000,
-     'theta': 45,
-     'phi': 0,
-     'voxel_size': [0.0005, 0.0005, 0.0005],
-     'target_rho': [1024, 1025, 1026],
-     'target_c': [1480, 1490, 1500],
-     'volume': np.array([[[0, 0, 0, 0],
-                          [0, 1, 2, 0],
-                          [0, 0, 0, 0]],
-                         [[0, 0, 0, 0],
-                          [0, 1, 0, 2],
-                          [0, 0, 0, 0]],
-                         [[0, 0, 0, 0],
-                          [0, 1, 0, 2],
-                          [0, 0, 0, 0]],
-                         [[0, 0, 0, 0],
-                          [0, 1, 0, 1],
-                          [0, 0, 0, 0]],
-                         [[0, 0, 0, 0],
-                          [0, 1, 0, 1],
-                          [0, 0, 0, 0]]])}
+# The benchmark sphere model
+m = rm.parameters('weakly scattering sphere')
+
+# make a 3d matrix of 0's and 1's and set to 1 for the sphere
+# and 0 for not the sphere
+m['voxel_size'] = (0.0001, 0.0001, 0.0001)  # [m]
+x = np.arange(-m['a'], m['a'], m['voxel_size'][0])
+(X, Y, Z) = np.meshgrid(x, x, x)
+
+m['volume'] = (np.sqrt(X**2 + Y**2 + Z**2) <= m['a']).astype(int)
+m['f'] = 38000
+m['theta'] = 90
+m['phi'] = 0
+m['rho'] = [m['medium_rho'], m['target_rho']]
+m['c'] = [m['medium_c'], m['target_c']]
+
+freqs = bmf['Frequency_kHz']*1e3
 
 pt = PTDWBAModel()
-pt.calculate_ts_single(**p)
+dwba_ts = []
+for f in freqs:
+    print(f/1e3)
+    m['f'] = f
+    # PTDWBAModel() doesn't yet support calling via the calculate_ts() call.
+    dwba_ts.append(pt.calculate_ts_single(**m))
 
-# Note that PTDWBAModel() doesn't yet support calling via the calculate_ts() call.
+plot_compare(freqs, dwba_ts, 'PT-DWBA',
+             freqs, bmf['Sphere_WeaklyScattering'], 'Benchmark',
+             'weakly scattering sphere')
 
-# The benchmark sphere model
-m = rm.parameters('weakly scattering sphere')
+# So, this PT_DWBA on a weakly scattering sphere is also different from the benchmark
+# TS values. Hmmmm. Now look at the difference between the PT-DWBA and MSS model runs...
+
+mss = MSSModel()
+mm = rm.parameters('weakly scattering sphere')
+mm['f'] = freqs
+mm['theta'] = 90.0
 
-# Add in the things that the PT-DWBA model needs
+mss_ts = mss.calculate_ts(mm)
 
-# Make up the voxel dataset based on the parameters in m
-voxel_size = ()  # same as used in Jech et al (2015)
+plot_compare(freqs, ts, 'PT-DWBA',
+             freqs, mss_ts, 'MSS',
+             'weakly scattering sphere')