Merge pull request mala-project#653 from franzpoeschel/align-ldos-openpmd

Align LDOS: Read from and write to openPMD

Author: RandomDefaultUser

examples/advanced/ex09_align_ldos.py

@@ -1,5 +1,7 @@
 import os
 
+import numpy
+
 import mala
 
 from mala.datahandling.data_repo import data_path
@@ -30,3 +32,31 @@
 ldos_aligner.align_ldos_to_ref(
     left_truncate=True, right_truncate_value=11, number_of_electrons=4
 )
+
+# The same thing works also with openPMD-based data
+
+try:
+    import openpmd_api
+    use_openpmd = True
+except ImportError:
+    use_openpmd = False
+
+if use_openpmd:
+    # initialize and add snapshots to workflow
+    ldos_aligner = mala.LDOSAligner(parameters)
+    ldos_aligner.clear_data()
+    ldos_aligner.add_snapshot("Be_snapshot0.out.h5",
+                              data_path, snapshot_type='openpmd')
+    ldos_aligner.add_snapshot("Be_snapshot1.out.h5",
+                              data_path, snapshot_type='openpmd')
+    ldos_aligner.add_snapshot("Be_snapshot2.out.h5",
+                              data_path, snapshot_type='openpmd')
+
+    # align and cut the snapshots from the left and right-hand sides
+    ldos_aligner.align_ldos_to_ref(
+        left_truncate=True, right_truncate_value=11, number_of_electrons=4
+    )
+
+    # A test that checks for data equivalence between the Numpy-based
+    # and openPMD-based implementations can be found under
+    # test/align_ldos_test.py.

mala/datahandling/ldos_aligner.py

@@ -85,8 +85,6 @@ def add_snapshot(
             snapshot_type=snapshot_type,
         )
 
