petra-catalogs 0.1.0__tar.gz

petra_catalogs-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Aaron D. Johnson, Javier Roulete, Katerina Chatziioannou, Michele Vallisneri, Kyle Gersbach, Chris Trejo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

petra_catalogs-0.1.0/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.3
+Name: petra-catalogs
+Version: 0.1.0
+Summary: From LISA global fit to catalogs
+Author: Aaron D. Johnson
+Requires-Python: >=3.11,<4.0
+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: numpy (>=2.2.3,<3.0.0)
+Requires-Dist: pandas (>=2.2.3,<3.0.0)
+Requires-Dist: scipy (>=1.15.2,<2.0.0)
+Description-Content-Type: text/markdown
+
+# petra
+ From the LISA global fit to catalogs of single sources
+

petra_catalogs-0.1.0/README.md

@@ -0,0 +1,2 @@
+# petra
+ From the LISA global fit to catalogs of single sources

petra_catalogs-0.1.0/petra/__init__.py

@@ -0,0 +1,2 @@
+from .make_catalog import make_catalog_mv_normal
+from .posterior_chain import PosteriorChain

petra_catalogs-0.1.0/petra/aux_distributions.py

@@ -0,0 +1,76 @@
+import numpy as np
+from typing import List, Tuple
+
+
+def mv_normal_aux_distribution(sample: np.ndarray,
+                               aux_parameters: Tuple[List[np.ndarray]],
+                               source_index) -> np.ndarray:
+    """
+    Vectorized: Compute the logpdf of a multivariate normal distribution evaluated at each source in sample.
+    Accepts source_index as an int or array-like.
+
+    Parameters:
+        sample: numpy array of shape (num_sources, num_params_per_source)
+        aux_parameters: Tuple of lists (means, cov_matrices) where each is of length num_distributions.
+        source_index: int or array-like indices specifying which distributions to evaluate.
+
+    Returns:
+        If source_index is a scalar, returns an array of shape (num_sources,).
+        Otherwise, returns an array of shape (len(source_index), num_sources).
+    """
+    means, cov_matrices = aux_parameters
+    # Ensure source_index is array-like.
+    source_indices = np.atleast_1d(source_index)
+    means = np.array(means)[source_indices]         # shape: (n, d)
+    cov_matrices = np.array(cov_matrices)[source_indices]  # shape: (n, d, d)
+    
+    num_sources, d = sample.shape
+    # Broadcast sample to match the number of distributions.
+    diff = sample[np.newaxis, :, :] - means[:, np.newaxis, :]  # shape: (n, num_sources, d)
+    inv_cov = np.linalg.inv(cov_matrices)  # shape: (n, d, d)
+    # Compute log determinant for each covariance matrix.
+    _, logdet = np.linalg.slogdet(cov_matrices)  # shape: (n,)
+    # Compute the Mahalanobis term over distributions and sources.
+    mahal = np.einsum('nsi, nij, nsj -> ns', diff, inv_cov, diff)  # shape: (n, num_sources)
+    logpdf = -0.5 * (d * np.log(2 * np.pi) + logdet[:, None] + mahal)  # shape: (n, num_sources)
+    
+    # Return result: squeeze out axis if only one distribution was requested.
+    if logpdf.shape[0] == 1:
+        return logpdf[0]
+    return logpdf
+
+
+def uni_normal_aux_distribution_single_parameter(sample: np.ndarray,
+                                                 aux_parameters: Tuple[List[np.ndarray]],
+                                                 source_index,
+                                                 single_parameter: int) -> np.ndarray:
+    """
+    Vectorized: Compute the logpdf of a univariate normal distribution on a single parameter for each source.
+    Accepts source_index as an int or array-like.
+    
+    Parameters:
+        sample: numpy array of shape (num_sources, num_params_per_source)
+        aux_parameters: Tuple of lists (means, stds) for the univariate case.
+        source_index: int or array-like indices specifying which distributions to evaluate.
+        single_parameter: the parameter index to process.
+    
+    Returns:
+        If source_index is a scalar, returns an array of shape (num_sources,).
+        Otherwise, returns an array of shape (len(source_index), num_sources).
+    """
+    values = sample[:, single_parameter]  # shape: (num_sources,)
+    means, stds = aux_parameters
+    source_indices = np.atleast_1d(source_index)
+    means = np.array(means)[source_indices]  # shape: (n,)
+    stds = np.array(stds)[source_indices]      # shape: (n,)
+    
+    # Compute logpdf in a vectorized way. Broadcasting: (n,1) vs (num_sources,)
+    logpdf = (
+        -0.5 * np.log(2 * np.pi)
+        - np.log(stds)[:, None]
+        - 0.5 * (((values - means[:, None]) / stds[:, None]) ** 2)
+    )  # shape: (n, num_sources)
+    
+    if logpdf.shape[0] == 1:
+        return logpdf[0]
+    return logpdf

