TIDE-HMS
TIDE-HMS is a Trajectory-Informed Hierarchical Multi-Scale model for predicting electrolyte density, dielectric constant, viscosity, diffusion coefficient, and ionic conductivity. The method transfers physical information from molecular-dynamics trajectories into a formula encoder through aligned supervision at the atom, functional-group, molecule, and formula scales. The resulting encoder can then predict electrolyte properties from composition and thermodynamic conditions without requiring a trajectory at inference time.
This repository contains only the final TIDE-HMS method: a hierarchical molecular graph encoder, a condition-aware mixture encoder, aligned four-scale trajectory pretraining with hard-freeze routing, and supervised five-property fine-tuning.
Model Framework
Overview
TIDE-HMS uses molecular trajectories only during pretraining. Given an electrolyte formula, the model first encodes each solvent and salt as a molecular graph, constructs hierarchical representations, and aggregates all components with composition and condition information. Four trajectory-derived physical signals supervise the corresponding representation levels. After pretraining, the trajectory heads are removed and the formula representation is fine-tuned for five-property regression.
1flowchart LR
2 A["Electrolyte formula<br/>SMILES, ratios, temperature, concentration"] --> B["Hierarchical molecular GNN"]
3 B --> C1["Atom state"]
4 B --> C2["Functional-group state"]
5 B --> C3["Molecule state"]
6 C3 --> D["Component type + ratio + condition encoding"]
7 D --> E["Transformer mixture encoder<br/>and gated pooling"]
8 E --> C4["Formula state"]
9
10 T["MD trajectory: 100 sampled frames"] --> T1["Li-donor coordination"]
11 T --> T2["Anion-donor fraction"]
12 T --> T3["Molecular MSD"]
13 T --> T4["Dynamic heterogeneity"]
14
15 T1 -. aligned supervision .-> C1
16 T2 -. aligned supervision .-> C2
17 T3 -. aligned supervision .-> C3
18 T4 -. aligned supervision .-> C4
19
20 C4 --> P["Five-property regression head"]
1. Hierarchical Molecular Representation
Each electrolyte component is converted from SMILES into a molecular graph. Atom features describe element type, degree, formal charge, ring membership, aromaticity, and hybridization; bond features describe bond order, ring membership, and conjugation. A three-layer GINE-style message-passing network produces node embeddings that are summarized into three molecular representation levels:
- Atom state: mean readout over all graph nodes, representing local atomic environments.
- Functional-group state: readout over donor, charged, or ring-associated atoms, representing chemically active local environments.
- Molecule state: joint projection of the atom state, functional-group state, and global molecular node statistics.
This hierarchy provides explicit representation states to which trajectory quantities at different physical scales can be aligned.
2. Formula and Condition Encoder
For a multi-component electrolyte, TIDE-HMS combines each molecule state with a solvent/salt type embedding and a continuous molar-ratio encoding. Temperature and salt concentration are encoded as formula-level conditions and added to every component token. A Transformer encoder models interactions among components, after which condition-aware gated pooling produces a single formula representation. The pooled representation is fused with the mixture-level molecule state and concatenated with the thermodynamic conditions.
The released architecture uses a 128-dimensional hierarchical GNN, three graph layers, one four-head Transformer layer, and a 256-dimensional property backbone.
3. Four-Scale Trajectory Supervision
TIDE-HMS samples 100 frames from each molecular-dynamics trajectory and extracts four frame-wise scalar targets:
| Target | Aligned representation | Physical quantity | Current definition |
|---|
atom | Atom state | Li-donor local coordination strength | Mean reciprocal-distance sum over the nearest donor atoms around Li |
functional_group | Functional-group state | Anion participation in the Li coordination shell | Anion-donor count divided by total anion- and solvent-donor count within the cutoff |
molecule | Molecule state | Translational motion of mobile molecules | Mean squared displacement of mobile molecular centers of mass |
formula | Formula state | Formula-level dynamic heterogeneity | Standard deviation of mobile-molecule squared displacements |
Each representation is connected to an independent trajectory prediction head that reconstructs its 100-frame target sequence. The four Huber losses are weighted equally by default.
4. Aligned Hard-Freeze Pretraining
During Stage I, TIDE-HMS optimizes one aligned scale at a time. When a trajectory target is processed, only its corresponding representation module and prediction head are trainable; all unrelated modules are temporarily frozen. For example, atom supervision updates the graph feature layers, message-passing layers, atom readout, atom projection, and atom trajectory head, whereas formula supervision updates the condition, ratio, type, Transformer, gated-pooling, formula projection, and formula trajectory head modules.
This hard-freeze routing prevents a trajectory target from being absorbed by an unrelated representation level and makes the target-to-state correspondence explicit. The four-scale pretrained backbone is subsequently transferred to property prediction.
5. Five-Property Fine-Tuning
During Stage II, the trajectory prediction heads are discarded and the complete pretrained backbone is connected to a multi-output regression head. The model jointly predicts:
- density;
- dielectric constant;
- viscosity;
- diffusion coefficient;
- ionic conductivity.
The default training schedule uses at most 600 pretraining epochs and 600 fine-tuning epochs. Both stages use validation-loss early stopping with a patience of 50 epochs and restore the best validation checkpoint.
Prediction Performance
The released checkpoints/models.pt was evaluated on a test set containing 537 electrolyte records. It achieved a mean coefficient of determination of R² = 0.9923 and a mean absolute percentage error of 3.12% across the five target properties.
| Property | Unit | R² | RMSE | MAE | MAPE |
|---|
| Density | g cm⁻³ | 0.9983 | 0.00419 | 0.00299 | 0.24% |
| Dielectric constant | dimensionless | 0.9882 | 1.1925 | 0.7320 | 3.50% |
| Viscosity | mPa s | 0.9843 | 0.6793 | 0.2468 | 5.46% |
| Diffusion coefficient | m² s⁻¹ | 0.9950 | 1.96 × 10⁻¹¹ | 1.28 × 10⁻¹¹ | 3.23% |
| Ionic conductivity | mS cm⁻¹ | 0.9959 | 0.4056 | 0.2485 | 3.16% |
The table reports rounded values from the release-checkpoint evaluation. Retraining the model with a different dataset version or a newly generated formula-level split may produce different metrics.
Formula-Level Split
Records are grouped by a stable formula identifier built from:
- canonical solvent SMILES and normalized solvent ratios;
- canonical salt SMILES and normalized salt ratios;
- salt concentration;
- component-order-independent sorting.
Temperature is excluded from the formula identifier. Therefore, all temperatures belonging to the same formula are assigned to the same train, validation, or test split. The default split ratio is approximately 70/20/10 by formula group, not by individual record.
The generated manifest is stored at data/splits/formula_split_seed42.json. Target transformation parameters and trajectory-target normalization statistics are calculated using training formulas only.
Repository Layout
1TIDE-HMS/
2├── configs/tide_hms.json
3├── data/raw/example_formula_data.json
4├── scripts/preprocess.py
5├── scripts/train.py
6├── scripts/predict.py
7└── tide_hms/
8 ├── data.py
9 ├── trajectory.py
10 ├── model.py
11 ├── training.py
12 └── metrics.py
Data Layout
Formula records use the JSON structure shown in data/raw/example_formula_data.json. Raw trajectories are expected at:
1data/traj/<case_name>/<case_name>.tpr
2data/traj/<case_name>/<case_name>.xtc
Trajectory directory names follow the existing convention case_<id>_<temperature>K_<replicate>, for example case_15_302.4K_1.
The public source bundle includes only a one-record format example. It does not include the full formula-property dataset, raw .tpr/.xtc trajectories, or precomputed four-scale trajectory caches used for model training. Reproducing the released training metrics therefore requires obtaining those assets separately and documenting their license, checksum, and download location. The released checkpoint can still be used for formula-property inference without trajectory files.
Installation
The required environment uses Python with PyTorch, RDKit, NumPy, MDAnalysis, and tqdm. CUDA is required for full training, while preprocessing and prediction can run on CPU.
Preprocessing
1python scripts/preprocess.py \
2 --source-json data/raw/formula_data.json \
3 --traj-root data/traj \
4 --cache-root data/cache \
5 --split-manifest data/splits/formula_split_seed42.json \
6 --n-frames 100 \
7 --seed 42 \
8 --target-device cuda:0
This command samples 100 frames, computes the four trajectory targets, writes traj_cache_100f.npz, and creates the formula-level split manifest.
Training
1python scripts/train.py \
2 --source-json data/raw/formula_data.json \
3 --traj-root data/traj \
4 --cache-root data/cache \
5 --output-dir outputs/tide_hms \
6 --split-manifest data/splits/formula_split_seed42.json \
7 --pretrain-epochs 600 \
8 --finetune-epochs 600 \
9 --pretrain-early-stopping-patience 50 \
10 --finetune-early-stopping-patience 50 \
11 --hidden-dim 256 \
12 --gnn-hidden-dim 128 \
13 --batch-size 32 \
14 --lr 1e-4 \
15 --device cuda:0
Training produces models.pt, metrics.json, and target_transform_params.json under the run directory. models.pt contains the model weights and the minimum architecture configuration required for loading. Dataset paths, output paths, optimizer settings, epoch counts, early-stopping settings, and runtime metadata are kept out of the public checkpoint.
Both stages use a maximum of 600 epochs and an early-stopping patience of 50 epochs. Validation is performed every 10 epochs by default, so five consecutive validation checks without improvement trigger early stopping. The best validation checkpoint is restored before the next stage or final save.
Prediction
1python scripts/predict.py \
2 --run-dir checkpoints \
3 --source-json data/raw/example_formula_data.json \
4 --target-transform-params configs/target_transform_params.json \
5 --device cpu
Configuration
configs/tide_hms.json records the canonical TIDE-HMS hyperparameters. Command-line options are the executable source of configuration for preprocessing and training.
