Update residual calculation to account for absorption and heterogeneous density

[btreeby]

The network was initially trained assuming constant density and no absorption. Initial tests show that the network may also generalise (even without re-training) to more general cases. This requires updating IterativeSolver.get_residual.

Absorption could possibly be implemented by adding an imaginary part to k_sq, while density gradients could be implemented using, e.g.,

rho_dx = self.Lap.dx(rho)

The additional parameters will need to be propagated through the code, e.g., by adding additional channels to the speed of sound input. This will need changing IterativeSolver.get_initials to only use the sound speed channel in the calculation of k_sq.

  k_sq, wavefield = self.get_initials(sos_maps)
  self.f.clear_states(wavefield)
  residual = self.get_residual(wavefield, k_sq)

To:

k_sq, wavefield, alpha, rho = self.get_initials(medium_maps)
self.f.clear_states(wavefield)
residual = self.get_residual(wavefield, k_sq, alpha, rho)

This will need to propagated in several places that call get_initials , get_residual, or use sos.