petra_catalogs-0.1.0/petra/cost_matrix.py

@@ -0,0 +1,54 @@
+import numpy as np
+from typing import Callable, List
+from functools import partial
+
+
+def create_compute_cost_matrix(aux_distribution: Callable, single_parameter: int = None) -> Callable:
+    """
+    Creates the function to compute the cost matrix for a given auxiliary parametric distribution.
+
+    Parameters
+    ----------
+    aux_distribution: The function to compute the cost matrix elements for the auxiliary distribution
+
+    Return
+    ------
+    compute_cost_matrix: The function to compute the cost matrix
+    """
+    if single_parameter is not None:
+        aux_distribution = partial(aux_distribution, single_parameter=single_parameter)
+
+    def compute_cost_matrix(sample: np.ndarray, aux_parameters: List[np.ndarray],
+                            prob_in_model: np.ndarray, num_distributions: int) -> np.ndarray:
+        """
+        Takes in a single sample and a set of parameters for all fitting distributions,
+        and returns the cost matrix.
+
+        Parameters
+        ----------
+        sample: A numpy array of shape (num_sources, num_params_per_source)
+        aux_parameters: A tuple of lists containing the auxiliary distribution parameters
+        prob_in_model: A numpy array of shape (num_distributions,) that contains the probability of each source being in the catalog
+        num_distributions: The number of distributions to use in the cost matrix
+
+        Return
+        ------
+        cost_matrix: A numpy array of shape (num_sources, num_distributions)
+        """
+        # Precompute logarithms.
+        with np.errstate(divide='ignore'):  # avoid divide by zero warnings, they are expected when we don't clip prob_in_model
+            log_prob = np.log(prob_in_model)  # shape: (num_distributions,)
+            log_prob_not = np.log1p(-prob_in_model)[np.newaxis, :]  # shape: (1, num_distributions)
+
+        # Vectorize over distribution indices.
+        distribution_indices = np.arange(num_distributions)
+        # Assume that aux_distribution is vectorized so that it returns an array of shape (num_distributions, num_sources)
+        cost_matrix = aux_distribution(sample, aux_parameters, distribution_indices)
+        cost_matrix = log_prob[:, np.newaxis] + cost_matrix
+
+        # Replace any NaN values with log(1 - prob_in_model)
+        cost_matrix = np.where(np.isnan(cost_matrix), log_prob_not, cost_matrix)
+
+        return cost_matrix
+
+    return compute_cost_matrix

petra_catalogs-0.1.0/petra/initialization.py

