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