class_docstr: Hierarchical radon model with contextual effects.
  
  Radon is a radioactive gas that enters homes through contact points with the
  ground. It is a carcinogen that is the primary cause of lung cancer in
  non-smokers. Radon levels vary greatly from household to household.
  
  The EPA did a study of radon levels in 80,000 houses. Two important predictors
  are: 1. Measurement in the basement or the first floor (radon higher in
  basements) 2. County uranium level (positive correlation with radon levels)
  We will build a hierarchical model of radon measurements in houses, in which
  the hierarchy is households within each county. We will incorporate a
  contextual effect corresponding to the mean floor at which the measurement
  was taken, by county.
  
  ```none
  
  county_effect_mean ~ Normal(loc=0, scale=1)
  county_effect_scale ~ Uniform(low=0, high=100)
  for i in range(num_counties):
    county_effect[i] ~ Normal(loc=county_effect_mean,
                              scale=county_effect_scale)
  for j in range(3):
    weight[j] ~ Normal(loc=0, scale=1)
  log_radon_scale ~ Uniform(low=0, high=100)
  
  for k in range(num_houses):
    log_radon[k] ~ Normal(
        loc=log_uranium * weight[1]           # effect of soil uranium
            + floor * weight[2]               # effect of floor
            + floor_by_county * weight[3]     # effect of mean floor by county
            + county_effect[county[k]],       # effect of county
        scale=log_radon_scale)
  ```
  
  This model is based on an example from [1] and is the same as the Stan model
  at
  <https://mc-stan.org/users/documentation/case-studies/radon.html
  #Correlations-among-levels>.
  
  Initializing this model with a `halfnormal` value for `prior_scale` will
  construct a modified version of this model in which the scales for the prior
  are constructed with a `HalfNormal` distribution instead of a `Uniform`
  distribution.
  
  ```none
  
  county_effect_scale ~ HalfNormal(scale=1.)
  log_radon_scale ~ HalfNormal(scale=1.)
  
  ```
  
  #### References
  
  [1] Gelman, A., & Hill, J. (2007). Data Analysis Using Regression and
      Multilevel/Hierarchical Models (1st ed.). Cambridge University Press.
  [2] Stan Development Team. 2018. Stan Modeling Language Users Guide and
      Reference Manual, Version 2.18.0. http://mc-stan.org
class_docstr_tok: ["Hierarchical","radon","model","with","contextual","effects",".","Radon","is","a","radioactive","gas","that","enters","homes","through","contact","points","with","the","ground",".","It","is","a","carcinogen","that","is","the","primary","cause","of","lung","cancer","in","non-smokers",".","Radon","levels","vary","greatly","from","household","to","household",".","The","EPA","did","a","study","of","radon","levels","in","80,000","houses",".","Two","important","predictors","are",":","1",".","Measurement","in","the","basement","or","the","first","floor","(","radon","higher","in","basements",")","2",".","County","uranium","level","(","positive","correlation","with","radon","levels",")","We","will","build","a","hierarchical","model","of","radon","measurements","in","houses",",","in","which","the","hierarchy","is","households","within","each","county",".","We","will","incorporate","a","contextual","effect","corresponding","to","the","mean","floor","at","which","the","measurement","was","taken",",","by","county",".","``","`","none","county_effect_mean","~","Normal","(","loc=0",",","scale=1",")","county_effect_scale","~","Uniform","(","low=0",",","high=100",")","for","i","in","range","(","num_counties",")",":","county_effect","[","i","]","~","Normal","(","loc=county_effect_mean",",","scale=county_effect_scale",")","for","j","in","range","(","3",")",":","weight","[","j","]","~","Normal","(","loc=0",",","scale=1",")","log_radon_scale","~","Uniform","(","low=0",",","high=100",")","for","k","in","range","(","num_houses",")",":","log_radon","[","k","]","~","Normal","(","loc=log_uranium","*","weight","[","1","]","#","effect","of","soil","uranium","+","floor","*","weight","[","2","]","#","effect","of","floor","+","floor_by_county","*","weight","[","3","]","#","effect","of","mean","floor","by","county","+","county_effect","[","county","[","k","]","]",",","#","effect","of","county","scale=log_radon_scale",")","``","`","This","model","is","based","on","an","example","from","[","1","]","and","is","the","same","as","the","Stan","model","at","<","https",":","\/\/mc-stan.org\/users\/documentation\/case-studies\/radon.html","#","Correlations-among-levels",">",".","Initializing","this","model","with","a","`","halfnormal","`","value","for","`","prior_scale","`","will","construct","a","modified","version","of","this","model","in","which","the","scales","for","the","prior","are","constructed","with","a","`","HalfNormal","`","distribution","instead","of","a","`","Uniform","`","distribution",".","``","`","none","county_effect_scale","~","HalfNormal","(","scale=1",".",")","log_radon_scale","~","HalfNormal","(","scale=1",".",")","``","`","#","#","#","#","References","[","1","]","Gelman",",","A.",",","&","Hill",",","J",".","(","2007",")",".","Data","Analysis","Using","Regression","and","Multilevel\/Hierarchical","Models","(","1st","ed.",")",".","Cambridge","University","Press",".","[","2","]","Stan","Development","Team",".","2018",".","Stan","Modeling","Language","Users","Guide","and","Reference","Manual",",","Version","2.18.0.","http",":","\/\/mc-stan.org"]
docstr: Construct the hierarchical radon model with contextual effects.
  
  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).
    test_log_uranium: Floating-point `Tensor` with shape `[num_test_points]`.
      Soil uranium measurements for the test set. Can be `None`, in which case
      test-related sample transformations are not computed.
    test_floor: Integer `Tensor` with shape `[num_test_points]`. Floor of the
      house on which the measurement was taken. Can be `None`, in which case
      test-related sample transformations are not computed.
    test_county: Integer `Tensor` with values in `range(0, num_counties)` of
      shape `[num_test_points]`. County in which the measurement was taken.
      Can be `None`, in which case test-related sample transformations are not
      computed.
    test_floor_by_county: Floating-point `Tensor` with shape
      `[num_test_points]`. Average floor on which the measurement was taken
      (calculated from the training set) for the county in which each house is
      located (the `Tensor` will have `num_counties` unique values). This
      represents the contextual effect. Can be `None`, in which case
      test-related sample transformations are not computed.
    test_log_radon: Floating-point `Tensor` with shape `[num_test_points]`.
      Radon measurement for each house (the dependent variable in the model).
      Can be `None`, in which case test-related sample transformations are not
      computed.
    prior_scale: String value. The default `uniform` value constructs the
      prior distribution's `county_effect_scale` and `log_radon_scale` with a
      `Uniform` distribution as in the original Stan model. A `halfnormal`
      value constructs the prior distribution's `county_effect_scale` and
      `log_radon_scale` with a `HalfNormal` distribution.
    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 any but not all of `test_*` inputs are None.
