<commit_before># utils.py
"""
Utilities module containing various useful
functions for use in other modules.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
import scipy.linalg as sl
import scipy.sparse as sps
import scipy.special as ss
from pkg_resources import Requirement, resource_filename
from scipy.integrate import odeint
from scipy.interpolate import interp1d

import enterprise
from enterprise import constants as const
from enterprise.signals.parameter import function
from enterprise.signals.gp_priors import powerlaw, turnover  # noqa: F401
from enterprise import signals as sigs  # noqa: F401
from enterprise.signals.gp_bases import (  # noqa: F401
    createfourierdesignmatrix_red,
    createfourierdesignmatrix_dm,
    createfourierdesignmatrix_env,
    createfourierdesignmatrix_ephem,
    createfourierdesignmatrix_eph,
)

try:
    from sksparse.cholmod import cholesky
except:
    print("You'll need sksparse for get_coefficients() and replicate()!")


def replicate(pta, ptac, p0, coefficients=False):
    """Create a replicated residuals conditioned on the data.
    Here pta is standard marginalized-likelihood PTA, and
    ptac is a hierarchical-likelihood version of pta with
    coefficients=True for all GPs. This function:
    
    - calls utils.get_coefficients(pta, p0) to get a realization
    of the GP coefficients conditioned on the data and on the
    hyperparameters in p0;
    - calls ptac.get_delay() to compute the resulting realized
    GPs at the toas;
    - adds measurement noise (including ECORR) consistent with
    the hyperparameters.
    
    To use this (pending further development), you need to set
    combine=False on the pta/ptac GPs, and method='sparse' on
    the ptac EcorrKernelNoise.
    
    Returns a list of replicated residuals, one list element
    per pulsar."""

    # GP delays
    if not coefficients:
        p0 = get_coefficients(pta, p0)

    ds = ptac.get_delay(params=p0)

    # note: the proper way to cache the Nmat computation is to give
    # a `sample` method to csc_matrix_alt and ndarray_alt, which
    # would then save the factorization in the instance

    nmats = ptac.get_ndiag(params=p0)
    for d, nmat in zip(ds, nmats):
        if isinstance(nmat, sps.csc_matrix):
            # add EFAC/EQUAD/ECORR noise
            # use xx' = I => (Lx)(Lx)' = LL' with LL' = PNP'
            # see https://scikit-sparse.readthedocs.io/en/latest/cholmod.html
            ch = cholesky(nmat)
            d[ch.P()] += ch.L() @ np.random.randn(len(d))
        elif isinstance(nmat, np.ndarray):
            # diagonal case, nmat will be ndarray_alt instance
            d += np.sqrt(nmat) * np.random.randn(len(d))
        else:
            raise NotImplementedError(
                "Cannot take Nmat factor; " "you may need to set the EcorrKernelNoise to 'sparse'."
            )
<commit_msg>Edited locally to satisfy GH black, but why doesn't mine fail locally?<commit_after># utils.py
"""
Utilities module containing various useful
functions for use in other modules.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
import scipy.linalg as sl
import scipy.sparse as sps
import scipy.special as ss
from pkg_resources import Requirement, resource_filename
from scipy.integrate import odeint
from scipy.interpolate import interp1d

import enterprise
from enterprise import constants as const
from enterprise.signals.parameter import function
from enterprise.signals.gp_priors import powerlaw, turnover  # noqa: F401
from enterprise import signals as sigs  # noqa: F401
from enterprise.signals.gp_bases import (  # noqa: F401
    createfourierdesignmatrix_red,
    createfourierdesignmatrix_dm,
    createfourierdesignmatrix_env,
    createfourierdesignmatrix_ephem,
    createfourierdesignmatrix_eph,
)

try:
    from sksparse.cholmod import cholesky
except:
    print("You'll need sksparse for get_coefficients() and replicate()!")


def replicate(pta, ptac, p0, coefficients=False):
    """Create a replicated residuals conditioned on the data.
    Here pta is standard marginalized-likelihood PTA, and
    ptac is a hierarchical-likelihood version of pta with
    coefficients=True for all GPs. This function:

    - calls utils.get_coefficients(pta, p0) to get a realization
    of the GP coefficients conditioned on the data and on the
    hyperparameters in p0;
    - calls ptac.get_delay() to compute the resulting realized
    GPs at the toas;
    - adds measurement noise (including ECORR) consistent with
    the hyperparameters.

    To use this (pending further development), you need to set
    combine=False on the pta/ptac GPs, and method='sparse' on
    the ptac EcorrKernelNoise.

    Returns a list of replicated residuals, one list element
    per pulsar."""

    # GP delays
    if not coefficients:
        p0 = get_coefficients(pta, p0)

    ds = ptac.get_delay(params=p0)

    # note: the proper way to cache the Nmat computation is to give
    # a `sample` method to csc_matrix_alt and ndarray_alt, which
    # would then save the factorization in the instance

    nmats = ptac.get_ndiag(params=p0)
    for d, nmat in zip(ds, nmats):
        if isinstance(nmat, sps.csc_matrix):
            # add EFAC/EQUAD/ECORR noise
            # use xx' = I => (Lx)(Lx)' = LL' with LL' = PNP'
            # see https://scikit-sparse.readthedocs.io/en/latest/cholmod.html
            ch = cholesky(nmat)
            d[ch.P()] += ch.L() @ np.random.randn(len(d))
        elif isinstance(nmat, np.ndarray):
            # diagonal case, nmat will be ndarray_alt instance
            d += np.sqrt(nmat) * np.random.randn(len(d))
        else:
            raise NotImplementedError(
                "Cannot take Nmat factor; " "you may need to set the EcorrKernelNoise to 'sparse'."
            )

{'fraction_non_alphanumeric': 0.0727735368956743, 'fraction_numerical': 0.005767599660729432, 'mean_word_length': 3.260115606936416, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '330353', 'n_tokens_mistral': 1853, 'n_tokens_neox': 1759, 'n_words': 617}