<reponame>garrettkatz/directional-fibers
"""
Measure global optimization performance of Levy function
"""

import sys, time
import numpy as np
import matplotlib.pyplot as pt
import multiprocessing as mp
import dfibers.traversal as tv
import dfibers.numerical_utilities as nu
import dfibers.logging_utilities as lu
import dfibers.fixed_points as fx
import dfibers.solvers as sv
import dfibers.examples.levy as lv
from mpl_toolkits.mplot3d import Axes3D

def run_trial(args):
    basename, sample, timeout = args
    stop_time = time.clock() + timeout

    logfile = open("%s_s%d.log"%(basename,sample),"w")

    # Set up fiber arguments
    np.random.seed()
    v = 20*np.random.rand(2,1) - 10 # random point in domain
    c = lv.f(v) # direction at that point
    c = c + 0.1*np.random.randn(2,1) # perturb for more variability
    fiber_kwargs = {
        "f": lv.f,
        "ef": lv.ef,
        "Df": lv.Df,
        "compute_step_amount": lambda trace: (0.0001, 0),
        "v": v,
        "c": c,
        "stop_time": stop_time,
        "terminate": lambda trace: (np.fabs(trace.x[:-1]) > 10).any(),
        "max_solve_iterations": 2**5,
    }

    solve_start = time.clock()

    # Run in one direction
    solution = sv.fiber_solver(
        logger=lu.Logger(logfile).plus_prefix("+: "),
        **fiber_kwargs)
    X1 = np.concatenate(solution["Fiber trace"].points, axis=1)
    V1 = solution["Fixed points"]
    z = solution["Fiber trace"].z_initial
    # print("Status: %s\n"%solution["Fiber trace"].status)    
    
    # Run in other direction (negate initial tangent)
    solution = sv.fiber_solver(
        z= -z,
        logger=lu.Logger(logfile).plus_prefix("-: "),
        **fiber_kwargs)
    X2 = np.concatenate(solution["Fiber trace"].points, axis=1)
    V2 = solution["Fixed points"]
    # print("Status: %s\n"%solution["Fiber trace"].status)    

    # Join fiber segments
    fiber = np.concatenate((np.fliplr(X1), X2), axis=1)

    # Union solutions
    fxpts = fx.sanitize_points(
        np.concatenate((V1, V2), axis=1),
        f = lv.f,
        ef = lv.ef,
        Df = lv.Df,
        duplicates = lambda V, v: (np.fabs(V - v) < 10**-6).all(axis=0),
    )

    # Save results
    with open("%s_s%d.npz"%(basename,sample), 'w') as rf: np.savez(rf, **{
        "fxpts": fxpts,
        "fiber": fiber,
        "runtime": time.clock() - solve_start })

    logfile.close()

def run_experiment(basename, num_samples, timeout, num_procs=0):

    pool_args = []
    for sample in range(num_samples):
        pool_args.append((basename, sample, timeout))

    if num_procs > 0:

        num_procs = min(num_procs, mp.cpu_count())
        print("using %d processes..."%num_procs)
        pool = mp.Pool(processes=num_procs)
        pool.map(run_trial, pool_args)
        pool.close()
        pool.join()

    else:

        for pa in pool_args: run_trial(pa)

def compile_results(basename, num_samples):
    
    L = []
    F = []
    runtimes = []
    for sample in range(num_samples):
        with open("%s_s%d.npz"%(basename,sample), 'r') as rf: data = dict(np.load(rf))
        fxpts = data["fxpts"]
        Fs = np.fabs(lv.f(fxpts)).max(axis=0)
        Ls = lv.levy(fxpts)
        within = (np.fabs(fxpts) < 10).all(axis=0)
        mean_within = Ls[within].mean() if within.any() else np.nan
        print("sample %d: %d secs, %d solns, mean %f, mean within %f, min %f"%(
            sample, data["runtime"], len(Ls), Ls.mean(), mean_within, Ls.min()))
        L.append(Ls)
        F.append(Fs)
        runtimes.append(data["runtime"])
    
    counts = np.array([len(Ls) for Ls in L])
    bests = np.array([Ls.min() for Ls in L])
    resids = np.array([Fs.max() for Fs in F])
    runtimes = np.array(runtimes)
    print("avg count = %d, avg best = %f, avg resid = %f, best best = %f"%(
        counts.mean(), bests.mean(), resids.mean(), bests.min()))
    return counts, bests, runtimes

