SOF-7776 Update: add WFN property

express/parsers/apps/espresso/formats/txt.py

@@ -701,6 +701,24 @@ def potential_profile(self, text):
         data = self._general_output_parser(text, **settings.REGEX["potential_profile"])
         return [[e[i] for e in data] for i in range(4)]
 
+    def wavefunction_amplitude(self, text):
+        """
+        Extracts wavefunction amplitude along z coordinate.
+
+        Example input:
+            #z (A)  Amplitude
+             -4.89        0.0012
+             -4.78        0.0034
+             -4.67        0.0067
+             -4.56        0.0123
+             -4.44        0.0234
+
+        Returns:
+            list[list[float]]
+        """
+        data = self._general_output_parser(text, **settings.REGEX["wavefunction_amplitude"])
+        return [[e[i] for e in data] for i in range(2)]
+
     def charge_density_profile(self, text):
         """
         Extracts total charge density along z.

express/parsers/apps/espresso/parser.py

@@ -274,9 +274,20 @@ def reaction_energies(self):
     def _get_esm_file(self):
         return find_file(".esm1", self.work_dir)
 
+    def _get_wavefunction_file(self):
+        return find_file(settings.WAVEFUNCTION_FILE, self.work_dir)
+
     def potential_profile(self):
         return self.txt_parser.potential_profile(self._get_file_content(self._get_esm_file()))
 
+    def wavefunction_amplitude(self):
+        # The data is [[coordinate], [value]]: for coordinates comes in alat units, we convert x data to angstroms
+        data = self.txt_parser.wavefunction_amplitude(self._get_file_content(self._get_wavefunction_file()))
+        lattice = self.final_lattice_vectors()
+        alat = lattice["vectors"]["alat"]
+        data[0] = [x * alat for x in data[0]]
+        return data
+
     def charge_density_profile(self):
         return self.txt_parser.charge_density_profile(self._get_file_content(self._get_esm_file()))
 

express/parsers/apps/espresso/settings.py

@@ -8,6 +8,7 @@
 AVERAGE_FILE = "avg.dat"
 HP_FILE = "__prefix__.Hubbard_parameters.dat"  # contains Hubbard U and V parameters
 HP_NN_FILE = "HUBBARD.dat"  # contains Hubbard V parameters for 6 nearest neighbors
+WAVEFUNCTION_FILE = "wf_r.dat"
 
 COMMON_REGEX = r"{0}\s+[=:<>]\s*([-+]?\d*\.?\d*([Ee][+-]?\d+)?)"
 DOUBLE_REGEX = GENERAL_REGEX.double_number
@@ -121,6 +122,12 @@
         "output_type": "float",
         "match_groups": [1, 2, 3, 4],
     },
+    "wavefunction_amplitude": {
+        "regex": r"^\s+({0})\s+({0})".format(DOUBLE_REGEX),
+        "occurrences": 0,
+        "output_type": "float",
+        "match_groups": [1, 2],
+    },
     "charge_density_profile": {
         "regex": r"^\s+({0})\s+({0})\s+{0}\s+{0}\s+{0}".format(DOUBLE_REGEX),
         "occurrences": 0,

express/properties/non_scalar/two_dimensional_plot/wavefunction_amplitude.py

@@ -0,0 +1,14 @@
+from express.properties.non_scalar.two_dimensional_plot import TwoDimensionalPlotProperty
+
+
+class WavefunctionAmplitude(TwoDimensionalPlotProperty):
+    """
+    Wavefunction amplitude.
+    """
+
+    def __init__(self, name, parser, *args, **kwargs):
+        super(WavefunctionAmplitude, self).__init__(name, parser, *args, **kwargs)
+        wavefunction_amplitude = self.parser.wavefunction_amplitude()
+        self.xDataArray = wavefunction_amplitude[0]
+        self.yDataSeries = [wavefunction_amplitude[1]]
+

express/settings.py

@@ -41,6 +41,9 @@
     "potential_profile": {
         "reference": "express.properties.non_scalar.two_dimensional_plot.potential_profile.PotentialProfile"
     },
+    "wavefunction_amplitude": {
+        "reference": "express.properties.non_scalar.two_dimensional_plot.wavefunction_amplitude.WavefunctionAmplitude"
+    },
     "charge_density_profile": {
         "reference": "express.properties.non_scalar.two_dimensional_plot.charge_density_profile.ChargeDensityProfile"
     },

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "munch==2.5.0",
     "pymatgen>=2023.8.10",
     "ase>=3.17.0",
-    "mat3ra-esse>=2024.1.25.post7",
+    "mat3ra-esse>=2026.1.23",
     "jarvis-tools>=2023.12.12",
     # To avoid module 'numpy.linalg._umath_linalg' has no attribute '_ilp64' in Colab
     "numpy>=1.24.4,<2",

tests/unit/properties/non_scalar/two_dimensional_plot/test_wavefunction_amplitude.py

@@ -0,0 +1,36 @@
+from tests.unit import UnitTestBase
+from express.properties.non_scalar.two_dimensional_plot.wavefunction_amplitude import WavefunctionAmplitude
+
+RAW_DATA_ALAT = [
+    [0.0, 0.0050251256, 0.0100502513, 0.0150753769, 0.0201005025],
+    [0.0000322091, 0.0000072134, -0.0000218274, -0.0000540398, -0.0000883573]
+]
+
+ALAT_ANGSTROM = 10.0
+
+CONVERTED_DATA_ANGSTROMS = [
+    [x * ALAT_ANGSTROM for x in RAW_DATA_ALAT[0]],
+    RAW_DATA_ALAT[1]
+]
+
+EXPECTED = {
+    "name": "wavefunction_amplitude",
+    "xDataArray": CONVERTED_DATA_ANGSTROMS[0],
+    "yDataSeries": [CONVERTED_DATA_ANGSTROMS[1]],
+    "xAxis": {"label": "coordinate", "units": "angstrom"},
+    "yAxis": {"label": "amplitude"},
+}
+
+
+class WavefunctionAmplitudeTest(UnitTestBase):
+    def setUp(self):
+        super().setUp()
+
+    def tearDown(self):
+        super().tearDown()
+
+    def test_wavefunction_amplitude(self):
+        parser = self.get_mocked_parser("wavefunction_amplitude", CONVERTED_DATA_ANGSTROMS)
+        property_ = WavefunctionAmplitude("wavefunction_amplitude", parser)
+        self.assertDeepAlmostEqual(property_.serialize_and_validate(), EXPECTED)
+