distributedmatrix.py 16.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329
"""distributedmatrix.py -- classes for distributed matrices

This file contains the python classes to use with the wrapper.
"""
import numpy as np
from functools import wraps
from .wrapper import Elpa

class ProcessorLayout:
    """Create rectangular processor layout for use with distributed matrices"""
    def __init__(self, comm):
        """Initialize processor layout.

        Args:
            comm: MPI communicator from mpi4py
        """
        nprocs = comm.Get_size()
        rank = comm.Get_rank()
        for np_cols in range(int(np.sqrt(nprocs)), 0, -1):
            if nprocs % np_cols == 0:
              break
        #if nprocs == 1:
        #    np_cols = 1
        np_rows = nprocs//np_cols
        # column major distribution of processors
        my_pcol = rank // np_rows
        my_prow = rank % np_rows
        self.np_cols, self.np_rows = np_cols, np_rows
        self.my_pcol, self.my_prow = my_pcol, my_prow
        self.comm = comm
        self.comm_f = comm.py2f()


class DistributedMatrix:
    """Class for generating a distributed block-cyclic matrix

    The data attribute contains the array in the correct size for the local
    processor.
    """
    def __init__(self, processor_layout, na, nev, nblk, dtype=np.float64):
        """Initialize distributed matrix for a given processor layout.

        Args:
            processor_layout (ProcessorLayout): has to be created from MPI
                communicator
            na (int): dimension of matrix
            nev (int): number of eigenvectors/eigenvalues to be computed
            nblk (int): block size of distributed matrix
            dtype: data type of matrix
        """
        self.na = na
        self.nev = nev
        self.nblk = nblk
        self.processor_layout = processor_layout

        # get local size
        self.na_rows = self.numroc(na, nblk, processor_layout.my_prow, 0,
                                   processor_layout.np_rows)
        self.na_cols = self.numroc(na, nblk, processor_layout.my_pcol, 0,
                                   processor_layout.np_cols)
        # create array
        self.data = np.empty((self.na_rows, self.na_cols),
                             dtype=dtype, order='F')

        self.elpa = None

    @classmethod
    def from_communicator(cls, comm, na, nev, nblk, dtype=np.float64):
        """Initialize distributed matrix from a MPI communicator.

        Args:
            comm: MPI communicator from mpi4py
            na (int): dimension of matrix
            nev (int): number of eigenvectors/eigenvalues to be computed
            nblk (int): block size of distributed matrix
            dtype: data type of matrix
        """
        processor_layout = ProcessorLayout(comm)
        return cls(processor_layout, na, nev, nblk, dtype)

    @classmethod
    def from_comm_world(cls, na, nev, nblk, dtype=np.float64):
        """Initialize distributed matrix from the MPI_COMM_WORLD communicator.

        Args:
            na (int): dimension of matrix
            nev (int): number of eigenvectors/eigenvalues to be computed
            nblk (int): block size of distributed matrix
            dtype: data type of matrix
        """
        from mpi4py import MPI
        comm = MPI.COMM_WORLD
        processor_layout = ProcessorLayout(comm)
        return cls(processor_layout, na, nev, nblk, dtype)

    @classmethod
    def like(cls, matrix):
        """Get a DistributedMatrix with the same parameters as matrix"""
        return cls(matrix.processor_layout, matrix.na, matrix.nev, matrix.nblk,
                   matrix.data.dtype)

    def get_local_index(self, global_row, global_col):
        """compute local row and column indices from global ones

        Returns a tuple of the local row and column indices
        """
        local_row = self.indxg2l(global_row, self.nblk,
                                 self.processor_layout.my_prow, 0,
                                 self.processor_layout.np_rows)
        local_col = self.indxg2l(global_col, self.nblk,
                                 self.processor_layout.my_pcol, 0,
                                 self.processor_layout.np_cols)
        return local_row, local_col

    def get_global_index(self, local_row, local_col):
        """compute global row and column indices from local ones

        Returns a tuple of the global row and column indices
        """
        global_row = self.indxl2g(local_row, self.nblk,
                                  self.processor_layout.my_prow, 0,
                                  self.processor_layout.np_rows)
        global_col = self.indxl2g(local_col, self.nblk,
                                  self.processor_layout.my_pcol, 0,
                                  self.processor_layout.np_cols)
        return global_row, global_col

    def is_local_index(self, global_row, global_col):
        """check if global index is stored on current processor"""
        return self.is_local_row(global_row) and self.is_local_col(global_col)

    def is_local_row(self, global_row):
        """check if global row is stored on this processor"""
        process_row = self.indxg2p(global_row, self.nblk,
                                   self.processor_layout.my_prow, 0,
                                   self.processor_layout.np_rows)
        return process_row == self.processor_layout.my_prow

    def is_local_col(self, global_col):
        process_col = self.indxg2p(global_col, self.nblk,
                                   self.processor_layout.my_pcol, 0,
                                   self.processor_layout.np_cols)
        return process_col == self.processor_layout.my_pcol

    @staticmethod
    def indxg2l(indxglob, nb, iproc, isrcproc, nprocs):
        """compute local index from global index indxglob

        original netlib scalapack source:

        .. code-block:: fortran

            INDXG2L = NB*((INDXGLOB-1)/(NB*NPROCS))+MOD(INDXGLOB-1,NB)+1
        """
        # adapt to python 0-based indexing
        return nb*(indxglob//(nb*nprocs)) + indxglob%nb

    @staticmethod
    def indxl2g(indxloc, nb, iproc, isrcproc, nprocs):
        """compute global index from local index indxloc

        original netlib scalapack source:

        .. code-block:: fortran

            INDXL2G = NPROCS*NB*((INDXLOC-1)/NB) + MOD(INDXLOC-1,NB) +
                      MOD(NPROCS+IPROC-ISRCPROC, NPROCS)*NB + 1
        """
        # adapt to python 0-based indexing
        return nprocs*nb*(indxloc//nb) + indxloc%nb + \
            ((nprocs+iproc-isrcproc)%nprocs)*nb

    @staticmethod
    def indxg2p(indxglob, nb, iproc, isrcproc, nprocs):
        """compute process coordinate for global index

        original netlib scalapack source:

        .. code-block:: fortran

            INDXG2P = MOD( ISRCPROC + (INDXGLOB - 1) / NB, NPROCS )
        """
        # adapt to python 0-based indexing
        return (isrcproc + indxglob // nb) % nprocs

    @staticmethod
    def numroc(n, nb, iproc, isrcproc, nprocs):
        """Get local dimensions of distributed block-cyclic matrix.

        Programmed after scalapack source (tools/numroc.f on netlib).
        """
        mydist = (nprocs + iproc - isrcproc) % nprocs
        nblocks = n // nb
        result = (nblocks // nprocs) * nb
        extrablks = nblocks % nprocs
        if mydist < extrablks:
            result += nb
        elif mydist == extrablks:
            result += n % nb
        return int(result)

    def _initialized_elpa(function):
        # wrapper to ensure one-time initialization of Elpa object
        @wraps(function)
        def wrapped_function(self):
            if self.elpa is None:
                self.elpa = Elpa.from_distributed_matrix(self)
            return function(self)
        return wrapped_function

    @_initialized_elpa
    def compute_eigenvectors(self):
        """Compute eigenvalues and eigenvectors

        The eigenvectors are stored in columns.
        This function returns a dictionary with entries 'eigenvalues' and
        'eigenvectors'.

        After computing the eigenvectors, the original content of the matrix is
        lost.
        """
        eigenvectors = DistributedMatrix.like(self)
        eigenvalues = np.zeros(self.na, dtype=np.float64)
        # call ELPA
        self.elpa.eigenvectors(self.data, eigenvalues, eigenvectors.data)
        return {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}

    @_initialized_elpa
    def compute_eigenvalues(self):
        """Compute only the eigenvalues.

        This function returns the eigenvalues as an array.

        After computing the eigenvalues, the original content of the matrix is
        lost.
        """
        eigenvalues = np.zeros(self.na, dtype=np.float64)
        # call ELPA
        self.elpa.eigenvalues(self.data, eigenvalues)
        return eigenvalues

    def set_data_from_global_matrix(self, matrix):
        """Set local part of the global matrix"""
        for local_row in range(self.na_rows):
            for local_col in range(self.na_cols):
                global_row, global_col = self.get_global_index(local_row,
                                                               local_col)
                self.data[local_row, local_col] = matrix[global_row,
                                                         global_col]

    def dot(self, vector):
        """Compute dot product of matrix with vector.

        This blocked implementation is much faster than the naive
        implementation.
        """
        if len(vector.shape) > 1 or vector.shape[0] != self.na:
            raise ValueError("Error: shape of vector {} incompatible to "
                             "matrix of size {:d}x{:d}.".format(
                                 vector.shape, self.na, self.na))
        from mpi4py import MPI
        summation = np.zeros_like(vector)
        # loop only over blocks here
        for local_row in range(0, self.na_rows, self.nblk):
            for local_col in range(0, self.na_cols, self.nblk):
                # do not go beyond the end of the matrix
                row_block_size = min(local_row + self.nblk,
                                     self.na_rows) - local_row
                col_block_size = min(local_col + self.nblk,
                                     self.na_cols) - local_col
                global_row, global_col = self.get_global_index(local_row,
                                                               local_col)
                # use numpy for faster dot product of local block
                summation[global_row:global_row+row_block_size] += \
                    np.dot(self.data[local_row:local_row + row_block_size,
                                     local_col:local_col + col_block_size],
                           vector[global_col:global_col+col_block_size])
        result = np.zeros_like(vector)
        self.processor_layout.comm.Allreduce(summation, result, op=MPI.SUM)
        return result

    def _dot_naive(self, vector):
        """Compute naive dot product of matrix with vector.

        Still in here as an example and for testing purposes.
        """
        from mpi4py import MPI
        summation = np.zeros_like(vector)
        for local_row in range(self.na_rows):
            for local_col in range(self.na_cols):
                global_row, global_col = self.get_global_index(local_row,
                                                               local_col)
                summation[global_row] += self.data[local_row, local_col] *\
                    vector[global_col]
        result = np.zeros_like(vector)
        self.processor_layout.comm.Allreduce(summation, result, op=MPI.SUM)
        return result

    def get_column(self, global_col):
        """Return global column"""
        from mpi4py import MPI
        column = np.zeros(self.na, dtype=self.data.dtype)
        temporary = np.zeros_like(column)
        if self.is_local_col(global_col):
            for global_row in range(self.na):
                if not self.is_local_row(global_row):
                    continue
                local_row, local_col = self.get_local_index(global_row,
                                                            global_col)
                temporary[global_row] = self.data[local_row, local_col]
        # this could be done more efficiently with a gather
        self.processor_layout.comm.Allreduce(temporary, column, op=MPI.SUM)
        return column

    def get_row(self, global_row):
        """Return global row"""
        from mpi4py import MPI
        row = np.zeros(self.na, dtype=self.data.dtype)
        temporary = np.zeros_like(row)
        if self.is_local_row(global_row):
            for global_col in range(self.na):
                if not self.is_local_col(global_col):
                    continue
                local_row, local_col = self.get_local_index(global_row,
                                                            global_col)
                temporary[global_col] = self.data[local_row, local_col]
        # this could be done more efficiently with a gather
        self.processor_layout.comm.Allreduce(temporary, row, op=MPI.SUM)
        return row
330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353

    def global_indices(self):
        """Return iterator over global indices of matrix.

        Use together with set_data_global_index and get_data_global_index.
        """
        for local_row in range(self.na_rows):
            for local_col in range(self.na_cols):
                yield self.get_global_index(local_row, local_col)

    def set_data_for_global_index(self, global_row, global_col, value):
        """Set value of matrix at global coordinates"""
        if self.is_local_index(global_row, global_col):
            local_row, local_col = self.get_local_index(global_row, global_col)
            self.data[local_row, local_col] = value

    def get_data_for_global_index(self, global_row, global_col):
        """Get value of matrix at global coordinates"""
        if self.is_local_index(global_row, global_col):
            local_row, local_col = self.get_local_index(global_row, global_col)
            return self.data[local_row, local_col]
        else:
            raise ValueError('Index out of bounds: global row {:d}, '
                             'global col {:d}'.format(global_row, global_col))
354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399

    def global_block_indices(self):
        """Return iterator over global indices of matrix blocks.

        Use together with set_block_global_index and get_block_global_index
        for more efficient loops.
        """
        for local_row in range(0, self.na_rows, self.nblk):
            for local_col in range(0, self.na_cols, self.nblk):
                # do not go beyond the end of the matrix
                row_block_size = min(local_row + self.nblk,
                                     self.na_rows) - local_row
                col_block_size = min(local_col + self.nblk,
                                     self.na_cols) - local_col
                global_row, global_col = self.get_global_index(local_row,
                                                               local_col)
                yield global_row, global_col, row_block_size, col_block_size

    def set_block_for_global_index(self, global_row, global_col,
                                   row_block_size, col_block_size, value):
        """Set value of block of matrix at global coordinates"""
        if self.is_local_index(global_row, global_col):
            local_row, local_col = self.get_local_index(global_row, global_col)
            if value.shape != (row_block_size, col_block_size):
                raise ValueError("value has the wrong shape. "
                                 "Expected: {}, found: {}."
                                 .format((row_block_size, col_block_size),
                                         value.shape)
                                 )
            self.data[local_row:local_row+row_block_size,
                      local_col:local_col+col_block_size] = value

    def get_block_for_global_index(self, global_row, global_col,
                                   row_block_size, col_block_size):
        """Get value of block of matrix at global coordinates"""
        if self.is_local_index(global_row, global_col):
            local_row, local_col = self.get_local_index(global_row, global_col)
            if local_row+row_block_size > self.na_rows or \
                    local_col+col_block_size > self.na_cols:
                raise ValueError("Block size wrong: exceeds dimensions of"
                                 " matrix.")
            return self.data[local_row:local_row+row_block_size,
                             local_col:local_col+col_block_size]
        else:
            raise ValueError('Index out of bounds: global row {:d}, '
                             'global col {:d}'.format(global_row, global_col))