Esempio n. 1
0
def brownian_motion(locs, innovation_noise, observation_noise):
  """Brownian Motion model.

  Args:
    locs: Array of loc parameters with np.nan value if loc is unobserved in
      shape (num_timesteps,)
    innovation_noise: Python `float`.
    observation_noise: Python `float`.

  Returns:
    model: `StanModel`.
  """

  code = """
  data {
    int<lower=0> num_timesteps;
    int<lower=0> num_observations;
    int<lower = 1, upper = num_timesteps> observation_indices[num_observations];
    vector[num_observations] observations;
    real<lower=0> innovation_noise;
    real<lower=0> observation_noise;
  }
  parameters {
    vector[num_timesteps] loc;
  }
  model {
    loc[1] ~ normal(0, innovation_noise);
    for (t in 2:num_timesteps){
      loc[t] ~ normal(loc[t-1], innovation_noise);
      }
    observations ~ normal(loc[observation_indices], observation_noise);
  }
  """

  stan_data = {
      'num_timesteps': len(locs),
      'num_observations': len(locs[np.isfinite(locs)]),
      'observation_indices': np.arange(1,
                                       len(locs) + 1)[np.isfinite(locs)],
      'observations': locs[np.isfinite(locs)],
      'innovation_noise': innovation_noise,
      'observation_noise': observation_noise
  }

  model = util.cached_stan_model(code)

  def _ext_identity(samples):
    """Extracts the values of all latent variables."""
    locs = util.get_columns(samples, r'^loc\.\d+$')
    return locs

  extract_fns = {'identity': _ext_identity}

  return stan_model.StanModel(
      extract_fns=extract_fns,
      sample_fn=util.make_sample_fn(model, data=stan_data),
  )
def stochastic_volatility(centered_returns):
  # pylint: disable=long-lines
  """Stochastic volatility model.

  This formulation is inspired by (a version in the Stan users' manual)[
  https://mc-stan.org/docs/2_21/stan-users-guide/stochastic-volatility-models.html].

  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.

  Returns:
    model: `StanModel`.
  """
  # pylint: enable=long-lines

  # This model is specified in 'noncentered' parameterization, in terms of
  # standardized residuals `log_volatilities_std`. We expect this form of the
  # model to mix more easily than a direct specification would. This makes
  # it valuable for obtaining ground truth, but caution should be used when
  # comparing performance of inference algorithms across parameterizations.
  code = """
  data {
    int<lower=0> num_timesteps;
    vector[num_timesteps] centered_returns;
  }
  parameters {
    real<lower=-1,upper=1> persistence;
    real mean_log_volatility;
    real<lower=0> white_noise_shock_scale;
    vector[num_timesteps] log_volatilities_std;
  }
  transformed parameters {
    vector[num_timesteps] log_volatilities = (
      log_volatilities_std * white_noise_shock_scale);
    log_volatilities[1] /= sqrt(1 - square(persistence));
    log_volatilities += mean_log_volatility;
    for (t in 2:num_timesteps)
      log_volatilities[t] += persistence * (
          log_volatilities[t - 1] - mean_log_volatility);
  }
  model {
    (persistence + 1) * 0.5 ~ beta(20, 1.5);
    white_noise_shock_scale ~ cauchy(0, 2);
    mean_log_volatility ~ cauchy(0, 5);
    log_volatilities_std ~ std_normal();

    centered_returns ~ normal(0, exp(log_volatilities / 2));
  }
  """

  stan_data = {
      'num_timesteps': len(centered_returns),
      'centered_returns': centered_returns
  }

  model = util.cached_stan_model(code)

  def _ext_identity(samples):
    """Extracts the values of all latent variables."""
    res = collections.OrderedDict()
    res['mean_log_volatility'] = util.get_columns(
        samples, r'^mean_log_volatility$')[:, 0]
    res['white_noise_shock_scale'] = util.get_columns(
        samples, r'^white_noise_shock_scale$')[:, 0]
    res['persistence_of_volatility'] = util.get_columns(
        samples, r'^persistence$')[:, 0]
    res['log_volatility'] = util.get_columns(
        samples, r'^log_volatilities\.\d+$',)
    return res

  extract_fns = {'identity': _ext_identity}

  return stan_model.StanModel(
      extract_fns=extract_fns,
      sample_fn=util.make_sample_fn(model, data=stan_data),
  )
