Initial development

.github/workflows/build.yaml

@@ -0,0 +1,38 @@
+name: build
+
+on: [push]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libgsl-dev libeigen3-dev
+          export CPLUS_INCLUDE_PATH=/usr/include/eigen3
+          python -m pip install --upgrade pip
+          pip install setuptools wheel numpy
+      - name: Install AGAMA
+        run: |
+          git clone https://github.com/GalacticDynamics-Oxford/Agama.git
+          cd Agama
+          python setup.py install --user --assume-yes
+          cd ..
+      - name: Build package
+        run: |
+          pip install -e .
+      - name: Run tests
+        run: |
+          pip install pytest
+          pytest

.gitignore

@@ -1,3 +1,5 @@
+examples/data/mock_dataset.txt
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]

README.md

@@ -0,0 +1,52 @@
+# StarStream
+
+[![version](https://img.shields.io/badge/version-0.1-blue.svg)](https://github.com/ybillchen/StarStream)
+[![license](https://img.shields.io/github/license/ybillchen/StarStream)](LICENSE)
+[![ADS](https://img.shields.io/badge/ADS-2025xxxxx-blue)](#)
+[![arXiv](https://img.shields.io/badge/arXiv-25xx.xxxxx-green)](#)
+
+An automatic detection algorithm for stellar streams using a physics-inspired stream model.
+
+The code is open source under a [BSD 3-Clause License](LICENSE), which allows you to redistribute and modify the code with moderate limitations. If you use this code for a publication, we kindly request you to cite the following original paper.
+
+- Y. Chen, O. Y. Gnedin, A. M. Price-Whelan, & C. Holm-Hansen (2025) *StarStream: Automatic detection algorithm for stellar streams*. [arXiv:25xx.xxxxx](https://arxiv.org/abs/25xx.xxxxx), [ADS link](#)
+
+## Install
+
+We have tested `StarStream` on `3.9 <= python <= 3.11`. However, lower or higher versions may also work. The prerequisites of this package are
+```
+numpy
+scipy
+astropy
+agama
+```
+
+To download the packge, `git clone` the source code from [GitHub](https://github.com/ybillchen/StarStream):
+```shell
+$ git clone https://github.com/ybillchen/StarStream.git
+```
+Next, `cd` the folder and use `pip` to install it:
+```shell
+$ cd StarStream/
+$ pip install -e .
+```
+The `-e` command allows you to make changes to the code.
+
+To check if the package is installed correctly, you may run the tests using `pytest` (make sure it's installed)
+```shell
+$ pytest
+```
+
+## Usage
+
+We provide [example notebooks](examples/) to demonstrate how to use apply this method to a mock dataset similar to *Gaia* DR3.
+
+
+## Contribute
+
+Feel free to dive in! [Raise an issue](https://github.com/ybillchen/StarStream/issues/new) or submit pull requests. We recommend you to contribute code following [GitHub flow](https://docs.github.com/en/get-started/quickstart/github-flow). 
+
+## Maintainers
+
+- [@ybillchen (Bill Chen)](https://github.com/ybillchen)
+

StarStream/__init__.py

@@ -0,0 +1,20 @@
+# Licensed under BSD-3-Clause License - see LICENSE
+
+from . import pdf
+from . import spray
+from . import optimize
+from . import utils
+__all__ = []
+__all__.append(pdf.__all__)
+__all__.append(spray.__all__)
+__all__.append(optimize.__all__)
+__all__.append(utils.__all__)
+
+from .pdf import *
+from .spray import *
+from .optimize import *
+from .utils import *
+from .version import __version__
+
+__name__ = "StarStream"
+__author__ = "Bill Chen"

StarStream/agama_stream.py

@@ -0,0 +1,140 @@
+# Licensed under BSD-3-Clause License - see LICENSE
+
+"""
+Agama implementation of Chen+25 particle spray algorithm
+"""
+
+import numpy as np
+import agama
+agama.setUnits(length=1, velocity=1, mass=1)
+
+def get_rot_mat(x, y, z, vx, vy, vz):
+    Lx = y * vz - z * vy
+    Ly = z * vx - x * vz
+    Lz = x * vy - y * vx
+    r = (x*x + y*y + z*z)**0.5
+    L = (Lx*Lx + Ly*Ly + Lz*Lz)**0.5
+    # rotation matrices transforming from the host to the satellite frame for each point on the trajectory
+    R = np.zeros((len(x), 3, 3))
+    R[:,0,0] = x/r
+    R[:,0,1] = y/r
+    R[:,0,2] = z/r
+    R[:,2,0] = Lx/L
+    R[:,2,1] = Ly/L
+    R[:,2,2] = Lz/L
+    R[:,1,0] = R[:,0,2] * R[:,2,1] - R[:,0,1] * R[:,2,2]
+    R[:,1,1] = R[:,0,0] * R[:,2,2] - R[:,0,2] * R[:,2,0]
+    R[:,1,2] = R[:,0,1] * R[:,2,0] - R[:,0,0] * R[:,2,1]
+    return R, L, r
+
+def get_d2Phi_dr2(pot_host, x, y, z):
+    # compute  the second derivative of potential by spherical radius
+    r = (x*x + y*y + z*z)**0.5
+    der = pot_host.forceDeriv(np.column_stack([x,y,z]))[1]
+    d2Phi_dr2 = -(x**2  * der[:,0] + y**2  * der[:,1] + z**2  * der[:,2] +
+                  2*x*y * der[:,3] + 2*y*z * der[:,4] + 2*z*x * der[:,5]) / r**2
+    return d2Phi_dr2
+
+def create_ic_chen25(rng, pot_host, orb_sat, mass_sat):
+    N = len(orb_sat)
+    x, y, z, vx, vy, vz = orb_sat.T
+    R, L, r = get_rot_mat(x, y, z, vx, vy, vz)
+    d2Phi_dr2 = get_d2Phi_dr2(pot_host, x, y, z)
+
+    # compute the tidal radius at this radius for each point on the trajectory
+    Omega = L / r**2
+    r_tidal = (agama.G * mass_sat / (Omega**2 - d2Phi_dr2))**(1./3)
+
+    # assign positions and velocities (in the satellite reference frame) of particles
+    r_tidal = np.repeat(r_tidal, 2)
+
+    mean = np.array([1.6, -30, 0, 1, 20, 0])
+
+    cov = np.array([
+        [0.1225,   0,   0, 0, -4.9,   0],
+        [     0, 529,   0, 0,    0,   0],
+        [     0,   0, 144, 0,    0,   0],
+        [     0,   0,   0, 0,    0,   0],
+        [  -4.9,   0,   0, 0,  400,   0],
+        [     0,   0,   0, 0,    0, 484],
+    ])
+
+    posvel = rng.multivariate_normal(mean, cov, size=2*N)
+
+    Dr = posvel[:, 0] * r_tidal
+    v_esc = np.sqrt(2 * agama.G * mass_sat / Dr)
+    Dv = posvel[:, 3] * v_esc
+
+    # convert degrees to radians
+    phi = posvel[:, 1] * np.pi / 180
+    theta = posvel[:, 2] * np.pi / 180
+    alpha = posvel[:, 4] * np.pi / 180
+    beta = posvel[:, 5] * np.pi / 180
+
+    dx = Dr * np.cos(theta) * np.cos(phi)
+    dy = Dr * np.cos(theta) * np.sin(phi)
+    dz = Dr * np.sin(theta)
+    dvx = Dv * np.cos(beta) * np.cos(alpha)
+    dvy = Dv * np.cos(beta) * np.sin(alpha)
+    dvz = Dv * np.sin(beta)
+
+    dq = np.column_stack([dx, dy, dz])
+    dp = np.column_stack([dvx, dvy, dvz])
+
+    ic_stream = np.tile(orb_sat, 2).reshape(2*N, 6)
+
+    # trailing arm
+    ic_stream[::2,0:3] += np.einsum("ni,nij->nj", dq[::2], R)
+    ic_stream[::2,3:6] += np.einsum("ni,nij->nj", dp[::2], R)
+
+    # leading arm
+    ic_stream[1::2,0:3] += np.einsum("ni,nij->nj", -dq[1::2], R)
+    ic_stream[1::2,3:6] += np.einsum("ni,nij->nj", -dp[1::2], R)
+
+    return ic_stream
+
+def integrate_prog(time_total, trajsize, pot_host, posvel_sat):
+    # integrate the orbit of the progenitor from its present-day posvel (at time t=0)
+    # back in time for an interval time_total, storing the trajectory at num_steps points
+    time_sat, orbit_sat = agama.orbit(potential=pot_host, ic=posvel_sat,
+        time=time_total, trajsize=trajsize)
+    if time_total < 0:
+        # reverse the arrays to make them increasing in time
+        time_sat  = time_sat [::-1]
+        orbit_sat = orbit_sat[::-1]
+    return time_sat, orbit_sat
+
+def create_stream(create_ic_method, rng, time_total, num_particles, pot_host, posvel_sat, mass_sat, 
+    pot_sat=None, nhalf_release=None, **kwargs):
+
+    if nhalf_release is None:
+        trajsize = num_particles//2
+    else:
+        trajsize = len(nhalf_release)
+    assert len(nhalf_release) == trajsize
+    time_sat, orbit_sat = integrate_prog(time_total, trajsize, pot_host, posvel_sat)
+
+    # at each point on the trajectory, create a pair of seed initial conditions
+    # for particles released at Lagrange points
+    if nhalf_release is None:
+        release_points = orbit_sat
+        time_seed = np.repeat(time_sat, 2)
+    else:
+        release_points = np.repeat(orbit_sat, nhalf_release, axis=0)
+        time_seed = np.repeat(time_sat, 2*nhalf_release)
+    ic_stream = create_ic_method(rng, pot_host, release_points, mass_sat, **kwargs)
+
+    if pot_sat is None:
+        pot_tot = pot_host
+    else:
+        # include the progenitor"s potential
+        traj = np.column_stack([time_sat, orbit_sat])
+        pot_traj = agama.Potential(potential=pot_sat, center=traj)
+        pot_tot = agama.Potential(pot_host, pot_traj)
+
+    dur = -time_seed if time_total<0 else time_total-time_seed
+    ic_stream = ic_stream[dur>0]
+    time_seed = time_seed[dur>0]
+    xv_stream = np.vstack(agama.orbit(potential=pot_tot,
+        ic=ic_stream, time=dur[dur>0], timestart=time_seed, trajsize=1)[:,1])
+    return time_sat, orbit_sat, xv_stream, ic_stream, time_seed

StarStream/optimize.py

@@ -0,0 +1,42 @@
+# Licensed under BSD-3-Clause License - see LICENSE
+
+import numpy as np
+
+__all__ = [
+    "optimize_mle"
+]
+
+def optimize_mle(data, err, pdf_signal, pdf_bg, frac_list=None):
+	"""
+	data: array-like (N, M): N data points with M dimensions
+	err: array-like (N, M): N data points with M dimensions
+	frac_list: list of fractions for evaluating lnL, default: None
+	pdf_signal, pdf_bg: KernelPDF objects
+	"""
+
+	prob_signal = pdf_signal.eval_pdf(data, err_eval=err)
+	prob_bg = pdf_bg.eval_pdf(data, err_eval=err)
+	prob_bg[prob_bg==0] = np.min(prob_bg[prob_bg!=0])
+
+	lnL_base = np.sum(np.log(prob_bg)) # no signal
+
+	if frac_list is None:
+		frac_list = 10.0**np.arange(-6.0, -1.0+0.001, 0.02)
+
+	lnL_list = []
+	for frac in frac_list:
+	    lnLi = np.log(prob_signal*frac + prob_bg*(1-frac))
+	    lnL_list.append(np.sum(lnLi))
+	lnL_list = np.array(lnL_list)
+
+	dlnL_list = lnL_list-lnL_base
+
+	fbest = frac_list[np.argmax(dlnL_list)]
+	dlnL_max = np.max(dlnL_list)
+
+	if dlnL_max <=0:
+		fbest = 0.0
+
+	prob = fbest*prob_signal / (fbest*prob_signal+(1-fbest)*prob_bg)
+
+	return fbest, prob, frac_list, dlnL_list