Merge pull request #128 from slimgroup/dft-fix

Fix DFT caching

Author: mloubout

Project.toml

@@ -1,7 +1,7 @@
 name = "JUDI"
 uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 authors = ["Philipp Witte, Mathias Louboutin"]
-version = "3.1.3"
+version = "3.1.4"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/pysource/fields.py

@@ -167,6 +167,24 @@ def lr_src_fields(model, weight, wavelet, empty_ws=False):
     return source_weight, wavelett
 
 
+def frequencies(freq, fdim=None):
+    """
+    Frequencies as a one dimensional Function
+
+    Parameters
+    ----------
+    freq: List or 1D array
+        List of frequencies
+    """
+    if freq is None:
+        return None, 0
+    nfreq = np.shape(freq)[0]
+    freq_dim = fdim or DefaultDimension(name='freq_dim', default_value=nfreq)
+    f = Function(name='f', dimensions=(freq_dim,), shape=(nfreq,))
+    f.data[:] = np.array(freq[:])
+    return f, nfreq
+
+
 def fourier_modes(u, freq):
     """
     On the fly DFT wavefield (frequency slices) and expression
@@ -182,10 +200,8 @@ def fourier_modes(u, freq):
         return None, None
 
     # Frequencies
-    nfreq = np.shape(freq)[0]
-    freq_dim = DefaultDimension(name='freq_dim', default_value=nfreq)
-    f = Function(name='f', dimensions=(freq_dim,), shape=(nfreq,))
-    f.data[:] = np.array(freq[:])
+    f, nfreq = frequencies(freq)
+    freq_dim = f.dimensions[0]
 
     dft_modes = []
     for wf in as_tuple(u):

src/pysource/operators.py

@@ -73,7 +73,7 @@ def __get__(self, obj, objtype):
 
 @memoized_func
 def forward_op(p_params, tti, visco, space_order, spacing, save, t_sub, fs, pt_src,
-               pt_rec, dft, dft_sub, ws, full_q):
+               pt_rec, nfreq, dft_sub, ws, full_q):
     """
     Low level forward operator creation, to be used through `propagator.py`
     Compute forward wavefield u = A(m)^{-1}*f and related quantities (u(xrcv))
