5
\$\begingroup\$

I've finally committed to moving to Python 3 from MATLAB. I'm currently porting my MATLAB code, and I'm sure I'm missing a lot of common best-practices.

For reference, I'm simulating a spring-loaded inverted pendulum (SLIP) model, which is commonly used to model human running. The system has hybrid dynamics, meaning the governing ODEs switch discretely (between stance and flight). Hence, the step function wraps solving the ODEs.
The simulation is eventually used in brute-force search (or in any case with a large number of rollouts), so performance suggestions are particularly useful.

I've taken a functional approach rather than making classes, partly because of issues I ran into with setting event properties, and partly for convenience. Below is the code with all the functions related to this specific model (other spring-mass models with different dynamics will be added eventually). Below that is a script to run everything.

import numpy as np
import scipy.integrate as integrate

def pMap(x,p):
    '''
    Wrapper function for step function, returning only x_next, and -1 if failed
    Essentially, the Poincare map.
    '''
    sol = step(x,p)
    return sol.y[:,-1], sol.failed

def step(x,p):
    '''
    Take one step from apex to apex/failure.
    returns a sol object from integrate.solve_ivp, with all phases
    '''

    # TODO: properly update sol object with all info, not just the trajectories

    # take one step (apex to apex)
    # the "step" function in MATLAB
    # x is the state vector, a list or np.array
    # p is a dict with all the parameters

    # set integration options

    x0 = x
    max_time = 5
    t0 = 0 # starting time

    # FLIGHT: simulate till touchdown
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: touchdownEvent(t,x,p)]
    for ev in events:
        ev.terminal = True
    sol = integrate.solve_ivp(fun=lambda t, x: flightDynamics(t, x, p),
    t_span = [t0, t0+max_time], y0 = x0, events=events, max_step=0.01)

    # STANCE: simulate till liftoff
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: liftoffEvent(t,x,p)]
    for ev in events:
        ev.terminal = True
    events[1].direction = 1 # only trigger when spring expands
    x0 = sol.y[:,-1]
    sol2 = integrate.solve_ivp(fun=lambda t, x: stanceDynamics(t, x, p),
    t_span = [sol.t[-1], sol.t[-1]+max_time], y0 = x0,
    events=events, max_step=0.0001)

    # FLIGHT: simulate till apex
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: apexEvent(t,x,p)]
    for ev in events:
        ev.terminal = True

    x0 = resetLeg(sol2.y[:,-1],p)
    sol3 = integrate.solve_ivp(fun=lambda t, x: flightDynamics(t, x, p),
    t_span = [sol2.t[-1], sol2.t[-1]+max_time], y0 = x0,
    events=events, max_step=0.01)

    # concatenate all solutions
    sol.t = np.concatenate((sol.t,sol2.t,sol3.t))
    sol.y = np.concatenate((sol.y,sol2.y,sol3.y),axis=1)
    sol.t_events += sol2.t_events + sol3.t_events

    # TODO: mark different phases
    for fail_idx in (0,2,4):
        if sol.t_events[fail_idx].size != 0: # if empty
            sol.failed = True
            break
    else:
        sol.failed = False
        # TODO: clean up the list

    return sol

def resetLeg(x,p):
    x[4] = x[0]+np.sin(p['aoa'])*p['resting_length']
    x[5] = x[1]-np.cos(p['aoa'])*p['resting_length']
    return x

def flightDynamics(t,x,p):
    # code in flight dynamics, xdot_ = f()
    return np.array([x[2], x[3], 0, -p['gravity'], x[2], x[3]])

def stanceDynamics(t, x,p):
    # stance dynamics
    # energy = computeTotalEnergy(x,p)
    # print(energy)
    alpha = np.arctan2(x[1]-x[5] , x[0]-x[4]) - np.pi/2.0
    leg_length = np.sqrt( (x[0]-x[4])**2 + (x[1]-x[5])**2 )
    xdotdot = -p["stiffness"]/p["mass"]*(p["resting_length"] -
                leg_length)*np.sin(alpha)
    ydotdot =  p["stiffness"]/p["mass"]*(p["resting_length"] -
                leg_length)*np.cos(alpha) - p["gravity"]
    return np.array([x[2], x[3], xdotdot, ydotdot, 0, 0])

def fallEvent(t,x,p):
    '''
    Event function to detect the body hitting the floor (failure)
    '''
    return x[1]
fallEvent.terminal = True
# TODO: direction

