utils.py 9.76 KB
Newer Older
1
2
3
4
from __future__ import print_function

""" Set of utilities."""

5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def extractFileNamefromUrl(url):
   """Extracts a file name from url. 

   :param url: url
   :type url: string
   :returns: file name 
   :rtype: string

   :Example:

   >>> from westpy import * 
   >>> extractFileNamefromUrl("http://www.west-code.org/database/gw100/xyz/CH4.xyz")
   """
   #
   fname = None 
   my_url = url[:-1] if url.endswith('/') else url
   if my_url.find('/'):
      fname = my_url.rsplit('/', 1)[1]
   return fname 
    

def download(url,fname=None):
27
28
29
30
   """Downloads a file from url. 

   :param url: url
   :type url: string
31
32
   :param fname: file name, optional 
   :type fname: string
33
34
35
36

   :Example:

   >>> from westpy import * 
37
   >>> download("http://www.west-code.org/database/gw100/xyz/CH4.xyz")
38
39
40
41

   .. note:: The file will be downloaded in the current directory. 
   """
   #
42
43
44
   if fname is None :
      fname = extractFileNamefromUrl(url)
   #
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
   from requests import get
   # open in binary mode
   with open(fname, "wb") as file:
       # get request
       response = get(url)
       # write to file
       file.write(response.content)
       #
       print("Downloaded file: ", fname, ", from url: ", url)


def bool2str( logical ):
   """Converts a boolean type into a string .TRUE. or .FALSE. . 

   :param logical: logical  
   :type logical: boolean
   :returns: .TRUE. or .FALSE.
   :rtype: string 

   :Example:

   >>> from westpy import * 
   >>> t = bool2str(True)
   >>> f = bool2str(False)
   >>> print(t,f) 
   .TRUE. .FALSE.
   """
   #
   if( logical ) : 
      return ".TRUE."
   else : 
      return ".FALSE."

Marco Govoni's avatar
Marco Govoni committed
78
def writeJsonFile(fname,data):
79
80
81
82
   """Writes data to file using the JSON format. 

   :param fname: file name
   :type fname: string
Marco Govoni's avatar
Marco Govoni committed
83
84
   :param data: data
   :type data: dict/list
85
86
87
88
89
90

   :Example:

   >>> from westpy import * 
   >>> data = {}
   >>> data["mass"] = 1.0
Marco Govoni's avatar
Marco Govoni committed
91
   >>> writeJsonFile("mass.json",data) 
92
93
94
95
96
97
98
99
100
101
102
103

   .. note:: The file will be generated in the current directory. 
   """
   #
   import json 
   #
   with open(fname, 'w') as file:
      json.dump(data, file, indent=2)
      #
      print("")
      print("File written : ", fname )  

Marco Govoni's avatar
Marco Govoni committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def readJsonFile(fname):
   """Reads data from file using the JSON format. 

   :param fname: file name
   :type fname: string
   :returns: data 
   :rtype: dict/list

   :Example:

   >>> from westpy import * 
   >>> data = readJsonFile("mass.json") 

   .. note:: The file will be read from the current directory. 
   """
   #
   import json 
   #
   with open(fname, 'r') as file:
      data = json.load(file)
      #
      print("")
      print("File read : ", fname )
   return data  

Marco Govoni's avatar
Marco Govoni committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def convertYaml2Json(fyml,fjson):
   """Converts the file from YAML to JSON. 

   :param fyml: Name of YAML file 
   :type fyml: string
   :param fjson: Name of JSON file 
   :type fjson: string

   :Example:

   >>> from westpy import * 
   >>> convertYaml2Json("file.yml","file.json") 

   .. note:: The file fjon will be created, fyml will not be overwritten. 
   """
   #
   import yaml, json 
   from westpy import writeJsonFile
   #
   data = yaml.load(open(fyml))
   writeJsonFile(fjson,data)
Marco Govoni's avatar
Marco Govoni committed
150

151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def listLinesWithKeyfromOnlineText(url,key):
   """List lines from text file located at url, with key.

   :param url: url
   :type url: string
   :param key: key word
   :type key: string
   :returns: list of lines
   :rtype: list

   :Example:

   >>> from westpy import * 
   >>> url = "http://www.quantum-simulation.org/potentials/sg15_oncv/upf/Si_ONCV_PBE-1.1.upf"
   >>> key = "z_valence"
   >>> l = listLinesWithKeyfromOnlineText(url,key)
   >>> print(l) 
   ['       z_valence="    4.00"'] 

   .. note:: Can be used to grep values from a UPF file.
   """
   #
   from urllib.request import urlopen
   import re
   data = urlopen(url) # parse the data
   greplist = []
   for line in data :
      if( key in str(line) ) : 
         greplist.append(line)
   return greplist

#
# list values from XML file located at url, with key  
#
def listValuesWithKeyFromOnlineXML(url,key):
   """List values from XML file located at url, with key.

   :param url: url
   :type url: string
   :param key: key word
   :type key: string
   :returns: list of values
   :rtype: list

   :Example:

   >>> from westpy import * 
   >>> url = "http://www.quantum-simulation.org/potentials/sg15_oncv/xml/Si_ONCV_PBE-1.1.xml"
   >>> key = "valence_charge"
   >>> l = listLinesWithKeyfromOnlineXML(url,key)
   >>> print(l) 
   ['4'] 

   .. note:: Can be used to grep values from a XML file.
   """
   #
   from urllib.request import urlopen
   import xml.etree.ElementTree as ET
   tree = ET.parse(urlopen(url)) # parse the data
   root = tree.getroot()
   xml_values = [str(xml_val.text).strip() for xml_val in root.iter(key)] #get values
   return xml_values