def logistic_regression(
    train_features,
    train_labels,
    test_features=None,
    test_labels=None,
):
  """Bayesian logistic regression with a Gaussian prior.

  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.

  Returns:
    model: `StanModel`.
  """

  code = """
  data {
    int<lower=0> num_train_points;
    int<lower=0> num_test_points;
    int<lower=0> num_features;
    matrix[num_train_points,num_features] train_features;
    int<lower=0,upper=1> train_labels[num_train_points];
    matrix[num_test_points,num_features] test_features;
    int<lower=0,upper=1> test_labels[num_test_points];
  }
  parameters {
    vector[num_features] weights;
  }
  model {
    {
      vector[num_train_points] logits;
      logits = train_features * weights;

      weights ~ normal(0, 1);
      train_labels ~ bernoulli_logit(logits);
    }
  }
  generated quantities {
    real test_nll;
    real per_example_test_nll[num_test_points];
    {
      vector[num_test_points] logits;
      logits = test_features * weights;

      test_nll = -bernoulli_logit_lpmf(test_labels | logits);
      for (i in 1:num_test_points) {
        per_example_test_nll[i] = -bernoulli_logit_lpmf(
            test_labels[i] | logits[i]);
      }
    }
  }
  """
  have_test = test_features is not None
  train_features = _add_bias(train_features)
  if have_test:
    test_features = _add_bias(test_features)
  else:
    # cmdstanpy can't handle zero-sized arrays at the moment:
    # https://github.com/stan-dev/cmdstanpy/issues/203
    test_features = train_features[:1]
    test_labels = train_labels[:1]
  stan_data = {
      'num_train_points': train_features.shape[0],
      'num_test_points': test_features.shape[0],
      'num_features': train_features.shape[1],
      'train_features': train_features,
      'train_labels': train_labels,
      'test_features': test_features,
      'test_labels': test_labels,
  }

  model = util.cached_stan_model(code)

  def _ext_identity(samples):
    return util.get_columns(samples, r'^weights\.\d+$')

  def _ext_test_nll(samples):
    return util.get_columns(samples, r'^test_nll$')[:, 0]

  def _ext_per_example_test_nll(samples):
    return util.get_columns(samples, r'^per_example_test_nll\.\d+$')

  extract_fns = {'identity': _ext_identity}
  if have_test:
    extract_fns['test_nll'] = _ext_test_nll
    extract_fns['per_example_test_nll'] = _ext_per_example_test_nll

  return stan_model.StanModel(
      extract_fns=extract_fns,
      sample_fn=util.make_sample_fn(model, data=stan_data),
  )
def radon_contextual_effects(num_counties, train_log_uranium, train_floor,
                             train_county, train_floor_by_county,
                             train_log_radon):
    """Heirarchical model of measured radon concentration in homes.

  The Stan model is cut and pasted from:
  https://mc-stan.org/users/documentation/case-studies/radon.html#Correlations-among-levels

  Args:
    num_counties: `int`, number of counties represented in the data.
    train_log_uranium: Floating-point `Tensor` with shape
      `[num_train_points]`. Soil uranium measurements.
    train_floor: Integer `Tensor` with shape `[num_train_points]`. Floor of
      the house on which the measurement was taken.
    train_county: Integer `Tensor` with values in `range(0, num_counties)` of
      shape `[num_train_points]`. County in which the measurement was taken.
    train_floor_by_county: Floating-point `Tensor` with shape
      `[num_train_points]`. Average floor on which the measurement was taken
      for the county in which each house is located (the `Tensor` will have
      `num_counties` unique values). This represents the contextual effect.
    train_log_radon: Floating-point `Tensor` with shape `[num_train_points]`.
      Radon measurement for each house (the dependent variable in the model).
  Returns:
    model: `StanModel`.
  """

    code = """
  data {
    int<lower=0> num_counties;
    int<lower=0> num_train;
    int<lower=0,upper=num_counties-1> county[num_train];
    vector[num_train] log_uranium;
    vector[num_train] which_floor;
    vector[num_train] floor_by_county;
    vector[num_train] log_radon;
  }
  parameters {
    vector[num_counties] county_effect;
    vector[3] weight;
    real county_effect_mean;
    real<lower=0,upper=100> county_effect_scale;
    real<lower=0,upper=100> log_radon_scale;
  }
  transformed parameters {
    vector[num_train] log_radon_mean;

    for (i in 1:num_train)
      log_radon_mean[i] <- county_effect[county[i] + 1]
                           + log_uranium[i] * weight[1]
                           + which_floor[i] * weight[2]
                           + floor_by_county[i] * weight[3];
  }
  model {
    county_effect_mean ~ normal(0, 1);
    county_effect ~ normal(county_effect_mean, county_effect_scale);
    weight ~ normal(0, 1);
    log_radon ~ normal(log_radon_mean, log_radon_scale);
  }
  """

    stan_data = {
        'num_train': train_log_radon.shape[0],
        'num_counties': num_counties,
        'county': np.array(train_county),
        'log_uranium': np.array(train_log_uranium),
        'floor_by_county': np.array(train_floor_by_county),
        'which_floor':
        np.array(train_floor),  # `floor` conflicts with a Stan fn
        'log_radon': np.array(train_log_radon)
    }

    model = util.cached_stan_model(code)

    def _ext_identity(samples):
        """Extracts the values of all latent variables."""
        res = collections.OrderedDict()
        res['county_effect_mean'] = util.get_columns(
            samples, r'^county_effect_mean$')[:, 0]
        res['county_effect_scale'] = util.get_columns(
            samples, r'^county_effect_scale$')[:, 0]
        res['county_effect'] = util.get_columns(samples,
                                                r'^county_effect\.\d+$')
        res['weight'] = util.get_columns(samples, r'^weight\.\d+$')
        res['log_radon_scale'] = (util.get_columns(samples,
                                                   r'^log_radon_scale$')[:, 0])
        return res

    extract_fns = {'identity': _ext_identity}

    return stan_model.StanModel(
        extract_fns=extract_fns,
        sample_fn=util.make_sample_fn(model, data=stan_data),
    )
