Feature/SOF-7776 WFN

assets/applications/executables/espresso/pp.x.yml

@@ -19,3 +19,12 @@ flavors:
       - standard_output
     applicationName: espresso
     executableName: pp.x
+
+  pp_wfn:
+    input:
+      - name: pp_wfn.in
+    results: []
+    monitors:
+      - standard_output
+    applicationName: espresso
+    executableName: pp.x

assets/applications/executables/python/python.yml

@@ -4,6 +4,7 @@ monitors:
 results:
   - file_content
   - 'workflow:pyml_predict'
+  - potential_profile
 flavors:
   hello_world:
     isDefault: true
@@ -34,6 +35,30 @@ flavors:
     applicationName: python
     executableName: python
 
+  extract_bands_fermi:
+    input:
+      - name: extract_bands_fermi.py
+      - name: requirements.txt
+        templateName: requirements_bands_fermi.txt
+    monitors:
+      - standard_output
+    applicationName: python
+    executableName: python
+
+  plot_wavefunction:
+    input:
+      - name: script.py
+        templateName: plot_wavefunction.py
+      - name: requirements.txt
+        templateName: requirements_plot_wavefunction.txt
+    results:
+      - potential_profile
+      - file_content
+    monitors:
+      - standard_output
+    applicationName: python
+    executableName: python
+
   'generic:post_processing:plot:matplotlib':
     input:
       - name: plot.py

assets/applications/input_files_templates/espresso/pp_wfn.j2.in

@@ -0,0 +1,16 @@
+&INPUTPP
+    outdir = {% raw %}'{{ JOB_WORK_DIR }}/outdir'{% endraw %}
+    prefix = '__prefix__'
+    plot_num = 7
+    kpoint = 1
+    kband = {% raw %}{{ KBAND_VALUE }}{% endraw %}
+    lsign = .true.
+/
+&PLOT
+    iflag = 1
+    fileout = 'wf_r.dat'
+    output_format = 0
+    x0(1) = 0.0, x0(2) = 0.0, x0(3) = 0.0
+    e1(1) = 0.0, e1(2) = 0.0, e1(3) = 1.0
+    nx = 200
+/

assets/applications/input_files_templates/python/espresso/extract_bands_fermi.pyi

@@ -0,0 +1,57 @@
+# ---------------------------------------------------------------- #
+# Extract band indices near Fermi energy from band_structure       #
+# This script expects fermi_energy and band_structure results      #
+# ---------------------------------------------------------------- #
+import json
+
+# The results below come as Munch objects, so we need to import Munch to handle them
+from munch import Munch
+
+# Data From Context
+# -----------------
+# fermi_energy: float (in eV) - from pw_scf result
+# band_structure: Munch object with band energies - from pw_bands result
+
+{% raw %}fermi_energy_value = {{ fermi_energy }}{% endraw %}
+{% raw %}band_structure_data = {{ band_structure }}{% endraw %}
+
+# Extract band energies at Gamma point (first k-point)
+# band_structure format from QE parser:
+# {
+#   "name": "band_structure",
+#   "xDataArray": [[kx, ky, kz], ...],  # k-points
+#   "yDataSeries": [[e1_k1, e1_k2, ...], [e2_k1, e2_k2, ...], ...]  # energies per band
+# }
+# yDataSeries[band_index][kpoint_index] = energy
+
+# Get energies at first k-point (Gamma, index 0) for all bands
+y_data = band_structure_data.get('yDataSeries', [])
+band_energies = [band_data[0] for band_data in y_data] if y_data else []
+
+# Find bands near Fermi energy (1-based indices as QE expects)
+valence_bands = [(i + 1, e) for i, e in enumerate(band_energies) if e <= fermi_energy_value]
+conduction_bands = [(i + 1, e) for i, e in enumerate(band_energies) if e > fermi_energy_value]
+
+if valence_bands:
+    valence_index, valence_energy = max(valence_bands, key=lambda x: x[1])
+else:
+    valence_index = 1
+    valence_energy = band_energies[0] if band_energies else 0.0
+
+if conduction_bands:
+    conduction_index, conduction_energy = min(conduction_bands, key=lambda x: x[1])
+else:
+    conduction_index = len(band_energies)
+    conduction_energy = band_energies[-1] if band_energies else 0.0
+
+result = {
+    "band_below_fermi": valence_index,
+    "band_above_fermi": conduction_index,
+    "fermi_energy": fermi_energy_value,
+    "valence_energy": valence_energy,
+    "conduction_energy": conduction_energy,
+    "total_bands": len(band_energies)
+}
+
+# Print to STDOUT for subsequent assignment unit
+print(json.dumps(result, indent=4))

assets/applications/input_files_templates/python/espresso/plot_wavefunction.pyi

@@ -0,0 +1,49 @@
+# ---------------------------------------------------------------- #
+#  Generate wavefunction plot from pp.x output                     #
+#  And saves static                                                #
+# ---------------------------------------------------------------- #
+
+import json
+
+import matplotlib
+import numpy as np
+
+matplotlib.use('Agg')  # Non-interactive backend
+import matplotlib.pyplot as plt
+
+# Load wavefunction data from pp.x output
+data = np.loadtxt('wf_r.dat')
+z = data[:, 0]
+psi_r = data[:, 1]
+
+# Calculate wavefunction amplitude
+psi_amplitude = np.abs(psi_r)
+
+# Create static PNG plot
+fig, ax = plt.subplots(figsize=(10, 6))
+ax.plot(z, psi_amplitude, 'b-', linewidth=2)
+ax.set_xlabel('Position z (Å)', fontsize=12)
+ax.set_ylabel('Wavefunction amplitude |ψ| (a.u.)', fontsize=12)
+ax.set_title('Wavefunction along z-axis', fontsize=14)
+ax.grid(True, alpha=0.3)
+plt.tight_layout()
+plt.savefig('wf_r.png', dpi=150, bbox_inches='tight')
+plt.close()
+
+# Create potential_profile JSON for platform rendering
+wavefunction_data = {
+    "name": "wavefunction_amplitude",
+    "xAxis": {
+        "label": "Position",
+        "units": "angstrom"
+    },
+    "xDataArray": z.tolist(),
+    "yAxis": {
+        "label": "Wavefunction Amplitude",
+        "units": "a.u."
+    },
+    "yDataSeries": [psi_amplitude.tolist()]
+}
+
+# Print JSON to STDOUT (will be captured as potential_profile result)
+print(json.dumps(wavefunction_data, indent=2))

assets/applications/input_files_templates/python/espresso/requirements_bands_fermi.j2.txt

@@ -0,0 +1,7 @@
+# ------------------------------------------------------------------ #
+#                                                                    #
+#  Python package requirements for extract_bands_fermi unit          #
+#                                                                    #
+# ------------------------------------------------------------------ #
+
+munch==2.5.0

assets/applications/input_files_templates/python/espresso/requirements_plot_wavefunction.j2.txt

@@ -0,0 +1,8 @@
+# ------------------------------------------------------------------ #
+#                                                                    #
+#  Python package requirements for plot_wavefunction unit            #
+#                                                                    #
+# ------------------------------------------------------------------ #
+
+numpy<2
+matplotlib

assets/applications/templates/espresso/pp.x.yml

@@ -9,3 +9,9 @@
   contextProviders: []
   applicationName: espresso
   executableName: pp.x
+
+- content: !readFile 'input_files_templates/espresso/pp_wfn.j2.in'
+  name: pp_wfn.in
+  contextProviders: []
+  applicationName: espresso
+  executableName: pp.x

assets/applications/templates/python/python.yml

@@ -16,6 +16,12 @@
   applicationName: python
   executableName: python
 
+- content: !readFile 'input_files_templates/python/espresso/requirements_bands_fermi.j2.txt'
+  name: requirements_bands_fermi.txt
+  contextProviders: []
+  applicationName: python
+  executableName: python
+
 - content: !readFile 'input_files_templates/python/generic/post_processing:plot:matplotlib.pyi'
   name: matplotlib_basic.py
   contextProviders: []
@@ -28,6 +34,24 @@
   applicationName: python
   executableName: python
 
+- content: !readFile 'input_files_templates/python/espresso/extract_bands_fermi.pyi'
+  name: extract_bands_fermi.py
+  contextProviders: []
+  applicationName: python
+  executableName: python
+
+- content: !readFile 'input_files_templates/python/espresso/plot_wavefunction.pyi'
+  name: plot_wavefunction.py
+  contextProviders: []
+  applicationName: python
+  executableName: python
+
+- content: !readFile 'input_files_templates/python/espresso/requirements_plot_wavefunction.j2.txt'
+  name: requirements_plot_wavefunction.txt
+  contextProviders: []
+  applicationName: python
+  executableName: python
+
 - content: !readFile 'input_files_templates/python/espresso_extract_kpoints.pyi'
   name: espresso_extract_kpoints.py
   contextProviders: []

