1import numpy as np
  2from scipy.ndimage.measurements import label
  3
  4class Forest(object):
  5    """ Forest can grow trees which eventually die."""
  6    def __init__(self, size=(150,150), p_sapling=0.0025):
  7        self.size = size
  8        self.trees = np.zeros(self.size, dtype=bool)
  9        self.p_sapling = p_sapling
 10        
 11    def __repr__(self):
 12        my_repr = "{}(size={})".format(self.__class__.__name__, self.size)
 13        return my_repr
 14    
 15    def __str__(self):
 16        return self.__class__.__name__
 17    
 18    @property
 19    def num_cells(self):
 20        """Number of cells available for growing trees"""
 21        return np.prod(self.size)
 22    
 23    @property
 24    def tree_fraction(self):
 25        """
 26        Fraction of trees
 27        """
 28        num_trees = self.trees.sum()
 29        return float(num_trees) / self.num_cells
 30    
 31    def _rand_bool(self, p):
 32        """
 33        Random boolean distributed according to p, less than p will be True
 34        """
 35        return np.random.uniform(size=self.trees.shape) < p
 36    
 37    def grow_trees(self):
 38        """
 39        Growing trees.
 40        """
 41        growth_sites = self._rand_bool(self.p_sapling)
 42        self.trees[growth_sites] = True    
 43        
 44    def advance_one_step(self):
 45        """
 46        Advance one step
 47        """
 48        self.grow_trees()
 49
 50class BurnableForest(Forest):
 51    """
 52    Burnable forest support fires
 53    """    
 54    def __init__(self, p_lightning=5.0e-6, **kwargs):
 55        super(BurnableForest, self).__init__(**kwargs)
 56        self.p_lightning = p_lightning        
 57        self.fires = np.zeros((self.size), dtype=bool)
 58    
 59    def advance_one_step(self):
 60        """
 61        Advance one step
 62        """
 63        super(BurnableForest, self).advance_one_step()
 64        self.start_fires()
 65        self.burn_trees()
 66        
 67    @property
 68    def fire_fraction(self):
 69        """
 70        Fraction of fires
 71        """
 72        num_fires = self.fires.sum()
 73        return float(num_fires) / self.num_cells
 74    
 75    def start_fires(self):
 76        """
 77        Start of fire.
 78        """
 79        lightning_strikes = (self._rand_bool(self.p_lightning) & 
 80            self.trees)
 81        self.fires[lightning_strikes] = True
 82    
 83    def burn_trees(self):    
 84        pass
 85    
 86class SlowBurnForest(BurnableForest):
 87    def burn_trees(self):
 88        """
 89        Burn trees.
 90        """
 91        fires = np.zeros((self.size[0] + 2, self.size[1] + 2), dtype=bool)
 92        fires[1:-1, 1:-1] = self.fires
 93        north = fires[:-2, 1:-1]
 94        south = fires[2:, 1:-1]
 95        east = fires[1:-1, :-2]
 96        west = fires[1:-1, 2:]
 97        new_fires = (north | south | east | west) & self.trees
 98        self.trees[self.fires] = False
 99        self.fires = new_fires
100
101class InstantBurnForest(BurnableForest):
102    def burn_trees(self):
103        # 起火点
104        strikes = self.fires
105        # 找到连通区域
106        groves, num_groves = label(self.trees)
107        fires = set(groves[strikes])
108        self.fires.fill(False)
109        # 将与着火点相连的区域都烧掉
110        for fire in fires:
111            self.fires[groves == fire] = True
112        self.trees[self.fires] = False
113        self.fires.fill(False)

 1forest = Forest()
 2sb_forest = SlowBurnForest()
 3ib_forest = InstantBurnForest()
 4
 5forests = [forest, sb_forest, ib_forest]
 6
 7tree_history = []
 8
 9for i in xrange(1500):
10    for fst in forests:
11        fst.advance_one_step()
12    tree_history.append(tuple(fst.tree_fraction for fst in forests))

1import matplotlib.pyplot as plt
2%matplotlib inline
3
4plt.figure(figsize=(10,6))
5
6plt.plot(tree_history)
7plt.legend([f.__str__() for f in forests])
8
9plt.show()