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 21.3 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
"""
__all__ = [
18
19
    'apply_equations',
    'tecplot_interpolate'
20
21
]
__author__ = 'Camilla D. K. Harris'
22
__email__ = 'cdha@umich.edu'
23

24
25
26
import os
import re

27
import h5py
28
import numpy as np
29
30
import tecplot

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

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

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

45
    Examples:
46
47
48
49
50
51
52
53
        ```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
54
        # tecplot.session.connect(port=7600)
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

        ## 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:
77
78
79
80
81
82
83
84
85
86
            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)
87
                if verbose:
88
                    print(eqn_str)
89
90
    if verbose:
        print('Successfully applied equations.')
91
92


93
def _shell_geometry(geometry_params: dict) -> dict:
94
95
    """Returns a dict containing points for the described shell geometry.
    """
96
97
    nlon = geometry_params['npoints'][0] # 360
    nlat = geometry_params['npoints'][1] # 179
98
99
    lons = np.linspace(0, 360, nlon, endpoint=False)
    dlat = 180/(nlat + 1)
100
    lats = np.linspace(-90.0+dlat, 90.0-dlat, nlat)
101
102
103
    print(f'lons: {lons}')
    print(f'lats: {lats}')

104
105
106
107
108
109
110
111
    latvals, lonvals = np.meshgrid(lats, lons)
    phvals = np.deg2rad(-1*lonvals + 90)
    thvals = np.deg2rad(90 - latvals)
    rhovals = geometry_params['radius'] * np.sin(thvals)
    xvals = rhovals * np.cos(phvals) + geometry_params['center'][0]
    yvals = rhovals * np.sin(phvals) + geometry_params['center'][1]
    zvals = (geometry_params['radius'] * np.cos(thvals)
             + geometry_params['center'][2])
112

113
    geometry_points = {
114
115
116
117
118
119
        'npoints': nlon * nlat,
        'latitude': latvals.flatten(),
        'longitude': lonvals.flatten(),
        'X': xvals.flatten(),
        'Y': yvals.flatten(),
        'Z': zvals.flatten()
120
121
    }
    return geometry_points
122
123


124
def _line_geometry(geometry_params: dict) -> dict:
125
126
    """Returns a dict containing points for the described line geometry.
    """
127
    geometry_points = {
128
129
130
131
132
133
134
135
136
137
138
139
140
        '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'])
141
142
    }
    return geometry_points
143
144


145
def _rectprism_geometry(geometry_params: dict) -> dict:
146
147
    """Returns a dict containing points for the described rectprism geometry.
    """
148
149
150
    npoints = (geometry_params['npoints'][0]
               * geometry_params['npoints'][1]
               * geometry_params['npoints'][2])
151
152
153
154
155
156
157
158
159
160
161
162
    vals = []
    for dim in range(3):
        minval = (geometry_params['center'][dim]
                  - geometry_params['halfwidths'][dim])
        maxval = (geometry_params['center'][dim]
                  + geometry_params['halfwidths'][dim])
        vals.append(np.linspace(
            minval,
            maxval,
            geometry_params['npoints'][dim]
        ))
    xvals, yvals, zvals = np.meshgrid(vals[0], vals[1], vals[2])
163
    geometry_points = {
164
165
166
167
        'npoints': npoints,
        'X': xvals.flatten(),
        'Y': yvals.flatten(),
        'Z': zvals.flatten(),
168
169
    }
    return geometry_points
170
171


172
def _trajectory_geometry(geometry_params: dict) -> dict:
173
    """Returns a dict containing points for the described trajectory geometry.
174
175

    Assumes format of trajectory file after SWMF SATELLITE command.