def plot_results(basename, num_samples, counts, bests, runtimes, timeout):

    ### Optimization order stats
    pt.figure(figsize=(5,4))
    pt.subplot(2,1,1)
    pt.plot(np.sort(bests), '-k.')
    pt.xlabel("Ordered samples")
    pt.ylabel("Best objective value")

    ##### Work complexity
    pt.subplot(2,1,2)
    terms = (runtimes < timeout)
    pt.plot(runtimes[terms], bests[terms], 'k+', markerfacecolor='none')
    pt.plot(runtimes[~terms], bests[~terms], 'ko', markerfacecolor='none')
    pt.legend(["terminated","timed out"])
    pt.xlabel("Runtime (seconds)")
    pt.ylabel("Best objective value")

    pt.tight_layout()
    pt.show()

    ### Fiber visuals

    pt.figure(figsize=(4,7))
    # objective fun
    X_surface, Y_surface = np.mgrid[-10:10:100j,-10:10:100j]
    L = lv.levy(np.array([X_surface.flatten(), Y_surface.flatten()])).reshape(X_surface.shape)
    ax_surface = pt.gcf().add_subplot(2,1,1,projection="3d")
    ax_surface.plot_surface(X_surface, Y_surface, L, linewidth=0, antialiased=False, color='gray')
    ax_surface.set_xlabel("v0")
    ax_surface.set_ylabel("v1")
    ax_surface.set_zlabel("levy(v)")
    ax_surface.view_init(azim=-80, elev=20)

    # fibers
    ax = pt.gcf().add_subplot(2,1,2)
    X_grid, Y_grid = np.mgrid[-10:10:60j,-10:10:60j]
    XY = np.array([X_grid.flatten(), Y_grid.flatten()])
    C_XY = lv.f(XY)
    ax.quiver(XY[0,:],XY[1,:],C_XY[0,:],C_XY[1,:],color=0.5*np.ones((1,3)),
        scale=10,units='xy',angles='xy')

    num_plot_samples = 3
    sort_idx = np.argsort(bests)
    plot_idx = [0] + list(np.random.permutation(num_samples)[:num_plot_samples-1])
    samples = sort_idx[plot_idx]
    # samples = [41,73,20] # all through global
    # samples = [41, 97, 11] # two through global
    # samples = [41, 49, 13] # two through global, one horiz not through
    # samples = [41, 46, 70] # one through global, one horiz
    # samples = [41, 96, 27] # two through global, one almost horiz
    samples = [41, 63, 28] # two through global, all interesting
    print("samples:")
    print(samples)
    for i,sample in enumerate(samples[::-1]):
        with open("%s_s%d.npz"%(basename,sample), 'r') as rf: data = dict(np.load(rf))
        fxpts = data["fxpts"]
        fiber = data["fiber"][:,::]
    
        L = lv.levy(fxpts).min()
        col = 0.5*float(num_plot_samples-i-1)/num_plot_samples
        print(sample,col)
        ax.plot(fiber[0],fiber[1],color=(col,col,col,1), linestyle='-', linewidth=1)
        pt.plot(fxpts[0],fxpts[1], 'o', color=(col,col,col,1))

    pt.xlabel("v0")
    pt.ylabel("v1",rotation=0)
    pt.yticks(np.linspace(-10,10,5))
    pt.xlim([-10,10])
    pt.ylim([-10,10])
    pt.tight_layout()
    pt.show()

if __name__ == "__main__":

    basename = "levy_opt"
    num_samples = 100
    num_plot_samples = 3
    timeout = 60*30
    num_procs = 10

    # run_experiment(basename, num_samples=num_samples, timeout=timeout, num_procs=num_procs)
    counts, bests, runtimes = compile_results(basename, num_samples)
    plot_results(basename, num_samples, counts, bests, runtimes, timeout)

"""
Measure global optimization performance of Levy function
"""

import sys, time
import numpy as np
import matplotlib.pyplot as pt
import multiprocessing as mp
import dfibers.traversal as tv
import dfibers.numerical_utilities as nu
import dfibers.logging_utilities as lu
import dfibers.fixed_points as fx
import dfibers.solvers as sv
import dfibers.examples.levy as lv
from mpl_toolkits.mplot3d import Axes3D

