Example for dynamically resizeable SoA (struct of arrays) class using numpy arrays

example_dynamic_soa_np.py

class Trajectory:
    def __init__(self, capacity):
        self._capacity = capacity
        self._size = 0
        self.position = np.empty(shape=(self._capacity, 2), dtype=np.float64)
        self.yaw = np.empty(shape=(self._capacity, 1), dtype=np.float64)
        self.yawrate = np.empty(shape=(self._capacity, 1), dtype=np.float64)

    def size(self):
        return self._size
    
    def capacity(self):
        return self._capacity
    
    def clear(self):
        self.resize(0)
        return self
    
    def resize(self, new_size):
        if new_size <= self._capacity:
            self._size = new_size
        else:
            delta = new_size - self._capacity
            new_capacity = int(math.ceil(self._capacity * 1.5)) + delta
            
            new_position = np.empty(shape=(new_capacity, 2), dtype=np.float64)
            new_yaw = np.empty(shape=(new_capacity, 1), dtype=np.float64)
            new_yawrate = np.empty(shape=(new_capacity, 1), dtype=np.float64)
            
            new_position[:self._capacity] = self.position
            new_yaw[:self._capacity] = self.yaw
            new_yawrate[:self._capacity] = self.yawrate
            
            self.position = new_position
            self.yaw = new_yaw
            self.yawrate = new_yawrate
            self._capacity = new_capacity
            self._size = new_size
            
        return self
    
    def shrink_to_fit(self):
        self.position = self.position[:self._size]
        self.yaw = self.yaw[:self._size]
        self.yawrate = self.yawrate[:self._size]
        self._capacity = self._size
        return self
    
    def concat(self, other_trajectory):
        my_size = self.size()
        other_size = other_trajectory.size()
        self.resize(my_size + other_size)
        self.position[my_size:my_size+other_size] = other_trajectory.position[:other_size]
        self.yaw[my_size:my_size+other_size] = other_trajectory.yaw[:other_size]
        self.yawrate[my_size:my_size+other_size] = other_trajectory.yawrate[:other_size]
        return self