@@ -0,0 +1,149 @@
+import numpy as np
+from petra.posterior_chain import PosteriorChain
+from petra.parametric_fits import uni_normal_fit_single_parameter
+from petra.aux_distributions import uni_normal_aux_distribution_single_parameter
+from petra.relabel import create_relabel_samples
+
+
+def relabel_by_sorting(posterior_chain: PosteriorChain, init_parameter_index: int = 0):
+    """
+    Relabel sources by simply sorting each row by the chosen parameter.
+    
+    NaNs are replaced by infinity so that they come last, then after sorting
+    we revert infinity back to NaN.
+    """
+    chain = posterior_chain.get_chain().copy()  # shape: (n_samples, n_sources, n_params)
+    n_samples, n_sources, _ = chain.shape
+
+    # Work on the parameter we want to sort on.
+    param = chain[:, :, init_parameter_index].copy()
+    # Replace NaNs with infinity so they appear at the end when sorting.
+    param[np.isnan(param)] = np.inf
+
+    # For each row, get the permutation that sorts the parameter values.
+    ordering = np.argsort(param, axis=1)
+    
+    # Reorder the entire chain.
+    reordered_chain = np.take_along_axis(chain, ordering[:, :, np.newaxis], axis=1)
+
+    # Convert any inf back to NaN in the parameter.
+    sorted_param = reordered_chain[:, :, init_parameter_index]
+    sorted_param[np.isinf(sorted_param)] = np.nan
+    reordered_chain[:, :, init_parameter_index] = sorted_param
+
+    return PosteriorChain(
+        reordered_chain,
+        posterior_chain.num_sources,
+        posterior_chain.num_params_per_source,
+        posterior_chain.trans_dimensional,
+        posterior_chain.prob_in_model,
+        posterior_chain.cost_dict,
+    )
+
+
+def relabel_by_histogram(posterior_chain: PosteriorChain, init_parameter_index: int = 0):
+    """
+    Assign labels to sources based on samples for one parameter,
+    ensuring that each original column (source) is used exactly once per row.
+    
+    The algorithm:
+      1. Extract the parameter values (for ordering) and replace NaNs with ∞ so they’re
+         chosen only after all finite values.
+      2. For each ordering column (i.e. for each source position):
+         a. Compute a histogram from all remaining finite values (across all rows).
+         b. Find the histogram’s mode.
+         c. For each row, choose the index (from the not-yet-assigned values) whose value is 
+            closest to the mode.
+         d. Mark that chosen value as “used” by setting it to NaN.
+      3. Reorder the full chain according to the computed ordering.
+      4. Convert any ∞ values in the sorted parameter back to NaN.
+    
+    Parameters
+    ----------
+    posterior_chain : PosteriorChain
+        Contains a 3D array of samples with shape (n_samples, n_sources, n_params).
+    init_parameter_index : int, optional
+        The parameter index used for sorting (default 0).
+    
+    Returns
+    -------
+    PosteriorChain
+        A new PosteriorChain with sources reordered so that each original column is used once per row.
+    """
+    # Get a copy of the full chain (shape: [n_samples, n_sources, n_params])
+    chain = posterior_chain.get_chain().copy()
+    n_samples, n_sources, _ = chain.shape
+
+    # Extract the parameter values we are sorting on.
+    # Work on a separate copy so that we can mark chosen entries.
+    param_values = chain[:, :, init_parameter_index].copy()
+    # Replace NaNs with infinity so they are assigned last.
+    param_values[np.isnan(param_values)] = np.inf
+
+    # Prepare an array to store the new ordering (each row should be a permutation of 0..n_sources-1).
+    reordered_indices = np.empty((n_samples, n_sources), dtype=int)
+
+    # Iterate for each new column in the ordering.
+    for new_col in range(n_sources):
+        # Compute the global histogram mode from all remaining finite values.
+        remaining = param_values[np.isfinite(param_values)]
+        if remaining.size > 0:
+            hist, edges = np.histogram(remaining, bins="auto")
+            centers = (edges[:-1] + edges[1:]) / 2
+            mode = centers[np.argmax(hist)]
+        else:
+            # Fallback if no finite values remain (should not occur given our assumptions).
+            mode = 0.0
+
+        # Now, for each row, pick the index (from values not yet chosen) that is closest to the mode.
+        for i in range(n_samples):
+            row = param_values[i, :]
+            # Find the positions that have not been chosen yet (i.e. are not NaN).
+            available = np.where(~np.isnan(row))[0]
+            if available.size > 0:
+                # Compute differences only over available values.
+                diffs = np.abs(row[available] - mode)
+                best_idx = available[np.argmin(diffs)]
+            else:
+                # Should not happen; every row originally has at least one finite value.
+                best_idx = 0
+            # Record the chosen index for this new column.
+            reordered_indices[i, new_col] = best_idx
+            # Mark the chosen value as used by setting it to NaN.
+            param_values[i, best_idx] = np.nan
+
+    # Reorder the entire chain using the computed ordering.
+    # This reorders along the source dimension.
+    reordered_chain = np.take_along_axis(chain, reordered_indices[:, :, np.newaxis], axis=1)
+
+    # In the sorted parameter, convert any ∞ values back to NaN.
+    sorted_param = reordered_chain[:, :, init_parameter_index]
+    sorted_param[np.isinf(sorted_param)] = np.nan
+    reordered_chain[:, :, init_parameter_index] = sorted_param
+
+    return PosteriorChain(
+        reordered_chain,
+        posterior_chain.num_sources,
+        posterior_chain.num_params_per_source,
+        posterior_chain.trans_dimensional,
+        posterior_chain.prob_in_model,
+        posterior_chain.cost_dict,
+    )
+
+
+def relabel_uni_normal_one_parameter(posterior_chain: PosteriorChain,
+                                     max_num_sources: int = None,
+                                     num_iterations: int = 20,
+                                     init_parameter_index: int = 0,
+                                     eps=1e-2):
+    # create single parameter function to relabel samples
+    relabel_samples = create_relabel_samples(uni_normal_fit_single_parameter,
+                                             uni_normal_aux_distribution_single_parameter,
+                                             single_parameter=init_parameter_index,
+                                             eps=eps)
+
+    return relabel_samples(
+        posterior_chain,
+        max_num_sources=max_num_sources,
+        num_iterations=num_iterations
+    )

petra_catalogs-0.1.0/petra/load_samples.py

