/[escript]/trunk/escriptcore/py_src/datamanager.py
ViewVC logotype

Contents of /trunk/escriptcore/py_src/datamanager.py

Parent Directory Parent Directory | Revision Log Revision Log


Revision 4508 - (show annotations)
Wed Jul 24 04:23:22 2013 UTC (6 years ago) by jfenwick
File MIME type: text/x-python
File size: 12523 byte(s)
Moving to escriptcore
1 ##############################################################################
2 #
3 # Copyright (c) 2003-2013 by University of Queensland
4 # http://www.uq.edu.au
5 #
6 # Primary Business: Queensland, Australia
7 # Licensed under the Open Software License version 3.0
8 # http://www.opensource.org/licenses/osl-3.0.php
9 #
10 # Development until 2012 by Earth Systems Science Computational Center (ESSCC)
11 # Development since 2012 by School of Earth Sciences
12 #
13 ##############################################################################
14
15 __copyright__="""Copyright (c) 2003-2013 by University of Queensland
16 http://www.uq.edu.au
17 Primary Business: Queensland, Australia"""
18 __license__="""Licensed under the Open Software License version 3.0
19 http://www.opensource.org/licenses/osl-3.0.php"""
20 __url__="https://launchpad.net/escript-finley"
21 __author__="Lutz Gross, Cihan Altinay"
22
23 """
24 an escript data import and export manager (still under development)
25
26 :var __author__: name of authors
27 :var __copyright__: copyrights
28 :var __license__: licence agreement
29 :var __url__: url entry point to documentation
30 """
31
32 import pickle
33 import os
34 import shutil
35 from . import util
36 from . import escriptcpp as esc
37
38 class DataManager(object):
39 """
40 Escript data import/export manager.
41
42 Example::
43
44 dm=DataManager(formats=[DataManager.RESTART,DataManager.VTK])
45 if dm.hasData():
46 dom = dm.getDomain()
47 time = dm.getValue("time")
48 dt = dm.getValue("dt")
49 T = dm.getValue("T")
50 u = dm.getValue("u")
51 else:
52 T = ...
53 u = ...
54 dm.addData(time=time,dt=dt,T=T,u=u) # add data and variables
55 dm.setTime(time) # set the simulation timestamp
56 dm.export() # write out data
57 """
58
59 RESTART, SILO, VISIT, VTK = list(range(4))
60
61 def __init__(self, formats=[RESTART], work_dir=".", restart_prefix="restart", do_restart=True):
62 """
63 Initialises the data manager. If do_restart is True and a restart
64 directory is found the contained data is loaded (hasData() returns True)
65 otherwise restart directories are removed (hasData() returns False).
66 Values are only written to disk when export() is called.
67
68 :param formats: A list of export file formats to use. Allowed values
69 are RESTART, SILO, VISIT, VTK.
70 :param work_dir: top-level directory where files are exported to
71 :param restart_prefix: prefix for restart directories. Will be used to
72 load restart files (if do_restart is True) and
73 store new restart files (if RESTART is used)
74 :param do_restart: whether to attempt to load restart files
75 """
76 self._metadata=""
77 self._md_schema=""
78 self._data={}
79 self._domain=None
80 self._meshlabels=["","",""]
81 self._meshunits=["","",""]
82 self._stamp={}
83 self._time=0.
84 self._restartdir=None
85 self._N=-1
86 self._checkpointfreq=1
87 self._myrank=esc.getMPIRankWorld()
88 self._exportformats=set(formats)
89 self._restartprefix=restart_prefix
90 self._workdir=work_dir
91 util.mkDir(self._workdir)
92 if self.VISIT in self._exportformats:
93 simFile=os.path.join(self._workdir, "escriptsim.sim2")
94 if not self.__initVisit(simFile, "Escript simulation"):
95 print("Warning: Could not initialize VisIt interface")
96 self._exportformats.remove(self.VISIT)
97 if self.RESTART in self._exportformats:
98 # find all restart directories
99 restart_folders = []
100 for f in os.listdir(self._workdir):
101 if f.startswith(self._restartprefix):
102 restart_folders.append(f)
103 # remove unneeded restart directories
104 if len(restart_folders)>0:
105 restart_folders.sort()
106 if do_restart:
107 self._restartdir=restart_folders[-1]
108 print(("Restart from "+os.path.join(self._workdir, self._restartdir)))
109 for f in restart_folders[:-1]:
110 self.__removeDirectory(f)
111 self.__loadState()
112 else:
113 for f in restart_folders:
114 self.__removeDirectory(f)
115
116 def addData(self, **data):
117 """
118 Adds 'escript.Data' objects and other data to be exported to this
119 manager.
120
121 :note: This method does not make copies of Data objects so
122 any modifications will be carried over until export() is called.
123 """
124 # if this is the first addition after a restart, clear data first
125 if self._restartdir != None:
126 self.__clearData()
127
128 for name,var in list(data.items()):
129 if hasattr(var, "getDomain"):
130 if self._domain==None:
131 self._domain=var.getDomain()
132 elif self._domain != var.getDomain():
133 raise ValueError("addData: Data must be on the same domain!")
134 self._data[name]=var
135 else:
136 self._stamp[name]=var
137
138 def setDomain(self, domain):
139 """
140 Sets the domain without adding data.
141 """
142 if self._domain==None:
143 self._domain = domain
144 elif self._domain != domain:
145 raise ValueError("setDomain: Domain already set!")
146
147
148 def hasData(self):
149 """
150 Returns True if the manager holds data for restart
151 """
152 return self._restartdir != None
153
154 def getDomain(self):
155 """
156 Returns the domain as recovered from restart files.
157 """
158 if not self.hasData():
159 raise ValueError("No restart data available")
160 return self._domain
161
162 def getValue(self, value_name):
163 """
164 Returns an 'escript.Data' object or other value that has been loaded
165 from restart files.
166 """
167 if not self.hasData():
168 raise ValueError("No restart data available")
169
170 if value_name in self._stamp:
171 return self._stamp[value_name]
172
173 ff=self.__getDumpFilename(value_name, self._restartdir)
174 var = esc.load(ff, self._domain)
175 #print("Value %s recovered from %s."%(value_name, ff))
176 return var
177
178 def getCycle(self):
179 """
180 Returns the export cycle (=number of times export() has been called)
181 """
182 return self._N
183
184 def setCheckpointFrequency(self, freq):
185 """
186 Sets the number of calls to export() before new restart files are
187 generated.
188 """
189 self._checkpointfreq=freq
190
191 def setTime(self, time):
192 """
193 Sets the simulation timestamp.
194 """
195 self._time = time
196
197 def setMeshLabels(self, x, y, z=""):
198 """
199 Sets labels for the mesh axes. These are currently only used by the
200 Silo exporter.
201 """
202 self._meshlabels=[x,y,z]
203
204 def setMeshUnits(self, x, y, z=""):
205 """
206 Sets units for the mesh axes. These are currently only used by the
207 Silo exporter.
208 """
209 self._meshunits=[x,y,z]
210
211 def setMetadataSchemaString(self, schema, metadata=""):
212 """
213 Sets metadata namespaces and the corresponding metadata.
214 Only used for the VTK file format at the moment.
215
216 :param schema: A dictionary that maps namespace prefixes to namespace
217 names, e.g. {'gml':'http://www.opengis.net/gml'}
218 :param metadata: The actual metadata string which will be enclosed in
219 '<MetaData>' tags.
220 """
221 self._metadata=metadata
222 ss=""
223 for i,p in list(schema.items()):
224 ss="%s xmlns:%s=\"%s\""%(ss, i, p)
225 self._md_schema=ss.strip()
226
227 def export(self):
228 """
229 Executes the actual data export. Depending on the formats parameter
230 used in the constructor all data added by addData() is written to disk
231 (RESTART,SILO,VTK) or made available through the VisIt simulation
232 interface (VISIT).
233 """
234
235 if self._domain == None:
236 print("Warning: DataManager.export() called but no domain set!")
237 return
238
239 self._N += 1
240 ds = None
241 nameprefix=os.path.join(self._workdir, "dataset.%04d"%(self._N))
242
243 for f in self._exportformats:
244 if f == self.SILO:
245 if ds == None:
246 ds=self.__createDataset()
247 ds.saveSilo(nameprefix)
248 elif f == self.VTK:
249 if ds == None:
250 ds=self.__createDataset()
251 ds.saveVTK(nameprefix)
252 elif f == self.VISIT:
253 from esys.weipa.weipacpp import visitPublishData
254 if ds == None:
255 ds=self.__createDataset()
256 visitPublishData(ds)
257 elif f == self.RESTART:
258 # only write checkpoint files with the requested frequency
259 if self._N % self._checkpointfreq==0:
260 self.__saveState()
261 else:
262 raise ValueError("export: Unknown export format "+str(f))
263
264 self.__clearData()
265
266 def __createDataset(self):
267 from esys.weipa.weipacpp import EscriptDataset
268 from esys.weipa import createDataset
269
270 ds = createDataset(self._domain, **self._data)
271 ds.setCycleAndTime(self._N, self._time)
272 ds.setMetadataSchemaString(self._md_schema, self._metadata)
273 ds.setMeshLabels(self._meshlabels[0], self._meshlabels[1], self._meshlabels[2])
274 ds.setMeshUnits(self._meshunits[0], self._meshunits[1], self._meshunits[2])
275 return ds
276
277 def __clearData(self):
278 #print("Clearing all data")
279 self._restartdir = None
280 self._domain = None
281 self._stamp = {}
282 self._data = {}
283
284 def __getStampFilename(self, dir_name):
285 return os.path.join(self._workdir, dir_name, "stamp.%d"%self._myrank)
286
287 def __getDumpFilename(self, data_name, dir_name):
288 return os.path.join(self._workdir, dir_name, "%s.nc"%data_name)
289
290 def __initVisit(self, simFile, comment=""):
291 """
292 Initialises the VisIt interface if available.
293
294 :param simFile: Name of the sim file to be generated which can be
295 loaded into a VisIt client
296 :param comment: A short description of this simulation
297 """
298 from esys.weipa.weipacpp import visitInitialize
299 return visitInitialize(simFile, comment)
300
301 def __loadState(self):
302 stamp_file=self.__getStampFilename(self._restartdir)
303 try:
304 self._stamp = pickle.load(open(stamp_file, "rb"))
305 self._N = int(self._restartdir[len(self._restartprefix)+1:])
306 except:
307 raise IOError("Could not load stamp file "+stamp_file)
308 # load domain
309 ff=self.__getDumpFilename("_domain",self._restartdir)
310 modname=self._stamp['__domainmodule']
311 clsname=self._stamp['__domainclass']
312 try:
313 domclass=__import__(modname, fromlist=[clsname])
314 self._domain = domclass.LoadMesh(ff)
315 except:
316 raise ImportError("Unable to load %s using %s.%s!"%(ff, modname, clsname))
317
318 def __saveState(self):
319 restartdir = "%s_%04d"%(self._restartprefix, self._N)
320 util.mkDir(os.path.join(self._workdir, restartdir))
321 stamp_file=self.__getStampFilename(restartdir)
322 self._stamp['__domainmodule']=self._domain.__module__
323 self._stamp['__domainclass']=type(self._domain).__name__
324 pickle.dump(self._stamp, open(stamp_file, "wb"))
325 ff=self.__getDumpFilename("_domain", restartdir)
326 self._domain.dump(ff)
327 for name, var in list(self._data.items()):
328 ff=self.__getDumpFilename(name, restartdir)
329 var.dump(ff)
330 print(("Restart files saved in "+os.path.join(self._workdir, restartdir)))
331 # keep only one restart directory
332 old_restartdir = "%s_%04d"%(self._restartprefix, self._N-self._checkpointfreq)
333 self.__removeDirectory(os.path.join(self._workdir, old_restartdir))
334
335 def __removeDirectory(self, path):
336 if self._myrank==0 and os.path.isdir(path):
337 shutil.rmtree(path, True)
338 #print("Removed restart directory %s."%path)
339 esc.MPIBarrierWorld()
340

  ViewVC Help
Powered by ViewVC 1.1.26