stsb3.core

AddBlock

Represents the result of adding two blocks toogether.

.model(...) is computed as x ~ p(x); y ~ p(y); x + y

Args:

• left (Block)
• right (Block)

_maybe_add_blocks

def _maybe_add_blocks(self, *args):

Adds parameters to prec and succ if they subclass Block.

Args:

• args: iterable of (name, parameter, bound)

_maybe_remove_blocks

None

_model

None

_transform

Defines a transform from a string argument.

Currently the following string arguments are supported:

• exp
• log
• logit
• invlogit
• tanh
• arctanh
• invlogit
• logit
• floor
• sin
• cos
• softplus
• diff (lowers time dimemsion by 1)
• logdiff (lowers time dimension by 1)

The resulting transform will be added to the transform stack iff it is not already at the top of the stack.

Args:

• arg (str): one of the above strings corresponding to function

Returns:

self (stsb.Block)

arctanh

x -> arctanh(x), i.e. x -> 0.5 log ((1 + x) / (1 - x))

clear_cache

Clears the block cache.

This method does not alter the cache mode.

cos

x -> cos x

diff

x -> x[1:] - x[:-1]

Note that this lowers the time dimension from T to T - 1.

exp

x -> exp(x)

floor

x -> x - [[x]], where [[.]] is the fractional part operator

invlogit

x -> 1 / (1 + exp(-x))

log

x -> log x

Block paths must be positive for valid output.

logdiff

x -> log x[1:] - log x[:-1]

Note that this lowers the time dimension from T to T - 1.

logit

x -> log(x / (1 - x))

model

def model(self, *args, **kwargs):

Draws a batch of samples from the block.

Args:

• args: optional positional arguments
• kwargs: optional keyword arguments

Returns:

• draws (torch.tensor) sampled values from the block

prec

Returns the predecessor nodes of self in the (implicit) compute graph

Returns:

_prec (list): list of predecessor nodes

sin

x -> sin x

softplus

x -> log(1 + exp(x))

succ

Returns the successor nodes of self in the (implicit) compute graph

Returns:

_succ (list): list of successor nodes

tanh

x -> tanh(x), i.e. x -> (exp(x) - exp(-x)) / (exp(x) + exp(-x))

Block

Base class that all STS blocks should subclass.

Defines a number of useful modeling constructs and methods, such as deterministic transformations.

_maybe_add_blocks

def _maybe_add_blocks(self, *args):

Adds parameters to prec and succ if they subclass Block.

Args:

• args: iterable of (name, parameter, bound)

_maybe_remove_blocks

None

_model

None

_transform

Defines a transform from a string argument.

Currently the following string arguments are supported:

• exp
• log
• logit
• invlogit
• tanh
• arctanh
• invlogit
• logit
• floor
• sin
• cos
• softplus
• diff (lowers time dimemsion by 1)
• logdiff (lowers time dimension by 1)

The resulting transform will be added to the transform stack iff it is not already at the top of the stack.

Args:

• arg (str): one of the above strings corresponding to function

Returns:

self (stsb.Block)

arctanh

x -> arctanh(x), i.e. x -> 0.5 log ((1 + x) / (1 - x))

clear_cache

Clears the block cache.

This method does not alter the cache mode.

cos

x -> cos x

diff

x -> x[1:] - x[:-1]

Note that this lowers the time dimension from T to T - 1.

exp

x -> exp(x)

floor

x -> x - [[x]], where [[.]] is the fractional part operator

invlogit

x -> 1 / (1 + exp(-x))

log

x -> log x

Block paths must be positive for valid output.

logdiff

x -> log x[1:] - log x[:-1]

Note that this lowers the time dimension from T to T - 1.

logit

x -> log(x / (1 - x))

model

def model(self, *args, **kwargs):

Draws a batch of samples from the block.

Args:

• args: optional positional arguments
• kwargs: optional keyword arguments

Returns:

• draws (torch.tensor) sampled values from the block