176
    """
177
178
    do_read = False
    trajectory_data = []
179
    with open(geometry_params['trajectory_data'], 'r') as trajectory_file:
180
181
182
183
184
185
186
187
        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
188
189
190
191
192
193
194
    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.'
        )
195
    geometry_points = {
196
197
198
199
200
201
202
        '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],
203
        'time': [((np.datetime64(
204
205
206
207
208
209
210
            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]}')
211
212
                   - np.datetime64('1970-01-01T00:00:00Z'))
                  / np.timedelta64(1, 's'))
213
                 for trajectory_point in trajectory_data]
214
    }
215
    return geometry_points
216
217


218
def _save_hdf5(filename, geometry_params, new_zone, variables) -> None:
219
220
    """Save the aux data and a subset of the variables in hdf5 format.
    """
221
222
223
224
225
226
227
228
229
    column_names = [var.name for var in variables]
    tp_data = []
    for var in variables:
        tp_data.append(new_zone.values(var)[:])
    tp_data_np = np.array(tp_data).transpose()
    with h5py.File(filename, 'w-') as h5_file:
        h5_file['data'] = tp_data_np
        h5_file['data'].attrs.update(geometry_params)
        h5_file['data'].attrs['names'] = column_names
230
231


232
def _save_csv(filename, geometry_params, new_zone, variables) -> None:
233
234
    """Save the aux data and a subset of the variables in plain-text format.
    """
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    aux_data = geometry_params.__repr__() + '\n'
    column_names = variables[0].name.__repr__()
    for var in variables[1:]:
        column_names += ' ' + var.name.__repr__()
    tp_data = []
    for var in variables:
        tp_data.append(new_zone.values(var)[:])
    tp_data_np = np.array(tp_data).transpose()
    np.savetxt(
        filename,
        tp_data_np,
        delimiter=' ',
        header=aux_data + column_names,
        comments=''
    )
250
251


252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def write_zone(
        tecplot_zone
        , write_as: str
        , filename: str
        , variables=None
        , verbose: bool = False
) -> None:
    """Writes a tecplot zone to various formats.

    Args:
        tecplot_zone (): The zone to save.
        write_as (str): Type of file to write to. Supported options are 'hdf5',
            'csv', 'tecplot_ascii', and 'tecplot_plt'.
        filename (str): Name of the file to write to.
        variables (): (Optional) Specify a subset of the dataset variables to
            save. This option may decrease the size of the output. Default
            behavior is to save all variables.
        verbose: (Optional) Print diagnostic information. Defaults to False.
    """
    if verbose:
        print('Saving variables:')
        for var in variables:
            print(var.name)
    aux_data = tecplot_zone.aux_data.as_dict()
    if verbose:
        print('Attaching auxiliary data:')
        print(aux_data.__repr__)
    ## save zone
    if verbose:
        print('Saving as:')
    if 'hdf5' in write_as:
        if verbose:
            print('hdf5')
        _save_hdf5(
            filename,
            aux_data,
            tecplot_zone,
            variables
        )
    elif 'csv' in write_as:
        _save_csv(
            filename,
            aux_data,
            tecplot_zone,
            variables
        )
    elif 'tecplot_ascii' in write_as:
        tecplot.data.save_tecplot_ascii(
            filename
            , zones=tecplot_zone
            , variables=variables
            , use_point_format=True
        )
    elif 'tecplot_plt' in write_as:
        tecplot.data.save_tecplot_plt(
            filename
            , zones=tecplot_zone
            , variables=variables
        )
    if verbose:
        print(f'Wrote {filename}')


315
def tecplot_interpolate(
316
317
318
319
320
321
322
323
        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
324
) -> None:
325
326
327
328
329
330
331
    """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',
332
            'csv', 'tecplot_ascii', and 'tecplot_plt'.
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
        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
350
351
            interpolate onto, excluding the north and south polar points.
            Defaults to (360,179).
352
353
354
355
356
357
358
359
360
        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).
