feat: a sympy backend (#442)

* feat: a sympy backend

* support sympy.Symbol as coordinates

* fix mypy errors

* Fix lite tests

* wrapper around lib

* fix pylint error

* implement copysign and sign

* convert bare coords to sympy symbols

* Add construction and conversion tests + fix minor bugs

* fix lite tests

* add more tests for conversion

* add tests for operators

* move things to one file

* assign lib as an object of SympyLib

* Add pi

* style: pre-commit fixes

* Add equality, non-equality, and conversion tests

* fix lite tests

* planar tests

* fix deltaphi tests

* Add a few spatial tests

* Add beta, gamma, Et, and Et2 tests

* More lorentz tests

* Test spatial methods

* evaluate the expressions numerically in planar tests + addition tests

* Evaluate expressions in spatial tests

* Numerical evaluation for more tests

* More lorentz tests

* test to_beta3

* common compute tests

* More compute tests

* spacelike and timelike sympy tests

* Fix implementation of SympyLib.sign

* Tests for scale

* Better english

* rm print statements

* final set of tests + xfail a few tests on numpy 2

* fix numpy 2 version

* a few tests work on numpy 2

* Add docs for the sympy backend

* Update docstrings

* Update src/vector/backends/sympy.py

* rm redundant code

* Update README.md

Co-authored-by: Jim Pivarski <[email protected]>

* Bump coverage value

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jim Pivarski <[email protected]>

.pre-commit-config.yaml

@@ -39,11 +39,13 @@ repos:
         additional_dependencies:
           - numpy
           - packaging
+          - sympy
 
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.6
     hooks:
       - id: codespell
+        exclude: ^docs/usage/intro.ipynb$
 
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: "v4.0.0-alpha.8"

README.md

@@ -26,10 +26,12 @@ Main features of Vector:
 - Uses names and conventions set by [ROOT](https://root.cern/)'s [TLorentzVector](https://root.cern.ch/doc/master/classTLorentzVector.html) and [Math::LorentzVector](https://root.cern.ch/doc/master/classROOT_1_1Math_1_1LorentzVector.html), as well as [scikit-hep/math](https://github.com/scikit-hep/scikit-hep/tree/master/skhep/math), [uproot-methods TLorentzVector](https://github.com/scikit-hep/uproot3-methods/blob/master/uproot3_methods/classes/TLorentzVector.py), [henryiii/hepvector](https://github.com/henryiii/hepvector), and [coffea.nanoevents.methods.vector](https://coffeateam.github.io/coffea/modules/coffea.nanoevents.methods.vector.html).
 - Implemented on a variety of backends:
   - pure Python objects
+  - [SymPy](https://www.sympy.org/en/index.html) vectors
   - NumPy arrays of vectors (as a [structured array](https://numpy.org/doc/stable/user/basics.rec.html) subclass)
   - [Awkward Arrays](https://awkward-array.org/) of vectors
   - potential for more: CuPy, TensorFlow, Torch, JAX...
-- NumPy/Awkward backends also implemented in [Numba](https://numba.pydata.org/) for JIT-compiled calculations on vectors.
+- Awkward backend also implemented in [Numba](https://numba.pydata.org/) for JIT-compiled calculations on vectors.
+- [JAX](https://awkward-array.org/doc/main/user-guide/how-to-specialize-differentiate-jax.html) and [Dask](https://dask-awkward.readthedocs.io/en/stable/) support through Awkward Arrays.
 - Distinction between geometrical vectors, which have a minimum of attribute and method names, and vectors representing momentum, which have synonyms like `pt` = `rho`, `energy` = `t`, `mass` = `tau`.
 
 ## Installation
@@ -58,7 +60,7 @@ The easiest way to create one or many vectors is with a helper function:
 
 ### Pure Python vectors
 
-You can directly use the `VectorObject` classes to construct object type vectors:
+You can directly use the `VectorObject` and `MomentumObject` classes to construct object type vectors:
 
 ```python
 vector.VectorObject2D(x=1.1, y=2.2)
@@ -159,6 +161,80 @@ vector.VectorObject4D.from_xyetatau(1.1, 2.2, 3.3, 4.4)
 
 and so on, for all combinations of azimuthal, longitudinal, and temporal coordinates, geometric and momentum-flavored.
 
+### SymPy vectors
+
+> **Note:** Operations on SymPy vectors are only 100% compatible with numeric vectors (Python, NumPy, and Awkward backends) if the vectors are positive time-like, that is, if `t**2 > x**2 + y**2 + z**2`. The space-like and negative time-like cases have different sign conventions.
+
+You can directly use the `VectorSympy` and `MomentumSympy` classes to construct object type vectors:
+
+```python
+import sympy
+
+x, y, z, t, px, py, pz, eta, tau = sympy.symbols(
+    "x y z t px py pz eta tau",
+    real=True,  # see sympy assumptions to add more restrictions on the symbols
+)
+vector.VectorSympy2D(x=x, y=y)
+vector.MomentumSympy3D(px=px, py=py, pz=pz)
+vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau)
+```
+
+and so on for every class.
+
+```python
+# Test instance type for any level of granularity.
+(
+    # is a vector or array of vectors
+    isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector),
+    # is 2D (not 3D or 4D)
+    isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector2D),
+    # is a sympy vector (not an array)
+    isinstance(vector.VectorSympy2D(x=x, y=y), vector.VectorSympy),
+    # has momentum synonyms
+    isinstance(vector.MomentumSympy2D(px=px, py=py), vector.Momentum),
+    # has transverse plane (2D, 3D, or 4D)
+    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Planar),
+    # has all spatial coordinates (3D or 4D)
+    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Spatial),
+    # has temporal coordinates (4D)
+    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Lorentz),
+    # azimuthal coordinate type
+    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).azimuthal, vector.AzimuthalXY),
+    # longitudinal coordinate type
+    isinstance(
+        vector.VectorSympy4D(x=x, y=y, z=z, t=t).longitudinal, vector.LongitudinalZ
+    ),
+    # temporal coordinate type
+    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).temporal, vector.TemporalT),
+)
+```
+
+Since `VectorSympy2D`, `VectorSympy3D`, `VectorSympy4D`, and their momentum equivalents operate on SymPy expressions, all of the normal SymPy methods and functions work on the results, coordinates, and the vectors.
+
+```python
+sympy.init_session()  # latex printing
+
+v1 = vector.VectorSympy2D(x=x, y=y)
+sympy.Eq(v1.rho, sympy.sqrt(x**2 + y**2))
+
+v2 = vector.VectorSympy4D(x=x, y=y, z=z, t=t)
+v2.to_rhophithetatau().tau
+
+values = {x: 3, y: 2, z: 1, t: 10}  # t**2 > x**2 + y**2 + z**2
+v2.is_timelike()
+v2.is_timelike().subs(values)
+
+v2.to_rhophithetatau().tau.subs(values).evalf()
+
+v2.boost(v2.to_beta3())
+v2.boost(v2.to_beta3()).t
+v2.boost(v2.to_beta3()).t.simplify()
+v2.boost(v2.to_beta3()).t.subs(values)
+v2.boost(v2.to_beta3()).t.subs(values).evalf()
+```
+
+All of the keyword arguments and rules that apply to `vector.obj` construction apply to `vector.VectorSympyND` and `vector.MomentumObjectND` objects.
+
 ### NumPy arrays of vectors
 
 You can directly use the `VectorNumpy` classes to construct object type vectors: