TACE 0.2.1__tar.gz → 0.2.2__tar.gz

{tace-0.2.1 → tace-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TACE
-Version: 0.2.1
+Version: 0.2.2
 Summary: TACE - ST and ICT based equivariant atomistic model.
 Author-email: Zemin Xu <xv_chana@163.com>
 Project-URL: Homepage, https://github.com/xvzemin/tace
@@ -14,10 +14,11 @@ Requires-Dist: hydra-core>=1.3
 Requires-Dist: matscipy
 Requires-Dist: ase
 Requires-Dist: numpy
-Requires-Dist: e3nn
+Requires-Dist: e3nn>=0.4.4
 Requires-Dist: scipy
 Requires-Dist: configargparse
 Requires-Dist: pandas
+Requires-Dist: matplotlib
 Requires-Dist: lmdb
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -25,6 +26,9 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff<0.17,>=0.16; extra == "dev"
 Provides-Extra: oeq
 Requires-Dist: openequivariance>=0.5.4; extra == "oeq"
+Provides-Extra: eqx
+Requires-Dist: ninja; extra == "eqx"
+Requires-Dist: filelock; extra == "eqx"
 Provides-Extra: torchsim
 Requires-Dist: torch-sim-atomistic>=0.6.1; extra == "torchsim"
 Provides-Extra: nvalchemi

{tace-0.2.1 → tace-0.2.2}/README.md

@@ -1,7 +1,7 @@
 [![TACE](https://img.shields.io/pypi/v/tace?style=for-the-badge&label=TACE)](https://pypi.org/project/tace/)
 [![Docs](https://img.shields.io/readthedocs/tace?style=for-the-badge&label=docs)](https://tace.readthedocs.io/en/latest/)
 [![License](https://img.shields.io/badge/License-MIT-yellow.svg?style=for-the-badge)](https://opensource.org/licenses/MIT)
-[![Matbench Discovery](https://img.shields.io/badge/Matbench%20Discovery-SOTA-brightgreen?style=for-the-badge)](https://matbench-discovery.materialsproject.org/)
+<!-- [![Matbench Discovery](https://img.shields.io/badge/Matbench%20Discovery-SOTA-brightgreen?style=for-the-badge)](https://matbench-discovery.materialsproject.org/) -->
 
 # Tensor Atomic/Edge Cluster Expansion (TACE/TECE)
 
@@ -9,30 +9,51 @@ TACE is designed with physical priors and strong inductive biases to enhance ext
 It performs Atomic Cluster Expansion and Edge Cluster Expansion based on spherical tensors 
 or irreducible Cartesian tensors, with an optional attention architecture.
 
-## Cartesian Architecture
+## O(3) Cartesian Architecture
 
 <img src="fig/cartesian_arch.png" width="100%" align="center">
 
-## Spherical/SO(2) Architecture
-
-The architecture of the spherical model is largely the same as that of the Cartesian space. 
-For details on the SO(2) component, please refer to our paper and code.
-
-## Wigner6j/O(2) Architecture
-
-The current implementation is still subject to change, and backward compatibility is not guaranteed. Please refer to our paper and code for details.
+```math
+(A\otimes_k^\delta B)_{\boldsymbol i\boldsymbol j}
+=\frac{1}{\sqrt{3^k}}\sum_{\boldsymbol a}
+A_{\boldsymbol i\boldsymbol a}B_{\boldsymbol a\boldsymbol j},
+\qquad l_3=l_1+l_2-2k,
+````
+
+```math
+(A\otimes_k^\epsilon B)_{\boldsymbol i w\boldsymbol j}
+=\frac{1}{\sqrt{2}}\frac{1}{\sqrt{3^k}}\sum_{\boldsymbol a,u,v}
+A_{\boldsymbol i u\boldsymbol a}\epsilon_{wuv}
+B_{\boldsymbol a v\boldsymbol j},
+\qquad l_3=l_1+l_2-2k-1.
+````
+
+## Time-reversal O(3) Spherical/O(2)/Magnetic Architecture
+
+The SO(2) implementation in TACE will be gradually replaced by an O(2) formulation 
+to ensure complete parity support for O(3). Backward compatibility will be maintained for the SO(2) implementation, 
+while backward compatibility for the O(2) implementation is not currently guaranteed.
+
+<img src="fig/o2.svg" width="100%" align="center">
+
+```math
+\left.(l,p,t)_{O(3)\times\mathbb Z_2^{\mathcal T}}
+\right\downarrow_{O(2)\times\mathbb Z_2^{\mathcal T}}
+=
+\underbrace{\bigl(0,p(-1)^l,t\bigr)}_{\text{1D},\,m=0}
+\oplus
+\underbrace{\bigoplus_{m=1}^{l}(m,0,t)}_{\text{2D},\,m>0}.
+````
 
 ## Documentation
 
 [TACE DOCS](https://tace.readthedocs.io/en/latest/index.html)
 
-## SOTA Foundation Model
+## Foundation Model
 
-[TACE FOUNDATION](https://github.com/xvzemin/tace-foundations)
+[HUGGING FACE](https://huggingface.co/xvzemin/tace-foundations/tree/main)
 
-Default Ranking on Matbench as of July 8, 2026
-
-<img src="fig/matbench_tece_rra.png" width="100%" align="center">
+[FOUNDATION MODEL DOCS](https://tace.readthedocs.io/en/latest/guide/foundation.html)
 
 ## Install, Train and Tutorial
 
@@ -76,23 +97,16 @@ Currently, the officially supported properties include:
 - Atomic stresses (conservative, predict only)
 - Atomic virials (conservative, predict only)
 - Absolute final collinear magmoms
-- Noncollinear magnetic forces (SOC, full O(3))
-
-For embedding property, we support:
-
-- fidelity_idx (different computational levels)
-- charges
-- total charge
-- electric field
-- initial noncollinear magmoms (SOC, full O(3))
+- Collinear magnetic forces
+- Noncollinear magnetic forces
 
 ## Plugins
 
 TACE currently supports the following plugin:
 
+- **mTACE**    (Magnetic, with and without Spin-Orbit Coupling)
 - **TACE-LES** (Latent Ewald Summation)
-- **TACE-QEq** (Lagrangian)
-- **mTACE**    (Magnetic, Spin-Orbit Coupling)
+- **TACE-QEq** (Lagrangian, under reconstruction)
 
 ## Interfaces
 
@@ -137,37 +151,37 @@ If you use TACE, please cite our papers:
 
 ```bibtex
 @misc{xu2026spectralspatialtensoratomiccluster,
-   title={Spectral/Spatial Tensor Atomic Cluster Expansion with Universal Embeddings in Cartesian Space}, 
-   author={Zemin Xu and Wenbo Xie and P. Hu},
-   year={2026},
-   eprint={2509.14961},
-   archivePrefix={arXiv},
-   primaryClass={stat.ML},
-   url={https://arxiv.org/abs/2509.14961}, 
+  title={Spectral/Spatial Tensor Atomic Cluster Expansion with Universal Embeddings in Cartesian Space}, 
+  author={Zemin Xu and Wenbo Xie and P. Hu},
+  year={2026},
+  eprint={2509.14961},
+  archivePrefix={arXiv},
+  primaryClass={stat.ML},
+  url={https://arxiv.org/abs/2509.14961}, 
 }
 
 @misc{xu2026edgeclusterexpansionradial,
-   title={Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials}, 
-   author={Zemin Xu and Wenbo Xie and P. Hu},
-   year={2026},
-   eprint={2607.10664},
-   archivePrefix={arXiv},
-   primaryClass={stat.ML},
-   url={https://arxiv.org/abs/2607.10664}, 
+  title={Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials}, 
+  author={Zemin Xu and Wenbo Xie and P. Hu},
+  year={2026},
+  eprint={2607.10664},
+  archivePrefix={arXiv},
+  primaryClass={stat.ML},
+  url={https://arxiv.org/abs/2607.10664}, 
 }
 ```
 
-If you use Local O(2) Frame or Generalized Wigner-6j Convlution, please cite our papers: 
+If you use Local O(2) Frame, please cite our papers: 
 
 ```bibtex
-@misc{xu2026completeo3interactionswigner6j,
-   title={Complete O(3) Interactions from Wigner-6j Recoupling to Local O(2) Frames}, 
-   author={Zemin Xu and Peijun Hu and Wenbo Xie},
-   year={2026},
-   eprint={2608.16592},
-   archivePrefix={arXiv},
-   primaryClass={physics.chem-ph},
-   url={https://arxiv.org/abs/2608.16592}, 
+@misc{xu2026lookingglassefficientparitycompletelearning,
+  title={Through the Looking-Glass: Efficient Parity-Complete Learning via Local $O(2)$ Frames}, 
+  author={Zemin Xu and Wenbo Xie},
+  year={2026},
+  eprint={2608.16592},
+  archivePrefix={arXiv},
+  primaryClass={physics.chem-ph},
+  url={https://arxiv.org/abs/2608.16592}, 
 }
 ```
 
@@ -175,11 +189,11 @@ If you use cartnn, Cartesian-3j, cMACE, cNequIP, cAllegro, please cite our paper
 
 ```bibtex
 @inproceedings{xu2026a,
-   title={A Cartesian-3j Framework for Machine Learning Interatomic Potentials},
-   author={Zemin Xu and Chenyu Wu and Wenbo Xie and Peijun Hu},
-   booktitle={Forty-third International Conference on Machine Learning},
-   year={2026},
-   url={https://openreview.net/forum?id=9ZWK6gneWq}
+  title={A Cartesian-3j Framework for Machine Learning Interatomic Potentials},
+  author={Zemin Xu and Chenyu Wu and Wenbo Xie and Peijun Hu},
+  booktitle={Forty-third International Conference on Machine Learning},
+  year={2026},
+  url={https://openreview.net/forum?id=9ZWK6gneWq}
 }
 ```
 
@@ -199,3 +213,5 @@ ruff format --check tace
 ## License
 
 The TACE code is published and distributed under the MIT License.
+EquivariantX (`eqx`) is licensed separately under
+[Creative Commons Attribution 4.0 International (CC BY 4.0)](eqx/LICENSE.md).

{tace-0.2.1 → tace-0.2.2}/TACE.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TACE
-Version: 0.2.1
+Version: 0.2.2
 Summary: TACE - ST and ICT based equivariant atomistic model.
 Author-email: Zemin Xu <xv_chana@163.com>
 Project-URL: Homepage, https://github.com/xvzemin/tace
@@ -14,10 +14,11 @@ Requires-Dist: hydra-core>=1.3
 Requires-Dist: matscipy
 Requires-Dist: ase
 Requires-Dist: numpy
-Requires-Dist: e3nn
+Requires-Dist: e3nn>=0.4.4
 Requires-Dist: scipy
 Requires-Dist: configargparse
 Requires-Dist: pandas
+Requires-Dist: matplotlib
 Requires-Dist: lmdb
 Requires-Dist: huggingface_hub
 Provides-Extra: dev
@@ -25,6 +26,9 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff<0.17,>=0.16; extra == "dev"
 Provides-Extra: oeq
 Requires-Dist: openequivariance>=0.5.4; extra == "oeq"
+Provides-Extra: eqx
+Requires-Dist: ninja; extra == "eqx"
+Requires-Dist: filelock; extra == "eqx"
 Provides-Extra: torchsim
 Requires-Dist: torch-sim-atomistic>=0.6.1; extra == "torchsim"
 Provides-Extra: nvalchemi

{tace-0.2.1 → tace-0.2.2}/TACE.egg-info/SOURCES.txt

@@ -7,30 +7,119 @@ TACE.egg-info/dependency_links.txt
 TACE.egg-info/entry_points.txt
 TACE.egg-info/requires.txt
 TACE.egg-info/top_level.txt
+eqx/LICENSE.md
 eqx/README.md
 eqx/__init__.py
+eqx/_layout.py
+eqx/utils.py
+eqx/ace/__init__.py
+eqx/ace/contraction.py
+eqx/ace/cuda.py
+eqx/ace/tace.py
+eqx/co2/__init__.py
+eqx/co2/basis.py
+eqx/co2/cartesian_harmonics.py
+eqx/co2/gate.py
+eqx/co2/irreps.py
+eqx/co2/linear.py
+eqx/co2/o3_tensor_product.py
+eqx/co2/restriction.py
+eqx/co2/spherical.py
+eqx/co2/tensor_product.py
 eqx/co3/__init__.py
+eqx/co3/basis.py
 eqx/co3/cartesian_harmonics.py
 eqx/co3/gate.py
 eqx/co3/irreps.py
-eqx/co3/layout.py
 eqx/co3/linear.py
-eqx/co3/projector.py
+eqx/co3/symmetric.py
 eqx/co3/tensor_product.py
-eqx/co3/utils.py
-eqx/docs/source/conf.py
+eqx/conv/__init__.py
+eqx/conv/angular.py
+eqx/conv/attention.py
+eqx/conv/codegen.py
+eqx/conv/contraction.py
+eqx/conv/edge.py
+eqx/conv/graph.py
+eqx/conv/program.py
+eqx/conv/radial.py
+eqx/conv/co3/__init__.py
+eqx/conv/co3/convolution.py
+eqx/conv/co3/linear.py
+eqx/conv/co3/polynomials.py
+eqx/conv/o2_o3/__init__.py
+eqx/conv/o2_o3/autotune.py
+eqx/conv/o2_o3/codegen.py
+eqx/conv/o2_o3/convolution.py
+eqx/conv/o2_o3/cuda.py
+eqx/conv/o2_o3/direction_codegen.py
+eqx/conv/o2_o3/geometry.py
+eqx/conv/o2_o3/schedule.py
+eqx/conv/o2_o3/transverse.py
+eqx/conv/o3/__init__.py
+eqx/conv/o3/codegen.py
+eqx/conv/o3/convolution.py
+eqx/conv/o3/cuda.py
+eqx/conv/o3/harmonics.py
+eqx/conv/uu_o2/__init__.py
+eqx/conv/uu_o2/convolution.py
+eqx/conv/uv_o2/__init__.py
+eqx/conv/uv_o2/convolution.py
+eqx/conv/uv_o2/transverse.py
+eqx/kernels/__init__.py
+eqx/kernels/channel_product.py
+eqx/kernels/codegen.py
+eqx/kernels/cuda.py
+eqx/kernels/cuda_graph.py
+eqx/kernels/layout.py
+eqx/kernels/quaternion.py
+eqx/kernels/recompute.py
+eqx/kernels/rotary.py
+eqx/kernels/rotation.py
+eqx/kernels/wigner.py
+eqx/kernels/csrc/runtime.cpp
+eqx/models/__init__.py
+eqx/models/convolution.py
+eqx/models/equflash/__init__.py
+eqx/models/equflash/conversion.py
+eqx/models/mace/__init__.py
+eqx/models/mace/conversion.py
+eqx/models/nequip/__init__.py
+eqx/models/nequip/conversion.py
+eqx/models/prophet/__init__.py
+eqx/models/prophet/conversion.py
+eqx/models/sevennet/__init__.py
+eqx/models/sevennet/conversion.py
+eqx/models/tace/__init__.py
+eqx/models/tace/tece_oam_rra/__init__.py
+eqx/models/tace/tece_oam_rra/bilinear_contraction.py
+eqx/models/tace/tece_oam_rra/bilinear_cuda.py
+eqx/models/tace/tece_oam_rra/cuda.py
+eqx/models/tace/tece_oam_rra/execution.py
+eqx/models/tace/tece_oam_rra/interaction.py
+eqx/models/tace/tece_oam_rra/product.py
+eqx/models/tace/tece_oam_rra/program.py
 eqx/o2/__init__.py
+eqx/o2/_clebsch_gordan.py
+eqx/o2/_layout.py
 eqx/o2/asymmetric_contraction.py
 eqx/o2/circular_harmonics.py
 eqx/o2/gate.py
 eqx/o2/irreps.py
 eqx/o2/linear.py
 eqx/o2/local_frame.py
+eqx/o2/o3_tensor_product.py
 eqx/o2/rotation_matrix.py
 eqx/o2/tensor_product.py
 eqx/o2/wigner.py
+eqx/o3/__init__.py
+eqx/o3/contraction.py
+eqx/o3/cuda.py
+eqx/o3/gate.py
+eqx/o3/linear.py
 tace/__init__.py
 tace/dataset/__init__.py
+tace/dataset/augmentation.py
 tace/dataset/dataloader.py
 tace/dataset/datamodule.py
 tace/dataset/element.py
@@ -75,15 +164,14 @@ tace/models/kspace.py
 tace/models/lammps.py
 tace/models/layout.py
 tace/models/linear.py
-tace/models/mag.py
 tace/models/mlp.py
 tace/models/normalizer.py
 tace/models/radial.py
 tace/models/s2.py
 tace/models/softmax.py
+tace/models/time_reversal.py
 tace/models/utils.py
 tace/models/zemin.py
-tace/models/_cart/__init__.py
 tace/models/_e3nn/__init__.py
 tace/models/_e3nn/base.py
 tace/models/_e3nn/basis_change.py
@@ -93,8 +181,8 @@ tace/models/_e3nn/edge.py
 tace/models/_e3nn/fused.py
 tace/models/_e3nn/inter.py
 tace/models/_e3nn/layer_norm.py
-tace/models/_e3nn/legacy_so2.py
 tace/models/_e3nn/les.py
+tace/models/_e3nn/magnetic.py
 tace/models/_e3nn/node.py
 tace/models/_e3nn/nonlinear.py
 tace/models/_e3nn/o2.py
@@ -106,6 +194,7 @@ tace/models/_e3nn/residual.py
 tace/models/_e3nn/scatter_norm.py
 tace/models/_e3nn/symmetric_contraction.py
 tace/models/_e3nn/tace.py
+tace/models/_e3nn/tece_oam_rra.py
 tace/models/_e3nn/ue.py
 tace/models/_e3nn/wigner6j.py
 tace/models/compile/__init__.py
@@ -154,12 +243,14 @@ tace/models/scf/longrange/utils.py
 tace/scripts/__init__.py
 tace/scripts/average.py
 tace/scripts/clean.py
-tace/scripts/convert.py
+tace/scripts/convert_cgtp.py
+tace/scripts/convert_lora.py
 tace/scripts/eval.py
 tace/scripts/export_eval.py
 tace/scripts/export_lammps.py
 tace/scripts/export_train.py
 tace/scripts/finetune.py
+tace/scripts/plot_diatom.py
 tace/scripts/split.py
 tace/scripts/train.py
 tace/scripts/update.py
@@ -200,12 +291,18 @@ tace/utils/lr_scheduler/wsd.py
 tace/utils/optimizer/__init__.py
 tace/utils/optimizer/hybrid_muon.py
 tace/utils/optimizer/soap.py
-tests/test_co3.py
+tests/test_cartesian.py
 tests/test_compile.py
 tests/test_eqt.py
+tests/test_eqx_integration.py
 tests/test_les.py
+tests/test_local_o2.py
 tests/test_loss.py
 tests/test_magnetic_optimizer.py
-tests/test_o2.py
+tests/test_product.py
+tests/test_radial.py
 tests/test_scf.py
+tests/test_statistics.py
+tests/test_time_reversal.py
+tests/test_torchsim.py
 tests/test_wigner6j.py

{tace-0.2.1 → tace-0.2.2}/TACE.egg-info/entry_points.txt

@@ -2,12 +2,14 @@
 tace-average = tace.scripts.average:main
 tace-clean = tace.scripts.clean:main
 tace-compile = tace.scripts.export_eval:main
-tace-convert = tace.scripts.convert:main
+tace-convert-cgtp = tace.scripts.convert_cgtp:main
+tace-convert-lora = tace.scripts.convert_lora:main
 tace-eval = tace.scripts.eval:main
 tace-export-eval = tace.scripts.export_eval:main
 tace-export-lammps = tace.scripts.export_lammps:main
 tace-export-train = tace.scripts.export_train:main
 tace-finetune = tace.scripts.finetune:main
+tace-plot-diatom = tace.scripts.plot_diatom:main
 tace-split = tace.scripts.split:main
 tace-train = tace.scripts.train:main
 tace-update = tace.scripts.update:main

{tace-0.2.1 → tace-0.2.2}/TACE.egg-info/requires.txt

@@ -6,10 +6,11 @@ hydra-core>=1.3
 matscipy
 ase
 numpy
-e3nn
+e3nn>=0.4.4
 scipy
 configargparse
 pandas
+matplotlib
 lmdb
 huggingface_hub
 
@@ -30,6 +31,10 @@ torch_dftd
 pytest
 ruff<0.17,>=0.16
 
+[eqx]
+ninja
+filelock
+
 [lammps12]
 cython
 cupy-cuda12x