@@ -0,0 +1,113 @@
+import numpy as np
+import os
+import glob
+from petra.utils import sort_by_number
+from petra.posterior_chain import PosteriorChain
+
+
+def load_samples_product_space(filepath, num_params_per_source, fill_value=np.nan):
+    """
+    Load chains from a file and reshape them into a 3D array
+    with dimensions (num_samples, num_sources, num_params)
+
+    Assumes the chain is organized as (param1_source1, param2_source1, ..., paramN_source1, paramN_source2, ...)
+    """
+    with open(filepath, "r") as f:
+        chain = np.loadtxt(f)
+    num_sources_vector = np.rint(chain[:, -5] + 1).astype(int)
+    num_sources = np.max(num_sources_vector)
+
+    only_samples_chain = (
+        np.zeros((chain.shape[0], num_sources * num_params_per_source)) + fill_value
+    )
+    for i in range(chain.shape[0]):
+        only_samples_chain[i, : num_sources_vector[i] * num_params_per_source] = chain[
+            i, : num_sources_vector[i] * num_params_per_source
+        ]
+    samples = only_samples_chain.reshape(-1, num_sources, num_params_per_source)
+    return PosteriorChain(
+        samples, num_sources, num_params_per_source, trans_dimensional=True
+    )
+
+
+def load_samples_fixed_num_sources(filepath, num_params_per_source, burn=0, thin=1):
+    """
+    Load chains from a file and reshape them into a 3D array
+    with dimensions (num_samples, num_sources, num_params)
+
+    Assumes the chain is organized as (param1_source1, param2_source1, ..., paramN_source1, paramN_source2, ...)
+    """
+    with open(filepath, "r") as f:
+        chain = np.loadtxt(f)
+    chain = chain[burn::thin, :-4]  # remove metadata columns
+    num_sources = (chain.shape[1]) // num_params_per_source
+    samples = chain.reshape(-1, num_sources, num_params_per_source)
+    return PosteriorChain(samples, num_sources, num_params_per_source)
+
+
+def load_samples_ucbmcmc(chain_folder, fill_value=np.nan, burn=0, thin=1, remove_low_numbers=False):
+    """
+    Load chains from files and reshape them into a 3D array
+    with dimensions (num_samples, num_sources, num_params).
+
+    Only files with more than 8 samples (i.e. more samples than the number
+    of dimensions in the model) are loaded.
+
+    Assumes the chain is organized as:
+      (param1_source1, param2_source1, ..., 
+       param1_source2, param2_source2, ..., 
+       ...,
+       paramN_source1, paramN_source2, ...)
+    """
+
+    # Find and sort all matching filepaths.
+    filepaths = sort_by_number(
+        glob.glob(os.path.join(chain_folder, "dimension_chain.dat.*"))
+    )
+    
+    # Build a list of valid files that have more samples than the model's 8 dimensions.
+    valid_files = []
+    for filepath in filepaths:
+        # Determine the number of sources from the filename.
+        nsources = int(filepath.split(".")[-1])
+        if nsources == 0:
+            continue  # Skip files with no sources
+        
+        # Count the total number of lines in the file.
+        with open(filepath, "r") as f:
+            total_lines = sum(1 for line in f)
+        # Calculate the number of samples by dividing by the number of sources.
+        # (Assumes that the file's total line count is an integer multiple of nsources.)
+        nsamples = total_lines // nsources
+        
+        # Only accept files with more than 8 samples.
+        if not remove_low_numbers:
+            valid_files.append((filepath, nsamples, nsources))
+        elif remove_low_numbers and nsamples > 8:
+            valid_files.append((filepath, nsamples, nsources))
+    
+    if not valid_files:
+        raise ValueError("No files with more than 8 samples found.")
+    
+    # Compute the total number of samples and maximum number of sources among valid files.
+    total_samples = sum(nsamples for (_, nsamples, _) in valid_files)
+    max_sources = max(nsources for (_, _, nsources) in valid_files)
+    
+    # Initialize the array to hold the samples.
+    samples = np.full((total_samples, max_sources, 8), fill_value)
+    
+    current_total = 0
+    # Loop over valid files and load their data.
+    for filepath, nsamples, nsources in valid_files:
+        # Load the chain data from the file.
+        chain = np.loadtxt(filepath)
+        # Reshape and place the data into the samples array.
+        samples[current_total: current_total + nsamples, :nsources, :] = chain.reshape((nsamples, nsources, 8))
+        current_total += nsamples
+    
+    return PosteriorChain(
+        samples[burn::thin],
+        max_sources,
+        trans_dimensional=True,
+        num_params_per_source=8,
+    )

petra_catalogs-0.1.0/petra/make_catalog.py

