TACE 0.2.0__tar.gz → 0.2.1__tar.gz

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

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: TACE
-Version: 0.2.0
+Version: 0.2.1
 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
 Requires-Python: >=3.9
 License-File: LICENSE.md
-Requires-Dist: torch<=2.13,>=2.4
+Requires-Dist: torch>=2.4
 Requires-Dist: torch_geometric>=2.4
 Requires-Dist: lightning
 Requires-Dist: omegaconf
@@ -20,10 +20,17 @@ Requires-Dist: configargparse
 Requires-Dist: pandas
 Requires-Dist: lmdb
 Requires-Dist: huggingface_hub
+Provides-Extra: dev
+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: torchsim
 Requires-Dist: torch-sim-atomistic>=0.6.1; extra == "torchsim"
+Provides-Extra: nvalchemi
+Requires-Dist: nvalchemi-toolkit>=0.1.0; extra == "nvalchemi"
+Requires-Dist: nvalchemi-toolkit-ops; extra == "nvalchemi"
+Provides-Extra: scf
 Provides-Extra: cueq12
 Requires-Dist: cuequivariance; extra == "cueq12"
 Requires-Dist: cuequivariance-torch; extra == "cueq12"

tace-0.2.1/README.md

@@ -0,0 +1,201 @@
+[![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/)
+
+# Tensor Atomic/Edge Cluster Expansion (TACE/TECE)
+
+TACE is designed with physical priors and strong inductive biases to enhance extrapolation capability. 
+It performs Atomic Cluster Expansion and Edge Cluster Expansion based on spherical tensors 
+or irreducible Cartesian tensors, with an optional attention architecture.
+
+## 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.
+
+## Documentation
+
+[TACE DOCS](https://tace.readthedocs.io/en/latest/index.html)
+
+## SOTA Foundation Model
+
+[TACE FOUNDATION](https://github.com/xvzemin/tace-foundations)
+
+Default Ranking on Matbench as of July 8, 2026
+
+<img src="fig/matbench_tece_rra.png" width="100%" align="center">
+
+## Install, Train and Tutorial
+
+The docs contain a complete tutorial. 
+
+We also provide complete input files and example scripts for training, ASE,
+TorchSim, LAMMPS, and other workflows in the
+[TACE examples](https://github.com/xvzemin/tace/tree/main/example).
+
+```bash
+# Minimal install and training example
+git clone https://github.com/xvzemin/tace
+cd tace
+pip install . # or pip install tace
+cd example/train
+tace-train -cn tace.yaml
+```
+
+## Fine-tuning
+
+- ✅ Full-parameter.
+
+- ✅ Freeze-parameter.
+
+- ✅ LoRA.
+
+## Overview
+
+Currently, the officially supported properties include:
+
+- Energy
+- Forces (conservative | direct)
+- Hessian (conservative, predict only)
+- Stress (conservative | direct)
+- Virials (conservative | direct)
+- Charges (lagrangian or uniform_distribution)
+- Dipole moment (conservative | direct)
+- Polarization (conservative, multi-value for PBC systems)
+- Polarizability (conservative | direct)
+- Born effective charges (conservative, under electric field)
+- 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))
+
+## Plugins
+
+TACE currently supports the following plugin:
+
+- **TACE-LES** (Latent Ewald Summation)
+- **TACE-QEq** (Lagrangian)
+- **mTACE**    (Magnetic, Spin-Orbit Coupling)
+
+## Interfaces
+
+- ✅ Supports integration with [ASE Calculator](https://wiki.fysik.dtu.dk/ase/).
+
+- ✅ Supports integration with [LAMMPS-ML-IAP](https://github.com/lammps/lammps).
+
+- ✅ Supports integration with [TorchSim](https://torchsim.github.io/torch-sim/).
+
+- ✅ Supports integration with [NVIDIA NValCHEMI](https://github.com/NVIDIA/nvalchemi-toolkit).
+
+- ✅ Supports integration with [OpenMM-ML](https://github.com/openmm/openmm-ml) (OpenMM-ML -> ASE -> TACE).
+
+- ✅ Supports integration with [USPEX](https://uspex-team.org/)
+  (USPEX -> LAMMPS-ML-IAP -> TACE) (Python=3.9).
+
+## Contact
+
+For bugs or feature requests, please use the
+[TACE issue](https://github.com/xvzemin/tace/issues).
+
+<!-- For usage discussions, you can join the TACE community through QQ or Discord.
+
+<table>
+  <tr>
+    <th>QQ</th>
+    <th>Discord</th>
+  </tr>
+  <tr>
+    <td align="center">
+      <img src="fig/qq.jpg" alt="TACE QQ QR code" width="260">
+    </td>
+    <td align="center">
+      <img src="fig/discord.jpg" alt="TACE Discord QR code" width="260">
+    </td>
+  </tr>
+</table> -->
+
+## Citing
+
+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}, 
+}
+
+@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}, 
+}
+```
+
+If you use Local O(2) Frame or Generalized Wigner-6j Convlution, 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}, 
+}
+```
+
+If you use cartnn, Cartesian-3j, cMACE, cNequIP, cAllegro, please cite our papers:
+
+```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}
+}
+```
+
+## Development
+
+Install the development tools and run the same formatting and lint checks used by
+pre-commit:
+
+```bash
+pip install -e ".[dev]"
+ruff check --fix tace
+ruff format tace
+ruff check tace
+ruff format --check tace
+```
+
+## License
+
+The TACE code is published and distributed under the MIT License.

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

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: TACE
-Version: 0.2.0
+Version: 0.2.1
 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
 Requires-Python: >=3.9
 License-File: LICENSE.md
-Requires-Dist: torch<=2.13,>=2.4
+Requires-Dist: torch>=2.4
 Requires-Dist: torch_geometric>=2.4
 Requires-Dist: lightning
 Requires-Dist: omegaconf
@@ -20,10 +20,17 @@ Requires-Dist: configargparse
 Requires-Dist: pandas
 Requires-Dist: lmdb
 Requires-Dist: huggingface_hub
+Provides-Extra: dev
+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: torchsim
 Requires-Dist: torch-sim-atomistic>=0.6.1; extra == "torchsim"
+Provides-Extra: nvalchemi
+Requires-Dist: nvalchemi-toolkit>=0.1.0; extra == "nvalchemi"
+Requires-Dist: nvalchemi-toolkit-ops; extra == "nvalchemi"
+Provides-Extra: scf
 Provides-Extra: cueq12
 Requires-Dist: cuequivariance; extra == "cueq12"
 Requires-Dist: cuequivariance-torch; extra == "cueq12"

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

@@ -1,5 +1,5 @@
 LICENSE.md
-README.rst
+README.md
 pyproject.toml
 TACE.egg-info/PKG-INFO
 TACE.egg-info/SOURCES.txt
@@ -7,6 +7,28 @@ TACE.egg-info/dependency_links.txt
 TACE.egg-info/entry_points.txt
 TACE.egg-info/requires.txt
 TACE.egg-info/top_level.txt
+eqx/README.md
+eqx/__init__.py
+eqx/co3/__init__.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/tensor_product.py
+eqx/co3/utils.py
+eqx/docs/source/conf.py
+eqx/o2/__init__.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/rotation_matrix.py
+eqx/o2/tensor_product.py
+eqx/o2/wigner.py
 tace/__init__.py
 tace/dataset/__init__.py
 tace/dataset/dataloader.py
@@ -26,8 +48,12 @@ tace/foundations/u_shift.py
 tace/interface/__init__.py
 tace/interface/ase/__init__.py
 tace/interface/ase/calculator.py
+tace/interface/ase/general_calculator.py
+tace/interface/ase/optimizer.py
 tace/interface/lammps/__init__.py
 tace/interface/lammps/mliap.py
+tace/interface/nvalchemi/__init__.py
+tace/interface/nvalchemi/wrapper.py
 tace/interface/openmm/__init__.py
 tace/interface/torchsim/__init__.py
 tace/interface/torchsim/torchsim.py
@@ -36,7 +62,7 @@ tace/interface/uspex26/relax.py
 tace/lightning/__init__.py
 tace/lightning/lit_model.py
 tace/lightning/lora.py
-tace/lightning/los_skip.py
+tace/lightning/loss_skip.py
 tace/lightning/torch_model.py
 tace/lightning/trainer.py
 tace/lightning/u_shift.py
@@ -45,43 +71,20 @@ tace/models/adapter.py
 tace/models/angular.py
 tace/models/blocks.py
 tace/models/ictd.py
+tace/models/kspace.py
 tace/models/lammps.py
 tace/models/layout.py
-tace/models/legacy.py
 tace/models/linear.py
+tace/models/mag.py
 tace/models/mlp.py
 tace/models/normalizer.py
-tace/models/precision.py
 tace/models/radial.py
 tace/models/s2.py
 tace/models/softmax.py
 tace/models/utils.py
 tace/models/zemin.py
 tace/models/_cart/__init__.py
-tace/models/_cart/base.py
-tace/models/_cart/basis_change.py
-tace/models/_cart/cartesian.py
-tace/models/_cart/default.py
-tace/models/_cart/dropout.py
-tace/models/_cart/edge.py
-tace/models/_cart/fused.py
-tace/models/_cart/inter.py
-tace/models/_cart/layer_norm.py
-tace/models/_cart/node.py
-tace/models/_cart/nonlinear.py
-tace/models/_cart/paths.py
-tace/models/_cart/prod.py
-tace/models/_cart/readout.py
-tace/models/_cart/representation.py
-tace/models/_cart/residual.py
-tace/models/_cart/tace.py
-tace/models/_cart/ue.py
-tace/models/_cue/__init__.py
-tace/models/_cue/_tp_scatter.py
-tace/models/_cue/_tp_uuu.py
-tace/models/_cue/paths.py
 tace/models/_e3nn/__init__.py
-tace/models/_e3nn/asymmetric_contraction.py
 tace/models/_e3nn/base.py
 tace/models/_e3nn/basis_change.py
 tace/models/_e3nn/default.py
@@ -90,8 +93,11 @@ 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/node.py
 tace/models/_e3nn/nonlinear.py
+tace/models/_e3nn/o2.py
 tace/models/_e3nn/paths.py
 tace/models/_e3nn/prod.py
 tace/models/_e3nn/readout.py
@@ -101,28 +107,50 @@ tace/models/_e3nn/scatter_norm.py
 tace/models/_e3nn/symmetric_contraction.py
 tace/models/_e3nn/tace.py
 tace/models/_e3nn/ue.py
-tace/models/_eqt/__init__.py
-tace/models/_eqt/_tp_uuu.py
-tace/models/_eqt/equitorch/__init__.py
-tace/models/_eqt/equitorch/irreps/__init__.py
-tace/models/_eqt/equitorch/nn/__init__.py
-tace/models/_eqt/equitorch/nn/sparse_product.py
-tace/models/_eqt/equitorch/nn/tensor_products.py
-tace/models/_eqt/equitorch/structs/__init__.py
-tace/models/_oeq/__init__.py
-tace/models/_oeq/_tp_scatter.py
+tace/models/_e3nn/wigner6j.py
 tace/models/compile/__init__.py
 tace/models/compile/aot.py
 tace/models/compile/compile.py
 tace/models/compile/tace.py
 tace/models/compile/wrapper.py
-tace/models/so2/__init__.py
-tace/models/so2/blocks.py
-tace/models/so2/rotation_matrix.py
-tace/models/so2/utils.py
-tace/models/so2/wigner.py
-tace/models/triton_ops/__init__.py
-tace/models/triton_ops/uu_so2_scatter.py
+tace/models/cue/__init__.py
+tace/models/cue/_tp_scatter.py
+tace/models/cue/_tp_uuu.py
+tace/models/cue/paths.py
+tace/models/eqt/__init__.py
+tace/models/eqt/_tp_uuu.py
+tace/models/eqt/equitorch/__init__.py
+tace/models/eqt/equitorch/irreps/__init__.py
+tace/models/eqt/equitorch/nn/__init__.py
+tace/models/eqt/equitorch/nn/sparse_product.py
+tace/models/eqt/equitorch/nn/tensor_products.py
+tace/models/eqt/equitorch/structs/__init__.py
+tace/models/oeq/__init__.py
+tace/models/oeq/_tp_scatter.py
+tace/models/scf/README.md
+tace/models/scf/__init__.py
+tace/models/scf/electrostatics.py
+tace/models/scf/energy_functional.py
+tace/models/scf/fixed_point.py
+tace/models/scf/local.py
+tace/models/scf/loss.py
+tace/models/scf/model.py
+tace/models/scf/qeq.py
+tace/models/scf/readout.py
+tace/models/scf/state.py
+tace/models/scf/utils.py
+tace/models/scf/longrange/LICENSE.md
+tace/models/scf/longrange/NOTICE.md
+tace/models/scf/longrange/__init__.py
+tace/models/scf/longrange/__version__.py
+tace/models/scf/longrange/energy.py
+tace/models/scf/longrange/features.py
+tace/models/scf/longrange/gto_utils.py
+tace/models/scf/longrange/kspace.py
+tace/models/scf/longrange/realspace_electrostatics.py
+tace/models/scf/longrange/realspace_grid_integrals.py
+tace/models/scf/longrange/slabs.py
+tace/models/scf/longrange/utils.py
 tace/scripts/__init__.py
 tace/scripts/average.py
 tace/scripts/clean.py
@@ -135,6 +163,10 @@ tace/scripts/finetune.py
 tace/scripts/split.py
 tace/scripts/train.py
 tace/scripts/update.py
+tace/scripts/utils/__init__.py
+tace/scripts/utils/check_equi.py
+tace/scripts/utils/check_soc.py
+tace/scripts/utils/convert_to_xyz.py
 tace/utils/__init__.py
 tace/utils/_global.py
 tace/utils/callbacks.py
@@ -168,7 +200,12 @@ tace/utils/lr_scheduler/wsd.py
 tace/utils/optimizer/__init__.py
 tace/utils/optimizer/hybrid_muon.py
 tace/utils/optimizer/soap.py
-test/test_aoti_single_system.py
-test/test_neighbour_list.py
-test/test_solid_harmonics.py
-test/test_wigner6j.py
+tests/test_co3.py
+tests/test_compile.py
+tests/test_eqt.py
+tests/test_les.py
+tests/test_loss.py
+tests/test_magnetic_optimizer.py
+tests/test_o2.py
+tests/test_scf.py
+tests/test_wigner6j.py

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

@@ -1,4 +1,4 @@
-torch<=2.13,>=2.4
+torch>=2.4
 torch_geometric>=2.4
 lightning
 omegaconf
@@ -26,6 +26,10 @@ cuequivariance-ops-torch-cu13
 [d3]
 torch_dftd
 
+[dev]
+pytest
+ruff<0.17,>=0.16
+
 [lammps12]
 cython
 cupy-cuda12x
@@ -39,6 +43,10 @@ ase
 matscipy
 vesin
 
+[nvalchemi]
+nvalchemi-toolkit>=0.1.0
+nvalchemi-toolkit-ops
+
 [oeq]
 openequivariance>=0.5.4
 
@@ -47,5 +55,7 @@ openmm
 openmm-ml
 openmmtorch
 
+[scf]
+
 [torchsim]
 torch-sim-atomistic>=0.6.1

{tace-0.2.0 → tace-0.2.1}/TACE.egg-info/top_level.txt

@@ -1 +1,2 @@
+eqx
 tace