Note: The default ITS GitLab runner is a shared resource and is subject to slowdowns during heavy usage.
You can run your own GitLab runner that is dedicated just to your group if you need to avoid processing delays.

tecplottools.py 16.4 KB
Newer Older
1
#!/usr/bin/env python
2
3
4
5
6
7
8
9
10
"""Tools for working with the Tecplot visualization software.

Requires an active Tecplot license and the pytecplot python package.
pytecplot ships with Tecplot 360 2017 R1 and later versions
but it is recommended that you install the latest version with
`pip install pytecplot`.
See the pytecplot documentation for more details about
[installation](https://www.tecplot.com/docs/pytecplot/install.html).
See also [TECPLOT](TECPLOT.markdown) for tips targeted to SWMF users.
11
12

Some useful references:
13
14
15
- [Tecplot User's Manual](download.tecplot.com/360/current/360_users_manual.pdf)
- [Tecplot Scripting Guide](download.tecplot.com/360/current/360_scripting_guide.pdf)
- [Pytecplot documentation](www.tecplot.com/docs/pytecplot/index.html)
16
17
18
19
20
"""
__all__ = [
    'apply_equations'
]
__author__ = 'Camilla D. K. Harris'
21
__email__ = 'cdha@umich.edu'
22

23
24
25
import os
import re

26
import numpy as np
27
28
import tecplot

29
def apply_equations(eqn_path: str, verbose: bool = False) -> None:
30
31
32
    """Apply an equations file in the Tecplot macro format to the active dataset

    Please reference the Tecplot User's Manual for more details on
33
    equation files and syntax. It is recommended to use this function with eqn
34
35
    files generated with the Tecplot GUI.
    See [TECPLOT](TECPLOT.markdown) for tips on using pytecplot.
36

37
38
    Args:
        eqn_file_path (str): The path to the equation macro file (typically with
39
            extension `.eqn`).
40
41
        verbose (bool): (Optional) Whether or not to print the equations as they
            are applied. Default behavior is silent.
42

43
    Examples:
44
45
46
47
48
49
50
51
        ```python
        import tecplot
        import swmfpy.tecplottools as tpt

        ## Uncomment this line if you are connecting to a running tecplot
        ## session. Be sure that the port number matches the port the GUI is
        ## listening to. See TECPLOT.markdown for tips on connecting to a
        ## running session or running your script seperately.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
52
        # tecplot.session.connect(port=7600)
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

        ## load a dataset
        dataset = tecplot.data.load_tecplot('./z=0_mhd_1_n00000000.plt')

        ## apply an equations file
        tpt.apply_equations('./gse_to_ephio.eqn', verbose= True)

        ## apply a frame style
        frame = tecplot.active_frame()
        frame.load_stylesheet('./density.sty')

        ## annotate with the zone name
        frame.add_text('&(ZONENAME[ACTIVEOFFSET=1])', position= (5, 95))

        ## save the image
        tecplot.export.save_png('./density.png', width= 1200, supersample= 8)
        ```
    """
    if verbose:
        print('Executing:')
    with open(eqn_path, 'r') as eqn_file:
        for line in eqn_file:
75
76
77
78
79
80
81
82
83
84
            if '$!alterdata' in line.lower():
                eqn_line = eqn_file.readline()
                try:
                    eqn_str = eqn_line.split("'")[1]
                except IndexError:
                    try:
                        eqn_str = eqn_line.split("\"")[1]
                    except:
                        raise ValueError(f'Unable to read equation: {eqn_line}')
                tecplot.data.operate.execute_equation(eqn_str)
85
                if verbose:
86
                    print(eqn_str)
87
88
    if verbose:
        print('Successfully applied equations.')
89
90


91
def _shell_geometry(geometry_params: dict) -> dict:
92
93
    """Returns a dict containing points for the described shell geometry.
    """
94
95
96
97
98
    npoints = geometry_params['npoints'][0]*geometry_params['npoints'][1]
    geometry_points = {
        'npoints': npoints
    }
    return geometry_points
99
100


101
def _line_geometry(geometry_params: dict) -> dict:
102
103
    """Returns a dict containing points for the described line geometry.
    """
