How to generate reverse time migration about deepwave HOT 6 OPEN

v.requires_grad_()
out = deepwave.scalar(v, ...)
loss = loss_fn(out[-1], observed_data)  # You may wish to mask direct arrivals here
loss.backward()
image = v.grad.detach()

import torch
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt
import deepwave
from deepwave import scalar, scalar_born
import numpy as np

device = torch.device('cuda:0' if torch.cuda.is_available()
                      else 'cpu')
ny = 2301
nx = 801
dx = 4.0
nt = 750
dt = 0.004
n_shots = 115
freq = 26

v_true = torch.ones(ny, nx) * 1500
v_true[:, -751:] = torch.from_file('marmousi_vp.bin',
                         size=ny * 751).reshape(ny, 751)

# Smooth v_true slightly to reduce multiples
v_true = torch.tensor(1 / gaussian_filter(1 / v_true.numpy(), 5))

smooth2 = (torch.tensor(1 / gaussian_filter(1 / v_true.numpy(), 5))
           .to(device))

v_true = v_true.to(device)

n_sources_per_shot = 1
d_source = 20  # 20 * 4m = 80m
first_source = 10 # 10 * 4m = 40m
source_depth = 2  # 2 * 4m = 8m

n_receivers_per_shot = 384
d_receiver = 6  # 6 * 4m = 24m
first_receiver = 0  # 0 * 4m = 0m
receiver_depth = 2  # 2 * 4m = 8m

peak_time = 1.5 / freq

source_locations = torch.zeros(n_shots, n_sources_per_shot, 2,
                               dtype=torch.long, device=device)
source_locations[..., 1] = source_depth
source_locations[:, 0, 0] = (torch.arange(n_shots) * d_source +
                             first_source)

receiver_locations = torch.zeros(n_shots, n_receivers_per_shot, 2,
                                 dtype=torch.long, device=device)
receiver_locations[..., 1] = receiver_depth
receiver_locations[:, :, 0] = (
    (torch.arange(n_receivers_per_shot) * d_receiver +
     first_receiver)
    .repeat(n_shots, 1)
)

source_amplitudes = (
    (deepwave.wavelets.ricker(freq, nt, dt, peak_time))
    .repeat(n_shots, n_sources_per_shot, 1).to(device)
)

# Use the same max_vel value for each propagation to ensure consistent
# internal time steps (for better data match)
observed_true_data = scalar(v_true, dx, dt, source_amplitudes=source_amplitudes,
                            source_locations=source_locations,
                            receiver_locations=receiver_locations,
                            pml_freq=freq,
                            max_vel=v_true.max().item())[-1]
observed_smooth2_data = scalar(smooth2, dx, dt, source_amplitudes=source_amplitudes,
                            source_locations=source_locations,
                            receiver_locations=receiver_locations,
                            pml_freq=freq,
                            max_vel=v_true.max().item())[-1]
observed_scattered_data = observed_true_data - observed_smooth2_data

scatter = torch.zeros_like(smooth2)
scatter.requires_grad_()

optimiser = torch.optim.LBFGS([scatter])
loss_fn = torch.nn.MSELoss()

n_epochs = 10

for epoch in range(n_epochs):
   def closure():
       optimiser.zero_grad()
       total_loss = 0
       for shotidx in range(n_shots):
           out = scalar_born(
               smooth2, scatter, dx, dt,
               source_amplitudes=source_amplitudes[shotidx:shotidx+1],
               source_locations=source_locations[shotidx:shotidx+1],
               receiver_locations=receiver_locations[shotidx:shotidx+1],
               pml_freq=freq,
               max_vel=v_true.max().item()
           )
           loss = 1e20 * loss_fn(out[-1][0], observed_scattered_data[shotidx])
           loss.backward()
           total_loss += loss.item()
       print(epoch, total_loss)
       return total_loss

   optimiser.step(closure)