fusion-energy / openmc_regular_mesh_plotter

A Python package for plotting OpenMC regular mesh tally results with underlying geometry from neutronics simulations.
MIT License
10 stars 1 forks source link

adding vtk writing #6

Closed shimwell closed 3 years ago

shimwell commented 3 years ago

perhaps a vtk writing ability can be added


def write_3d_mesh_tally_to_vtk(
    xs: np.linspace,
    ys: np.linspace,
    zs: np.linspace,
    tally_data: List[float],
    error_data: Optional[List[float]] = None,
    outfile: Optional[str] = "3d_mesh_tally_data.vtk",
    tally_label: Optional[str] = "3d_mesh_tally_data",
) -> str:
    """Converts regular 3d data into a vtk file for visualizing the data.
    Programs that can visualize vtk files include Paraview
    https://www.paraview.org/ and VisIt
    https://wci.llnl.gov/simulation/computer-codes/visit

    Arguments:
        xs: A numpy array containing evenly spaced numbers from the lowest x
            coordinate value to the highest x coordinate value.
        ys: A numpy array containing evenly spaced numbers from the lowest y
            coordinate value to the highest y coordinate value.
        zs: A numpy array containing evenly spaced numbers from the lowest z
            coordinate value to the highest z coordinate value.
        tally_data: A list of data values to assign to the vtk dataset.
        error_data: A list of error data values to assign to the vtk dataset.
        outfile: The filename of the output vtk file.
        tally_label: The name to assign to the dataset in the vtk file.

    Returns:
        str: the filename of the file produced
    """

    vtk_box = vtk.vtkRectilinearGrid()

    vtk_box.SetDimensions(len(xs), len(ys), len(zs))

    vtk_x_array = vtk.vtkDoubleArray()
    vtk_x_array.SetName("x-coords")
    vtk_x_array.SetArray(xs, len(xs), True)
    vtk_box.SetXCoordinates(vtk_x_array)

    vtk_y_array = vtk.vtkDoubleArray()
    vtk_y_array.SetName("y-coords")
    vtk_y_array.SetArray(ys, len(ys), True)
    vtk_box.SetYCoordinates(vtk_y_array)

    vtk_z_array = vtk.vtkDoubleArray()
    vtk_z_array.SetName("z-coords")
    vtk_z_array.SetArray(zs, len(zs), True)
    vtk_box.SetZCoordinates(vtk_z_array)

    tally = np.array(tally_data)
    tally_data = vtk.vtkDoubleArray()
    tally_data.SetName(tally_label)
    tally_data.SetArray(tally, tally.size, True)

    if error_data is not None:
        error = np.array(error_data)
        error_data = vtk.vtkDoubleArray()
        error_data.SetName("error_tag")
        error_data.SetArray(error, error.size, True)

    vtk_box.GetCellData().AddArray(tally_data)
    vtk_box.GetCellData().AddArray(error_data)

    writer = vtk.vtkRectilinearGridWriter()

    writer.SetFileName(outfile)

    writer.SetInputData(vtk_box)

    print("Writing %s" % outfile)

    writer.Write()

    return outfile
shimwell commented 3 years ago

This is being solved in another package https://github.com/fusion-energy/openmc_mesh_tally_to_vtk