Merge pull request #132 from adrn/ptypes-class-level

Access physical types for parameters at the class level

docs/dynamics/mockstreams.rst

@@ -85,8 +85,9 @@ potential and integrate the orbit of the progenitor system. We'll use a
 spherical NFW potential, a progenitor mass :math:`m=10^4~{\rm M}_\odot` and
 initial conditions that place the progenitor on a mildly eccentric orbit:
 
-    >>> pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc,
-    ...                                units=galactic)
+    >>> pot = gp.NFWPotential.from_circular_velocity(v_c=250*u.km/u.s,
+    ...                                              r_s=10*u.kpc,
+    ...                                              units=galactic)
     >>> prog_mass = 1E4*u.Msun
     >>> prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc,
     ...                                 vel=[75, 150, 30.]*u.km/u.s)
@@ -103,7 +104,7 @@ initial conditions that place the progenitor on a mildly eccentric orbit:
     import gala.dynamics as gd
     from gala.units import galactic
 
-    pot = gp.NFWPotential.from_circular_velocity(v_c=175*u.km/u.s, r_s=10*u.kpc,
+    pot = gp.NFWPotential.from_circular_velocity(v_c=250*u.km/u.s, r_s=10*u.kpc,
                                                  units=galactic)
     prog_mass = 1E4*u.Msun
     prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc,
@@ -200,4 +201,3 @@ API
 .. automodapi:: gala.dynamics.mockstream
     :no-heading:
     :headings: ^^
-

docs/potential/index.rst

@@ -263,8 +263,6 @@ API
 ===
 
 .. automodapi:: gala.potential.potential
-    :skip: SphericalNFWPotential
-    :skip: FlattenedNFWPotential
 
 .. automodapi:: gala.potential.frame.builtin
 

gala/dynamics/mockstream/core.py

@@ -15,7 +15,8 @@
            'dissolved_fardal_stream']
 
 def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp,
-                release_every=1, Integrator=DOPRI853Integrator, Integrator_kwargs=dict(),
+                release_every=1, Integrator=DOPRI853Integrator,
+                Integrator_kwargs=dict(),
                 snapshot_filename=None, seed=None):
     """
     Generate a mock stellar stream in the specified potential with a
@@ -31,15 +32,15 @@ def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp,
         A single mass or an array of masses if the progenitor mass evolves
         with time.
     k_mean : `numpy.ndarray`
-        Array of mean :math:`k` values (see Fardal et al. 2015). These are used to determine
-        the exact prescription for generating the mock stream. The components are for:
-        :math:`(R,\phi,z,v_R,v_\phi,v_z)`. If 1D, assumed constant in time. If 2D, time axis
-        is axis 0.
+        Array of mean :math:`k` values (see Fardal et al. 2015). These are used
+        to determine the exact prescription for generating the mock stream. The
+        components are for: :math:`(R,\phi,z,v_R,v_\phi,v_z)`. If 1D, assumed
+        constant in time. If 2D, time axis is axis 0.
     k_disp : `numpy.ndarray`
-        Array of :math:`k` value dispersions (see Fardal et al. 2015). These are used to determine
-        the exact prescription for generating the mock stream. The components are for:
-        :math:`(R,\phi,z,v_R,v_\phi,v_z)`. If 1D, assumed constant in time. If 2D, time axis
-        is axis 0.
+        Array of :math:`k` value dispersions (see Fardal et al. 2015). These are
+        used to determine the exact prescription for generating the mock stream.
+        The components are for: :math:`(R,\phi,z,v_R,v_\phi,v_z)`. If 1D,
+        assumed constant in time. If 2D, time axis is axis 0.
     release_every : int (optional)
         Release particles at the Lagrange points every X timesteps.
     Integrator : `~gala.integrate.Integrator` (optional)
@@ -60,11 +61,12 @@ def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp,
     """
 
     if isinstance(hamiltonian, CPotentialBase):
-        warnings.warn("This function now expects a `Hamiltonian` instance instead of "
-                      "a `PotentialBase` subclass instance. If you are using a "
-                      "static reference frame, you just need to pass your "
-                      "potential object in to the Hamiltonian constructor to use, e.g., "
-                      "Hamiltonian(potential).", DeprecationWarning)
+        warnings.warn("This function now expects a `Hamiltonian` instance "
+                      "instead of a `PotentialBase` subclass instance. If you "
+                      "are using a static reference frame, you just need to "
+                      "pass your potential object in to the Hamiltonian "
+                      "constructor to use, e.g., Hamiltonian(potential).",
+                      DeprecationWarning)
 
         hamiltonian = Hamiltonian(hamiltonian)
 
