Merge pull request #20 from TomMonks/dev

PATCH: 0.3.3

setup.py

@@ -1,8 +1,8 @@
 import setuptools
 
-# patch ed - build error due to imporing numpy into __init__
+# patched - build error due to imporing numpy into __init__
 # from sim_tools import __version__
-VERSION = "0.3.2"
+VERSION = "0.3.3"
 
 
 # Read in the requirements.txt file

sim_tools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.3.2'
+__version__ = '0.3.3'
 __author__ = 'Thomas Monks'
 
 from . import datasets, distributions, time_dependent, ovs

sim_tools/distributions.py

@@ -813,7 +813,13 @@ def sample(self, size: Optional[int] = None) -> float | np.ndarray:
             Number of samples to return. If integer then
             numpy array returned.
         """
-        return self.rng.choice(self.values, p=self.probabilities, size=size).item()
+        sample =  self.rng.choice(self.values, p=self.probabilities, size=size)
+
+        if size is None:
+            return sample.item()
+        else:
+            return sample
+
 
 
 class TruncatedDistribution(Distribution):

tests/test_dists.py

@@ -123,10 +123,10 @@ def test_discrete():
     assert type(d.sample()) == int
 
 
-def test_discrete():
-    d = dists.Discrete(values=[1, 2, 3], freq=[95, 3, 2], random_seed=SEED_1)
-    assert type(d.sample()) == int
 
+def test_discrete_multiple():
+    d = dists.Discrete(values=[1, 2, 3], freq=[95, 3, 2], random_seed=SEED_1)
+    assert len(d.sample(size=100)) == 100
 
 def test_truncated_type():
     d1 = dists.Normal(10, 1, random_seed=SEED_1)