QuAcK/tests/molecule.py

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!")