Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

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

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

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

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

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

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

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

# coding: utf-8 

# Copyright (c) Pymatgen Development Team. 

# Distributed under the terms of the MIT License. 

"""Wrapper for netCDF readers.""" 

from __future__ import unicode_literals, division, print_function 

 

import os.path 

 

from monty.dev import requires, deprecated 

from monty.collections import AttrDict 

from monty.functools import lazy_property 

from pymatgen.core.units import ArrayWithUnit 

from pymatgen.core.structure import Structure 

 

import logging 

logger = logging.getLogger(__name__) 

 

 

__author__ = "Matteo Giantomassi" 

__copyright__ = "Copyright 2013, The Materials Project" 

__version__ = "0.1" 

__maintainer__ = "Matteo Giantomassi" 

__email__ = "gmatteo at gmail.com" 

__status__ = "Development" 

__date__ = "$Feb 21, 2013M$" 

 

__all__ = [ 

"as_ncreader", 

"as_etsfreader", 

"NetcdfReader", 

"ETSF_Reader", 

"structure_from_ncdata", 

] 

 

try: 

import netCDF4 

except ImportError: 

netCDF4 = None 

 

 

def _asreader(file, cls): 

closeit = False 

if not isinstance(file, cls): 

file, closeit = cls(file), True 

return file, closeit 

 

 

def as_ncreader(file): 

""" 

Convert file into a NetcdfReader instance. 

Returns reader, closeit where closeit is set to True 

if we have to close the file before leaving the procedure. 

""" 

return _asreader(file, NetcdfReader) 

 

 

def as_etsfreader(file): 

return _asreader(file, ETSF_Reader) 

 

 

class NetcdfReaderError(Exception): 

"""Base error class for NetcdfReader""" 

 

 

class NO_DEFAULT(object): 

"""Signal that read_value should raise an Error""" 

 

 

class NetcdfReader(object): 

""" 

Wraps and extends netCDF4.Dataset. Read only mode. Supports with statements. 

 

Additional documentation available at: 

http://netcdf4-python.googlecode.com/svn/trunk/docs/netCDF4-module.html 

""" 

Error = NetcdfReaderError 

 

@requires(netCDF4 is not None, "netCDF4 must be installed to use this class") 

def __init__(self, path): 

"""Open the Netcdf file specified by path (read mode).""" 

self.path = os.path.abspath(path) 

 

try: 

self.rootgrp = netCDF4.Dataset(self.path, mode="r") 

except Exception as exc: 

raise self.Error("In file %s: %s" % (self.path, str(exc))) 

 

self.ngroups = len(list(self.walk_tree())) 

 

#self.path2group = collections.OrderedDict() 

#for children in self.walk_tree(): 

# for child in children: 

# #print(child.group, child.path) 

# self.path2group[child.path] = child.group 

 

def __enter__(self): 

"""Activated when used in the with statement.""" 

return self 

 

def __exit__(self, type, value, traceback): 

"""Activated at the end of the with statement. It automatically closes the file.""" 

self.rootgrp.close() 

 

def close(self): 

try: 

self.rootgrp.close() 

except Exception as exc: 

logger.warning("Exception %s while trying to close %s" % (exc, self.path)) 

 

#@staticmethod 

#def pathjoin(*args): 

# return "/".join(args) 

 

def walk_tree(self, top=None): 

""" 

Navigate all the groups in the file starting from top. 

If top is None, the root group is used. 

""" 

if top is None: 

top = self.rootgrp 

 

values = top.groups.values() 

yield values 

for value in top.groups.values(): 

for children in self.walk_tree(value): 

yield children 

 

def print_tree(self): 

for children in self.walk_tree(): 

for child in children: 

print(child) 

 

def read_dimvalue(self, dimname, path="/"): 

"""Returns the value of a dimension.""" 

dim = self._read_dimensions(dimname, path=path)[0] 

return len(dim) 

 

def read_varnames(self, path="/"): 

"""List of variable names stored in the group specified by path.""" 

if path == "/": 

return self.rootgrp.variables.keys() 

else: 

group = self.path2group[path] 

return group.variables.keys() 

 

def read_value(self, varname, path="/", cmode=None, default=NO_DEFAULT): 

""" 

Returns the values of variable with name varname in the group specified by path. 

 

Args: 

varname: Name of the variable 

path: path to the group. 

cmode: if cmode=="c", a complex ndarrays is constructed and returned 

(netcdf does not provide native support from complex datatype). 

default: read_value returns default if varname is not present. 

 

Returns: 

numpy array if varname represents an array, scalar otherwise. 

""" 

try: 

var = self.read_variable(varname, path=path) 

except self.Error: 

if default is NO_DEFAULT: raise 

