Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: a sympy backend #442

Merged
merged 45 commits into from
Apr 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
45 commits
Select commit Hold shift + click to select a range
6ebf2e7
feat: a sympy backend
Saransh-cpp Mar 21, 2024
33effcf
support sympy.Symbol as coordinates
Saransh-cpp Mar 25, 2024
ca2b2cc
fix mypy errors
Saransh-cpp Mar 25, 2024
0b8ea42
Fix lite tests
Saransh-cpp Mar 25, 2024
e33fd94
wrapper around lib
Saransh-cpp Mar 26, 2024
e1bd38a
fix pylint error
Saransh-cpp Mar 26, 2024
2f79a05
implement copysign and sign
Saransh-cpp Mar 26, 2024
c50462e
convert bare coords to sympy symbols
Saransh-cpp Mar 27, 2024
99311ff
Add construction and conversion tests + fix minor bugs
Saransh-cpp Apr 4, 2024
2c53d44
fix lite tests
Saransh-cpp Apr 4, 2024
83bcd22
add more tests for conversion
Saransh-cpp Apr 5, 2024
3825f1f
add tests for operators
Saransh-cpp Apr 5, 2024
dee3828
move things to one file
Saransh-cpp Apr 5, 2024
3af081e
assign lib as an object of SympyLib
Saransh-cpp Apr 5, 2024
ad74d47
Add pi
Saransh-cpp Apr 5, 2024
506fb09
style: pre-commit fixes
pre-commit-ci[bot] Apr 5, 2024
2531c25
Add equality, non-equality, and conversion tests
Saransh-cpp Apr 8, 2024
4590f76
fix lite tests
Saransh-cpp Apr 8, 2024
395a192
planar tests
Saransh-cpp Apr 8, 2024
55815a9
fix deltaphi tests
Saransh-cpp Apr 8, 2024
773f973
Add a few spatial tests
Saransh-cpp Apr 9, 2024
2d8f4a0
Add beta, gamma, Et, and Et2 tests
Saransh-cpp Apr 10, 2024
fb30546
More lorentz tests
Saransh-cpp Apr 12, 2024
ebf42f6
Test spatial methods
Saransh-cpp Apr 15, 2024
fb457d0
evaluate the expressions numerically in planar tests + addition tests
Saransh-cpp Apr 16, 2024
218318d
Evaluate expressions in spatial tests
Saransh-cpp Apr 16, 2024
eaf1f85
Numerical evaluation for more tests
Saransh-cpp Apr 16, 2024
a632370
More lorentz tests
Saransh-cpp Apr 16, 2024
e4c6f85
test to_beta3
Saransh-cpp Apr 16, 2024
69f8b24
common compute tests
Saransh-cpp Apr 17, 2024
d86d9a8
More compute tests
Saransh-cpp Apr 18, 2024
49261b2
spacelike and timelike sympy tests
Saransh-cpp Apr 18, 2024
5c28a36
Fix implementation of SympyLib.sign
Saransh-cpp Apr 18, 2024
ac56df0
Tests for scale
Saransh-cpp Apr 18, 2024
655bc26
Better english
Saransh-cpp Apr 18, 2024
f8e94dd
rm print statements
Saransh-cpp Apr 19, 2024
e7930a6
final set of tests + xfail a few tests on numpy 2
Saransh-cpp Apr 19, 2024
a05ec9b
fix numpy 2 version
Saransh-cpp Apr 19, 2024
2fef21b
a few tests work on numpy 2
Saransh-cpp Apr 19, 2024
1192cd8
Add docs for the sympy backend
Saransh-cpp Apr 22, 2024
62a48ed
Update docstrings
Saransh-cpp Apr 22, 2024
cca8a96
Update src/vector/backends/sympy.py
Saransh-cpp Apr 22, 2024
169d834
rm redundant code
Saransh-cpp Apr 22, 2024
4bc4407
Update README.md
Saransh-cpp Apr 23, 2024
2a29f81
Bump coverage value
Saransh-cpp Apr 23, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
80 changes: 78 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
Loading