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 19 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
    'apply_equations',
19
    'bracketify',
20
21
    'write_zone',
    'interpolate_zone_to_geometry'
22
23
]
__author__ = 'Camilla D. K. Harris'
24
__email__ = 'cdha@umich.edu'
25

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

30

31
def _shell_geometry(geometry_params: dict) -> dict:
32
33
    """Returns a dict containing points for the described shell geometry.
    """
34
35
    nlon = geometry_params['npoints'][0] # 360
    nlat = geometry_params['npoints'][1] # 179
36
37
    lons = np.linspace(0, 360, nlon, endpoint=False)
    dlat = 180/(nlat + 1)
38
    lats = np.linspace(-90.0+dlat, 90.0-dlat, nlat)
39
40
41
    print(f'lons: {lons}')
    print(f'lats: {lats}')

42
43
44
45
46
47
48
49
    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])
50

51
    geometry_points = {
52
53
54
55
56
57
        'npoints': nlon * nlat,
        'latitude': latvals.flatten(),
        'longitude': lonvals.flatten(),
        'X': xvals.flatten(),
        'Y': yvals.flatten(),
        'Z': zvals.flatten()
58
59
    }
    return geometry_points
60
61


62
def _line_geometry(geometry_params: dict) -> dict:
63
64
    """Returns a dict containing points for the described line geometry.
    """
65
    geometry_points = {
66
67
68
69
70
71
72
73
74
75
76
77
78
        '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'])
79
80
    }
    return geometry_points
81
82


83
def _rectprism_geometry(geometry_params: dict) -> dict:
84
85
    """Returns a dict containing points for the described rectprism geometry.
    """
86
87
88
    npoints = (geometry_params['npoints'][0]
               * geometry_params['npoints'][1]
               * geometry_params['npoints'][2])
89
90
91
92
93
94
95
96
97
98
99
100
    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])
101
    geometry_points = {
102
103
104
105
        'npoints': npoints,
        'X': xvals.flatten(),
        'Y': yvals.flatten(),
        'Z': zvals.flatten(),
106
107
    }
    return geometry_points
108
109


110
def _trajectory_geometry(geometry_params: dict) -> dict:
111
    """Returns a dict containing points for the described trajectory geometry.
112
113

    Assumes format of trajectory file after SWMF SATELLITE command.
114
    """
115
116
    do_read = False
    trajectory_data = []
117
    with open(geometry_params['trajectory_data'], 'r') as trajectory_file:
118
119
120
121
122
123
124
125
        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
126
127
128
129
130
131
132
    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.'
        )
133
    geometry_points = {
134
135
136
137
138
139
140
        '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],
141
        'time': [((np.datetime64(
142
143
144
145
146
147
148
            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]}')
149
150
                   - np.datetime64('1970-01-01T00:00:00Z'))
                  / np.timedelta64(1, 's'))
151
                 for trajectory_point in trajectory_data]
152
    }
153
    return geometry_points
154
155


156
def _save_hdf5(filename, geometry_params, new_zone, variables) -> None:
157
158
    """Save the aux data and a subset of the variables in hdf5 format.
    """
159
160
161
162
163
164
165
166
167
    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
168
169


170
def _save_csv(filename, geometry_params, new_zone, variables) -> None:
171
172
    """Save the aux data and a subset of the variables in plain-text format.
    """
173
174
175
    aux_data = geometry_params.__repr__() + '\n'
    column_names = variables[0].name.__repr__()
    for var in variables[1:]:
176
        column_names += ',' + var.name.__repr__()
177
178
179
180
181
182
183
    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,
184
        delimiter=',',
185
186
187
        header=aux_data + column_names,
        comments=''
    )
188
189


