Babu09/crystal-se3-autoencoder

See also: MEIDNet Prism - the current project of the same author: inverse design of crystalline materials from target properties (npj Computational Materials, 2026; model Babu09/MEIDNet).

SE(3)-Equivariant Hybrid Graph Autoencoder for Crystal Structures

This repository contains an implementation of a Hybrid Graph Autoencoder for Crystal Structures with SE(3)-Equivariance, improved lattice prediction, and enforced translational invariance.

Quick Start Guide

Option 1: Using the Web Interface

1. Visit our Hugging Face Space: Crystal Structure Reconstruction Demo
2. Upload your CIF file
3. Get the reconstructed structure instantly!

Option 2: Local Installation

1. Clone the repository:

   git clone https://huggingface.co/Babu09/crystal-se3-autoencoder
   cd crystal-se3-autoencoder
   

2. Create a Python environment (recommended):

   # Using conda
   conda create -n crystal_env python=3.8
   conda activate crystal_env
   
   # Or using venv
   python -m venv crystal_env
   # On Windows
   .\crystal_env\Scripts\activate
   # On Linux/Mac
   source crystal_env/bin/activate
   

3. Install dependencies:

   pip install -r requirements.txt
   

4. Test the model:

   import torch
   from h2_e3_n4 import GraphAutoencoderHybridSE3, reconstruct_and_write
   
   # Initialize model
   model = GraphAutoencoderHybridSE3(
       max_sites=20,
       num_species=56,
       latent_dim=128,
       node_hidden_dim=128,
       species_embedding_dim=64
   )
   
   # Load pre-trained weights
   model.load_state_dict(torch.load('hybrid_graph_autoencoder_se3.pt'))
   model.eval()
   
   # Reconstruct a structure
   reconstruct_and_write(model, 'input.cif', 'reconstructed.cif')
   

Example: Silicon Crystal Reconstruction

Here's an example of reconstructing a silicon crystal structure:

# Save this as Si.cif
data_Si
_symmetry_space_group_name_H-M   'P 1'
_cell_length_a   5.44370237
_cell_length_b   5.44370237
_cell_length_c   5.44370237
_cell_angle_alpha   90.00000000
_cell_angle_beta   90.00000000
_cell_angle_gamma   90.00000000
_symmetry_Int_Tables_number   1
_chemical_formula_structural   Si
_chemical_formula_sum   Si8
_cell_volume   161.31810739
_cell_formula_units_Z   8
loop_
 _symmetry_equiv_pos_site_id
 _symmetry_equiv_pos_as_xyz
  1  'x, y, z'
loop_
 _atom_type_symbol
 _atom_type_oxidation_number
  Si0+  0.0
loop_
 _atom_site_type_symbol
 _atom_site_label
 _atom_site_symmetry_multiplicity
 _atom_site_fract_x
 _atom_site_fract_y
 _atom_site_fract_z
 _atom_site_occupancy
  Si0+  Si0  1  0.75000000  0.75000000  0.25000000  1
  Si0+  Si1  1  0.00000000  0.50000000  0.50000000  1
  Si0+  Si2  1  0.75000000  0.25000000  0.75000000  1
  Si0+  Si3  1  0.00000000  0.00000000  0.00000000  1
  Si0+  Si4  1  0.25000000  0.75000000  0.75000000  1
  Si0+  Si5  1  0.50000000  0.50000000  0.00000000  1
  Si0+  Si6  1  0.25000000  0.25000000  0.25000000  1
  Si0+  Si7  1  0.50000000  0.00000000  0.50000000  1

Running the reconstruction:

reconstruct_and_write(model, 'Si.cif', 'Si_reconstructed.cif')

The model will generate a reconstructed structure that preserves:

• Crystal symmetry
• Atomic positions
• Chemical composition
• Lattice parameters (within ~0.3% accuracy)

Model Features

• SE(3)-equivariant layers in both encoder and decoder (EGNN-style)
• Translation invariance through center subtraction in encoder
• Improved lattice prediction without sigmoid activation
• Supports visualization with t-SNE plots
• Handles crystal structures with up to 20 sites
• Supports 56 different chemical species

Supported Elements

The model supports 56 chemical elements:

SPECIES_LIST_56 = [
    "Ag", "Al", "As", "Au", "B", "Ba", "Be", "Bi", "Ca", "Cd", "Co", "Cr", "Cs",
    "Cu", "F", "Fe", "Ga", "Ge", "Hf", "Hg", "In", "Ir", "K", "La", "Li", "Mg",
    "Mn", "Mo", "N", "Na", "Nb", "Ni", "O", "Os", "Pb", "Pd", "Pt", "Rb", "Re",
    "Rh", "Ru", "S", "Sb", "Sc", "Si", "Sn", "Sr", "Ta", "Te", "Ti", "Tl", "V",
    "W", "Y", "Zn", "Zr"
]

Troubleshooting

1. ImportError: No module named 'h2_e3_n4'

   # Make sure you're in the correct directory
   cd crystal-se3-autoencoder
   
   # Install the package in editable mode
   pip install -e .
   

2. CUDA out of memory

   • The model supports CPU inference if GPU memory is limited
   • Move the model to CPU:

     model = model.cpu()
     

3. File not found error

   • Make sure your CIF file is in the correct directory
   • Use absolute paths if needed:

     reconstruct_and_write(model, '/path/to/input.cif', '/path/to/output.cif')
     

Citation

If you use this code in your research, please cite:

@article{your-paper,
    title={SE(3)-Equivariant Hybrid Graph Autoencoder for Crystal Structures},
    author={[Your Name]},
    year={2025}
}

License

MIT License

Contact

For questions and support, please open an issue on the Hugging Face repository.