Coverage for subcell_pipeline/simulation/readdy/data_structures.py: 62%

1"""Data structures for ReaDDy simulations.""" 

2 

3from typing import Optional 

4 

5import numpy as np 

6 

7 

8class TopologyData: 

9 """Data class representing a ReaDDy topology of connected particles.""" 

10 

11 uid: int 

12 """Unique ID of the topology from ReaDDy.""" 

13 

14 type_name: str 

15 """ReaDDy type name of the topology.""" 

16 

17 particle_ids: list[int] 

18 """List of unique IDs of each particle in the topology.""" 

19 

20 def __init__(self, uid: int, type_name: str, particle_ids: list[int]): 

21 self.uid = uid 

22 self.type_name = type_name 

23 self.particle_ids = particle_ids 

24 

25 def __str__(self) -> str: 

26 return ( 

27 "Topology(\n" 

28 f" id = {self.uid}\n" 

29 f" type_name = {self.type_name}\n" 

30 f" particles = {self.particle_ids}\n" 

31 ")" 

32 ) 

33 

34 

35class ParticleData: 

36 """Data class representing a ReaDDy particle.""" 

37 

38 uid: int 

39 """Unique ID of the particle from ReaDDy.""" 

40 

41 type_name: str 

42 """ReaDDy type name of the particle.""" 

43 

44 position: np.ndarray 

45 """The x,y,z position of the particle.""" 

46 

47 neighbor_ids: list[int] 

48 """List of unique IDs of each neighbor particle connected by an edge.""" 

49 

50 def __init__( 

51 self, uid: int, type_name: str, position: np.ndarray, neighbor_ids: list[int] 

52 ): 

53 self.uid = uid 

54 self.type_name = type_name 

55 self.position = position 

56 self.neighbor_ids = neighbor_ids 

57 

58 def __str__(self) -> str: 

59 return ( 

60 f"Particle(\n" 

61 f" id = {self.uid}\n" 

62 f" type_name = {self.type_name}\n" 

63 f" position = {self.position}\n" 

64 f" neighbors = {self.neighbor_ids}\n" 

65 ")" 

66 ) 

67 

68 

69class FrameData: 

70 """Data class representing one ReaDDy timestep.""" 

71 

72 time: float 

73 """Current time of the simulation for this frame.""" 

74 

75 topologies: dict[int, TopologyData] 

76 """Mapping of topology ID to a TopologyData for each topology.""" 

77 

78 particles: dict[int, ParticleData] 

79 """Mapping of particle ID to a ParticleData for each particle.""" 

80 

81 edge_ids: list[list[int]] 

82 """List of edges, each is a list of IDs of the two connected particles.""" 

83 

84 def __init__( 

85 self, 

86 time: float, 

87 topologies: Optional[dict[int, TopologyData]] = None, 

88 particles: Optional[dict[int, ParticleData]] = None, 

89 edge_ids: Optional[list[list[int]]] = None, 

90 ): 

91 self.time = time 

92 self.topologies = topologies if topologies is not None else {} 

93 self.particles = particles if particles is not None else {} 

94 self.edge_ids = edge_ids if edge_ids is not None else [] 

95 

96 def __str__(self) -> str: 

97 top_str = "\n" 

98 for top_id in self.topologies: 

99 top_str += f"{top_id} : \n{self.topologies[top_id]}\n" 

100 p_str = "\n" 

101 for p_id in self.particles: 

102 p_str += f"{p_id} : \n{self.particles[p_id]}\n" 

103 return ( 

104 f"Frame(\n" 

105 f" time={self.time}\n" 

106 f" topologies=\n{top_str}\n\n" 

107 f" particles=\n{p_str}\n\n" 

108 ")" 

109 )