Add radial velocity component to the fit_map rotation curve

eddy/rotationmap.py

@@ -125,6 +125,11 @@ def fit_map(self, p0, params, r_min=None, r_max=None, optimize=True,
         low spatial resolution data, or for rotation maps made via the
         intensity-weighted average velocity.
 
+        A radial velocity component, ``v_r(r) = vr_100 * (r / 100)^vr_q``, can
+        be added on top of the rotation curve (whether Keplerian or power-law)
+        by setting ``params['vr_100']`` (and ``params['vr_q']``), e.g. to
+        capture a disk wind or infall signature.
+
         .. _Rosenfeld et al. (2013): https://ui.adsabs.harvard.edu/abs/2013ApJ...774...16R/abstract
 
         Args:
@@ -1211,7 +1216,7 @@ def verify_params_dictionary(self, params):
         fitting and fill in defaults for any that are missing. Sets the
         rotation velocity function (``'vfunc'``) based on whether a power-law
         profile or Keplerian profile is requested, and flags whether a vortex
-        component should be included.
+        or radial velocity component should be included.
 
         Args:
             params (dict): Dictionary of model parameters, as passed to
@@ -1246,6 +1251,10 @@ def verify_params_dictionary(self, params):
         else:
             params['vortex'] = False
 
+        # Radial velocity component.
+
+        params['vradial'] = params['vr_100'] is not None
+
         # Deprojection properties.
 
         params['z_func'] = params.pop('z_func', None)
@@ -1602,6 +1611,11 @@ def _vpow(self, rvals, tvals, zvals, params):
         vpow = (rvals * params['dist'] / 100.)**params['vp_q']
         return params['vp_100'] * vpow
 
+    def _vrad(self, rvals, params):
+        """Power-law radial velocity profile."""
+        vrad = (rvals * params['dist'] / 100.)**params['vr_q']
+        return params['vr_100'] * vrad
+
     def _vpow_pressure(self, rvals, tvals, zvals, params):
         """Power-law rotation with pressure term."""
         vpow = self._vpow(rvals, tvals, zvals, params)
@@ -1766,17 +1780,22 @@ def _make_model(self, params):
             rvals = 1.0 / (2.0 / rvals - 1.0 / a)
             #tvals = jnp.arcsin(xvals / rvals)
 
-        # Calculate the velocity profile and project. This includes an
-        # additional component from the vortex.
+        # Calculate the velocity profile and project. This includes
+        # additional components from the vortex and radial velocity.
 
         vphi = params['vfunc'](rvals, tvals, zvals, params)
         vphi_proj = self._proj_vphi(vphi, tvals, params)
         if params['vortex']:
             vvor_proj = self._make_model_vortex(rvals, tvals, params)
         else:
             vvor_proj = 0.0
+        if params['vradial']:
+            vrad_proj = self._proj_vrad(self._vrad(rvals, params),
+                                         tvals, params)
+        else:
+            vrad_proj = 0.0
 
-        v0 = vphi_proj + vvor_proj + params['vlsr']
+        v0 = vphi_proj + vvor_proj + vrad_proj + params['vlsr']
 
         # Convolve if necessary.
 

tests/test_rotationmap.py

@@ -81,3 +81,26 @@ def test_fit_map_returns_percentiles(hd163296_rotationmap):
     assert np.all(np.isfinite(pcts))
 
 
+def test_make_model_includes_radial_velocity(hd163296_rotationmap):
+    """``vr_100``/``vr_q`` must perturb ``_make_model``'s output. Guards
+    against a repeat of the regression where the radial velocity term was
+    silently dropped from the sky-projected model."""
+    base_params = {
+        'x0': 0.0, 'y0': 0.0, 'PA': 312.0, 'mstar': 2.0, 'vlsr': 5.7e3,
+        'inc': 46.7, 'dist': 101.0,
+    }
+
+    no_rad = hd163296_rotationmap.verify_params_dictionary(dict(base_params))
+    assert no_rad['vradial'] is False
+    model_no_rad = hd163296_rotationmap._make_model(no_rad)
+
+    with_rad = hd163296_rotationmap.verify_params_dictionary(
+        dict(base_params, vr_100=500.0, vr_q=0.0))
+    assert with_rad['vradial'] is True
+    model_with_rad = hd163296_rotationmap._make_model(with_rad)
+
+    assert not np.array_equal(np.asarray(model_no_rad),
+                               np.asarray(model_with_rad))
+    assert np.all(np.isfinite(model_with_rad))
+
+