You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
57 lines
1.2 KiB
Python
57 lines
1.2 KiB
Python
import numpy as np
|
|
import sys
|
|
from refine import refine as ref
|
|
|
|
|
|
def get_p(pn, pdir, pprefix):
|
|
|
|
p = np.load(pdir + pprefix + "0" + ".npy")
|
|
for i in range(1, pn):
|
|
p = np.concatenate((p, np.load(pdir + pprefix + str(i) + ".npy")), axis=0)
|
|
|
|
return p
|
|
|
|
|
|
def get_k(pn, kdir, kprefix):
|
|
|
|
k = (np.load(kdir + kprefix + "0" + ".npy"))[1:-1, :, :]
|
|
for i in range(1, pn):
|
|
k = np.concatenate(
|
|
(k, (np.load(kdir + kprefix + str(i) + ".npy"))[1:-1, :, :]), axis=0
|
|
)
|
|
return ref(k, 2, 2, 2)
|
|
|
|
|
|
def kef(P, K, i, j, k, pbc):
|
|
# tx=2*K[:,:,i]*K[:,:,i+1]/(K[:,:,i]+K[:,:,i+1])
|
|
# ty=2*K[:,j,:]*K[:,j+1,:]/(K[:,j,:]+K[:,j+1,:])
|
|
tz = 2 * K[k, :, :] * K[k + 1, :, :] / (K[k, :, :] + K[k + 1, :, :])
|
|
|
|
# qx=tx*(P[:,:,i+1]-P[:,:,i])
|
|
# qy=ty*(P[:,j+1,:]-P[:,j,:])
|
|
qz = -tz * (P[k + 1, :, :] - P[k, :, :])
|
|
|
|
kz = qz.sum() * (K.shape[0] + 1) / (pbc * K.shape[1] * K.shape[2])
|
|
return kz
|
|
|
|
|
|
def test(pn, kdir, pdir, kprefix, pprefix):
|
|
|
|
K = get_k(pn, kdir, kprefix)
|
|
print(K.shape)
|
|
P = get_p(pn, pdir, pprefix)
|
|
print(P)
|
|
print(P.shape)
|
|
print(kef(P, K, 1, 1, 1, 1000))
|
|
return
|
|
|
|
|
|
pn = 4
|
|
kdir = "./test/"
|
|
pdir = "./test/"
|
|
kprefix = "k"
|
|
pprefix = "P"
|
|
|
|
|
|
test(pn, kdir, pdir, kprefix, pprefix)
|