docstr_tok: ["Construct","the","hierarchical","radon","model","with","contextual","effects",".","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",")",".","test_log_uranium",":","Floating-point","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Soil","uranium","measurements","for","the","test","set",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_floor",":","Integer","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Floor","of","the","house","on","which","the","measurement","was","taken",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_county",":","Integer","`","Tensor","`","with","values","in","`","range","(","0",",","num_counties",")","`","of","shape","`","[","num_test_points","]","`",".","County","in","which","the","measurement","was","taken",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_floor_by_county",":","Floating-point","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Average","floor","on","which","the","measurement","was","taken","(","calculated","from","the","training","set",")","for","the","county","in","which","each","house","is","located","(","the","`","Tensor","`","will","have","`","num_counties","`","unique","values",")",".","This","represents","the","contextual","effect",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_log_radon",":","Floating-point","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Radon","measurement","for","each","house","(","the","dependent","variable","in","the","model",")",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","prior_scale",":","String","value",".","The","default","`","uniform","`","value","constructs","the","prior","distribution","'s","`","county_effect_scale","`","and","`","log_radon_scale","`","with","a","`","Uniform","`","distribution","as","in","the","original","Stan","model",".","A","`","halfnormal","`","value","constructs","the","prior","distribution","'s","`","county_effect_scale","`","and","`","log_radon_scale","`","with","a","`","HalfNormal","`","distribution",".","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","any","but","not","all","of","`","test_","*","`","inputs","are","None","."]
code: def __init__(
      self,
      num_counties,
      train_log_uranium,
      train_floor,
      train_county,
      train_floor_by_county,
      train_log_radon,
      test_log_uranium=None,
      test_floor=None,
      test_county=None,
      test_floor_by_county=None,
      test_log_radon=None,
      prior_scale="uniform",
      name="radon_contextual_effects",
      pretty_name="Radon Contextual Effects",
  ):
      """Construct the hierarchical radon model with contextual effects.
  
      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).
        test_log_uranium: Floating-point `Tensor` with shape `[num_test_points]`.
          Soil uranium measurements for the test set. Can be `None`, in which case
          test-related sample transformations are not computed.
        test_floor: Integer `Tensor` with shape `[num_test_points]`. Floor of the
          house on which the measurement was taken. Can be `None`, in which case
          test-related sample transformations are not computed.
        test_county: Integer `Tensor` with values in `range(0, num_counties)` of
          shape `[num_test_points]`. County in which the measurement was taken.
          Can be `None`, in which case test-related sample transformations are not
          computed.
        test_floor_by_county: Floating-point `Tensor` with shape
          `[num_test_points]`. Average floor on which the measurement was taken
          (calculated from the training set) for the county in which each house is
          located (the `Tensor` will have `num_counties` unique values). This
          represents the contextual effect. Can be `None`, in which case
          test-related sample transformations are not computed.
        test_log_radon: Floating-point `Tensor` with shape `[num_test_points]`.
          Radon measurement for each house (the dependent variable in the model).
          Can be `None`, in which case test-related sample transformations are not
          computed.
        prior_scale: String value. The default `uniform` value constructs the
          prior distribution's `county_effect_scale` and `log_radon_scale` with a
          `Uniform` distribution as in the original Stan model. A `halfnormal`
          value constructs the prior distribution's `county_effect_scale` and
          `log_radon_scale` with a `HalfNormal` distribution.
        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 any but not all of `test_*` inputs are None.
      """
      with tf.name_scope(name):
          test_data = (
              test_log_uranium,
              test_floor,
              test_county,
              test_floor_by_county,
              test_log_radon,
          )
          test_data_present = (d is not None for d in test_data)
          self._have_test = all(test_data_present)
          if not self._have_test and any(test_data_present):
              raise ValueError(
                  "Test set values must all be specified or all `None`. Got:"
                  "`test_log_uranium`={}, `test_floor`={}, `test_county`={},"
                  "`test_floor_by_county`={}, `test_log_radon`={}`".format(
                      *test_data
                  )
              )
  
          dtype = train_log_radon.dtype
          self._prior_dist = make_radon_prior(
              num_counties, dtype=dtype, prior_scale=prior_scale
          )
          self._train_log_likelihood_fn = functools.partial(
              radon_log_likelihood_fn,
              log_uranium=train_log_uranium,
              floor=train_floor,
              county=train_county,
              floor_by_county=train_floor_by_county,
              log_radon=train_log_radon,
          )
  
          sample_transformations = {
              "identity": model.Model.SampleTransformation(
                  fn=lambda params: params,
                  pretty_name="Identity",
                  dtype=self._prior_dist.dtype,
              )
          }
  
          if self._have_test:
              test_log_likelihood_fn = functools.partial(
                  radon_log_likelihood_fn,
                  log_uranium=test_log_uranium,
                  floor=test_floor,
                  county=test_county,
                  floor_by_county=test_floor_by_county,
                  log_radon=test_log_radon,
              )
  
              sample_transformations[
                  "test_nll"
              ] = model.Model.SampleTransformation(
                  fn=test_log_likelihood_fn,
                  pretty_name="Test NLL",
              )
              sample_transformations[
                  "per_example_test_nll"
              ] = model.Model.SampleTransformation(
                  fn=functools.partial(
                      test_log_likelihood_fn, reduce_sum=False
                  ),
                  pretty_name="Per-example Test NLL",
              )
  
          self._train_log_uranium = train_log_uranium
          self._train_floor = train_floor
          self._train_county = train_county
          self._train_floor_by_county = train_floor_by_county
          self._test_log_uranium = test_log_uranium
          self._test_floor = test_floor
          self._test_county = test_county
          self._test_floor_by_county = test_floor_by_county
          self._num_counties = num_counties
          self._prior_scale = prior_scale
  
      super(RadonContextualEffects, self).__init__(
          default_event_space_bijector={
              "county_effect_mean": tfb.Identity(),
              "county_effect_scale": tfb.Sigmoid(
                  low=tf.zeros([], dtype=dtype), high=100.0
              )
              if self._prior_scale == "uniform"
              else tfb.Softplus(),
              "county_effect": tfb.Identity(),
              "weight": tfb.Identity(),
              "log_radon_scale": tfb.Sigmoid(
                  low=tf.zeros([], dtype=dtype), high=100.0
              )
              if self._prior_scale == "uniform"
              else tfb.Softplus(),
          },
          event_shape=self._prior_dist.event_shape,
          dtype=self._prior_dist.dtype,
          name=name,
          pretty_name=pretty_name,
          sample_transformations=sample_transformations,
      )
  
imports: import functools
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import dtype_util
  from inference_gym.internal import data
  from inference_gym.targets import bayesian_model
  from inference_gym.targets import model
  from inference_gym.targets.ground_truth import radon_contextual_effects_indiana
  from inference_gym.targets.ground_truth import radon_contextual_effects_indiana_halfnormal
  from inference_gym.targets.ground_truth import radon_contextual_effects_minnesota
  from inference_gym.targets.ground_truth import radon_contextual_effects_minnesota_halfnormal
label_desc: Use image node_id 1 to create a new RadonContextualEffects object from inherited base classes: bayesian_model with example: obj = RadonContextualEffects(num_counties, train_log_uranium, train_floor, train_county, train_floor_by_county, train_log_radon, test_log_uranium, test_floor, test_county, test_floor_by_county, test_log_radon, prior_scale, name, pretty_name)