@@ -0,0 +1,54 @@
+from petra.posterior_chain import PosteriorChain
+from petra.relabel import create_relabel_samples
+from petra.aux_distributions import mv_normal_aux_distribution
+from petra.parametric_fits import mv_normal_fit
+from petra.initialization import relabel_uni_normal_one_parameter, relabel_by_histogram
+
+
+def relabel_mv_normal(posterior_chain: PosteriorChain,
+                      max_num_sources: int = None,
+                      num_iterations: int = 20,
+                      eps=1e-2):
+
+    relabel_samples = create_relabel_samples(mv_normal_fit,
+                                             mv_normal_aux_distribution,
+                                             eps=eps)
+
+    return relabel_samples(
+        posterior_chain,
+        max_num_sources=max_num_sources,
+        num_iterations=num_iterations
+    )
+
+
+def make_catalog_mv_normal(posterior_chain: PosteriorChain,
+                           max_num_sources: int,
+                           num_iterations: int = 200,
+                           initialization_param_index: int = None,
+                           shuffle_seed: int = None):
+    # shuffle the entries
+    posterior_chain.randomize_entries(shuffle_seed)
+
+    # make sure that posterior_chain has the right shape
+    if posterior_chain.num_sources > max_num_sources:
+        posterior_chain.expand_chain(max_num_sources)
+
+    # initialize here:
+    if initialization_param_index is not None:
+        print("Initializing with univariate normal distribution.")
+        # initial_posterior_chain = relabel_uni_normal_one_parameter(posterior_chain,
+        #                                                            max_num_sources=max_num_sources,
+        #                                                            num_iterations=init_num_iterations,
+        #                                                            init_parameter_index=initialization_param_index)
+        initial_posterior_chain = relabel_by_histogram(posterior_chain,
+                                                       init_parameter_index=initialization_param_index)
+    else:
+        initial_posterior_chain = posterior_chain
+
+    # relabel using mv normal
+    print("Relabeling with multivariate normal distribution.")
+    relabeled_chain = relabel_mv_normal(initial_posterior_chain,
+                                        max_num_sources=max_num_sources,
+                                        num_iterations=num_iterations)
+    
+    return relabeled_chain

petra_catalogs-0.1.0/petra/parametric_fits.py

@@ -0,0 +1,131 @@
+import pandas as pd
+import numpy as np
+from functools import partial
+from typing import Callable
+from petra.utils import find_prob_in_model
+
+
+def create_parametric_fit(fit_function, single_parameter=None):
+    """
+    Create a function to fit a parametric distribution to the chain of samples. Extensible to any parametric distribution.
+
+    Parameters
+    ----------
+    fit_function: Callable function with the signature (chain, max_num_sources, single_parameter (optional))
+        The function to fit the parametric distribution to the chain of samples
+
+    Return
+    ------
+    parametric_fit: Callable function with the signature (chain, max_num_sources)
+        The function to fit the parametric distribution to the chain of samples
+    """
+
+    if single_parameter is not None:
+        fit_function = partial(fit_function, fit_parameter=single_parameter)
+
+    def parametric_fit(chain, max_num_sources):
+        """
+        Fit a parametric distribution to the chain of samples .
+
+        Parameters
+        ----------
+        chain: A numpy array of shape (num_samples, num_entries, num_params_per_source)
+        max_num_sources: The maximum number of sources to consider in the catalog
+
+        Return
+        ------
+        fit: The fit to the chain of samples
+        """
+        return fit_function(chain, max_num_sources)
+
+    return parametric_fit
+
+
+def mv_normal_fit(chain, max_num_sources):
+    """
+    Get a MV Gaussian fit (mu, cov) to use in the cost matrix. This function takes in a chain of samples and returns the
+    mean and covariance matrix. NaN values are removed by pandas before the calculation.
+
+    Parameters
+    ----------
+    chain: A numpy array of shape (num_samples, num_entries, num_params_per_source)
+    max_num_sources: The maximum number of sources to consider in the catalog
+
+    Return
+    -------
+    means: A list of length max_num_sources of numpy arrays of shape (num_params_per_source,)
+    cov_matrices: A list of length max_num_sources of numpy arrays of shape (num_params_per_source, num_params_per_source)
+    """
+
+    means = []
+    cov_matrices = []
+    for source in range(max_num_sources):
+        sample_i = chain[:, source, :]  # shape: (num_samples, num_params)
+        valid = ~np.isnan(sample_i).any(axis=1)
+        valid_samples = sample_i[valid]
+        if valid_samples.shape[0] < 8:
+            print('Fewer than 8 values in source index {}. Appending a standard normal distribution.'.format(source))
+            means.append(np.zeros(chain.shape[2]))
+            cov_matrices.append(np.identity(chain.shape[2]))
+            continue
+        df = pd.DataFrame(chain[:, source, :])
+        means.append(np.array(df.dropna().mean()))
+        cov_matrices.append(np.array(df.dropna().cov()))
+    return np.array(means), np.array(cov_matrices)
+
+
+def uni_normal_fit_single_parameter(chain, max_num_sources, fit_parameter):
+    """
+    Get a univariate Gaussian fit (mu, std) to use in the cost matrix. This function takes in a chain of samples and returns the
+    mean and standard deviation. NaN values are removed by pandas before the calculation.
+
+    Parameters
+    ----------
+    chain: A numpy array of shape (num_samples, num_entries, num_params_per_source)
+    max_num_sources: The maximum number of sources to consider in the catalog
+    fit_parameter: The index of the parameter to fit
+
+    Return
+    ------
+    means: A list of length max_num_sources of numpy arrays of shape (num_params_per_source,)
+    stds: A list of length max_num_sources of numpy arrays of shape (num_params_per_source,)
+    """
+    means = []
+    stds = []
+    for source in range(max_num_sources):
+        sample_i = chain[:, source, fit_parameter]  # shape: (num_samples, num_params)
+        valid = np.where(~np.isnan(sample_i))
+        valid_samples = sample_i[valid]
+        if valid_samples.shape[0] < 8:
+            print('Fewer than 8 values in source index {}. Appending a standard normal distribution.'.format(source))
+            means.append(0)
+            stds.append(1)
+            continue
+        else:
+            df = pd.DataFrame(chain[:, source, fit_parameter])
+            mean = np.nanmean(df)
+            std = np.nanstd(df)
+            means.append(mean)
+            stds.append(std)
+    return np.array(means), np.array(stds)
+
+
+def update_parametric_fit_and_prob_in_model(posterior_chain, max_num_sources, parametric_fit_function: Callable, eps=1e-2):
+    """
+    Update the parametric fit and the probability of each source being in the model.
+
+    Parameters
+    ----------
+    chain: A numpy array of shape (num_samples, num_entries, num_params_per_source)
+    max_num_sources: The maximum number of sources to consider in the catalog
+    fit_function: A function that fits a parametric distribution to the chain of samples
+    prob_in_model: A numpy array of shape (max_num_sources,) that contains the probability of each source being in the model
+
+    Return
+    ------
+    parametric_fit: The fit to the chain of samples
+    prob_in_model: A numpy array of shape (max_num_sources,)
+    """
+    aux_params = parametric_fit_function(posterior_chain.get_chain(), max_num_sources)
+    prob_in_model = find_prob_in_model(posterior_chain.get_chain(), max_num_sources, eps=eps)
+    return aux_params, prob_in_model

