main.py

import streamlit as st
import os
from ase.io import read
from CifFile import ReadCif
import torch
from models.master import create_model
from process import process_ase

import gc
from io import BytesIO, StringIO
import argparse


torch.set_grad_enabled(False)


def main(local):
    

    model = create_model()
    st.title("CartNet Thermal Ellipsoid Prediction")
    st.image('fig/frontpage.png')

    st.markdown("""
    <h3 align="center">
    🌐 <a href="https://imatge-upc.github.io/CartNet/" target="_blank">Project</a> |
    📄 <a href="https://pubs.rsc.org/en/content/articlelanding/2024/dd/d4dd00352g" target="_blank">Paper</a> | 
    🐙 <a href="https://github.com/imatge-upc/CartNet" target="_blank">GitHub</a>
    </h3>
    """, unsafe_allow_html=True)

    st.markdown("""
                CartNet is a graph neural network specifically designed for predicting Anisotropic Displacement Parameters (ADPs) in crystal structures. The model has been trained on over 220,000 molecular crystal structures from the Cambridge Structural Database (CSD), making it highly accurate and robust for ADP prediction tasks. CartNet addresses the computational challenges of traditional methods by encoding the full 3D geometry of atomic structures into a Cartesian reference frame, bypassing the need for unit cell encoding. The model incorporates innovative features, including a neighbour equalization technique to enhance interaction detection and a Cholesky-based output layer to ensure valid ADP predictions. Additionally, it introduces a rotational SO(3) data augmentation technique to improve generalization across different crystal structure orientations, making the model highly efficient and accurate in predicting ADPs while significantly reducing computational costs.
    """)

    uploaded_file = st.file_uploader("Upload a CIF file", type=["cif"], accept_multiple_files=False)

    total_sctructures = None
    success = 0

    if uploaded_file is not None:
        try:
            filename = str(uploaded_file.name)
            file = BytesIO(uploaded_file.getbuffer())
            cif = ReadCif(file)

            if len(cif.keys())>1:
                st.warning("⚠️ **Warning**: Found " + str(len(cif.keys())) + " blocks in the CIF file. We will process all of them and export as separate CIF files.")
                total_sctructures = len(cif.keys())
            
            st.markdown(f"### CIF file: {filename}")
            for key in cif.keys():
                st.markdown(f"### Structure: {key}")
                try:
                    block = "data_"+str(key)+"\n"+ cif[key].printsection()
                    atoms = read(StringIO(block), format="cif")

                    if len(atoms.positions) > 1000 and not local:
                        st.error("""
                        ⚠️ **Warning**: The structure is too large. Please upload a smaller one or use the [local implementation of CartNet Web App](https://github.com/alexsoleg/cartnet-streamlit/).
                        """)
                        continue
                    
                    cif_data = cif[key]
                    if "_diffrn_ambient_temperature" in cif_data.keys():
                        temperature = float(cif_data["_diffrn_ambient_temperature"].split("(")[0])
                    elif "_cell_measurement_temperature" in cif_data.keys():
                        temperature = float(cif_data["_cell_measurement_temperature"].split("(")[0])
                    else:
                        st.error("Temperature not found in the CIF file. \
                                        Please provide a temperature in the field _diffrn_ambient_temperature o in the field _cell_measurement_temperature from the CIF file.")
                        continue
                    st.success("CIF file successfully read.")
                except Exception as e:
                    st.error(f"Error: {e}")
                    st.error(f"We couldn't find any structure for the block {key}. Please make sure the CIF is compatible with ASE. If the error message is a blank line, it means ASE didn't found any coordinates.")
                    
                    continue

                cif_file = process_ase(atoms, temperature, model)
                
                cif_file = BytesIO(cif_file.getvalue().encode())
                st.download_button(
                    label="Download processed CIF file",
                    data=cif_file,
                    file_name=f"output_{key}.cif",
                    mime="text/plain",
                    key=f"download_button_{key}"
                )

                gc.collect()
                success += 1
            if total_sctructures is not None:
                st.warning(f"A total of {total_sctructures} data blocks read and a total of {success} structures has been read and the ADP successfully predicted")
            gc.collect()
        except Exception as e:
            st.error(f"An error occurred while reading the CIF file: {e}")
    
    if not local:
        st.warning("""
        ⚠️ **Warning**: This online web application is designed for structures with up to 1000 atoms in the unit cell. For larger structures, please use the [local implementation of CartNet Web App](https://huggingface.co/spaces/alexsoleg/cartnet-demo/blob/main/README_offline.md).
        """)

    st.markdown("""
    📌 **Notes**: 
                
    - We use [ASE library](https://wiki.fysik.dtu.dk/ase/) for reading the CIF files, please make sure your CIF file is compatible.   
    - While CartNet excels at predicting ADPs for most crystal structures, it is currently optimized for ordered structures without solvents. For solvents or disordered structures, the performance may be slightly reduced.
    - Hydrogens are not predicted by CartNet. To avoid errors while reading the output file, the Hydrogens are included in the output CIF file as very small isotropic ADPs.
    - If CartNet cannot predict the ADP or produces unexpected results for a specific atom, an error will be displayed for that atom. The atom will then be ignored, and its ADPs will be saved as very small isotropic values.
    
    """)

    

    st.markdown("""
    ### How to cite

    If you use CartNet in your research, please cite our paper:

    ```bibtex
    @Article{D4DD00352G,
    author ="Solé, Àlex and Mosella-Montoro, Albert and Cardona, Joan and Gómez-Coca, Silvia and Aravena, Daniel and Ruiz, Eliseo and Ruiz-Hidalgo, Javier",
    title  ="A Cartesian encoding graph neural network for crystal structure property prediction: application to thermal ellipsoid estimation",
    journal  ="Digital Discovery",
    year  ="2025",
    pages  ="-",
    publisher  ="RSC",
    doi  ="10.1039/D4DD00352G",
    url  ="http://dx.doi.org/10.1039/D4DD00352G",}
    ```
    """)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--local', action='store_true')
    args = parser.parse_args()
    main(args.local)