chore: untrack oldCode directory

import jax

import jax.scipy.stats.norm as jax_norm

import jax.scipy.stats.uniform as jax_uniform

import jax.scipy.stats.truncnorm as jax_truncnorm

import numpy as np

class JAXDistribution:

    """Base class for JAX-compatible distributions with both JAX logpdf and numpy rvs."""

    pass

class JAXUniform(JAXDistribution):

    """JAX-compatible uniform distribution.

    Initialization expects `loc` (lower bound) and `high` (upper bound).

    The class computes `scale = high - loc` internally so YAML can be

    provided as `loc` and `high` (more user-friendly than `loc` and `scale`).

    """

    def __init__(self, loc, high):

        self.loc = loc

        self.high = high

        self.scale = float(self.high - self.loc)

    def logpdf(self, x):

        """JAX-compatible log-pdf."""

        return jax_uniform.logpdf(x, loc=self.loc, scale=self.scale)

    def ppf(self, x):

        """JAX-compatible ppf."""

        return jax_uniform.ppf(x, loc=self.loc, scale=self.scale)

    def rvs(self):

        """JAX-based random sample using `jax.random.uniform`.

        Returns a Python float to match previous `.rvs()` behaviour used in

        host-side code (e.g., initialization of walker positions).

        """

        key = jax.random.PRNGKey(int(np.random.randint(0, 2**31)))

        return float(jax.random.uniform(key, minval=self.loc, maxval=self.high))

class JAXNorm(JAXDistribution):

    """JAX-compatible normal distribution."""

    def __init__(self, loc, scale):

        self.loc = loc

        self.scale = scale

    def logpdf(self, x):

        """JAX-compatible log-pdf."""

        return jax_norm.logpdf(x, loc=self.loc, scale=self.scale)

    def ppf(self, x):

        """JAX-compatible ppf."""

        return jax_norm.ppf(x, loc=self.loc, scale=self.scale)

    def rvs(self):

        """JAX-based random sample."""

        key = jax.random.PRNGKey(int(np.random.randint(0, 2**31)))

        return float(self.loc + self.scale * jax.random.normal(key))

class JAXTruncnorm(JAXDistribution):

    """JAX-compatible truncated normal distribution."""

    def __init__(self, a, b, loc, scale):

        self.a = a

        self.b = b

        self.loc = loc

        self.scale = scale

    def logpdf(self, x):

        """JAX-compatible log-pdf."""

        return jax_truncnorm.logpdf(x, self.a, self.b, loc=self.loc, scale=self.scale)

    def ppf(self, x):

        """JAX-compatible ppf."""

        # return jax_truncnorm.ppf(x, loc=self.loc, scale=self.scale)

        raise NotImplementedError("JAXTruncnorm does not have a ppf method implemented.")

    def rvs(self):

        """JAX-based random sample using inverse transform sampling."""

        raise NotImplementedError("The rvs method for JAXTruncnorm required a scipy.stats dependency which has been removed. Walker initialization with truncated normal priors is currently not supported.")

def get_jax_logpdf(dist):

    """

    Returns a JAX-compatible log-pdf function for a given distribution.

    Assumes all distributions are JAX-compatible (JAXDistribution subclasses or float/int).

    Args:

        dist: A JAXDistribution or float/int (fixed parameter).

    Returns:

        A callable that computes the log-pdf using JAX, or the scalar value if fixed.

    """

    if isinstance(dist, (float, int)):

        return dist

    # All distributions should be JAX distributions at this point

    return dist.logpdf

    # If we reach here, it's likely an error - but provide informative message

    raise TypeError(f"Expected JAXDistribution or float/int, got {type(dist).__name__}. "

                    "Ensure distributions are created as JAX distributions in the YAML config.")

#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

Created on Thu May 15 11:41:17 2025

@author: gas

