ocelot 0.3.1a0__py3-none-any.whl

ocelot/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.4.0'

ocelot/calculate/__init__.py

@@ -0,0 +1 @@
+# Todo decide what to make top-level here

ocelot/calculate/calculate.py

@@ -0,0 +1,207 @@
+"""A set of functions for calculating typical cluster parameters.
+
+Todo: error treatment here could be made more bayesian
+"""
+
+from typing import Optional
+from astropy.coordinates import SkyCoord
+
+import numpy as np
+import pandas as pd
+
+from .constants import mas_per_yr_to_rad_per_s, pc_to_m, deg_to_rad
+
+
+def _handle_ra_discontinuity(ra_data, middle_ras_raise_error=True):
+    """Tries to detect when the ras in a field cross the (0, 360) ra discontinuity and returns corrected results. Will
+    raise an error if ras are all over the place (which will happen e.g. at very high declinations) in which you
+    ought to instead switch to a method free of spherical distortions.
+
+    Args:
+        ra_data (pd.Series or np.ndarray): data on ras.
+        middle_ras_raise_error (bool): whether or not a cluster having right ascensions in all ranges [0, 90), [90, 270]
+            and (270, 360] raises an error. The error here indicates that this cluster has extreme spherical
+            discontinuities (e.g. it's near a coordinate pole) and that the mean ra and mean dec will be inaccurate.
+            Default: True
+
+    Returns:
+        ra_data but corrected for distortions. If values are both <90 and >270, the new ra data will be in the range
+            (-90, 90).
+
+    """
+    # Firstly, check that the ras are valid ras
+    if np.any(ra_data >= 360) or np.any(ra_data < 0):
+        raise ValueError(
+            "at least one input ra value was invalid! Ras must be in the range [0, 360)."
+        )
+
+    # Next, grab all the locations of dodgy friends and check that all three aren't ever in use at the same time
+    low_ra = ra_data < 90
+    high_ra = ra_data > 270
+    middle_ra = np.logical_not(np.logical_or(low_ra, high_ra))
+
+    # Proceed if we have both low and high ras
+    if np.any(low_ra) and np.any(high_ra):
+        # Stop if we have middle too (would imply stars everywhere or an extreme dec value)
+        if np.any(middle_ra) and middle_ras_raise_error:
+            raise ValueError(
+                "ra values are in all three ranges: [0, 90), [90, 270] and (270, 360). This cluster can't "
+                "be processed by this function! Spherical distortions must be removed first."
+            )
+
+        # Otherwise, apply the discontinuity removal
+        else:
+            # Make a copy so nothing weird happens
+            ra_data = ra_data.copy()
+
+            # And remove the distortion for all high numbers
+            ra_data[high_ra] = ra_data[high_ra] - 360
+
+    return ra_data
+
+
+def mean_radius(
+    data_gaia: pd.DataFrame,
+    membership_probabilities: Optional[np.ndarray] = None,
+    already_inferred_parameters: Optional[dict] = None,
+    key_ra: str = "ra",
+    key_ra_error: str = "ra_error",
+    key_dec: str = "dec",
+    key_dec_error: str = "dec_error",
+    distance_to_use: str = "inverse_parallax",
+    middle_ras_raise_error: bool = True,
+    **kwargs,
+):
+    """Produces various radius statistics on a given cluster, finding its sky location and three radii: the core, tidal
+    and 50% radius.
+
+    Done in a very basic, frequentist way, whereby means are weighted based on the membership probabilities (if
+    specified).
+
+    N.B. unlike the above functions, errors do *not change the mean* as this would potentially bias the
+    estimates towards being dominated by large, centrally-located stars within clusters (that have generally lower
+    velocities.) Hence, estimates here will be less precise but hopefully more accurate.
+
+    Todo: add error estimation to this function (hard)
+
+    Todo: add galactic l, b to the output of this function
+
+    Args:
+        data_gaia (pd.DataFrame): Gaia data for the cluster in the standard format (e.g. as in DR2.)
+        membership_probabilities (optional, np.ndarray): membership probabilities for the cluster. When specified,
+            they can increase or decrease the effect of certain stars on the mean.
+        already_inferred_parameters (optional, dict): a parameter dictionary of the mean distance and proper motion.
+            Otherwise, this function calculates a version.
+        key_ra (str): Gaia parameter name.
+        key_ra_error (str): Gaia parameter name.
+        key_dec (str): Gaia parameter name.
+        key_dec_error (str): Gaia parameter name.
+        distance_to_use (str): which already inferred distance to use to convert angular radii to parsecs.
+            Default: "inverse_parallax"
+        middle_ras_raise_error (bool): whether or not a cluster having right ascensions in all ranges [0, 90), [90, 270]
+            and (270, 360] raises an error. The error here indicates that this cluster has extreme spherical
+            discontinuities (e.g. it's near a coordinate pole) and that the mean ra and mean dec will be inaccurate.
+            Default: True
+
+    Returns:
+        a dict, formatted with:
+            {
+            # Position
+            "ra": ra of the cluster
+            "ra_error": error on the above
+            "dec": dec of the cluster
+            "dec_error": error on the above
+
+            # Angular radii
+            "ang_radius_50": median ang. distance from the center, i.e.angular radius of the cluster with 50% of members
+            "ang_radius_50_error": error on the above
+            "ang_radius_c": angular King's core radius of the cluster
+            "ang_radius_c_error": error on the above
+            "ang_radius_t": maximum angular distance from the center, i.e. angular King's tidal radius of the cluster
+            "ang_radius_t_error": error on the above
+
+            # Parsec radii
+            "radius_50": median distance from the center, i.e.radius of the cluster with 50% of members
+            "radius_50_error": error on the above
+            "radius_c": King's core radius of the cluster
+            "radius_c_error": error on the above
+            "radius_t": maximum distance from the center, i.e. King's tidal radius of the cluster
+            "radius_t_error": error on the above
+            }
+
+    """
+    inferred_parameters = {}
+    sqrt_n_stars = np.sqrt(data_gaia.shape[0])
+
+    # Grab the distances if they aren't specified - we'll need them in a moment!
+    if already_inferred_parameters is None:
+        already_inferred_parameters = mean_distance(data_gaia, membership_probabilities)
+
+    # Estimate the ra, dec of the cluster as the weighted mean
+    ra_data = _handle_ra_discontinuity(
+        data_gaia[key_ra], middle_ras_raise_error=middle_ras_raise_error
+    )
+
+    inferred_parameters["ra"] = np.average(ra_data, weights=membership_probabilities)
+    inferred_parameters["ra_std"] = _weighted_standard_deviation(
+        ra_data, membership_probabilities
+    )
+    inferred_parameters["ra_error"] = inferred_parameters["ra_std"] / sqrt_n_stars
+
+    if inferred_parameters["ra"] < 0:
+        inferred_parameters["ra"] += 360
+
+    inferred_parameters["dec"] = np.average(
+        data_gaia[key_dec], weights=membership_probabilities
+    )
+    inferred_parameters["dec_std"] = _weighted_standard_deviation(
+        data_gaia[key_dec], membership_probabilities
+    )
+    inferred_parameters["dec_error"] = inferred_parameters["dec_std"] / sqrt_n_stars
+
+    inferred_parameters["ang_dispersion"] = np.sqrt(
+        inferred_parameters["ra_std"] ** 2 + inferred_parameters["dec_std"] ** 2
+    )
+
+    # Calculate how far every star in the cluster is from the central point
+    center_skycoord = SkyCoord(
+        ra=inferred_parameters["ra"], dec=inferred_parameters["dec"], unit="deg"
+    )
+    star_skycoords = SkyCoord(
+        ra=data_gaia[key_ra].to_numpy(), dec=data_gaia[key_dec].to_numpy(), unit="deg"
+    )
+
+    distances_from_center = center_skycoord.separation(star_skycoords).degree
+
+    # And say something about the radii in this case
+    inferred_parameters["ang_radius_50"] = np.median(distances_from_center)
+    inferred_parameters["ang_radius_50_error"] = np.nan
+
+    inferred_parameters["ang_radius_c"] = np.nan
+    inferred_parameters["ang_radius_c_error"] = np.nan
+
+    inferred_parameters["ang_radius_t"] = np.max(distances_from_center)
+    inferred_parameters["ang_radius_t_error"] = np.nan
+
+    # Convert the angular distances into parsecs
+    inferred_parameters["radius_50"] = (
+        np.tan(inferred_parameters["ang_radius_50"] * deg_to_rad)
+        * already_inferred_parameters[distance_to_use]
+    )
+    inferred_parameters["radius_50_error"] = np.nan
+    inferred_parameters["radius_c"] = np.nan
+    inferred_parameters["radius_c_error"] = np.nan
+    inferred_parameters["radius_t"] = (
+        np.tan(inferred_parameters["ang_radius_t"] * deg_to_rad)
+        * already_inferred_parameters[distance_to_use]
+    )
+    inferred_parameters["radius_t_error"] = np.nan
+
+    return inferred_parameters
+
+
+def all_statistics():
+    """
+    """
+    # Todo refactor this (and other high-level calculation methods)
+    pass

