tecplottools.py 22.7 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
32
33
34
35
def _get_variable_names(variables):
    """For getting the names of a group of Tecplot variables"""
    return [var.name for var in variables]


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

45
46
47
48
49
50
51
52
    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])
53

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


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


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


113
def _trajectory_geometry(geometry_params: dict) -> dict:
114
    """Returns a dict containing points for the described trajectory geometry.
115
116

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


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


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

    This is helpful for accessing Tecplot variables.

266
267
268
269
    Args:
        variable_name (str): A string which may contain the meta-characters * ?
        [ or ].

270
271
272
273
274
275
276
277
278
279
280
281
282
    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.
    """
283
284
285
286
287
288
289
    translation = {
        '[':'[[]',
        ']':'[]]',
        '*':'[*]',
        '?':'[?]'
    }
    return variable_name.translate(str.maketrans(translation))
290
291


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

    Args:
303
        tecplot_dataset (): The dataset to save.
304
305
306
307
308
309
310
311
        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.
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345

    Examples:
        ```python
        import tecplot
        import swmfpy.tecplottools as tpt

        ## load a dataset and configure the layout
        dataset = tecplot.data.load_tecplot(
            '3d__mhd_1_n00000100.plt')
        frame = tecplot.active_frame()
        frame.plot_type = tecplot.constant.PlotType.Cartesian3D
        plot = frame.plot()

        ## set the vector variables
        plot.vector.u_variable = dataset.variable(4)
        plot.vector.v_variable = dataset.variable(5)
        plot.vector.w_variable = dataset.variable(6)

        ## seed and extract a streamtrace
        plot.streamtraces.add(
            seed_point=(1.5, 1.0, 0.0),
            stream_type=tecplot.constant.Streamtrace.VolumeLine
        )
        streamtrace_zones = plot.streamtraces.extract()
        new_zone = next(streamtrace_zones)

        ## write the new zone to hdf5 format
        tpt.write_zone(
            tecplot_dataset=dataset,
            tecplot_zone=new_zone,
            write_as='hdf5',
            filename='streamtrace.h5'
        )
        ```
346
    """
347
    if verbose and variables:
348
        print('Saving variables:')
349
        print(_get_variable_names(variables).__repr__())
350
351
352
    aux_data = tecplot_zone.aux_data.as_dict()
    if verbose:
        print('Attaching auxiliary data:')
353
        print(aux_data.__repr__())
354
        print('Saving as:')
355
    ## save zone
356
357
358
    if 'hdf5' in write_as:
        if verbose:
            print('hdf5')
359
360
        if not variables:
            variables = list(tecplot_dataset.variables())
361
362
363
364
365
366
367
        _save_hdf5(
            filename,
            aux_data,
            tecplot_zone,
            variables
        )
    elif 'csv' in write_as:
368
369
        if not variables:
            variables = list(tecplot_dataset.variables())
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
        _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
        )
389
390
    else:
        raise ValueError(f'File type {write_as} not supported!')
391
392
393
394
    if verbose:
        print(f'Wrote {filename}')


395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
def _assign_geometry_defaults(
        geometry: str,
        default_params: dict,
        geometry_params: dict
):
    """Checks parameters with defaults and assigns them.

    If the parameters are already set nothing will change.

    Args:
        geometry (str): String identifying the geometry to look for.
        default_params (dict): Dictionary of the default parameters.
        geomatry_params (dict): Dictionary in which to look for and set
            parameters.
    """
    if geometry in geometry_params['geometry']:
        for key, value in default_params.items():
            geometry_params[key] = geometry_params.get(
                key,
                value
            )
    return geometry_params


def _check_geometry_requirements(
        geometry_requirements: dict,
        geometry_params: dict
):
    """Checks that the required kwargs for the given geometry have been set.
    """
    if geometry_params['geometry'] not in geometry_requirements:
        raise ValueError(f'Geometry {geometry_params["geometry"]} '
                         'not supported!')
    for param in geometry_requirements[geometry_params['geometry']]:
        if param not in geometry_params:
            raise TypeError(
                f'Geometry {geometry_params["geometry"]} '
                f'requires argument {param}!')

434
435
436
def interpolate_zone_to_geometry(
        dataset
        , source_zone
437
        , geometry: str
438
        , variables: list = None
439
440
        , verbose: bool = False
        , **kwargs
441
):
442
443
444
    """Interpolates Tecplot binary data onto various geometries.

    Args:
445
446
        dataset: The loaded Tecplot dataset.
        source_zone: The Tecplot zone to interpolate onto the geometry.
447
        geometry (str): Type of geometry for interpolation. Supported geometries
448
449
450
451
            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.
452
453
454
455
456
457
458
459
        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
460
461
            interpolate onto, excluding the north and south polar points.
            Defaults to (360,179).
462
463
464
465
466
467
468
469
470
        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).
471
        halfwidths (array-like): Argument for the 'rectprism' geometry. Contains