petra_catalogs-0.1.0/petra/posterior_chain.py

@@ -0,0 +1,70 @@
+import numpy as np
+from dataclasses import dataclass
+
+
+@dataclass
+class PosteriorChain:
+    """
+    A dataclass to store the chains, number of sources, and number of parameters per source
+
+    Parameters
+    ----------
+    chain: np.ndarray
+        The chain of samples
+    num_sources: int
+        The number of sources
+    num_params_per_source: int
+        The number of parameters per source
+    trans_dimensional: bool
+        Whether the chain has variable number of sources
+    prob_in_model: ndarray
+        Probability of a source being in the model
+    cost_dict: dict
+        A dictionary of total cost for the labeling
+    """
+
+    chain: np.ndarray
+    num_sources: int
+    num_params_per_source: int
+    trans_dimensional: bool = False
+    prob_in_model: np.ndarray = None
+    cost_dict: dict = None
+
+    def __post_init__(self):
+        if self.cost_dict is None:
+            self.cost_dict = {}
+        self.chain = self.chain.reshape(-1, self.num_sources, self.num_params_per_source)
+
+    def __repr__(self):
+        return self.chain.view().__repr__()
+
+    @property
+    def shape(self):
+        return self.chain.shape
+
+    def __getitem__(self, index):
+        return self.chain[index]
+
+    def __setitem__(self, index, value):
+        self.chain[index] = value
+
+    def get_chain(self, burn=0, thin=1):
+        return self.chain[burn::thin]
+
+    def expand_chain(self, max_num_sources):
+        """
+        Expand the second dimension of the chain to the maximum number of sources.
+        """
+        expanded_chain = np.zeros((self.chain.shape[0], max_num_sources, self.num_params_per_source)) + np.nan
+        expanded_chain[:, :self.num_sources, :] = self.chain
+        return PosteriorChain(expanded_chain, max_num_sources, self.num_params_per_source, True, self.prob_in_model, self.cost_dict)
+    
+    def randomize_entries(self, seed=None):
+        """
+        Randomize the entries (second dimension) of the chain without repetition.
+        """
+        # Shuffle along axis=1 (the second dimension) for each index in the first dimension
+        if seed is not None:
+            np.random.seed(seed)
+        for i in range(self.chain.shape[0]):
+            np.random.shuffle(self.chain[i])

petra_catalogs-0.1.0/petra/relabel.py

