Merge pull request #2182 from devitocodes/custom-fd-fix-2

dsl: prevent aggregation for symbolic coefficients

.github/workflows/examples-mpi.yml

@@ -24,7 +24,17 @@ on:
 jobs:
   build:
     name: Examples with mpi
-    runs-on: ["self-hosted", "mpi", "examples"]
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        mpi: ['openmpi', 'intelmpi']
+
+        include:
+        - mpi: openmpi
+          mpiarg: "-n 2 --oversubscribe"
+
+        - mpi: intelmpi
+          mpiarg: "-n 4"
 
     env:
       DEVITO_MPI: "1"
@@ -38,18 +48,18 @@ jobs:
     - name: Checkout devito
       uses: actions/checkout@v3
 
-    - name: Set environment
-      run: |
-        source $ENVHOME/devito-cpu-mpi/bin/activate
-        echo "PATH=$PATH" >> $GITHUB_ENV
-        echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH" >> $GITHUB_ENV
+    - name: Setup MPI
+      uses: mpi4py/setup-mpi@v1
+      with:
+        mpi: ${{ matrix.mpi }}
 
     - name: Install dependencies
       run: |
         pip install --upgrade pip
         pip install -e .[extras,mpi,tests]
 
     - name: Test mpi notebooks
+      continue-on-error: true
       run : |
         ./scripts/create_ipyparallel_mpi_profile.sh
         ipcluster start --profile=mpi --engines=mpi -n 4 --daemonize
@@ -60,11 +70,11 @@ jobs:
 
     - name: Test seismic examples
       run: |
-        mpirun -n 4 pytest examples/seismic/tti/tti_example.py
-        mpirun -n 4 pytest examples/seismic/elastic/elastic_example.py
-        mpirun -n 4 pytest examples/seismic/viscoacoustic/viscoacoustic_example.py
-        mpirun -n 4 pytest examples/seismic/viscoelastic/viscoelastic_example.py
+        mpirun ${{ matrix.mpiarg }} pytest examples/seismic/tti/tti_example.py
+        mpirun ${{ matrix.mpiarg }} pytest examples/seismic/elastic/elastic_example.py
+        mpirun ${{ matrix.mpiarg }} pytest examples/seismic/viscoacoustic/viscoacoustic_example.py
+        mpirun ${{ matrix.mpiarg }} pytest examples/seismic/viscoelastic/viscoelastic_example.py
 
     - name: Test fwi examples with mpi
       run: |
-        mpirun -n 4 python examples/seismic/inversion/fwi.py
+        mpirun ${{ matrix.mpiarg }} python examples/seismic/inversion/fwi.py

devito/parameters.py

@@ -244,7 +244,7 @@ def __enter__(self, condition=True, **params):
     def __exit__(self, exc_type, exc_val, exc_tb):
         for k, v in self.params.items():
             try:
-                configuration.update(k, self.previous[k])
+                configuration[k] = self.previous[k]
             except ValueError:
                 # E.g., `platform` and `compiler` will end up here
                 super(Parameters, configuration).__setitem__(k, self.previous[k])

devito/passes/equations/linearity.py

@@ -128,6 +128,11 @@ def _(expr, mapper, nn_derivs=None):
     if len(with_derivs) > 1:
         return expr
 
+    # Cannot factorize derivatives with symbolic coefficients since
+    # they may have different coefficient values at evaluation
+    if any(d._uses_symbolic_coefficients for w in with_derivs for d in w[1]):
+        return expr
+
     try:
         with_deriv, derivs, others = with_derivs.pop(0)
     except IndexError:

tests/test_symbolic_coefficients.py

@@ -6,7 +6,7 @@
                     Dimension, solve, Operator, NODE)
 from devito.finite_differences import Differentiable
 from devito.tools import as_tuple
-from devito.passes.equations.linearity import factorize_derivatives
+from devito.passes.equations.linearity import factorize_derivatives, aggregate_coeffs
 
 _PRECISION = 9
 
@@ -359,3 +359,18 @@ def test_collect_w_custom_coeffs(self):
         assert collected == expr
         assert collected.is_Add
         Operator([Eq(p.forward, expr)])(time_M=2)  # noqa
+
+    def test_aggregate_w_custom_coeffs(self):
+        grid = Grid(shape=(11, 11, 11))
+        q = TimeFunction(name='q', grid=grid, space_order=8, time_order=2,
+                         coefficients='symbolic')
+
+        expr = 0.5 * q.dx2
+        aggregated = aggregate_coeffs(expr, {})
+
+        assert aggregated == expr
+        assert aggregated.is_Mul
+        assert aggregated.args[0] == .5
+        assert aggregated.args[1] == q.dx2
+
+        Operator([Eq(q.forward, expr)])(time_M=2)  # noqa