Optimally allocating a resource with time-varying demand and cost

Question

I'm working on the following DP which finds the optimal way to allocate a resource. At each time step I can either allocate (0.2 resources) at cost C or not in which case the storage is reduced by the demand. The cost C depends on which time step I do the allocation.

I need to add quite a few more complex features, but this is the bare working version in Python 2.7:

from math import floor
import numpy as np
import time      


def discretise(x,levels = 40,Hmax = 2, Hmin = 0):
    # discretise the levels of storage
    factor = levels / (Hmax-Hmin)
    return floor(x*factor)/factor

# function to track the storage evolution at the stages
def record(D,key,item):
    """Records the value of item in the dictionary d with k"""
    if key in D.keys():
        D[key].append(item)
    else:
        D[key] = [item]

    return D

# Memoization decorator i found online that works
def memoize(f):
    cache = {}

    def memoizedFunction(*args):
        if args not in cache:
            cache[args] = f(*args)
        return cache[args]

    memoizedFunction.cache = cache
    return memoizedFunction

# define demands and cost (Energy)
D = [0.1, 0.1, 0.13, 0.2, 0.4, 0.3, 0.1]
Energy  = [2.0, 5.0,  9.0, 3.0, 7.0, 4.0, 6.0]



@memoize
def val(h,i=0):
    if h < 0 or h > 2: # h < Hmin or h > Hmax:
        return 1000000000
    elif i == T:  # i = T
        return -h*5 # h*m the linear value of h at i = T
    else:
        #           pump off         pump on    
        record(H,i,discretise(h))  

        no_pump = val(discretise(h-D[i]),i+1) # don't allocate
        pump = Energy[i] + val(discretise(h+0.2-D[i]),i+1) # allocate


        v = min(no_pump  , pump   )

        # record optimal policy 
        if v == no_pump:
            pos = 0
        else:
            pos = 1

        record(X,i+1,pos)
        record(V,i+1,v)


        return v


T = 7
H = {}
V = {}
X = {}
ts = time.time()
for ii in xrange(0,20):
    print val(ii/10.0)

te = time.time()

print te - ts

How can I get this code nicer / more pythonic?

Answer 1

Styling

I don't think PEP8 actually says anything about this, but I prefer to have spaces after commas in just about anything. The comma is a way of separating two items, but sometimes it just isn't enough and other times it just looks squashed without the space. You should not, however, put spaces around the = sign in all cases. From PEP8:

Don't use spaces around the = sign when used to indicate a keyword argument or a default parameter value.

You break that rule in the definition of discretise().

Again from PEP8:

Use the function naming rules: lowercase with words separated by underscores as necessary to improve readability.

That is under the title Method Names and Instance Variables. Classes should be PascalCase, but your variables should be as defined above. Your variables, T, H, V, and X should have more descriptive names besides the fact that they should be lowercase. I realize that you use v in val(), but if you gave more descriptive names to your variables, there wouldn't be a clash.

v = min(no_pump  , pump   )

What are you doing with all that whitespace? It should look like this:

v = min(no_pump, pump)

You have many cases of things such as i+1, ii/10.0, etc. PEP8 doesn't actually say anything against that, but my opinion is that you should use whitespace around the operators unless it makes something look confusing. For example, a*b + c*d would look confusing if there were spaces around the asterisks, but in the case of val(ii/10.0), for example, I would prefer it to be val(ii / 10.0).

The logic

def record(D,key,item):
    if key in D.keys():
        D[key].append(item)
    else:
        D[key] = [item]

That function could be simplified to this:

def record(D, key, item):
    D.setdefault(key, []).append(item)

Take a look at the docs. If you wanted to leave it how it was, though, you could still take out the .keys() part of the if statement. Turning a dictionary into an iterator, which in does, turns it into an iterator of its keys.

def memoize(*args):
    cache = {}

    def memoizedFunction(*args):
        if args not in cache:
            cache[args] = f(*args)
        return cache[args]

    memoizedFunction.cache = cache
    return memoizedFunction

Using memoizedFunction.cache = cache isn't helping you here. You aren't even using it. You might be better off defining a class:

class Memoize:
    def __init__(self, function):
        self.function = function
        self.cache = {}

    def __call__(self, *args):
        if args not in self.cache:
            self.cache[args] = self.function(*args)
        return self.cache[args]

I changed the name to Memoize to comply with PEP8, but that just means that you need to change @memoize to @Memoize.

        #           pump off         pump on    

It isn't clear to me what that comment is saying. The pump on part goes beyond the end of the next line. If you meant it to be a comment on a different line, you should move it.

v = min(no_pump  , pump   )

# record optimal policy 
if v == no_pump:
    pos = 0
else:
    pos = 1

Here would be a good use of the cmp() function:

v = cmp(no_pump, pump)
if v <= 0:
    pos = 0
else:
    pos = 1

ts = time.time()
...
te = time.time()

I can tell from the last line, print te - ts that those are meant to stand for time start and time end, but why not spell it out? In this case it's fairly easy to tell what they are for, but it's a good habit to make your variable names very clear. I would define them as start_time and end_time.