def touchdownEvent(t,x,p):
    '''
    Event function for foot touchdown (transition to stance)
    '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
    return x[5]
touchdownEvent.terminal = True # no longer actually necessary...
# direction

def liftoffEvent(t,x,p):
    '''
    Event function to reach maximum spring extension (transition to flight)
    '''
    return ((x[0]-x[4])**2 + (x[1]-x[5])**2) - p["resting_length"]**2
liftoffEvent.terminal = True
liftoffEvent.direction = 1

def apexEvent(t,x,p):
    '''
    Event function to reach apex
    '''
    return x[3]
apexEvent.terminal = True

def computeTotalEnergy(x,p):
    # TODO: make this accept a trajectory, and output parts as well
    return (p["mass"]/2*(x[2]**2+x[3]**2) +
    p["gravity"]*p["mass"]*(x[1]) +
    p["stiffness"]/2*
    (p["resting_length"]-np.sqrt((x[0]-x[4])**2 + (x[1]-x[5])**2))**2)

Script to run an example step:

from slip import *
import numpy as np
import matplotlib.pyplot as plt

p = {'mass':80.0, 'stiffness':8200.0, 'resting_length':1.0, 'gravity':9.81,
'aoa':1/5*np.pi} # aoa stands for angle_of_attack
x0 = [0, 0.85, 5.5, 0, 0, 0]
x0 = resetLeg(x0,p)
p['total_energy'] = computeTotalEnergy(x0,p)
sol = step(x0,p)

plt.plot(sol.y[0],sol.y[1], color='orange')
plt.show()

I'm aware my docstrings are not up to spec, and that's something on my todo list already. If you're interested, the code repo is here, and the code I'm porting was used for this paper.

\$\endgroup\$
  • \$\begingroup\$ Welcome to Code Review! For what Python version did you write this? \$\endgroup\$ – Mast Apr 18 at 12:36
  • \$\begingroup\$ Thanks, edited the title: python3 \$\endgroup\$ – Steve Heim Apr 18 at 12:51
  • 1
    \$\begingroup\$ The first code block is contained in the file slip.py, right? And what is aoa? \$\endgroup\$ – Graipher Apr 18 at 13:07
  • 1
    \$\begingroup\$ Yes, the first codeblock is in slip.py. aoa is shorthand for "angle of attack", which started being really long... \$\endgroup\$ – Steve Heim Apr 18 at 13:33
  • 1
    \$\begingroup\$ What is x? It looks like some vector in a configuration space? I suggest that you explain what the components of x represent. \$\endgroup\$ – 200_success Apr 18 at 18:11
5
\$\begingroup\$

Style

Your code looks quite good in general, but there is always something to nitpick on, isn't it?

Python comes with an official Style Guide, often just called PEP8. It has quite an extensive collection of style best practices you should follow in Python.

I think one of the most often cited "rules" is to use snake_case for variable and function names. You seem to follow that partly, since your variable names are written in snake_case while your function names use camelCase. I personally prefer snake case because it blends more nicely with the rest of the Python and NumPy/SciPy names.

Another rule (appropriately titled Whitespace in Expressions and Statements - Pet Peeves), I personally consider as more "essential" than the previous one, is to put a space after you have written a ,, e.g. while listing function arguments. It helps to avoid visual clutter IMHO. E.g.

p = {'mass':80.0, 'stiffness':8200.0, 'resting_length':1.0, 'gravity':9.81,
'aoa':1/5*np.pi} # aoa stands for angle_of_attack
x0 = [0, 0.85, 5.5, 0, 0, 0]
x0 = resetLeg(x0,p)
p['total_energy'] = computeTotalEnergy(x0,p)
sol = step(x0,p)

plt.plot(sol.y[0],sol.y[1], color='orange')
plt.show()

would become

p = {'mass': 80.0, 'stiffness': 8200.0, 'resting_length': 1.0, 'gravity': 9.81,
     'aoa': 1/5*np.pi} # aoa stands for angle_of_attack
x0 = [0, 0.85, 5.5, 0, 0, 0]
x0 = resetLeg(x0, p)
p['total_energy'] = computeTotalEnergy(x0, p)
sol = step(x0, p)

plt.plot(sol.y[0], sol.y[1], color='orange')
plt.show()

As you can see I also added some indentation to the parameter dictionary defintion to make it more clear which parts belong together. This effect can also be nicely shown in the body of step, where you could go from

sol2 = integrate.solve_ivp(fun=lambda t, x: stanceDynamics(t, x, p),
t_span = [sol.t[-1], sol.t[-1]+max_time], y0 = x0,
events=events, max_step=0.0001)

to

sol2 = integrate.solve_ivp(
    fun=lambda t, x: stanceDynamics(t, x, p),
    t_span=[sol.t[-1], sol.t[-1]+max_time],
    y0=x0, events=events, max_step=0.0001
)

I think it's hard to argue that this does not look more easy to read.

Since you do not try to rename the integrate submodule from SciPy to something shorter, there is no need to use the from ... import ... as ... syntax. Instead you could just use from scipy import integrate.

The last thing I would like to talk about here in this section is documentation. You have already mentioned that this is on your to-do list, and it definitely should be. The official Style Guide has a nice concise section on that topic, too, but since your specifically working in the scientific Python stack, I think it might be worth pointing you to the documentation style used in NumPy and SciPy. Using this docstyle for step would yield something like:

def step(x, p):
    '''
    Take one step from apex to apex/failure.

    Parameters
    ----------
    x : array_like
        state vector with 6 elements
    p : dict
        physical parameters used in the differential equations

    Returns
    -------
    object
        a solution object from scipy.integrate.solve_ivp with all phases
    '''
    # your code here

Elaborating further on this topic would clearly be to much for this review, so please read the ressources linked above to get more information.

Performance

Solving differential equations is not my area of expertise, so there are probably way more fundamental changes other members of the community can come up with in order to help you. I will nevertheless try to give you some insights I gained while working with your code.

I did some profiling on your code using the cProfile module from Python, and, as one might already has expected, most of the time is spent on numerically solving the ODEs in SciPy (I think I saw some Runge-Kutta pop op here and there). Furthermore, the profiler told me that solving the stance dynamics takes most of the time outside of the SciPy backend out of the three solutions. That was in line with my expectations since that's the call to integrate.solve_ivp where max_step is two orders of magnitude smaller. With your other parameters that lead to around 13610 calls to stanceDynamics (through the lambda expression of course). So that was my first point to look at.

The first micro-optimization I came up with was to replace leg_length = np.sqrt((x[0]-x[4])**2 + (x[1]-x[5])**2) with the well known hypot function found in Python's math module as well as in NumPy to form leg_length = np.hypot(x[0]-x[4], x[1]-x[5]). The effects where minuscule, but that was also within my expectations.

So what to do next? I recently got to know (read: no expert here!) a Python package called Numba which allows just-in-time compilation of Python code in order to improve performance, especially if Python code is called within tight loops. Great plus: it's quite easy to use (most of the time) and works on Python code as well as with NumPy. With a little help Numba seems reasonable capable to deduce the datatype of parameters you put into a function on it's own. I went from your original code

def stanceDynamics(t, x, p):
    alpha = np.arctan2(x[1]-x[5] , x[0]-x[4]) - np.pi/2.0
    leg_length = np.sqrt((x[0]-x[4])**2 + (x[1]-x[5])**2)
    xdotdot = -p["stiffness"]/p["mass"]*(p["resting_length"] - leg_length) * np.sin(alpha)
    ydotdot =  p["stiffness"]/p["mass"]*(p["resting_length"] - leg_length)*np.cos(alpha) - p["gravity"]
    return np.array([x[2], x[3], xdotdot, ydotdot, 0, 0])

to the following

import numba as nb

@nb.jit
def stanceDynamics2(t, x, stiffness, mass, resting_length, gravity):
    alpha = np.arctan2(x[1] - x[5], x[0] - x[4]) - np.pi / 2.0
    leg_length = np.hypot(x[0]-x[4], x[1]-x[5])
    xdotdot = -stiffness / mass * (resting_length - leg_length) * np.sin(alpha)
    ydotdot = stiffness / mass * (resting_length - leg_length) * np.cos(alpha) - gravity
    return np.array([x[2], x[3], xdotdot, ydotdot, 0, 0])

As you can see, the dictionary vanished from the function signature and was replaced by the plain values in order to help Numba deduce the type (and save some dictionary lookups). You can see the result of this when looking at the profiler output (10 repetitions for each variant):

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   136100    2.985    0.000    3.550    0.000 slip.py:102(stanceDynamics)
   136100    0.565    0.000    0.565    0.000 slip.py:115(stanceDynamics2)

In the overall timing averaged over ten runs that gives you a performance gain of about 25-30% (or about ~240ms in absolute values) on my not so powerful laptop1 here. Your results may vary.

As I said, there is possibly a lot more to gain here, but I would consider this at least as a respectable achievement for the amount of effort needed to implement those changes.


1 Intel Core i5 460M, Python 3.7.3 [MSC v.1915 64 bit (AMD64)] :: Anaconda, Inc. on win32

\$\endgroup\$
  • \$\begingroup\$ Two things I would suggest: (1) create named arguments that default to all the functions that are called, and initialize them to the dotted values def stanceDynamics(..., atan2=np.arctan2, hypot=np.hypot, sin=np.sin, cos=np.cos, array=np.array); and (2) aggressively pull out common sub-expressions into local variables: anything that is used more than once should move into a local. \$\endgroup\$ – Austin Hastings Apr 18 at 21:42
  • \$\begingroup\$ @AustinHastings: Would these hints apply to the "standard" version or to Numba? I can't see any real difference with (1) (standard & Numba) and (2) seems to make it worse with Numba. \$\endgroup\$ – AlexV Apr 18 at 22:23
  • \$\begingroup\$ they are for the standard version. I don't have experience with Numba. \$\endgroup\$ – Austin Hastings Apr 18 at 22:33
6
\$\begingroup\$

Inspired by @Alex's answer, I did some profiling and simple timing of your code, and some changes that seemed "obvious" to me.

My environment for this is a 32-bit laptop that is old enough to have a driver's license, running python 3.7 on Windows. Note: Because of this, it's quite likely that your performance will differ from mine. Verify these results, it's important!

I downloaded scipy for this, but I don't have any real knowledge of what you're doing. Also, at this point it's pretty late, so I'm not as smart as I was a few hours ago. Forgive me if I'm incoherent.

I created modified versions of your main function, step. I'll describe the changes, and the result, then past in the code at the bottom for you to try to reproduce my results.

]$ python test.py -t
Start: (N/A) 1555646668200
Done orig: (1466) 1555646669666
Done const True: (1404) 1555646671070
Done const_lookup True: (1373) 1555646672443
Done const_lookup_cse True: (1310) 1555646673753
Done const_lookup_cse2 True: (1264) 1555646675017
Done const_lookup_cse2_ary2 False: (1217) 1555646676234

The numbers in parens are elapsed milliseconds to run the step_ function, so lower is better. The trailing numbers are current-time-in-ms, which you should ignore. The boolean value indicates whether the returned result matches the original step result -- note that _ary2 is False.

The names get longer because I was adding more changes, except for the _ary version, which I suppressed because it didn't work right. (I left it in, so you can see a failed example. But see the comments on _ary2 first.)

The functions are named step_XXX where XXX is whatever appears in the table above (e.g., step_const_lookup_cse). The parts of the name indicate what I tried to do in that iteration.

orig (aka, step)

This is the original. Nothing to see here, move along.

step_const

In this version, I moved all the event and dynamics functions inside the step function as nested functions. I unpacked the p[] dictionary into variables (with uppercased names, so I called them constants) that don't change during the execution.

For example: GRAVITY = p["gravity"]

I then rewrote the various helpers to use these names instead of dictionary lookups.

const_lookup

In this version, I tried to remove lookups where possible. I did this by storing names into named-arguments with default values. Thus, for example, I would have a function with a parameter list like (..., array=np.array) so that the array named parameter could be used instead of a lookup of np.array each time.

const_lookup_cse

This was my first attempt at Common-Subexpression-Elimination. In general, I tried to eliminate repeated computations of the same result by storing results in local variables.

const_lookup_cse2

This is an extended version of _cse above. There aren't many changes, but not for lack of trying.

One thing I tried was to "unpack" the x[] array into variables x0, x1, ..., x5 in functions that used them. This totally didn't work for me, but it might be worth trying for you. That is, instead of expressions like x[1] - x[5] I used x1 - x5. I was surprised at how bad the performance was, so I quickly undid that change and want for some simpler ones.

const_lookup_cse2_ary2

I used the cProfile profiler to investigate the performance of various functions. There is a lot of time spent in library code, which is "good" in the sense that there's work being done and it's not your code.

The dominant functions are related to the stance_dynamics call. One of the non-obvious culprits is np.array so I figured I would try to come up with some alternatives to calling that function.

In my _ary version, which I suppressed, I created a ring of 4 np.arrays that I would alternate through in sequence. Sadly, this performed worse than just constructing the array from scratch each time. Yikes! Also, the results returned were not the same as the original function.

So I tried a different approach, _ary2 in which I held a single np.array for each function that returns one, and reset the values inside the array. This gets improved performance, but again, the results are not the same.

Now, here's the important part: I have no understanding of what you are doing. So it may be that the different results don't matter. If, for example, the only time the returned np.array is important is immediately after it is returned, then the overall results might be valid even though the details are different. So I left this in. On the other than, if the integration solver is referring back to previous values, and by overwriting the contents of the array I'm screwing those up, then this is an invalid change.

But it's a decent speed-up, so I'm leaving that to you.

One thing I noticed reading the docs for the solver, there's an option to vectorize the function. In that case, instead of an array(n) it takes and returns an array(n, k). In my ignorance, I don't know if that means it would be a "growing" array, with new values being appended, or some other thing. That's up to you to figure out. But if it would let you append the results to the end of an existing array, it might give a similar performance boost.

Files

Here's the test driver, test.py, and then the module slip.py:

test.py

import sys
import time

#import matplotlib.pyplot as plt
import numpy as np
from slip import *

_last_time_ms = 0
def timestamp(msg):
    global _last_time_ms
    time_ms = int(round(time.time() * 1000))
    delta = 'N/A' if _last_time_ms == 0 else time_ms - _last_time_ms
    print(f"{msg}: ({delta}) {time_ms}")
    _last_time_ms = time_ms


p = {'mass':80.0, 'stiffness':8200.0, 'resting_length':1.0, 'gravity':9.81,
'aoa':1/5*np.pi}
# x0 = [-np.sin(p['aoa'])*p['resting_length'],1.1*p['resting_length'], # body position
#     5*p['resting_length'],0.0, # velocities
#     0.0,0.0] # foot position

def get_x0():
    x0 = [0, 0.85, 5.5, 0, 0, 0]
    x0 = resetLeg(x0, p)
    p['total_energy'] = computeTotalEnergy(x0, p)
    return x0

def sol_close(a, b):
    return np.allclose(a.t, b.t) and np.allclose(a.y, b.y)

TIMING = ('-t' in sys.argv)

if TIMING:
    run_all = True

    timestamp("Start")
    sol = step(get_x0(), p)
    timestamp("Done orig")

    if run_all:

        x = step_const(get_x0(), p)
        same = sol_close(sol, x)
        timestamp(f"Done const {same}")

        x = step_const_lookup(get_x0(), p)
        same = sol_close(sol, x)
        timestamp(f"Done const_lookup {same}")

        x = step_const_lookup_cse(get_x0(), p)
        same = sol_close(sol, x)
        timestamp(f"Done const_lookup_cse {same}")

        x = step_const_lookup_cse2(get_x0(), p)
        same = sol_close(sol, x)
        timestamp(f"Done const_lookup_cse2 {same}")

        # Doesn't work. Also, runs slower.
        #x = step_const_lookup_cse2_ary(get_x0(), p)
        #same = sol_close(sol, x)
        #timestamp(f"Done const_lookup_cse2_ary {same}")

        x = step_const_lookup_cse2_ary2(get_x0(), p)
        same = sol_close(sol, x)
        timestamp(f"Done const_lookup_cse2_ary2 {same}")

else:
    import cProfile
    import pstats

    statsfile = 'step_profile'

    cProfile.run('step_const_lookup_cse2(get_x0(), p)', statsfile)

    p = pstats.Stats(statsfile)
    p.strip_dirs().sort_stats('cumtime').print_stats()

# plt.plot(sol.t,sol.y[0])
# plt.plot(sol.t,sol.y[1], color='green')
# plt.plot(sol.t,sol.y[0])
# plt.plot(sol.y[0],sol.y[1], color='orange')
# plt.show()

slip.py

import math
import numpy as np
import scipy.integrate as integrate

def pMap(x,p):
    '''
    Wrapper function for step, returning only x_next, and -1 if failed
    '''
    sol = step(x,p)
    return sol.y[:,-1], sol.failed

def step_const(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x):
        ''' code in flight dynamics, xdot_ = f() '''
        return np.array([x[2], x[3], 0, -GRAVITY, x[2], x[3]])

    def liftoff_event(t, x, RESTING_LENGTH_SQ=RESTING_LENGTH**2):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return ((x[0]-x[4])**2 + (x[1]-x[5])**2) - RESTING_LENGTH_SQ
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        alpha = np.arctan2(x[1]-x[5], x[0]-x[4]) - HALF_PI
        leg_length = np.sqrt((x[0] - x[4]) ** 2 + (x[1] - x[5]) ** 2)
        xdotdot = -SPECIFIC_STIFFNESS * (RESTING_LENGTH - leg_length) \
                * np.sin(alpha)
        ydotdot = SPECIFIC_STIFFNESS * (RESTING_LENGTH - leg_length) \
                * np.cos(alpha) - GRAVITY
        return np.array([x[2], x[3], xdotdot, ydotdot, 0, 0])

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    APEX_EVENTS = (fall_event, apex_event)
    FLIGHT_EVENTS = (fall_event, touchdown_event)
    STANCE_EVENTS = (fall_event, liftoff_event)


    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    # FLIGHT: simulate till touchdown
    sol = integrate.solve_ivp(
            events=FLIGHT_EVENTS,
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    sol2 = integrate.solve_ivp(
            events=STANCE_EVENTS,
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[sol.t[-1], sol.t[-1] + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    sol3 = integrate.solve_ivp(
            events=APEX_EVENTS,
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[sol2.t[-1], sol2.t[-1] + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    sol.t = np.concatenate((sol.t, sol2.t, sol3.t))
    sol.y = np.concatenate((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    sol.failed = any(sol.t_events[i].size != 0 for i in (0, 2, 4))
    return sol

def step_const_lookup(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x, array=np.array):
        ''' code in flight dynamics, xdot_ = f() '''
        return array([x[2], x[3], 0, -GRAVITY, x[2], x[3]])

    def liftoff_event(t, x, RESTING_LENGTH_SQ=RESTING_LENGTH**2):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return ((x[0]-x[4])**2 + (x[1]-x[5])**2) - RESTING_LENGTH_SQ
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x, atan2=np.arctan2, sqrt=np.sqrt, sin=np.sin,
            cos=np.cos, array=np.array):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        alpha = atan2(x[1]-x[5], x[0]-x[4]) - HALF_PI
        leg_length = sqrt((x[0] - x[4]) ** 2 + (x[1] - x[5]) ** 2)
        xdotdot = -SPECIFIC_STIFFNESS * (RESTING_LENGTH - leg_length) \
                * sin(alpha)
        ydotdot = SPECIFIC_STIFFNESS * (RESTING_LENGTH - leg_length) \
                * cos(alpha) - GRAVITY
        return array([x[2], x[3], xdotdot, ydotdot, 0, 0])

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    APEX_EVENTS = (fall_event, apex_event)
    FLIGHT_EVENTS = (fall_event, touchdown_event)
    STANCE_EVENTS = (fall_event, liftoff_event)


    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    solve_ivp = integrate.solve_ivp
    ''' Cache lookup '''

    # FLIGHT: simulate till touchdown
    sol = solve_ivp(
            events=FLIGHT_EVENTS,
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    sol2 = solve_ivp(
            events=STANCE_EVENTS,
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[sol.t[-1], sol.t[-1] + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    sol3 = solve_ivp(
            events=APEX_EVENTS,
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[sol2.t[-1], sol2.t[-1] + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    concat = np.concatenate
    sol.t = concat((sol.t, sol2.t, sol3.t))
    sol.y = concat((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    ste = sol.t_events
    sol.failed = any(ste[i].size != 0 for i in (0, 2, 4))
    return sol

def step_const_lookup_cse(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x, array=np.array):
        ''' code in flight dynamics, xdot_ = f() '''
        return array([x[2], x[3], 0, -GRAVITY, x[2], x[3]])

    def liftoff_event(t, x, RESTING_LENGTH_SQ=RESTING_LENGTH**2):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return ((x[0]-x[4])**2 + (x[1]-x[5])**2) - RESTING_LENGTH_SQ
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x, atan2=np.arctan2, hypot=math.hypot, sqrt=np.sqrt,
            cos=np.cos, sin=np.sin, array=np.array):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        alpha = atan2(x[1] - x[5], x[0] - x[4]) - HALF_PI
        #leg_length = sqrt((x[0] - x[4]) ** 2 + (x[1] - x[5]) ** 2)
        stiff_x_leg = SPECIFIC_STIFFNESS * (RESTING_LENGTH -
                                            hypot(x[0] - x[4], x[1] - x[5]))
        xdotdot = -stiff_x_leg * sin(alpha)
        ydotdot = stiff_x_leg * cos(alpha) - GRAVITY
        return array([x[2], x[3], xdotdot, ydotdot, 0, 0])

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    solve_ivp = integrate.solve_ivp
    ''' Cache lookup '''

    # FLIGHT: simulate till touchdown
    sol = solve_ivp(
            events=(fall_event, touchdown_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    last_t = sol.t[-1]
    sol2 = solve_ivp(
            events=(fall_event, liftoff_event),
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    last_t = sol2.t[-1]
    sol3 = solve_ivp(
            events=(fall_event, apex_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    concat = np.concatenate
    sol.t = concat((sol.t, sol2.t, sol3.t))
    sol.y = concat((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    ste = sol.t_events
    sol.failed = any(ste[i].size != 0 for i in (0, 2, 4))
    return sol

def step_const_lookup_cse2(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x, array=np.array):
        ''' code in flight dynamics, xdot_ = f() '''
        return array([x[2], x[3], 0, -GRAVITY, x[2], x[3]])

    def liftoff_event(t, x, hypot=math.hypot):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return hypot(x[0] - x[4], x[1] - x[5]) - RESTING_LENGTH
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x, atan2=np.arctan2, hypot=math.hypot, sqrt=np.sqrt,
            cos=np.cos, sin=np.sin, array=np.array):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        x15 = x[1] - x[5]
        x04 = x[0] - x[4]
        alpha = atan2(x15, x04) - HALF_PI
        stiff_x_leg = SPECIFIC_STIFFNESS * (RESTING_LENGTH - hypot(x15, x04))
        xdotdot = -stiff_x_leg * sin(alpha)
        ydotdot = stiff_x_leg * cos(alpha) - GRAVITY
        return array([x[2], x[3], xdotdot, ydotdot, 0, 0])

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    solve_ivp = integrate.solve_ivp
    ''' Cache lookup '''

    # FLIGHT: simulate till touchdown
    sol = solve_ivp(
            events=(fall_event, touchdown_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    last_t = sol.t[-1]
    sol2 = solve_ivp(
            events=(fall_event, liftoff_event),
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    last_t = sol2.t[-1]
    sol3 = solve_ivp(
            events=(fall_event, apex_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    concat = np.concatenate
    sol.t = concat((sol.t, sol2.t, sol3.t))
    sol.y = concat((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    ste = sol.t_events
    sol.failed = any(ste[i].size != 0 for i in (0, 2, 4))
    return sol

# NB: I pasted the code in two parts, and this is the seam.

def step_const_lookup_cse2_ary(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    _arrays = [np.array([0]*6) for _ in range(4)]
    _arrays.append(0)

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x, arrays=_arrays):
        ''' code in flight dynamics, xdot_ = f() '''
        i = arrays[-1]
        arrays[-1] += 1
        a = arrays[i & 3]
        a[4] = a[0] = x[2]
        a[5] = a[1] = x[3]
        a[2] = 0
        a[3] = -GRAVITY
        return a

    def liftoff_event(t, x, hypot=math.hypot):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return hypot(x[0] - x[4], x[1] - x[5]) - RESTING_LENGTH
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x, atan2=np.arctan2, hypot=math.hypot, sqrt=np.sqrt,
            cos=np.cos, sin=np.sin, arrays=_arrays):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        x15 = x[1] - x[5]
        x04 = x[0] - x[4]
        alpha = atan2(x15, x04) - HALF_PI
        stiff_x_leg = SPECIFIC_STIFFNESS * (RESTING_LENGTH - hypot(x15, x04))
        xdotdot = -stiff_x_leg * sin(alpha)
        ydotdot = stiff_x_leg * cos(alpha) - GRAVITY
        i = arrays[-1]
        arrays[-1] += 1
        a = arrays[i & 3]
        a[0] = x[2]
        a[1] = x[3]
        a[2] = xdotdot
        a[3] = ydotdot
        a[4] = a[5] = 0
        return a

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    solve_ivp = integrate.solve_ivp
    ''' Cache lookup '''

    # FLIGHT: simulate till touchdown
    sol = solve_ivp(
            events=(fall_event, touchdown_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    last_t = sol.t[-1]
    sol2 = solve_ivp(
            events=(fall_event, liftoff_event),
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    last_t = sol2.t[-1]
    sol3 = solve_ivp(
            events=(fall_event, apex_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    concat = np.concatenate
    sol.t = concat((sol.t, sol2.t, sol3.t))
    sol.y = concat((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    ste = sol.t_events
    sol.failed = any(ste[i].size != 0 for i in (0, 2, 4))
    return sol

def step_const_lookup_cse2_ary2(x, p):
    AOA = p['aoa']
    GRAVITY = p['gravity']
    MASS = p['mass']
    RESTING_LENGTH = p['resting_length']
    STIFFNESS = p['stiffness']
    TOTAL_ENERGY = p['total_energy']

    SPECIFIC_STIFFNESS = STIFFNESS / MASS  # FIXME: Is this name right?

    # Note: not taken from p[]
    HALF_PI = np.pi / 2.0
    MAX_TIME = 5

    # Function definitions: specific to the constants provided in `p`

    def apex_event(t, x):
        ''' Event function to reach apex '''
        return x[3]
    apex_event.terminal = True

    def fall_event(t, x):
        ''' Event function to detect the body hitting the floor (failure)
        '''
        return x[1]
    fall_event.terminal = True

    def flight_dynamics(t, x, a=np.zeros(6)):
        ''' code in flight dynamics, xdot_ = f() '''
        a[4] = a[0] = x[2]
        a[5] = a[1] = x[3]
        a[2] = 0
        a[3] = -GRAVITY
        return a

    def liftoff_event(t, x, hypot=math.hypot):
        ''' Event function for maximum spring extension (transition to flight)
        '''
        return hypot(x[0] - x[4], x[1] - x[5]) - RESTING_LENGTH
    liftoff_event.terminal = True
    liftoff_event.direction = 1

    def stance_dynamics(t, x, atan2=np.arctan2, hypot=math.hypot, sqrt=np.sqrt,
            cos=np.cos, sin=np.sin, a=np.zeros(6)):
        # energy = computeTotalEnergy(x,p)
        # print(energy)
        x15 = x[1] - x[5]
        x04 = x[0] - x[4]
        alpha = atan2(x15, x04) - HALF_PI
        stiff_x_leg = SPECIFIC_STIFFNESS * (RESTING_LENGTH - hypot(x15, x04))
        xdotdot = -stiff_x_leg * sin(alpha)
        ydotdot = stiff_x_leg * cos(alpha) - GRAVITY
        a[0] = x[2]
        a[1] = x[3]
        a[2] = xdotdot
        a[3] = ydotdot
        a[4] = a[5] = 0
        return a

    def touchdown_event(t, x):
        ''' Event function for foot touchdown (transition to stance)
        '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
        return x[5]
    touchdown_event.terminal = True

    t0 = 0
    ''' Starting time '''
    x0 = x
    ''' Starting state '''

    solve_ivp = integrate.solve_ivp
    ''' Cache lookup '''

    # FLIGHT: simulate till touchdown
    sol = solve_ivp(
            events=(fall_event, touchdown_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[t0, t0 + MAX_TIME],
            y0=x0,
    )

    # STANCE: simulate till liftoff
    x0 = sol.y[:, -1]
    last_t = sol.t[-1]
    sol2 = solve_ivp(
            events=(fall_event, liftoff_event),
            fun=stance_dynamics,
            max_step=0.0001,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # FLIGHT: simulate till apex
    x0 = resetLeg(sol2.y[:, -1], p)
    last_t = sol2.t[-1]
    sol3 = solve_ivp(
            events=(fall_event, apex_event),
            fun=flight_dynamics,
            max_step=0.01,
            t_span=[last_t, last_t + MAX_TIME],
            y0=x0,
    )

    # concatenate all solutions
    concat = np.concatenate
    sol.t = concat((sol.t, sol2.t, sol3.t))
    sol.y = concat((sol.y, sol2.y, sol3.y), axis=1)
    sol.t_events += sol2.t_events + sol3.t_events
    ste = sol.t_events
    sol.failed = any(ste[i].size != 0 for i in (0, 2, 4))
    return sol

def step(x,p):
    '''
    Take one step from apex to apex/failure.
    returns a sol object from integrate.solve_ivp, with all phases
    '''

    # TODO: properly update sol object with all info, not just the trajectories

    # take one step (apex to apex)
    # the "step" function in MATLAB
    # x is the state vector, a list or np.array
    # p is a dict with all the parameters

    # set integration options

    x0 = x
    max_time = 5
    t0 = 0 # starting time

    # FLIGHT: simulate till touchdown
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: touchdownEvent(t,x,p)]
    for ev in events:
        ev.terminal = True
    sol = integrate.solve_ivp(fun=lambda t, x: flightDynamics(t, x, p),
    t_span = [t0, t0+max_time], y0 = x0, events=events, max_step=0.01)

    # STANCE: simulate till liftoff
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: liftoffEvent(t,x,p)]
    for ev in events:
        ev.terminal = True
    events[1].direction = 1 # only trigger when spring expands
    x0 = sol.y[:,-1]
    sol2 = integrate.solve_ivp(fun=lambda t, x: stanceDynamics(t, x, p),
    t_span = [sol.t[-1], sol.t[-1]+max_time], y0 = x0,
    events=events, max_step=0.0001)

    # FLIGHT: simulate till apex
    events = [lambda t,x: fallEvent(t,x,p), lambda t,x: apexEvent(t,x,p)]
    for ev in events:
        ev.terminal = True

    x0 = resetLeg(sol2.y[:,-1], p)
    sol3 = integrate.solve_ivp(fun=lambda t, x: flightDynamics(t, x, p),
    t_span = [sol2.t[-1], sol2.t[-1]+max_time], y0 = x0,
    events=events, max_step=0.01)

    # concatenate all solutions
    sol.t = np.concatenate((sol.t,sol2.t,sol3.t))
    sol.y = np.concatenate((sol.y,sol2.y,sol3.y),axis=1)
    sol.t_events += sol2.t_events + sol3.t_events

    # TODO: mark different phases
    for fail_idx in (0,2,4):
        if sol.t_events[fail_idx].size != 0: # if empty
            sol.failed = True
            break
    else:
        sol.failed = False
        # TODO: clean up the list

    return sol

