/usr/lib/python3-escript-mpi/esys/escriptcore/datamanager.py is in python3-escript-mpi 5.1-5.
This file is owned by root:root, with mode 0o644.
The actual contents of the file can be viewed below.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 | ##############################################################################
#
# Copyright (c) 2003-2017 by The University of Queensland
# http://www.uq.edu.au
#
# Primary Business: Queensland, Australia
# Licensed under the Apache License, version 2.0
# http://www.apache.org/licenses/LICENSE-2.0
#
# Development until 2012 by Earth Systems Science Computational Center (ESSCC)
# Development 2012-2013 by School of Earth Sciences
# Development from 2014 by Centre for Geoscience Computing (GeoComp)
#
##############################################################################
from __future__ import print_function, division
__copyright__="""Copyright (c) 2003-2017 by The University of Queensland
http://www.uq.edu.au
Primary Business: Queensland, Australia"""
__license__="""Licensed under the Apache License, version 2.0
http://www.apache.org/licenses/LICENSE-2.0"""
__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 pickle
import os
import shutil
from . import util
from . import escriptcpp as esc
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 = list(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=esc.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 sorted(data.items(), key=lambda x: x[0]):
if hasattr(var, "getDomain"):
if self._domain is 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 is 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 = esc.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 sorted(list(schema.items()), key=lambda x: x[0]):
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 is 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 is None:
ds=self.__createDataset()
ds.saveSilo(nameprefix)
elif f == self.VTK:
if ds is None:
ds=self.__createDataset()
ds.saveVTK(nameprefix)
elif f == self.VISIT:
from esys.weipa.weipacpp import visitPublishData
if ds is 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 = pickle.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__
pickle.dump(self._stamp, open(stamp_file, "wb"))
ff=self.__getDumpFilename("_domain", restartdir)
self._domain.dump(ff)
for name, var in sorted(list(self._data.items()), key=lambda x: x[0]):
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)
esc.MPIBarrierWorld()
|