"""The WaveBlocks Project
IOM plugin providing functions for handling energy data.
@author: R. Bourquin
@copyright: Copyright (C) 2010, 2011, 2012, 2013 R. Bourquin
@license: Modified BSD License
"""
import numpy as np
[docs]def add_energy(self, parameters, timeslots=None, blockid=0, key=("kin", "pot")):
r"""
:param parameters: A :py:class:`ParameterProvider` instance containing
at least the key `ncomponents`.
:param timeslots: The number of time slots we need. Can be set to ``None``
to get automatically growing datasets.
:param blockid: The ID of the data block to operate on.
:param key: Specify which energies to save. All are independent.
:type key: Tuple of valid identifier strings that are ``kin``, ``pot`` and ``tot``.
Default is ``("kin", "pot")``.
"""
N = parameters["ncomponents"]
if timeslots is None:
T = 0
Ts = None
csTs = 64
else:
T = timeslots
Ts = timeslots
csTs = min(64, Ts)
# Check that the "observables" group is present
grp_ob = self._srf[self._prefixb + str(blockid)].require_group("observables")
# Add a new group for energies
grp_en = grp_ob.create_group("energies")
# Add all requested data sets
if "kin" in key and "kinetic" not in grp_en.keys():
daset_tgek = grp_en.create_dataset("timegrid_kin", (T,), dtype=np.integer, chunks=True, maxshape=(Ts,), fillvalue=-1)
grp_en.create_dataset("kinetic", (T, N), dtype=np.floating, chunks=(csTs, N), maxshape=(Ts, N))
daset_tgek.attrs["pointer"] = 0
if "pot" in key and "potential" not in grp_en.keys():
daset_tgep = grp_en.create_dataset("timegrid_pot", (T,), dtype=np.integer, chunks=True, maxshape=(Ts,), fillvalue=-1)
grp_en.create_dataset("potential", (T, N), dtype=np.floating, chunks=(csTs, N), maxshape=(Ts, N))
daset_tgep.attrs["pointer"] = 0
if "tot" in key and "total" not in grp_en.keys():
daset_tget = grp_en.create_dataset("timegrid_tot", (T,), dtype=np.integer, chunks=True, maxshape=(Ts,), fillvalue=-1)
grp_en.create_dataset("total", (T, 1), dtype=np.floating, chunks=(csTs, 1), maxshape=(Ts, 1))
daset_tget.attrs["pointer"] = 0
[docs]def delete_energy(self, blockid=0):
r"""Remove the stored energies
:param blockid: The ID of the data block to operate on.
"""
try:
del self._srf[self._prefixb + str(blockid) + "/observables/energies"]
# Check if there are other children, if not remove the whole node.
if len(self._srf[self._prefixb + str(blockid) + "/observables"].keys()) == 0:
del self._srf[self._prefixb + str(blockid) + "/observables"]
except KeyError:
pass
[docs]def has_energy(self, blockid=0, key=("kin", "pot")):
r"""Ask if the specified data block has the desired data tensor.
:param blockid: The ID of the data block to operate on.
:param key: Specify which energies to save. All are independent.
:type key: Tuple of valid identifier strings that are ``kin``, ``pot`` and ``tot``.
Default is ``("kin", "pot")``.
"""
r = True
r &= ("observables" in self._srf[self._prefixb + str(blockid)].keys())
if r is True:
r &= ("energies" in self._srf[self._prefixb + str(blockid)]["observables"].keys())
if r and "kin" in key:
r &= ("kinetic" in self._srf[self._prefixb + str(blockid)]["observables/energies"].keys())
if r and "pot" in key:
r &= ("potential" in self._srf[self._prefixb + str(blockid)]["observables/energies"].keys())
if r and "tot" in key:
r &= ("total" in self._srf[self._prefixb + str(blockid)]["observables/energies"].keys())
return r
[docs]def save_energy(self, energies, timestep=None, blockid=0, key=("kin", "pot")):
r"""Save the kinetic and potential energies to a file.
:param energies: A tuple containing the energies. The order is important,
it has to match the order in the ``key`` argument. Per default
the order has to be :math:`(E_\text{kin}, E_\text{pot})`.
:param timestep: The timestep at which we save the data.
:param blockid: The ID of the data block to operate on.
:param key: Specify which energies to save. All are independent.
:type key: Tuple of valid identifier strings that are ``kin``, ``pot`` and ``tot``.
Default is ``("kin", "pot")``.
"""
for item, datum in zip(key, energies):
if item == "kin":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_kin"
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/kinetic"
elif item == "pot":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_pot"
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/potential"
elif item == "tot":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_tot"
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/total"
timeslot = self._srf[pathtg].attrs["pointer"]
# TODO: remove np.array
energy = np.real(np.array(datum))
# Write the data
self.must_resize(pathd, timeslot)
self._srf[pathd][timeslot, :] = energy
# Write the timestep to which the stored values belong into the timegrid
self.must_resize(pathtg, timeslot)
self._srf[pathtg][timeslot] = timestep
# Update the pointer
self._srf[pathtg].attrs["pointer"] += 1
[docs]def load_energy_timegrid(self, blockid=0, key=("kin", "pot")):
r"""Load the time grid for specified energies.
:param blockid: The ID of the data block to operate on.
:param key: Specify which energy timegrids to load. All are independent.
:type key: Tuple of valid identifier strings that are ``kin``, ``pot`` and ``tot``.
Default is ``("kin", "pot")``.
"""
tg = []
for item in key:
if item == "kin":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_kin"
tg.append(self._srf[pathtg][:])
elif item == "pot":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_pot"
tg.append(self._srf[pathtg][:])
elif item == "tot":
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_tot"
tg.append(self._srf[pathtg][:])
if len(tg) == 1:
return tg[0]
else:
return tuple(tg)
[docs]def load_energy(self, timestep=None, split=False, blockid=0, key=("kin", "pot")):
r"""Load the energy data.
:param timestep: Load only the data of this timestep.
:param split: Split the array into arrays for each component.
:param blockid: The ID of the data block to operate on.
:param key: Specify which energies to save. All are independent.
:type key: Tuple of valid identifier strings that are ``kin``, ``pot`` and ``tot``.
Default is ``("kin", "pot")``.
"""
result = []
for item in key:
if item == "kin":
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/kinetic"
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_ek"
elif item == "pot":
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/potential"
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_ep"
elif item == "tot":
pathd = "/" + self._prefixb + str(blockid) + "/observables/energies/total"
pathtg = "/" + self._prefixb + str(blockid) + "/observables/energies/timegrid_et"
if timestep is not None:
index = self.find_timestep_index(pathtg, timestep)
axis = 0
else:
index = slice(None)
axis = 1
if split is True:
energy = self.split_data(self._srf[pathd][index, ...], axis)
else:
energy = self._srf[pathd][index, ...]
result.append(energy)
if len(result) == 1:
return result[0]
else:
return tuple(result)