def resetLeg(x,p):
    x[4] = x[0]+np.sin(p['aoa'])*p['resting_length']
    x[5] = x[1]-np.cos(p['aoa'])*p['resting_length']
    return x

def stanceDynamics(t, x,p):
    # stance dynamics
    # energy = computeTotalEnergy(x,p)
    # print(energy)
    alpha = np.arctan2(x[1]-x[5] , x[0]-x[4]) - np.pi/2.0
    leg_length = np.sqrt( (x[0]-x[4])**2 + (x[1]-x[5])**2 )
    xdotdot = -p["stiffness"]/p["mass"]*(p["resting_length"] -
                leg_length)*np.sin(alpha)
    ydotdot =  p["stiffness"]/p["mass"]*(p["resting_length"] -
                leg_length)*np.cos(alpha) - p["gravity"]
    return np.array([x[2], x[3], xdotdot, ydotdot, 0, 0])

def fallEvent(t,x,p):
    '''
    Event function to detect the body hitting the floor (failure)
    '''
    return x[1]
fallEvent.terminal = True
# TODO: direction

def touchdownEvent(t,x,p):
    '''
    Event function for foot touchdown (transition to stance)
    '''
        # x[1]- np.cos(p["aoa"])*p["resting_length"] (which is = x[5])
    return x[5]
