Initial import of cluster simulation code

Author: emilyhunt

.gitignore

@@ -3,6 +3,7 @@
 *.png
 .vscode/
 *.code-workspace
+data/
 
 ## GITHUB PYTHON .gitignore FILE
 # Byte-compiled / optimized / DLL files
@@ -92,7 +93,7 @@ ipython_config.py
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
-# .python-version
+.python-version
 
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.

docs/CNAME

@@ -1,18 +1,18 @@
 [project]
-name = "ocelot"  # Required
-version = "0.3.1-alpha"  # Required
-description = "A toolbox for working with observations of star clusters."  # Optional
+name = "ocelot" # Required
+version = "0.3.1-alpha" # Required
+description = "A toolbox for working with observations of star clusters." # Optional
 readme = "README.md" # Optional
 requires-python = ">=3.8"
-license = {file = "LICENSE"}
-keywords = ["astronomy", "star cluster", "threading", "development"]  # Optional
+license = { file = "LICENSE" }
+keywords = ["astronomy", "star cluster", "threading", "development"] # Optional
 authors = [
-  {name = "Emily Hunt", email = "emily.hunt.physics@gmail.com" } # Optional
+  { name = "Emily Hunt", email = "emily.hunt.physics@gmail.com" }, # Optional
 ]
 maintainers = [
-  {name = "Emily Hunt", email = "emily.hunt.physics@gmail.com" } # Optional
+  { name = "Emily Hunt", email = "emily.hunt.physics@gmail.com" }, # Optional
 ]