assets/workflows/subworkflows/categories.yml

@@ -30,6 +30,7 @@ categories:
     - remove-all-results
     - set-io-unit-filenames
     - variable-cell_relaxation
+    - wfn_plot
   application:
     - espresso
     - nwchem
@@ -209,6 +210,10 @@ entities:
     categories:
       - espresso
       - python
+  - filename: espresso/extract_bands_fermi.json
+    categories:
+      - espresso
+      - python
   - filename: espresso/fixed_cell_relaxation.json
     categories:
       - atomic_forces
@@ -303,10 +308,19 @@ entities:
       - espresso
       - phonon_dispersions
       - phonon_dos
+  - filename: espresso/plot_wavefunction.json
+    categories:
+      - espresso
+      - file_content
+      - potential_profile
+      - python
   - filename: espresso/post_processor.json
     categories:
       - espresso
       - shell
+  - filename: espresso/pp_wfn.json
+    categories:
+      - espresso
   - filename: espresso/pre_processor.json
     categories:
       - espresso
@@ -335,6 +349,18 @@ entities:
       - total_energy
       - total_energy_contributions
       - total_force
+  - filename: espresso/total_energy_with_bands.json
+    categories:
+      - atomic_forces
+      - band_structure
+      - espresso
+      - fermi_energy
+      - pressure
+      - stress_tensor
+      - total_energy
+      - total_energy_contributions
+      - total_force
+      - wfn_plot
   - filename: espresso/total_energy.json
     categories:
       - atomic_forces
@@ -378,6 +404,9 @@ entities:
       - nwchem
       - total_energy
       - total_energy_contributions
+  - filename: python/extract_bands_fermi.json
+    categories:
+      - python
   - filename: python/ml/classification_tail.json
     categories:
       - creates-predictions-csv-during-predict-phase
@@ -404,6 +433,11 @@ entities:
       - pyml:workflow-type-setter
       - python
       - set-io-unit-filenames
+  - filename: python/plot_wavefunction.json
+    categories:
+      - file_content
+      - potential_profile
+      - python
   - filename: python/python_script.json
     categories:
       - python

assets/workflows/subworkflows/espresso/extract_bands_fermi.yml

@@ -0,0 +1,39 @@
+# Note: This subworkflow extracts band indices near Fermi energy
+# Expects fermi_energy and band_structure to be passed via unitConfigs
+# Outputs: KBAND_VALUE_BELOW_EF, KBAND_VALUE_ABOVE_EF, and KBAND_VALUE (selected)
+application:
+  name: python
+  version: 3.10.13
+method:
+  name: UnknownMethod
+model:
+  name: UnknownModel
+name: Extract Bands Near Fermi
+units:
+  - config:
+      execName: python
+      flavorName: extract_bands_fermi
+      name: extract_bands_fermi
+      flowchartId: extract-band-fermi
+    type: executionBuilder
+  - config:
+      name: Store Band Below EF
+      operand: KBAND_VALUE_BELOW_EF
+      value: "json.loads(STDOUT)['band_below_fermi']"
+      input:
+        - name: STDOUT
+          scope: extract-band-fermi
+    type: assignment
+  - config:
+      name: Store Band Above EF
+      operand: KBAND_VALUE_ABOVE_EF
+      value: "json.loads(STDOUT)['band_above_fermi']"
+      input:
+        - name: STDOUT
+          scope: extract-band-fermi
+    type: assignment
+  - config:
+      name: Select Band
+      operand: KBAND_VALUE
+      value: "KBAND_VALUE_BELOW_EF"
+    type: assignment

assets/workflows/subworkflows/espresso/plot_wavefunction.yml

@@ -0,0 +1,20 @@
+application:
+  name: python
+  version: 3.10.13
+method:
+  name: UnknownMethod
+model:
+  name: UnknownModel
+name: Plot Wavefunction
+properties:
+  - file_content
+units:
+  - config:
+      execName: python
+      flavorName: plot_wavefunction
+      name: plot WFN
+      results:
+        - name: file_content
+          basename: wf_r.png
+          filetype: image
+    type: executionBuilder

assets/workflows/subworkflows/espresso/pp_wfn.yml

@@ -0,0 +1,15 @@
+application:
+  name: espresso
+  version: "6.3"
+method:
+  name: PseudopotentialMethod
+model:
+  name: DFTModel
+name: PP Wavefunction
+units:
+  - config:
+      execName: pp.x
+      flavorName: pp_wfn
+      name: pp_wfn
+      flowchartId: pp-wfn
+    type: executionBuilder