File: wrapper.pyx

package info (click to toggle)
elpa 2022.11.001-4
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 33,636 kB
  • sloc: f90: 61,114; ansic: 39,478; sh: 4,795; cpp: 4,291; makefile: 923; asm: 834; python: 635; perl: 141
file content (249 lines) | stat: -rw-r--r-- 10,613 bytes parent folder | download | duplicates (3)
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
"""wrapper.pyx -- python wrapper for ELPA

This file contains the cython part of the wrapper.
"""
cimport numpy as np
import numpy as np
import sys

if 'mpi4py.MPI' in sys.modules.keys():
    raise NotImplementedError('Please load the pyelpa module before mpi4py, '
                              'otherwise there will be MPI problems.')

# import the function definitions from the ELPA header
cdef import from "<elpa/elpa.h>":
    cdef struct elpa_struct:
        pass
    ctypedef elpa_struct *elpa_t
    int elpa_init(int api_version)
    void elpa_uninit(int *error)
    elpa_t elpa_allocate(int *error)
    void elpa_deallocate(elpa_t handle, int *error)
    int elpa_setup(elpa_t handle)
    void elpa_set_integer(elpa_t handle, const char *name, int value, int *error)
    void elpa_get_integer(elpa_t handle, const char *name, int *value, int *error)
    void elpa_set_double(elpa_t handle, const char *name, double value, int *error)
    void elpa_get_double(elpa_t handle, const char *name, double *value, int *error)
    void elpa_eigenvectors_a_h_a_d(elpa_t handle, double *a, double *ev, double *q, int *error)
    void elpa_eigenvectors_a_h_a_f(elpa_t handle, float *a, float *ev, float *q, int *error)
    void elpa_eigenvectors_a_h_a_dc(elpa_t handle, double complex *a, double *ev, double complex *q, int *error)
    void elpa_eigenvectors_a_h_a_fc(elpa_t handle, float complex *a, float *ev, float complex *q, int *error)
    void elpa_eigenvalues_a_h_a_d(elpa_t handle, double *a, double *ev, int *error)
    void elpa_eigenvalues_a_h_a_f(elpa_t handle, float *a, float *ev, int *error)
    void elpa_eigenvalues_a_h_a_dc(elpa_t handle, double complex *a, double *ev, int *error)
    void elpa_eigenvalues_a_h_a_fc(elpa_t handle, float complex *a, float *ev, int *error)
    int ELPA_OK
    int ELPA_SOLVER_2STAGE