Marco Govoni's avatar
Marco Govoni committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235


def gaussian(x, mu, sig):
   """return normal distribution at point x.

   :math:`f(x;\\mu,\\sigma) = \\frac{1}{\\sigma\sqrt{2\\pi}}e^{-\\frac{(x-\\mu)^2}{2\\sigma^2}}`

   :param x: x
   :type x: float
   :param mu: :math:`\\mu`
   :type mu: float
   :param sigma: :math:`\\sigma`
   :type sigma: float
   :returns: :math:`f(x;\\mu,\\sigma)`
   :rtype: float

   :Example:

   >>> from westpy import * 
   >>> gaussian(1.0,2.0,3.0)
   """
   import numpy as np
   return 1./(np.sqrt(2.*np.pi)*sig)*np.exp(-np.power((x - mu)/sig, 2.)/2)
Marco Govoni's avatar
Marco Govoni committed
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323


def _putline(*args):
    """
    Generate a line to be written to a cube file where 
    the first field is an int and the remaining fields are floats.
    
    params:
        *args: first arg is formatted as int and remaining as floats
    
    returns: formatted string to be written to file with trailing newline
    """
    s = "{0:^ 8d}".format(args[0])
    s += "".join("{0:< 12.6f}".format(arg) for arg in args[1:])
    return s + "\n"

def _getline(cube):
    """
    Read a line from cube file where first field is an int 
    and the remaining fields are floats.
    
    params:
        cube: file object of the cube file
    
    returns: (int, list<float>)
    """
    l = cube.readline().strip().split()
    return int(l[0]), map(float, l[1:])

def read_cube(fname):
    """ 
    Read cube file into numpy array
    
    params:
        fname: filename of cube file
        
    returns: (data: np.array, metadata: dict)
    """
    import numpy as np
    meta = {}
    with open(fname, 'r') as cube:
        cube.readline(); cube.readline()  # ignore comments
        natm, meta['org'] = _getline(cube)
        nx, meta['xvec'] = _getline(cube)
        ny, meta['yvec'] = _getline(cube)
        nz, meta['zvec'] = _getline(cube)
        meta['atoms'] = [_getline(cube) for i in range(natm)]
        data = np.zeros((nx*ny*nz))
        idx = 0
        for line in cube:
            for val in line.strip().split():
                data[idx] = float(val)
                idx += 1
    data = np.reshape(data, (nx, ny, nz))
    return data, meta


def read_imcube(rfname, ifname = ""):
    """
    Convenience function to read in two cube files at once, 
    where one contains the real part and the other contains the 
    imag part. If only one filename given, other filename is inferred.
    
    params:
        rfname: filename of cube file of real part
        ifname: optional, filename of cube file of imag part
        
    returns: np.array (real part + j*imag part)
    """
    import numpy as np
    ifname = ifname or rfname.replace('real', 'imag')
    _debug("reading from files", rfname, "and", ifname)
    re, im = read_cube(rfname), read_cube(ifname)
    fin = np.zeros(re[0].shape, dtype='complex128')
    if re[1] != im[1]:
        _debug("warning: meta data mismatch, real part metadata retained")
    fin += re[0]
    fin += 1j*im[0]
    return fin, re[1]



def write_cube(data, meta, fname):
    """
    Write volumetric data to cube file along
    
    params:
        data: volumetric data consisting real values
Marco Govoni's avatar
Marco Govoni committed
324
325
326
327
        meta: dict containing metadata with following keys: 
              - atoms: list of atoms in the form (mass, [position])
              - org: origin
              - xvec,yvec,zvec: lattice vector basis
Marco Govoni's avatar
Marco Govoni committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
        fname: filename of cubefile (existing files overwritten)
    
    returns: None
    """
    with open(fname, "w") as cube:
        # first two lines are comments
        cube.write(" Cubefile created by cubetools.py\n  source: none\n")
        natm = len(meta['atoms'])
        nx, ny, nz = data.shape
        cube.write(_putline(natm, *meta['org'])) # 3rd line #atoms and origin
        cube.write(_putline(nx, *meta['xvec']))
        cube.write(_putline(ny, *meta['yvec']))
        cube.write(_putline(nz, *meta['zvec']))
        for atom_mass, atom_pos in meta['atoms']:
            cube.write(_putline(atom_mass, *atom_pos)) #skip the newline
        for i in range(nx):
            for j in range(ny):
                for k in range(nz):
                    if (i or j or k) and k%6==0:
                        cube.write("\n")
                    cube.write(" {0: .5E}".format(data[i,j,k]))

def write_imcube(data, meta, rfname, ifname=""):
    """
    Convenience function to write two cube files from complex valued 
    volumetric data, one for the real part and one for the imaginary part.
    Data about atoms, origin and lattice vectors are kept same for both.
    If only one filename given, other filename is inferred.
    
    params: 
        data: volumetric data consisting complex values
        meta: dict containing metadata with following keys
Marco Govoni's avatar
Marco Govoni committed
360
361
362
              - atoms: list of atoms in the form (mass, [position])
              - org: origin
              - xvec,yvec,zvec: lattice vector basis
Marco Govoni's avatar
Marco Govoni committed
363
364
365
366
367
368
369
370
371
        rfname: filename of cube file containing real part
        ifname: optional, filename of cube file containing imag part
        
    returns: None
    """
    ifname = ifname or rfname.replace('real', 'imag')
    _debug("writing data to files", rfname, "and", ifname)
    write_cube(data.real, meta, rfname)
    write_cube(data.imag, meta, ifname)