├── .github └── workflows │ └── ci-tests.yaml ├── .gitignore ├── LICENSE ├── README.md ├── devito4pytorch ├── __init__.py └── devito_wrapper.py ├── setup.cfg ├── setup.py └── tests ├── adjoint_born_test.py ├── forward_born_test.py └── forward_modeling_test.py /.github/workflows/ci-tests.yaml: -------------------------------------------------------------------------------- 1 | name: CI-tests 2 | 3 | on: 4 | # Trigger the workflow on push or pull request, 5 | # but only for the master branch 6 | push: 7 | branches: 8 | - master 9 | pull_request: 10 | branches: 11 | - master 12 | 13 | jobs: 14 | tutorials: 15 | name: Run all tests 16 | runs-on: ubuntu-latest 17 | 18 | env: 19 | DEVITO_ARCH: gcc-7 20 | DEVITO_LANGUAGE: "openmp" 21 | DEVITO_BACKEND: "core" 22 | PYTHON_VERSION: "3.7" 23 | OMP_NUM_THREADS: 2 24 | 25 | strategy: 26 | # Prevent all build to stop if a single one fails 27 | fail-fast: false 28 | 29 | steps: 30 | - name: Checkout Devito4PyTorch 31 | uses: actions/checkout@v1 32 | 33 | 34 | - name: Set up Python 3.7 35 | uses: actions/setup-python@v1 36 | with: 37 | python-version: 3.7 38 | 39 | - name: Install dependencies 40 | if: matrix.name != 'pytest-docker-py36-gcc-omp' 41 | run: | 42 | pip install --upgrade pip 43 | pip install -e . 44 | pip install flake8 codecov pytest-cov pytest pytest-runner 45 | 46 | - name: Lint with flake8 47 | run: | 48 | flake8 --builtins=ArgumentError . 49 | 50 | - name: Test with pytest 51 | run: | 52 | pytest --cov --cov-config=.coveragerc --cov-report=xml tests 53 | 54 | - name: Upload coverage to Codecov 55 | if: matrix.name != 'pytest-docker-py36-gcc-omp' 56 | uses: codecov/codecov-action@v1.0.6 57 | with: 58 | token: ${{ secrets.CODECOV_TOKEN }} 59 | name: ${{ matrix.name }} -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 SLIM group @ Georgia Institute of Technology 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![CI Status](https://github.com/slimgroup/Devito4PyTorch/workflows/CI-tests/badge.svg)](https://github.com/slimgroup/Devito4PyTorch/actions?query=workflow%3ACI-tests) 2 | [![Code Coverage](https://codecov.io/gh/slimgroup/Devito4PyTorch/branch/master/graph/badge.svg)](https://codecov.io/gh/slimgroup/Devito4PyTorch) 3 | [![DOI](https://zenodo.org/badge/259481291.svg)](https://zenodo.org/badge/latestdoi/259481291) 4 | 5 | ## Devito4PyTorch 6 | 7 | Devito4PyTorch integrates [Devito](https://www.devitoproject.org/) into PyTorch via defining highly optimized finite difference kernels as PyTorch "layers". Devito4PyTorch is an extension of Devito, a symbolic finite-difference domain specific language that provides a high-level interface to the definition of partial differential equations (PDE), such as wave equation. During backpropagation, Devito4PyTorch calls Devito's adjoint PDE solvers, thus making it possible to backpropagate efficiently through the composition of PDE solvers and neural networks. 8 | 9 | ## Installation 10 | 11 | You can install the package with `pip` via 12 | 13 | ```bash 14 | pip install git+https://github.com/slimgroup/Devito4PyTorch 15 | ``` 16 | 17 | Or if you want a developper version, you can clone and install the package as 18 | 19 | ```bash 20 | git clone https://github.com/slimgroup/Devito4PyTorch 21 | cd Devito4PyTorch 22 | pip install -e . 23 | ``` 24 | 25 | # GPU requirement 26 | 27 | If you wish to run experiments using a GPU, you will need to install `cudatoolkit`. To do so, you can use conda as: 28 | 29 | ```bash 30 | conda install cudatoolkit=10.1 -c pytorch 31 | ``` 32 | -------------------------------------------------------------------------------- /devito4pytorch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/slimgroup/Devito4PyTorch/38e225f2dd41205d7a5175d6bff0f1aed95bca8c/devito4pytorch/__init__.py -------------------------------------------------------------------------------- /devito4pytorch/devito_wrapper.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from devito.builtins import initialize_function 4 | from devito import Function 5 | 6 | 7 | class ForwardBorn(torch.autograd.Function): 8 | """ 9 | Wrapping the forward-born operator 10 | """ 11 | @staticmethod 12 | def forward(ctx, input, model, geometry, solver, device): 13 | 14 | ctx.model = model 15 | ctx.geometry = geometry 16 | ctx.solver = solver 17 | ctx.device = device 18 | 19 | # Prepare input 20 | input = torch.nn.ReplicationPad2d((ctx.model.nbl))(input) 21 | input = input.detach().cpu().numpy() 22 | 23 | # Linearized forward modeling 24 | d_lin = ctx.solver.jacobian(input[0, 0, :, :])[0].data 25 | 26 | return torch.from_numpy(np.array(d_lin)).to(ctx.device) 27 | 28 | @staticmethod 29 | def backward(ctx, grad_output): 30 | 31 | grad_output = grad_output.detach().cpu().numpy() 32 | rec = ctx.geometry.rec 33 | rec.data[:] = grad_output[:] 34 | 35 | # Adjoint linearized modeling 36 | u0 = ctx.solver.forward(save=True)[1] 37 | g = ctx.solver.jacobian_adjoint(rec, u0)[0].data 38 | 39 | # Remove padding 40 | nb = ctx.model.nbl 41 | g = torch.from_numpy(np.array(g[nb:-nb, nb:-nb])).to(ctx.device) 42 | 43 | return g.view(1, 1, g.shape[0], g.shape[1]), None, None, None, None 44 | 45 | 46 | class AdjointBorn(torch.autograd.Function): 47 | """ 48 | Wrapping the adjoint-born operator 49 | """ 50 | @staticmethod 51 | def forward(ctx, input, model, geometry, solver, device): 52 | 53 | ctx.model = model 54 | ctx.geometry = geometry 55 | ctx.solver = solver 56 | ctx.device = device 57 | 58 | # Adjoint born modeling 59 | input = input.detach().cpu().numpy() 60 | rec = ctx.geometry.rec 61 | rec.data[:] = input[:] 62 | 63 | u0 = ctx.solver.forward(save=True)[1] 64 | g = ctx.solver.jacobian_adjoint(rec, u=u0)[0].data 65 | 66 | # Remove padding 67 | nb = ctx.model.nbl 68 | g = torch.from_numpy(np.array(g[nb:-nb, nb:-nb])).to(ctx.device) 69 | 70 | return g.view(1, 1, g.shape[0], g.shape[1]) 71 | 72 | @staticmethod 73 | def backward(ctx, grad_output): 74 | 75 | grad_output = torch.nn.ReplicationPad2d((ctx.model.nbl))(grad_output) 76 | grad_output = grad_output.detach().cpu().numpy()[0, 0, :, :] 77 | 78 | # Linearized forward modeling 79 | d_lin = ctx.solver.jacobian(grad_output)[0].data 80 | 81 | return (torch.from_numpy(np.array(d_lin)).to(ctx.device), None, 82 | None, None, None) 83 | 84 | 85 | class ForwardModeling(torch.autograd.Function): 86 | """ 87 | Wrapping forward-modeling operator 88 | """ 89 | @staticmethod 90 | def forward(ctx, input, model, geometry, solver, device): 91 | 92 | ctx.model = model 93 | ctx.geometry = geometry 94 | ctx.solver = solver 95 | ctx.device = device 96 | 97 | # Prepare input 98 | input = input[0, 0, ...].detach().cpu().numpy() 99 | vp = Function(name='vp', grid=ctx.model.grid, 100 | space_order=ctx.model.space_order) 101 | initialize_function(vp, input**(-0.5), ctx.model.nbl) 102 | ctx.model.vp = vp 103 | 104 | # Nonlinear forward modeling 105 | d_nonlin, ctx.u0 = ctx.solver.forward(save=True)[:2] 106 | 107 | return torch.from_numpy(np.array(d_nonlin.data)).to(ctx.device) 108 | 109 | @staticmethod 110 | def backward(ctx, grad_output): 111 | 112 | grad_output = grad_output.detach().cpu().numpy() 113 | rec = ctx.geometry.rec 114 | rec.data[:] = grad_output[:] 115 | 116 | g = ctx.solver.jacobian_adjoint(rec, u=ctx.u0)[0].data 117 | 118 | # Remove padding 119 | nb = ctx.model.nbl 120 | g = torch.from_numpy(np.array(g[nb:-nb, nb:-nb])).to(ctx.device) 121 | 122 | return g.view(1, 1, g.shape[0], g.shape[1]), None, None, None, None 123 | 124 | 125 | class AdjointModeling(torch.autograd.Function): 126 | """ 127 | Wrapping adjoint-modeling operator 128 | """ 129 | @staticmethod 130 | def forward(ctx, input, model, geometry, solver, device): 131 | 132 | ctx.model = model 133 | ctx.geometry = geometry 134 | ctx.solver = solver 135 | ctx.device = device 136 | 137 | # Prepare input 138 | input = input[0, 0, ...].detach().cpu().numpy() 139 | vp = Function(name='vp', grid=ctx.model.grid, 140 | space_order=ctx.model.space_order) 141 | initialize_function(vp, input**(-0.5), ctx.model.nbl) 142 | ctx.model.vp = vp 143 | 144 | # Nonlinear forward modeling 145 | d_nonlin, ctx.u0 = ctx.solver.forward(save=True, 146 | vp=ctx.model.vp)[:2] 147 | 148 | return torch.from_numpy(np.array(d_nonlin.data)).to(ctx.device) 149 | 150 | @staticmethod 151 | def backward(ctx, grad_output): 152 | 153 | grad_output = grad_output.detach().cpu().numpy() 154 | rec = ctx.geometry.rec 155 | rec.data[:] = grad_output[:] 156 | 157 | g = ctx.solver.jacobian_adjoint(rec, u=ctx.u0)[0].data 158 | 159 | # Remove padding 160 | nb = ctx.model.nbl 161 | g = torch.from_numpy(np.array(g[nb:-nb, nb:-nb])).to(ctx.device) 162 | 163 | return g.view(1, 1, g.shape[0], g.shape[1]), None, None, None, None 164 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 90 3 | ignore = F403,E226,E731,W503,F405,E722,E741,W504,W605 4 | exclude = .github 5 | 6 | [tool:pytest] 7 | python_files = *_test.py 8 | addopts = --durations=20 --maxfail=5 9 | 10 | [aliases] 11 | test = pytest -v 12 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r") as fh: 4 | long_description = fh.read() 5 | 6 | reqs = ['torch', "devito @ git+https://github.com/devitocodes/devito@master", 7 | "torchvision"] 8 | setuptools.setup( 9 | name="devito4pytorch", 10 | version="0.0.5", 11 | author="Ali Siahkoohi", 12 | author_email="alisk@gatech.edu", 13 | description="Integrating Devito into PyTorch", 14 | long_description=long_description, 15 | long_description_content_type="text/markdown", 16 | url="https://github.com/slimgroup/Devito4PyTorch", 17 | license='MIT', 18 | install_requires=reqs, 19 | packages=setuptools.find_packages() 20 | ) 21 | -------------------------------------------------------------------------------- /tests/adjoint_born_test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from devito4pytorch import devito_wrapper 4 | from devito import gaussian_smooth 5 | from examples.seismic import demo_model, setup_geometry 6 | from examples.seismic.acoustic import AcousticWaveSolver 7 | 8 | if not torch.cuda.is_available(): 9 | device = torch.device('cpu') 10 | torch.set_default_tensor_type('torch.FloatTensor') 11 | else: 12 | device = torch.device('cuda') 13 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 14 | 15 | 16 | class AdjointBornLayer(torch.nn.Module): 17 | """ 18 | Creates a wrapped forward-born operator 19 | """ 20 | def __init__(self, model, geometry, device): 21 | super(AdjointBornLayer, self).__init__() 22 | self.forward_born = devito_wrapper.AdjointBorn() 23 | self.model = model 24 | self.geometry = geometry 25 | self.device = device 26 | self.solver = AcousticWaveSolver(self.model, self.geometry, 27 | space_order=8) 28 | 29 | def forward(self, x): 30 | return self.forward_born.apply(x, self.model, self.geometry, 31 | self.solver, self.device) 32 | 33 | 34 | def test_adjoint_born(): 35 | 36 | tn = 1000. 37 | shape = (101, 101) 38 | nbl = 40 39 | model = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 40 | spacing=(10., 10.), nbl=nbl, nlayers=2, 41 | space_order=8) 42 | model0 = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 43 | spacing=(10., 10.), nbl=nbl, nlayers=2, 44 | space_order=8) 45 | 46 | gaussian_smooth(model0.vp, sigma=(1, 1)) 47 | geometry0 = setup_geometry(model0, tn) 48 | 49 | # Pure Devito 50 | solver0 = AcousticWaveSolver(model0, geometry0, space_order=8) 51 | dm = model.vp.data**(-2) - model0.vp.data**(-2) 52 | 53 | grad_devito = np.array(solver0.jacobian(dm)[0].data) 54 | 55 | # Devito4PyTorch 56 | dm = torch.from_numpy(np.array(dm[nbl:-nbl, nbl:-nbl])).to(device) 57 | 58 | d_est = torch.zeros(geometry0.rec.data.shape, requires_grad=True, 59 | device=device) 60 | adjoint_born = AdjointBornLayer(model0, geometry0, device) 61 | 62 | loss = 0.5*torch.norm(adjoint_born(d_est) - dm)**2 63 | 64 | # Negative gradient to be compared with linearized data computed above 65 | grad = -torch.autograd.grad(loss, d_est, create_graph=False)[0] 66 | 67 | # Test 68 | rel_err = np.linalg.norm(grad.cpu().numpy() 69 | - grad_devito) / np.linalg.norm(grad_devito) 70 | assert np.isclose(rel_err, 0., atol=1.e-6) 71 | -------------------------------------------------------------------------------- /tests/forward_born_test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from devito4pytorch import devito_wrapper 4 | from devito import gaussian_smooth 5 | from examples.seismic import demo_model, setup_geometry 6 | from examples.seismic.acoustic import AcousticWaveSolver 7 | 8 | if not torch.cuda.is_available(): 9 | device = torch.device('cpu') 10 | torch.set_default_tensor_type('torch.FloatTensor') 11 | else: 12 | device = torch.device('cuda') 13 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 14 | 15 | 16 | class ForwardBornLayer(torch.nn.Module): 17 | """ 18 | Creates a wrapped forward-born operator 19 | """ 20 | def __init__(self, model, geometry, device): 21 | super(ForwardBornLayer, self).__init__() 22 | self.forward_born = devito_wrapper.ForwardBorn() 23 | self.model = model 24 | self.geometry = geometry 25 | self.device = device 26 | self.solver = AcousticWaveSolver(self.model, self.geometry, 27 | space_order=8) 28 | 29 | def forward(self, x): 30 | return self.forward_born.apply(x, self.model, self.geometry, 31 | self.solver, self.device) 32 | 33 | 34 | def test_forward_born(): 35 | 36 | tn = 1000. 37 | shape = (101, 101) 38 | nbl = 40 39 | model = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 40 | spacing=(10., 10.), nbl=nbl, nlayers=2, 41 | space_order=8) 42 | model0 = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 43 | spacing=(10., 10.), nbl=nbl, nlayers=2, 44 | space_order=8) 45 | 46 | gaussian_smooth(model0.vp, sigma=(1, 1)) 47 | geometry0 = setup_geometry(model0, tn) 48 | geometry = setup_geometry(model, tn) 49 | 50 | # Pure Devito 51 | solver = AcousticWaveSolver(model, geometry, space_order=8) 52 | solver0 = AcousticWaveSolver(model0, geometry0, space_order=8) 53 | 54 | d = solver.forward(vp=model.vp)[0] 55 | d0, u0 = solver0.forward(save=True, vp=model0.vp)[:2] 56 | d_lin = d.data - d0.data 57 | 58 | rec = geometry0.rec 59 | rec.data[:] = -d_lin[:] 60 | grad_devito = solver0.jacobian_adjoint(rec, u0)[0].data 61 | grad_devito = np.array(grad_devito)[nbl:-nbl, nbl:-nbl] 62 | 63 | # Devito4PyTorch 64 | d_lin = torch.from_numpy(np.array(d_lin)).to(device) 65 | 66 | forward_born = ForwardBornLayer(model0, geometry0, device) 67 | dm_est = torch.zeros([1, 1, shape[0], shape[1]], requires_grad=True, 68 | device=device) 69 | 70 | loss = 0.5*torch.norm(forward_born(dm_est) - d_lin)**2 71 | grad = torch.autograd.grad(loss, dm_est, create_graph=False)[0] 72 | 73 | # Test 74 | rel_err = np.linalg.norm(grad.cpu().numpy() 75 | - grad_devito) / np.linalg.norm(grad_devito) 76 | assert np.isclose(rel_err, 0., atol=1.e-6) 77 | -------------------------------------------------------------------------------- /tests/forward_modeling_test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from devito4pytorch import devito_wrapper 4 | from devito import gaussian_smooth 5 | from examples.seismic import demo_model, setup_geometry 6 | from examples.seismic.acoustic import AcousticWaveSolver 7 | 8 | if not torch.cuda.is_available(): 9 | device = torch.device('cpu') 10 | torch.set_default_tensor_type('torch.FloatTensor') 11 | else: 12 | device = torch.device('cuda') 13 | torch.set_default_tensor_type('torch.cuda.FloatTensor') 14 | 15 | 16 | class ForwardModelingLayer(torch.nn.Module): 17 | """ 18 | Creates a wrapped forward-modeling operator 19 | """ 20 | def __init__(self, model, geometry, device): 21 | super(ForwardModelingLayer, self).__init__() 22 | self.forward_modeling = devito_wrapper.ForwardModeling() 23 | self.model = model 24 | self.geometry = geometry 25 | self.device = device 26 | self.solver = AcousticWaveSolver(self.model, self.geometry, 27 | space_order=8) 28 | 29 | def forward(self, x): 30 | return self.forward_modeling.apply(x, self.model, self.geometry, 31 | self.solver, self.device) 32 | 33 | 34 | def test_forward(): 35 | 36 | tn = 1000. 37 | shape = (101, 101) 38 | nbl = 40 39 | model = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 40 | spacing=(10., 10.), nbl=nbl, nlayers=2, 41 | space_order=8) 42 | model0 = demo_model('layers-isotropic', origin=(0., 0.), shape=shape, 43 | spacing=(10., 10.), nbl=nbl, nlayers=2, 44 | space_order=8) 45 | 46 | gaussian_smooth(model0.vp, sigma=(1, 1)) 47 | geometry0 = setup_geometry(model0, tn) 48 | geometry = setup_geometry(model, tn) 49 | 50 | # Pure Devito 51 | solver = AcousticWaveSolver(model, geometry, space_order=8) 52 | solver0 = AcousticWaveSolver(model0, geometry0, space_order=8) 53 | 54 | d = solver.forward()[0] 55 | d0, u0 = solver0.forward(save=True, vp=model0.vp)[:2] 56 | residual = d.data - d0.data 57 | 58 | rec = geometry0.rec 59 | rec.data[:] = -residual[:] 60 | grad_devito = solver0.jacobian_adjoint(rec, u0)[0].data 61 | grad_devito = np.array(grad_devito)[nbl:-nbl, nbl:-nbl] 62 | 63 | # Devito4PyTorch 64 | d = torch.from_numpy(np.array(d.data)).to(device) 65 | 66 | forward_modeling = ForwardModelingLayer(model0, geometry0, device) 67 | 68 | m0 = np.array(model0.vp.data**(-2))[nbl:-nbl, nbl:-nbl] 69 | m0 = torch.Tensor(m0).unsqueeze(0).unsqueeze(0).to(device) 70 | m0.requires_grad = True 71 | 72 | loss = 0.5*torch.norm(forward_modeling(m0) - d)**2 73 | grad = torch.autograd.grad(loss, m0, create_graph=False)[0] 74 | 75 | # Test 76 | rel_err = np.linalg.norm(grad.cpu().numpy() 77 | - grad_devito) / np.linalg.norm(grad_devito) 78 | assert np.isclose(rel_err, 0., atol=1.e-6) 79 | --------------------------------------------------------------------------------