Fix JLD2 ext

ext/FluxJUDIExt.jl

@@ -1,9 +1,8 @@
 module FluxJUDIExt
 
-isdefined(Base, :get_extension) ? (using JUDI) : (using ..JUDI)
-using Flux
+import JUDI: LazyPropagation, judiVector, eval_prop
+isdefined(Base, :get_extension) ? (using Flux) : (using ..Flux)
 
-Flux.Zygote.unbroadcast(x::AbstractArray, x̄::LazyPropagation) = Zygote.unbroadcast(x, eval_prop(x̄))
 Flux.cpu(x::LazyPropagation) = Flux.cpu(eval_prop(x))
 Flux.gpu(x::LazyPropagation) = Flux.gpu(eval_prop(x))
 Flux.CUDA.cu(F::LazyPropagation) = Flux.CUDA.cu(eval_prop(F))

ext/JLD2JUDIExt.jl

@@ -1,12 +1,11 @@
 module JLD2JUDIExt
 
-isdefined(Base, :get_extension) ? (using JUDI) : (using ..JUDI)
-using JLD2
+using JUDI
+isdefined(Base, :get_extension) ? (using JLD2) : (using ..JLD2)
+
+JLD2.rconvert(::Type{Geometry}, x::JLD2.ReconstructedMutable{N, FN, NT}) where {N, FN, NT} = Geometry([JUDI.tof32(getproperty(x, f)) for f in FN]...)
+JUDI.Geometry(x::JLD2.ReconstructedMutable{N, FN, NT}) where {N, FN, NT} = Geometry([JUDI.tof32(getproperty(x, f)) for f in FN]...)
 
-function JLD2.rconvert(::Type{Geometry}, x::JLD2.ReconstructedMutable{N, FN, NT}) where {N, FN, NT}
-    args = [JUDI.tof32(getproperty(x, f)) for f in FN]
-    return Geometry(args...)
-end
 
 function JUDI.tof32(x::JLD2.ReconstructedStatic{N, FN, NT}) where {N, FN, NT}
     # Drop "typed" signature

ext/ZygoteJUDIExt.jl

@@ -1,7 +1,7 @@
 module ZygoteJUDIExt
 
-isdefined(Base, :get_extension) ? (using JUDI) : (using ..JUDI)
-using Zygote
+import JUDI: LazyPropagation, judiVector, eval_prop
+isdefined(Base, :get_extension) ? (using Zygote) : (using ..Zygote)
 
 Zygote.unbroadcast(x::AbstractArray, x̄::LazyPropagation) = Zygote.unbroadcast(x, eval_prop(x̄))
 

src/JUDI.jl

@@ -187,23 +187,22 @@ function __init__()
 
     # Optional dependencies
     @static if !isdefined(Base, :get_extension)
-
         # JLD2 compat for loading older version of JUDI types
-        @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3b3d" begin
+        @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819" begin
+            @info "JLD2 compat enabled"
             include("../ext/JLD2JUDIExt.jl")
-            using  JLD2JUDIExt
         end
 
         # Additional Zygote compat if in use
         @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
+            @info "Zygote compat enabled"
             include("../ext/ZygoteJUDIExt.jl")
-            using ZygoteJUDIExt
         end
 
         # Additional Flux compat if in use
         @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" begin
+            @info "Flux compat enabled"
             include("../ext/FluxJUDIExt.jl")
-            using FluxJUDIExt
         end
     end
 

src/TimeModeling/Types/GeometryStructure.jl

@@ -73,7 +73,7 @@ get_nsrc(S::SeisCon) = length(S)
 get_nsrc(S::Vector{SeisCon}) = length(S)
 get_nsrc(S::SeisBlock) = length(unique(get_header(S, "FieldRecord")))
 
-n_samples(g::GeometryOOC, nsrc::Integer) = sum([g.nrec[j]*get_nt(g, j) for j=1:nsrc])
+n_samples(g::GeometryOOC, nsrc::Integer) = sum(g.nrec .* get_nt(g))
 n_samples(g::GeometryIC, nsrc::Integer) = sum([length(g.xloc[j])*get_nt(g, j) for j=1:nsrc])
 n_samples(g::Geometry) = n_samples(g, get_nsrc(g))
 

src/TimeModeling/Utils/auxiliaryFunctions.jl

@@ -810,6 +810,8 @@ filter_none(x) = x
 
 """
     _maybe_pad_t0(q, qGeom, data, dataGeom)
+
+Pad zeros for data with non-zero t0, usually from a segy file so that time axis and array size match for the source and data.
 """
 function _maybe_pad_t0(qIn::Matrix{T}, qGeom::Geometry, dObserved::Matrix{T}, dataGeom::Geometry) where T<:Number
     if size(dObserved, 1) != size(qIn, 1)
@@ -823,4 +825,7 @@ function _maybe_pad_t0(qIn::Matrix{T}, qGeom::Geometry, dObserved::Matrix{T}, da
         end
     end
     return qIn, dObserved
-end
+end
+
+_maybe_pad_t0(qIn::judiVector{T, AT}, dObserved::judiVector{T, AT}) where{T<:Number, AT} =
+    _maybe_pad_t0(qIn.data, qIn.geometry, dObserved.data, dObserved.geometry)

src/pysource/fields.py

@@ -6,6 +6,11 @@
 from devito.data.allocators import ExternalAllocator
 from devito.tools import as_tuple
 
+try:
+    from devitopro import *
+except ImportError:
+    pass
+
 
 def wavefield(model, space_order, save=False, nt=None, fw=True, name='', t_sub=1):
     """

src/pysource/models.py

@@ -7,6 +7,12 @@
 from devito.data.allocators import ExternalAllocator
 from devito.tools import as_tuple, memoized_func
 
+try:
+    from devitopro import *
+except ImportError:
+    pass
+
+
 __all__ = ['Model']
 
 

src/pysource/operators.py

@@ -13,6 +13,11 @@
 from sensitivity import grad_expr, lin_src
 from utils import opt_op
 
+try:
+    from devitopro import *
+except ImportError:
+    pass
+
 
 def name(model):
     if model.is_tti:

test/test_issues.jl

@@ -28,6 +28,7 @@ end
         # Load file with old judiVector type and julia <1.7 StepRangeLen
         @load "$(datapath)backward_comp.jld" dat
 
+        @show dat.geometry
         @test typeof(dat) == judiVector{Float32, Matrix{Float32}}
         @test typeof(dat.geometry) == GeometryIC{Float32}
         @test typeof(dat.geometry.xloc) == Vector{Vector{Float32}}