104
    geometry_points = {
105
106
107
108
109
110
111
112
113
114
115
116
117
        'npoints': geometry_params['npoints'],
        'X': np.linspace(
            geometry_params['r1'][0],
            geometry_params['r2'][0],
            geometry_params['npoints']),
        'Y': np.linspace(
            geometry_params['r1'][1],
            geometry_params['r2'][1],
            geometry_params['npoints']),
        'Z': np.linspace(
            geometry_params['r1'][2],
            geometry_params['r2'][2],
            geometry_params['npoints'])
118
119
    }
    return geometry_points
120
121


122
def _rectprism_geometry(geometry_params: dict) -> dict:
123
124
    """Returns a dict containing points for the described rectprism geometry.
    """
125
126
127
128
129
130
131
    npoints = (geometry_params['npoints'][0]
               * geometry_params['npoints'][1]
               * geometry_params['npoints'][2])
    geometry_points = {
        'npoints': npoints
    }
    return geometry_points
132
133


134
def _trajectory_geometry(geometry_params: dict) -> dict:
135
    """Returns a dict containing points for the described trajectory geometry.
136
137

    Assumes format of trajectory file after SWMF SATELLITE command.
138
    """
139
140
141
142
143
144
145
146
147
148
149
    do_read = False
    trajectory_data = []
    with open('r', geometry_params['trajectory_file']) as trajectory_file:
        for line in trajectory_file:
            if do_read:
                if len(line.split()) == 10:
                    trajectory_data.append(line.split())
                else:
                    do_read = False
            elif '#START' in line:
                do_read = True
150
151
152
153
154
155
156
    try:
        assert len(trajectory_data) >= 1
    except:
        raise ValueError(
            'No points could be read from the trajectory file. Consult the '
            'SWMF documentation for advice on trajectory format.'
        )
157
    geometry_points = {
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
        'npoints': len(trajectory_data),
        'X': [float(trajectory_point[7])
              for trajectory_point in trajectory_data],
        'Y': [float(trajectory_point[8])
              for trajectory_point in trajectory_data],
        'Z': [float(trajectory_point[9])
              for trajectory_point in trajectory_data],
        'time': [np.datetime64(
            f'{trajectory_point[0]}'
            f'-{trajectory_point[1]}'
            f'-{trajectory_point[2]}'
            f'T{trajectory_point[3]}'
            f':{trajectory_point[4]}'
            f':{trajectory_point[5]}'
            f'.{trajectory_point[6]}')
                 for trajectory_point in trajectory_data]
174
    }
175
    return geometry_points
176
177


178
def _save_hdf5() -> None:
179
180
181
182
    """Save the aux data and a subset of the variables in hdf5 format.
    """


183
def _save_csv() -> None:
184
185
186
187
188
    """Save the aux data and a subset of the variables in plain-text format.
    """


