Page MenuHomec4science

scatter.py
No OneTemporary

File Metadata

Created
Thu, Oct 10, 16:23

scatter.py

# Copyright (C) 2018-2020 by the RROMPy authors
#
# This file is part of RROMPy.
#
# RROMPy is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# RROMPy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with RROMPy. If not, see <http://www.gnu.org/licenses/>.
#
from rrompy.utilities.base.types import Tuple, ListAny
from .base import poolSize, poolRank, forcedSerial, forceSerialExecution
__all__ = ['indicesScatter', 'listScatter']
def partitionIndices(a:int, n:int, j:int,
return_sizes : bool = False) -> ListAny:
aeach, nextra = a // n, a % n
lb = aeach * j + min(j, nextra)
idx = list(range(lb, lb + aeach + (j < nextra)))
if return_sizes: return idx, [aeach + 1] * nextra + [aeach] * (n - nextra)
return idx
def indicesScatter(a:int, force : bool = True,
return_sizes : bool = False) -> ListAny:
"""Split indices among parallel cores."""
size, rank = (1, 0) if forcedSerial() else (poolSize(), poolRank())
if force: forceSerialExecution(True)
return partitionIndices(a, size, rank, return_sizes)
def listScatter(obj:ListAny, force : bool = True,
return_sizes : bool = False) -> Tuple[ListAny, ListAny]:
"""Split list among parallel cores."""
idx = indicesScatter(len(obj), force, return_sizes)
if not return_sizes: idx = (idx,)
try:
return (obj[idx[0]],) + idx
except:
return ([obj[i] for i in idx[0]],) + idx

Event Timeline