Package: polytensor
.. PolyTensor documentation master file, created by
sphinx-quickstart on Fri Dec 22 09:52:54 2023.
You can adapt this file completely to your liking, but it should at least
contain the roottoctreedirective... image:: ./docs/source/_static/icon/moonrabbit.svg
:align: center
:width: 100
:alt: polytensor logopolytensorpolytensoris a python package for CUDA-accelerated, parallel polynomial evaluation and regression.$$f(x) = c + \sum_{i=0}^n a_i x_i + \sum_{i < j}^n a_{i,j} x_i x_j + ... $$Why?Evaluating standard, non-matrix polynomials in a CUDA-accelerated, parallel fashion has never been easier! Polytensor is a PyTorch-based package for computing millions of polynomials in parallel on a CUDA GPU. We offer two flavors, sparse-vanilla and dense-rocky-road. My work on quantum-inspired energy models requires computing all kinds of polynomials for optimization and dynamic simulations, and I wanted a clean way of computing the energy function for these models in parallel.Quick StartTo use the latest stable version ofpolytensor, install it usingpipfrom the command line:.. code-block:: console$ pip install polytensorFor the latest development version, install it directly from this repo:.. code-block:: console$ python -m venv .venv
$ source .venv/bin/activate
$ (.venv) python -m pip install git+https://github.com/btrainwilson/polytensor.gitOr, if you want to developpolytensor, install it in editable mode:.. code-block:: console$ git clone git+https://github.com/btrainwilson/polytensor.git
$ python -m pip install -e polytensorUsageExamplesAll of the following examples assume that you have importedpolytensor:.. code-block:: pythonimport polytensorPolynomialThe easiest way to begin is by creating a polynomial with a list of terms and coefficients. Consider the following polynomial 



$$f(x) = x_0 + 2 x_1 + 3 x_0 x_1 + 5 x_1^2$$


Each term is specified by a list of indeces and a value. For example, the coefficient $3$ is associated with the term $3 x_0 x_1$, the coefficient $1$ is associated with the term $x_0$, etc.

.. code-block:: python

    terms = {
      (0,)   : 1.0,     # 1.0 * x_0
      (1,)   : 2.0,     # 2.0 * x_1
      (0, 1) : 3.0,     # 3.0 * x_0 * x_1
      (1, 1) : 5.0,     # 5.0 * x_1^2
    }

We can also create random polynomials using the ``polytensor.generators`` module. For example, the following expression is a random polynomial with 10 variables where the coefficients are sampled with a Gaussian $\\mathcal{N}(0,1)$.

$$f(x) = \\sum_{i=0}^{10} a_i x_i + \\sum_{i < j}^{10} a_{i,j} x_i x_j + \\sum_{i < j < k}^{10} a_{i,j, k} x_i x_j x_k + ... : a_s \\sim \\mathcal{N}(0, 1)$$

where $s$ is the term in the polynomial, e.g., $s = (i, j, k)$.


.. code-block:: python

    import torch

    num_vars = 10

    # Create a random polynomial with 10 variables and 5 terms per degree
    num_per_degree = [num_vars, 5, 5, 5]

    # Function to sample coefficients
    sample_fn = lambda: torch.randn(1)


    # Create coefficients for a random polynomial with 10 variables and 5 terms per degree up to degree 4
    terms = polytensor.generators.coeffPUBORandomSampler(
        n=num_vars, num_terms=num_per_degree,sample_fn=sample_fn
        )

Given these coefficients, we can create a polynomial using either a sparse representation or a dense representation. The sparse representation is more efficient for polynomials with fewer terms, while the dense representation is more efficient for polynomials with more terms.


Sparse PolynomialsUnder the hood, the terms remain in their dictionary definition, where the keys are tuples of indeces and the values are the coefficients. For example, the following code creates a sparse polynomial with the coefficients from the previous example... code-block:: pythonterms = {
  (0,)   : 1.0,     # 1.0 * x_0
  (1,)   : 2.0,     # 2.0 * x_1
  (0, 1) : 3.0,     # 3.0 * x_0 * x_1
  (1, 1) : 5.0,     # 5.0 * x_1^2
}

poly = polytensor.SparsePolynomial(terms)

x = torch.Tensor([1.0, 2.0])

# Evaluate the polynomial at x
y_p = poly(x)

# Which is equivalent to
y_s = 0.0
for term, v in terms.items():
    y_s = y_s + v * torch.prod(x[..., term])

assert np.allclose(y_p.detach().cpu().numpy(), y_s.detach().cpu().numpy())In fact, the loop above is exactly how the polynomial is evaluated. TheSparsePolynomialclass is a wrapper around the dictionary of terms and coefficients. The__call__method loops through the terms and evaluates the polynomial at the given point :math:x. Now, we consider dense polynomials.Dense PolynomialsAt a glance, the ``DensePolynomial`` stores the terms in a list of dense ``torch.Tensor`` s, one tensor for each degree, where the indeces of the tensor are the term indeces and the tensor element is the coefficient. The ``DensePolynomial`` class exploits the ``einsum`` function in ``torch`` to evaluate the polynomial using the dense tensors. 


Sparse vs Dense RepresentationWhen to use the sparse representation? The sparse representation is more efficient than the dense representation when the number of terms :math:Nis small compared to the number of possible terms, i.e.,$$N << n^d$$Forpolytensor.DensePolynomial, The number of terms in the tensor for degree :math:dis :math:n^dwhere :math:nis the number of variables in the polynomial. The einsum computation using this representation is way faster than the sparse enumeration if the number of terms is similar to the size of the tensors. Under the hood ofpolytensor.DensePolynomial,torch.einsumexploits CUDA acceleration to parallelize the computation. However, if the number of terms in the polynomial is nowhere close to the number of terms in the dense tensor representation, then most of the terms in the dense tensors will be :math:0and the sparse polynomial is a better representation. For example, if your polynomial has :math:100terms, most of which are quadratic or linear, then a dense representation is likely more efficient. However, if those 100 terms are distributed throughout 6 degree monomials, then a sparse representation is more efficient.ContributingWe welcome contributions!To set up the test environment (.tenv virtual environment), run the following commands:.. code-block:: console$ git clone git+https://github.com/btrainwilson/polytensor.git
$ cd polytensor
$ make .tenv
$ source .tenv/bin/activateThis will handle installing the development dependencies and setting up the virtual environment.TestingTo run the tests, run the following command:

.. code-block:: console

    $ make test

If everything is set up properly, the tests should pass with green text at the bottom. 

DocumentationTo build the documentation, run the following command:.. code-block:: console$ make docThis will build the documentation in thedocs/builddirectory.
To view the documentation,.. code-block:: console$ make serveand navigate tolocalhost:8018in your browser.Pull RequestsTo submit a contribution, fork the repo and submit a pull request with your changes. We will review the request by running our test suites to ensure the interface is not broken and then check for code cleanliness and correctness. To increase the chances of accepting a PR, build a unit test in the test/ directory as a part of the PR.