├── requirements.txt ├── occamypy ├── __version__.py ├── cupy │ ├── __init__.py │ ├── operator │ │ ├── __init__.py │ │ └── signal.py │ └── vector.py ├── numpy │ ├── __init__.py │ ├── operator │ │ ├── __init__.py │ │ ├── transform.py │ │ ├── pylops_interface.py │ │ └── signal.py │ └── vector.py ├── torch │ ├── __init__.py │ ├── operator │ │ ├── __init__.py │ │ ├── transform.py │ │ └── signal.py │ ├── back_utils.py │ ├── autograd.py │ └── vector.py ├── vector │ ├── __init__.py │ └── axis_info.py ├── dask │ ├── __init__.py │ └── utils.py ├── utils │ ├── __init__.py │ ├── logger.py │ ├── backend.py │ ├── plot.py │ ├── os.py │ └── sep.py ├── __init__.py ├── problem │ ├── __init__.py │ └── base.py ├── solver │ └── __init__.py └── operator │ ├── __init__.py │ ├── matrix.py │ ├── linear.py │ ├── signal.py │ └── derivative.py ├── readme_img └── logo192.png ├── tutorials ├── data │ ├── python.npy │ ├── monarch.npy │ ├── ricker20.npy │ ├── Marmousi2Vp.bin │ ├── reflectivity1D.npy │ └── shepp_logan_phantom.npy ├── devitoseismic │ ├── acoustic │ │ ├── __init__.py │ │ ├── acoustic_example.py │ │ ├── acoustic_time_update_nb.ipynb │ │ ├── operators.py │ │ └── wavesolver.py │ ├── __init__.py │ ├── plotting.py │ ├── utils.py │ └── source.py └── born_devito.py ├── .gitignore ├── envs ├── devel.yml ├── env.yml ├── env_cupy.yml ├── devel_cupy.yml └── swung.yml ├── LICENSE.md ├── setup.py ├── CONTRIBUTING.md ├── CHANGELOG.md └── README.md /requirements.txt: -------------------------------------------------------------------------------- 1 | . 2 | -------------------------------------------------------------------------------- /occamypy/__version__.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.2.0" 2 | -------------------------------------------------------------------------------- /readme_img/logo192.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/readme_img/logo192.png -------------------------------------------------------------------------------- /tutorials/data/python.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/python.npy -------------------------------------------------------------------------------- /tutorials/data/monarch.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/monarch.npy -------------------------------------------------------------------------------- /tutorials/data/ricker20.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/ricker20.npy -------------------------------------------------------------------------------- /tutorials/data/Marmousi2Vp.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/Marmousi2Vp.bin -------------------------------------------------------------------------------- /tutorials/data/reflectivity1D.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/reflectivity1D.npy -------------------------------------------------------------------------------- /occamypy/cupy/__init__.py: -------------------------------------------------------------------------------- 1 | from .vector import * 2 | from .operator import * 3 | 4 | __all__ = [ 5 | "VectorCupy", 6 | ] 7 | -------------------------------------------------------------------------------- /tutorials/data/shepp_logan_phantom.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fpicetti/occamypy/HEAD/tutorials/data/shepp_logan_phantom.npy -------------------------------------------------------------------------------- /tutorials/devitoseismic/acoustic/__init__.py: -------------------------------------------------------------------------------- 1 | from .operators import * # noqa 2 | from .wavesolver import * # noqa 3 | from .acoustic_example import * # noqa 4 | -------------------------------------------------------------------------------- /occamypy/cupy/operator/__init__.py: -------------------------------------------------------------------------------- 1 | from .signal import * 2 | 3 | __all__ = [ 4 | "GaussianFilter", 5 | "ConvND", 6 | "Padding", 7 | "ZeroPad", 8 | ] 9 | -------------------------------------------------------------------------------- /occamypy/numpy/__init__.py: -------------------------------------------------------------------------------- 1 | from .vector import VectorNumpy 2 | from .operator import * 3 | from .operator import pylops_interface 4 | 5 | __all__ = [ 6 | "VectorNumpy", 7 | "pylops_interface", 8 | ] 9 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/__init__.py: -------------------------------------------------------------------------------- 1 | from .model import * # noqa 2 | from .source import * # noqa 3 | from .plotting import * # noqa 4 | from .preset_models import * # noqa 5 | from .utils import * # noqa 6 | -------------------------------------------------------------------------------- /occamypy/torch/__init__.py: -------------------------------------------------------------------------------- 1 | from .operator import * 2 | from .vector import * 3 | from .autograd import * 4 | 5 | __all__ = [ 6 | "VectorTorch", 7 | "VectorAD", 8 | "AutogradFunction", 9 | ] 10 | -------------------------------------------------------------------------------- /occamypy/numpy/operator/__init__.py: -------------------------------------------------------------------------------- 1 | from .signal import * 2 | from .transform import * 3 | 4 | __all__ = [ 5 | "ConvND", 6 | "GaussianFilter", 7 | "Padding", 8 | "ZeroPad", 9 | "FFT", 10 | ] 11 | -------------------------------------------------------------------------------- /occamypy/torch/operator/__init__.py: -------------------------------------------------------------------------------- 1 | from .signal import * 2 | from .transform import * 3 | 4 | __all__ = [ 5 | "ConvND", 6 | "GaussianFilter", 7 | "Padding", 8 | "ZeroPad", 9 | "FFT", 10 | ] 11 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | *.pyc 3 | /occamypy.egg-info 4 | /build 5 | /dist 6 | .DS_Store 7 | .ipynb_checkpoints/ 8 | dask-worker-space 9 | /devel 10 | publishing.md 11 | tutorials/devito 12 | pykonal-0.2.3b3 13 | **__pycache__ -------------------------------------------------------------------------------- /occamypy/vector/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import * 2 | from .out_core import * 3 | from .axis_info import AxInfo 4 | 5 | __all__ = [ 6 | "Vector", 7 | "superVector", 8 | "VectorSet", 9 | "VectorOC", 10 | "AxInfo", 11 | ] 12 | -------------------------------------------------------------------------------- /occamypy/dask/__init__.py: -------------------------------------------------------------------------------- 1 | from .vector import * 2 | from .operator import * 3 | from .utils import * 4 | 5 | __all__ = [ 6 | "DaskClient", 7 | "DaskVector", 8 | "DaskOperator", 9 | "DaskSpread", 10 | "DaskCollect", 11 | ] 12 | -------------------------------------------------------------------------------- /occamypy/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .backend import * 2 | from .logger import * 3 | from .os import * 4 | from .sep import * 5 | 6 | __all__ = [ 7 | "Logger", 8 | "read_file", 9 | "write_file", 10 | "RunShellCmd", 11 | "hashfile", 12 | "mkdir", 13 | "rand_name", 14 | "CUPY_ENABLED", 15 | "ZERO", 16 | "get_backend", 17 | ] 18 | -------------------------------------------------------------------------------- /occamypy/__init__.py: -------------------------------------------------------------------------------- 1 | from .__version__ import __version__ 2 | from .vector import * 3 | from .operator import * 4 | from .numpy import * 5 | from .problem import * 6 | from .solver import * 7 | from .dask import * 8 | from .utils import * 9 | from .torch import * 10 | from .utils import plot 11 | from .utils import backend 12 | 13 | if CUPY_ENABLED: 14 | from .cupy import * 15 | 16 | -------------------------------------------------------------------------------- /envs/devel.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | dependencies: 5 | - jupyter 6 | - numpy>=1.15 7 | - scipy>=1.4 8 | - numba 9 | - matplotlib 10 | - imageio 11 | - dask 12 | - dask-jobqueue 13 | - dask-kubernetes 14 | - bokeh>=2.1.1 15 | - pylops 16 | - pip 17 | - pip: 18 | - gputil 19 | - torch>=1.7 # here for automatic cuda installation 20 | - functorch -------------------------------------------------------------------------------- /envs/env.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | dependencies: 5 | - jupyter 6 | - numpy>=1.15 7 | - scipy>=1.4 8 | - numba 9 | - matplotlib 10 | - imageio 11 | - dask 12 | - dask-jobqueue 13 | - dask-kubernetes 14 | - bokeh>=2.1.1 15 | - pylops 16 | - pip 17 | - pip: 18 | - gputil 19 | - torch>=1.7 # here for automatic cuda installation 20 | - occamypy 21 | -------------------------------------------------------------------------------- /occamypy/problem/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import * 2 | from .linear import * 3 | from .nonlinear import * 4 | 5 | __all__ = [ 6 | "Bounds", 7 | "Problem", 8 | "LeastSquares", 9 | "LeastSquaresSymmetric", 10 | "LeastSquaresRegularized", 11 | "Lasso", 12 | "GeneralizedLasso", 13 | "NonlinearLeastSquares", 14 | "NonlinearLeastSquaresRegularized", 15 | "VarProRegularized", 16 | ] 17 | -------------------------------------------------------------------------------- /envs/env_cupy.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | - rapidsai 5 | dependencies: 6 | - jupyter 7 | - numpy>=1.15 8 | - scipy>=1.4 9 | - numba 10 | - matplotlib 11 | - imageio 12 | - dask 13 | - dask-jobqueue 14 | - dask-kubernetes 15 | - bokeh>=2.1.1 16 | - pylops 17 | - cupy 18 | - cusignal 19 | - pip 20 | - pip: 21 | - gputil 22 | - torch>=1.7 # here for automatic cuda installation 23 | - occamypy 24 | -------------------------------------------------------------------------------- /envs/devel_cupy.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | - rapidsai 5 | dependencies: 6 | - jupyter 7 | - numpy>=1.15 8 | - scipy>=1.4 9 | - numba 10 | - matplotlib 11 | - imageio 12 | - dask 13 | - dask-jobqueue 14 | - dask-kubernetes 15 | - bokeh>=2.1.1 16 | - pylops 17 | - cupy # if cuda driver is correctly compiled 18 | - cusignal # to use cupy signal 19 | - pip 20 | - pip: 21 | - gputil 22 | - torch>=1.7 # here for automatic cuda installation 23 | - functorch -------------------------------------------------------------------------------- /envs/swung.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | dependencies: 5 | - jupyter 6 | - numpy>=1.15 7 | - scipy>=1.4 8 | - numba 9 | - matplotlib 10 | - imageio 11 | - dask 12 | - dask-jobqueue 13 | - dask-kubernetes 14 | - bokeh>=2.1.1 15 | - pylops 16 | - cython # for the tomography example 17 | - pip 18 | - pip: 19 | - gputil 20 | - torch>=1.7 # here for automatic cuda installation 21 | - functorch 22 | - git+https://github.com/devitocodes/devito.git # for the LS-RTM example 23 | - git+https://github.com/malcolmw/pykonal@0.2.3b3 # for the tomography example 24 | - occamypy -------------------------------------------------------------------------------- /occamypy/solver/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import * 2 | from .stopper import * 3 | from .linear import * 4 | from .nonlinear import * 5 | from .sparsity import * 6 | from .stepper import * 7 | from .mcmc import * 8 | 9 | __all__ = [ 10 | "Solver", 11 | "BasicStopper", 12 | "SamplingStopper", 13 | "CG", 14 | "SD", 15 | "LSQR", 16 | "CGsym", 17 | "NLCG", 18 | "LBFGS", 19 | "LBFGSB", 20 | "TNewton", 21 | "MCMC", 22 | "ISTA", 23 | "FISTA", 24 | "ISTC", 25 | "SplitBregman", 26 | "Stepper", 27 | "CvSrchStep", 28 | "ParabolicStep", 29 | "ParabolicStepConst", 30 | "StrongWolfe", 31 | ] 32 | -------------------------------------------------------------------------------- /occamypy/operator/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import * 2 | from .linear import * 3 | from .nonlinear import * 4 | from .derivative import * 5 | from .matrix import * 6 | from .signal import * 7 | 8 | __all__ = [ 9 | "Operator", 10 | "Vstack", 11 | "Hstack", 12 | "Dstack", 13 | "Chain", 14 | "Zero", 15 | "Identity", 16 | "Scaling", 17 | "Diagonal", 18 | "Matrix", 19 | "NonlinearOperator", 20 | "NonlinearComb", 21 | "NonlinearSum", 22 | "NonlinearVstack", 23 | "VarProOperator", 24 | "FirstDerivative", 25 | "SecondDerivative", 26 | "Gradient", 27 | "Laplacian", 28 | "FFT", 29 | "ConvND", 30 | "GaussianFilter", 31 | "Padding", 32 | "ZeroPad", 33 | ] 34 | -------------------------------------------------------------------------------- /occamypy/utils/logger.py: -------------------------------------------------------------------------------- 1 | class Logger: 2 | """System logger class to store messages to a log file""" 3 | def __init__(self, file, bufsize: int = 1): 4 | """ 5 | Logger constructor 6 | 7 | Args: 8 | file: path/to/file to write the log 9 | bufsize: buffer size 10 | """ 11 | self.file = open(file, "a", bufsize) 12 | 13 | def __del__(self): 14 | """Close the log file""" 15 | self.file.close() 16 | return 17 | 18 | def addToLog(self, msg: str): 19 | """Write a message to log file 20 | 21 | Args: 22 | msg: message to be added to log file 23 | """ 24 | self.file.write(msg + "\n") 25 | return 26 | -------------------------------------------------------------------------------- /occamypy/vector/axis_info.py: -------------------------------------------------------------------------------- 1 | from typing import NamedTuple 2 | 3 | import numpy as np 4 | 5 | 6 | class AxInfo(NamedTuple): 7 | """ 8 | Store information about vectors' axis 9 | 10 | Attributes: 11 | N: number of samples along the axis 12 | o: axis origin value (i.e., value the first sample) 13 | d: sampling/discretization step 14 | l: label of the axis 15 | last: value of the last sample 16 | """ 17 | n: int = 1 18 | o: float = 0. 19 | d: float = 1. 20 | l: str = "undefined" 21 | 22 | def to_string(self, ax: int = 1): 23 | """ 24 | Create a description of the axis 25 | 26 | Args: 27 | ax: axis number for printing (for SEPlib compatibility) 28 | 29 | Returns: string 30 | 31 | """ 32 | return "n%s=%s o%s=%s d%s=%s label%s='%s'\n" % (ax, self.n, ax, self.o, ax, self.d, ax, self.l) 33 | 34 | def plot(self): 35 | """ 36 | Create a np.ndarray of the axis, useful for plotting 37 | 38 | """ 39 | return np.arange(self.n) * self.d + self.o 40 | 41 | @property 42 | def last(self): 43 | return self.o + (self.n - 1) * self.d 44 | -------------------------------------------------------------------------------- /occamypy/utils/backend.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy 4 | import torch 5 | 6 | try: 7 | import cupy 8 | except ModuleNotFoundError: 9 | pass 10 | 11 | __all__ = [ 12 | "set_seed_everywhere", 13 | "get_backend", 14 | "get_vector_type", 15 | ] 16 | 17 | 18 | def set_seed_everywhere(seed=0): 19 | """Set random seed on every computation engine""" 20 | numpy.random.seed(seed) 21 | torch.manual_seed(seed) 22 | torch.cuda.manual_seed(seed) 23 | os.environ["CUPY_SEED"] = str(seed) 24 | 25 | 26 | def get_backend(vector): 27 | """ 28 | Get the vector content backend 29 | Args: 30 | vector: vector to analyze 31 | 32 | Returns: package (one of numpy, cupy, torch) 33 | """ 34 | if vector.whoami == "VectorTorch": 35 | backend = torch 36 | elif vector.whoami == "VectorNumpy": 37 | backend = numpy 38 | elif vector.whoami == "VectorCupy": 39 | backend = cupy 40 | 41 | return backend 42 | 43 | 44 | def get_vector_type(vector): 45 | """ 46 | Get the vector content original classs 47 | Args: 48 | vector: vector to analyze 49 | 50 | Returns: array class (numpy.ndarray, cupy.ndarray, torch.Tensor) 51 | """ 52 | if vector.whoami == "VectorTorch": 53 | return torch.Tensor 54 | elif vector.whoami == "VectorNumpy": 55 | return numpy.ndarray 56 | elif vector.whoami == "VectorCupy": 57 | return cupy.ndarray 58 | -------------------------------------------------------------------------------- /occamypy/torch/back_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from GPUtil import getFirstAvailable 4 | 5 | __all__ = [ 6 | "set_backends", 7 | "get_device", 8 | "get_device_name", 9 | ] 10 | 11 | 12 | def set_backends(): 13 | """Set the GPU backend to enable reproducibility""" 14 | torch.backends.cudnn.enabled = True 15 | torch.backends.cudnn.benchmark = True 16 | torch.backends.cudnn.deterministic = True 17 | 18 | 19 | def get_device(devID: int = None) -> torch.device: 20 | """ 21 | Get the computation device 22 | 23 | Args: 24 | devID: device id to be used (None for CPU, -1 for max free memory) 25 | 26 | Returns: torch.device object 27 | """ 28 | if devID is None: 29 | dev = "cpu" 30 | elif devID == -1: 31 | dev = getFirstAvailable(order='memory')[0] 32 | else: 33 | dev = int(devID) 34 | if dev > torch.cuda.device_count(): 35 | dev = "cpu" 36 | raise UserWarning("The selected device is not available, switched to CPU.") 37 | 38 | return torch.device(dev) 39 | 40 | 41 | def get_device_name(devID: int = None) -> str: 42 | """ 43 | Get the device name as a nice string 44 | 45 | Args: 46 | devID: device ID for torch 47 | """ 48 | if devID is None or isinstance(torch.cuda.get_device_name(devID), torch.NoneType): 49 | return "CPU" 50 | else: 51 | return "GPU %d - %s" % (devID, torch.cuda.get_device_name(devID)) 52 | -------------------------------------------------------------------------------- /occamypy/operator/matrix.py: -------------------------------------------------------------------------------- 1 | from occamypy.vector.base import Vector 2 | from occamypy.operator.base import Operator 3 | from occamypy.utils.backend import get_backend, get_vector_type 4 | 5 | 6 | class Matrix(Operator): 7 | """ 8 | Linear Operator build upon an explicit matrix 9 | 10 | Attributes: 11 | matrix: Vector array that contains the matrix 12 | """ 13 | def __init__(self, matrix: Vector, domain: Vector, range: Vector, outcore=False): 14 | """ 15 | Matrix constructor 16 | 17 | Args: 18 | matrix: vector that contains the matrix 19 | domain: domain vector 20 | range: range vector 21 | outcore: whether to use out-of-core SEPlib operators 22 | """ 23 | if not (type(domain) == type(range) == type(matrix)): 24 | raise TypeError("ERROR! Domain, Range and Matrix have to be the same vector type") 25 | 26 | if matrix.shape[1] != domain.size: 27 | raise ValueError 28 | if matrix.shape[0] != range.size: 29 | raise ValueError 30 | 31 | super(Matrix, self).__init__(domain=domain, range=range, name="Matrix") 32 | self.backend = get_backend(matrix) 33 | self.matrix_type = get_vector_type(matrix) 34 | 35 | self.matrix = matrix 36 | self.outcore = outcore 37 | 38 | def forward(self, add, model, data): 39 | self.checkDomainRange(model, data) 40 | if not add: 41 | data.zero() 42 | data[:] += self.backend.matmul(self.matrix[:], model[:].flatten()).reshape(data.shape) 43 | return 44 | 45 | def adjoint(self, add, model, data): 46 | self.checkDomainRange(model, data) 47 | if not add: 48 | model.zero() 49 | model[:] += self.backend.matmul(self.matrix.hermitian()[:], data[:].flatten()).reshape(model.shape) 50 | return 51 | 52 | def getNdArray(self): 53 | """Get the matrix vector""" 54 | return self.matrix.getNdArray() 55 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 2 | 3 | * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 4 | * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 5 | * Neither the name of Stanford University, nor of Politecnico of Milan, nor the name of Stanford Exploration Project, nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. 6 | * If the software is used to develop scientific or technical material that is published in any peer-reviewed papers, conference abstracts or similar publications, the recipient agrees to acknowledge the authors of the source code in a manner consistent with industry practice by citing the provided references. 7 | * The authors would appreciate being notified of any errors found in the supplied code by opening a GIT issue or by emailing: ebiondi@caltech.edu or francesco.picetti@polimi.it or barnier@gmail.com or bob@sep.stanford.edu or sfarris@stanford.edu 8 | 9 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 10 | 11 | REFERENCES LIST: 12 | 13 | E. Biondi, G. Barnier, R. G. Clapp, F. Picetti, and S. Farris, 2020, An object-oriented optimization framework for large-scale inverse problems: SIAM Journal of Scientific Computing (submitted) -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from setuptools import setup, find_packages 3 | 4 | NAME = "occamypy" 5 | DESCRIPTION = "OccamyPy. An object-oriented optimization library for small- and large-scale problems." 6 | URL = "https://github.com/fpicetti/occamypy" 7 | EMAIL = "francesco.picetti@polimi.it" 8 | AUTHOR = "Ettore Biondi, Guillame Barnier, Robert Clapp, Francesco Picetti, Stuart Farris" 9 | REQUIRES_PYTHON = ">=3.6.0" 10 | PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) 11 | 12 | 13 | def load_readme(): 14 | readme_path = os.path.join(PROJECT_ROOT, "README.md") 15 | with open(readme_path, encoding="utf-8") as f: 16 | return f"\n{f.read()}" 17 | 18 | 19 | def load_version(): 20 | context = {} 21 | with open(os.path.join(PROJECT_ROOT, "occamypy", "__version__.py")) as f: 22 | exec(f.read(), context) 23 | return context["__version__"] 24 | 25 | 26 | setup(name='occamypy', 27 | version=load_version(), 28 | url="https://github.com/fpicetti/occamypy", 29 | description='An Object-Oriented Optimization Framework for Large-Scale Inverse Problems', 30 | long_description=load_readme(), 31 | long_description_content_type='text/markdown', 32 | keywords=['algebra', 'inverse problems', 'large-scale optimization'], 33 | classifiers=[ 34 | 'Development Status :: 4 - Beta', 35 | 'Intended Audience :: Developers', 36 | 'Intended Audience :: Science/Research', 37 | 'Natural Language :: English', 38 | 'Programming Language :: Python :: 3.6', 39 | 'Programming Language :: Python :: 3.7', 40 | 'Programming Language :: Python :: 3.8', 41 | 'Topic :: Scientific/Engineering :: Mathematics', 42 | 'Operating System :: Unix' 43 | ], 44 | author=AUTHOR, 45 | author_email=EMAIL, 46 | install_requires=['numpy', 47 | 'scipy', 48 | 'numba', 49 | 'torch>=1.7.0', 50 | 'dask', 51 | 'dask-jobqueue', 52 | 'dask-kubernetes', 53 | 'matplotlib', 54 | 'gputil', 55 | 'imageio', 56 | ], 57 | packages=find_packages(), 58 | zip_safe=True) 59 | -------------------------------------------------------------------------------- /occamypy/utils/plot.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import numpy as np 4 | from imageio import mimsave 5 | 6 | 7 | def float2png(in_content: np.ndarray) -> np.ndarray: 8 | """ 9 | Convert a np.ndarray dynamics to [0,255] uint8 10 | 11 | Args: 12 | in_content: float image to be converted 13 | 14 | Returns: 15 | image in the png dynamics 16 | """ 17 | in_min = np.min(in_content) 18 | in_max = np.max(in_content) 19 | in_content = 255 * (in_content - in_min) / (in_max - in_min) # x in [0,255] 20 | return in_content.astype(np.uint8) 21 | 22 | 23 | def vector2gif(in_content: np.ndarray, filename: str, transpose: bool = False, clip: float = 100., frame_axis=-1, 24 | fps: int = 25) -> None: 25 | """ 26 | Save a 3D vector to an animated GIF file. 27 | 28 | Args: 29 | in_content: 3D numpy.array 30 | filename: path to file to be written with extension 31 | transpose: whether to transpose the image axis [False]. Notice that sepvector have the axis swapped w.r.t. numpy 32 | clip: percentile for clipping the data 33 | frame_axis: frame axis of data [-1] 34 | fps: number of frames per second [25] 35 | """ 36 | if in_content.ndim != 3: 37 | raise ValueError("in_content has to be a 3D vector") 38 | 39 | if clip != 100.: 40 | clip_val = np.percentile(np.absolute(in_content), clip) 41 | in_content = np.clip(in_content, -clip_val, clip_val) 42 | 43 | in_content = float2png(in_content) 44 | 45 | if frame_axis != 0: # need to bring the dim. to first dim 46 | in_content = np.swapaxes(in_content, frame_axis, 0) 47 | if transpose: 48 | frames = [in_content[_].T for _ in range(in_content.shape[0])] 49 | else: 50 | frames = [in_content[_] for _ in range(in_content.shape[0])] 51 | 52 | mimsave(filename, frames, **{'fps': fps}) 53 | 54 | 55 | def clim(in_content: np.ndarray, ratio: float = 95) -> Tuple[float, float]: 56 | """ 57 | Compute the clim tuple for plotting an array 58 | 59 | Args: 60 | in_content: array to be analyzed 61 | ratio: percentile value for clipping 62 | 63 | Returns: 64 | tuple (-c, c) of clipping values 65 | """ 66 | c = np.percentile(np.absolute(in_content), ratio) 67 | return -c, c 68 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | Contributions are welcome and greatly appreciated! 4 | 5 | ### Report Bugs 6 | 7 | Report bugs at https://github.com/fpicetti/occamypy/issues with tag *bug* including: 8 | 9 | * Your operating system name and version, and computing facilies details (e.g., a cluster). 10 | * Your Python environment variables, packages, etc. 11 | * Minimum working example. 12 | 13 | ### Propose New Objects and Features 14 | 15 | Open an issue at https://github.com/fpicetti/occamypy/issues 16 | with tag *enhancement* including: 17 | * Detailed explanation about what you would like, how it would work, etc. 18 | * Keep the scope as narrow as possible, to make it easier to implement. 19 | 20 | ### Implement Your Own Codes 21 | 22 | You are super-welcome to add your own tutorial, discuss and improve the documentation, 23 | and code your own operators, problems and solvers. Also, check the *enhancement* issue tag! 24 | We would love to expand our users pool to different areas of scientific computing and engineering! 25 | 26 | ## Getting Started to contribute 27 | 28 | Ready to contribute? 29 | 30 | 1. Fork the `occamypy` repo. 31 | 32 | 2. Clone your fork locally: 33 | ``` 34 | git clone https://github.com/your_name_here/occamypy.git 35 | ``` 36 | 37 | 3. Follow the installation instructions for *developers* that you find 38 | in the README.md or in the online documentation. 39 | Ensure that you are able to *pass all the tests before moving forward*. 40 | 41 | 4. Create a branch for local development: 42 | ``` 43 | git checkout -b name-of-your-branch 44 | ``` 45 | Now you can make your changes locally. 46 | 47 | 5. Commit your changes and push your branch to GitHub: 48 | ``` 49 | git add . 50 | git commit -m "Your detailed description of your changes." 51 | git push origin name-of-your-branch 52 | ``` 53 | Remember to add ``-u`` when pushing the branch for the first time. 54 | 55 | 6. Submit a pull request through the GitHub website. 56 | 57 | 58 | ### Pull Request Guidelines 59 | 60 | Before you submit a pull request, check that it meets these guidelines: 61 | 62 | 1. The pull request should include new tests for all the core routines that have been developed. 63 | 2. If the pull request adds functionality, please add a tutorial notebook or example script. 64 | Remember to write a docstring, keeping in mind that these codes will be read mainly by humans. 65 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # 0.2.0 2 | * Added a `name` attribute to Operator, Problem, and Solver classes (useful for printing/logging) 3 | * Changed to version v0.2.0 after v0.1.5 as it is more appropriate (hello developers!) 4 | * Added CG comparison to [Devito-based LS-RTM](tutorials/2D%20LS-RTM%20with%20devito,%20dask,%20and%20regularizers.ipynb) tutorial 5 | 6 | # 0.1.5 7 | * Moved `set_seed_everywhere` in the main utils submodule 8 | * Signal processing operators are called directly from occamypy init (no cupy/torch submodule needed) 9 | * Removed `snr` logic from `rand` and `randn`; now they accept distribution descriptors 10 | * A lot of cleanings in the documentation, names, and methods 11 | * Added [PyLops](tutorials/PyLops%20and%20OccamyPy%20together.ipynb) tutorial 12 | * Added Traveltime tomography tutorials in [1D](tutorials/1D%20Travel-time%20tomography%20using%20the%20Eikonal%20equation.ipynb), [2D](tutorials/2D%20Travel-time%20tomography%20using%20the%20Eikonal%20equation.ipynb), and [Marmousi](tutorials/Traveltime%20tomography%20for%20Seismic%20Exploration.ipynb) 13 | * Added [Devito-based LS-RTM](tutorials/2D%20LS-RTM%20with%20devito,%20dask,%20and%20regularizers.ipynb) tutorial 14 | * Added [autograd](occamypy/torch/autograd.py) torch submodule to cast linear operators to torch automatic differentiation engine (see the [tutorial](tutorials/2D%20LS-RTM%20with%20devito%20and%20Automatic%20Differentiation.ipynb)) 15 | * Added a future work tutorial on [automatically differentiated operators](tutorials/Automatic%20Differentiation%20for%20nonlinear%20operators.ipynb) 16 | 17 | # 0.1.4 18 | * Added support for F-contiguous arrays 19 | * Added [PyLops](https://pylops.readthedocs.io/en/stable/) interface [operators](ea8505947c926e376a6def40b1fccfbadf3940d2) 20 | * Added plot utilities 21 | * Added [`AxInfo`](tutorials/AxInfo%20-%20exploit%20physical%20vectors.ipynb) class for handling physical vectors 22 | * Added Padding operators different from ZeroPad 23 | * Improvements on VectorTorch methods and attributes 24 | * Added FISTA solver wrapper 25 | 26 | # 0.1.3 27 | * Fix circular imports 28 | 29 | # 0.1.2 30 | * Added `__getitem__()` method to vector class 31 | * Added PyTorch FFT operators 32 | * Fix convolution in PyTorch 33 | * Added a number of utilities 34 | 35 | # 0.1.1 36 | * Derivative operators are now agnostic to the computation engine 37 | * Added Dask Blocky Operator 38 | * fixed `rand()` in VectorNumpy 39 | * added kwargs for Dask Operators 40 | 41 | # 0.1.0 42 | * First official release. 43 | -------------------------------------------------------------------------------- /occamypy/numpy/operator/transform.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from occamypy.operator.base import Operator 4 | from occamypy.numpy.vector import VectorNumpy 5 | 6 | 7 | class FFT(Operator): 8 | """N-dimensional Fast Fourier Transform for complex input""" 9 | 10 | def __init__(self, domain, axes=None, nfft=None, sampling=None): 11 | """ 12 | FFT (numpy) constructor 13 | 14 | Args: 15 | domain: domain vector 16 | axes: index of axes on which the FFT is computed 17 | nfft: number of frequency bins for each axis 18 | sampling: sampling step on each axis 19 | """ 20 | if axes is None: 21 | axes = tuple(range(domain.ndim)) 22 | elif not isinstance(axes, tuple) and domain.ndim == 1: 23 | axes = (axes,) 24 | if nfft is None: 25 | nfft = domain.shape 26 | elif not isinstance(nfft, tuple) and domain.ndim == 1: 27 | nfft = (nfft,) 28 | if sampling is None: 29 | sampling = tuple([1.] * domain.ndim) 30 | elif not isinstance(sampling, tuple) and domain.ndim == 1: 31 | sampling = (sampling,) 32 | 33 | if len(axes) != len(nfft) != len(sampling): 34 | raise ValueError('axes, nffts, and sampling must have same number of elements') 35 | 36 | self.axes = axes 37 | self.nfft = nfft 38 | self.sampling = sampling 39 | 40 | self.fs = [np.fft.fftfreq(n, d=s) for n, s in zip(nfft, sampling)] 41 | 42 | dims_fft = np.asarray(domain.shape) 43 | for a, n in zip(self.axes, self.nfft): 44 | dims_fft[a] = n 45 | 46 | super(FFT, self).__init__(domain=VectorNumpy(np.zeros(domain.shape, dtype=complex)), 47 | range=VectorNumpy(np.zeros(shape=dims_fft, dtype=complex)), 48 | name="FFT") 49 | 50 | def forward(self, add, model, data): 51 | self.checkDomainRange(model, data) 52 | if not add: 53 | data.zero() 54 | data[:] += np.fft.fftn(model[:], s=self.nfft, axes=self.axes, norm='ortho') 55 | return 56 | 57 | def adjoint(self, add, model, data): 58 | self.checkDomainRange(model, data) 59 | if not add: 60 | model.zero() 61 | # here we need to separate the computation and use np.take for handling nfft > model.shape 62 | temp = np.fft.ifftn(data[:], s=self.nfft, axes=self.axes, norm='ortho') 63 | for a in self.axes: 64 | temp = np.take(temp, range(self.domain.shape[a]), axis=a) 65 | model[:] += temp 66 | return 67 | -------------------------------------------------------------------------------- /occamypy/torch/operator/transform.py: -------------------------------------------------------------------------------- 1 | from itertools import product 2 | 3 | import torch 4 | import torch.fft as fft 5 | from numpy.fft import fftfreq 6 | 7 | from occamypy.operator.base import Operator 8 | from occamypy.torch.back_utils import set_backends 9 | from occamypy.torch.vector import VectorTorch 10 | 11 | set_backends() 12 | 13 | __all__ = [ 14 | "FFT" 15 | ] 16 | 17 | 18 | class FFT(Operator): 19 | """N-dimensional Fast Fourier Transform for complex input""" 20 | 21 | def __init__(self, domain, axes=None, nfft=None, sampling=None): 22 | """ 23 | FFT (torch) constructor 24 | 25 | Args: 26 | domain: domain vector 27 | axes: dimension along which FFT is computed (all by default) 28 | nfft: number of samples in Fourier Transform for each direction (same as domain by default) 29 | sampling: sampling steps on each axis (1. by default) 30 | """ 31 | if axes is None: 32 | axes = tuple(range(domain.ndim)) 33 | elif not isinstance(axes, tuple) and domain.ndim == 1: 34 | axes = (axes,) 35 | if nfft is None: 36 | nfft = domain.shape 37 | elif not isinstance(nfft, tuple) and domain.ndim == 1: 38 | nfft = (nfft,) 39 | if sampling is None: 40 | sampling = tuple([1.] * domain.ndim) 41 | elif not isinstance(sampling, tuple) and domain.ndim == 1: 42 | sampling = (sampling,) 43 | 44 | if len(axes) != len(nfft) != len(sampling): 45 | raise ValueError('axes, nffts, and sampling must have same number of elements') 46 | 47 | self.axes = axes 48 | self.nfft = nfft 49 | self.sampling = sampling 50 | self.fs = [fftfreq(n, d=s) for n, s in zip(self.nfft, self.sampling)] 51 | 52 | dims_fft = list(domain.shape) 53 | for a, n in zip(self.axes, self.nfft): 54 | dims_fft[a] = n 55 | self.inner_idx = [torch.arange(0, domain.shape[i]) for i in range(len(dims_fft))] 56 | 57 | super(FFT, self).__init__(domain=VectorTorch(torch.zeros(domain.shape).type(torch.double)), 58 | range=VectorTorch(torch.zeros(dims_fft).type(torch.complex128)), 59 | name="FFT") 60 | 61 | def forward(self, add, model, data): 62 | self.checkDomainRange(model, data) 63 | if not add: 64 | data.zero() 65 | data[:] += fft.fftn(model.getNdArray(), s=self.nfft, dim=self.axes, norm='ortho') 66 | return 67 | 68 | def adjoint(self, add, model, data): 69 | self.checkDomainRange(model, data) 70 | if not add: 71 | model.zero() 72 | # compute IFFT 73 | x = fft.ifftn(data.getNdArray(), s=self.nfft, dim=self.axes, norm='ortho').type(model.getNdArray().dtype) 74 | # handle nfft > model.shape 75 | x = torch.Tensor([x[coord] for coord in product(*self.inner_idx)]).reshape(self.domain.shape).to(model.device) 76 | model[:] += x 77 | return 78 | -------------------------------------------------------------------------------- /occamypy/numpy/operator/pylops_interface.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from occamypy.operator.base import Operator 4 | 5 | try: 6 | import pylops 7 | except ImportError: 8 | raise UserWarning("PyLops is not installed. To use this feature please run: pip install pylops") 9 | 10 | __all__ = [ 11 | "ToPylops", 12 | "FromPylops", 13 | ] 14 | 15 | 16 | class FromPylops(Operator): 17 | """Cast a pylops.LinearOperator to occamypy.Operator""" 18 | 19 | def __init__(self, domain, range, op): 20 | """ 21 | FromPylops constructor 22 | 23 | Args: 24 | domain: domain vector 25 | range: range vector 26 | op: pylops LinearOperator 27 | """ 28 | if not isinstance(op, pylops.LinearOperator): 29 | raise TypeError("op has to be a pylops.LinearOperator") 30 | if op.shape[0] != range.size: 31 | raise ValueError("Range and operator rows mismatch") 32 | if op.shape[1] != domain.size: 33 | raise ValueError("Domain and operator columns mismatch") 34 | 35 | self.op = op 36 | 37 | super(FromPylops, self).__init__(domain, range, name=op.__str__().replace("<", "").replace(">", "")) 38 | 39 | def forward(self, add, model, data): 40 | self.checkDomainRange(model, data) 41 | if not add: 42 | data.zero() 43 | x = model.getNdArray().ravel() 44 | y = self.op.matvec(x) 45 | data.getNdArray()[:] += y.reshape(data.shape) 46 | return 47 | 48 | def adjoint(self, add, model, data): 49 | self.checkDomainRange(model, data) 50 | if not add: 51 | model.zero() 52 | y = data.getNdArray().ravel() 53 | x = self.op.rmatvec(y) 54 | model[:] += x.reshape(model.shape) 55 | return 56 | 57 | 58 | class ToPylops(pylops.LinearOperator): 59 | """Cast an numpy-based occamypy.Operator to pylops.LinearOperator""" 60 | 61 | def __init__(self, op: Operator): 62 | """ 63 | ToPylops constructor 64 | 65 | Args: 66 | op: occamypy.Operator 67 | """ 68 | super(ToPylops, self).__init__(explicit=False) 69 | self.shape = (op.range.size, op.domain.size) 70 | self.dtype = op.domain.getNdArray().dtype 71 | 72 | if not isinstance(op, Operator): 73 | raise TypeError("op has to be an Operator") 74 | self.op = op 75 | 76 | # these are just temporary vectors, used by forward and adjoint computations 77 | self.domain = op.domain.clone() 78 | self.range = op.range.clone() 79 | self._name = op.name 80 | 81 | def _matvec(self, x: np.ndarray) -> np.ndarray: 82 | x_ = self.domain.clone() 83 | x_[:] = x.reshape(self.domain.shape).astype(self.dtype) 84 | y_ = self.range.clone() 85 | self.op.forward(False, x_, y_) 86 | return y_.getNdArray().ravel() 87 | 88 | def _rmatvec(self, y: np.ndarray) -> np.ndarray: 89 | y_ = self.range.clone() 90 | y_[:] = y.reshape(self.range.shape).astype(self.dtype) 91 | x_ = self.domain.clone() 92 | self.op.adjoint(False, x_, y_) 93 | return x_.getNdArray().ravel() 94 | -------------------------------------------------------------------------------- /occamypy/operator/linear.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from occamypy.operator.base import Operator 3 | 4 | 5 | class Zero(Operator): 6 | """Zero matrix operator; useful for Jacobian matrices that are zeros""" 7 | 8 | def __init__(self, domain, range): 9 | """ 10 | Zero constructor 11 | 12 | Args: 13 | domain: domain vector 14 | range: range vector 15 | """ 16 | super(Zero, self).__init__(domain, range, name="Zero") 17 | 18 | def forward(self, add, model, data): 19 | self.checkDomainRange(model, data) 20 | if not add: 21 | data.zero() 22 | 23 | def adjoint(self, add, model, data): 24 | self.checkDomainRange(model, data) 25 | if not add: 26 | model.zero() 27 | 28 | 29 | class Identity(Operator): 30 | """Identity operator""" 31 | 32 | def __init__(self, domain): 33 | """ 34 | Identity constructor 35 | 36 | Args: 37 | domain: domain vector 38 | """ 39 | super(Identity, self).__init__(domain, domain, name="Identity") 40 | 41 | def forward(self, add, model, data): 42 | self.checkDomainRange(model, data) 43 | if add: 44 | data.scaleAdd(model) 45 | else: 46 | data.copy(model) 47 | 48 | def adjoint(self, add, model, data): 49 | self.checkDomainRange(model, data) 50 | if add: 51 | model.scaleAdd(data) 52 | else: 53 | model.copy(data) 54 | 55 | 56 | class Scaling(Operator): 57 | """scalar multiplication operator""" 58 | 59 | def __init__(self, domain, scalar): 60 | """ 61 | Scaling constructor 62 | 63 | Args: 64 | domain: domain vector 65 | scalar: scaling coefficient 66 | """ 67 | super(Scaling, self).__init__(domain, domain, name="Scaling") 68 | if not np.isscalar(scalar): 69 | raise ValueError('scalar has to be (indeed) a scalar variable') 70 | self.scalar = scalar 71 | 72 | def forward(self, add, model, data): 73 | self.checkDomainRange(model, data) 74 | data.scaleAdd(model, 1. if add else 0., self.scalar) 75 | 76 | def adjoint(self, add, model, data): 77 | self.checkDomainRange(model, data) 78 | model.scaleAdd(data, 1. if add else 0., self.scalar) 79 | 80 | 81 | class Diagonal(Operator): 82 | """Diagonal operator for performing element-wise multiplication""" 83 | 84 | def __init__(self, diag): 85 | """ 86 | Diagonal constructor 87 | 88 | Args: 89 | diag: vector to be stored on the diagonal 90 | """ 91 | super(Diagonal, self).__init__(diag, diag, name="Diagonal") 92 | self.diag = diag 93 | 94 | def forward(self, add, model, data): 95 | self.checkDomainRange(model, data) 96 | data.scaleAdd(model, 1. if add else 0.) 97 | data.multiply(self.diag) 98 | 99 | def adjoint(self, add, model, data): 100 | self.checkDomainRange(model, data) 101 | model.scaleAdd(data, 1. if add else 0.) 102 | model.multiply(self.diag) 103 | -------------------------------------------------------------------------------- /occamypy/operator/signal.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Union 2 | 3 | from occamypy.utils.os import CUPY_ENABLED 4 | 5 | from occamypy.numpy import operator as np 6 | from occamypy.torch import operator as tc 7 | if CUPY_ENABLED: 8 | from occamypy.cupy import operator as cp 9 | 10 | 11 | __all__ = [ 12 | "FFT", 13 | "ConvND", 14 | "GaussianFilter", 15 | "Padding", 16 | "ZeroPad", 17 | ] 18 | 19 | 20 | def FFT(domain, axes=None, nfft=None, sampling=None): 21 | """ 22 | N-dimensional Fast Fourier Transform 23 | 24 | Args: 25 | domain: domain vector 26 | axes: dimension along which FFT is computed (all by default) 27 | nfft: number of samples in Fourier Transform for each direction (same as domain by default) 28 | sampling: sampling steps on each axis (1. by default) 29 | """ 30 | if domain.whoami == "VectorNumpy": 31 | return np.FFT(domain=domain, axes=axes, nfft=nfft, sampling=sampling) 32 | elif domain.whoami == "VectorTorch": 33 | return tc.FFT(domain=domain, axes=axes, nfft=nfft, sampling=sampling) 34 | elif domain.whoami == "VectorCupy": 35 | raise NotImplementedError("FFT operator for VectorCupy is not implemented yet") 36 | else: 37 | raise TypeError("Domain vector not recognized") 38 | 39 | 40 | def ConvND(domain, kernel, method: str = 'auto'): 41 | """ 42 | ND convolution square operator in the domain space 43 | 44 | Args: 45 | domain: domain vector 46 | kernel: kernel vector 47 | method: how to compute the convolution [auto, direct, fft] 48 | """ 49 | if domain.whoami == "VectorNumpy": 50 | return np.ConvND(domain=domain, kernel=kernel, method=method) 51 | elif domain.whoami == "VectorTorch": 52 | return tc.ConvND(domain=domain, kernel=kernel) 53 | elif domain.whoami == "VectorCupy": 54 | return cp.ConvND(domain=domain, kernel=kernel, method=method) 55 | else: 56 | raise TypeError("Domain vector not recognized") 57 | 58 | 59 | def GaussianFilter(domain, sigma: Tuple[float]): 60 | """ 61 | Gaussian smoothing operator 62 | 63 | Args: 64 | domain: domain vector 65 | sigma: standard deviation along the domain directions 66 | """ 67 | if domain.whoami == "VectorNumpy": 68 | return np.GaussianFilter(domain=domain, sigma=sigma) 69 | elif domain.whoami == "VectorTorch": 70 | return tc.GaussianFilter(domain=domain, sigma=sigma) 71 | elif domain.whoami == "VectorCupy": 72 | return cp.GaussianFilter(domain=domain, sigma=sigma) 73 | else: 74 | raise TypeError("Domain vector not recognized") 75 | 76 | 77 | def Padding(domain, pad: Union[int, Tuple[int]], mode: str = "constant"): 78 | """ 79 | Padding operator 80 | 81 | Notes: 82 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 83 | pad=((2,2), (3,3)) 84 | 85 | Args: 86 | domain: domain vector 87 | pad: number of samples to be added at each end of the dimension, for each dimension 88 | mode: padding mode 89 | """ 90 | if domain.whoami == "VectorNumpy": 91 | return np.Padding(domain=domain, pad=pad, mode=mode) 92 | elif domain.whoami == "VectorTorch": 93 | return tc.Padding(domain=domain, pad=pad, mode=mode) 94 | elif domain.whoami == "VectorCupy": 95 | return cp.Padding(domain=domain, pad=pad, mode=mode) 96 | else: 97 | raise TypeError("Domain vector not recognized") 98 | 99 | 100 | def ZeroPad(domain, pad): 101 | """ 102 | Zero-Padding operator 103 | 104 | Notes: 105 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 106 | pad=((2,2), (3,3)) 107 | 108 | Args: 109 | domain: domain vector 110 | pad: number of samples to be added at each end of the dimension, for each dimension 111 | """ 112 | return Padding(domain=domain, pad=pad) 113 | -------------------------------------------------------------------------------- /occamypy/utils/os.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import os 3 | import random 4 | import string 5 | import subprocess 6 | import sys 7 | from importlib.util import find_spec 8 | 9 | import numpy as np 10 | 11 | from occamypy.utils.logger import Logger 12 | 13 | CUPY_ENABLED = find_spec("cupy") is not None 14 | 15 | debug = False # Debug flag for printing screen output of RunShellCmd as it runs commands 16 | debug_log = None # File where debug outputs are written if requested (it must be a logger object) 17 | BUF_SIZE = 8388608 # read binary files in 64Mb chunks! 18 | DEVNULL = open(os.devnull, 'wb') 19 | 20 | # Check for avoid Overflow or Underflow 21 | ZERO = 10 ** (np.floor(np.log10(np.abs(float(np.finfo(np.float64).tiny)))) + 2) 22 | 23 | 24 | def mkdir(directory): 25 | try: 26 | if not os.path.exists(directory): 27 | os.makedirs(directory) 28 | except OSError: 29 | print('Error: Creating directory. ' + directory) 30 | 31 | 32 | def rand_name(N: int = 6) -> str: 33 | """Get a random sequence of N letters and numbers 34 | 35 | Args: 36 | N: desired length of the string 37 | Returns: 38 | string of length N 39 | """ 40 | return "".join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(N)) 41 | 42 | 43 | def hashfile(binfile): 44 | """Function hashing a binary file. It uses a BUF_SIZE to partially store file in memory 45 | and do not completely load the file into the RAM""" 46 | md5 = hashlib.md5() 47 | with open(binfile, 'rb') as fid: 48 | while True: 49 | data = fid.read(BUF_SIZE) 50 | if not data: 51 | break 52 | md5.update(data) 53 | return md5.hexdigest() 54 | 55 | 56 | def RunShellCmd(cmd, print_cmd=False, print_output=False, synch=True, check_code=True, get_stat=True, get_output=True): 57 | """Function to run a Shell command through python, return code and """ 58 | # Overwrites any previous definition (when used within other programs) 59 | global debug, debug_log 60 | # Running command synchronously or asynchronously? 61 | if synch: 62 | if debug: 63 | print_cmd = True 64 | print_output = True 65 | # Printing command to be run if requested 66 | info = "RunShellCmd running: \'%s\'" % cmd 67 | if isinstance(debug_log, Logger): 68 | debug_log.addToLog(info) 69 | if print_cmd: 70 | print(info) 71 | # Starting the process (Using PIPE to streaming output) 72 | proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, 73 | stderr=subprocess.STDOUT, shell=True, universal_newlines=True) 74 | # Creating Stdout to save command output 75 | stdout = [] 76 | # Streaming the stdout to screen if requested 77 | while True: 78 | line = proc.stdout.readline() 79 | if line == '' and proc.poll() is not None: 80 | break 81 | else: 82 | stdout.append(line) 83 | line = line.rstrip() 84 | if line != '': 85 | # Print to debug file? 86 | if isinstance(debug_log, Logger): 87 | debug_log.addToLog(line) 88 | # Print to screen? 89 | if print_output: 90 | print(line) 91 | sys.stdout.flush() 92 | proc.stdout.flush() 93 | stdout = ''.join(stdout) 94 | else: 95 | # Running process asynchronously (Avoiding PIPE and removing standard output) 96 | global DEVNULL 97 | proc = subprocess.Popen([cmd], stdout=DEVNULL, shell=True, universal_newlines=True) 98 | return proc, "Running command asynchronously, returning process" 99 | # Command has finished, checking error code and returning requested variables 100 | err_code = proc.poll() 101 | return_var = [] 102 | # Returning error code or status 103 | if get_stat: 104 | return_var.append(err_code) 105 | # Returning output 106 | if get_output: 107 | return_var.append(stdout) 108 | # Checking error code 109 | if check_code and err_code != 0: 110 | # Writing error code to debug file if any 111 | info = "ERROR! Command failed: %s; Error code: %s" % (cmd, err_code) 112 | if isinstance(debug_log, Logger): 113 | debug_log.addToLog(info) 114 | raise SystemError("ERROR! Command failed: %s; Error code: %s; Output: %s" % (cmd, err_code, stdout)) 115 | # Returning 116 | return return_var 117 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/acoustic/acoustic_example.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | 4 | from devito.logger import info 5 | from devito import Constant, Function, smooth, norm 6 | from . import AcousticWaveSolver 7 | from .. import demo_model, setup_geometry, seismic_args 8 | 9 | 10 | def acoustic_setup(shape=(50, 50, 50), spacing=(15.0, 15.0, 15.0), 11 | tn=500., kernel='OT2', space_order=4, nbl=10, 12 | preset='layers-isotropic', fs=False, **kwargs): 13 | model = demo_model(preset, space_order=space_order, shape=shape, nbl=nbl, 14 | dtype=kwargs.pop('dtype', np.float32), spacing=spacing, 15 | fs=fs, **kwargs) 16 | 17 | # Source and receiver geometries 18 | geometry = setup_geometry(model, tn) 19 | 20 | # Create solver object to provide relevant operators 21 | solver = AcousticWaveSolver(model, geometry, kernel=kernel, 22 | space_order=space_order, **kwargs) 23 | return solver 24 | 25 | 26 | def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0, 27 | space_order=4, kernel='OT2', nbl=40, full_run=False, fs=False, 28 | autotune=False, preset='layers-isotropic', checkpointing=False, **kwargs): 29 | 30 | solver = acoustic_setup(shape=shape, spacing=spacing, nbl=nbl, tn=tn, 31 | space_order=space_order, kernel=kernel, fs=fs, 32 | preset=preset, **kwargs) 33 | 34 | info("Applying Forward") 35 | # Whether or not we save the whole time history. We only need the full wavefield 36 | # with 'save=True' if we compute the gradient without checkpointing, if we use 37 | # checkpointing, PyRevolve will take care of the time history 38 | save = full_run and not checkpointing 39 | # Define receiver geometry (spread across x, just below surface) 40 | rec, u, summary = solver.forward(save=save, autotune=autotune) 41 | 42 | if preset == 'constant': 43 | # With a new m as Constant 44 | v0 = Constant(name="v", value=2.0, dtype=np.float32) 45 | solver.forward(save=save, vp=v0) 46 | # With a new vp as a scalar value 47 | solver.forward(save=save, vp=2.0) 48 | 49 | if not full_run: 50 | return summary.gflopss, summary.oi, summary.timings, [rec, u.data] 51 | 52 | # Smooth velocity 53 | initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order) 54 | smooth(initial_vp, solver.model.vp) 55 | dm = np.float32(initial_vp.data ** (-2) - solver.model.vp.data ** (-2)) 56 | 57 | info("Applying Adjoint") 58 | solver.adjoint(rec, autotune=autotune) 59 | info("Applying Born") 60 | solver.jacobian(dm, autotune=autotune) 61 | info("Applying Gradient") 62 | solver.jacobian_adjoint(rec, u, autotune=autotune, checkpointing=checkpointing) 63 | return summary.gflopss, summary.oi, summary.timings, [rec, u.data] 64 | 65 | 66 | @pytest.mark.parametrize('shape', [(101,), (51, 51), (16, 16, 16)]) 67 | @pytest.mark.parametrize('k', ['OT2', 'OT4']) 68 | def test_isoacoustic_stability(shape, k): 69 | spacing = tuple([20]*len(shape)) 70 | _, _, _, [rec, _] = run(shape=shape, spacing=spacing, tn=20000.0, nbl=0, kernel=k) 71 | assert np.isfinite(norm(rec)) 72 | 73 | 74 | @pytest.mark.parametrize('fs, normrec, dtype', [(True, 369.955, np.float32), 75 | (False, 459.1678, np.float64)]) 76 | def test_isoacoustic(fs, normrec, dtype): 77 | _, _, _, [rec, _] = run(fs=fs, dtype=dtype) 78 | assert np.isclose(norm(rec), normrec, rtol=1e-3, atol=0) 79 | 80 | 81 | if __name__ == "__main__": 82 | description = ("Example script for a set of acoustic operators.") 83 | parser = seismic_args(description) 84 | parser.add_argument('--fs', dest='fs', default=False, action='store_true', 85 | help="Whether or not to use a freesurface") 86 | parser.add_argument("-k", dest="kernel", default='OT2', 87 | choices=['OT2', 'OT4'], 88 | help="Choice of finite-difference kernel") 89 | args = parser.parse_args() 90 | 91 | # 3D preset parameters 92 | ndim = args.ndim 93 | shape = args.shape[:args.ndim] 94 | spacing = tuple(ndim * [15.0]) 95 | tn = args.tn if args.tn > 0 else (750. if ndim < 3 else 1250.) 96 | 97 | preset = 'constant-isotropic' if args.constant else 'layers-isotropic' 98 | run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn, fs=args.fs, 99 | space_order=args.space_order, preset=preset, kernel=args.kernel, 100 | autotune=args.autotune, opt=args.opt, full_run=args.full, 101 | checkpointing=args.checkpointing, dtype=args.dtype) 102 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | try: 3 | import matplotlib as mpl 4 | import matplotlib.pyplot as plt 5 | from matplotlib import cm 6 | from mpl_toolkits.axes_grid1 import make_axes_locatable 7 | 8 | mpl.rc('font', size=16) 9 | mpl.rc('figure', figsize=(8, 6)) 10 | except: 11 | plt = None 12 | cm = None 13 | 14 | 15 | def plot_perturbation(model, model1, colorbar=True): 16 | """ 17 | Plot a two-dimensional velocity perturbation from two devitoseismic `Model` 18 | objects. 19 | 20 | Parameters 21 | ---------- 22 | model : Model 23 | The first velocity model. 24 | model1 : Model 25 | The second velocity model. 26 | colorbar : bool 27 | Option to plot the colorbar. 28 | """ 29 | domain_size = 1.e-3 * np.array(model.domain_size) 30 | extent = [model.origin[0], model.origin[0] + domain_size[0], 31 | model.origin[1] + domain_size[1], model.origin[1]] 32 | dv = np.transpose(model.vp.csg_nonlinear) - np.transpose(model1.vp.csg_nonlinear) 33 | 34 | plot = plt.imshow(dv, animated=True, cmap=cm.jet, 35 | vmin=min(dv.reshape(-1)), vmax=max(dv.reshape(-1)), 36 | extent=extent) 37 | plt.xlabel('X position (km)') 38 | plt.ylabel('Depth (km)') 39 | 40 | # Create aligned colorbar on the right 41 | if colorbar: 42 | ax = plt.gca() 43 | divider = make_axes_locatable(ax) 44 | cax = divider.append_axes("right", size="5%", pad=0.05) 45 | cbar = plt.colorbar(plot, cax=cax) 46 | cbar.set_label('Velocity perturbation (km/s)') 47 | plt.show() 48 | 49 | 50 | def plot_velocity(model, source=None, receiver=None, colorbar=True, cmap="jet"): 51 | """ 52 | Plot a two-dimensional velocity field from a devitoseismic `Model` 53 | object. Optionally also includes point markers for sources and receivers. 54 | 55 | Parameters 56 | ---------- 57 | model : Model 58 | Object that holds the velocity model. 59 | source : array_like or float 60 | Coordinates of the source point. 61 | receiver : array_like or float 62 | Coordinates of the receiver points. 63 | colorbar : bool 64 | Option to plot the colorbar. 65 | """ 66 | domain_size = 1.e-3 * np.array(model.domain_size) 67 | extent = [model.origin[0], model.origin[0] + domain_size[0], 68 | model.origin[1] + domain_size[1], model.origin[1]] 69 | 70 | slices = tuple(slice(model.nbl, -model.nbl) for _ in range(2)) 71 | if getattr(model, 'vp', None) is not None: 72 | field = model.vp.data[slices] 73 | else: 74 | field = model.lam.data[slices] 75 | plot = plt.imshow(np.transpose(field), animated=True, cmap=cmap, 76 | vmin=np.min(field), vmax=np.max(field), 77 | extent=extent) 78 | plt.xlabel('X position (km)') 79 | plt.ylabel('Depth (km)') 80 | 81 | # Plot source points, if provided 82 | if receiver is not None: 83 | plt.scatter(1e-3*receiver[:, 0], 1e-3*receiver[:, 1], 84 | s=25, c='green', marker='D') 85 | 86 | # Plot receiver points, if provided 87 | if source is not None: 88 | plt.scatter(1e-3*source[:, 0], 1e-3*source[:, 1], 89 | s=25, c='red', marker='o') 90 | 91 | # Ensure axis limits 92 | plt.xlim(model.origin[0], model.origin[0] + domain_size[0]) 93 | plt.ylim(model.origin[1] + domain_size[1], model.origin[1]) 94 | 95 | # Create aligned colorbar on the right 96 | if colorbar: 97 | ax = plt.gca() 98 | divider = make_axes_locatable(ax) 99 | cax = divider.append_axes("right", size="5%", pad=0.05) 100 | cbar = plt.colorbar(plot, cax=cax) 101 | cbar.set_label('Velocity (km/s)') 102 | plt.show() 103 | 104 | 105 | def plot_shotrecord(rec, model, t0, tn, colorbar=True): 106 | """ 107 | Plot a shot record (receiver values over time). 108 | 109 | Parameters 110 | ---------- 111 | rec : 112 | Receiver data with shape (time, points). 113 | model : Model 114 | object that holds the velocity model. 115 | t0 : int 116 | Start of time dimension to plot. 117 | tn : int 118 | End of time dimension to plot. 119 | """ 120 | scale = np.max(rec) / 10. 121 | extent = [model.origin[0], model.origin[0] + 1e-3*model.domain_size[0], 122 | 1e-3*tn, t0] 123 | 124 | plot = plt.imshow(rec, vmin=-scale, vmax=scale, cmap=cm.gray, extent=extent) 125 | plt.xlabel('X position (km)') 126 | plt.ylabel('Time (s)') 127 | 128 | # Create aligned colorbar on the right 129 | if colorbar: 130 | ax = plt.gca() 131 | divider = make_axes_locatable(ax) 132 | cax = divider.append_axes("right", size="5%", pad=0.05) 133 | plt.colorbar(plot, cax=cax) 134 | plt.show() 135 | 136 | 137 | def plot_image(data, vmin=None, vmax=None, colorbar=True, cmap="gray"): 138 | """ 139 | Plot image data, such as RTM images or FWI gradients. 140 | 141 | Parameters 142 | ---------- 143 | data : ndarray 144 | Image data to plot. 145 | cmap : str 146 | Choice of colormap. Defaults to gray scale for images as a 147 | devitoseismic convention. 148 | """ 149 | plot = plt.imshow(np.transpose(data), 150 | vmin=vmin or 0.9 * np.min(data), 151 | vmax=vmax or 1.1 * np.max(data), 152 | cmap=cmap) 153 | 154 | # Create aligned colorbar on the right 155 | if colorbar: 156 | ax = plt.gca() 157 | divider = make_axes_locatable(ax) 158 | cax = divider.append_axes("right", size="5%", pad=0.05) 159 | plt.colorbar(plot, cax=cax) 160 | plt.show() 161 | -------------------------------------------------------------------------------- /tutorials/born_devito.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import occamypy as o 3 | from typing import Tuple, List 4 | from devitoseismic import AcquisitionGeometry, demo_model, SeismicModel 5 | from devitoseismic.acoustic import AcousticWaveSolver 6 | import devito 7 | 8 | devito.configuration['log-level'] = 'ERROR' 9 | 10 | 11 | def create_models(args: dict) -> Tuple[SeismicModel, SeismicModel, SeismicModel]: 12 | hard = demo_model('layers-isotropic', origin=(0., 0.), 13 | shape=args["shape"], spacing=args["spacing"], 14 | nbl=args["nbl"], grid=None, nlayers=2) 15 | smooth = demo_model('layers-isotropic', origin=(0., 0.), 16 | shape=args["shape"], spacing=args["spacing"], 17 | nbl=args["nbl"], grid=hard.grid, nlayers=2) 18 | 19 | devito.gaussian_smooth(smooth.vp, sigma=args["filter_sigma"]) 20 | 21 | water = demo_model('layers-isotropic', origin=(0., 0.), 22 | shape=args["shape"], spacing=args["spacing"], 23 | nbl=args["nbl"], grid=hard.grid, nlayers=1) 24 | 25 | return hard, smooth, water 26 | 27 | 28 | def build_src_coordinates(x: float, z: float) -> np.ndarray: 29 | src = np.empty((1, 2), dtype=np.float32) 30 | src[0, :] = x 31 | src[0, -1] = z 32 | return src 33 | 34 | 35 | def build_rec_coordinates(model: SeismicModel, args: dict) -> np.ndarray: 36 | """Receivers equispaced on the whole domain""" 37 | rec = np.empty((args["nreceivers"], 2)) 38 | rec[:, 0] = np.linspace(0, model.domain_size[0], num=args["nreceivers"]) 39 | rec[:, 1] = args["rec_depth"] 40 | 41 | return rec 42 | 43 | 44 | def direct_arrival_mask(data: o.Vector, rec_pos: np.ndarray, src_pos: np.ndarray, 45 | vel_sep: float = 1500., offset: float = 0.) -> o.Vector: 46 | dt = data.ax_info[0].d 47 | 48 | direct = np.sqrt(np.sum((src_pos - rec_pos) ** 2, axis=1)) / vel_sep 49 | direct += offset 50 | 51 | mask = data.clone().zero() 52 | 53 | iwin = np.round(direct / dt).astype(int) 54 | for i in range(rec_pos.shape[0]): 55 | mask[iwin[i]:, i] = 1. 56 | 57 | return mask 58 | 59 | 60 | def depth_compensation_mask(model: o.Vector, z_pow: float = 1., depth_axis: int = 1) -> o.Vector: 61 | mask = np.ones(model.shape) 62 | 63 | if z_pow != 0.: 64 | for z in range(mask.shape[depth_axis]): 65 | 66 | if depth_axis != 0: 67 | mask = np.swapaxes(mask, 0, depth_axis) 68 | 69 | mask[z] *= pow(z, z_pow) 70 | 71 | if depth_axis != 0: 72 | mask = np.swapaxes(mask, 0, depth_axis) 73 | 74 | mask = model.__class__(mask, ax_info=model.ax_info) 75 | 76 | return mask 77 | 78 | 79 | def _propagate_shot(model: SeismicModel, rec_pos: np.ndarray, src_pos: np.ndarray, param: dict) -> o.VectorNumpy: 80 | geometry = AcquisitionGeometry(model, rec_pos, src_pos, **param) 81 | solver = AcousticWaveSolver(model, geometry, **param) 82 | 83 | devito.clear_cache() 84 | 85 | # propagate (source -> receiver data) 86 | data = o.VectorNumpy(solver.forward()[0].data.__array__()) 87 | 88 | data.ax_info = [o.AxInfo(geometry.nt, geometry.t0, geometry.dt / 1000, "time [s]"), 89 | o.AxInfo(geometry.nrec, float(rec_pos[0][0]), float(rec_pos[1][0] - rec_pos[0][0]), "rec pos x [m]")] 90 | 91 | devito.clear_cache() 92 | return data 93 | 94 | 95 | def propagate_shots(model: SeismicModel, rec_pos: np.ndarray, src_pos: List[np.ndarray], param: dict): 96 | if len(src_pos) == 1: 97 | return _propagate_shot(model=model, rec_pos=rec_pos, src_pos=src_pos[0], param=param) 98 | else: 99 | return o.superVector([_propagate_shot(model=model, rec_pos=rec_pos, src_pos=s, param=param) for s in src_pos]) 100 | 101 | 102 | class BornSingleSource(o.Operator): 103 | 104 | def __init__(self, velocity: SeismicModel, src_pos: np.ndarray, rec_pos: np.ndarray, args: dict): 105 | 106 | # store params 107 | self.src_pos = src_pos 108 | self.rec_pos = rec_pos 109 | self.nbl = args["nbl"] 110 | 111 | # build geometry and acoustic solver 112 | self.geometry = AcquisitionGeometry(velocity, rec_pos, src_pos, **args) 113 | self.solver = AcousticWaveSolver(velocity, self.geometry, **args) 114 | 115 | # allocate vectors 116 | self.velocity = o.VectorNumpy(velocity.vp.data.__array__()) 117 | self.velocity.ax_info = [ 118 | o.AxInfo(velocity.vp.shape[0], velocity.origin[0] - self.nbl * velocity.spacing[0], velocity.spacing[0], 119 | "x [m]"), 120 | o.AxInfo(velocity.vp.shape[1], velocity.origin[1] - self.nbl * velocity.spacing[1], velocity.spacing[1], 121 | "z [m]")] 122 | 123 | csg = o.VectorNumpy((self.geometry.nt, self.geometry.nrec)) 124 | csg.ax_info = [o.AxInfo(self.geometry.nt, self.geometry.t0, self.geometry.dt / 1000, "time [s]"), 125 | o.AxInfo(self.geometry.nrec, float(rec_pos[0][0]), float(rec_pos[1][0] - rec_pos[0][0]), 126 | "rec pos x [m]")] 127 | 128 | super(BornSingleSource, self).__init__(self.velocity, csg) 129 | self.name = "DeviBorn" 130 | # store source wavefield 131 | self.src_wfld = self.solver.forward(save=True)[1] 132 | 133 | def forward(self, add, model, data): 134 | """Modeling function: image -> residual data""" 135 | self.checkDomainRange(model, data) 136 | if not add: 137 | data.zero() 138 | 139 | recs = self.solver.jacobian(dmin=model[:])[0] 140 | data[:] += recs.data.__array__() 141 | 142 | return 143 | 144 | def adjoint(self, add, model, data): 145 | """Adjoint function: data -> image""" 146 | self.checkDomainRange(model, data) 147 | if not add: 148 | model.zero() 149 | 150 | recs = self.geometry.rec.copy() 151 | recs.data[:] = data[:] 152 | 153 | img = self.solver.gradient(rec=recs, u=self.src_wfld)[0] 154 | model[:] += img.data.__array__() 155 | 156 | return 157 | -------------------------------------------------------------------------------- /occamypy/numpy/operator/signal.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | 3 | import numpy as np 4 | from scipy.ndimage import gaussian_filter 5 | from scipy.signal import convolve, correlate 6 | 7 | from occamypy.operator.base import Operator, Dstack 8 | from occamypy.numpy.vector import VectorNumpy 9 | from occamypy.vector.base import superVector 10 | 11 | 12 | class GaussianFilter(Operator): 13 | """Gaussian smoothing operator using scipy smoothing""" 14 | 15 | def __init__(self, domain, sigma): 16 | """ 17 | GaussianFilter (numpy) constructor 18 | 19 | Args: 20 | domain: domain vector 21 | sigma: standard deviation along the domain directions 22 | """ 23 | self.sigma = sigma 24 | self.scaling = np.sqrt(np.prod(np.array(self.sigma) / np.pi)) # in order to have the max amplitude 1 25 | super(GaussianFilter, self).__init__(domain, domain, name="GaussFilt") 26 | 27 | def forward(self, add, model, data): 28 | self.checkDomainRange(model, data) 29 | if not add: 30 | data.zero() 31 | # Getting Ndarrays 32 | model_arr = model.getNdArray() 33 | data_arr = data.getNdArray() 34 | data_arr[:] += self.scaling * gaussian_filter(model_arr, sigma=self.sigma) 35 | return 36 | 37 | def adjoint(self, add, model, data): 38 | self.forward(add, data, model) 39 | return 40 | 41 | 42 | class ConvND(Operator): 43 | """ND convolution square operator in the domain space""" 44 | 45 | def __init__(self, domain: VectorNumpy, kernel: Union[VectorNumpy, np.ndarray], method='auto'): 46 | """ 47 | ConvND (numpy) constructor 48 | 49 | Args: 50 | domain: domain vector 51 | kernel: kernel vector 52 | method: how to compute the convolution [auto, direct, fft] 53 | """ 54 | if isinstance(kernel, VectorNumpy): 55 | self.kernel = kernel.clone().getNdArray() 56 | elif isinstance(kernel, np.ndarray): 57 | self.kernel = kernel.copy() 58 | else: 59 | raise ValueError("kernel has to be either a vector or a numpy.ndarray") 60 | 61 | # Padding array to avoid edge effects 62 | pad_width = [] 63 | for len_filt in self.kernel.shape: 64 | half_len = int(len_filt / 2) 65 | if np.mod(len_filt, 2): 66 | padding = (half_len, half_len) 67 | else: 68 | padding = (half_len, half_len - 1) 69 | pad_width.append(padding) 70 | self.kernel = np.pad(self.kernel, pad_width, mode='constant') 71 | 72 | if len(domain.shape) != len(self.kernel.shape): 73 | raise ValueError("Domain and kernel number of dimensions mismatch") 74 | 75 | if method not in ["auto", "direct", "fft"]: 76 | raise ValueError("method has to be auto, direct or fft") 77 | self.method = method 78 | 79 | super(ConvND, self).__init__(domain=domain, range=domain, name="Convolve") 80 | 81 | def forward(self, add, model, data): 82 | self.checkDomainRange(model, data) 83 | if not add: 84 | data.zero() 85 | modelNd = model.getNdArray() 86 | dataNd = data.getNdArray()[:] 87 | dataNd += convolve(modelNd, self.kernel, mode='same', method=self.method) 88 | return 89 | 90 | def adjoint(self, add, model, data): 91 | self.checkDomainRange(model, data) 92 | if not add: 93 | model.zero() 94 | modelNd = model.getNdArray() 95 | dataNd = data.getNdArray()[:] 96 | modelNd += correlate(dataNd, self.kernel, mode='same', method=self.method) 97 | return 98 | 99 | 100 | def Padding(domain, pad, mode: str = "constant"): 101 | """ 102 | Padding operator 103 | 104 | Notes: 105 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 106 | pad=((2,2), (3,3)) 107 | 108 | Args: 109 | domain: domain vector 110 | pad: number of samples to be added at each end of the dimension, for each dimension 111 | mode: padding mode (see https://numpy.org/doc/stable/reference/generated/numpy.pad.html) 112 | """ 113 | if isinstance(domain, VectorNumpy): 114 | return _Padding(domain, pad, mode) 115 | elif isinstance(domain, superVector): 116 | return Dstack([_Padding(v, pad, mode) for v in domain.vecs]) 117 | else: 118 | raise ValueError("ERROR! Provided domain has to be either vector or superVector") 119 | 120 | 121 | def ZeroPad(domain, pad): 122 | """ 123 | Zero-Padding operator 124 | 125 | Notes: 126 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 127 | pad=((2,2), (3,3)) 128 | 129 | Args: 130 | domain: domain vector 131 | pad: number of samples to be added at each end of the dimension, for each dimension 132 | """ 133 | return Padding(domain, pad, mode="constant") 134 | 135 | 136 | def _pad_VectorNumpy(vec, pad): 137 | if not isinstance(vec, VectorNumpy): 138 | raise ValueError("ERROR! Provided vector must be a VectorNumpy") 139 | if len(vec.shape) != len(pad): 140 | raise ValueError("Dimensions of vector and padding mismatch!") 141 | 142 | vec_new_shape = tuple(np.asarray(vec.shape) + [sum(pad[_]) for _ in range(len(pad))]) 143 | return VectorNumpy(np.empty(vec_new_shape, dtype=vec.getNdArray().dtype)) 144 | 145 | 146 | class _Padding(Operator): 147 | 148 | def __init__(self, domain: VectorNumpy, pad, mode: str = "constant"): 149 | 150 | self.dims = domain.shape 151 | pad = [(pad, pad)] * len(self.dims) if isinstance(pad, int) else list(pad) 152 | if (np.array(pad) < 0).any(): 153 | raise ValueError('Padding must be positive or zero') 154 | self.pad = pad 155 | self.mode = mode 156 | super(_Padding, self).__init__(domain, _pad_VectorNumpy(domain, self.pad), name="Padding") 157 | 158 | def forward(self, add, model, data): 159 | """Pad the domain""" 160 | self.checkDomainRange(model, data) 161 | if add: 162 | temp = data.clone() 163 | y = np.pad(model.arr, self.pad, mode=self.mode) 164 | data.arr = y 165 | if add: 166 | data.scaleAdd(temp, 1., 1.) 167 | return 168 | 169 | def adjoint(self, add, model, data): 170 | """Extract original subsequence""" 171 | self.checkDomainRange(model, data) 172 | if add: 173 | temp = model.clone() 174 | x = data.clone().arr 175 | for ax, pad in enumerate(self.pad): 176 | x = np.take(x, pad[0] + np.arange(self.dims[ax]), axis=ax) 177 | model.arr = x 178 | if add: 179 | model.scaleAdd(temp, 1., 1.) 180 | return 181 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/acoustic/acoustic_time_update_nb.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# Time update for *Cerjan* damped constant density acoustics with forward approximation for $\\partial_t$\n", 8 | "\n", 9 | "We show a derivation for the time update expression used for the constant density acoustic solver. You can compare the end result of this derivation in the last equation line below with lines 58-59 in the file ```examples/seismic/acoustic/operators.py```." 10 | ] 11 | }, 12 | { 13 | "cell_type": "markdown", 14 | "metadata": {}, 15 | "source": [ 16 | "## Table of symbols\n", 17 | "\n", 18 | "| Symbol         | Description | Dimensionality | \n", 19 | "| :--- | :--- | :--- |\n", 20 | "| $\\delta t$ | Temporal sampling interval | constant |\n", 21 | "| $m(x,y,z)$ | slowness squared | function of space |\n", 22 | "| $\\eta(x,y,z)$ | Damping coefficient | function of space |\n", 23 | "| $u(t,x,y,z)$ | Pressure wavefield | function of time and space |\n", 24 | "| $q(t,x,y,z)$ | Source term | function of time, localized in space |\n", 25 | "| $\\partial_{t}$ | first derivative wrt $t$ | time |\n", 26 | "| $\\partial_{tt}$ | second derivative wrt $t$ | time |\n", 27 | "| $\\nabla^2$ | Laplacian operator | space |" 28 | ] 29 | }, 30 | { 31 | "cell_type": "markdown", 32 | "metadata": {}, 33 | "source": [ 34 | "## A word about notation \n", 35 | "\n", 36 | "For clarity in the following derivation we will drop the space notatation for certain variables:\n", 37 | "- $m(x,y,z) \\rightarrow m$\n", 38 | "- $\\eta(x,y,z) \\rightarrow \\eta$\n", 39 | "- $u(t,x,y,z) \\rightarrow u(t)$\n", 40 | "- $q(t,x,y,z) \\rightarrow q(t)$" 41 | ] 42 | }, 43 | { 44 | "cell_type": "markdown", 45 | "metadata": {}, 46 | "source": [ 47 | "## The time update equation\n", 48 | "\n", 49 | "To implement the Devito modeling operator we define the equation used to update the pressure wavefield as a function of time. What follows is a bit of algebra using the wave equation and finite difference approximations to time derivatives to express the pressure wavefield forward in time $u(t+\\delta t)$ as a function of the current $u(t)$ and previous $u(t-\\delta t)$ pressure wavefields.\n", 50 | "\n", 51 | "#### 1. First order numerical derivative (forward):\n", 52 | "The first order accurate forward approximation to the first time derivative involves two wavefields: $u(t-\\delta t)$, and $u(t)$. We can use this expression as is. \n", 53 | "\n", 54 | "$$\n", 55 | "\\partial_{t}\\ u(t) = \\frac{u(t+\\delta t) - u(t)}{\\delta t}\n", 56 | "$$\n", 57 | "
\n", 58 | "\n", 59 | "#### 2. Second order numerical derivative:\n", 60 | "The second order accurate centered approximation to the second time derivative involves three wavefields: $u(t-\\delta t)$, $u(t)$, and $u(t+\\delta t)$. \n", 61 | "\n", 62 | "$$\n", 63 | "\\partial_{tt}\\ u(t) = \\frac{u(t+\\delta t) - 2\\ u(t) + u(t-\\delta t)}{\\delta t^2}\n", 64 | "$$\n", 65 | "
\n", 66 | "\n", 67 | "#### 3. Second order time update:\n", 68 | "In order to advance our finite difference solution in time, we solve for $u(t+\\delta t)$.\n", 69 | "\n", 70 | "$$\n", 71 | "u(t+\\delta t) = \\delta t^2\\ \\partial_{tt}\\ u(t) + 2\\ u(t) - u(t-\\delta t)\n", 72 | "$$\n", 73 | "
\n", 74 | "\n", 75 | "#### 4. Damped wave equation:\n", 76 | "\n", 77 | "Our *Cerjan* (reference below) damped wave equation, which we solve for $\\partial_{tt}$:\n", 78 | "\n", 79 | "$$\n", 80 | "\\begin{aligned}\n", 81 | "m\\ \\partial_{tt}\\ u(t) + \\eta\\ \\partial_{t}\\ u(t) &= \\nabla^2 u(t) + q(t) \\\\[10pt]\n", 82 | "\\partial_{tt}\\ u(t) &=\n", 83 | " \\frac{1}{m} \\Bigr[ \\nabla^2 u(t) + q(t) - \\eta\\ \\partial_{t}\\ u(t) \\Bigr]\n", 84 | "\\end{aligned}\n", 85 | "$$\n", 86 | "\n", 87 | "#### 5. Time update:\n", 88 | "Next we plug in the expression for $\\partial_{tt}\\ u$ (from the wave equation) and $\\partial_{t}\\ u$ (from the numerical derivative) into the the time update expression for $u(t+\\delta t)$ from step 3.\n", 89 | "\n", 90 | "$$\n", 91 | "\\begin{aligned}\n", 92 | "u(t+\\delta t) &=\n", 93 | " \\frac{\\delta t^2}{m} \\Bigr[ \\nabla^2 u(t) + q(t)\n", 94 | " - \\frac{\\eta}{\\delta t} \\bigr\\{ u(t+\\delta t) - u(t) \\bigr\\} \\Bigr]\n", 95 | " + 2\\ u(t) - u(t-\\delta t)\n", 96 | "\\end{aligned}\n", 97 | "$$\n", 98 | "\n", 99 | "#### 6. Simplify:\n", 100 | "\n", 101 | "Finally we simplify this expression to the form used in the Devito ```Operator```.\n", 102 | "\n", 103 | "$$\n", 104 | "\\begin{aligned}\n", 105 | "\\left(1 + \\frac{\\delta t\\ \\eta}{m}\\right) u(t+\\delta t) &= \n", 106 | " \\frac{\\delta t^2}{m} \\Bigr\\{ \\nabla^2 u(t) + q(t) \\Bigr\\}\n", 107 | " + \\frac{\\delta t\\ \\eta}{m}\\ u(t) + 2\\ u(t) - u(t-\\delta t) \\\\[15pt]\n", 108 | "u(t+\\delta t) &=\n", 109 | " \\left( \\frac{1}{m+\\delta t\\ \\eta} \\right) \\Bigr[\n", 110 | " \\delta t^2 \\Bigr\\{ \\nabla^2 u(t) + q(t) \\Bigr\\}\n", 111 | " + \\delta t\\ \\eta\\ u(t) + m\\ \\left[2\\ u(t) - u(t-\\delta t) \\right]\n", 112 | " \\Bigr]\n", 113 | "\\end{aligned}\n", 114 | "$$\n", 115 | "\n", 116 | "\n", 117 | "#### 7. Compare:\n", 118 | "\n", 119 | "Please compare the last equation above with [lines 58-59 in examples/seismic/acoustic/operators.py](https://github.com/devitocodes/devito/blob/master/examples/seismic/acoustic/operators.py#L58-L59)\n", 120 | "\n", 121 | "```\n", 122 | "eq_time = ((lap + q) * s**2 + s * damp * field +\n", 123 | " m * (2 * field - prev))/(s * damp + m)```" 124 | ] 125 | }, 126 | { 127 | "cell_type": "markdown", 128 | "metadata": {}, 129 | "source": [ 130 | "## References\n", 131 | "\n", 132 | "- **A nonreflecting boundary condition for discrete acoustic and elastic wave equations** (1985)\n", 133 | "
Charles Cerjan, Dan Kosloft. Ronnie Kosloff, and Moshe Resheq\n", 134 | "
Geophysics, Vol. 50, No. 4\n", 135 | "
https://library.seg.org/doi/pdfplus/10.1190/segam2016-13878451.1" 136 | ] 137 | }, 138 | { 139 | "cell_type": "code", 140 | "execution_count": null, 141 | "metadata": {}, 142 | "outputs": [], 143 | "source": [] 144 | } 145 | ], 146 | "metadata": { 147 | "kernelspec": { 148 | "display_name": "Python 3", 149 | "language": "python", 150 | "name": "python3" 151 | }, 152 | "language_info": { 153 | "codemirror_mode": { 154 | "name": "ipython", 155 | "version": 3 156 | }, 157 | "file_extension": ".py", 158 | "mimetype": "text/x-python", 159 | "name": "python", 160 | "nbconvert_exporter": "python", 161 | "pygments_lexer": "ipython3", 162 | "version": "3.8.2" 163 | } 164 | }, 165 | "nbformat": 4, 166 | "nbformat_minor": 4 167 | } 168 | -------------------------------------------------------------------------------- /occamypy/cupy/operator/signal.py: -------------------------------------------------------------------------------- 1 | import cupy as cp 2 | import numpy as np 3 | from cupyx.scipy.ndimage import gaussian_filter 4 | 5 | try: 6 | from cusignal.convolution import convolve, correlate 7 | except ModuleNotFoundError: 8 | raise ModuleNotFoundError("cuSIGNAL is not installed. Please install it") 9 | 10 | from occamypy.vector.base import Vector, superVector 11 | from occamypy.operator.base import Operator, Dstack 12 | from occamypy.cupy.vector import VectorCupy 13 | 14 | 15 | class GaussianFilter(Operator): 16 | """Gaussian smoothing operator using scipy smoothing""" 17 | 18 | def __init__(self, domain, sigma): 19 | """ 20 | GaussianFilter (cupy) constructor 21 | 22 | Args: 23 | domain: domain vector 24 | sigma: standard deviation along the domain directions 25 | """ 26 | self.sigma = sigma 27 | self.scaling = np.sqrt(np.prod(np.array(self.sigma) / cp.pi)) # in order to have the max amplitude 1 28 | 29 | super(GaussianFilter, self).__init__(domain=domain, range=domain, name="GaussFilt") 30 | 31 | def forward(self, add, model, data): 32 | self.checkDomainRange(model, data) 33 | if not add: 34 | data.zero() 35 | # Getting Ndarrays 36 | model_arr = model.getNdArray() 37 | data_arr = data.getNdArray() 38 | data_arr[:] += self.scaling * gaussian_filter(model_arr, sigma=self.sigma) 39 | return 40 | 41 | def adjoint(self, add, model, data): 42 | self.forward(add, data, model) 43 | return 44 | 45 | 46 | class ConvND(Operator): 47 | """ND convolution square operator in the domain space""" 48 | 49 | def __init__(self, domain, kernel, method='auto'): 50 | """ 51 | ConvND (cupy) constructor 52 | 53 | Args: 54 | domain: domain vector 55 | kernel: kernel vector 56 | method: how to compute the convolution [auto, direct, fft] 57 | """ 58 | if isinstance(kernel, Vector): 59 | self.kernel = kernel.clone().getNdArray() 60 | elif isinstance(kernel, cp.ndarray): 61 | self.kernel = kernel.copy() 62 | else: 63 | raise ValueError("kernel has to be either a vector or a cupy.ndarray") 64 | 65 | # Padding array to avoid edge effects 66 | pad_width = [] 67 | for len_filt in self.kernel.shape: 68 | half_len = int(len_filt / 2) 69 | if np.mod(len_filt, 2): 70 | padding = (half_len, half_len) 71 | else: 72 | padding = (half_len, half_len - 1) 73 | pad_width.append(padding) 74 | self.kernel = cp.padding(self.kernel, pad_width, mode='constant') 75 | 76 | if len(domain.shape()) != len(self.kernel.shape): 77 | raise ValueError("Domain and kernel number of dimensions mismatch") 78 | 79 | if method not in ["auto", "direct", "fft"]: 80 | raise ValueError("method has to be auto, direct or fft") 81 | self.method = method 82 | 83 | super(ConvND, self).__init__(domain=domain, range=domain, name="Convolve") 84 | 85 | def forward(self, add, model, data): 86 | self.checkDomainRange(model, data) 87 | if not add: 88 | data.zero() 89 | modelNd = model.getNdArray() 90 | dataNd = data.getNdArray()[:] 91 | dataNd += convolve(modelNd, self.kernel) 92 | return 93 | 94 | def adjoint(self, add, model, data): 95 | self.checkDomainRange(model, data) 96 | if not add: 97 | model.zero() 98 | modelNd = model.getNdArray() 99 | dataNd = data.getNdArray()[:] 100 | modelNd += correlate(dataNd, self.kernel) 101 | return 102 | 103 | 104 | def Padding(domain, pad, mode: str = "constant"): 105 | """ 106 | Padding operator 107 | 108 | Notes: 109 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 110 | pad=((2,2), (3,3)) 111 | 112 | Args: 113 | domain: domain vector 114 | pad: number of samples to be added at each end of the dimension, for each dimension 115 | mode: padding mode (see https://numpy.org/doc/stable/reference/generated/numpy.pad.html) 116 | """ 117 | if isinstance(domain, VectorCupy): 118 | return _Padding(domain, pad, mode) 119 | elif isinstance(domain, superVector): 120 | # TODO add the possibility to have different padding for each sub-vector 121 | return Dstack([_Padding(v, pad, mode) for v in domain.vecs]) 122 | else: 123 | raise ValueError("ERROR! Provided domain has to be either vector or superVector") 124 | 125 | 126 | def ZeroPad(domain, pad): 127 | """ 128 | Zero-Padding operator 129 | 130 | Notes: 131 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 132 | pad=((2,2), (3,3)) 133 | 134 | Args: 135 | domain: domain vector 136 | pad: number of samples to be added at each end of the dimension, for each dimension 137 | """ 138 | return Padding(domain, pad, mode="constant") 139 | 140 | 141 | def _pad_VectorCupy(vec, pad): 142 | if not isinstance(vec, VectorCupy): 143 | raise ValueError("ERROR! Provided vector has to be a VectorCupy") 144 | if len(vec.shape) != len(pad): 145 | raise ValueError("Dimensions of vector and padding mismatch!") 146 | 147 | vec_new_shape = tuple(cp.asarray(vec.shape) + [sum(pad[_]) for _ in range(len(pad))]) 148 | if isinstance(vec, VectorCupy): 149 | return VectorCupy(cp.empty(vec_new_shape, dtype=vec.getNdArray().dtype)) 150 | else: 151 | raise ValueError("ERROR! For now only vectorCupy is supported!") 152 | 153 | 154 | class _Padding(Operator): 155 | 156 | def __init__(self, domain, pad, mode: str = "constant"): 157 | 158 | if isinstance(domain, VectorCupy): 159 | self.dims = domain.shape 160 | pad = [(pad, pad)] * len(self.dims) if isinstance(pad, int) else list(pad) 161 | if (cp.array(pad) < 0).any(): 162 | raise ValueError('Padding must be positive or zero') 163 | self.pad = pad 164 | self.mode = mode 165 | super(_Padding, self).__init__(domain, _pad_VectorCupy(domain, self.pad), name="Padding") 166 | 167 | def forward(self, add, model, data): 168 | """Pad the domain""" 169 | self.checkDomainRange(model, data) 170 | if add: 171 | temp = data.clone() 172 | y = cp.padding(model.getNdArray(), self.pad, mode=self.mode) 173 | data.getNdArray()[:] = y 174 | if add: 175 | data.scaleAdd(temp, 1., 1.) 176 | return 177 | 178 | def adjoint(self, add, model, data): 179 | """Extract original subsequence""" 180 | self.checkDomainRange(model, data) 181 | if add: 182 | temp = model.clone() 183 | x = data.clone().getNdArray() 184 | for ax, pad in enumerate(self.pad): 185 | x = cp.take(x, pad[0] + cp.arange(self.dims[ax]), axis=ax) 186 | model.arr = x 187 | if add: 188 | model.scaleAdd(temp, 1., 1.) 189 | return 190 | -------------------------------------------------------------------------------- /occamypy/utils/sep.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | 4 | import numpy as np 5 | 6 | from occamypy.utils.os import RunShellCmd 7 | from occamypy.vector.axis_info import AxInfo 8 | 9 | # Assigning datapath 10 | HOME = os.environ["HOME"] 11 | datapath = None 12 | # Checking environment definition first 13 | if "DATAPATH" in os.environ: 14 | datapath = os.environ["DATAPATH"] 15 | # Checking local directory 16 | elif os.path.isfile('.datapath'): 17 | out = (RunShellCmd("cat .datapath | head -N 1", check_code=False, get_stat=False)[0]).rstrip() 18 | datapath = out.split("=")[1] 19 | # Checking whether the local host has a datapath 20 | else: 21 | if os.path.isfile(HOME + "/.datapath"): 22 | out = RunShellCmd("cat $HOME/.datapath | grep $HOST", check_code=False, get_stat=False)[0] 23 | if len(out) == 0: 24 | out = (RunShellCmd("cat $HOME/.datapath | head -N 1", check_code=False, get_stat=False)[0]).rstrip() 25 | datapath = out.split("=")[1] 26 | 27 | # Checking if datapath was found 28 | _already_warned = False 29 | if datapath is None: 30 | if os.path.isdir("/tmp/"): 31 | datapath = "/tmp/" 32 | if not _already_warned: 33 | print("WARNING! DATAPATH not found. The folder /tmp will be used to write binary files") 34 | _already_warned = True 35 | else: 36 | raise IOError("SEP datapath not found\n Set env variable DATAPATH to a folder to write binary files") 37 | 38 | 39 | def rm_file(filename): 40 | """File to remove header and binary files""" 41 | binfile = get_binary(filename) 42 | if os.path.isfile(filename): 43 | os.remove(filename) 44 | if os.path.isfile(binfile): 45 | os.remove(binfile) 46 | return 47 | 48 | 49 | def get_par(filename, par): 50 | """ Function to obtain a header parameter within the passed header file""" 51 | info = None 52 | # Checking if label is requested 53 | if "label" in par: 54 | reg_prog = re.compile("%s=(\'(.*?)\'|\"(.*?)\")" % par) 55 | else: 56 | reg_prog = re.compile("%s=([^\s]+)" % par) 57 | if not os.path.isfile(filename): 58 | raise OSError("ERROR! No %s file found!" % filename) 59 | for line in reversed(open(filename).readlines()): 60 | if info is None: 61 | find = reg_prog.search(line) 62 | if find: 63 | info = find.group(1) 64 | if info is None: 65 | raise IOError("%s parameter not found in file %s" % (filename, par)) 66 | # Removing extra characters from found parameter 67 | if info is not None: 68 | info = info.replace('"', '') 69 | info = info.replace('\'', '') 70 | return info 71 | 72 | 73 | def get_binary(filename): 74 | """ Function to obtain binary file associated with a given header file""" 75 | return get_par(filename, "in") 76 | 77 | 78 | def get_axes(filename): 79 | """Function returning all axis information related to a header file""" 80 | axes = [] 81 | # Currently handling maximum 7 axis 82 | for iaxis in range(7): 83 | # Obtaining number of elements within each axis 84 | try: 85 | axis_n = int(get_par(filename, par="n%s" % (iaxis + 1))) 86 | except IOError as exc: 87 | if iaxis == 0: 88 | print(exc.args) 89 | print("ERROR! First axis parameters must be found! Returning None") 90 | return None 91 | else: 92 | # Default value for an unset axis 93 | axis_n = 1 94 | # Obtaining origin of each axis 95 | try: 96 | axis_o = float(get_par(filename, par="o%s" % (iaxis + 1))) 97 | except IOError as exc: 98 | if iaxis == 0: 99 | print(exc.args) 100 | print("ERROR! First axis parameters must be found! Returning None") 101 | return None 102 | else: 103 | # Default value for an unset axis 104 | axis_o = 0.0 105 | # Obtaining sampling of each axis 106 | try: 107 | axis_d = float(get_par(filename, par="d%s" % (iaxis + 1))) 108 | except IOError as exc: 109 | if iaxis == 0: 110 | print(exc.args) 111 | print("ERROR! First axis parameters must be found! Returning None") 112 | return None 113 | else: 114 | # Default value for an unset axis 115 | axis_d = 0.0 116 | # Obtaining label of each axis 117 | try: 118 | axis_lab = get_par(filename, par="label%s" % (iaxis + 1)) 119 | except IOError as exc: 120 | # Default value for an unset axis 121 | axis_lab = "undefined" 122 | axes.append(AxInfo(int(axis_n), float(axis_o), float(axis_d), axis_lab)) 123 | return axes 124 | 125 | 126 | def get_num_axes(filename): 127 | """Function to obtain number of axes in a header file""" 128 | # Obtaining elements in each dimensions 129 | axis_info = get_axes(filename) 130 | axis_elements = [ax.n for ax in axis_info] 131 | index = [i for i, nelements in enumerate(axis_elements) if int(nelements) > 1] 132 | if index: 133 | n_axes = index[-1] + 1 134 | else: 135 | n_axes = 1 136 | return n_axes 137 | 138 | 139 | def read_file(filename, formatting='>f', mem_order="C"): 140 | """Function for reading header files""" 141 | axis_info = get_axes(filename) 142 | n_axis = get_num_axes(filename) 143 | shape = [ax.n for ax in axis_info] 144 | shape = shape[:n_axis] 145 | if mem_order == "C": 146 | shape = tuple(reversed(shape)) 147 | elif mem_order != "F": 148 | raise ValueError("ERROR! %s not an supported array order" % mem_order) 149 | fid = open(get_binary(filename), 'r+b') 150 | # Default formatting big-ending floating point number 151 | data = np.fromfile(fid, dtype=formatting) 152 | # Reshaping array and forcing memory continuity 153 | if mem_order == "C": 154 | data = np.ascontiguousarray(np.reshape(data, shape, order=mem_order)) 155 | else: 156 | data = np.asfortranarray(np.reshape(data, shape, order=mem_order)) 157 | fid.close() 158 | return [data, axis_info] 159 | 160 | 161 | def write_file(filename, data, axis_info=None, formatting='>f'): 162 | """Function for writing header files""" 163 | global datapath 164 | # write binary file 165 | binfile = datapath + filename.split('/')[-1] + '@' 166 | with open(binfile, 'w+b') as fid: 167 | # Default formatting big-ending floating point number 168 | if np.isfortran(data): # Forcing column-wise binary writing 169 | data.flatten('F').astype(formatting).tofile(fid) 170 | else: 171 | data.astype(formatting).tofile(fid) 172 | fid.close() 173 | 174 | # If axis_info is not provided all the present axis are set to d=1.0 o=0.0 label='undefined' 175 | if axis_info is None: 176 | naxis = data.shape 177 | if not np.isfortran(data): 178 | naxis = tuple(reversed(naxis)) # If C last axis is the "fastest" 179 | axis_info = [AxInfo(naxis[ii]) for ii in range(len(naxis))] 180 | 181 | # writing header/pointer file 182 | with open(filename, 'w') as fid: 183 | # Writing axis info 184 | for ii, ax_info in enumerate(axis_info): 185 | fid.write(ax_info.to_string(ii+1)) 186 | fid.write("in='%s'\n" % binfile) 187 | fid.write("data_format='xdr_float'\n") 188 | fid.write("esize=4\n") 189 | fid.close() 190 | return 191 | 192 | -------------------------------------------------------------------------------- /occamypy/torch/autograd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from occamypy.torch.vector import VectorTorch 4 | 5 | __all__ = [ 6 | "AutogradFunction", 7 | "VectorAD", 8 | ] 9 | 10 | 11 | class VectorAD(VectorTorch): 12 | """ 13 | VectorTorch child which allows tensors to be atteched to the graph (requires_grad=True) 14 | 15 | Notes: 16 | tensors are stored in C-contiguous memory 17 | """ 18 | 19 | def __init__(self, in_content, device: int = None, *args, **kwargs): 20 | """ 21 | VectorAD constructor 22 | 23 | Args: 24 | in_content: Vector, np.ndarray, torch.Tensor or tuple 25 | device: computation device (None for CPU, -1 for least used GPU) 26 | *args: list of arguments for Vector construction 27 | **kwargs: dict of arguments for Vector construction 28 | """ 29 | if isinstance(in_content, VectorTorch): 30 | super(VectorAD, self).__init__(in_content=in_content[:], *args, **kwargs) 31 | self.getNdArray().requires_grad = True 32 | else: 33 | super(VectorAD, self).__init__(in_content=in_content, device=device, *args, **kwargs) 34 | if not self.getNdArray().requires_grad: 35 | self.getNdArray().requires_grad = True 36 | 37 | @property 38 | def requires_grad(self): 39 | return self.getNdArray().requires_grad 40 | 41 | @property 42 | def grad(self): 43 | return self.getNdArray().grad 44 | 45 | def backward(self, *args, **kwargs): 46 | return self.getNdArray().backward(*args, **kwargs) 47 | 48 | def max(self): 49 | max = self.getNdArray().max() 50 | return max 51 | 52 | def min(self): 53 | min = self.getNdArray().min() 54 | return min 55 | 56 | def norm(self, N=2): 57 | norm = torch.linalg.norm(self.getNdArray().flatten(), ord=N) 58 | return norm 59 | 60 | def zero(self): 61 | self[:] = torch.zeros_like(self.getNdArray()) 62 | return self 63 | 64 | def set(self, val: float or int): 65 | self[:] = val * torch.ones_like(self[:]) 66 | return self 67 | 68 | 69 | class _Function(torch.autograd.Function): 70 | 71 | @staticmethod 72 | def forward(ctx, model: torch.Tensor, vec_class, fwd_fn, adj_fn, op_dev: torch.device = None, 73 | torch_comp: bool = True) -> torch.Tensor: 74 | ctx.fwd = fwd_fn 75 | ctx.adj = adj_fn 76 | ctx.op_dev = op_dev 77 | ctx.torch_comp = torch_comp 78 | ctx.vec_class = vec_class 79 | 80 | if ctx.torch_comp: 81 | # operator is torch-compatible, so we move the domain to the same device of the operator 82 | data = ctx.fwd(ctx.vec_class(model.clone().to(ctx.op_dev)))[:].to(model.device) 83 | else: 84 | # operator is not torch-compatible, so we move the domain to CPU to the same vec_class of the operator 85 | data = ctx.fwd(ctx.vec_class(model.detach().cpu().numpy())) 86 | data = torch.from_numpy(data[:]).to(model.device) 87 | return data 88 | 89 | @staticmethod 90 | def backward(ctx, data: torch.Tensor): 91 | 92 | if ctx.torch_comp: 93 | # operator is torch-compatible, so we move the domain to the same device of the operator 94 | model = ctx.adj(ctx.vec_class(data.to(ctx.op_dev)))[:].to(data.device) 95 | else: 96 | # operator is not torch-compatible, so we move the domain to CPU 97 | model = ctx.adj(data.detach().cpu().numpy()) 98 | model = torch.from_numpy(model[:]).to(data.device) 99 | return model, None, None, None, None, None 100 | 101 | 102 | class AutogradFunction: 103 | """ 104 | Cast a Operator to a Autograd Function to be used in the torch graph 105 | 106 | Examples: 107 | T = ToAutograd(my_operator) 108 | 109 | T(tensor) -> tensor 110 | 111 | T(vector) -> tensor 112 | 113 | T * vector -> vector 114 | """ 115 | 116 | def __init__(self, operator): 117 | """ 118 | AutogradFunction constructor 119 | 120 | Args: 121 | operator: linear operator, can be based on any vector backend 122 | """ 123 | # check if operator is torch compatible 124 | _is_vectorAD = isinstance(operator.domain, VectorAD) and isinstance(operator.range, VectorAD) 125 | _is_vectorTorch = isinstance(operator.domain, VectorTorch) and isinstance(operator.range, VectorTorch) 126 | self.vec_class = operator.range.__class__ 127 | 128 | if _is_vectorAD or _is_vectorTorch: 129 | self.torch_comp = True 130 | else: 131 | self.torch_comp = False 132 | 133 | # check the device on which the operator runs 134 | if self.torch_comp: 135 | try: 136 | self.op_dev = operator.device 137 | except AttributeError: 138 | self.op_dev = operator.domain.device 139 | else: # is not a torch based operator 140 | self.op_dev = "cpu" 141 | 142 | self.fwd = lambda x: operator * x 143 | self.adj = lambda x: operator.T * self.vec_class(x) 144 | self.function = _Function.apply 145 | 146 | def __mul__(self, other): # occamypy notation: self * vector -> vector 147 | out = self(other) 148 | return other.__class__(out) 149 | 150 | def __call__(self, other): # torch notation: self(tensor, vector) -> tensor 151 | if isinstance(other, torch.Tensor): 152 | return self.apply(other) 153 | else: 154 | return self.apply(other.getNdArray()) 155 | 156 | def apply(self, model: torch.Tensor) -> torch.Tensor: 157 | return self.function(model, self.vec_class, 158 | self.fwd, self.adj, 159 | self.op_dev, self.torch_comp, ) 160 | 161 | 162 | if __name__ == "__main__": 163 | import occamypy as o 164 | 165 | S = torch.nn.Sigmoid() 166 | 167 | # use with VectorTorch (requires_grad=False) 168 | x = o.VectorTorch(torch.ones(2)) 169 | T = AutogradFunction(o.Scaling(x, 2)) 170 | y_ = T(x) 171 | sig_y_ = S(T(x)) 172 | y_sig = T(S(x[:])) 173 | y = T * x 174 | del y, y_, sig_y_, y_sig 175 | 176 | # use with VectorAD (requires_grad=True) to wrap learnable tensors 177 | x = VectorAD(torch.ones(2)) 178 | T = AutogradFunction(o.Scaling(x, 2)) 179 | y_ = T(x) 180 | sig_y_ = S(T(x)) 181 | y_sig = T(S(x[:])) 182 | y = T * x 183 | del y, y_, sig_y_, y_sig 184 | 185 | # now try a numpy-based operator on a tensor with requires_grad=False 186 | x = VectorTorch(torch.zeros(21)) 187 | x[10] = 1 188 | T = AutogradFunction(o.GaussianFilter(o.VectorNumpy((x.size,)), 1)) 189 | y_ = T(x) 190 | sig_y_ = S(T(x)) 191 | y_sig = T(S(x[:])) 192 | y = T * x 193 | del y, y_, sig_y_, y_sig 194 | 195 | # now try a numpy-based operator on a tensor with requires_grad=True 196 | x = o.VectorAD(x) 197 | T = o.AutogradFunction(o.GaussianFilter(o.VectorNumpy((x.size,)), 1)) 198 | y_ = T(x) 199 | sig_y_ = S(T(x)) 200 | y_sig = T(S(x[:])) 201 | y = T * x 202 | 203 | # now try a numpy-based operator on a tensor with requires_grad=True on GPU 204 | # like: we have a tensor that comes from a GPU (cnn) and we want to apply a numpy-based operator like devito 205 | x = o.VectorAD(torch.Tensor([0., 1., 0.]), device=0) 206 | T = o.AutogradFunction(o.GaussianFilter(o.VectorNumpy((x.size,)), 1)) 207 | y_ = T(x) 208 | sig_y_ = S(T(x)) 209 | y_sig = T(S(x[:])) 210 | y = T * x 211 | del y, y_, sig_y_, y_sig 212 | 213 | print(0) 214 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![occamypy](readme_img/logo192.png) 2 | 3 | # OccamyPy: an object-oriented optimization framework for small- and large-scale problems 4 | 5 | We present an object-oriented optimization framework that can be employed to solve 6 | small- and large-scale problems based on the concept of vectors and operators. 7 | By using such a strategy, we implement different iterative optimization algorithms 8 | that can be used in combination with architecture-independent vectors and operators, 9 | allowing the minimization of single-machine or cluster-based problems with a unique codebase. 10 | We implement a Python library following the described structure with a user-friendly interface. 11 | We demonstrate its flexibility and scalability on multiple inverse problems, 12 | where convex and non-convex objective functions are optimized with different iterative algorithms. 13 | 14 | ### Installation 15 | Preferred way is through Python Package Index: 16 | ```bash 17 | pip install occamypy 18 | ``` 19 | In order to have Cupy-based vectors and operators, you should install also [Cupy](https://docs.cupy.dev/en/stable/install.html#install-cupy) and [cuSIGNAL](https://github.com/rapidsai/cusignal#installation). 20 | They are not included in this installation as they are dependent on the target CUDA device and compiler. 21 | 22 | As this library strongly relies on Numpy, we suggest installing OccamyPy in a conda environment like [this](./envs/env.yml) with: 23 | ```bash 24 | conda env create -n MYENV -f env.yml 25 | ``` 26 | 27 | ### History 28 | This library was initially developed at 29 | [Stanford Exploration Project](http://zapad.stanford.edu/ettore88/python-solver) 30 | for solving large scale seismic problems. 31 | Inspired by Equinor's [PyLops](https://github.com/equinor/pylops) 32 | we publish this library as our contribution to scientific community. 33 | 34 | ## How it works 35 | This framework allows for the definition of linear and non-linear mapping functions that 36 | operate on abstract vector objects that can be defined to use 37 | heterogeneous computational resources, from personal laptops to HPC environments. 38 | 39 | - **vector** class: this is the building block for handling data. It contains the required 40 | mathematical operations such as norm, scaling, dot-product, sum, point-wise multiplication. 41 | These methods can be implemented using existing libraries (e.g., Numpy, Cupy, PyTorch) or 42 | user-defined ones (e.g., [SEPLib](http://sepwww.stanford.edu/doku.php?id=sep:software:seplib)). 43 | See the [`vector`](./occamypy/vector) subpackage for details and implementations. 44 | 45 | - **operator** class: a mapping function between a `domain` vector and a `range` vector. 46 | It can be linear and non-linear. 47 | Linear operators require the definition of both the forward and adjoint functions; 48 | non-linear operators require the forward mapping and its Jacobian operator. 49 | See the [`operator`](./occamypy/operator) subpackage for details and implementations. 50 | 51 | - **problem** class: it represents the objective function related to an optimization problem. 52 | Defined upon operators (e.g., modeling and regularization) and vectors (observed data, priors). 53 | It contains the methods for objective function and gradient computation, as our solvers are mainly gradient based. 54 | See the [`problem`](./occamypy/problem) subpackage for details and implementations. 55 | 56 | - **solver** class: it aims at finding the solution to a problem by employing methods 57 | defined within the vector, operator and problem classes. 58 | Additionally, it allows to restart an optimization method from an intermediate result 59 | written as serialized objects on permanent computer memory. 60 | We have a number of linear and nonlinear solver, along with some stepper algorithms. 61 | See the [`solver`](./occamypy/solver) subpackage for details and implementations. 62 | Solvers come with a Logger object that we found helpful for saving large-scale inversions. Check it out in the tutorials! 63 | 64 | ### Features at a glance 65 | 66 | | vector engines | operators | problems | solvers | 67 | |-|-|-|-| 68 | | numpy | linear | least squares | Conjugate Gradient | 69 | | cupy | nonlinear | symmetric least squares | Steepest Descent | 70 | | torch | distributed | L2-reg least squares | LSQR | 71 | | | | LASSO | symmetric Conjugate Gradient | 72 | | | | generalized LASSO | nonlinear Conjugate Gradient | 73 | | | | nonlinear least squares | L-BFGS | 74 | | | | L2-reg nonlinear least squares | L-BFGS-B | 75 | | | | regularized Variable Projection | Truncated Newton | 76 | | | | | Markov Chain Monte Carlo | 77 | | | | | ISTA and Fast-ISTA | 78 | | | | | ISTC (ISTA with cooling) | 79 | | | | | Split-Bregman | 80 | 81 | ### Scalability 82 | The main objective of the described framework and implemented library is to solve large-scale inverse problems. 83 | Any vector and operator can be split into blocks to be distributed to multiple nodes. 84 | This is achieved via custom [Dask](https://dask.org/) vector and operator classes. 85 | See the [`dask`](./occamypy/dask) subpackage for details and implementations. 86 | 87 | ### Tutorials 88 | We provide some [tutorials](./tutorials) that demonstrate the flexibility of occamypy. 89 | Please refer to them as a good starting point for developing your own code. 90 | If you have a good application example, contact us! We will be happy to see OccamyPy in action. 91 | 92 | Check out the [tutorial](https://curvenote.com/@swung/transform-2022-occamypy-an-oo-optimizaton-library/overview) we gave at SWUNG's Transform 2022! 93 | 94 | ### Contributing 95 | Follow the following instructions and read carefully the [CONTRIBUTING](CONTRIBUTING.md) file before getting started. 96 | 97 | We have a lot of ideas that might be helpful to scientists! 98 | We are currently working on: 99 | * wrapping linear operators to PyTorch optimizers: see the [LS-RTM tutorial](./tutorials/2D%20LS-RTM%20with%20devito%20and%20Automatic%20Differentiation.ipynb)! 100 | This can be useful for using neural networks and operators (i.e., deep priors and physical modeling). 101 | * using PyTorch's [functorch](https://github.com/pytorch/functorch) library to compute the Jacobian-vector product of nonlinear operators: see the [first step](tutorials/Automatic%20Differentiation%20for%20nonlinear%20operators.ipynb)! 102 | * implement computation-demanding operators natively in OccamyPy, so that they can be used on CPU/GPU and HPC clusters. 103 | 104 | ### Authors 105 | - [Ettore Biondi](https://github.com/biondiettore) 106 | - [Guillame Barnier](https://github.com/gbarnier) 107 | - [Robert Clapp](http://zapad.stanford.edu/bob) 108 | - [Francesco Picetti](https://github.com/fpicetti) 109 | - [Stuart Farris](http://zapad.stanford.edu/sfarris) 110 | 111 | ### Citation 112 | ``` 113 | @article{biondi2021object, 114 | title = {An object-oriented optimization framework for large-scale inverse problems}, 115 | author = {Ettore Biondi and Guillaume Barnier and Robert G. Clapp and Francesco Picetti and Stuart Farris}, 116 | journal = {Computers & Geosciences}, 117 | volume = {154}, 118 | pages = {104790}, 119 | year = {2021}, 120 | doi = {https://doi.org/10.1016/j.cageo.2021.104790}, 121 | } 122 | ``` 123 | 124 | ### Publications using OccamyPy 125 | 126 | * E. Biondi, G. Barnier, R. G. Clapp, F. Picetti, and S. Farris. "Object-Oriented Optimization for Small- and Large-Scale Seismic Inversion Procedures", in _European Association of Geophysicists and Engineers (EAGE) Workshop on High Performance Computing for Upstream_, 2021. [link](https://doi.org/10.3997/2214-4609.202181003). 127 | * E. Biondi, G. Barnier, R. G. Clapp, F. Picetti, and S. Farris. "Object-oriented optimization for large-scale seismic inversion of ocean-bottom-node pressure data", in _International Conference on Parallel Computational Fluid Dynamics (ParCFD)_, 2021. [link](https://parcfd2020.sciencesconf.org/345756). 128 | 129 | If you have one to add, reach us out! 130 | -------------------------------------------------------------------------------- /occamypy/torch/operator/signal.py: -------------------------------------------------------------------------------- 1 | from itertools import accumulate, product 2 | from typing import Union, List, Tuple 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from occamypy.operator.base import Operator, Dstack 8 | from occamypy.torch.back_utils import set_backends 9 | from occamypy.torch.vector import VectorTorch 10 | from occamypy.vector.base import superVector 11 | 12 | set_backends() 13 | 14 | 15 | def _gaussian_kernel1d(sigma: float, order: int = 0, truncate: float = 4.) -> torch.Tensor: 16 | """Computes a 1-D Gaussian convolution kernel""" 17 | 18 | radius = int(truncate * sigma + 0.5) 19 | if order < 0: 20 | raise ValueError('order must be non-negative') 21 | exponent_range = torch.arange(order + 1) 22 | sigma2 = sigma * sigma 23 | x = torch.arange(-radius, radius+1) 24 | phi_x = torch.exp(-0.5 / sigma2 * x ** 2) 25 | phi_x = phi_x / phi_x.sum() 26 | 27 | if order == 0: 28 | return phi_x 29 | else: 30 | # f(x) = q(x) * phi(x) = q(x) * exp(p(x)) 31 | # f'(x) = (q'(x) + q(x) * p'(x)) * phi(x) 32 | # p'(x) = -1 / sigma ** 2 33 | # Implement q'(x) + q(x) * p'(x) as a matrix operator and apply to the 34 | # coefficients of q(x) 35 | q = torch.zeros(order + 1) 36 | q[0] = 1 37 | D = torch.diag(exponent_range[1:], 1) # D @ q(x) = q'(x) 38 | P = torch.diag(torch.ones(order)/-sigma2, -1) # P @ q(x) = q(x) * p'(x) 39 | Q_deriv = D + P 40 | for _ in range(order): 41 | q = Q_deriv.dot(q) 42 | q = (x[:, None] ** exponent_range).dot(q) 43 | return q * phi_x 44 | 45 | 46 | class ConvND(Operator): 47 | """ND convolution square operator in the domain space""" 48 | 49 | def __init__(self, domain: VectorTorch, kernel: Union[VectorTorch, torch.Tensor]): 50 | """ 51 | ConvND (torch) constructor 52 | 53 | Args: 54 | domain: domain vector 55 | kernel: kernel vector or tensor 56 | """ 57 | if isinstance(kernel, VectorTorch): 58 | self.kernel = kernel.getNdArray().clone() 59 | elif isinstance(kernel, torch.Tensor): 60 | self.kernel = kernel.clone() 61 | 62 | if domain.ndim != self.kernel.ndim: 63 | raise ValueError("Domain and kernel number of dimensions mismatch") 64 | 65 | if domain.device != self.kernel.device: 66 | raise ValueError("Domain and kernel has to live in the same device") 67 | 68 | self.kernel_size = tuple(self.kernel.shape) 69 | self.pad_size = tuple([k // 2 for k in self.kernel_size]) 70 | 71 | if domain.ndim == 1: 72 | corr = torch.nn.functional.conv1d 73 | elif domain.ndim == 2: 74 | corr = torch.nn.functional.conv2d 75 | elif domain.ndim == 3: 76 | corr = torch.nn.functional.conv3d 77 | else: 78 | raise ValueError 79 | 80 | # torch.nn functions require batch and channel dimensions 81 | self.conv = lambda x: corr(x.unsqueeze(0).unsqueeze(0), 82 | torch.flip(self.kernel, dims=tuple(range(self.kernel.ndim))).unsqueeze(0).unsqueeze(0), 83 | padding=self.pad_size).flatten(end_dim=2) 84 | 85 | self.corr = lambda x: corr(x.unsqueeze(0).unsqueeze(0), 86 | self.kernel.unsqueeze(0).unsqueeze(0), 87 | padding=self.pad_size).flatten(end_dim=2) 88 | 89 | super(ConvND, self).__init__(domain, domain, name="Convolve") 90 | 91 | def forward(self, add, model, data): 92 | self.checkDomainRange(model, data) 93 | if not add: 94 | data.zero() 95 | data[:] += self.conv(model.getNdArray()) 96 | return 97 | 98 | def adjoint(self, add, model, data): 99 | self.checkDomainRange(model, data) 100 | if not add: 101 | model.zero() 102 | model[:] += self.corr(data.getNdArray()) 103 | return 104 | 105 | 106 | class GaussianFilter(ConvND): 107 | """Gaussian smoothing operator""" 108 | 109 | def __init__(self, domain, sigma): 110 | """ 111 | GaussianFilter (torch) constructor 112 | 113 | Args: 114 | domain: domain vector 115 | sigma: standard deviation along the domain directions 116 | """ 117 | self.sigma = [sigma] if isinstance(sigma, (float, int)) else sigma 118 | if not isinstance(self.sigma, (list, tuple)): 119 | raise TypeError("sigma has to be either a list or a tuple") 120 | self.scaling = np.sqrt(np.prod(np.array(self.sigma) / np.pi)) 121 | kernels = [_gaussian_kernel1d(s) for s in self.sigma] 122 | 123 | self.kernel = [*accumulate(kernels, lambda a, b: torch.outer(a.flatten(), b))][-1] 124 | self.kernel.reshape([k.numel() for k in kernels]) 125 | super(GaussianFilter, self).__init__(domain, self.kernel.to(domain.device)) 126 | self.name = "GaussFilt" 127 | 128 | 129 | def Padding(domain, pad, mode="constant"): 130 | """ 131 | Padding operator 132 | 133 | Notes: 134 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 135 | pad=((2,2), (3,3)) 136 | 137 | Args: 138 | domain: domain vector 139 | pad: number of samples to be added at each end of the dimension, for each dimension 140 | mode: padding mode (see https://pytorch.org/docs/1.10/generated/torch.nn.functional.pad.html) 141 | """ 142 | if isinstance(domain, VectorTorch): 143 | return _Padding(domain, pad, mode) 144 | elif isinstance(domain, superVector): 145 | return Dstack([_Padding(v, pad, mode) for v in domain.vecs]) 146 | else: 147 | raise ValueError("ERROR! Provided domain has to be either vector or superVector") 148 | 149 | 150 | def ZeroPad(domain, pad): 151 | """ 152 | Zero-Padding operator 153 | 154 | Notes: 155 | To pad 2 values to each side of the first dim, and 3 values to each side of the second dim, use: 156 | pad=((2,2), (3,3)) 157 | 158 | Args: 159 | domain: domain vector 160 | pad: number of samples to be added at each end of the dimension, for each dimension 161 | """ 162 | return Padding(domain, pad, mode="constant") 163 | 164 | 165 | class _Padding(Operator): 166 | 167 | def __init__(self, domain: VectorTorch, pad: Union[int, Tuple[int]], mode: str = "constant"): 168 | 169 | nd = domain.ndim 170 | 171 | if isinstance(pad, int): 172 | pad = [int(pad), int(pad)] * nd 173 | else: 174 | pad = np.asarray(pad).ravel().tolist() 175 | if len(pad) != 2 * nd: 176 | raise ValueError("len(pad) has to be 2*nd") 177 | 178 | if (np.array(pad) < 0).any(): 179 | raise ValueError('Padding must be positive or zero') 180 | 181 | self.pad = list(pad) 182 | 183 | self.padded_shape = tuple(np.asarray(domain.shape) + [self.pad[i] + self.pad[i + 1] for i in range(0, 2 * nd, 2)]) 184 | 185 | super(_Padding, self).__init__(domain, VectorTorch(self.padded_shape, device=domain.device.index), name="Paddding") 186 | 187 | self.inner_idx = [list(torch.arange(start=self.pad[0:-1:2][i], end=self.range.shape[i]-pad[1::2][i])) for i in range(nd)] 188 | self.mode = mode 189 | 190 | def forward(self, add, model, data): 191 | """Pad the domain""" 192 | self.checkDomainRange(model, data) 193 | if not add: 194 | data.zero() 195 | # torch counts the axes in reverse order 196 | if self.mode == "constant": 197 | padded = torch.nn.functional.pad(model.getNdArray(), self.pad[::-1], mode=self.mode) 198 | else: 199 | padded = torch.nn.functional.pad(model.getNdArray()[None], self.pad[::-1], mode=self.mode).squeeze(0) 200 | data[:] += padded 201 | return 202 | 203 | def adjoint(self, add, model, data): 204 | """Extract original subsequence""" 205 | self.checkDomainRange(model, data) 206 | if not add: 207 | model.zero() 208 | x = torch.Tensor([data[coord] for coord in product(*self.inner_idx)]).reshape(self.domain.shape).to(model.device) 209 | model[:] += x 210 | return 211 | -------------------------------------------------------------------------------- /occamypy/numpy/vector.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | from sys import version_info 3 | 4 | import numpy as np 5 | 6 | from occamypy.vector.base import Vector 7 | 8 | 9 | class VectorNumpy(Vector): 10 | """Vector class based on numpy.ndarray""" 11 | 12 | def __init__(self, in_content, *args, **kwargs): 13 | """ 14 | VectorNumpy constructor 15 | 16 | Args: 17 | in_content: numpy.ndarray, tuple or path_to_file to load a numpy.ndarray 18 | *args: list of arguments for Vector construction 19 | **kwargs: dict of arguments for Vector construction 20 | """ 21 | if isinstance(in_content, str): # Header file name passed to constructor 22 | self.arr = np.load(in_content, allow_pickle=True) 23 | elif isinstance(in_content, np.ndarray): # Numpy array passed to constructor 24 | # if np.isfortran(in_content): # for seplib compatibility 25 | # raise TypeError('Input array not a C contiguous array!') 26 | self.arr = np.array(in_content, copy=False) 27 | elif isinstance(in_content, tuple): # Tuple size passed to constructor 28 | self.arr = np.zeros(in_content) 29 | else: # Not supported type 30 | raise ValueError("ERROR! Input variable not currently supported!") 31 | 32 | super(VectorNumpy, self).__init__(*args, **kwargs) 33 | # Number of elements per axis (tuple). Checking also the memory order 34 | self.shape = self.arr.shape # If fortran the first axis is the "fastest" 35 | self.ndim = self.arr.ndim # Number of axes integer 36 | self.size = self.arr.size # Total number of elements 37 | 38 | def getNdArray(self): 39 | return self.arr 40 | 41 | def norm(self, N=2): 42 | return np.linalg.norm(self.getNdArray().ravel(), ord=N) 43 | 44 | def zero(self): 45 | self.getNdArray().fill(0) 46 | return self 47 | 48 | def max(self): 49 | return self.getNdArray().max() 50 | 51 | def min(self): 52 | return self.getNdArray().min() 53 | 54 | def set(self, val): 55 | self.getNdArray().fill(val) 56 | return self 57 | 58 | def scale(self, sc): 59 | self.getNdArray()[:] *= sc 60 | return self 61 | 62 | def addbias(self, bias): 63 | self.getNdArray()[:] += bias 64 | return self 65 | 66 | def rand(self, low: float = -1., high: float = 1.): 67 | self.arr = np.random.uniform(low=low, high=high, size=self.shape) 68 | return self 69 | 70 | def randn(self, mean: float = 0., std: float = 1.): 71 | self.arr = np.random.normal(loc=mean, scale=std, size=self.shape) 72 | return self 73 | 74 | def clone(self): 75 | vec_clone = deepcopy(self) # Deep clone of vector 76 | # Checking if a vector space was provided 77 | if vec_clone.getNdArray().size == 0: # this is the shape of np.ndarray! 78 | vec_clone.arr = np.zeros(vec_clone.shape, dtype=self.getNdArray().dtype) 79 | return vec_clone 80 | 81 | def cloneSpace(self): 82 | vec_space = VectorNumpy(np.empty(0, dtype=self.getNdArray().dtype)) 83 | vec_space.ax_info = self.ax_info 84 | # Cloning space of input vector 85 | vec_space.ndim = self.ndim 86 | vec_space.shape = self.shape 87 | vec_space.size = self.size 88 | return vec_space 89 | 90 | def checkSame(self, other): 91 | return self.shape == other.shape 92 | 93 | def abs(self): 94 | self.getNdArray()[:] = np.abs(self.getNdArray()) 95 | return self 96 | 97 | def sign(self): 98 | self.getNdArray()[:] = np.sign(self.getNdArray()) 99 | return self 100 | 101 | def reciprocal(self): 102 | self.getNdArray()[:] = 1. / self.getNdArray() 103 | return self 104 | 105 | def maximum(self, other): 106 | if np.isscalar(other): 107 | self.getNdArray()[:] = np.maximum(self.getNdArray(), other) 108 | return self 109 | elif isinstance(other, VectorNumpy): 110 | if not self.checkSame(other): 111 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 112 | self.getNdArray()[:] = np.maximum(self.getNdArray(), other.getNdArray()) 113 | return self 114 | else: 115 | raise TypeError("Provided input has to be either a scalar or a %s!" % self.whoami) 116 | 117 | def conj(self): 118 | self.getNdArray()[:] = np.conjugate(self.getNdArray()) 119 | return self 120 | 121 | def transpose(self): 122 | other = VectorNumpy(tuple(reversed(self.shape))) 123 | other[:] = self.getNdArray().T 124 | return other 125 | 126 | def pow(self, power): 127 | self.getNdArray()[:] = self.getNdArray() ** power 128 | return self 129 | 130 | def real(self): 131 | self.getNdArray()[:] = self.getNdArray().real 132 | return self 133 | 134 | def imag(self, ): 135 | self.getNdArray()[:] = self.getNdArray().imag 136 | return self 137 | 138 | def copy(self, other): 139 | # Checking whether the input is a vector or not 140 | if not isinstance(other, VectorNumpy): 141 | raise TypeError("Provided input vector not a %s!" % self.whoami) 142 | # Checking dimensionality 143 | if not self.checkSame(other): 144 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 145 | # Element-wise copy of the input array 146 | self.getNdArray()[:] = other.getNdArray() 147 | return self 148 | 149 | def scaleAdd(self, other, sc1=1.0, sc2=1.0): 150 | # Checking whether the input is a vector or not 151 | if not isinstance(other, VectorNumpy): 152 | raise TypeError("Provided input vector not a %s!" % self.whoami) 153 | # Checking dimensionality 154 | if not self.checkSame(other): 155 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 156 | # Performing scaling and addition 157 | self.getNdArray()[:] = sc1 * self.getNdArray() + sc2 * other.getNdArray() 158 | return self 159 | 160 | def dot(self, other): 161 | # Checking whether the input is a vector or not 162 | if not isinstance(other, VectorNumpy): 163 | raise TypeError("Provided input vector not a %s!" % self.whoami) 164 | # Checking size (must have same number of elements) 165 | if self.size != other.size: 166 | raise ValueError("Vector size mismatching: self = %d; other = %d" % (self.size, other.size)) 167 | # Checking dimensionality 168 | if not self.checkSame(other): 169 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 170 | return np.vdot(self.getNdArray().ravel(), other.getNdArray().ravel()) 171 | 172 | def multiply(self, other): 173 | # Checking whether the input is a vector or not 174 | if not isinstance(other, VectorNumpy): 175 | raise TypeError("Provided input vector not a %s!" % self.whoami) 176 | # Checking size (must have same number of elements) 177 | if self.size != other.size: 178 | raise ValueError("Vector size mismatching: self = %s; other = %s" % (self.size, other.size)) 179 | # Checking dimensionality 180 | if not self.checkSame(other): 181 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 182 | # Performing element-wise multiplication 183 | self.getNdArray()[:] = np.multiply(self.getNdArray(), other.getNdArray()) 184 | return self 185 | 186 | def isDifferent(self, other): 187 | # Checking whether the input is a vector or not 188 | if not isinstance(other, VectorNumpy): 189 | raise TypeError("Provided input vector not a %s!" % self.whoami) 190 | # Using Hash table for python2 and numpy built-in function array_equal otherwise 191 | if version_info[0] == 2: 192 | # First make both array buffers read-only 193 | self.arr.flags.writeable = False 194 | other.arr.flags.writeable = False 195 | chcksum1 = hash(self.getNdArray().data) 196 | chcksum2 = hash(other.getNdArray().data) 197 | # Remake array buffers writable 198 | self.arr.flags.writeable = True 199 | other.arr.flags.writeable = True 200 | isDiff = (chcksum1 != chcksum2) 201 | else: 202 | isDiff = (not np.array_equal(self.getNdArray(), other.getNdArray())) 203 | return isDiff 204 | 205 | def clip(self, low, high): 206 | if not isinstance(low, VectorNumpy): 207 | raise TypeError("Provided input low vector not a %s!" % self.whoami) 208 | if not isinstance(high, VectorNumpy): 209 | raise TypeError("Provided input high vector not a %s!" % self.whoami) 210 | self.getNdArray()[:] = np.minimum(np.maximum(low.getNdArray(), self.getNdArray()), high.getNdArray()) 211 | return self 212 | 213 | def plot(self): 214 | return self.getNdArray() 215 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from argparse import Action, ArgumentError, ArgumentParser 3 | 4 | from devito import error, configuration, warning 5 | from devito.tools import Pickable 6 | 7 | from .source import * 8 | 9 | __all__ = ['AcquisitionGeometry', 'setup_geometry', 'seismic_args'] 10 | 11 | 12 | def setup_geometry(model, tn, f0=0.010): 13 | # Source and receiver geometries 14 | src_coordinates = np.empty((1, model.dim)) 15 | src_coordinates[0, :] = np.array(model.domain_size) * .5 16 | if model.dim > 1: 17 | src_coordinates[0, -1] = model.origin[-1] + model.spacing[-1] 18 | 19 | rec_coordinates = setup_rec_coords(model) 20 | 21 | geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, 22 | t0=0.0, tn=tn, src_type='Ricker', f0=f0) 23 | 24 | return geometry 25 | 26 | 27 | def setup_rec_coords(model): 28 | nrecx = model.shape[0] 29 | recx = np.linspace(model.origin[0], model.domain_size[0], nrecx) 30 | 31 | if model.dim == 1: 32 | return recx.reshape((nrecx, 1)) 33 | elif model.dim == 2: 34 | rec_coordinates = np.empty((nrecx, model.dim)) 35 | rec_coordinates[:, 0] = recx 36 | rec_coordinates[:, -1] = model.origin[-1] + 2 * model.spacing[-1] 37 | return rec_coordinates 38 | else: 39 | nrecy = model.shape[1] 40 | recy = np.linspace(model.origin[1], model.domain_size[1], nrecy) 41 | rec_coordinates = np.empty((nrecx*nrecy, model.dim)) 42 | rec_coordinates[:, 0] = np.array([recx[i] for i in range(nrecx) 43 | for j in range(nrecy)]) 44 | rec_coordinates[:, 1] = np.array([recy[j] for i in range(nrecx) 45 | for j in range(nrecy)]) 46 | rec_coordinates[:, -1] = model.origin[-1] + 2 * model.spacing[-1] 47 | return rec_coordinates 48 | 49 | 50 | class AcquisitionGeometry(Pickable): 51 | """ 52 | Encapsulate the geometry of an acquisition: 53 | - receiver positions and number 54 | - source positions and number 55 | 56 | In practice this would only point to a segy file with the 57 | necessary information 58 | """ 59 | 60 | def __init__(self, model, rec_positions, src_positions, t0, tn, **kwargs): 61 | """ 62 | In practice would be __init__(segyfile) and all below parameters 63 | would come from a segy_read (at property call rather than at init) 64 | """ 65 | src_positions = np.reshape(src_positions, (-1, model.dim)) 66 | rec_positions = np.reshape(rec_positions, (-1, model.dim)) 67 | self.rec_positions = rec_positions 68 | self._nrec = rec_positions.shape[0] 69 | self.src_positions = src_positions 70 | self._nsrc = src_positions.shape[0] 71 | self._src_type = kwargs.get('src_type') 72 | assert (self.src_type in sources or self.src_type is None) 73 | self._f0 = kwargs.get('f0') 74 | self._a = kwargs.get('a', None) 75 | self._t0w = kwargs.get('t0w', None) 76 | if self._src_type is not None and self._f0 is None: 77 | error("Peak frequency must be provided in KH" + 78 | " for source of type %s" % self._src_type) 79 | 80 | self._grid = model.grid 81 | self._model = model 82 | self._dt = model.critical_dt 83 | self._t0 = t0 84 | self._tn = tn 85 | 86 | def resample(self, dt): 87 | self._dt = dt 88 | return self 89 | 90 | @property 91 | def time_axis(self): 92 | return TimeAxis(start=self.t0, stop=self.tn, step=self.dt) 93 | 94 | @property 95 | def src_type(self): 96 | return self._src_type 97 | 98 | @property 99 | def grid(self): 100 | return self._grid 101 | 102 | @property 103 | def model(self): 104 | warning("Model is kept for backward compatibility but should not be" 105 | "obtained from the geometry") 106 | return self._model 107 | 108 | @property 109 | def f0(self): 110 | return self._f0 111 | 112 | @property 113 | def tn(self): 114 | return self._tn 115 | 116 | @property 117 | def t0(self): 118 | return self._t0 119 | 120 | @property 121 | def dt(self): 122 | return self._dt 123 | 124 | @property 125 | def nt(self): 126 | return self.time_axis.num 127 | 128 | @property 129 | def nrec(self): 130 | return self._nrec 131 | 132 | @property 133 | def nsrc(self): 134 | return self._nsrc 135 | 136 | @property 137 | def dtype(self): 138 | return self.grid.dtype 139 | 140 | @property 141 | def rec(self): 142 | return self.new_rec() 143 | 144 | def new_rec(self, name='rec'): 145 | return Receiver(name=name, grid=self.grid, 146 | time_range=self.time_axis, npoint=self.nrec, 147 | coordinates=self.rec_positions) 148 | 149 | @property 150 | def adj_src(self): 151 | if self.src_type is None: 152 | warning("No source type defined, returning uninitiallized (zero) shot record") 153 | return self.new_rec() 154 | adj_src = sources[self.src_type](name='rec', grid=self.grid, f0=self.f0, 155 | time_range=self.time_axis, npoint=self.nrec, 156 | coordinates=self.rec_positions, 157 | t0=self._t0w, a=self._a) 158 | # Revert time axis to have a proper shot record and not compute on zeros 159 | for i in range(self.nrec): 160 | adj_src.data[:, i] = adj_src.wavelet[::-1] 161 | return adj_src 162 | 163 | @property 164 | def src(self): 165 | return self.new_src() 166 | 167 | def new_src(self, name='src', src_type='self'): 168 | if self.src_type is None or src_type is None: 169 | warning("No surce type defined, returning uninistiallized (zero) source") 170 | return PointSource(name=name, grid=self.grid, 171 | time_range=self.time_axis, npoint=self.nsrc, 172 | coordinates=self.src_positions) 173 | else: 174 | return sources[self.src_type](name=name, grid=self.grid, f0=self.f0, 175 | time_range=self.time_axis, npoint=self.nsrc, 176 | coordinates=self.src_positions, 177 | t0=self._t0w, a=self._a) 178 | 179 | _pickle_args = ['grid', 'rec_positions', 'src_positions', 't0', 'tn'] 180 | _pickle_kwargs = ['f0', 'src_type'] 181 | 182 | 183 | sources = {'Wavelet': WaveletSource, 'Ricker': RickerSource, 'Gabor': GaborSource} 184 | 185 | 186 | def seismic_args(description): 187 | """ 188 | Command line options for the devitoseismic examples 189 | """ 190 | 191 | class _dtype_store(Action): 192 | def __call__(self, parser, args, values, option_string=None): 193 | values = {'float32': np.float32, 'float64': np.float64}[values] 194 | setattr(args, self.dest, values) 195 | 196 | class _opt_action(Action): 197 | def __call__(self, parser, args, values, option_string=None): 198 | try: 199 | # E.g., `('advanced', {'par-tile': True})` 200 | values = eval(values) 201 | if not isinstance(values, tuple) and len(values) >= 1: 202 | raise ArgumentError(self, ("Invalid choice `%s` (`opt` must be " 203 | "either str or tuple)" % str(values))) 204 | opt = values[0] 205 | except NameError: 206 | # E.g. `'advanced'` 207 | opt = values 208 | if opt not in configuration._accepted['opt']: 209 | raise ArgumentError(self, ("Invalid choice `%s` (choose from %s)" 210 | % (opt, str(configuration._accepted['opt'])))) 211 | setattr(args, self.dest, values) 212 | 213 | parser = ArgumentParser(description=description) 214 | parser.add_argument("-nd", dest="ndim", default=3, type=int, 215 | help="Number of dimensions") 216 | parser.add_argument("-d", "--shape", default=(51, 51, 51), type=int, nargs="+", 217 | help="Number of grid points along each axis") 218 | parser.add_argument('-f', '--full', default=False, action='store_true', 219 | help="Execute all operators and store forward wavefield") 220 | parser.add_argument("-so", "--space_order", default=4, 221 | type=int, help="Space order of the simulation") 222 | parser.add_argument("--nbl", default=40, 223 | type=int, help="Number of boundary layers around the domain") 224 | parser.add_argument("--constant", default=False, action='store_true', 225 | help="Constant velocity model, default is a two layer model") 226 | parser.add_argument("--checkpointing", default=False, action='store_true', 227 | help="Use checkpointing, default is false") 228 | parser.add_argument("-opt", default="advanced", action=_opt_action, 229 | help="Performance optimization level") 230 | parser.add_argument('-a', '--autotune', default='off', 231 | choices=(configuration._accepted['autotuning']), 232 | help="Operator auto-tuning mode") 233 | parser.add_argument("-tn", "--tn", default=0, 234 | type=float, help="Simulation time in millisecond") 235 | parser.add_argument("-dtype", action=_dtype_store, dest="dtype", default=np.float32, 236 | choices=['float32', 'float64']) 237 | return parser 238 | -------------------------------------------------------------------------------- /occamypy/problem/base.py: -------------------------------------------------------------------------------- 1 | __all__ = [ 2 | "Bounds", 3 | "Problem", 4 | ] 5 | 6 | 7 | class Bounds: 8 | """ 9 | Class used to enforce boundary constraints during the inversion 10 | 11 | Methods: 12 | apply(in_content): apply bounds to input vector 13 | """ 14 | 15 | def __init__(self, minBound=None, maxBound=None): 16 | """ 17 | Bounds constructor 18 | 19 | Args: 20 | minBound: vector containing minimum values of the domain vector 21 | maxBound: vector containing maximum values of the domain vector 22 | """ 23 | self.minBound = minBound 24 | self.maxBound = maxBound 25 | if minBound is not None: 26 | self.minBound = minBound.clone() 27 | if maxBound is not None: 28 | self.maxBound = maxBound.clone() 29 | # If only the lower bound was provided we use the opposite of the lower bound to clip the values 30 | if self.minBound is not None and self.maxBound is None: 31 | self.minBound.scale(-1.0) 32 | 33 | def apply(self, in_content): 34 | """ 35 | Apply bounds to the input vector 36 | 37 | Args: 38 | in_content: vector to be processed 39 | """ 40 | if self.minBound is not None and self.maxBound is None: 41 | if not in_content.checkSame(self.minBound): 42 | raise ValueError("Input vector not consistent with bound space") 43 | in_content.scale(-1.0) 44 | in_content.clip(in_content, self.minBound) 45 | in_content.scale(-1.0) 46 | elif self.minBound is None and self.maxBound is not None: 47 | if not in_content.checkSame(self.maxBound): 48 | raise ValueError("Input vector not consistent with bound space") 49 | in_content.clip(in_content, self.maxBound) 50 | elif self.minBound is not None and self.maxBound is not None: 51 | if (not (in_content.checkSame(self.minBound) and in_content.checkSame( 52 | self.maxBound))): 53 | raise ValueError("Input vector not consistent with bound space") 54 | in_content.clip(self.minBound, self.maxBound) 55 | return 56 | 57 | 58 | class Problem: 59 | """Base problem class""" 60 | 61 | def __init__(self, model, data, minBound=None, maxBound=None, boundProj=None, name: str = "Problem"): 62 | """ 63 | Problem constructor 64 | 65 | Args: 66 | model: model vector to be optimized 67 | data: data vector 68 | minBound: vector containing minimum values of the domain vector 69 | maxBound: vector containing maximum values of the domain vector 70 | boundProj: class with a function "apply(input_vec)" to project input_vec onto some convex set 71 | name: problem name 72 | """ 73 | self.minBound = minBound 74 | self.maxBound = maxBound 75 | self.boundProj = boundProj 76 | 77 | if minBound is not None or maxBound is not None: 78 | # Simple box bounds 79 | self.bounds = Bounds(minBound, maxBound) # Setting the bounds of the problem (if necessary) 80 | elif boundProj is not None: 81 | # Projection operator onto the bounds 82 | self.bounds = boundProj 83 | 84 | # use model and data as pointers, not clone 85 | self.model = model 86 | self.data = data 87 | 88 | # Setting common variables 89 | self.pert_model = model.clone().zero() 90 | self.res = data.clone().zero() 91 | self.pert_res = self.res.clone() 92 | 93 | self.obj_updated = False 94 | self.res_updated = False 95 | self.grad_updated = False 96 | self.pert_res_updated = False 97 | self.fevals = 0 98 | self.gevals = 0 99 | self.counter = 0 100 | self.linear = False 101 | self.name = str(name) 102 | 103 | def __str__(self): 104 | return self.name 105 | 106 | def setDefaults(self): 107 | """Default common variables for any inverse problem""" 108 | self.obj_updated = False 109 | self.res_updated = False 110 | self.grad_updated = False 111 | self.pert_res_updated = False 112 | self.fevals = 0 113 | self.gevals = 0 114 | self.counter = 0 115 | return 116 | 117 | def set_model(self, model): 118 | """ 119 | Setting internal domain vector 120 | 121 | Args: 122 | model: domain vector to be copied 123 | """ 124 | if model.isDifferent(self.model): 125 | self.model.copy(model) 126 | self.obj_updated = False 127 | self.res_updated = False 128 | self.grad_updated = False 129 | self.pert_res_updated = False 130 | 131 | def set_residual(self, residual): 132 | """ 133 | Setting internal residual vector 134 | 135 | Args: 136 | residual: residual vector to be copied 137 | """ 138 | # Useful for linear inversion (to avoid residual computation) 139 | if self.res.isDifferent(residual): 140 | self.res.copy(residual) 141 | # If residuals have changed, recompute gradient and objective function value 142 | self.grad_updated = False 143 | self.obj_updated = False 144 | self.res_updated = True 145 | return 146 | 147 | def get_model(self): 148 | """Geet the domain vector""" 149 | return self.model 150 | 151 | def get_pert_model(self): 152 | """Get the model perturbation vector""" 153 | return self.pert_model 154 | 155 | def get_rnorm(self, model): 156 | """ 157 | Compute the residual vector norm 158 | 159 | Args: 160 | model: domain vector 161 | """ 162 | self.get_res(model) 163 | return self.get_res(model).norm() 164 | 165 | def get_gnorm(self, model): 166 | """ 167 | Compute the gradient vector norm 168 | 169 | Args: 170 | model: domain vector 171 | """ 172 | return self.get_grad(model).norm() 173 | 174 | def get_obj(self, model): 175 | """Compute the objective function 176 | 177 | Args: 178 | model: domain vector 179 | """ 180 | self.set_model(model) 181 | if not self.obj_updated: 182 | self.res = self.get_res(self.model) 183 | self.obj = self.obj_func(self.res) 184 | self.obj_updated = True 185 | return self.obj 186 | 187 | def get_res(self, model): 188 | """ 189 | Compute the residual vector 190 | 191 | Args: 192 | model: domain vector 193 | """ 194 | self.set_model(model) 195 | if not self.res_updated: 196 | self.fevals += 1 197 | self.res = self.res_func(self.model) 198 | self.res_updated = True 199 | return self.res 200 | 201 | def get_grad(self, model): 202 | """ 203 | Compute the gradient vector 204 | 205 | Args: 206 | model: domain vector 207 | """ 208 | self.set_model(model) 209 | if not self.grad_updated: 210 | self.res = self.get_res(self.model) 211 | self.grad = self.grad_func(self.model, self.res) 212 | self.gevals += 1 213 | if self.linear: 214 | self.fevals += 1 215 | self.grad_updated = True 216 | return self.grad 217 | 218 | def get_pert_res(self, model, pert_model): 219 | """Compute the perturbation residual vector (i.e., application of the Jacobian to model perturbation vector) 220 | 221 | Args: 222 | model: domain vector 223 | pert_model: domain perturbation vector 224 | """ 225 | self.set_model(model) 226 | if not self.pert_res_updated or pert_model.isDifferent(self.pert_model): 227 | self.pert_model.copy(pert_model) 228 | self.pert_res = self.pert_res_func(self.model, self.pert_model) 229 | if self.linear: 230 | self.fevals += 1 231 | self.pert_res_updated = True 232 | return self.pert_res 233 | 234 | def get_fevals(self): 235 | """Get the number of objective function evalutions""" 236 | return self.fevals 237 | 238 | def get_gevals(self): 239 | """Get the number of gradient evalutions""" 240 | return self.gevals 241 | 242 | def obj_func(self, residual): 243 | """ 244 | Compute the objective function 245 | 246 | Args: 247 | residual: residual vector 248 | Returns: 249 | objective function value 250 | """ 251 | raise NotImplementedError("Implement objf for problem in the derived class!") 252 | 253 | def res_func(self, model): 254 | """ 255 | Compute the residual vector 256 | 257 | Args: 258 | model: domain vector 259 | 260 | Returns: residual vector based on the domain 261 | """ 262 | raise NotImplementedError("Implement resf for problem in the derived class!") 263 | 264 | def pert_res_func(self, model, pert_model): 265 | """ 266 | Compute the residual vector 267 | Args: 268 | model: domain vector 269 | pert_model: domain perturbation vector 270 | 271 | Returns: residual vector 272 | """ 273 | raise NotImplementedError("Implement pert_res_func for problem in the derived class!") 274 | 275 | def grad_func(self, model, residual): 276 | """ 277 | Compute the gradient vector from the residual (i.e., g = A' r = A'(Am - d)) 278 | Args: 279 | model: domain vector 280 | residual: residual vector 281 | 282 | Returns: gradient vector 283 | """ 284 | raise NotImplementedError("Implement gradf for problem in the derived class!") 285 | -------------------------------------------------------------------------------- /occamypy/torch/vector.py: -------------------------------------------------------------------------------- 1 | from math import sqrt 2 | 3 | import torch 4 | from numpy import ndarray 5 | 6 | from occamypy.vector.base import Vector 7 | from occamypy.torch.back_utils import get_device, get_device_name 8 | 9 | 10 | class VectorTorch(Vector): 11 | """ 12 | Vector class based on torch.Tensor 13 | 14 | Notes: 15 | tensors are stored in C-contiguous memory 16 | """ 17 | def __init__(self, in_content, device: int = None, *args, **kwargs): 18 | """ 19 | VectorTorch constructor 20 | 21 | Args: 22 | in_content: Vector, np.ndarray, torch.Tensor or tuple 23 | device: computation device (None for CPU, -1 for least used GPU) 24 | *args: list of arguments for Vector construction 25 | **kwargs: dict of arguments for Vector construction 26 | """ 27 | super(VectorTorch, self).__init__(*args, **kwargs) 28 | 29 | if isinstance(in_content, Vector): 30 | try: 31 | self.arr = torch.from_numpy(in_content.getNdArray()).contiguous() 32 | self.ax_info = in_content.ax_info 33 | except: 34 | raise UserWarning("Torch cannot handle the input array type") 35 | elif isinstance(in_content, ndarray): # Numpy array passed to constructor 36 | self.arr = torch.from_numpy(in_content).contiguous() 37 | elif isinstance(in_content, torch.Tensor): # Tensor passed to constructor 38 | self.arr = in_content.contiguous() 39 | elif isinstance(in_content, tuple): # Tuple size passed to constructor 40 | self.arr = torch.zeros(in_content) 41 | else: # Not supported type 42 | raise ValueError("ERROR! Input variable not currently supported!") 43 | 44 | self.setDevice(device) 45 | 46 | self.shape = tuple(self.arr.shape) # Number of elements per axis (tuple) 47 | self.ndim = self.arr.ndim # Number of axes (integer) 48 | self.size = self.arr.numel() # Total number of elements (integer) 49 | 50 | def _check_same_device(self, other): 51 | if not isinstance(self, VectorTorch): 52 | raise TypeError("The self vector has to be a VectorTorch") 53 | if not isinstance(other, VectorTorch): 54 | raise TypeError("The other vector has to be a VectorTorch") 55 | answer = self.device == other.device 56 | if not answer: 57 | raise Warning('The two vectors live in different devices: %s - %s' % (self.device, other.device)) 58 | return answer 59 | 60 | @property 61 | def device(self): 62 | try: 63 | return self.arr.device 64 | except AttributeError: 65 | return None 66 | 67 | def setDevice(self, devID): 68 | if isinstance(devID, int): 69 | self.arr = self.arr.to(get_device(devID)) 70 | elif isinstance(devID, torch.device): 71 | self.arr = self.arr.to(devID) 72 | else: 73 | ValueError("Device type not understood") 74 | 75 | @ property 76 | def deviceName(self): 77 | return get_device_name(self.device.index) 78 | 79 | def getNdArray(self): 80 | return self.arr 81 | 82 | def norm(self, N=2): 83 | norm = torch.linalg.norm(self.getNdArray().flatten(), ord=N) 84 | return norm.item() 85 | 86 | def zero(self): 87 | self.set(0) 88 | return self 89 | 90 | def max(self): 91 | max = self.getNdArray().max() 92 | return max.item() 93 | 94 | def min(self): 95 | min = self.getNdArray().min() 96 | return min.item() 97 | 98 | def set(self, val: float or int): 99 | self.getNdArray().fill_(val) 100 | return self 101 | 102 | def scale(self, sc: float or int): 103 | self.getNdArray()[:] *= sc 104 | return self 105 | 106 | def addbias(self, bias: float or int): 107 | self.getNdArray()[:] += bias 108 | return self 109 | 110 | def rand(self, low: float = -1., high: float = 1.): 111 | self.arr.uniform_(low, high) 112 | return self 113 | 114 | def randn(self, mean: float = 0., std: float = 1.): 115 | self.arr.normal_(mean, std) 116 | return self 117 | 118 | def clone(self): 119 | # If self is a Space vector, it is empty and it has only the shape, size and ndim attributes 120 | if self.getNdArray().numel() == 0: # this is the shape of tensor! 121 | vec_clone = VectorTorch(torch.zeros(self.shape).type(self.getNdArray().dtype), ax_info=self.ax_info.copy()) 122 | 123 | else: # self is a data vector, just clone 124 | vec_clone = VectorTorch(self.getNdArray().clone(), ax_info=self.ax_info.copy()) 125 | 126 | vec_clone.setDevice(self.device.index) 127 | return vec_clone 128 | 129 | def cloneSpace(self): 130 | vec_space = self.__class__(torch.empty(0).type(self.getNdArray().dtype)) 131 | vec_space.setDevice(self.device.index) 132 | vec_space.ax_info = self.ax_info 133 | # Cloning space of input vector 134 | vec_space.ndim = self.ndim 135 | vec_space.shape = self.shape 136 | vec_space.size = self.size 137 | return vec_space 138 | 139 | def checkSame(self, other): 140 | return self.shape == other.shape 141 | 142 | def abs(self): 143 | self.getNdArray().abs_() 144 | return self 145 | 146 | def sign(self): 147 | self.getNdArray().sign_() 148 | return self 149 | 150 | def reciprocal(self): 151 | self.getNdArray().reciprocal_() 152 | return self 153 | 154 | def maximum(self, other): 155 | if isinstance(other, (int, float)): 156 | # create a vector filled with the scalar value 157 | other = self.clone().set(other) 158 | 159 | if not self.checkSame(other): 160 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 161 | self.getNdArray()[:] = torch.maximum(self.getNdArray(), other.getNdArray()) 162 | return self 163 | 164 | def conj(self): 165 | self.getNdArray()[:] = torch.conj(self.getNdArray()) 166 | return self 167 | 168 | def transpose(self): 169 | other = VectorTorch(self.getNdArray().T) 170 | other.getNdArray()[:] = other.getNdArray()[:].contiguous() 171 | return other 172 | 173 | def pow(self, power): 174 | self.getNdArray()[:].pow_(power) 175 | return self 176 | 177 | def real(self): 178 | self.getNdArray()[:] = torch.real(self.getNdArray()) 179 | return self 180 | 181 | def imag(self): 182 | self.getNdArray()[:] = torch.imag(self.getNdArray()) 183 | return self 184 | 185 | def copy(self, other): 186 | # Checking whether the input is a vector or not 187 | if not isinstance(other, VectorTorch): 188 | raise TypeError("Provided input vector not a %s!" % self.whoami) 189 | # Checking dimensionality 190 | if not self.checkSame(other): 191 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 192 | # Element-wise copy of the input array 193 | self.getNdArray()[:].copy_(other.getNdArray()) 194 | return self 195 | 196 | def scaleAdd(self, other, sc1=1.0, sc2=1.0): 197 | # Checking whether the input is a vector or not 198 | if not isinstance(other, VectorTorch): 199 | raise TypeError("Provided input vector not a %s!" % self.whoami) 200 | # Checking dimensionality 201 | if not self.checkSame(other): 202 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 203 | # Performing scaling and addition 204 | self.scale(sc1) 205 | self.getNdArray()[:] += sc2 * other.getNdArray() 206 | return self 207 | 208 | def dot(self, other): 209 | # Checking whether the input is a vector or not 210 | if not isinstance(other, VectorTorch): 211 | raise TypeError("Provided input vector not a %s!" % self.whoami) 212 | # Checking size (must have same number of elements) 213 | if self.size != other.size: 214 | raise ValueError("Vector size mismatching: self = %d; other = %d" % (self.size, other.size)) 215 | # Checking dimensionality 216 | if not self.checkSame(other): 217 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 218 | return torch.vdot(self.getNdArray().flatten(), other.getNdArray().flatten()) 219 | 220 | def multiply(self, other): 221 | # Checking whether the input is a vector or not 222 | if not isinstance(other, VectorTorch): 223 | raise TypeError("Provided input vector not a %s!" % self.whoami) 224 | # Checking size (must have same number of elements) 225 | if self.size != other.size: 226 | raise ValueError("Vector size mismatching: self = %s; other = %s" % (self.size, other.size)) 227 | # Checking dimensionality 228 | if not self.checkSame(other): 229 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 230 | # Performing element-wise multiplication 231 | self.getNdArray()[:].multiply_(other.getNdArray()) 232 | return self 233 | 234 | def isDifferent(self, other): 235 | # Checking whether the input is a vector or not 236 | if not isinstance(other, VectorTorch): 237 | raise TypeError("Provided input vector not a %s!" % self.whoami) 238 | isDiff = not torch.equal(self.getNdArray(), other.getNdArray()) 239 | return isDiff 240 | 241 | def clip(self, low, high): 242 | if not isinstance(low, VectorTorch): 243 | raise TypeError("Provided input low vector not a %s!" % self.whoami) 244 | if not isinstance(high, VectorTorch): 245 | raise TypeError("Provided input high vector not a %s!" % self.whoami) 246 | self.getNdArray()[:] = torch.minimum(torch.maximum(low.getNdArray(), self.getNdArray()), high.getNdArray()) 247 | return self 248 | 249 | def plot(self): 250 | return self.getNdArray().detach().cpu().numpy() 251 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/acoustic/operators.py: -------------------------------------------------------------------------------- 1 | from devito import Eq, Operator, Function, TimeFunction, Inc, solve, sign 2 | from devito.symbolics import retrieve_functions, INT 3 | from .. import PointSource, Receiver 4 | 5 | 6 | def freesurface(model, eq): 7 | """ 8 | Generate the stencil that mirrors the field as a free surface modeling for 9 | the acoustic wave equation. 10 | 11 | Parameters 12 | ---------- 13 | model : Model 14 | Physical model. 15 | eq : Eq 16 | Time-stepping stencil (time update) to mirror at the freesurface. 17 | """ 18 | lhs, rhs = eq.evaluate.args 19 | # Get vertical dimension and corresponding subdimension 20 | zfs = model.grid.subdomains['fsdomain'].dimensions[-1] 21 | z = zfs.parent 22 | 23 | # Functions present in the stencil 24 | funcs = retrieve_functions(rhs) 25 | mapper = {} 26 | # Antisymmetric mirror at negative indices 27 | # TODO: Make a proper "mirror_indices" tool function 28 | for f in funcs: 29 | zind = f.indices[-1] 30 | if (zind - z).as_coeff_Mul()[0] < 0: 31 | s = sign(zind.subs({z: zfs, z.spacing: 1})) 32 | mapper.update({f: s * f.subs({zind: INT(abs(zind))})}) 33 | return Eq(lhs, rhs.subs(mapper), subdomain=model.grid.subdomains['fsdomain']) 34 | 35 | 36 | def laplacian(field, model, kernel): 37 | """ 38 | Spatial discretization for the isotropic acoustic wave equation. For a 4th 39 | order in time formulation, the 4th order time derivative is replaced by a 40 | double laplacian: 41 | H = (laplacian + s**2/12 laplacian(1/m*laplacian)) 42 | 43 | Parameters 44 | ---------- 45 | field : TimeFunction 46 | The computed solution. 47 | model : Model 48 | Physical model. 49 | """ 50 | if kernel not in ['OT2', 'OT4']: 51 | raise ValueError("Unrecognized kernel") 52 | s = model.grid.time_dim.spacing 53 | biharmonic = field.biharmonic(1/model.m) if kernel == 'OT4' else 0 54 | return field.laplace + s**2/12 * biharmonic 55 | 56 | 57 | def iso_stencil(field, model, kernel, **kwargs): 58 | """ 59 | Stencil for the acoustic isotropic wave-equation: 60 | u.dt2 - H + damp*u.dt = 0. 61 | 62 | Parameters 63 | ---------- 64 | field : TimeFunction 65 | The computed solution. 66 | model : Model 67 | Physical model. 68 | kernel : str, optional 69 | Type of discretization, 'OT2' or 'OT4'. 70 | q : TimeFunction, Function or float 71 | Full-space/time source of the wave-equation. 72 | forward : bool, optional 73 | Whether to propagate forward (True) or backward (False) in time. 74 | """ 75 | # Forward or backward 76 | forward = kwargs.get('forward', True) 77 | # Define time step to be updated 78 | unext = field.forward if forward else field.backward 79 | udt = field.dt if forward else field.dt.T 80 | # Get the spacial FD 81 | lap = laplacian(field, model, kernel) 82 | # Get source 83 | q = kwargs.get('q', 0) 84 | # Define PDE and update rule 85 | eq_time = solve(model.m * field.dt2 - lap - q + model.damp * udt, unext) 86 | 87 | # Time-stepping stencil. 88 | eqns = [Eq(unext, eq_time, subdomain=model.grid.subdomains['physdomain'])] 89 | 90 | # Add free surface 91 | if model.fs: 92 | eqns.append(freesurface(model, Eq(unext, eq_time))) 93 | return eqns 94 | 95 | 96 | def ForwardOperator(model, geometry, space_order=4, 97 | save=False, kernel='OT2', **kwargs): 98 | """ 99 | Construct a forward modelling operator in an acoustic medium. 100 | 101 | Parameters 102 | ---------- 103 | model : Model 104 | Object containing the physical parameters. 105 | geometry : AcquisitionGeometry 106 | Geometry object that contains the source (SparseTimeFunction) and 107 | receivers (SparseTimeFunction) and their position. 108 | space_order : int, optional 109 | Space discretization order. 110 | save : int or Buffer, optional 111 | Saving flag, True saves all time steps. False saves three timesteps. 112 | Defaults to False. 113 | kernel : str, optional 114 | Type of discretization, 'OT2' or 'OT4'. 115 | """ 116 | m = model.m 117 | 118 | # Create symbols for forward wavefield, source and receivers 119 | u = TimeFunction(name='u', grid=model.grid, 120 | save=geometry.nt if save else None, 121 | time_order=2, space_order=space_order) 122 | src = PointSource(name='src', grid=geometry.grid, time_range=geometry.time_axis, 123 | npoint=geometry.nsrc) 124 | 125 | rec = Receiver(name='rec', grid=geometry.grid, time_range=geometry.time_axis, 126 | npoint=geometry.nrec) 127 | 128 | s = model.grid.stepping_dim.spacing 129 | eqn = iso_stencil(u, model, kernel) 130 | 131 | # Construct expression to inject source values 132 | src_term = src.inject(field=u.forward, expr=src * s**2 / m) 133 | 134 | # Create interpolation expression for receivers 135 | rec_term = rec.interpolate(expr=u) 136 | 137 | # Substitute spacing terms to reduce flops 138 | return Operator(eqn + src_term + rec_term, subs=model.spacing_map, 139 | name='Forward', **kwargs) 140 | 141 | 142 | def AdjointOperator(model, geometry, space_order=4, 143 | kernel='OT2', **kwargs): 144 | """ 145 | Construct an adjoint modelling operator in an acoustic media. 146 | 147 | Parameters 148 | ---------- 149 | model : Model 150 | Object containing the physical parameters. 151 | geometry : AcquisitionGeometry 152 | Geometry object that contains the source (SparseTimeFunction) and 153 | receivers (SparseTimeFunction) and their position. 154 | space_order : int, optional 155 | Space discretization order. 156 | kernel : str, optional 157 | Type of discretization, 'OT2' or 'OT4'. 158 | """ 159 | m = model.m 160 | 161 | v = TimeFunction(name='v', grid=model.grid, save=None, 162 | time_order=2, space_order=space_order) 163 | srca = PointSource(name='srca', grid=model.grid, time_range=geometry.time_axis, 164 | npoint=geometry.nsrc) 165 | rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, 166 | npoint=geometry.nrec) 167 | 168 | s = model.grid.stepping_dim.spacing 169 | eqn = iso_stencil(v, model, kernel, forward=False) 170 | 171 | # Construct expression to inject receiver values 172 | receivers = rec.inject(field=v.backward, expr=rec * s**2 / m) 173 | 174 | # Create interpolation expression for the adjoint-source 175 | source_a = srca.interpolate(expr=v) 176 | 177 | # Substitute spacing terms to reduce flops 178 | return Operator(eqn + receivers + source_a, subs=model.spacing_map, 179 | name='Adjoint', **kwargs) 180 | 181 | 182 | def GradientOperator(model, geometry, space_order=4, save=True, 183 | kernel='OT2', **kwargs): 184 | """ 185 | Construct a gradient operator in an acoustic media. 186 | 187 | Parameters 188 | ---------- 189 | model : Model 190 | Object containing the physical parameters. 191 | geometry : AcquisitionGeometry 192 | Geometry object that contains the source (SparseTimeFunction) and 193 | receivers (SparseTimeFunction) and their position. 194 | space_order : int, optional 195 | Space discretization order. 196 | save : int or Buffer, optional 197 | Option to store the entire (unrolled) wavefield. 198 | kernel : str, optional 199 | Type of discretization, centered or shifted. 200 | """ 201 | m = model.m 202 | 203 | # Gradient symbol and wavefield symbols 204 | grad = Function(name='grad', grid=model.grid) 205 | u = TimeFunction(name='u', grid=model.grid, save=geometry.nt if save 206 | else None, time_order=2, space_order=space_order) 207 | v = TimeFunction(name='v', grid=model.grid, save=None, 208 | time_order=2, space_order=space_order) 209 | rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, 210 | npoint=geometry.nrec) 211 | 212 | s = model.grid.stepping_dim.spacing 213 | eqn = iso_stencil(v, model, kernel, forward=False) 214 | 215 | if kernel == 'OT2': 216 | gradient_update = Inc(grad, - u * v.dt2) 217 | elif kernel == 'OT4': 218 | gradient_update = Inc(grad, - u * v.dt2 - s**2 / 12.0 * u.biharmonic(m**(-2)) * v) 219 | # Add expression for receiver injection 220 | receivers = rec.inject(field=v.backward, expr=rec * s**2 / m) 221 | 222 | # Substitute spacing terms to reduce flops 223 | return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map, 224 | name='Gradient', **kwargs) 225 | 226 | 227 | def BornOperator(model, geometry, space_order=4, 228 | kernel='OT2', **kwargs): 229 | """ 230 | Construct an Linearized Born operator in an acoustic media. 231 | 232 | Parameters 233 | ---------- 234 | model : Model 235 | Object containing the physical parameters. 236 | geometry : AcquisitionGeometry 237 | Geometry object that contains the source (SparseTimeFunction) and 238 | receivers (SparseTimeFunction) and their position. 239 | space_order : int, optional 240 | Space discretization order. 241 | kernel : str, optional 242 | Type of discretization, centered or shifted. 243 | """ 244 | m = model.m 245 | 246 | # Create source and receiver symbols 247 | src = Receiver(name='src', grid=model.grid, time_range=geometry.time_axis, 248 | npoint=geometry.nsrc) 249 | 250 | rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, 251 | npoint=geometry.nrec) 252 | 253 | # Create wavefields and a dm field 254 | u = TimeFunction(name="u", grid=model.grid, save=None, 255 | time_order=2, space_order=space_order) 256 | U = TimeFunction(name="U", grid=model.grid, save=None, 257 | time_order=2, space_order=space_order) 258 | dm = Function(name="dm", grid=model.grid, space_order=0) 259 | 260 | s = model.grid.stepping_dim.spacing 261 | eqn1 = iso_stencil(u, model, kernel) 262 | eqn2 = iso_stencil(U, model, kernel, q=-dm*u.dt2) 263 | 264 | # Add source term expression for u 265 | source = src.inject(field=u.forward, expr=src * s**2 / m) 266 | 267 | # Create receiver interpolation expression from U 268 | receivers = rec.interpolate(expr=U) 269 | 270 | # Substitute spacing terms to reduce flops 271 | return Operator(eqn1 + source + eqn2 + receivers, subs=model.spacing_map, 272 | name='Born', **kwargs) 273 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/acoustic/wavesolver.py: -------------------------------------------------------------------------------- 1 | from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator 2 | from devito.tools import memoized_meth 3 | from .operators import ( 4 | ForwardOperator, AdjointOperator, GradientOperator, BornOperator 5 | ) 6 | from pyrevolve import Revolver 7 | 8 | 9 | class AcousticWaveSolver(object): 10 | """ 11 | Solver object that provides operators for devitoseismic inversion problems 12 | and encapsulates the time and space discretization for a given problem 13 | setup. 14 | 15 | Parameters 16 | ---------- 17 | model : Model 18 | Physical model with domain parameters. 19 | geometry : AcquisitionGeometry 20 | Geometry object that contains the source (SparseTimeFunction) and 21 | receivers (SparseTimeFunction) and their position. 22 | kernel : str, optional 23 | Type of discretization, centered or shifted. 24 | space_order: int, optional 25 | Order of the spatial stencil discretisation. Defaults to 4. 26 | """ 27 | def __init__(self, model, geometry, kernel='OT2', space_order=4, **kwargs): 28 | self.model = model 29 | self.model._initialize_bcs(bcs="damp") 30 | self.geometry = geometry 31 | 32 | assert self.model.grid == geometry.grid 33 | 34 | self.space_order = space_order 35 | self.kernel = kernel 36 | 37 | # Cache compiler options 38 | self._kwargs = kwargs 39 | 40 | @property 41 | def dt(self): 42 | # Time step can be \sqrt{3}=1.73 bigger with 4th order 43 | if self.kernel == 'OT4': 44 | return self.model.dtype(1.73 * self.model.critical_dt) 45 | return self.model.critical_dt 46 | 47 | @memoized_meth 48 | def op_fwd(self, save=None): 49 | """Cached operator for forward runs with buffered wavefield""" 50 | return ForwardOperator(self.model, save=save, geometry=self.geometry, 51 | kernel=self.kernel, space_order=self.space_order, 52 | **self._kwargs) 53 | 54 | @memoized_meth 55 | def op_adj(self): 56 | """Cached operator for adjoint runs""" 57 | return AdjointOperator(self.model, save=None, geometry=self.geometry, 58 | kernel=self.kernel, space_order=self.space_order, 59 | **self._kwargs) 60 | 61 | @memoized_meth 62 | def op_grad(self, save=True): 63 | """Cached operator for gradient runs""" 64 | return GradientOperator(self.model, save=save, geometry=self.geometry, 65 | kernel=self.kernel, space_order=self.space_order, 66 | **self._kwargs) 67 | 68 | @memoized_meth 69 | def op_born(self): 70 | """Cached operator for born runs""" 71 | return BornOperator(self.model, save=None, geometry=self.geometry, 72 | kernel=self.kernel, space_order=self.space_order, 73 | **self._kwargs) 74 | 75 | def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs): 76 | """ 77 | Forward modelling function that creates the necessary 78 | data objects for running a forward modelling operator. 79 | 80 | Parameters 81 | ---------- 82 | src : SparseTimeFunction or array_like, optional 83 | Time series data for the injected source term. 84 | rec : SparseTimeFunction or array_like, optional 85 | The interpolated receiver data. 86 | u : TimeFunction, optional 87 | Stores the computed wavefield. 88 | model : Model, optional 89 | Object containing the physical parameters. 90 | vp : Function or float, optional 91 | The time-constant velocity. 92 | save : bool, optional 93 | Whether or not to save the entire (unrolled) wavefield. 94 | 95 | Returns 96 | ------- 97 | Receiver, wavefield and performance summary 98 | """ 99 | # Source term is read-only, so re-use the default 100 | src = src or self.geometry.src 101 | # Create a new receiver object to store the result 102 | rec = rec or self.geometry.rec 103 | 104 | # Create the forward wavefield if not provided 105 | u = u or TimeFunction(name='u', grid=self.model.grid, 106 | save=self.geometry.nt if save else None, 107 | time_order=2, space_order=self.space_order) 108 | 109 | model = model or self.model 110 | # Pick vp from model unless explicitly provided 111 | kwargs.update(model.physical_params(**kwargs)) 112 | 113 | # Execute operator and return wavefield and receiver data 114 | summary = self.op_fwd(save).apply(src=src, rec=rec, u=u, 115 | dt=kwargs.pop('dt', self.dt), **kwargs) 116 | 117 | return rec, u, summary 118 | 119 | def adjoint(self, rec, srca=None, v=None, model=None, **kwargs): 120 | """ 121 | Adjoint modelling function that creates the necessary 122 | data objects for running an adjoint modelling operator. 123 | 124 | Parameters 125 | ---------- 126 | rec : SparseTimeFunction or array-like 127 | The receiver data. Please note that 128 | these act as the source term in the adjoint run. 129 | srca : SparseTimeFunction or array-like 130 | The resulting data for the interpolated at the 131 | original source location. 132 | v: TimeFunction, optional 133 | The computed wavefield. 134 | model : Model, optional 135 | Object containing the physical parameters. 136 | vp : Function or float, optional 137 | The time-constant velocity. 138 | 139 | Returns 140 | ------- 141 | Adjoint source, wavefield and performance summary. 142 | """ 143 | # Create a new adjoint source and receiver symbol 144 | srca = srca or self.geometry.new_src(name='srca', src_type=None) 145 | 146 | # Create the adjoint wavefield if not provided 147 | v = v or TimeFunction(name='v', grid=self.model.grid, 148 | time_order=2, space_order=self.space_order) 149 | 150 | model = model or self.model 151 | # Pick vp from model unless explicitly provided 152 | kwargs.update(model.physical_params(**kwargs)) 153 | 154 | # Execute operator and return wavefield and receiver data 155 | summary = self.op_adj().apply(srca=srca, rec=rec, v=v, 156 | dt=kwargs.pop('dt', self.dt), **kwargs) 157 | return srca, v, summary 158 | 159 | def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None, 160 | checkpointing=False, **kwargs): 161 | """ 162 | Gradient modelling function for computing the adjoint of the 163 | Linearized Born modelling function, ie. the action of the 164 | Jacobian adjoint on an input data. 165 | 166 | Parameters 167 | ---------- 168 | rec : SparseTimeFunction 169 | Receiver data. 170 | u : TimeFunction 171 | Full wavefield `u` (created with save=True). 172 | v : TimeFunction, optional 173 | Stores the computed wavefield. 174 | grad : Function, optional 175 | Stores the gradient field. 176 | model : Model, optional 177 | Object containing the physical parameters. 178 | vp : Function or float, optional 179 | The time-constant velocity. 180 | 181 | Returns 182 | ------- 183 | Gradient field and performance summary. 184 | """ 185 | dt = kwargs.pop('dt', self.dt) 186 | # Gradient symbol 187 | grad = grad or Function(name='grad', grid=self.model.grid) 188 | 189 | # Create the forward wavefield 190 | v = v or TimeFunction(name='v', grid=self.model.grid, 191 | time_order=2, space_order=self.space_order) 192 | 193 | model = model or self.model 194 | # Pick vp from model unless explicitly provided 195 | kwargs.update(model.physical_params(**kwargs)) 196 | 197 | if checkpointing: 198 | u = TimeFunction(name='u', grid=self.model.grid, 199 | time_order=2, space_order=self.space_order) 200 | cp = DevitoCheckpoint([u]) 201 | n_checkpoints = None 202 | wrap_fw = CheckpointOperator(self.op_fwd(save=False), 203 | src=src or self.geometry.src, 204 | u=u, dt=dt, **kwargs) 205 | wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v, 206 | rec=rec, dt=dt, grad=grad, **kwargs) 207 | 208 | # Run forward 209 | wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0] - 2) 210 | wrp.apply_forward() 211 | summary = wrp.apply_reverse() 212 | else: 213 | summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt, 214 | **kwargs) 215 | return grad, summary 216 | 217 | def jacobian(self, dmin, src=None, rec=None, u=None, U=None, model=None, **kwargs): 218 | """ 219 | Linearized Born modelling function that creates the necessary 220 | data objects for running an adjoint modelling operator. 221 | 222 | Parameters 223 | ---------- 224 | src : SparseTimeFunction or array_like, optional 225 | Time series data for the injected source term. 226 | rec : SparseTimeFunction or array_like, optional 227 | The interpolated receiver data. 228 | u : TimeFunction, optional 229 | The forward wavefield. 230 | U : TimeFunction, optional 231 | The linearized wavefield. 232 | model : Model, optional 233 | Object containing the physical parameters. 234 | vp : Function or float, optional 235 | The time-constant velocity. 236 | """ 237 | # Source term is read-only, so re-use the default 238 | src = src or self.geometry.src 239 | # Create a new receiver object to store the result 240 | rec = rec or self.geometry.rec 241 | 242 | # Create the forward wavefields u and U if not provided 243 | u = u or TimeFunction(name='u', grid=self.model.grid, 244 | time_order=2, space_order=self.space_order) 245 | U = U or TimeFunction(name='U', grid=self.model.grid, 246 | time_order=2, space_order=self.space_order) 247 | 248 | model = model or self.model 249 | # Pick vp from model unless explicitly provided 250 | kwargs.update(model.physical_params(**kwargs)) 251 | 252 | # Execute operator and return wavefield and receiver data 253 | summary = self.op_born().apply(dm=dmin, u=u, U=U, src=src, rec=rec, 254 | dt=kwargs.pop('dt', self.dt), **kwargs) 255 | return rec, u, U, summary 256 | 257 | # Backward compatibility 258 | born = jacobian 259 | gradient = jacobian_adjoint 260 | -------------------------------------------------------------------------------- /occamypy/cupy/vector.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | from sys import version_info 3 | 4 | import cupy as cp 5 | import numpy as np 6 | from GPUtil import getGPUs, getFirstAvailable 7 | 8 | from occamypy.vector.base import Vector 9 | from occamypy.numpy.vector import VectorNumpy 10 | 11 | 12 | class VectorCupy(Vector): 13 | """Vector class based on cupy.ndarray""" 14 | 15 | def __init__(self, in_content, device=None, *args, **kwargs): 16 | """ 17 | VectorCupy constructor 18 | 19 | Args: 20 | in_content: numpy.ndarray, cupy.ndarray, tuple or VectorNumpy 21 | device: computation device (None for CPU, -1 for least used GPU) 22 | *args: list of arguments for Vector construction 23 | **kwargs: dict of arguments for Vector construction 24 | """ 25 | if isinstance(in_content, cp.ndarray) or isinstance(in_content, np.ndarray): 26 | if cp.isfortran(in_content): 27 | raise TypeError('Input array not a C contiguous array!') 28 | self.arr = cp.array(in_content, copy=False) 29 | elif isinstance(in_content, tuple): # Tuple size passed to constructor 30 | # self.arr = cp.zeros(tuple(reversed(in_content))) 31 | self.arr = cp.zeros(in_content) 32 | elif isinstance(in_content, VectorNumpy): 33 | self.arr = in_content.getNdArray().copy() 34 | self.ax_info = in_content.ax_info 35 | else: # Not supported type 36 | raise ValueError("ERROR! Input variable not currently supported!") 37 | 38 | self.setDevice(device) 39 | 40 | super(VectorCupy, self).__init__(*args, **kwargs) 41 | self.shape = self.arr.shape # Number of elements per axis (tuple) 42 | self.ndim = self.arr.ndim # Number of axes (integer) 43 | self.size = self.arr.size # Total number of elements (integer) 44 | 45 | def _check_same_device(self, other): 46 | if not isinstance(self, VectorCupy): 47 | raise TypeError("self vector has to be a VectorCupy") 48 | if not isinstance(other, VectorCupy): 49 | raise TypeError("other vector has to be a VectorCupy") 50 | answer = self.device == other.device 51 | if not answer: 52 | raise Warning('The two vectors live in different devices: %s - %s' % (self.device, other.device)) 53 | return answer 54 | 55 | @property 56 | def device(self): 57 | try: 58 | return self.arr.device 59 | except AttributeError: 60 | return None 61 | 62 | def setDevice(self, devID=0): 63 | if devID is not None: # Move to GPU 64 | if devID == -1: 65 | devID = getFirstAvailable(order='memory')[0] 66 | try: 67 | with cp.cuda.Device(devID): 68 | self.arr = cp.asarray(self.arr) 69 | except AttributeError: 70 | self.arr.device = None 71 | else: # move to CPU 72 | if self.device is None: # already on CPU 73 | pass 74 | else: 75 | self.arr = cp.asnumpy(self.arr) 76 | 77 | @property 78 | def deviceName(self): 79 | if self.device is None: 80 | return "CPU" 81 | else: 82 | name = getGPUs()[self.device.id].name 83 | return "GPU %d - %s" % (self.device.id, name) 84 | 85 | def getNdArray(self): 86 | return self.arr 87 | 88 | def norm(self, N=2): 89 | return cp.linalg.norm(self.getNdArray().ravel(), ord=N) 90 | 91 | def zero(self): 92 | self.getNdArray().fill(0) 93 | return self 94 | 95 | def max(self): 96 | return self.getNdArray().max() 97 | 98 | def min(self): 99 | return self.getNdArray().min() 100 | 101 | def set(self, val): 102 | self.getNdArray().fill(val) 103 | return self 104 | 105 | def scale(self, sc): 106 | self.getNdArray()[:] *= sc 107 | return self 108 | 109 | def addbias(self, bias): 110 | self.getNdArray()[:] += bias 111 | return self 112 | 113 | def rand(self, low: float = -1., high: float = 1.): 114 | self.zero() 115 | self[:] += cp.random.uniform(low=low, high=high, size=self.shape) 116 | return self 117 | 118 | def randn(self, mean: float = 0., std: float = 1.): 119 | self.zero() 120 | self[:] += cp.random.normal(loc=mean, scale=std, size=self.shape) 121 | return self 122 | 123 | def clone(self): 124 | vec_clone = deepcopy(self) # Deep clone of vector 125 | # Checking if a vector space was provided 126 | if vec_clone.getNdArray().size == 0: 127 | vec_clone.arr = cp.zeros(vec_clone.shape, dtype=self.getNdArray().dtype) 128 | return vec_clone 129 | 130 | def cloneSpace(self): 131 | arr = cp.empty(0, dtype=self.getNdArray().dtype) 132 | vec_space = VectorCupy(arr) 133 | vec_space.ax_info = self.ax_info 134 | # Cloning space of input vector 135 | vec_space.shape = self.shape 136 | vec_space.ndim = self.ndim 137 | vec_space.size = self.size 138 | return vec_space 139 | 140 | def checkSame(self, other): 141 | return self.shape == other.shape 142 | 143 | def abs(self): 144 | self.getNdArray()[:] = cp.abs(self.getNdArray()) 145 | return self 146 | 147 | def sign(self): 148 | self.getNdArray()[:] = cp.sign(self.getNdArray()) 149 | return self 150 | 151 | def reciprocal(self): 152 | self.getNdArray()[:] = 1. / self.getNdArray() 153 | return self 154 | 155 | def maximum(self, other): 156 | if cp.isscalar(other): 157 | self.getNdArray()[:] = cp.maximum(self.getNdArray(), other) 158 | return self 159 | elif isinstance(other, VectorCupy): 160 | if not self.checkSame(other): 161 | raise ValueError('Dimensionality not equal: self = %s; other = %s' % (self.shape, other.shape)) 162 | if not self._check_same_device(other): 163 | raise ValueError('Provided input has to live in the same device') 164 | self.getNdArray()[:] = cp.maximum(self.getNdArray(), other.getNdArray()) 165 | return self 166 | else: 167 | raise TypeError("Provided input has to be either a scalar or a %s!" % self.whoami) 168 | 169 | def conj(self): 170 | self.getNdArray()[:] = cp.conj(self.getNdArray()) 171 | return self 172 | 173 | def pow(self, power): 174 | self.getNdArray()[:] = self.getNdArray() ** power 175 | return self 176 | 177 | def real(self): 178 | self.getNdArray()[:] = self.getNdArray().real 179 | return self 180 | 181 | def imag(self,): 182 | self.getNdArray()[:] = self.getNdArray().imag 183 | return self 184 | 185 | def copy(self, other): 186 | # Checking whether the input is a vector or not 187 | if not isinstance(other, VectorCupy): 188 | raise TypeError("Provided input vector not a %s!" % self.whoami) 189 | # Checking dimensionality 190 | if not self.checkSame(other): 191 | raise ValueError("Dimensionality not equal: self = %s; other = %s" % (self.shape, other.shape)) 192 | # Element-wise copy of the input array 193 | self.getNdArray()[:] = other.getNdArray() 194 | return self 195 | 196 | def scaleAdd(self, other, sc1=1.0, sc2=1.0): 197 | # Checking whether the input is a vector or not 198 | if not isinstance(other, VectorCupy): 199 | raise TypeError("Provided input vector not a %s!" % self.whoami) 200 | # Checking dimensionality 201 | if not self.checkSame(other): 202 | raise ValueError("Dimensionality not equal: self = %s; other = %s" % (self.shape, other.shape)) 203 | # Performing scaling and addition 204 | if not self._check_same_device(other): 205 | raise ValueError('Provided input has to live in the same device') 206 | self.getNdArray()[:] = sc1 * self.getNdArray() + sc2 * other.getNdArray() 207 | return self 208 | 209 | def dot(self, other): 210 | # Checking whether the input is a vector or not 211 | if not isinstance(other, VectorCupy): 212 | raise TypeError("Provided input vector not a %s!" % self.whoami) 213 | # Checking size (must have same number of elements) 214 | if self.size != other.size: 215 | raise ValueError("Vector size mismatching: self = %d; other = %d" % (self.size, other.size)) 216 | # Checking dimensionality 217 | if not self.checkSame(other): 218 | raise ValueError("Dimensionality not equal: self = %s; other = %s" % (self.shape, other.shape)) 219 | if not self._check_same_device(other): 220 | raise ValueError('Provided input has to live in the same device') 221 | return cp.vdot(self.getNdArray().flatten(), other.getNdArray().flatten()) 222 | 223 | def multiply(self, other): 224 | # Checking whether the input is a vector or not 225 | if not isinstance(other, VectorCupy): 226 | raise TypeError("Provided input vector not a %s!" % self.whoami) 227 | # Checking size (must have same number of elements) 228 | if self.size != other.size: 229 | raise ValueError("Vector size mismatching: self = %d; other = %d" % (self.size, other.size)) 230 | # Checking dimensionality 231 | if not self.checkSame(other): 232 | raise ValueError("Dimensionality not equal: self = %s; other = %s" % (self.shape, other.shape)) 233 | # Performing element-wise multiplication 234 | if not self._check_same_device(other): 235 | raise ValueError('Provided input has to live in the same device') 236 | self.getNdArray()[:] = cp.multiply(self.getNdArray(), other.getNdArray()) 237 | return self 238 | 239 | def isDifferent(self, other): 240 | # Checking whether the input is a vector or not 241 | if not isinstance(other, VectorCupy): 242 | raise TypeError("Provided input vector not a %s!" % self.whoami) 243 | if not self._check_same_device(other): 244 | raise ValueError('Provided input has to live in the same device') 245 | # Using Hash table for python2 and numpy built-in function array_equal otherwise 246 | if version_info[0] == 2: 247 | # First make both array buffers read-only 248 | self.arr.flags.writeable = False 249 | other.arr.flags.writeable = False 250 | chcksum1 = hash(self.getNdArray().data) 251 | chcksum2 = hash(other.getNdArray().data) 252 | # Remake array buffers writable 253 | self.arr.flags.writeable = True 254 | other.arr.flags.writeable = True 255 | isDiff = (chcksum1 != chcksum2) 256 | else: 257 | isDiff = (not cp.equal(self.getNdArray(), other.getNdArray()).all()) 258 | return isDiff 259 | 260 | def clip(self, low, high): 261 | if not isinstance(low, VectorCupy): 262 | raise TypeError("Provided input low vector not a %s!" % self.whoami) 263 | if not isinstance(high, VectorCupy): 264 | raise TypeError("Provided input high vector not a %s!" % self.whoami) 265 | self.getNdArray()[:] = cp.minimum(cp.maximum(low.getNdArray(), self.getNdArray()), high.getNdArray()) 266 | return self 267 | 268 | def plot(self): 269 | return self.getNdArray().get() 270 | -------------------------------------------------------------------------------- /occamypy/dask/utils.py: -------------------------------------------------------------------------------- 1 | import atexit 2 | import json 3 | import os 4 | import random 5 | import socket 6 | import subprocess 7 | from time import time 8 | # Adding following import for fixing compatibility issue with python3.9 9 | # (https://github.com/dask/distributed/issues/4168) 10 | import multiprocessing.popen_spawn_posix 11 | 12 | DEVNULL = open(os.devnull, "wb") 13 | import dask.distributed as daskD 14 | from dask_jobqueue import PBSCluster, LSFCluster, SLURMCluster 15 | from dask_kubernetes import KubeCluster, make_pod_spec 16 | 17 | 18 | def get_tcp_info(filename): 19 | """ 20 | Obtain the scheduler TCP information 21 | 22 | Args: 23 | filename: file to read 24 | 25 | Returns: 26 | TCP address 27 | """ 28 | tcp_info = None 29 | with open(filename) as json_file: 30 | data = json.load(json_file) 31 | if "address" in data: 32 | tcp_info = data["address"] 33 | return tcp_info 34 | 35 | 36 | def create_hostnames(machine_names, Nworkers): 37 | """ 38 | Create hostnames variables (i.e., list of IP addresses) from machine names and number of wokers per machine 39 | 40 | Args: 41 | machine_names: list of host names 42 | Nworkers: list of client workers 43 | 44 | Returns: 45 | list of host IP addresses 46 | """ 47 | ip_adds = [] 48 | for host in machine_names: 49 | ip_adds.append(socket.gethostbyname(host)) 50 | 51 | if len(Nworkers) != len(ip_adds): 52 | raise ValueError("Lenght of number of workers (%s) not consistent with number of machines available (%s)" 53 | % (len(Nworkers), len(ip_adds))) 54 | 55 | hostnames = [] 56 | for idx, ip in enumerate(ip_adds): 57 | hostnames += [ip] * Nworkers[idx] 58 | return hostnames 59 | 60 | 61 | def client_startup(cluster, n_jobs: int, total_workers: int): 62 | """ 63 | Start a dask client 64 | 65 | Args: 66 | cluster: Dask cluster 67 | n_jobs: number of jobs to submit to the cluster 68 | total_workers: number of total workers in the cluster 69 | 70 | Returns: 71 | DaskClient instance, list of workers ID 72 | """ 73 | if n_jobs <= 0: 74 | raise ValueError("n_jobs must equal or greater than 1!") 75 | if isinstance(cluster, daskD.LocalCluster) or isinstance(cluster, KubeCluster): 76 | cluster.scale(n_jobs) 77 | else: 78 | cluster.scale(jobs=n_jobs) 79 | # Creating dask Client 80 | client = daskD.Client(cluster) 81 | workers = 0 82 | t0 = time() 83 | while workers < total_workers: 84 | workers = len(client.get_worker_logs().keys()) 85 | # If the number of workers is not reached in 5 minutes raise exception 86 | if time() - t0 > 300.0: 87 | raise SystemError("Dask could not start the requested workers within 5 minutes!" 88 | "Try different n_jobs.") 89 | WorkerIds = list(client.get_worker_logs().keys()) 90 | return client, WorkerIds 91 | 92 | 93 | class DaskClient: 94 | """ 95 | Dask Client to be used with Dask vectors and operators 96 | 97 | Notes: 98 | The Kubernetes pods are created using the Docker image "ettore88/occamypy:devel". 99 | To change the image to be use, provide the item image within the kube_params dictionary. 100 | """ 101 | 102 | def __init__(self, **kwargs): 103 | """ 104 | DaskClient constructor. 105 | 106 | Args: 107 | 1) Cluster with shared file system and ssh capability 108 | 109 | hostnames (list) [None]: host names or IP addresses of the machines that the user wants to use in their cluster/client (First hostname will be running the scheduler!) 110 | scheduler_file_prefix (str): prefix to used to create dask scheduler-file. 111 | logging (bool) [True]: whether to log scheduler and worker stdout to files within dask_logs folder 112 | Must be a mounted path on all the machines. Necessary if hostnames are provided [$HOME/scheduler-] 113 | 114 | 2) Local cluster 115 | local_params (dict) [None]: Local Cluster options (see help(LocalCluster) for help) 116 | n_wrks (int) [1]: number of workers to start 117 | 118 | 3) PBS cluster 119 | pbs_params (dict) [None]: PBS Cluster options (see help(PBSCluster) for help) 120 | n_jobs (int): number of jobs to be submitted to the cluster 121 | n_wrks (int) [1]: number of workers per job 122 | 123 | 4) LSF cluster 124 | lfs_params (dict) [None]: LSF Cluster options (see help(LSFCluster) for help) 125 | n_jobs (int): number of jobs to be submitted to the cluster 126 | n_wrks (int) [1]: number of workers per job 127 | 128 | 5) SLURM cluster 129 | slurm_params (dict) [None]: SLURM Cluster options (see help(SLURMCluster) for help) 130 | n_jobs (int): number of jobs to be submitted to the cluster 131 | n_wrks (int) [1]: number of workers per job 132 | 133 | 6) Kubernetes cluster 134 | kube_params (dict): KubeCluster options 135 | (see help(KubeCluster) and help(make_pod_spec) for help) [None] 136 | n_wrks (int) [1]: number of workers per job 137 | """ 138 | hostnames = kwargs.get("hostnames", None) 139 | local_params = kwargs.get("local_params", None) 140 | pbs_params = kwargs.get("pbs_params", None) 141 | lsf_params = kwargs.get("lsf_params", None) 142 | slurm_params = kwargs.get("slurm_params", None) 143 | kube_params = kwargs.get("kube_params", None) 144 | logging = kwargs.get("logging", True) 145 | 146 | ClusterInit = None 147 | cluster_params = None 148 | if local_params: 149 | cluster_params = local_params 150 | ClusterInit = daskD.LocalCluster 151 | elif pbs_params: 152 | cluster_params = pbs_params 153 | ClusterInit = PBSCluster 154 | elif lsf_params: 155 | cluster_params = lsf_params 156 | ClusterInit = LSFCluster 157 | elif slurm_params: 158 | cluster_params = slurm_params 159 | ClusterInit = SLURMCluster 160 | # Checking interface to be used 161 | if hostnames: 162 | if not isinstance(hostnames, list): 163 | raise ValueError("User must provide a list with host names") 164 | scheduler_file_prefix = kwargs.get("scheduler_file_prefix", os.path.expanduser("~") + "/scheduler-") 165 | # Random port number 166 | self.port = ''.join(["1"] + [str(random.randint(0, 9)) for _ in range(3)]) 167 | # Creating logging interface 168 | stdout_scheduler = DEVNULL 169 | stdout_workers = [DEVNULL] * len(hostnames) 170 | if logging: 171 | # Creating logging folder 172 | try: 173 | os.mkdir("dask_logs") 174 | except OSError: 175 | pass 176 | stdout_scheduler = open("dask_logs/dask-scheduler.log", "w") 177 | stdout_workers = [open("dask_logs/dask-worker-%s.log" % (ii + 1), "w") for ii in range(len(hostnames))] 178 | # Starting scheduler 179 | scheduler_file = "%s%s" % (scheduler_file_prefix, self.port) + ".json" 180 | cmd = ["ssh"] + [hostnames[0]] + \ 181 | ["dask-scheduler"] + ["--scheduler-file"] + [scheduler_file] + \ 182 | ["--port"] + [self.port] 183 | self.scheduler_proc = subprocess.Popen(cmd, stdout=stdout_scheduler, stderr=subprocess.STDOUT) 184 | # Checking if scheduler has started and getting tpc information 185 | t0 = time() 186 | while True: 187 | if os.path.isfile(scheduler_file): 188 | if get_tcp_info(scheduler_file): break 189 | # If the dask scheduler is not started in 5 minutes raise exception 190 | if time() - t0 > 300.0: 191 | raise SystemError("Dask could not start scheduler! Try different first host name.") 192 | # Creating dask Client 193 | self.client = daskD.Client(scheduler_file=scheduler_file) 194 | # Starting workers on all the other hosts 195 | self.worker_procs = [] 196 | worker_ips = [] 197 | for ii, hostname in enumerate(hostnames): 198 | cmd = ["ssh"] + [hostname] + ["dask-worker"] + ["--scheduler-file"] + [scheduler_file] 199 | # Starting worker 200 | self.worker_procs.append(subprocess.Popen(cmd, stdout=stdout_workers[ii], stderr=subprocess.STDOUT)) 201 | # Obtaining IP address of host for the started worker (necessary to resort workers) 202 | worker_ips.append( 203 | subprocess.check_output( 204 | ["ssh"] + [hostname] + ["hostname -i"] + ["| awk '{print $1}'"]).rstrip().decode("utf-8")) 205 | # Waiting until all the requested workers are up and running 206 | workers = 0 207 | requested = len(hostnames) 208 | t0 = time() 209 | while workers < requested: 210 | workers = len(self.client.get_worker_logs().keys()) 211 | # If the number of workers is not reached in 5 minutes raise exception 212 | if time() - t0 > 300.0: 213 | raise SystemError( 214 | "Dask could not start the requested workers within 5 minutes! Try different hostnames.") 215 | # Resorting worker IDs according to user-provided list 216 | self.WorkerIds = [] 217 | wrkIds = list(self.client.get_worker_logs().keys()) # Unsorted workers ids 218 | wrk_ips = [idw.split(":")[1][2:] for idw in wrkIds] # Unsorted ip addresses 219 | for ip in worker_ips: 220 | idx = wrk_ips.index(ip) 221 | self.WorkerIds.append(wrkIds[idx]) 222 | wrkIds.pop(idx) 223 | wrk_ips.pop(idx) 224 | elif kube_params: 225 | n_wrks = kwargs.get("n_wrks") 226 | if "image" not in kube_params: 227 | kube_params.update({"image": 'ettore88/occamypy:devel'}) 228 | pod_spec = make_pod_spec(**kube_params) 229 | self.cluster = KubeCluster(pod_spec, deploy_mode="remote") 230 | self.client, self.WorkerIds = client_startup(self.cluster, n_wrks, n_wrks) 231 | elif ClusterInit: 232 | n_wrks = kwargs.get("n_wrks", 1) 233 | if n_wrks <= 0: 234 | raise ValueError("n_wrks must equal or greater than 1!") 235 | if "local_params" in kwargs: 236 | # Starting local cluster 237 | n_jobs = n_wrks 238 | n_wrks = 1 239 | else: 240 | # Starting scheduler-based clusters 241 | n_jobs = kwargs.get("n_jobs") 242 | if n_jobs <= 0: 243 | raise ValueError("n_jobs must equal or greater than 1!") 244 | cluster_params.update({"processes": n_wrks}) 245 | if n_wrks > 1: 246 | # forcing nanny to be true (otherwise, dask-worker command will fail) 247 | cluster_params.update({"nanny": True}) 248 | self.cluster = ClusterInit(**cluster_params) 249 | self.client, self.WorkerIds = client_startup(self.cluster, n_jobs, n_jobs * n_wrks) 250 | else: 251 | raise ValueError("Either hostnames or local_params or pbs/lsf/slurm_params or kube_params must be " 252 | "provided!") 253 | 254 | self.num_workers = len(self.WorkerIds) 255 | self.dashboard_link = self.cluster.dashboard_link 256 | 257 | # Closing dask processes 258 | atexit.register(self.client.shutdown) 259 | -------------------------------------------------------------------------------- /tutorials/devitoseismic/source.py: -------------------------------------------------------------------------------- 1 | from scipy import interpolate 2 | from cached_property import cached_property 3 | import numpy as np 4 | try: 5 | import matplotlib.pyplot as plt 6 | except: 7 | plt = None 8 | 9 | from devito.types import SparseTimeFunction 10 | 11 | __all__ = ['PointSource', 'Receiver', 'Shot', 'WaveletSource', 12 | 'RickerSource', 'GaborSource', 'DGaussSource', 'TimeAxis'] 13 | 14 | 15 | class TimeAxis(object): 16 | """ 17 | Data object to store the TimeAxis. Exactly three of the four key arguments 18 | must be prescribed. Because of remainder values it is not possible to create 19 | a TimeAxis that exactly adhears to the inputs therefore start, stop, step 20 | and num values should be taken from the TimeAxis object rather than relying 21 | upon the input values. 22 | 23 | The four possible cases are: 24 | start is None: start = step*(1 - num) + stop 25 | step is None: step = (stop - start)/(num - 1) 26 | num is None: num = ceil((stop - start + step)/step); 27 | because of remainder stop = step*(num - 1) + start 28 | stop is None: stop = step*(num - 1) + start 29 | 30 | Parameters 31 | ---------- 32 | start : float, optional 33 | Start of time axis. 34 | step : float, optional 35 | Time interval. 36 | num : int, optional 37 | Number of values (Note: this is the number of intervals + 1). 38 | Stop value is reset to correct for remainder. 39 | stop : float, optional 40 | End time. 41 | """ 42 | def __init__(self, start=None, step=None, num=None, stop=None): 43 | try: 44 | if start is None: 45 | start = step*(1 - num) + stop 46 | elif step is None: 47 | step = (stop - start)/(num - 1) 48 | elif num is None: 49 | num = int(np.ceil((stop - start + step)/step)) 50 | stop = step*(num - 1) + start 51 | elif stop is None: 52 | stop = step*(num - 1) + start 53 | else: 54 | raise ValueError("Only three of start, step, num and stop may be set") 55 | except: 56 | raise ValueError("Three of args start, step, num and stop may be set") 57 | 58 | if not isinstance(num, int): 59 | raise TypeError("input argument must be of type int") 60 | 61 | self.start = start 62 | self.stop = stop 63 | self.step = step 64 | self.num = num 65 | 66 | def __str__(self): 67 | return "TimeAxis: start=%g, stop=%g, step=%g, num=%g" % \ 68 | (self.start, self.stop, self.step, self.num) 69 | 70 | def _rebuild(self): 71 | return TimeAxis(start=self.start, stop=self.stop, num=self.num) 72 | 73 | @cached_property 74 | def time_values(self): 75 | return np.linspace(self.start, self.stop, self.num) 76 | 77 | 78 | class PointSource(SparseTimeFunction): 79 | """Symbolic data object for a set of sparse point sources 80 | 81 | Parameters 82 | ---------- 83 | name : str 84 | Name of the symbol representing this source. 85 | grid : Grid 86 | The computational domain. 87 | time_range : TimeAxis 88 | TimeAxis(start, step, num) object. 89 | npoint : int, optional 90 | Number of sparse points represented by this source. 91 | data : ndarray, optional 92 | Data values to initialise point data. 93 | coordinates : ndarray, optional 94 | Point coordinates for this source. 95 | space_order : int, optional 96 | Space discretization order. 97 | time_order : int, optional 98 | Time discretization order (defaults to 2). 99 | dtype : data-type, optional 100 | Data type of the buffered data. 101 | dimension : Dimension, optional 102 | Represents the number of points in this source. 103 | """ 104 | 105 | @classmethod 106 | def __args_setup__(cls, *args, **kwargs): 107 | kwargs['nt'] = kwargs['time_range'].num 108 | 109 | # Either `npoint` or `coordinates` must be provided 110 | npoint = kwargs.get('npoint') 111 | if npoint is None: 112 | coordinates = kwargs.get('coordinates', kwargs.get('coordinates_data')) 113 | if coordinates is None: 114 | raise TypeError("Need either `npoint` or `coordinates`") 115 | kwargs['npoint'] = coordinates.shape[0] 116 | 117 | return args, kwargs 118 | 119 | def __init_finalize__(self, *args, **kwargs): 120 | time_range = kwargs.pop('time_range') 121 | data = kwargs.pop('data', None) 122 | 123 | kwargs.setdefault('time_order', 2) 124 | super(PointSource, self).__init_finalize__(*args, **kwargs) 125 | 126 | self._time_range = time_range._rebuild() 127 | 128 | # If provided, copy initial data into the allocated buffer 129 | if data is not None: 130 | self.data[:] = data 131 | 132 | @cached_property 133 | def time_values(self): 134 | return self._time_range.time_values 135 | 136 | @property 137 | def time_range(self): 138 | return self._time_range 139 | 140 | def resample(self, dt=None, num=None, rtol=1e-5, order=3): 141 | # Only one of dt or num may be set. 142 | if dt is None: 143 | assert num is not None 144 | else: 145 | assert num is None 146 | 147 | start, stop = self._time_range.start, self._time_range.stop 148 | dt0 = self._time_range.step 149 | 150 | if dt is None: 151 | new_time_range = TimeAxis(start=start, stop=stop, num=num) 152 | dt = new_time_range.step 153 | else: 154 | new_time_range = TimeAxis(start=start, stop=stop, step=dt) 155 | 156 | if np.isclose(dt, dt0): 157 | return self 158 | 159 | nsamples, ntraces = self.data.shape 160 | 161 | new_traces = np.zeros((new_time_range.num, ntraces)) 162 | 163 | for i in range(ntraces): 164 | tck = interpolate.splrep(self._time_range.time_values, 165 | self.data[:, i], k=order) 166 | new_traces[:, i] = interpolate.splev(new_time_range.time_values, tck) 167 | 168 | # Return new object 169 | return PointSource(name=self.name, grid=self.grid, data=new_traces, 170 | time_range=new_time_range, coordinates=self.coordinates.data) 171 | 172 | # Pickling support 173 | _pickle_kwargs = SparseTimeFunction._pickle_kwargs + ['time_range'] 174 | _pickle_kwargs.remove('nt') # `nt` is inferred from `time_range` 175 | 176 | 177 | Receiver = PointSource 178 | Shot = PointSource 179 | 180 | 181 | class WaveletSource(PointSource): 182 | 183 | """ 184 | Abstract base class for symbolic objects that encapsulate a set of 185 | sources with a pre-defined source signal wavelet. 186 | 187 | Parameters 188 | ---------- 189 | name : str 190 | Name for the resulting symbol. 191 | grid : Grid 192 | The computational domain. 193 | f0 : float 194 | Peak frequency for Ricker wavelet in kHz. 195 | time_values : TimeAxis 196 | Discretized values of time in ms. 197 | a : float, optional 198 | Amplitude of the wavelet (defaults to 1). 199 | t0 : float, optional 200 | Firing time (defaults to 1 / f0) 201 | """ 202 | 203 | @classmethod 204 | def __args_setup__(cls, *args, **kwargs): 205 | kwargs.setdefault('npoint', 1) 206 | 207 | return super(WaveletSource, cls).__args_setup__(*args, **kwargs) 208 | 209 | def __init_finalize__(self, *args, **kwargs): 210 | super(WaveletSource, self).__init_finalize__(*args, **kwargs) 211 | 212 | self.f0 = kwargs.get('f0') 213 | self.a = kwargs.get('a') 214 | self.t0 = kwargs.get('t0') 215 | for p in range(kwargs['npoint']): 216 | self.data[:, p] = self.wavelet 217 | 218 | @property 219 | def wavelet(self): 220 | """ 221 | Return a wavelet with a peak frequency ``f0`` at time ``t0``. 222 | """ 223 | raise NotImplementedError('Wavelet not defined') 224 | 225 | def show(self, idx=0, wavelet=None): 226 | """ 227 | Plot the wavelet of the specified source. 228 | 229 | Parameters 230 | ---------- 231 | idx : int 232 | Index of the source point for which to plot wavelet. 233 | wavelet : ndarray or callable 234 | Prescribed wavelet instead of one from this symbol. 235 | """ 236 | wavelet = wavelet or self.data[:, idx] 237 | plt.figure() 238 | plt.plot(self.time_values, wavelet) 239 | plt.xlabel('Time (ms)') 240 | plt.ylabel('Amplitude') 241 | plt.tick_params() 242 | plt.show() 243 | 244 | # Pickling support 245 | _pickle_kwargs = PointSource._pickle_kwargs + ['f0', 'a', 'f0'] 246 | 247 | 248 | class RickerSource(WaveletSource): 249 | 250 | """ 251 | Symbolic object that encapsulate a set of sources with a 252 | pre-defined Ricker wavelet: 253 | 254 | http://subsurfwiki.org/wiki/Ricker_wavelet 255 | 256 | Parameters 257 | ---------- 258 | name : str 259 | Name for the resulting symbol. 260 | grid : Grid 261 | The computational domain. 262 | f0 : float 263 | Peak frequency for Ricker wavelet in kHz. 264 | time : TimeAxis 265 | Discretized values of time in ms. 266 | 267 | Returns 268 | ---------- 269 | op Ricker wavelet. 270 | """ 271 | 272 | @property 273 | def wavelet(self): 274 | t0 = self.t0 or 1 / self.f0 275 | a = self.a or 1 276 | r = (np.pi * self.f0 * (self.time_values - t0)) 277 | return a * (1-2.*r**2)*np.exp(-r**2) 278 | 279 | 280 | class GaborSource(WaveletSource): 281 | 282 | """ 283 | Symbolic object that encapsulate a set of sources with a 284 | pre-defined Gabor wavelet: 285 | 286 | https://en.wikipedia.org/wiki/Gabor_wavelet 287 | 288 | Parameters 289 | ---------- 290 | name : str 291 | Name for the resulting symbol. 292 | grid : Grid 293 | defining the computational domain. 294 | f0 : float 295 | Peak frequency for Ricker wavelet in kHz. 296 | time : TimeAxis 297 | Discretized values of time in ms. 298 | 299 | Returns 300 | ------- 301 | op Gabor wavelet. 302 | """ 303 | 304 | @property 305 | def wavelet(self): 306 | agauss = 0.5 * self.f0 307 | tcut = self.t0 or 1.5 / agauss 308 | s = (self.time_values - tcut) * agauss 309 | a = self.a or 1 310 | return a * np.exp(-2*s**2) * np.cos(2 * np.pi * s) 311 | 312 | 313 | class DGaussSource(WaveletSource): 314 | 315 | """ 316 | Symbolic object that encapsulate a set of sources with a 317 | pre-defined 1st derivative wavelet of a Gaussian Source. 318 | 319 | Notes 320 | ----- 321 | For visualizing the second or third order derivative 322 | of Gaussian wavelets, the convention is to use the 323 | negative of the normalized derivative. In the case 324 | of the second derivative, scaling by -1 produces a 325 | wavelet with its main lobe in the positive y direction. 326 | This scaling also makes the Gaussian wavelet resemble 327 | the Mexican hat, or Ricker, wavelet. The validity of 328 | the wavelet is not affected by the -1 scaling factor. 329 | 330 | Parameters 331 | ---------- 332 | name : str 333 | Name for the resulting symbol. 334 | grid : Grid 335 | The computational domain. 336 | f0 : float 337 | Peak frequency for wavelet in kHz. 338 | time : TimeAxis 339 | Discretized values of time in ms. 340 | 341 | Returns 342 | ------- 343 | The 1st order derivative of the Gaussian wavelet. 344 | """ 345 | 346 | @property 347 | def wavelet(self): 348 | t0 = self.t0 or 1 / self.f0 349 | a = self.a or 1 350 | time = (self.time_values - t0) 351 | return -2 * a * time * np.exp(- a * time**2) 352 | -------------------------------------------------------------------------------- /occamypy/operator/derivative.py: -------------------------------------------------------------------------------- 1 | from occamypy.utils import get_backend 2 | from occamypy.operator.base import Operator, Vstack 3 | 4 | 5 | class FirstDerivative(Operator): 6 | r""" 7 | First Derivative with a stencil 8 | 9 | 1) 2nd order centered: 10 | 11 | .. math:: 12 | y[i] = 0.5 (x[i+1] - x[i-1]) / dx 13 | 14 | 2) 1st order forward: 15 | 16 | .. math:: 17 | y[i] = (x[i+1] - x[i]) / dx 18 | 19 | 1) 1st order backward: 20 | 21 | .. math:: 22 | y[i] = 0.5 (x[i] - x[i-1]) / dx 23 | """ 24 | def __init__(self, domain, sampling=1., axis=0, stencil='centered'): 25 | """ 26 | FirstDerivative costructor 27 | 28 | Args: 29 | domain: domain vector 30 | sampling: sampling step along the differentiation axis 31 | axis: axis along which to compute the derivative [0] 32 | stencil: derivative kind (centered, forward, backward) 33 | """ 34 | self.sampling = sampling 35 | self.dims = domain.getNdArray().shape 36 | self.axis = axis if axis >= 0 else len(self.dims) + axis 37 | self.stencil = stencil 38 | 39 | if self.stencil == 'centered': 40 | self.forward = self._forwardC 41 | self.adjoint = self._adjointC 42 | elif self.stencil == 'backward': 43 | self.forward = self._forwardB 44 | self.adjoint = self._adjointB 45 | elif self.stencil == 'forward': 46 | self.forward = self._forwardF 47 | self.adjoint = self._adjointF 48 | else: 49 | raise ValueError("Derivative stencil must be centered, forward or backward") 50 | 51 | self.backend = get_backend(domain) 52 | 53 | super(FirstDerivative, self).__init__(domain, domain, name="Der1") 54 | 55 | def _forwardF(self, add, model, data): 56 | self.checkDomainRange(model, data) 57 | if not add: 58 | data.zero() 59 | # Getting Ndarrays 60 | x = model.clone().getNdArray() 61 | if self.axis > 0: # need to bring the dim. to derive to first dim 62 | x = self.backend.swapaxes(x, self.axis, 0) 63 | y = self.backend.zeros_like(x) 64 | 65 | y[:-1] = (x[1:] - x[:-1]) / self.sampling 66 | if self.axis > 0: # reset axis order 67 | y = self.backend.swapaxes(y, 0, self.axis) 68 | data[:] += y 69 | return 70 | 71 | def _adjointF(self, add, model, data): 72 | self.checkDomainRange(model, data) 73 | if not add: 74 | model.zero() 75 | # Getting Ndarrays 76 | y = data.clone().getNdArray() 77 | if self.axis > 0: # need to bring the dim. to derive to first dim 78 | y = self.backend.swapaxes(y, self.axis, 0) 79 | x = self.backend.zeros_like(y) 80 | 81 | x[:-1] -= y[:-1] / self.sampling 82 | x[1:] += y[:-1] / self.sampling 83 | 84 | if self.axis > 0: 85 | x = self.backend.swapaxes(x, 0, self.axis) 86 | model[:] += x 87 | return 88 | 89 | def _forwardC(self, add, model, data): 90 | self.checkDomainRange(model, data) 91 | if not add: 92 | data.zero() 93 | # Getting Ndarrays 94 | x = model.clone().getNdArray() 95 | if self.axis > 0: # need to bring the dim. to derive to first dim 96 | x = self.backend.swapaxes(x, self.axis, 0) 97 | y = self.backend.zeros_like(x) 98 | 99 | y[1:-1] = (.5 * x[2:] - 0.5 * x[:-2]) / self.sampling 100 | if self.axis > 0: # reset axis order 101 | y = self.backend.swapaxes(y, 0, self.axis) 102 | data[:] += y 103 | return 104 | 105 | def _adjointC(self, add, model, data): 106 | self.checkDomainRange(model, data) 107 | if not add: 108 | model.zero() 109 | # Getting Ndarrays 110 | y = data.clone().getNdArray() 111 | if self.axis > 0: # need to bring the dim. to derive to first dim 112 | y = self.backend.swapaxes(y, self.axis, 0) 113 | x = self.backend.zeros_like(y) 114 | 115 | x[:-2] -= 0.5 * y[1:-1] / self.sampling 116 | x[2:] += 0.5 * y[1:-1] / self.sampling 117 | 118 | if self.axis > 0: 119 | x = self.backend.swapaxes(x, 0, self.axis) 120 | model[:] += x 121 | return 122 | 123 | def _forwardB(self, add, model, data): 124 | self.checkDomainRange(model, data) 125 | if not add: 126 | data.zero() 127 | # Getting Ndarrays 128 | x = model.clone().getNdArray() 129 | if self.axis > 0: # need to bring the dim. to derive to first dim 130 | x = self.backend.swapaxes(x, self.axis, 0) 131 | y = self.backend.zeros_like(x) 132 | 133 | y[1:] = (x[1:] - x[:-1]) / self.sampling 134 | if self.axis > 0: # reset axis order 135 | y = self.backend.swapaxes(y, 0, self.axis) 136 | data[:] += y 137 | return 138 | 139 | def _adjointB(self, add, model, data): 140 | self.checkDomainRange(model, data) 141 | if not add: 142 | model.zero() 143 | # Getting Ndarrays 144 | y = data.clone().getNdArray() 145 | if self.axis > 0: # need to bring the dim. to derive to first dim 146 | y = self.backend.swapaxes(y, self.axis, 0) 147 | x = self.backend.zeros_like(y) 148 | 149 | x[:-1] -= y[1:] / self.sampling 150 | x[1:] += y[1:] / self.sampling 151 | 152 | if self.axis > 0: 153 | x = self.backend.swapaxes(x, 0, self.axis) 154 | model[:] += x 155 | return 156 | 157 | 158 | class SecondDerivative(Operator): 159 | r""" 160 | Compute 2nd order second derivative 161 | 162 | .. math:: 163 | y[i] = (x[i+1] - 2x[i] + x[i-1]) / dx^2 164 | """ 165 | def __init__(self, domain, sampling=1., axis=0): 166 | """ 167 | SecondDerivative constructor 168 | 169 | Args: 170 | domain: domain vector 171 | sampling: sampling step along the differentiation axis 172 | axis: axis along which to compute the derivative 173 | """ 174 | self.sampling = sampling 175 | self.data_tmp = domain.clone().zero() 176 | self.dims = domain.getNdArray().shape 177 | self.axis = axis if axis >= 0 else len(self.dims) + axis 178 | 179 | self.backend = get_backend(domain) 180 | 181 | super(SecondDerivative, self).__init__(domain, domain, name="Der2") 182 | 183 | def forward(self, add, model, data): 184 | self.checkDomainRange(model, data) 185 | if not add: 186 | data.zero() 187 | 188 | # Getting Ndarrays 189 | x = model.clone().getNdArray() 190 | if self.axis > 0: # need to bring the dim. to derive to first dim 191 | x = self.backend.swapaxes(x, self.axis, 0) 192 | y = self.backend.zeros_like(x) 193 | 194 | y[1:-1] = (x[0:-2] - 2 * x[1:-1] + x[2:]) / self.sampling ** 2 195 | 196 | if self.axis > 0: # reset axis order 197 | y = self.backend.swapaxes(y, 0, self.axis) 198 | data[:] += y 199 | return 200 | 201 | def adjoint(self, add, model, data): 202 | self.checkDomainRange(model, data) 203 | if not add: 204 | model.zero() 205 | 206 | # Getting numpy arrays 207 | y = data.clone().getNdArray() 208 | if self.axis > 0: # need to bring the dim. to derive to first dim 209 | y = self.backend.swapaxes(y, self.axis, 0) 210 | x = self.backend.zeros_like(y) 211 | 212 | x[0:-2] += (y[1:-1]) / self.sampling ** 2 213 | x[1:-1] -= (2 * y[1:-1]) / self.sampling ** 2 214 | x[2:] += (y[1:-1]) / self.sampling ** 2 215 | 216 | if self.axis > 0: 217 | x = self.backend.swapaxes(x, 0, self.axis) 218 | model[:] += x 219 | return 220 | 221 | 222 | class Gradient(Operator): 223 | """N-Dimensional Gradient operator""" 224 | 225 | def __init__(self, domain, sampling=None, stencil=None): 226 | """ 227 | Gradient constructor 228 | 229 | Args: 230 | domain: domain vector 231 | sampling: sampling steps 232 | stencil: stencil kind for each direction 233 | """ 234 | self.dims = domain.getNdArray().shape 235 | self.sampling = sampling if sampling is not None else tuple([1] * len(self.dims)) 236 | 237 | if stencil is None: 238 | self.stencil = tuple(['centered'] * len(self.dims)) 239 | elif isinstance(stencil, str): 240 | self.stencil = tuple([stencil] * len(self.dims)) 241 | elif isinstance(stencil, tuple) or isinstance(stencil, list): 242 | self.stencil = stencil 243 | if len(self.sampling) == 0: 244 | raise ValueError("Provide at least one sampling item") 245 | if len(self.stencil) == 0: 246 | raise ValueError("Provide at least one stencil item") 247 | if len(self.sampling) != len(self.stencil): 248 | raise ValueError("There is something wrong with the dimensions") 249 | 250 | self.op = Vstack([FirstDerivative(domain, sampling=self.sampling[d], axis=d) 251 | for d in range(len(self.dims))]) 252 | super(Gradient, self).__init__(domain=self.op.domain, range=self.op.range, name="Gradient") 253 | 254 | def __str__(self): 255 | return "Gradient" 256 | 257 | def forward(self, add, model, data): 258 | return self.op.forward(add, model, data) 259 | 260 | def adjoint(self, add, model, data): 261 | return self.op.adjoint(add, model, data) 262 | 263 | def merge_directions(self, grad_vector, iso=True): 264 | """ 265 | Merge the different directional contributes, using the L2 norm (iso=True) or the simple sum (iso=False) 266 | """ 267 | self.range.checkSame(grad_vector) 268 | 269 | data = self.domain.clone() 270 | 271 | if iso: 272 | for v in grad_vector.vecs: 273 | data.scaleAdd(v.clone().pow(2), 1., 1.) 274 | data.pow(.5) 275 | else: 276 | for v in grad_vector.vecs: 277 | data.scaleAdd(v, 1., 1.) 278 | 279 | return data 280 | 281 | 282 | class Laplacian(Operator): 283 | """ 284 | Laplacian operator. 285 | 286 | Notes: 287 | The input parameters are tailored for >2D, but it works also for 1D. 288 | """ 289 | 290 | def __init__(self, domain, axis=None, weights=None, sampling=None): 291 | """ 292 | Laplacian constructor 293 | 294 | Args: 295 | domain: domain vector 296 | axis: axes along which to compute the derivative 297 | weights: scalar weights for each axis 298 | sampling: sampling steps for each axis 299 | """ 300 | self.dims = domain.getNdArray().shape 301 | self.axis = axis if axis is not None else tuple(range(len(self.dims))) 302 | self.sampling = sampling if sampling is not None else tuple([1] * len(self.dims)) 303 | self.weights = weights if weights is not None else tuple([1] * len(self.dims)) 304 | 305 | if not (len(self.axis) == len(self.weights) == len(self.sampling)): 306 | raise ValueError("There is something wrong with the dimensions") 307 | 308 | self.data_tmp = domain.clone().zero() 309 | 310 | self.op = self.weights[0] * SecondDerivative(domain, sampling=self.sampling[0], axis=self.axis[0]) 311 | for d in range(1, len(self.axis)): 312 | self.op += self.weights[d] * SecondDerivative(domain, sampling=self.sampling[d], axis=self.axis[d]) 313 | super(Laplacian, self).__init__(domain, domain, name="Laplacian") 314 | 315 | def forward(self, add, model, data): 316 | return self.op.forward(add, model, data) 317 | 318 | def adjoint(self, add, model, data): 319 | return self.op.adjoint(add, model, data) 320 | --------------------------------------------------------------------------------