updated code, readme and image

README_streamlit.md → README_offline.md

@@ -1,19 +1,33 @@
 # CartNet Streamlit Web App
 
-![Pipeline](fig/pipeline.png)
+![Pipeline](fig/frontpage.png)
 
-### CartNet online demo available at: [CartNet Web App](https://cartnet-adp-estimation.streamlit.app)
+<h3 align="center">
+  🌐 <a href="https://imatge-upc.github.io/CartNet/" target="_blank">Project</a> | 
+  🐙 <a href="https://github.com/imatge-upc/CartNet" target="_blank">GitHub</a> | 
+  <!-- 📃 <a href="https://imatge-upc.github.io/CartNet/static/pdfs/CartNet.pdf" target="_blank">Paper</a> |  -->
+  🤗 <a href="https://huggingface.co/spaces/alexsoleg/cartnet-demo" target="_blank">Demo</a>
+</h3>
 
 
 CartNet is a graph neural network specifically designed to predict 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.
 
 This repository contains a web application based on the official implementation of CartNet, which can be found at [imatge-upc/CartNet](https://github.com/imatge-upc/CartNet).
 
-⚠️ **Warning**: The online web application can only process systems with less than 300 atoms in the unit cell. For large systems, please use the local application.
+⚠️ **Warning**: The online web application can only process systems with less than 1000 atoms in the unit cell. For larger systems, please use the local application.
 
 ## Local Application
 ### Installation of the local application
 
+Clone the following repository:
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com)
+git lfs install
+
+git clone https://huggingface.co/spaces/alexsoleg/cartnet-demo
+```
+
 To set up the local application, you need to install the dependencies listed in `requirements.txt`. You can do this by running the following command:
 
 ```bash
@@ -37,7 +51,7 @@ python predict.py input.cif output.cif
 Or, if you prefer, you can use the browser app on your local machine without the atom number limitation by running:
 
 ```bash
-streamlit run main_local.py
+streamlit run main.py -- --local
 ```
 
 ## How to cite

main.py

@@ -2,23 +2,22 @@ import streamlit as st
 import os
 from ase.io import read
 from CifFile import ReadCif
-from torch_geometric.data import Data, Batch
 import torch
 from models.master import create_model
-from process import process_data
-from utils import radius_graph_pbc
+from process import process_ase
+
 import gc
 from io import BytesIO, StringIO
+import argparse
+
 
-MEAN_TEMP = torch.tensor(192.1785) #training temp mean
-STD_TEMP = torch.tensor(81.2135) #training temp std
+torch.set_grad_enabled(False)
 
 
-@torch.no_grad()
-def main():
+def main(local):
     model = create_model()
     st.title("CartNet Thermal Ellipsoid Prediction")
-    st.image('fig/pipeline.png')
+    st.image('fig/frontpage.png')
 
     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.
@@ -46,7 +45,7 @@ def main():
                     block = "data_"+str(key)+"\n"+ cif[key].printsection()
                     atoms = read(StringIO(block), format="cif")
 
