基本的にJaxはほとんどトランスフォーマーの構築に特化されているが、要するに行列というかテンソルの計算が専門なので、意外に原子の露出順位づけなどに使える。むしろ、既製のパッケージより速い。

20 Feb, 2026

import jax
import jax.numpy as jnp

def exposure_proxy(patch_xyz, env_xyz, patch_q,
                   r_cut=8.0, s=0.5, C0=10.0):
    # patch_xyz: [Np, 3]
    # env_xyz:   [Ne, 3] 
    d = jnp.linalg.norm(patch_xyz[:, None, :] - env_xyz[None, :, :], axis=-1) 
    w = jax.nn.sigmoid((r_cut - d) / s)
    cn = jnp.sum(w, axis=1)  
    E  = jnp.exp(-cn / C0)   
    B  = jnp.sum(jnp.maximum(patch_q, 0.0) * E)
    return B, E, cn

こんな感じ。