Coverage for HARK/simulator.py: 89%
1350 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-02 05:14 +0000
1"""
2A module with classes and functions for automated simulation of HARK.AgentType
3models from a human- and machine-readable model specification.
4"""
6from dataclasses import dataclass, field
7from copy import copy, deepcopy
8import numpy as np
9from numba import njit
10from sympy.utilities.lambdify import lambdify
11from sympy import symbols, IndexedBase
12from typing import Callable
13from HARK.utilities import NullFunc
14from HARK.distributions import Distribution
15from scipy.sparse import csr_matrix
16from scipy.sparse.linalg import eigs
17from itertools import product
18import importlib.resources
19import yaml
21# Prevent pre-commit from removing sympy
22x = symbols("x")
23del x
24y = IndexedBase("y")
25del y
28@dataclass(kw_only=True)
29class ModelEvent:
30 """
31 Class for representing "events" that happen to agents in the course of their
32 model. These might be statements of dynamics, realization of a random shock,
33 or the evaluation of a function (potentially a control or other solution-
34 based object). This is a superclass for types of events defined below.
36 Parameters
37 ----------
38 description : str
39 Text description of this model event.
40 statement : str
41 The line of the model statement that this event corresponds to.
42 parameters : dict
43 Dictionary of objects that are static / universal within this event.
44 assigns : list[str]
45 List of names of variables that this event assigns values for.
46 needs : list[str]
47 List of names of variables that this event requires to be run.
48 data : dict
49 Dictionary of current variable values within this event.
50 common : bool
51 Indicator for whether the variables assigned in this event are commonly
52 held across all agents, rather than idiosyncratic.
53 N : int
54 Number of agents currently in this event.
55 """
57 statement: str = field(default="")
58 parameters: dict = field(default_factory=dict)
59 description: str = field(default="")
60 assigns: list[str] = field(default_factory=list, repr=False)
61 needs: list = field(default_factory=list, repr=False)
62 data: dict = field(default_factory=dict, repr=False)
63 common: bool = field(default=False, repr=False)
64 N: int = field(default=1, repr=False)
66 def run(self):
67 """
68 This method should be filled in by each subclass.
69 """
70 pass
72 def reset(self):
73 self.data = {}
75 def assign(self, output):
76 if len(self.assigns) > 1:
77 assert len(self.assigns) == len(output)
78 for j in range(len(self.assigns)):
79 var = self.assigns[j]
80 if type(output[j]) is not np.ndarray:
81 output[j] = np.array([output[j]])
82 self.data[var] = output[j]
83 else:
84 var = self.assigns[0]
85 if type(output) is not np.ndarray:
86 output = np.array([output])
87 self.data[var] = output
89 def expand_information(self, origins, probs, atoms, which=None):
90 """
91 This method is only called internally when a RandomEvent or MarkovEvent
92 runs its quasi_run() method. It expands the set of of "probability blobs"
93 by applying a random realization event. All extant blobs for which the
94 shock applies are replicated for each atom in the random event, with the
95 probability mass divided among the replicates.
97 Parameters
98 ----------
99 origins : np.array
100 Array that tracks which arrival state space node each blob originated
101 from. This is expanded into origins_new, which is returned.
102 probs : np.array
103 Vector of probabilities of each of the random possibilities.
104 atoms : [np.array]
105 List of arrays with realization values for the distribution. Each
106 array corresponds to one variable that is assigned by this event.
107 which : np.array or None
108 If given, a Boolean array indicating which of the pre-existing blobs
109 is affected by the given probabilities and atoms. By default, all
110 blobs are assumed to be affected.
112 Returns
113 -------
114 origins_new : np.array
115 Expanded boolean array of indicating the arrival state space node that
116 each blob originated from.
117 """
118 K = probs.size
119 N = self.N
120 if which is None:
121 which = np.ones(N, dtype=bool)
122 other = np.logical_not(which)
123 M = np.sum(which) # how many blobs are we affecting?
124 MX = N - M # how many blobs are we not affecting?
126 # Update probabilities of outcomes
127 pmv_old = np.reshape(self.data["pmv_"][which], (M, 1))
128 pmv_new = (pmv_old * np.reshape(probs, (1, K))).flatten()
129 self.data["pmv_"] = np.concatenate((self.data["pmv_"][other], pmv_new))
131 # Replicate the pre-existing data for each atom
132 for var in self.data.keys():
133 if (var == "pmv_") or (var in self.assigns):
134 continue # don't double expand pmv, and don't touch assigned variables
135 data_old = np.reshape(self.data[var][which], (M, 1))
136 data_new = np.tile(data_old, (1, K)).flatten()
137 self.data[var] = np.concatenate((self.data[var][other], data_new))
139 # If any of the assigned variables don't exist yet, add dummy versions
140 # of them. This section exists so that the code works with "partial events"
141 # on both the first pass and subsequent passes.
142 for j in range(len(self.assigns)):
143 var = self.assigns[j]
144 if var in self.data.keys():
145 continue
146 self.data[var] = np.zeros(N, dtype=atoms[j].dtype)
147 # Zeros are just dummy values
149 # Add the new random variables to the simulation data. This generates
150 # replicates for the affected blobs and leaves the others untouched,
151 # still with their dummy values. They will be altered on later passes.
152 for j in range(len(self.assigns)):
153 var = self.assigns[j]
154 data_new = np.tile(np.reshape(atoms[j], (1, K)), (M, 1)).flatten()
155 self.data[var] = np.concatenate((self.data[var][other], data_new))
157 # Expand the origins array to account for the new replicates
158 origins_new = np.tile(np.reshape(origins[which], (M, 1)), (1, K)).flatten()
159 origins_new = np.concatenate((origins[other], origins_new))
160 self.N = MX + M * K
162 # Send the new origins array back to the calling process
163 return origins_new
165 def add_idiosyncratic_bernoulli_info(self, origins, probs):
166 """
167 Special method for adding Bernoulli outcomes to the information set when
168 probabilities are idiosyncratic to each agent. All extant blobs are duplicated
169 with the appropriate probability
171 Parameters
172 ----------
173 origins : np.array
174 Array that tracks which arrival state space node each blob originated
175 from. This is expanded into origins_new, which is returned.
176 probs : np.array
177 Vector of probabilities of drawing True for each blob.
179 Returns
180 -------
181 origins_new : np.array
182 Expanded boolean array of indicating the arrival state space node that
183 each blob originated from.
184 """
185 N = self.N
187 # # Update probabilities of outcomes, replicating each one
188 pmv_old = np.reshape(self.data["pmv_"], (N, 1))
189 P = np.reshape(probs, (N, 1))
190 PX = np.concatenate([1.0 - P, P], axis=1)
191 pmv_new = (pmv_old * PX).flatten()
192 self.data["pmv_"] = pmv_new
194 # Replicate the pre-existing data for each atom
195 for var in self.data.keys():
196 if (var == "pmv_") or (var in self.assigns):
197 continue # don't double expand pmv, and don't touch assigned variables
198 data_old = np.reshape(self.data[var], (N, 1))
199 data_new = np.tile(data_old, (1, 2)).flatten()
200 self.data[var] = data_new
202 # Add the (one and only) new random variable to the simulation data
203 var = self.assigns[0]
204 data_new = np.tile(np.array([[0, 1]]), (N, 1)).flatten()
205 self.data[var] = data_new
207 # Expand the origins array to account for the new replicates
208 origins_new = np.tile(np.reshape(origins, (N, 1)), (1, 2)).flatten()
209 self.N = N * 2
211 # Send the new origins array back to the calling process
212 return origins_new
215@dataclass(kw_only=True)
216class DynamicEvent(ModelEvent):
217 """
218 Class for representing model dynamics for an agent, consisting of an expression
219 to be evaluated and variables to which the results are assigned.
221 Parameters
222 ----------
223 expr : Callable
224 Function or expression to be evaluated for the assigned variables.
225 args : list[str]
226 Ordered list of argument names for the expression.
227 """
229 expr: Callable = field(default_factory=NullFunc, repr=False)
230 args: list[str] = field(default_factory=list, repr=False)
232 def evaluate(self):
233 temp_dict = self.data.copy()
234 temp_dict.update(self.parameters)
235 args = (temp_dict[arg] for arg in self.args)
236 out = self.expr(*args)
237 return out
239 def run(self):
240 self.assign(self.evaluate())
242 def quasi_run(self, origins, norm=None):
243 self.run()
244 return origins
247@dataclass(kw_only=True)
248class RandomEvent(ModelEvent):
249 """
250 Class for representing the realization of random variables for an agent,
251 consisting of a shock distribution and variables to which the results are assigned.
253 Parameters
254 ----------
255 dstn : Distribution
256 Distribution of one or more random variables that are drawn from during
257 this event and assigned to the corresponding variables.
258 """
260 dstn: Distribution = field(default_factory=Distribution, repr=False)
262 def reset(self):
263 self.dstn.reset()
264 ModelEvent.reset(self)
266 def draw(self):
267 out = np.empty((len(self.assigns), self.N))
268 if not self.common:
269 out[:, :] = self.dstn.draw(self.N)
270 else:
271 out[:, :] = self.dstn.draw(1)
272 if len(self.assigns) == 1:
273 out = out.flatten()
274 return out
276 def run(self):
277 self.assign(self.draw())
279 def quasi_run(self, origins, norm=None):
280 # Get distribution
281 atoms = self.dstn.atoms
282 probs = self.dstn.pmv.copy()
284 # Apply Harmenberg normalization if applicable
285 try:
286 harm_idx = self.assigns.index(norm)
287 probs *= atoms[harm_idx]
288 except:
289 pass
291 # Expand the set of simulated blobs
292 origins_new = self.expand_information(origins, probs, atoms)
293 return origins_new
296@dataclass(kw_only=True)
297class RandomIndexedEvent(RandomEvent):
298 """
299 Class for representing the realization of random variables for an agent,
300 consisting of a list of shock distributions, and index for the list, and the
301 variables to which the results are assigned.
303 Parameters
304 ----------
305 dstn : [Distribution]
306 List of distributions of one or more random variables that are drawn
307 from during this event and assigned to the corresponding variables.
308 index : str
309 Name of the index that is used to choose a distribution for each agent.
310 """
312 index: str = field(default="", repr=False)
313 dstn: list[Distribution] = field(default_factory=list, repr=False)
315 def draw(self):
316 idx = self.data[self.index]
317 K = len(self.assigns)
318 out = np.empty((K, self.N))
319 out.fill(np.nan)
321 if self.common:
322 k = idx[0] # this will behave badly if index is not itself common
323 out[:, :] = self.dstn[k].draw(1)
324 return out
326 for k in range(len(self.dstn)):
327 these = idx == k
328 if not np.any(these):
329 continue
330 out[:, these] = self.dstn[k].draw(np.sum(these))
331 if K == 1:
332 out = out.flatten()
333 return out
335 def reset(self):
336 for k in range(len(self.dstn)):
337 self.dstn[k].reset()
338 ModelEvent.reset(self)
340 def quasi_run(self, origins, norm=None):
341 origins_new = origins.copy()
342 J = len(self.dstn)
344 for j in range(J):
345 idx = self.data[self.index]
346 these = idx == j
348 # Get distribution
349 atoms = self.dstn[j].atoms
350 probs = self.dstn[j].pmv.copy()
352 # Apply Harmenberg normalization if applicable
353 try:
354 harm_idx = self.assigns.index(norm)
355 probs *= atoms[harm_idx]
356 except:
357 pass
359 # Expand the set of simulated blobs
360 origins_new = self.expand_information(
361 origins_new, probs, atoms, which=these
362 )
364 # Return the altered origins array
365 return origins_new
368@dataclass(kw_only=True)
369class MarkovEvent(ModelEvent):
370 """
371 Class for representing the realization of a Markov draw for an agent, in which
372 a Markov probabilities (array, vector, or a single float) is used to determine
373 the realization of some discrete outcome. If the probabilities are a 2D array,
374 it represents a Markov matrix (rows sum to 1), and there must be an index; if
375 the probabilities are a vector, it should be a stochastic vector; if it's a
376 single float, it represents a Bernoulli probability.
377 """
379 probs: str = field(default="", repr=False)
380 index: str = field(default="", repr=False)
381 N: int = field(default=1, repr=False)
382 seed: int = field(default=0, repr=False)
383 # seed is overwritten when each period is created
385 def __post_init__(self):
386 self.reset_rng()
388 def reset(self):
389 self.reset_rng()
390 ModelEvent.reset(self)
392 def reset_rng(self):
393 self.RNG = np.random.RandomState(self.seed)
395 def draw(self):
396 # Initialize the output
397 out = -np.ones(self.N, dtype=int)
398 if self.probs in self.parameters:
399 probs = self.parameters[self.probs]
400 probs_are_param = True
401 else:
402 probs = self.data[self.probs]
403 probs_are_param = False
405 # Make the base draw(s)
406 if self.common:
407 X = self.RNG.rand(1)
408 else:
409 X = self.RNG.rand(self.N)
411 if self.index: # it's a Markov matrix
412 idx = self.data[self.index]
413 J = probs.shape[0]
414 for j in range(J):
415 these = idx == j
416 if not np.any(these):
417 continue
418 P = np.cumsum(probs[j, :])
419 if self.common:
420 out[:] = np.searchsorted(P, X[0]) # only one value of X!
421 else:
422 out[these] = np.searchsorted(P, X[these])
423 return out
425 if (isinstance(probs, np.ndarray)) and (
426 probs_are_param
427 ): # it's a stochastic vector
428 P = np.cumsum(probs)
429 if self.common:
430 out[:] = np.searchsorted(P, X[0])
431 return out
432 else:
433 return np.searchsorted(P, X)
435 # Otherwise, this is just a Bernoulli RV
436 P = probs
437 if self.common:
438 out[:] = X < P
439 return out
440 else:
441 return X < P # basic Bernoulli
443 def run(self):
444 self.assign(self.draw())
446 def quasi_run(self, origins, norm=None):
447 if self.probs in self.parameters:
448 probs = self.parameters[self.probs]
449 probs_are_param = True
450 else:
451 probs = self.data[self.probs]
452 probs_are_param = False
454 # If it's a Markov matrix:
455 if self.index:
456 K = probs.shape[0]
457 atoms = np.array([np.arange(probs.shape[1], dtype=int)])
458 origins_new = origins.copy()
459 for k in range(K):
460 idx = self.data[self.index]
461 these = idx == k
462 probs_temp = probs[k, :]
463 origins_new = self.expand_information(
464 origins_new, probs_temp, atoms, which=these
465 )
466 return origins_new
468 # If it's a stochastic vector:
469 if (isinstance(probs, np.ndarray)) and (probs_are_param):
470 atoms = np.array([np.arange(probs.shape[0], dtype=int)])
471 origins_new = self.expand_information(origins, probs, atoms)
472 return origins_new
474 # Otherwise, this is just a Bernoulli RV, but it might have idiosyncratic probability
475 if probs_are_param:
476 P = probs
477 atoms = np.array([[False, True]])
478 origins_new = self.expand_information(origins, np.array([1 - P, P]), atoms)
479 return origins_new
481 # Final case: probability is idiosyncratic Bernoulli
482 origins_new = self.add_idiosyncratic_bernoulli_info(origins, probs)
483 return origins_new
486@dataclass(kw_only=True)
487class EvaluationEvent(ModelEvent):
488 """
489 Class for representing the evaluation of a model function. This might be from
490 the solution of the model (like a policy function or decision rule) or just
491 a non-algebraic function used in the model. This looks a lot like DynamicEvent.
493 Parameters
494 ----------
495 func : Callable
496 Model function that is evaluated in this event, with the output assigned
497 to the appropriate variables.
498 """
500 func: Callable = field(default_factory=NullFunc, repr=False)
501 arguments: list[str] = field(default_factory=list, repr=False)
503 def evaluate(self):
504 temp_dict = self.data.copy()
505 temp_dict.update(self.parameters)
506 args_temp = (temp_dict[arg] for arg in self.arguments)
507 out = self.func(*args_temp)
508 return out
510 def run(self):
511 self.assign(self.evaluate())
513 def quasi_run(self, origins, norm=None):
514 self.run()
515 return origins
518@dataclass(kw_only=True)
519class SimBlock:
520 """
521 Class for representing a "block" of a simulated model, which might be a whole
522 period or a "stage" within a period.
524 Parameters
525 ----------
526 description : str
527 Textual description of what happens in this simulated block.
528 statement : str
529 Verbatim model statement that was used to create this block.
530 content : dict
531 Dictionary of objects that are constant / universal within the block.
532 This includes both traditional numeric parameters as well as functions.
533 arrival : list[str]
534 List of inbound states: information available at the *start* of the block.
535 events: list[ModelEvent]
536 Ordered list of events that happen during the block.
537 data: dict
538 Dictionary that stores current variable values.
539 N : int
540 Number of idiosyncratic agents in this block.
541 """
543 statement: str = field(default="", repr=False)
544 content: dict = field(default_factory=dict)
545 description: str = field(default="", repr=False)
546 arrival: list[str] = field(default_factory=list, repr=False)
547 events: list[ModelEvent] = field(default_factory=list, repr=False)
548 data: dict = field(default_factory=dict, repr=False)
549 N: int = field(default=1, repr=False)
551 def run(self):
552 """
553 Run this simulated block by running each of its events in order.
554 """
555 for j in range(len(self.events)):
556 event = self.events[j]
557 for k in range(len(event.assigns)):
558 var = event.assigns[k]
559 if var in event.data.keys():
560 del event.data[var]
561 for k in range(len(event.needs)):
562 var = event.needs[k]
563 event.data[var] = self.data[var]
564 event.N = self.N
565 event.run()
566 for k in range(len(event.assigns)):
567 var = event.assigns[k]
568 self.data[var] = event.data[var]
570 def reset(self):
571 """
572 Reset the simulated block by resetting each of its events.
573 """
574 self.data = {}
575 for j in range(len(self.events)):
576 self.events[j].reset()
578 def distribute_content(self):
579 """
580 Fill in parameters, functions, and distributions to each event.
581 """
582 for event in self.events:
583 for param in event.parameters.keys():
584 try:
585 event.parameters[param] = self.content[param]
586 except:
587 raise ValueError(
588 "Could not distribute the parameter called " + param + "!"
589 )
590 if (type(event) is RandomEvent) or (type(event) is RandomIndexedEvent):
591 try:
592 event.dstn = self.content[event._dstn_name]
593 except:
594 raise ValueError(
595 "Could not find a distribution called " + event._dstn_name + "!"
596 )
597 if type(event) is EvaluationEvent:
598 try:
599 event.func = self.content[event._func_name]
600 except:
601 raise ValueError(
602 "Could not find a function called " + event._func_name + "!"
603 )
605 def make_transition_matrices(self, grid_specs, twist=None, norm=None):
606 """
607 Construct a transition matrix for this block, moving from a discretized
608 grid of arrival variables to a discretized grid of end-of-block variables.
609 User specifies how the grids of pre-states should be built. Output is
610 stored in attributes of self as follows:
612 - matrices : A dictionary of arrays that cast from the arrival state space
613 to the grid of outcome variables. Doing np.dot(dstn, matrices[var])
614 will yield the discretized distribution of that outcome variable.
615 - grids : A dictionary of discretized grids for outcome variables. Doing
616 np.dot(np.dot(dstn, matrices[var]), grids[var]) yields the *average*
617 of that outcome in the population.
619 Parameters
620 ----------
621 grid_specs : dict
622 Dictionary of dictionaries of grid specifications. For now, these have
623 at most a minimum value, a maximum value, a number of nodes, and a poly-
624 nomial order. They are equispaced if a min and max are specified, and
625 polynomially spaced with the specified order > 0 if provided. Otherwise,
626 they are set at 0,..,N if only N is provided.
627 twist : dict or None
628 Mapping from end-of-period (continuation) variables to successor's
629 arrival variables. When this is specified, additional output is created
630 for the "full period" arrival-to-arrival transition matrix.
631 norm : str or None
632 Name of the shock variable by which to normalize for Harmenberg
633 aggregation. By default, no normalization happens.
635 Returns
636 -------
637 None
638 """
639 # Initialize dictionaries of input and output grids
640 arrival_N = len(self.arrival)
641 completed = arrival_N * [False]
642 grids_in = {}
643 grids_out = {}
644 if arrival_N == 0: # should only be for initializer block
645 dummy_grid = np.array([0])
646 grids_in["_dummy"] = dummy_grid
648 # Construct a grid for each requested variable
649 continuous_grid_out_bool = []
650 grid_orders = {}
651 for var in grid_specs.keys():
652 spec = grid_specs[var]
653 try:
654 idx = self.arrival.index(var)
655 completed[idx] = True
656 is_arrival = True
657 except:
658 is_arrival = False
659 if ("min" in spec) and ("max" in spec):
660 Q = spec["order"] if "order" in spec else 1.0
661 bot = spec["min"]
662 top = spec["max"]
663 N = spec["N"]
664 new_grid = np.linspace(0.0, 1.0, N) ** Q * (top - bot) + bot
665 is_cont = True
666 grid_orders[var] = Q
667 elif "N" in spec:
668 new_grid = np.arange(spec["N"], dtype=int)
669 is_cont = False
670 grid_orders[var] = -1
671 else:
672 new_grid = None # could not make grid, construct later
673 is_cont = False
674 grid_orders[var] = None
676 if is_arrival:
677 grids_in[var] = new_grid
678 else:
679 grids_out[var] = new_grid
680 continuous_grid_out_bool.append(is_cont)
682 # Verify that specifications were passed for all arrival variables
683 for j in range(len(self.arrival)):
684 if not completed[j]:
685 raise ValueError(
686 "No grid specification was provided for " + self.arrival[var] + "!"
687 )
689 # If an intertemporal twist was specified, make result grids for continuation variables.
690 # This overrides any grids for these variables that were explicitly specified
691 if twist is not None:
692 for cont_var in twist.keys():
693 arr_var = twist[cont_var]
694 if cont_var not in list(grids_out.keys()):
695 is_cont = grids_in[arr_var].dtype is np.dtype(np.float64)
696 continuous_grid_out_bool.append(is_cont)
697 grids_out[cont_var] = copy(grids_in[arr_var])
698 grid_orders[cont_var] = grid_orders[arr_var]
699 grid_out_is_continuous = np.array(continuous_grid_out_bool)
701 # Make meshes of all the arrival grids, which will be the initial simulation data
702 if arrival_N > 0:
703 state_meshes = np.meshgrid(
704 *[grids_in[k] for k in self.arrival], indexing="ij"
705 )
706 else: # this only happens in the initializer block
707 state_meshes = [dummy_grid.copy()]
708 state_init = {
709 self.arrival[k]: state_meshes[k].flatten() for k in range(arrival_N)
710 }
711 N_orig = state_meshes[0].size
712 self.N = N_orig
713 mesh_tuples = [
714 [state_init[self.arrival[k]][n] for k in range(arrival_N)]
715 for n in range(self.N)
716 ]
718 # Make the initial vector of probability masses
719 state_init["pmv_"] = np.ones(self.N)
721 # Initialize the array of arrival states
722 origin_array = np.arange(self.N, dtype=int)
724 # Reset the block's state and give it the initial state data
725 self.reset()
726 self.data.update(state_init)
728 # Loop through each event in order and quasi-simulate it
729 for j in range(len(self.events)):
730 event = self.events[j]
731 event.data = self.data # Give event *all* data directly
732 event.N = self.N
733 origin_array = event.quasi_run(origin_array, norm=norm)
734 self.N = self.data["pmv_"].size
736 # Add survival to output if mortality is in the model
737 if "dead" in self.data.keys():
738 grids_out["dead"] = None
740 # Get continuation variable names, making sure they're in the same order
741 # as named by the arrival variables. This should maybe be done in the
742 # simulator when it's initialized.
743 if twist is not None:
744 cont_vars_orig = list(twist.keys())
745 temp_dict = {twist[var]: var for var in cont_vars_orig}
746 cont_vars = []
747 for var in self.arrival:
748 cont_vars.append(temp_dict[var])
749 if "dead" in self.data.keys():
750 cont_vars.append("dead")
751 grid_out_is_continuous = np.concatenate(
752 (grid_out_is_continuous, [False])
753 )
754 else:
755 cont_vars = list(grids_out.keys()) # all outcomes are arrival vars
756 D = len(cont_vars)
758 # Now project the final results onto the output or result grids
759 N = self.N
760 J = state_meshes[0].size
761 matrices_out = {}
762 cont_idx = {}
763 cont_alpha = {}
764 cont_M = {}
765 cont_discrete = {}
766 k = 0
767 for var in grids_out.keys():
768 if var not in self.data.keys():
769 raise ValueError(
770 "Variable " + var + " does not exist but a grid was specified!"
771 )
772 grid = grids_out[var]
773 vals = self.data[var]
774 pmv = self.data["pmv_"]
775 M = grid.size if grid is not None else 0
777 # Semi-hacky fix to deal with omitted arrival variables
778 if (M == 1) and (vals.dtype is np.dtype(np.float64)):
779 grid = grid.astype(float)
780 grids_out[var] = grid
781 grid_out_is_continuous[k] = True
783 if grid_out_is_continuous[k]:
784 # Split the final values among discrete gridpoints on the interior.
785 # NB: This will only work properly if the grid is equispaced
786 if M > 1:
787 Q = grid_orders[var]
788 if var in cont_vars:
789 trans_matrix, cont_idx[var], cont_alpha[var] = (
790 aggregate_blobs_onto_polynomial_grid_alt(
791 vals, pmv, origin_array, grid, J, Q
792 )
793 )
794 cont_M[var] = M
795 cont_discrete[var] = False
796 else:
797 trans_matrix = aggregate_blobs_onto_polynomial_grid(
798 vals, pmv, origin_array, grid, J, Q
799 )
800 else: # Skip if the grid is a dummy with only one value.
801 trans_matrix = np.ones((J, M))
802 if var in cont_vars:
803 cont_idx[var] = np.zeros(N, dtype=int)
804 cont_alpha[var] = np.zeros(N)
805 cont_M[var] = M
806 cont_discrete[var] = False
808 else: # Grid is discrete, can use simpler method
809 if grid is None:
810 M = np.max(vals.astype(int))
811 if var == "dead":
812 M = 2
813 grid = np.arange(M, dtype=int)
814 grids_out[var] = grid
815 M = grid.size
816 vals = vals.astype(int)
817 trans_matrix = aggregate_blobs_onto_discrete_grid(
818 vals, pmv, origin_array, M, J
819 )
820 if var in cont_vars:
821 cont_idx[var] = vals
822 cont_alpha[var] = np.zeros(N)
823 cont_M[var] = M
824 cont_discrete[var] = True
826 # Store the transition matrix for this variable
827 matrices_out[var] = trans_matrix
828 k += 1
830 # Construct an overall transition matrix from arrival to continuation variables.
831 # If this is the initializer block, the "arrival" variable is just the initial
832 # dummy state, and the "continuation" variables are actually the arrival variables
833 # for ordinary blocks/periods.
835 # Count the number of non-trivial dimensions. A continuation dimension
836 # is non-trivial if it is both continuous and has more than one grid node.
837 C = 0
838 shape = [N_orig]
839 trivial = []
840 for var in cont_vars:
841 shape.append(cont_M[var])
842 if (not cont_discrete[var]) and (cont_M[var] > 1):
843 C += 1
844 trivial.append(False)
845 else:
846 trivial.append(True)
847 trivial = np.array(trivial)
849 # Make a binary array of offsets from the base index
850 bin_array_base = np.array(list(product([0, 1], repeat=C)))
851 bin_array = np.empty((2**C, D), dtype=int)
852 some_zeros = np.zeros(2**C, dtype=int)
853 c = 0
854 for d in range(D):
855 bin_array[:, d] = some_zeros if trivial[d] else bin_array_base[:, c]
856 c += not trivial[d]
858 # Make a vector of dimensional offsets from the base index
859 dim_offsets = np.ones(D, dtype=int)
860 for d in range(D - 1):
861 dim_offsets[d] = np.prod(shape[(d + 2) :])
862 dim_offsets_X = np.tile(dim_offsets, (2**C, 1))
863 offsets = np.sum(bin_array * dim_offsets_X, axis=1)
865 # Make combined arrays of indices and alphas
866 index_array = np.empty((N, D), dtype=int)
867 alpha_array = np.empty((N, D, 2))
868 for d in range(D):
869 var = cont_vars[d]
870 index_array[:, d] = cont_idx[var]
871 alpha_array[:, d, 0] = 1.0 - cont_alpha[var]
872 alpha_array[:, d, 1] = cont_alpha[var]
873 idx_array = np.dot(index_array, dim_offsets)
875 # Make the master transition array
876 blank = np.zeros(np.array((N_orig, np.prod(shape[1:]))))
877 master_trans_array_X = calc_overall_trans_probs(
878 blank, idx_array, alpha_array, bin_array, offsets, pmv, origin_array
879 )
881 # Condition on survival if relevant
882 if "dead" in self.data.keys():
883 master_trans_array_X = np.reshape(master_trans_array_X, (N_orig, N_orig, 2))
884 survival_probs = np.reshape(matrices_out["dead"][:, 0], [N_orig, 1])
885 master_trans_array_X = master_trans_array_X[..., 0] / survival_probs
887 # Reshape the transition matrix depending on what kind of block this is
888 if arrival_N == 0:
889 # If this is the initializer block, the "transition" matrix is really
890 # just the initial distribution of states at model birth; flatten it.
891 master_init_array = master_trans_array_X.flatten()
892 else:
893 # In an ordinary period, reshape the transition array so it's square.
894 master_trans_array = np.reshape(master_trans_array_X, (N_orig, N_orig))
896 # Store the results as attributes of self
897 grids = {}
898 grids.update(grids_in)
899 grids.update(grids_out)
900 self.grids = grids
901 self.matrices = matrices_out
902 self.mesh = mesh_tuples
903 if twist is not None:
904 self.trans_array = master_trans_array
905 if arrival_N == 0:
906 self.init_dstn = master_init_array
909@dataclass(kw_only=True)
910class AgentSimulator:
911 """
912 A class for representing an entire simulator structure for an AgentType.
913 It includes a sequence of SimBlocks representing periods of the model, which
914 could be built from the information on an AgentType instance.
916 Parameters
917 ----------
918 name : str
919 Short name of this model.s
920 description : str
921 Textual description of what happens in this simulated block.
922 statement : str
923 Verbatim model statement that was used to create this simulator.
924 comments : dict
925 Dictionary of comments or descriptions for various model objects.
926 parameters : list[str]
927 List of parameter names used in the model.
928 distributions : list[str]
929 List of distribution names used in the model.
930 functions : list[str]
931 List of function names used in the model.
932 common: list[str]
933 Names of variables that are common across idiosyncratic agents.
934 types: dict
935 Dictionary of data types for all variables in the model.
936 N_agents: int
937 Number of idiosyncratic agents in this simulation.
938 T_total: int
939 Total number of periods in these agents' model.
940 T_sim: int
941 Maximum number of periods that will be simulated, determining the size
942 of the history arrays.
943 T_age: int
944 Period after which to automatically terminate an agent if they would
945 survive past this period.
946 stop_dead : bool
947 Whether simulated agents who draw dead=True should actually cease acting.
948 Default is True. Setting to False allows "cohort-style" simulation that
949 will generate many agents that survive to old ages. In most cases, T_sim
950 should not exceed T_age, unless the user really does want multiple succ-
951 essive cohorts to be born and fully simulated.
952 replace_dead : bool
953 Whether simulated agents who are marked as dead should be replaced with
954 newborns (default True) or simply cease acting without replacement (False).
955 The latter option is useful for models with state-dependent mortality,
956 to allow "cohort-style" simulation with the correct distribution of states
957 for survivors at each age. Setting to False has no effect if stop_dead is True.
958 periods: list[SimBlock]
959 Ordered list of simulation blocks, each representing a period.
960 twist : dict
961 Dictionary that maps period t-1 variables to period t variables, as a
962 relabeling "between" periods.
963 initializer : SimBlock
964 A special simulated block that should have *no* arrival variables, because
965 it represents the initialization of "newborn" agents.
966 data : dict
967 Dictionary that holds *current* values of model variables.
968 track_vars : list[str]
969 List of names of variables whose history should be tracked in the simulation.
970 history : dict
971 Dictionary that holds the histories of tracked variables.
972 """
974 name: str = field(default="")
975 description: str = field(default="")
976 statement: str = field(default="", repr=False)
977 comments: dict = field(default_factory=dict, repr=False)
978 parameters: list[str] = field(default_factory=list, repr=False)
979 distributions: list[str] = field(default_factory=list, repr=False)
980 functions: list[str] = field(default_factory=list, repr=False)
981 common: list[str] = field(default_factory=list, repr=False)
982 types: dict = field(default_factory=dict, repr=False)
983 N_agents: int = field(default=1)
984 T_total: int = field(default=1, repr=False)
985 T_sim: int = field(default=1)
986 T_age: int = field(default=0, repr=False)
987 stop_dead: bool = field(default=True)
988 replace_dead: bool = field(default=True)
989 periods: list[SimBlock] = field(default_factory=list, repr=False)
990 twist: dict = field(default_factory=dict, repr=False)
991 data: dict = field(default_factory=dict, repr=False)
992 initializer: field(default_factory=SimBlock, repr=False)
993 track_vars: list[str] = field(default_factory=list, repr=False)
994 history: dict = field(default_factory=dict, repr=False)
996 def simulate(self, T=None):
997 """
998 Simulates the model for T periods, including replacing dead agents as
999 warranted and storing tracked variables in the history. If T is not
1000 specified, the agents are simulated for the entire T_sim periods.
1001 This is the primary user-facing simulation method.
1002 """
1003 if T is None:
1004 T = self.T_sim - self.t_sim # All remaining simulated periods
1005 if (T + self.t_sim) > self.T_sim:
1006 raise ValueError("Can't simulate more than T_sim periods!")
1008 # Execute the simulation loop for T periods
1009 for t in range(T):
1010 # Do the ordinary work for simulating a period
1011 self.sim_one_period()
1013 # Mark agents who have reached maximum allowable age
1014 if "dead" in self.data.keys() and self.T_age > 0:
1015 too_old = self.t_age == self.T_age
1016 self.data["dead"][too_old] = True
1018 # Record tracked variables and advance age
1019 self.store_tracked_vars()
1020 self.advance_age()
1022 # Handle death and replacement depending on simulation style
1023 if "dead" in self.data.keys() and self.stop_dead:
1024 self.mark_dead_agents()
1025 self.t_sim += 1
1027 def reset(self):
1028 """
1029 Completely reset this simulator back to its original state so that it
1030 can be run from scratch. This should allow it to generate the same results
1031 every single time the simulator is run (if nothing changes).
1032 """
1033 N = self.N_agents
1034 T = self.T_sim
1035 self.t_sim = 0 # Time index for the simulation
1037 # Reset the variable data and history arrays
1038 self.clear_data()
1039 self.history = {}
1040 for var in self.track_vars:
1041 self.history[var] = np.empty((T, N), dtype=self.types[var])
1043 # Reset all of the blocks / periods
1044 self.initializer.reset()
1045 for t in range(len(self.periods)):
1046 self.periods[t].reset()
1048 # Specify all agents as "newborns" assigned to the initializer block
1049 self.t_seq_bool_array = np.zeros((self.T_total, N), dtype=bool)
1050 self.t_age = -np.ones(N, dtype=int)
1052 def clear_data(self, skip=None):
1053 """
1054 Reset all current data arrays back to blank, other than those designated
1055 to be skipped, if any.
1057 Parameters
1058 ----------
1059 skip : [str] or None
1060 Names of variables *not* to be cleared from data. Default is None.
1062 Returns
1063 -------
1064 None
1065 """
1066 if skip is None:
1067 skip = []
1068 N = self.N_agents
1069 # self.data = {}
1070 for var in self.types.keys():
1071 if var in skip:
1072 continue
1073 this_type = self.types[var]
1074 if this_type is float:
1075 self.data[var] = np.full((N,), np.nan)
1076 elif this_type is bool:
1077 self.data[var] = np.zeros((N,), dtype=bool)
1078 elif this_type is int:
1079 self.data[var] = np.zeros((N,), dtype=np.int32)
1080 elif this_type is complex:
1081 self.data[var] = np.full((N,), np.nan, dtype=complex)
1082 else:
1083 raise ValueError(
1084 "Type "
1085 + str(this_type)
1086 + " of variable "
1087 + var
1088 + " was not recognized!"
1089 )
1091 def mark_dead_agents(self):
1092 """
1093 Looks at the special data field "dead" and marks those agents for replacement.
1094 If no variable called "dead" has been defined, this is skipped.
1095 """
1096 who_died = self.data["dead"]
1097 self.t_seq_bool_array[:, who_died] = False
1098 self.t_age[who_died] = -1
1100 def create_newborns(self):
1101 """
1102 Calls the initializer to generate newborns where needed.
1103 """
1104 # Skip this step if there are no newborns
1105 newborns = self.t_age == -1
1106 if not np.any(newborns):
1107 return
1109 # Generate initial arrival variables
1110 N = np.sum(newborns)
1111 self.initializer.data = {} # by definition
1112 self.initializer.N = N
1113 self.initializer.run()
1115 # Set the initial arrival data for newborns and clear other variables
1116 init_arrival = self.periods[0].arrival
1117 for var in self.types:
1118 self.data[var][newborns] = (
1119 self.initializer.data[var]
1120 if var in init_arrival
1121 else np.empty(N, dtype=self.types[var])
1122 )
1124 # Set newborns' period to 0
1125 self.t_age[newborns] = 0
1126 self.t_seq_bool_array[0, newborns] = True
1128 def store_tracked_vars(self):
1129 """
1130 Record current values of requested variables in the history dictionary.
1131 """
1132 for var in self.track_vars:
1133 self.history[var][self.t_sim, :] = self.data[var]
1135 def advance_age(self):
1136 """
1137 Increments age for all agents, altering t_age and t_age_bool. Agents in
1138 the last period of the sequence will be assigned to the initial period.
1139 In a lifecycle model, those agents should be marked as dead and replaced
1140 in short order.
1141 """
1142 alive = self.t_age >= 0 # Don't age the dead
1143 self.t_age[alive] += 1
1144 X = self.t_seq_bool_array # For shorter typing on next line
1145 self.t_seq_bool_array[:, alive] = np.concatenate(
1146 (X[-1:, alive], X[:-1, alive]), axis=0
1147 )
1149 def sim_one_period(self):
1150 """
1151 Simulates one period of the model by advancing all agents one period.
1152 This includes creating newborns, but it does NOT include eliminating
1153 dead agents nor storing tracked results in the history. This method
1154 should usually not be called by a user, instead using simulate(1) if
1155 you want to run the model for exactly one period.
1156 """
1157 # Use the "twist" information to advance last period's end-of-period
1158 # information/values to be the arrival variables for this period. Then, for
1159 # each variable other than those brought in with the twist, wipe it clean.
1160 keepers = []
1161 for var_tm1 in self.twist:
1162 var_t = self.twist[var_tm1]
1163 keepers.append(var_t)
1164 self.data[var_t] = self.data[var_tm1].copy()
1165 self.clear_data(skip=keepers)
1167 # Create newborns first so the arrival vars exist. This should be done in
1168 # the first simulated period (t_sim=0) or if decedents should be replaced.
1169 if self.replace_dead or self.t_sim == 0:
1170 self.create_newborns()
1172 # Loop through ages and run the model on the appropriately aged agents
1173 for t in range(self.T_total):
1174 these = self.t_seq_bool_array[t, :]
1175 if not np.any(these):
1176 continue # Skip any "empty ages"
1177 this_period = self.periods[t]
1179 data_temp = {var: self.data[var][these] for var in this_period.arrival}
1180 this_period.data = data_temp
1181 this_period.N = np.sum(these)
1182 this_period.run()
1184 # Extract all of the variables from this period and write it to data
1185 for var in this_period.data.keys():
1186 self.data[var][these] = this_period.data[var]
1188 # Put time information into the data dictionary
1189 self.data["t_age"] = self.t_age.copy()
1190 self.data["t_seq"] = np.argmax(self.t_seq_bool_array, axis=0).astype(int)
1192 def make_transition_matrices(
1193 self, grid_specs, norm=None, fake_news_timing=False, for_t=None
1194 ):
1195 """
1196 Build Markov-style transition matrices for each period of the model, as
1197 well as the initial distribution of arrival variables for newborns.
1198 Stores results to the attributes of self as follows:
1200 - trans_arrays : List of Markov matrices for transitioning from the arrival
1201 state space in period t to the arrival state space in t+1.
1202 This transition includes death (and replacement).
1203 - newborn_dstn : Stochastic vector as a NumPy array, representing the distribution
1204 of arrival states for "newborns" who were just initialized.
1205 - state_grids : Nested list of tuples representing the arrival state space for
1206 each period. Each element corresponds to the discretized arrival
1207 state space point with the same index in trans_arrays (and
1208 newborn_dstn). Arrival states are ordered within a tuple in the
1209 same order as the model file. Linked from period[t].mesh.
1210 - outcome_arrays : List of dictionaries of arrays that cast from the arrival
1211 state space to the grid of outcome variables, for each period.
1212 Doing np.dot(state_dstn, outcome_arrays[t][var]) will yield
1213 the discretized distribution of that outcome variable. Linked
1214 from periods[t].matrices.
1215 - outcome_grids : List of dictionaries of discretized outcomes in each period.
1216 Keys are names of outcome variables, and entries are vectors
1217 of discretized values that the outcome variable can take on.
1218 Doing np.dot(np.dot(state_dstn, outcome_arrays[var]), outcome_grids[var])
1219 yields the *average* of that outcome in the population. Linked
1220 from periods[t].grids.
1222 Parameters
1223 ----------
1224 grid_specs : dict
1225 Dictionary of dictionaries with specifications for discretized grids
1226 of all variables of interest. If any arrival variables are omitted,
1227 they will be given a default trivial grid with one node at 0. This
1228 should only be done if that arrival variable is closely tied to the
1229 Harmenberg normalizing variable; see below. A grid specification must
1230 include a number of gridpoints N, and should also include a min and
1231 max if the variable is continuous. If the variable is discrete, the
1232 grid values are assumed to be 0,..,N.
1233 norm : str or None
1234 Name of the variable for which Harmenberg normalization should be
1235 applied, if any. This should be a variable that is directly drawn
1236 from a distribution, not a "downstream" variable.
1237 fake_news_timing : bool
1238 Indicator for whether this call is part of the "fake news" algorithm
1239 for constructing sequence space Jacobians (SSJs). This should only
1240 ever be set to True in that situation, which affects how mortality
1241 is handled between periods. In short, the simulator usually assumes
1242 that "newborns" start with t_seq=0, but during the fake news algorithm,
1243 that is not the case.
1244 for_t : list or None
1245 Optional list of time indices for which the matrices should be built.
1246 When not specified, all periods are constructed. The most common use
1247 for this arg is during the "fake news" algorithm for lifecycle models.
1249 Returns
1250 -------
1251 None
1252 """
1253 # Sort grid specifications into those needed by the initializer vs those
1254 # used by other blocks (ordinary periods)
1255 arrival = self.periods[0].arrival
1256 arrival_N = len(arrival)
1257 check_bool = np.zeros(arrival_N, dtype=bool)
1258 grid_specs_init_orig = {}
1259 grid_specs_other = {}
1260 for name in grid_specs.keys():
1261 if name in arrival:
1262 idx = arrival.index(name)
1263 check_bool[idx] = True
1264 grid_specs_init_orig[name] = copy(grid_specs[name])
1265 grid_specs_other[name] = copy(grid_specs[name])
1267 # Build the dictionary of arrival variables, making sure it's in the
1268 # same order as named self.arrival. For any arrival grids that are
1269 # not specified, make a dummy specification.
1270 grid_specs_init = {}
1271 for n in range(arrival_N):
1272 name = arrival[n]
1273 if check_bool[n]:
1274 grid_specs_init[name] = grid_specs_init_orig[name]
1275 continue
1276 dummy_grid_spec = {"N": 1}
1277 grid_specs_init[name] = dummy_grid_spec
1278 grid_specs_other[name] = dummy_grid_spec
1280 # Make the initial state distribution for newborns
1281 self.initializer.make_transition_matrices(grid_specs_init)
1282 self.newborn_dstn = self.initializer.init_dstn
1283 K = self.newborn_dstn.size
1285 # Make the period-by-period transition matrices
1286 these_t = range(len(self.periods)) if for_t is None else for_t
1287 for t in these_t:
1288 block = self.periods[t]
1289 block.make_transition_matrices(
1290 grid_specs_other, twist=self.twist, norm=norm
1291 )
1292 block.reset()
1294 # Extract the master transition matrices into a single list
1295 p2p_trans_arrays = [block.trans_array for block in self.periods]
1297 # Apply agent replacement to the last period of the model, representing
1298 # newborns filling in for decedents. This will usually only do anything
1299 # at all in "one period infinite horizon" models. If this is part of the
1300 # fake news algorithm for constructing SSJs, then replace decedents with
1301 # newborns in *all* periods, because model timing is funny in this case.
1302 if fake_news_timing:
1303 T_set = np.arange(len(self.periods)).tolist()
1304 else:
1305 T_set = [-1]
1306 newborn_dstn = np.reshape(self.newborn_dstn, (1, K))
1307 for t in T_set:
1308 if "dead" not in self.periods[t].matrices.keys():
1309 continue
1310 death_prbs = self.periods[t].matrices["dead"][:, 1]
1311 p2p_trans_arrays[t] *= np.tile(np.reshape(1 - death_prbs, (K, 1)), (1, K))
1312 p2p_trans_arrays[t] += np.reshape(death_prbs, (K, 1)) * newborn_dstn
1314 # Store the transition arrays as attributes of self
1315 self.trans_arrays = p2p_trans_arrays
1317 # Build and store lists of state meshes, outcome arrays, and outcome grids
1318 self.state_grids = [self.periods[t].mesh for t in range(len(self.periods))]
1319 self.outcome_grids = [self.periods[t].grids for t in range(len(self.periods))]
1320 self.outcome_arrays = [
1321 self.periods[t].matrices for t in range(len(self.periods))
1322 ]
1324 def find_steady_state(self):
1325 """
1326 Calculates the steady state distribution of arrival states for a "one period
1327 infinite horizon" model, storing the result to the attribute steady_state_dstn.
1328 Should only be run after make_transition_matrices(), and only if T_total = 1
1329 and the model is infinite horizon.
1330 """
1331 if self.T_total != 1:
1332 raise ValueError(
1333 "This method currently only works with one period infinite horizon problems."
1334 )
1336 # Find the eigenvector associated with the largest eigenvalue of the
1337 # infinite horizon transition matrix. The largest eigenvalue *should*
1338 # be 1 for any Markov matrix, but double check to be sure.
1339 trans_T = csr_matrix(self.trans_arrays[0].transpose())
1340 v, V = eigs(trans_T, k=1)
1341 if not np.isclose(v[0], 1.0):
1342 raise ValueError(
1343 "The largest eigenvalue of the transition matrix isn't close to 1!"
1344 )
1346 # Normalize that eigenvector and make sure its real, then store it
1347 D = V[:, 0]
1348 SS_dstn = (D / np.sum(D)).real
1349 self.steady_state_dstn = SS_dstn
1351 def get_long_run_average(self, var):
1352 """
1353 Calculate and return the long run / steady state population average of
1354 one named variable. Should only be run after find_steady_state().
1356 Parameters
1357 ----------
1358 var : str
1359 Name of the variable for which to calculate the long run average.
1361 Returns
1362 -------
1363 var_mean : float
1364 Long run / steady state population average of the variable.
1365 """
1366 if not hasattr(self, "steady_state_dstn"):
1367 raise ValueError("This method can only be run after find_steady_state()!")
1369 dstn = self.steady_state_dstn
1370 array = self.outcome_arrays[0][var]
1371 grid = self.outcome_grids[0][var]
1373 var_dstn = np.dot(dstn, array)
1374 var_mean = np.dot(var_dstn, grid)
1375 return var_mean
1377 def simulate_cohort_by_grids(
1378 self,
1379 outcomes,
1380 T_max=None,
1381 calc_dstn=False,
1382 calc_avg=True,
1383 from_dstn=None,
1384 ):
1385 """
1386 Generate a simulated "cohort style" history for this type of agents using
1387 discretized grid methods. Can only be run after running make_transition_matrices().
1388 Starting from the distribution of states at birth, the population is moved
1389 forward in time via the transition matrices, and the distribution and/or
1390 average of specified outcomes are stored in the dictionary attributes
1391 history_dstn and history_avg respectively.
1393 Parameters
1394 ----------
1395 outcomes : str or [str]
1396 Names of one or more outcome variables to be tracked during the grid
1397 simulation. Each named variable should have an outcome grid specified
1398 when make_transition_matrices() was called, whether explicitly or
1399 implicitly. The existence of these grids is checked as a first step.
1400 T_max : int or None
1401 If specified, the number of periods of the model to actually generate
1402 output for. If not specified, all periods are run.
1403 calc_dstn : bool
1404 Whether outcome distributions should be stored in the dictionary
1405 attribute history_dstn. The default is False.
1406 calc_avg : bool
1407 Whether outcome averages should be stored in the dictionary attribute
1408 history_avg. The default is True.
1409 from_dstn : np.array or None
1410 Optional initial distribution of arrival states. If not specified, the
1411 newborn distribution in the initializer is assumed to be used.
1413 Returns
1414 -------
1415 None
1416 """
1417 # First, verify that newborn and transition matrices exist for all periods
1418 if not hasattr(self, "newborn_dstn"):
1419 raise ValueError(
1420 "The newborn state distribution does not exist; make_transition_matrices() must be run before grid simulations!"
1421 )
1422 if T_max is None:
1423 T_max = self.T_total
1424 T_max = np.minimum(T_max, self.T_total)
1425 if not hasattr(self, "trans_arrays"):
1426 raise ValueError(
1427 "The transition arrays do not exist; make_transition_matrices() must be run before grid simulations!"
1428 )
1429 if len(self.trans_arrays) < T_max:
1430 raise ValueError(
1431 "There are somehow fewer elements of trans_array than there should be!"
1432 )
1433 if not (calc_dstn or calc_avg):
1434 return # No work actually requested, we're done here
1436 # Initialize generated output as requested
1437 if isinstance(outcomes, str):
1438 outcomes = [outcomes]
1439 if calc_dstn:
1440 history_dstn = {}
1441 for name in outcomes: # List will be concatenated to array at end
1442 history_dstn[name] = [] # if all distributions are same size
1443 if calc_avg:
1444 history_avg = {}
1445 for name in outcomes:
1446 history_avg[name] = np.empty(T_max)
1448 # Initialize the state distribution
1449 current_dstn = (
1450 self.newborn_dstn.copy() if from_dstn is None else from_dstn.copy()
1451 )
1452 state_dstn_by_age = []
1454 # Loop over requested periods of this agent type's model
1455 for t in range(T_max):
1456 state_dstn_by_age.append(current_dstn)
1458 # Calculate outcome distributions and averages as requested
1459 for name in outcomes:
1460 this_outcome = self.periods[t].matrices[name].transpose()
1461 this_dstn = np.dot(this_outcome, current_dstn)
1462 if calc_dstn:
1463 history_dstn[name].append(this_dstn)
1464 if calc_avg:
1465 this_grid = self.periods[t].grids[name]
1466 history_avg[name][t] = np.dot(this_dstn, this_grid)
1468 # Advance the distribution to the next period
1469 current_dstn = np.dot(self.trans_arrays[t].transpose(), current_dstn)
1471 # Reshape the distribution histories if possible
1472 if calc_dstn:
1473 for name in outcomes:
1474 dstn_sizes = np.array([dstn.size for dstn in history_dstn[name]])
1475 if np.all(dstn_sizes == dstn_sizes[0]):
1476 history_dstn[name] = np.stack(history_dstn[name], axis=1)
1478 # Store results as attributes of self
1479 self.state_dstn_by_age = state_dstn_by_age
1480 if calc_dstn:
1481 self.history_dstn = history_dstn
1482 if calc_avg:
1483 self.history_avg = history_avg
1485 def describe_model(self, display=True):
1486 """
1487 Convenience method that prints model information to screen.
1488 """
1489 # Make a twist statement
1490 twist_statement = ""
1491 for var_tm1 in self.twist.keys():
1492 var_t = self.twist[var_tm1]
1493 new_line = var_tm1 + "[t-1] <---> " + var_t + "[t]\n"
1494 twist_statement += new_line
1496 # Assemble the overall model statement
1497 output = ""
1498 output += "----------------------------------\n"
1499 output += "%%%%% INITIALIZATION AT BIRTH %%%%\n"
1500 output += "----------------------------------\n"
1501 output += self.initializer.statement
1502 output += "----------------------------------\n"
1503 output += "%%%% DYNAMICS WITHIN PERIOD t %%%%\n"
1504 output += "----------------------------------\n"
1505 output += self.statement
1506 output += "----------------------------------\n"
1507 output += "%%%%%%% RELABELING / TWIST %%%%%%%\n"
1508 output += "----------------------------------\n"
1509 output += twist_statement
1510 output += "-----------------------------------"
1512 # Return or print the output
1513 if display:
1514 print(output)
1515 return
1516 else:
1517 return output
1519 def describe_symbols(self, display=True):
1520 """
1521 Convenience method that prints symbol information to screen.
1522 """
1523 # Get names and types
1524 symbols_lines = []
1525 comments = []
1526 for key in self.comments.keys():
1527 comments.append(self.comments[key])
1529 # Get type of object
1530 if key in self.types.keys():
1531 this_type = str(self.types[key].__name__)
1532 elif key in self.distributions:
1533 this_type = "dstn"
1534 elif key in self.parameters:
1535 this_type = "param"
1536 elif key in self.functions:
1537 this_type = "func"
1539 # Add tags
1540 if key in self.common:
1541 this_type += ", common"
1542 # if key in self.solution:
1543 # this_type += ', solution'
1544 this_line = key + " (" + this_type + ")"
1545 symbols_lines.append(this_line)
1547 # Add comments, aligned
1548 symbols_text = ""
1549 longest = np.max([len(this) for this in symbols_lines])
1550 for j in range(len(symbols_lines)):
1551 line = symbols_lines[j]
1552 comment = comments[j]
1553 L = len(line)
1554 pad = (longest + 1) - L
1555 symbols_text += line + pad * " " + ": " + comment + "\n"
1557 # Return or print the output
1558 output = symbols_text
1559 if display:
1560 print(output)
1561 return
1562 else:
1563 return output
1565 def describe(self, symbols=True, model=True, display=True):
1566 """
1567 Convenience method for showing all information about the model.
1568 """
1569 # Asssemble the requested output
1570 output = self.name + ": " + self.description + "\n"
1571 if symbols or model:
1572 output += "\n"
1573 if symbols:
1574 output += "----------------------------------\n"
1575 output += "%%%%%%%%%%%%% SYMBOLS %%%%%%%%%%%%\n"
1576 output += "----------------------------------\n"
1577 output += self.describe_symbols(display=False)
1578 if model:
1579 output += self.describe_model(display=False)
1580 if symbols and not model:
1581 output += "----------------------------------"
1583 # Return or print the output
1584 if display:
1585 print(output)
1586 return
1587 else:
1588 return output
1591def make_simulator_from_agent(agent, stop_dead=True, replace_dead=True, common=None):
1592 """
1593 Build an AgentSimulator instance based on an AgentType instance. The AgentType
1594 should have its model attribute defined so that it can be parsed and translated
1595 into the simulator structure. The names of objects in the model statement
1596 should correspond to attributes of the AgentType.
1598 Parameters
1599 ----------
1600 agent : AgentType
1601 Agents for whom a new simulator is to be constructed.
1602 stop_dead : bool
1603 Whether simulated agents who draw dead=True should actually cease acting.
1604 Default is True. Setting to False allows "cohort-style" simulation that
1605 will generate many agents that survive to old ages. In most cases, T_sim
1606 should not exceed T_age, unless the user really does want multiple succ-
1607 essive cohorts to be born and fully simulated.
1608 replace_dead : bool
1609 Whether simulated agents who are marked as dead should be replaced with
1610 newborns (default True) or simply cease acting without replacement (False).
1611 The latter option is useful for models with state-dependent mortality,
1612 to allow "cohort-style" simulation with the correct distribution of states
1613 for survivors at each age. Setting False has no effect if stop_dead is True.
1614 common : [str] or None
1615 List of random variables that should be treated as commonly shared across
1616 all agents, rather than idiosyncratically drawn. If this is provided, it
1617 will override the model defaults.
1619 Returns
1620 -------
1621 new_simulator : AgentSimulator
1622 A simulator structure based on the agents.
1623 """
1624 # Read the model statement into a dictionary, and get names of attributes
1625 if hasattr(agent, "model_statement"): # look for a custom model statement
1626 model_statement = copy(agent.model_statement)
1627 else: # otherwise use the default model file
1628 with importlib.resources.open_text("HARK.models", agent.model_file) as f:
1629 model_statement = f.read()
1630 f.close()
1631 model = yaml.safe_load(model_statement)
1632 time_vary = agent.time_vary
1633 time_inv = agent.time_inv
1634 cycles = agent.cycles
1635 T_age = agent.T_age
1636 comments = {}
1637 RNG = agent.RNG # this is only for generating seeds for MarkovEvents
1639 # Extract basic fields from the model
1640 try:
1641 model_name = model["name"]
1642 except:
1643 model_name = "DEFAULT_NAME"
1644 try:
1645 description = model["description"]
1646 except:
1647 description = "(no description provided)"
1648 try:
1649 variables = model["symbols"]["variables"]
1650 except:
1651 variables = []
1652 try:
1653 twist = model["twist"]
1654 except:
1655 twist = {}
1656 if common is None:
1657 try:
1658 common = model["symbols"]["common"]
1659 except:
1660 common = []
1662 # Extract arrival variable names that were explicitly listed
1663 try:
1664 arrival = model["symbols"]["arrival"]
1665 except:
1666 arrival = []
1668 # Make a dictionary of declared data types and add comments
1669 types = {}
1670 for var_line in variables: # Loop through declared variables
1671 var_name, var_type, flags, desc = parse_declaration_for_parts(var_line)
1672 if var_type is not None:
1673 try:
1674 var_type = eval(var_type)
1675 except:
1676 raise ValueError(
1677 "Couldn't understand type "
1678 + var_type
1679 + " for declared variable "
1680 + var_name
1681 + "!"
1682 )
1683 else:
1684 var_type = float
1685 types[var_name] = var_type
1686 comments[var_name] = desc
1687 if ("arrival" in flags) and (var_name not in arrival):
1688 arrival.append(var_name)
1689 if ("common" in flags) and (var_name not in common):
1690 common.append(var_name)
1692 # Make a blank "template" period with structure but no data
1693 template_period, information, offset, solution, block_comments = (
1694 make_template_block(model, arrival, common)
1695 )
1696 comments.update(block_comments)
1698 # Make the agent initializer, without parameter values (etc)
1699 initializer, init_info = make_initializer(model, arrival, common)
1701 # Extract basic fields from the template period and model
1702 statement = template_period.statement
1703 content = template_period.content
1705 # Get the names of parameters, functions, and distributions
1706 parameters = []
1707 functions = []
1708 distributions = []
1709 for key in information.keys():
1710 val = information[key]
1711 if val is None:
1712 parameters.append(key)
1713 elif type(val) is NullFunc:
1714 functions.append(key)
1715 elif type(val) is Distribution:
1716 distributions.append(key)
1718 # Loop through variables that appear in the model block but were undeclared
1719 for var in information.keys():
1720 if var in types.keys():
1721 continue
1722 this = information[var]
1723 if (this is None) or (type(this) is Distribution) or (type(this) is NullFunc):
1724 continue
1725 types[var] = float
1726 comments[var] = ""
1727 if "dead" in types.keys():
1728 types["dead"] = bool
1729 comments["dead"] = "whether agent died this period"
1730 types["t_seq"] = int
1731 types["t_age"] = int
1732 comments["t_seq"] = "which period of the sequence the agent is on"
1733 comments["t_age"] = "how many periods the agent has already lived for"
1735 # Make a dictionary for the initializer and distribute information
1736 init_dict = {}
1737 for name in init_info.keys():
1738 try:
1739 init_dict[name] = getattr(agent, name)
1740 except:
1741 raise ValueError(
1742 "Couldn't get a value for initializer object " + name + "!"
1743 )
1744 initializer.content = init_dict
1745 initializer.distribute_content()
1747 # Make a dictionary of time-invariant parameters
1748 time_inv_dict = {}
1749 for name in content:
1750 if name in time_inv:
1751 try:
1752 time_inv_dict[name] = getattr(agent, name)
1753 except:
1754 raise ValueError(
1755 "Couldn't get a value for time-invariant object " + name + "!"
1756 )
1758 # Create a list of periods, pulling appropriate data from the agent for each one
1759 T_seq = len(agent.solution) # Number of periods in the solution sequence
1760 periods = []
1761 T_cycle = agent.T_cycle
1762 t_cycle = 0
1763 for t in range(T_seq):
1764 # Make a fresh copy of the template period
1765 new_period = deepcopy(template_period)
1767 # Make sure each period's events have unique seeds; this is only for MarkovEvents
1768 for event in new_period.events:
1769 if hasattr(event, "seed"):
1770 event.seed = RNG.integers(0, 2**31 - 1)
1772 # Make the parameter dictionary for this period
1773 new_param_dict = deepcopy(time_inv_dict)
1774 for name in content:
1775 if name in solution:
1776 new_param_dict[name] = getattr(agent.solution[t], name)
1777 elif name in time_vary:
1778 s = (t_cycle - 1) if name in offset else t_cycle
1779 try:
1780 new_param_dict[name] = getattr(agent, name)[s]
1781 except:
1782 raise ValueError(
1783 "Couldn't get a value for time-varying object "
1784 + name
1785 + " at time index "
1786 + str(s)
1787 + "!"
1788 )
1789 elif name in time_inv:
1790 continue
1791 else:
1792 raise ValueError(
1793 "The object called "
1794 + name
1795 + " is not named in time_inv nor time_vary!"
1796 )
1798 # Fill in content for this period, then add it to the list
1799 new_period.content = new_param_dict
1800 new_period.distribute_content()
1801 periods.append(new_period)
1803 # Advance time according to the cycle
1804 t_cycle += 1
1805 if t_cycle == T_cycle:
1806 t_cycle = 0
1808 # Calculate maximum age
1809 if T_age is None:
1810 T_age = 0
1811 if cycles > 0:
1812 T_age_max = T_seq - 1
1813 T_age = np.minimum(T_age_max, T_age)
1814 try:
1815 T_sim = agent.T_sim
1816 except:
1817 T_sim = 0 # very boring default!
1819 # Make and return the new simulator
1820 new_simulator = AgentSimulator(
1821 name=model_name,
1822 description=description,
1823 statement=statement,
1824 comments=comments,
1825 parameters=parameters,
1826 functions=functions,
1827 distributions=distributions,
1828 common=common,
1829 types=types,
1830 N_agents=agent.AgentCount,
1831 T_total=T_seq,
1832 T_sim=T_sim,
1833 T_age=T_age,
1834 stop_dead=stop_dead,
1835 replace_dead=replace_dead,
1836 periods=periods,
1837 twist=twist,
1838 initializer=initializer,
1839 track_vars=agent.track_vars,
1840 )
1841 new_simulator.solution = solution # this is for use by SSJ constructor
1842 return new_simulator
1845def make_template_block(model, arrival=None, common=None):
1846 """
1847 Construct a new SimBlock object as a "template" of the model block. It has
1848 events and reference information, but no values filled in.
1850 Parameters
1851 ----------
1852 model : dict
1853 Dictionary with model block information, probably read in as a yaml.
1854 arrival : [str] or None
1855 List of arrival variables that were flagged or explicitly listed.
1856 common : [str] or None
1857 List of variables that are common or shared across all agents, rather
1858 than idiosyncratically drawn.
1860 Returns
1861 -------
1862 template_block : SimBlock
1863 A "template" of this model block, with no parameters (etc) on it.
1864 info : dict
1865 Dictionary of model objects that were referenced within the block. Keys
1866 are object names and entries reveal what kind of object they are:
1867 - None --> parameter
1868 - 0 --> outcome/data variable (including arrival variables)
1869 - NullFunc --> function
1870 - Distribution --> distribution
1871 offset : [str]
1872 List of object names that are offset in time by one period.
1873 solution : [str]
1874 List of object names that are part of the model solution.
1875 comments : dict
1876 Dictionary of comments included with declared functions, distributions,
1877 and parameters.
1878 """
1879 if arrival is None:
1880 arrival = []
1881 if common is None:
1882 common = []
1884 # Extract explicitly listed metadata
1885 try:
1886 name = model["name"]
1887 except:
1888 name = "DEFAULT_NAME"
1889 try:
1890 offset = model["symbols"]["offset"]
1891 except:
1892 offset = []
1893 try:
1894 solution = model["symbols"]["solution"]
1895 except:
1896 solution = []
1898 # Extract parameters, functions, and distributions
1899 comments = {}
1900 parameters = {}
1901 if "parameters" in model["symbols"].keys():
1902 param_lines = model["symbols"]["parameters"]
1903 for line in param_lines:
1904 param_name, datatype, flags, desc = parse_declaration_for_parts(line)
1905 parameters[param_name] = None
1906 comments[param_name] = desc
1907 # TODO: what to do with parameter types?
1908 if ("offset" in flags) and (param_name not in offset):
1909 offset.append(param_name)
1910 if ("solution" in flags) and (param_name not in solution):
1911 solution.append(param_name)
1913 functions = {}
1914 if "functions" in model["symbols"].keys():
1915 func_lines = model["symbols"]["functions"]
1916 for line in func_lines:
1917 func_name, datatype, flags, desc = parse_declaration_for_parts(line)
1918 if (datatype is not None) and (datatype != "func"):
1919 raise ValueError(
1920 func_name
1921 + " was declared as a function, but given a different datatype!"
1922 )
1923 functions[func_name] = NullFunc()
1924 comments[func_name] = desc
1925 if ("offset" in flags) and (func_name not in offset):
1926 offset.append(func_name)
1927 if ("solution" in flags) and (func_name not in solution):
1928 solution.append(func_name)
1930 distributions = {}
1931 if "distributions" in model["symbols"].keys():
1932 dstn_lines = model["symbols"]["distributions"]
1933 for line in dstn_lines:
1934 dstn_name, datatype, flags, desc = parse_declaration_for_parts(line)
1935 if (datatype is not None) and (datatype != "dstn"):
1936 raise ValueError(
1937 dstn_name
1938 + " was declared as a distribution, but given a different datatype!"
1939 )
1940 distributions[dstn_name] = Distribution()
1941 comments[dstn_name] = desc
1942 if ("offset" in flags) and (dstn_name not in offset):
1943 offset.append(dstn_name)
1944 if ("solution" in flags) and (dstn_name not in solution):
1945 solution.append(dstn_name)
1947 # Combine those dictionaries into a single "information" dictionary, which
1948 # represents objects available *at that point* in the dynamic block
1949 content = parameters.copy()
1950 content.update(functions)
1951 content.update(distributions)
1952 info = deepcopy(content)
1953 for var in arrival:
1954 info[var] = 0 # Mark as a state variable
1956 # Parse the model dynamics
1957 dynamics = format_block_statement(model["dynamics"])
1959 # Make the list of ordered events
1960 events = []
1961 names_used_in_dynamics = []
1962 for line in dynamics:
1963 # Make the new event and add it to the list
1964 new_event, names_used = make_new_event(line, info)
1965 events.append(new_event)
1966 names_used_in_dynamics += names_used
1968 # Add newly assigned variables to the information set
1969 for var in new_event.assigns:
1970 if var in info.keys():
1971 raise ValueError(var + " is assigned, but already exists!")
1972 info[var] = 0
1974 # If any assigned variables are common, mark the event as common
1975 for var in new_event.assigns:
1976 if var in common:
1977 new_event.common = True
1978 break # No need to check further
1980 # Remove content that is never referenced within the dynamics
1981 delete_these = []
1982 for name in content.keys():
1983 if name not in names_used_in_dynamics:
1984 delete_these.append(name)
1985 for name in delete_these:
1986 del content[name]
1988 # Make a single string model statement
1989 statement = ""
1990 longest = np.max([len(event.statement) for event in events])
1991 for event in events:
1992 this_statement = event.statement
1993 L = len(this_statement)
1994 pad = (longest + 1) - L
1995 statement += this_statement + pad * " " + ": " + event.description + "\n"
1997 # Make a description for the template block
1998 if name is None:
1999 description = "template block for unnamed block"
2000 else:
2001 description = "template block for " + name
2003 # Make and return the new SimBlock
2004 template_block = SimBlock(
2005 description=description,
2006 arrival=arrival,
2007 content=content,
2008 statement=statement,
2009 events=events,
2010 )
2011 return template_block, info, offset, solution, comments
2014def make_initializer(model, arrival=None, common=None):
2015 """
2016 Construct a new SimBlock object to be the agent initializer, based on the
2017 model dictionary. It has structure and events, but no parameters (etc).
2019 Parameters
2020 ----------
2021 model : dict
2022 Dictionary with model initializer information, probably read in as a yaml.
2023 arrival : [str]
2024 List of arrival variables that were flagged or explicitly listed.
2026 Returns
2027 -------
2028 initializer : SimBlock
2029 A "template" of this model block, with no parameters (etc) on it.
2030 init_requires : dict
2031 Dictionary of model objects that are needed by the initializer to run.
2032 Keys are object names and entries reveal what kind of object they are:
2033 - None --> parameter
2034 - 0 --> outcome variable (these should include all arrival variables)
2035 - NullFunc --> function
2036 - Distribution --> distribution
2037 """
2038 if arrival is None:
2039 arrival = []
2040 if common is None:
2041 common = []
2042 try:
2043 name = model["name"]
2044 except:
2045 name = "DEFAULT_NAME"
2047 # Extract parameters, functions, and distributions
2048 parameters = {}
2049 if "parameters" in model["symbols"].keys():
2050 param_lines = model["symbols"]["parameters"]
2051 for line in param_lines:
2052 param_name, datatype, flags, desc = parse_declaration_for_parts(line)
2053 parameters[param_name] = None
2055 functions = {}
2056 if "functions" in model["symbols"].keys():
2057 func_lines = model["symbols"]["functions"]
2058 for line in func_lines:
2059 func_name, datatype, flags, desc = parse_declaration_for_parts(line)
2060 if (datatype is not None) and (datatype != "func"):
2061 raise ValueError(
2062 func_name
2063 + " was declared as a function, but given a different datatype!"
2064 )
2065 functions[func_name] = NullFunc()
2067 distributions = {}
2068 if "distributions" in model["symbols"].keys():
2069 dstn_lines = model["symbols"]["distributions"]
2070 for line in dstn_lines:
2071 dstn_name, datatype, flags, desc = parse_declaration_for_parts(line)
2072 if (datatype is not None) and (datatype != "dstn"):
2073 raise ValueError(
2074 dstn_name
2075 + " was declared as a distribution, but given a different datatype!"
2076 )
2077 distributions[dstn_name] = Distribution()
2079 # Combine those dictionaries into a single "information" dictionary
2080 content = parameters.copy()
2081 content.update(functions)
2082 content.update(distributions)
2083 info = deepcopy(content)
2085 # Parse the initialization routine
2086 initialize = format_block_statement(model["initialize"])
2088 # Make the list of ordered events
2089 events = []
2090 names_used_in_initialize = [] # this doesn't actually get used
2091 for line in initialize:
2092 # Make the new event and add it to the list
2093 new_event, names_used = make_new_event(line, info)
2094 events.append(new_event)
2095 names_used_in_initialize += names_used
2097 # Add newly assigned variables to the information set
2098 for var in new_event.assigns:
2099 if var in info.keys():
2100 raise ValueError(var + " is assigned, but already exists!")
2101 info[var] = 0
2103 # If any assigned variables are common, mark the event as common
2104 for var in new_event.assigns:
2105 if var in common:
2106 new_event.common = True
2107 break # No need to check further
2109 # Verify that all arrival variables were created in the initializer
2110 for var in arrival:
2111 if var not in info.keys():
2112 raise ValueError(
2113 "The arrival variable " + var + " was not set in the initialize block!"
2114 )
2116 # Make a blank dictionary with information the initializer needs
2117 init_requires = {}
2118 for event in events:
2119 for var in event.parameters.keys():
2120 if var not in init_requires.keys():
2121 try:
2122 init_requires[var] = parameters[var]
2123 except:
2124 raise ValueError(
2125 var
2126 + " was referenced in initialize, but not declared as a parameter!"
2127 )
2128 if type(event) is RandomEvent:
2129 try:
2130 dstn_name = event._dstn_name
2131 init_requires[dstn_name] = distributions[dstn_name]
2132 except:
2133 raise ValueError(
2134 dstn_name
2135 + " was referenced in initialize, but not declared as a distribution!"
2136 )
2137 if type(event) is EvaluationEvent:
2138 try:
2139 func_name = event._func_name
2140 init_requires[dstn_name] = functions[func_name]
2141 except:
2142 raise ValueError(
2143 func_name
2144 + " was referenced in initialize, but not declared as a function!"
2145 )
2147 # Make a single string initializer statement
2148 statement = ""
2149 longest = np.max([len(event.statement) for event in events])
2150 for event in events:
2151 this_statement = event.statement
2152 L = len(this_statement)
2153 pad = (longest + 1) - L
2154 statement += this_statement + pad * " " + ": " + event.description + "\n"
2156 # Make and return the new SimBlock
2157 initializer = SimBlock(
2158 description="agent initializer for " + name,
2159 content=init_requires,
2160 statement=statement,
2161 events=events,
2162 )
2163 return initializer, init_requires
2166def make_new_event(statement, info):
2167 """
2168 Makes a "blank" version of a model event based on a statement line. Determines
2169 which objects are needed vs assigned vs parameters / information from context.
2171 Parameters
2172 ----------
2173 statement : str
2174 One line of a model statement, which will be turned into an event.
2175 info : dict
2176 Empty dictionary of model information that already exists. Consists of
2177 arrival variables, already assigned variables, parameters, and functions.
2178 Typing of each is based on the kind of "empty" object.
2180 Returns
2181 -------
2182 new_event : ModelEvent
2183 A new model event with values and information missing, but structure set.
2184 names_used : [str]
2185 List of names of objects used in this expression.
2186 """
2187 # First determine what kind of event this is
2188 has_eq = "=" in statement
2189 has_tld = "~" in statement
2190 has_amp = "@" in statement
2191 has_brc = ("{" in statement) and ("}" in statement)
2192 has_brk = ("[" in statement) and ("]" in statement)
2193 event_type = None
2194 if has_eq:
2195 if has_tld:
2196 raise ValueError("A statement line can't have both an = and a ~!")
2197 if has_amp:
2198 event_type = EvaluationEvent
2199 else:
2200 event_type = DynamicEvent
2201 if has_tld:
2202 if has_brc:
2203 event_type = MarkovEvent
2204 elif has_brk:
2205 event_type = RandomIndexedEvent
2206 else:
2207 event_type = RandomEvent
2208 if event_type is None:
2209 raise ValueError("Statement line was not any valid type!")
2211 # Now make and return an appropriate event for that type
2212 if event_type is DynamicEvent:
2213 event_maker = make_new_dynamic
2214 if event_type is RandomEvent:
2215 event_maker = make_new_random
2216 if event_type is RandomIndexedEvent:
2217 event_maker = make_new_random_indexed
2218 if event_type is MarkovEvent:
2219 event_maker = make_new_markov
2220 if event_type is EvaluationEvent:
2221 event_maker = make_new_evaluation
2223 new_event, names_used = event_maker(statement, info)
2224 return new_event, names_used
2227def make_new_dynamic(statement, info):
2228 """
2229 Construct a new instance of DynamicEvent based on the given model statement
2230 line and a blank dictionary of parameters. The statement should already be
2231 verified to be a valid dynamic statement: it has an = but no ~ or @.
2233 Parameters
2234 ----------
2235 statement : str
2236 One line dynamics statement, which will be turned into a DynamicEvent.
2237 info : dict
2238 Empty dictionary of available information.
2240 Returns
2241 -------
2242 new_dynamic : DynamicEvent
2243 A new dynamic event with values and information missing, but structure set.
2244 names_used : [str]
2245 List of names of objects used in this expression.
2246 """
2247 # Cut the statement up into its LHS, RHS, and description
2248 lhs, rhs, description = parse_line_for_parts(statement, "=")
2250 # Parse the LHS (assignment) to get assigned variables
2251 assigns = parse_assignment(lhs)
2253 # Parse the RHS (dynamic statement) to extract object names used
2254 obj_names, is_indexed = extract_var_names_from_expr(rhs)
2256 # Allocate each variable to needed dynamic variables or parameters
2257 needs = []
2258 parameters = {}
2259 for j in range(len(obj_names)):
2260 var = obj_names[j]
2261 if var not in info.keys():
2262 raise ValueError(
2263 var + " is used in a dynamic expression, but does not (yet) exist!"
2264 )
2265 val = info[var]
2266 if type(val) is NullFunc:
2267 raise ValueError(
2268 var + " is used in a dynamic expression, but it's a function!"
2269 )
2270 if type(val) is Distribution:
2271 raise ValueError(
2272 var + " is used in a dynamic expression, but it's a distribution!"
2273 )
2274 if val is None:
2275 parameters[var] = None
2276 else:
2277 needs.append(var)
2279 # Declare a SymPy symbol for each variable used; these are temporary
2280 _args = []
2281 for j in range(len(obj_names)):
2282 _var = obj_names[j]
2283 if is_indexed[j]:
2284 exec(_var + " = IndexedBase('" + _var + "')")
2285 else:
2286 exec(_var + " = symbols('" + _var + "')")
2287 _args.append(symbols(_var))
2289 # Make a SymPy expression, then lambdify it
2290 sympy_expr = symbols(rhs)
2291 expr = lambdify(_args, sympy_expr)
2293 # Make an overall list of object names referenced in this event
2294 names_used = assigns + obj_names
2296 # Make and return the new dynamic event
2297 new_dynamic = DynamicEvent(
2298 description=description,
2299 statement=lhs + " = " + rhs,
2300 assigns=assigns,
2301 needs=needs,
2302 parameters=parameters,
2303 expr=expr,
2304 args=obj_names,
2305 )
2306 return new_dynamic, names_used
2309def make_new_random(statement, info):
2310 """
2311 Make a new random variable realization event based on the given model statement
2312 line and a blank dictionary of parameters. The statement should already be
2313 verified to be a valid random statement: it has a ~ but no = or [].
2315 Parameters
2316 ----------
2317 statement : str
2318 One line of the model statement, which will be turned into a random event.
2319 info : dict
2320 Empty dictionary of available information.
2322 Returns
2323 -------
2324 new_random : RandomEvent
2325 A new random event with values and information missing, but structure set.
2326 names_used : [str]
2327 List of names of objects used in this expression.
2328 """
2329 # Cut the statement up into its LHS, RHS, and description
2330 lhs, rhs, description = parse_line_for_parts(statement, "~")
2332 # Parse the LHS (assignment) to get assigned variables
2333 assigns = parse_assignment(lhs)
2335 # Verify that the RHS is actually a distribution
2336 if type(info[rhs]) is not Distribution:
2337 raise ValueError(
2338 rhs + " was treated as a distribution, but not declared as one!"
2339 )
2341 # Make an overall list of object names referenced in this event
2342 names_used = assigns + [rhs]
2344 # Make and return the new random event
2345 new_random = RandomEvent(
2346 description=description,
2347 statement=lhs + " ~ " + rhs,
2348 assigns=assigns,
2349 needs=[],
2350 parameters={},
2351 dstn=info[rhs],
2352 )
2353 new_random._dstn_name = rhs
2354 return new_random, names_used
2357def make_new_random_indexed(statement, info):
2358 """
2359 Make a new indexed random variable realization event based on the given model
2360 statement line and a blank dictionary of parameters. The statement should
2361 already be verified to be a valid random statement: it has a ~ and [].
2363 Parameters
2364 ----------
2365 statement : str
2366 One line of the model statement, which will be turned into a random event.
2367 info : dict
2368 Empty dictionary of available information.
2370 Returns
2371 -------
2372 new_random_indexed : RandomEvent
2373 A new random indexed event with values and information missing, but structure set.
2374 names_used : [str]
2375 List of names of objects used in this expression.
2376 """
2377 # Cut the statement up into its LHS, RHS, and description
2378 lhs, rhs, description = parse_line_for_parts(statement, "~")
2380 # Parse the LHS (assignment) to get assigned variables
2381 assigns = parse_assignment(lhs)
2383 # Split the RHS into the distribution and the index
2384 dstn, index = parse_random_indexed(rhs)
2386 # Verify that the RHS is actually a distribution
2387 if type(info[dstn]) is not Distribution:
2388 raise ValueError(
2389 dstn + " was treated as a distribution, but not declared as one!"
2390 )
2392 # Make an overall list of object names referenced in this event
2393 names_used = assigns + [dstn, index]
2395 # Make and return the new random indexed event
2396 new_random_indexed = RandomIndexedEvent(
2397 description=description,
2398 statement=lhs + " ~ " + rhs,
2399 assigns=assigns,
2400 needs=[index],
2401 parameters={},
2402 index=index,
2403 )
2404 new_random_indexed._dstn_name = dstn
2405 return new_random_indexed, names_used
2408def make_new_markov(statement, info):
2409 """
2410 Make a new Markov-type event based on the given model statement line and a
2411 blank dictionary of parameters. The statement should already be verified to
2412 be a valid Markov statement: it has a ~ and {} and maybe (). This can represent
2413 a Markov matrix transition event, a draw from a discrete index, or just a
2414 Bernoulli random variable. If a Bernoulli event, the "probabilties" can be
2415 idiosyncratic data.
2417 Parameters
2418 ----------
2419 statement : str
2420 One line of the model statement, which will be turned into a random event.
2421 info : dict
2422 Empty dictionary of available information.
2424 Returns
2425 -------
2426 new_markov : MarkovEvent
2427 A new Markov draw event with values and information missing, but structure set.
2428 names_used : [str]
2429 List of names of objects used in this expression.
2430 """
2431 # Cut the statement up into its LHS, RHS, and description
2432 lhs, rhs, description = parse_line_for_parts(statement, "~")
2434 # Parse the LHS (assignment) to get assigned variables
2435 assigns = parse_assignment(lhs)
2437 # Parse the RHS (Markov statement) for the array and index
2438 probs, index = parse_markov(rhs)
2439 if index is None:
2440 needs = []
2441 else:
2442 needs = [index]
2444 # Determine whether probs is an idiosyncratic variable or a parameter, and
2445 # set up the event to grab it appropriately
2446 if info[probs] is None:
2447 parameters = {probs: None}
2448 else:
2449 needs += [probs]
2450 parameters = {}
2452 # Make an overall list of object names referenced in this event
2453 names_used = assigns + needs + [probs]
2455 # Make and return the new Markov event
2456 new_markov = MarkovEvent(
2457 description=description,
2458 statement=lhs + " ~ " + rhs,
2459 assigns=assigns,
2460 needs=needs,
2461 parameters=parameters,
2462 probs=probs,
2463 index=index,
2464 )
2465 return new_markov, names_used
2468def make_new_evaluation(statement, info):
2469 """
2470 Make a new function evaluation event based the given model statement line
2471 and a blank dictionary of parameters. The statement should already be verified
2472 to be a valid evaluation statement: it has an @ and an = but no ~.
2474 Parameters
2475 ----------
2476 statement : str
2477 One line of the model statement, which will be turned into an eval event.
2478 info : dict
2479 Empty dictionary of available information.
2481 Returns
2482 -------
2483 new_evaluation : EvaluationEvent
2484 A new evaluation event with values and information missing, but structure set.
2485 names_used : [str]
2486 List of names of objects used in this expression.
2487 """
2488 # Cut the statement up into its LHS, RHS, and description
2489 lhs, rhs, description = parse_line_for_parts(statement, "=")
2491 # Parse the LHS (assignment) to get assigned variables
2492 assigns = parse_assignment(lhs)
2494 # Parse the RHS (evaluation) for the function and its arguments
2495 func, arguments = parse_evaluation(rhs)
2497 # Allocate each variable to needed dynamic variables or parameters
2498 needs = []
2499 parameters = {}
2500 for j in range(len(arguments)):
2501 var = arguments[j]
2502 if var not in info.keys():
2503 raise ValueError(
2504 var + " is used in an evaluation statement, but does not (yet) exist!"
2505 )
2506 val = info[var]
2507 if type(val) is NullFunc:
2508 raise ValueError(
2509 var
2510 + " is used as an argument an evaluation statement, but it's a function!"
2511 )
2512 if type(val) is Distribution:
2513 raise ValueError(
2514 var + " is used in an evaluation statement, but it's a distribution!"
2515 )
2516 if val is None:
2517 parameters[var] = None
2518 else:
2519 needs.append(var)
2521 # Make an overall list of object names referenced in this event
2522 names_used = assigns + arguments + [func]
2524 # Make and return the new evaluation event
2525 new_evaluation = EvaluationEvent(
2526 description=description,
2527 statement=lhs + " = " + rhs,
2528 assigns=assigns,
2529 needs=needs,
2530 parameters=parameters,
2531 arguments=arguments,
2532 func=info[func],
2533 )
2534 new_evaluation._func_name = func
2535 return new_evaluation, names_used
2538def look_for_char_and_remove(phrase, symb):
2539 """
2540 Check whether a symbol appears in a string, and remove it if it does.
2542 Parameters
2543 ----------
2544 phrase : str
2545 String to be searched for a symbol.
2546 symb : char
2547 Single character to be searched for.
2549 Returns
2550 -------
2551 out : str
2552 Possibly shortened input phrase.
2553 found : bool
2554 Whether the symbol was found and removed.
2555 """
2556 found = symb in phrase
2557 out = phrase.replace(symb, "")
2558 return out, found
2561def parse_declaration_for_parts(line):
2562 """
2563 Split a declaration line from a model file into the object's name, its datatype,
2564 any metadata flags, and any provided comment or description.
2566 Parameters
2567 ----------
2568 line : str
2569 Line of to be parsed into the object name, object type, and a comment or description.
2571 Returns
2572 -------
2573 name : str
2574 Name of the object.
2575 datatype : str or None
2576 Provided datatype string, in parentheses, if any.
2577 flags : [str]
2578 List of metadata flags that were detected. These include ! for a variable
2579 that is in arrival, * for any non-variable that's part of the solution,
2580 + for any object that is offset in time, and & for a common random variable.
2582 desc : str
2583 Comment or description, after //, if any.
2584 """
2585 flags = []
2586 check_for_flags = {"offset": "+", "arrival": "!", "solution": "*", "common": "&"}
2588 # First, separate off the comment or description, if any
2589 slashes = line.find("\\")
2590 desc = "" if slashes == -1 else line[(slashes + 2) :].strip()
2591 rem = line if slashes == -1 else line[:slashes].strip()
2593 # Now look for bracketing parentheses declaring a datatype
2594 lp = rem.find("(")
2595 if lp > -1:
2596 rp = rem.find(")")
2597 if rp == -1:
2598 raise ValueError("Unclosed parentheses on object declaration line!")
2599 datatype = rem[(lp + 1) : rp].strip()
2600 leftover = rem[:lp].strip()
2601 else:
2602 datatype = None
2603 leftover = rem
2605 # What's left over should be the object name plus any flags
2606 for key in check_for_flags.keys():
2607 symb = check_for_flags[key]
2608 leftover, found = look_for_char_and_remove(leftover, symb)
2609 if found:
2610 flags.append(key)
2612 # Remove any remaining spaces, and that *should* be the name
2613 name = leftover.replace(" ", "")
2614 # TODO: Check for valid name formatting based on characters.
2616 return name, datatype, flags, desc
2619def parse_line_for_parts(statement, symb):
2620 """
2621 Split one line of a model statement into its LHS, RHS, and description. The
2622 description is everything following \\, while the LHS and RHS are determined
2623 by a special symbol.
2625 Parameters
2626 ----------
2627 statement : str
2628 One line of a model statement, which will be parsed for its parts.
2629 symb : char
2630 The character that represents the divide between LHS and RHS
2632 Returns
2633 -------
2634 lhs : str
2635 The left-hand (assignment) side of the expression.
2636 rhs : str
2637 The right-hand (evaluation) side of the expression.
2638 desc : str
2639 The provided description of the expression.
2640 """
2641 eq = statement.find(symb)
2642 lhs = statement[:eq].replace(" ", "")
2643 not_lhs = statement[(eq + 1) :]
2644 comment = not_lhs.find("\\")
2645 desc = "" if comment == -1 else not_lhs[(comment + 2) :].strip()
2646 rhs = not_lhs if comment == -1 else not_lhs[:comment]
2647 rhs = rhs.replace(" ", "")
2648 return lhs, rhs, desc
2651def parse_assignment(lhs):
2652 """
2653 Get ordered list of assigned variables from the LHS of a model line.
2655 Parameters
2656 ----------
2657 lhs : str
2658 Left-hand side of a model expression
2660 Returns
2661 -------
2662 assigns : List[str]
2663 List of variable names that are assigned in this model line.
2664 """
2665 if lhs[0] == "(":
2666 if not lhs[-1] == ")":
2667 raise ValueError("Parentheses on assignment was not closed!")
2668 assigns = []
2669 pos = 0
2670 while pos != -1:
2671 pos += 1
2672 end = lhs.find(",", pos)
2673 var = lhs[pos:end]
2674 if var != "":
2675 assigns.append(var)
2676 pos = end
2677 else:
2678 assigns = [lhs]
2679 return assigns
2682def extract_var_names_from_expr(expression):
2683 """
2684 Parse the RHS of a dynamic model statement to get variable names used in it.
2686 Parameters
2687 ----------
2688 expression : str
2689 RHS of a model statement to be parsed for variable names.
2691 Returns
2692 -------
2693 var_names : List[str]
2694 List of variable names used in the expression. These *should* be dynamic
2695 variables and parameters, but not functions.
2696 indexed : List[bool]
2697 Indicators for whether each variable seems to be used with indexing.
2698 """
2699 var_names = []
2700 indexed = []
2701 math_symbols = "+-/*^%.(),[]{}<>"
2702 digits = "01234567890"
2703 cur = ""
2704 for j in range(len(expression)):
2705 c = expression[j]
2706 if (c in math_symbols) or ((c in digits) and cur == ""):
2707 if cur == "":
2708 continue
2709 if cur in var_names:
2710 cur = ""
2711 continue
2712 var_names.append(cur)
2713 if c == "[":
2714 indexed.append(True)
2715 else:
2716 indexed.append(False)
2717 cur = ""
2718 else:
2719 cur += c
2720 if cur != "" and cur not in var_names:
2721 var_names.append(cur)
2722 indexed.append(False) # final symbol couldn't possibly be indexed
2723 return var_names, indexed
2726def parse_evaluation(expression):
2727 """
2728 Separate a function evaluation expression into the function that is called
2729 and the variable inputs that are passed to it.
2731 Parameters
2732 ----------
2733 expression : str
2734 RHS of a function evaluation model statement, which will be parsed for
2735 the function and its inputs.
2737 Returns
2738 -------
2739 func_name : str
2740 Name of the function that will be called in this event.
2741 arg_names : List[str]
2742 List of arguments of the function.
2743 """
2744 # Get the name of the function: what's to the left of the @
2745 amp = expression.find("@")
2746 func_name = expression[:amp]
2748 # Check for parentheses formatting
2749 rem = expression[(amp + 1) :]
2750 if not rem[0] == "(":
2751 raise ValueError(
2752 "The @ in a function evaluation statement must be followed by (!"
2753 )
2754 if not rem[-1] == ")":
2755 raise ValueError("A function evaluation statement must end in )!")
2756 rem = rem[1:-1]
2758 # Parse what's inside the parentheses for argument names
2759 arg_names = []
2760 pos = 0
2761 go = True
2762 while go:
2763 end = rem.find(",", pos)
2764 if end > -1:
2765 arg = rem[pos:end]
2766 else:
2767 arg = rem[pos:]
2768 go = False
2769 if arg != "":
2770 arg_names.append(arg)
2771 pos = end + 1
2773 return func_name, arg_names
2776def parse_markov(expression):
2777 """
2778 Separate a Markov draw declaration into the array of probabilities and the
2779 index for idiosyncratic values.
2781 Parameters
2782 ----------
2783 expression : str
2784 RHS of a function evaluation model statement, which will be parsed for
2785 the probabilities name and index name.
2787 Returns
2788 -------
2789 probs : str
2790 Name of the probabilities object in this statement.
2791 index : str
2792 Name of the indexing variable in this statement.
2793 """
2794 # Get the name of the probabilitie
2795 lb = expression.find("{") # this *should* be 0
2796 rb = expression.find("}")
2797 if lb == -1 or rb == -1 or rb < (lb + 2):
2798 raise ValueError("A Markov assignment must have an {array}!")
2799 probs = expression[(lb + 1) : rb]
2801 # Get the name of the index, if any
2802 x = rb + 1
2803 lp = expression.find("(", x)
2804 rp = expression.find(")", x)
2805 if lp == -1 and rp == -1: # no index present at all
2806 return probs, None
2807 if lp == -1 or rp == -1 or rp < (lp + 2):
2808 raise ValueError("Improper Markov formatting: should be {probs}(index)!")
2809 index = expression[(lp + 1) : rp]
2811 return probs, index
2814def parse_random_indexed(expression):
2815 """
2816 Separate an indexed random variable assignment into the distribution and
2817 the index for it.
2819 Parameters
2820 ----------
2821 expression : str
2822 RHS of a function evaluation model statement, which will be parsed for
2823 the distribution name and index name.
2825 Returns
2826 -------
2827 dstn : str
2828 Name of the distribution in this statement.
2829 index : str
2830 Name of the indexing variable in this statement.
2831 """
2832 # Get the name of the index
2833 lb = expression.find("[")
2834 rb = expression.find("]")
2835 if lb == -1 or rb == -1 or rb < (lb + 2):
2836 raise ValueError("An indexed random variable assignment must have an [index]!")
2837 index = expression[(lb + 1) : rb]
2839 # Get the name of the distribution
2840 dstn = expression[:lb]
2842 return dstn, index
2845def format_block_statement(statement):
2846 """
2847 Ensure that a string stagement of a model block (maybe a period, maybe an
2848 initializer) is formatted as a list of strings, one statement per entry.
2850 Parameters
2851 ----------
2852 statement : str
2853 A model statement, which might be for a block or an initializer. The
2854 statement might be formatted as a list or as a single string.
2856 Returns
2857 -------
2858 block_statements: [str]
2859 A list of model statements, one per entry.
2860 """
2861 if type(statement) is str:
2862 if statement.find("\n") > -1:
2863 block_statements = []
2864 pos = 0
2865 end = statement.find("\n", pos)
2866 while end > -1:
2867 new_line = statement[pos:end]
2868 block_statements.append(new_line)
2869 pos = end + 1
2870 end = statement.find("\n", pos)
2871 else:
2872 block_statements = [statement.copy()]
2873 if type(statement) is list:
2874 for line in statement:
2875 if type(line) is not str:
2876 raise ValueError("The model statement somehow includes a non-string!")
2877 block_statements = statement.copy()
2878 return block_statements
2881@njit
2882def aggregate_blobs_onto_polynomial_grid(
2883 vals, pmv, origins, grid, J, Q
2884): # pragma: no cover
2885 """
2886 Numba-compatible helper function for casting "probability blobs" onto a discretized
2887 grid of outcome values, based on their origin in the arrival state space. This
2888 version is for non-continuation variables, returning only the probability array
2889 mapping from arrival states to the outcome variable.
2890 """
2891 bot = grid[0]
2892 top = grid[-1]
2893 M = grid.size
2894 Mm1 = M - 1
2895 N = pmv.size
2896 scale = 1.0 / (top - bot)
2897 order = 1.0 / Q
2898 diffs = grid[1:] - grid[:-1]
2900 probs = np.zeros((J, M))
2902 for n in range(N):
2903 x = vals[n]
2904 jj = origins[n]
2905 p = pmv[n]
2906 if (x > bot) and (x < top):
2907 ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
2908 temp = (x - grid[ii]) / diffs[ii]
2909 probs[jj, ii] += (1.0 - temp) * p
2910 probs[jj, ii + 1] += temp * p
2911 elif x <= bot:
2912 probs[jj, 0] += p
2913 else:
2914 probs[jj, -1] += p
2915 return probs
2918@njit
2919def aggregate_blobs_onto_polynomial_grid_alt(
2920 vals, pmv, origins, grid, J, Q
2921): # pragma: no cover
2922 """
2923 Numba-compatible helper function for casting "probability blobs" onto a discretized
2924 grid of outcome values, based on their origin in the arrival state space. This
2925 version is for ncontinuation variables, returning the probability array mapping
2926 from arrival states to the outcome variable, the index in the outcome variable grid
2927 for each blob, and the alpha weighting between gridpoints.
2928 """
2929 bot = grid[0]
2930 top = grid[-1]
2931 M = grid.size
2932 Mm1 = M - 1
2933 N = pmv.size
2934 scale = 1.0 / (top - bot)
2935 order = 1.0 / Q
2936 diffs = grid[1:] - grid[:-1]
2938 probs = np.zeros((J, M))
2939 idx = np.empty(N, dtype=np.dtype(np.int32))
2940 alpha = np.empty(N)
2942 for n in range(N):
2943 x = vals[n]
2944 jj = origins[n]
2945 p = pmv[n]
2946 if (x > bot) and (x < top):
2947 ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
2948 temp = (x - grid[ii]) / diffs[ii]
2949 probs[jj, ii] += (1.0 - temp) * p
2950 probs[jj, ii + 1] += temp * p
2951 alpha[n] = temp
2952 idx[n] = ii
2953 elif x <= bot:
2954 probs[jj, 0] += p
2955 alpha[n] = 0.0
2956 idx[n] = 0
2957 else:
2958 probs[jj, -1] += p
2959 alpha[n] = 1.0
2960 idx[n] = M - 2
2961 return probs, idx, alpha
2964@njit
2965def aggregate_blobs_onto_discrete_grid(vals, pmv, origins, M, J): # pragma: no cover
2966 """
2967 Numba-compatible helper function for allocating "probability blobs" to a grid
2968 over a discrete state-- the state itself is truly discrete.
2969 """
2970 out = np.zeros((J, M))
2971 N = pmv.size
2972 for n in range(N):
2973 ii = vals[n]
2974 jj = origins[n]
2975 p = pmv[n]
2976 out[jj, ii] += p
2977 return out
2980@njit
2981def calc_overall_trans_probs(
2982 out, idx, alpha, binary, offset, pmv, origins
2983): # pragma: no cover
2984 """
2985 Numba-compatible helper function for combining transition probabilities from
2986 the arrival state space to *multiple* continuation variables into a single
2987 unified transition matrix.
2988 """
2989 N = alpha.shape[0]
2990 B = binary.shape[0]
2991 D = binary.shape[1]
2992 for n in range(N):
2993 ii = origins[n]
2994 jj_base = idx[n]
2995 p = pmv[n]
2996 for b in range(B):
2997 adj = offset[b]
2998 P = p
2999 for d in range(D):
3000 k = binary[b, d]
3001 P *= alpha[n, d, k]
3002 jj = jj_base + adj
3003 out[ii, jj] += P
3004 return out