ocelot/calculate/distance.py

@@ -0,0 +1 @@
+"""Functions for working with cluster distances."""

ocelot/calculate/generic.py

@@ -0,0 +1,97 @@
+"""Generic calculation utilities used across the module."""
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from typing import Optional, Union
+from ocelot.util.check import _check_matching_lengths_of_non_nones
+from scipy.stats import directional_stats
+
+
+def _weighted_standard_deviation(x: ArrayLike, weights: Optional[ArrayLike] = None):
+    """Computes weighted standard deviation. Uses method from
+    https://stackoverflow.com/a/52655244/12709989.
+    """
+    # Todo: not sure that this deals with small numbers of points correctly!
+    #   See: unit test fails when v. few points used
+    return np.sqrt(np.cov(x, aweights=weights))
+
+
+def standard_error(
+    standard_deviation: Union[ArrayLike[Union[float, int]], float, int],
+    number_of_measurements: Union[ArrayLike[int], int],
+) -> float:
+    """Calculates the standard error on the mean of some parameter given the standard
+    deviation.
+
+    Parameters
+    ----------
+    standard_deviation : array-like, float, or int
+        Standard deviation(s)
+    number_of_measurements : array-like of ints, int
+        Number of measurements used to find standard deviation.
+
+    Returns
+    -------
+    standard_error : float
+    """
+    return standard_deviation / np.sqrt(number_of_measurements)
+
+
+def mean_and_deviation(
+    values: ArrayLike,
+    weights: Optional[ArrayLike] = None,
+) -> tuple[float]:
+    """Calculates the mean and standard deviation of some set of values.
+
+    Parameters
+    ----------
+    values : array-like
+        Values to calculate mean and standard deviation of.
+    weights : array-like, optional
+        Array of weights to use to compute a weighted mean and average.
+
+    Returns
+    -------
+    mean : float
+        Mean of values.
+    std : float
+        Standard deviation of values.
+    """
+    _check_matching_lengths_of_non_nones(values, weights)
+
+    return (
+        np.average(values, weights=weights),
+        _weighted_standard_deviation(values, weights),
+    )
+
+
+def lonlat_to_unitvec(longitudes: NDArray, latitudes: NDArray):
+    """Converts longitudes and latitudes to unit vectors on a unit sphere. Uses method 
+    at https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates.
+    Assumes that latitudes is in the range [-pi / 2, pi / 2] as is common in 
+    astronomical unit systems.
+    """
+    x = np.cos(latitudes) * np.cos(longitudes)
+    y = np.cos(latitudes) * np.sin(longitudes)
+    z = np.sin(latitudes)
+    return np.column_stack((x, y, z))
+
+
+def unitvec_to_lonlat(unit_vectors: NDArray):
+    """Converts unit vectors on a unit sphere to longitudes and latitudes. See 
+    `lonlat_to_unitvec` for more details.
+    """
+    x, y, z = [column.ravel() for column in np.hsplit(unit_vectors, 3)]
+    longitudes = np.arctan2(y, x)
+    latitudes  = np.arcsin(z / np.sqrt(x**2 + y**2 + z**2))
+    return longitudes, latitudes
+
+
+def spherical_mean(longitudes: ArrayLike, latitudes: ArrayLike):
+    """Calculates the spherical mean of angular positions."""
+    longitudes = np.asarray_chkfinite(longitudes)
+    latitudes = np.asarray_chkfinite(latitudes)
+
+    unit_vectors = lonlat_to_unitvec(longitudes, latitudes)
+    mean_unit_vector = directional_stats(unit_vectors).mean_direction
+    mean_lon, mean_lat = unitvec_to_lonlat(mean_unit_vector)
+    return mean_lon[0], mean_lat[0]

