mirror of
https://github.com/pfloos/quack
synced 2024-09-27 20:11:05 +02:00
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
|
|
import json
|
|
import sqlite3
|
|
|
|
class Molecule:
|
|
def __init__(self, name, multiplicity, geometry, energies):
|
|
self.name = name
|
|
self.multiplicity = multiplicity
|
|
self.geometry = geometry # List of tuples (atom, x, y, z)
|
|
self.energies = energies # Dictionary of dictionaries: {method: {basis: energy}}
|
|
|
|
def get_energy(self, method, basis):
|
|
"""Retrieve energy for a specific method and basis set."""
|
|
return self.energies.get(method, {}).get(basis, None)
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"name": self.name,
|
|
"multiplicity": self.multiplicity,
|
|
"geometry": self.geometry,
|
|
"energies": self.energies,
|
|
}
|
|
|
|
@staticmethod
|
|
def from_dict(data):
|
|
return Molecule(
|
|
name=data["name"],
|
|
multiplicity=data["multiplicity"],
|
|
geometry=data["geometry"],
|
|
energies=data["energies"]
|
|
)
|
|
|
|
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):
|
|
conn = sqlite3.connect(db_name)
|
|
cursor = conn.cursor()
|
|
cursor.execute('''CREATE TABLE IF NOT EXISTS molecules
|
|
(name TEXT, multiplicity INTEGER, geometry TEXT, energies TEXT)''')
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def add_molecule_to_db(db_name, molecule):
|
|
conn = sqlite3.connect(db_name)
|
|
cursor = conn.cursor()
|
|
geometry_str = json.dumps(molecule.geometry)
|
|
energies_str = json.dumps(molecule.energies)
|
|
cursor.execute("INSERT INTO molecules VALUES (?, ?, ?, ?)",
|
|
(molecule.name, molecule.multiplicity, geometry_str, energies_str))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_molecules_from_db(db_name):
|
|
conn = sqlite3.connect(db_name)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT name, multiplicity, geometry, energies FROM molecules")
|
|
rows = cursor.fetchall()
|
|
molecules = []
|
|
for row in rows:
|
|
name, multiplicity, geometry_str, energies_str = row
|
|
geometry = json.loads(geometry_str)
|
|
energies = json.loads(energies_str) # energies is a dictionary of dictionaries
|
|
molecules.append(Molecule(name, multiplicity, geometry, energies))
|
|
conn.close()
|
|
return molecules
|
|
|
|
def write_geometry_to_xyz(molecule, filename):
|
|
with open(filename, 'w') as f:
|
|
# First line: number of atoms
|
|
f.write(f"{len(molecule.geometry)}\n")
|
|
# Second line: empty comment line
|
|
f.write("\n")
|
|
# Remaining lines: atom positions
|
|
for atom, x, y, z in molecule.geometry:
|
|
f.write(f"{atom} {x:.6f} {y:.6f} {z:.6f}\n")
|
|
|
|
|