[BUG] Numerical instabilities when medium.alpha_coeff is close to 1

[BruceAbrams]

Describe the bug
Setting the attenuation exponent value in the range between 0.95 and 1.03 causes the output of kspaceFirstOrder2D to have NaNs. This issue is not present in MATLAB. I'm trying to understand if that's a known problem, or if I'm doing something wrong.

To Reproduce
Here's sample code:

import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import RegularGridInterpolator
from kwave.data import Vector
from kwave.kgrid import kWaveGrid
from kwave.kmedium import kWaveMedium
from kwave.ksensor import kSensor
from kwave.ksource import kSource
from kwave.kspaceFirstOrder2D import kspaceFirstOrder2D
from kwave.options.simulation_execution_options import SimulationExecutionOptions
from kwave.options.simulation_options import SimulationOptions
from kwave.utils.mapgen import make_cart_circle
from scipy.signal import gausspulse

from kwave.utils.conversion import cart2grid

from skimage.data import shepp_logan_phantom
from skimage.transform import rescale

def soundSpeedPhantom2D(Xi, Yi, image=None, rotAngle=2.85*np.pi):
    """
    INPUTS:
        Xi, Yi - meshgrid of points over which sound speed is defined
        option - 1 for breast CT; 2 for breast MRI
    OUTPUTS:
        C - sound speed map of grind [m/s]
        c_bkgnd - background spind speed map [m/s]
    """

    if image is not None:
        img = image
    else:
        img = np.asarray(Image.open('breast_mri.jpg'))
        
    img = img/np.max(img)
    Ny, Nx = np.shape(img)
    dx = 0.00015 # Grid spacing [m]
    dy = dx
    x = (np.arange(Nx) - (Nx - 1) / 2.0) * dx
    y = (np.arange(Ny) - (Ny - 1) / 2.0) * dy

    # Create sound speed imnage [m/s]
    c_min = 1420
    c_max = 1640
    c = img * (c_max-c_min) + c_min
    c_bckgd = c[0,0]

    # Put sound speed map on input meshgrid
    R = np.sqrt(Xi**2 + Yi**2)
    T = np.arctan2(Yi, Xi) - rotAngle
    
    interp = RegularGridInterpolator((y,x), c, bounds_error=False, fill_value=c_bckgd)
    points = np.stack((R * np.sin(T), R*np.cos(T)), axis=-1)
    C = interp(points)

    return C, c_bckgd

image = shepp_logan_phantom()
image = rescale(image, scale=4, mode='reflect', channel_axis=None)
[Ny, Nx] = np.shape(image)

dxi = 0.15e-3
xmax = 120e-3
xi = np.linspace(-xmax, xmax, Nx)
zi = xi
Nxi = len(xi)
Nzi = len(zi)

Xi, Zi = np.meshgrid(xi, zi)
C, c_bckgd = soundSpeedPhantom2D(Xi, Zi, image=image, rotAngle=0)

# Create k-Wave Grid

N = Vector([Nzi, Nxi])
dx = Vector([dxi, dxi])
kgrid = kWaveGrid(N, dx)

t_end = 1.3 * Nzi * dxi / np.min(C)
cfl = 0.3 

kgrid.makeTime(C, cfl=cfl, t_end=t_end)
t_array = kgrid.t_array
dt = kgrid.dt

fs = 1.0/dt

# Attenuation map

atten_bkgnd = 0.0025             # dB/(MHz^y cm)
sos2atten = 6e-2                 # conversion factor
atten_varying = sos2atten * np.abs(C - C[0, 0])
atten = atten_bkgnd + atten_varying

y_atten = 1.01                   # power-law exponent, unstable around 0.95 - 1.03

# Create Transducer Ring

circle_radius = 110e-3           # 110 mm
numElements = 512 

circle_rad_pixels = int(np.floor(circle_radius / dxi))


focus_pos = Vector([0, 0])  # [m]

element_pos = make_cart_circle(circle_radius, numElements, focus_pos, plot_circle=False)

sensor_mask,_,_ = cart2grid(kgrid, element_pos)

rho0 = 1000.0                                # kg/m^3

medium_sound_speed = C                       # m/s
medium_density = rho0 * np.ones_like(C)      # kg/m^3
medium_alpha_coeff = atten                   # dB/(MHz^y cm)
medium_alpha_power = y_atten                 # exponent

alpha_mode = 'no_dispersion'                 # ignore velocity dispersion

medium = kWaveMedium(
    sound_speed=medium_sound_speed,
    density=medium_density,
    alpha_power=medium_alpha_power,
    alpha_coeff=medium_alpha_coeff,
    alpha_mode=alpha_mode
)

# Define Properties of the Tone Burst Used to Drive Transducer
fracBW = 0.75                         # fractional bandwidth
fTx = 1.0e6                           # transmit frequency [Hz]

tc = gausspulse('cutoff', fc=fTx, bw = fracBW, bwr=-6, tpr=-80) #fracBW @ -6dB, cutoff time at -80dB
t = np.arange(-np.ceil(tc/kgrid.dt), np.ceil(tc/kgrid.dt+1)) * kgrid.dt
src_amplitude = 1.0e1
tx_signal = src_amplitude * gausspulse(t=t, fc=fTx, bw=fracBW, bwr=-6)
tx_signal= tx_signal[np.newaxis, :] #need to add an extra dimension, so that k-Wave doesn't complain

emission_samples = tx_signal.size

# Define a Binary Sensor Mask

sensor = kSensor(mask=sensor_mask)

simulation_options = SimulationOptions(
    pml_inside=False,
    smooth_p0=False,
    save_to_disk=True,
    data_cast="single",
    pml_alpha=10
)
execution_options = SimulationExecutionOptions(is_gpu_simulation=True, show_sim_log=False)

tx_elmt_idx = 0
source = kSource()
pos = element_pos[:,tx_elmt_idx].reshape(-1,1)
sensor_mask2,_,_ = cart2grid(kgrid, pos)
source.p_mask = sensor_mask2

source.p = tx_signal
source.p_mode = 'dirichlet'

sensor_data_unordered = kspaceFirstOrder2D(
    kgrid=kgrid,
    source=source,
    sensor=sensor,
    medium=medium,
    simulation_options=simulation_options,
    execution_options=execution_options
)

sensor_data = sensor_data_unordered['p'].T

Plotting sensor_data, we get the following:

Expected behavior
The output of sensor_data should not have any NaNs and should look something like this

• OS: Linux