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

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

Parent Directory Parent Directory | Revision Log Revision Log | View Patch Patch

revision 3095 by gross, Wed Aug 4 03:26:50 2010 UTC revision 3096 by caltinay, Thu Aug 19 01:40:52 2010 UTC
# Line 1  Line 1 
 # -*- coding: utf-8 -*-  
1  ########################################################  ########################################################
2  #  #
3  # Copyright (c) 2003-2010 by University of Queensland  # Copyright (c) 2003-2010 by University of Queensland
# Line 18  Primary Business: Queensland, Australia" Line 17  Primary Business: Queensland, Australia"
17  __license__="""Licensed under the Open Software License version 3.0  __license__="""Licensed under the Open Software License version 3.0
18  http://www.opensource.org/licenses/osl-3.0.php"""  http://www.opensource.org/licenses/osl-3.0.php"""
19  __url__="https://launchpad.net/escript-finley"  __url__="https://launchpad.net/escript-finley"
20    __author__="Lutz Gross, Cihan Altinay"
21    
22  """  """
23  a simple restart manager (still under testing)  an escript data import and export manager (still under development)
24    
25  :var __author__: name of author  :var __author__: name of authors
26  :var __copyright__: copyrights  :var __copyright__: copyrights
27  :var __license__: licence agreement  :var __license__: licence agreement
28  :var __url__: url entry point on documentation  :var __url__: url entry point to documentation
 :var __version__: version  
 :var __date__: date of the version  
29  """  """
30    
 __author__="Lutz Gross, l.gross@uq.edu.au"  
   
 import util  
