
# This is file is generated from interpolate.py.mako
# Do not edit this file
from pysph.sph.equation import Equation
from compyle.api import declare
from pysph.sph.wc.linalg import gj_solve, augmented_matrix


class EvaluateUhat(Equation):
    def _get_helpers_(self):
        return [gj_solve, augmented_matrix]

    def __init__(self, dest, sources, dim=1):
        self.dim = dim

        super(EvaluateUhat, self).__init__(dest, sources)

    def initialize(self, d_idx, d_uho, d_Buh, d_vho, d_Bvh, d_who, d_Bwh):
        i = declare('int')

        for i in range(3):
            d_uho[4*d_idx+i] = 0.0
            d_Buh[4*d_idx+i] = 0.0
            d_vho[4*d_idx+i] = 0.0
            d_Bvh[4*d_idx+i] = 0.0
            d_who[4*d_idx+i] = 0.0
            d_Bwh[4*d_idx+i] = 0.0

    def loop(self, d_idx, d_h, s_h, s_x, s_y, s_z, d_x, d_y, d_z, s_rho,
             s_m, s_idx, XIJ, DWIJ,
             WIJ, s_uhat, d_Buh, s_vhat, d_Bvh, s_what, d_Bwh):
        Vj = s_m[s_idx] / s_rho[s_idx]
        uhj = s_uhat[s_idx]
        vhj = s_vhat[s_idx]
        whj = s_what[s_idx]
        i4 = declare('int')
        i4 = 4*d_idx

        d_Buh[i4+0] += uhj * WIJ * Vj
        d_Buh[i4+1] += uhj * DWIJ[0] * Vj
        d_Buh[i4+2] += uhj * DWIJ[1] * Vj
        d_Buh[i4+3] += uhj * DWIJ[2] * Vj
        d_Bvh[i4+0] += vhj * WIJ * Vj
        d_Bvh[i4+1] += vhj * DWIJ[0] * Vj
        d_Bvh[i4+2] += vhj * DWIJ[1] * Vj
        d_Bvh[i4+3] += vhj * DWIJ[2] * Vj
        d_Bwh[i4+0] += whj * WIJ * Vj
        d_Bwh[i4+1] += whj * DWIJ[0] * Vj
        d_Bwh[i4+2] += whj * DWIJ[1] * Vj
        d_Bwh[i4+3] += whj * DWIJ[2] * Vj

    def post_loop(self, d_idx, d_A, d_uho, d_Buh, d_vho, d_Bvh, d_who, d_Bwh):
        a_mat = declare('matrix(16)')
        aug_mat = declare('matrix(20)')
        b_uh = declare('matrix(4)')
        res_uh = declare('matrix(4)')
        b_vh = declare('matrix(4)')
        res_vh = declare('matrix(4)')
        b_wh = declare('matrix(4)')
        res_wh = declare('matrix(4)')
        i, n, i16, i4 = declare('int', 4)
        i16 = 16*d_idx
        i4 = 4*d_idx

        for i in range(16):
            a_mat[i] = d_A[i16+i]
        for i in range(20):
            aug_mat[i] = 0.0
        for i in range(4):
            b_uh[i] = d_Buh[i4+i]
            res_uh[i] = 0.0
            b_vh[i] = d_Bvh[i4+i]
            res_vh[i] = 0.0
            b_wh[i] = d_Bwh[i4+i]
            res_wh[i] = 0.0

        n = self.dim + 1
        augmented_matrix(a_mat, b_uh, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_uh)
        for i in range(4):
            d_uho[i4+i] = res_uh[i]
        augmented_matrix(a_mat, b_vh, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_vh)
        for i in range(4):
            d_vho[i4+i] = res_vh[i]
        augmented_matrix(a_mat, b_wh, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_wh)
        for i in range(4):
            d_who[i4+i] = res_wh[i]


class ExtrapolateUhat(Equation):
    def initialize(self, d_idx, d_uhat, d_vhat, d_what):
        d_uhat[d_idx] = 0.0
        d_vhat[d_idx] = 0.0
        d_what[d_idx] = 0.0

    def loop(self, d_idx, d_uhat, d_uho, d_vhat, d_vho, d_what, d_who,
             d_disp, d_xn, d_yn, d_zn):

        delx = 2 * d_disp[d_idx] * d_xn[d_idx]
        dely = 2 * d_disp[d_idx] * d_yn[d_idx]
        delz = 2 * d_disp[d_idx] * d_zn[d_idx]
        d_uhat[d_idx] = -1. * (
             d_uho[4*d_idx+0]
             - delx * d_uho[4*d_idx+1]
             - dely * d_uho[4*d_idx+2]
             - delz * d_uho[4*d_idx+3]
                      )

        d_vhat[d_idx] = (
             d_vho[4*d_idx+0]
             - delx * d_vho[4*d_idx+1]
             - dely * d_vho[4*d_idx+2]
             - delz * d_vho[4*d_idx+3]
                      )

        d_what[d_idx] = (
             d_who[4*d_idx+0]
             - delx * d_who[4*d_idx+1]
             - dely * d_who[4*d_idx+2]
             - delz * d_who[4*d_idx+3]
                      )


