better free surface

setup.sh

@@ -2,15 +2,13 @@
 
 # Install devito
 pip install --upgrade pip
-pip install --user git+https://github.com/devitocodes/[email protected]
-pip install --user -r docker/devito_requirements.txt
+pip install -U --user "devito[tests,extras]"
 
 # Set devito enviroment variables
 export DEVITO_ARCH="gcc"
-export DEVITO_OPENMP="0"
+export DEVITO_LANGUAGE="openmp"
 
 # Point PyCall to correct Python version
-export PYTHON=$(which python)
-julia -e 'Pkg.build("PyCall")'
+PYTHON=$(which python) julia -e 'using Pkg; Pkg.build("PyCall")'
 
 

src/JUDI.jl

@@ -65,6 +65,7 @@ const pyut = PyNULL()
 const devito = PyNULL()
 
 set_devito_config(key::String, val::String) = set!(devito."configuration", key, val)
+set_devito_config(key::String, val::Bool) = set!(devito."configuration", key, val)
 
 # Create a lock for pycall FOR THREAD/TASK SAFETY
 # See discussion at
@@ -178,6 +179,9 @@ function __init__()
     # Initialize lock at session start
     PYLOCK[] = ReentrantLock()
 
+    # Prevent autopadding to use external allocator
+    set_devito_config("autopadding", false)
+
     # Make sure there is no conflict for the cuda init thread with CUDA.jl
     if get(ENV, "DEVITO_PLATFORM", "") == "nvidiaX"
         @info "Initializing openacc/openmp offloading"

src/pysource/fields.py

@@ -7,9 +7,9 @@
 from devito.tools import as_tuple
 
 try:
-    from devitopro import *
+    import devitopro as dvp  # noqa
 except ImportError:
-    pass
+    import devito as dvp
 
 
 def wavefield(model, space_order, save=False, nt=None, fw=True, name='', t_sub=1):
@@ -139,9 +139,9 @@ def wavefield_subsampled(model, u, nt, t_sub, space_order=8):
         return None
     wf_s = []
     for wf in as_tuple(u):
-        usave = TimeFunction(name='us_%s' % wf.name, grid=model.grid, time_order=2,
-                             space_order=space_order, time_dim=time_subsampled,
-                             save=int(nsave))
+        usave = dvp.TimeFunction(name='us_%s' % wf.name, grid=model.grid, time_order=2,
+                                 space_order=space_order, time_dim=time_subsampled,
+                                 save=nsave)
         wf_s.append(usave)
     return wf_s
 

src/pysource/fields_exprs.py

@@ -2,7 +2,7 @@
 
 from devito import Inc, Eq, ConditionalDimension, cos, sin, sign
 from devito.tools import as_tuple
-from devito.symbolics import retrieve_functions, INT
+from devito.symbolics import retrieve_functions, INT, retrieve_derivatives
 
 from fields import (wavefield_subsampled, lr_src_fields, fourier_modes, norm_holder,
                     illumination)
@@ -101,7 +101,7 @@ def extended_rec(model, wavelet, v):
     return [Inc(ws, model.grid.time_dim.spacing * wf * wt)]
 
 
-def freesurface(model, eq):
+def freesurface(model, eq, mfuncs=None):
     """
     Generate the stencil that mirrors the field as a free surface modeling for
     the acoustic wave equation
@@ -112,30 +112,33 @@ def freesurface(model, eq):
         Physical model
     eq: Eq or List of Eq
         Equation to apply mirror to
+    mfuncs: List of Functions
+        List of functions to mirror (default=None). Mirrors all functions if not provided
     """
     fs_eq = []
+    fsdomain = model.grid.subdomains['fsdomain']
     for eq_i in eq:
         for p in eq_i._flatten:
-            lhs, rhs = p.evaluate.args
             # Add modulo replacements to to rhs
+            lhs, rhs = p.lhs, p.rhs
             zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
             z = zfs.parent
 
