/[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 3132 - (show annotations)
Thu Sep 2 01:48:31 2010 UTC (9 years ago) by caltinay
File MIME type: text/x-python
File size: 11807 byte(s)
Removed unused variable.

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 self._N += 1
204 nameprefix=os.path.join(self._workdir, "dataset%04d"%(self._N))
205
206 from esys.weipa.weipacpp import EscriptDataset
207 ds=EscriptDataset()
208 ds.setDomain(self._domain)
209 ds.setCycleAndTime(self._N, self._time)
210 ds.setMetadataSchemaString(self._metadata, self._md_schema)
211 #ds.setMeshLabels("x","y","z") # FIXME
212 #ds.setMeshUnits("cm","cm","cm") # FIXME
213 for n,d in idata.items():
214 ds.addData(d, n, "") #FIXME: data units are not supported yet
215
216 for f in self._exportformats:
217 if f == self.SILO:
218 filename=nameprefix+".silo"
219 ds.saveSilo(filename)
220 elif f == self.VTK:
221 filename=nameprefix+".vtu"
222 ds.saveVTK(filename)
223 elif f == self.VISIT:
224 from esys.weipa.weipacpp import visitPublishData
225 visitPublishData(ds)
226 elif f == self.RESTART:
227 self.__saveState()
228 else:
229 raise ValueError, "export: Unknown export format "+f
230
231 self.__clearData()
232
233 def __clearData(self):
234 #print "Clearing all data"
235 self._restartdir = None
236 self._domain = None
237 self._stamp = {}
238 self._data = {}
239
240 def __getStampFilename(self, dir_name):
241 return os.path.join(self._workdir, dir_name, "stamp.%d"%self._myrank)
242
243 def __getDumpFilename(self, data_name, dir_name):
244 return os.path.join(self._workdir, dir_name, "%s.nc"%data_name)
245
246 def __initVisit(self, simFile, comment=""):
247 """
248 Initialises the VisIt interface if available.
249
250 :param simFile: Name of the sim file to be generated which can be
251 loaded into a VisIt client
252 :param comment: A short description of this simulation
253 """
254 from esys.weipa.weipacpp import visitInitialize
255 return visitInitialize(simFile, comment)
256
257 def __interpolateData(self):
258 # (Reduced)Solution is not directly supported so interpolate to a
259 # different function space
260 from esys.escript import Solution, ReducedSolution
261 from esys.escript import ContinuousFunction, ReducedContinuousFunction
262 from esys.escript.util import interpolate
263 new_data={}
264 for n,d in self._data.items():
265 if not d.isEmpty():
266 fs=d.getFunctionSpace()
267 domain=fs.getDomain()
268 if fs == Solution(domain):
269 new_data[n]=interpolate(d, ContinuousFunction(domain))
270 elif fs == ReducedSolution(domain):
271 new_data[n]=interpolate(d, ReducedContinuousFunction(domain))
272 else:
273 new_data[n]=d
274 return new_data
275
276 def __loadState(self):
277 stamp_file=self.__getStampFilename(self._restartdir)
278 try:
279 self._stamp = cPickle.load(open(stamp_file, "rb"))
280 self._N = int(self._restartdir[len(self._restartprefix)+1:])
281 except:
282 raise IOError, "Could not load stamp file "+stamp_file
283 # load domain
284 ff=self.__getDumpFilename("_domain",self._restartdir)
285 modname=self._stamp['__domainmodule']
286 clsname=self._stamp['__domainclass']
287 try:
288 domclass=__import__(modname, fromlist=[clsname])
289 self._domain = domclass.LoadMesh(ff)
290 except:
291 raise ImportError, "Unable to load %s using %s.%s!"%(ff, modname, clsname)
292
293 def __saveState(self):
294 restartdir = "%s_%04d"%(self._restartprefix, self._N)
295 util.mkDir(os.path.join(self._workdir, restartdir))
296 stamp_file=self.__getStampFilename(restartdir)
297 self._stamp['__domainmodule']=self._domain.__module__
298 self._stamp['__domainclass']=type(self._domain).__name__
299 cPickle.dump(self._stamp, open(stamp_file, "wb"))
300 ff=self.__getDumpFilename("_domain", restartdir)
301 self._domain.dump(ff)
302 for name, var in self._data.items():
303 ff=self.__getDumpFilename(name, restartdir)
304 var.dump(ff)
305 print "Restart files saved in "+os.path.join(self._workdir, restartdir)
306 # keep only one restart directory
307 old_restartdir = "%s_%04d"%(self._restartprefix, self._N-1)
308 self.__removeDirectory(os.path.join(self._workdir, old_restartdir))
309
310 def __removeDirectory(self, path):
311 if self._myrank==0 and os.path.isdir(path):
312 shutil.rmtree(path, True)
313 #print "Removed restart directory %s."%path
314 MPIBarrierWorld()
315

  ViewVC Help
Powered by ViewVC 1.1.26