@@ -93,16 +95,25 @@ def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp,
     # ------------------------------------------------------------------------
 
     if prog_orbit.t[1] < prog_orbit.t[0]:
-        raise ValueError("Progenitor orbit goes backwards in time. Streams can only "
-                         "be generated on orbits that run forwards. Hint: you can "
-                         "reverse the orbit with prog_orbit[::-1], but make sure the array"
-                         "of k_mean values is ordered correctly.")
+        raise ValueError("Progenitor orbit goes backwards in time. Streams can "
+                         "only be generated on orbits that run forwards. Hint: "
+                         "you can reverse the orbit with prog_orbit[::-1], but "
+                         "make sure the array of k_mean values is ordered "
+                         "correctly.")
 
     c_w = np.squeeze(prog_orbit.w(hamiltonian.units)).T # transpose for Cython funcs
     prog_w = np.ascontiguousarray(c_w)
     prog_t = np.ascontiguousarray(prog_orbit.t.decompose(hamiltonian.units).value)
-    if hasattr(prog_mass, 'unit'):
-        prog_mass = prog_mass.decompose(hamiltonian.units).value
+    if not hasattr(prog_mass, 'unit'):
+        prog_mass = prog_mass * hamiltonian.units['mass']
+
+    if not prog_mass.isscalar:
+        if len(prog_mass) != prog_orbit.ntimes:
+            raise ValueError("If passing in an array of progenitor masses, it "
+                             "must have the same length as the number of "
+                             "timesteps in the input orbit.")
+
+    prog_mass = prog_mass.decompose(hamiltonian.units).value
 
     if Integrator == LeapfrogIntegrator:
         stream_w = _mock_stream_leapfrog(hamiltonian, t=prog_t, prog_w=prog_w,

gala/potential/common.py

@@ -12,6 +12,7 @@
 from ..units import UnitSystem, DimensionlessUnitSystem
 
 class CommonBase(object):
+    _physical_types = None
 
     def _validate_units(self, units):
 
@@ -25,13 +26,17 @@ def _validate_units(self, units):
         return units
 
     @classmethod
-    def _prepare_parameters(cls, parameters, ptypes, units):
+    def _prepare_parameters(cls, parameters, units):
+        ptypes = cls._physical_types
+        if ptypes is None:
+            ptypes = dict()
+
         pars = OrderedDict()
-        for k,v in parameters.items():
+        for k, v in parameters.items():
             if hasattr(v, 'unit'):
                 pars[k] = v.decompose(units)
 
-            elif k in ptypes:
+            elif k in ptypes and ptypes[k]: # treat empty string as u.one
                 # HACK TODO: remove when fix potentials that ask for scale velocity
                 if ptypes[k] == 'speed':
                     pars[k] = v * units['length']/units['time']

gala/potential/frame/cframe.pyx

@@ -81,14 +81,9 @@ cdef class CFrameWrapper:
 
 class CFrameBase(FrameBase):
 
-    def __init__(self, Wrapper, parameters, units, parameter_physical_types=None, ndim=3):
+    def __init__(self, Wrapper, parameters, units, ndim=3):
         self.units = self._validate_units(units)
-
-        if parameter_physical_types is None:
-            parameter_physical_types = dict()
-        self._ptypes = parameter_physical_types
-
-        self.parameters = self._prepare_parameters(parameters, self._ptypes, self.units)
+        self.parameters = self._prepare_parameters(parameters, self.units)
         self.c_parameters = np.ravel([v.value for v in self.parameters.values()])
         self.c_instance = Wrapper(*self.c_parameters)
 

gala/potential/hamiltonian/tests/test_with_frame_potential.py

@@ -99,7 +99,8 @@ def to_rotating_frame(omega, w, t=None):
 # ----------------------------------------------------------------------------
 
 class TestWithPotentialStaticFrame(_TestBase):
-    obj = Hamiltonian(NFWPotential.from_circular_velocity(v_c=0.2, r_s=20., units=galactic),
+    obj = Hamiltonian(NFWPotential.from_circular_velocity(v_c=0.2, r_s=20.,
+                                                          units=galactic),
                       StaticFrame(units=galactic))
 
     @pytest.mark.skip("Not implemented")