FEX-DM

This repository contains the code implementation for the paper "Identifying stochastic dynamics via Finite Expression Methods".

FEX-DM is a machine learning two-stage framework that models stochastic differential equations (OU process, Double Well Potential, etc).

Table of Contents

• Overview
• Requirements
• Installation
• Usage
• Citation
• License
• Contact

Overview

FEX-DM is a two-stage framework for learning stochastic differential equations (SDEs):

1. Stage 1 (Deterministic): Learn the drift term using Finite Expression (FEX) methods
2. Stage 2 (Stochastic): Learn the diffusion term using neural networks

The framework can handle various SDE models including Ornstein-Uhlenbeck (OU) processes, Double Well Potential, and more.

Requirements

• Python 3.8+
• PyTorch
• NumPy
• SciPy
• Matplotlib
• SymPy
• FAISS (for nearest neighbor search)

Installation

1. Clone the repository:

git clone <repository-url>
cd FEX-DM

2. Set up environment:

Stage 0: Set up env

Run config.py in the src folder. The required packages will be automatically installed.

python src/config.py

Usage

Stage 1: Deterministic Training

Train the FEX model to learn the drift term:

python src/1stage_deterministic.py

Stage 2: Stochastic Training

Train the neural network to learn the diffusion term:

python src/2stage_stochastic_time_independent.py

Citation

If you use FEX-DM, please cite the following paper:

@article{LIANG2026114756,
title = {Identifying stochastic dynamics via finite expression methods},
journal = {Journal of Computational Physics},
volume = {554},
pages = {114756},
year = {2026},
issn = {0021-9991},
doi = {https://doi.org/10.1016/j.jcp.2026.114756},
url = {https://www.sciencedirect.com/science/article/pii/S0021999126001063},
author = {Senwei Liang and Chunmei Wang and Xingjian Xu},

License

This project is licensed under the MIT License - see the LICENSE file for details.

Copyright (c) 2025 Xingjian Xu and Senwei Liang

Contact

For questions, or issues:

• Email: [email protected], [email protected]
• Issues: Please use the GitHub Issues page for bug reports