Adding test for ibcdfo with jax

install/install_ibcdfo.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-git clone --recurse-submodules -b develop https://github.com/POptUS/IBCDFO.git
+git clone --recurse-submodules -b updates_manifold_sampling https://github.com/POptUS/IBCDFO.git
 pushd IBCDFO/minq/py/minq5/
 export PYTHONPATH="$PYTHONPATH:$(pwd)"
 echo "PYTHONPATH=$PYTHONPATH" >> $GITHUB_ENV

libensemble/gen_funcs/aposmm_localopt_support.py

@@ -9,6 +9,7 @@
     "run_local_tao",
     "run_local_dfols",
     "run_local_ibcdfo_pounders",
+    "run_local_ibcdfo_manifold_sampling",
     "run_local_scipy_opt",
     "run_external_localopt",
 ]
@@ -26,7 +27,7 @@ class APOSMMException(Exception):
     """Raised for any exception in APOSMM"""
 
 
-optimizer_list = ["petsc", "nlopt", "dfols", "scipy", "ibcdfo", "external"]
+optimizer_list = ["petsc", "nlopt", "dfols", "scipy", "ibcdfo_pounders", "ibcdfo_manifold_sampling", "external"]
 optimizers = libensemble.gen_funcs.rc.aposmm_optimizers
 
 if optimizers is not None:
@@ -42,8 +43,10 @@ class APOSMMException(Exception):
         import nlopt  # noqa: F401
     if "dfols" in optimizers:
         import dfols  # noqa: F401
-    if "ibcdfo" in optimizers:
-        from ibcdfo import pounders  # noqa: F401
+    if "ibcdfo_pounders" in optimizers:
+        from ibcdfo.pounders import pounders  # noqa: F401
+    if "ibcdfo_manifold_sampling" in optimizers:
+        from ibcdfo.manifold_sampling import manifold_sampling_primal  # noqa: F401
     if "scipy" in optimizers:
         from scipy import optimize as sp_opt  # noqa: F401
     if "external_localopt" in optimizers:
@@ -79,6 +82,7 @@ class LocalOptInterfacer(object):
     - PETSc/TAO [``'pounders'``, ``'blmvm'``, ``'nm'``]
     - SciPy [``'scipy_Nelder-Mead'``, ``'scipy_COBYLA'``, ``'scipy_BFGS'``]
     - DFOLS [``'dfols'``]
+    - IBCDFO [``'pounders'``, ``'manifold_sampling_primal'``]
     - External local optimizer [``'external_localopt'``] (which use files to pass/receive ``x/f`` values)
     """
 
@@ -123,6 +127,8 @@ def __init__(self, user_specs, x0, f0, grad0=None):
             run_local_opt = run_local_dfols
         elif user_specs["localopt_method"] in ["ibcdfo_pounders"]:
             run_local_opt = run_local_ibcdfo_pounders
+        elif user_specs["localopt_method"] in ["ibcdfo_manifold_sampling"]:
+            run_local_opt = run_local_ibcdfo_manifold_sampling
         elif user_specs["localopt_method"] in ["external_localopt"]:
             run_local_opt = run_external_localopt
         else:
@@ -417,6 +423,60 @@ def run_local_dfols(user_specs, comm_queue, x0, f0, child_can_read, parent_can_r
     finish_queue(x_opt, opt_flag, comm_queue, parent_can_read, user_specs)
 
 
+def run_local_ibcdfo_manifold_sampling(user_specs, comm_queue, x0, f0, child_can_read, parent_can_read):
+    """
+    Runs a IBCDFO local optimization run starting at ``x0``, governed by the
+    parameters in ``user_specs``.
+
+    Although IBCDFO methods can receive previous evaluations, few other methods
+    support that, so APOSMM assumes the first point will be re-evaluated (but
+    not be sent back to the manager).
+    """
+    n = len(x0)
+    # Define bound constraints (lower <= x <= upper)
+    lb = np.zeros(n)
+    ub = np.ones(n)
+
+    # Set random seed (for reproducibility)
+    np.random.seed(0)
+
+    # dist_to_bound = min(min(ub - x0), min(x0 - lb))
+    # assert dist_to_bound > np.finfo(np.float64).eps, "The distance to the boundary is too small"
+
+    run_max_eval = user_specs.get("run_max_eval", 100 * (n + 1))
+    # g_tol = 1e-8
+    # delta_0 = 0.5 * dist_to_bound
+    # m = len(f0)
+    subprob_switch = "linprog"
+
+    [X, F, hF, xkin, flag] = manifold_sampling_primal(
+        user_specs["hfun"],
+        lambda x: scipy_dfols_callback_fun(x, comm_queue, child_can_read, parent_can_read, user_specs),
+        x0,
+        lb,
+        ub,
+        run_max_eval,
+        subprob_switch,
+    )
+
+    assert flag >= 0 or flag == -6, "IBCDFO errored"
+
+    x_opt = X[xkin]
+
+    if flag > 0:
+        opt_flag = 1
+    else:
+        print(
+            "[APOSMM] The IBCDFO run started from " + str(x0) + " stopped with an exit "
+            "flag of " + str(flag) + ". No point from this run will be "
+            "ruled as a minimum! APOSMM may start a new run from some point "
+            "in this run."
+        )
+        opt_flag = 0
+
+    finish_queue(x_opt, opt_flag, comm_queue, parent_can_read, user_specs)
+
+
 def run_local_ibcdfo_pounders(user_specs, comm_queue, x0, f0, child_can_read, parent_can_read):
     """
     Runs a IBCDFO local optimization run starting at ``x0``, governed by the