return default 

 

if cmode is None: 

# scalar or array 

# getValue is not portable! 

try: 

return var.getValue()[0] if not var.shape else var[:] 

except IndexError: 

return var.getValue() if not var.shape else var[:] 

 

else: 

assert var.shape[-1] == 2 

if cmode == "c": 

return var[...,0] + 1j*var[...,1] 

else: 

raise ValueError("Wrong value for cmode %s" % cmode) 

 

def read_variable(self, varname, path="/"): 

"""Returns the variable with name varname in the group specified by path.""" 

return self._read_variables(varname, path=path)[0] 

 

def _read_dimensions(self, *dimnames, **kwargs): 

path = kwargs.get("path", "/") 

try: 

if path == "/": 

return [self.rootgrp.dimensions[dname] for dname in dimnames] 

else: 

group = self.path2group[path] 

return [group.dimensions[dname] for dname in dimnames] 

 

except KeyError: 

raise self.Error("In file %s:\ndimnames %s, kwargs %s" % (self.path, dimnames, kwargs)) 

 

def _read_variables(self, *varnames, **kwargs): 

path = kwargs.get("path", "/") 

try: 

if path == "/": 

return [self.rootgrp.variables[vname] for vname in varnames] 

else: 

group = self.path2group[path] 

return [group.variables[vname] for vname in varnames] 

 

except KeyError: 

raise self.Error("In file %s:\nvarnames %s, kwargs %s" % (self.path, varnames, kwargs)) 

 

def read_keys(self, keys, dict_cls=AttrDict, path="/"): 

""" 

Read a list of variables/dimensions from file. If a key is not present the corresponding 

entry in the output dictionary is set to None. 

""" 

od = dict_cls() 

for k in keys: 

try: 

# Try to read a variable. 

od[k] = self.read_value(k, path=path) 

except self.Error: 

try: 

# Try to read a dimension. 

od[k] = self.read_dimvalue(k, path=path) 

except self.Error: 

od[k] = None 

 

return od 

 

 

class ETSF_Reader(NetcdfReader): 

""" 

This object reads data from a file written according to the ETSF-IO specifications. 

 

We assume that the netcdf file contains at least the crystallographic section. 

""" 

@lazy_property 

def chemical_symbols(self): 

"""Chemical symbols char [number of atom species][symbol length].""" 

charr = self.read_value("chemical_symbols") 

symbols = [] 

for v in charr: 

symbols.append("".join(c for c in v)) 

 

#symbols = ["".join(str(c)) for symb in symbols for c in symb] 

#symbols = [s.decode("ascii") for s in symbols] 

#chemical_symbols = [str("".join(s)) for s in symbols] 

#print(symbols) 

return symbols 

 

def typeidx_from_symbol(self, symbol): 

"""Returns the type index from the chemical symbol. Note python convention.""" 

return self.chemical_symbols.index(symbol) 

 

def read_structure(self, cls=Structure): 

"""Returns the crystalline structure.""" 

if self.ngroups != 1: 

raise NotImplementedError("In file %s: ngroups != 1" % self.path) 

 

return structure_from_ncdata(self, cls=cls) 

 

 

def structure_from_ncdata(ncdata, site_properties=None, cls=Structure): 

""" 

Reads and returns a pymatgen structure from a NetCDF file 

containing crystallographic data in the ETSF-IO format. 

 

Args: 

ncdata: filename or NetcdfReader instance. 

site_properties: Dictionary with site properties. 

cls: The Structure class to instanciate. 

""" 

ncdata, closeit = as_ncreader(ncdata) 

 

# TODO check whether atomic units are used 

lattice = ArrayWithUnit(ncdata.read_value("primitive_vectors"), "bohr").to("ang") 

 

red_coords = ncdata.read_value("reduced_atom_positions") 

natom = len(red_coords) 

 

znucl_type = ncdata.read_value("atomic_numbers") 

 

# type_atom[0:natom] --> index Between 1 and number of atom species 

type_atom = ncdata.read_value("atom_species") 

 

# Fortran to C index and float --> int conversion. 

species = natom * [None] 

for atom in range(natom): 

type_idx = type_atom[atom] - 1 

species[atom] = int(znucl_type[type_idx]) 

 

d = {} 

if site_properties is not None: 

for prop in site_properties: 

d[property] = ncdata.read_value(prop) 

 

structure = cls(lattice, species, red_coords, site_properties=d) 

 

# Quick and dirty hack. 

# I need an abipy structure since I need to_abivars and other methods. 

try: 

from abipy.core.structure import Structure as AbipyStructure 

structure.__class__ = AbipyStructure 

except ImportError: 

pass 

 

if closeit: 

ncdata.close() 

 

return structure