ocelot/calculate/motion.py

@@ -0,0 +1,11 @@
+"""Tools for calculating values based on proper motions and/or velocities."""
+
+
+def radial_velocity(velocities, with_frame_correction=False):
+    # Todo
+    pass
+
+
+def proper_motion(proper_motions, with_frame_correction=False):
+    # Todo
+    pass

ocelot/calculate/position.py

@@ -0,0 +1,64 @@
+"""Different methods for calculating the center of a cluster in spherical coordinates. 
+(This is actually oddly difficult, thanks to how spheres work. Damn spheres.)
+"""
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ocelot.calculate.generic import spherical_mean
+
+
+def mean_position(
+    longitudes: ArrayLike, latitudes: ArrayLike, degrees=True
+) -> tuple[float]:
+    """Calculates the spherical mean of angular positions, specified as longitudes and
+    latitudes. This uses directional statistics to do so in a way that is aware of
+    discontinuities, such as the fact that 0° = 360°.
+
+    Parameters
+    ----------
+    longitudes : array-like
+        Array of longitudinal positions of stars in your cluster (e.g. right ascensions
+        or galactic longitudes.) Assumed to be in the range [0°, 360°].
+    latitudes : array-like
+        Array of latitudinal positions of stars in your cluster (e.g. declinations or 
+        galactic latitudes.) Assumed to be in the range [-90°, 90°].
+    degrees : bool
+        Whether longitudes and latitudes are in degrees, and whether to return an answer
+        in degrees. Defaults to True. If False, longitudes and latitudes are assumed to
+        be in radians, with ranges [0, 2π] and [-π/2, π/2] respectively.
+
+    Returns
+    -------
+    mean_longitude : float
+    mean_latitude : float
+
+    Notes
+    -----
+    This function explicitly assumes that your star cluster *has* a well-defined mean
+    position. Some configurations (such as points uniformly distributed in at least one
+    axis of a sphere) will not have a meaningful mean position.
+
+    Internally, this function uses `scipy.stats.directional_stats`, with a definition
+    taken from [1]. See [2] for more background.
+
+    References
+    ----------
+    [1] Mardia, Jupp. (2000). Directional Statistics (p. 163). Wiley.
+    [2] https://en.wikipedia.org/wiki/Directional_statistics
+    """
+    if degrees:
+        longitudes, latitudes = np.radians(longitudes), np.radians(latitudes)
+    mean_lon, mean_lat = spherical_mean(longitudes, latitudes)
+    if degrees:
+        return np.degrees(mean_lon), np.degrees(mean_lat)
+    return mean_lon, mean_lat
+
+
+def mode_position():
+    """Attempts to find the mode of a star cluster's 2D on-sky distribution. This is a 
+    better estimator than the mean position for clusters that are assymmetric, which is
+    often the case for clusters with assymetric tidal tails (e.g. due to one side being
+    more easily detected than the other.)
+    """
+    # Todo
+    pass

