1# This file is part of QuTiP: Quantum Toolbox in Python.
2#
3#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
4#    All rights reserved.
5#
6#    Redistribution and use in source and binary forms, with or without
7#    modification, are permitted provided that the following conditions are
8#    met:
9#
10#    1. Redistributions of source code must retain the above copyright notice,
11#       this list of conditions and the following disclaimer.
12#
13#    2. Redistributions in binary form must reproduce the above copyright
14#       notice, this list of conditions and the following disclaimer in the
15#       documentation and/or other materials provided with the distribution.
16#
17#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
18#       of its contributors may be used to endorse or promote products derived
19#       from this software without specific prior written permission.
20#
21#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
22#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
23#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
24#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
25#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
26#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
27#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
28#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
29#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
30#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
31#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
32###############################################################################
33
34__all__ = ['file_data_store', 'file_data_read', 'qsave', 'qload']
35
36import pickle
37import numpy as np
38import sys
39from qutip.qobj import Qobj
40from qutip.solver import Result
41
42
43# -----------------------------------------------------------------------------
44# Write matrix data to a file
45#
46def file_data_store(filename, data, numtype="complex", numformat="decimal",
47                    sep=","):
48    """Stores a matrix of data to a file to be read by an external program.
49
50    Parameters
51    ----------
52    filename : str
53        Name of data file to be stored, including extension.
54    data: array_like
55        Data to be written to file.
56    numtype : str {'complex, 'real'}
57        Type of numerical data.
58    numformat : str {'decimal','exp'}
59        Format for written data.
60    sep : str
61        Single-character field seperator.  Usually a tab, space, comma,
62        or semicolon.
63
64    """
65    if filename is None or data is None:
66        raise ValueError("filename or data is unspecified")
67
68    M, N = np.shape(data)
69
70    f = open(filename, "w")
71
72    f.write("# Generated by QuTiP: %dx%d %s matrix " % (M, N, numtype) +
73            "in %s format ['%s' separated values].\n" % (numformat, sep))
74
75    if numtype == "complex":
76
77        if numformat == "exp":
78
79            for m in range(M):
80                for n in range(N):
81                    if np.imag(data[m, n]) >= 0.0:
82                        f.write("%.10e+%.10ej" % (np.real(data[m, n]),
83                                                  np.imag(data[m, n])))
84                    else:
85                        f.write("%.10e%.10ej" % (np.real(data[m, n]),
86                                                 np.imag(data[m, n])))
87                    if n != N - 1:
88                        f.write(sep)
89                f.write("\n")
90
91        elif numformat == "decimal":
92
93            for m in range(M):
94                for n in range(N):
95                    if np.imag(data[m, n]) >= 0.0:
96                        f.write("%.10f+%.10fj" % (np.real(data[m, n]),
97                                                  np.imag(data[m, n])))
98                    else:
99                        f.write("%.10f%.10fj" % (np.real(data[m, n]),
100                                                 np.imag(data[m, n])))
101                    if n != N - 1:
102                        f.write(sep)
103                f.write("\n")
104
105        else:
106            raise ValueError("Illegal numformat value (should be " +
107                             "'exp' or 'decimal')")
108
109    elif numtype == "real":
110
111        if numformat == "exp":
112
113            for m in range(M):
114                for n in range(N):
115                    f.write("%.10e" % (np.real(data[m, n])))
116                    if n != N - 1:
117                        f.write(sep)
118                f.write("\n")
119
120        elif numformat == "decimal":
121
122            for m in range(M):
123                for n in range(N):
124                    f.write("%.10f" % (np.real(data[m, n])))
125                    if n != N - 1:
126                        f.write(sep)
127                f.write("\n")
128
129        else:
130            raise ValueError("Illegal numformat value (should be " +
131                             "'exp' or 'decimal')")
132
133    else:
134        raise ValueError("Illegal numtype value (should be " +
135                         "'complex' or 'real')")
136
137    f.close()
138
139
140# -----------------------------------------------------------------------------
141# Read matrix data from a file
142#
143def file_data_read(filename, sep=None):
144    """Retrieves an array of data from the requested file.
145
146    Parameters
147    ----------
148    filename : str
149        Name of file containing reqested data.
150    sep : str
151        Seperator used to store data.
152
153    Returns
154    -------
155    data : array_like
156        Data from selected file.
157
158    """
159    if filename is None:
160        raise ValueError("filename is unspecified")
161
162    f = open(filename, "r")
163
164    #
165    # first count lines and numbers of
166    #
167    M = N = 0
168    for line in f:
169        # skip comment lines
170        if line[0] == '#' or line[0] == '%':
171            continue
172        # find delim
173        if N == 0 and sep is None:
174            if len(line.rstrip().split(",")) > 1:
175                sep = ","
176            elif len(line.rstrip().split(";")) > 1:
177                sep = ";"
178            elif len(line.rstrip().split(":")) > 1:
179                sep = ":"
180            elif len(line.rstrip().split("|")) > 1:
181                sep = "|"
182            elif len(line.rstrip().split()) > 1:
183                # sepical case for a mix of white space deliminators
184                sep = None
185            else:
186                raise ValueError("Unrecognized column deliminator")
187        # split the line
188        line_vec = line.split(sep)
189        n = len(line_vec)
190        if N == 0 and n > 0:
191            N = n
192            # check type
193            if ("j" in line_vec[0]) or ("i" in line_vec[0]):
194                numtype = "complex"
195            else:
196                numtype = "np.real"
197
198            # check format
199            if ("e" in line_vec[0]) or ("E" in line_vec[0]):
200                numformat = "exp"
201            else:
202                numformat = "decimal"
203
204        elif N != n:
205            raise ValueError("Badly formatted data file: " +
206                             "unequal number of columns")
207        M += 1
208
209    #
210    # read data and store in a matrix
211    #
212    f.seek(0)
213
214    if numtype == "complex":
215        data = np.zeros((M, N), dtype="complex")
216        m = n = 0
217        for line in f:
218            # skip comment lines
219            if line[0] == '#' or line[0] == '%':
220                continue
221            n = 0
222            for item in line.rstrip().split(sep):
223                data[m, n] = complex(item)
224                n += 1
225            m += 1
226
227    else:
228        data = np.zeros((M, N), dtype="float")
229        m = n = 0
230        for line in f:
231            # skip comment lines
232            if line[0] == '#' or line[0] == '%':
233                continue
234            n = 0
235            for item in line.rstrip().split(sep):
236                data[m, n] = float(item)
237                n += 1
238            m += 1
239
240    f.close()
241
242    return data
243
244
245def qsave(data, name='qutip_data'):
246    """
247    Saves given data to file named 'filename.qu' in current directory.
248
249    Parameters
250    ----------
251    data : instance/array_like
252        Input Python object to be stored.
253    filename : str
254        Name of output data file.
255
256    """
257    # open the file for writing
258    fileObject = open(name + '.qu', 'wb')
259    # this writes the object a to the file named 'filename.qu'
260    pickle.dump(data, fileObject)
261    fileObject.close()
262
263
264def qload(name):
265    """
266    Loads data file from file named 'filename.qu' in current directory.
267
268    Parameters
269    ----------
270    name : str
271        Name of data file to be loaded.
272
273    Returns
274    -------
275    qobject : instance / array_like
276        Object retrieved from requested file.
277
278    """
279    with open(name + ".qu", "rb") as fileObject:
280        if sys.version_info >= (3, 0):
281            out = pickle.load(fileObject, encoding='latin1')
282        else:
283            out = pickle.load(fileObject)
284    if isinstance(out, Qobj):  # for quantum objects
285        print('Loaded Qobj object:')
286        str1 = "Quantum object: " + "dims = " + str(out.dims) \
287            + ", shape = " + str(out.shape) + ", type = " + out.type
288        if out.type == 'oper' or out.type == 'super':
289            str1 += ", isHerm = " + str(out.isherm) + "\n"
290        else:
291            str1 += "\n"
292        print(str1)
293    elif isinstance(out, Result):
294        print('Loaded Result object:')
295        print(out)
296    else:
297        print("Loaded " + str(type(out).__name__) + " object.")
298    return out
299