def test_normal_context(self): with _DrawValuesContext() as context0: assert context0.parent is None context0.drawn_vars['root_test'] = 1 with _DrawValuesContext() as context1: assert id(context1.drawn_vars) == id(context0.drawn_vars) assert context1.parent == context0 with _DrawValuesContext() as context2: assert id(context2.drawn_vars) == id(context0.drawn_vars) assert context2.parent == context1 context2.drawn_vars['leaf_test'] = 2 assert context1.drawn_vars['leaf_test'] == 2 context1.drawn_vars['root_test'] = 3 assert context0.drawn_vars['root_test'] == 3 assert context0.drawn_vars['leaf_test'] == 2
def test_mixed_contexts(): modelA = Model() modelB = Model() with raises((ValueError, TypeError)): modelcontext(None) with modelA: with modelB: assert Model.get_context() == modelB assert modelcontext(None) == modelB dvc = _DrawValuesContext() with dvc: assert Model.get_context() == modelB assert modelcontext(None) == modelB assert _DrawValuesContext.get_context() == dvc dvcb = _DrawValuesContextBlocker() with dvcb: assert _DrawValuesContext.get_context() == dvcb assert _DrawValuesContextBlocker.get_context() == dvcb assert _DrawValuesContext.get_context() == dvc assert _DrawValuesContextBlocker.get_context() is dvc assert Model.get_context() == modelB assert modelcontext(None) == modelB assert _DrawValuesContext.get_context(error_if_none=False) is None with raises(TypeError): _DrawValuesContext.get_context() assert Model.get_context() == modelB assert modelcontext(None) == modelB assert Model.get_context() == modelA assert modelcontext(None) == modelA assert Model.get_context(error_if_none=False) is None with raises(TypeError): Model.get_context(error_if_none=True) with raises((ValueError, TypeError)): modelcontext(None)
def random(self, point=None, size=None): """Sample from this distribution conditional on a given set of values. Parameters ---------- point: dict, optional Dict of variable values on which random values are to be conditioned (uses default point if not specified). size: int, optional Desired size of random sample (returns one sample if not specified). Returns ------- array """ with _DrawValuesContext(): (states, ) = draw_values([self.states], point=point, size=size) # This is a terrible thing to have to do here, but it's better than # having to (know to) update `Distribution.shape` when/if dimensions # change (e.g. when sampling new state sequences). bcast_comps = np.broadcast( states, *[dist.random(point=point) for dist in self.comp_dists]) self_shape = bcast_comps.shape if size: # `draw_values` will not honor the `size` parameter if its arguments # don't contain random variables, so, when our `self.states` are # constants, we have to broadcast `states` so that it matches `size + # self.shape`. expanded_states = np.broadcast_to( states, tuple(np.atleast_1d(size)) + self_shape) else: expanded_states = np.broadcast_to(states, self_shape) samples = np.empty(expanded_states.shape) for i, dist in enumerate(self.comp_dists): # We want to sample from only the parts of our component # distributions that are active given the states. # This is only really relevant when the component distributions # change over the state space (e.g. Poisson means that change # over time). # We could always sample such components over the entire space # (e.g. time), but, for spaces with large dimension, that would # be extremely costly and wasteful. i_idx = np.where(expanded_states == i) i_size = len(i_idx[0]) if i_size > 0: subset_kwargs = distribution_subset_args( dist, expanded_states.shape, i_idx) state_dist = dist.dist(**subset_kwargs) sample = state_dist.random(point=point) samples[i_idx] = sample return samples
def test_blocking_context(self): with _DrawValuesContext() as context0: assert context0.parent is None context0.drawn_vars['root_test'] = 1 with _DrawValuesContext() as context1: assert id(context1.drawn_vars) == id(context0.drawn_vars) assert context1.parent == context0 with _DrawValuesContextBlocker() as blocker: assert id(blocker.drawn_vars) != id(context0.drawn_vars) assert blocker.parent is None blocker.drawn_vars['root_test'] = 2 with _DrawValuesContext() as context2: assert id(context2.drawn_vars) == id(blocker.drawn_vars) assert context2.parent == blocker context2.drawn_vars['root_test'] = 3 context2.drawn_vars['leaf_test'] = 4 assert blocker.drawn_vars['root_test'] == 3 assert 'leaf_test' not in context1.drawn_vars assert context0.drawn_vars['root_test'] == 1
def random(self, point=None, size=None): """Sample from this distribution conditional on a given set of values. Parameters ---------- point: dict, optional Dict of variable values on which random values are to be conditioned (uses default point if not specified). size: int, optional Desired size of random sample (returns one sample if not specified). Returns ------- array """ with _DrawValuesContext() as draw_context: # TODO FIXME: Very, very lame... term_smpl = draw_context.drawn_vars.get((self.states, 1), None) if term_smpl is not None: point[self.states.name] = term_smpl # `draw_values` is inconsistent and will not use the `size` # parameter if the variables aren't random variables. if hasattr(self.states, "distribution"): (states,) = draw_values([self.states], point=point, size=size) else: states = pm.Constant.dist(self.states).random(point=point, size=size) # states = states.T samples = np.empty(states.shape) for i, dist in enumerate(self.comp_dists): # We want to sample from only the parts of our component # distributions that are active given the states. # This is only really relevant when the component distributions # change over the state space (e.g. Poisson means that change # over time). # We could always sample such components over the entire space # (e.g. time), but, for spaces with large dimension, that would # be extremely costly and wasteful. i_idx = np.where(states == i) i_size = len(i_idx[0]) if i_size > 0: subset_args = distribution_subset_args( dist, states.shape, i_idx, point=point ) samples[i_idx] = dist.dist(*subset_args).random(point=point) return samples
def random(self, point=None, size=None): """Sample from this distribution conditional on a given set of values. Parameters ---------- point: dict, optional Dict of variable values on which random values are to be conditioned (uses default point if not specified). size: int, optional Desired size of random sample (returns one sample if not specified). Returns ------- array """ with _DrawValuesContext() as draw_context: terms = [self.gamma_0, self.Gammas] gamma_0, Gamma = draw_values(terms, point=point) # Sample state 0 in each state sequence state_n = pm.Categorical.dist(gamma_0, shape=self.shape[:-1]).random( point=point, size=size ) state_shape = state_n.shape N = self.shape[-1] states = np.empty(state_shape + (N,), dtype=self.dtype) unif_samples = np.random.uniform(size=states.shape) # Make sure we have a transition matrix for each element in a state # sequence Gamma = np.broadcast_to(Gamma, tuple(states.shape) + Gamma.shape[-2:]) # Slices across each independent/replication dimension slices = [slice(None, d) for d in state_shape] slices = tuple(np.ogrid[slices]) for n in range(0, N): gamma_t = Gamma[..., n, :, :] gamma_t = gamma_t[slices + (state_n,)] state_n = vsearchsorted(gamma_t.cumsum(axis=-1), unif_samples[..., n]) states[..., n] = state_n return states
def draw_values(self) -> list[np.ndarray]: vars = self.vars trace = self.trace samples = self.samples # size = self.size params = dict(enumerate(vars)) with _DrawValuesContext() as context: self.init() self.make_graph() drawn = context.drawn_vars # Init givens and the stack of nodes to try to `_draw_value` from givens = { p.name: (p, v) for (p, samples), v in drawn.items() if getattr(p, "name", None) is not None } stack = list( self.leaf_nodes.values()) # A queue would be more appropriate while stack: next_ = stack.pop(0) if (next_, samples) in drawn: # If the node already has a givens value, skip it continue elif isinstance(next_, (Constant, SharedVariable)): # If the node is a aesara.tensor.TensorConstant or a # aesara.tensor.sharedvar.SharedVariable, its value will be # available automatically in _compile_aesara_function so # we can skip it. Furthermore, if this node was treated as a # TensorVariable that should be compiled by aesara in # _compile_aesara_function, it would raise a `TypeError: # ('Constants not allowed in param list', ...)` for # TensorConstant, and a `TypeError: Cannot use a shared # variable (...) as explicit input` for SharedVariable. # ObservedRV and MultiObservedRV instances are ViewOPs # of TensorConstants or SharedVariables, we must add them # to the stack or risk evaluating deterministics with the # wrong values (issue #3354) stack.extend([ node for node in self.named_nodes_parents[next_] if isinstance(node, (ObservedRV, MultiObservedRV)) and (node, samples) not in drawn ]) continue else: # If the node does not have a givens value, try to draw it. # The named node's children givens values must also be taken # into account. children = self.named_nodes_children[next_] temp_givens = [givens[k] for k in givens if k in children] try: # This may fail for autotransformed RVs, which don't # have the random method value = self.draw_value(next_, trace=trace, givens=temp_givens) assert isinstance(value, np.ndarray) givens[next_.name] = (next_, value) drawn[(next_, samples)] = value except aesara.graph.fg.MissingInputError: # The node failed, so we must add the node's parents to # the stack of nodes to try to draw from. We exclude the # nodes in the `params` list. stack.extend([ node for node in self.named_nodes_parents[next_] if node is not None and (node, samples) not in drawn ]) # the below makes sure the graph is evaluated in order # test_distributions_random::TestDrawValues::test_draw_order fails without it # The remaining params that must be drawn are all hashable to_eval: set[int] = set() missing_inputs: set[int] = {j for j, p in self.symbolic_params} while to_eval or missing_inputs: if to_eval == missing_inputs: raise ValueError("Cannot resolve inputs for {}".format( [get_var_name(trace.varnames[j]) for j in to_eval])) to_eval = set(missing_inputs) missing_inputs = set() for param_idx in to_eval: param = vars[param_idx] drawn = context.drawn_vars if (param, samples) in drawn: self.evaluated[param_idx] = drawn[(param, samples)] else: try: if param in self.named_nodes_children: for node in self.named_nodes_children[param]: if node.name not in givens and ( node, samples) in drawn: givens[node.name] = ( node, drawn[(node, samples)], ) value = self.draw_value(param, trace=self.trace, givens=givens.values()) assert isinstance(value, np.ndarray) self.evaluated[param_idx] = drawn[( param, samples)] = value givens[param.name] = (param, value) except aesara.graph.fg.MissingInputError: missing_inputs.add(param_idx) return [self.evaluated[j] for j in params]
def random(self, point=None, size=None): """ Draw random values from defined Mixture distribution. Parameters ---------- point: dict, optional Dict of variable values on which random values are to be conditioned (uses default point if not specified). size: int, optional Desired size of random sample (returns one sample if not specified). Returns ------- array """ # Convert size to tuple size = to_tuple(size) # Draw mixture weights and infer the comp_dists shapes with _DrawValuesContext() as draw_context: # We first need to check w and comp_tmp shapes and re compute size w = draw_values([self.w], point=point, size=size)[0] comp_dist_shapes, broadcast_shape = self.infer_comp_dist_shapes(point=point) # When size is not None, it's hard to tell the w parameter shape if size is not None and w.shape[: len(size)] == size: w_shape = w.shape[len(size) :] else: w_shape = w.shape # Try to determine parameter shape and dist_shape if self.comp_is_distribution: param_shape = np.broadcast(np.empty(w_shape), np.empty(broadcast_shape)).shape else: param_shape = np.broadcast(np.empty(w_shape), np.empty(broadcast_shape + (1,))).shape if np.asarray(self.shape).size != 0: dist_shape = np.broadcast(np.empty(self.shape), np.empty(param_shape[:-1])).shape else: dist_shape = param_shape[:-1] # Try to determine the size that must be used to get the mixture # components (i.e. get random choices using w). # 1. There must be size independent choices based on w. # 2. There must also be independent draws for each non singleton axis # of w. # 3. There must also be independent draws for each dimension added by # self.shape with respect to the w.ndim. These usually correspond to # observed variables with batch shapes wsh = (1,) * (len(dist_shape) - len(w_shape) + 1) + w_shape[:-1] psh = (1,) * (len(dist_shape) - len(param_shape) + 1) + param_shape[:-1] w_sample_size = [] # Loop through the dist_shape to get the conditions 2 and 3 first for i in range(len(dist_shape)): if dist_shape[i] != psh[i] and wsh[i] == 1: # self.shape[i] is a non singleton dimension (usually caused by # observed data) sh = dist_shape[i] else: sh = wsh[i] w_sample_size.append(sh) if size is not None and w_sample_size[: len(size)] != size: w_sample_size = size + tuple(w_sample_size) # Broadcast w to the w_sample_size (add a singleton last axis for the # mixture components) w = broadcast_distribution_samples([w, np.empty(w_sample_size + (1,))], size=size)[0] # Semiflatten the mixture weights. The last axis is the number of # mixture mixture components, and the rest is all about size, # dist_shape and broadcasting w_ = np.reshape(w, (-1, w.shape[-1])) w_samples = random_choice(p=w_, size=None) # w's shape already includes size # Now we broadcast the chosen components to the dist_shape w_samples = np.reshape(w_samples, w.shape[:-1]) if size is not None and dist_shape[: len(size)] != size: w_samples = np.broadcast_to(w_samples, size + dist_shape) else: w_samples = np.broadcast_to(w_samples, dist_shape) # When size is not None, maybe dist_shape partially overlaps with size if size is not None: if size == dist_shape: size = None elif size[-len(dist_shape) :] == dist_shape: size = size[: len(size) - len(dist_shape)] # We get an integer _size instead of a tuple size for drawing the # mixture, then we just reshape the output if size is None: _size = None else: _size = int(np.prod(size)) # Compute the total size of the mixture's random call with size if _size is not None: output_size = int(_size * np.prod(dist_shape) * param_shape[-1]) else: output_size = int(np.prod(dist_shape) * param_shape[-1]) # Get the size we need for the mixture's random call if self.comp_is_distribution: mixture_size = int(output_size // np.prod(broadcast_shape)) else: mixture_size = int(output_size // (np.prod(broadcast_shape) * param_shape[-1])) if mixture_size == 1 and _size is None: mixture_size = None # Sample from the mixture with draw_context: mixed_samples = self._comp_samples( point=point, size=mixture_size, broadcast_shape=broadcast_shape, comp_dist_shapes=comp_dist_shapes, ) # Test that the mixture has the same number of "samples" as w if w_samples.size != (mixed_samples.size // w.shape[-1]): raise ValueError( "Inconsistent number of samples from the " "mixture and mixture weights. Drew {} mixture " "weights elements, and {} samples from the " "mixture components.".format(w_samples.size, mixed_samples.size // w.shape[-1]) ) # Semiflatten the mixture to be able to zip it with w_samples w_samples = w_samples.flatten() mixed_samples = np.reshape(mixed_samples, (-1, w.shape[-1])) # Select the samples from the mixture samples = np.array([mixed[choice] for choice, mixed in zip(w_samples, mixed_samples)]) # Reshape the samples to the correct output shape if size is None: samples = np.reshape(samples, dist_shape) else: samples = np.reshape(samples, size + dist_shape) return samples