diff --git a/pygeos-tools/examples/solvers/gravity/adjoint-test/README.md b/pygeos-tools/examples/solvers/gravity/adjoint-test/README.md new file mode 100644 index 00000000..469f101f --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/adjoint-test/README.md @@ -0,0 +1,27 @@ +# Adjoint Test for GEOS Gravity modeling + + +## Typical Usage + +### Rectangular Prism Anomaly Meshed with Tetrahedra + +```bash +python gravityAdjointTest.py --xml ../data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml --n_model 10368 --n_data 1271 +``` + +### Rectangular Prism Anomaly Meshed with Hexahedra: + +```bash +python gravityAdjointTest.py --xml ../data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml --n_model 1728 --n_data 1271 +``` + +## Typical Output + +```text +=== Adjoint Test Summary === + = 7.585858080970467e-07 + = 7.585858080970469e-07 +Passed = True +Error = 1.2332569442944734e-16 +``` + diff --git a/pygeos-tools/examples/solvers/gravity/adjoint-test/gravityAdjointTest.py b/pygeos-tools/examples/solvers/gravity/adjoint-test/gravityAdjointTest.py new file mode 100644 index 00000000..ed8314e8 --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/adjoint-test/gravityAdjointTest.py @@ -0,0 +1,84 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ + +import argparse +import numpy as np +import sys +from mpi4py import MPI + +from geos.pygeos_tools.input import XML +from geos.pygeos_tools.solvers import GravityLinearOpSolver +from geos.pygeos_tools.solvers.InversionUtils import adjointTest + +__doc__ = """ +This is an example of how to set up and run an adjoint test for GEOS simulation using the GravitySolver. +""" + + +def parse_args(): + """ + Parse command-line arguments. + + Returns + ------- + argparse.Namespace + Parsed arguments including XML input, dimensions and scaling factor. + """ + parser = argparse.ArgumentParser( description="Run a GEOS gravity adjoint test." ) + parser.add_argument( '--xml', type=str, required=True, help="Input XML file for GEOSX." ) + parser.add_argument( '--n_model', type=int, required=True, help="Number of model samples." ) + parser.add_argument( '--n_data', type=int, required=True, help="Number of data samples." ) + parser.add_argument( '--scale', type=float, default=1.0, help="Scaling factor." ) + return parser.parse_args() + + +def main(): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + verbose = ( rank == 0 ) + + args = parse_args() + xml_file = args.xml + scale = args.scale + nm = args.n_model + nd = args.n_data + + # Load XML configuration + xml = XML( xml_file ) + + # Generate random fields + np.random.seed( 42 ) # Set random seed for reproducibility + if rank == 0: + x = np.random.rand( nm ) + y = np.random.rand( nd ) + else: + x = None + y = None + + x = comm.bcast( x, root=0 ) + y = comm.bcast( y, root=0 ) + + # Initialize solver as a linear operator + solver = GravityLinearOpSolver( rank=rank, xml=xml, nm=args.n_model, nd=args.n_data, scaleData=scale ) + + # Run the test + passed, error = adjointTest( solver._matvec, solver._rmatvec, x, y, flag_verbose=verbose ) + + solver.finalize() + + if not passed: + sys.exit( 1 ) + + +if __name__ == "__main__": + main() diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/dm.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/dm.npy new file mode 100644 index 00000000..4e1695ee Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/dm.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m0.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m0.npy new file mode 100644 index 00000000..78ee3c76 Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m0.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m_true.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m_true.npy new file mode 100644 index 00000000..29a954ba Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/m_true.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml new file mode 100644 index 00000000..fc39ee68 --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml @@ -0,0 +1,2625 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/dm.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/dm.npy new file mode 100644 index 00000000..d374ba64 Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/dm.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m0.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m0.npy new file mode 100644 index 00000000..31f9f7ba Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m0.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m_true.npy b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m_true.npy new file mode 100644 index 00000000..c0262740 Binary files /dev/null and b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/m_true.npy differ diff --git a/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml new file mode 100644 index 00000000..5a40c18e --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml @@ -0,0 +1,2625 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/pygeos-tools/examples/solvers/gravity/gradient-test/README.md b/pygeos-tools/examples/solvers/gravity/gradient-test/README.md new file mode 100644 index 00000000..4795b34b --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/gradient-test/README.md @@ -0,0 +1,55 @@ +# Gradient Test for GEOS Gravity modeling + + +## Typical Usage + +### Rectangular Prism Anomaly Meshed with Tetrahedra + +```bash +python gravityGradientTest.py --xml ../data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml --m_true ../data/rectangularPrism_C3D4/m_true.npy --m0 ../data/rectangularPrism_C3D4/m0.npy --dm ../data/rectangularPrism_C3D4/dm.npy +``` + +### Rectangular Prism Anomaly Meshed with Hexahedra: + +```bash +python gravityGradientTest.py --xml ../data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml --m_true ../data/rectangularPrism_C3D8/m_true.npy --m0 ../data/rectangularPrism_C3D8/m0.npy --dm ../data/rectangularPrism_C3D8/dm.npy +``` + + +## Typical Output showing the expected quadratic error convergence up to machine precision +```text + h h2 e0 e1 +5.5556e-01 3.0864e-01 6.8482e-04 1.4697e-07 +3.0864e-01 9.5260e-02 3.8049e-04 4.5360e-08 +1.7147e-01 2.9401e-02 2.1140e-04 1.4000e-08 +9.5260e-02 9.0744e-03 1.1745e-04 4.3210e-09 +5.2922e-02 2.8008e-03 6.5249e-05 1.3336e-09 +2.9401e-02 8.6443e-04 3.6250e-05 4.1162e-10 +1.6334e-02 2.6680e-04 2.0139e-05 1.2704e-10 +9.0744e-03 8.2346e-05 1.1188e-05 3.9214e-11 +5.0414e-03 2.5415e-05 6.2157e-06 1.2104e-11 +2.8008e-03 7.8442e-06 3.4532e-06 3.7367e-12 +1.5560e-03 2.4211e-06 1.9184e-06 1.1550e-12 +8.6443e-04 7.4724e-07 1.0658e-06 3.5841e-13 +4.8024e-04 2.3063e-07 5.9211e-07 1.1246e-13 +2.6680e-04 7.1182e-08 3.2895e-07 3.5991e-14 +1.4822e-04 2.1970e-08 1.8275e-07 1.2474e-14 +8.2346e-05 6.7808e-09 1.0153e-07 6.0212e-15 +4.5748e-05 2.0928e-09 5.6404e-08 3.2005e-15 +2.5415e-05 6.4594e-10 3.1336e-08 2.5012e-15 +1.4120e-05 1.9936e-10 1.7409e-08 2.8565e-15 +7.8442e-06 6.1532e-11 9.6715e-09 2.3721e-15 +4.3579e-06 1.8991e-11 5.3731e-09 2.8261e-15 +2.4211e-06 5.8615e-12 2.9850e-09 2.9543e-15 +1.3450e-06 1.8091e-12 1.6583e-09 2.4471e-15 +7.4724e-07 5.5837e-13 9.2130e-10 2.8471e-15 +4.1513e-07 1.7234e-13 5.1183e-10 2.1396e-15 + +=== Gradient Test Summary === +Passed: True +Slope: [ 1.99999993 1.99999983 1.99999957 1.9999963 1.99999456 1.99997363 + 1.99988094 1.99984551 1.99960117 1.99755535 1.99079356 1.97189443 + 1.93835375 1.80267853 1.2392288 1.07519441 0.41943078 -0.22598208 + 0.31615789 -0.29794799 -0.07551366 0.32048653 -0.25757919 0.48605047] +History saved to: grad_test.txt +``` diff --git a/pygeos-tools/examples/solvers/gravity/gradient-test/gravityGradientTest.py b/pygeos-tools/examples/solvers/gravity/gradient-test/gravityGradientTest.py new file mode 100644 index 00000000..ca68499b --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/gradient-test/gravityGradientTest.py @@ -0,0 +1,100 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ + +import argparse +import numpy as np +import sys +from mpi4py import MPI + +from geos.pygeos_tools.input import XML +from geos.pygeos_tools.solvers import GravityLinearOpSolver +from geos.pygeos_tools.solvers.InversionUtils import gradientTest, plotGradientTestFromFile + +__doc__ = """ +This is an example of how to set up and run a gradient test for GEOS simulation using the GravitySolver. +""" + + +def parse_args(): + """ + Parse command-line arguments. + + Returns + ------- + argparse.Namespace + Parsed arguments including XML input, models, and scaling factor. + """ + parser = argparse.ArgumentParser( description="Run a GEOS gravity simulation and gradient test." ) + parser.add_argument( '--xml', type=str, required=True, help="Input XML file for GEOSX." ) + parser.add_argument( '--m_true', type=str, required=True, help="True model file (.npy)." ) + parser.add_argument( '--m0', type=str, required=True, help="Initial model file (.npy)." ) + parser.add_argument( '--dm', type=str, required=True, help="Perturbation model file (.npy)." ) + parser.add_argument( '--scale', type=float, default=1.0, help="Scaling factor for the loss." ) + parser.add_argument( '--history', type=str, default='grad_test.txt', help="Output file for gradient test history." ) + parser.add_argument( '--plot', action='store_true', help="Plot the gradient test results." ) + return parser.parse_args() + + +def main(): + + rank = MPI.COMM_WORLD.Get_rank() + verbose = ( rank == 0 ) + + args = parse_args() + xml_file = args.xml + scale = args.scale + + # Load XML configuration + xml = XML( xml_file ) + + # Load models from .npy files + try: + m_true = np.load( args.m_true ) + m0 = np.load( args.m0 ) + dm = np.load( args.dm ) + except FileNotFoundError as e: + raise RuntimeError( f"Could not load input model file: {e}" ) + + # Initialize solver as a linear operator + solver = GravityLinearOpSolver( rank=rank, xml=xml, nm=m_true.size, scaleData=scale ) + + # Generate observed data + uobs = solver.getData( m_true ) + + # Run gradient test + passed, slope, history = gradientTest( lambda m: solver.getLoss( m, uobs, scale=scale ), + lambda m: solver.getGradient( m, uobs, scale=scale ), + m0, + dm, + history_filename=args.history, + flag_verbose=verbose ) + + if verbose: + print( "\n=== Gradient Test Summary ===" ) + print( f"Passed: {passed}" ) + print( f"Slope: {slope if slope.size > 0 else 'N/A'}" ) + print( f"History saved to: {args.history}" ) + + if args.plot: + plotGradientTestFromFile( args.history, save_path="grad_test.png", style={ "color": "red", "marker": "o" } ) + + # Finalize GEOS + solver.finalize() + + # Exit with non-zero code if test failed + if not passed: + sys.exit( 1 ) + + +if __name__ == "__main__": + main() diff --git a/pygeos-tools/examples/solvers/gravity/modeling/README.md b/pygeos-tools/examples/solvers/gravity/modeling/README.md new file mode 100644 index 00000000..24c0dc8f --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/modeling/README.md @@ -0,0 +1,25 @@ +# Gravity Modeling + + +## Typical Usage + +### Rectangular Prism Anomaly Meshed with Tetrahedra + +```bash +python gravityModeling.py --xml ../data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml --m_true ../data/rectangularPrism_C3D4/m_true.npy +``` +or equivalently: +```bash +python gravityModelingLinearOp.py --xml ../data/rectangularPrism_C3D4/rectangularPrism_C3D4.xml --m_true ../data/rectangularPrism_C3D4/m_true.npy +``` + +### Rectangular Prism Anomaly Meshed with Hexahedra: + +```bash +python gravityModeling.py --xml ../data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml --m_true ../data/rectangularPrism_C3D8/m_true.npy +``` +or equivalently: +```bash +python gravityModelingLinearOp.py --xml ../data/rectangularPrism_C3D8/rectangularPrism_C3D8.xml --m_true ../data/rectangularPrism_C3D8/m_true.npy +``` + diff --git a/pygeos-tools/examples/solvers/gravity/modeling/gravityModeling.py b/pygeos-tools/examples/solvers/gravity/modeling/gravityModeling.py new file mode 100644 index 00000000..169e30b1 --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/modeling/gravityModeling.py @@ -0,0 +1,62 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ +import argparse +import copy + +from mpi4py import MPI + +import numpy as np +from geos.pygeos_tools.input import XML +from geos.pygeos_tools.solvers import GravitySolver +from pygeosx import run, COMPLETED + +__doc__ = """ +“This is an example of how to set up and run your GEOS simulation using the GravitySolver. +""" + + +def parse_args(): + """Get arguments + + Returns: + argument '--xml': Input xml file for GEOSX + """ + parser = argparse.ArgumentParser( description="Gravity simulation example" ) + parser.add_argument( '--xml', type=str, required=True, help="Input xml file for GEOS" ) + parser.add_argument( '--m_true', type=str, required=True, help="True model (.npy)" ) + + args, _ = parser.parse_known_args() + return args + + +def main(): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + args = parse_args() + xmlfile = args.xml + xml = XML( xmlfile ) + + solver = GravitySolver() + solver.initialize( rank=rank, xml=xml ) + + density = np.load( args.m_true ) + + gz = solver.modeling( density ) + print( f"gz min={np.min(gz)}, max={np.max(gz)}" ) + + solver.finalize() + + +if __name__ == "__main__": + main() diff --git a/pygeos-tools/examples/solvers/gravity/modeling/gravityModelingLinearOp.py b/pygeos-tools/examples/solvers/gravity/modeling/gravityModelingLinearOp.py new file mode 100644 index 00000000..9995a05a --- /dev/null +++ b/pygeos-tools/examples/solvers/gravity/modeling/gravityModelingLinearOp.py @@ -0,0 +1,63 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ +import argparse +import copy + +from mpi4py import MPI + +import numpy as np +from geos.pygeos_tools.input import XML +from geos.pygeos_tools.solvers import GravityLinearOpSolver +from pygeosx import run, COMPLETED + +__doc__ = """ +“This is an example of how to set up and run your GEOS simulation using the GravitySolver. +""" + + +def parse_args(): + """Get arguments + + Returns: + argument '--xml': Input xml file for GEOSX + """ + parser = argparse.ArgumentParser( description="Gravity simulation example" ) + parser.add_argument( '--xml', type=str, required=True, help="Input xml file for GEOS" ) + parser.add_argument( '--m_true', type=str, required=True, help="True model (.npy)" ) + + args, _ = parser.parse_known_args() + return args + + +def main(): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + args = parse_args() + xmlfile = args.xml + xml = XML( xmlfile ) + + m_true = np.load( args.m_true ) + + solver = GravityLinearOpSolver( rank=rank, xml=xml, nm=m_true.size ) + + print( f"Density min={np.min(m_true)}, max={np.max(m_true)}", flush=True ) + + gz = solver.getData( m_true ) + print( f"gz min={np.min(gz)}, max={np.max(gz)}" ) + + solver.finalize() + + +if __name__ == "__main__": + main() diff --git a/pygeos-tools/src/geos/pygeos_tools/solvers/GravityLinearOpSolver.py b/pygeos-tools/src/geos/pygeos_tools/solvers/GravityLinearOpSolver.py new file mode 100644 index 00000000..7ff0e5c2 --- /dev/null +++ b/pygeos-tools/src/geos/pygeos_tools/solvers/GravityLinearOpSolver.py @@ -0,0 +1,191 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ + +import numpy as np +from numpy.typing import DTypeLike +from typing import Optional, Tuple +from scipy.sparse.linalg import LinearOperator + +from geos.pygeos_tools.input.Xml import XML +from geos.pygeos_tools.solvers.InversionUtils import computeResidual, computeL2Loss +from geos.pygeos_tools.solvers.GravitySolver import GravitySolver + + +class GravityLinearOpSolver( LinearOperator ): + """ + A SciPy-compatible linear operator that wraps GEOS-based forward and adjoint gravity modeling. + + This class provides a bridge between GEOS's GravitySolver and SciPy's optimization and inversion + tools by exposing the forward and adjoint operations as a LinearOperator. It supports: + + - Forward modeling via `_matvec` + - Adjoint modeling via `_rmatvec` + - Loss and gradient computation for inversion workflows + + Parameters + ---------- + rank : int, optional + MPI rank of the current process (default is 0). + scaleData : float, optional + Scaling factor applied to the forward model output (default is 1.0). + xml : object, optional + Parsed XML configuration for GEOS. + nm : int + Number of model parameters (must be > 0). + nd : int, optional + Number of data points. If not provided, it will be inferred on first forward call. + dtype : np.dtype, optional + Data type for the operator (default is np.float64). + """ + + def __init__( self, + rank: int = 0, + scaleData: float = 1.0, + xml: Optional[ XML ] = None, + nm: int = 0, + nd: Optional[ int ] = None, + dtype: DTypeLike = np.float64 ) -> None: + + if nm <= 0: + raise ValueError( f"Invalid dimensions: nm={nm}, must be a positive integer." ) + + self.rank = rank + self.scaleData = scaleData + self.xml = xml + self.nm = nm + self.nd = nd + self.dtype = dtype + self._shape_initialized = False + self.solver = GravitySolver( solverType="GravityFE" ) + + if self.nd is not None: + super().__init__( dtype=self.dtype, shape=( self.nd, self.nm ) ) + self._shape_initialized = True + + def _initialize_shape( self, nd: int ) -> None: + if not self._shape_initialized: + super().__init__( dtype=self.dtype, shape=( nd, self.nm ) ) + self._shape_initialized = True + + def _resetSolver( self ) -> None: + self.solver.initialize( rank=self.rank, xml=self.xml ) + + def _matvec( self, x: np.ndarray ) -> np.ndarray: + self._resetSolver() + y = self.solver.modeling( x, scale_data=self.scaleData ) + self._initialize_shape( nd=y.size ) + return y + + def _rmatvec( self, y: np.ndarray ) -> np.ndarray: + self._resetSolver() + self._initialize_shape( nd=y.size ) + return self.solver.adjoint( self.nm, y ) + + def getData( self, model: np.ndarray ) -> np.ndarray: + """ + Run forward modeling on the given model. + + Parameters + ---------- + model : np.ndarray + Model parameters. + + Returns + ------- + np.ndarray + Predicted data. + """ + return self._matvec( model ) + + def getLoss( self, model: np.ndarray, dObs: np.ndarray, scale: float = 1.0 ) -> float: + """ + Compute the L2 loss between predicted and observed data. + + Parameters + ---------- + model : np.ndarray + Model parameters. + dObs : np.ndarray + Observed data. + scale : float, optional + Scaling factor for the loss (default is 1.0). + + Returns + ------- + float + Scaled L2 loss. + """ + dPred = self._matvec( model ) + return computeL2Loss( dPred, dObs, scale ) + + def getGradient( self, model: np.ndarray, dObs: np.ndarray, scale: float = 1.0 ) -> np.ndarray: + """ + Compute the gradient of the loss with respect to the model. + + Parameters + ---------- + model : np.ndarray + Model parameters. + dObs : np.ndarray + Observed data. + scale : float, optional + Scaling factor for the gradient (default is 1.0). + + Returns + ------- + np.ndarray + Gradient vector. + """ + dPred = self._matvec( model ) + residue = computeResidual( dPred, dObs ) + return self._rmatvec( residue ) * scale + + def getLossAndGradient( self, + model: np.ndarray, + dObs: np.ndarray, + scale: float = 1.0 ) -> Tuple[ float, np.ndarray ]: + """ + Compute both the loss and its gradient. + + Parameters + ---------- + model : np.ndarray + Model parameters. + dObs : np.ndarray + Observed data. + scale : float, optional + Scaling factor (default is 1.0). + + Returns + ------- + Tuple[float, np.ndarray] + Tuple containing the loss and the gradient. + """ + dPred = self._matvec( model ) + residue = computeResidual( dPred, dObs ) + loss = computeL2Loss( dPred, dObs, scale ) + gradient = self._rmatvec( residue ) * scale + return loss, gradient + + def finalize( self ) -> None: + """ + Finalize the GEOS solver if it was initialized. + """ + if self.solver.alreadyInitialized: + self.solver.finalize() + + def __enter__( self ) -> "GravityLinearOpSolver": + return self + + def __exit__( self, excType, excValue, traceback ) -> None: + self.finalize() diff --git a/pygeos-tools/src/geos/pygeos_tools/solvers/GravitySolver.py b/pygeos-tools/src/geos/pygeos_tools/solvers/GravitySolver.py new file mode 100644 index 00000000..f4b59aee --- /dev/null +++ b/pygeos-tools/src/geos/pygeos_tools/solvers/GravitySolver.py @@ -0,0 +1,156 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ + +import numpy as np +import numpy.typing as npt +from typing import Dict, List, Optional +from typing_extensions import Self +from geos.pygeos_tools.solvers.Solver import Solver +from geos.pygeos_tools.wrapper import get_matching_wrapper_path, set_wrapper_to_value, allgather_wrapper +from pygeosx import run, COMPLETED +import copy + +__doc__ = """ +GravitySolver class inherits from Solver class. + +This adds accessor methods for gravity modeling and adjoint. +""" + + +class GravitySolver( Solver ): + """ + Gravity solver object containing methods to run gravity simulations in GEOS + + Attributes + ----------- + The ones inherited from Solver class + """ + + def __init__( self: Self, solverType: str = "GravityFE", **kwargs ): + """ + Parameters + ----------- + solverType : str + The solver used in the XML file + """ + super().__init__( solverType, **kwargs ) + + def getGzAtStations( self: Self ) -> npt.NDArray: + """ + Get gz values at station coordinates + + Returns + ------ + numpy Array : Array containing gz values at all stations coordinates + """ + return self.getGeosWrapperByName( "gzAtStations", [ "Solvers" ] ) + + def getDensityModel( self: Self, filterGhost: bool = False, **kwargs ) -> Optional[np.ndarray]: + """ + Get the density values + WARNING: this function aims to work in the specific case of having only 1 CellElementRegion in your XML file + and that this CellElementRegion contains only one cellBlock. + + Parameters + ----------- + densityName : str + Name of density array in GEOS + filterGhost : bool + Filter the ghost ranks + + Returns + ------- + Numpy Array : Array containing the density values + """ + density = self.getSolverFieldWithPrefix( "mediumDensity", **kwargs ) + + if density is not None: + if filterGhost: + density_filtered = self.filterGhostRankFor1RegionWith1CellBlock( density, **kwargs ) + if density_filtered is not None: + return density_filtered + else: + print( "getDensityModel->filterGhostRank: No ghostRank was found.", flush=True ) + return None + else: + return density + else: + print( "getDensityModel: No density was found.", flush=True ) + return None + + def updateDensityModel( self: Self, density: npt.NDArray ) -> None: + """ + Update density values in GEOS + + Parameters + ----------- + density : array + New values for the density + """ + self.setGeosWrapperValueByName( "mediumDensity", value=density, filters=[ self.discretization ] ) + + def setMode( self: Self, mode: str ) -> None: + mode_key = get_matching_wrapper_path( self.geosx, [ 'mode' ] ) + self.geosx.get_wrapper( mode_key ).set_value( mode ) + + def modeling( self: Self, model: npt.NDArray, scale_data: float = 1.0 ) -> npt.NDArray: + + # Make sure we are in modeling mode. + self.setMode( "modeling" ) + self.applyInitialConditions() + + # Send model. + self.updateDensityModel( model ) + + # Run. + while run() != COMPLETED: + pass + + # Get vertical component gz at all stations. + gz = copy.deepcopy( self.getGzAtStations() ) * scale_data + + return gz + + def adjoint( self: Self, nm: int, residue: npt.NDArray ) -> npt.NDArray: + + # Make sure we are in adjoint mode. + self.setMode( "adjoint" ) + self.applyInitialConditions() + + # Send the residues to Geos. + self.setGeosWrapperValueByName( "residue", value=residue ) + + # Run + while run() != COMPLETED: + pass + + # Retrieve adjoint: + # * If any, gather all subdomains. + # * Make sure to remove ghosts. + try: + localToGlobal_key = get_matching_wrapper_path( + self.geosx, [ self.discretization, 'elementSubRegions', 'localToGlobalMap' ] ) + ghost_key = get_matching_wrapper_path( self.geosx, + [ self.discretization, 'elementSubRegions', 'ghostRank' ] ) + adjoint_key = get_matching_wrapper_path( self.geosx, + [ self.discretization, 'elementSubRegions', 'adjoint' ] ) + except Exception as e: + raise RuntimeError( f"Failed to resolve wrapper paths: {e}" ) + + adjoint_field = copy.deepcopy( allgather_wrapper( self.geosx, adjoint_key, ghost_key=ghost_key ) ) + + localToGlobal_noGhost = allgather_wrapper( self.geosx, localToGlobal_key, ghost_key=ghost_key ).astype( int ) + adjoint = np.zeros( nm ) + adjoint[ localToGlobal_noGhost ] = adjoint_field + + return adjoint diff --git a/pygeos-tools/src/geos/pygeos_tools/solvers/InversionUtils.py b/pygeos-tools/src/geos/pygeos_tools/solvers/InversionUtils.py new file mode 100644 index 00000000..9b034b02 --- /dev/null +++ b/pygeos-tools/src/geos/pygeos_tools/solvers/InversionUtils.py @@ -0,0 +1,307 @@ +# ------------------------------------------------------------------------------------------------------------ +# SPDX-License-Identifier: LGPL-2.1-only +# +# Copyright (c) 2016-2024 Lawrence Livermore National Security LLC +# Copyright (c) 2018-2024 TotalEnergies +# Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University +# Copyright (c) 2023-2024 Chevron +# Copyright (c) 2019- GEOS/GEOSX Contributors +# All rights reserved +# +# See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details. +# ------------------------------------------------------------------------------------------------------------ + +import numpy as np +import importlib.util +import warnings +from math import isclose +from typing import Callable, Optional, Tuple + + +def computeResidual( predicted: np.ndarray, observed: np.ndarray ) -> np.ndarray: + """ + Compute the residual vector between predicted and observed data. + + Parameters + ---------- + predicted : np.ndarray + Predicted data vector. + observed : np.ndarray + Observed data vector. + + Returns + ------- + np.ndarray + Residual vector (predicted - observed). + """ + return predicted - observed + + +def computeL2Loss( predicted: np.ndarray, observed: np.ndarray, scale: float = 1.0 ) -> float: + """ + Compute the L2 loss (squared L2 norm) between predicted and observed data. + + Parameters + ---------- + predicted : np.ndarray + Predicted data vector. + observed : np.ndarray + Observed data vector. + scale : float, optional + Scaling factor applied to the loss (default is 1.0). + + Returns + ------- + float + Scaled L2 loss value. + """ + residual = computeResidual( predicted, observed ) + loss = 0.5 * np.linalg.norm( residual )**2 + + return float(scale * loss) + + +def gradientTest( f: Callable[ [ np.ndarray ], float ], + g: Callable[ [ np.ndarray ], np.ndarray ], + m0: np.ndarray, + dm: np.ndarray, + flag_verbose: bool = True, + nk: int = 25, + dk: float = 1.8, + history_filename: Optional[ str ] = None ) -> Tuple[ bool, np.ndarray, np.ndarray ]: + r""" + Tests that the gradient g of an objective function f satisfies the Taylor expansion: + f(m0 + h·Δm) ≈ f(m0) + h·g(m0)ᵀ·Δm + O(h²) + + Parameters + ---------- + f : function + Objective function of the form m->f(m). + g : function + Gradient function of the form m->g(m). + m0 : :obj:`np.ndarray` + Reference model. + dm : :obj:`np.ndarray` + Model perturbation. + flag_verbose : :obj:`logical`, optional + Print convergence report. + nk : :obj:`int`, optional + Number of steps (values of h). + dk : :obj:`float`, optional + Ratio between two consecutive steps. + + Returns + ------- + passed : :obj:`logical` + True if the gradient test passed, False otherwise. + slope : :obj:`np.ndarray` + Slope of the error curve at every step (should be 2. if the jacobian is correct). + history: :obj:`np.ndarray` + Full convergence history (steps and errors), useful for further display. + + """ + f0 = f( m0 ) + g0_dm = np.dot( g( m0 ), dm ) + + e0 = np.zeros( nk ) + e1 = np.zeros( nk ) + h = np.zeros( nk ) + + for k in range( nk ): + if flag_verbose: + print( f"Gradient test: Refinement {k} / {nk}", flush=True ) + + h[ k ] = pow( dk, -( k + 1 ) ) + f1 = f( m0 + h[ k ] * dm ) + + e0[ k ] = abs( f1 - f0 ) / abs( f0 ) + e1[ k ] = abs( f1 - f0 - h[ k ] * g0_dm ) / abs( f0 ) + + if flag_verbose: + print( " h h2 e0 e1", flush=True ) + for k in range( nk ): + print( f'{h[k]:.4e} {h[k]*h[k]:.4e} {e0[k]:.4e} {e1[k]:.4e}', flush=True ) + + # Keep meaningful part of the curve, and estimate its slope... + # Should be 2 since quadratic behaviour is expected. + ind = np.where( ( 1e-15 < e1 ) & ( e1 < 1e-2 ) )[ 0 ] + if ind.size > 1: + slope = np.diff( np.log( e1[ ind ] ) ) / np.diff( np.log( h[ ind ] ) ) + mean_slope = np.mean( slope ) + passed = bool( isclose( mean_slope, 2.0, abs_tol=0.1 ) or np.count_nonzero( slope > 1.9 ) > 2 + or np.all( e1 < 1e-15 ) ) + else: + slope = np.array( [] ) + passed = False + + history = np.vstack( ( h, h * h, e0, e1 ) ) + + if history_filename is not None: + np.savetxt( history_filename, history ) + + return passed, slope, history + + +def plotGradientTest( h: np.ndarray, + e1: np.ndarray, + save_path: Optional[ str ] = None, + style: Optional[ dict ] = None ) -> Optional[ "matplotlib.figure.Figure" ]: + """ + Plot the gradient test error curve and reference trends. + + This function visualizes the convergence behavior of a gradient test by plotting + the first-order Taylor error (e1) against the step size (h) on a log-log scale. + It also overlays reference lines for linear and quadratic convergence trends. + + Parameters + ---------- + h : np.ndarray + Array of step sizes used in the gradient test. + e1 : np.ndarray + Array of first-order Taylor error values. + save_path : str, optional + If provided, the plot will be saved to this file path (e.g., "plot.png"). + style : dict, optional + Dictionary of matplotlib style arguments (e.g., color, marker, linestyle) to customize the error curve. + + Returns + ------- + matplotlib.figure.Figure or None + The matplotlib Figure object if plotting is successful, otherwise None if matplotlib is not available or data is invalid. + """ + + if importlib.util.find_spec( "matplotlib" ) is None: + warnings.warn( "matplotlib is not installed." ) + return None + + import matplotlib.pyplot as plt + + # Filter out invalid values + h = np.asarray( h ) + e1 = np.asarray( e1 ) + valid = np.isfinite( h ) & np.isfinite( e1 ) & ( h > 0 ) & ( e1 > 0 ) + + if not np.any( valid ): + warnings.warn( "No valid data points for plotting." ) + return None + + h_valid = h[ valid ] + e1_valid = e1[ valid ] + + try: + fig, ax = plt.subplots() + except Exception as e: + print( f"Subplot creation failed: {e}" ) + fig = plt.figure() + ax = fig.add_subplot( 1, 1, 1 ) + + style = style or {} + try: + ax.loglog( h_valid, e1_valid, label="Error", **style ) + ax.loglog( h_valid, h_valid, label="Linear trend", linestyle='--' ) + ax.loglog( h_valid, np.power( h_valid, 2 ), label="Quadratic trend", linestyle=':' ) + except Exception as e: + print( f"Plotting failed: {e}" ) + return None + + ax.grid( True ) + ax.legend() + ax.set_title( "Gradient test" ) + ax.set_xlabel( "h" ) + ax.set_ylabel( "Error" ) + + if save_path: + try: + fig.savefig( save_path, bbox_inches='tight' ) + print( f"Plot saved to {save_path}" ) + except Exception as e: + print( f"Saving plot failed: {e}" ) + return None + + return fig + + +def plotGradientTestFromFile( filename: str, + column_h: int = 0, + column_error: int = 3, + **kwargs ) -> Optional[ "matplotlib.figure.Figure" ]: + """ + Load gradient test history from a file and plot the error curve. + + This function reads a saved gradient test history file (typically generated by `gradientTest`) + and plots the first-order Taylor error against the step size using `plotGradientTest`. + + Parameters + ---------- + filename : str + Path to the file containing the gradient test history (as a 2D array). + column_h : int, optional + Index of the column containing step sizes (default is 0). + column_error : int, optional + Index of the column containing first-order Taylor errors (default is 3). + **kwargs : dict + Additional keyword arguments passed to `plotGradientTest`, such as `style` or `save_path`. + + Returns + ------- + matplotlib.figure.Figure or None + The matplotlib Figure object if plotting is successful, otherwise None if matplotlib is not available. + """ + history = np.loadtxt( filename ) + h = history[ column_h, : ] + e1 = history[ column_error, : ] + return plotGradientTest( h, e1, **kwargs ) + + +def adjointTest( A: Callable, + AT: Callable, + x: np.ndarray, + y: np.ndarray, + tol: float = 1e-13, + flag_verbose: bool = True ) -> Tuple[ bool, float ]: + """ + Perform an adjoint test to verify that ≈ . + + Parameters + ---------- + A : function + Function that computes the forward operation A(x). + AT : function + Function that computes the adjoint operation A^T(y). + x : np.ndarray + Input vector for the forward operation. + y : np.ndarray + Input vector for the adjoint operation. + tol : float, optional + Tolerance for the error to consider the test passed. + flag_verbose : bool, optional + If True, prints the test summary. + + Returns + ------- + passed : bool + True if the adjoint test passed, False otherwise. + error : float + Relative error between the inner products. + """ + Ax = A( x ) + ATy = AT( y ) + + IP1 = np.dot( Ax.T, y ) + IP2 = np.dot( x.T, ATy ) + + nAx = np.linalg.norm( Ax, 2 ) + nx = np.linalg.norm( x, 2 ) + nATy = np.linalg.norm( ATy, 2 ) + ny = np.linalg.norm( y, 2 ) + + error = abs( IP1 - IP2 ) / max( nAx * ny, nATy * nx ) + passed = error < tol + if flag_verbose: + print( "\n=== Adjoint Test Summary ===" ) + print( f" = {IP1}" ) + print( f" = {IP2}" ) + print( f"Passed = {passed}" ) + print( f"Error = {error}" ) + + return passed, error diff --git a/pygeos-tools/src/geos/pygeos_tools/solvers/Solver.py b/pygeos-tools/src/geos/pygeos_tools/solvers/Solver.py index e562836d..c28ff3a8 100644 --- a/pygeos-tools/src/geos/pygeos_tools/solvers/Solver.py +++ b/pygeos-tools/src/geos/pygeos_tools/solvers/Solver.py @@ -143,8 +143,8 @@ def initialize( self: Self, rank: int = 0, xml: XML = None ) -> None: raise ValueError( f"The solver type '{self.type}' does not exist in your XML '{self.xml.filename}'." ) + geosState: int = self._getGEOSState() if not self.alreadyInitialized: - geosState: int = self._getGEOSState() if geosState == GEOS_STATE.UNINITIALIZED.value: self.geosx = pygeosx.initialize( rank, self.geosxArgs.getCommandLine() ) self.alreadyInitialized = True @@ -168,6 +168,10 @@ def initialize( self: Self, rank: int = 0, xml: XML = None ) -> None: self.updateOutputs() self.updateTimeVariables() + else: + if geosState == GEOS_STATE.COMPLETED.value: + self.geosx = pygeosx.reinit( self.geosxArgs.getCommandLine() ) + """ Accessors from pygeosx and xml """ diff --git a/pygeos-tools/src/geos/pygeos_tools/solvers/__init__.py b/pygeos-tools/src/geos/pygeos_tools/solvers/__init__.py index 6918bc50..b28890ce 100644 --- a/pygeos-tools/src/geos/pygeos_tools/solvers/__init__.py +++ b/pygeos-tools/src/geos/pygeos_tools/solvers/__init__.py @@ -18,3 +18,5 @@ from geos.pygeos_tools.solvers.ReservoirSolver import ReservoirSolver from geos.pygeos_tools.solvers.Solver import Solver from geos.pygeos_tools.solvers.WaveSolver import WaveSolver +from geos.pygeos_tools.solvers.GravitySolver import GravitySolver +from geos.pygeos_tools.solvers.GravityLinearOpSolver import GravityLinearOpSolver