prec

Returns the predecessor nodes of self in the (implicit) compute graph

Returns:

_prec (list): list of predecessor nodes

sin

x -> sin x

softplus

x -> log(1 + exp(x))

succ

Returns the successor nodes of self in the (implicit) compute graph

Returns:

_succ (list): list of successor nodes

tanh

x -> tanh(x), i.e. x -> (exp(x) - exp(-x)) / (exp(x) + exp(-x))

LikelihoodCollection

def __init__(
    self,
    dgp,
    name2lik,
    name=None,
    t0=None,
    t1=None,
):

A NoiseBlock subclass that facilitates running inference against multiple NoiseBlock objects simultaneously.

Args:

• dgp (callable): a callable that, when called, calls all NoiseBlock objects in the correct order. This callable might be created with a call to core._collect.
• name2lik (dict): structure {str: Block}, name of the block mapping to the block object.

See NoiseBlock for definitions of other parameters.

_fit_autoguide

None

_maybe_add_blocks

def _maybe_add_blocks(self, *args):

Adds parameters to prec and succ if they subclass Block.

Args:

• args: iterable of (name, parameter, bound)

_maybe_remove_blocks

None

_model

None

_transform

Defines a transform from a string argument.

Currently the following string arguments are supported:

• exp
• log
• logit
• invlogit
• tanh
• arctanh
• invlogit
• logit
• floor
• sin
• cos
• softplus
• diff (lowers time dimemsion by 1)
• logdiff (lowers time dimension by 1)

The resulting transform will be added to the transform stack iff it is not already at the top of the stack.

Args:

• arg (str): one of the above strings corresponding to function

Returns:

self (stsb.Block)

arctanh

x -> arctanh(x), i.e. x -> 0.5 log ((1 + x) / (1 - x))

clear_cache

Clears the block cache.

This method does not alter the cache mode.

cos

x -> cos x

diff

x -> x[1:] - x[:-1]

Note that this lowers the time dimension from T to T - 1.

exp

x -> exp(x)

fit

def fit(self, method="nf_block_ar", method_kwargs=dict(), verbosity=0.01):

Fits a guide (variational posterior) to the model.

Wraps multiple Pyro implementations of variational inference. To minimize noise in the estimation you should follow the Pyro guidelines about marginalizing out discrete latent rvs, etc.

Args:

• method (str): one of “advi”, “low_rank”, or “nf_block_ar”.
  • "advi": fits a diagonal normal distribution in unconstrained latent space
  • "low_rank": fits a low-rank multivariate normal in unconstrained latent space. Unlike the diagonal normal, this guide can capture some nonlocal dependence in latent rvs.
  • "nf_block_ar": fits a normalizing flow block autoregressive neural density estimator in unconstrained latent space. This method uses two stacked block autoregressive NNs. See the Pyro docs for more details about this.
• method_kwargs (dict): optional keyword arguments to pass to Pyro’s inference capabilities. If no keyword arguments are specified, sane defaults will be passed instead. Some arguments could include:
  • "niter": number of iterations to run optimization (default 1000)
  • "lr": the learning rate (default 0.01)
  • "loss": the loss function to use (default "Trace_ELBO")
  • "optim": the optimizer to use (default "AdamW")
• verbosity (float): status messages are printed every int(1.0 / verbosity) iterations

floor

x -> x - [[x]], where [[.]] is the fractional part operator

invlogit

x -> 1 / (1 + exp(-x))

log

x -> log x

Block paths must be positive for valid output.

logdiff

x -> log x[1:] - log x[:-1]

Note that this lowers the time dimension from T to T - 1.

logit

x -> log(x / (1 - x))

model

def model(self, *args, **kwargs):

Draws a batch of samples from the block.

Args:

• args: optional positional arguments
• kwargs: optional keyword arguments

Returns:

• draws (torch.tensor) sampled values from the block

posterior_predictive

def posterior_predictive(
    self,
    nsamples=1,
):

