Gradient is always zero when running FWI with OBC data

Hi, I am currently running a Full Waveform Inversion (FWI) with LBFGS using OBC data with the JUDI.jl package, but I have encountered an issue where the gradient is consistently returning as zero. This is causing the optimization process to not converge as expected.

Code Snippet:

using Statistics, Random, LinearAlgebra
using JUDI, NLopt, HDF5, SegyIO, PyPlot

n, d, o, m0 = read(h5open("obcmodel.h5", "r"), "n", "d", "o", "m0")

# Ensure the data types are correct
n = convert(Vector{Int64}, reshape(n, 2))
d = convert(Vector{Float64}, reshape(d, 2))
o = convert(Vector{Float32}, reshape(o, 2))
m0 = convert(Matrix{Float32}, m0)
model0 = Model((n[1], n[2]), (d[1], d[2]), (o[1], o[2]), m0)

# Bound constraints
vmin = 1.4f0
vmax = 6.5f0

block = segy_read("obc.segy")
d_obs = judiVector(block)

# Set up wavelet and source vector
src_geometry = Geometry(block; key = "source", segy_depth_key = "SourceDepth")
wavelet = ricker_wavelet(src_geometry.t[1], src_geometry.dt[1], 0.008f0) # 8 Hz wavelet
q = judiVector(src_geometry, wavelet)

# optimization parameters
fevals = 100    # allow 10 function evalutations
batchsize = 2
count = 0
fhistory = zeros(Float32, fevals+1)

# NLopt objective function
println("No.  ", "fval         ", "norm(gradient)")
function f!(x, grad)

    # Update model
    model0.m .= x

    # Select batch and calculate gradient
    i = randperm(d_obs.nsrc)[1:batchsize]
    fval, gradient = fwi_objective(model0, q[i], d_obs[i])

    # Reset gradient in water column to zero
    gradient = reshape(gradient, model0.n); gradient[:, 1:21] .= 0f0
    grad[1:end] .= gradient[1:end]

    global count; count += 1
    println(count, "    ", fval, "    ", norm(grad))
    fhistory[count] = fval
    return convert(Float64, fval)
end

# Optimization parameters
opt = Opt(:LD_LBFGS, prod(model0.n))
min_objective!(opt, f!)
mmin = (1f0 ./ vmax)^2   # bound constraints on slowness squared
mmax = (1f0 ./ vmin)^2
lower_bounds!(opt, mmin); upper_bounds!(opt, mmax)
maxeval!(opt, fevals)
(minf, minx, ret) = optimize(opt, vec(copy(model0.m)))

# Save results, function values and elapsed time
h5open("result_2D_obc_lbfgs.h5", "w") do file
    write(file, "x", sqrt.(1f0 ./ reshape(minx, model0.n)), "fhistory", fhistory)
end

Actual Result: The gradient is always reported as zero, and the optimization does not make any progress:

No.  fval         norm(gradient)
1    1.4131794e-7    0.0

I have verified that the SEGY file is correctly formatted and contains the expected data:

block = segy_read("/home/whz/.julia/environments/v1.7/datacube_in.segy")
d_obs = judiVector(block)
block.traceheaders

BinaryTraceHeader: TraceNumWithinLine: 0 TraceNumWithinFile: 1 FieldRecord: 0 TraceNumber: 1 EnergySourcePoint: 0 CDP: 0 CDPTrace: 0 TraceIDCode: 0 NSummedTraces: 0 NStackedTraces: 0 DataUse: 0 Offset: 0 RecGroupElevation: 0 SourceSurfaceElevation: 0 SourceDepth: 0 RecDatumElevation: 0 SourceDatumElevation: 0 SourceWaterDepth: 0 GroupWaterDepth: 0 ElevationScalar: 0 RecSourceScalar: 0 SourceX: 0 SourceY: 0 GroupX: 198 GroupY: 0 CoordUnits: 0 WeatheringVelocity: 0 SubWeatheringVelocity: 0 UpholeTimeSource: 0 UpholeTimeGroup: 0 StaticCorrectionSource: 0 StaticCorrectionGroup: 0 TotalStaticApplied: 0 LagTimeA: 0 LagTimeB: 0 DelayRecordingTime: 0 MuteTimeStart: 0 MuteTimeEnd: 0 ns: 10000 dt: 1000 GainType: 0 InstrumentGainConstant: 0 InstrumntInitialGain: 0 Correlated: 0 SweepFrequencyStart: 0 SweepFrequencyEnd: 0 SweepLength: 0 SweepType: 0 SweepTraceTaperLengthStart: 0 SweepTraceTaperLengthEnd: 0 TaperType: 0 AliasFilterFrequency: 0 AliasFilterSlope: 0 NotchFilterFrequency: 0 NotchFilterSlope: 0 LowCutFrequency: 0 HighCutFrequency: 0 LowCutSlope: 0 HighCutSlope: 0 Year: 2024 DayOfYear: 154 HourOfDay: 22 MinuteOfHour: 3 SecondOfMinute: 25 TimeCode: 0 TraceWeightingFactor: 0 GeophoneGroupNumberRoll: 0 GeophoneGroupNumberTraceStart: 0 GeophoneGroupNumberTraceEnd: 0 GapSize: 0 OverTravel: 0 CDPX: 0 CDPY: 0 Inline3D: 0 Crossline3D: 0 ShotPoint: 0 ShotPointScalar: 0 TraceValueMeasurmentUnit: 0 TransductionConstnatMantissa: 0 TransductionConstantPower: 0 TransductionUnit: 0 TraceIdentifier: 0 ScalarTraceHeader: 0 SourceType: 0 SourceEnergyDirectionMantissa: 0 SourceEnergyDirectionExponent: 0 SourceMeasurmentMantissa: 0 SourceMeasurementExponent: 0 SourceMeasurmentUnit: 0 Unassigned1: 0 Unassigned2: 0

I am unsure if there might be an issue with how the headers are being interpreted or used in the FWI process. Is there something specific in the SEGY file or trace headers that I should be looking out for when using OBC data with JUDI.jl? I would greatly appreciate any guidance or suggestions on how to resolve this issue.

slimgroup / JUDI.jl

Gradient is always zero when running FWI with OBC data #259