"""

import numpy as np

import matplotlib.pyplot as plt; plt.close('all')

def plotGLS(gls,freq=False,ticks=[],subplots=1,axes=np.nan,plots=np.nan,timeZoom=np.array([0]),perColor='black',label=None):

    if subplots>0:

        fig, axes = plt.subplots(subplots)

        if subplots>1:

            ax=axes.copy(); ax=ax[0]

        else:

            ax=axes

    else:

        fig,axes=plots

        ax=axes.copy(); ax=ax[-subplots]

    if freq is False:

        x=np.log10(1./gls['freq'])

        ax.plot(x,gls['power'],color=perColor,linewidth=0.5,label=label)

        # ax.plot([x.min(),x.max()],np.repeat(np.var(gls['y'])/len(gls['y']),2),'k--',linewidth=0.5)

        if len(ticks)==0:

            xarray1=int(min(x)); xarray2=int(max(x))+1

            xarray=np.asarray(range(xarray1,xarray2+1))

        else:

            xarray=np.log10(ticks)

            xarray1=np.min(ticks); xarray2=np.max(ticks)

        #xarray=np.log10(np.array([30,32,34,36,38]))

        ax.xaxis.set_ticks(xarray)

        ax.set_xticklabels(np.round(10.**xarray*100)/100)

        ticks=np.linspace(2,9,num=8)*10.**xarray1

        xarray=ticks.copy()

#        for i in range(0,(xarray2-xarray1+1)):

        for i in range(0,int(np.floor(xarray2-xarray1+1))):

            xarray=np.append(xarray,ticks*10.**i)

        xarray=np.log10(xarray)

        ax.xaxis.set_ticks(xarray*100, minor = True)

        ax.set_xlabel('Period [d]')

        if len(timeZoom)==2:

            ax.set_xlim(np.log10(timeZoom))

        else:

            ax.set_xlim([min(x),max(x)])

        '''

        x=-gls.freq

        ax.plot(x,gls.power,color='black')

        xarray1=int(min(np.log10(1./gls.freq))); xarray2=int(max(np.log10(1./gls.freq)))+1

        xarray=10**np.asarray(range(xarray1,xarray2))

        ax.xaxis.set_ticks(-1./xarray)

        ax.set_xticklabels(xarray)

        ticks=np.linspace(2,9,num=8)*10.**xarray1

        xarray=ticks.copy()

        for i in range(0,(xarray2-xarray1+1)):

            xarray=np.append(xarray,ticks*10.**i)

        ax.xaxis.set_ticks(-1./xarray, minor = True)

        '''

        ax.xaxis.set_tick_params(top='on',direction='in',which='both')

        # ax.relim()      # make sure all the data fits

        # ax.autoscale()  # auto-scale

    else:

        x=gls['freq']

        ax.plot(x,gls['power'],color=perColor,linewidth=0.5,label=label)

        # ax.plot([x.min(),x.max()],np.repeat(np.var(gls['y'])/len(gls['y']),2),'k--',linewidth=0.5)

        ax.set_xlabel('f [d$^{-1}$]')

        xarray=ax.get_xticks()

        xarray=xarray[np.where(xarray>=0)[0]]

        step=(xarray[1]-xarray[0])/10

        num=(max(xarray)-min(xarray))/step+1

        xarray=np.linspace(min(xarray),max(xarray),num=int(np.round(num)))

        ax.xaxis.set_ticks(xarray, minor = True)

        ax.tick_params(axis='x',which='minor',bottom='on')

        #ax.relim()      # make sure all the data fits

        #ax.autoscale()  # auto-scale

        #new_tick_locations = 1./np.array([2,3,5,10,100,1000])

        if len(timeZoom)==2:

            ax.set_xlim(1/np.flipud(timeZoom))

        if np.size(ticks)>1:

            ax2 = ax.twiny()

            new_tick_locations =1./ticks

            #def tick_function(X):

            #    V = 1./X

            #    return ["%i" % z for z in V]

            new_tick_labels=[str(x) for x in ticks]

            ax2.set_xlim(ax.get_xlim())

            ax2.set_xticks(new_tick_locations)

            #ax2.set_xticklabels(tick_function(new_tick_locations), rotation=90)

            ax2.set_xticklabels(new_tick_labels, rotation=90)

            ax2.set_xlabel("Period [d]")

            ax2.xaxis.set_tick_params(top='on',direction='in',which='major')

            ax.xaxis.set_tick_params(top='off',direction='in',which='both')

        else:

            ax.xaxis.set_tick_params(top='on',direction='in',which='both')

    ax.yaxis.set_tick_params(right='on',direction='in',which='both')

    ylim=np.array(ax.get_ylim()); ylim[0]=-ylim[1]*.05; ylim[1]=ylim[1]*1.1

    #ylim=np.array([0,np.max(np.append(gls.power))])

    ax.set_ylim(ylim)

    ax.set_ylabel(r'units$^2$/d$^{-1}$')

    # ax.set_xscale('log'); ax.set_yscale('log')

    return(fig,axes)

def plotFAP(gls,fapColors,xlim,ax):

    for i in range(0,gls['faps'].size):

        ax.plot(xlim,[gls['fapLevels'][i],gls['fapLevels'][i]],color=fapColors[i],linewidth=0.5,label='FAP='+str(gls['faps'][i])+'%')

    ax.legend(fontsize='small')

    return

def getNormalizedPeriodogram(t,y,dy=None,freqs=None,faps=None,n_bootstrap=1000,normalization='psd'):

    from astropy.timeseries import LombScargle

    import numpy as np

    if faps is None: faps=np.array([5,1,0.1])# in %

    ls = LombScargle(t, y)

    if freqs is None:

        freqs,power=ls.autopower()

    else:

        power=ls.power(freqs)

    spw=LombScargle(t, np.ones(len(t))).power(freqs)

    # fLimit=.5/np.min(np.diff(t))

    # integralSupport=np.where(freqs<fLimit)[0]

    # zeta=1/(1*np.trapz(spw[integralSupport],freqs[integralSupport]))

    # power=power*zeta #from power spectrum to power spectral density

    if normalization=='psd':

        rescale=np.var(y)/np.trapz(power,freqs)

        power=power*rescale

    else:

        rescale=1

    fapLevels=ls.false_alarm_level(faps/100, method='bootstrap', method_kwds={'n_bootstraps':n_bootstrap})*rescale

    return {'freq':freqs,'power':power,'spw':spw,'faps':faps,'fapLevels':fapLevels}

 No newline at end of file

#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

Created on Mon Feb 17 16:30:47 2025

@author: gas

"""

