# Data-driven ergodic control from equilibrium samples This repository contains the code used to study data-driven solutions of an ergodic stochastic optimal-control problem for overdamped Langevin dynamics. The method estimates the principal eigen-pair of $$ (\sigma \ell-\mathcal{L})\varphi=\lambda\varphi, $$ from sampled data. The estimated positive eigenfunction determines the optimal bias potential $$ V_{\mathrm{bias}}=-\frac{2}{\beta}\log\varphi. $$ The repository reproduces numerical studies for 1D double-well, 2D lemon-slice, and 2D three-hole benchmark potentials, together with an alanine-dipeptide application. ## Contents The archive contains the following main components: - **Generator eigen-pair estimation:** rejection-sampled and temperature-reweighted data are used to approximate the generator eigenproblem. - **Basis representations:** random Fourier features (RFF) and tensor-product B-spline bases are compared against finite-element reference solutions. - **Neural Cole--Hopf (NCH):** multi-seed calculations assess sensitivity to neural-network initialization for two-dimensional systems. - **Temperature reweighting:** estimation of target-temperature eigen-pairs from data generated at a single reference inverse temperature. - **Closed-loop validation:** learned biases are tested in independent fixed-horizon first-passage-time simulations. - **Alanine dipeptide:** eigen-pairs inferred from equilibrium MD and reweighted metadynamics data are compared and validated in closed loop. ## Data availability The accompanying Zenodo archive contains the source code, generated numerical data, and publication figures. The project should have the following top-level structure after extraction: ```text Eigenvalue_Optimal_Control/ ├── tool/ │ ├── System.py │ ├── eigen_pro.py │ ├── EigenRunner.py │ ├── FEM_Solver.py │ ├── ClosedLoopFPT.py │ ├── NCH.py │ └── Plot_Functions.py │ ├── scripts/ ├── Data/ ├── Fig/ │ ├── parameters.json ├── parameters_ls.json ├── parameters_th.json │ ├── me_rj.py ├── me_reweighted.py ├── me_fem.py ├── me_ad.py ├── convergence_sigma.py ├── fpt.py ├── fpt_AD.py ├── nch.py ├── plot.py │ ├── README.md ├── LICENSE ├── environment_fem.yml ├── requirement.txt └── requirement_cluster.txt ``` Extract code.zip, Data.zip, and Fig.zip into one common project directory: ```bash cd /Eigenvalue_Optimal_Control mv code/tool . mv code/scripts . mv code/parameters.json . mv code/parameters_ls.json . mv code/parameters_th.json . mv code/me_rj.py . mv code/me_reweighted.py . mv code/me_fem.py . mv code/me_ad.py . mv code/convergence_sigma.py . mv code/fpt.py . mv code/fpt_AD.py . mv code/nch.py . mv code/plot.py . ``` The supplied scripts use paths beginning with `Data/` and `Fig/`. These directories should therefore remain directly inside the repository root. The plotting driver automatically creates `Fig/` if it does not already exist. ## Repository structure After extracting the archives into a common project directory, the relevant structure is ```text tool/ Core models, bases, generator estimators, and runners System.py Benchmark potentials and dynamics eigen_pro.py Eigenproblem and reweighting routines EigenRunner.py Parallel parameter-sweep framework FEM_Solver.py FEM reference ClosedLoopFPT.py Fixed-horizon first-passage-time utilities NCH.py Shared NCH multi-seed base class Plot_Functions.py Data loaders abd publication-figure functions parameters*.json Sweep settings for the benchmark systems me_rj.py Direct rejection-sampling eigen-pair calculations me_reweighted.py Temperature-reweighted eigen-pair calculations me_fem.py FEM reference calculations me_ad.py Alanine-dipeptide eigen-pair calculations convergence_sigma.py Small-sigma convergence calculations fpt.py Benchmark fixed-horizon FPT calculations fpt_AD.py Alanine-dipeptide fixed-horizon FPT calculations nch.py NCH robustness scripts plot.py Batch script that reproduces manuscript figures scripts/ Local wrapper scripts and example HPC submission files Data/ Generated eigen-pair, FPT, and reference data ├── Convergence/ │ ├── DW/ │ ├── LS/ │ ├── TH/ │ └── AD/ ├── DW/ │ ├── FEM/ │ ├── RFF/ │ └── Bspline/ ├── LS/ │ ├── FEM/ │ ├── RFF/ │ └── Bspline/ ├── TH/ │ ├── FEM/ │ ├── RFF/ │ └── Bspline/ ├── MultiSeed/ └── AD/ ├── input/ │ ├── MD/ │ ├── MetaD/ │ ├── ala2_pot_app.pkl │ └── ala2_diff.pkl ├── MD/ └── MetaD/ Fig/ Generated figures ``` Here `DW`, `LS`, and `TH` denote the double-well, lemon-slice, and three-hole benchmark systems, respectively. Subdirectories such as `rejection`, `reweight`, `fix_m`, `fix_sigma`, and `fpt` distinguish the sampling protocol and numerical experiment. `Data/Convergence/` contains the small-sigma results used for the convergence analysis. The `Data/AD/input/` directory contains input molecular data and effective-model files, whereas `Data/AD/MD/` and `Data/AD/MetaD/` contain numerical results generated from those inputs. ## Software requirements Three software environments are used for different parts of the study: - `requirement_cluster.txt`: benchmark eigen-pair, temperature-reweighting, and first-passage-time calculations performed on the HPC cluster; - `requirement.txt`: macOS post-processing, manuscript figures, NCH calculations, and molecular-data analysis; - `environment_fem.yml`: finite-element reference calculations only. ### HPC-Cluster environment The benchmark calculations were run with Python 3.9.18. Load an appropriate Python version on the cluster, then create and activate a virtual environment: ```bash python3.9 -m venv .venv-py39 source .venv-py39/bin/activate python -m pip install --upgrade pip python -m pip install -r requirement_cluster.txt ``` The tested core package versions are: ```text numpy==1.26.4 scipy==1.13.1 sympy==1.14.0 joblib==1.5.3 scikit-learn==1.3.2 matplotlib ``` ### Local macOS environment Plotting, NCH calculations, alanine-dipeptide analysis and figure generation. were run with Python 3.11.9. ```bash python -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -r requirements.txt ``` The tested package versions are: ```text numpy==2.4.4 scipy==1.17.1 sympy==1.14.0 matplotlib==3.10.9 scikit-learn==1.9.0 mdtraj==1.11.1.post2 jax==0.10.2 jaxlib==0.10.2 ``` Some saved files contain serialized scikit-learn objects. For best compatibility, use the environment corresponding to the calculation that produced the data. ### FEM environment The FEM reference calculations require a separate FEniCSx environment. ```bash conda env create -f environment_fem.yml conda activate fenicsx python -m pip install gmsh==4.15.2 ``` This environment is required **only** for `me_fem.py`. ## Parameters and job-array indices The parameter files ```text parameters.json parameters_ls.json parameters_th.json ``` define the numerical sweeps for DW, LS, and TH. Important entries include: - **beta_list:** inverse temperatures; - **sigma_list:** running-cost prefactors; - **sigma_list_0:** small-sigma values used in the convergence study. - **m_list:** data-set sizes; - **ndim:** number of state-space dimension; - **nf:** number of random Fourier features; - **nsim:** number of independent realizations; - **tol:** Relative singular-value threshold used when whitening the basis representation. - **dt_rw:** Time step used to simulate trajectories for the temperature-reweighting calculations. - **beta0:** Reference inverse temperature at which reweighting trajectories are generated. - **DATA_SEED:** base random seed for reproducible sampled dara; ### `fix_m` jobs In `fix_m` mode, the data-set size is fixed and the job index selects one pair (beta, sigma): ```text beta_index = job_id // len(sigma_list) sigma_index = job_id % len(sigma_list) ``` ### `fix_sigma` jobs In `fix_sigma` mode, the cost prefactor is fixed and the job index selects one pair (beta, m): ```text beta_index = job_id // len(m_list) m_index = job_id % len(m_list) ``` Thus, job 0 selects the first entry in both lists. The job index advances through all values of sigma or m before advancing to the next value of beta. The fixed-horizon FPT calculations use the same (beta, sigma) principle. The precise parameter subsets used for FPT are defined in fpt.py. The job index can be passed directly: ```bash SYSTEM=DW python fpt.py 0 ``` On Slurm, the scripts read `SLURM_ARRAY_TASK_ID` automatically. ## Reproducing the manuscript figures The archived numerical results are sufficient to reproduce the manuscript figures without rerunning the eigen-pair or closed-loop simulations. Generate one figure: ```bash python plot.py 1 ``` Generate selected figures: ```bash python plot.py 1 2 3 4 5 6 ``` Generate every figure: ```bash python plot.py all ``` or simply ```bash python plot.py ``` The plotting driver loads data **only for the requested figures**. It is therefore not necessary to have the complete numerical archive available when reproducing only a subset of the manuscript figures. The current figure numbering is: 1. double-well illustrative example; 2. benchmark potential comparison; 3. RFF-gEDMD eigenvalue calculations; 4. B-spline-gEDMD eigenvalue calculations; 5. Small-sigma convergence and equilibrium-average recovery. 6. double-well RFF potential reconstruction; 7. double-well B-spline potential reconstruction; 8. lemon-slice potential comparison; 9. three-hole potential comparison; 10. neural Cole--Hopf potential comparison; 11. benchmark first-passage-time comparison; 12. temperature-reweighting results; 13. alanine-dipeptide free-energy comparison; 14. alanine-dipeptide spectral and potential comparison, including small-sigma convergence. 15. alanine-dipeptide first-passage-time comparison. Generated figures are written to `Fig/` ## Reproducing benchmark eigen-pair calculations The following commands are required only when the numerical results are to be recomputed from scratch. They are **not required** for reproducing figures from the archived numerical data. ### FEM reference data Activate the FEM environment and run ```bash conda activate fenicsx python me_fem.py DW python me_fem.py LS python me_fem.py TH ``` ### Direct rejection-sampling gEDMD Run one RFF calculation: ```bash SYSTEM=DW BASIS=RFF MODE=fix_m SIZE_DATA=2000 python me_rj.py 0 ``` Run one B-spline calculation: ```bash SYSTEM=DW BASIS=BSP MODE=fix_m SIZE_DATA=10000 python me_rj.py 0 ``` For `fix_sigma` calculations, omit `SIZE_DATA`: ```bash SYSTEM=DW BASIS=RFF MODE=fix_sigma python me_rj.py 0 ``` Replace `DW` with `LS` or `TH` as needed. ### Temperature-reweighted RFF gEDMD Run one reweighted calculation: ```bash SYSTEM=DW python me_reweighted.py 0 ``` Replace `DW` with `LS` or `TH`. ### Optional local wrappers for computing eigen-pair Run Direct-sampling sweeps sequentially: ```bash bash scripts/run_me_rj.sh DW RFF fix_sigma bash scripts/run_me_rj.sh LS RFF fix_sigma bash scripts/run_me_rj.sh TH RFF fix_sigma bash scripts/run_me_rj.sh DW RFF fix_m 2000 bash scripts/run_me_rj.sh DW RFF fix_m 10000 bash scripts/run_me_rj.sh LS RFF fix_m 500000 bash scripts/run_me_rj.sh TH RFF fix_m 500000 ``` For B-splines: ```bash bash scripts/run_me_rj.sh DW BSP fix_sigma bash scripts/run_me_rj.sh LS BSP fix_sigma bash scripts/run_me_rj.sh TH BSP fix_sigma bash scripts/run_me_rj.sh DW BSP fix_m 10000 bash scripts/run_me_rj.sh LS BSP fix_m 500000 bash scripts/run_me_rj.sh TH BSP fix_m 500000 ``` Run the temperature-reweighted sweeps sequentially: ```bash bash scripts/run_me_reweighted.sh DW bash scripts/run_me_reweighted.sh LS bash scripts/run_me_reweighted.sh TH ``` ### Small-sigma congvergence data The convergence study uses `sigma_list_0` and RFF-gEDMD in `fix_m` mode only. The fixed settings are: - `DW`: `SIZE_DATA=10_000` and `gamma=0.3` - `LS`: `SIZE_DATA=500_000` and `gamma=0.3` - `TH`: `SIZE_DATA=500_000` and `gamma=0.5` The three-hole calculation additionally uses an enlarged domain and 200 RFF features. Run one local calculation: ```bash SYSTEM=DW python convergence_sigma.py 0 SYSTEM=LS python convergence_sigma.py 0 SYSTEM=TH python convergence_sigma.py 0 ``` To run all benchmark convergence calculations sequentially: ```bash bash scripts/run_convergence.sh ``` ### Fixed-horizon first-passage-time calculations Benchmark FPT calculations require the corresponding direct-sampling eigen-pair data because they load the reconstructed controller from `Data/`. The required order is: ```text FEM reference data ↓ Direct-sampling eigen-pair data: me_rj.py ↓ Benchmark FPT data: fpt.py ↓ Figures: plot.py ``` Run one FPT job: ```bash SYSTEM=DW python fpt.py 0 ``` Use `SYSTEM=LS` or `SYSTEM=TH` for the other benchmark systems. Optional sequential wrappers are provided: ```bash bash scripts/run_fpt_dw.sh bash scripts/run_fpt_ls.sh bash scripts/run_fpt_th.sh ``` To run all benchmark FPT calculations sequentially: ```bash bash scripts/run_all_fpt.sh ``` The wrappers are optional conveniences for local execution; they are not required when using Slurm arrays. ## HPC job-array submission Example Slurm submission scripts are provided under `scripts/`. Before submitting, adjust the project directory, virtual-environment activation, partition, runtime, memory, CPU allocation, and email settings for your HPC system. Create a directory for job logs: ```bash mkdir -p logs ``` ### Combined benchmark eigen-pair and reweighting array submits the full benchmark eigen-pair and temperature-reweighting calculations with: ```bash sbatch scripts/submit_all_me.slurm ``` The script uses 16 experiment groups with 64 parameter jobs each: ```text 16 experiment groups × 64 parameter jobs = 1024 tasks ``` The global Slurm array index is mapped to ```text GROUP = global_id // 64 JOB_ID = global_id % 64 ``` where `JOB_ID` is the local parameter index passed to the Python calculation. The 16 groups are: ```text Group 0 DW RFF fix_sigma Group 1 LS RFF fix_sigma Group 2 TH RFF fix_sigma Group 3 DW RFF fix_m, size_data = 2,000 Group 4 DW RFF fix_m, size_data = 10,000 Group 5 LS RFF fix_m, size_data = 500,000 Group 6 TH RFF fix_m, size_data = 500,000 Group 7 DW BSP fix_sigma Group 8 LS BSP fix_sigma Group 9 TH BSP fix_sigma Group 10 DW BSP fix_m, size_data = 10,000 Group 11 LS BSP fix_m, size_data = 500,000 Group 12 TH BSP fix_m, size_data = 500,000 Group 13 DW temperature reweighting Group 14 LS temperature reweighting Group 15 TH temperature reweighting ``` Within each Slurm task, the script resets ```text SLURM_ARRAY_TASK_ID ``` to the local job index `0,...,63` before calling `me_rj.py` or `me_reweighted.py`. This is necessary because these Python scripts read the Slurm array index directly. The supplied script currently uses ```text #SBATCH --array=0-1023%10 ``` so at most ten array tasks run concurrently. It also sets ```text N_JOBS = SLURM_CPUS_PER_TASK ``` so each Python calculation can use the CPUs assigned to that Slurm task. ### Benchmark convergence calculations Submit the convergence calculations for all three benchmark systems with: ```bash sbatch scripts/submit_convergence.slurm ``` The array contains: ```text 3 systems × 8 beta values × 7 small-sigma values = 168 tasks ``` The global array index is mapped as: ```text GROUP = global_id // 56 JOB_ID = global_id % 56 ``` The task ranges are: ```text 0–55 DW 56–111 LS 112–167 TH ``` The submission script resets `SLURM_ARRAY_TASK_ID` to the local `JOB_ID` because convergence_sigma.py reads that value as its parameter index. ### Benchmark FPT calculations Submit FPT jobs **only after the required direct-sampling eigen-pair calculations have been generated**: ```bash sbatch scripts/submit_all_fpt.slurm ``` The FPT array contains 42 selected \((\beta,\sigma)\) combinations for each of the three benchmark systems, giving 126 tasks in total: ```text global tasks 0--41 -> DW, local jobs 0--41 global tasks 42--83 -> LS, local jobs 0--41 global tasks 84--125 -> TH, local jobs 0--41 ``` The Slurm script maps each global task to its system and local FPT job index automatically. ## Neural Cole--Hopf robustness calculations The neural Cole--Hopf calculations approximate the positive eigenfunction as $$ \varphi_\theta(x)=\exp[-W_\theta(x)]. $$ Run the multi-seed calculations with: ```bash python nch.py ``` and the archived multi-seed results are stored under ```text Data/MultiSeed/ ``` The saved results contain the numerical configuration, evaluation grid, seed-wise eigenvalues, reconstructed fields, training histories, and trained network parameters. ## Alanine-dipeptide application The alanine-dipeptide calculations compare eigen-pairs inferred from equilibrium MD data and reweighted MetaD data at $T=300\,\mathrm{K}$. Required input data are placed under: ```text Data/AD/input/MD/ Data/AD/input/MetaD/ ``` The input archive includes MD and MetaD samples, MetaD weights, and fitted effective-potential and diffusion-model files. ### Eigen-pair calculations Run the two data sources separately: ```bash python me_ad.py MD python me_ad.py MetaD ``` The MD calculation can be memory intensive. Limit the number of workers when needed: ```bash N_JOBS=1 python me_ad.py MD ``` ### Alanine-dipeptide convergence The convergence calculation uses: ```text SIGMA_KEY=sigma_list_0 ``` to select the small-sigma list. Submit the convergence calculation with: ```bash sbatch scripts/submit_ad_convergence.slurm ``` The job runs the MD and MetaD calculations sequentially: ```bash SIGMA_KEY=sigma_list_0 N_JOBS=4 python me_ad.py MD SIGMA_KEY=sigma_list_0 N_JOBS=4 python me_ad.py MetaD ``` Reduce `N_JOBS` if the node does not provide sufficient memory. The results are written under: ```text Data/Convergence/AD/MD/ Data/Convergence/AD/MetaD/ ``` ### Alanine-dipeptide FPT calculations Run one FPT job for each source: ```bash SOURCE=MD python fpt_AD.py 0 SOURCE=MetaD python fpt_AD.py 0 ``` Run all sigma values sequentially: ```bash bash scripts/run_fpt_ad.sh MD bash scripts/run_fpt_ad.sh MetaD ``` The alanine-dipeptide FPT calculations are computationally expensive and may require a day or longer, depending on the hardware. ## Figure dependencies When using the archived data, figures can be generated directly with: ```bash python plot.py FIGURE_NUMBER ``` When regenerating results from scratch, first run the calculation that provides the required data: -`1-4,6-9`: Direct-sampling eigen-pair calculations `me_rj.py` and run `python plot.py 1 2 3 4 6 7 8 9` -`5`: Small-sigma benchmark convergence `convergence_sigma.py` and run `python plot.py 5` -`10`: Neural Cole--Hopf calculations `nch.py` and run `python plot.py 10` -`11`: benchmark FPT calculations `fpt.py` and run `python plot.py 11` -`12`: Temperature-reweighted eigen-pair calculations `me_reweighted.py` and run `python plot.py 12` -`13`: Alanine-dipeptide input data and run `python plot.py 13` -`14`: Alanine-dipeptide eigen-pairs and small-sigma convergence `me_ad.py` and run `python plot.py 14` -`15`: Alanine-dipeptide FPT calculations `fpt_AD.py` and run `python plot.py 15` For example, to regenerate Figure 5 from scratch: ```bash bash scripts/run_convergence.sh python plot.py 5 ``` To regenerate Figure 11 from scratch: ```bash bash scripts/run_me_rj.sh DW BSP fix_m 10000 bash scripts/run_me_rj.sh LS BSP fix_m 500000 bash scripts/run_me_rj.sh TH BSP fix_m 500000 bash scripts/run_all_fpt.sh python plot.py 11 ``` ## Recommended full regeneration order ```text 1. FEM reference calculations 2. Direct rejection-sampling eigen-pair calculations 3. Small-sigma convergence calculations 4. Benchmark FPT calculations 5. Temperature-reweighting calculations 6. Neural Cole--Hopf calculations 7. Alanine-dipeptide eigen-pair and FPT calculations 8. Manuscript figures ``` For figure reproduction from the Zenodo archive, only the final step is needed: ```bash python plot.py ```