class CopyUhatFromGhost(Equation):
    def initialize_pair(self,
                        d_idx, d_uhat, s_uhat, d_vhat, s_vhat, d_what, s_what):
        d_uhat[d_idx] = -1.0 * s_uhat[d_idx]

        d_vhat[d_idx] = s_vhat[d_idx]

        d_what[d_idx] = s_what[d_idx]


class EvaluateU(Equation):
    def _get_helpers_(self):
        return [gj_solve, augmented_matrix]

    def __init__(self, dest, sources, dim=1):
        self.dim = dim

        super(EvaluateU, self).__init__(dest, sources)

    def initialize(self, d_idx, d_uo, d_Bu, d_vo, d_Bv, d_wo, d_Bw):
        i = declare('int')

        for i in range(3):
            d_uo[4*d_idx+i] = 0.0
            d_Bu[4*d_idx+i] = 0.0
            d_vo[4*d_idx+i] = 0.0
            d_Bv[4*d_idx+i] = 0.0
            d_wo[4*d_idx+i] = 0.0
            d_Bw[4*d_idx+i] = 0.0

    def loop(self, d_idx, d_h, s_h, s_x, s_y, s_z, d_x, d_y, d_z, s_rho,
             s_m, s_idx, XIJ, DWIJ,
             WIJ, s_u, d_Bu, s_v, d_Bv, s_w, d_Bw):
        Vj = s_m[s_idx] / s_rho[s_idx]
        uj = s_u[s_idx]
        vj = s_v[s_idx]
        wj = s_w[s_idx]
        i4 = declare('int')
        i4 = 4*d_idx

        d_Bu[i4+0] += uj * WIJ * Vj
        d_Bu[i4+1] += uj * DWIJ[0] * Vj
        d_Bu[i4+2] += uj * DWIJ[1] * Vj
        d_Bu[i4+3] += uj * DWIJ[2] * Vj
        d_Bv[i4+0] += vj * WIJ * Vj
        d_Bv[i4+1] += vj * DWIJ[0] * Vj
        d_Bv[i4+2] += vj * DWIJ[1] * Vj
        d_Bv[i4+3] += vj * DWIJ[2] * Vj
        d_Bw[i4+0] += wj * WIJ * Vj
        d_Bw[i4+1] += wj * DWIJ[0] * Vj
        d_Bw[i4+2] += wj * DWIJ[1] * Vj
        d_Bw[i4+3] += wj * DWIJ[2] * Vj

    def post_loop(self, d_idx, d_A, d_uo, d_Bu, d_vo, d_Bv, d_wo, d_Bw):
        a_mat = declare('matrix(16)')
        aug_mat = declare('matrix(20)')
        b_u = declare('matrix(4)')
        res_u = declare('matrix(4)')
        b_v = declare('matrix(4)')
        res_v = declare('matrix(4)')
        b_w = declare('matrix(4)')
        res_w = declare('matrix(4)')
        i, n, i16, i4 = declare('int', 4)
        i16 = 16*d_idx
        i4 = 4*d_idx

        for i in range(16):
            a_mat[i] = d_A[i16+i]
        for i in range(20):
            aug_mat[i] = 0.0
        for i in range(4):
            b_u[i] = d_Bu[i4+i]
            res_u[i] = 0.0
            b_v[i] = d_Bv[i4+i]
            res_v[i] = 0.0
            b_w[i] = d_Bw[i4+i]
            res_w[i] = 0.0

        n = self.dim + 1
        augmented_matrix(a_mat, b_u, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_u)
        for i in range(4):
            d_uo[i4+i] = res_u[i]
        augmented_matrix(a_mat, b_v, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_v)
        for i in range(4):
            d_vo[i4+i] = res_v[i]
        augmented_matrix(a_mat, b_w, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_w)
        for i in range(4):
            d_wo[i4+i] = res_w[i]


class ExtrapolateU(Equation):
    def initialize(self, d_idx, d_u, d_v, d_w):
        d_u[d_idx] = 0.0
        d_v[d_idx] = 0.0
        d_w[d_idx] = 0.0

    def loop(self, d_idx, d_u, d_uo, d_v, d_vo, d_w, d_wo,
             d_disp, d_xn, d_yn, d_zn):

        delx = 2 * d_disp[d_idx] * d_xn[d_idx]
        dely = 2 * d_disp[d_idx] * d_yn[d_idx]
        delz = 2 * d_disp[d_idx] * d_zn[d_idx]
        d_u[d_idx] = -1. * (
             d_uo[4*d_idx+0]
             - delx * d_uo[4*d_idx+1]
             - dely * d_uo[4*d_idx+2]
             - delz * d_uo[4*d_idx+3]
                      )

        d_v[d_idx] = (
             d_vo[4*d_idx+0]
             - delx * d_vo[4*d_idx+1]
             - dely * d_vo[4*d_idx+2]
             - delz * d_vo[4*d_idx+3]
                      )

        d_w[d_idx] = (
             d_wo[4*d_idx+0]
             - delx * d_wo[4*d_idx+1]
             - dely * d_wo[4*d_idx+2]
             - delz * d_wo[4*d_idx+3]
                      )


