Calculate middle values faster

Question

I want to improve the performance of this code:

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy
import random
import math

random.seed(42)

class TrainModel:
    def __init__(cls, scheduling_interval, trains):
        cls.scheduling_interval = scheduling_interval
        cls.trains = trains


class Train:
    def __init__( cls, legs):
        cls.legs = legs


class Leg: # Teilfahrt
    def __init__(cls, travel_time, earliest_departure_time, latest_departure_time,
                    power_profile):
        cls.travel_time = travel_time # Fahrzeit
        cls.earliest_departure_time = earliest_departure_time # Früheste Abfraht
        cls.latest_departure_time = latest_departure_time # Späteste Abfahrt
        cls.power_profile = power_profile # Leistungsprofil

# Test model
scheduling_interval = 60*24 # one day in minutes
train_number = 10
trains = []
for i in xrange(train_number):
    legs = []
    leg_number = random.randint(5,10)
    for j in xrange(leg_number):
        travel_time = random.randint(5,40)
        earliest_departure_time = random.randint(0,scheduling_interval-100)
        latest_departure_time = earliest_departure_time+7
        power_profile = numpy.random.rand(travel_time*60+1)
        legs.append(Leg(travel_time, earliest_departure_time, latest_departure_time, power_profile))
    train = Train(legs)
    trains.append(train)

train_model = TrainModel(scheduling_interval, trains)

def calc_middle(t, profile):
    x_min = int(min((901-1)*(t+1), len(profile)-1))
    sum_m = sum(a for a in profile[(901-1)*t+1:x_min] if a > 0)
    sum_m += max(0.5 * profile[(901-1)*t], 0)
    sum_m += max(0.5 * profile[x_min], 0)
    return sum_m

storage = {}

# Slow loops profiling shows that most of the time is spend in the sum function of calc_middle
number_of_middle_intervals = int(math.ceil((train_model.scheduling_interval*60+1)/901))
for t in xrange(len(train_model.trains)):
    for r in xrange(len(train_model.trains[t].legs)):
        current_train = train_model.trains[t].legs[r]
        for j in xrange(current_train.latest_departure_time - current_train.earliest_departure_time + 1):
            full_profile = numpy.zeros(train_model.scheduling_interval*60+1)
            dim_start = (current_train.earliest_departure_time+j)*60
            dim_end = (current_train.earliest_departure_time+j+current_train.travel_time)*60+1
            full_profile[dim_start:dim_end] = full_profile[dim_start:dim_end] + current_train.power_profile.clip(0)
            power_value_per_middle_intervall = []
            for i in xrange(number_of_middle_intervals):
                x_min = int(min((901-1)*(i+1), len(full_profile)-1))
                x_start = (901-1)*i
                if x_min < dim_start or x_start > dim_end:
                    tmp_sum = 0
                else:
                    tmp_sum = calc_middle(i, full_profile)
                power_value = tmp_sum

                power_value_per_middle_intervall.append(power_value)
            storage[t,r,j] = power_value_per_middle_intervall

The loops does the following:

It loops through all trains and all legs of each train. Then for each departure time of the current leg it creates a big array full_profile, which will hold for all time points in the scheduling interval the power consumption. It will put the small array current_train.power_profile at the specific position.

Then it needs to calculate the middle values of 901 consecutive points, by the formula of calc_middle. (this means sum up all the 899 inner points and add 0.5 times the other outer 2 points, the x_min is needed in case the last interval is shorter than 901 points). This needs to be done for all middle intervals, but you can improve the computation with the condition if x_min < dim_start or x_start > dim_end: because you know that the middle interval will be zero then.

The next departure time of the same leg, will have the position, where the small array is inserted be shifted by 60 indexes, this means that the middle intervals which are affected might change. In general the big array is much bigger than the small array.

Note the big array is always scheduling_interval*60+1 big and the small array travel_time*60+1.

EDIT

I included now a sample test model. I wrote it that you can vary the model size by increasing the train_number. The parameters are all realistic, except the train_number is now 10, but it can be up to 24000 in general.

Answer 1

Running a profiler over your code shows that most of the time is spent in this line in calc_middle():

sum_m = sum(a for a in profile[(901-1)*t+1:x_min] if a > 0)

So if we want to make this go faster, we need to speed up this line. Some suggestions:

1. Pre-filter the negative elements with numpy.clip().

   The same array is coming through this function many times – wouldn’t it be cheaper to filter out negative elements just once?

   We always get full_profile being passed in here, so add

   full_profile = full_profile.clip(0)
   

   before entering the for i in xrange(number_of_middle_intervals) loop.

   This lets you reduce calc_middle to:

   def calc_middle(t, profile):
       x_min = int(min((901-1)*(t+1), len(profile)-1))
       profile = profile.clip(0)
       return sum(profile[(901-1)*t:x_min+1])
   

   There's a nice little saving.

