File: test_rotateZ.py

package info (click to toggle)
python-vector 1.6.3-1
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 2,984 kB
  • sloc: python: 40,200; makefile: 13
file content (35 lines) | stat: -rw-r--r-- 1,251 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/vector for details.

from __future__ import annotations

import pytest

import vector

sympy = pytest.importorskip("sympy")

pytestmark = pytest.mark.sympy

x, y, rho, phi = sympy.symbols("x y rho phi")
values = {x: 1, y: 0, rho: 1, phi: 0}


def test_xy():
    vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y))
    assert vec.rotateZ(1).x == x * sympy.cos(1) - y * sympy.sin(1)
    assert vec.rotateZ(1).y == x * sympy.sin(1) + y * sympy.cos(1)
    assert vec.rotateZ(0.1).x.subs(values) == pytest.approx(0.9950041652780258)
    assert vec.rotateZ(0.1).y.subs(values) == pytest.approx(0.09983341664682815)


def test_rhophi():
    vec = vector.VectorSympy2D(
        azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi)
    )
    assert vec.rotateZ(1).rho == rho
    assert vec.rotateZ(1).phi == sympy.Mod(phi + 1 + sympy.pi, 2 * sympy.pi) - sympy.pi
    assert vec.rotateZ(0.1).rho.subs(values) == pytest.approx(1)
    assert vec.rotateZ(0.1).phi.subs(values) == pytest.approx(0.1)