/[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 3109 - (show annotations)
Thu Aug 26 04:33:18 2010 UTC (9 years ago) by caltinay
File MIME type: text/x-python
File size: 11556 byte(s)
Added accessor for the simulation cycle and log which directory was used to
restart.

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

  ViewVC Help
Powered by ViewVC 1.1.26