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.
65 lines
1.0 KiB
Python
65 lines
1.0 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)
|
|
|
|
|
|
|
|
|
|
|