escript/py_src/datamanager.py

########################################################
#
# Copyright (c) 2003-2010 by University of Queensland
# Earth Systems Science Computational Center (ESSCC)
# http://www.uq.edu.au/esscc
#
# Primary Business: Queensland, Australia
# Licensed under the Open Software License version 3.0
# http://www.opensource.org/licenses/osl-3.0.php
#
########################################################

__copyright__="""Copyright (c) 2003-2010 by University of Queensland
Earth Systems Science Computational Center (ESSCC)
http://www.uq.edu.au/esscc
Primary Business: Queensland, Australia"""
__license__="""Licensed under the Open Software License version 3.0
http://www.opensource.org/licenses/osl-3.0.php"""
__url__="https://launchpad.net/escript-finley"
__author__="Lutz Gross, Cihan Altinay"

"""
an escript data import and export manager (still under development)

:var __author__: name of authors
:var __copyright__: copyrights
:var __license__: licence agreement
:var __url__: url entry point to documentation
"""

import cPickle
import os
import shutil
import util
from esys.escript import getMPIRankWorld, MPIBarrierWorld, load

class DataManager(object):
    """
    Escript data import/export manager.

    Example::

        dm=DataManager(formats=[DataManager.RESTART,DataManager.VTK])
        if dm.hasData():
            dom = dm.getDomain()
            time = dm.getValue("time")
            dt = dm.getValue("dt")
            T = dm.getValue("T")
            u = dm.getValue("u")
        else:
            T = ...
            u = ...
        dm.addData(time=time,dt=dt,T=T,u=u) # add data and variables
        dm.setTime(time)                    # set the simulation timestamp
        dm.export()                         # write out data
    """

    RESTART, SILO, VISIT, VTK = range(4)

    def __init__(self, formats=[RESTART], work_dir=".", restart_prefix="restart", do_restart=True):
        """
        Initialises the data manager. If do_restart is True and a restart
        directory is found the contained data is loaded (hasData() returns True)
        otherwise restart directories are removed (hasData() returns False).
        Values are only written to disk when export() is called.

        :param formats: A list of export file formats to use. Allowed values
                        are RESTART, SILO, VISIT, VTK.
        :param work_dir: top-level directory where files are exported to
        :param restart_prefix: prefix for restart directories. Will be used to
                               load restart files (if do_restart is True) and
                               store new restart files (if RESTART is used)
        :param do_restart: whether to attempt to load restart files
        """
        self._metadata=""
        self._md_schema=""
        self._data={}
        self._domain=None
        self._meshlabels=["","",""]
        self._meshunits=["","",""]
        self._stamp={}
        self._time=0.
        self._restartdir=None
        self._N=-1
        self._checkpointfreq=1
        self._myrank=getMPIRankWorld()
        self._exportformats=set(formats)
        self._restartprefix=restart_prefix
        self._workdir=work_dir
        util.mkDir(self._workdir)
        if self.VISIT in self._exportformats:
            simFile=os.path.join(self._workdir, "escriptsim.sim2")
            if not self.__initVisit(simFile, "Escript simulation"):
                print("Warning: Could not initialize VisIt interface")
                self._exportformats.remove(self.VISIT)
        if self.RESTART in self._exportformats:
            # find all restart directories
            restart_folders = []
            for f in os.listdir(self._workdir):
                if f.startswith(self._restartprefix):
                    restart_folders.append(f)
            # remove unneeded restart directories
            if len(restart_folders)>0:
                restart_folders.sort()
                if do_restart:
                    self._restartdir=restart_folders[-1]
                    print("Restart from "+os.path.join(self._workdir, self._restartdir))
                    for f in restart_folders[:-1]:
                        self.__removeDirectory(f)
                    self.__loadState()
                else:
                    for f in restart_folders:
                        self.__removeDirectory(f)

    def addData(self, **data):
        """
        Adds 'escript.Data' objects and other data to be exported to this
        manager.

        :note: This method does not make copies of Data objects so
               any modifications will be carried over until export() is called.
        """
        # if this is the first addition after a restart, clear data first
        if self._restartdir != None:
            self.__clearData()

        for name,var in data.items():
            if hasattr(var, "getDomain"):
                if self._domain==None:
                    self._domain=var.getDomain()
                elif self._domain != var.getDomain():
                    raise ValueError("addData: Data must be on the same domain!")
                self._data[name]=var
            else:
                self._stamp[name]=var

    def setDomain(self, domain):
        """
        Sets the domain without adding data.
        """
        if self._domain==None:
            self._domain = domain
        elif self._domain != domain:
            raise ValueError("setDomain: Domain already set!")


    def hasData(self):
        """
        Returns True if the manager holds data for restart
        """
        return self._restartdir != None

    def getDomain(self):
        """
        Returns the domain as recovered from restart files.
        """
        if not self.hasData():
            raise ValueError("No restart data available")
        return self._domain 

    def getValue(self, value_name):
        """
        Returns an 'escript.Data' object or other value that has been loaded
        from restart files.
        """
        if not self.hasData():
            raise ValueError("No restart data available")

        if value_name in self._stamp:
            return self._stamp[value_name]

        ff=self.__getDumpFilename(value_name, self._restartdir)
        var = load(ff, self._domain)
        #print("Value %s recovered from %s."%(value_name, ff))
        return var 

    def getCycle(self):
        """
        Returns the export cycle (=number of times export() has been called)
        """
        return self._N

    def setCheckpointFrequency(self, freq):
        """
        Sets the number of calls to export() before new restart files are
        generated.
        """
        self._checkpointfreq=freq

    def setTime(self, time):
        """
        Sets the simulation timestamp.
        """
        self._time = time

    def setMeshLabels(self, x, y, z=""):
        """
        Sets labels for the mesh axes. These are currently only used by the
        Silo exporter.
        """
        self._meshlabels=[x,y,z]

    def setMeshUnits(self, x, y, z=""):
        """
        Sets units for the mesh axes. These are currently only used by the
        Silo exporter.
        """
        self._meshunits=[x,y,z]

    def setMetadataSchemaString(self, schema, metadata=""):
        """
        Sets metadata namespaces and the corresponding metadata.
        Only used for the VTK file format at the moment.

        :param schema: A dictionary that maps namespace prefixes to namespace
                       names, e.g. {'gml':'http://www.opengis.net/gml'}
        :param metadata: The actual metadata string which will be enclosed in
                         '<MetaData>' tags.
        """
        self._metadata=metadata
        ss=""
        for i,p in schema.items():
            ss="%s xmlns:%s=\"%s\""%(ss, i, p)
        self._md_schema=ss.strip()

    def export(self):
        """
        Executes the actual data export. Depending on the formats parameter
        used in the constructor all data added by addData() is written to disk
        (RESTART,SILO,VTK) or made available through the VisIt simulation
        interface (VISIT).
        """

        if self._domain == None:
            print("Warning: DataManager.export() called but no domain set!")
            return

        self._N += 1
        ds = None
        nameprefix=os.path.join(self._workdir, "dataset.%04d"%(self._N))

        for f in self._exportformats:
            if f == self.SILO:
                if ds == None:
                    ds=self.__createDataset()
                ds.saveSilo(nameprefix)
            elif f == self.VTK:
                if ds == None:
                    ds=self.__createDataset()
                ds.saveVTK(nameprefix)
            elif f == self.VISIT:
                from esys.weipa.weipacpp import visitPublishData
                if ds == None:
                    ds=self.__createDataset()
                visitPublishData(ds)
            elif f == self.RESTART:
                # only write checkpoint files with the requested frequency
                if self._N % self._checkpointfreq==0:
                    self.__saveState()
            else:
                raise ValueError("export: Unknown export format "+str(f))

        self.__clearData()

    def __createDataset(self):
        from esys.weipa.weipacpp import EscriptDataset
        from esys.weipa import createDataset

        ds = createDataset(self._domain, self._data)
        ds.setCycleAndTime(self._N, self._time)
        ds.setMetadataSchemaString(self._md_schema, self._metadata)
        ds.setMeshLabels(self._meshlabels[0], self._meshlabels[1], self._meshlabels[2])
        ds.setMeshUnits(self._meshunits[0], self._meshunits[1], self._meshunits[2])
        return ds
    
    def __clearData(self):
        #print("Clearing all data")
        self._restartdir = None
        self._domain = None
        self._stamp = {}
        self._data = {}

    def __getStampFilename(self, dir_name):
        return os.path.join(self._workdir, dir_name, "stamp.%d"%self._myrank)

    def __getDumpFilename(self, data_name, dir_name):
        return os.path.join(self._workdir, dir_name, "%s.nc"%data_name)

    def __initVisit(self, simFile, comment=""):
        """
        Initialises the VisIt interface if available.

        :param simFile: Name of the sim file to be generated which can be
                        loaded into a VisIt client
        :param comment: A short description of this simulation
        """
        from esys.weipa.weipacpp import visitInitialize
        return visitInitialize(simFile, comment)

    def __loadState(self):
        stamp_file=self.__getStampFilename(self._restartdir)
        try:
            self._stamp = cPickle.load(open(stamp_file, "rb"))
            self._N = int(self._restartdir[len(self._restartprefix)+1:])
        except:
            raise IOError("Could not load stamp file "+stamp_file)
        # load domain
        ff=self.__getDumpFilename("_domain",self._restartdir)
        modname=self._stamp['__domainmodule']
        clsname=self._stamp['__domainclass']
        try:
            domclass=__import__(modname, fromlist=[clsname])
            self._domain = domclass.LoadMesh(ff)
        except:
            raise ImportError("Unable to load %s using %s.%s!"%(ff, modname, clsname))

    def __saveState(self):
        restartdir = "%s_%04d"%(self._restartprefix, self._N)
        util.mkDir(os.path.join(self._workdir, restartdir))
        stamp_file=self.__getStampFilename(restartdir)
        self._stamp['__domainmodule']=self._domain.__module__
        self._stamp['__domainclass']=type(self._domain).__name__
        cPickle.dump(self._stamp, open(stamp_file, "wb"))
        ff=self.__getDumpFilename("_domain", restartdir)
        self._domain.dump(ff)
        for name, var in self._data.items():
            ff=self.__getDumpFilename(name, restartdir)
            var.dump(ff)
        print("Restart files saved in "+os.path.join(self._workdir, restartdir))
        # keep only one restart directory
        old_restartdir = "%s_%04d"%(self._restartprefix, self._N-self._checkpointfreq)
        self.__removeDirectory(os.path.join(self._workdir, old_restartdir))

    def __removeDirectory(self, path):
        if self._myrank==0 and os.path.isdir(path):
            shutil.rmtree(path, True)
            #print("Removed restart directory %s."%path)
        MPIBarrierWorld()

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