Source code for

import warnings

import dataclasses
import enum
import numpy as np
from typing import Optional

from import CenteredGrid
from import EvaluationOptions
from import Transform
from .grid_modules import RegularGrid, CustomGrid, Sections
from .grid_modules.topography import Topography

[docs] @dataclasses.dataclass class Grid:
[docs] class GridTypes(enum.Flag): OCTREE = 2 ** 0 DENSE = 2 ** 1 CUSTOM = 2 ** 2 TOPOGRAPHY = 2 ** 3 SECTIONS = 2 ** 4 CENTERED = 2 ** 5 NONE = 2 ** 10
# ? What should we do with the extent? _octree_grid: Optional[RegularGrid] = None _dense_grid: Optional[RegularGrid] = None _custom_grid: Optional[CustomGrid] = None _topography: Optional[Topography] = None _sections: Optional[Sections] = None _centered_grid: Optional[CenteredGrid] = None values: np.ndarray = np.empty((0, 3)) length: np.ndarray = np.empty(0) _active_grids = GridTypes.NONE _transform: Optional[Transform] = None _octree_levels: int = -1
[docs] def __init__(self, extent=None, resolution=None): # Init basic grid empty if extent is not None and resolution is not None: self.dense_grid = RegularGrid(extent, resolution)
@classmethod def init_octree_grid(cls, extent, octree_levels): grid = cls() grid._octree_grid = RegularGrid( extent=extent, resolution=np.array([2 ** octree_levels] * 3), ) grid.active_grids |= grid.GridTypes.OCTREE grid._update_values() return grid @classmethod def init_dense_grid(cls, extent, resolution): return cls(extent, resolution) def __str__(self): active_grid_types_str = [g_type for g_type in self.GridTypes if self.active_grids & g_type] grid_summary = [f"{g_type} (active: {getattr(self, g_type + '_grid_active')}): {len(getattr(self, g_type + '_grid').values)} points" for g_type in active_grid_types_str] grid_summary_str = "\n".join(grid_summary) return f"Grid Object:\n{grid_summary_str}" @property def transform(self) -> Transform: if self.dense_grid is not None: return self.dense_grid.transform elif self.octree_grid is not None: return self.octree_grid.transform else: return Transform.init_neutral() @transform.setter def transform(self, value: Transform): self._transform = value @property def extent(self): if self.dense_grid is not None: return self.dense_grid.extent elif self.octree_grid is not None: return self.octree_grid.extent else: raise AttributeError('Extent is not defined') @property def corner_min(self): return self.extent[::2] @property def corner_max(self): return self.extent[1::2] @property def bounding_box(self): extents = self.extent # Define 3D points of the bounding box corners based on extents bounding_box_points = np.array([[extents[0], extents[2], extents[4]], # min x, min y, min z [extents[0], extents[2], extents[5]], # min x, min y, max z [extents[0], extents[3], extents[4]], # min x, max y, min z [extents[0], extents[3], extents[5]], # min x, max y, max z [extents[1], extents[2], extents[4]], # max x, min y, min z [extents[1], extents[2], extents[5]], # max x, min y, max z [extents[1], extents[3], extents[4]], # max x, max y, min z [extents[1], extents[3], extents[5]]]) # max x, max y, max z return bounding_box_points @property def active_grids(self): return self._active_grids @active_grids.setter def active_grids(self, value): self._active_grids = value self._update_values() @property def dense_grid(self) -> RegularGrid: return self._dense_grid @dense_grid.setter def dense_grid(self, value): self._dense_grid = value self.active_grids |= self.GridTypes.DENSE self._update_values() @property def octree_grid(self): return self._octree_grid @octree_grid.setter def octree_grid(self, value): raise AttributeError('Octree grid is not allowed to be set directly. Use init_octree_grid instead') def set_octree_grid(self, regular_grid: RegularGrid, evaluation_options: EvaluationOptions): regular_grid_resolution = regular_grid.resolution # Check all directions has the same res if not np.all(regular_grid_resolution == regular_grid_resolution[0]): raise AttributeError('Octree resolution must be isotropic') octree_levels = int(np.log2(regular_grid_resolution[0])) # Check if octrree levels are the same if octree_levels != evaluation_options.number_octree_levels: raise AttributeError('Regular grid resolution does not match octree levels. Resolution must be 2^n') self._octree_grid = regular_grid self.active_grids |= self.GridTypes.OCTREE self._update_values() def set_octree_grid_by_levels(self, octree_levels: int, evaluation_options: EvaluationOptions, extent: Optional[np.ndarray] = None): if extent is None: extent = self.extent self._octree_grid = RegularGrid( extent=extent, resolution=np.array([2 ** octree_levels] * 3), ) evaluation_options.number_octree_levels = octree_levels self.active_grids |= self.GridTypes.OCTREE self._update_values() @property def octree_levels(self): return self._octree_levels @octree_levels.setter def octree_levels(self, value): raise AttributeError('Octree levels are not allowed to be set directly. Use set_octree_grid instead') @property def custom_grid(self): return self._custom_grid @custom_grid.setter def custom_grid(self, value): self._custom_grid = value self.active_grids |= self.GridTypes.CUSTOM self._update_values() @property def topography(self): return self._topography @topography.setter def topography(self, value): self._topography = value self.active_grids |= self.GridTypes.TOPOGRAPHY self._update_values() @property def sections(self): return self._sections @sections.setter def sections(self, value): self._sections = value self.active_grids |= self.GridTypes.SECTIONS self._update_values() @property def centered_grid(self): return self._centered_grid @centered_grid.setter def centered_grid(self, value): self._centered_grid = value self.active_grids |= self.GridTypes.CENTERED self._update_values() @property def regular_grid(self): dense_grid_exists_and_active = self.dense_grid is not None and self.GridTypes.DENSE in self.active_grids octree_grid_exists_and_active = self.octree_grid is not None and self.GridTypes.OCTREE in self.active_grids if dense_grid_exists_and_active and octree_grid_exists_and_active: raise AttributeError('Both dense_grid and octree_grid are active. This is not possible.') elif self.dense_grid is not None: return self.dense_grid elif self.octree_grid is not None: return self.octree_grid else: return None # noinspection t def _update_values(self): values = [] if self.GridTypes.OCTREE in self.active_grids: if self.octree_grid is None: raise AttributeError('Octree grid is active but not defined') values.append(self.octree_grid.values) if self.GridTypes.DENSE in self.active_grids: if self.dense_grid is None: raise AttributeError('Dense grid is active but not defined') values.append(self.dense_grid.values) if self.GridTypes.CUSTOM in self.active_grids: if self.custom_grid is None: raise AttributeError('Custom grid is active but not defined') values.append(self.custom_grid.values) if self.GridTypes.TOPOGRAPHY in self.active_grids: if self.topography is None: raise AttributeError('Topography grid is active but not defined') values.append(self.topography.values) if self.GridTypes.SECTIONS in self.active_grids: if self.sections is None: raise AttributeError('Sections grid is active but not defined') values.append(self.sections.values) if self.GridTypes.CENTERED in self.active_grids: if self.centered_grid is None: raise AttributeError('Centered grid is active but not defined') values.append(self.centered_grid.values) self.values = np.concatenate(values) return self.values def get_section_args(self, section_name: str): # TODO: This method should be part of the sections # assert type(section_name) is str, 'Only one section type can be retrieved' l0, l1 = self.get_grid_args('sections') where = np.where(self.sections.names == section_name)[0][0] return l0 + self.sections.length[where], l0 + self.sections.length[where + 1]