LUMOS

This is the official code release for LUMOS, a data-physics dual-driven framework for fluorescent molecule design.

Our paper: Multi-objective fluorescent molecule design with a data-physics dual-driven generative framework.

Installation

Running example scripts in the paper relies on three separate Conda environments due to dependency constraints:

• LUMOS (main): model training/inference, data processing, and core utilities
• GPU4PySCF: GPU-accelerated TD-DFT calculations
• ADMET-AI: ADMET property prediction (for cell permeability optimization)

Our configuration includes Ubuntu 22.04 (GNU/Linux x86_64), NVIDIA A100-SXM4-80GB, CUDA 12 and Anaconda 24.9.1.

LUMOS main environment

You can create the main environment from the provided YAML:

conda env create -f environments/lumos.yml
# you can also create environment using conda-lock
# https://github.com/conda/conda-lock
conda-lock install -n lumos environments/lumos_conda-lock.yml

GPU4PySCF and ADMET-AI environments

Please follow the official instructions in their repositories:

• GPU4PySCF: https://github.com/pyscf/gpu4pyscf
• ADMET-AI: https://github.com/swansonk14/admet_ai

We also provide optional YAML files in environments folder for reproducibility.

Data and parameters

Training data and pretrained model checkpoints are hosted on Zenodo:

# checkpoints
curl -L -o checkpoints.zip https://zenodo.org/records/18295513/files/checkpoints.zip # ~324 MB
unzip checkpoints.zip
rm checkpoints.zip # optional

# data
curl -L -o data.zip https://zenodo.org/records/18295513/files/data.zip # ~5.7 MB
unzip data.zip
rm data.zip # optional

By default, the code expects the following directory layout:

./checkpoints/...
./data/...

Reproducing the paper results

We provide Python scripts to reproduce the main results reported in the paper, including:

• Representation learning: reconstruction, dimensionality reduction, and similarity analysis
• Neural network (NN) property prediction: training and inference for the attentive graph predictor (AGP) and the latent surrogate predictor (LSP)
• Physics-anchored property prediction: high-throughput TD-DFT workflow and hybrid model training/inference
• De novo generation: prompt-conditioned generation and gradient-guided generation
• Molecular optimization: global optimization, fragment optimization and cell-permeability optimization

All commands below assume you have downloaded the Zenodo assets, activated the LUMOS main environment, and switched to the repository root.

Representation learning

For reconstruction accuracy test:

bash tools/reconstruct_test.sh /path/to/input.csv /path/to/output.pkl

The input CSV must contain a smiles column. The script writes reconstruction results to output.pkl, including output SMILES, input SMILES and corresponding latent representations.

For dimensionality reduction and similarity analysis, you need to get the latent representations for molecules first:

bash experiments/dim_reduction/inferz.sh /path/to/input.csv /path/to/save_dir/
# expected output: /path/to/save_dir/latents.npz

This will save the latent representations under /path/to/save_dir/. Then run the analysis notebook main.ipynb for t-SNE visualization and similarity analysis.

Neural network (NN) property prediction

To train the NN predictors, use the script train_pred_model.sh.

bash scripts/train_pred_model.sh 

For a full list of CLI arguments, run:

python main_train_unified_pred_model.py --help

Before running inference, update the absolute paths in infer_AGP.json and infer_LSP.json, for example, in infer_LSP.json:

{
    "use_latent": true,
    "use_unified": true,
    "batch_size": 512,
    "vae_config_path": "/your/path/to/lumos/checkpoints/autoencoder/config.pkl",
    "vae_params_path": "/your/path/to/lumos/checkpoints/autoencoder/params.pkl",
    "predictor_config_path": "/your/path/to/lumos/checkpoints/latent_surrogate_predictor_light/config.pkl",
    "predictor_params_path": "/your/path/to/lumos/checkpoints/latent_surrogate_predictor_light/best_params.pkl"
}

The vae_*_path and predictor_*_path fields should point to the corresponding configuration files and model checkpoints on your machine.

To run inference with the NN predictors, use the script infer_pred_model.sh.

bash scripts/infer_pred_model.sh /path/to/infer_config.json /path/to/input_smiles.txt /path/to/output.csv
# expected output: /path/to/output.csv

