# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Dict, List, Union, Callable
from pathlib import Path
import inspect
import torch
import numpy as np
from modulus.sym.domain.inferencer import PointVTKInferencer
from modulus.sym.domain.constraint import Constraint
from modulus.sym.graph import Graph
from modulus.sym.key import Key
from modulus.sym.node import Node
from modulus.sym.distributed import DistributedManager
from modulus.sym.utils.io import InferencerPlotter
from modulus.sym.utils.io.vtk import var_to_polyvtk, VTKBase, VTKUniformGrid
from modulus.sym.dataset import DictInferencePointwiseDataset
[docs]class VoxelInferencer(PointVTKInferencer):
"""
Inferencer for creating volex representations.
This inferencer works bu creating a uniform mesh of voxels and masking out the ones defined by a callable function.
The result is a voxel based representation of any complex geometery at any resolution.
Parameters
----------
bounds : List[List[int]]
List of domain bounds to form uniform rectangular domain
npoints : List[int]
Resolution of voxels in each domain
nodes : List[Node]
List of Modulus Nodes to unroll graph with.
output_names : List[str]
List of desired outputs.
export_map : Dict[str, List[str]], optional
Export map dictionary with keys that are VTK variables names and values that are lists of output variables. Will use 1:1 mapping if none is provided, by default None
invar : Dict[str, np.array], optional
Dictionary of additional numpy arrays as input, by default {}
mask_fn : Union[Callable, None], optional
Masking function to remove points from inferencing, by default None
mask_value : float, optional
Value to assign masked points, by default Nan
plotter : Plotter, optional
Modulus `Plotter` for showing results in tensorboard., by default None
requires_grad : bool, optional
If automatic differentiation is needed for computing results., by default True
log_iter : bool, optional
Save results to different file each call, by default False
"""
def __init__(
self,
bounds: List[List[int]],
npoints: List[int],
nodes: List[Node],
output_names: List[str],
export_map: Union[None, Dict[str, List[str]]] = None,
invar: Dict[str, np.array] = {}, # Additional inputs
batch_size: int = 1024,
mask_fn: Union[Callable, None] = None,
mask_value: float = np.nan,
plotter=None,
requires_grad: bool = False,
log_iter: bool = False,
model=None,
):
# No export map means one to one with outvars
self.npoints = npoints
if export_map is None:
export_map = {name: name for name in output_names}
coords = ["x", "y", "z"]
input_vtk_map = {coords[i]: coords[i] for i in range(len(bounds))}
# Create uniform grid dataset
vtk_obj = VTKUniformGrid(
bounds=bounds,
npoints=npoints,
export_map=export_map,
)
super().__init__(
vtk_obj,
nodes,
input_vtk_map=input_vtk_map,
output_names=output_names,
invar=invar, # Additional inputs
batch_size=batch_size,
mask_fn=mask_fn,
mask_value=mask_value,
plotter=plotter,
requires_grad=requires_grad,
log_iter=log_iter,
model=model,
)
def _write_results(
self, invar, predvar, name, results_dir, writer, save_filetypes, step
):
# Save batch to vtk/np files
if "np" in save_filetypes:
# Reshape into grid numpy arrays [cin, xdim, ydim, zdim]
np_vars = {}
for key, value in {**invar, **predvar}.items():
shape = self.npoints + [value.shape[1]]
np_vars[key] = np.moveaxis(np.reshape(value, (shape)), -1, 0)
np.savez(results_dir + name, np_vars)
if "vtk" in save_filetypes:
self.vtk_obj.file_dir = Path(results_dir)
self.vtk_obj.file_name = Path(name).stem
if self.log_iter:
self.vtk_obj.var_to_vtk(data_vars={**invar, **predvar}, step=step)
else:
self.vtk_obj.var_to_vtk(data_vars={**invar, **predvar})
# Add tensorboard plots
if self.plotter is not None:
self.plotter._add_figures(
"Inferencers", name, results_dir, writer, step, invar, predvar
)