|  | 
|  | 1 | +# cython: language_level=3 | 
|  | 2 | +from cpython.mem cimport PyMem_Malloc, PyMem_Free | 
|  | 3 | + | 
|  | 4 | +cdef extern from "tdnac.c": | 
|  | 5 | +    void TD_NAC(int istep, int nst, int nbasis, int norb, int nocc, int nvirt, double dt, \ | 
|  | 6 | +        int *orb_ini, int *orb_final, double **nacme, double **ao_overlap, \ | 
|  | 7 | +        double **mo_coef_old, double **mo_coef_new, double ***ci_coef_old, double ***ci_coef_new) | 
|  | 8 | + | 
|  | 9 | +def wf_overlap(qm, molecule, istep_py, dt_py): | 
|  | 10 | +    cdef: | 
|  | 11 | +        int *orb_ini | 
|  | 12 | +        int *orb_final | 
|  | 13 | +        double **nacme | 
|  | 14 | +        double **ao_overlap | 
|  | 15 | +        double **mo_coef_old | 
|  | 16 | +        double **mo_coef_new | 
|  | 17 | +        double ***ci_coef_old | 
|  | 18 | +        double ***ci_coef_new | 
|  | 19 | + | 
|  | 20 | +        int istep, ist, nst, ibasis, jbasis, iorb, jorb, nbasis, norb, nocc, nvirt | 
|  | 21 | +        double dt | 
|  | 22 | + | 
|  | 23 | +    # Assign size variables | 
|  | 24 | +    dt = dt_py | 
|  | 25 | +    istep = istep_py | 
|  | 26 | +    nst = molecule.nst | 
|  | 27 | +    nbasis = qm.nbasis | 
|  | 28 | +    norb = qm.norb | 
|  | 29 | +    nocc = qm.nocc | 
|  | 30 | +    nvirt = qm.nvirt | 
|  | 31 | + | 
|  | 32 | +    # Allocate NACME variables | 
|  | 33 | +    orb_ini = <int*> PyMem_Malloc(1 * sizeof(int)) | 
|  | 34 | +    orb_final = <int*> PyMem_Malloc(1 * sizeof(int)) | 
|  | 35 | + | 
|  | 36 | +    nacme = <double**> PyMem_Malloc(nst * sizeof(double*)) | 
|  | 37 | + | 
|  | 38 | +    ao_overlap = <double**> PyMem_Malloc(nbasis * sizeof(double*)) | 
|  | 39 | +    mo_coef_old = <double**> PyMem_Malloc(norb * sizeof(double*)) | 
|  | 40 | +    mo_coef_new = <double**> PyMem_Malloc(norb * sizeof(double*)) | 
|  | 41 | + | 
|  | 42 | +    for ist in range(nst): | 
|  | 43 | +        nacme[ist] = <double*> PyMem_Malloc(nst * sizeof(double)) | 
|  | 44 | + | 
|  | 45 | +    for ibasis in range(nbasis): | 
|  | 46 | +        ao_overlap[ibasis] = <double*> PyMem_Malloc(nbasis * sizeof(double)) | 
|  | 47 | + | 
|  | 48 | +    for iorb in range(norb): | 
|  | 49 | +        mo_coef_old[iorb] = <double*> PyMem_Malloc(nbasis * sizeof(double)) | 
|  | 50 | +        mo_coef_new[iorb] = <double*> PyMem_Malloc(nbasis * sizeof(double)) | 
|  | 51 | + | 
|  | 52 | +    ci_coef_old = <double***> PyMem_Malloc(nst * sizeof(double**)) | 
|  | 53 | +    ci_coef_new = <double***> PyMem_Malloc(nst * sizeof(double**)) | 
|  | 54 | + | 
|  | 55 | +    for ist in range(nst): | 
|  | 56 | +        ci_coef_old[ist] = <double**> PyMem_Malloc(nocc * sizeof(double*)) | 
|  | 57 | +        ci_coef_new[ist] = <double**> PyMem_Malloc(nocc * sizeof(double*)) | 
|  | 58 | + | 
|  | 59 | +    for ist in range(nst): | 
|  | 60 | +        for iorb in range(nocc): | 
|  | 61 | +            ci_coef_old[ist][iorb] = <double*> PyMem_Malloc(nvirt * sizeof(double)) | 
|  | 62 | +            ci_coef_new[ist][iorb] = <double*> PyMem_Malloc(nvirt * sizeof(double)) | 
|  | 63 | + | 
|  | 64 | +    # Assign NACME variables from python to C | 
|  | 65 | +    orb_ini[0] = qm.orb_ini[0] | 
|  | 66 | +    orb_final[0] = qm.orb_final[0] | 
|  | 67 | + | 
|  | 68 | +    for ist in range(nst): | 
|  | 69 | +        for jst in range(nst): | 
|  | 70 | +            nacme[ist][jst] = 0. | 
|  | 71 | + | 
|  | 72 | +    for ibasis in range(nbasis): | 
|  | 73 | +        for jbasis in range(nbasis): | 
|  | 74 | +            ao_overlap[ibasis][jbasis] = qm.ao_overlap[ibasis, jbasis] | 
|  | 75 | +     | 
|  | 76 | +    for iorb in range(norb): | 
|  | 77 | +        for ibasis in range(nbasis): | 
|  | 78 | +            mo_coef_old[iorb][ibasis] = qm.mo_coef_old[iorb, ibasis] | 
|  | 79 | +            mo_coef_new[iorb][ibasis] = qm.mo_coef_new[iorb, ibasis] | 
|  | 80 | + | 
|  | 81 | +    for ist in range(nst): | 
|  | 82 | +        for iorb in range(nocc): | 
|  | 83 | +            for jorb in range(nvirt): | 
|  | 84 | +                ci_coef_old[ist][iorb][jorb] = qm.ci_coef_old[ist, iorb, jorb] | 
|  | 85 | +                ci_coef_new[ist][iorb][jorb] = qm.ci_coef_new[ist, iorb, jorb] | 
|  | 86 | + | 
|  | 87 | +    # Calculate TDNAC term for CIoverlap | 
|  | 88 | +    TD_NAC(istep, nst, nbasis, norb, nocc, nvirt, dt, orb_ini, orb_final, nacme, \ | 
|  | 89 | +        ao_overlap, mo_coef_old, mo_coef_new, ci_coef_old, ci_coef_new) | 
|  | 90 | + | 
|  | 91 | +    # Assign NACME variables from C to python | 
|  | 92 | +    for ist in range(nst): | 
|  | 93 | +        for jst in range(nst): | 
|  | 94 | +             molecule.nacme[ist, jst] = nacme[ist][jst] | 
|  | 95 | + | 
|  | 96 | +    for iorb in range(norb): | 
|  | 97 | +        for ibasis in range(nbasis): | 
|  | 98 | +            qm.mo_coef_old[iorb, ibasis] = mo_coef_new[iorb][ibasis] | 
|  | 99 | + | 
|  | 100 | +    for ist in range(nst): | 
|  | 101 | +        for iorb in range(nocc): | 
|  | 102 | +            for jorb in range(nvirt): | 
|  | 103 | +                qm.ci_coef_old[ist, iorb, jorb] = ci_coef_new[ist][iorb][jorb] | 
|  | 104 | + | 
|  | 105 | +    # Deallocate NACME variables | 
|  | 106 | +    PyMem_Free(orb_ini) | 
|  | 107 | +    PyMem_Free(orb_final) | 
|  | 108 | + | 
|  | 109 | +    for ist in range(nst): | 
|  | 110 | +        PyMem_Free(nacme[ist]) | 
|  | 111 | + | 
|  | 112 | +    PyMem_Free(nacme) | 
|  | 113 | + | 
|  | 114 | +    for ibasis in range(nbasis): | 
|  | 115 | +        PyMem_Free(ao_overlap[ibasis]) | 
|  | 116 | + | 
|  | 117 | +    for iorb in range(norb): | 
|  | 118 | +        PyMem_Free(mo_coef_old[iorb]) | 
|  | 119 | +        PyMem_Free(mo_coef_new[iorb]) | 
|  | 120 | + | 
|  | 121 | +    PyMem_Free(ao_overlap) | 
|  | 122 | +    PyMem_Free(mo_coef_old) | 
|  | 123 | +    PyMem_Free(mo_coef_new) | 
|  | 124 | + | 
|  | 125 | +    for ist in range(nst): | 
|  | 126 | +        for iorb in range(nocc): | 
|  | 127 | +            PyMem_Free(ci_coef_old[ist][iorb]) | 
|  | 128 | +            PyMem_Free(ci_coef_new[ist][iorb]) | 
|  | 129 | + | 
|  | 130 | +    for ist in range(nst): | 
|  | 131 | +        PyMem_Free(ci_coef_old[ist]) | 
|  | 132 | +        PyMem_Free(ci_coef_new[ist]) | 
|  | 133 | + | 
|  | 134 | +    PyMem_Free(ci_coef_old) | 
|  | 135 | +    PyMem_Free(ci_coef_new) | 
|  | 136 | + | 
|  | 137 | + | 
0 commit comments