190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
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
233
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
def _add_variable_value(dataset, variable_name: str, zone, values):
    """Adds and populates a new variable to a zone in a dataset."""
    dataset.add_variable(variable_name)
    zone.values(bracketify(variable_name))[:] = values


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

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

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

    Examples:
        ```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.
        # tecplot.session.connect(port=7600)

        ## 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:
            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)
                if verbose:
                    print(eqn_str)
    if verbose:
        print('Successfully applied equations.')


def bracketify(variable_name: str) -> str:
    """Surrounds square brackets with more brackets.

    This is helpful for accessing Tecplot variables.

263
264
265
266
    Args:
        variable_name (str): A string which may contain the meta-characters * ?
        [ or ].

267
268
269
270
271
272
273
274
275
276
277
278
279
    Examples:
        In a dataset which contains the variable 'X [R]',
        ```print(dataset.variable_names)
        >>> ['X [R]', ... ]```
        The following will fail:
        ```print(dataset.variable('X [R]').name)
        >>> TecplotPatternMatchWarning: no variables found matching: "X [R]" For
        a literal match, the meta-characters: * ? [ ] must be wrapped in square-
        brackets. For example, "[?]" matches the character "?"...```
        However,
        ```print(dataset.variable(tpt.bracketify('X [R]')).name)```
        will succeed.
    """
280
281
282
283
284
285
286
    translation = {
        '[':'[[]',
        ']':'[]]',
        '*':'[*]',
        '?':'[?]'
    }
    return variable_name.translate(str.maketrans(translation))
287
288


289
def write_zone(
290
291
        tecplot_dataset
        , tecplot_zone
292
293
294
295
296
297
298
299
        , write_as: str
        , filename: str
        , variables=None
        , verbose: bool = False
) -> None:
    """Writes a tecplot zone to various formats.

    Args:
300
        tecplot_dataset (): The dataset to save.
301
302
303
304
305
306
307
308
309
        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.
    """
310
    if verbose and variables:
311
312
313
314
315
316
        print('Saving variables:')
        for var in variables:
            print(var.name)
    aux_data = tecplot_zone.aux_data.as_dict()
    if verbose:
        print('Attaching auxiliary data:')
317
        print(aux_data.__repr__())
318
319
320
321
322
323
    ## save zone
    if verbose:
        print('Saving as:')
    if 'hdf5' in write_as:
        if verbose:
            print('hdf5')
324
325
        if not variables:
            variables = list(tecplot_dataset.variables())
326
327
328
329
330
331
332
        _save_hdf5(
            filename,
            aux_data,
            tecplot_zone,
            variables
        )
    elif 'csv' in write_as:
333
334
        if not variables:
            variables = list(tecplot_dataset.variables())
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
        _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
        )
354
355
    else:
        raise ValueError(f'File type {write_as} not supported!')
356
357
358
359
    if verbose:
        print(f'Wrote {filename}')


360
361
362
def interpolate_zone_to_geometry(
        dataset
        , source_zone
363
        , geometry: str
364
        , variables: list = None
365
366
        , verbose: bool = False
        , **kwargs
367
):
368
369
370
    """Interpolates Tecplot binary data onto various geometries.

    Args:
371
372
        dataset: The loaded Tecplot dataset.
        source_zone: The Tecplot zone to interpolate onto the geometry.
373
        geometry (str): Type of geometry for interpolation. Supported geometries
374
375
376
377
            are 'shell', 'line', 'rectprism', or 'trajectory'. See below for the
            required keyword arguments for each geometry.
        variables (list): (Optional) Subset of variables to interpolate. Default
            behavior is to interpolate all variables.
378
379
380
381
382
383
384
385
        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
386
387
            interpolate onto, excluding the north and south polar points.
            Defaults to (360,179).
388
389
390
391
392
393
394
395
396
        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).
397
        halfwidths (array-like): Argument for the 'rectprism' geometry. Contains
398
399
400
401
402
403
404
405
406
407
408
409
410
            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.
    """
411
412
413
414
415
    if verbose:
        print('Collecting parameters')

    ## collect the geometry parameters
    geometry_params = {
416
        'geometry':geometry
417
418
419
420
421
    }
    geometry_params.update(kwargs)

    if verbose:
        print('Adding defaults')
422
423
    ## assign defaults for shell and rectprism
    if 'shell' in geometry_params['geometry']:
424
425
426
427
428
429
430
431
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
        geometry_params['npoints'] = geometry_params.get(
            'npoints'
            , (359, 181)
        )
432
    elif 'rectprism' in geometry_params['geometry']:
433
434
435
436
        geometry_params['center'] = geometry_params.get(
            'center'
            , (0.0, 0.0, 0.0)
        )
437
438
439

    ## check that we support the geometry
    geometry_param_names = {
440
441
442
443
        'shell': ('radius',),
        'line': ('r1', 'r2', 'npoints'),
        'rectprism': ('halfwidths', 'npoints'),
        'trajectory': ('trajectory_data', 'trajectory_format')
444
    }
445
446
447
    if geometry_params['geometry'] not in geometry_param_names:
        raise ValueError(f'Geometry {geometry_params["geometry"]} '
                         'not supported!')
448
    ## check that we've gotten the right /required/ geometry arguments
449
    for param in geometry_param_names[geometry_params['geometry']]:
450
451
        if param not in geometry_params:
            raise TypeError(
452
                f'Geometry {geometry_params["geometry"]} '
453
454
455
456
457
                f'requires argument {param}!')

    ## describe the interpolation we're about to do on the data
    if verbose:
        print('Geometry to be interpolated:')
458
        for key, value in geometry_params.items():
459
460
461
462
463
            print(f'\t{key}: {value.__repr__()}')

    ## describe the loaded tecplot data
    if verbose:
        print('Loaded tecplot data with variables:')
464
        print(dataset.variable_names)
465
466

    ## create geometry zone
467
    if 'shell' in geometry_params['geometry']:
468
        geometry_points = _shell_geometry(geometry_params)
469
    elif 'line' in geometry_params['geometry']:
470
        geometry_points = _line_geometry(geometry_params)
471
    elif 'rectprism' in geometry_params['geometry']:
472
        geometry_points = _rectprism_geometry(geometry_params)
473
    elif 'trajectory' in geometry_params['geometry']:
474
475
        if 'batsrus' in geometry_params['trajectory_format']:
            geometry_points = _trajectory_geometry(geometry_params)
476

477
    if ('trajectory' in geometry_params['geometry']
478
            and 'tecplot' in geometry_params['trajectory_format']):
479
        dataset = tecplot.data.load_tecplot(
480
            filenames=geometry_params['trajectory_data']
481
            , read_data_option=tecplot.constant.ReadDataOption.Append
482
        )
483
        dataset.zone(-1).name = 'geometry'
484
    else:
485
        dataset.add_ordered_zone(
486
487
488
            'geometry'
            , geometry_points['npoints']
        )
489
        for i, direction in zip((0, 1, 2), ('X', 'Y', 'Z')):
490
491
            dataset.zone('geometry').values(i)[:] = \
                geometry_points[direction][:]
492

493
    ## interpolate variables on to the geometry
494
    if verbose and variables:
495
496
497
        print('Interpolating variables:')
        for var in variables:
            print(var.name)
498
499
500
    ## dataset.variables('...') will return a generator of variables.
    ## This call will break if `variables` is not recast as a list before
    ## passing it to the function. Why?????
501
    tecplot.data.operate.interpolate_linear(
502
503
504
        destination_zone=dataset.zone('geometry'),
        source_zones=source_zone,
        variables=variables
505
    )
506

507
508
509
510
511
512
513
    ## add variables for shell and trajectory cases
    if 'shell' in  geometry_params['geometry']:
        _add_variable_value(
            dataset,
            'latitude [deg]',
            dataset.zone('geometry'),
            geometry_points['latitude']
514
        )
515
516
517
518
519
        _add_variable_value(
            dataset,
            'longitude [deg]',
            dataset.zone('geometry'),
            geometry_points['longitude']
520
        )
521
522
523
524
525
526
527
    if ('trajectory' in geometry_params['geometry']
            and 'batsrus' in geometry_params['trajectory_format']):
        _add_variable_value(
            dataset,
            'time',
            dataset.zone('geometry'),
            geometry_points['time']
528
        )
529
530
531
532
533

        ## add auxiliary data
    dataset.zone('geometry').aux_data.update(geometry_params)

    return dataset.zone('geometry')