[tet.py] Fix bug when Ns < mpi_size in gradient methods.

USTC-TNS · Nov 19, 2023 · 7b60f9b · 7b60f9b
1 parent 6490602
commit 7b60f9b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/CHANGELOG.org b/CHANGELOG.org
@@ -20,6 +20,7 @@
 *** Removed
 + *tetragono*: Remove ansatz product state support, which may be added again in the future, but it require many updates.
 *** Fixed
++ *tetragono*: Fix bug when total sampling number less than process number in sampling gradient.
 
 ** [[https://github.com/USTC-TNS/TAT/compare/v0.3.7...v0.3.9][v0.3.9]] - 2023-09-07
 

diff --git a/tetragono/tetragono/sampling_lattice/gradient.py b/tetragono/tetragono/sampling_lattice/gradient.py
@@ -236,9 +236,10 @@ def gradient_descent(
                             configuration_pool.append((possibility, configuration))
                         show(f"sampling {sampling_step}/{sampling_total_step}, energy={observer.energy}")
                 # Save configuration
-                new_configurations = configuration.export_configuration()
-                sampling_configurations.resize(new_configurations.shape, refcheck=False)
-                np.copyto(sampling_configurations, new_configurations)
+                if mpi_rank < sampling_total_step:
+                    new_configurations = configuration.export_configuration()
+                    sampling_configurations.resize(new_configurations.shape, refcheck=False)
+                    np.copyto(sampling_configurations, new_configurations)
             showln(f"sampling done, total_step={sampling_total_step}, energy={observer.energy}")
             if sampling_method == "direct":
                 showln(f"direct sampling stability is {observer.stability}")