def tecplot_interpolate(
189
190
191
192
193
194
195
196
        tecplot_plt_file_path: str
        , geometry: str
        , write_as: str
        , filename: str = None
        , tecplot_equation_file_path: str = None
        , tecplot_variable_pattern: str = None
        , verbose: bool = False
        , **kwargs
197
) -> None:
198
199
200
201
202
203
204
    """Interpolates Tecplot binary data onto various geometries.

    Args:
        tecplot_plt_file_path (str): Path to the tecplot binary file.
        geometry (str): Type of geometry for interpolation. Supported geometries
            are 'shell', 'line', 'rectprism', or 'trajectory'.
        write_as (str): Type of file to write to. Supported options are 'hdf5',
205
            'csv', 'tecplot_ascii', and 'tecplot_plt'.
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
        filename (str): (Optional) Name of the file to write to. Defaults to a
            concatenation of the tecplot file name and the geometry type.
        tecplot_equation_file_path (str): (Optional) Path to an equation file to
            be applied to the tecplot dataset before interpolation. Defaults to
            no equations.
        tecplot_variable_pattern (str): (Optional) Regex-style variable name
            pattern used to save a subset of the variables. This option may be
            used to decrease the size of the hdf5 output. Default behavior is to
            save all variables.
        verbose: (Optional) Print diagnostic information. Defaults to False.

    Keyword Args:
        center (array-like): Argument for the 'shell' geometry. Contains the X,
            Y, and Z positions of the shell. Defaults to (0,0,0).
        radius (float): Argument for the 'shell' geometry. Required.
        npoints (array-like): Argument for the 'shell' geometry. Contains the
            number of points in the azimuthal and polar directions to
            interpolate onto. Defaults to (359,181).
        r1 (array-like): Argument for the 'line' geometry. Contains the X, Y,
            and Z positions of the point where the line starts. Required.
        r2 (array-like): Argument for the 'line' geometry. Contains the X, Y,
            and Z positions of the point where the line ends. Required.
        npoints (int): Argument for the 'line' geometry. The number of points
            along the line to interpolate onto. Required.
        center (array-like): Argument for the 'rectprism' geometry. Contains the
            X, Y, and Z positions of the center of the rectangular prism.
            Defaults to (0,0,0).
233
        halfwidths (array-like): Argument for the 'rectprism' geometry. Contains
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
            the half widths of the rectangular prism in the X, Y, and Z
            directions. Required.
        npoints (array-like): Argument for the 'rectprism' geometry. Contains
            the number of points in the X, Y, and Z directions to interpolate
            onto. Required.
        trajectory_data (str): Argument for the 'trajectory' geometry. The path
            to the ASCII trajectory data file. Required.
        trajectory_format (str): Argument for the 'trajectory' geometry. The
            format of the trajectory data file. Supported formats are 'tecplot'
            (data is a tecplot zone with 3D positional variables and 'time') and
            'batsrus' (data is formatted as described for the #SATELLITE
            command, see SWMF documentation). Required.

    Examples:
        ```tecplot_interpolate(
            tecplot_plt_file_path='./path/to/data.plt'
            ,geometry='shell'
            ,write_as='tecplot_ascii'
            ,center=[0.0, 0.0, 0.0]
            ,radius=1.01
        )

        tecplot_interpolate(
            tecplot_plt_file_path='./path/to/data.plt'
            ,geometry='line'
            ,write_as='tecplot_ascii'
            ,tecplot_equation_file_path='./path/to/equations.eqn'
            ,tecplot_variable_pattern='B.*|E.*'
            ,r1=[1.0, 0.0, 0.0]
            ,r2=[6.0, 0.0, 0.0]
            ,npoints=101
        )
        ```
    """
268
269
270
271
272
273
274
275
276
277
278
279
280
    if verbose:
        print('Collecting parameters')

    ## collect the geometry parameters
    geometry_params = {
        'kind':geometry
    }
    geometry_params.update(kwargs)

    ## assign defaults for shell
    if verbose:
        print('Adding defaults')
    if 'shell' in geometry_params['kind']:
281
282
283
284
285
286
287
288
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
        geometry_params['npoints'] = geometry_params.get(
            'npoints'
            , (359, 181)
        )
289
    elif 'rectprism' in geometry_params['kind']:
290
291
292
293
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
294
295
296

    ## check that we support the geometry
    geometry_param_names = {
297
298
299
300
        'shell': ('radius',),
        'line': ('r1', 'r2', 'npoints'),
        'rectprism': ('halfwidths', 'npoints'),
        'trajectory': ('trajectory_data', 'trajectory_format')
301
302
303
304
305
306
307
308
309
310
311
312
313
314
    }
    if geometry_params['kind'] not in geometry_param_names:
        raise ValueError(f'Geometry {geometry_params["kind"]} not supported!')

    ## check that we've gotten the right /required/ geometry arguments
    for param in geometry_param_names[geometry_params['kind']]:
        if param not in geometry_params:
            raise TypeError(
                f'Geometry {geometry_params["kind"]} '
                f'requires argument {param}!')

    ## check that we support the file type to save as
    file_types = (
        'hdf5'
315
316
317
        , 'csv'
        , 'tecplot_ascii'
        , 'tecplot_plt'
318
319
320
321
322
323
324
    )
    if write_as not in file_types:
        raise ValueError(f'File type {write_as} not supported!')

    ## describe the interpolation we're about to do on the data
    if verbose:
        print('Geometry to be interpolated:')
325
        for key, value in geometry_params.items():
326
327
328
329
            print(f'\t{key}: {value.__repr__()}')

    ## check whether we are using equations
    ## check that the equations file is there
330
    use_equations = not tecplot_equation_file_path is None
331
    if use_equations:
332
333
        equations_file = open(tecplot_equation_file_path, 'r')
        equations_file.close()
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
        if verbose:
            print('Applying equations from file:')
            print(tecplot_equation_file_path)
    else:
        if verbose:
            print('Not applying any equations')

    ## check patterns
    if not (tecplot_variable_pattern is None) and verbose:
        print(f'Applying pattern {tecplot_variable_pattern} to variables')

    ## check that the tecplot file is there
    if not os.path.exists(tecplot_plt_file_path):
        raise FileNotFoundError(
            f'Tecplot file does not exist: {tecplot_plt_file_path}')
