Merge pull request #24 from TomHilder/developer

Reconcile developer and master branches. Updates include partial port of `spiralmoments`, as well as dev option smooth_box and various small changes.

.github/workflows/Tests.yml

@@ -13,9 +13,9 @@ jobs:
         python-version: ['3.6', '3.7', '3.8', '3.9', '3.10']
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

.github/workflows/draft-pdf.yml

@@ -6,7 +6,7 @@ jobs:
     name: Paper Draft
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Build draft PDF
         uses: openjournals/openjournals-draft-action@master
         with:

.gitignore

@@ -23,6 +23,9 @@
 # Mac OS
 .DS_Store
 
+# vs code config
+.vscode/*
+
 # Notebooks extras
 /.ipynb_checkpoints
 

docs/API.rst

@@ -17,6 +17,35 @@ Wakeflow module
    :undoc-members:
    :show-inheritance:
 
+Spiralmoments User API
+======================
+
+Sub-module contents
+-------------------
+
+.. automodule:: spiralmoments
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Spiralmoments sub-module
+------------------------
+
+.. automodule:: spiralmoments.spirals
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: spiralmoments.height
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: spiralmoments.phi
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Wakeflow Dev API
 ================
 

pyproject.toml

@@ -42,7 +42,7 @@ testpaths = [
 ]
 
 [tool.bumpver]
-current_version = "1.2.0"
+current_version = "1.3.2"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

setup.cfg

@@ -11,6 +11,6 @@ platforms = unix, linux, osx, windows
 [options]
 package_dir=
     =src
-packages=wakeflow
+packages=wakeflow,spiralmoments
 python_requires = >=3.6
 include_package_data = True

src/spiralmoments/__init__.py

@@ -0,0 +1,12 @@
+# __init__.py
+# Written by Thomas Hilder
+
+"""
+Spiralmoments allows users to map and project spiral-shapes and velocity fields in 3 dimensions to predict peak-velocity maps
+and line-of-sight spiral shapes.
+"""
+
+# spiral shape manipulation stuff
+from .spirals import Spiral
+from .height  import HeightFunctions
+from .phi     import PhiFunctions

src/spiralmoments/height.py

@@ -0,0 +1,63 @@
+# height.py
+# Written by Thomas Hilder
+
+"""
+Contains the HeightFunctions class, which allows users to generate functions that return the disk height as a function of radius, using the parameterisation of their choice.
+"""
+
+import numpy as np
+
+from typing import Callable
+
+class HeightFunctions():
+    """
+    Class containing methods to generate functions that allow users to get disk height as a function of radius (most-likely used for the height of the emitting layer for a certain line). 
+    """
+
+    @staticmethod
+    def powerlaw(z_ref: float, r_ref: float, power_index: float) -> Callable:
+        """Power-law height profile height = z_ref * (r / r_ref)**power_index.
+        
+        Parameters
+        ----------
+        z_ref : disk height at r_ref
+        r_ref : reference radius
+        power_index : power-law index for the height profile
+        
+        returns : Callable function that will return the height value at a provided radius, for the power-law profile.
+        """
+        # return power-law function
+        return lambda r : z_ref * np.power(r / r_ref, power_index)
+
+    @staticmethod
+    def powerlaw_tapered(z_ref: float, r_ref: float, power_index: float, r_taper: float, power_index_taper: float) -> Callable:
+        """Power-law height profile height = z_red * (r / r_ref)**power_intex * exp(- (r / r_taper)**power_index_taper).
+        
+        Parameters
+        ----------
+        z_ref : disk height at r_ref
+        r_ref : reference radius
+        power_index : power-law index for the height profile
+        r_taper : taper radius
+        power_index_taper : power-law index for the exponential taper
+        
+        returns : Callable function that will return the height value at a provided radius, for the tapered power-law profile.
+        """
+        # get regular power-law
+        powerlaw = HeightFunctions.powerlaw(z_ref, r_ref, power_index)
+        # return modification with taper
+        return lambda r : powerlaw(r) * np.exp(
+            -1 * np.power(
+                r / r_taper, power_index_taper
+                )
+        )
+
+    @staticmethod
+    def interpolate_tabulated() -> Callable:
+        """Height function defined by interpolating over user-specified values. NOT IMPLEMENTED.
+        
+        Parameters
+        ----------
+        """
+        # TODO: implement method
+        return NotImplementedError

src/spiralmoments/lindblad.py

@@ -0,0 +1,39 @@
+# lindblad.py
+# Written by Thomas Hilder
+
+"""
+Contains a function to calculate the Lindblad resonance locations
+"""
+
+# TODO: add a function that calculates the resonant positions for an arbitrary Omega function (will involve root finding)
+
+import numpy as np
+
+from typing import Tuple
+
+def lindblad_resonances(m: int, r_0: float) -> Tuple[float, float]:
+    """Calculates the positions of the mth Lindblad resonance position in the disk, 
+    for a perturber located at r_0. Assumes Keplerian rotation.
+    
+    Parameters
+    ----------
+    m : quantum number associated with the Lindblad resonances (azimuthal wave number)
+    r_0 : orbital radius of the perturbing body
+    
+    returns : Tuple containing the inner and outer Lindblad radii respectively.
+    """
+
+    # convert m to a float to make sure we don't do integer division
+    m_float = float(m)
+    # inner resonance
+    r_inner = np.power(
+        1 - np.reciprocal(m_float),
+        2. / 3.
+    )
+    # outer resonance
+    r_outer = np.power(
+        1 + np.reciprocal(m_float),
+        2. / 3.
+    )
+    # return results scaled by r_0
+    return r_inner * r_0, r_outer * r_0

src/spiralmoments/phi.py

@@ -0,0 +1,138 @@
+# phi.py
+# Written by Thomas Hilder
+
+"""
+Contains the PhiFunctions class, which allows users to generate functions that return the azimuthal position of a spiral as a function of radius, using the parameterisation of their choice. Includes options to use physical spirals from the Lin-Shu dispersion relation, including the planet wake shape (Ogilvie and Lubow 2002), as well as parameterised spirals such as logarithmic or power-law spirals.
+"""
+
+import numpy as np
+
+from scipy.integrate import quad
+from typing import Callable
+from .lindblad import lindblad_resonances
+
+class PhiFunctions():
+    """
+    Class containing methods to generate functions that allow users to get spiral shape functions \phi(r) using either simple or physically-motivated parameterisations. In the latter case, \phi(r) can be generated from arbitrary rotation and sound speed curves.
+    """
+
+    @staticmethod
+    def log(b: float, c: float) -> Callable:
+        """Logarithmic spiral, defined by phi = b * log(r) + c. It has constant pitch angle arctan(1 / b).
+        
+        Parameters
+        ----------
+        b : spiral scaling factor (pitch angle = arctan(1 / b))
+        c : constant offset in phi
+        
+        returns : Callable function that will return the phi value at a provided radius, given the defined spiral shape.
+        """
+        return lambda r : b * np.log(r) + c
+
+    @staticmethod
+    def powerlaw(a: float, c: float, power_index: float) -> Callable:
+        """Power-law spiral, defined by phi = a * r**power_index + c. The pitch-angle evolves as a power-law. TODO: add form for pitch-angle
+        
+        Parameters
+        ----------
+        a : spiral scaling factor
+        c : constant offset in phi
+        
+        returns : Callable function that will return the phi value at a provided radius, given the defined spiral shape.
+        """
+        return lambda r : a * np.power(r, power_index) + c
+
+    @staticmethod
+    def powerlaw_plus_log(a: float, b: float, c: float, power_index: float) -> Callable:
+        """Composite spiral with both a power-law and a logarithmic component. Defined by phi = a * r**power_index + b * log(r) + c. Corresponds to
+        a spiral that evolves with a pitch-angle with both a constant and a power-law component. TODO: add form for pitch-angle
+        
+        Parameters
+        ----------
+        a : power-law spiral scaling factor
+        b : log spiral scaling factor (pitch angle approaches arctan(1 / b) as r approaches zero)
+        c : constant offset in phi
+        
+        returns : Callable function that will return the phi value at a provided radius, given the defined spiral shape.
+        """
+        return lambda r : a * np.power(r, power_index) + b * np.log(r) + c
+
+    @staticmethod
+    def planet_wake_powerlaw_disk(phi_planet: float, r_planet: float, hr_planet: float, index_cs: float, rotation_sign: int) -> Callable:
+        """Ogilvie and Lubow (2002) analytical planet wake shape for a power-law disk with Keplerian rotation.
+        
+        Parameters
+        ----------
+        phi_planet : azimuthal location of the planet in radians
+        r_planet : orbital radius of the planet
+        hr_planet : disk aspect ratio H/r at the planet orbital radius r=r_planet. H is the pressure scale height H = cs / Omega.
+        index_cs : power law index for the radial power-law parameterisation of the sound speed in the disk cs = cs_planet * (r / r_planet)**-index_cs
+        rotation_sign : sign of the rotation. +1 for clockwise and -1 for anticlockwise.
+        
+        returns : Callable function that will return the phi value at a provided radius, given the defined planet wake shape.
+        """
+        # scale r by the planet radius
+        rr = lambda r : r / r_planet
+        # calculate the position of the wake relative to the planet
+        phi_wake = lambda r : np.reciprocal(float(hr_planet)) * (
+            np.power(rr(r), index_cs - 0.5) / (index_cs - 0.5) -
+            np.power(rr(r), index_cs + 1.0) / (index_cs + 1.0) -
+            3 / (
+                (2 * index_cs - 1.0) * (index_cs + 1.0)
+            )
+        )
+        # return function for the absolute position of the wake
+        return lambda r : phi_planet + rotation_sign * np.sign(r - r_planet) * phi_wake(r)
+
+    @staticmethod
+    def planet_wake_numerical(phi_planet: float, r_planet: float, Omega_func: float, cs_func: float, rotation_sign: int) -> Callable:
+        """Planet wake form given in Rafikov (2002), defined in terms of the rotation curve Omega(r) and the sounds speed cs(r), calculated numerically.
+        
+        Parameters
+        ----------
+        phi_planet : azimuthal location of the planet in radians
+        r_planet : orbital radius of the planet
+        Omega_func : function that takes only r as an argument and returns the rotation Omega at r
+        cs_func : function that takes only r as an argument and returns the sound speed cs at r
+        rotation_sign : sign of the rotation. +1 for clockwise and -1 for anticlockwise.
+        
+        returns : Callable function that will numerically evaluate and return the phi value at a provided radius, given the defined planet wake shape.
+        """
+        # get a function for the integrand
+        integrand = lambda r : (Omega_func(r) - Omega_func(r_planet)) / cs_func(r)
+        # get a function that returns the result of the integral from r_planet to r
+        integral = lambda r : quad(integrand, r_planet, r)[0]
+        # return function for absolute position of the wake
+        return np.vectorize(
+            lambda r : phi_planet + rotation_sign * np.sign(r - r_planet) * integral(r)
+            )
+
+    @staticmethod
+    def m_mode_spiral_numerical(phi_0: float, r_0: float, m: int, Omega_func: float, cs_func: float, rotation_sign: int) -> Callable:
+        """Individual spiral mode with azimuthal wave number m, defined in terms of the rotation curve Omega(r) and the sounds speed cs(r). Keep in mind that there are evanescent regions interior to the locations of the Lindblad radii. TODO: add a useful reference here.
+        
+        Parameters
+        ----------
+        phi_0 : azimuthal launching location for the spiral wave
+        r_0 : radial launching location for the spiral wave
+        m : azimuthal wave number of the spiral wave
+        Omega_func : function that takes only r as an argument and returns the rotation Omega at r
+        cs_func : function that takes only r as an argument and returns the sound speed cs at r
+        rotation_sign : sign of the rotation. +1 for clockwise and -1 for anticlockwise.
+        
+        returns : Callable function that will numerically evaluate and return the phi value at a provided radius, given the origin and mode of the wave.
+        """
+        # get a function to calculate the integrand (we use nan_to_num to ignore wave propagation in evanescent region)
+        integrand = lambda r : (
+            np.reciprocal(cs_func(r))
+        ) * (
+            np.nan_to_num(
+                np.sqrt((Omega_func(r) - Omega_func(r_0))**2 - (Omega_func(r) / m)**2)
+            )
+        )
+        # get a function that returns the result of the integral from r_0 to r
+        integral = lambda r : quad(integrand, r_0, r)[0]
+        # return function for absolute position of the spiral
+        return np.vectorize(
+            lambda r : phi_0 - rotation_sign * integral(r)
+            )

src/spiralmoments/rotations.py

@@ -0,0 +1,48 @@
+# rotations.py
+# Written by Thomas Hilder
+
+"""
+Contains the rotation_matrix function used to perform 3D rotations of vector fields.
+"""
+
+import numpy as np
+
+def rotation_matrix(angle: float, axis='x') -> np.ndarray:
+    """Function for rotation matrices
+    
+    Parameters
+    ----------
+    angle : angle (in radians) around which vector field should be rotated
+    axis : axis around which vector field should be rotated, either 'x', 'y', or 'z'
+    
+    returns : numpy.ndarray with shape (3,3) representing the rotation matrix.
+    """
+
+    # get phi in [0, 2pi)
+    angle = np.mod(angle, 2*np.pi)
+
+    # get corresponding matrix for chosen axis
+    if axis == "x":
+        arr = [
+            [1,             0,              0],
+            [0, np.cos(angle), -np.sin(angle)],
+            [0, np.sin(angle),  np.cos(angle)]
+        ]
+    elif axis == "y":
+        arr = [
+            [ np.cos(angle), 0, np.sin(angle)],
+            [ 0,             1,             0],
+            [-np.sin(angle), 0, np.cos(angle)]
+        ]
+    elif axis == "z":
+        arr = [
+            [np.cos(angle), -np.sin(angle), 0],
+            [np.sin(angle),  np.cos(angle), 0],
+            [0,                          0, 1]
+        ]
+    # if bogus axis choice throw error
+    else:
+        raise ValueError("axis must be x, y or z")
+    # return array
+    return np.array(arr)
+