diff --git a/mcbackend/backends/numpy.py b/mcbackend/backends/numpy.py index cb37f06..6245b4c 100644 --- a/mcbackend/backends/numpy.py +++ b/mcbackend/backends/numpy.py @@ -23,7 +23,7 @@ def grow_append( length = len(target) if length == draw_idx: # Grow the array by 10 % - ngrow = math.ceil(0.1 * length) + ngrow = max(10, math.ceil(0.1 * length)) if rigid[vn]: extension = numpy.empty((ngrow,) + numpy.shape(v)) else: @@ -52,7 +52,7 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non and grow the allocated memory by 10 % when needed. Exceptions are variables with non-rigid shapes (indicated by 0 in the shape tuple) where the correct amount of memory cannot be pre-allocated. - In these cases, and when ``preallocate == 0`` object arrays are used. + In these cases object arrays are used. """ self._var_is_rigid: Dict[str, bool] = {} self._samples: Dict[str, numpy.ndarray] = {} @@ -68,11 +68,11 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non for var in variables: rigid = is_rigid(var.shape) and not var.undefined_ndim and var.dtype != "str" rigid_dict[var.name] = rigid - if preallocate > 0 and rigid: + if rigid: reserve = (preallocate, *var.shape) target_dict[var.name] = numpy.empty(reserve, var.dtype) else: - target_dict[var.name] = numpy.array([None] * preallocate) + target_dict[var.name] = numpy.array([None] * preallocate, dtype=object) super().__init__(cmeta, rmeta) diff --git a/mcbackend/test_backend_numpy.py b/mcbackend/test_backend_numpy.py index f3e49af..86850a3 100644 --- a/mcbackend/test_backend_numpy.py +++ b/mcbackend/test_backend_numpy.py @@ -2,6 +2,7 @@ import hagelkorn import numpy +import pytest from mcbackend.backends.numpy import NumPyBackend, NumPyChain, NumPyRun from mcbackend.core import RunMeta @@ -43,8 +44,9 @@ def test_targets(self): assert chain._samples["changeling"].dtype == object pass - def test_growing(self): - imb = NumPyBackend(preallocate=15) + @pytest.mark.parametrize("preallocate", [0, 75]) + def test_growing(self, preallocate): + imb = NumPyBackend(preallocate=preallocate) rm = RunMeta( rid=hagelkorn.random(), variables=[ @@ -62,19 +64,27 @@ def test_growing(self): ) run = imb.init_run(rm) chain = run.init_chain(0) - assert chain._samples["A"].shape == (15, 2) - assert chain._samples["B"].shape == (15,) - for _ in range(22): + assert chain._samples["A"].shape == (preallocate, 2) + assert chain._samples["B"].shape == (preallocate,) + for _ in range(130): draw = { "A": numpy.random.uniform(size=(2,)), "B": numpy.random.uniform(size=(random.randint(0, 10),)), } chain.append(draw) - # Growth: 15 → 17 → 19 → 21 → 24 - assert chain._samples["A"].shape == (24, 2) - assert chain._samples["B"].shape == (24,) - assert chain.get_draws("A").shape == (22, 2) - assert chain.get_draws("B").shape == (22,) + # NB: Growth algorithm adds max(10, ceil(0.1*length)) + if preallocate == 75: + # 75 → 85 → 95 → 105 → 116 → 128 → 141 + assert chain._samples["A"].shape == (141, 2) + assert chain._samples["B"].shape == (141,) + elif preallocate == 0: + # 10 → 20 → ... → 90 → 100 → 110 → 121 → 134 + assert chain._samples["A"].shape == (134, 2) + assert chain._samples["B"].shape == (134,) + else: + assert False, f"Missing test for {preallocate=}" + assert chain.get_draws("A").shape == (130, 2) + assert chain.get_draws("B").shape == (130,) pass