2. Use numpy.sum().

   Since the slice is now a pure numpy array, rather than a Python generator expression, we can call numpy.sum() here for extra speed:

   def calc_middle(t, profile):
       x_min = int(min((901-1)*(t+1), len(profile)-1))
       profile = profile.clip(0)
       return numpy.sum(profile[(901-1)*t:x_min+1])
   

---

Reprofiling the new code highlights another line which is chewing up cycles: computing x_min, e.g.:

x_min = int(min((901-1)*(t+1), len(profile)-1))

You can make a few tweaks for speed here:

• Reuse values where possible
• Don't unnecessarily cast to int()
• Only compute this value if you actually need to

Knowing that this line is quite a sink, here’s a rewrite of the innermost for loop that tries to speed it up:

x_start = 900 * i
if x_start > dim_end:
    tmp_sum = 0
else:
    x_min = min(901 * i, len(full_profile)-1)
    if x_min < dim_start:
        tmp_sum = 0
    else:
        tmp_sum = numpy.sum(full_profile[x_start:x_min+1])

---

There are probably more optimisations to be made; this is just based on an hour or so poking around before bed. In informal testing, it takes the time to run with train_number=24 from ~2.4s to ~0.6s. Better, but probably not yet suitable for 24000 trains.

---

---

I thought about it a bit more this morning. Here are a few more things I found. None of these are as significant, but they’re useful savings.

• If you compute x_start as:

  x_start = 0
  for i in xrange(number_of_middle_intervals):
      # do stuff
      x_start += 900
  

  rather than

  for i in xrange(number_of_middle_intervals):
      x_start = 900 * i
  

  I get a small saving. Saves about 0.2s when running with train_number=200 (from 4.69s to 4.5s).

• Since x_start is increasing with every iteration of the loop, when you get x_start > dim_end, you know that this will always be true for the current value of the j-loop. So you can fill in the rest of power_value_per_middle_interval with zeroes:

  if x_start > dim_end:
      remaining_length = len(power_value_per_middle_intervall) - number_of_middle_intervals
      power_value_per_middle_intervall += [0] * remaining_length
      break
  

  4.5s ~> 4.12s.

Answer 2

Try to keep everything in functions. Local variables aren't just nicer to work with; they're also marginally faster.

Split model generation into a new function.

Consider calc_middle. I change this under alexwlchan's suggestions of hoisting out clip to

def calc_middle(i, profile):
    start = 900 * i
    end = start + 900
    return (
        profile[start] / 2 +
        profile[start+1:end].sum() +
        profile[end] / 2
    )

Then we merge it into the loop:

power_per_interval = numpy.zeros(number_of_middle_intervals)
for i in xrange(number_of_middle_intervals):
    start = 900 * i
    end = start + 900

    if dim_start <= end and start <= dim_end:
        power_value = (
            full_profile[start] / 2 +
            full_profile[start+1:end].sum() +
            full_profile[end] / 2
        )
        power_per_interval[i] = power_value
storage[t,r,j] = power_per_interval

Then we can remove the if by cropping the xrange:

idx_start = dim_start // 900 - 1
idx_end   = dim_end   // 900 + 1
for i in xrange(idx_start, idx_end):
    ...

We can clean things up a bit above:

dim_start = (current_train.earliest_departure_time+j)*60
dim_end   = dim_start + current_train.travel_time * 60 + 1

full_profile = numpy.zeros(train_model.scheduling_interval*60+1)
full_profile[dim_start:dim_end] = current_train.power_profile.clip(0)

Then avoid some mess by using full_profile = current_train.power_profile.clip(0) and offsetting the indexes with appropriate clipping.

dim_start = (current_train.earliest_departure_time+j)*60
dim_end   = dim_start + current_train.travel_time * 60 + 1

full_profile = current_train.power_profile.clip(0)

power_per_interval = numpy.zeros(number_of_middle_intervals)
idx_start = dim_start // 900 - 1
idx_end   = dim_end   // 900 + 1
for i in xrange(idx_start, idx_end):
    start = 900 * i - dim_start
    end = start + 900

    power_value = 0
    if 0 <= start < len(full_profile):
        power_value += full_profile[start] / 2
    if 0 <= end:
        power_value += full_profile[max(0, start+1):end].sum()
    if 0 <= end <= len(full_profile):
        power_value += full_profile[end] / 2

    power_per_interval[i] = power_value

storage[t,r,j] = power_per_interval

This took some care, but it works OK. Then we can remove the sum by using a single cumulative sum outside the loop:

current_train = train_model.trains[t].legs[r]
full_profile = current_train.power_profile.clip(0)
cumulative = numpy.add.accumulate(full_profile)

