Classroom exercise: energy calculation¶

import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline

density =  np.array([0, 0, 3, 5, 8, 4, 2, 1])
fig, ax = plt.subplots()
ax.bar(np.arange(len(density)), density)
ax.xrange=[-0.5, len(density)-0.5]
ax.set_ylabel("Particle count $n_i$")
ax.set_xlabel("Position $i$")

%%bash
mkdir -p diffusion
touch diffusion/__init__.py

%%writefile diffusion/model.py
def energy(density, coeff=1.0):
    """ 
    Energy associated with the diffusion model

    Parameters
    ----------

    density: array of positive integers
        Number of particles at each position i in the array
    coeff: float
        Diffusion coefficient.
    """
  # implementation goes here

%%writefile diffusion/test_model.py
from .model import energy
def test_energy():
    """ Optional description for reporting """
    # Test something

%%bash
cd diffusion
pytest

%%writefile diffusion/model.py
"""  Simplistic 1-dimensional diffusion model """

def energy(density):
    """ 
    Energy associated with the diffusion model
    
    :Parameters:
    
    density: array of positive integers
      Number of particles at each position i in the array/geometry
    """
    from numpy import array, any, sum

    # Make sure input is an numpy array
    density = array(density)

    # ...of the right kind (integer). Unless it is zero length, 
    #    in which case type does not matter.
      
    if density.dtype.kind != 'i' and len(density) > 0:
        raise TypeError("Density should be a array of *integers*.")
    # and the right values (positive or null)
    if any(density < 0):
        raise ValueError("Density should be an array of *positive* integers.")
    if density.ndim != 1:
        raise ValueError("Density should be an a *1-dimensional*" + 
                         "array of positive integers.")
    
    return sum(density * (density - 1))

%%writefile diffusion/test_model.py
""" Unit tests for a diffusion model """

from pytest import raises
from .model import energy

def test_energy_fails_on_non_integer_density():
    with raises(TypeError) as exception: 
        energy([1.0, 2, 3])
    
def test_energy_fails_on_negative_density():
    with raises(ValueError) as exception: energy(
            [-1, 2, 3])
        
def test_energy_fails_ndimensional_density():
    with raises(ValueError) as exception: energy(
            [[1, 2, 3], [3, 4, 5]])

def test_zero_energy_cases():
    # Zero energy at zero density
    densities = [ [], [0], [0, 0, 0] ]
    for density in densities: 
        assert energy(density) == 0

def test_derivative():
    from numpy.random import randint

    # Loop over vectors of different sizes (but not empty)
    for vector_size in randint(1, 1000, size=30): 

        # Create random density of size N
        density = randint(50, size=vector_size)

        # will do derivative at this index
        element_index = randint(vector_size)

        # modified densities
        density_plus_one = density.copy()
        density_plus_one[element_index] += 1

        # Compute and check result
        # d(n^2-1)/dn = 2n
        expected = (2.0 * density[element_index] 
                    if density[element_index] > 0 
                    else 0 )
        actual = energy(density_plus_one) - energy(density) 
        assert expected == actual

def test_derivative_no_self_energy():
    """ If particle is alone, then its participation to energy is zero """
    from numpy import array

    density = array([1, 0, 1, 10, 15, 0])
    density_plus_one = density.copy()
    density[1] += 1 

    expected = 0
    actual = energy(density_plus_one) - energy(density) 
    assert expected == actual

%%bash
cd diffusion
pytest

%%bash
cd diffusion
pytest --cov="diffusion"

%%bash
cd diffusion
pytest --cov="diffusion" --cov-report html