Draws from the posterior predictive distribution of the graph with self as the root

Args:

• nsamples (int): number of samples to draw

Returns:

samples (torch.tensor)

prec

Returns the predecessor nodes of self in the (implicit) compute graph

Returns:

_prec (list): list of predecessor nodes

prior_predictive

def prior_predictive(
    self,
    nsamples=1,
):

Draws from the prior predictive distribution of the graph with self as the root

Args:

• nsamples (int): number of samples to draw

Returns:

samples (torch.tensor)

sample

def sample(
    self,
    nsamples=100,
    thin=0.1,
    burnin=500,
):

Sample from the model’s posterior using the Pyro implementation of the No-U Turn Sampler

This could take a very long time for long time series. It is recommended to use .fit(...) instead.

Args:

• nsamples (int): number of desired samples after burn in and thinning
• thin (float): every int(1.0 / thin) sample is kept
• burnin (int): samples[burnin:] are kept

sin

x -> sin x

softplus

x -> log(1 + exp(x))

succ

Returns the successor nodes of self in the (implicit) compute graph

Returns:

_succ (list): list of successor nodes

tanh

x -> tanh(x), i.e. x -> (exp(x) - exp(-x)) / (exp(x) + exp(-x))

MultiplyBlock

Represents the result of multiplying two blocks toogether.

.model(...) is computed as x ~ p(x); y ~ p(y); x * y

Args:

• left (Block)
• right (Block)

_maybe_add_blocks

def _maybe_add_blocks(self, *args):

Adds parameters to prec and succ if they subclass Block.

Args:

• args: iterable of (name, parameter, bound)

_maybe_remove_blocks

None

_model

None

_transform

Defines a transform from a string argument.

Currently the following string arguments are supported:

• exp
• log
• logit
• invlogit
• tanh
• arctanh
• invlogit
• logit
• floor
• sin
• cos
• softplus
• diff (lowers time dimemsion by 1)
• logdiff (lowers time dimension by 1)

The resulting transform will be added to the transform stack iff it is not already at the top of the stack.

Args:

• arg (str): one of the above strings corresponding to function

Returns:

self (stsb.Block)

arctanh

x -> arctanh(x), i.e. x -> 0.5 log ((1 + x) / (1 - x))

clear_cache

Clears the block cache.

This method does not alter the cache mode.

cos

x -> cos x

diff

x -> x[1:] - x[:-1]

Note that this lowers the time dimension from T to T - 1.

exp

x -> exp(x)

floor

x -> x - [[x]], where [[.]] is the fractional part operator

invlogit

x -> 1 / (1 + exp(-x))

log

x -> log x

Block paths must be positive for valid output.

logdiff

x -> log x[1:] - log x[:-1]

Note that this lowers the time dimension from T to T - 1.

logit

x -> log(x / (1 - x))

model

def model(self, *args, **kwargs):

Draws a batch of samples from the block.

Args:

• args: optional positional arguments
• kwargs: optional keyword arguments

Returns:

• draws (torch.tensor) sampled values from the block

prec

Returns the predecessor nodes of self in the (implicit) compute graph

Returns:

_prec (list): list of predecessor nodes

sin

x -> sin x

softplus

x -> log(1 + exp(x))

succ

Returns the successor nodes of self in the (implicit) compute graph

Returns:

_succ (list): list of successor nodes

tanh

x -> tanh(x), i.e. x -> (exp(x) - exp(-x)) / (exp(x) + exp(-x))

NoiseBlock

Base class for all likelihood function-type blocks

Implements a number of inference wrappers to Pyro implementations.

Args:

• dgp (Block): the latent data-generating process for which self serves as a likelihood function
• data (None || torch.tensor): the observed data. If data is None, then using the noise block is equivalent to drawing from the prior of a state space model
• name (None || str): a unique name of the block. If name is None, a unique name will be automatically generated

For other argument documentation, see Block

_fit_autoguide

None

