Skip to content

Commit

Permalink
fixed nc compilation with complex numbers
Browse files Browse the repository at this point in the history
Signed-off-by: Nick Papior <nickpapior@gmail.com>
  • Loading branch information
zerothi committed Nov 13, 2024
1 parent 6b65ab8 commit a9a5f0f
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 30 deletions.
24 changes: 19 additions & 5 deletions src/sisl/physics/_matrix_phase.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,10 @@ from sisl._core._dtypes cimport (
)

from ._matrix_utils cimport (
_f_matrix_box_nc,
_f_matrix_box_so,
_matrix_box_nc,
_matrix_box_nc_cmplx,
_matrix_box_nc_real,
_matrix_box_so_cmplx,
_matrix_box_so_real,
)
Expand Down Expand Up @@ -326,9 +328,15 @@ def _phase_csr_nc(ints_st[::1] ptr,
cdef ints_st r, rr, ind, s, s_idx, c

cdef complexs_st ph
cdef _f_matrix_box_nc func
cdef numerics_st *d
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

with nogil:
if p_opt == -1:
for r in range(nr):
Expand Down Expand Up @@ -356,7 +364,7 @@ def _phase_csr_nc(ints_st[::1] ptr,
s_idx = _index_sorted(tmp, c)

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[v_ptr[rr] + s_idx] += M[0]
v[v_ptr[rr] + s_idx+1] += M[1]
v[v_ptr[rr+1] + s_idx] += M[2]
Expand All @@ -374,7 +382,7 @@ def _phase_csr_nc(ints_st[::1] ptr,
s_idx = _index_sorted(tmp, c)

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[v_ptr[rr] + s_idx] += M[0]
v[v_ptr[rr] + s_idx+1] += M[1]
v[v_ptr[rr+1] + s_idx] += M[2]
Expand Down Expand Up @@ -406,9 +414,15 @@ def _phase_array_nc(ints_st[::1] ptr,
cdef ints_st r, rr, ind, s, c

cdef complexs_st ph
cdef _f_matrix_box_nc func
cdef numerics_st *d
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

with nogil:
if p_opt == -1:
for r in range(nr):
Expand All @@ -430,7 +444,7 @@ def _phase_array_nc(ints_st[::1] ptr,
ph = phases[ind]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[rr, c] += M[0]
v[rr, c + 1] += M[1]
v[rr + 1, c] += M[2]
Expand All @@ -445,7 +459,7 @@ def _phase_array_nc(ints_st[::1] ptr,
ph = phases[s]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[rr, c] += M[0]
v[rr, c + 1] += M[1]
v[rr + 1, c] += M[2]
Expand Down
40 changes: 27 additions & 13 deletions src/sisl/physics/_matrix_phase3.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@ from sisl._core._dtypes cimport (
)

from ._matrix_utils cimport (
_f_matrix_box_nc,
_f_matrix_box_so,
_matrix_box_nc,
_matrix_box_nc_cmplx,
_matrix_box_nc_real,
_matrix_box_so_cmplx,
_matrix_box_so_real,
)
Expand Down Expand Up @@ -175,8 +177,14 @@ def _phase3_csr_nc(ints_st[::1] ptr,
cdef ints_st r, rr, ind, s, c
cdef ints_st s_idx
cdef numerics_st *d
cdef _f_matrix_box_nc func
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

with nogil:
if p_opt == 0:
for r in range(nr):
Expand All @@ -188,21 +196,21 @@ def _phase3_csr_nc(ints_st[::1] ptr,
d = &D[ind, 0]

ph = phases[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vx[v_ptr[rr] + s_idx] += M[0]
Vx[v_ptr[rr] + s_idx+1] += M[1]
Vx[v_ptr[rr+1] + s_idx] += M[2]
Vx[v_ptr[rr+1] + s_idx+1] += M[3]

ph = phases[ind, 1]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vy[v_ptr[rr] + s_idx] += M[0]
Vy[v_ptr[rr] + s_idx+1] += M[1]
Vy[v_ptr[rr+1] + s_idx] += M[2]
Vy[v_ptr[rr+1] + s_idx+1] += M[3]

ph = phases[ind, 2]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vz[v_ptr[rr] + s_idx] += M[0]
Vz[v_ptr[rr] + s_idx+1] += M[1]
Vz[v_ptr[rr+1] + s_idx] += M[2]
Expand All @@ -220,21 +228,21 @@ def _phase3_csr_nc(ints_st[::1] ptr,
d = &D[ind, 0]

ph = phases[s, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vx[v_ptr[rr] + s_idx] += M[0]
Vx[v_ptr[rr] + s_idx+1] += M[1]
Vx[v_ptr[rr+1] + s_idx] += M[2]
Vx[v_ptr[rr+1] + s_idx+1] += M[3]

ph = phases[s, 1]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vy[v_ptr[rr] + s_idx] += M[0]
Vy[v_ptr[rr] + s_idx+1] += M[1]
Vy[v_ptr[rr+1] + s_idx] += M[2]
Vy[v_ptr[rr+1] + s_idx+1] += M[3]

ph = phases[s, 2]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vz[v_ptr[rr] + s_idx] += M[0]
Vz[v_ptr[rr] + s_idx+1] += M[1]
Vz[v_ptr[rr+1] + s_idx] += M[2]
Expand Down Expand Up @@ -266,8 +274,14 @@ def _phase3_array_nc(ints_st[::1] ptr,
cdef ints_st r, rr, ind, s, c
cdef ints_st s_idx
cdef numerics_st *d
cdef _f_matrix_box_nc func
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

with nogil:
if p_opt == 0:
for r in range(nr):
Expand All @@ -278,21 +292,21 @@ def _phase3_array_nc(ints_st[::1] ptr,
d = &D[ind, 0]

ph = phases[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vx[rr, c] += M[0]
Vx[rr, c+1] += M[1]
Vx[rr+1, c] += M[2]
Vx[rr+1, c+1] += M[3]

ph = phases[ind, 1]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vy[rr, c] += M[0]
Vy[rr, c+1] += M[1]
Vy[rr+1, c] += M[2]
Vy[rr+1, c+1] += M[3]

ph = phases[ind, 2]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vz[rr, c] += M[0]
Vz[rr, c+1] += M[1]
Vz[rr+1, c] += M[2]
Expand All @@ -308,21 +322,21 @@ def _phase3_array_nc(ints_st[::1] ptr,
d = &D[ind, 0]

ph = phases[s, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vx[rr, c] += M[0]
Vx[rr, c+1] += M[1]
Vx[rr+1, c] += M[2]
Vx[rr+1, c+1] += M[3]

ph = phases[s, 1]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vy[rr, c] += M[0]
Vy[rr, c+1] += M[1]
Vy[rr+1, c] += M[2]
Vy[rr+1, c+1] += M[3]

ph = phases[s, 2]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
Vz[rr, c] += M[0]
Vz[rr, c+1] += M[1]
Vz[rr+1, c] += M[2]
Expand Down
24 changes: 19 additions & 5 deletions src/sisl/physics/_matrix_phase_sc.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,10 @@ from sisl._core._sparse cimport ncol2ptr_nc
from sisl._indices cimport _index_sorted

from ._matrix_utils cimport (
_f_matrix_box_nc,
_f_matrix_box_so,
_matrix_box_nc,
_matrix_box_nc_cmplx,
_matrix_box_nc_real,
_matrix_box_so_cmplx,
_matrix_box_so_real,
)
Expand Down Expand Up @@ -181,9 +183,15 @@ def _phase_sc_csr_nc(ints_st[::1] ptr,

cdef ints_st r, rr, cind, c, nz, ind
cdef complexs_st ph
cdef _f_matrix_box_nc func
cdef numerics_st *d
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

# We have to do it manually due to the double elements per matrix element
ncol2ptr_nc(nr, ncol, v_ptr, 2)

Expand Down Expand Up @@ -222,7 +230,7 @@ def _phase_sc_csr_nc(ints_st[::1] ptr,
ph = phases[ind]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[v_ptr[rr] + cind] = M[0]
v_col[v_ptr[rr] + cind] = c
v[v_ptr[rr] + cind+1] = M[1]
Expand All @@ -246,7 +254,7 @@ def _phase_sc_csr_nc(ints_st[::1] ptr,
ph = phases[col[ind] / nr]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)

v[v_ptr[rr] + cind] = M[0]
v_col[v_ptr[rr] + cind] = c
Expand Down Expand Up @@ -283,8 +291,14 @@ def _phase_sc_array_nc(ints_st[::1] ptr,
cdef complexs_st ph
cdef ints_st r, rr, c, nz, ind
cdef numerics_st *d
cdef _f_matrix_box_nc func
cdef complexs_st *M = [0, 0, 0, 0]

if numerics_st in complexs_st:
func = _matrix_box_nc_cmplx
else:
func = _matrix_box_nc_real

with nogil:
if p_opt == -1:
for r in range(nr):
Expand All @@ -305,7 +319,7 @@ def _phase_sc_array_nc(ints_st[::1] ptr,
ph = phases[ind]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[rr, c] = M[0]
v[rr, c+1] = M[1]
v[rr+1, c] = M[2]
Expand All @@ -319,7 +333,7 @@ def _phase_sc_array_nc(ints_st[::1] ptr,
ph = phases[col[ind] / nr]

d = &D[ind, 0]
_matrix_box_nc(d, ph, M)
func(d, ph, M)
v[rr, c] = M[0]
v[rr, c+1] = M[1]
v[rr+1, c] = M[2]
Expand Down
16 changes: 12 additions & 4 deletions src/sisl/physics/_matrix_utils.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,21 @@ ctypedef fused _internal_complexs_st:
float complex
double complex

ctypedef void(*_f_matrix_box_so)(const numerics_st *data,
ctypedef void(*_f_matrix_box_nc)(const numerics_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil

cdef void _matrix_box_nc(const numerics_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil
cdef void _matrix_box_nc_real(const reals_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil

cdef void _matrix_box_nc_cmplx(const _internal_complexs_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil

ctypedef void(*_f_matrix_box_so)(const numerics_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil

cdef void _matrix_box_so_real(const reals_st *data,
const complexs_st phase,
Expand Down
18 changes: 15 additions & 3 deletions src/sisl/physics/_matrix_utils.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,27 @@ M[3] == spin[1, 1]
@cython.wraparound(False)
@cython.initializedcheck(False)
@cython.cdivision(True)
cdef inline void _matrix_box_nc(const numerics_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil:
cdef inline void _matrix_box_nc_real(const reals_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil:
M[0] = <complexs_st> (data[0] * phase)
M[1] = <complexs_st> ((data[2] + 1j * data[3]) * phase)
M[2] = <complexs_st> ((data[2] + 1j * data[3]).conjugate() * phase)
M[3] = <complexs_st> (data[1] * phase)


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
@cython.cdivision(True)
cdef inline void _matrix_box_nc_cmplx(const _internal_complexs_st *data,
const complexs_st phase,
complexs_st *M) noexcept nogil:
M[0] = <complexs_st> (data[0] * phase)
M[1] = <complexs_st> (data[2] * phase)
M[2] = <complexs_st> (data[2].conjugate() * phase)
M[3] = <complexs_st> (data[1] * phase)


@cython.boundscheck(False)
@cython.wraparound(False)
Expand Down

0 comments on commit a9a5f0f

Please sign in to comment.