ocelot/calculate/profile.py

@@ -0,0 +1,225 @@
+# Todo: refactor to new structure (e.g. class-based, mirroring scipy.stats maybe?)
+
+from typing import Union, Optional
+
+import numpy as np
+
+
+def king_surface_density(
+    r_values: Union[float, np.ndarray],
+    r_core: float,
+    r_tidal: float,
+    normalise: bool = False,
+):
+    """Computes the King surface density (King 1962, equation 14) given the three parameters. Can take vectorised input.
+    Will return the surface density per square unit expected at a distance r_values from the core of the cluster.
+
+    Valid only for:
+        0 <= r < r_tidal (0 elsewhere - this is checked internally)
+        0 < r_core < r_tidal (raises an error if not the case)
+
+    Args:
+        r_values (float or np.ndarray): r values to compute the surface density at.
+        r_core (float): the core radius of the cluster.
+        r_tidal (float): the tidal radius of the cluster.
+        normalise (bool): whether or not to return the normalised King surface density profile, calculated
+            numerically.
+            Default: False
+
+    Returns:
+        a float or array of floats of the surface density for the cluster.
+
+    """
+    _check_core_and_tidal_radii(r_core, r_tidal)
+    r_values = _check_r_values(r_values)
+
+    # Constants
+    rt_rc = r_tidal / r_core
+    a = (1 + rt_rc**2) ** -0.5
+
+    # Compute normalisation constant if desired
+    if normalise:
+        term_1 = r_core * np.arctan(rt_rc)
+        term_2 = -2 * a * r_core * np.log(rt_rc + 1 / a)
+        term_3 = r_tidal * a**2
+        normalisation_constant = 1 / (term_1 + term_2 + term_3)
+    else:
+        normalisation_constant = 1.0
+
+    # Work out where  0 <= r < r_tidal
+    r_valid = np.logical_and(r_values >= 0, r_values < r_tidal)
+    r_invalid = np.invert(r_valid)
+
+    # Compute result
+    reduced_r_values = (1 + (r_values[r_valid] / r_core) ** 2) ** (-0.5)
+    result = np.empty(r_values.shape)
+    result[r_valid] = normalisation_constant * (reduced_r_values - a) ** 2
+    result[r_invalid] = 0
+
+    return result
+
+
+def _check_r_values(r_values):
+    # Convert to a np array that's at least 1d, and check that it isn't bad
+    r_values = np.atleast_1d(r_values)
+    if not np.isfinite(r_values).all():
+        raise ValueError("invalid r_values (e.g. nan or inf) are not allowed!")
+    if np.any(r_values < 0):
+        raise ValueError("all r_values must be positive or zero.")
+    return r_values
+
+
+def _check_core_and_tidal_radii(r_core, r_tidal):
+    if not np.isfinite(r_core).all() or not np.isfinite(r_tidal).all():
+        raise ValueError("input parameters must be finite!")
+    if r_tidal < r_core or r_core < 0 or r_tidal < 0:
+        raise ValueError("parameters must satisfy 0 < r_core < r_tidal")
+
+
+def king_surface_density_fast(r_values: np.ndarray, r_core: float, r_tidal: float):
+    """Computes the King surface density (King 1962, equation 14) given the three parameters. Can take vectorised input.
+    Fast version intended for use with random sampling - it has no checks and cannot be normalised! Be careful!
+
+    Valid only for:
+        0 <= r < r_tidal (NOT CHECKED in this function!)
+        0 < r_core < r_tidal (NOT CHECKED in this function!)
+
+    Args:
+        r_values (np.ndarray): r values to compute the surface density at. Must be a numpy array!
+        r_core (float): the core radius of the cluster.
+        r_tidal (float): the tidal radius of the cluster.
+
+    Returns:
+        a float or array of floats of the surface density for the cluster.
+
+    """
+    # Constants
+    rt_rc = r_tidal / r_core
+    a = (1 + rt_rc**2) ** -0.5
+
+    # Compute result
+    return ((1 + (r_values / r_core) ** 2) ** (-0.5) - a) ** 2
+
+
+def sample_2d_king_profile(
+    r_core: float,
+    r_tidal: float,
+    n_samples: int,
+    seed=None,
+    oversampling_factor: float = 10,
+    return_generator: bool = False,
+):
+    """Samples a 2D King profile to return n_samples sample radii.
+
+    Valid only for:
+        0 < r_core < r_tidal
+
+    Args:
+        r_core (float): the core radius of the cluster.
+        r_tidal (float): the tidal radius of the cluster.
+        n_samples (int): the number of samples to generate.
+        seed (int, optional): the seed of the random number generator. Default: None.
+        oversampling_factor (float): how many times n_samples to generate each step, which helps to make sure that
+            enough samples are quickly generated in just one or two loops. Default: 10.
+        return_generator (bool): whether or not to return the numpy random number generator created. Default: False
+
+    Returns:
+        an array of sample radii of size n_samples, plus the random generator if return_generator==True.
+    """
+    _check_core_and_tidal_radii(r_core, r_tidal)
+
+    r_samples = np.empty(n_samples, dtype=float)
+    generator = np.random.default_rng(seed=seed)
+    completed_samples = 0
+
+    max_value = king_surface_density_fast(np.zeros(1), r_core, r_tidal)[0]
+
+    while completed_samples < n_samples:
+        remaining_samples = n_samples - completed_samples
+        samples_to_generate = int(
+            np.clip(remaining_samples * oversampling_factor, 10, np.inf)
+        )
+
+        # Generate some initial radius samples
+        test_r_values = generator.uniform(high=r_tidal, size=samples_to_generate)
+        test_king_values = king_surface_density_fast(test_r_values, r_core, r_tidal)
+        test_mcmc_values = generator.uniform(size=samples_to_generate, high=max_value)
+
+        # See which & how many are valid and save them!
+        valid_test_samples = test_king_values < test_mcmc_values
+        n_valid_test_samples = np.count_nonzero(valid_test_samples)
+
+        if n_valid_test_samples > remaining_samples:
+            r_samples[completed_samples:] = test_r_values[valid_test_samples][
+                :remaining_samples
+            ]
+            break
+
+        r_samples[
+            completed_samples : completed_samples + n_valid_test_samples
+        ] = test_r_values[valid_test_samples]
+        completed_samples += n_valid_test_samples
+
+    if return_generator:
+        return r_samples, generator
+    else:
+        return r_samples
+
+
+def sample_1d_king_profile(
+    r_core: float,
+    r_tidal: float,
+    n_samples: int,
+    seed: Optional[int] = None,
+    oversampling_factor: float = 10,
+):
+    """Samples a 1D King profile (e.g. useful to get line of sight distances from the center of a cluster.) Uses a
+    little trick - assumes that we're looking at the cluster side-on and removes the not-line-of-sight axis as if we
+    were looking at it from the front. (This is because I cba to work out a 1D profile and more to the point, I couldn't
+    fucking find one)
+
+    Todo: change this to using the strip density g(x), which I *think* could do this better...
+
+    Valid only for:
+        0 < r_core < r_tidal
+
+    Args:
+        r_core (float): the core radius of the cluster.
+        r_tidal (float): the tidal radius of the cluster.
+        n_samples (int): the number of samples to generate.
+        seed (int, optional): the seed of the random number generator. Default: None.
+        oversampling_factor (float): how many times n_samples to generate each step, which helps to make sure that
+            enough samples are quickly generated in just one or two loops. Default: 10.
+
+    Returns:
+        an array of sample radii of size n_samples
+    """
+    r_samples, generator = sample_2d_king_profile(
+        r_core,
+        r_tidal,
+        n_samples,
+        seed=seed,
+        oversampling_factor=oversampling_factor,
+        return_generator=True,
+    )
+
+    # Now, deproject this to 1D by giving every value an angle and then finding the 1D radius with the cosine
+    random_angles = generator.uniform(high=2 * np.pi, size=n_samples)
+
+    return r_samples * np.cos(random_angles)
+
+
+def king_number_density(r, r_core, r_tidal, k=1):
+    """Calculates the King1962 number density (eqn 18 in the paper.)
+    
+    Unnormalised by default (i.e. k=1.)
+    """
+    x = (r / r_core)**2
+    x_t = (r_tidal / r_core)**2
+
+    term_1 = np.log(1 + x)
+    term_2 = - 4 * ((1 + x)**(0.5) - 1) / (1 + x_t)**(0.5)
+    term_3 = x / (1 + x_t)
+    
+    return np.pi * r_core**2 * k * (term_1 + term_2 + term_3)
+

ocelot/cluster/__init__.py

@@ -0,0 +1,11 @@
+from . import epsilon
+from .nearest_neighbor import precalculate_nn_distances
+from .preprocess import cut_dataset, rescale_dataset, recenter_dataset
+from .resample import generate_gaia_covariance_matrix, resample_gaia_astrometry
+import warnings
+
+warnings.warn(
+    "ocelot.cluster API will change soon. Most objects will move, change, or be "
+    "deprecated.",
+    DeprecationWarning,
+)