# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Try to reconfigure OpSim for other cadence
#
# *Evenly distribute and include the galactic plane*

# +
from OpSim import OpSim
from astropy.coordinates import SkyCoord
from astropy import units
import numpy as np

import matplotlib.pyplot as plt

# %matplotlib notebook

filters = ['u_', 'g_', 'r_', 'i_', 'z_', 'y_']
# -

OpS = OpSim()
OpS.dbFile = '/Users/ageller/WORK/LSST/onGitHub/EBLSST/input/db/minion_1016_sqlite.db' #for the OpSim database   
OpS.getAllOpSimFields()

print(OpS.Nobs[1])

print(np.max(OpS.Nobs[np.where(OpS.Nobs < 5000)]))

# +
coords = SkyCoord(OpS.RA, OpS.Dec, unit=(units.degree, units.degree),frame='icrs')

RAwrap = coords.ra.wrap_at(180.*units.degree).degree
Decwrap = coords.dec.wrap_at(180.*units.degree).degree

f, ax = plt.subplots(subplot_kw={'projection': "mollweide"}, figsize=(8,5))
ax.grid(True)
ax.set_xlabel("RA",fontsize=16)
ax.set_ylabel("Dec",fontsize=16)
mlw = ax.scatter(np.array(RAwrap).ravel()*np.pi/180., 
                 np.array(Decwrap).ravel()*np.pi/180., 
                 c=OpS.Nobs, #probably want this in a given filter, but that takes a long time to gather
                 cmap='viridis_r', s = 4, vmin=0, vmax=1000)
cbar = f.colorbar(mlw, shrink=0.7)
cbar.set_label(r'Nobs')

# -

# ### Take only the fields that are observed and get their dates

observed = np.where(OpS.Nobs > 0)
primary = np.where(OpS.Nobs > 800)
print(len(observed[0]))
print(len(primary[0]))

# *There will be this many observations in each of these fields*
#
# *Neglect the fields with the very large number of obs -- is this a bug or because of overlap and/or a planned revisit?*

# +
Nfields = len(observed[0])
Nobs = OpS.Nobs[observed]

bad = np.where(Nobs > 5000)
good = np.where(Nobs < 5000)
#print(bad, Nobs[bad])
Nobs[bad] = np.max(Nobs[good])
#print(Nobs[bad])

NtotalObs = np.sum(Nobs)
print("Nfields, NtotalObs, NObsPerField = ",Nfields, NtotalObs, NtotalObs/Nfields)
# -

# ### Let's try to construct a CDF of the dt inbetween observations

#this takes a long time to run through...
OpSimi = OpS.fieldID[primary].astype("int")
OpS.dt = np.array([{} for i in OpSimi])
OpSimi = OpSimi[0:2] #take only the first two for testing
for i in OpSimi: 
    print(i)
    dt = {}
    OpS.setDates(i, filters)
    for f in filters:
        dt[f] = np.diff(OpS.obsDates[i][f])
    OpS.dt[i] = dt

i = OpSimi[0]
print(OpS.NobsDates[i])
print(OpS.Nobs)
#print(OpS.dt[i][filters[0]])

# *Construct the CDF.  Ideally we wouldn't bin it, but I think this will be easier to do given that there are about 3000 fields and we'd want to run this in MPI on Quest.*

# +
f,ax = plt.subplots(len(filters),1, figsize=(5,15), sharex=True)
Nbins = 100
dist = {}
for a,f in zip(ax, filters):
    pdf = np.zeros(Nbins)
    for i in OpSimi: #sum up all the dt's in that filter across all fields
        p,bin_edges = np.histogram(np.log10(OpS.dt[i][f]), bins=Nbins)
        pdf += p
    bins = bin_edges[:-1] + np.diff(bin_edges)/2.
    cdf = np.cumsum(pdf)
    a.step(bins,pdf/np.max(pdf),linestyle='--')
    a.step(bins,cdf/np.max(cdf))
    a.set_ylabel(f)
    
    dist[f] = {}
    dist[f]['bins']= bins
    dist[f]['cdf']= cdf/np.max(cdf)
    dist[f]['pdf']= pdf/np.max(pdf)
    
ax[-1].set_xlabel('log10(dt [days])')
plt.subplots_adjust(hspace=0)
# -

print(dist[filters[1]])

# ### Files from Quest

import os
import pickle

# +
f,ax = plt.subplots(len(filters),1, figsize=(5,15), sharex=True)

#d = 'primary_output_files/' #primary only take into account the fields with most obs, but using this would overcount total number of available observations
d = 'output_files/' #uses all fields with >0 observations
files = os.listdir(d)
IDs = []
bins = {}
pdf = {}
cdf = {}
Nobs = {}
for i,file in enumerate(files):
    dist = pickle.load( open( d+file, "rb"))
    for f in filters:
        if (i == 0):
            bins[f] = dist[f]['bins']
            pdf[f] = dist[f]['pdf']
            Nobs[f] = dist[f]['Nobs']
        if (i > 0):
            xx = np.where(np.logical_or(np.isinf(dist[f]['pdf']), np.isnan(dist[f]['pdf'])))[0]
            if (len(xx) > 0):
                #print(file, xx, dist[f])
                dist[f]['pdf'][xx] = np.zeros_like(xx)
            pdf[f] += dist[f]['pdf']
            Nobs[f] += dist[f]['Nobs']

outDist = {}
for a,f in zip(ax, filters):
    Nobs[f] /= float(len(files))
    
    cdf[f] = np.cumsum(pdf[f])
    
    pdf[f] /= np.max(pdf[f])
    cdf[f] /= np.max(cdf[f])

    outDist[f] = {}
    outDist[f]['Nobs'] = Nobs[f]
    outDist[f]['bins'] = bins[f]
    outDist[f]['cdf'] = cdf[f]

    a.step(bins[f],pdf[f],linestyle='--')
    a.step(bins[f],cdf[f])
    a.set_ylabel(f)


ax[-1].set_xlabel('log10(dt [days])')
plt.subplots_adjust(hspace=0)

#pickle.dump( outDist, open( "OpSim_primary_dtDist.pickle", "wb" ) )
pickle.dump( outDist, open( "OpSim_observed_dtDist.pickle", "wb" ) )
#print(cdf)
print(Nobs)
# -

# *Test drawing observations*

# +

f,ax = plt.subplots(len(filters),1, figsize=(5,15), sharex=True)
for f,a in zip(filters,ax):
    N = round(Nobs[f])
    dt = []
    for i in range(N):
        y = np.random.random()
        dt.append(10.**np.interp(y, cdf[f], bins[f]))
    dates = np.cumsum(np.array(dt))

    _ = a.hist(np.log10(dt), bins=50)
    a.set_ylabel(f)


ax[-1].set_xlabel('log10(dt [days])')
plt.subplots_adjust(hspace=0)
# -



StarcoderdataJupyterScriptsDedupFiltered