Add a pseudo diffusivity to the model

diffusion/diffusion_model.py

@@ -1,23 +1,27 @@
 """Model diffusion through random motion of particles."""
 
 import agentpy as ap
+import numpy as np
 
+BASE_WEIGHT = 10
+
+directions = [
+    (0, 0),
+    (1, 0),
+    (0, 1),
+    (-1, 0),
+    (0, -1),
+]
+weights = np.ones(len(directions), dtype=int) * BASE_WEIGHT
 
-class Particle(ap.Agent):
 
-    direction = [
-        (0, 0),
-        (1, 0),
-        (0, 1),
-        (-1, 0),
-        (0, -1),
-    ]
+class Particle(ap.Agent):
 
     def setup(self):
         self.displacement = [0, 0]
 
     def set_random_displacement(self):
-        self.displacement = self.model.random.choice(Particle.direction)
+        self.displacement = self.model.random.choices(directions, weights=weights).pop()
 
 
 class DiffusionModel(ap.Model):
@@ -29,8 +33,12 @@ def setup(self):
         positions = (self.p.initial_location,) * self.p.agents
         self.grid.add_agents(self.agents, positions=positions)
 
+        self.set_diffusivity()
         self.histogram = None
 
+    def set_diffusivity(self):
+        weights[1:] = self.p.diffusivity
+
     def update(self):
         self.histogram = self.grid.apply(len, field="agents")
         self.agents.set_random_displacement()

examples/config-big.yaml

@@ -5,4 +5,5 @@ DiffusionModel:
   n_cols: 100
   n_rows: 100
   initial_location: [50, 50]
+  diffusivity: 2
   bmi_version: "2.0"

examples/config.yaml

@@ -5,4 +5,5 @@ DiffusionModel:
   n_cols: 10
   n_rows: 10
   initial_location: [5, 5]
+  diffusivity: 7
   bmi_version: "2.0"

examples/explore-model.ipynb

@@ -61,6 +61,7 @@
     "The model uses 100 agents.\n",
     "Initially, all the agents are placed at position `[5,5]` on a 10 x 10 grid.\n",
     "By default, the model runs for 20 time steps.\n",
+    "It has a pseudo diffusivity of 7.\n",
     "\n",
     "Initialize the model, then call its `setup` and `update` methods to get a view of the initial conditions."
    ]

tests/diffusion_model_test.py

@@ -12,6 +12,7 @@
     "n_cols": 6,
     "n_rows": 5,
     "initial_location": [2, 2],
+    "diffusivity": 5,
     "bmi_version": "2.0",
 }
 
@@ -32,6 +33,7 @@ def test_model_has_parameters():
     assert m.p.steps == PARAMETERS["steps"]
     assert m.p.n_cols == PARAMETERS["n_cols"]
     assert m.p.n_rows == PARAMETERS["n_rows"]
+    assert m.p.diffusivity == PARAMETERS["diffusivity"]
 
 
 def test_model_setup():
@@ -53,3 +55,10 @@ def test_particle_displacement():
     final_positions = list(m.grid.positions.values())
     for i in range(len(m.agents)):
         assert math.dist(initial_positions[i], final_positions[i]) <= 1.0
+
+
+def test_zero_diffusivity():
+    PARAMETERS["diffusivity"] = 0
+    m = DiffusionModel(PARAMETERS)
+    m.run(steps=1)
+    assert m.histogram[tuple(PARAMETERS["initial_location"])] == PARAMETERS["agents"]