Esempio n. 5
0
def log_gaussian_cox_process(
    train_locations,
    train_extents,
    train_counts,
):
  """Log-Gaussian Cox Process model.

  Args:
    train_locations: Float `Tensor` with shape `[num_train_points, D]`. Training
      set locations where counts were measured.
    train_extents: Float `Tensor` with shape `[num_train_points]`. Training set
      location extents, must be positive.
    train_counts: Integer `Tensor` with shape `[num_train_points]`. Training set
      counts, must be positive.

  Returns:
    model: `StanModel`.
  """

  code = """
  data {
    int<lower=0> num_points;
    int<lower=0> num_features;
    vector[num_features] locations[num_points];
    real<lower=0> extents[num_points];
    int<lower=0> counts[num_points];
  }
  transformed data {
    vector[num_points] loc;
    real mean_log_intensity;
    {
      mean_log_intensity = 0;
      for (i in 1:num_points) {
        mean_log_intensity += (
          log(counts[i]) - log(extents[i])) / num_points;
      }
      for (i in 1:num_points) loc[i] = mean_log_intensity;  // otherwise nan!
    }
  }
  parameters {
    real<lower=0> amplitude;
    real<lower=0> length_scale;
    vector[num_points] log_intensity;
  }
  model {
    {
      matrix[num_points, num_points] L_K;
      matrix[num_points, num_points] K = gp_matern32_cov(
          locations, amplitude + .001, length_scale + .001);
      for (i in 1:num_points) K[i,i] += 1e-6;  // GP jitter
      L_K = cholesky_decompose(K);

      amplitude ~ lognormal(-1., .5);
      length_scale ~ lognormal(-1., 1.);
      log_intensity ~ multi_normal_cholesky(loc, L_K);
      for (i in 1:num_points) {
        counts[i] ~ poisson_log(
          log(extents[i]) + log_intensity[i]);
      }
    }
  }
  """

  num_points = train_locations.shape[0]
  num_features = train_locations.shape[1]
  stan_data = {
      'num_points': num_points,
      'num_features': num_features,
      'locations': train_locations,
      'extents': train_extents,
      'counts': train_counts,
  }

  model = util.cached_stan_model(code)

  def _ext_identity(samples):
    """Extract all the parameters."""
    res = collections.OrderedDict()
    res['amplitude'] = util.get_columns(
        samples,
        r'^amplitude$',
    )[:, 0]
    res['length_scale'] = util.get_columns(
        samples,
        r'^length_scale$',
    )[:, 0]
    res['log_intensity'] = util.get_columns(
        samples,
        r'^log_intensity\.\d+$',
    )
    return res

  extract_fns = {'identity': _ext_identity}

  return stan_model.StanModel(
      extract_fns=extract_fns,
      sample_fn=util.make_sample_fn(model, data=stan_data),
  )
