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

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

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3139 - (show annotations)
Thu Sep 2 06:37:43 2010 UTC (9 years, 1 month ago) by caltinay
File MIME type: text/x-python
File size: 12347 byte(s)
Fixed a typo.

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._myrank=getMPIRankWorld()
86 self._exportformats=set(formats)
87 self._restartprefix=restart_prefix
88 self._workdir=work_dir
89 util.mkDir(self._workdir)
90 if self.VISIT in self._exportformats:
91 simFile=os.path.join(self._workdir, "escriptsim.sim2")
92 if not self.__initVisit(simFile, "Escript simulation"):
93 print "Warning: Could not initialize VisIt interface"
94 self._exportformats.remove(self.VISIT)
95 # find all restart directories
96 restart_folders = []
97 for f in os.listdir(self._workdir):
98 if f.startswith(self._restartprefix):
99 restart_folders.append(f)
100 # remove unneeded restart directories
101 if len(restart_folders)>0:
102 restart_folders.sort()
103 if do_restart:
104 self._restartdir=restart_folders[-1]
105 print "Restart from "+os.path.join(self._workdir, self._restartdir)
106 for f in restart_folders[:-1]:
107 self.__removeDirectory(f)
108 self.__loadState()
109 else:
110 for f in restart_folders:
111 self.__removeDirectory(f)
112
113 def addData(self, **data):
114 """
115 Adds 'escript.Data' objects and other data to be exported to this
116 manager.
117
118 :note: This method does not make copies of Data objects so
119 any modifications will be carried over until export() is called.
120 """
121 # if this is the first addition after a restart, clear data first
122 if self._restartdir != None:
123 self.__clearData()
124
125 for name,var in data.items():
126 if hasattr(var, "getDomain"):
127 if self._domain==None:
128 self._domain=var.getDomain()
129 elif self._domain != var.getDomain():
130 raise ValueError, "addData: Data must be on the same domain!"
131 self._data[name]=var
132 else:
133 self._stamp[name]=var
134
135 def hasData(self):
136 """
137 Returns True if the manager holds data for restart
138 """
139 return self._restartdir != None
140
141 def getDomain(self):
142 """
143 Returns the domain as recovered from restart files.
144 """
145 if not self.hasData():
146 raise ValueError, "No restart data available"
147 return self._domain
148
149 def getValue(self, value_name):
150 """
151 Returns an 'escript.Data' object or other value that has been loaded
152 from restart files.
153 """
154 if not self.hasData():
155 raise ValueError, "No restart data available"
156
157 if value_name in self._stamp:
158 return self._stamp[value_name]
159
160 ff=self.__getDumpFilename(value_name, self._restartdir)
161 var = load(ff, self._domain)
162 #print "Value %s recovered from %s."%(value_name, ff)
163 return var
164
165 def getCycle(self):
166 """
167 Returns the export cycle (=number of times export() has been called)
168 """
169 return self._N
170
171 def setTime(self, time):
172 """
173 Sets the simulation timestamp.
174 """
175 self._time = time
176
177 def setMeshLabels(self, x, y, z=""):
178 """
179 Sets labels for the mesh axes. These are currently only used by the
180 Silo exporter.
181 """
182 self._meshlabels=[x,y,z]
183
184 def setMeshUnits(self, x, y, z=""):
185 """
186 Sets units for the mesh axes. These are currently only used by the
187 Silo exporter.
188 """
189 self._meshunits=[x,y,z]
190
191 def setMetadataSchemaString(self, schema, metadata=""):
192 """
193 Sets metadata namespaces and the corresponding metadata.
194 Only used for the VTK file format at the moment.
195
196 :param schema: A dictionary that maps namespace prefixes to namespace
197 names, e.g. {'gml':'http://www.opengis.net/gml'}
198 :param metadata: The actual metadata string which will be enclosed in
199 '<MetaData>' tags.
200 """
201 self._metadata="<MetaData>"+metadata+"</MetaData>"
202 ss=""
203 for i,p in schema.items():
204 ss="%s xmlns:%s=\"%s\""%(ss, i, p)
205 self._md_schema=ss.strip()
206
207 def export(self):
208 """
209 Executes the actual data export. Depending on the formats parameter
210 used in the constructor all data added by addData() is written to disk
211 (RESTART,SILO,VTK) or made available through the VisIt simulation
212 interface (VISIT).
213 """
214
215 if self._domain == None:
216 return
217
218 idata = self.__interpolateData()
219 self._N += 1
220 nameprefix=os.path.join(self._workdir, "dataset%04d"%(self._N))
221
222 from esys.weipa.weipacpp import EscriptDataset
223 ds=EscriptDataset()
224 ds.setDomain(self._domain)
225 ds.setCycleAndTime(self._N, self._time)
226 ds.setMetadataSchemaString(self._md_schema, self._metadata)
227 ds.setMeshLabels(self._meshlabels[0], self._meshlabels[1], self._meshlabels[2])
228 ds.setMeshUnits(self._meshunits[0], self._meshunits[1], self._meshunits[2])
229 for n,d in idata.items():
230 ds.addData(d, n, "") #FIXME: data units are not supported yet
231
232 for f in self._exportformats:
233 if f == self.SILO:
234 filename=nameprefix+".silo"
235 ds.saveSilo(filename)
236 elif f == self.VTK:
237 filename=nameprefix+".vtu"
238 ds.saveVTK(filename)
239 elif f == self.VISIT:
240 from esys.weipa.weipacpp import visitPublishData
241 visitPublishData(ds)
242 elif f == self.RESTART:
243 self.__saveState()
244 else:
245 raise ValueError, "export: Unknown export format "+f
246
247 self.__clearData()
248
249 def __clearData(self):
250 #print "Clearing all data"
251 self._restartdir = None
252 self._domain = None
253 self._stamp = {}
254 self._data = {}
255
256 def __getStampFilename(self, dir_name):
257 return os.path.join(self._workdir, dir_name, "stamp.%d"%self._myrank)
258
259 def __getDumpFilename(self, data_name, dir_name):
260 return os.path.join(self._workdir, dir_name, "%s.nc"%data_name)
261
262 def __initVisit(self, simFile, comment=""):
263 """
264 Initialises the VisIt interface if available.
265
266 :param simFile: Name of the sim file to be generated which can be
267 loaded into a VisIt client
268 :param comment: A short description of this simulation
269 """
270 from esys.weipa.weipacpp import visitInitialize
271 return visitInitialize(simFile, comment)
272
273 def __interpolateData(self):
274 # (Reduced)Solution is not directly supported so interpolate to a
275 # different function space
276 from esys.escript import Solution, ReducedSolution
277 from esys.escript import ContinuousFunction, ReducedContinuousFunction
278 from esys.escript.util import interpolate
279 new_data={}
280 for n,d in self._data.items():
281 if not d.isEmpty():
282 fs=d.getFunctionSpace()
283 domain=fs.getDomain()
284 if fs == Solution(domain):
285 new_data[n]=interpolate(d, ContinuousFunction(domain))
286 elif fs == ReducedSolution(domain):
287 new_data[n]=interpolate(d, ReducedContinuousFunction(domain))
288 else:
289 new_data[n]=d
290 return new_data
291
292 def __loadState(self):
293 stamp_file=self.__getStampFilename(self._restartdir)
294 try:
295 self._stamp = cPickle.load(open(stamp_file, "rb"))
296 self._N = int(self._restartdir[len(self._restartprefix)+1:])
297 except:
298 raise IOError, "Could not load stamp file "+stamp_file
299 # load domain
300 ff=self.__getDumpFilename("_domain",self._restartdir)
301 modname=self._stamp['__domainmodule']
302 clsname=self._stamp['__domainclass']
303 try:
304 domclass=__import__(modname, fromlist=[clsname])
305 self._domain = domclass.LoadMesh(ff)
306 except:
307 raise ImportError, "Unable to load %s using %s.%s!"%(ff, modname, clsname)
308
309 def __saveState(self):
310 restartdir = "%s_%04d"%(self._restartprefix, self._N)
311 util.mkDir(os.path.join(self._workdir, restartdir))
312 stamp_file=self.__getStampFilename(restartdir)
313 self._stamp['__domainmodule']=self._domain.__module__
314 self._stamp['__domainclass']=type(self._domain).__name__
315 cPickle.dump(self._stamp, open(stamp_file, "wb"))
316 ff=self.__getDumpFilename("_domain", restartdir)
317 self._domain.dump(ff)
318 for name, var in self._data.items():
319 ff=self.__getDumpFilename(name, restartdir)
320 var.dump(ff)
321 print "Restart files saved in "+os.path.join(self._workdir, restartdir)
322 # keep only one restart directory
323 old_restartdir = "%s_%04d"%(self._restartprefix, self._N-1)
324 self.__removeDirectory(os.path.join(self._workdir, old_restartdir))
325
326 def __removeDirectory(self, path):
327 if self._myrank==0 and os.path.isdir(path):
328 shutil.rmtree(path, True)
329 #print "Removed restart directory %s."%path
330 MPIBarrierWorld()
331

  ViewVC Help
Powered by ViewVC 1.1.26