OkadaTorch: Differentiable PyTorch Okada Model

• OkadaTorch is a differentiable PyTorch implementation of the Okada model, translating analytical solutions for elastic dislocations into automatic gradient-computable operations.
• It leverages vectorized tensor operations to compute precise surface displacements and strains for both point and finite rectangular faults.
• The framework supports gradient-based inversion, Bayesian inference, and sensitivity analysis, integrating seamlessly with modern scientific machine learning workflows.

Searching arXiv for the specified paper and closely related references. OkadaTorch is a PyTorch implementation of the Okada model in which the entire computation is differentiable, enabling gradients with respect to inputs to be computed using automatic differentiation. It is presented as a direct translation of the original Okada model into PyTorch together with a wrapper interface for efficiently computing gradients and Hessians with respect to either observation station coordinates or fault parameters (Someya et al., 23 Jul 2025). Within the stated scope of a 3D, homogeneous, isotropic elastic half-space with a traction-free surface, it reproduces the functionality of the standard subroutines SPOINT, SRECTF, DC3D0, and DC3D for point and finite rectangular dislocation sources at or below the surface, while making the full computational graph available to PyTorch autograd (Someya et al., 23 Jul 2025).

1. Position within elastic dislocation modeling

The Okada model is a widely used analytical solution for displacements and strains caused by a point or rectangular dislocation source in a 3D elastic half-space (Someya et al., 23 Jul 2025). OkadaTorch specifically translates the classic Okada analytical solutions for elastic deformation due to shear and tensile dislocations into PyTorch tensor operations, preserving the original subroutine structure while exposing differentiability with respect to model inputs (Someya et al., 23 Jul 2025).

The physical setting is the standard half-space formulation: a finite rectangular dislocation embedded in an elastic half-space, with the Earth modeled as homogeneous, isotropic, linearly elastic, and bounded by a traction-free surface (Someya et al., 23 Jul 2025). The source is parameterized by a rectangular fault with length $L$ along strike and width $W$ down dip, orientation given by strike $\theta$, dip $\delta$, and rake $\lambda$, and depth $d$ defined at the top edge or fault center depending on configuration (Someya et al., 23 Jul 2025). Material dependence is expressed through Poisson’s ratio $\nu$, with Okada’s parameter

$\alpha = \frac{1 - 2\nu}{1 - \nu},$

while shear modulus $\mu$ determines stress scaling but is not required for displacements in the standard formulation (Someya et al., 23 Jul 2025).

A central motivation is methodological rather than geophysical novelty. The implementation was developed to support first- and second-order optimization, sensitivity analysis, Bayesian inference with HMC/NUTS, and end-to-end scientific machine learning models without hand-derived symbolic derivatives (Someya et al., 23 Jul 2025). This suggests that its main contribution lies in computational differentiation and workflow integration rather than in altering the underlying elastic theory.

2. Fault geometry, coordinates, and kinematic parameterization

OkadaTorch adopts a right-handed global coordinate system with $x=\text{east}$, $W$0, and $W$1; at the surface $W$2, and below the surface $W$3 in its global convention (Someya et al., 23 Jul 2025). Observation stations are specified by coordinates $W$4, and the source geometry is referenced to a fault location $W$5 with configurable origin semantics (Someya et al., 23 Jul 2025).

The horizontal transformation from global to strike-aligned coordinates is

$W$6

$W$7

with unit vectors

$W$8

where $W$9 is along strike and $\theta$0 is to the right of strike (Someya et al., 23 Jul 2025). Introducing $\theta$1 so that depth is positive downward, the dip-plane transformation is

$\theta$2

$\theta$3

with $\theta$4 in Okada’s local coordinates (Someya et al., 23 Jul 2025).

For a finite rectangle, the along-strike and down-dip coordinates are denoted $\theta$5 and $\theta$6 (Someya et al., 23 Jul 2025). In common conventions, $\theta$7, $\theta$8, $\theta$9, and $\delta$0, but OkadaTorch allows different origins via fault_origin (Someya et al., 23 Jul 2025). At the wrapper level, fault_origin selects whether the fault location refers to the top-left corner ("topleft", default) or the rectangle center ("center") (Someya et al., 23 Jul 2025).

Slip is decomposed into three components: strike-slip (DISL1), dip-slip (DISL2), and tensile/opening (DISL3) (Someya et al., 23 Jul 2025). The rake $\delta$1 determines this decomposition, and displacements scale with slip (Someya et al., 23 Jul 2025). In the wrapper interface, the parameter dictionary contains scalar tensors for "x_fault", "y_fault", "depth", "length", "width", "strike", "dip", "rake", and "slip" (Someya et al., 23 Jul 2025).

3. Mathematical structure of the forward model

OkadaTorch follows the analytical structure of Okada’s displacement and strain fields as superpositions of Green’s functions for the three slip modes (Someya et al., 23 Jul 2025). In fault-local coordinates, a finite rectangular source is evaluated by summing signed contributions from its four corners. The representative displacement expression given for surface points is

