def __init__(self, centered_returns, name='stochastic_volatility', pretty_name='Stochastic Volatility'): """Construct the stochastic volatility model. Args: centered_returns: Float `Tensor` of shape `[num_timesteps]` giving the mean-adjusted return (change in asset price, minus the average change) observed at each step. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. """ with tf.name_scope(name): num_timesteps = ps.size0(centered_returns) if tf.is_tensor(num_timesteps): raise ValueError( 'Returns series length must be static, but saw ' 'shape {}.'.format(centered_returns.shape)) self._prior_dist = tfd.JointDistributionCoroutine( functools.partial(stochastic_volatility_prior_fn, num_timesteps=num_timesteps)) self._log_likelihood_fn = functools.partial( stochastic_volatility_log_likelihood_fn, centered_returns=centered_returns) def _ext_identity(params): res = collections.OrderedDict() res['persistence_of_volatility'] = params[0] res['mean_log_volatility'] = params[1] res['white_noise_shock_scale'] = params[2] res['log_volatility'] = tf.stack(params[3], axis=-1) return res sample_transformations = { 'identity': model.Model.SampleTransformation( fn=_ext_identity, pretty_name='Identity', ) } super(StochasticVolatility, self).__init__( default_event_space_bijector=(tfb.Sigmoid( -1., 1.), tfb.Identity(), tfb.Softplus()) + ((tfb.Identity(), ) * num_timesteps, ), event_shape=self._prior_dist.event_shape, dtype=self._prior_dist.dtype, name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )
def sample_from_banana(): def banana_model(): x0 = yield tfd.JointDistributionCoroutine.Root(tfd.Normal(0., 10.)) _ = yield tfd.Normal(0.03 * (tf.square(x0) - 100.), 1.) banana = tfd.JointDistributionCoroutine(banana_model) kernel = tfp.mcmc.NoUTurnSampler(banana.log_prob, step_size=0.35) trace_fn = lambda _, pkr: pkr.log_accept_ratio return tfp.mcmc.sample_chain(50, [0., 0.], kernel=kernel, trace_fn=trace_fn, seed=seed)[1]
def stochastic_volatility_prior_fn(num_timesteps): """Generative process for the stochastic volatility model.""" persistence_of_volatility = yield Root( tfb.Shift(-1.)(tfb.Scale(2.)(tfd.Beta( concentration1=20., concentration0=1.5, name='persistence_of_volatility')))) mean_log_volatility = yield Root( tfd.Cauchy(loc=0., scale=5., name='mean_log_volatility')) white_noise_shock_scale = yield Root( tfd.HalfCauchy(loc=0., scale=2., name='white_noise_shock_scale')) _ = yield tfd.JointDistributionCoroutine(functools.partial( autoregressive_series_fn, num_timesteps=num_timesteps, mean=mean_log_volatility, noise_scale=white_noise_shock_scale, persistence=persistence_of_volatility), name='log_volatility')
def __init__(self, train_features, train_labels, test_features=None, test_labels=None, name='logistic_regression', pretty_name='Logistic Regression'): """Construct the logistic regression model. Args: train_features: Floating-point `Tensor` with shape `[num_train_points, num_features]`. Training features. train_labels: Integer `Tensor` with shape `[num_train_points]`. Training labels. test_features: Floating-point `Tensor` with shape `[num_test_points, num_features]`. Testing features. Can be `None`, in which case test-related sample transformations are not computed. test_labels: Integer `Tensor` with shape `[num_test_points]`. Testing labels. Can be `None`, in which case test-related sample transformations are not computed. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. """ with tf.name_scope(name): train_features = _add_bias(train_features) train_labels = tf.convert_to_tensor(train_labels) num_features = int(train_features.shape[1]) root = tfd.JointDistributionCoroutine.Root zero = tf.zeros(num_features) one = tf.ones(num_features) def model_fn(features): weights = yield root(tfd.Independent(tfd.Normal(zero, one), 1)) logits = tf.einsum('nd,...d->...n', features, weights) yield tfd.Independent(tfd.Bernoulli(logits=logits), 1) train_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, features=train_features)) sample_transformations = { 'identity': bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: params, pretty_name='Identity', ) } if test_features is not None and test_labels is not None: test_features = _add_bias(test_features) test_labels = tf.convert_to_tensor(test_labels) test_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, features=test_features)) def _get_label_dist(weights): # TODO(b/150897904): The seed does nothing since the model is fully # conditioned. distributions, _ = test_joint_dist.sample_distributions( value=[weights, test_labels], seed=42) return distributions[-1] sample_transformations['test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=lambda weights: -( # pylint: disable=g-long-lambda _get_label_dist(weights).log_prob(test_labels)), pretty_name='Test NLL', )) sample_transformations['per_example_test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=lambda weights: -( # pylint: disable=g-long-lambda _get_label_dist(weights).distribution.log_prob( test_labels)), pretty_name='Per-example Test NLL', )) self._train_joint_dist = train_joint_dist self._train_labels = train_labels super(LogisticRegression, self).__init__( default_event_space_bijector=tfb.Identity(), event_shape=train_joint_dist.event_shape[0], dtype=train_joint_dist.dtype[0], name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )
def __init__(self, train_features, train_labels, test_features=None, test_labels=None, name='sparse_logistic_regression', pretty_name='Sparse Logistic Regression'): """Construct the sparse logistic regression model. Args: train_features: Floating-point `Tensor` with shape `[num_train_points, num_features]`. Training features. train_labels: Integer `Tensor` with shape `[num_train_points]`. Training labels. test_features: Floating-point `Tensor` with shape `[num_test_points, num_features]`. Testing features. Can be `None`, in which case test-related sample transformations are not computed. test_labels: Integer `Tensor` with shape `[num_test_points]`. Testing labels. Can be `None`, in which case test-related sample transformations are not computed. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. Raises: ValueError: If `test_features` and `test_labels` are either not both `None` or not both specified. """ with tf.name_scope(name): train_features = _add_bias(train_features) train_labels = tf.convert_to_tensor(train_labels) num_features = int(train_features.shape[1]) root = tfd.JointDistributionCoroutine.Root zero = tf.zeros(num_features) one = tf.ones(num_features) half = tf.fill([num_features], 0.5) def model_fn(features): """Model definition.""" unscaled_weights = yield root( tfd.Independent(tfd.Normal(zero, one), 1, name='unscaled_weights')) local_scales = yield root( tfd.Independent(tfd.Gamma(half, half), 1, name='local_scales')) global_scale = yield root( tfd.Gamma(0.5, 0.5, name='global_scale')) weights = unscaled_weights * local_scales * global_scale[ ..., tf.newaxis] logits = tf.einsum('nd,...d->...n', features, weights) yield tfd.Independent(tfd.Bernoulli(logits=logits), 1, name='labels') train_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, features=train_features)) sample_transformations = { 'identity': bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: params, pretty_name='Identity', ) } if (test_features is not None) != (test_labels is not None): raise ValueError( '`test_features` and `test_labels` must either both ' 'be `None` or both specified. Got: test_features={}, ' 'test_labels={}'.format(test_features, test_labels)) if test_features is not None and test_labels is not None: test_features = _add_bias(test_features) test_labels = tf.convert_to_tensor(test_labels) test_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, features=test_features)) def _get_label_dist(params): # TODO(b/150897904): The seed does nothing since the model is fully # conditioned. distributions, _ = test_joint_dist.sample_distributions( value=self._dict_to_tuple(params) + (test_labels, ), seed=42) return distributions[-1] sample_transformations['test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: -( # pylint: disable=g-long-lambda _get_label_dist(params).log_prob(test_labels)), pretty_name='Test NLL', )) sample_transformations['per_example_test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: -( # pylint: disable=g-long-lambda _get_label_dist(params).distribution.log_prob( test_labels)), pretty_name='Per-example Test NLL', )) self._train_joint_dist = train_joint_dist self._train_labels = train_labels super(SparseLogisticRegression, self).__init__( default_event_space_bijector=self._tuple_to_dict( (tfb.Identity(), tfb.Exp(), tfb.Exp())), event_shape=self._tuple_to_dict(train_joint_dist.event_shape[:-1]), dtype=self._tuple_to_dict(train_joint_dist.dtype[:-1]), name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )
def __init__( self, train_student_ids, train_question_ids, train_correct, test_student_ids=None, test_question_ids=None, test_correct=None, name='item_response_theory', pretty_name='Item-Response Theory', ): """Construct the item-response theory model. This models a set of students answering a set of questions, and being scored whether they get the question correct or not. Each student is associated with a scalar `student_ability`, and each question is associated with a scalar `question_difficulty`. Additionally, a scalar `mean_student_ability` is shared between all the students. This corresponds to the [1PL item-response theory](1) model. The data are encoded into three parallel arrays per set. I.e. `*_correct[i] == 1` means that student `*_student_ids[i]` answered question `*_question_ids[i]` correctly; `*_correct[i] == 0` means they didn't. Args: train_student_ids: integer `tensor` with shape `[num_train_points]`. training student ids, ranging from 0 to `num_students`. train_question_ids: integer `tensor` with shape `[num_train_points]`. training question ids, ranging from 0 to `num_questions`. train_correct: integer `tensor` with shape `[num_train_points]`. whether the student in the training set answered the question correctly, either 0 or 1. test_student_ids: Integer `Tensor` with shape `[num_test_points]`. Testing student ids, ranging from 0 to `num_students`. Can be `None`, in which case test-related sample transformations are not computed. test_question_ids: Integer `Tensor` with shape `[num_test_points]`. Testing question ids, ranging from 0 to `num_questions`. Can be `None`, in which case test-related sample transformations are not computed. test_correct: Integer `Tensor` with shape `[num_test_points]`. Whether the student in the testing set answered the question correctly, either 0 or 1. Can be `None`, in which case test-related sample transformations are not computed. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. Raises: ValueError: If `test_student_ids`, `test_question_ids` or `test_correct` are not either all `None` or are all specified. ValueError: If the parallel arrays are not all of the same size. #### References 1. https://en.wikipedia.org/wiki/Item_response_theory """ with tf.name_scope(name): test_data_present = ( e is not None for e in [test_student_ids, test_question_ids, test_correct]) self._have_test = all(test_data_present) if not self._have_test and any(test_data_present): raise ValueError( '`test_student_ids`, `test_question_ids` and ' '`test_correct` must either all be `None` or ' 'all be specified. Got: test_student_ids={}, ' 'test_question_ids={}, test_correct={}'.format( test_student_ids, test_question_ids, test_correct)) if not (train_student_ids.shape[0] == train_question_ids.shape[0] == train_correct.shape[0]): raise ValueError( '`train_student_ids`, `train_question_ids` and ' '`train_correct` must all have the same length. ' 'Got: {} {} {}'.format(train_student_ids.shape[0], train_question_ids.shape[0], train_correct.shape[0])) max_student_id = train_student_ids.max() max_question_id = train_question_ids.max() if self._have_test: max_student_id = max(max_student_id, test_student_ids.max()) max_question_id = max(max_question_id, test_question_ids.max()) self._num_students = max_student_id + 1 self._num_questions = max_question_id + 1 # TODO(siege): Make it an option to use a sparse encoding. The dense # encoding is only efficient when the dataset is not very sparse to begin # with. train_dense_y, train_y_mask = self._sparse_to_dense( train_student_ids, train_question_ids, train_correct, ) root = tfd.JointDistributionCoroutine.Root def model_fn(dense_y, y_mask): """Model definition.""" mean_student_ability = yield root( tfd.Normal(0.75, 1., name='mean_student_ability')) student_ability = yield root( tfd.Sample( tfd.Normal(0., 1.), dense_y.shape[0], name='student_ability', )) question_difficulty = yield root( tfd.Sample( tfd.Normal(0., 1.), dense_y.shape[1], name='question_difficulty', )) logits = (mean_student_ability[..., tf.newaxis, tf.newaxis] + student_ability[..., tf.newaxis] - question_difficulty[..., tf.newaxis, :]) # TODO(b/150949917): Use a more dedicated masking functionality. masked_logits = logits * y_mask - 1e10 * (1 - y_mask) yield tfd.Independent(tfd.Bernoulli(masked_logits), 2, name='correct') train_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, train_dense_y, train_y_mask)) dtype = self._tuple_to_dict(train_joint_dist.dtype[:-1]) sample_transformations = { 'identity': bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: params, pretty_name='Identity', dtype=dtype, ) } if self._have_test: if not (test_student_ids.shape[0] == test_question_ids.shape[0] == test_correct.shape[0]): raise ValueError( '`test_student_ids`, `test_question_ids` and ' '`test_correct` must all have the same length. ' 'Got: {} {} {}'.format(test_student_ids.shape[0], test_question_ids.shape[0], test_correct.shape[0])) test_dense_y, test_y_mask = self._sparse_to_dense( test_student_ids, test_question_ids, test_correct, ) test_joint_dist = tfd.JointDistributionCoroutine( functools.partial(model_fn, test_dense_y, test_y_mask)) def _get_label_dist(params): # TODO(b/150897904): The seed does nothing since the model is fully # conditioned. distributions, _ = test_joint_dist.sample_distributions( value=self._dict_to_tuple(params) + (test_dense_y, ), seed=42) return distributions[-1] sample_transformations['test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=lambda params: -( # pylint: disable=g-long-lambda _get_label_dist(params).log_prob(test_dense_y)), pretty_name='Test NLL', )) def _per_example_test_nll(params): """Computes per-example test NLL.""" dense_nll = _get_label_dist(params).distribution.log_prob( test_dense_y) return self._dense_to_sparse(test_student_ids, test_question_ids, dense_nll) sample_transformations['per_example_test_nll'] = ( bayesian_model.BayesianModel.SampleTransformation( fn=_per_example_test_nll, pretty_name='Per-example Test NLL', )) self._train_joint_dist = train_joint_dist self._train_correct = train_correct self._train_student_ids = train_student_ids self._train_question_ids = train_question_ids self._test_student_ids = test_student_ids self._test_question_ids = test_question_ids self._evidence_val = train_dense_y super(ItemResponseTheory, self).__init__( default_event_space_bijector=self._tuple_to_dict( (tfb.Identity(), tfb.Identity(), tfb.Identity())), event_shape=self._tuple_to_dict(train_joint_dist.event_shape[:-1]), dtype=dtype, name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )