Coverage for src/arcade_collection/convert/convert_to_contours.py: 100%

1import tarfile 

2 

3import numpy as np 

4from skimage import measure 

5 

6from arcade_collection.output.extract_tick_json import extract_tick_json 

7from arcade_collection.output.get_location_voxels import get_location_voxels 

8 

9ROTATIONS = { 

10 "top": (0, 1, 2), 

11 "side1": (0, 2, 1), 

12 "side2": (2, 1, 0), 

13} 

14"""Axis rotations for different contour views.""" 

15 

16 

17def convert_to_contours( 

18 series_key: str, 

19 locations_tar: tarfile.TarFile, 

20 frame: int, 

21 regions: list[str], 

22 box: tuple[int, int, int], 

23 indices: dict[str, list[int]], 

24) -> dict[str, dict[str, dict[int, list]]]: 

25 """ 

26 Convert data to iso-valued contours. 

27 

28 Contours are calculated using "marching squares" method. Note that these 

29 contours follow the iso-values, which means that a single "square" voxel 

30 will produce a diamond-shaped contour. For the exact outline of a set of 

31 voxels, consider using ``extract_voxel_contours`` from the 

32 ``abm_shape_collection`` package. 

33 

34 Parameters 

35 ---------- 

36 series_key 

37 Simulation series key. 

38 locations_tar 

39 Archive of location data. 

40 frame 

41 Frame number. 

42 regions 

43 List of regions. 

44 box 

45 Size of bounding box. 

46 indices 

47 Map of view to slice indices. 

48 

49 Returns 

50 ------- 

51 : 

52 Map of region, view, and index to contours. 

53 """ 

54 

55 locations = extract_tick_json(locations_tar, series_key, frame, "LOCATIONS") 

56 

57 contours: dict[str, dict[str, dict[int, list]]] = { 

58 region: {view: {} for view in indices} for region in regions 

59 } 

60 

61 for location in locations: 

62 for region in regions: 

63 array = np.zeros(box) 

64 voxels = get_location_voxels(location, region if region != "DEFAULT" else None) 

65 

66 if len(voxels) == 0: 

67 continue 

68 

69 array[tuple(np.transpose(voxels))] = 1 

70 

71 for view in indices: 

72 array_rotated = np.moveaxis(array, [0, 1, 2], ROTATIONS[view]) 

73 

74 for index in indices[view]: 

75 array_slice = array_rotated[:, :, index] 

76 

77 if np.sum(array_slice) == 0: 

78 continue 

79 

80 if index not in contours[region][view]: 

81 contours[region][view][index] = [] 

82 

83 array_contours = [ 

84 contour.tolist() for contour in measure.find_contours(array_slice) 

85 ] 

86 contours[region][view][index].extend(array_contours) 

87 

88 return contours