fastatomstruct 1.5.0__cp314-cp314-win_amd64.whl

fastatomstruct/__init__.py

@@ -0,0 +1,103 @@
+from typing import Any, Callable, Dict, List
+from time import time
+
+import ase
+import numpy as np
+from .fastatomstruct import *
+
+rank = None
+
+
+def is_root() -> bool:
+    """Check if the current process is the root process."""
+    if rank is None:
+        return True
+    return rank == 0
+
+
+def chunks(indices: List[int], n: int):
+    """Yield successive chunks from list of atoms."""
+    for i in range(0, len(indices), n):
+        yield indices[i : i + n]
+
+
+def ipar(
+    func: Callable, atoms: List[ase.Atoms], *args: List[Any], **kwargs: Dict[Any, Any]
+) -> List[Any]:
+    """Image-based parallelization.
+
+    The calculation of many structural quantities is parallelized over atoms.
+    However, image-based parallelization (or a mixture of both) can be more efficient
+    in some cases. This function makes using this parallelization layer quite easy.
+    You can just use it as a simple wrapper around other functions implemented in
+    `fastatomstruct` (see example below). **This function needs a working installation
+    of the MPI4Py package!** You can set the `RAYON_NUM_THREADS` enviroment variable
+    to control the number of threads.
+
+    Arguments:
+        func (Callable): Some function from the `fastatomstruct` package
+        atoms (List[ase.Atoms]): List of ASE configurations
+
+    Returns:
+        List of results; type depends on the output of func
+
+    Examples
+    --------
+
+    The exemplary file "Sb-1.00-300-100.traj" `can be found here <https://zivgitlab.uni-muenster.de/ag-salinga/fastatomstruct/-/raw/master/tests/structures/Sb-1.00-300-100.traj>`__.
+    The code below needs to be run e.g. with `mpirun -n 2 test.py`. In this case,
+    two processes will be used.
+
+    .. plot::
+
+       >>> import matplotlib.pyplot as plt
+       >>> import fastatomstruct as fs
+       >>> from ase import io
+       >>> if fs.is_root():
+       >>>     atoms = io.read("Sb-1.00-300-100.traj", index="::200")
+       >>> q = fs.ipar(fs.q_tetrahedral, atoms, fs.CutoffMode.Fixed(3.2), 3)
+       >>> if fs.is_root():
+       >>>     plt.hist(q[0], bins=25, alpha=0.5)
+       >>>     plt.hist(q[1], bins=25, color="C3", alpha=0.5)
+    """
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    global rank
+    rank = comm.Get_rank()
+
+    if rank == 0:
+        print("\nImage parallelization report\n" + "----------------------------")
+
+    if rank == 0:
+        res_full, t0_f = [], time()
+
+    for i, c in enumerate(chunks(list(range(len(atoms))), comm.size)):
+        while len(c) < comm.size:
+            c.append(None)
+
+        if rank == 0:
+            per = (len(c) - np.count_nonzero([ci is None for ci in c])) / len(c)
+            t0 = time()
+        if c[rank] is not None:
+            a = atoms[c[rank]]
+            res = func(a, *args, **kwargs)
+        else:
+            res = None
+
+        res = comm.gather(res, root=0)
+        if rank == 0:
+            t = time() - t0
+            print(f"Chunk {i + 1}: {per * 100:03.2f}% | {t:.2f} s")
+            res_full.extend(res)
+
+    if rank == 0:
+        for i in range(len(res_full) - 1, -1, -1):
+            if res_full[i] is None:
+                del res_full[i]
+        print("----------------------------")
+        print(f"Total: {(time() - t0_f):.2f} s\n")
+    else:
+        res_full = None
+
+    return res_full

fastatomstruct/__init__.pyi

@@ -0,0 +1,60 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import ase
+from typing import Any, Callable, Dict, List
+from .fastatomstruct import *
+
+"""
+This type stub file was generated by pyright.
+"""
+rank = ...
+
+def is_root() -> bool:
+    """Check if the current process is the root process."""
+    ...
+
+def chunks(indices: List[int], n: int):
+    """Yield successive chunks from list of atoms."""
+    ...
+
+def ipar(
+    func: Callable, atoms: List[ase.Atoms], *args: List[Any], **kwargs: Dict[Any, Any]
+) -> List[Any]:
+    """Image-based parallelization.
+
+    The calculation of many structural quantities is parallelized over atoms.
+    However, image-based parallelization (or a mixture of both) can be more efficient
+    in some cases. This function makes using this parallelization layer quite easy.
+    You can just use it as a simple wrapper around other functions implemented in
+    `fastatomstruct` (see example below). **This function needs a working installation
+    of the MPI4Py package!** You can set the `RAYON_NUM_THREADS` enviroment variable
+    to control the number of threads.
+
+    Arguments:
+        func (Callable): Some function from the `fastatomstruct` package
+        atoms (List[ase.Atoms]): List of ASE configurations
+
+    Returns:
+        List of results; type depends on the output of func
+
+    Examples:
+
+    The exemplary file "Sb-1.00-300-100.traj" `can be found here <https://zivgitlab.uni-muenster.de/ag-salinga/fastatomstruct/-/raw/master/python/test/structures/Sb-1.00-300-100.traj>`__.
+    The code below needs to be run e.g. with `mpirun -n 2 test.py`. In this case,
+    two processes will be used.
+
+    .. plot::
+
+       >>> import matplotlib.pyplot as plt
+       >>> import fastatomstruct as fs
+       >>> from ase import io
+       >>> if fs.is_root():
+       >>>     atoms = io.read("Sb-1.00-300-100.traj", index="::200")
+       >>> q = fs.ipar(fs.q_tetrahedral, atoms, fs.CutoffMode.Fixed(3.2), 3)
+       >>> if fs.is_root():
+       >>>     plt.hist(q[0], bins=25, alpha=0.5)
+       >>>     plt.hist(q[1], bins=25, color="C3", alpha=0.5)
+    """
+    ...