Ggthjy

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 laptop¹ 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.

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

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