initial commit

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -0,0 +1,295 @@
+import time
+import json
+import random
+import base64
+from io import BytesIO
+from fire import Fire
+
+import streamlit as st
+from ase.atoms import Atoms
+from ase.build import bulk
+from ase.io import write
+from chemeleon import Chemeleon
+from chemeleon.visualize import Visualizer
+
+from utils import dict_to_atoms
+
+# Constants
+TIMESTEPS = 1000
+TRAJECTORY_STEPS = 100
+DEFAULT_NUM_SAMPLES = 3
+DEMO = False
+
+# Set page configuration
+st.set_page_config(page_title="Chemeleon", layout="wide")
+
+# Hide Streamlit's default menu and footer for a cleaner look
+hide_streamlit_style = """
+            <style>
+            #MainMenu {visibility: hidden;}
+            footer {visibility: hidden;}
+            </style>
+            """
+st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
+
+def demo_generator_structures(num_atoms, text_input, num_samples):
+    """
+    Generate crystal structures for demonstration purposes.
+    """
+    elements = random.choices(["Si", "Ge", "C", "Na", "Cl"], k=num_samples)
+    random_elements = random.choices(elements, k=num_atoms)
+
+    for step in range(TIMESTEPS):
+        time.sleep(0.001)
+        random_atoms = Atoms(
+            "Li",
+            positions=[[random.random() * 5 for _ in range(3)]],
+        )
+        atoms_list = [bulk(element, "fcc", a=5.43) for element in random_elements]
+        new_atoms_list = []
+
+        for atoms in atoms_list:
+            # Adding random atoms to each bulk structure
+            combined_atoms = atoms + random_atoms
+            new_atoms_list.append(combined_atoms)
+
+        yield new_atoms_list
+
+
+def generator_structures_chemeleon(
+    num_atoms, test_input, num_samples, use_client=False
+):
+    """
+    Generate crystal structures based on the given number of atoms and input text.
+    """
+    if use_client:
+        response = client(
+            url="https://8000-01j80snre5xdhq828s1q5brs0m.cloudspaces.litng.ai/predict",
+            n_samples=num_samples,
+            n_atoms=num_atoms,
+            text_input=test_input,
+        )
+
+        for line in response.iter_lines():
+            output = json.loads(line)["output"]
+            atom_dict = json.loads(output)
+            atoms_list = [dict_to_atoms(atoms_dict) for atoms_dict in atom_dict]
+            yield atoms_list
+    else:
+        chemeleon = Chemeleon.load_general_text_model()
+        for atoms_list in chemeleon.sample(
+            text_input=test_input,
+            n_atoms=num_atoms,
+            n_samples=num_samples,
+            stream=True,
+        ):
+            yield atoms_list
+
+
+def visualize_structure(atoms):
+    """
+    Visualize the given atomic structure using Plotly.
+    """
+    visualizer = Visualizer([atoms], atomic_size=0.6, resolution=20)
+    fig = visualizer.view()
+    return fig
+
+
+def visualize_trajectory(atoms_list):
+    """
+    Visualize the given atomic structure trajectory using Plotly.
+    """
+    visualizer = Visualizer(atoms_list, atomic_size=0.6, resolution=20)
+    fig = visualizer.view_trajectory(duration=1000)
+    return fig
+
+
+# Main application function
+def main(use_client=False):
+    # Initialize session state
+    if "structures" not in st.session_state:
+        st.session_state.structures = []
+    if "trajectory" not in st.session_state:
+        st.session_state.trajectory = []
+    if "progress_in_generating" not in st.session_state:
+        st.session_state["progress_in_generating"] = False
+
+    # Sidebar for user inputs
+    with st.sidebar:
+        st.image("assets/logo_static.jpg", width=200)
+        st.markdown(
+            """
+            <h1 style='text-align: center; color: #4CAF50;'>Chemeleon</h1>
+            <h3 style='text-align: center;'>A text-guided diffusion model for crystal structure generation</h3>
+            """,
+            unsafe_allow_html=True,
+        )
+        st.markdown("---")
+        description = st.text_input(
+            "Input your text prompt to generate crystal structures",
+            "A Crystal Structure of LiMnO4 with orthorhombic symmetry",
+            help="Examples: 'LiMnO4' or 'A Crystal Structure of BaTiO3 with cubic symmetry'",
+        )
+        num_atoms = st.slider(
+            "🔢 Number of Atoms:",
+            min_value=1,
+            max_value=20,
+            value=6,
+            help="Select the number of atoms in the unit cell.",
+        )
+        num_samples = st.number_input(
+            "🧪 Number of Samples:",
+            min_value=1,
+            max_value=5,
+            value=DEFAULT_NUM_SAMPLES,
+            step=1,
+            help="Determine how many structure samples to generate.",
+        )
+
+    # Generate Structures when button is clicked
+    if st.session_state["progress_in_generating"]:
+        # Clear previous structures
+        st.session_state.structures = []
+        st.session_state.trajectory = []
+
+        # Initialize progress bar in the sidebar
+        progress_placeholder = st.empty()
+        progress_bar = progress_placeholder.progress(0)
+
+        # Initialize loading animation
+        image_placeholder = st.empty()
+        with st.spinner("Generating structures..."):
+            with image_placeholder:
+                data_url = base64.b64encode(
+                    open("assets/logo.gif", "rb").read()
+                ).decode()
+                image_placeholder.markdown(
+                    f'<img src="data:image/gif;base64,{data_url}" width=100>',
+                    unsafe_allow_html=True,
+                )
+
+        # Generate structures
+        trajectory = []
+        if DEMO:
+            generator = demo_generator_structures(num_atoms, description, num_samples)
+        else:
+            generator = generator_structures_chemeleon(
+                num_atoms, description, num_samples, use_client
+            )
+        for step, atoms_list in enumerate(generator):
+            progress_bar.progress((step + 1) / TIMESTEPS)
+            if step % TRAJECTORY_STEPS == 0 or step == TIMESTEPS - 1:
+                st.session_state.structures = atoms_list
+                trajectory.append(atoms_list)
+
+        st.session_state.trajectory = trajectory
+
+        # Remove the progress bar
+        progress_placeholder.empty()
+
+        # Remove the loading animation
+        image_placeholder.empty()
+
+        # Reset the progress state
+        st.session_state["progress_in_generating"] = False
+
+        # Display success message
+        st.sidebar.success("✨ Structures generated successfully!")
+
+    with st.sidebar:
+        if st.button(
+            "Generate Structures 🚀",
+            disabled=st.session_state["progress_in_generating"],
+        ):
+            st.session_state["progress_in_generating"] = True
+            st.rerun()
+
+    # Check if structures are generated
+    if st.session_state.structures:
+        # Tabs for visualization
+        tabs = st.tabs(["Structure Visualization", "Trajectory Analysis"])
+
+        # Structure Visualization Tab
+        with tabs[0]:
+            col1, col2 = st.columns([1, 3])
+            with col1:
+                st.session_state.selected_sample_index = (
+                    st.radio(
+                        "Select Sample",
+                        options=list(range(1, num_samples + 1)),
+                        index=0,
+                        help="Choose which sample to visualize.",
+                    )
+                    - 1
+                )  # Adjust for zero-based indexing
+                # Download file
+                atoms = st.session_state.structures[
+                    st.session_state.selected_sample_index
+                ]
+                buffer = BytesIO()
+                write(buffer, atoms, format="cif")
+                buffer.seek(0)
+                st.download_button(
+                    label="Download CIF File",
+                    data=buffer,
+                    file_name=f"{str(atoms.symbols)}.cif",
+                    mime="chemical/cif",
+                )
+            with col2:
+                atoms = st.session_state.structures[
+                    st.session_state.selected_sample_index
+                ]
+                fig = visualize_structure(atoms)
+                st.plotly_chart(fig, use_container_width=True)
+
+        # Trajectory Analysis Tab
+        with tabs[1]:
+            if st.session_state.trajectory:
+                trajectory = [
+                    traj[st.session_state.selected_sample_index]
+                    for traj in st.session_state.trajectory
+                ]
+                tabs_2 = st.tabs(["Animation", "Step View"])
+                # Animation
+                with tabs_2[0]:
+                    fig = visualize_trajectory(trajectory)
+                    st.plotly_chart(fig, use_container_width=True)
+                # Slider
+                with tabs_2[1]:
+                    trajectory_index = st.slider(
+                        "Select Trajectory Step",
+                        min_value=0,
+                        max_value=len(trajectory) - 1,
+                        value=0,
+                        step=1,
+                        help="Navigate through different steps of the structure generation.",
+                    )
+                    selected_atoms = trajectory[trajectory_index]
+                    trajectory_fig = visualize_structure(selected_atoms)
+                    st.plotly_chart(trajectory_fig, use_container_width=True)
+            else:
+                st.info("No trajectory data available.")
+
+    # Footer
+    st.markdown(
+        """
+        <div style="text-align: center; color: grey; margin-top: 50px;">
+            <p style="font-size: 14px; margin: 0;">
+                Developed by 
+                <a href="https://hspark1212.github.io" target="_blank">Hyunsoo Park</a>, 
+                as a part of <a href="https://github.com/wmd-group" target="_blank">Materials Design Group</a> 
+                at Imperial College London
+            </p>
+            <p>
+                <a href="https://chemrxiv.org/engage/chemrxiv/article-details/6728e27cf9980725cf118177" target="_blank">Research Paper</a> | 
+                <a href="https://github.com/hspark1212/chemeleon" target="_blank">Repository</a>
+            </p>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
+
+
+if __name__ == "__main__":
+    Fire(main)  # Usage example: streamlit run app/streamlit_app.py -- --use_client=True

assets/style.css

@@ -0,0 +1,4 @@
+/* In your assets/style.css */
+#logo-image {
+    transition: all 0.5s ease-in-out;
+}

requirements.txt

@@ -0,0 +1 @@
+chemeleon

utils.py

@@ -0,0 +1,44 @@
+import numpy as np
+from ase import Atoms
+import plotly.graph_objects as go
+
+empty_fig = go.Figure(
+    data=[
+        go.Scatter(
+            x=[],
+            y=[],
+            mode="markers",
+            marker=dict(color="rgba(0,0,0,0)"),
+        )
+    ],
+    layout=go.Layout(
+        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        template="plotly_white",
+    ),
+)
+
+
+def atoms_to_dict(atoms):
+    """
+    Converts an ASE Atoms object to a dictionary for storing.
+    """
+    return {
+        "symbols": atoms.get_chemical_symbols(),
+        "positions": atoms.get_positions().tolist(),
+        "cell": atoms.get_cell().tolist(),
+        "pbc": atoms.get_pbc().tolist(),
+    }
+
+
+def dict_to_atoms(data):
+    """
+    Converts a dictionary back to an ASE Atoms object.
+    """
+    atoms = Atoms(
+        symbols=data["symbols"],
+        positions=np.array(data["positions"]),
+        cell=np.array(data["cell"]),
+        pbc=np.array(data["pbc"]),
+    )
+    return atoms