Basic plotting and forest creation functions for the forest fire model in CMSE 201 Spring 2017

create_and_plot_forest.py

```python
import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np

def make_forest( height, width, density ):
    '''
    This function will take in a width, height, and density specified by the user and use those values
    to generate a forest (2-d numpy array) matching those characteristics. A 0 will represent an empty
    space, a 1 will represent a tree
    
    Inputs:
    ======
    height (int) : A whole number representing the desired height of the forest
    
    width (int) : A whole number representing the desired width of the forest
    
    density (float) : A value between 0 and 1 which expresses the percent of spaces in the array to populate
        with trees.
        
    Outputs:
    ======
    new_forest (2d ndarray) : Forest as a 2d numpy array populated with trees at a given density
    '''
    
    # Make empty forest
    new_forest = np.zeros( shape=(width, height) )
    
    # Put down trees
    for i in range(width):
        for j in range(height):
            a_random_num = np.random.random()
            if a_random_num < density:
                new_forest[i][j] = 1
    
    return new_forest

def display_forest( forest, state_colors ):
    '''
    This function will nicely visualize an input forest. The function also takes as input a list of
    colors which correspond to the colors desired for each state, and will raise an error if the number
    of states and the number of colors provided do not match.
    
    Inputs:
    ======
    forest (2d ndarray) : A forest with some number of states, preferably starting at 0 and counting up
        by one
        
    state_colors (list) : A list a colors which will assign a particular color to each state. As an
        example: if my forest has two states with "empty space" being 0 and "tree" being 1, then we might
        pass in ['white', 'green']
        
    Outputs:
    ======
    None : This function just plots your forest, it doesn't return anything
    '''
    
    # Check for a good number of colors
    num_states = len(np.unique(forest))
    if len(state_colors) != num_states:
        raise ValueError("Number of colors does not equal number of states")
    
    # Make a custom colormap
    cmap = colors.ListedColormap(state_colors)
    
    # Plot the forest
    plt.imshow(forest, cmap=cmap, interpolation='none')
    plt.axis('off')
    plt.show()
    
    return None

# Simple example of generating and plotting
ff = make_forest(25,20,.7)
colors_list = ['white', 'green']
display_forest(ff,colors_list)
```

![Image of a forest plotted with white squares for empty and green squares for trees][1]

[1]: http://d.pr/i/BNcS+