_maybe_add_blocks

def _maybe_add_blocks(self, *args):

Adds parameters to prec and succ if they subclass Block.

Args:

• args: iterable of (name, parameter, bound)

_maybe_remove_blocks

None

_model

None

_transform

Defines a transform from a string argument.

Currently the following string arguments are supported:

• exp
• log
• logit
• invlogit
• tanh
• arctanh
• invlogit
• logit
• floor
• sin
• cos
• softplus
• diff (lowers time dimemsion by 1)
• logdiff (lowers time dimension by 1)

The resulting transform will be added to the transform stack iff it is not already at the top of the stack.

Args:

• arg (str): one of the above strings corresponding to function

Returns:

self (stsb.Block)

arctanh

x -> arctanh(x), i.e. x -> 0.5 log ((1 + x) / (1 - x))

clear_cache

Clears the block cache.

This method does not alter the cache mode.

cos

x -> cos x

diff

x -> x[1:] - x[:-1]

Note that this lowers the time dimension from T to T - 1.

exp

x -> exp(x)

fit

def fit(self, method="nf_block_ar", method_kwargs=dict(), verbosity=0.01):

Fits a guide (variational posterior) to the model.

Wraps multiple Pyro implementations of variational inference. To minimize noise in the estimation you should follow the Pyro guidelines about marginalizing out discrete latent rvs, etc.

Args:

• method (str): one of “advi”, “low_rank”, or “nf_block_ar”.
  • "advi": fits a diagonal normal distribution in unconstrained latent space
  • "low_rank": fits a low-rank multivariate normal in unconstrained latent space. Unlike the diagonal normal, this guide can capture some nonlocal dependence in latent rvs.
  • "nf_block_ar": fits a normalizing flow block autoregressive neural density estimator in unconstrained latent space. This method uses two stacked block autoregressive NNs. See the Pyro docs for more details about this.
• method_kwargs (dict): optional keyword arguments to pass to Pyro’s inference capabilities. If no keyword arguments are specified, sane defaults will be passed instead. Some arguments could include:
  • "niter": number of iterations to run optimization (default 1000)
  • "lr": the learning rate (default 0.01)
  • "loss": the loss function to use (default "Trace_ELBO")
  • "optim": the optimizer to use (default "AdamW")
• verbosity (float): status messages are printed every int(1.0 / verbosity) iterations

floor

x -> x - [[x]], where [[.]] is the fractional part operator

invlogit

x -> 1 / (1 + exp(-x))

log

x -> log x

Block paths must be positive for valid output.

logdiff

x -> log x[1:] - log x[:-1]

Note that this lowers the time dimension from T to T - 1.

logit

x -> log(x / (1 - x))

model

def model(self, *args, **kwargs):

Draws a batch of samples from the block.

Args:

• args: optional positional arguments
• kwargs: optional keyword arguments

Returns:

• draws (torch.tensor) sampled values from the block

posterior_predictive

def posterior_predictive(
    self,
    nsamples=1,
):

Draws from the posterior predictive distribution of the graph with self as the root

Args:

• nsamples (int): number of samples to draw

Returns:

samples (torch.tensor)

prec

Returns the predecessor nodes of self in the (implicit) compute graph

Returns:

_prec (list): list of predecessor nodes

prior_predictive

def prior_predictive(
    self,
    nsamples=1,
):

Draws from the prior predictive distribution of the graph with self as the root

Args:

• nsamples (int): number of samples to draw

Returns:

samples (torch.tensor)

sample

def sample(
    self,
    nsamples=100,
    thin=0.1,
    burnin=500,
):

Sample from the model’s posterior using the Pyro implementation of the No-U Turn Sampler

This could take a very long time for long time series. It is recommended to use .fit(...) instead.

Args:

• nsamples (int): number of desired samples after burn in and thinning
• thin (float): every int(1.0 / thin) sample is kept
• burnin (int): samples[burnin:] are kept

sin

x -> sin x