$\delta$2

where the superscripts denote strike-slip, dip-slip, and tensile components, respectively (Someya et al., 23 Jul 2025).

The constituent kernels are built from elementary terms involving

$\delta$3

together with $\delta$4 and $\delta$5, as well as dip-dependent and $\delta$6-dependent coefficients (Someya et al., 23 Jul 2025). A representative strike-slip kernel is written as

$\delta$7

with analogous expressions for dip-slip and tensile components (Someya et al., 23 Jul 2025). The full closed forms are attributed to Okada (1985) for surface points and Okada (1992) for subsurface points (Someya et al., 23 Jul 2025).

After evaluating local components, displacements are rotated back to global coordinates according to

$\delta$8

with $\delta$9 handled consistently with the sign convention $\lambda$0 up in OkadaTorch (Someya et al., 23 Jul 2025).

The strain tensor is defined as

$\lambda$1

and the translated subroutines compute all spatial derivatives $\lambda$2; symmetrization can then be applied if strictly symmetric strain components are required (Someya et al., 23 Jul 2025).

4. Differentiable programming design

A defining feature of OkadaTorch is that all primitive operations—such as addition, multiplication, square roots, logarithms, and arctangents—are implemented as PyTorch tensor operations, so autograd tracks the entire computational graph (Someya et al., 23 Jul 2025). No Python loops over stations are required, and vectorized tensors enable efficient execution on CPU or GPU (Someya et al., 23 Jul 2025).

If is_degree=True, which is the default, angular parameters are converted through differentiable scale factors $\lambda$3, so the rotation matrices remain differentiable functions of strike and dip (Someya et al., 23 Jul 2025). The wrapper exposes derivative computation in two forms. OkadaWrapper.gradient uses torch.func.jacfwd to compute derivatives of outputs with respect to a single argument such as "x", "depth", or "strike" (Someya et al., 23 Jul 2025). OkadaWrapper.hessian computes second derivatives with respect to two arguments by nesting jacfwd, allowing both diagonal and mixed entries (Someya et al., 23 Jul 2025). For scalar losses, ordinary reverse-mode autodiff via loss.backward() computes gradients with respect to any parameter whose tensor has requires_grad=True, and double backward is supported for second derivatives (Someya et al., 23 Jul 2025).

The paper also notes the internal use of torch.func—formerly functorch—in PyTorch 2.x for efficient vectorization through vmap during derivative evaluation (Someya et al., 23 Jul 2025). The broader PyTorch framework is described in the cited systems literature (Paszke et al., 2019). A plausible implication is that the implementation is intended to align geophysical forward modeling with standard differentiable programming idioms already established in machine learning and scientific computing.

The treatment of non-differentiable behavior is conservative. The core algorithms avoid branching on the differentiated quantities, while functions such as atan2 and log are described as differentiable almost everywhere, with the usual caveats near singular points (Someya et al., 23 Jul 2025).

5. Interfaces, outputs, and computational workflow

OkadaTorch provides two interface layers: direct translations of the original subroutines and a higher-level wrapper (Someya et al., 23 Jul 2025). The direct layer includes SPOINT for surface point sources, SRECTF for surface rectangular sources, DC3D0 for subsurface point sources, and DC3D for subsurface rectangular sources (Someya et al., 23 Jul 2025). Compared with the original Fortran, the translated functions add the keyword arguments compute_strain and is_degree (Someya et al., 23 Jul 2025).

For DC3D, the documented call signature is $d$4 where ALPHA = (1 − 2ν)/(1 − ν), out is [ux, uy, uz] or displacements plus strains if requested, each with the shape of X, Y, Z, and IRET is the status integer inherited from the original code (Someya et al., 23 Jul 2025). The values AL1, AL2, AW1, and AW2 correspond to the fault-plane extent parameters of Okada’s DC3D interface and appear explicitly in the core functions (Someya et al., 23 Jul 2025).

The high-level OkadaWrapper provides:

• compute(coords, params, compute_strain=True, is_degree=True, fault_origin="topleft", nu=0.25)
• gradient(coords, params, arg=..., compute_strain=True, is_degree=True, fault_origin="topleft", nu=0.25)
• hessian(coords, params, arg1=..., arg2=..., compute_strain=True, is_degree=True, fault_origin="topleft", nu=0.25) (Someya et al., 23 Jul 2025)

The coords argument is a dictionary containing tensors "x", "y", and optionally "z", all of identical shape and interpreted as batched stations (Someya et al., 23 Jul 2025). The params argument is a dictionary of scalar torch tensors, each optionally carrying requires_grad=True (Someya et al., 23 Jul 2025). Units are required to be consistent for coordinates and geometry, for example kilometers throughout, while slip is specified in meters; outputs are then displacements in meters and strains in displacement per unit length, such as m/km (Someya et al., 23 Jul 2025).

If compute_strain=True, the wrapper returns a list of 12 tensors:

$\lambda$4

from which strain components such as

$\lambda$5

