Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import gradio as gr | |
| import groq | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import py3Dmol | |
| from pymatgen.core.structure import Structure | |
| from pymatgen.io.cif import CifWriter | |
| from pymatgen.io.xyz import XYZ | |
| from pymatgen.symmetry.analyzer import SpacegroupAnalyzer | |
| import requests | |
| from dotenv import load_dotenv | |
| import tempfile | |
| import base64 | |
| # Load environment variables | |
| load_dotenv() | |
| # Retrieve the API key from the environment variable | |
| groq_api_key = os.getenv("GROQ_API_KEY") | |
| if not groq_api_key: | |
| raise ValueError("GROQ_API_KEY environment variable not set") | |
| # Initialize Groq client | |
| client = groq.Groq(api_key=groq_api_key) | |
| # Function to generate crystal structures | |
| def generate_crystal_structure(material_name): | |
| """Generate a crystal structure for a given material.""" | |
| # This is a simplified version - in a real application, you would use a database or API | |
| # to get real crystal structure data for materials | |
| # Dictionary of common materials and their crystal structures | |
| material_structures = { | |
| "Silicon": { | |
| "spacegroup": "Fd-3m", | |
| "lattice": [[0, 5.431/2, 5.431/2], [5.431/2, 0, 5.431/2], [5.431/2, 5.431/2, 0]], | |
| "species": ["Si"], | |
| "coords": [[0, 0, 0]] | |
| }, | |
| "Titanium Dioxide": { | |
| "spacegroup": "P42/mnm", | |
| "lattice": [[4.594, 0, 0], [0, 4.594, 0], [0, 0, 2.959]], | |
| "species": ["Ti", "O", "O"], | |
| "coords": [[0, 0, 0], [0.3053, 0.3053, 0], [0.3053, 0.6947, 0.5]] | |
| }, | |
| "Graphene": { | |
| "spacegroup": "P6/mmm", | |
| "lattice": [[2.46, 0, 0], [-2.46/2, 2.46*np.sqrt(3)/2, 0], [0, 0, 15]], | |
| "species": ["C", "C"], | |
| "coords": [[0, 0, 0], [1/3, 2/3, 0]] | |
| }, | |
| "Copper": { | |
| "spacegroup": "Fm-3m", | |
| "lattice": [[3.615, 0, 0], [0, 3.615, 0], [0, 0, 3.615]], | |
| "species": ["Cu"], | |
| "coords": [[0, 0, 0]] | |
| }, | |
| "Aluminum": { | |
| "spacegroup": "Fm-3m", | |
| "lattice": [[4.05, 0, 0], [0, 4.05, 0], [0, 0, 4.05]], | |
| "species": ["Al"], | |
| "coords": [[0, 0, 0]] | |
| }, | |
| "Gold": { | |
| "spacegroup": "Fm-3m", | |
| "lattice": [[4.078, 0, 0], [0, 4.078, 0], [0, 0, 4.078]], | |
| "species": ["Au"], | |
| "coords": [[0, 0, 0]] | |
| }, | |
| "Diamond": { | |
| "spacegroup": "Fd-3m", | |
| "lattice": [[0, 3.567/2, 3.567/2], [3.567/2, 0, 3.567/2], [3.567/2, 3.567/2, 0]], | |
| "species": ["C"], | |
| "coords": [[0, 0, 0]] | |
| }, | |
| "Graphite": { | |
| "spacegroup": "P63/mmc", | |
| "lattice": [[2.46, 0, 0], [-1.23, 2.13, 0], [0, 0, 6.71]], | |
| "species": ["C", "C", "C", "C"], | |
| "coords": [[0, 0, 0], [0, 0, 0.5], [1/3, 2/3, 0], [2/3, 1/3, 0.5]] | |
| } | |
| } | |
| # Try to match the material name with our database (case insensitive) | |
| for known_material, structure_data in material_structures.items(): | |
| if material_name.lower() in known_material.lower(): | |
| structure = Structure.from_spacegroup( | |
| structure_data["spacegroup"], | |
| lattice=structure_data["lattice"], | |
| species=structure_data["species"], | |
| coords=structure_data["coords"] | |
| ) | |
| return structure | |
| # If material not found, create a generic structure | |
| return Structure.from_spacegroup( | |
| "Pm-3m", | |
| lattice=[[4.0, 0, 0], [0, 4.0, 0], [0, 0, 4.0]], | |
| species=["X"], | |
| coords=[[0, 0, 0]] | |
| ) | |
| # Function to get material recommendations from LLM | |
| def get_material_recommendations(query): | |
| """Get material recommendations from the LLM based on user query.""" | |
| system_prompt = """You are a materials science expert. Your task is to recommend the 3 best materials for a specific application or with certain properties based on the user's query. | |
| For each material, provide: | |
| 1. Material name | |
| 2. Chemical formula | |
| 3. Key properties relevant to the application | |
| 4. Why it's suitable for the application | |
| 5. Any limitations or considerations | |
| Format your response as a JSON object with the following structure: | |
| { | |
| "materials": [ | |
| { | |
| "name": "Material Name", | |
| "formula": "Chemical Formula", | |
| "properties": "Key properties relevant to the application", | |
| "suitability": "Why it's suitable for the application", | |
| "limitations": "Any limitations or considerations" | |
| }, | |
| // Second material | |
| // Third material | |
| ] | |
| } | |
| Ensure your response is strictly in this JSON format with no additional text.""" | |
| try: | |
| completion = client.chat.completions.create( | |
| model="deepseek-r1", | |
| messages=[ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": query} | |
| ], | |
| temperature=0.2, | |
| max_tokens=1000 | |
| ) | |
| response_text = completion.choices[0].message.content | |
| # Extract JSON from the response | |
| try: | |
| # Try to parse the entire response as JSON | |
| recommendations = json.loads(response_text) | |
| except json.JSONDecodeError: | |
| # If that fails, try to extract JSON using string manipulation | |
| json_start = response_text.find('{') | |
| json_end = response_text.rfind('}') + 1 | |
| if json_start >= 0 and json_end > json_start: | |
| json_str = response_text[json_start:json_end] | |
| recommendations = json.loads(json_str) | |
| else: | |
| raise ValueError("Could not extract valid JSON from LLM response") | |
| return recommendations | |
| except Exception as e: | |
| return {"error": str(e)} | |
| # Function to get crystal structure information | |
| def get_crystal_structure_info(material_name): | |
| """Get crystal structure information for a given material.""" | |
| try: | |
| # Generate the crystal structure | |
| structure = generate_crystal_structure(material_name) | |
| # Create a CIF file | |
| cif_writer = CifWriter(structure) | |
| with tempfile.NamedTemporaryFile(suffix='.cif', delete=False) as temp_cif: | |
| cif_writer.write_file(temp_cif.name) | |
| cif_path = temp_cif.name | |
| with open(cif_path, 'r') as f: | |
| cif_content = f.read() | |
| # Create an XYZ file | |
| with tempfile.NamedTemporaryFile(suffix='.xyz', delete=False) as temp_xyz: | |
| xyz_path = temp_xyz.name | |
| # Convert structure to XYZ format | |
| ase_atoms = structure.to_ase() | |
| from ase.io import write as ase_write | |
| ase_write(xyz_path, ase_atoms, format='xyz') | |
| with open(xyz_path, 'r') as f: | |
| xyz_content = f.read() | |
| # Get space group information | |
| analyzer = SpacegroupAnalyzer(structure) | |
| spacegroup = analyzer.get_space_group_symbol() | |
| # Generate 3D visualization | |
| view = py3Dmol.view(width=500, height=400) | |
| view.addModel(xyz_content, 'xyz') | |
| view.setStyle({'sphere': {'colorscheme': 'Jmol', 'scale': 0.3}, | |
| 'stick': {'radius': 0.2}}) | |
| view.zoomTo() | |
| view.spin(True) | |
| view.setBackgroundColor('white') | |
| view.render() | |
| # Convert the view to HTML | |
| html_str = view._make_html() | |
| # Clean up temporary files | |
| os.unlink(cif_path) | |
| os.unlink(xyz_path) | |
| return { | |
| "material_name": material_name, | |
| "formula": structure.composition.reduced_formula, | |
| "space_group": spacegroup, | |
| "num_atoms": len(structure), | |
| "lattice_parameters": { | |
| "a": structure.lattice.a, | |
| "b": structure.lattice.b, | |
| "c": structure.lattice.c, | |
| "alpha": structure.lattice.alpha, | |
| "beta": structure.lattice.beta, | |
| "gamma": structure.lattice.gamma | |
| }, | |
| "cif_content": cif_content, | |
| "xyz_content": xyz_content, | |
| "visualization_html": html_str | |
| } | |
| except Exception as e: | |
| return {"error": str(e)} | |
| # Function to create a downloadable file | |
| def create_downloadable_file(content, filename): | |
| """Create a downloadable file with the given content.""" | |
| with open(filename, 'w') as f: | |
| f.write(content) | |
| return filename | |
| # Gradio interface | |
| def process_query(query): | |
| """Process the user query and return material recommendations and crystal structure visualization.""" | |
| try: | |
| # Get material recommendations from LLM | |
| recommendations = get_material_recommendations(query) | |
| if "error" in recommendations: | |
| return f"Error getting recommendations: {recommendations['error']}", None, None, None | |
| # Format the recommendations as text | |
| recommendation_text = "# Material Recommendations\n\n" | |
| for i, material in enumerate(recommendations.get("materials", [])): | |
| recommendation_text += f"## {i+1}. {material.get('name', 'Unknown')}\n" | |
| recommendation_text += f"**Formula:** {material.get('formula', 'N/A')}\n\n" | |
| recommendation_text += f"**Properties:** {material.get('properties', 'N/A')}\n\n" | |
| recommendation_text += f"**Suitability:** {material.get('suitability', 'N/A')}\n\n" | |
| recommendation_text += f"**Limitations:** {material.get('limitations', 'N/A')}\n\n" | |
| # Get crystal structure information for the first recommended material | |
| cif_file = None | |
| xyz_file = None | |
| if recommendations.get("materials") and len(recommendations.get("materials")) > 0: | |
| first_material = recommendations["materials"][0]["name"] | |
| structure_info = get_crystal_structure_info(first_material) | |
| if "error" in structure_info: | |
| structure_html = f"<p>Error getting crystal structure: {structure_info['error']}</p>" | |
| else: | |
| # Add crystal structure information to the recommendation text | |
| recommendation_text += f"# Crystal Structure of {structure_info['material_name']}\n\n" | |
| recommendation_text += f"**Formula:** {structure_info['formula']}\n\n" | |
| recommendation_text += f"**Space Group:** {structure_info['space_group']}\n\n" | |
| recommendation_text += f"**Number of Atoms:** {structure_info['num_atoms']}\n\n" | |
| recommendation_text += "**Lattice Parameters:**\n" | |
| recommendation_text += f"a = {structure_info['lattice_parameters']['a']:.4f} Å, " | |
| recommendation_text += f"b = {structure_info['lattice_parameters']['b']:.4f} Å, " | |
| recommendation_text += f"c = {structure_info['lattice_parameters']['c']:.4f} Å\n" | |
| recommendation_text += f"α = {structure_info['lattice_parameters']['alpha']:.2f}°, " | |
| recommendation_text += f"β = {structure_info['lattice_parameters']['beta']:.2f}°, " | |
| recommendation_text += f"γ = {structure_info['lattice_parameters']['gamma']:.2f}°\n\n" | |
| # Create HTML for the 3D visualization | |
| structure_html = structure_info['visualization_html'] | |
| # Create downloadable files | |
| cif_file = create_downloadable_file(structure_info['cif_content'], f"{structure_info['material_name'].replace(' ', '_')}.cif") | |
| xyz_file = create_downloadable_file(structure_info['xyz_content'], f"{structure_info['material_name'].replace(' ', '_')}.xyz") | |
| else: | |
| structure_html = "<p>No materials recommended to visualize.</p>" | |
| return recommendation_text, structure_html, cif_file, xyz_file | |
| except Exception as e: | |
| return f"Error processing query: {str(e)}", None, None, None | |
| # Create the Gradio interface | |
| with gr.Blocks(title="Material Science Expert") as demo: | |
| gr.Markdown("# Material Science Expert") | |
| gr.Markdown("Ask for a material for a specific application or with certain properties.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| query_input = gr.Textbox( | |
| label="Your Query", | |
| placeholder="I need a material with high thermal conductivity for electronics cooling.", | |
| lines=3 | |
| ) | |
| submit_btn = gr.Button("Get Recommendations") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| recommendations_output = gr.Markdown(label="Material Recommendations") | |
| with gr.Column(scale=1): | |
| structure_output = gr.HTML(label="Crystal Structure Visualization") | |
| with gr.Row(): | |
| with gr.Column(): | |
| cif_file_output = gr.File(label="Download CIF File") | |
| xyz_file_output = gr.File(label="Download XYZ File") | |
| submit_btn.click( | |
| fn=process_query, | |
| inputs=[query_input], | |
| outputs=[recommendations_output, structure_output, cif_file_output, xyz_file_output] | |
| ) | |
| gr.Markdown(""" | |
| ## Example Queries: | |
| - I need a material with high thermal conductivity for electronics cooling. | |
| - What are the best materials for solar cell applications? | |
| - Recommend materials with high strength-to-weight ratio for aerospace applications. | |
| - I need a transparent conductive material for touchscreens. | |
| """) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0") | |