gel.scripts.get_displacement_from_gpr

  1#!/usr/bin/env python3
  2"""Within the FM-TRACK environment, perform GPR-interpolation with text
  3input and output.
  4
  5For all arguments, run `get_displacement_from_gpr --help`
  6"""
  7import argparse
  8import numpy as np
  9import os
 10import pickle
 11
 12
 13def my_optimizer(obj_func, initial_theta, bounds):
 14    """Function sometimes used for kernel hyperparameter optimization"""
 15    import scipy
 16    opt_res = scipy.optimize.minimize(
 17        obj_func, initial_theta, method="L-BFGS-B", jac=True,
 18        bounds=bounds, options={"maxiter" : 3, "disp" : True})
 19    theta_opt, func_min = opt_res.x, opt_res.fun
 20    return theta_opt, func_min
 21
 22
 23_cache = dict()
 24def get_predicted_u(gpr_path, vertices):
 25    """Evaluates GPR model from directory.
 26
 27    * `gpr_path`: str path to directory with files:
 28        * gp_U_cleaned.sav: FM-Track GPR model of x1-axis displacements
 29        * gp_V_cleaned.sav: FM-Track GPR model of x2-axis displacements
 30        * gp_W_cleaned.sav: FM-Track GPR model of x3-axis displacements
 31        * scaler_cleaned.sav: FM-Track GPR model pre-process component
 32    * `vertices`: (N,3) array of float coordinates in hydrogel
 33
 34    Returns:
 35    * (N,3) array of float displacements at coordinates
 36    """
 37    # loads GPR models
 38    if gpr_path not in _cache:
 39        gp_U = pickle.load(open(os.path.join(gpr_path, 'gp_U_cleaned.sav'), 'rb'))
 40        gp_V = pickle.load(open(os.path.join(gpr_path, 'gp_V_cleaned.sav'), 'rb'))
 41        gp_W = pickle.load(open(os.path.join(gpr_path, 'gp_W_cleaned.sav'), 'rb'))
 42        scaler = pickle.load(open(os.path.join(gpr_path, 'scaler_cleaned.sav'),'rb'))
 43        _cache[gpr_path] = (gp_U, gp_V, gp_W, scaler)
 44    else:
 45        (gp_U, gp_V, gp_W, scaler) = _cache[gpr_path]
 46
 47    input_arr = scaler.transform(vertices)
 48
 49    N = 310000
 50    if len(input_arr) < N:
 51        u = gp_U.predict(input_arr)
 52        v = gp_V.predict(input_arr)
 53        w = gp_W.predict(input_arr)
 54        disp = np.column_stack((u,v,w))
 55    else:
 56        # Must stage it
 57        disp = np.zeros_like(input_arr)
 58        for batch_i in range(int(np.ceil(len(input_arr)/N))):
 59            u = gp_U.predict(input_arr[N*batch_i:N*(batch_i+1)])
 60            v = gp_V.predict(input_arr[N*batch_i:N*(batch_i+1)])
 61            w = gp_W.predict(input_arr[N*batch_i:N*(batch_i+1)])
 62            disp[N*batch_i:N*(batch_i+1)] = np.column_stack((u,v,w))
 63
 64    return disp
 65
 66
 67def get_u_from_gpr_main():
 68    """The function invoked by the command. Parses arguments and passes
 69    to `get_predicted_u`.
 70    """
 71    parser = argparse.ArgumentParser(
 72        description="Evaluate GPR model at provided vertices, for use "
 73        "in FM-Track environment"
 74    )
 75    parser.add_argument(
 76        "-v",
 77        "--vertices-file",
 78        type=str,
 79        metavar="VERTICES_FILE"
 80    )
 81    parser.add_argument(
 82        "-d",
 83        "--dest",
 84        type=str,
 85        metavar="DEST"
 86    )
 87    parser.add_argument(
 88        "-g",
 89        "--gpr-dir",
 90        type=str,
 91        metavar="GPR_DIR"
 92    )
 93    args = parser.parse_args()
 94
 95    vertices = np.loadtxt(args.vertices_file)
 96    u = get_predicted_u(args.gpr_dir, vertices)
 97    np.savetxt(args.dest, u)
 98
 99
100if __name__=="__main__":
101    get_u_from_gpr_main()