cdef class Elpa:
    """Wrapper for ELPA C interface.

    Provides routines for initialization, deinitialization, setting and getting
    properties and for calling the eigenvectors and eigenvalues routines.
    The routines eigenvectors and eigenvalues select the right ELPA routine to
    call depending on the argument type.
    """
    cdef elpa_t handle

    def __init__(self):
        """Run initialization and allocation of handle"""
        if elpa_init(20171201) != ELPA_OK:
            raise RuntimeError("ELPA API version not supported")
        cdef int error
        handle = elpa_allocate(&error)
        self.handle = handle

    def set_integer(self, description, int value):
        """Wraps elpa_set_integer"""
        cdef int error
        if isinstance(description, unicode):
            # encode to ascii for passing to C
            description = (<unicode>description).encode('ascii')
        cdef const char* c_string = description
        elpa_set_integer(<elpa_t>self.handle, description, value, &error)

    def get_integer(self, description):
        """Wraps elpa_get_integer"""
        cdef int error
        if isinstance(description, unicode):
            # encode to ascii for passing to C
            description = (<unicode>description).encode('ascii')
        cdef const char* c_string = description
        cdef int tmp
        elpa_get_integer(<elpa_t>self.handle, c_string, &tmp, &error)
        return tmp

    def set_double(self, description, double value):
        """Wraps elpa_set_double"""
        cdef int error
        if isinstance(description, unicode):
            # encode to ascii for passing to C
            description = (<unicode>description).encode('ascii')
        cdef const char* c_string = description
        elpa_set_double(<elpa_t>self.handle, description, value, &error)

    def get_double(self, description):
        """Wraps elpa_get_double"""
        cdef int error
        if isinstance(description, unicode):
            # encode to ascii for passing to C
            description = (<unicode>description).encode('ascii')
        cdef const char* c_string = description
        cdef double tmp
        elpa_get_double(<elpa_t>self.handle, c_string, &tmp, &error)
        return tmp

    def setup(self):
        """call setup function"""
        elpa_setup(<elpa_t>self.handle)

    def __del__(self):
        """Deallocation of handle and deinitialization"""
        cdef int error
        elpa_deallocate(<elpa_t>self.handle, &error)
        elpa_uninit(&error)

    def eigenvectors_d(self,
                       np.ndarray[np.float64_t, ndim=2] a,
                       np.ndarray[np.float64_t, ndim=1] ev,
                       np.ndarray[np.float64_t, ndim=2] q):
        cdef int error
        elpa_eigenvectors_a_h_a_d(<elpa_t>self.handle, <np.float64_t *>a.data,
                            <np.float64_t *>ev.data, <np.float64_t *>q.data,
                            <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvectors_f(self,
                       np.ndarray[np.float32_t, ndim=2] a,
                       np.ndarray[np.float32_t, ndim=1] ev,
                       np.ndarray[np.float32_t, ndim=2] q):
        cdef int error
        elpa_eigenvectors_a_h_a_f(<elpa_t>self.handle, <np.float32_t *>a.data,
                            <np.float32_t *>ev.data, <np.float32_t *>q.data,
                            <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvectors_dc(self,
                        np.ndarray[np.complex128_t, ndim=2] a,
                        np.ndarray[np.float64_t, ndim=1] ev,
                        np.ndarray[np.complex128_t, ndim=2] q):
        cdef int error
        elpa_eigenvectors_a_h_a_dc(<elpa_t>self.handle, <np.complex128_t *>a.data,
                             <np.float64_t *>ev.data, <np.complex128_t *>q.data,
                             <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvectors_fc(self,
                        np.ndarray[np.complex64_t, ndim=2] a,
                        np.ndarray[np.float32_t, ndim=1] ev,
                        np.ndarray[np.complex64_t, ndim=2] q):
        cdef int error
        elpa_eigenvectors_a_h_a_fc(<elpa_t>self.handle, <np.complex64_t *>a.data,
                             <np.float32_t *>ev.data, <np.complex64_t *>q.data,
                             <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvectors(self, a, ev, q):
        """Compute eigenvalues and eigenvectors.

        The data type of a is tested and the corresponding ELPA routine called

        Args:
            a (DistributedMatrix): problem matrix
            ev (numpy.ndarray): array of size a.na to store eigenvalues
            q (DistributedMatrix): store eigenvectors
        """
        if a.dtype == np.float64:
            self.eigenvectors_d(a, ev, q)
        elif a.dtype == np.float32:
            self.eigenvectors_f(a, ev, q)
        elif a.dtype == np.complex128:
            self.eigenvectors_dc(a, ev, q)
        elif a.dtype == np.complex64:
            self.eigenvectors_fc(a, ev, q)
        else:
            raise TypeError("Type not known.")

    def eigenvalues_d(self,
                       np.ndarray[np.float64_t, ndim=2] a,
                       np.ndarray[np.float64_t, ndim=1] ev):
        cdef int error
        elpa_eigenvalues_a_h_a_d(<elpa_t>self.handle, <np.float64_t *>a.data,
                           <np.float64_t *>ev.data, <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvalues_f(self,
                       np.ndarray[np.float32_t, ndim=2] a,
                       np.ndarray[np.float32_t, ndim=1] ev):
        cdef int error
        elpa_eigenvalues_a_h_a_f(<elpa_t>self.handle, <np.float32_t *>a.data,
                           <np.float32_t *>ev.data, <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvalues_dc(self,
                        np.ndarray[np.complex128_t, ndim=2] a,
                        np.ndarray[np.float64_t, ndim=1] ev):
        cdef int error
        elpa_eigenvalues_a_h_a_dc(<elpa_t>self.handle, <np.complex128_t *>a.data,
                            <np.float64_t *>ev.data, <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvalues_fc(self,
                        np.ndarray[np.complex64_t, ndim=2] a,
                        np.ndarray[np.float32_t, ndim=1] ev):
        cdef int error
        elpa_eigenvalues_a_h_a_fc(<elpa_t>self.handle, <np.complex64_t *>a.data,
                            <np.float32_t *>ev.data, <int*>&error)
        if error != ELPA_OK:
            raise RuntimeError("ELPA returned error value {:d}.".format(error))

    def eigenvalues(self, a, ev):
        """Compute eigenvalues.

        The data type of a is tested and the corresponding ELPA routine called

        Args:
            a (DistributedMatrix): problem matrix
            ev (numpy.ndarray): array of size a.na to store eigenvalues
        """
        if a.dtype == np.float64:
            self.eigenvalues_d(a, ev)
        elif a.dtype == np.float32:
            self.eigenvalues_f(a, ev)
        elif a.dtype == np.complex128:
            self.eigenvalues_dc(a, ev)
        elif a.dtype == np.complex64:
            self.eigenvalues_fc(a, ev)
        else:
            raise TypeError("Type not known.")

    @classmethod
    def from_distributed_matrix(cls, a):
        """Initialize ELPA with values from a distributed matrix

        Args:
            a (DistributedMatrix): matrix to get values from
        """
        self = cls()
        # Set parameters the matrix and it's MPI distribution
        self.set_integer("mpi_comm_parent", <int>a.processor_layout.comm_f)
        self.set_integer("na", <int>a.na)
        self.set_integer("nev", <int>a.nev)
        self.set_integer("local_nrows", <int>a.na_rows)
        self.set_integer("local_ncols", <int>a.na_cols)
        self.set_integer("nblk", <int>a.nblk)
        self.set_integer("process_row", <int>a.processor_layout.my_prow)
        self.set_integer("process_col", <int>a.processor_layout.my_pcol)
        # Setup
        self.setup()
        # if desired, set tunable run-time options
        self.set_integer("solver", ELPA_SOLVER_2STAGE)
        return self