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

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

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3344 - (show annotations)
Thu Nov 11 23:26:52 2010 UTC (8 years, 2 months ago) by caltinay
File MIME type: text/x-python
File size: 12447 byte(s)
Phew!
-escript, finley, and dudley now uses weipa's saveVTK implementation
-moved tests from finley to weipa accordingly; dudley still to do
-rebaselined all test files
-fixed a few issues in weipa.saveVTK, e.g. saving metadata without schema
-added a deprecation warning to esys.escript.util.saveVTK
-todo: change doco, tests and other places to use weipa.saveVTK


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

  ViewVC Help
Powered by ViewVC 1.1.26