A JAX-accelerated Python framework for seismic source inversion and uncertainty quantification using JMA intensity.
This tool is a high-performance Python framework designed for Bayesian inversion of earthquake epicenters and moment magnitudes (Mw) specifically tailored for the Japanese JMA instrumental seismic intensity (Shindo) scale. Following the probabilistic inverse theory of Tarantola (2005), it fully accounts for uncertainties in both the measured data (seismic intensity) and the theoretical model (ground motion prediction equation, GMPE). The system evaluates the posterior Probability Density Function (PDF) for the epicenter location and magnitude via an exhaustive 3D grid search, ensuring no local minima are missed. By leveraging JAX for massive parallelization, modern it allows for near-instantaneous evaluation on both CPU and GPU, making it ideal for both real-time seismology and the processing of large-scale historical earthquake catalogs. Key Features:
- HPC Ready: Fully vectorized backend using JAX (XLA) with seamless support for CPU/GPU/TPU acceleration.
- Bayesian Framework: Complete 3D PDF evaluation accounting for both observational uncertainties and modeling errors.
- Scientific GMPE: Implements the Morikawa & Fujiwara (2013) Ground Motion Prediction Equation for forward computation.
- Automatic Site Effects: Integrated workflow for automated VS30 retrieval from a high-resolution J-SHIS-derived SQL database for Japan (for sites in Japan without direct VS30 measurements).
- Historical & Modern Data: Support for both recent instrumental records and macroseismic (historical) observations.
- User-Friendly: Simple ASCII input/output, PEP8 compliant, and structured for researchers and Python-beginners.
The inversion follows the probabilistic inverse theory and total error budget accounting for both observational and theoretical uncertainties as described in Tarantola (2005). The JMA intensity prediction is following Morikawa and Fujiwara (2013).
Tarantola, A. (2005, Chapter 7.1). Inverse Problem Theory and Methods for Model Parameter Estimation, Society for Industrial and Applied Mathematics, Philadelphia, USA.
Morikawa, N., Fujiwara, H. (2013). A New Ground Motion Prediction Equation for Japan Applicable up to M9 Mega-Earthquake, J. Disaster Res., 8(5), 878-888. https://doi.org/10.20965/jdr.2013.p0878
If missing VS30 values are detected, the system automatically interfaces with an optimized SQLite subset of the J-SHIS-derived database for Japan (Hallo, 2026):
Hallo, M. (2026). Research Dataset: Optimized Site Parameters (Vs30) for Seismic Hazard Analysis in Japan (derived from J-SHIS) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.19379171
The computational engine is engineered for maximum throughput by bypassing standard Python execution loops in favor of XLA (Accelerated Linear Algebra):
- JIT Compilation: Every critical path, from the Forward GMPE evaluation to the Likelihood summation, is Just-In-Time (JIT) compiled. This transforms Python code into optimized machine code tailored for the specific hardware (CPU or GPU).
- Multi-Device Scaling: Leveraging JAX allows the 3D grid-search to be offloaded to GPU/TPU without code changes. On multi-core CPUs, it utilizes all available threads via vectorized operations rather than standard multiprocessing.
- Vectorized Grid Search: Instead of iterative loops, the framework uses nested vectorization (
vmap). The 3D parameter space is treated as a high-dimensional tensor, allowing the hardware to evaluate thousands of potential epicenters and magnitudes in a single clock cycle. - SQLite Spatial Indexing: Automatic retrieval is powered by a high-performance SQLite backend. This allows for rapid spatial lookups within a processed database, ensuring that even datasets with missing site data are enriched with industry-standard values.
- PEP8 & Linting: The codebase is strictly linted using Flake8, ensuring high maintainability and integration readiness for enterprise environments.
- Modular I/O: Designed with a decoupled architecture, allowing the core inversion engine to be wrapped into automated APIs or triggered via CLI as part of a larger pipeline.
The official software version is archived on Zenodo:
Cite as: Hallo, M. (2026). JAX-accelerated Bayesian earthquake location from JMA seismic intensity (v1.0). Zenodo. https://doi.org/10.5281/zenodo.19603409
utils/— Directory containing supporting Python modules and SQL interfacelocation_jma_intensity.py— Main execution script for the Bayesian inversionconfig.py— Configuration module to define inversion parameters and search areaplotting.py— Python module for high-quality visualization of results and PDFsINPUT.txt— Example input file with observed JMA intensity data and VS30 valuesrequirements.txt— Pip requirements file for automated installation of dependencies
Python: Version 3.12 or higher
Libraries: jax, numpy, matplotlib, pandas, geopandas, sqlalchemy, requests, shapely, psutil
Install dependencies via pip:
pip install -r requirements.txt- Prepare your
INPUT.txtinput file (see the detailed format in the comments ofconfig.py) - Set up the search area and inversion parameters in the
config.pyfile - Execute the inversion:
python location_jma_intensity.py - Check the
resultsdirectory for output figures and the summary text file
The computation process is monitored, and the tool informs the user in real-time about the progress. See the example below:
[*] Read input file
[*] Prepare input data
[*] Found 34 invalid Vs30 values
[*] Check or download SQLite database
[*] Initialize SQL engine
[*] Extracting missing Vs30 data from SQL
[+] SUCCESS: Vs30 values assigned from SQL database
[+] SUCCESS: Updated file saved to: UPDATED_INPUT.txt
[*] Vectorize data for JAX
[*] Memory Usage Forecast
Grid dimensions: 1001 x 1001 x 201
Number of stations: 163
Final 3D matrix (RAM): 0.81 GB
JAX peak per slice: 0.65 GB
Available System RAM: 5.97 GB
[*] SUCCESS: Memory check passed
[*] Evaluate probability in 3D model space (Parallel-JAX)
Progress: 30.3%
Progress: 60.2%
Progress: 90.0%
Progress: 100.0%
[*] SUCCESS: 3D PDF evaluation finished
[*] Estimating ML/MAP solution
[*] Estimating solution uncertainties from marginal PDFs
[*] Compute Posterior Mean (PM) solution
[*] Save results to a text file
[+] SUCCESS: Results saved to: results/20260416_144319_loc.txt
[*] Evaluating misfit for the ML/MAP solution
[*] Plot results
[+] SUCCESS: Saved to results/20260416_144319_loc_map.png
[+] SUCCESS: Saved to results/20260416_144319_loc_pdf_cross_section.png
[+] SUCCESS: Saved to results/20260416_144319_loc_pdf_marginal.png
[+] SUCCESS: Saved to results/20260416_144319_loc_misfit.png
[*] SUCCESS: All done
Regarding the results, the figure below illustrates the output 3D posterior Probability Density Function (PDF) with orthogonal slices (Mw, N-S, E-W) passing through the Maximum Likelihood (ML) solution.
The tool saves results into a structured text file. It contains precise coordinates for the Maximum Likelihood (ML) or Maximum A Posteriori (MAP) solutions (which are identical in this case), including standard deviations (1σ) and double standard deviations (2σ) to quantify spatial uncertainty. Additionally, the output includes the Posterior Mean (PM) solution, which is particularly useful for characterizing the location in cases of highly asymmetric posterior PDFs.
# SOLUTION FOR THE EARTHQUAKE EPICENTER LOCATION
# --------------------------------------------------------------------
# Maximum Likelihood (ML) / Maximum a Posteriori (MAP)
# Latitude, Longitude, Mw, Easting, Northing, E_sigma, N_sigma, Mw_sigma, E_2sigma, N_2sigma, Mw_2sigma [km]
35.05460 135.66053 6.36 5.550 -5.000 4.442 6.517 0.216 8.885 13.035 0.431
# --------------------------------------------------------------------
# Posterior Mean solution (PM)
# Latitude, Longitude, Mw, Easting, Northing [km]
35.06636 135.65568 6.46 5.099 -3.699
Copyright (C) 2026 Kyoto University
This program is published under the GNU General Public License (GNU GPL).
This program is free software: you can modify it and/or redistribute it or any derivative version under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
This code is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY. We would like to kindly ask you to acknowledge the authors and don't remove their names from the code.
You should have received copy of the GNU General Public License along with this program. If not, see http://www.gnu.org/licenses/.