/[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 3106 - (show annotations)
Wed Aug 25 08:15:38 2010 UTC (8 years, 10 months ago) by caltinay
File MIME type: text/x-python
File size: 11329 byte(s)
Fixed export when 'work_dir' is set to sth other than the current dir.

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

  ViewVC Help
Powered by ViewVC 1.1.26