NDH_PythonTools/loadmat.py at main · nholschuh/NDH_PythonTools · GitHub

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
import hdf5storage
import mat73
import scipy
import numpy as np

################## NDH Tools self imports
###########################################################
from .read_h5 import read_h5
###########################################################

def loadmat(fn, varnames=None ,debug_flag = 0, force_method=0):
    """
    % (C) Nick Holschuh - Amherst College -- 2024 ([email protected])
    %
    % This function uses existing tools to intelligently load .mat files
    %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % The inputs are:
    %
    %     fn -- the .mat file to load (full path)
    %     varnames -- default=None, this allows you to load just a subset of variables.
    %     force_method -- requirue the use of mat73 (1), read_h5 (2), or scipy (3). If 0,
    %                     the code will attempt all methods until one works (or it fails)
    %
    %%%%%%%%%%%%%%%
    % The outputs are:
    %
    %      data -- Ideally, a dictionary with keys corresponding to variables in the .mat file
    %              I think, sometimes, it produces a tuple that you have to index into, although
    %              I've tried to prevent that.
    %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    """
    import scipy.io

    if force_method == 0:
        ############################ Try statement, that will catch if all else fails
        try:
            if varnames == None:

                #################### Try for options 1, 2, and 3, no defined variables
                try:
                    try:
                        if debug_flag == 1:
                            print('Trying Method 1 with no described variables')
                        data = mat73.loadmat(fn)
                    except:
                        if debug_flag == 1:
                            print('Trying Method 2 with no described variables')
                        data = read_h5(fn)
                        data = data[0]
                        for var_opts in data.keys():
                            if type(data[var_opts]) == type(np.array([])):
                                if len(data[var_opts].dtype) > 1:
                                    data[var_opts] = data[var_opts]
                                    temp = data[var_opts]['real'] + 1j * data[var_opts]['imag']
                                    data[var_opts] = temp.astype(np.complex64)
                except:
                    if debug_flag == 1:
                        print('Trying Method 3 with no described variables')
                    data = scipy.io.loadmat(fn,squeeze_me=True)

            else:
                #################### Attempt with defined variables
                try:
                    try:
                        if debug_flag == 1:
                            print('Trying Method 1 with defined variables')
                        data = scipy.io.loadmat(fn,variable_names=varnames,squeeze_me=True)
                    except:
                        data = read_h5(fn,varnames)
                        data = data[0]
                        for var_opts in data.keys():
                            if type(data[var_opts]) == type(np.array([])):
                                if len(data[var_opts].dtype) > 1:
                                    data[var_opts] = data[var_opts]
                                    temp = data[var_opts]['real'] + 1j * data[var_opts]['imag']
                                    data[var_opts] = temp.astype(np.complex64)
                except:
                    if debug_flag == 1:
                        print('Abandoning the goal of loading specific variables')
                    data = mat73.loadmat(fn)
                    print('You couldn''t load just the variables you asked for, but loaded the whole file instead')

        ############################ The catch all option when everything fails
        except:
            if debug_flag == 1:
                print('Something is wrong with this .mat file')
            data = {}

    ################################ Forcing a particular method for testing
    elif force_method == 1:
        data = mat73.loadmat(fn)

    elif force_method == 2:
        if varnames == None:
            data = read_h5(fn)
        else:
            data = read_h5(fn,varnames)

        data = data[0]
        for var_opts in data.keys():
            if type(data[var_opts]) == type(np.array([])):
                if len(data[var_opts].dtype) > 1:
                    data[var_opts] = data[var_opts]
                    temp = data[var_opts]['real'] + 1j * data[var_opts]['imag']
                    data[var_opts] = temp.astype(np.complex64)

    elif force_method == 3:
        if varnames == None:
            data = scipy.io.loadmat(fn,variable_names=varnames,squeeze_me=True)
        else:
            data = scipy.io.loadmat(fn,squeeze_me=True)

    ################################ Forcing a particular method for testing
    if isinstance(data,tuple) == 1:
        data = data[0]

    return data