Skip to content

Commit 9b3bd6e

Browse files
nstarmannstarman
committed
⬆️ dep-bump(python): Enable python 3.14 support
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
1 parent d1a031b commit 9b3bd6e

142 files changed

Lines changed: 2386 additions & 2727 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.copier-answers.yml‎

Lines changed: 0 additions & 12 deletions
This file was deleted.

‎README.md‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,14 +33,12 @@ import galax.dynamics as gd
3333
import galax.potential as gp
3434

3535
w = gc.PhaseSpaceCoordinate(
36-
q=u.Quantity([8, 0, 0], "kpc"),
37-
p=u.Quantity([0, 220, 0], "km/s"),
38-
t=u.Quantity(0, "Myr"),
36+
t=u.Q(0, "Myr"), q=u.Q([8, 0, 0], "kpc"), p=u.Q([0, 220, 0], "km/s")
3937
)
4038

4139
pot = gp.MilkyWayPotential()
4240

43-
orbit = gd.evaluate_orbit(pot, w, u.Quantity(jnp.linspace(0, 1, 100), "Gyr"))
41+
orbit = gd.evaluate_orbit(pot, w, u.Q(jnp.linspace(0, 1, 100), "Gyr"))
4442
print(orbit)
4543
# Orbit(
4644
# q=<CartesianPos3D: (x, y, z) [kpc]
@@ -51,7 +49,7 @@ print(orbit)
5149
# [[ 0. 0.225 0. ]
5250
# ...
5351
# [ 0.018 0.23 0. ]]>,
54-
# t=Quantity(Array([0., ..., 1000.], dtype=float64), unit='Myr')
52+
# t=Q([0., ..., 1000.], 'Myr')
5553
# )
5654

5755
orbit_sph = orbit.vconvert(cx.vecs.LonLatSphericalPos)

‎docs/getting_started.rst‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,11 @@ that enable fast calculations of computed or derived quantities. For example,
7171
we could compute the potential energy or the acceleration at a Cartesian
7272
position near the Sun::
7373

74-
>>> xyz = u.Quantity([-8., 0, 0], "kpc")
74+
>>> xyz = u.Q([-8., 0, 0], "kpc")
7575
>>> mw.potential(xyz, t=0).uconvert("kpc2 / Myr2")
76-
Quantity(Array(-0.16440296, dtype=float64), unit='kpc2 / Myr2')
76+
Q(-0.16440296, 'kpc2 / Myr2')
7777
>>> mw.acceleration(xyz, t=0)
78-
Quantity(Array([ 0.00702262, -0. , -0. ], dtype=float64), unit='kpc / Myr2')
78+
Q([ 0.00702262, -0. , -0. ], 'kpc / Myr2')
7979

8080
The values that are returned by most methods in :mod:`galax` are provided as
8181
Astropy :class:`~astropy.units.Quantity` objects, which represent numerical data
@@ -84,9 +84,9 @@ re-represented in any equivalent units, so, for example, we could display the
8484
energy or acceleration in other units::
8585

8686
>>> mw.potential(xyz, t=0).uconvert("kpc2/Myr2")
87-
Quantity(Array(-0.16440296, dtype=float64), unit='kpc2 / Myr2')
87+
Q(-0.16440296, 'kpc2 / Myr2')
8888
>>> mw.acceleration(xyz, t=0)
89-
Quantity(Array([ 0.00702262, -0. , -0. ], dtype=float64), unit='kpc / Myr2')
89+
Q([ 0.00702262, -0. , -0. ], 'kpc / Myr2')
9090

9191
Now that we have a potential model, if we want to compute an orbit, we need to
9292
specify a set of initial conditions to initialize the numerical orbit
@@ -98,8 +98,8 @@ velocity vectors. As an example orbit, we will use a position and velocity that
9898
is close to the Sun's Galactocentric position and velocity::
9999

100100
>>> import galax.coordinates as gc
101-
>>> psp = gc.PhaseSpacePosition(q=u.Quantity([-8.1, 0, 0.02], "kpc"),
102-
... p=u.Quantity([13, 245, 8.], "km/s"))
101+
>>> psp = gc.PhaseSpacePosition(q=u.Q([-8.1, 0, 0.02], "kpc"),
102+
... p=u.Q([13, 245, 8.], "km/s"))
103103