-            if lhs.is_TimeFunction:
-                so = lhs.space_order
-                z0= S.Zero
-                zfs = CustomDimension(name="zfs", symbolic_min=1, symbolic_max=so,
-                                      symbolic_size=so)
-                fs_eq.extend([Eq(lhs.subs({z: z0 - zfs}), -lhs.subs({z: z0 + zfs})),
-                              Eq(lhs.subs({z: z.symbolic_min + z0}), 0)])
-    #         mapper = {}
-    #         for f in funcs:
-    #             zind = f.indices[-1]
-    #             if (zind - z).as_coeff_Mul()[0] < 0:
-    #                 s = sign((zind - z.symbolic_min).subs({z: zfs, z.spacing: 1}))
-    #                 mapper.update({f: s * f.subs({zind: INT(abs(zind))})})
-    #         fs_eq.append(Eq(lhs, sign(lhs.indices[-1]-z.symbolic_min) * rhs.subs(mapper),
-    #                         subdomain=model.grid.subdomains['fsdomain']))
+            Dz = {d: d.evaluate for d in retrieve_derivatives(rhs) if z in d.dims}
+
+            funcs = {f for f in retrieve_functions(Dz.values()) if f.indices[-1] is not z}
+            if mfuncs:
+                funcs = {f for f in funcs if f.function in mfuncs}
+
+            mapper = {}
+            for f in funcs:
+                zind = f.indices[-1]
+                if (zind - z).as_coeff_Mul()[0] < 0:
+                    s = sign(zind._subs(z.spacing, 1))
+                    mapper.update({f: s * f._subs(zind, INT(abs(zind)))})
+            dzmapper = {d: v.subs(mapper) for d, v in Dz.items()}
+            fs_eq.append(Eq(lhs, rhs.subs(dzmapper), subdomain=fsdomain))
+            fs_eq.append(Eq(lhs._subs(z, 0), 0, subdomain=fsdomain))
     return fs_eq
 
 

src/pysource/kernels.py

@@ -58,11 +58,11 @@ def acoustic_kernel(model, u, fw=True, q=None):
     damp = model.damp
     stencil = solve(wmr * u.dt2 + damp * udt - ulaplace - q, u_n)
 
-    # if 'nofsdomain' in model.grid.subdomains:
-    #     pde = [Eq(u_n, stencil, subdomain=model.grid.subdomains['nofsdomain'])]
-    #     pde += freesurface(model, pde)
-    # else:
-    pde = [Eq(u_n, stencil)]
+    if 'nofsdomain' in model.grid.subdomains:
+        pde = [Eq(u_n, stencil, subdomain=model.grid.subdomains['nofsdomain'])]
+        pde += freesurface(model, pde, (u,))
+    else:
+        pde = [Eq(u_n, stencil)]
 
     return pde
 
@@ -166,7 +166,7 @@ def tti_kernel(model, u1, u2, fw=True, q=None):
         # Water at free surface, no anisotrpy
         acout_ttip = [Eq(u1_n, stencilp.subs(model.zero_thomsen))]
         acout_ttir = [Eq(u2_n, stencilr.subs(model.zero_thomsen))]
-        pdea = freesurface(model, acout_ttip) + freesurface(model, acout_ttir)
+        pdea = freesurface(model, (acout_ttip, acout_ttir), (u1, u2))
         # Standard PDE in subsurface
         first_stencil = Eq(u1_n, stencilp, subdomain=model.grid.subdomains['nofsdomain'])
         second_stencil = Eq(u2_n, stencilr, subdomain=model.grid.subdomains['nofsdomain'])

src/pysource/models.py

@@ -8,7 +8,7 @@
 from devito.tools import as_tuple, memoized_func
 
 try:
-    from devitopro import *
+    from devitopro import *  # noqa
 except ImportError:
     pass
 

src/pysource/operators.py

@@ -14,7 +14,7 @@
 from utils import opt_op
 
 try:
-    from devitopro import *
+    from devitopro import *  # noqa
 except ImportError:
     pass
 

src/pysource/sensitivity.py

@@ -42,12 +42,12 @@ def grad_expr(gradm, u, v, model, w=None, freq=None, dft_sub=None, ic="as"):
     """
     ic_func = ic_dict[func_name(freq=freq, ic=ic)]
     expr = ic_func(as_tuple(u), as_tuple(v), model, freq=freq, factor=dft_sub, w=w)
-    # if model.fs and ic in ["fwi", "isic"]:
-    #     # Only need `fs` processing for isic for the spatial derivatives.
-    #     eq_g = [Eq(gradm, gradm - expr, subdomain=model.grid.subdomains['nofsdomain'])]
-    #     eq_g += freesurface(model, eq_g)
-    # else:
-    eq_g = [Eq(gradm, gradm - expr)]
+    if model.fs and ic in ["fwi", "isic"]:
+        # Only need `fs` processing for isic for the spatial derivatives.
+        eq_g = [Eq(gradm, gradm - expr, subdomain=model.grid.subdomains['nofsdomain'])]
+        eq_g += freesurface(model, eq_g, (*u, *v))
+    else:
+        eq_g = [Eq(gradm, gradm - expr)]
     return eq_g