def item_response_theory(
    train_student_ids,
    train_question_ids,
    train_correct,
    test_student_ids=None,
    test_question_ids=None,
    test_correct=None,
):
    """One-parameter logistic item-response theory (IRT) model.

  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.

  Returns:
    target: `StanModel`.
  """

    code = """
  data {
    int<lower=0> num_students;
    int<lower=0> num_questions;
    int<lower=0> num_train_pairs;
    int<lower=0> num_test_pairs;
    int<lower=1,upper=num_students> train_student_ids[num_train_pairs];
    int<lower=1,upper=num_questions> train_question_ids[num_train_pairs];
    int<lower=0,upper=1> train_responses[num_train_pairs];
    int<lower=1,upper=num_students> test_student_ids[num_test_pairs];
    int<lower=1,upper=num_questions> test_question_ids[num_test_pairs];
    int<lower=0,upper=1> test_responses[num_test_pairs];
  }
  parameters {
    real mean_student_ability;
    vector[num_students] student_ability;
    vector[num_questions] question_difficulty;
  }
  model {
    {
      mean_student_ability ~ normal(0.75, 1);
      student_ability ~ normal(0, 1);
      question_difficulty ~ normal(0, 1);

      for (i in 1:num_train_pairs) {
        real pair_logit;
        pair_logit = (
            mean_student_ability + student_ability[train_student_ids[i]] -
            question_difficulty[train_question_ids[i]]
        );
        train_responses[i] ~ bernoulli_logit(pair_logit);
      }
    }
  }
  generated quantities {
    real test_nll = 0.;
    real per_example_test_nll[num_test_pairs];
    {
      for (i in 1:num_test_pairs) {
        real pair_logit;
        pair_logit = (
            mean_student_ability + student_ability[test_student_ids[i]] -
            question_difficulty[test_question_ids[i]]
        );
        per_example_test_nll[i] = -bernoulli_logit_lpmf(test_responses[i] | pair_logit);
      }
      test_nll = sum(per_example_test_nll);
    }
  }
  """

    have_test = test_student_ids is not None
    # cmdstanpy can't handle zero-sized arrays at the moment:
    # https://github.com/stan-dev/cmdstanpy/issues/203
    if not have_test:
        test_student_ids = train_student_ids[:1]
        test_question_ids = train_question_ids[:1]
        test_correct = train_correct[:1]
    stan_data = {
        'num_train_pairs':
        train_student_ids.shape[0],
        'num_test_pairs':
        test_student_ids.shape[0],
        'num_students':
        max(int(train_student_ids.max()), int(test_student_ids.max())) + 1,
        'num_questions':
        max(int(train_question_ids.max()), int(test_question_ids.max())) + 1,
        'train_student_ids':
        train_student_ids + 1,  # N.B. Stan arrays are 1-indexed.
        'train_question_ids':
        train_question_ids + 1,
        'train_responses':
        train_correct,
        'test_student_ids':
        test_student_ids + 1,
        'test_question_ids':
        test_question_ids + 1,
        'test_responses':
        test_correct,
    }

    model = util.cached_stan_model(code)

    def _ext_identity(samples):
        """Extracts all the parameters."""
        res = collections.OrderedDict()
        res['mean_student_ability'] = util.get_columns(
            samples,
            r'^mean_student_ability$',
        )[:, 0]
        res['student_ability'] = util.get_columns(
            samples,
            r'^student_ability\.\d+$',
        )
        res['question_difficulty'] = util.get_columns(
            samples,
            r'^question_difficulty\.\d+$',
        )
        return res

    def _ext_test_nll(samples):
        return util.get_columns(samples, r'^test_nll$')[:, 0]

    def _ext_per_example_test_nll(samples):
        return util.get_columns(samples, r'^per_example_test_nll\.\d+$')

    extract_fns = {'identity': _ext_identity}
    if have_test:
        extract_fns['test_nll'] = _ext_test_nll
        extract_fns['per_example_test_nll'] = _ext_per_example_test_nll

    return stan_model.StanModel(
        extract_fns=extract_fns,
        sample_fn=util.make_sample_fn(model, data=stan_data),
    )