Sound speed profile from impulse response

Reference

This tutorial is adapted from Example D presented in:

• Mandar Chitre, “Differentiable Ocean Acoustic Propagation Modeling,” in OCEANS 2023 IEEE/MTS – Limerick, 5-8 June 2023.

Problem statement

Consider a setup with a 1 kHz acoustic source at 1 km depth that sends a broadband pulse once every week. A receiver 10 km away at a depth of 800 m measures the impulse response from the received broadband pulse. We assume that we have an initial sound speed profile measurement with a CTD at the start of the experiment. The sound speed profile changes over the weeks of the experiment, and we wish to track the changes using the measured impulse response every week.

Key idea

Since the sound speed profile is an unknown function of depth, we model it using a small 3-layer neural network. We initialize the parameters by training the neural network on the known/estimated sound speed profile from the previous week. We use the neural network to provide sound speed estimates to RaySolver, thus effectively creating a model that combines a numerical physics-based differential equation solver with a data-driven neural network. The hybrid model returns a vector of predicted delays of the first few multipath arrivals in the impulse response. We minimize a loss function that measures the difference between the predictions and measurements using gradient descent and automatic differentiation. This essentially trains the neural network to approximate the sound speed profile via a loss function that utilizes the propagation model. At the end of each week’s training, we get a revised estimate of the sound speed profile.

Dataset

To illustrate this idea, let us generate a synthetic dataset:

using UnderwaterAcoustics
using AcousticRayTracers

# define Munk sound speed profile
struct MunkSSP <: UnderwaterAcoustics.DepthDependent end

function (ssp::MunkSSP)(pos)
  ϵ = 0.00737
  z̃ = 2.0 * (-pos.z - 1300.0) / 1300.0
  1500.0 * (1.0 + ϵ * (z̃ - 1 + exp(-z̃)))
end

# simulate acoustic propagation using RaySolver to generate impulse response
function model(h, r, d1, d2, c, n)
  env = UnderwaterEnvironment(bathymetry=h, seabed=Rock, soundspeed=c)
  pm = RaySolver(env)
  tx = AcousticSource(0.0, -d1, 1000.0)
  rx = AcousticReceiver(r, -d2)
  arr = arrivals(pm, tx, rx)
  t = getfield.(arr, :t)        # get a.t for each arrival a in arr
  (t[2:n] .- t[1]) .* 1000      # relative time in milliseconds
end

# estimated sound speed profile from previous week
prev_depths = [0, 500.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0]
prev_soundspeeds = [1518.6, 1494.5, 1490.6, 1505.8, 1530.3, 1556.7, 1583.6]

# generate data with first 6 multipath delays
data = model(5000.0, 10000.0, 1000.0, 1800.0, MunkSSP(), 7)

6-element Vector{Float64}:
  178.42786295390846
 1406.6963669822208
 2245.763511823167
 3008.5299257637885
 4014.4245725479186
 6395.490526963947

The previous week’s sound speed profile is shown in the plot above as solid circles, whereas the true sound speed profile is shown as a blue line.

Neural network based propagation model

We don’t know the sound speed profile, so we will use a small 3-layer neural network to model it:

using SimpleChains
using StaticArrays
using Random: MersenneTwister

# neural network based sound speed profile data type
struct NeuralSSP{T1,T2} <: UnderwaterAcoustics.DepthDependent
  model::T1                 # neural network structure
  ps::T2                    # neural network parameters
  max_depth::Float64        # used for scaling the input
  min_soundspeed::Float64   # used for scaling the output
  max_soundspeed::Float64   # used for scaling the output
end

# scale input, pass through the neural network, and scale output
function (ssp::NeuralSSP)(pos)
  z̄ = -pos.z / ssp.max_depth
  c̄ = only(ssp.model(SA[z̄], ssp.ps))
  (c̄ + 1) / 2 * (ssp.max_soundspeed - ssp.min_soundspeed) + ssp.min_soundspeed
end

# define the neural network
ssp = let mlp = SimpleChain(
    static(1),
    TurboDense(tanh, 3),
    TurboDense(tanh, 3),
    TurboDense(tanh, 1)
  )
  ps = SimpleChains.init_params(mlp; rng=MersenneTwister(0))
  NeuralSSP(mlp, ps, 5000.0, 1450.0, 1600.0)
end

We initialize the neural network parameters by training the neural network on the known/estimated sound speed profile from the previous week:

c̄ = @. (prev_soundspeeds - ssp.min_soundspeed) / (ssp.max_soundspeed - ssp.min_soundspeed) * 2 - 1
z̄ = @. prev_depths / ssp.max_depth
G = SimpleChains.alloc_threaded_grad(ssp.model)
L = SimpleChains.add_loss(ssp.model, SquaredLoss(transpose(c̄)))
SimpleChains.train!(G, ssp.ps, L, transpose(collect(z̄)), SimpleChains.ADAM(0.1), 10000)

We plot the model output against the training data to check that it fits well:

plot(ssp, -5000, 0; xlabel="Sound speed (m/s)", size=(350, 400))
scatter!(prev_soundspeeds, -prev_depths)

Now we define a loss function that computes the difference between the predictions of RaySolver using the sound speed given by the NeuralSSP, and the measured impulse response:

# compute loss using predictions of RaySolver using the NeuralSSP
function loss(ps)
  ssp1 = NeuralSSP(ssp.model, ps, 5000.0, 1450.0, 1600.0)
  pred = model(5000.0, 10000.0, 1000.0, 1800.0, ssp1, length(data)+1)
  sum(abs2, pred - data)
end

and perform gradient descent to minimize the loss function:

import ForwardDiff

η = 1e-6            # learning rate
ps = copy(ssp.ps)   # start with the previous week’s parameters

# this will take a few minutes to run...
for i ∈ 1:20
  g = ForwardDiff.gradient(loss, ps)
  ps .-= η .* g     # gradient descent
end

We plot the trained neural sound speed profile against the ground truth:

let ssp = NeuralSSP(ssp.model, ps, 5000.0, 1450.0, 1600.0)
  plot(MunkSSP(), -5000, 0; xlabel="Sound speed (m/s)", size=(350, 400))
  scatter!(prev_soundspeeds, -prev_depths)
  plot!(ssp, -5000, 0)
end

and see that it indeed fits quite well!