104104
By convention, I typically use the variable ``w`` to represent phase-space
105105
positions, so here ``psp`` is meant to imply "initial conditions." Note that,
@@ -136,15 +136,8 @@ phase-space positions at times::
136136

137137
>>> orbit
138138
Orbit(
139-
q=CartesianPos3D(
140-
x=Quantity([...], unit='kpc'),
141-
...
142-
),
143-
p=CartesianVel3D(...),
144-
t=Quantity([...], unit='Myr'),
145-
frame=SimulationFrame(),
146-
interpolant=None
147-
)
139+
q=CartesianPos3D( x=Q([-8.1 , ...], 'kpc'), ... ), p=CartesianVel3D(...),
140+
t=Q([...], 'Myr'), frame=SimulationFrame(), interpolant=None )
148141

149142
:class:`~galax.dynamics.Orbit` objects have many of their own useful methods for
150143
performing common tasks, like plotting an orbit::
@@ -164,9 +157,9 @@ performing common tasks, like plotting an orbit::
164157
import galax.potential as gp
165158

166159
mw = gp.MilkyWayPotential()
167-
psp = gc.PhaseSpacePosition(q=u.Quantity([-8.1, 0, 0.02], "kpc"),
168-
p=u.Quantity([13, 245, 8.], "km/s"))
169-
ts = u.Quantity(jnp.arange(0, 2_000, step=1), "Myr")
160+
psp = gc.PhaseSpacePosition(q=u.Q([-8.1, 0, 0.02], "kpc"),
161+
p=u.Q([13, 245, 8.], "km/s"))
162+
ts = u.Q(jnp.arange(0, 2_000, step=1), "Myr")
170163
orbit = gd.evaluate_orbit(psp, ts)
171164

172165
orbit.plot(['x', 'y'])

‎pyproject.toml‎

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,17 +27,17 @@
2727
"Typing :: Typed",
2828
]
2929
dependencies = [
30-
"astropy>=7.0",
31-
"beartype==0.21.0", # TODO: figure out why 0.22.0 breaks things
32-
"coordinax>=0.23.2",
30+
"astropy>=7.0.0",
31+
"beartype>=0.22.9", # TODO: figure out why 0.22.0 breaks things
32+
"coordinax>=0.23.3",
3333
"dataclassish>=0.8.0",
3434
"diffrax>=0.7",
3535
"diffraxtra>=1.5.2",
36-
"equinox>=0.13.2",
36+
"equinox>=0.13.8",
3737
"interpax>=0.3.11",
3838
"is-annotated>=1.0",
39-
"jax>=0.6.0",
40-
"jaxlib>=0.6.0",
39+
"jax>=0.7.2",
40+
"jaxlib>=0.7.2",
4141
"jaxtyping>=0.3.3",
4242
"oncequinox>=0.1.0",
4343
"optimistix>=0.0.11",
@@ -46,15 +46,16 @@
4646
"optype>=0.14.0; python_version >= '3.12'",
4747
"packaging>=24.1",
4848
"plotting_backends>=0.1",
49-
"plum-dispatch>=2.5.8",
50-
"quax>=0.2.1",
51-
"quaxed>=0.10.2",
49+
"plum>=2.8.0; python_version < '3.14'",
50+
"plum>=2.9.0; python_version >= '3.14'",
51+
"quax>=0.3.2",
52+
"quaxed>=0.10.5",
5253
"tfp-nightly[jax]>=0.25.0",
5354
"typing-extensions>=4.13.2",
54-
"unxt>=1.10.3",
55+
"unxt>=1.11.2",
5556
"wadler_lindig>=0.1.7",
5657
"xmmutablemap>=0.2.1",
57-
"zeroth>=1.0",
58+
"zeroth>=1.2.0",
5859
]
5960