-                    if len(atoms.positions) > 1000:
+                    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/).
                         """)
@@ -68,39 +67,7 @@ def main():
                     
                     continue
 
-                data = Data()
-                data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)
-                
-                data.pos = torch.tensor(atoms.positions, dtype=torch.float32)
-                data.temperature_og = torch.tensor([temperature], dtype=torch.float32)
-                data.temperature = (data.temperature_og - MEAN_TEMP) / STD_TEMP
-                data.cell = torch.tensor(atoms.cell.array, dtype=torch.float32).unsqueeze(0)
-
-                data.pbc = torch.tensor([True, True, True])
-                data.natoms = len(atoms)
-
-                del atoms
-                gc.collect()
-                batch = Batch.from_data_list([data])
-                
-
-                edge_index, _, _, edge_attr = radius_graph_pbc(batch, 5.0, 64)
-                del batch
-                gc.collect()
-                data.cart_dist = torch.norm(edge_attr, dim=-1)
-                data.cart_dir = torch.nn.functional.normalize(edge_attr, dim=-1)
-                data.edge_index = edge_index
-                data.non_H_mask = data.x != 1
-                delattr(data, "pbc")
-                delattr(data, "natoms")
-                batch = Batch.from_data_list([data])
-                del data, edge_index, edge_attr
-                gc.collect()
-
-                st.success("Graph successfully created.")
-
-                cif_file = process_data(batch, model)
-                st.success("ADPs successfully predicted.")
+                cif_file = process_ase(atoms, temperature, model)
                 
                 cif_file = BytesIO(cif_file.getvalue().encode())
                 st.download_button(
@@ -118,10 +85,11 @@ def main():
             gc.collect()
         except Exception as e:
             st.error(f"An error occurred while reading the CIF file: {e}")
-
-    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://github.com/alexsoleg/cartnet-streamlit/).
-    """)
+    
+    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://github.com/alexsoleg/cartnet-streamlit/).
+        """)
 
     st.warning("""
     ⚠️ **Warning**: We use [ASE library](https://wiki.fysik.dtu.dk/ase/) for reading the CIF files, please make sure your CIF file is compatible.
@@ -150,6 +118,9 @@ def main():
     """)
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--local', action='store_true')
+    args = parser.parse_args()
+    main(args.local)
 
 

main_local.py

@@ -1,139 +0,0 @@
-import streamlit as st
-import os
-from ase.io import read
-from CifFile import ReadCif
-from torch_geometric.data import Data, Batch
-import torch
-from models.master import create_model
-from process import process_data
-from utils import radius_graph_pbc
-from io import BytesIO, StringIO
-import gc
-
-MEAN_TEMP = torch.tensor(192.1785) #training temp mean
-STD_TEMP = torch.tensor(81.2135) #training temp std
-
-
-@torch.no_grad()
-def main():
-    model = create_model()
-    st.title("CartNet Thermal Ellipsoid Prediction")
-    st.image('fig/pipeline.png')
-
-    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)
-
-    
-    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.")
-            
-            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")
-                    
-                    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
-
-                data = Data()
-                data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)
-                
-                data.pos = torch.tensor(atoms.positions, dtype=torch.float32)
-                data.temperature_og = torch.tensor([temperature], dtype=torch.float32)
-                data.temperature = (data.temperature_og - MEAN_TEMP) / STD_TEMP
-                data.cell = torch.tensor(atoms.cell.array, dtype=torch.float32).unsqueeze(0)
-
-                data.pbc = torch.tensor([True, True, True])
-                data.natoms = len(atoms)
-
-                del atoms
-                gc.collect()
-                batch = Batch.from_data_list([data])
-                
-
-                edge_index, _, _, edge_attr = radius_graph_pbc(batch, 5.0, 64)
-                del batch
-                gc.collect()
-                data.cart_dist = torch.norm(edge_attr, dim=-1)
-                data.cart_dir = torch.nn.functional.normalize(edge_attr, dim=-1)
-                data.edge_index = edge_index
-                data.non_H_mask = data.x != 1
-                delattr(data, "pbc")
-                delattr(data, "natoms")
-                batch = Batch.from_data_list([data])
-                del data, edge_index, edge_attr
-                gc.collect()
-
-                st.success("Graph successfully created.")
-
-                cif_file = process_data(batch, model)
-                st.success("ADPs successfully predicted.")
-                
-                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()
-            
-            gc.collect()
-        except Exception as e:
-            st.error(f"An error occurred while reading the CIF file: {e}")
-
-    
-    st.markdown("""
-    📌 The official implementation of the paper with all experiments can be found at [CartNet GitHub Repository](https://github.com/imatge-upc/CartNet).
-    """)
-
-    st.warning("""
-    ⚠️ **Warning**: We use [ASE library](https://wiki.fysik.dtu.dk/ase/) for reading the CIF files, please make sure your CIF file is compatible.
-    """)
-
-    st.markdown("""
-    ### How to cite
-
-    If you use CartNet in your research, please cite our paper:
-
-    ```bibtex
-    @article{your_paper_citation,
-    title={Title of the Paper},
-    author={Author1 and Author2 and Author3},
-    journal={Journal Name},
-    year={2023},
-    volume={XX},
-    number={YY},
-    pages={ZZZ}
-    }
-    ```
-    """)
-
-if __name__ == "__main__":
-    main()

process.py

@@ -1,10 +1,51 @@
+import streamlit as st
 import torch
+from torch_geometric.data import Data, Batch
 from ase.io import write
 from ase import Atoms
 import gc
 from io import BytesIO, StringIO
+from utils import radius_graph_pbc
+
+MEAN_TEMP = torch.tensor(192.1785) #training temp mean
+STD_TEMP = torch.tensor(81.2135) #training temp std
+
+def process_ase(atoms, temperature, model):
+    data = Data()
+    data.x = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int32)
+    
+    data.pos = torch.tensor(atoms.positions, dtype=torch.float32)
+    data.temperature_og = torch.tensor([temperature], dtype=torch.float32)
+    data.temperature = (data.temperature_og - MEAN_TEMP) / STD_TEMP
+    data.cell = torch.tensor(atoms.cell.array, dtype=torch.float32).unsqueeze(0)
+
+    data.pbc = torch.tensor([True, True, True])
+    data.natoms = len(atoms)
+
+    del atoms
+    gc.collect()
+    batch = Batch.from_data_list([data])
+    
+
+    edge_index, _, _, edge_attr = radius_graph_pbc(batch, 5.0, 64)
+    del batch
+    gc.collect()
+    data.cart_dist = torch.norm(edge_attr, dim=-1)
+    data.cart_dir = torch.nn.functional.normalize(edge_attr, dim=-1)
+    data.edge_index = edge_index
+    data.non_H_mask = data.x != 1
+    delattr(data, "pbc")
+    delattr(data, "natoms")
+    batch = Batch.from_data_list([data])
+    del data, edge_index, edge_attr
+    gc.collect()
+
+    st.success("Graph successfully created.")
+
+    cif_file = process_data(batch, model)
+    st.success("ADPs successfully predicted.")
+    return cif_file
 
-@torch.no_grad()
 def process_data(batch, model, output_file="output.cif"):
     atoms = batch.x.numpy().astype(int)  # Atomic numbers
     positions = batch.pos.numpy()  # Atomic positions
@@ -72,7 +113,7 @@ def process_data(batch, model, output_file="output.cif"):
             element_count[element] = 0
         element_count[element] += 1
         label = f"{element}{element_count[element]}"
-        u_iso = torch.trace(adps[indices[i]]).mean() if element != 'H' else 0.01
+        u_iso = torch.trace(adps[indices[i]]).mean() if element != 'H' else 0.0001
         type = "Uani" if element != 'H' else "Uiso"
         cif_file.write(f"{label} {element} {frac_pos[0]} {frac_pos[1]} {frac_pos[2]} {u_iso} {type}\n")