-classifiers = [  # Optional
+classifiers = [ # Optional
   "Development Status :: 3 - Alpha",
   "Intended Audience :: Developers",
   "License :: OSI Approved :: MIT License",
@@ -26,21 +26,36 @@ classifiers = [  # Optional
 ]
 dependencies = [
   "numpy>1.21.0,<3.0",
+  "numba<1.0",
   "matplotlib>3.4.0,<4.0",
   "scikit-learn>0.24.0,<2.0",
   "scipy>1.0.0,<2.0",
   "pandas>1.0.0,<3.0",
+  "pyarrow",
   "astropy>4.0.0,<7.0",
   "healpy>1.13.0,<2.0",
+  "galpy>1.9,<2.0",
+  "imf @ git+https://github.com/keflavich/imf",
+  "dustmaps",
 ]
 
 
 [project.optional-dependencies]
-dev = ["check-manifest", "pytest", "mkdocs-material[imaging]", "mkdocstrings[python]>=0.18"]
-docs = ["mkdocs-material[imaging]", "mkdocstrings[python]>=0.18"]  # For building docs on GitHub
+dev = [
+  "check-manifest",
+  "pytest",
+  "mkdocs-material[imaging]",
+  "mkdocstrings[python]>=0.18",
+  "ruff",
+  "ezpadova @ git+https://github.com/mfouesneau/ezpadova",
+]
+docs = [
+  "mkdocs-material[imaging]",
+  "mkdocstrings[python]>=0.18",
+] # For building docs on GitHub
 test = ["pytest"]
 
-[project.urls]  # Optional
+[project.urls] # Optional
 "Homepage" = "https://ocelot-docs.org"
 "Bug Reports" = "https://github.com/emilyhunt/ocelot/issues"
 "Source" = "https://github.com/emilyhunt/ocelot"
@@ -50,9 +65,7 @@ test = ["pytest"]
 
 [tool.pytest.ini_options]
 pythonpath = "src"
-addopts = [
-    "--import-mode=importlib",
-]
+addopts = ["--import-mode=importlib"]
 
 [build-system]
 requires = ["setuptools>=43.0.0", "wheel"]

ocelot_data.zip

@@ -0,0 +1,3 @@
+# Scripts
+
+The following scripts are intended for various development purposes.

pyproject.toml

@@ -0,0 +1,56 @@
+"""Script to download isochrones from PARSEC with ezpadova. Intended for development
+use only. Downloads to ../data/.... .
+"""
+
+import ezpadova
+import pandas as pd
+import numpy as np
+import time
+from pathlib import Path
+
+
+outdir = Path("../data/isochrones/PARSEC_v1.2S")
+outdir.mkdir(exist_ok=True, parents=True)
+
+minimum_age = 5.0
+break_point = 8.0
+maximum_age = 10.15
+age_spacing = 0.01
+
+minimum_metallicity = -2.2
+maximum_metallicity = 0.5
+metallicity_spacing = 0.05
+
+photometric_system = "YBC_tab_mag_odfnew/tab_mag_gaiaEDR3.dat"
+
+
+metallicities = np.linspace(
+    minimum_metallicity,
+    maximum_metallicity,
+    num=int((maximum_metallicity - minimum_metallicity) / metallicity_spacing + 1),
+)
+for a_metallicity in metallicities:
+    print(f"Working on [M/H]={a_metallicity:.3f}")
+    # Download in two batches to get around size limits
+    isochrones_first_half = ezpadova.get_isochrones(
+        logage=(minimum_age, break_point, age_spacing),
+        MH=(a_metallicity, a_metallicity, 0.0),
+        photsys_file=photometric_system,
+        # track_parsec="parsec_CAF09_v1.2S",
+        # track_colibri="parsec_CAF09_v1.2S_NOV13",
+    )
+    isochrones_second_half = ezpadova.get_isochrones(
+        logage=(break_point + age_spacing, maximum_age, age_spacing),
+        MH=(a_metallicity, a_metallicity, 0.0),
+        photsys_file=photometric_system,
+        # track_parsec="parsec_CAF09_v1.2S",
+        # track_colibri="parsec_CAF09_v1.2S_NOV13",
+    )
+
+    isochrones = pd.concat(
+        [isochrones_first_half, isochrones_second_half], ignore_index=True
+    )
+    isochrones.to_parquet(outdir / f"mh={a_metallicity:.3f}.parquet")
+
+    print("Sleeping to avoid rate limit...")
+    time.sleep(5)

scripts/README.md

@@ -1 +1,8 @@
-__version__ = '0.4.0'
+import importlib.util
+from pathlib import Path
+
+__version__ = "0.4.0"
+
+# Set a library path, helping other modules to find the data directory etc.
+_module_path = Path(importlib.util.find_spec("ocelot").origin).parent
+_data_path = _module_path / "data"

scripts/download_parsec_isochrones.py

@@ -0,0 +1,2 @@
+"""A package for simulating star clusters."""
+from .cluster import Cluster, ClusterParameters  # noqa: F401

src/ocelot/__init__.py

@@ -0,0 +1,91 @@
+"""Functions for dealing with astrometry-related things."""
+
+from __future__ import annotations  # Necessary to type hint without cyclic import
+import numpy as np
+from ocelot.simulate.uncertainties import apply_gaia_astrometric_uncertainties
+from ocelot.calculate.profile import sample_1d_king_profile
+from ocelot.util.random import points_on_sphere
+import ocelot.simulate.cluster
+from astropy.coordinates import SkyCoord, CartesianRepresentation, CartesianDifferential
+from astropy import units as u
+from scipy.stats import multivariate_normal
+
+
+def generate_star_positions(cluster: ocelot.simulate.cluster.Cluster):
+    """Generates positions of member stars in polar coordinates relative to the center
+    of the cluster.
+    """
+    radii = sample_1d_king_profile(
+        cluster.parameters.r_core,
+        cluster.parameters.r_tidal,
+        cluster.stars,
+        seed=cluster.random_seed,
+    )
+    phis, thetas = points_on_sphere(
+        len(radii), phi_symmetric=False, seed=cluster.random_seed
+    )
+    x_values, y_values, z_values = spherical_to_cartesian(radii, thetas, phis)
+    return CartesianRepresentation(x_values, y_values, z_values, unit=u.pc)
+
+
+def spherical_to_cartesian(radii, thetas, phis):
+    """Converts from spherical to Cartesian coordinates. See:
+    https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates
+
+    Assumes radii ∈ [0, inf), thetas ∈ [0, pi], phis ∈ [0, 2pi)
+    """
+    x_values = radii * np.sin(thetas) * np.cos(phis)
+    y_values = radii * np.sin(thetas) * np.sin(phis)
+    z_values = radii * np.cos(thetas)
+    return x_values, y_values, z_values
+
+
+def generate_star_velocities(cluster: ocelot.simulate.cluster.Cluster):
+    """Generates the velocities of stars in a cluster."""
+    distribution = multivariate_normal(
+        mean=np.zeros(3),
+        cov=np.identity(3) * cluster.parameters.velocity_dispersion_1d,
+        seed=cluster.random_generator,
+    )
+    v_x, v_y, v_z = distribution.rvs(cluster.stars).T.reshape(3, -1)  # We also reshape to make sure a size-1 cluster is handled correctly
+    return CartesianDifferential(d_x=v_x, d_y=v_y, d_z=v_z, unit=u.m / u.s)
+
+
+def generate_true_star_astrometry(cluster: ocelot.simulate.cluster.Cluster):
+    """Generates the true values of cluster astrometry (not affected by errors)."""
+    positions = generate_star_positions(cluster)
+    velocities = generate_star_velocities(cluster)
+
+    # Do coordinate frame stuff to get final values (dont look astropy devs, dont tell
+    # me what I can and cant do, i live in a lawless realm, this works so it works)
+    cluster_center = cluster.parameters.get_position_as_skycoord(
+        frame="galactocentric", with_zeroed_proper_motions=True
+    ).cartesian
+    cluster_differential = cluster_center.differentials["s"]
+
+    final_positions = positions + cluster_center
+    final_velocities = velocities + cluster_differential
+
+    final_coords = SkyCoord(
+        CartesianRepresentation(final_positions, differentials=final_velocities),
+        frame="galactocentric",
+    ).transform_to("icrs")
+    final_coords_galactic = final_coords.transform_to("galactic")
+
+    # Assign these values to cluster df
+    cluster.cluster["ra"] = final_coords.ra.value
+    cluster.cluster["dec"] = final_coords.dec.value
+    cluster.cluster["l"] = final_coords_galactic.l.value
+    cluster.cluster["b"] = final_coords_galactic.b.value
+    cluster.cluster["pmra"] = np.nan  # Gets set later
+    cluster.cluster["pmdec"] = np.nan  # Gets set later
+    cluster.cluster["pmra_true"] = final_coords.pm_ra_cosdec.value
+    cluster.cluster["pmdec_true"] = final_coords.pm_dec.value
+    cluster.cluster["radial_velocity_true"] = final_coords.radial_velocity.value
+    cluster.cluster["parallax_true"] = 1000 / final_coords.distance.value
+
+
+def generate_cluster_astrometry(cluster: ocelot.simulate.cluster.Cluster):
+    """Generates the astrometry of clusters."""
+    generate_true_star_astrometry(cluster)
+    apply_gaia_astrometric_uncertainties(cluster)