361
        halfwidths (array-like): Argument for the 'rectprism' geometry. Contains
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
            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
        )
        ```
    """
396
397
398
399
400
401
402
403
404
405
406
407
408
    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']:
409
410
411
412
413
414
415
416
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
        geometry_params['npoints'] = geometry_params.get(
            'npoints'
            , (359, 181)
        )
417
    elif 'rectprism' in geometry_params['kind']:
418
419
420
421
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
422
423
424

    ## check that we support the geometry
    geometry_param_names = {
425
426
427
428
        'shell': ('radius',),
        'line': ('r1', 'r2', 'npoints'),
        'rectprism': ('halfwidths', 'npoints'),
        'trajectory': ('trajectory_data', 'trajectory_format')
429
430
431
432
433
434
435
436
437
438
439
440
441
442
    }
    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'
443
444
445
        , 'csv'
        , 'tecplot_ascii'
        , 'tecplot_plt'
446
447
448
449
450
451
452
    )
    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:')
453
        for key, value in geometry_params.items():
454
455
456
457
            print(f'\t{key}: {value.__repr__()}')

    ## check whether we are using equations
    ## check that the equations file is there
458
    use_equations = not tecplot_equation_file_path is None
459
    if use_equations:
460
461
        equations_file = open(tecplot_equation_file_path, 'r')
        equations_file.close()
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
        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}')
477

478
479
480
481
482
483
484
485
486
487
488
489
490
    ## 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')
491
    apply_equations(tecplot_equation_file_path)
492
493
494
495
496
497
    if verbose:
        print('Variables after equations:')
        print(batsrus.variable_names)

    ## create geometry zone
    if 'shell' in geometry_params['kind']:
498
        geometry_points = _shell_geometry(geometry_params)
499
    elif 'line' in geometry_params['kind']:
500
        geometry_points = _line_geometry(geometry_params)
501
    elif 'rectprism' in geometry_params['kind']:
502
        geometry_points = _rectprism_geometry(geometry_params)
503
    elif 'trajectory' in geometry_params['kind']:
504
505
        if 'batsrus' in geometry_params['trajectory_format']:
            geometry_points = _trajectory_geometry(geometry_params)
506
507

    source_zone = list(batsrus.zones())
508
    if ('trajectory' in geometry_params['kind']
509
            and 'tecplot' in geometry_params['trajectory_format']):
510
        batsrus = tecplot.data.load_tecplot(
511
            filenames=geometry_params['trajectory_data']
512
            , read_data_option=tecplot.constant.ReadDataOption.Append
513
        )
514
        new_zone = batsrus.zone(-1)
515
516
517
518
519
520
        new_zone.name = 'geometry'
    else:
        new_zone = batsrus.add_ordered_zone(
            'geometry'
            , geometry_points['npoints']
        )
521
522
        for i, direction in zip((0, 1, 2), ('X', 'Y', 'Z')):
            new_zone.values(i)[:] = geometry_points[direction][:]
523

524
525
526
    ## interpolate variables on to the geometry
    if verbose:
        print('Interpolating variables:')
527
    positions = list(batsrus.variables('*[[]R[]]'))
528
529
530
531
532
533
    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
534
535
        , source_zones=source_zone
        , variables=variables
536
537
538
    )
    ## add variables for shell and trajectory cases
    if 'shell' in  geometry_params['kind']:
539
540
541
542
543
        batsrus.add_variable('latitude [deg]')
        new_zone.values('latitude [[]deg[]]')[:] = geometry_points['latitude']
        batsrus.add_variable('longitude [deg]')
        new_zone.values('longitude [[]deg[]]')[:] = geometry_points['longitude']
        variables = variables + list(batsrus.variables('*itude [[]deg[]]'))
544
545
    if 'trajectory' in geometry_params['kind']:
        batsrus.add_variable('time')
546
547
548
        if 'batsrus' in geometry_params['trajectory_format']:
            new_zone.values('time')[:] = geometry_points['time']
        variables = variables + [batsrus.variable('time')]
549
550
551
552

    ## add auxiliary data
    new_zone.aux_data.update(geometry_params)
    if ('trajectory' in geometry_params['kind']
553
            and 'pandas' in geometry_params['trajectory_format']):
554
555
556
557
558
        new_zone.aux_data.update(
            {'trajectory_data': type(geometry_params['trajectory_data'])}
        )

    ## construct default filename
559
    no_file_name = False
560
    if filename is None:
561
        no_file_name = True
562
        filename = tecplot_plt_file_path[:-4] + f'_{geometry_params["kind"]}'
563

564
565
    ## save zone
    if 'hdf5' in write_as:
566
567
        if no_file_name:
            filename += '.h5'
568
569
570
571
572
573
        _save_hdf5(
            filename,
            geometry_params,
            new_zone,
            positions + variables
        )
574
    elif 'csv' in write_as:
575
576
        if no_file_name:
            filename += '.csv'
577
578
579
580
581
582
        _save_csv(
            filename,
            geometry_params,
            new_zone,
            positions + variables
        )
583
    elif 'tecplot_ascii' in write_as:
584
585
        if no_file_name:
            filename += '.dat'
586
587
        tecplot.data.save_tecplot_ascii(
            filename
588
            , zones=new_zone
589
590
            , variables=positions + variables
            , use_point_format=True
591
592
        )
    elif 'tecplot_plt' in write_as:
593
594
        if no_file_name:
            filename += '.plt'
595
596
        tecplot.data.save_tecplot_plt(
            filename
597
            , zones=new_zone
598
            , variables=positions + variables
599
        )
600
601
    if verbose:
        print(f'Wrote {filename}')