examples/notebooks/writing_a_new_distribution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Writing a new Distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n",
    "#\n",
    "#   Licensed under the Apache License, Version 2.0 (the \"License\").\n",
    "#   You may not use this file except in compliance with the License.\n",
    "#   A copy of the License is located at\n",
    "#\n",
    "#       http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "#   or in the \"license\" file accompanying this file. This file is distributed\n",
    "#   on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n",
    "#   express or implied. See the License for the specific language governing\n",
    "#   permissions and limitations under the License.\n",
    "# ==============================================================================\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To write and and use a new Distribution class in MXFusion, fill out the Distribution interface and either the Univariate or Multivariate interface, depending on the type of distribution you are creating.\n",
    "\n",
    "There are 4 primary methods to fill out for a Distribution in MXFusion:\n",
    "* ```__init__```  - This is the constructor for the Distribution. It takes in any parameters the distribution needs. It also defines names for the input variable\\[s\\] that the distribution takes and the output variable\\[s\\] it produces.\n",
    "* ```log_pdf``` - This method returns the logarithm of probabilistic density function for the distribution. This is called during Inference time as necessary to perform the Inference algorithm.\n",
    "* ```draw_samples``` - This method returns drawn samples from the distribution. This is called during Inference time as necessary to perform the Inference algorithm.\n",
    "* ```define_variable``` - This is used to generate random variables drawn from the Distribution used during model definition.\n",
    "\n",
    "```log_pdf``` and ```draw_samples``` are implemented using MXNet functions to compute on the input variables, which at Inference time are MXNet arrays or MXNet symbolic variables.\n",
    "\n",
    "This notebook will take the Normal distribution as a reference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## File Structure\n",
    "Code for distributions lives in the mxfusion/components/distributions directory.\n",
    "\n",
    "If you're implementing  the *FancyNew* distribution then you should create a file called mxfusion/components/distributions/fancy_new.py for the class to live in. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interface Implementation\n",
    "Since this example is for a Univariate Normal distribution, our class extends the UnivatiateDistribution class.\n",
    "\n",
    "The Normal distribution's constructor takes in objects for its ```mean``` and ```variance```, specifications for data type and context, and a random number generator if not the default. \n",
    "\n",
    "In addition, a distribution can take in additional parameters used for calculating that aren't inputs. We refer to these additional parameters as the Distribution's ```attributes```. The difference between an input and an attribute is primarily that inputs are dynamic at inference time, while attributes are static throughout a given inference run. \n",
    "\n",
    "In this case, ```minibatch_ratio``` is a static attribute, as it doesn't change for a given minibatch size during inference. \n",
    "\n",
    "The mean and variance can be either Variables or MXNet arrays if they are constants.\n",
    "\n",
    "As mentioned above, you define names for the input and output variable\\[s\\] for the distribution here. These names are used when printing and generally inspecting the model, so give meaningful names. We prefer names like ```mean``` and ```variance``` to ones like ```location``` and ```scale``` or greek letters like ```mew``` and ```sigma```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Normal(UnivariateDistribution):\n",
    "    \"\"\"\n",
    "    The one-dimensional normal distribution. The normal distribution can be defined over a scalar random variable or an array of random variables. In case of an array of random variables, the mean and variance are broadcasted to the shape of the output random variable (array).\n",
    "\n",
    "    :param mean: Mean of the normal distribution.\n",
    "    :type mean: Variable\n",
    "    :param variance: Variance of the normal distribution.\n",
    "    :type variance: Variable\n",
    "    :param rand_gen: the random generator (default: MXNetRandomGenerator)\n",
    "    :type rand_gen: RandomGenerator\n",
    "    :param dtype: the data type for float point numbers\n",
    "    :type dtype: numpy.float32 or numpy.float64\n",
    "    :param ctx: the mxnet context (default: None/current context)\n",
    "    :type ctx: None or mxnet.cpu or mxnet.gpu\n",
    "    \"\"\"\n",
    "    def __init__(self, mean, variance, rand_gen=None, minibatch_ratio=1.,\n",
    "                 dtype=None, ctx=None):\n",
    "        self.minibatch_ratio = minibatch_ratio\n",
    "        if not isinstance(mean, Variable):\n",
    "            mean = Variable(value=mean)\n",
    "        if not isinstance(variance, Variable):\n",
    "            variance = Variable(value=variance)\n",
    "\n",
    "        inputs = [('mean', mean), ('variance', variance)]\n",
    "        input_names = ['mean', 'variance']\n",
    "        output_names = ['random_variable']\n",
    "        super(Normal, self).__init__(inputs=inputs, outputs=None,\n",
    "                                     input_names=input_names,\n",
    "                                     output_names=output_names,\n",
    "                                     rand_gen=rand_gen, dtype=dtype, ctx=ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your distribution's ```__init__``` function only takes in parameters that get passed onto its super constructor,  you don't need to implement ```replicate_self```. If it does take additional parameters (as the Normal distribution does for minibatch_ratio), those parameters need to be copied over to the replicant Distribution before returning, as below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    def replicate_self(self, attribute_map=None):\n",
    "        \"\"\"\n",
    "        Replicates this Factor, using new inputs, outputs, and a new uuid.\n",
    "        Used during model replication to functionally replicate a factor into a new graph.\n",
    "        :param inputs: new input variables of the factor\n",
    "        :type inputs: a dict of {'name' : Variable} or None\n",
    "        :param outputs: new output variables of the factor.\n",
    "        :type outputs: a dict of {'name' : Variable} or None\n",
    "        \"\"\"\n",
    "        replicant = super(Normal, self).replicate_self(attribute_map)\n",
    "        replicant.minibatch_ratio = self.minibatch_ratio\n",
    "        return replicant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```log_pdf``` and ```draw_samples``` are relatively straightforward for implementation. These are the meaningful parts of the Distribution that you're implementing, putting the math into code using MXNet operators for the compute. \n",
    "\n",
    "If it's a distribution that isn't well documented on Wikipedia, please add a link to a paper or other resource that explains what it's doing and why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    def log_pdf(self, mean, variance, random_variable, F=None):\n",
    "        \"\"\"\n",
    "        Computes the logarithm of the probability density function (PDF) of the normal distribution.\n",
    "\n",
    "        :param mean: the mean of the normal distribution\n",
    "        :type mean: MXNet NDArray or MXNet Symbol\n",
    "        :param variance: the variance of the normal distributions\n",
    "        :type variance: MXNet NDArray or MXNet Symbol\n",
    "        :param random_variable: the random variable of the normal distribution\n",
    "        :type random_variable: MXNet NDArray or MXNet Symbol\n",
    "        :param F: the MXNet computation mode (mxnet.symbol or mxnet.ndarray)\n",
    "        :returns: log pdf of the distribution\n",
    "        :rtypes: MXNet NDArray or MXNet Symbol\n",
    "        \"\"\"\n",
    "        F = get_default_MXNet_mode() if F is None else F\n",
    "        logvar = np.log(2 * np.pi) / -2 + F.log(variance) / -2\n",
    "        logL = F.broadcast_add(logvar, F.broadcast_div(F.square(\n",
    "            F.broadcast_minus(random_variable, mean)), -2 * variance))\n",
    "        return logL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    def draw_samples(self, mean, variance, rv_shape, num_samples=1, F=None):\n",
    "        \"\"\"\n",
    "        Draw samples from the normal distribution.\n",
    "\n",
    "        :param mean: the mean of the normal distribution\n",
    "        :type mean: MXNet NDArray or MXNet Symbol\n",
    "        :param variance: the variance of the normal distributions\n",
    "        :type variance: MXNet NDArray or MXNet Symbol\n",
    "        :param rv_shape: the shape of each sample\n",
    "        :type rv_shape: tuple\n",
    "        :param num_samples: the number of drawn samples (default: one)\n",
    "        :int num_samples: int\n",
    "        :param F: the MXNet computation mode (mxnet.symbol or mxnet.ndarray)\n",
    "        :returns: a set samples of the normal distribution\n",
    "        :rtypes: MXNet NDArray or MXNet Symbol\n",
    "        \"\"\"\n",
    "        F = get_default_MXNet_mode() if F is None else F\n",
    "        out_shape = (num_samples,) + rv_shape\n",
    "        return F.broadcast_add(F.broadcast_mul(self._rand_gen.sample_normal(\n",
    "            shape=out_shape, dtype=self.dtype, ctx=self.ctx),\n",
    "            F.sqrt(variance)), mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``` define_variable ``` is just a helper function for end users. All it does is take in parameters for the distribution, create a distribution based on those parameters, then return the output variables of that distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    @staticmethod\n",
    "    def define_variable(mean=0., variance=1., shape=None, rand_gen=None,\n",
    "                        minibatch_ratio=1., dtype=None, ctx=None):\n",
    "        \"\"\"\n",
    "        Creates and returns a random variable drawn from a normal distribution.\n",
    "\n",
    "        :param mean: Mean of the distribution.\n",
    "        :param variance: Variance of the distribution.\n",
    "        :param shape: the shape of the random variable(s)\n",
    "        :type shape: tuple or [tuple]\n",
    "        :param rand_gen: the random generator (default: MXNetRandomGenerator)\n",
    "        :type rand_gen: RandomGenerator\n",
    "        :param dtype: the data type for float point numbers\n",
    "        :type dtype: numpy.float32 or numpy.float64\n",
    "        :param ctx: the mxnet context (default: None/current context)\n",
    "        :type ctx: None or mxnet.cpu or mxnet.gpu\n",
    "        :returns: the random variables drawn from the normal distribution.\n",
    "        :rtypes: Variable\n",
    "        \"\"\"\n",
    "        normal = Normal(mean=mean, variance=variance, rand_gen=rand_gen,\n",
    "                        dtype=dtype, ctx=ctx)\n",
    "        normal._generate_outputs(shape=shape)\n",
    "        return normal.random_variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using your new distribution\n",
    "At this point, you should be ready to start testing your new Distribution's functionality by importing it like any other MXFusion component."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing\n",
    "Before submitting your new code as a pull request, please write unit tests that verify it works as expected. This should include numerical checks against edge cases. See the existing test cases for the Normal or Categorical distributions for example tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}