6061
[project.optional-dependencies]
@@ -115,7 +116,7 @@
115116
"pytest-codspeed>=3.2.0",
116117
"pytest-cov>=5",
117118
"pytest>=8.3",
118-
"sybil>=8.0.0",
119+
"sybil>=9.0.0",
119120
"pytest-env>=1.1.5",
120121
]
121122
test-mpl = ["pytest-mpl>=0.17.0"]
@@ -221,6 +222,7 @@ ignore-words-list = "Hart" # in dynamics/_src/cluster/relax_time.py
221222
"--showlocals",
222223
"--strict-config",
223224
"--strict-markers",
225+
"-p no:doctest", # using sybil
224226
"-ra",
225227
]
226228
filterwarnings = [
@@ -340,3 +342,14 @@ ignore-words-list = "Hart" # in dynamics/_src/cluster/relax_time.py
340342

341343
[tool.ruff.lint.pydocstyle]
342344
convention = "numpy"
345+
346+
[tool.unxt]
347+
quantity.repr.include_params = false
348+
quantity.repr.named_unit = false
349+
quantity.repr.short_arrays = "compact"
350+
quantity.repr.use_short_name = true
351+
352+
quantity.str.include_params = false
353+
quantity.str.named_unit = false
354+
quantity.str.short_arrays = false
355+
quantity.str.use_short_name = true

‎src/galax/_interop/galax_interop_astropy/coordinates.py‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,7 @@ def from_(
3030
>>> import galax.coordinates as gc
3131
3232
>>> vec = coord.SphericalRepresentation(
33-
... lon=u.Quantity(10, u.deg), lat=u.Quantity(34, u.deg),
34-
... distance=u.Quantity(3, u.kpc),
33+
... lon=10*u.deg, lat=34*u.deg, distance=3*u.kpc,
3534
... differentials=coord.SphericalCosLatDifferential(
3635
... d_lon_coslat=1*u.deg/u.Myr, d_lat=1*u.deg/u.Myr,
3736
... d_distance=1*u.kpc/u.Myr) )
@@ -77,8 +76,7 @@ def from_(
7776
>>> import galax.coordinates as gc
7877
7978
>>> vec = coord.SphericalRepresentation(
80-
... lon=u.Quantity(10, u.deg), lat=u.Quantity(34, u.deg),
81-
... distance=u.Quantity(3, u.kpc),
79+
... lon=10*u.deg, lat=34*u.deg, distance=3*u.kpc,
8280
... differentials=coord.SphericalCosLatDifferential(
8381
... d_lon_coslat=1*u.deg/u.Myr, d_lat=1*u.deg/u.Myr,
8482
... d_distance=1*u.kpc/u.Myr) )
@@ -87,13 +85,10 @@ def from_(
8785
(10., 34., 3.)
8886
(has differentials w.r.t.: 's')>
8987
90-
>>> gc.PhaseSpaceCoordinate.from_(vec, u.Quantity(0, "Myr"))
88+
>>> gc.PhaseSpaceCoordinate.from_(vec, 0*u.Myr)
9189
PhaseSpaceCoordinate(
92-
q=LonLatSphericalPos( lon=..., lat=..., distance=... ),
93-
p=LonCosLatSphericalVel( lon_coslat=..., lat=..., distance=... ),
94-
t=Quantity(0., unit='Myr'),
95-
frame=SimulationFrame()
96-
)
90+
q=LonLatSphericalPos(...), p=LonCosLatSphericalVel(...),
91+
t=Q(0., 'Myr'), frame=SimulationFrame() )
9792
9893
"""
9994
if "s" not in vec.differentials:
@@ -128,8 +123,7 @@ def from_(
128123
>>> import galax.coordinates as gc
129124
130125
>>> vec = coord.SphericalRepresentation(
131-
... lon=u.Quantity(10, u.deg), lat=u.Quantity(34, u.deg),
132-
... distance=u.Quantity(3, u.kpc))
126+
... lon=10*u.deg, lat=34*u.deg, distance=3*u.kpc)
133127
>>> dif = coord.SphericalCosLatDifferential(
134128
... d_lon_coslat=1*u.deg/u.Myr, d_lat=1*u.deg/u.Myr,
135129
... d_distance=1*u.kpc/u.Myr)
@@ -174,22 +168,18 @@ def from_(
174168
>>> import galax.coordinates as gc
175169
176170
>>> vec = coord.SphericalRepresentation(
177-
... lon=u.Quantity(10, u.deg), lat=u.Quantity(34, u.deg),
178-
... distance=u.Quantity(3, u.kpc))
171+
... lon=10*u.deg, lat=34*u.deg, distance=3*u.kpc)
179172
>>> dif = coord.SphericalCosLatDifferential(
180173
... d_lon_coslat=1*u.deg/u.Myr, d_lat=1*u.deg/u.Myr,
181174
... d_distance=1*u.kpc/u.Myr)
182175
>>> vec
183176
<SphericalRepresentation (lon, lat, distance) in (deg, deg, kpc)
184177
(10., 34., 3.)>
185178
186-
>>> gc.PhaseSpaceCoordinate.from_(vec, dif, u.Quantity(0, "Myr"))
179+
>>> gc.PhaseSpaceCoordinate.from_(vec, dif, 0*u.Myr)
187180
PhaseSpaceCoordinate(
188-
q=LonLatSphericalPos( lon=..., lat=..., distance=... ),
189-
p=LonCosLatSphericalVel( lon_coslat=..., lat=..., distance=... ),
190-
t=Quantity(0., unit='Myr'),
191-
frame=SimulationFrame()
192-
)
181+
q=LonLatSphericalPos(...), p=LonCosLatSphericalVel(...),
182+
t=Q(0., 'Myr'), frame=SimulationFrame() )
193183
194184
"""
195185
return gc.PhaseSpaceCoordinate(q=vec.without_differentials(), p=dif, t=t)

‎src/galax/_interop/galax_interop_astropy/dynamics.py‎

Lines changed: 17 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -48,56 +48,42 @@ def evaluate_orbit(
4848
We can then integrate an initial phase-space position in this potential to
4949
get an orbit:
5050
51-
>>> w0 = gc.PhaseSpaceCoordinate(q=u.Quantity([10., 0., 0.], "kpc"),
52-
... p=u.Quantity([0., 200, 0.], "km/s"),
53-
... t=u.Quantity(-100, "Myr"))
54-
>>> ts = u.Quantity(np.linspace(0., 1., 4), "Gyr")
51+
>>> w0 = gc.PhaseSpaceCoordinate(t=u.Q(-100, "Myr"),
52+
... q=u.Q([10., 0., 0.], "kpc"), p=u.Q([0., 200, 0.], "km/s"))
53+
>>> ts = u.Q(np.linspace(0., 1., 4), "Gyr")
5554
5655
>>> orbit = gd.evaluate_orbit(potential, w0, ts)
5756
>>> orbit
5857
Orbit(
59-
q=CartesianPos3D(...), p=CartesianVel3D(...),
60-
t=Quantity([...], unit='Myr'),
61-
frame=SimulationFrame(),
62-
interpolant=None
63-
)
58+
q=CartesianPos3D(...), p=CartesianVel3D(...), t=Q([...], 'Myr'),
59+
frame=SimulationFrame(), interpolant=None )
6460
65-
>>> ts = u.Quantity(np.linspace(0., 1., 10), "Gyr")
61+
>>> ts = u.Q(np.linspace(0., 1., 10), "Gyr")
6662
>>> orbit = gd.evaluate_orbit(potential, w0, ts)
6763
>>> orbit
6864
Orbit(
69-
q=CartesianPos3D(...), p=CartesianVel3D(...),
70-
t=Quantity([...], unit='Myr'),
71-
frame=SimulationFrame(),
72-
interpolant=None
73-
)
65+
q=CartesianPos3D(...), p=CartesianVel3D(...), t=Q( [...], 'Myr' ),
66+
frame=SimulationFrame(), interpolant=None )
7467
7568
We can also integrate a batch of orbits at once:
7669
77-
>>> w0 = gc.PhaseSpaceCoordinate(q=u.Quantity([[10., 0, 0], [10., 0, 0]], "kpc"),
78-
... p=u.Quantity([[0, 200, 0], [0, 220, 0]], "km/s"),
79-
... t=u.Quantity([-100, -150], "Myr"))
70+
>>> w0 = gc.PhaseSpaceCoordinate(q=u.Q([[10., 0, 0], [10., 0, 0]], "kpc"),
71+
... p=u.Q([[0, 200, 0], [0, 220, 0]], "km/s"),
72+
... t=u.Q([-100, -150], "Myr"))
8073
>>> orbit = gd.evaluate_orbit(potential, w0, ts)
8174
>>> orbit
8275
Orbit(
83-
q=CartesianPos3D(
84-
x=Quantity([...], unit='kpc'),
85-
...
86-
),
87-
p=CartesianVel3D(...),
88-
t=Quantity([...], unit='Myr'),
89-
frame=SimulationFrame(),
90-
interpolant=None
91-
)
76+
q=CartesianPos3D(...), p=CartesianVel3D(...), t=Q( [...], 'Myr' ),
77+
frame=SimulationFrame(), interpolant=None )
9278
9379
:class:`~galax.dynamics.PhaseSpaceCoordinate` has a ``t`` argument for the
9480
time at which the position is given. As noted earlier, this can be used to
9581
integrate from a different time than the initial time of the position:
9682
97-
>>> w0 = gc.PhaseSpaceCoordinate(q=u.Quantity([10., 0., 0.], "kpc"),
98-
... p=u.Quantity([0., 200, 0.], "km/s"),
99-
... t=u.Quantity(0, "Myr"))
100-
>>> ts = u.Quantity(np.linspace(0.3, 1.0, 8), "Gyr")
83+
>>> w0 = gc.PhaseSpaceCoordinate(q=u.Q([10., 0., 0.], "kpc"),
84+
... p=u.Q([0., 200, 0.], "km/s"),
85+
... t=u.Q(0, "Myr"))
86+
>>> ts = u.Q(np.linspace(0.3, 1.0, 8), "Gyr")
10187
>>> orbit = gd.evaluate_orbit(potential, w0, ts)
10288
>>> orbit.q[0] # doctest: +SKIP
10389
Array([ 9.779, -0.3102, 0. ], dtype=float64)

‎src/galax/_interop/galax_interop_astropy/potential.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def parse_to_xyz_t(
4747
# =============================================================================
4848

4949

50-
@dispatch(precedence=1)
50+
@dispatch(precedence=1) # type: ignore[call-overload,misc]
5151
def potential(
5252
pot: gp.AbstractPotential,
5353
xyz: Real[APYQuantity, "*#batch 3"],
@@ -63,7 +63,7 @@ def potential(
6363
return gp.potential(pot, convert(xyz, FastQ), convert(t, FastQ))
6464

6565

66-
@dispatch(precedence=1)
66+
@dispatch(precedence=1) # type: ignore[call-overload,misc]
6767
def gradient(
6868
pot: gp.AbstractPotential,
6969
xyz: Real[APYQuantity, "*#batch 3"],
@@ -79,7 +79,7 @@ def gradient(
7979
return gp.gradient(pot, convert(xyz, FastQ), convert(t, FastQ))
8080

8181

82-
@dispatch(precedence=1)
82+
@dispatch(precedence=1) # type: ignore[call-overload,misc]
8383
def density(
8484
pot: gp.AbstractPotential,
8585
xyz: Real[APYQuantity, "*#batch 3"],
@@ -95,7 +95,7 @@ def density(
9595
return gp.density(pot, convert(xyz, FastQ), convert(t, FastQ))
9696

9797

98-
@dispatch(precedence=1)
98+
@dispatch(precedence=1) # type: ignore[call-overload,misc]
9999
def hessian(
100100
pot: gp.AbstractPotential,
101101
xyz: Real[APYQuantity, "*#batch 3"],

‎src/galax/_interop/galax_interop_gala/coordinates.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,9 @@ def gala_psp_to_galax_psp(obj: galad.PhaseSpacePosition, /) -> gc.PhaseSpacePosi
4141
>>> galax_w = convert(gala_w, gcx.PhaseSpacePosition)
4242
>>> galax_w
4343
PhaseSpacePosition(
44-
q=CartesianPos3D( ... ),
45-
p=CartesianVel3D( ... ),
46-
frame=SimulationFrame()
44+
q=CartesianPos3D(x=Q(1., 'kpc'), y=Q(2., 'kpc'), z=Q(3., 'kpc')),
45+
p=CartesianVel3D(x=Q(4., 'km / s'), y=Q(5., 'km / s'), z=Q(6., 'km / s')),
46+
frame=SimulationFrame()
4747
)
4848
"""
4949
return gc.PhaseSpacePosition(q=obj.pos, p=obj.vel, frame=gc.frames.simulation_frame)
@@ -129,8 +129,7 @@ def galax_psp_to_gala_psp(obj: gc.PhaseSpaceCoordinate, /) -> galad.PhaseSpacePo
129129
it to a :class:`gala.dynamics.PhaseSpacePosition`.
130130
131131
>>> galax_w = gcx.PhaseSpaceCoordinate(
132-
... q=[1, 2, 3] * u.kpc, p=[4, 5, 6] * u.km / u.s, t=2 * u.Myr
133-
... )
132+
... q=[1, 2, 3] * u.kpc, p=[4, 5, 6] * u.km / u.s, t=2 * u.Myr )
134133
135134
>>> with catch_warnings(action="ignore"):
136135
... gala_w = convert(galax_w, gd.PhaseSpacePosition)

0 commit comments

Comments
 (0)