parameters: (grid: Grid, filled_fraction=0.25)
return_statement: return grid.with_values(filled_t)
docstring: Treats channel dimensions as batch dimensions.
docstring_summary: Treats channel dimensions as batch dimensions.
docstring_tokens: Treats channel dimensions as batch dimensions .
function: def discretize(grid: Grid, filled_fraction=0.25):
      """ Treats channel dimensions as batch dimensions. """
      import numpy as np
      data = math.reshaped_native(grid.values, [grid.shape.non_spatial, grid.shape.spatial])
      ranked_idx = np.argsort(data, axis=-1)
      filled_idx = ranked_idx[:, int(round(grid.shape.spatial.volume * (1 - filled_fraction))):]
      filled = np.zeros_like(data)
      np.put_along_axis(filled, filled_idx, 1, axis=-1)
      filled_t = math.reshaped_tensor(filled, [grid.shape.non_spatial, grid.shape.spatial])
      return grid.with_values(filled_t)
function_tokens: def discretize ( grid : Grid , filled_fraction = 0.25 ) : import numpy as np data = math . reshaped_native ( grid . values , [ grid . shape . non_spatial , grid . shape . spatial ] ) ranked_idx = np . argsort ( data , axis = - 1 ) filled_idx = ranked_idx [ : , int ( round ( grid . shape . spatial . volume * ( 1 - filled_fraction ) ) ) : ] filled = np . zeros_like ( data ) np . put_along_axis ( filled , filled_idx , 1 , axis = - 1 ) filled_t = math . reshaped_tensor ( filled , [ grid . shape . non_spatial , grid . shape . spatial ] ) return grid . with_values ( filled_t )