Skip to content

Commit

Permalink
starting to work on #53; adding there atoms [e,x,y,z] array into Grid…
Browse files Browse the repository at this point in the history
…Utils and HighLevel
  • Loading branch information
ondrejkrejci committed Feb 20, 2023
1 parent 6c18e93 commit 0524753
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 14 deletions.
30 changes: 17 additions & 13 deletions ppafm/GridUtils.py
Original file line number Diff line number Diff line change
Expand Up @@ -337,14 +337,16 @@ def saveWSxM_3D( prefix, data, extent, slices=None ):

#================ Npy

def saveNpy(fname, data, lvec , head=None):
def saveNpy(fname, data, lvec , atoms, head=None):
np.save(fname+'.npy', data)
np.save(fname+'_vec.npy',lvec)
np.save(fname+'_atoms.npy',np.array(atoms,dtype=float))

def loadNpy(fname):
data = np.load(fname+'.npy')
lvec = np.load(fname+'_vec.npy')
return data.copy(), lvec; #necessary for being 'C_CONTINUOS'
atoms = np.load(fname+'_atoms.npy')
return data.copy(), lvec, atoms; #necessary for being 'C_CONTINUOS'

# =============== Vector Field

Expand All @@ -371,24 +373,23 @@ def loadVecFieldNpy( fname, FF = None ):
Fy = np.load(fname+'_y.npy' )
Fz = np.load(fname+'_z.npy' )
lvec = np.load(fname+'_vec.npy' )
atoms = np.load(fname+'atoms.npy')
FF = packVecGrid( Fx, Fy, Fz, FF )
del Fx,Fy,Fz
return FF, lvec
return FF, lvec, atoms

def saveVecFieldXsf( fname, FF, lvec, head = XSF_HEAD_DEFAULT ):
saveXSF(fname+'_x.xsf', FF[:,:,:,0], lvec, head )
saveXSF(fname+'_y.xsf', FF[:,:,:,1], lvec, head )
saveXSF(fname+'_z.xsf', FF[:,:,:,2], lvec, head )

def saveVecFieldNpy( fname, FF, lvec , head = XSF_HEAD_DEFAULT ):
def saveVecFieldNpy( fname, FF, lvec , atoms, head = XSF_HEAD_DEFAULT ):
np.save(fname+'_x.npy', FF[:,:,:,0] )
np.save(fname+'_y.npy', FF[:,:,:,1] )
np.save(fname+'_z.npy', FF[:,:,:,2] )
np.save(fname+'_vec.npy', lvec )
if (head != XSF_HEAD_DEFAULT ):
print("saving atoms")
tmp0=head[0]; q=np.zeros(len(tmp0)); #head: [e,[x,y,z],lvec]
np.save(fname+'_atoms.npy',[tmp0,head[1][0],head[1][1],head[1][2],q]) #atoms: [e, x, y, z, q]
print("saving atoms") # atoms: [e, x, y, z];
np.save(fname+'_atoms.npy',np.array(atoms,dtype=float)) # at the moment we did not switched npz, I will do it later, maybe #

def limit_vec_field( FF, Fmax=100.0 ):
'''
Expand All @@ -400,14 +401,14 @@ def limit_vec_field( FF, Fmax=100.0 ):
FF[:,:,:,1].flat[mask] *= Fmax/FR[mask]
FF[:,:,:,2].flat[mask] *= Fmax/FR[mask]

def save_vec_field(fname, data, lvec, data_format="xsf", head = XSF_HEAD_DEFAULT ):
def save_vec_field(fname, data, lvec, data_format="xsf", head = XSF_HEAD_DEFAULT, atoms = [[0.],[0.],[0.],[0.]] ):
'''
Saving scalar fields into xsf, or npy
'''
if (data_format=="xsf"):
saveVecFieldXsf(fname, data, lvec, head = head )
elif (data_format=="npy"):
saveVecFieldNpy(fname, data, lvec, head = head )
saveVecFieldNpy(fname, data, lvec, atoms, head = head )
else:
print("I cannot save this format!")

Expand All @@ -419,7 +420,7 @@ def load_vec_field(fname, data_format="xsf"):
if (data_format=="xsf"):
data, lvec, ndim, head =loadVecFieldXsf(fname)
elif (data_format=="npy"):
data, lvec = loadVecFieldNpy(fname)
data, lvec, atoms = loadVecFieldNpy(fname)
ndim = np.delete(data.shape,3)
else:
print("I cannot load this format!")
Expand All @@ -428,14 +429,16 @@ def load_vec_field(fname, data_format="xsf"):

# =============== Scalar Fields

def save_scal_field(fname, data, lvec, data_format="xsf", head = XSF_HEAD_DEFAULT ):
def save_scal_field(fname, data, lvec, data_format="xsf", head = XSF_HEAD_DEFAULT , atoms = [[0.],[0.],[0.],[0.]]):
'''
Saving scalar fields into xsf, or npy
'''
if (data_format=="xsf"):
saveXSF(fname+".xsf", data, lvec, head = head)
elif (data_format=="npy"):
saveNpy(fname, data, lvec, head = head)
attmp = np.array(atoms,dtype=float);
np.save(fname+'_atoms.npy',attmp) # at the moment we did not switched npz, I will do it later, maybe #
else:
print("I cannot save this format!")

Expand All @@ -444,10 +447,11 @@ def load_scal_field(fname, data_format="xsf"):
'''
Loading scalar fields into xsf, or npy
'''
atoms = [[0.],[0.],[0.],[0.]] #
if (data_format=="xsf"):
data, lvec, ndim, head =loadXSF(fname+".xsf")
elif (data_format=="npy"):
data, lvec = loadNpy(fname)
data, lvec, atoms = loadNpy(fname)
ndim = data.shape
elif (data_format=="cube"):
data,lvec, ndim, head = loadCUBE(fname+".cube")
Expand Down
2 changes: 1 addition & 1 deletion ppafm/HighLevel.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ def computeLJ( geomFile, speciesFile, save_format=None, computeVpot=False, Fmax=
# --- save to files ?
if save_format is not None:
if(verbose>0): print("computeLJ Save ", save_format)
GU.save_vec_field( 'FF'+ffModel, FF, lvec, data_format=save_format, head=atomstring )
GU.save_vec_field( 'FF'+ffModel, FF, lvec, data_format=save_format, head=atomstring, atoms = atoms )
if computeVpot:
GU.save_scal_field( 'E'+ffModel, V, lvec, data_format=save_format, head=atomstring )
if(verbose>0): print("<<<END: computeLJ()")
Expand Down

0 comments on commit 0524753

Please sign in to comment.