@@ -86,7 +86,7 @@ def forward_op(p_params, tti, visco, space_order, spacing, save, t_sub, fs, pt_s
     scords = np.ones((1, ndim)) if pt_src else None
     rcords = np.ones((1, ndim)) if pt_rec else None
     wavelet = np.ones((nt, 1))
-    freq_list = np.ones((2,)) if dft else None
+    freq_list = np.ones((nfreq,)) if nfreq > 0 else None
     q = wavefield(model, 0, save=True, nt=nt, name="qwf") if full_q else 0
     wsrc = Function(name='src_weight', grid=model.grid, space_order=0) if ws else None
 
@@ -120,7 +120,7 @@ def forward_op(p_params, tti, visco, space_order, spacing, save, t_sub, fs, pt_s
 
 @memoized_func
 def adjoint_op(p_params, tti, visco, space_order, spacing, save, nv_weights, fs,
-               pt_src, pt_rec, dft, dft_sub, ws, full_q):
+               pt_src, pt_rec, nfreq, dft_sub, ws, full_q):
     """
     Low level adjoint operator creation, to be used through `propagators.py`
     Compute adjoint wavefield v = adjoint(F(m))*y
@@ -134,7 +134,7 @@ def adjoint_op(p_params, tti, visco, space_order, spacing, save, nv_weights, fs,
     scords = np.ones((1, ndim)) if pt_src else None
     rcords = np.ones((1, ndim)) if pt_rec else None
     wavelet = np.ones((nt, 1))
-    freq_list = np.ones((2,)) if dft else None
+    freq_list = np.ones((nfreq,)) if nfreq > 0 else None
     q = wavefield(model, 0, save=True, nt=nt, fw=False, name="qwf") if full_q else 0
 
     # Setting adjoint wavefield
@@ -167,7 +167,7 @@ def adjoint_op(p_params, tti, visco, space_order, spacing, save, nv_weights, fs,
 
 @memoized_func
 def born_op(p_params, tti, visco, space_order, spacing, save, pt_src,
-            pt_rec, fs, t_sub, ws, dft, dft_sub, isic, nlind):
+            pt_rec, fs, t_sub, ws, nfreq, dft_sub, isic, nlind):
     """
     Low level born operator creation, to be used through `interface.py`
     Compute linearized wavefield U = J(m)* δ m
@@ -181,7 +181,7 @@ def born_op(p_params, tti, visco, space_order, spacing, save, pt_src,
     wavelet = np.ones((nt, 1))
     scords = np.ones((1, ndim)) if pt_src else None
     rcords = np.ones((1, ndim)) if pt_rec else None
-    freq_list = np.ones((2,)) if dft else None
+    freq_list = np.ones((nfreq,)) if nfreq > 0 else None
     wsrc = Function(name='src_weight', grid=model.grid, space_order=0) if ws else None
     f0 = Constant('f0')
 
@@ -220,7 +220,7 @@ def born_op(p_params, tti, visco, space_order, spacing, save, pt_src,
 
 @memoized_func
 def adjoint_born_op(p_params, tti, visco, space_order, spacing, pt_rec, fs, w,
-                    save, t_sub, dft, dft_sub, isic):
+                    save, t_sub, nfreq, dft_sub, isic):
     """
     Low level gradient operator creation, to be used through `propagators.py`
     Compute the action of the adjoint Jacobian onto a residual J'* δ d.
@@ -231,10 +231,11 @@ def adjoint_born_op(p_params, tti, visco, space_order, spacing, pt_rec, fs, w,
     ndim = len(spacing)
     residual = np.ones((nt, 1))
     rcords = np.ones((1, ndim)) if pt_rec else None
-    freq_list = np.ones((2,)) if dft else None
+    freq_list = np.ones((nfreq,)) if nfreq > 0 else None
     # Setting adjoint wavefieldgradient
     v = wavefield(model, space_order, fw=False)
-    u = forward_wavefield(model, space_order, save=save, nt=nt, dft=dft, t_sub=t_sub)
+    u = forward_wavefield(model, space_order, save=save, nt=nt,
+                          dft=nfreq > 0, t_sub=t_sub)
 
     # Set up PDE expression and rearrange
     pde = wave_kernel(model, v, fw=False, f0=Constant('f0'))

src/pysource/propagators.py

@@ -1,10 +1,10 @@
 from kernels import wave_kernel
 from geom_utils import src_rec, geom_expr
 from fields import (fourier_modes, wavefield, lr_src_fields,
-                    wavefield_subsampled, norm_holder)
+                    wavefield_subsampled, norm_holder, frequencies)
 from fields_exprs import extented_src
 from sensitivity import grad_expr
-from utils import weight_fun, opt_op, fields_kwargs
+from utils import weight_fun, opt_op, fields_kwargs, nfreq
 from operators import forward_op, adjoint_op, born_op, adjoint_born_op
 
 from devito import Operator, Function, Constant
@@ -32,7 +32,7 @@ def forward(model, src_coords, rcv_coords, wavelet, space_order=8, save=False,
     op = forward_op(model.physical_parameters, model.is_tti, model.is_viscoacoustic,
                     space_order, model.spacing, save, t_sub, model.fs,
                     src_coords is not None, rcv_coords is not None,
-                    freq_list is not None, dft_sub, ws is not None, qwf is not None)
+                    nfreq(freq_list), dft_sub, ws is not None, qwf is not None)
 
     # Make kwargs
     kw = {'dt': model.critical_dt}
@@ -83,7 +83,7 @@ def adjoint(model, y, src_coords, rcv_coords, space_order=8, qwf=None, dft_sub=N
     op = adjoint_op(model.physical_parameters, model.is_tti, model.is_viscoacoustic,
                     space_order, model.spacing, save, nv_weights, model.fs,
                     src_coords is not None, rcv_coords is not None,
-                    freq_list is not None, dft_sub, ws is not None, qwf is not None)
+                    nfreq(freq_list), dft_sub, ws is not None, qwf is not None)
 
     # On-the-fly Fourier
     dft_modes, fr = fourier_modes(v, freq_list)
@@ -133,12 +133,13 @@ def gradient(model, residual, rcv_coords, u, return_op=False, space_order=8,
     # Create operator and run
     op = adjoint_born_op(model.physical_parameters, model.is_tti, model.is_viscoacoustic,
                          space_order, model.spacing, rcv_coords is not None, model.fs, w,
-                         not return_op, t_sub, freq is not None, dft_sub, isic)
+                         not return_op, t_sub, nfreq(freq), dft_sub, isic)
 
     # Update kwargs
     kw = {'dt': model.critical_dt}
+    f, _factor = frequencies(freq)
     f0q = Constant('f0', value=f0) if model.is_viscoacoustic else None
-    kw.update(fields_kwargs(src, u, v, gradm, f0q))
+    kw.update(fields_kwargs(src, u, v, gradm, f0q, f))
     kw.update(model.physical_params())
 
     if return_op:
@@ -172,7 +173,7 @@ def born(model, src_coords, rcv_coords, wavelet, space_order=8, save=False,
     op = born_op(model.physical_parameters, model.is_tti, model.is_viscoacoustic,
                  space_order, model.spacing, save,
                  src_coords is not None, rcv_coords is not None, model.fs, t_sub,
-                 ws is not None, freq_list is not None, dft_sub, isic, nlind)
+                 ws is not None, nfreq(freq_list), dft_sub, isic, nlind)
 
     # Make kwargs
     kw = {'dt': model.critical_dt}

src/pysource/sensitivity.py

@@ -4,6 +4,7 @@
 from devito import Eq
 from devito.tools import as_tuple
 
+from fields import frequencies
 from fields_exprs import sub_time, freesurface
 from FD_utils import divs, grads
 
@@ -89,17 +90,15 @@ def crosscorr_freq(u, v, model, freq=None, dft_sub=None, **kwargs):
     tsave, factor = sub_time(time, dft_sub)
     expr = 0
 
+    fdim = as_tuple(u)[0][0].dimensions[0]
+    f, nfreq = frequencies(freq, fdim=fdim)
+    omega_t = 2*np.pi*f*tsave*factor*dt
+    # Gradient weighting is (2*np.pi*f)**2/nt
+    w = -(2*np.pi*f)**2/time.symbolic_max
+
     for uu, vv in zip(u, v):
         ufr, ufi = uu
-        # Frequencies
-        nfreq = np.shape(freq)[0]
-        fdim = ufr.dimensions[0]
-        omega_t = lambda f: 2*np.pi*f*tsave*factor*dt
-        # Gradient weighting is (2*np.pi*f)**2/nt
-        w = lambda f: -(2*np.pi*f)**2/time.symbolic_max
-        expr += sum(w(freq[ff])*(ufr._subs(fdim, ff)*cos(omega_t(freq[ff])) -
-                                 ufi._subs(fdim, ff)*sin(omega_t(freq[ff])))
-                    for ff in range(nfreq))*vv
+        expr += w*(ufr*cos(omega_t) - ufi*sin(omega_t))*vv
     return expr
 
 
@@ -139,21 +138,18 @@ def isic_freq(u, v, model, **kwargs):
     time = model.grid.time_dim
     dt = time.spacing
     tsave, factor = sub_time(time, kwargs.get('factor'))
+    fdim = as_tuple(u)[0][0].dimensions[0]
+    f, nfreq = frequencies(freq, fdim=fdim)
+    omega_t = 2*np.pi*f*tsave*factor*dt
+    w = -(2*np.pi*f)**2/time.symbolic_max
+    w2 = factor / time.symbolic_max
+
     expr = 0
     for uu, vv in zip(u, v):
         ufr, ufi = uu
-        # Frequencies
-        nfreq = np.shape(freq)[0]
-        fdim = ufr.dimensions[0]
-        omega_t = lambda f: 2*np.pi*f*tsave*factor*dt
-        # Gradient weighting is (2*np.pi*f)**2/nt
-        w = lambda f: -(2*np.pi*f)**2/time.symbolic_max
-        w2 = factor / time.symbolic_max
-        for ff in range(nfreq):
-            cwt, swt = cos(omega_t(freq[ff])), sin(omega_t(freq[ff]))
-            ufrf, ufif = ufr._subs(fdim, ff), ufi._subs(fdim, ff)
-            idftu = (ufrf * cwt - ufif * swt)
-            expr += w(freq[ff]) * idftu * vv * model.m - w2 * inner_grad(idftu, vv)
+        cwt, swt = cos(omega_t), sin(omega_t)
+        idftu = (ufr * cwt - ufi * swt)
+        expr += w * idftu * vv * model.m - w2 * inner_grad(idftu, vv)
     return expr
 
 

src/pysource/utils.py

@@ -115,6 +115,13 @@ def opt_op(model):
     return ('advanced', opts)
 
 
+def nfreq(freq_list):
+    """
+    Check number of on-the-fly DFT frequencies.
+    """
+    return 0 if freq_list is None else np.shape(freq_list)[0]
+
+
 def fields_kwargs(*args):
     """
     Creates a dictionary of {f.name: f} for any field argument that is not None