@@ -0,0 +1,146 @@
+import numpy as np
+from scipy import optimize
+from typing import Callable
+
+from petra.utils import fill_missing_indices
+from petra.posterior_chain import PosteriorChain
+from petra.parametric_fits import update_parametric_fit_and_prob_in_model, create_parametric_fit
+from petra.cost_matrix import create_compute_cost_matrix
+
+
+def relabel_samples_one_iteration(chain, aux_parameters, prob_in_model, max_num_sources, compute_cost_matrix: Callable):
+    """
+    Takes in a chain of samples and a MV Gaussian fit and returns the sorted chain
+
+    Parameters
+    ----------
+    chain: A numpy array of shape (num_chains, num_sources, num_params_per_source)
+    parameters: A set of parameters that fit the aux distribution to the chain
+    prob_in_model: A numpy array of shape (max_num_sources,) that contains the probability of each source being in the model
+    max_num_sources: The maximum number of sources to consider in the catalog (can be less than, equal to, or greater than the number of entries in the chain)
+
+    Return
+    ------
+    relabeled_chain: A numpy array of shape (num_chains, num_sources, num_params_per_source)
+    cost: Cost of the optimal assignment
+    """
+    # TODO(Aaron): Parallelize this for loop over samples
+    relabeled_list = []
+    total_cost = 0
+    for sample in chain:
+        cost_matrix = compute_cost_matrix(sample, aux_parameters, prob_in_model, max_num_sources)
+        total_entries = max(len(sample), max_num_sources)  # pick the larger of number of aux distributions or sample entries
+        row_ind, col_ind = optimize.linear_sum_assignment(cost_matrix, maximize=True)  # solve the linear sum assignment problem
+        total_cost -= cost_matrix[row_ind, col_ind].sum()
+        # TODO(Aaron): Is the following line still necessary? It shouldn't change the result either way.
+        relabeled_list.append(fill_missing_indices(total_entries, col_ind))  # sum assignment only outputs 1 number for each labeling distribution input
+    relabeled_array = np.array(relabeled_list)
+
+    relabeled_samples = np.array([row[m] for row, m in zip(chain, relabeled_array)])  # reorder the original samples
+    total_cost /= len(chain)
+    return relabeled_samples, total_cost
+
+
+def relabel_posterior_chain_one_iteration(posterior_chain: PosteriorChain, aux_parameters, prob_in_model, max_num_sources, compute_cost_matrix: Callable):
+    """
+    Takes in a chain of samples and a MV Gaussian fit and returns the sorted chain
+
+    Parameters
+    ----------
+    chain: A PosteriorChain object
+    parameters: A set of parameters that fit the aux distribution to the chain
+    prob_in_model: A numpy array of shape (max_num_sources,) that contains the probability of each source being in the model
+    max_num_sources: The maximum number of sources to consider in the catalog (can be less than, equal to, or greater than the number of entries in the chain)
+
+    Return
+    ------
+    relabeled_chain: A PosteriorChain object
+    cost: Cost of the optimal assignment
+    """
+    cost_dict = posterior_chain.cost_dict
+    relabeled_chain, total_cost = relabel_samples_one_iteration(posterior_chain.get_chain(), aux_parameters, prob_in_model, max_num_sources, compute_cost_matrix)
+    if cost_dict is None:
+        cost_dict = {}
+    cost_dict[max_num_sources] = total_cost
+    return PosteriorChain(relabeled_chain, posterior_chain.num_sources, posterior_chain.num_params_per_source, prob_in_model=prob_in_model, cost_dict=cost_dict)
+
+
+def create_relabel_samples(parametric_fit_function: Callable,
+                           aux_distribution: Callable,
+                           single_parameter: int = None,
+                           eps: float = 1e-2):
+
+    # make all the necessary pieces
+    param_fit = create_parametric_fit(parametric_fit_function, single_parameter=single_parameter)
+    compute_cost_matrix = create_compute_cost_matrix(aux_distribution, single_parameter=single_parameter)
+
+    def relabel_samples(posterior_chain: PosteriorChain,
+                        max_num_sources: int = None,
+                        num_iterations: int = 200):
+        """
+        Takes in a posterior_chain and relabels the samples using a multivariate normal auxiliary fitting distribution
+
+        Parameters
+        ----------
+        posterior_chain: A PosteriorChain object
+        max_num_sources: The maximum number of sources to consider in the catalog (can be less than, equal to, or greater than the number of entries in the chain)
+        num_iterations: How many tries to use before stopping the algorithm
+        init_parameter_index: Which parameter index to use in the initialization procedures
+        shuffle_entries: Whether to shuffle the entries in the chain at the start
+        shuffle_seed: Seed for the random number generator
+
+        Return
+        ------
+        relabeled_chain: A PosteriorChain object
+        cost: Cost of the optimal assignment
+        """
+
+        # if shuffle_entries:
+        #     posterior_chain.randomize_entries(shuffle_seed)
+        if max_num_sources is None:
+            max_num_sources = posterior_chain.num_sources
+        if max_num_sources > posterior_chain.num_sources:
+            # increase the width of the posterior chain to the maximum number of sources
+            posterior_chain = posterior_chain.expand_chain(max_num_sources)
+        if max_num_sources < posterior_chain.num_sources:
+            raise ValueError("The maximum number of sources cannot be less than the number of entries in the chain.")
+
+        print()
+        print(f"Sorting the posterior chain:\n\tMaximum number of iterations: {num_iterations}\n\tMaximum number of source labels:{max_num_sources}\n")
+
+        # set up the for loop
+        old_posterior_chain = posterior_chain
+        old_parametric_fit, old_prob_in_model = update_parametric_fit_and_prob_in_model(posterior_chain, max_num_sources, param_fit, eps=eps)
+        old_cost_of_assignment = 0
+
+        for iteration in range(num_iterations):
+            # get the new values
+            new_posterior_chain = relabel_posterior_chain_one_iteration(old_posterior_chain, old_parametric_fit, old_prob_in_model, max_num_sources, compute_cost_matrix)
+            new_parametric_fit, new_prob_in_model = update_parametric_fit_and_prob_in_model(new_posterior_chain, max_num_sources, param_fit, eps=eps)
+            new_cost_of_assignment = new_posterior_chain.cost_dict[max_num_sources]
+
+            # print out the results
+            delta_cost_of_assignment = new_cost_of_assignment - old_cost_of_assignment
+            print(f"Iteration {iteration + 1}: Difference in cost of assignment is {delta_cost_of_assignment} with total cost of {new_cost_of_assignment}.")
+            print(f"\tProbabilities in model: {new_prob_in_model}")
+
+            # break if converged
+            if (delta_cost_of_assignment == 0):
+                print(f"Stopped after {iteration + 1} iterations because the cost didn't change from the previous iteration.")
+                new_parametric_fit = old_parametric_fit
+                new_prob_in_model = old_prob_in_model
+                new_posterior_chain = old_posterior_chain
+                break
+
+            # update the old values to the new values
+            old_posterior_chain = new_posterior_chain
+            old_parametric_fit = new_parametric_fit
+            old_prob_in_model = new_prob_in_model
+            old_cost_of_assignment = new_cost_of_assignment
+
+        if iteration == num_iterations - 1:
+            print(f"Final cost of assignment: {new_cost_of_assignment} after the maximum number ({num_iterations}) of iterations.")
+
+        return new_posterior_chain
+
+    return relabel_samples

