"""The WaveBlocks Project

This script loads and transforms data from a datafile
to make them suitable as initial values.

@author: R. Bourquin
@copyright: Copyright (C) 2012 R. Bourquin
@license: Modified BSD License
"""

from WaveBlocksND import IOManager
from WaveBlocksND import BlockFactory


__all__ = ["load_from_file"]


def load_from_file(filepath, blockid=0, timestep=0):

    IOM = IOManager()
    IOM.open_file(filepath)

    # Check if we have data
    tg = IOM.load_wavepacket_timegrid(blockid=blockid)
    if not timestep in tg:
        raise ValueError("No data for timestep "+str(timestep))

    # Load data and assemble packet
    BF = BlockFactory()

    # Basis shapes
    BS_descr = IOM.load_wavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in BS_descr.iteritems():
        BS[ahash] = BF.create_basis_shape(descr)

    # Create a packet
    wpd = IOM.load_wavepacket_description(blockid=blockid)
    HAWP = BF.create_wavepacket(wpd)

    # Data
    ha, ci = IOM.load_wavepacket_coefficients(blockid=blockid, timestep=timestep, get_hashes=True)
    Pi = IOM.load_wavepacket_parameters(blockid=blockid, timestep=timestep)

    HAWP.set_parameters(Pi)
    HAWP.set_basis_shapes([ BS[int(h)] for h in ha ])
    HAWP.set_coefficients(ci)

    # Reformat data
    C = []

    for n in xrange(HAWP.get_number_components()):
        B = HAWP.get_basis_shapes(component=n)
        cn = HAWP.get_coefficients(component=n)
        l = []
        for i in xrange(B.get_basis_size()):
            l.append( (B[i], cn[i,0]) )
        C.append(l)

    return Pi, C

@@ -167,16 +167,22 @@
 (C) 2012
+, 2013
  R. Bour
@@ -186,16 +186,16 @@
 ourquin%0A
-
 @license
@@ -378,18 +378,887 @@
 mestep=0
-):
+, sizeK=None):%0A    r%22%22%22Utility script to load wavepacket parameters and coefficients%0A    from another simulation result in a form suitable for the input%0A    configuration of a new simulation. This is (mainly) used%0A    to start simulations with previously computed eigenstates.%0A%0A    :param filepath: The path to the %60.hdf5%60 file from which data will be read.%0A    :param blockid: The %60datablock%60 from which to read the data.%0A                    Default is the block with %60blockid=0%60.%0A    :param timestep: Load the data corresponding to the given %60timestep%60.%0A                     The default timestep is %600%60.%0A    :param sizeK: Load at most 'sizeK' many coefficients. Note that the order%0A                  is defined by the linearization mapping :math:%60%5Cmu%60 of the%0A                  packet's current basis shape. We then pick the first %60sizeK%60%0A                  ones.%0A    %22%22%22
 %0A%0A    IO
@@ -2426,16 +2426,16 @@
 i,0%5D) )%0A
-
         
@@ -2447,16 +2447,218 @@
 end(l)%0A%0A
+    if sizeK is not None:%0A        # We load at most 'sizeK' coefficients.%0A        # Note that this does NOT specify which%0A        # ones in terms of multi-indices.%0A        C = %5Bci%5B:sizeK%5D for ci in C%5D%0A%0A
     retu

Improve the load_from_file parameter and coefficients loader script