472
473
474
475
476
477
478
479
480
481
482
483
            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.
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540

    Examples:
        ```python
        import tecplot
        import swmfpy.tecplottools as tpt

        tecplot.session.connect(port=7600)

        ## Load a dataset and configure the layout.
        dataset = tecplot.data.load_tecplot('3d__mhd_1_n00000100.plt')

        ## Create a new zone with the specified geometry
        ## and interpolate data onto it.

        ## geometry: shell
        tpt.interpolate_zone_to_geometry(
            dataset=dataset,
            source_zone=dataset.zone(0),
            geometry='shell',
            radius=1.5,
            npoints=(4,3)
        )

        ## geometry: line
        tpt.interpolate_zone_to_geometry(
            dataset=dataset,
            source_zone=dataset.zone(0),
            geometry='line',
            r1=[1.0, 0.0, 0.0],
            r2=[3.0, 0.0, 0.0],
            npoints=100
        )

        ## geometry: rectangular prism
        new_zone = tpt.interpolate_zone_to_geometry(
            dataset=dataset,
            source_zone=dataset.zone(0),
            geometry='rectprism',
            center=[0.0, 0.0, 0.0],
            halfwidths=[2.0, 2.0, 2.0],
            npoints=[5, 5, 5]
        )

        ## geometry: spacecraft trajectory as specified for the
        ## BATSRUS #SATELLITE command
        tpt.interpolate_zone_to_geometry(
            dataset=dataset,
            source_zone=dataset.zone(0),
            geometry='trajectory',
            trajectory_format='batsrus',
            trajectory_data='./test_data/satellite_e4.dat'
        )

        ## The new zones are labeled with the name of the geometry and can be
        ## manipulated in the Tecplot GUI.
        ```

541
    """
542
543
544
545
546
    if verbose:
        print('Collecting parameters')

    ## collect the geometry parameters
    geometry_params = {
547
        'geometry':geometry
548
549
550
551
552
    }
    geometry_params.update(kwargs)

    if verbose:
        print('Adding defaults')
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
    geometry_params = _assign_geometry_defaults(
        'shell',
        {
            'center':(0.0, 0.0, 0.0),
            'npoints':(359, 181)
        },
        geometry_params
    )
    geometry_params = _assign_geometry_defaults(
        'rectprism',
        {
            'center':(0.0, 0.0, 0.0),
        },
        geometry_params
    )
568

569
570
571
572
573
574
575
576
577
    _check_geometry_requirements(
        {
            'shell': ('radius',),
            'line': ('r1', 'r2', 'npoints'),
            'rectprism': ('halfwidths', 'npoints'),
            'trajectory': ('trajectory_data', 'trajectory_format')
        },
        geometry_params
    )
578
579

    if verbose:
580
        ## describe the interpolation we're about to do on the data
581
        print('Geometry to be interpolated:')
582
        for key, value in geometry_params.items():
583
584
            print(f'\t{key}: {value.__repr__()}')

585
        ## describe the loaded tecplot data
586
        print('Loaded tecplot data with variables:')
587
        print(dataset.variable_names)
588
589

    ## create geometry zone
590
    if 'shell' in geometry_params['geometry']:
591
        geometry_points = _shell_geometry(geometry_params)
592
    elif 'line' in geometry_params['geometry']:
593
        geometry_points = _line_geometry(geometry_params)
594
    elif 'rectprism' in geometry_params['geometry']:
595
        geometry_points = _rectprism_geometry(geometry_params)
596
    elif 'trajectory' in geometry_params['geometry']:
597
598
        if 'batsrus' in geometry_params['trajectory_format']:
            geometry_points = _trajectory_geometry(geometry_params)
599
600
601
602
603
604
605
606
607
608
609
610
611
            dataset.add_ordered_zone(
                geometry_params['geometry']
                , geometry_points['npoints']
            )
            for i, direction in zip((0, 1, 2), ('X', 'Y', 'Z')):
                dataset.zone(geometry_params['geometry']).values(i)[:] = \
                     geometry_points[direction][:]
        elif 'tecplot' in geometry_params['trajectory_format']:
            dataset = tecplot.data.load_tecplot(
                filenames=geometry_params['trajectory_data']
                , read_data_option=tecplot.constant.ReadDataOption.Append
            )
            dataset.zone(-1).name = geometry_params['geometry']
612

613
    ## interpolate variables on to the geometry
614
    if verbose and variables:
615
        print('Interpolating variables:')
616
617
        print(_get_variable_names(variables).__repr__())

618
619
620
    ## 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?????
621
    tecplot.data.operate.interpolate_linear(
622
        destination_zone=dataset.zone(geometry_params['geometry']),
623
624
        source_zones=source_zone,
        variables=variables
625
    )
626

627
628
629
630
631
    ## add variables for shell and trajectory cases
    if 'shell' in  geometry_params['geometry']:
        _add_variable_value(
            dataset,
            'latitude [deg]',
632
            dataset.zone(geometry_params['geometry']),
633
            geometry_points['latitude']
634
        )
635
636
637
        _add_variable_value(
            dataset,
            'longitude [deg]',
638
            dataset.zone(geometry_params['geometry']),
639
            geometry_points['longitude']
640
        )
641
642
643
644
645
    if ('trajectory' in geometry_params['geometry']
            and 'batsrus' in geometry_params['trajectory_format']):
        _add_variable_value(
            dataset,
            'time',
646
            dataset.zone(geometry_params['geometry']),
647
            geometry_points['time']
648
        )
649
        geometry_params['time_seconds_since'] = '1970-01-01T00:00:00Z'
650

651
652
    ## add auxiliary data
    dataset.zone(geometry_params['geometry']).aux_data.update(geometry_params)
653

654
    return dataset.zone(geometry_params['geometry'])