The input TXT file must contain two columns (whitespace- or tab-separated), corresponding to the fluorophore SMILES and the solvent SMILES:

[SMILES-1]  [SOLVENT_SMILES-1]
[SMILES-2]  [SOLVENT_SMILES-2]
......
[SMILES-N]  [SOLVENT_SMILES-N]

Physics-anchored property prediction

To run the high-throughput TD-DFT pipeline, please update the absolute path in pyscf_python_path.json:

{
    "python_path": "/your/path/to/gpu4pyscf/bin/python"
}

Then run the script parallel_ris_xtb.py:

python tools/parallel_ris_xtb.py --n_nodes 1 --node_id 0 --data_path /path/to/input.csv --save_path /path/to/save_dir/
# expected output: /path/to/save_dir/ris_*

Note that the input CSV must contain columns smiles, solvent and idx:

• smiles: SMILES of the fluorophores
• solvent: SMILES of the solvent molecules (see the supported solvents are in sol_name2smiles.json)
• idx: user-defined integer IDs used to keep outputs in a consistent order

To train the hybrid model, use the script train_delta_learning.sh.

bash scripts/train_delta_learning.sh
# for CLI arguments
python main_train_delta_learning.py --help

To run inference with the hybrid model, please update the absolute paths in infer_hybrid.json first, and use the script infer_delta_learning.sh.

bash scripts/infer_delta_learning.sh /path/to/input.csv /path/to/output.csv

Hybrid model inference requires precomputed TD-DFT results (absorption and emission wavelengths). You can use collect_pyscf.ipynb to collect the PySCF outputs and assemble the input CSV for inference.

De novo generation

Before running the generation experiments, please update the absolute paths in experiment_config.json first:

{
    "dit_config_path": "/your/path/to/lumos/checkpoints/diffusion_transformer/config.pkl",
    "dit_params_path": "/your/path/to/lumos/checkpoints/diffusion_transformer/params_step1000000.pkl",
    "vae_config_path": "/your/path/to/lumos/checkpoints/autoencoder/config.pkl",
    "vae_params_path": "/your/path/to/lumos/checkpoints/autoencoder/params.pkl"
}

Here, vae_*_path and dit_*_path should point to the corresponding autoencoder and diffusion transformer configuration files and checkpoints on your machine.

For prompt-conditioned generation, follow the notebooks in prompt_conditioned_generation:

• conditional_generation.ipynb: single-prompt generation
• conditional_generation_two_inputs.ipynb: two-prompt generation

For gradient-guided generation, please check the absolute path in infer_LSP.json and run the script gen.sh:

# label: abs / emi / loge / plqy
bash experiments/gradient_guided_generation/gen.sh [label]

Molecular optimization

The configuration files, scripts and post analysis notebooks for three molecular optimization cases reported in our paper are provided under cases.

To run the test cases:

# global optimization
python cases/case_molopt/diff_evo.py --config_path cases/case_molopt/config.json --save_path cases/case_molopt

# fragment optimization
python cases/case_fragopt/diff_evo.py --config_path cases/case_fragopt/config.json --save_path cases/case_fragopt

# cell permeability optimization
python cases/case_admet/diff_evo.py --config_path cases/case_admet/config.json --save_path cases/case_admet

# expected outputs (per case directory): log.txt & diffusion_evo_opt.pkl

NOTE: before running the scripts, please update the absolute paths in corresponding config.json file first.

Training the foundation models

The full training pipeline including three steps:

• Training autoencoder model: following train_ae.sh
• Running inference with autoencoder model (to get training data for diffusion transformer): following infer_ae.sh
• Training diffusion transformer model: following train_dit.sh

Citation

@misc{li2026multiobjectivefluorescentmoleculedesign,
      title={Multi-objective fluorescent molecule design with a data-physics dual-driven generative framework}, 
      author={Yanheng Li and Zhichen Pu and Lijiang Yang and Zehao Zhou and Yi Qin Gao},
      year={2026},
      eprint={2601.13564},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2601.13564}, 
}

Contact

For questions or further information, please contact grlyh@pku.edu.cn.