31  import cPickle  import cPickle
32  import os  import os
33    import shutil
34    import util
35  from esys.escript import getMPIRankWorld, MPIBarrierWorld, load  from esys.escript import getMPIRankWorld, MPIBarrierWorld, load
36    
37  class RestartManager(object):  class DataManager(object):
38     """      """
39     this is a simple restart manager      Escript data import/export manager.
40    
41     usage      Example::
42        
43     rm=RestartManager(prefix="myrestart")          dm=DataManager(formats=[DataManager.RESTART,DataManager.VTK])
44     if rm.hasData():          if dm.hasData():
45          dom=rm.getDomain()              dom = dm.getDomain()
46          t=rm.getStamp("t")              time = dm.getValue("time")
47          dt=rm.getStamp("stemp")              dt = dm.getValue("dt")
48          T=rm.getValue("T")              T = dm.getValue("T")
49          u=rm.getValues("d")              u = dm.getValue("u")
50     else:          else:
51          T=...              T = ...
52          u=...              u = ...
53                dm.addData(time=time,dt=dt,T=T,u=u) # add data and variables
54     rm.setStamp(t=t,step=dt) # this creates a new restart directory          dm.setTime(time)                    # set the simulation timestamp
55     rm.checkIn(T=T,d=u)      # adds values to the restart data base          dm.export()                         # write out data
56     """      """
57     def __init__(self, new=False, prefix="restart", work_dir=".", myrank=getMPIRankWorld()):  
58         """      RESTART, SILO, VISIT, VTK = range(4)
59         initializes the restart manager. If new is True a new restart directory is created  
60         otherwise the last available directory is used for restart.      def __init__(self, formats=[RESTART], work_dir=".", restart_prefix="restart", do_restart=True):
61                  """
62         :param new: fresh start. Otherwise the last available restart file will be used.          Initialises the data manager. If do_restart is True and a restart
63         :type new: `bool`          directory is found the contained data is loaded (hasData() returns True)
64         :param prefix: directory name prefix          otherwise restart directories are removed (hasData() returns False).
65         :type prefix: `str`          Values are only written to disk when export() is called.
66         :param work_dir: directory where restart files are created  
67         :type prefix: `str`          :param formats: A list of export file formats to use. Allowed values
68         """                          are RESTART, SILO, VISIT, VTK.
69         self.PREFIX=prefix          :param work_dir: top-level directory where files are exported to
70         self.MYRANK=myrank          :param restart_prefix: prefix for restart directories. Will be used to
71         self.WORK_DIR=work_dir                                 load restart files (if do_restart is True) and
72         self.NEW_RESTART_DIR=None                                 store new restart files (if RESTART is used)
73         util.mkDir(self.WORK_DIR)          :param do_restart: whether to attempt to load restart files
74                  """
75         restart_files=[]          self._metadata=""
76                  self._md_schema=""
77         for f in os.listdir(self.WORK_DIR):          self._data={}
78             if f.startswith(self.PREFIX): restart_files.append(f)          self._domain=None
79         restart_files.sort()          self._stamp={}
80                      self._time=0.
81         if len(restart_files)==0 or new:          self._restartdir=None
82            for f in restart_files: self.__removeRestartDirectory(f)          self._N=-1
83            self.RESTART_DIR=None          self._myrank=getMPIRankWorld()
84            self.N=-1          self._exportformats=set(formats)
85            self.__STAMP=None          self._restartprefix=restart_prefix
86                      self._workdir=work_dir
87         else:          util.mkDir(self._workdir)
88            self.RESTART_DIR=restart_files[-1]          if self.VISIT in self._exportformats:
89            for f in restart_files[:-1]: self.__removeRestartDirectory(f)              simFile=os.path.join(self._workdir, "escriptsim.sim2")
90            print "Restart from file %s."%self.RESTART_DIR              if not self.__initVisit(simFile, "Escript simulation"):
91            self.N=int(self.RESTART_DIR[len(self.PREFIX)+1:])                  print "Warning: Could not initialize VisIt interface"
92            self.__STAMP=cPickle.load(open(self.__getStampFileName(self.RESTART_DIR),"rb"))                  self._exportformats.remove(self.VISIT)
93         self.DOMAIN = None          # find all restart directories
94                  restart_folders = []
95     def __removeRestartDirectory(self, dir_name):          for f in os.listdir(self._workdir):
96         if self.MYRANK==0 and os.path.isdir(dir_name):              if f.startswith(self._restartprefix):
97             for root, dirs, files in os.walk(dir_name, topdown=False):                  restart_folders.append(f)
98                 for name in files: os.remove(os.path.join(root,name))          # remove unneeded restart directories
99                 for name in dirs: os.remove(os.path.join(root,name))          if len(restart_folders)>0:
100             os.rmdir(dir_name)              restart_folders.sort()
101             print "Restart files %s have been removed."%dir_name              if do_restart:
102         MPIBarrierWorld()                  self._restartdir=restart_folders[-1]
103                    for f in restart_folders[:-1]:
104     def __getStampFileName(self, dir_name):                      self.__removeDirectory(f)
105           return os.path.join(self.WORK_DIR,dir_name,"stamp.%d"%self.MYRANK)                  self.__loadState()
106     def __getValueFileName(self, value_name, dir_name):              else:
107           return os.path.join(self.WORK_DIR,dir_name,"%s.nc"%value_name)                  for f in restart_folders:
108                        self.__removeDirectory(f)
109     def hasData(self):  
110         """      def addData(self, **data):
111         returns True if manager holds data for restart          """
112         """          Adds 'escript.Data' objects and other data to be exported to this
113         return not self.RESTART_DIR == None          manager.
114    
115     def setStamp(self, **values):          :note: This method does not make copies of Data objects so
116         """                 any modifications will be carried over until export() is called.
117         sets a time stap defined by a set of pickable values (e.g. float, int, str)          """
118         A new restart directory is created and the directory before the last is removed.          # if this is the first addition after a restart, clear data first
119         """          if self._restartdir != None:
120         # complete last saver:              self.__clearData()
121         if not self.NEW_RESTART_DIR == None:  
122            if self.RESTART_DIR!= None: self.__removeRestartDirectory(self.RESTART_DIR)          for name,var in data.items():
123            self.RESTART_DIR=self.NEW_RESTART_DIR              if hasattr(var, "getDomain"):
124            self.DOMAIN=self.NEW_DOMAIN                  if self._domain==None:
125            self.N+=1                      self._domain=var.getDomain()
126         # create a new restart directory                  elif self._domain != var.getDomain():
127         self.NEW_RESTART_DIR="%s_%d"%(self.PREFIX,self.N+1)                      raise ValueError, "addData: Data must be on the same domain!"
128         self.NEW_DOMAIN = None                  self._data[name]=var
129         util.mkDir(self.NEW_RESTART_DIR)              else:
130         cPickle.dump(values,open(self.__getStampFileName(self.NEW_RESTART_DIR),"wb"))                  self._stamp[name]=var
131    
132     def getStamp(self,value_name):      def hasData(self):
133         """          """
134         returns the value of stamp paramater with name value_name          Returns True if the manager holds data for restart
135                  """
136         :param value_name: requested value          return self._restartdir != None
137         :type value_name: `str`  
138         :return: requested value from latest restart file.      def getDomain(self):
139         """          """
140         if self.__STAMP == None:          Returns the domain as recovered from restart files.
141            raise ValueError,"no restart data available."          """
142         if not self.__STAMP.has_key(value_name):          if not self.hasData():
143             raise ValueError,"no stamp parameter %s."%value_name              raise ValueError, "No restart data available"
144         return self.__STAMP[value_name]          return self._domain
145    
146     def checkIn(self,**values):      def getValue(self, value_name):
147         """          """
148         writes 'escript.Data' objects to restart dierctory.          Returns an 'escript.Data' object or other value that has been loaded
149         """          from restart files.
150         if self.NEW_RESTART_DIR == None:          """
151             raise ValueError,"No time stamp set."          if not self.hasData():
152         for i,v in values.items():              raise ValueError, "No restart data available"
153           self.checkinDomain(v.getDomain())  
154           ff=self.__getValueFileName(i, self.NEW_RESTART_DIR)          if value_name in self._stamp:
155           v.dump(ff)              return self._stamp[value_name]
156           print "%s dumped to file %s."%(i,ff)  
157            ff=self.__getDumpFilename(value_name, self._restartdir)
158     def getValue(self,value_name):          var = load(ff, self._domain)
159         """          #print "Value %s recovered from %s."%(value_name, ff)
160         returns the value of value_name in the latest restart file          return var
161          
162         :param value_name: requested value      def setTime(self, time):
163         :type value_name: `str`          """
164         :return: requested value from latest restart file.          Sets the simulation timestamp.
165         """          """
166         if self.RESTART_DIR == None:          self._time = time
167            raise ValueError,"no restart data available."  
168         domain=self.getDomain()      def setMetadataSchemaString(self, schema, metadata=""):
169         if domain == None:          """
170        raise ValueError,"no domain found for restart."          Do not use this yet.
171         ff=self.__getValueFileName(value_name, self.RESTART_DIR)          """
172         v=load(ff, domain)          self._metadata="<MetaData>"+metadata+"</MetaData>"
173         print "Value %s recovered from %s."%(value_name, ff)          schema=""
174         return v          for i,p in self._md_schema.items():
175     def checkinDomain(self, domain):              schema="%s xmlns:%s=\"%s\""%(schema, i, p)
176         """          self._md_schema=schema.strip()
177         writes the domain to current restart directory  
178         """      def export(self):
179         if self.NEW_RESTART_DIR == None:          """
180             raise ValueError,"No time stamp set."          Executes the actual data export. Depending on the formats parameter
181         if self.NEW_DOMAIN == None:          used in the constructor all data added by addData() is written to disk
182             ff=self.__getValueFileName("__mesh_finley",self.NEW_RESTART_DIR) # finley extension added to mark domain type          (RESTART,SILO,VTK) or made available through the VisIt simulation
183             domain.dump(ff)          interface (VISIT).
184             print "Domain dumped to file %s."%(ff,)          """
185             self.NEW_DOMAIN=domain  
186            if self._domain == None:
187     def getDomain(self):              return
188         """  
189         recovers the domain from the current restart directory          idata = self.__interpolateData()
190         """  
191         if self.RESTART_DIR == None:          self._N += 1
192            raise ValueError,"no restart data available."          nameprefix=os.path.join(self._workdir, "dataset%04d"%(self._N))
193         if self.DOMAIN == None:  
194        for f in os.listdir(os.path.join(self.WORK_DIR, self.RESTART_DIR)):          for f in self._exportformats:
195                 if f.startswith("__mesh_finley"):              if f == self.SILO:
196                   ff=self.__getValueFileName("__mesh_finley",self.RESTART_DIR)                  from esys.weipa.weipacpp import _saveSilo
197                   import esys.finley                  filename=nameprefix+".silo"
198                   self.DOMAIN=esys.finley.LoadMesh(ff)                  _saveSilo(filename, self._N, self._time, self._domain, idata)
199                   print "Domain recovered from file %s."%ff              elif f == self.VTK:
200         return self.DOMAIN                  from esys.weipa.weipacpp import _saveVTK
201                    filename=nameprefix+".vtu"
202                    _saveVTK(filename, self._N, self._time, self._domain,
203                            idata, self._metadata, self._md_schema)
204                elif f == self.VISIT:
205                    from esys.weipa.weipacpp import _visitPublishData
206                    _visitPublishData(self._N, self._time, self._domain, idata)
207                elif f == self.RESTART:
208                    self.__saveState()
209                else:
210                    raise ValueError, "export: Unknown export format "+f
211    
212            self.__clearData()
213    
214        def __clearData(self):
215            #print "Clearing all data"
216            self._restartdir = None
217            self._domain = None
218            self._stamp = {}
219            self._data = {}
220    
221        def __getStampFilename(self, dir_name):
222            return os.path.join(self._workdir, dir_name, "stamp.%d"%self._myrank)
223    
224        def __getDumpFilename(self, data_name, dir_name):
225            return os.path.join(self._workdir, dir_name, "%s.nc"%data_name)
226    
227        def __initVisit(self, simFile, comment=""):
228            """
229            Initialises the VisIt interface if available.
230    
231            :param simFile: Name of the sim file to be generated which can be
232                            loaded into a VisIt client
233            :param comment: A short description of this simulation
234            """
235            from esys.weipa.weipacpp import _visitInitialize
236            return _visitInitialize(simFile, comment)
237    
238        def __interpolateData(self):
239            # (Reduced)Solution is not directly supported so interpolate to
240            # different function space
241            from esys.escript import Solution, ReducedSolution
242            from esys.escript import ContinuousFunction, ReducedContinuousFunction
243            from esys.escript.util import interpolate
244            new_data={}
245            for n,d in self._data.items():
246                if not d.isEmpty():
247                    fs=d.getFunctionSpace()
248                    domain=fs.getDomain()
249                    if fs == Solution(domain):
250                        new_data[n]=interpolate(d, ContinuousFunction(domain))
251                    elif fs == ReducedSolution(domain):
252                        new_data[n]=interpolate(d, ReducedContinuousFunction(domain))
253                    else:
254                        new_data[n]=d
255            return new_data
256    
257        def __loadState(self):
258            stamp_file=self.__getStampFilename(self._restartdir)
259            try:
260                self._stamp = cPickle.load(open(stamp_file, "rb"))
261                self._N = int(self._restartdir[len(self._restartprefix)+1:])
262            except:
263                raise IOError, "Could not load stamp file "+stamp_file
264            # load domain
265            path=os.path.join(self._workdir, self._restartdir)
266            for f in os.listdir(path):
267                if f.startswith("__mesh_finley"):
268                    ff=self.__getDumpFilename("__mesh_finley",self._restartdir)
269                    import esys.finley
270                    self._domain = esys.finley.LoadMesh(ff)
271                    #print "Domain recovered from file %s."%ff
272                    break
273    
274        def __saveState(self):
275            restartdir = "%s_%04d"%(self._restartprefix, self._N)
276            util.mkDir(restartdir)
277            stamp_file=self.__getStampFilename(restartdir)
278            cPickle.dump(self._stamp, open(stamp_file, "wb"))
279            ff=self.__getDumpFilename("__mesh_finley", restartdir)
280            self._domain.dump(ff)
281            #print "Domain dumped to %s."%(ff)
282            for name, var in self._data.items():
283                ff=self.__getDumpFilename(name, restartdir)
284                var.dump(ff)
285                #print "%s dumped to %s."%(name, ff)
286            # keep only one restart directory
287            old_restartdir = "%s_%04d"%(self._restartprefix, self._N-1)
288            self.__removeDirectory(old_restartdir)
289    
290        def __removeDirectory(self, path):
291            if self._myrank==0 and os.path.isdir(path):
292                shutil.rmtree(path, True)
293                #print "Removed restart directory %s."%path
294            MPIBarrierWorld()
295    

Legend:
Removed from v.3095  
changed lines
  Added in v.3096

  ViewVC Help
Powered by ViewVC 1.1.26