GalacticDynamics / galax

Galactic and Gravitational Dynamics in Python (+ GPU and autodiff)
MIT License
33 stars 6 forks source link

Time is needed in phasespaceposition? #501

Open adrn opened 5 days ago

adrn commented 5 days ago

Noticed this when writing an example:

pot = gp.PlummerPotential(
    m_tot=Quantity(5e11, "Msun"),
    b=Quantity(2.5, "kpc"),
    units=unxt.unitsystem("galactic"),
)

psp = gc.PhaseSpacePosition(
    Quantity([1.0, 2.0, 1.0], "kpc"), Quantity([0.0, 200.0, 0.0], "km/s")
)
pot.potential(psp)

This errors with:

AttributeError                            Traceback (most recent call last)
Cell In[36], line 1
----> 1 pot.potential(psp)

    [... skipping hidden 2 frame]

File ~/projects/scratch/.venv/lib/python3.12/site-packages/jaxtyping/_decorator.py:483, in jaxtyped.<locals>.wrapped_fn_impl(args, kwargs, bound, memos)
    480             raise TypeCheckError(msg) from e
    482 # Actually call the function.
--> 483 out = fn(*args, **kwargs)
    485 if full_signature.return_annotation is not inspect.Signature.empty:
    486     # Now type-check the return value. We need to include the
    487     # parameters in the type-checking here in case there are any
   (...)
    498     # checking of the parameters. Unfortunately there doesn't seem
    499     # to be a way around that, so c'est la vie.
    500     kwargs[output_name] = out

File ~/projects/scratch/.venv/lib/python3.12/site-packages/galax/potential/_src/base.py:135, in AbstractBasePotential.potential(self, *args, **kwargs)
    129 \"\"\"Compute the potential energy at the given position(s).
    130 
    131 See :func:`~galax.potential.potential` for details.
    132 \"\"\"
    133 from .funcs import potential
--> 135 return potential(self, *args, **kwargs)

    [... skipping hidden 2 frame]

File ~/projects/scratch/.venv/lib/python3.12/site-packages/jaxtyping/_decorator.py:483, in jaxtyped.<locals>.wrapped_fn_impl(args, kwargs, bound, memos)
    480             raise TypeCheckError(msg) from e
    482 # Actually call the function.
--> 483 out = fn(*args, **kwargs)
    485 if full_signature.return_annotation is not inspect.Signature.empty:
    486     # Now type-check the return value. We need to include the
    487     # parameters in the type-checking here in case there are any
   (...)
    498     # checking of the parameters. Unfortunately there doesn't seem
    499     # to be a way around that, so c'est la vie.
    500     kwargs[output_name] = out

File ~/projects/scratch/.venv/lib/python3.12/site-packages/galax/potential/_src/funcs.py:97, in potential(pot, pspt)
     44 \"\"\"Compute the potential energy at the given position(s).
     45 
     46 Parameters
   (...)
     94 Quantity[...](Array(-1.20227527, dtype=float64), unit='kpc2 / Myr2')
     95 \"\"\"
     96 q = parse_to_quantity(pspt.q, units=pot.units)
---> 97 return pot._potential(q, pspt.t)

    [... skipping hidden 13 frame]

File ~/projects/scratch/.venv/lib/python3.12/site-packages/galax/potential/_src/builtin/builtin.py:801, in PlummerPotential._potential(self, q, t)
    796 @partial(jax.jit, inline=True)
    797 def _potential(
    798     self, q: gt.BatchQVec3, t: gt.BatchableRealQScalar, /
    799 ) -> gt.SpecificEnergyBatchScalar:
    800     r2 = jnp.linalg.vector_norm(q, axis=-1) ** 2
--> 801     return -self.constants[\"G\"] * self.m_tot(t) / jnp.sqrt(r2 + self.b(t) ** 2)

    [... skipping hidden 13 frame]

File ~/projects/scratch/.venv/lib/python3.12/site-packages/galax/potential/_src/params/core.py:156, in ConstantParameter.__call__(self, t, **_)
    139 @partial(jax.jit, inline=True)
    140 def __call__(self, t: BatchableRealQScalar = t0, **_: Any) -> FloatQAnyShape:
    141     \"\"\"Return the constant parameter value.
    142 
    143     Parameters
   (...)
    154         The constant parameter value.
    155     \"\"\"
--> 156     return expand_batch_dims(self.value, t.ndim)

AttributeError: 'NoneType' object has no attribute 'ndim'"

but I think this should work from the user's perspective?

nstarman commented 5 days ago

RN we don't have a mechanism to check whether a potential is time dependent. All are assumed to be. It should error if the potential has time dependence and the coordinate does not include a time. But I agree if it isn't time dependent then it should work.