Coverage for src/abm_shape_collection/construct_mesh_from_coeffs.py: 100%

20 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2024-09-25 19:34 +0000

1import numpy as np 

2import pandas as pd 

3from aicsshparam import shtools 

4from vtk import ( # pylint: disable=no-name-in-module 

5 vtkPolyData, 

6 vtkTransform, 

7 vtkTransformPolyDataFilter, 

8) 

9 

10 

11def construct_mesh_from_coeffs( 

12 coeffs: pd.DataFrame, 

13 order: int, 

14 prefix: str = "", 

15 suffix: str = "", 

16 scale: float = 1.0, 

17) -> vtkPolyData: 

18 """ 

19 Construct a mesh from spherical harmonic coefficients. 

20 

21 Parameters 

22 ---------- 

23 coeffs 

24 Spherical harmonic coefficients. 

25 order 

26 Order of the spherical harmonics coefficient parametrization. 

27 prefix 

28 Prefix string for all coefficient columns. 

29 suffix 

30 Suffix string for all coefficient columns. 

31 scale 

32 Scale factor for mesh points. 

33 

34 Returns 

35 ------- 

36 : 

37 Mesh object. 

38 """ 

39 

40 coeffs_map = np.zeros((2, order + 1, order + 1), dtype=np.float32) 

41 

42 for lix in range(order + 1): 

43 for mix in range(order + 1): 

44 coeffs_map[0, lix, mix] = coeffs[f"{prefix}shcoeffs_L{lix}M{mix}C{suffix}"] 

45 coeffs_map[1, lix, mix] = coeffs[f"{prefix}shcoeffs_L{lix}M{mix}S{suffix}"] 

46 

47 mesh, _ = shtools.get_reconstruction_from_coeffs(coeffs_map) 

48 

49 if scale != 1.0: 

50 transform = vtkTransform() 

51 transform.Scale((scale, scale, scale)) 

52 

53 transform_filter = vtkTransformPolyDataFilter() 

54 transform_filter.SetInputData(mesh) 

55 transform_filter.SetTransform(transform) 

56 transform_filter.Update() 

57 

58 mesh = transform_filter.GetOutput(0) 

59 

60 return mesh