def run_trial(args):
    basename, sample, timeout = args
    stop_time = time.clock() + timeout

    logfile = open("%s_s%d.log"%(basename,sample),"w")

    # Set up fiber arguments
    np.random.seed()
    v = 20*np.random.rand(2,1) - 10 # random point in domain
    c = lv.f(v) # direction at that point
    c = c + 0.1*np.random.randn(2,1) # perturb for more variability
    fiber_kwargs = {
        "f": lv.f,
        "ef": lv.ef,
        "Df": lv.Df,
        "compute_step_amount": lambda trace: (0.0001, 0),
        "v": v,
        "c": c,
        "stop_time": stop_time,
        "terminate": lambda trace: (np.fabs(trace.x[:-1]) > 10).any(),
        "max_solve_iterations": 2**5,
    }

    solve_start = time.clock()

    # Run in one direction
    solution = sv.fiber_solver(
        logger=lu.Logger(logfile).plus_prefix("+: "),
        **fiber_kwargs)
    X1 = np.concatenate(solution["Fiber trace"].points, axis=1)
    V1 = solution["Fixed points"]
    z = solution["Fiber trace"].z_initial
    # print("Status: %s\n"%solution["Fiber trace"].status)    
    
    # Run in other direction (negate initial tangent)
    solution = sv.fiber_solver(
        z= -z,
        logger=lu.Logger(logfile).plus_prefix("-: "),
        **fiber_kwargs)
    X2 = np.concatenate(solution["Fiber trace"].points, axis=1)
    V2 = solution["Fixed points"]
    # print("Status: %s\n"%solution["Fiber trace"].status)    

    # Join fiber segments
    fiber = np.concatenate((np.fliplr(X1), X2), axis=1)

    # Union solutions
    fxpts = fx.sanitize_points(
        np.concatenate((V1, V2), axis=1),
        f = lv.f,
        ef = lv.ef,
        Df = lv.Df,
        duplicates = lambda V, v: (np.fabs(V - v) < 10**-6).all(axis=0),
    )

    # Save results
    with open("%s_s%d.npz"%(basename,sample), 'w') as rf: np.savez(rf, **{
        "fxpts": fxpts,
        "fiber": fiber,
        "runtime": time.clock() - solve_start })

    logfile.close()

def run_experiment(basename, num_samples, timeout, num_procs=0):

    pool_args = []
    for sample in range(num_samples):
        pool_args.append((basename, sample, timeout))

    if num_procs > 0:

        num_procs = min(num_procs, mp.cpu_count())
        print("using %d processes..."%num_procs)
        pool = mp.Pool(processes=num_procs)
        pool.map(run_trial, pool_args)
        pool.close()
        pool.join()

    else:

        for pa in pool_args: run_trial(pa)

def compile_results(basename, num_samples):
    
    L = []
    F = []
    runtimes = []
    for sample in range(num_samples):
        with open("%s_s%d.npz"%(basename,sample), 'r') as rf: data = dict(np.load(rf))
        fxpts = data["fxpts"]
        Fs = np.fabs(lv.f(fxpts)).max(axis=0)
        Ls = lv.levy(fxpts)
        within = (np.fabs(fxpts) < 10).all(axis=0)
        mean_within = Ls[within].mean() if within.any() else np.nan
        print("sample %d: %d secs, %d solns, mean %f, mean within %f, min %f"%(
            sample, data["runtime"], len(Ls), Ls.mean(), mean_within, Ls.min()))
        L.append(Ls)
        F.append(Fs)
        runtimes.append(data["runtime"])
    
    counts = np.array([len(Ls) for Ls in L])
    bests = np.array([Ls.min() for Ls in L])
    resids = np.array([Fs.max() for Fs in F])
    runtimes = np.array(runtimes)
    print("avg count = %d, avg best = %f, avg resid = %f, best best = %f"%(
        counts.mean(), bests.mean(), resids.mean(), bests.min()))
    return counts, bests, runtimes