349

350
351
352
353
354
355
356
357
358
359
360
361
362
    ## load the tecplot data
    if verbose:
        print('Loading tecplot data')
    batsrus = tecplot.data.load_tecplot(tecplot_plt_file_path)

    ## describe the loaded tecplot data
    if verbose:
        print('Loaded tecplot data with variables:')
        print(batsrus.variable_names)

    ## apply equations
    if verbose:
        print('Applying equations to data')
363
    apply_equations(tecplot_equation_file_path)
364
365
366
367
368
369
    if verbose:
        print('Variables after equations:')
        print(batsrus.variable_names)

    ## create geometry zone
    if 'shell' in geometry_params['kind']:
370
        geometry_points = _shell_geometry(geometry_params)
371
    elif 'line' in geometry_params['kind']:
372
        geometry_points = _line_geometry(geometry_params)
373
    elif 'rectprism' in geometry_params['kind']:
374
        geometry_points = _rectprism_geometry(geometry_params)
375
    elif 'trajectory' in geometry_params['kind']:
376
377
        if 'batsrus' in geometry_params['trajectory_format']:
            geometry_points = _trajectory_geometry(geometry_params)
378
379

    source_zone = list(batsrus.zones())
380
    if ('trajectory' in geometry_params['kind']
381
            and 'tecplot' in geometry_params['trajectory_format']):
382
        batsrus = tecplot.data.load_tecplot(
383
384
            filenames=geometry_params['trajectory_file']
            , read_data_option=tecplot.constant.ReadDataOption.Append
385
386
387
388
389
390
391
392
        )
        new_zone = batsrus.zones(-1)
        new_zone.name = 'geometry'
    else:
        new_zone = batsrus.add_ordered_zone(
            'geometry'
            , geometry_points['npoints']
        )
393
394
        for i, direction in zip((0, 1, 2), ('X', 'Y', 'Z')):
            new_zone.values(i)[:] = geometry_points[direction][:]
395

396
397
398
    ## interpolate variables on to the geometry
    if verbose:
        print('Interpolating variables:')
399
    positions = list(batsrus.variables('*[[]R[]]'))
400
401
402
403
404
405
    variables = list(batsrus.variables(re.compile(tecplot_variable_pattern)))
    if verbose:
        for var in variables:
            print(var.name)
    tecplot.data.operate.interpolate_linear(
        destination_zone=new_zone
406
407
        , source_zones=source_zone
        , variables=variables
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    )
    ## add variables for shell and trajectory cases
    if 'shell' in  geometry_params['kind']:
        batsrus.add_variable('polar')
        new_zone.values('polar')[:] = geometry_points['polar']
        batsrus.add_variable('azimuthal')
        new_zone.values('azimuthal')[:] = geometry_points['azimuthal']
    if 'trajectory' in geometry_params['kind']:
        batsrus.add_variable('time')
        new_zone.values('time')[:] = geometry_points['time']

    ## add auxiliary data
    new_zone.aux_data.update(geometry_params)
    if ('trajectory' in geometry_params['kind']
422
            and 'pandas' in geometry_params['trajectory_format']):
423
424
425
426
427
        new_zone.aux_data.update(
            {'trajectory_data': type(geometry_params['trajectory_data'])}
        )

    ## construct default filename
428
    if filename is None:
429
        filename = tecplot_plt_file_path[:-4] + f'_{geometry_params["kind"]}'
430

431
432
    ## save zone
    if 'hdf5' in write_as:
433
        filename += '.h5'
434
435
        _save_hdf5()
    elif 'csv' in write_as:
436
        filename += '.csv'
437
        _save_csv()
438
    elif 'tecplot_ascii' in write_as:
439
        filename += '.dat'
440
441
        tecplot.data.save_tecplot_ascii(
            filename
442
            , zones=new_zone
443
444
            , variables=positions + variables
            , use_point_format=True
445
446
        )
    elif 'tecplot_plt' in write_as:
447
        filename += '.plt'
448
449
        tecplot.data.save_tecplot_plt(
            filename
450
            , zones=new_zone
451
        )
452
453
    if verbose:
        print(f'Wrote {filename}')