persistent_tasmanian

Required: Tasmanian, pypackaging, scikit-build

Note that Tasmanian can be pip installed, but currently must use either venv or –user install.

E.g: pip install scikit-build packaging Tasmanian --user

A persistent generator using the uncertainty quantification capabilities in Tasmanian.

persistent_tasmanian.get_2D_duplicate_indices(x, y, x_ord=<MagicMock name='mock()' id='140379695138192'>, y_ord=<MagicMock name='mock()' id='140379695512400'>, tol=1e-12)

Finds the row indices of a 2D numpy array x which overlap with y. If x_ord (resp. y_ord) is empty, then x (resp. y) must be lexicographically sorted. Otherwise, x[x_ord] (resp. y[y_ord]) must be lexicographically sorted.Complexity is O(x.shape[0] + y.shape[0]).

persistent_tasmanian.get_2D_insert_indices(x, y, x_ord=<MagicMock name='mock()' id='140379695482896'>, y_ord=<MagicMock name='mock()' id='140379695017808'>, tol=1e-12)

Finds the row indices in a 2D numpy array x for which the sorted values of y can be inserted into. If x_ord (resp. y_ord) is empty, then x (resp. y) must be lexicographically sorted. Otherwise, x[x_ord] (resp. y[y_ord]) must be lexicographically sorted. Complexity is O(x.shape[0] + y.shape[0]).

persistent_tasmanian.get_H0(gen_specs, refined_pts, refined_ord, queued_pts, queued_ids, tol=1e-12)

For runs following the first one, get the history array H0 based on the ordering in refined_pts

persistent_tasmanian.get_sparse_grid_specs(user_specs, sim_f, num_dims, num_outputs=1, mode='batched')

Helper function that generates the simulator, generator, and allocator specs as well as the persis_info dictionary to ensure that they are compatible with the custom generators in this script. The outputs should be used in the main libE() call.

INPUTS:
user_specs (dict)a dictionary of user specs that is needed in the generator specs;

expects the key ‘tasmanian_init’ whose value is a 0-argument lambda that initializes an appropriate Tasmanian sparse grid object.

sim_f (func)a lambda function that takes in generator outputs (simulator inputs)

and returns simulator outputs.

num_dims (int) : number of model inputs.

num_outputs (int) : number of model outputs.

mode (string) : can either be ‘batched’ or ‘async’.

OUTPUTS:

sim_specs (dict) : a dictionary of simulation specs and also one of the inputs of libE().

gen_specs (dict) : a dictionary of generator specs and also one of the inputs of libE().

alloc_specs (dict) : a dictionary of allocation specs and also one of the inputs of libE().

persis_info (dict)a dictionary containing common information that is passed to all

workers and also one of the inputs of libE().

persistent_tasmanian.get_state(queued_pts, queued_ids, id_offset, new_points=<MagicMock name='mock()' id='140379695264144'>, completed_points=<MagicMock name='mock()' id='140379695376272'>, tol=1e-12)

Creates the data to be sent and updates the state arrays and scalars if new information (new_points or compeleted_points) arrives. Ensures that the output state arrays remain sorted if the input state arrays are already sorted.

persistent_tasmanian.lex_le(x, y, tol=1e-12)

Returns True if x <= y lexicographically up to some tolerance.

persistent_tasmanian.sparse_grid_async(H, persis_info, gen_specs, libE_info)

Implements asynchronous construction for a Tasmanian sparse grid, using the logic in the dynamic Tasmanian model construction function: sparse grid dynamic example

persistent_tasmanian.sparse_grid_batched(H, persis_info, gen_specs, libE_info)

Implements batched construction for a Tasmanian sparse grid, using the loop described in Tasmanian Example 09: sparse grid example