for j in xrange(current_train.latest_departure_time - current_train.earliest_departure_time + 1):
    dim_start = (current_train.earliest_departure_time+j)*60
    dim_end   = dim_start + current_train.travel_time * 60 + 1

    power_per_interval = numpy.zeros(number_of_middle_intervals)
    idx_start = dim_start // 900 - 1
    idx_end   = dim_end   // 900 + 1
    for i in xrange(idx_start, idx_end):
        start = 900 * i - dim_start
        end = start + 900

        power_value = 0
        if 0 <= start < len(full_profile):
            power_value += full_profile[start] / 2

        if 0 <= start < len(cumulative):
            power_value -= cumulative[start]
        elif len(cumulative) <= start:
            power_value -= cumulative[-1]

        if 0 <= end - 1 < len(cumulative):
            power_value += cumulative[end - 1]
        elif len(cumulative) <= end - 1:
            power_value += cumulative[-1]

        if 0 <= end <= len(full_profile):
            power_value += full_profile[end] / 2

        power_per_interval[i] = power_value

    storage[t,r,j] = power_per_interval

We can now vectorize the inner loop, but it actually slows us down since the loop only runs 2-3 times.

It's likely this can be improved upon if you're willing to change the storage to a dictionary of 2D-arrays of size number_of_middle_intervals by latest_departure_time - earliest_departure_time + 1, but that's not really to spec.

By this point, the actual analysis stage takes under 0.01 seconds on my computer. For 1000 trains, the program runs in about 0.62 seconds; your original takes about 36 (or about 58x as long).

24000 trains takes 15 seconds. It's not amazingly fast, and if needed I can make it go faster, but it's at least usable.

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy
import random
import math

class TrainModel:
    def __init__(cls, scheduling_interval, trains):
        cls.scheduling_interval = scheduling_interval
        cls.trains = trains


class Train:
    def __init__( cls, legs):
        cls.legs = legs


class Leg: # Teilfahrt
    def __init__(cls, travel_time, earliest_departure_time, latest_departure_time,
                    power_profile):
        cls.travel_time = travel_time # Fahrzeit
        cls.earliest_departure_time = earliest_departure_time # Früheste Abfraht
        cls.latest_departure_time = latest_departure_time # Späteste Abfahrt
        cls.power_profile = power_profile # Leistungsprofil


def generate_test_data(train_number):
    random.seed(42)
    numpy.random.seed(234)

    # Test model
    scheduling_interval = 60 * 24 # one day in minutes
    trains = []
    for i in xrange(train_number):
        legs = []
        leg_number = random.randint(5,10)
        for j in xrange(leg_number):
            travel_time = random.randint(5,40)
            earliest_departure_time = random.randint(0,scheduling_interval-100)
            latest_departure_time = earliest_departure_time+7
            power_profile = numpy.random.rand(travel_time*60+1)
            legs.append(Leg(travel_time, earliest_departure_time, latest_departure_time, power_profile))
        train = Train(legs)
        trains.append(train)

    return TrainModel(scheduling_interval, trains)


def generate_storage(train_model):
    storage = {}

    number_of_middle_intervals = int(math.ceil((train_model.scheduling_interval*60+1)/901))
    for t, train in enumerate(train_model.trains):
        for l, current_train in enumerate(train.legs):
            full_profile = current_train.power_profile.clip(0)
            cumulative = numpy.add.accumulate(full_profile)

            times = xrange(current_train.earliest_departure_time, current_train.latest_departure_time + 1)
            for j, dim_start in enumerate(times):
                dim_start *= 60
                dim_end   = dim_start + current_train.travel_time * 60 + 1

                power_per_interval = numpy.zeros(number_of_middle_intervals)
                idx_start = dim_start // 900 - 1
                idx_end   = dim_end   // 900 + 1
                for i in xrange(idx_start, idx_end):
                    start = 900 * i - dim_start
                    end = start + 900

                    power_value = 0
                    if 0 <= start < len(full_profile):
                        power_value += full_profile[start] / 2
                        power_value -= cumulative[start]
                    elif len(cumulative) <= start:
                        power_value -= cumulative[-1]

                    if 0 <= end - 1 < len(cumulative):
                        power_value += cumulative[end - 1]
                    elif len(cumulative) <= end - 1:
                        power_value += cumulative[-1]

                    if 0 <= end <= len(full_profile):
                        power_value += full_profile[end] / 2

                    power_per_interval[i] = power_value

                storage[t,l,j] = power_per_interval

    return storage

import time

if __name__ == "__main__":
    s = time.time()
    model = generate_test_data(24000)
    print(time.time() - s)
    s = time.time()
    generate_storage(model)
    print(time.time() - s)