@@ -447,7 +507,7 @@ def run_local_ibcdfo_pounders(user_specs, comm_queue, x0, f0, child_can_read, pa
     else:
         Options = None
 
-    [X, F, hF, flag, xkin] = pounders.pounders(
+    [X, F, hF, flag, xkin] = pounders(
         lambda x: scipy_dfols_callback_fun(x, comm_queue, child_can_read, parent_can_read, user_specs),
         x0,
         n,

libensemble/gen_funcs/persistent_aposmm.py

@@ -739,7 +739,7 @@ def initialize_children(user_specs):
         "nm",
     ]:
         fields_to_pass = ["x_on_cube", "f"]
-    elif user_specs["localopt_method"] in ["pounders", "ibcdfo_pounders", "dfols"]:
+    elif user_specs["localopt_method"] in ["pounders", "ibcdfo_pounders", "ibcdfo_manifold_sampling", "dfols"]:
         fields_to_pass = ["x_on_cube", "fvec"]
     else:
         raise NotImplementedError(f"Unknown local optimization method {user_specs['localopt_method']}.")

libensemble/tests/regression_tests/declare_hfun_and_combine_model_with_jax.py

@@ -0,0 +1,50 @@
+# This declares the hfun for Test_compare_pounder_pounders_with_jax.py and
+# then used jax to combine the quadratic models of each component of the
+# inputs to the hfun.
+#
+# For other general use cases of pounders on smooth hfuns, only the hfun below
+# needs to be changed (and combinemodels_jax can be given to pounders)
+
+
+import jax
+import numpy
+
+jax.config.update("jax_enable_x64", True)
+
+
+def hfun(z):
+    res = z[0] * z[1] - z[2] ** 2
+    return res
+
+
+@jax.jit
+def hfun_d(z, zd):
+    resd = jax.jvp(hfun, (z,), (zd,))
+    return resd
+
+
+@jax.jit
+def hfun_dd(z, zd, zdt, zdd):
+    _, resdd = jax.jvp(hfun_d, (z, zd), (zdt, zdd))
+    return resdd
+
+
+def G_combine(Cres, Gres):
+    n, m = Gres.shape
+    G = numpy.zeros(n)
+    for i in range(n):
+        _, G[i] = hfun_d(Cres, Gres[i, :])
+    return G
+
+
+def H_combine(Cres, Gres, Hres):
+    n, _, m = Hres.shape
+    H = numpy.zeros((n, n))
+    for i in range(n):
+        for j in range(n):
+            _, H[i, j] = hfun_dd(Cres, Gres[i, :], Gres[j, :], Hres[i, j, :])
+    return H
+
+
+def combinemodels_jax(Cres, Gres, Hres):
+    return G_combine(Cres, Gres), H_combine(Cres, Gres, Hres)

libensemble/tests/regression_tests/test_persistent_aposmm_ibcdfo_manifold_sampling.py

@@ -0,0 +1,129 @@
+"""
+Runs libEnsemble with APOSMM+IBCDFO on two test problems. Only a single
+optimization run is being performed for the below setup.
+
+The first case uses POUNDERS to solve the chwirut least-squares problem. For
+this case, all chwirut 214 residual calculations for a given point are
+performed as a single simulation evaluation.
+
+The second case uses the generalized POUNDERS to minimize normalized beamline
+emittance. The "beamline simulation" is a synthetic polynomial test function
+that takes in 4 variables and returning 3 outputs. These outputs represent
+position <x>, momentum <p_x>, and the correlation between them <x p_x>.
+
+These values are then mapped to the normalized emittance <x> <p_x> - <x p_x>.
+
+Execute via one of the following commands:
+   mpiexec -np 3 python test_persistent_aposmm_ibcdfo_pounders.py
+   python test_persistent_aposmm_ibcdfo_pounders.py --nworkers 2
+Both will run with 1 manager, 1 worker running APOSMM+IBCDFO, and 1 worker
+doing the simulation evaluations.
+"""
+
+# Do not change these lines - they are parsed by run-tests.sh
+# TESTSUITE_COMMS: local mpi
+# TESTSUITE_NPROCS: 3
+
+import multiprocessing
+import sys
+
+import numpy as np
+
+import libensemble.gen_funcs
+from libensemble.libE import libE
+
+libensemble.gen_funcs.rc.aposmm_optimizers = "ibcdfo_manifold_sampling"
+
+from libensemble.alloc_funcs.persistent_aposmm_alloc import persistent_aposmm_alloc as alloc_f
+from libensemble.gen_funcs.persistent_aposmm import aposmm as gen_f
+from libensemble.tools import add_unique_random_streams, parse_args, save_libE_output
+
+try:
+    from ibcdfo.manifold_sampling import manifold_sampling_primal  # noqa: F401
+    from ibcdfo.manifold_sampling.h_examples import pw_maximum as hfun
+
+except ModuleNotFoundError:
+    sys.exit("Please 'pip install ibcdfo'")
+
+try:
+    from minqsw import minqsw  # noqa: F401
+
+except ModuleNotFoundError:
+    sys.exit("Ensure https://github.com/POptUS/minq has been cloned and that minq/py/minq5/ is on the PYTHONPATH")
+
+
+def synthetic_beamline_mapping(H, _, sim_specs):
+    x = H["x"][0]
+    assert len(x) == 4, "Assuming 4 inputs to this function"
+    y = np.zeros(3)  # Synthetic beamline outputs
+    y[0] = x[0] ** 2 + 1.0
+    y[1] = x[1] ** 2 + 2.0
+    y[2] = x[2] * x[3] + 0.5
+
+    Out = np.zeros(1, dtype=sim_specs["out"])
+    Out["fvec"] = y
+    Out["f"] = np.max(y)
+    return Out
+
+
+# Main block is necessary only when using local comms with spawn start method (default on macOS and Windows).
+if __name__ == "__main__":
+    multiprocessing.set_start_method("fork", force=True)
+
+    nworkers, is_manager, libE_specs, _ = parse_args()
+
+    assert nworkers == 2, "This test is just for two workers"
+
+    m = 3
+    n = 4
+    sim_f = synthetic_beamline_mapping
+
+    sim_specs = {
+        "sim_f": sim_f,
+        "in": ["x"],
+        "out": [("f", float), ("fvec", float, m)],
+    }
+
+    gen_out = [
+        ("x", float, n),
+        ("x_on_cube", float, n),
+        ("sim_id", int),
+        ("local_min", bool),
+        ("local_pt", bool),
+        ("started_run", bool),
+    ]
+
+    gen_specs = {
+        "gen_f": gen_f,
+        "persis_in": ["f", "fvec"] + [n[0] for n in gen_out],
+        "out": gen_out,
+        "user": {
+            "initial_sample_size": 1,
+            "stop_after_k_runs": 1,
+            "max_active_runs": 1,
+            "sample_points": np.atleast_2d(0.1 * (np.arange(n) + 1)),
+            "localopt_method": "ibcdfo_manifold_sampling",
+            "run_max_eval": 100 * (n + 1),
+            "components": m,
+            "lb": -1 * np.ones(n),
+            "ub": np.ones(n),
+        },
+    }
+
+    gen_specs["user"]["hfun"] = hfun
+
+    alloc_specs = {"alloc_f": alloc_f}
+
+    persis_info = add_unique_random_streams({}, nworkers + 1)
+
+    exit_criteria = {"sim_max": 500}
+
+    # Perform the run
+    H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs)
+
+    if is_manager:
+        assert np.min(H["f"]) == 2.0, "The best is 2"
+        assert persis_info[1].get("run_order"), "Run_order should have been given back"
+        assert flag == 0
+
+        save_libE_output(H, persis_info, __file__, nworkers)

libensemble/tests/regression_tests/test_persistent_aposmm_ibcdfo_pounders.py

@@ -33,7 +33,7 @@
 from libensemble.libE import libE
 from libensemble.sim_funcs.chwirut1 import chwirut_eval
 
-libensemble.gen_funcs.rc.aposmm_optimizers = "ibcdfo"
+libensemble.gen_funcs.rc.aposmm_optimizers = "ibcdfo_pounders"
 
 from libensemble.alloc_funcs.persistent_aposmm_alloc import persistent_aposmm_alloc as alloc_f
 from libensemble.gen_funcs.persistent_aposmm import aposmm as gen_f