""" a Structure is a combination of a Geometry with a material (and an optional mesh order) """
from typing import Dict, List, Tuple, Union, overload
import numpy as np
from pydantic.v1 import Field
from .base_model import BaseModel
from .geometries import Geometry2D, Geometry3D
from .materials import Material
DEFAULT_MESH_ORDER = 5
[docs]def Structure(
*,
material: Material,
geometry: Union[Geometry2D, Geometry3D],
mesh_order: int = DEFAULT_MESH_ORDER,
):
kwargs = {
"material": material,
"geometry": geometry,
"mesh_order": mesh_order,
}
if isinstance(geometry, Geometry2D):
return Structure2D(**kwargs)
else:
return Structure3D(**kwargs)
[docs]class Structure2D(BaseModel):
"""a `Structure2D` is an association between a `Geometry2D` and a `Material`"""
material: Material = Field(description="the material of the structure")
geometry: Geometry2D = Field(description="the geometry of the structure")
mesh_order: int = Field(
default=DEFAULT_MESH_ORDER, description="the mesh order of the structure"
)
def _visualize(self):
color = self.material.meta.get("color", None)
return self.geometry._visualize(color=color)
[docs]class Structure3D(BaseModel):
"""a `Structure3D` is an association between a `Geometry3D` and a `Material`"""
material: Material = Field(description="the material of the structure")
geometry: Geometry3D = Field(description="the geometry of the structure")
mesh_order: int = Field(
default=DEFAULT_MESH_ORDER, description="the mesh order of the structure"
)
def _project(self, z) -> List[Structure2D]:
geometry_2d = self.geometry._project(z)
structs = []
for geom in geometry_2d:
struct = Structure2D(
material=self.material,
geometry=geom,
mesh_order=self.mesh_order,
)
structs.append(struct)
return structs
def _lumadd(self, sim, env, unit=1e-6, xyz="yzx"):
material_name = self.material._lumadd(sim, env, unit)
self.geometry._lumadd(sim, material_name, self.mesh_order, unit, xyz)
def _trimesh(self, color=None, scale=None):
return self.geometry._trimesh(
color=(color or self.material.meta.get("color")),
scale=scale,
)
def _visualize(self, scale=None):
return self._trimesh(scale=scale).show()
def _visualize_structures(structures: List[Structure3D], scale=None):
"""easily visualize a collection (list) of `Structure3D` objects"""
from trimesh.scene import Scene # fmt: skip
from trimesh.transformations import rotation_matrix # fmt: skip
scene = Scene(
geometry=[s._trimesh(scale=scale) for s in _sort_structures(structures)]
)
scene.apply_transform(rotation_matrix(np.pi - np.pi / 6, (0, 1, 0)))
return scene.show()
@overload
def _sort_structures(structures: List[Structure3D]) -> List[Structure3D]:
...
@overload
def _sort_structures(structures: List[Structure2D]) -> List[Structure2D]:
...
def _sort_structures(
structures: Union[List[Structure3D], List[Structure2D]]
) -> Union[List[Structure2D], List[Structure3D]]:
struct_info = [(s.mesh_order, -i, s) for i, s in enumerate(structures)]
sorted_struct_info = sorted(struct_info, key=lambda I: (I[0], I[1]), reverse=True)
return [s for _, _, s in sorted_struct_info] # type: ignore
@overload
def _classify_structures_by_mesh_order_and_material(
structures: List[Structure3D], materials: Dict[Material, int]
) -> Dict[Tuple[int, int], List[Structure3D]]:
...
@overload
def _classify_structures_by_mesh_order_and_material(
structures: List[Structure2D], materials: Dict[Material, int]
) -> Dict[Tuple[int, int], List[Structure2D]]:
...
def _classify_structures_by_mesh_order_and_material(
structures: Union[List[Structure3D], List[Structure2D]],
materials: Dict[Material, int],
) -> Union[
Dict[Tuple[int, int], List[Structure2D]], Dict[Tuple[int, int], List[Structure3D]]
]:
structures = _sort_structures(structures)
structures_dict = {}
for structure in structures:
mo = structure.mesh_order
mat = materials[structure.material]
if (mo, mat) not in structures_dict:
structures_dict[mo, mat] = []
structures_dict[mo, mat].append(structure)
return structures_dict