def plot_results(basename, num_samples, counts, bests, runtimes, timeout):

    ### Optimization order stats
    pt.figure(figsize=(5,4))
    pt.subplot(2,1,1)
    pt.plot(np.sort(bests), '-k.')
    pt.xlabel("Ordered samples")
    pt.ylabel("Best objective value")

    ##### Work complexity
    pt.subplot(2,1,2)
    terms = (runtimes < timeout)
    pt.plot(runtimes[terms], bests[terms], 'k+', markerfacecolor='none')
    pt.plot(runtimes[~terms], bests[~terms], 'ko', markerfacecolor='none')
    pt.legend(["terminated","timed out"])
    pt.xlabel("Runtime (seconds)")
    pt.ylabel("Best objective value")

    pt.tight_layout()
    pt.show()

    ### Fiber visuals

    pt.figure(figsize=(4,7))
    # objective fun
    X_surface, Y_surface = np.mgrid[-10:10:100j,-10:10:100j]
    L = lv.levy(np.array([X_surface.flatten(), Y_surface.flatten()])).reshape(X_surface.shape)
    ax_surface = pt.gcf().add_subplot(2,1,1,projection="3d")
    ax_surface.plot_surface(X_surface, Y_surface, L, linewidth=0, antialiased=False, color='gray')
    ax_surface.set_xlabel("v0")
    ax_surface.set_ylabel("v1")
    ax_surface.set_zlabel("levy(v)")
    ax_surface.view_init(azim=-80, elev=20)

    # fibers
    ax = pt.gcf().add_subplot(2,1,2)
    X_grid, Y_grid = np.mgrid[-10:10:60j,-10:10:60j]
    XY = np.array([X_grid.flatten(), Y_grid.flatten()])
    C_XY = lv.f(XY)
    ax.quiver(XY[0,:],XY[1,:],C_XY[0,:],C_XY[1,:],color=0.5*np.ones((1,3)),
        scale=10,units='xy',angles='xy')

    num_plot_samples = 3
    sort_idx = np.argsort(bests)
    plot_idx = [0] + list(np.random.permutation(num_samples)[:num_plot_samples-1])
    samples = sort_idx[plot_idx]
    # samples = [41,73,20] # all through global
    # samples = [41, 97, 11] # two through global
    # samples = [41, 49, 13] # two through global, one horiz not through
    # samples = [41, 46, 70] # one through global, one horiz
    # samples = [41, 96, 27] # two through global, one almost horiz
    samples = [41, 63, 28] # two through global, all interesting
    print("samples:")
    print(samples)
    for i,sample in enumerate(samples[::-1]):
        with open("%s_s%d.npz"%(basename,sample), 'r') as rf: data = dict(np.load(rf))
        fxpts = data["fxpts"]
        fiber = data["fiber"][:,::]
    
        L = lv.levy(fxpts).min()
        col = 0.5*float(num_plot_samples-i-1)/num_plot_samples
        print(sample,col)
        ax.plot(fiber[0],fiber[1],color=(col,col,col,1), linestyle='-', linewidth=1)
        pt.plot(fxpts[0],fxpts[1], 'o', color=(col,col,col,1))

    pt.xlabel("v0")
    pt.ylabel("v1",rotation=0)
    pt.yticks(np.linspace(-10,10,5))
    pt.xlim([-10,10])
    pt.ylim([-10,10])
    pt.tight_layout()
    pt.show()

if __name__ == "__main__":

    basename = "levy_opt"
    num_samples = 100
    num_plot_samples = 3
    timeout = 60*30
    num_procs = 10

    # run_experiment(basename, num_samples=num_samples, timeout=timeout, num_procs=num_procs)
    counts, bests, runtimes = compile_results(basename, num_samples)
    plot_results(basename, num_samples, counts, bests, runtimes, timeout)
Measure global optimization performance of Levy function # Set up fiber arguments # random point in domain # direction at that point # perturb for more variability # Run in one direction # print("Status: %s\n"%solution["Fiber trace"].status) # Run in other direction (negate initial tangent) # print("Status: %s\n"%solution["Fiber trace"].status) # Join fiber segments # Union solutions # Save results ### Optimization order stats ##### Work complexity ### Fiber visuals # objective fun # fibers # samples = [41,73,20] # all through global # samples = [41, 97, 11] # two through global # samples = [41, 49, 13] # two through global, one horiz not through # samples = [41, 46, 70] # one through global, one horiz # samples = [41, 96, 27] # two through global, one almost horiz # two through global, all interesting # run_experiment(basename, num_samples=num_samples, timeout=timeout, num_procs=num_procs)