def run(path, dataFile, configFile, restart=False):

    import matplotlib

    matplotlib.use('PDF')

    import jax, sys, os, socket, h5py

    import jax.numpy as jnp

    from jax import jit, random

    from pytransit import QuadraticModel, EclipseModel

    # from snappyko import Orbit

    from meepmeep import Orbit

    import numpy as np

    # import jax.numpy as np

    import matplotlib.pyplot as plt

    from matplotlib.backends.backend_pdf import PdfPages

    import time as ttt

    import importlib

    # np.seterr(divide = 'ignore')

# =============================================================================

    hostname=socket.gethostname()

    if hostname in 'lor hon thX280 zephyrus':

        sys.path.append('/dati/software/pyLib/')

        # sys.path.append('/dati/software/l1periodogram/l1periodogram_codes/')

        os.chdir('/dati/ricerca/cheops')

    elif 'r3' in hostname:

        sys.path.append('/home/gscandariato/software/pyLib')

        os.chdir('/mnt/beegfs/gscandariato/cheops')

    else:

        # sys.path.append('/u/gscandariato/software/pyLib')

        # os.chdir('/beegfs/gscandariato/cheops')

        sys.path.append('/content/drive/MyDrive/dati/software/pyLib/')

        os.chdir('/content/drive/MyDrive/dati/ricerca/cheops/')

    # =============================================================================

    # Ensure the local package root is importable on any host

    # file is at <repo_root>/PyAETNA/runPyAETNA.py

    # add repo_root (parent of the package dir) so `import PyAETNA...` works

    file_dir = os.path.dirname(os.path.abspath(__file__))

    repo_root = os.path.dirname(file_dir)

    if repo_root not in sys.path:

        sys.path.insert(0, repo_root)

    # (alternatively export PYTHONPATH to include repo_root on the cluster)

    # os.environ["OMP_NUM_THREADS"] = "6"

    # os.environ["OPENBLAS_NUM_THREADS"] = "6"

    # os.environ["XLA_FLAGS"]="--xla_force_host_platform_device_count=8"    

    # sys.path.append('./pyTools')

    from PyAETNA.utils import PyAETNA_MCMCFunctions as gsMCMC# import runSampler, plotMCMC, computeBE, planetModel

    importlib.reload(gsMCMC)

    os.chdir(path)

    if not os.path.isdir('PyAETNAplots'): os.mkdir('PyAETNAplots')

    if not os.path.isdir('PyAETNAouts'): os.mkdir('PyAETNAouts')

    gsMCMC.logFile='PyAETNAouts/'+configFile+'_'+dataFile+'.log'

    if restart:

        file = open(gsMCMC.logFile, "w"); file.close()

        try:

            os.remove('PyAETNAouts/'+configFile+'_'+dataFile+'.h5')

        except:

            pass

    #jax.config.update("jax_platforms", "cpu") #TODO: set a keyword in the yaml fiel to choose the device

    file = open(gsMCMC.logFile, "a")

    print("\nJAX version {}".format(jax.__version__),file=file)

    #print('NODE_LIST',node_list,file=file)

    print('\nJAX devices:', jax.devices(),file=file)

    print('\nJAX env info:', jax.print_environment_info(return_string=True),file=file)

    print('\nJAX default backend:', jax.default_backend(),file=file)

    file.close()

    from utils.PyAETNA_orbitModel import circ_mandelagol_pyTransit as planetModel1

    from utils.PyAETNA_orbitModel import circ_mandelagol_pyTransit_LongCadence as planetModel2

    # lnTEST=gsMCMC.lnprob_vectorized_TEST()

    # import time

    # time.sleep(5)

    ### RUNNING THE SAMPLER

    gsMCMC.planetModel=[planetModel1,planetModel2]

    model=gsMCMC.runSampler(configFile,dataFile,outDir='PyAETNAouts/')

    # return #comment this line if you want to plot the results of the sampling

    ### PLOTTING

    flog=open(gsMCMC.logFile,'a'); flog.write('\n'+ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime())+' Plotting...\n'); flog.close() 

    pdfName='PyAETNAplots/'+configFile+'_'+dataFile+'_.pdf'; pdf = PdfPages(pdfName)    

    pdfName='PyAETNAplots/'+configFile+'_'+dataFile+'_BestFit.pdf'; pdfBestFit = PdfPages(pdfName)    

    # with jax.default_device(jax.devices('cpu')[0]):

        # time,yres,drv,flag=gsMCMC.plotMCMC(configFile,dataFile,pdf,pdfBestFit,outDir='PyAETNAouts/')    

    time,yres,drv,flag=gsMCMC.plotMCMC(configFile,dataFile,pdf,pdfBestFit,outDir='PyAETNAouts/')    

    pdf.close(); pdfBestFit.close()

    flog=open(gsMCMC.logFile,'a'); flog.write('\n'+ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime())+' Plotting done.\n'); flog.close()    

    ### COMPUTING BE    

    gsMCMC.computeBE(configFile,dataFile,outDir='PyAETNAouts/')    

    flog=open(gsMCMC.logFile,'a'); flog.write('\n'+ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime())+' Finished!!!\n'); flog.close()

    # samplerPool.close()

    os.chdir('..')

    return

if __name__ == "__main__":

    import sys

    try:

        path=sys.argv[1]

        dataFile=sys.argv[2]

        configFile=sys.argv[3]

        restart=sys.argv[4]

    except:

        # path='V1298Tau'

        # dataFile='V1298Tau_201_sa.sub.fits.dat'

        # configFile='testV4_noPlanet_tinyGP'

        path='/dati/ricerca/cheops/testPyAETNA/'

        #'CHEOPS_ALL_test.dat' for the whole dataset, 'CHEOPS_ALL.dat' for a smaller subset (for faster testing)

        # dataFile='CHEOPS_ALL_test.dat'

        dataFile='CHEOPS_ALL_semi.dat'

        # dataFile='CHEOPS_ALL.dat'

        #'testV4_ALL_noPlanet_noGP' for lighter model with no GP

        # 'testV4_ALL_noPlanet_tinyGP' for a more complex model with GP computed with tinyGP/jax, good for testing GPU

        configFile='testV4_ALL_noPlanet_noGP'

        restart=True

    run(path, dataFile, configFile, restart=restart)