add initial version of QE simulations (#16)

* add initial version of QE simulations

Author: jacobfilik

database/tables.sql

@@ -146,9 +146,17 @@ CREATE TABLE fdmnes_simulation (
 
 COMMENT ON TABLE fdmnes_simulation IS 'Specific information for a FDMNES simulation';
 
+CREATE TYPE qe_edge_enum AS ENUM('k', 'l1', 'l2', 'l23');
+
+CREATE TYPE conductivity_enum AS ENUM('metallic', 'semiconductor', 'insulator');
+
 CREATE TABLE qe_simulation (
     simulation_id INTEGER PRIMARY KEY,
     simulation_type_id INTEGER GENERATED ALWAYS AS (3) STORED,
+    crystal_structure_id INTEGER NOT NULL,
+    absorbing_atom INTEGER NOT NULL,
+    edge qe_edge_enum NOT NULL,
+    conductivity conductivity_enum NOT NULL,
     FOREIGN KEY(simulation_id, simulation_type_id) REFERENCES simulation (id,simulation_type_id)
 );
 

web-conexs-api/src/slurm_submission_service/crud.py

@@ -2,12 +2,17 @@
 
 from sqlmodel import or_, select
 
-from web_conexs_api.jobfilebuilders import build_fdmnes_inputfile, build_orca_input_file
+from web_conexs_api.jobfilebuilders import (
+    build_fdmnes_inputfile,
+    build_orca_input_file,
+    build_qe_inputfile,
+)
 from web_conexs_api.models.models import (
     CrystalStructure,
     FdmnesSimulation,
     MolecularStructure,
     OrcaSimulation,
+    QESimulation,
     Simulation,
     SimulationStatus,
 )
@@ -91,3 +96,19 @@ def update_simulation(session, simulation: Simulation):
     session.commit()
     session.refresh(simulation)
     return simulation
+
+
+def get_qe_simulation(session, id) -> QESimulation:
+    simulation = session.get(QESimulation, id)
+
+    if simulation:
+        return simulation
+    else:
+        raise RuntimeError("FDMNES simulation not found")
+
+
+def get_qe_jobfile(session, id):
+    qe_simulation = get_qe_simulation(session, id)
+    structure = get_crystal_structure(session, qe_simulation.crystal_structure_id)
+
+    return build_qe_inputfile(qe_simulation, structure)

web-conexs-api/src/slurm_submission_service/submitter.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import os
+import shutil
 import sys
 import time
 import uuid
@@ -12,12 +13,14 @@
 import requests
 
 from web_conexs_api.database import get_session
+from web_conexs_api.jobfilebuilders import build_qe_xspectra_input
 from web_conexs_api.models.models import OrcaCalculation, Simulation, SimulationStatus
 
 from .crud import (
     get_active_simulations,
     get_fdmnes_jobfile,
     get_orca_jobfile_with_technique,
+    get_qe_jobfile,
     get_submitted_simulations,
     update_simulation,
 )
@@ -32,10 +35,14 @@
 SLURM_USER = os.environ.get("SLURM_USER")
 SLURM_API = os.environ.get("SLURM_API")
 SLURM_PARTITION = os.environ.get("SLURM_PARTITION")
+SLURM_RESPONSE_KEY = os.environ.get("SLURM_RESPONSE_KEY", "account")
 
 ORCA_IMAGE = os.environ.get("ORCA_IMAGE")
 FDMNES_IMAGE = os.environ.get("FDMNES_IMAGE")
+QE_IMAGE = os.environ.get("QE_IMAGE")
 CONTAINER_IMAGE_DIR = os.environ.get("CONTAINER_IMAGE_DIR")
+PP_DIR = os.environ.get("PP_DIR")
+WFC_DIR = os.environ.get("WFC_DIR")
 
 JOB_RUNNING = "RUNNING"
 JOB_COMPLETED = "COMPLETED"
@@ -225,6 +232,8 @@ def run_update():
                 submit_orca(session, sim)
             if sim.simulation_type_id == 2:
                 submit_fdmnes(session, sim)
+            if sim.simulation_type_id == 3:
+                submit_qe(session, sim)
 
                 # orca = get_orca_jobfile(session, sim.id)
                 # determine workspace
@@ -256,7 +265,8 @@ def run_update():
         job_map = {}
 
         for j in jobs:
-            if j["account"] == SLURM_USER:
+            # NEED TO CHANGE TO USER
+            if j[SLURM_RESPONSE_KEY] == SLURM_USER:
                 job_map[j["job_id"]] = {"state": j["job_state"][0]}
 
         if len(active) != 0:
@@ -290,6 +300,60 @@ def run_update():
                 update_simulation(session, a)
 
 
+def submit_qe(session, sim: Simulation):
+    jobfile, absorbing_atom, abs_edge, pp, pp_abs = get_qe_jobfile(session, sim.id)
+    application_name = "qe"
+    user = sim.person.identifier
+    uid = uuid.uuid4()
+
+    job_name = application_name + "-" + str(uid)
+    working_dir = ROOT_DIR + user + "/" + job_name
+    cluster_dir = CLUSTER_ROOT_DIR + user + "/" + job_name
+
+    Path(working_dir).mkdir(parents=True)
+
+    with open(working_dir / Path("job.inp"), "w+") as f:
+        f.write(jobfile)
+
+    for pp_filename in pp:
+        shutil.copy(os.path.join(PP_DIR, pp_filename), working_dir)
+
+    xspectra_input_file = working_dir / Path("xspectra_input.inp")
+
+    core_file = Path(pp_abs).with_suffix(".wfc")
+
+    wffile = WFC_DIR / core_file
+    shutil.copy(wffile, working_dir)
+
+    xspectra_jobfile = build_qe_xspectra_input(absorbing_atom, abs_edge, str(core_file))
+
+    with open(xspectra_input_file, "w+") as ff:
+        ff.write(xspectra_jobfile)
+
+    qe_sif = os.path.join(CONTAINER_IMAGE_DIR, QE_IMAGE)
+
+    script = (
+        "#!/bin/bash\n"
+        + f"singularity exec {qe_sif} mpirun"
+        + f" -np {sim.n_cores} pw.x -in job.inp > result.pwo \n"
+        + f"singularity exec {qe_sif} mpirun"
+        + f" -np {sim.n_cores} xspectra.x -in xspectra_input.inp"
+    )
+
+    try:
+        job_id = build_job_and_run(
+            script, user, working_dir, job_name, sim.n_cores, sim.memory, cluster_dir
+        )
+        sim.job_id = job_id
+        sim.working_directory = working_dir
+        sim.submission_date = datetime.datetime.now()
+        sim.status = SimulationStatus.submitted
+        update_simulation(session, sim)
+    except Exception:
+        sim.status = SimulationStatus.failed
+        update_simulation(session, sim)
+
+
 def test_read():
     with contextmanager(get_session)() as session:
         sims = get_submitted_simulations(session)

web-conexs-api/src/web_conexs_api/app.py

@@ -1,12 +1,13 @@
 from fastapi import FastAPI
 from fastapi_pagination import add_pagination
 
-from .routers import crystals, fdmnes, molecules, orca, simulations, user
+from .routers import crystals, fdmnes, molecules, orca, qe, simulations, user
 
 app = FastAPI()
 
 app.include_router(orca.router)
 app.include_router(fdmnes.router)
+app.include_router(qe.router)
 app.include_router(molecules.router)
 app.include_router(crystals.router)
 app.include_router(simulations.router)

web-conexs-api/src/web_conexs_api/crud.py

@@ -8,7 +8,11 @@
 from fastapi_pagination.ext.sqlalchemy import paginate
 from sqlmodel import Session, and_, select
 
-from .jobfilebuilders import build_fdmnes_inputfile, build_orca_input_file
+from .jobfilebuilders import (
+    build_fdmnes_inputfile,
+    build_orca_input_file,
+    build_qe_inputfile,
+)
 from .models.models import (
     CrystalStructure,
     CrystalStructureInput,
@@ -20,6 +24,8 @@
     OrcaSimulationInput,
     Person,
     PersonInput,
+    QESimulation,
+    QESimulationInput,
     Simulation,
 )
 
@@ -339,3 +345,82 @@ def get_user(session, user_id):
         person = Person(id=None, identifier=user_id, admin=False)
 
     return person
+
+
+def submit_qe_simulation(
+    qe_input: QESimulationInput, session: Session, user_id: str
+) -> QESimulation:
+    person = get_or_create_person(session, user_id)
+
+    smodel = {
+        "person_id": person.id,
+        "simulation_type_id": 3,
+        "request_date": datetime.datetime.now(),
+    }
+
+    simulation = Simulation.model_validate(smodel)
+    qe = QESimulation.model_validate(qe_input)
+    qe.simulation = simulation
+
+    session.add(qe)
+    session.commit()
+    session.refresh(qe)
+
+    return qe
+
+    # get_qe_xas,
+
+
+def get_qe_simulation(session, id, user_id) -> QESimulation:
+    statement = (
+        select(QESimulation)
+        .join(Simulation)
+        .join(Person)
+        .where(and_(Person.identifier == user_id, QESimulation.simulation_id == id))
+    )
+
+    simulation = session.exec(statement).first()
+
+    if simulation:
+        return simulation
+    else:
+        raise HTTPException(status_code=404, detail=f"No simulation with id={id}")
+
+
+def get_qe_jobfile(session, id, user_id):
+    qe_simulation = get_qe_simulation(session, id, user_id)
+    structure = get_crystal_structure(
+        session, qe_simulation.crystal_structure_id, user_id
+    )
+
+    return build_qe_inputfile(qe_simulation, structure)
+
+
+def get_qe_output(session, id, user_id):
+    qe_simulation = get_qe_simulation(session, id, user_id)
+
+    wd = qe_simulation.simulation.working_directory
+
+    result_file = wd + "/result.pwo"
+
+    with open(result_file) as fh:
+        file = fh.read()
+
+        return file
+
+
+def get_qe_xas(session, id, user_id):
+    qe_simulation = get_qe_simulation(session, id, user_id)
+
+    wd = qe_simulation.simulation.working_directory
+
+    result_file = wd + "/xanes.dat"
+
+    if not os.path.isfile(result_file):
+        raise HTTPException(status_code=404, detail="Item not found")
+
+    out = np.loadtxt(result_file)
+
+    output = {"energy": out[:, 0].tolist(), "xas": out[:, 1].tolist()}
+
+    return output