persistent_tasmanian.py

  1"""
  2A persistent generator using the uncertainty quantification capabilities in
  3`Tasmanian <https://tasmanian.ornl.gov/>`_.
  4"""
  5
  6import numpy as np
  7from libensemble.message_numbers import STOP_TAG, PERSIS_STOP, FINISHED_PERSISTENT_GEN_TAG, EVAL_GEN_TAG
  8from libensemble.tools.persistent_support import PersistentSupport
  9from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens as allocf
 10from libensemble.tools import parse_args
 11
 12
 13def lex_le(x, y, tol=1e-12):
 14    """
 15    Returns True if x <= y lexicographically up to some tolerance.
 16    """
 17    cmp = np.fabs(x - y) > tol
 18    ind = np.argmax(cmp)
 19    if not cmp[ind]:
 20        return True
 21    return x[ind] <= y[ind]
 22
 23
 24def get_2D_insert_indices(x, y, x_ord=np.empty(0, dtype="int"), y_ord=np.empty(0, dtype="int"), tol=1e-12):
 25    """
 26    Finds the row indices in a 2D numpy array `x` for which the sorted values of `y` can be inserted
 27    into. If `x_ord` (resp. `y_ord`) is empty, then `x` (resp. `y`) must be lexicographically
 28    sorted. Otherwise, `x[x_ord]` (resp. `y[y_ord]`) must be lexicographically sorted. Complexity is
 29    O(x.shape[0] + y.shape[0]).
 30    """
 31    assert len(x.shape) == 2
 32    assert len(y.shape) == 2
 33    if x.size == 0:
 34        return np.zeros(y.shape[0], dtype="int")
 35    else:
 36        if x_ord.size == 0:
 37            x_ord = np.arange(x.shape[0], dtype="int")
 38        if y_ord.size == 0:
 39            y_ord = np.arange(y.shape[0], dtype="int")
 40        x_ptr = 0
 41        y_ptr = 0
 42        out_ord = np.empty(0, dtype="int")
 43        while y_ptr < y.shape[0]:
 44            # The case where y[k] <= max of x[k:end, :]
 45            xk = x[x_ord[x_ptr], :]
 46            yk = y[y_ord[y_ptr], :]
 47            if lex_le(yk, xk, tol=tol):
 48                out_ord = np.append(out_ord, x_ord[x_ptr])
 49                y_ptr += 1
 50            else:
 51                x_ptr += 1
 52                # The edge case where y[k] is the largest of all elements of x.
 53                if x_ptr >= x_ord.shape[0]:
 54                    for i in range(y_ptr, y_ord.shape[0], 1):
 55                        out_ord = np.append(out_ord, x_ord.shape[0])
 56                        y_ptr += 1
 57                    break
 58        return out_ord
 59
 60
 61def get_2D_duplicate_indices(x, y, x_ord=np.empty(0, dtype="int"), y_ord=np.empty(0, dtype="int"), tol=1e-12):
 62    """
 63    Finds the row indices of a 2D numpy array `x` which overlap with `y`. If `x_ord` (resp. `y_ord`)
 64    is empty, then `x` (resp. `y`) must be lexicographically sorted. Otherwise, `x[x_ord]` (resp.
 65    `y[y_ord]`) must be lexicographically sorted.Complexity is O(x.shape[0] + y.shape[0]).
 66    """
 67    assert len(x.shape) == 2
 68    assert len(y.shape) == 2
 69    if x.size == 0:
 70        return np.empty(0, dtype="int")
 71    else:
 72        if x_ord.size == 0:
 73            x_ord = np.arange(x.shape[0], dtype="int")
 74        if y_ord.size == 0:
 75            y_ord = np.arange(y.shape[0], dtype="int")
 76        x_ptr = 0
 77        y_ptr = 0
 78        out_ord = np.empty(0, dtype="int")
 79        while y_ptr < y.shape[0] and x_ptr < x.shape[0]:
 80            # The case where y[k] <= max of x[k:end, :]
 81            xk = x[x_ord[x_ptr], :]
 82            yk = y[y_ord[y_ptr], :]
 83            if all(np.fabs(yk - xk) <= tol):
 84                out_ord = np.append(out_ord, x_ord[x_ptr])
 85                x_ptr += 1
 86            elif lex_le(xk, yk, tol=tol):
 87                x_ptr += 1
 88            else:
 89                y_ptr += 1
 90        return out_ord
 91
 92
 93def get_state(queued_pts, queued_ids, id_offset, new_points=np.array([]), completed_points=np.array([]), tol=1e-12):
 94    """
 95    Creates the data to be sent and updates the state arrays and scalars if new information
 96    (new_points or compeleted_points) arrives. Ensures that the output state arrays remain sorted if
 97    the input state arrays are already sorted.
 98    """
 99    if new_points.size > 0:
100        new_points_ord = np.lexsort(np.rot90(new_points))
101        new_points_ids = id_offset + np.arange(new_points.shape[0])
102        id_offset += new_points.shape[0]
103        insert_idx = get_2D_insert_indices(queued_pts, new_points, y_ord=new_points_ord, tol=tol)
104        queued_pts = np.insert(queued_pts, insert_idx, new_points[new_points_ord], axis=0)
105        queued_ids = np.insert(queued_ids, insert_idx, new_points_ids[new_points_ord], axis=0)
106
107    if completed_points.size > 0:
108        completed_ord = np.lexsort(np.rot90(completed_points))
109        delete_ind = get_2D_duplicate_indices(queued_pts, completed_points, y_ord=completed_ord, tol=tol)
110        queued_pts = np.delete(queued_pts, delete_ind, axis=0)
111        queued_ids = np.delete(queued_ids, delete_ind, axis=0)
112
113    return queued_pts, queued_ids, id_offset
114
115
116def get_H0(gen_specs, refined_pts, refined_ord, queued_pts, queued_ids, tol=1e-12):
117    """
118    For runs following the first one, get the history array H0 based on the ordering in `refined_pts`
119    """
120
121    def approx_eq(x, y):
122        return np.argmax(np.fabs(x - y)) <= tol
123
124    num_ids = queued_ids.shape[0]
125    H0 = np.zeros(num_ids, dtype=gen_specs["out"])
126    refined_priority = np.flip(np.arange(refined_pts.shape[0], dtype="int"))
127    rptr = 0
128    for qptr in range(num_ids):
129        while not approx_eq(refined_pts[refined_ord[rptr]], queued_pts[qptr]):
130            rptr += 1
131        assert rptr <= refined_pts.shape[0]
132        H0["x"][qptr] = queued_pts[qptr]
133        H0["sim_id"][qptr] = queued_ids[qptr]
134        H0["priority"][qptr] = refined_priority[refined_ord[rptr]]
135    return H0
136
137
138# ========================
139# Main generator functions
140# ========================
141
142
143def sparse_grid_batched(H, persis_info, gen_specs, libE_info):
144    """
145    Implements batched construction for a Tasmanian sparse grid,
146    using the loop described in Tasmanian Example 09:
147    `sparse grid example <https://github.com/ORNL/TASMANIAN/blob/master/InterfacePython/example_sparse_grids_09.py>`_
148
149    """
150    U = gen_specs["user"]
151    ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
152    grid = U["tasmanian_init"]()  # initialize the grid
153    allowed_refinements = [
154        "setAnisotropicRefinement",
155        "getAnisotropicRefinement",
156        "setSurplusRefinement",
157        "getSurplusRefinement",
158        "none",
159    ]
160    assert (
161        "refinement" in U and U["refinement"] in allowed_refinements
162    ), "Must provide a gen_specs['user']['refinement'] in: {}".format(allowed_refinements)
163
164    while grid.getNumNeeded() > 0:
165        aPoints = grid.getNeededPoints()
166
167        H0 = np.zeros(len(aPoints), dtype=gen_specs["out"])
168        H0["x"] = aPoints
169
170        # Receive values from manager
171        tag, Work, calc_in = ps.send_recv(H0)
172        if tag in [STOP_TAG, PERSIS_STOP]:
173            break
174        aModelValues = calc_in["f"]
175
176        # Update surrogate on grid
177        t = aModelValues.reshape((aModelValues.shape[0], grid.getNumOutputs()))
178        t = t.flatten()
179        t = np.atleast_2d(t).T
180        grid.loadNeededPoints(t)
181
182        if "tasmanian_checkpoint_file" in U:
183            grid.write(U["tasmanian_checkpoint_file"])
184
185        # set refinement, using user['refinement'] to pick the refinement strategy
186        if U["refinement"] in ["setAnisotropicRefinement", "getAnisotropicRefinement"]:
187            assert "sType" in U
188            assert "iMinGrowth" in U
189            assert "iOutput" in U
190            grid.setAnisotropicRefinement(U["sType"], U["iMinGrowth"], U["iOutput"])
191        elif U["refinement"] in ["setSurplusRefinement", "getSurplusRefinement"]:
192            assert "fTolerance" in U
193            assert "iOutput" in U
194            assert "sCriteria" in U
195            grid.setSurplusRefinement(U["fTolerance"], U["iOutput"], U["sCriteria"])
196
197    return H0, persis_info, FINISHED_PERSISTENT_GEN_TAG
198
199
200def sparse_grid_async(H, persis_info, gen_specs, libE_info):
201    """
202    Implements asynchronous construction for a Tasmanian sparse grid,
203    using the logic in the dynamic Tasmanian model construction function:
204    `sparse grid dynamic example <https://github.com/ORNL/TASMANIAN/blob/master/Addons/tsgConstructSurrogate.hpp>`_
205
206    """
207    U = gen_specs["user"]
208    ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
209    grid = U["tasmanian_init"]()  # initialize the grid
210    allowed_refinements = ["getCandidateConstructionPoints", "getCandidateConstructionPointsSurplus"]
211    assert (
212        "refinement" in U and U["refinement"] in allowed_refinements
213    ), "Must provide a gen_specs['user']['refinement'] in: {}".format(allowed_refinements)
214    tol = U["_match_tolerance"] if "_match_tolerance" in U else 1.0e-12
215
216    # Choose the refinement function based on U['refinement'].
217    if U["refinement"] == "getCandidateConstructionPoints":
218        assert "sType" in U
219        assert "liAnisotropicWeightsOrOutput" in U
220    if U["refinement"] == "getCandidateConstructionPointsSurplus":
221        assert "fTolerance" in U
222        assert "sRefinementType" in U
223
224    def get_refined_points(g, U):
225        if U["refinement"] == "getCandidateConstructionPoints":
226            return g.getCandidateConstructionPoints(U["sType"], U["liAnisotropicWeightsOrOutput"])
227        else:
228            assert U["refinement"] == "getCandidateConstructionPointsSurplus"
229            return g.getCandidateConstructionPointsSurplus(U["fTolerance"], U["sRefinementType"])
230        # else:
231        #     raise ValueError("Unknown refinement string")
232
233    # Asynchronous helper and state variables.
234    num_dims = grid.getNumDimensions()
235    num_completed = 0
236    offset = 0
237    queued_pts = np.empty((0, num_dims), dtype="float")
238    queued_ids = np.empty(0, dtype="int")
239
240    # First run.
241    grid.beginConstruction()
242    init_pts = get_refined_points(grid, U)
243    queued_pts, queued_ids, offset = get_state(queued_pts, queued_ids, offset, new_points=init_pts, tol=tol)
244    H0 = np.zeros(init_pts.shape[0], dtype=gen_specs["out"])
245    H0["x"] = init_pts
246    H0["sim_id"] = np.arange(init_pts.shape[0], dtype="int")
247    H0["priority"] = np.flip(H0["sim_id"])
248    tag, Work, calc_in = ps.send_recv(H0)
249
250    # Subsequent runs.
251    while tag not in [STOP_TAG, PERSIS_STOP]:
252
253        # Parse the points returned by the allocator.
254        num_completed += calc_in["x"].shape[0]
255        queued_pts, queued_ids, offset = get_state(
256            queued_pts, queued_ids, offset, completed_points=calc_in["x"], tol=tol
257        )
258
259        # Compute the next batch of points (if they exist).
260        new_pts = np.empty((0, num_dims), dtype="float")
261        refined_pts = np.empty((0, num_dims), dtype="float")
262        refined_ord = np.empty(0, dtype="int")
263        if grid.getNumLoaded() < 1000 or num_completed > 0.2 * grid.getNumLoaded():
264            # A copy is needed because the data in the calc_in arrays are not contiguous.
265            grid.loadConstructedPoint(np.copy(calc_in["x"]), np.copy(calc_in["f"]))
266            if "tasmanian_checkpoint_file" in U:
267                grid.write(U["tasmanian_checkpoint_file"])
268            refined_pts = get_refined_points(grid, U)
269            # If the refined points are empty, then there is a stopping condition internal to the
270            # Tasmanian sparse grid that is being triggered by the loaded points.
271            if refined_pts.size == 0:
272                break
273            refined_ord = np.lexsort(np.rot90(refined_pts))
274            delete_ind = get_2D_duplicate_indices(refined_pts, queued_pts, x_ord=refined_ord, tol=tol)
275            new_pts = np.delete(refined_pts, delete_ind, axis=0)
276
277        if new_pts.shape[0] > 0:
278            # Update the state variables with the refined points and update the queue in the allocator.
279            num_completed = 0
280            queued_pts, queued_ids, offset = get_state(queued_pts, queued_ids, offset, new_points=new_pts, tol=tol)
281            H0 = get_H0(gen_specs, refined_pts, refined_ord, queued_pts, queued_ids, tol=tol)
282            tag, Work, calc_in = ps.send_recv(H0)
283        else:
284            tag, Work, calc_in = ps.recv()
285
286    return [], persis_info, FINISHED_PERSISTENT_GEN_TAG
287
288
289def get_sparse_grid_specs(user_specs, sim_f, num_dims, num_outputs=1, mode="batched"):
290    """
291    Helper function that generates the simulator, generator, and allocator specs as well as the
292    persis_info dictionary to ensure that they are compatible with the custom generators in this
293    script. The outputs should be used in the main libE() call.
294
295    INPUTS:
296        user_specs  (dict)   : a dictionary of user specs that is needed in the generator specs;
297                               expects the key 'tasmanian_init' whose value is a 0-argument lambda
298                               that initializes an appropriate Tasmanian sparse grid object.
299
300        sim_f       (func)   : a lambda function that takes in generator outputs (simulator inputs)
301                               and returns simulator outputs.
302
303        num_dims    (int)    : number of model inputs.
304
305        num_outputs (int)    : number of model outputs.
306
307        mode        (string) : can either be 'batched' or 'async'.
308
309    OUTPUTS:
310        sim_specs   (dict) : a dictionary of simulation specs and also one of the inputs of libE().
311
312        gen_specs   (dict) : a dictionary of generator specs and also one of the inputs of libE().
313
314        alloc_specs (dict) : a dictionary of allocation specs and also one of the inputs of libE().
315
316        persis_info (dict) : a dictionary containing common information that is passed to all
317                             workers and also one of the inputs of libE().
318
319    """
320
321    assert "tasmanian_init" in user_specs
322    assert mode in ["batched", "async"]
323
324    sim_specs = {
325        "sim_f": sim_f,
326        "in": ["x"],
327    }
328    gen_out = [
329        ("x", float, (num_dims,)),
330        ("sim_id", int),
331        ("priority", int),
332    ]
333    gen_specs = {
334        "persis_in": [t[0] for t in gen_out] + ["f"],
335        "out": gen_out,
336        "user": user_specs,
337    }
338    alloc_specs = {
339        "alloc_f": allocf,
340        "user": {},
341    }
342
343    if mode == "batched":
344        gen_specs["gen_f"] = sparse_grid_batched
345        sim_specs["out"] = [("f", float, (num_outputs,))]
346    if mode == "async":
347        gen_specs["gen_f"] = sparse_grid_async
348        sim_specs["out"] = [("x", float, (num_dims,)), ("f", float, (num_outputs,))]
349        alloc_specs["user"]["active_recv_gen"] = True
350        alloc_specs["user"]["async_return"] = True
351
352    nworkers, _, _, _ = parse_args()
353    persis_info = {}
354    for i in range(nworkers + 1):
355        persis_info[i] = {"worker_num": i}
356
357    return sim_specs, gen_specs, alloc_specs, persis_info