kim-jax 1.3.0

A package for developing knowledge informed inverse or forward mapping using ensemble learning

Knowledge-Informed Mapping (KIM) Toolkit

Table of Contents

• Overview
• Installation
• Documentation
• Examples
• License
• Acknowledgments
• How to Cite
• Contacts

Overview

KIM is a Knowledge-Informed Mapping toolkit in Python to optimize the development of the mapping $ƒ$ from a vector of inputs $\mathbf{X}$ to a vector of outputs $\mathbf{Y}$. KIM mainly builds on the methodology development of deep learning-based inverse mapping in Jiang et al. 2023 and Wang et al. 2025. It involves two key procedures: (1) an exploratory data analysis using information theory to identify the dependency between $\mathbf{X}$ and $\mathbf{Y}$ and filter out both insignificant and redundant inputs through global sensitivity analysis and conditional independence testing; and (2) ensemble learning of $ƒ$ using neural networks to account for its structural uncertainty. KIM is mostly rewritten in JAX and also supports basic parallel computing on CPU cores for statistical significance test and ensemble learning by using Joblib.

Installation

1. Install Miniconda if you don't have it already.

2. Open a terminal and create a new virtual environment named kim with Python 3.11:

   conda create --name kim python=3.11
   

3. Activate the newly created virtual environment:

   conda activate kim
   

4. Install the package using pip within the activated virtual environment:

   pip install kim-jax
   

5. (Optional) Download the git repo to get the example jupyter notebooks:

   git clone https://github.com/PeishiJiang/KIM.git
   

Documentation

The official documentation is hosted on the package website. Please refer to Math behind KIM for a complete description of the theory behind the package.

Examples

We provide one tutorial case and two real cases of applying KIM to performing inverse modeling by using Jupyter notebook to illustrate the package usage.

Case 0: Emulating a multivariate nonlinear system. We develop forward mappings to emulate three predictand driven by four predictors via a nonlinear system.

Case 1: Calibrating a cloud chamber model. We develop inverse mappings to estimate two key parameters, i.e., wall fluxes ($\lambda_w$) and collision processes ($\lambda_c$) of a cloud chamber model from synthetic observations.

Case 2: Calibrating an integrated hydrological model. We develop inverse mappings to estimate eight parameters of the Advanced Terrestrial Simulator (ATS) from the streamflow observations at the outlet of Coal Creek watershed, CO, USA.

License

Distributed under the Simplified BSD License. See LICENSE for more information.

Acknowledgements

This work was funded by the Laboratory Directed Research and Development Program at Pacific Northwest National Laboratory and the IDEAS-Watersheds project.

How to Cite

The repository is under review. We will provide a complete citation upon the acceptance of the repo/paper.

Contacts

Peishi Jiang ([e-mail])

---

Go to Top