can be assembled (Someya et al., 23 Jul 2025). These derivatives are themselves differentiable with respect to parameters, so quantities such as $\lambda$6 or second derivatives can be obtained through the same wrapper methods (Someya et al., 23 Jul 2025).

6. Performance, validation, and numerical behavior

The implementation is vectorized over stations: the core subroutines accept tensor fields X, Y, and Z, eliminating Python loops and permitting efficient handling of large grids on CPU or GPU (Someya et al., 23 Jul 2025). Its reported computational complexity is roughly $\lambda$7, with the wrapper applying the four-corner summation internally for each rectangular fault (Someya et al., 23 Jul 2025). GPU acceleration follows directly from the use of PyTorch tensor operations; moving inputs and parameters to CUDA tensors is sufficient to execute on GPU (Someya et al., 23 Jul 2025). The paper also notes a memory-compute trade-off: setting compute_strain=False avoids computing and storing intermediate variables needed only for derivatives, reducing compute time and memory when only displacements are needed (Someya et al., 23 Jul 2025).

Validation is described primarily in qualitative and workflow terms. The core is stated to be a direct translation of the standard Fortran DC3D family, and the wrapper is said to demonstrate physically consistent displacement and strain fields at the surface (Someya et al., 23 Jul 2025). The paper shows forward fields for a representative rectangular fault, including $\lambda$8, $\lambda$9, $d$0, and all nine $d$1 fields; first derivatives with respect to each parameter, such as $d$2, illustrating expected sensitivities including linearity in slip; and a synthetic inversion example in which Adam optimization recovers parameters close to ground truth with an interpretable loss history (Someya et al., 23 Jul 2025).

At the same time, the reported evidence has stated limits. Unit tests and formal accuracy metrics versus legacy codes are not detailed in the paper (Someya et al., 23 Jul 2025). This suggests that the validation emphasis is on direct translation fidelity, qualitative field behavior, and end-to-end inversion demonstrations rather than on a published benchmark suite.

The numerical cautions closely mirror those of the classical Okada solution. Because the kernels involve logarithms, arctangents, and $d$3-like terms, evaluations very near a fault edge or at zero distance can be ill-conditioned (Someya et al., 23 Jul 2025). The implementation preserves the original subroutine logic, including the IRET status code, and is described as robust for standard use, but users are advised to avoid placing stations exactly on fault edges and to exercise care with very shallow sources and observation points directly above the upper edge (Someya et al., 23 Jul 2025). For gradient-based inversion, regularization or parameter bounds such as minimum depth or minimum width can help when gradients become large near singular configurations (Someya et al., 23 Jul 2025).

7. Applications, scope, and limitations

The differentiable design enables several use cases explicitly identified in the paper. These include gradient-based optimization of fault parameters from GNSS or InSAR displacements using torch.optim, gradient-informed Bayesian inference such as HMC/NUTS, sensitivity analysis and Fisher/Hessian-based uncertainty quantification including Laplace approximations around optima, and integration with scientific machine learning pipelines in which neural networks output slip distributions and a differentiable Okada layer computes the resulting elastic fields (Someya et al., 23 Jul 2025). The paper further notes possible coupling with differentiable tsunami or crustal flow models (Someya et al., 23 Jul 2025).

A concise summary of these use cases is given below.

Use case	Mechanism in OkadaTorch	Stated purpose
Gradient-based inversion	loss.backward() or wrapper derivatives	Fault parameter optimization
Bayesian inference	Exact model gradients	HMC/NUTS workflows
Uncertainty analysis	Gradients and Hessians	Fisher/Hessian and Laplace approximations
SciML integration	End-to-end backpropagation through physics layer	Neural-network-driven slip or coupled models

The scope is nonetheless restricted by the assumptions of the underlying Okada model. OkadaTorch assumes a homogeneous, isotropic, linear elastic half-space with a traction-free surface (Someya et al., 23 Jul 2025). It does not include layering, topography, viscoelastic or poroelastic effects, rate/state friction, or inelastic processes (Someya et al., 23 Jul 2025). For layered media or time-dependent rheologies, different forward models are required (Someya et al., 23 Jul 2025).

The paper also notes that classical identifiability issues remain: trade-offs such as those between depth and slip, or between geometry and rake, are not removed by differentiability, and gradient-based methods can converge to local minima without good initialization and appropriate regularization or priors (Someya et al., 23 Jul 2025). Future extensions mentioned include multi-layered media, topography, smoothing parameterizations for slip distributions, and richer priors for Bayesian workflows (Someya et al., 23 Jul 2025).

Practically, the code is available in a public repository together with examples, documentation, and example notebooks for reproducing forward and derivative fields and the synthetic inversion (Someya et al., 23 Jul 2025). The repository is intended for use with PyTorch, with PyTorch 2.x recommended for torch.func.jacfwd (Someya et al., 23 Jul 2025). Publication was permitted by NIED, and the programs are stated to differ from those on the NIED site (Someya et al., 23 Jul 2025).