import os import json import sqlite3 from utils import print_col class Molecule: def __init__(self, name, multiplicity, geometry, properties): self.name = name self.multiplicity = multiplicity self.geometry = geometry self.properties = properties def to_dict(self): return { "name": self.name, "multiplicity": self.multiplicity, "geometry": self.geometry, "properties": self.properties, } @staticmethod def from_dict(data): return Molecule( name=data["name"], multiplicity=data["multiplicity"], geometry=data["geometry"], properties=data["properties"] ) def save_molecules_to_json(molecules, filename): with open(filename, 'w') as f: json_data = [molecule.to_dict() for molecule in molecules] json.dump(json_data, f, indent=4) def load_molecules_from_json(filename): with open(filename, 'r') as f: json_data = json.load(f) return [Molecule.from_dict(data) for data in json_data] def create_database(db_name): if os.path.exists(db_name): conn = sqlite3.connect(db_name) cursor = conn.cursor() # Check if the table already exists cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='molecules';") table_exists = cursor.fetchone() if table_exists: print_col(f"Database '{db_name}' already exists and table 'molecules' is already created.", "yellow") else: # Create the table if it does not exist cursor.execute('''CREATE TABLE molecules (name TEXT, multiplicity INTEGER, geometry TEXT, properties TEXT)''') conn.commit() print_col(f"Table 'molecules' created in existing database '{db_name}' successfully.", "green") conn.close() else: # Create the database and table conn = sqlite3.connect(db_name) cursor = conn.cursor() cursor.execute('''CREATE TABLE molecules (name TEXT, multiplicity INTEGER, geometry TEXT, properties TEXT)''') conn.commit() conn.close() print_col(f"Database '{db_name}' created and table 'molecules' added successfully.", "green") def add_molecule_to_db(db_name, molecule): conn = sqlite3.connect(db_name) cursor = conn.cursor() # Convert geometry and properties to JSON strings geometry_str = json.dumps(molecule.geometry) energies_str = json.dumps(molecule.properties) # Check if the molecule already exists cursor.execute("SELECT COUNT(*) FROM molecules WHERE name = ?", (molecule.name,)) count = cursor.fetchone()[0] if count > 0: print_col(f"Molecule '{molecule.name}' already exists in {db_name}.", "yellow") else: # Insert the molecule if it does not exist cursor.execute("INSERT INTO molecules (name, multiplicity, geometry, properties) VALUES (?, ?, ?, ?)", (molecule.name, molecule.multiplicity, geometry_str, energies_str)) conn.commit() print_col(f"'{molecule.name}' added to {db_name} successfully.", "green") conn.close() def remove_database(db_name): if os.path.exists(db_name): os.remove(db_name) print_col(f"Database '{db_name}' removed successfully.", "red") else: print_col(f"Database '{db_name}' does not exist.", "red") def get_molecules_from_db(db_name): conn = sqlite3.connect(db_name) cursor = conn.cursor() cursor.execute("SELECT name, multiplicity, geometry, properties FROM molecules") rows = cursor.fetchall() molecules = [] for row in rows: name, multiplicity, geometry_str, energies_str = row geometry = json.loads(geometry_str) properties = json.loads(energies_str) molecules.append(Molecule(name, multiplicity, geometry, properties)) conn.close() return molecules def generate_xyz(elements, filename="output.xyz", verbose=False): """ Generate an XYZ file from a list of elements. Parameters: elements (list): A list of dictionaries, where each dictionary represents an atom with its element and x, y, z coordinates. filename (str): The name of the output XYZ file. Default is 'output.xyz'. """ # Get the number of atoms num_atoms = len(elements) # Open the file in write mode with open(filename, 'w') as f: # Write the number of atoms f.write(f"{num_atoms}\n") # Write a comment line (can be left blank or customized) f.write("XYZ file generated by generate_xyz function\n") # Write the element and coordinates for atom in elements: element = atom['element'] x = atom['x'] y = atom['y'] z = atom['z'] f.write(f"{element} {x:.6f} {y:.6f} {z:.6f}\n") if(verbose): print(f"XYZ file '{filename}' generated successfully!")