Sym v1.2.0

Source code for modulus.sym.domain.inferencer.voxel

# 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 import InferencerPlotter
from 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 )
© Copyright 2023, NVIDIA Modulus Team. Last updated on Jan 25, 2024.