mirror of
https://github.com/pfloos/quack
synced 2024-09-27 20:11:05 +02:00
106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
|
|
import json
|
|
import sqlite3
|
|
|
|
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):
|
|
conn = sqlite3.connect(db_name)
|
|
cursor = conn.cursor()
|
|
cursor.execute('''CREATE TABLE IF NOT EXISTS molecules
|
|
(name TEXT, multiplicity INTEGER, geometry TEXT, properties 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.properties)
|
|
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, 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!")
|
|
|
|
|