Package: probmods
ProbMods: Probabilistic ModelsAn interface to facilitate rapid development and deployment of probabilistic modelsContents:InstallationManualWhat is the Problem?Basic PrinciplesModels and InstancesAutomatic Argument Casting:@castParameters Without VarzDetails of Model InstantiationAutomatic Model Instantiation:@instancemethodDescription of ModelsPrior and Posterior Methods:@priormethodand@posteriormethodModel FittingTransformedAutomatic Model TestsInstallationSeethe instructions here.
Then simplypip install probmodsManualWhat is the Problem?Suppose that we have implemented a probabilistic model, and we would like to
apply it to data.
A typical way to do this would be to write a script roughly along the following
lines:importlab.tensorflowasBimporttensorflowastffromstheno.tensorflowimportEQ,GPfromvarzimportVars,minimise_l_bfgs_b# Increase regularisation for `float32`s.B.epsilon=1e-6# Initial model parameters:init_noise=1e-2init_variance=1init_length_scale=1defprior(vs):"""Construct the prior of the model."""ps=vs.struct# Dynamic parameter structvariance=ps.variance.positive(init_variance)length_scale=ps.length_scale.positive(init_length_scale)noise=ps.noise.positive(init_noise)returnGP(variance*EQ().stretch(length_scale)),noisedefsample(vs,x):"""Sample from the prior."""f,noise=prior(vs)returnf(x,noise).sample()# Create a variable container.vs=Vars(tf.float32)# Generate data by sampling from the prior.x=B.linspace(tf.float32,0,10,100)y=sample(vs,x)defobjective(vs):f,noise=prior(vs)return-f(x,noise).logpdf(y)# Fit the model.minimise_l_bfgs_b(objective,vs,trace=True,jit=True)defposterior(vs):"""Construct the posterior."""f,noise=prior(vs)post=f|(f(x,noise),y)returnpost,noisedefpredict(vs,x):"""Make predictions at new input locations."""f,noise=posterior(vs)returnf(x).marginals()# Make predictions.mean,var=predict(vs,B.linspace(tf.float32,10,15,100))In the example, we sample data from a Gaussian process usingStheno, learn hyperparameters for the
Gaussian process, and finally make predictions for new input locations.
Several aspects are not completely satisfactory:There is not one model object which you can conveniently pass around.
Instead, the initial model parameters are defined as global variables, and
the model is implemented with many functions (prior,sample,objective,posterior,predict, ...) which all depend on
each other.
This is not a convenient interface.
If you wanted to share your model with someone else, so they could apply it
to their data, they would not be able to just insert your model in existing
code, but they would have to take your whole script.
Moreover, they even might have to adjust it  appropriately, for what if they
wanted to sample from the posterior or fit the posterior to some data?Since the script uses TensorFlow as a backend, you have to be careful
to convert everything to TensorFlow tensors of the appropriate data type.The script is not easily extensible.
What if you wanted to add in data normalisation by subtracting the mean
and dividing by the standard deviation?
Among other things, you would have to keep track of those parameters
and appropriately transform the predictions back to the original domain.
What if, in addition, your data was positive, so you would want to also
employ a log-transform?
The code now starts to become spaghetti-like.The packageprobmodsaims to solve all of these problems.
Below is how the above model could be implemented withprobmods:importlab.tensorflowasBimportnumpyasnpimporttensorflowastffromsthenoimportEQ,GPfromprobmodsimportModel,instancemethod,cast,Transformed# Increase regularisation for `float32`s.B.epsilon=1e-6classGPModel(Model):def__init__(self,init_variance,init_length_scale,init_noise):self.init_variance=init_varianceself.init_length_scale=init_length_scaleself.init_noise=init_noisedef__prior__(self):"""Construct the prior of the model."""variance=self.ps.variance.positive(self.init_variance)length_scale=self.ps.length_scale.positive(self.init_length_scale)self.f=GP(variance*EQ().stretch(length_scale))self.noise=self.ps.noise.positive(self.init_noise)def__noiseless__(self):"""Transform the model into a noiseless one."""self.noise=0@castdef__condition__(self,x,y):"""Condition the model on data."""self.f=self.f|(self.f(x,self.noise),y)@instancemethod@castdeflogpdf(self,x,y):"""Compute the log-pdf."""returnself.f(x).logpdf(y)@instancemethod@castdefpredict(self,x):"""Make predictions at new input locations."""returnself.f(x,self.noise).marginals()@instancemethod@castdefsample(self,x):"""Sample at new input locations."""returnself.f(x,self.noise).sample()# This model object can easily be inserted anywhere in existing code and be# passed around.model=Transformed(tf.float32,GPModel(1,1,1e-1),transform="normalise+positive",)# Generate data by sampling from the prior.x=np.linspace(0,10,100)y=model.sample(x)# Fit the model and print the learned parameters.model.fit(x,y,trace=True,jit=True)model.vs.print()# Make predictions.x_new=np.linspace(10,15,100)posterior=model.condition(x,y)mean,var=posterior.predict(x_new)# But we can go on....# ...to sample from the posterior.y_new=posterior.sample(x_new)# Or to train the model parameters by _fitting the posterior_ to new data!posterior.fit(x_new,y_new,trace=True,jit=True)# Or to condition on even more data!posterior=posterior.condition(x_new,y_new)# Or to make more preditions!mean,var=posterior.predict(x_new)Basic PrinciplesModels and InstancesModelsare functions fromlearnable parameterstoinstances of models.
Aninstance of a model, or simply aninstance, is an object with concrete
values for it's all parameters and which can do things like sample or compute
a log-pdf.
Moreover, models can be parametrised bynon-learnabe parameters, like
initial values for learnable parameters or parameters which define the structure
of the model, like other submodels.
In this sense, models can interpreted asconfigurationsortemplatesfor instances of models.For example, consider>>>model=GPModel(init_variance=1,init_length_scale=1,init_noise=1e-2)Heremodelis a model parametrised with initial values for learnable
parameters.
We can try to sample from it, but this runs into an error, because, although the
initial values for the parameters ofmodelare set, the actual parameters ofmodeland the internals, such as the Gaussian process, are not yet created:>>>x=np.linspace(0,10,10)>>>model.sample(x)AttributeError:Parameterstructnotavailable.Important assumption:It is assumed that models likemodelcan safely be copied usingcopy.copy,
which  performs ashallow copy.
This means that the constructor of a model should not do much more than
model configuration through setting attributes.
If a shallow copy is not appropriate, you should implementmodel.__copy__.The objectmodelacts like a function from parameters to instances of model.
To demonstrate this, we first need to create parameters.probmodsusesVarzmanage parameters:>>>fromvarzimportVars>>>parameters=Vars(tf.float32)The objectparameterswill create and keep track of all parameters whichmodelwill use.
We can feedparameterstomodelto get an instance, which we can then
sample from.>>>instance=model(parameters)>>>instance.sample(x)array([[0.58702797],[0.40569574],[0.42661083],[1.1435565],[0.02888119],[-1.8267081],[-0.5065604],[0.7860895],[-0.32402134],[-2.4540234]],dtype=float32)You can check whether a model is instantiated or not with the propertyinstantiated:>>>model.instantiatedFalse>>>instance.instantiatedTrueRepresenting models as functions from parameters to instances has a number of
benefits:Transparancy:You can always construct an instance simply by calling it
with parameters. As can sometimes be the case with objects, you do not
need to call particular methods in a particular sequence or worry about
other side effects.Efficiency:Functions can be compiled using a JIT, which eliminates Python's
overhead and can create extremely efficient implementations:@jitdefpplp(parameters,x1,y1,x2,y2):"""Compute the log-pdf of `(x1, y1)` given `(x2, y2)`."""posterior=model.condition(x2,y2)returnposterior(parameters).logpdf(x1,y1)Composability:Models can easily be used as components in bigger models.Automatic Argument Casting:@castAlthough the internal variables ofinstanceare TensorFlow tensors,
you can simply feed a NumPy array toinstance.sample.
Furthermore, the output ofinstance.sample(x)is a NumPy array, rather than a
TensorFlow tensor:>>>x=np.linspace(0,10,10)>>>instance.sample(x)array([[0.58702797],[0.40569574],[0.42661083],[1.1435565],[0.02888119],[-1.8267081],[-0.5065604],[0.7860895],[-0.32402134],[-2.4540234]],dtype=float32)This behaviour is due to the@castdecorator, which automatically
converts NumPy arguments  to the right framework (in this case, TensorFlow) and
the right data type (in this case,tf.float32).
Moreover, ifonlyNumPy arguments were given,@castthen also converts
back to the result to NumPy.
For example, if we were to pass a TensorFlow tensor, we would get a TensorFlow
tensor back:>>>instance.sample(tf.constant(x,dtype=tf.float32))<tf.Tensor:shape=(10,1),dtype=float32,numpy=array([[0.37403315],[-1.423271],[-0.60986364],[0.94153786],[2.247231],[2.799852],[2.8013554],[1.9703895],[0.6884947],[-0.47107112]],dtype=float32)>Parameters Without VarzAlthough the package is integrated withVarzto make parameter management as painless as possible, you are not forced to use
Varz.
If you do not want to use Varz, you should give the appropriate parameters
when you callmodelto instantiate it; these parameters which will then be
passed to__prior__.
Here's howGPModelcould be modified to work in this way:...classGPModel(Model):@castdef__prior__(self,variance,length_scale,noise):"""Construct the prior of the model."""self.f=GP(variance*EQ().stretch(length_scale))self.noise=noise...Note that specifying the initial values of the parameters in the constructor is
not necessary anymore, because all parameter values are given to__prior__upon instantiation.>>>model=GPModel()>>>instance=model(1,1,0.1)>>>instance.fGP(0,EQ()>1)>>>instance.noise0.1Details of Model InstantiationWhenmodelis instantiated by calling it asmodel(*args, **kw_args),
the following happens:First of all, the model iscopiedto safely allow mutation of the copy:instance=copy.copy(model)This copy is ashallow copy.
If a shallow copy is not appropriate, then you should implementinstance.__copy__.If the first argument tomodelwas a variable container of typevarz.Varsor a parameter struct of typevarz.spec.Struct,instance.ps(short for
parameters) is set to extract parameters from it.
If no such argument was given,instance.pswill extract parameters frommodel.vs, ifmodel.vsexists.
If alsomodel.vsdoes not exist,instance.pswill remain unavailable.
Whatever case happens,instance.dtypewill reflect the data type
of the parameters or arguments with whichmodelwas instantiated (argsandkw_args).The prior is constructed:instance.__prior__(*args,**kw_args)The argumentsargsand keyword argumentskw_argsare those given
to the model to instantiate it:model(*args, **kw_args).
Callinginstance.__prior__()mutatesinstance, but that's fine, becauseinstanceis a copy of the original, so no harm done.
The implementation ofinstance.__prior__can access learnable parameters
throughinstance.ps.For every previousmodel = model.condition(x, y)ormodel = model.noiselesscall, the corresponding operations are performed
oninstancein the same order:instance.__condition__(x,y)instance.__noiseless__()We're done! The resultinstanceis returned.instanceis populated with
parameters, has constructed its prior, and has done any potential
conditioning, so it is ready to be used e.g. for sampling:instance.sample(x).Automatic Model Instantiation:@instancemethodFrom what we've seen so far, you can create a model and sample from it in the
following way:# Create a model.model=GPModel(1,1,1e-2)# Sample from it at inputs `x`.parameters=Vars(tf.float32)instance=model(parameters)sample=instance.sample(x)This pattern is slightly suboptimal in two ways:You need to constantly carry around a variable containerparameters.You need to not forget to instantiate the model (callingmodel(parameters))
before doing an operation like sampling.The decorator@instancemethodis designed to help with these issues.
If you decorate a method with@instancemethod, then that indicates that
that method can only be called oninstances of modelsrather thanmodels.
If you call an@instancemethodwithout instantiatingmodel, then
the decorator will automatically instantiate the model with the variable
containermodel.vs, assuming that it is available.
That is, ifmodel.sampleis an@instancemethod, thenmodel.sample(x)automatically translates tomodel(model.vs).sample(x).model.vsdoes not automatically contain a variable container:
you will need to assign it one.# Create a model.model=GPModel(1,1,1e-2)# Assign a variable container.model.vs=Vars(tf.float32)# Sample from the model.sample=model.sample(x)Description of ModelsTheModelclass offers the following properties:PropertyDescriptionmodel.vsA variable container which will be used to automatically instantiate the model when an@instancemethodis called uninstantiated. You need to explicitly assign a variable container tomodel.vs.model.psOnce the model is instantiated,model.ps(orself.psfrom within the class) can be used to initialise constrained variables.model.psis not available for uninstantiated models. As an example, after instantiation,self.ps.parameter_group.matrix.orthogonal(shape=(5, 5))returns a randomly initialised orthogonal matrix of shape(5, 5)namedparameter_group.matrix.psbehaves like a nested struct, dictionary, or list. SeeVarzfor more details.model.instantiatedTrueifmodelis instantiated andFalseotherwise.model.priorTrueifmodelis notconditioned. Throws an exception ifmodelis not instantiated.model.posteriorTrueifmodelisconditioned. Throws an exception ifmodelis not instantiated.model.dtypeIf the model is instantiated, this return the data type ofmodel.ps. If the model is not instantiated, this attempts to returns the data type ofmodel.vs. If neithermodel.psnormodel.vsis available, the data type is automatically determined from the arguments tomodel.__prior__.model.num_outputsA convenience property which can be set to the number of outputs of the model.When you subclassModel, you can implement the following methods:MethodDescription__prior__(self, *args, *kw_args)Construct the prior of the model.__condition__(self, x, y)The prior was previously constructed. Update the model by conditioning on(x, y). You may want to use@convert. You can either return the conditioned model or mutate the current model and return nothing.__noiseless__(self)Remove noise from the current model. You can either return the noiseless model or mutate the current model and return nothing.logpdf(self, x, y)Compute the logpdf for(x, y). This needs to be an@instancemethodand you may want to use@convert.sample(self, x)Sample at inputsx. This needs to be an@instancemethodand you may want to use@convert.predict(self)Predict at inputsx. The default implementation samples and computes the mean and variance of these samples, but you can override this implementation. This needs to be an@instancemethodand you may want to use@convert.For reference, we again show the implementation ofGPModelhere:fromsthenoimportEQ,GPfromprobmodsimportModel,instancemethod,castclassGPModel(Model):def__init__(self,init_variance,init_length_scale,init_noise):self.init_variance=init_varianceself.init_length_scale=init_length_scaleself.init_noise=init_noisedef__prior__(self):"""Construct the prior of the model."""variance=self.ps.variance.positive(self.init_variance)length_scale=self.ps.length_scale.positive(self.init_length_scale)self.f=GP(variance*EQ().stretch(length_scale))self.noise=self.ps.noise.positive(self.init_noise)def__noiseless__(self):"""Transform the model into a noiseless one."""self.noise=0@castdef__condition__(self,x,y):"""Condition the model on data."""self.f=self.f|(self.f(x,self.noise),y)@instancemethod@castdeflogpdf(self,x,y):"""Compute the log-pdf."""returnself.f(x).logpdf(y)@instancemethod@castdefpredict(self,x):"""Make predictions at new input locations."""returnself.f(x,self.noise).marginals()@instancemethod@castdefsample(self,x):"""Sample at new input locations."""returnself.f(x,self.noise).sample()Prior and Posterior Methods:@priormethodand@posteriormethodIt might be that the implementation of an operation, like sampling, is
different for the prior and posterior.
You can use the decorators@priormethodand@posteriormethodto provide
different implementations for the prior and posterior.
These decorators will also automatically instantiate the model, so there is
no need for an additional@instancemethod.Example:fromprobmodsimportModel,priormethod,posteriormethodclassMyModel(Model):def__prior__(self):passdef__condition__(self):pass@priormethoddefsample(self):return"sample from the prior"@posteriormethoddefsample(self):return"sample from the posterior">>>model=MyModel()>>>model.sample()'sample from the prior'>>>model.condition().sample()'sample from the posterior'Important note:The decorators@priormethodand@posteriormethodshould always be theoutermostones.TransformedThe package offers an implementation of one model:Transformed.Transformedtakes an existing model and transforms the output of the model,
e.g., to deal with positive data or to normalise data.Example:model=Transformed(tf.float32,GPModel(1,1,1e-2),transform="normalise+positive",)The first argumenttf.float32indicates the data type of the parameters
that you would like to use.Transformedthen automatically creates a variable container and assigns
it tomodel.vs.
The second and third arguments are the model to transform and the specification
of the transform.
The following transformations are possible:TransformationDescription"normalise"Subtract the mean and divide by the standard deviation. The mean to subtract and the standard deviation to divide by are computed from the data to which the transform is first applied; these values are then remembered."positive"Perform a log-transform. This is handy for positive data."squishing"Perform a transform which suppresses tails. This is handy for heavy-tailed data.You can combine transforms by joining the strings with a,or+.
For example,"normalise+positive"first applies a log-transform and then
normalises the data.
For a more detailed description of, please seeprobmods.bijection.parse.Finally, the optional keyword argumentlearn_transformcan be set toTrueorFalse(default) which specifies whether the parameters of the data
transform should be learned.Model FittingTo fit a model, you can just callmodel.fit(x, y).
The default implementation simply maximisesmodel.logpdf(x, y).
Seeprobmods.model.fitfor a description of the arguments tofit.If you want to provide a custom fitting procedure for your model,
then you can implement a method forfit:fromprobmodsimportfit@fit.dispatchdeffit(model:GPModel,x,y):...# Custom fitting procedure.Note that this will only apply tomodels which are of the typeGPModel.
For example, this will not apply toTransformed(dtype, GPModel(...)).
To implement a fitting procedure for a transformed version ofGPModel, the
following is possible:fromprobmodsimportfit,Transformed@fit.dispatchdeffit(model:Transformed[GPModel],x,y):...# Custom fitting procedure.Automatic Model TestsThe functionprobmods.test.check_modelcan be used to perform some basic
assertions on a model.
See the documentation for more details.