snc-sf 1.1.1__tar.gz

snc_sf-1.1.1/LICENSE.md

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, Ilija Medan
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

snc_sf-1.1.1/PKG-INFO

@@ -0,0 +1,61 @@
+Metadata-Version: 2.3
+Name: snc-sf
+Version: 1.1.1
+Summary: 
+Author: imedan
+Author-email: ilija.medan@gmail.com
+Requires-Python: >=3.11
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: astropy (>=6.1.0)
+Requires-Dist: astroquery (>=0.4.7)
+Requires-Dist: gaiaunlimited (>=0.3.2,<0.4.0)
+Requires-Dist: healpy (>=1.17.1)
+Requires-Dist: jax (>=0.4.27)
+Requires-Dist: jaxlib (>=0.4.27)
+Requires-Dist: matplotlib (>=3.8.0)
+Requires-Dist: notebook (>=7.0.0)
+Requires-Dist: numpy (>=2.0.0)
+Requires-Dist: numpyro (>=0.14.0)
+Requires-Dist: polars (>=1.0.0)
+Requires-Dist: pyarrow (>=14.0.0)
+Requires-Dist: scipy (>=1.13.0)
+Requires-Dist: sdss-semaphore (>=0.2.4,<0.3.0)
+Requires-Dist: tqdm (>=4.60.0)
+Description-Content-Type: text/markdown
+
+# SDSS-V Solar Neighborhood Census Selection Function and Subpopulation Probabilities
+
+This repository provides the code to calculate the selection function for the SDSS-V Solar Neighborhood Census (SNC) relative to the Gaia Catalog of Nearby Stars ([GCNS](https://ui.adsabs.harvard.edu/abs/2021A%26A...649A...6G/abstract)). With this selection function, the code allows for the forward modeling of subpopulation probabilities across the HR diagram. A use case would be selecting all stars in the SNC with [Fe/H] < -1 and the forward model would evaluate the likely probability of selecting [Fe/H] < -1 stars from the GCNS in discrete bins across the HR diagram.
+
+## Installation
+
+The code can be installed with `pip`:
+```
+git clone https://github.com/imedan/snc_sf
+cd snc_sf
+pip install .
+```
+or in a fresh virtual environment with `poetry` to fully replicate the development environment
+```
+git clone https://github.com/imedan/snc_sf
+cd snc_sf
+conda create -n "snc_sf" python=3.11
+conda activate snc_sf
+pip install poetry
+poetry install
+```
+
+To fully utilize the code, SDSS-V and GCNS data is need. The SDSS-V DR19 `astra` summary data can be accessed [here](https://data.sdss.org/sas/dr19/spectro/astra/0.6.0/summary/). The required GCNS data will automatically be downloaded when first initializing a `snc_sf.selection_function.SNCSelectionFunction()` object.
+
+## Examples
+
+A number of examples using the code are located in [`notebooks/`](notebooks/). These examples are fully explained in the paper accompanying this work. In summary, the examples included related to the following:
+- [`dr19_obs.ipynb`](notebooks/dr19_obs.ipynb): Constructs the base SNC 100 pc dataset and creates the plots in Figure 1 of the paper.
+- [`SF_example_plots.ipynb`](notebooks/SF_example_plots.ipynb): Example of the selecrtion function for different regions on the sky, as shown in Figure 2 of the paper.
+- [`forward_model_GCNS_mock_proof.py`](notebooks/forward_model_GCNS_mock_proof.py): Script that replicates the example with mock data in Section 4.1 and Figure 3 in the paper.
+- [`forward_model_halpha_ex.py`](notebooks/forward_model_halpha_ex.py): Script that replicates the example examing the distribution of H-alpha emitters across the HR diagram in Section 4.2 and Figure 4 in the paper.
+- [`forward_model_mass_dens.py`](notebooks/forward_model_mass_dens.py): Script that replicates the example that calculates the stellar number desnity in bins of metallicity and mass in Section 4.3 in the paper.
+- [`appendix_example.ipynb`](notebooks/appendix_example.ipynb): Example to demonstrate basic code usage as illustrated in the Appendix of the paper.

snc_sf-1.1.1/README.md

@@ -0,0 +1,33 @@
+# SDSS-V Solar Neighborhood Census Selection Function and Subpopulation Probabilities
+
+This repository provides the code to calculate the selection function for the SDSS-V Solar Neighborhood Census (SNC) relative to the Gaia Catalog of Nearby Stars ([GCNS](https://ui.adsabs.harvard.edu/abs/2021A%26A...649A...6G/abstract)). With this selection function, the code allows for the forward modeling of subpopulation probabilities across the HR diagram. A use case would be selecting all stars in the SNC with [Fe/H] < -1 and the forward model would evaluate the likely probability of selecting [Fe/H] < -1 stars from the GCNS in discrete bins across the HR diagram.
+
+## Installation
+
+The code can be installed with `pip`:
+```
+git clone https://github.com/imedan/snc_sf
+cd snc_sf
+pip install .
+```
+or in a fresh virtual environment with `poetry` to fully replicate the development environment
+```
+git clone https://github.com/imedan/snc_sf
+cd snc_sf
+conda create -n "snc_sf" python=3.11
+conda activate snc_sf
+pip install poetry
+poetry install
+```
+
+To fully utilize the code, SDSS-V and GCNS data is need. The SDSS-V DR19 `astra` summary data can be accessed [here](https://data.sdss.org/sas/dr19/spectro/astra/0.6.0/summary/). The required GCNS data will automatically be downloaded when first initializing a `snc_sf.selection_function.SNCSelectionFunction()` object.
+
+## Examples
+
+A number of examples using the code are located in [`notebooks/`](notebooks/). These examples are fully explained in the paper accompanying this work. In summary, the examples included related to the following:
+- [`dr19_obs.ipynb`](notebooks/dr19_obs.ipynb): Constructs the base SNC 100 pc dataset and creates the plots in Figure 1 of the paper.
+- [`SF_example_plots.ipynb`](notebooks/SF_example_plots.ipynb): Example of the selecrtion function for different regions on the sky, as shown in Figure 2 of the paper.
+- [`forward_model_GCNS_mock_proof.py`](notebooks/forward_model_GCNS_mock_proof.py): Script that replicates the example with mock data in Section 4.1 and Figure 3 in the paper.
+- [`forward_model_halpha_ex.py`](notebooks/forward_model_halpha_ex.py): Script that replicates the example examing the distribution of H-alpha emitters across the HR diagram in Section 4.2 and Figure 4 in the paper.
+- [`forward_model_mass_dens.py`](notebooks/forward_model_mass_dens.py): Script that replicates the example that calculates the stellar number desnity in bins of metallicity and mass in Section 4.3 in the paper.
+- [`appendix_example.ipynb`](notebooks/appendix_example.ipynb): Example to demonstrate basic code usage as illustrated in the Appendix of the paper.

snc_sf-1.1.1/pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "snc-sf"
+version = "1.1.1"
+description = ""
+authors = [
+    {name = "imedan",email = "ilija.medan@gmail.com"}
+]
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "astropy (>=6.1.0)",
+    "numpy (>=2.0.0)",
+    "matplotlib (>=3.8.0)",
+    "healpy (>=1.17.1)",
+    "tqdm (>=4.60.0)",
+    "sdss-semaphore (>=0.2.4,<0.3.0)",
+    "gaiaunlimited (>=0.3.2,<0.4.0)",
+    "polars (>=1.0.0)",
+    "jax (>=0.4.27)",
+    "jaxlib (>=0.4.27)",
+    "notebook (>=7.0.0)",
+    "numpyro (>=0.14.0)",
+    "pyarrow (>=14.0.0)",
+    "astroquery (>=0.4.7)",
+    "scipy (>=1.13.0)",
+]
+
+[tool.poetry]
+packages = [{include = "snc_sf", from = "src"}]
+
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.group.dev.dependencies]
+pytest = ">=8.0.0"

snc_sf-1.1.1/src/snc_sf/optimize.py

@@ -0,0 +1,106 @@
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import logsumexp
+from jax.nn import sigmoid
+from jax.experimental.sparse import BCOO
+try:
+    from jaxlib._jax import ArrayImpl
+except ModuleNotFoundError:
+    from jaxlib.xla_extension import ArrayImpl
+
+
+@jax.jit
+def compute_single_loglike(p: ArrayImpl, A_jk_bcoo: BCOO,
+                           S: ArrayImpl, idx_mod: ArrayImpl) -> float:
+    """
+    Compute the log likelihood for a single posterior draw
+    of the selection function
+
+    Parameters
+    ----------
+    p: jaxlib._jax.ArrayImpl
+        Probability of a subpopulation of GCNS stars being in bins
+        across the HR diagram. This is a 1D array raveled
+        from a 2D array of G - RP vs M_G
+    
+    A_jk_bcoo: jax.experimental.sparse.BCOO
+        Sparse matrix of the sum of the selection function values
+        for the intersection of each selection function bin with
+        each HR diagram bin. Should be of shape
+        (1D ravel of selection bins, 1D ravel of HR diagram bins)
+    
+    S: jaxlib._jax.ArrayImpl
+        Selection function values for the obserbed data in a subpopulation.
+
+    idx_mod: jaxlib._jax.ArrayImpl
+        The HR diagram 1D raveled indexes on the HR diagram of the observed
+        data in the subpopulation
+    
+    Returns
+    --------
+    log_like: float
+        The poisson point process log likelihood
+    """
+    Z = A_jk_bcoo @ p
+    mask = (S > 0)
+    p_obs = p[idx_mod] * S
+    safe_p_obs = jnp.where(mask, jnp.maximum(p_obs, 1e-300), 1.0)
+    log_like_obs = jnp.sum(jnp.log(safe_p_obs))
+    log_like_norm = jnp.sum(Z)
+    log_like = log_like_obs - log_like_norm
+    return log_like
+
+
+def objective_jax_multi(theta: ArrayImpl, S_data: ArrayImpl, A_jks_bcoo: tuple,
+                        idx_mod_data: ArrayImpl) -> float:
+    """
+    Compute the total log likelihood
+
+    Parameters
+    ----------
+    theta: jaxlib._jax.ArrayImpl
+        The sigmoid of theta gives the probability of a subpopulation
+        of GCNS stars being in bins
+        across the HR diagram. This is a 1D array raveled
+        from a 2D array of G - RP vs M_G.
+
+    S_data: jaxlib._jax.ArrayImpl
+        Selection function values for the obserbed data in a subpopulation.
+    
+    A_jks_bcoo: tuple
+        Sparse matrix of the sum of the selection function values
+        for the intersection of each selection function bin with
+        each HR diagram bin. Should be of shape
+        (1D ravel of selection bins, 1D ravel of HR diagram bins). This should
+        be a tuple where each index is for a different posterior draw
+
+    idx_mod_data: jaxlib._jax.ArrayImpl
+        The HR diagram 1D raveled indexes on the HR diagram of the observed
+        data in the subpopulation
+    
+    Returns
+    --------
+    neg_log_like: float
+        The total poisson point process negative log likelihood
+    """
+    p = sigmoid(theta)
+    p = jnp.nan_to_num(p, nan=0, posinf=1, neginf=0)
+
+    log_like_samples = []
+    for i in range(len(A_jks_bcoo)):
+        log_like_samples.append(compute_single_loglike(p, A_jks_bcoo[i],
+                                                       S_data[:, i], idx_mod_data))
+    log_like_samples = jnp.stack(log_like_samples)
+
+    log_like_samples = jnp.where(jnp.isfinite(log_like_samples), log_like_samples, -1e10)
+
+    neg_log_like = -1.0 * (logsumexp(log_like_samples) - jnp.log(len(A_jks_bcoo)))
+    return neg_log_like
+
+
+@jax.jit
+def sigmoid_inv(y: ArrayImpl):
+    """
+    Inverse of the sigmoid
+    """
+    return jnp.log(y / (1 - y))