-        if snapshot_type != "numpy":
-            raise Exception("Snapshot type must be numpy for LDOS alignment")
 
     def align_ldos_to_ref(
         self,
@@ -136,6 +134,9 @@ def align_ldos_to_ref(
             comm = get_comm()
             rank = comm.rank
             size = comm.size
+            import sys
+            print(
+                "[mala.LDOSAligner.align_ldos_to_ref] Warning: MPI support is experimental. Use with caution.")
         else:
             comm = None
             rank = 0
@@ -146,13 +147,23 @@ def align_ldos_to_ref(
             snapshot_ref = self.parameters.snapshot_directories_list[
                 reference_index
             ]
-            ldos_ref = np.load(
-                os.path.join(
-                    snapshot_ref.output_npy_directory,
-                    snapshot_ref.output_npy_file,
-                ),
-                mmap_mode="r",
-            )
+            if snapshot_ref.snapshot_type == 'numpy':
+                ldos_ref = self.target_calculator.read_from_numpy_file(
+                    os.path.join(
+                        snapshot_ref.output_npy_directory,
+                        snapshot_ref.output_npy_file,
+                    )
+                )
+            elif snapshot_ref.snapshot_type == 'openpmd':
+                ldos_ref = self.target_calculator.read_from_openpmd_file(
+                    os.path.join(
+                        snapshot_ref.output_npy_directory,
+                        snapshot_ref.output_npy_file,
+                    )
+                )
+            else:
+                raise Exception(
+                    "Unknown snapshot type '{snapshot_ref.snapshot_type}'.")
 
             # get the mean
             n_target = ldos_ref.shape[-1]
@@ -205,13 +216,23 @@ def align_ldos_to_ref(
         for idx in local_snapshots:
             snapshot = self.parameters.snapshot_directories_list[idx]
             print(f"Aligning snapshot {idx+1} of {N_snapshots}")
-            ldos = np.load(
-                os.path.join(
-                    snapshot.output_npy_directory,
-                    snapshot.output_npy_file,
-                ),
-                mmap_mode="r",
-            )
+            if snapshot_ref.snapshot_type == 'numpy':
+                ldos = self.target_calculator.read_from_numpy_file(
+                    os.path.join(
+                        snapshot.output_npy_directory,
+                        snapshot.output_npy_file,
+                    )
+                )
+            elif snapshot_ref.snapshot_type == 'openpmd':
+                ldos = self.target_calculator.read_from_openpmd_file(
+                    os.path.join(
+                        snapshot.output_npy_directory,
+                        snapshot.output_npy_file,
+                    )
+                )
+            else:
+                raise Exception(
+                    "Unknown snapshot type '{snapshot_ref.snapshot_type}'.")
 
             # get the mean
             nx = ldos.shape[0]
@@ -279,30 +300,34 @@ def align_ldos_to_ref(
                 snapshot.output_npy_directory, save_path_ext
             )
             save_name = snapshot.output_npy_file
-
-            stripped_output_file_name = snapshot.output_npy_file.replace(
-                ".out", ""
-            )
-            ldos_shift_info_save_name = stripped_output_file_name.replace(
-                ".npy", ".ldos_shift.info.json"
-            )
-
+            target_name = os.path.join(save_path, save_name)
             os.makedirs(save_path, exist_ok=True)
 
-            if "*" in save_name:
-                save_name = save_name.replace("*", str(idx))
-                ldos_shift_info_save_name.replace("*", str(idx))
+            self.ldos_parameters.ldos_gridsize = ldos_shifted.shape[-1]
 
-            target_name = os.path.join(save_path, save_name)
+            if snapshot.snapshot_type == 'numpy':
+                stripped_output_file_name = snapshot.output_npy_file.replace(
+                    ".out", ""
+                )
+                ldos_shift_info_save_name = stripped_output_file_name.replace(
+                    ".npy", ".ldos_shift.info.json"
+                )
 
-            self.target_calculator.write_to_numpy_file(
-                target_name, ldos_shifted
-            )
+                if "*" in save_name:
+                    save_name = save_name.replace("*", str(idx))
+                    ldos_shift_info_save_name.replace("*", str(idx))
 
-            with open(
-                os.path.join(save_path, ldos_shift_info_save_name), "w"
-            ) as f:
-                json.dump(ldos_shift_info, f, indent=2)
+                self.target_calculator.write_to_numpy_file(
+                    target_name, ldos_shifted
+                )
+
+                with open(
+                    os.path.join(save_path, ldos_shift_info_save_name), "w"
+                ) as f:
+                    json.dump(ldos_shift_info, f, indent=2)
+            else:
+                self.target_calculator.write_to_openpmd_file(
+                    target_name, ldos_shifted, additional_attributes=ldos_shift_info)
 
         barrier()
 

test/splitting_test.py test/align_ldos_test.pytest/splitting_test.py renamed to test/align_ldos_test.py

@@ -83,3 +83,56 @@ def test_dos_splitting(self):
                 )
             )
         )
+
+    # This test is based on examples/advanced/ex09_align_ldos.py, but adds a
+    # test to check that the aligned data is equivalent between the numpy-based
+    # and the openPMD-based implementations.
+    def test_ldos_alignment(self):
+        parameters = mala.Parameters()
+        parameters.targets.ldos_gridoffset_ev = -5
+        parameters.targets.ldos_gridsize = 11
+        parameters.targets.ldos_gridspacing_ev = 2.5
+
+        # initialize and add snapshots to workflow
+        ldos_aligner = mala.LDOSAligner(parameters)
+        ldos_aligner.clear_data()
+        ldos_aligner.add_snapshot("Be_snapshot0.out.npy", data_path)
+        ldos_aligner.add_snapshot("Be_snapshot1.out.npy", data_path)
+        ldos_aligner.add_snapshot("Be_snapshot2.out.npy", data_path)
+
+        # align and cut the snapshots from the left and right-hand sides
+        ldos_aligner.align_ldos_to_ref(
+            left_truncate=True, right_truncate_value=11, number_of_electrons=4
+        )
+
+        try:
+            import openpmd_api
+            use_openpmd = True
+        except ImportError:
+            use_openpmd = False
+
+        if use_openpmd:
+            # initialize and add snapshots to workflow
+            ldos_aligner = mala.LDOSAligner(parameters)
+            ldos_aligner.clear_data()
+            ldos_aligner.add_snapshot("Be_snapshot0.out.h5",
+                                      data_path, snapshot_type='openpmd')
+            ldos_aligner.add_snapshot("Be_snapshot1.out.h5",
+                                      data_path, snapshot_type='openpmd')
+            ldos_aligner.add_snapshot("Be_snapshot2.out.h5",
+                                      data_path, snapshot_type='openpmd')
+
+            # align and cut the snapshots from the left and right-hand sides
+            ldos_aligner.align_ldos_to_ref(
+                left_truncate=True, right_truncate_value=11, number_of_electrons=4
+            )
+
+            parameters = mala.Parameters()
+            data_handler = mala.DataHandler(parameters)
+            for i in range(1, 4):
+                data_openpmd = data_handler.target_calculator.read_from_openpmd_file(
+                    f"{data_path}/aligned/Be_snapshot0.out.h5")
+                data_numpy = data_handler.target_calculator.read_from_numpy_file(
+                    f"{data_path}/aligned/Be_snapshot0.out.npy")
+                if not np.allclose(data_numpy, data_openpmd):
+                    raise Exception("Inconsistency in snapshot", i)