petra_catalogs-0.1.0/petra/utils.py

@@ -0,0 +1,61 @@
+import numpy as np
+
+
+def find_prob_in_model(chain, max_num_sources, eps=1e-2):
+    """
+    Find the number of samples for each entry in the chain and divide by the total number of samples.
+
+    Parameters
+    ----------
+    chains: A numpy array of shape (num_samples, num_entries, num_params_per_source)
+    max_num_sources: The maximum number of sources to consider in the catalog
+    eps: A small number to avoid division by zero
+
+    Return
+    ------
+    prob_in_model: A numpy array of shape (max_num_sources,)
+    """
+    num_samples = chain.shape[0]
+    prob_in_model = np.zeros(max_num_sources)
+    for i in range(max_num_sources):
+        prob_in_model[i] = np.sum(~np.isnan(chain[:, i, 0])) / num_samples
+    # clip the probabilities to avoid division by zero in np.log
+    prob_in_model = np.clip(prob_in_model, eps, 1 - eps)
+    return prob_in_model
+
+
+def fill_missing_indices(total_sources, given_indices):
+    # Step 1: Generate a list of all indices from 0 to total_sources - 1
+    all_indices = np.arange(total_sources)
+    # Step 2: Convert given_indices to a set for faster operations
+    given_indices_set = set(given_indices)
+    # Step 3: Filter out the given indices from all_indices to get missing indices
+    missing_indices = [index for index in all_indices if index not in given_indices_set]
+    # Step 4: Combine the given indices with the missing indices
+    filled_indices = list(given_indices) + missing_indices
+
+    return np.array(filled_indices)
+
+
+def count_swaps(arr):
+    """
+    Counts the number of swaps needed to sort an array.
+
+    Args:
+        arr (np.array): The input array.
+
+    Returns:
+        int: The number of swaps.
+    """
+    sorted_arr = np.sort(arr)
+    swaps = np.sum(arr != sorted_arr)
+    return swaps // 2
+
+
+def sort_by_number(filenames):
+    # Extract the number after the dot and convert it to an integer
+    def extract_number(filename):
+        return int(filename.split(".")[-1])
+
+    # Sort the filenames using the extracted number
+    return sorted(filenames, key=extract_number)

petra_catalogs-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[tool.poetry]
+name = "petra-catalogs"
+version = "0.1.0"
+description = "From LISA global fit to catalogs"
+authors = ["Aaron D. Johnson", "Javier Roulete", "Katerina Chatziioannou", "Michele Vallisneri", "Kyle Gersbach", "Chris Trejo"]
+readme = "README.md"
+packages = [{ include = "petra"}]
+
+[tool.poetry.dependencies]
+python = "^3.11"
+numpy = "^2.2.3"
+scipy = "^1.15.2"
+pandas = "^2.2.3"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"