Coverage for src/arcade_collection/input/generate_setup_file.py: 100%

1from __future__ import annotations 

2 

3import xml.etree.ElementTree as ET 

4from typing import TYPE_CHECKING 

5 

6import numpy as np 

7 

8if TYPE_CHECKING: 

9 import pandas as pd 

10 

11DEFAULT_POPULATION_ID = "X" 

12"""Default population ID used in setup file.""" 

13 

14 

15def generate_setup_file( 

16 samples: pd.DataFrame, margins: tuple[int, int, int], terms: list[str] 

17) -> str: 

18 """ 

19 Create ARCADE setup file from samples, margins, and CPM Hamiltonian terms. 

20 

21 Initial number of cells is determined by number of unique ids in samples. 

22 Regions are included if samples contains valid regions. 

23 

24 Parameters 

25 ---------- 

26 samples 

27 Sample cell ids and coordinates. 

28 margins 

29 Margin size in x, y, and z directions. 

30 terms 

31 List of Potts Hamiltonian terms for setup file. 

32 

33 Returns 

34 ------- 

35 : 

36 Contents of ARCADE setup file. 

37 """ 

38 

39 init = len(samples["id"].unique()) 

40 bounds = calculate_sample_bounds(samples, margins) 

41 regions = ( 

42 samples["region"].unique() 

43 if "region" in samples.columns and not samples["region"].isna().all() 

44 else None 

45 ) 

46 return make_setup_file(init, bounds, terms, regions) 

47 

48 

49def calculate_sample_bounds( 

50 samples: pd.DataFrame, margins: tuple[int, int, int] 

51) -> tuple[int, int, int]: 

52 """ 

53 Calculate transformed sample bounds including margin. 

54 

55 Parameters 

56 ---------- 

57 samples 

58 Sample cell ids and coordinates. 

59 margins 

60 Margin size in x, y, and z directions. 

61 

62 Returns 

63 ------- 

64 : 

65 Bounds in x, y, and z directions. 

66 """ 

67 

68 mins = (min(samples.x), min(samples.y), min(samples.z)) 

69 maxs = (max(samples.x), max(samples.y), max(samples.z)) 

70 

71 bound_x, bound_y, bound_z = np.subtract(maxs, mins) + np.multiply(2, margins) + 3 

72 

73 return (bound_x, bound_y, bound_z) 

74 

75 

76def make_setup_file( 

77 init: int, 

78 bounds: tuple[int, int, int], 

79 terms: list[str], 

80 regions: list[str] | None = None, 

81) -> str: 

82 """ 

83 Create ARCADE setup file. 

84 

85 Parameters 

86 ---------- 

87 init 

88 Number of initial cells. 

89 bounds 

90 Bounds in x, y, and z directions. 

91 regions 

92 List of regions. 

93 terms 

94 List of Potts Hamiltonian terms for setup file. 

95 

96 Returns 

97 ------- 

98 : 

99 Contents of ARCADE setup file. 

100 """ 

101 

102 root = ET.fromstring("<set></set>") # noqa: S314 

103 series = ET.SubElement( 

104 root, 

105 "series", 

106 { 

107 "name": "ARCADE", 

108 "interval": "1", 

109 "start": "0", 

110 "end": "0", 

111 "dt": "1", 

112 "ds": "1", 

113 "ticks": "1", 

114 "length": str(int(bounds[0])), 

115 "width": str(int(bounds[1])), 

116 "height": str(int(bounds[2])), 

117 }, 

118 ) 

119 

120 potts = ET.SubElement(series, "potts") 

121 for term in terms: 

122 ET.SubElement(potts, "potts.term", {"id": term}) 

123 

124 agents = ET.SubElement(series, "agents") 

125 populations = ET.SubElement(agents, "populations") 

126 population = ET.SubElement(populations, "population", {"id": "X", "init": str(init)}) 

127 

128 if regions is not None: 

129 for region in regions: 

130 ET.SubElement(population, "population.region", {"id": region}) 

131 

132 ET.indent(root, space=" ", level=0) 

133 return ET.tostring(root, encoding="unicode")