touchdownEvent.terminal = True # no longer actually necessary...
# direction

def liftoffEvent(t,x,p):
    '''
    Event function to reach maximum spring extension (transition to flight)
    '''
    return ((x[0]-x[4])**2 + (x[1]-x[5])**2) - p["resting_length"]**2
liftoffEvent.terminal = True
liftoffEvent.direction = 1

def apexEvent(t,x,p):
    '''
    Event function to reach apex
    '''
    return x[3]
apexEvent.terminal = True

def computeTotalEnergy(x,p):
    # TODO: make this accept a trajectory, and output parts as well
    return (p["mass"]/2*(x[2]**2+x[3]**2) +
    p["gravity"]*p["mass"]*(x[1]) +
    p["stiffness"]/2*
    (p["resting_length"]-np.sqrt((x[0]-x[4])**2 + (x[1]-x[5])**2))**2)

def flightDynamics(t, x, p):
    ''' code in flight dynamics, xdot_ = f() '''
    return np.array([x[2], x[3], 0, -p["gravity"], x[2], x[3]])

### Functions for Viability
def map2e(x, p):
    '''
    map an apex state to its dimensionless normalized height
    TODO: make this accept trajectories
    '''
    assert(np.isclose(x[3],0))
    potential_energy = p['mass']*p['gravity']*x[1]
    kinetic_energy = p['mass']/2*x[3]**2
    return potential_energy/(potential_energy+kinetic_energy)