softplus

x -> log(1 + exp(x))

succ

Returns the successor nodes of self in the (implicit) compute graph

Returns:

_succ (list): list of successor nodes

tanh

x -> tanh(x), i.e. x -> (exp(x) - exp(-x)) / (exp(x) + exp(-x))

_add_fns_to_repr

None

_apply_fns

None

_closure_init

def closure_init(init_fn):

Generates a complete constructor given only the implementation-specific portion

_collect

Generates a single callable model from a collection of NoiseBlock instances.

The callable model dgp returns a dictionary {model1_name: model1_returnval,...}

_forecast_is_marginalized_var

None

_forecast_replaceattrs

None

_forecast_setattrs

None

_generic_init

def _generic_init(self, init_fn, name=None, t0=0, t1=2, size=1, **kwargs)

Generic portion of init function for dynamically-created blocks

Args:

• init_fn (callable): defines the implementation-specific portion of the constructor

For other argument documentation, see Block

_is_block

None

_is_observable_block

None

_is_pyro_dist

None

_leaf_arg_error

None

_make_2d

None

_make_id

None

_noblock_leaf_arg_error

None

_obj_name_to_definite

None

_obj_name_to_definite_fob

None

_obj_name_to_definite_likelihood

None

_obj_name_to_definite_season

None

collect

def collect(*liks, name="likelihood_collection")

Factory function that generates a unified collection of multiple NoiseBlock subclasses.

Args:

• *liks: each element of liks must subclass NoiseBlock
• name (str): the name of the returned LikelihoodCollection.

construct_init

def construct_init(fn_addr_param):

Constructs an implementation-specific constructor given a parameter specification dictionary.

Args:

• fn_addr_params (dict): {function_address: parameters}. parametersis adict` with structure

      {
          "expand": bool,
          "domain": domain defined in constants,
          "default": Block || torch.tensor || pyro.distributions
      }

forecast

def forecast(dgp, samples, *args, Nt=1, nsamples=1, **kwargs):

Forecasts the root node of the DGP forward in time.

Args:

• dgp (Block): the root node to forecast forward
• samples (dict): {semantic site name: value} The value tensors should have shape (m, n, T), where m is the number of samples, n is the batch size, and T is the length of the time series
• *args: any additional positional arguments to pass to dgp.model
• Nt (int): number of timesteps for which to generate forecast. Forecast is generated from t1 + 1 to t1 + 1 + Nt.
• nsamples (int): number of samples to draw from the forecast distribution
• design_tensors (Dict[str, torch.Tensor]):
• **kwargs: any additional keyword arguments to pass to dgp.model

name_to_definite

def name_to_definite(skeleton, *names,):

Makes the names associated with the block skeleton into torch.tensors according to the skeleton’s current interpretation.

Args:

• skeleton (Block-like): an incompletely-constructed block. The block will be incompletely constructed because the definition of this method is part of the block’s definition.
• *names (List[str]): names to make definite

Returns:

• tuple of torch.tensor corresponding to the passed names.

redefine

def redefine(
    block,
    attribute,
    obj,
):

Redefines an attribute of a block to the passed object

Args:

• block (Block)
• attribute (str)
• obj (Block || torch.tensor || pyro.distributions)

register_address_component

def register_address_component(
    name,
    expand,
    domain=None,
):

Registers a new functional address component and metadata.

Addresses in stsb3 look like x/y-z or y-z, where x is a context coomponent, y is the name of the rv, and z describes its function. register_address_component allows for runtime definition of z address components for use in new blocks or structure search algorithms.

Args:

• name (str): the name. E.g., already-defined names include loc, scale, and amplitude
• expand (bool): whether objects whose address contains this component can be expanded using Blocks
• domain (None || tuple): if expand, must be a tuple in constants.DOMAINS (i.e., one of (-inf, inf), (0, inf), or (0, 1))

This function is not safe – if name already exists in constants, this will overwrite its definition and properties.