wavelet_miniapp.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. r"""
  2. This script is miniapp that acts as a sim code. It uses `vtkRTAnalyticSource`
  3. to generate a temporal dataset which are treated as the simulation generate
  4. data. One can pass command line arguments to control the parameters of the
  5. generated data. Additional command line arguments can be passed to provide a
  6. Catalyst co-processing script for in situ processing.
  7. This is a good demo for integration Catalyst support in any Python-based
  8. simulation code. Look for use of the 'bridge' module from 'paraview.catalyst'
  9. in the code to see how.
  10. """
  11. import math, argparse, time, os.path
  12. #----------------------------------------------------------------
  13. # parse command line arguments
  14. parser = argparse.ArgumentParser(\
  15. description="Wavelet MiniApp for Catalyst testing")
  16. parser.add_argument("-t", "--timesteps", type=int,
  17. help="number of timesteps to run the miniapp for (default: 100)", default=100)
  18. parser.add_argument("--size", type=int,
  19. help="number of samples in each coordinate direction (default: 101)", default=101)
  20. parser.add_argument("-s", "--script", type=str, action="append",
  21. help="path(s) to the Catalyst script(s) to use for in situ processing. Can be a "
  22. ".py file or a Python package zip or directory",
  23. required=True)
  24. parser.add_argument("--script-version", type=int,
  25. help="choose Catalyst analysis script version explicitly, otherwise it "
  26. "will be determined automatically. When specifying multiple scripts, this "
  27. "setting applies to all scripts.", default=0)
  28. parser.add_argument("-d", "--delay", type=float,
  29. help="delay (in seconds) between timesteps (default: 0.0)", default=0.0)
  30. parser.add_argument("-c", "--channel", type=str,
  31. help="Catalyst channel name (default: input)", default="input")
  32. #----------------------------------------------------------------
  33. # A helper function that creates a VTK dataset per timestep/per rank.
  34. def create_dataset(timestep, args, piece, npieces):
  35. # We'll use vtkRTAnalyticSource to generate our dataset
  36. # to keep things simple.
  37. from vtkmodules.vtkImagingCore import vtkRTAnalyticSource
  38. wavelet = vtkRTAnalyticSource()
  39. ext = (args.size - 1) // 2
  40. wavelet.SetWholeExtent(-ext, ext, -ext, ext, -ext, ext)
  41. wholeExtent = wavelet.GetWholeExtent()
  42. # put in some variation in the point data that changes with timestep
  43. wavelet.SetMaximum(255+200*math.sin(timestep))
  44. # using 'UpdatePiece' lets us generate a subextent based on the
  45. # 'piece' and 'npieces'; thus works seamlessly in distributed and
  46. # non-distributed modes
  47. wavelet.UpdatePiece(piece, npieces, 0)
  48. # typically, here you'll have some adaptor code that converts your
  49. # simulation data into a vtkDataObject subclass. In this example,
  50. # there's nothing to do since we're directly generating a
  51. # vtkDataObject.
  52. dataset = wavelet.GetOutputDataObject(0)
  53. return (dataset, wholeExtent)
  54. #----------------------------------------------------------------
  55. # Here's our simulation main loop
  56. def main(args):
  57. """The main loop"""
  58. # initialize Catalyst
  59. from paraview.catalyst import bridge
  60. from paraview import print_info, print_warning
  61. bridge.initialize()
  62. # add analysis script
  63. for script in args.script:
  64. bridge.add_pipeline(script, args.script_version)
  65. # Some MPI related stuff to figure out if we're running with MPI
  66. # and if so, on how many ranks.
  67. try:
  68. from mpi4py import MPI
  69. comm = MPI.COMM_WORLD
  70. rank = comm.Get_rank()
  71. num_ranks = comm.Get_size()
  72. except ImportError:
  73. print_warning("missing mpi4py, running in serial (non-distributed) mode")
  74. rank = 0
  75. num_ranks = 1
  76. numsteps = args.timesteps
  77. for step in range(numsteps):
  78. if args.delay > 0:
  79. time.sleep(args.delay)
  80. # assume simulation time starts at 0
  81. timevalue = step/float(numsteps)
  82. if rank == 0:
  83. print_info("timestep: {0}/{1}; timevalue: {2}".format(step+1, numsteps, timevalue))
  84. dataset, wholeExtent = create_dataset(step, args, rank, num_ranks)
  85. # "perform" coprocessing. results are outputted only if
  86. # the passed in script says we should at timevalue/step
  87. bridge.coprocess(timevalue, step, dataset, name=args.channel, wholeExtent=wholeExtent)
  88. del dataset
  89. del wholeExtent
  90. # finalize Catalyst
  91. bridge.finalize()
  92. if __name__ == "__main__":
  93. args = parser.parse_args()
  94. main(args)