# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['efax',
 'efax._src',
 'efax._src.distributions',
 'efax._src.distributions.cmvn',
 'efax._src.distributions.multivariate_normal',
 'efax._src.scipy_replacement']

package_data = \
{'': ['*']}

install_requires = \
['jax>=0.3',
 'numpy>=1.22',
 'scipy>=1.4,<2.0',
 'tensorflow_probability>=0.15',
 'tjax>=0.19.1']

setup_kwargs = {
    'name': 'efax',
    'version': '1.6.1',
    'description': 'Exponential families for JAX',
    'long_description': "=================================\nEFAX: Exponential Families in JAX\n=================================\n.. image:: https://badge.fury.io/py/efax.svg\n    :target: https://badge.fury.io/py/efax\n\n.. role:: bash(code)\n    :language: bash\n\n.. role:: python(code)\n   :language: python\n\nThis library provides a set of tools for working with *exponential family distributions* in the\ndifferential programming library `JAX <https://github.com/google/jax/>`_.\n\nThe *exponential families* are an important class of probability distributions that include the\nnormal, gamma, beta, exponential, Poisson, binomial, and Bernoulli distributions.\nFor an explanation of the fundamental ideas behind this library, see our `overview on exponential\nfamilies <https://github.com/NeilGirdhar/efax/blob/master/expfam.pdf>`_.\n\nThe main motivation for using EFAX over a library like tensorflow-probability or the basic functions\nin JAX is that EFAX provides the two most important parametrizations for each exponential family—the\nnatural and expectation parametrizations—and a uniform interface to efficient implementations of the\nmain functions used in machine learning.  An example of why this matters is that the most efficient\nway to implement cross entropy between X and Y relies on X being in the expectation parametrization\nand Y in the natural parametrization.\n\nFramework\n=========\nRepresentation\n--------------\nEFAX has a single base class for its objects: :python:`Parametrization` whose type encodes the\ndistribution family.\n\nEach parametrization object has a shape, and so it can store any number of distributions.\nOperations on these objects are vectorized.\nThis is unlike SciPy where each distribution is represented by a single object, and so a thousand\ndistributions need a thousand objects, and corresponding calls to functions that operate on them.\n\nAll parametrization objects are dataclasses using :python:`tjax.dataclass`.  These dataclasses are\na modification of Python's dataclasses to support JAX's “PyTree” type registration.\n\nEach of the fields of a parametrization object stores a parameter over a specified support.\nSome parameters are marked as “fixed”, which means that they are fixed with respect to the\nexponential family.  An example of a fixed parameter is the failure number of the negative binomial\ndistribution.\n\nFor example:\n\n.. code:: python\n\n    @dataclass\n    class MultivariateNormalNP(NaturalParametrization['MultivariateNormalEP']):\n        mean_times_precision: RealArray = distribution_parameter(VectorSupport())\n        negative_half_precision: RealArray = distribution_parameter(SymmetricMatrixSupport())\n\nIn this case, we see that there are two natural parameters for the multivariate normal distribution.\nObjects of this type can hold any number of distributions:  if such an object :python:`x` has shape\n:python:`s`, then the shape of\n:python:`x.mean_times_precision` is :python:`(*s, n)` and the shape of\n:python:`x.negative_half_precision` is :python:`(*s, n, n)`.\n\nParametrizations\n----------------\nEach exponential family distribution has two special parametrizations: the natural and the\nexpectation parametrization.  (These are described in the overview pdf.)\nConsequently, every distribution has at least two base classes, one inheriting from\n:python:`NaturalParametrization` and one from :python:`ExpectationParametrization`.\n\nThe motivation for the natural parametrization is combining and scaling independent predictive\nevidence.  In the natural parametrization, these operations correspond to scaling and addition.\n\nThe motivation for the expectation parametrization is combining independent observations into the\nmaximum likelihood distribution that could have produced them.  In the expectation parametrization,\nthis is an expected value.\n\nEFAX provides conversions between the two parametrizations through the\n:python:`NaturalParametrization.to_exp` and :python:`ExpectationParametrization.to_nat` methods.\n\nImportant methods\n-----------------\nEFAX aims to provide the main methods used in machine learning.\n\nEvery :python:`Parametrization` has methods to flatten and unflatten the parameters into a single\narray: :python:`flattened` and :python:`unflattened`.  Typically, array-valued signals in a machine\nlearning model would be unflattened into a distribution object, operated on, and then flattened\nbefore being sent back to the model.  Flattening is careful with distributions with symmetric (or Hermitian) matrix-valued parameters.  It only stores the upper triangular elements.\n\nEvery :python:`NaturalParametrization` has methods:\n\n- :python:`sufficient_statistics` to produce the sufficient statistics given an observation (used in\n  maximum likelihood estimation),\n- :python:`pdf`, which is the density or mass function,\n- :python:`fisher_information`, which is the Fisher information matrix, and\n- :python:`entropy`, which is the Shannon entropy.\n\nEvery :python:`ExpectationParametrization` has methods:\n\n- :python:`cross_entropy` that is an efficient cross entropy armed with a numerically optimized\n  custom JAX gradient.  This is possible because the gradient of the cross entropy is the difference\n  of expectation parameters plus the expected carrier measure.\n\nNumerical optimization\n----------------------\nBecause of the nature of the log-normalizer and carrier measure, some methods for some distributions\nrequire numerical optimization.  These are the conversion from expectation parameters to natural\nones, the entropy, and the cross entropy.\n\nDistributions\n=============\nEFAX supports the following distributions:\n\n- Bernoulli\n- beta\n- chi\n- chi-square\n- Dirichlet\n- exponential\n- gamma\n- geometric\n- logarithmic\n- multinomial\n- negative binomial\n- normal\n\n  - univariate real\n  - univariate complex\n  - multivariate real\n\n    - with arbitrary variance\n    - with diagonal variance\n    - with isotropic variance\n    - with unit variance\n\n  - multivariate complex\n\n    - with unit variance and zero pseudo-variance\n    - circularly symmetric\n- Poisson\n- Rayleigh\n- von Mises-Fisher\n- Weibull\n\nUsage\n=====\nBasic usage\n-----------\nA basic use of the two parametrizations:\n\n.. code:: python\n\n    from jax import numpy as jnp\n\n    from efax import BernoulliEP, BernoulliNP\n\n    # p is the expectation parameters of three Bernoulli distributions having probabilities 0.4, 0.5,\n    # and 0.6.\n    p = BernoulliEP(jnp.array([0.4, 0.5, 0.6]))\n\n    # q is the natural parameters of three Bernoulli distributions having log-odds 0, which is\n    # probability 0.5.\n    q = BernoulliNP(jnp.zeros(3))\n\n    print(p.cross_entropy(q))\n    # [0.6931472 0.6931472 0.6931472]\n\n    # q2 is natural parameters of Bernoulli distributions having a probability of 0.3.\n    p2 = BernoulliEP(0.3 * jnp.ones(3))\n    q2 = p2.to_nat()\n\n    print(p.cross_entropy(q2))\n    # [0.6955941  0.78032386 0.86505365]\n    # A Bernoulli distribution with probability 0.3 predicts a Bernoulli observation with probability\n    # 0.4 better than the other observations.\n\nOptimization\n------------\nUsing the cross entropy to iteratively optimize a prediction is simple:\n\n.. code:: python\n\n    import jax.numpy as jnp\n    from jax import grad, jit, lax\n    from jax.tree_util import tree_map, tree_reduce\n    from tjax import BooleanNumeric, RealArray, RealNumeric\n\n    from efax import BernoulliEP, BernoulliNP\n\n\n    def cross_entropy_loss(p: BernoulliEP, q: BernoulliNP) -> RealNumeric:\n        return jnp.sum(p.cross_entropy(q))\n\n\n    gce = jit(grad(cross_entropy_loss, 1))\n\n\n    def apply(x: RealArray, x_bar: RealArray) -> RealArray:\n        return x - 1e-4 * x_bar\n\n\n    def body_fun(q: BernoulliNP) -> BernoulliNP:\n        q_bar = gce(some_p, q)\n        return tree_map(apply, q, q_bar)\n\n\n    def cond_fun(q: BernoulliNP) -> BooleanNumeric:\n        q_bar = gce(some_p, q)\n        total = tree_reduce(jnp.sum,\n                            tree_map(lambda x: jnp.sum(jnp.square(x)), q_bar))\n        return total > 1e-6\n\n\n    # some_p are expectation parameters of a Bernoulli distribution corresponding\n    # to probabilities 0.3, 0.4, and 0.7.\n    some_p = BernoulliEP(jnp.array([0.3, 0.4, 0.7]))\n\n    # some_q are natural parameters of a Bernoulli distribution corresponding to\n    # log-odds 0, which is probability 0.5.\n    some_q = BernoulliNP(jnp.zeros(3))\n\n    # Optimize the predictive distribution iteratively, and output the natural parameters of the\n    # prediction.\n    optimized_q = lax.while_loop(cond_fun, body_fun, some_q)\n    print(optimized_q)\n    # BernoulliNP\n    #     log_odds=Jax Array (3,) float32\n    #            -0.8440     -0.4047      0.8440\n\n    # Compare with the true value.\n    print(some_p.to_nat())\n    # BernoulliNP\n    #     log_odds=Jax Array (3,) float32\n    #            -0.8473     -0.4055      0.8473\n\n    # Print optimized natural parameters as expectation parameters.\n    print(optimized_q.to_exp())\n    # BernoulliEP\n    #     probability=Jax Array (3,) float32\n    #             0.3007      0.4002      0.6993\n\nContribution guidelines\n=======================\n\nContributions are welcome!\n\nIt's not hard to add a new distribution.  The steps are:\n\n- Create an issue for the new distribution.\n\n- Solve for or research the equations needed to fill the blanks in the overview pdf, and put them in\n  the issue.  I'll add them to the pdf for you.\n\n- Implement the natural and expectation parametrizations, either:\n\n  - directly like in the Bernoulli distribution, or\n  - as a transformation of an existing exponential family like the Rayleigh distribution.\n\n- Implement the conversion from the expectation to the natural parametrization.  If this has no\n  analytical solution, then there's a mixin that implements a numerical solution.  This can be seen\n  in the Dirichlet distribution.\n\n- Add the new distribution to the tests by adding it to `create_info <https://github.com/NeilGirdhar/efax/blob/master/tests/create_info.py>`_.)\n\nImplementation should respect PEP8.\nThe tests can be run using :bash:`pytest . -n auto`.\nThere are a few tools to clean and check the source:\n\n- :bash:`isort .`\n\n- :bash:`pylint efax`\n\n- :bash:`flake8 efax`\n\n- :bash:`mypy efax`\n",
    'author': 'Neil Girdhar',
    'author_email': 'mistersheik@gmail.com',
    'maintainer': None,
    'maintainer_email': None,
    'url': 'https://github.com/NeilGirdhar/efax',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.8,<3.11',
}


setup(**setup_kwargs)