class CopyUFromGhost(Equation):
    def initialize_pair(self,
                        d_idx, d_u, s_u, d_v, s_v, d_w, s_w):
        d_u[d_idx] = -1.0 * s_u[d_idx]

        d_v[d_idx] = s_v[d_idx]

        d_w[d_idx] = s_w[d_idx]


class EvaluateP(Equation):
    def _get_helpers_(self):
        return [gj_solve, augmented_matrix]

    def __init__(self, dest, sources, dim=1):
        self.dim = dim

        super(EvaluateP, self).__init__(dest, sources)

    def initialize(self, d_idx, d_po, d_Bp):
        i = declare('int')

        for i in range(3):
            d_po[4*d_idx+i] = 0.0
            d_Bp[4*d_idx+i] = 0.0

    def loop(self, d_idx, d_h, s_h, s_x, s_y, s_z, d_x, d_y, d_z, s_rho,
             s_m, s_idx, XIJ, DWIJ,
             WIJ, s_p, d_Bp):
        Vj = s_m[s_idx] / s_rho[s_idx]
        pj = s_p[s_idx]
        i4 = declare('int')
        i4 = 4*d_idx

        d_Bp[i4+0] += pj * WIJ * Vj
        d_Bp[i4+1] += pj * DWIJ[0] * Vj
        d_Bp[i4+2] += pj * DWIJ[1] * Vj
        d_Bp[i4+3] += pj * DWIJ[2] * Vj

    def post_loop(self, d_idx, d_A, d_po, d_Bp):
        a_mat = declare('matrix(16)')
        aug_mat = declare('matrix(20)')
        b_p = declare('matrix(4)')
        res_p = declare('matrix(4)')
        i, n, i16, i4 = declare('int', 4)
        i16 = 16*d_idx
        i4 = 4*d_idx

        for i in range(16):
            a_mat[i] = d_A[i16+i]
        for i in range(20):
            aug_mat[i] = 0.0
        for i in range(4):
            b_p[i] = d_Bp[i4+i]
            res_p[i] = 0.0

        n = self.dim + 1
        augmented_matrix(a_mat, b_p, n, 1, 4, aug_mat)
        gj_solve(aug_mat, n, 1, res_p)
        for i in range(4):
            d_po[i4+i] = res_p[i]


class ExtrapolateP(Equation):
    def initialize(self, d_idx, d_p):
        d_p[d_idx] = 0.0

    def loop(self, d_idx, d_p, d_po,
             d_disp, d_xn, d_yn, d_zn):

        delx = 2 * d_disp[d_idx] * d_xn[d_idx]
        dely = 2 * d_disp[d_idx] * d_yn[d_idx]
        delz = 2 * d_disp[d_idx] * d_zn[d_idx]
        d_p[d_idx] = (
             d_po[4*d_idx+0]
             - delx * d_po[4*d_idx+1]
             - dely * d_po[4*d_idx+2]
             - delz * d_po[4*d_idx+3]
                      )


class CopyPFromGhost(Equation):
    def initialize_pair(self,
                        d_idx, d_p, s_p):
        d_p[d_idx] = s_p[d_idx]


class UpdateMomentMatrix(Equation):
    def __init__(self, dest, sources, dim=1):
        self.dim = dim

        super(UpdateMomentMatrix, self).__init__(dest, sources)

    def initialize(self, d_idx, d_A):
        i, j = declare('int', 2)

        for i in range(4):
            for j in range(4):
                d_A[16*d_idx + j+4*i] = 0.0

    def loop(self, d_idx, s_idx, d_h, s_h, s_x, s_y, s_z, d_x, d_y,
             d_z, s_rho, s_m, d_A, XIJ, WIJ, DWIJ):
        Vj = s_m[s_idx] / s_rho[s_idx]
        i16 = declare('int')
        i16 = 16*d_idx
        d_A[i16+0] += WIJ * Vj

        d_A[i16+1] += -XIJ[0] * WIJ * Vj
        d_A[i16+2] += -XIJ[1] * WIJ * Vj
        d_A[i16+3] += -XIJ[2] * WIJ * Vj

        d_A[i16+4] += DWIJ[0] * Vj
        d_A[i16+8] += DWIJ[1] * Vj
        d_A[i16+12] += DWIJ[2] * Vj

        d_A[i16+5] += -XIJ[0] * DWIJ[0] * Vj
        d_A[i16+6] += -XIJ[1] * DWIJ[0] * Vj
        d_A[i16+7] += -XIJ[2] * DWIJ[0] * Vj

        d_A[i16+9] += -XIJ[0] * DWIJ[1] * Vj
        d_A[i16+10] += -XIJ[1] * DWIJ[1] * Vj
        d_A[i16+11] += -XIJ[2] * DWIJ[1] * Vj

        d_A[i16+13] += -XIJ[0] * DWIJ[2] * Vj
        d_A[i16+14] += -XIJ[1] * DWIJ[2] * Vj
        d_A[i16+15] += -XIJ[2] * DWIJ[2] * Vj
