test.py

import numpy as np

from constants import UNIT_CUBE, IHAT, JHAT, KHAT
from math_utils import unit, centroid

from geometry.hex import Hex
from geometry.quad import Quad
from geometry.point import Point
from mesh.mesh_abc import Mesh
from space.lookup import QuantizedLookup, KDTreeLookup
from foam_definitions import *

from time import perf_counter
from functools import wraps
from grading.propagation import propagate_grading_elements
from grading.elements import GradingSegment

def benchmark(func_):

  @wraps(func_)
  def wrapper(*args, **kwargs):

    s = perf_counter()
    res = func_(*args, **kwargs)
    e = perf_counter()
    print(f"{func_.__name__} took {e - s:.3f} seconds: args={args}, kwargs={kwargs}")
    return res
  
  return wrapper

class MyMesh(Mesh):
  def write():
    return None
  
def gen_cube_tensor(max_ax=5):

  cube_array = np.array(UNIT_CUBE, np.float64)
  max_ax = np.arange(0, max_ax)
  X, Y, Z = np.meshgrid(max_ax, max_ax, max_ax, indexing="ij")
  offsets = np.stack([X, Y, Z], axis=-1).reshape(-1, 3)
  cube_tensor = offsets[:, np.newaxis, :] + cube_array.copy()[np.newaxis, :, :]
  return cube_tensor

@benchmark
def test_lookup(max_ax=5):

  cube_tensor = gen_cube_tensor(max_ax)
  vectors = cube_tensor.reshape((-1,3))
  qtz = QuantizedLookup(vectors)
  kdt = KDTreeLookup(vectors)
  return qtz, kdt

@benchmark
def test_grid(max_ax=5):

  cube_tensor = gen_cube_tensor(max_ax)
  geoms = map(lambda coors: Hex(coors), cube_tensor)
  mesh = MyMesh()
  added: list[Hex] = []
  for geom in geoms:
    reg = mesh.add(geom)
    added.append(reg)

  return mesh

@benchmark
def test_grid_enhanced(max_ax = 5):
  cube_tensor = gen_cube_tensor(max_ax)
  geoms = map(lambda coors: Hex(coors), cube_tensor)
  mesh, geoms = MyMesh.from_iterable(geoms)
  return mesh, geoms

@benchmark
def test_point():
  p = Point((0, 1, 2))
  return p

@benchmark
def test_quad():

  q = Quad(UNIT_CUBE[:4])
  return q

@benchmark
def test_hex():

  h = Hex(UNIT_CUBE)
  h.copy().transform.rotate(0.5, (1,1,0))
  return h

@benchmark
def test_propagation(max_ax):

  cube_tensor = gen_cube_tensor(max_ax)
  geoms = map(lambda coors: Hex(coors), cube_tensor)
  mesh, geoms = MyMesh.from_iterable(geoms)
  geoms[0]._tf.e1 = GradingSegment(1, 1, 1)

  propagate_grading_elements(mesh)
  count = 0
  for geom in geoms:
    if geom._tf.e1 is not None:
      count += 1
  print("Succesfully propagated", count)
  
  return mesh

if __name__ == '__main__':
  print(f'Executing {__file__} ...')

  from space.vectorspace import VectorSpace

  vs = VectorSpace(UNIT_CUBE)
  print(vs.data, vs._lookup.index_map)
  vs[1] = (0,)*3
  print(vs.data, vs._lookup.index_map)
  
  from geometry.composite.flat.hgrid import CartesianGrid

  cg = CartesianGrid()
  print(cg)