def map2x(x,p,e):
    '''
    map a desired dimensionless height `e` to it's state-vector
    '''
    if 'total_energy' not in p:
        print('WARNING: you did not initialize your parameters with '
        'total energy. You really should do this...')

    assert(np.isclose(x[3],0)) # check that we are at apex

    x_new = x
    x_new[1] = p['total_energy']*e/p['mass']/p['gravity']
    x_new[2] = np.sqrt(p['total_energy']*(1-e)/p['mass']*2)
    x_new[3] = 0.0 # shouldn't be necessary, but avoids errors accumulating
    return x_new

def mapSA2xp_height_angle(state_action,x,p):
    '''
    Specifically map state_actions to x and p
    '''
    p['aoa'] = state_action[1]
    x = map2x(x,p,state_action[0])
    return x,p
\$\endgroup\$
  • \$\begingroup\$ Thanks! I like the nested functions, I'll need to check if I can do that without sacrificing other functionality (not shown here), but I think most if not all of it is actually fine this way. Also, in regards to unpacking, is this improving performance due to not passing the variables around to subfunctions, or not looking up a dict repeatedly? My understanding is that dictionary look-ups are very fast, so I would not expect that to result in a speedup. \$\endgroup\$ – Steve Heim Apr 19 at 17:34
  • \$\begingroup\$ My unpacking was of the x array, and I hoped a single unpack would be faster than several index operations. But it wasn't, which is why having hard numbers and making incremental changes is important. \$\endgroup\$ – Austin Hastings Apr 19 at 18:45
  • 1
    \$\begingroup\$ Ahh, you meant unpacking the p dictionary into constants. It's true that dictionary lookups are fast. But fast compared to what? Looking something up in a dictionary held in a variable means looking up the variable, then passing the key to the get item method of that variable. Whereas looking up a separately named variable means looking up the variable, done! \$\endgroup\$ – Austin Hastings Apr 19 at 18:50

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.