Simple Python script to calculate a free-falling body using general relativity

freeFallGRT.py

import copy
import math
import numpy as np

'''
    This program simply calculates a particle in free fall by assuming a number
    of constant-velocity patches.
    The formula used is the (approcimate) time delation of General relativity,
    but actually it is the same as the principle of least action from clasical mechanics
    See http://scienceblogs.de/hier-wohnen-drachen/2017/01/13/der-freie-fall-und-die-maximale-eigenzeit

    Using
    Pure Python/Numpy implementation of the Nelder-Mead algorithm.
    Reference: https://en.wikipedia.org/wiki/Nelder%E2%80%93Mead_method
'''

#g=9.81
g=10.
c=299792458
c2=c*c
Tmax=1
hmax=5
## Number of T-values in the array
N = 25
deltaH=hmax/(N+1.)

def nelder_mead(f, x_start,
                step=0.1, no_improve_thr=10e-35,
                no_improv_break=50, max_iter=0,
                alpha=1., gamma=2., rho=-0.5, sigma=0.5):
    '''
        @param f (function): function to optimize, must return a scalar score
            and operate over a numpy array of the same dimensions as x_start
        @param x_start (numpy array): initial position
        @param step (float): look-around radius in initial step
        @no_improv_thr,  no_improv_break (float, int): break after no_improv_break iterations with
            an improvement lower than no_improv_thr
        @max_iter (int): always break after this number of iterations.
            Set it to 0 to loop indefinitely.
        @alpha, gamma, rho, sigma (floats): parameters of the algorithm
            (see Wikipedia page for reference)

        return: tuple (best parameter array, best score)
    '''

    # init
    dim = len(x_start)
    prev_best = f(x_start)
    no_improv = 0
    res = [[x_start, prev_best]]

    for i in range(dim):
        x = copy.copy(x_start)
        x[i] = x[i] + step
        score = f(x)
        res.append([x, score])

    # simplex iter
    iters = 0
    while 1:
        # order
        res.sort(key=lambda x: x[1])
        best = res[0][1]

        # break after max_iter
        if max_iter and iters >= max_iter:
            return res[0]
        iters += 1

        # break after no_improv_break iterations with no improvement
#        print '...best so far:', best

        if best < prev_best - no_improve_thr:
            no_improv = 0
            prev_best = best
        else:
            no_improv += 1

        if no_improv >= no_improv_break:
            return res[0]

        # centroid
        x0 = [0.] * dim
        for tup in res[:-1]:
            for i, c in enumerate(tup[0]):
                x0[i] += c / (len(res)-1)

        # reflection
        xr = x0 + alpha*(x0 - res[-1][0])
        rscore = f(xr)
        if res[0][1] <= rscore < res[-2][1]:
            del res[-1]
            res.append([xr, rscore])
            continue

        # expansion
        if rscore < res[0][1]:
            xe = x0 + gamma*(x0 - res[-1][0])
            escore = f(xe)
            if escore < rscore:
                del res[-1]
                res.append([xe, escore])
                continue
            else:
                del res[-1]
                res.append([xr, rscore])
                continue

        # contraction
        xc = x0 + rho*(x0 - res[-1][0])
        cscore = f(xc)
        if cscore < res[-1][1]:
            del res[-1]
            res.append([xc, cscore])
            continue

        # reduction
        x1 = res[0][0]
        nres = []
        for tup in res:
            redx = x1 + sigma*(tup[0] - x1)
            score = f(redx)
            nres.append([redx, score])
        res = nres

def eta(t,h):
    return g*h/c2 - (deltaH**2/t**2)/c2/2.


## T(i) is the time at which the height i is reached
def totaltime(T):
    tausum=0
    for i in range(N):
        cheight= hmax * (1.-(i+0.5)/(N+1.))
        if (i==0):
            deltaT=T[i]
        else:
            deltaT=T[i]-T[i-1]

        tausum += deltaT*eta(deltaT, cheight)
    cheight= hmax * (0.5/(N+1.))
    deltaT=1-T[i]
    tausum += deltaT*eta(1-T[N-1], cheight)
    return -tausum
if __name__ == "__main__":
    # test

    T=np.zeros(N)
    for i in range(N):
#        T[i] = math.sqrt(2*deltaH*(i+1)/g)
        T[i] = (i+1.)/(N+1.)
    print T
    totaltime(T)
    solution= nelder_mead(totaltime, T)
    print hmax,0
    for i in range(N):
        print hmax-(i+1)*deltaH, solution[0][i]
    print 0,1