python class for dynamically resizeable struct of arrays (SoA) with numpy arrays

soa_np.py

# MIT License
# 
# Copyright (c) 2022 xaedes
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


class SoA:
    def __init__(self, fields, capacity=1):
        self._capacity = capacity
        self._size = 0
        self.fields = dict(fields)
        self.arrays = dict([
            (name, np.empty(self._capacity, dtype))
            for name, dtype in self.fields.items()
        ])
        
    def __getattr__(self, name):
        return self.arrays[name]
                    
    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_arrays = dict([
                (name, np.empty(new_capacity, dtype))
                for name, dtype in self.fields.items()
            ])
            
            for name in self.fields.keys():
                new_arrays[name][:self._capacity] = self.arrays[name][:]
                self.arrays[name] = new_arrays[name]
                
            self._capacity = new_capacity
            self._size = new_size
            
        return self
    
    def shrink_to_fit(self):
        for name in self.fields.keys():
            self.arrays[name] = self.arrays[name][:self._size]
        self._capacity = self._size
        return self
    
    def concat(self, other):
        my_size = self.size()
        other_size = other.size()
        self.resize(my_size + other_size)
        for name in self.fields.keys():
            self.arrays[name][my_size:my_size+other_size] = other.arrays[name][:other_size]
        return self
    
    def append(self, **kwargs):
        k = self.size()
        self.resize(k + 1)
        for name, val in kwargs.items():
            self.arrays[name][k] = val
        return self
    
    def append0(self, **kwargs):
        k = self.size()
        self.resize(k + 1)
        for name in self.fields.keys():
            self.arrays[name][k] = 0
        for name, val in kwargs.items():
            self.arrays[name][k] = val
        return self    
        
class Trajectory(SoA):
    def __init__(self, capacity=1):
        fields = {
            "position": np.dtype((np.float64, (2,))),
            "yaw": np.dtype((np.float64, (1,))),
            "yawrate": np.dtype((np.float64, (1,)))
        }
        super().__init__(fields, capacity)