Robochimps/hcn-iresnet-jacobi

HCN iResNet coordinate transformation — Jacobi coordinates

Trained iResNet normalising-flow bijection that maps HCN Jacobi internal coordinates (r, R, γ) to dimensionless flow coordinates (q̃₁, q̃₂, q̃₃) optimised for variational nuclear-motion calculations.

This model was published in:

Yahya Saleh, Álvaro Fernández Corral, Emil Vogt, Armin Iske, Jochen Küpper, and Andrey Yachmenev. Computing Excited States of Molecules Using Normalizing Flows. J. Chem. Theory Comput. 2025, 21 (10), 5221–5229. https://doi.org/10.1021/acs.jctc.5c00590

What this is

The coordinate transformation is an invertible residual network (iResNet) that learns a bijection from the physical Jacobi coordinate space of HCN to a latent space. The latent space improves the accuracy of variational vibrational energy level calculations.

Input / output

	Symbol	Unit	Description
Input	r	Å	C–N bond length
Input	R	Å	H-to-CoM(CN) distance
Input	γ	rad	Angle between CN bond vector and H–CoM(CN) vector (γ = 0 → linear HCN, γ = π → linear NCH)
Output	q̃₁, q̃₂, q̃₃	dimensionless	iResNet flow coordinates

The model operates on raw physical coordinates — no displacement from a reference geometry is applied.

Physical coordinate ranges

Coordinate	Range
r (C–N)	[0.8996, 1.8521] Å
R (H–CoM)	[0.5297, 3.7037] Å
γ	[0, π] rad

Architecture

• Type: invertible residual network (iResNet)
• Blocks: 10 residual blocks
• Dense layers per block: 3 (widths 8 → 8 → 3)
• Activation: LipSwish ( x·σ(x)/1.1 )
• Lipschitz constraint: SVD spectral clipping at 0.9 per layer
• Output scaling: per-DOF vector xmax = [10.0, 10.0, 0.5]
• Inversion: 30-step fixed-point iteration

All arithmetic is in float64. Both the forward and inverse passes are fully JAX-traceable (jax.jit, jax.grad, jax.hessian, jax.vmap).

Repository contents

File	Description
config.json	Architecture hyperparameters (model_type: "iresnet")
flax_model.msgpack	Trained parameters {"params": ...}, serialised with flax.serialization (float64)

The repository contains no code. The architecture is defined once, for all molecules, in pyhami as the flax module pyhami.core.flax_iresnet.IResNetFlow. msgpack is plain data, so loading the weights cannot execute code.

Loading and using the model

With pyhami (recommended):

import jax.numpy as jnp
from pyhami.molecules.hcn.coords import HCNIResNetTransformation

transf  = HCNIResNetTransformation()                 # downloads config + weights once, then cached
q_phys  = jnp.array([1.065, 2.014, 0.0])   # (r [Å], R [Å], γ [rad]) — linear HCN
q_tilde = transf.forward(q_phys)   # → flow coords
q_back  = transf.inverse(q_tilde)  # → physical coords

# fine-tuning: parameters become JAX pytree leaves
transf = HCNIResNetTransformation(trainable=True)

The raw files can also be read directly with flax:

import json
from flax.serialization import msgpack_restore
from huggingface_hub import hf_hub_download

config    = json.load(open(hf_hub_download("Robochimps/hcn-iresnet-jacobi", "config.json")))
variables = msgpack_restore(open(hf_hub_download("Robochimps/hcn-iresnet-jacobi", "flax_model.msgpack"), "rb").read())

Citation

If you use this model in your work, please cite:

@article{saleh2025flows,
  author  = {Saleh, Yahya and Fern{\'a}ndez Corral, {\'A}lvaro and Vogt, Emil
             and Iske, Armin and K{\"u}pper, Jochen and Yachmenev, Andrey},
  title   = {Computing Excited States of Molecules Using Normalizing Flows},
  journal = {Journal of Chemical Theory and Computation},
  year    = {2025},
  volume  = {21},
  number  = {10},
  pages   = {5221--5229},
  doi     = {10.1021/acs.jctc.5c00590},
}