Fix GoParallelTuple tuple truncation (\#120) and Tensor2 uninitialized phonon checks (\#47)

cellconstructor/ForceTensor.py

@@ -95,6 +95,15 @@ def SetupFromPhonons(self, phonons):
                 The dynamical matrix from which you want to setup the tensor
         """
 
+        # Check if the phonon has been initialized
+        if phonons.structure is None or len(phonons.q_tot) == 0:
+            ERR = """
+Error, the phonon object is not initialized.
+       Please load a dynamical matrix or provide a structure before
+       calling SetupFromPhonons.
+"""
+            raise ValueError(ERR)
+
         # Check if the supercell is correct
         ERR = """
 Error, the supercell of the phonon object is {}.

cellconstructor/Settings.py

@@ -247,6 +247,30 @@ def GoParallel(function, list_of_inputs, reduce_op = None, timer=None):
         computing_list = list_of_inputs[start:end]
         t2 = time.time()
 
+        my_len = len(computing_list)
+        if __PARALLEL_TYPE__ == "mpi4py":
+            comm = mpi4py.MPI.COMM_WORLD
+            min_len = comm.allreduce(my_len, op=mpi4py.MPI.MIN)
+            if min_len == 0:
+                if rank == 0:
+                    raise IndexError(
+                        "Some MPI ranks have no work items "
+                        "(%d inputs across %d ranks). "
+                        "Reduce the number of MPI ranks or increase "
+                        "the number of inputs so every rank has at "
+                        "least one element."
+                        % (len(list_of_inputs), n_proc))
+                comm.Abort(1)
+            elif my_len == 0:
+                raise IndexError("GoParallel rank %d: no work items" % rank)
+        else:
+            if my_len == 0:
+                raise IndexError(
+                    "Rank %d has no work items (%d inputs across %d ranks). "
+                    "Reduce the number of MPI ranks or increase the number "
+                    "of inputs so every rank has at least one element."
+                    % (rank, len(list_of_inputs), n_proc))
+
         if timer is not None:
             timer.add_timer("broadcast", t2 - t1)
 
@@ -393,6 +417,33 @@ def GoParallelTuple(function, list_of_inputs, reduce_op = None):
         for i in range(rank, len(list_of_inputs), n_proc):
             computing_list.append(list_of_inputs[i])
 
+        # Synchronize empty-work check across all MPI ranks so no
+        # rank raises alone (which would hang the others at the
+        # collective call that follows).
+        my_len = len(computing_list)
+        if __PARALLEL_TYPE__ == "mpi4py":
+            comm = mpi4py.MPI.COMM_WORLD
+            min_len = comm.allreduce(my_len, op=mpi4py.MPI.MIN)
+            if min_len == 0:
+                if rank == 0:
+                    raise IndexError(
+                        "Some MPI ranks have no work items "
+                        "(%d inputs across %d ranks). "
+                        "Reduce the number of MPI ranks or increase "
+                        "the number of inputs so every rank has at "
+                        "least one element."
+                        % (len(list_of_inputs), n_proc))
+                comm.Abort(1)
+            elif my_len == 0:
+                raise IndexError("GoParallelTuple rank %d: no work items" % rank)
+        else:
+            if my_len == 0:
+                raise IndexError(
+                    "Rank %d has no work items (%d inputs across %d ranks). "
+                    "Reduce the number of MPI ranks or increase the number "
+                    "of inputs so every rank has at least one element."
+                    % (rank, len(list_of_inputs), n_proc))
+
         #print("Rank {} is computing {} elements".format(rank, len(computing_list)))
 
         # Work! TODO: THIS IS VERY MEMORY HEAVY
@@ -428,18 +479,26 @@ def GoParallelTuple(function, list_of_inputs, reduce_op = None):
                 raise NotImplementedError("Error, not implemented {}".format(__PARALLEL_TYPE__))
 
 
-            result = results[0]
+            # `results` is a list of allgathered lists, one per tuple
+            # element: [[el0_from_rank0, el0_from_rank1, ...],
+            #           [el1_from_rank0, el1_from_rank1, ...], ...].
+            # Reduce each element across ranks and build the final
+            # result list.
+            final_result = []
             for j in range(len(results)):
-                # Perform the last reduction
+                reduced = results[j][0]
                 if reduce_op == "+":
-                    for i in range(1,len(results[j])):
-                        result[j]+= results[j][i]
+                    for i in range(1, len(results[j])):
+                        reduced += results[j][i]
                 elif reduce_op == "*":
-                    for i in range(1,len(results[j])):
-                        result[j]*= results[j][i]
+                    for i in range(1, len(results[j])):
+                        reduced *= results[j][i]
+                final_result.append(reduced)
 
-            return result 
+            return final_result
         else:
+            if __PARALLEL_TYPE__ == "serial":
+                return results
             raise NotImplementedError("Error, for now parallelization with MPI implemented only with reduction")
     else:
         raise NotImplementedError("Something went wrong: {}".format(__PARALLEL_TYPE__))

tests/TestParallel/test_goparalleltuple.py

@@ -0,0 +1,80 @@
+import sys
+import os
+
+import numpy as np
+import pytest
+
+import cellconstructor as CC
+import cellconstructor.Settings
+
+
+# Detect mpi4py availability (used by skipif decorators below).
+_has_mpi4py = False
+try:
+    import mpi4py  # noqa: F401
+    _has_mpi4py = True
+except ImportError:
+    pass
+
+
+def test_goparalleltuple_serial_returns_all_elements():
+    def func(x):
+        return [np.array([1.0, 2.0]), np.array([[3.0, 4.0], [5.0, 6.0]])]
+
+    result = CC.Settings.GoParallelTuple(func, [[1, 2]], "+")
+    assert len(result) == 2, \
+        "GoParallelTuple serial returned %d elements, expected 2" % len(result)
+    np.testing.assert_array_equal(result[0], np.array([1.0, 2.0]))
+    np.testing.assert_array_equal(result[1], np.array([[3.0, 4.0], [5.0, 6.0]]))
+
+
+def test_goparalleltuple_serial_reduction_sum():
+    def func(x):
+        return [np.array([x[0]]), np.array([x[1]])]
+
+    result = CC.Settings.GoParallelTuple(func, [[1, 10], [2, 20], [3, 30]], "+")
+    assert len(result) == 2, \
+        "GoParallelTuple serial reduction returned %d elements, expected 2" % len(result)
+    np.testing.assert_array_equal(result[0], np.array([6]))   # 1+2+3
+    np.testing.assert_array_equal(result[1], np.array([60]))  # 10+20+30
+
+
+@pytest.mark.skipif(_has_mpi4py,
+                    reason="no-reduction not implemented for mpi4py path")
+def test_goparalleltuple_serial_no_reduction():
+    def func(x):
+        return [np.array([x]), np.array([x * 10])]
+
+    result = CC.Settings.GoParallelTuple(func, [[1], [2]], reduce_op=None)
+    assert len(result) == 2, \
+        "GoParallelTuple serial no reduction returned %d elements, expected 2" % len(result)
+
+
+
+@pytest.mark.skipif(not _has_mpi4py, reason="mpi4py not installed")
+def test_goparalleltuple_mpi4py_returns_all_elements():
+    """
+    Regression test for issue #120.
+
+    GoParallelTuple with mpi4py must return all elements of the function's
+    return tuple, not just the first. Even with a single MPI process the
+    allgather path is entered and the original bug dropped elements 2+.
+    """
+    def func(x):
+        return [np.array([1.0, 2.0]), np.array([[3.0, 4.0], [5.0, 6.0]])]
+
+    result = CC.Settings.GoParallelTuple(func, [[1, 2]], "+")
+    msg = "GoParallelTuple mpi4py returned %d elements, expected 2 (issue #120)" \
+          % len(result)
+    assert len(result) == 2, msg
+    np.testing.assert_array_equal(result[0], np.array([1.0, 2.0]))
+    np.testing.assert_array_equal(result[1], np.array([[3.0, 4.0], [5.0, 6.0]]))
+
+
+if __name__ == "__main__":
+    test_goparalleltuple_serial_returns_all_elements()
+    test_goparalleltuple_serial_reduction_sum()
+    test_goparalleltuple_serial_no_reduction()
+    if _has_mpi4py:
+        test_goparalleltuple_mpi4py_returns_all_elements()
+    print("ok")

tests/TestTensor2/test_setup_from_phonons.py

@@ -0,0 +1,57 @@
+import numpy as np
+import pytest
+
+import cellconstructor as CC
+import cellconstructor.ForceTensor
+import cellconstructor.Phonons
+
+
+def _make_test_structure():
+    s = CC.Structure.Structure(2)
+    s.atoms = ["Si", "Si"]
+    s.coords[0, :] = [0.0, 0.0, 0.0]
+    s.coords[1, :] = [1.3575, 1.3575, 1.3575]
+    s.unit_cell = np.eye(3) * 5.43
+    s.has_unit_cell = True
+    s.masses = {"Si": 28.0855 / CC.Units.MASS_RY_TO_UMA}
+    return s
+
+
+def test_setup_from_uninitialized_phonons():
+    """
+    Passing an empty Phonons() object (no structure loaded) to
+    SetupFromPhonons must raise a clear error, not crash cryptically.
+    """
+    uc = _make_test_structure()
+    sc = uc.generate_supercell((2, 2, 2))
+    tensor = CC.ForceTensor.Tensor2(uc, sc, (2, 2, 2))
+
+    empty_phon = CC.Phonons.Phonons()
+
+    with pytest.raises(ValueError,
+                       match="not initialized|no structure"):
+        tensor.SetupFromPhonons(empty_phon)
+
+
+def test_setup_wrong_supercell():
+    """
+    Passing a phonon with a different supercell to SetupFromPhonons
+    must raise an error about supercell mismatch.
+    """
+    uc = _make_test_structure()
+    sc_222 = uc.generate_supercell((2, 2, 2))
+    tensor = CC.ForceTensor.Tensor2(uc, sc_222, (2, 2, 2))
+
+    # Create a phonon on a 1x1x1 cell (smaller supercell)
+    phon = CC.Phonons.Phonons(uc, nqirr=1)
+    assert list(phon.GetSupercell()) == [1, 1, 1]
+
+    with pytest.raises((ValueError, AssertionError, IOError),
+                       match="supercell"):
+        tensor.SetupFromPhonons(phon)
+
+
+if __name__ == "__main__":
+    test_setup_from_uninitialized_phonons()
+    test_setup_wrong_supercell()
+    print("ok")