<commit_before>import sys
sys.path.append('/home/jwalker/dynamics/python/atmos-tools')
sys.path.append('/home/jwalker/dynamics/python/atmos-read')

import numpy as np
import xray
import pandas as pd
import matplotlib.pyplot as plt
import collections

import atmos as atm
import merra
import indices
import utils

# ----------------------------------------------------------------------
yearstr = '1979-2014'
onset_nm = 'CHP_MFC'
ndays = 5
lon1, lon2 = 60, 100
plev = 200
daynm, yearnm, latname, lonname = 'dayrel', 'year', 'YDim', 'XDim'
datadir = atm.homedir() + 'datastore/merra/analysis/'
savedir = 'figs/'
filenm = 'merra_ubudget%d_dailyrel_%s_ndays%d_%dE-%dE_%s.nc'
filenm = datadir + filenm % (plev, onset_nm, ndays, lon1, lon2, yearstr)
files = {}
files['ubudget'] = filenm
varnms = ['U', 'V']
for nm in varnms:
    filenm = datadir + 'merra_%s%d_dailyrel_%s_%s.nc'
    files[nm] = filenm % (nm, plev, onset_nm, yearstr)
    filenm = datadir + 'merra_%s_sector_%dE-%dE_dailyrel_%s_%s.nc'
    files[nm + '_latp'] = filenm % (nm, lon1, lon2, onset_nm, yearstr)

# ----------------------------------------------------------------------
# Read data from each year

# Zonal momentum budget components
with xray.open_dataset(files['ubudget']) as ubudget:
    ubudget.load()


# Scaling factor for all terms in momentum budget
scale = 1e-4
ubudget.attrs['comp_units'] = '%.0e m/s2' % scale
ubudget = ubudget / scale

# Read other lat-lon variables and smooth with rolling mean
for nm in varnms:
    varnm = '%s%d' % (nm, plev)
    with xray.open_dataset(files[nm]) as ds:
        var = ds[varnm].load()
    daydim = atm.get_coord(var, coord_name=daynm, return_type='dim')
    ubudget[nm] = atm.rolling_mean(var, ndays, axis=daydim, center=True)

# Read other lat-pres variables and smooth with rolling mean
data_latp = xray.Dataset()
for nm in varnms:
    varnm = nm + '_latp'<commit_msg>Added multiple plevels<commit_after>import sys
sys.path.append('/home/jwalker/dynamics/python/atmos-tools')
sys.path.append('/home/jwalker/dynamics/python/atmos-read')

import numpy as np
import xray
import pandas as pd
import matplotlib.pyplot as plt
import collections

import atmos as atm
import merra
import indices
import utils

# ----------------------------------------------------------------------
yearstr = '1979-2014'
onset_nm = 'CHP_MFC'
ndays = 5
lon1, lon2 = 60, 100
plevs = [1000,925,850,775,700,600,500,400,300,250,200,150,100,70,50,30,20]
daynm, yearnm, latname, lonname = 'dayrel', 'year', 'YDim', 'XDim'
datadir = atm.homedir() + 'datastore/merra/analysis/'
savedir = 'figs/'
filenm = datadir + 'merra_ubudget%d_dailyrel_%s_ndays%d_%dE-%dE_%s.nc'
filenm = datadir + filenm % (plev, onset_nm, ndays, lon1, lon2, yearstr)
files = {}
files['ubudget'] = [filenm % (plev, onset_nm, ndays, lon1, lon2, yearstr)
                    for plev in plevs]
varnms = ['U', 'V']
for nm in varnms:
    filenm = datadir + 'merra_%s%d_dailyrel_%s_%s.nc'
    files[nm] = filenm % (nm, plev, onset_nm, yearstr)
    filenm = datadir + 'merra_%s_sector_%dE-%dE_dailyrel_%s_%s.nc'
    files[nm + '_latp'] = filenm % (nm, lon1, lon2, onset_nm, yearstr)

# ----------------------------------------------------------------------
# Read data from each year

# Zonal momentum budget components
ubudget = xray.Dataset()
pname, pdim = 'Height', 1
for i, plev in enumerate(plevs):
    filenm = files['ubudget'][i]
    with xray.open_dataset(filenm) as ds:
        ds.load()
    for nm in ds.data_vars:
        ds[nm] = atm.expand_dims(ds[nm], pname, plev, axis=pdim)
    if i == 0:
        ubudget = ds
    else:
        ubudget = xray.concat([ubudget, ds], dim=pname)


# Scaling factor for all terms in momentum budget
scale = 1e-4
ubudget.attrs['comp_units'] = '%.0e m/s2' % scale
ubudget = ubudget / scale

# Read other lat-lon variables and smooth with rolling mean
for nm in varnms:
    varnm = '%s%d' % (nm, plev)
    with xray.open_dataset(files[nm]) as ds:
        var = ds[varnm].load()
    daydim = atm.get_coord(var, coord_name=daynm, return_type='dim')
    ubudget[nm] = atm.rolling_mean(var, ndays, axis=daydim, center=True)

# Read other lat-pres variables and smooth with rolling mean
data_latp = xray.Dataset()
for nm in varnms:
    varnm = nm + '_latp'
{'fraction_non_alphanumeric': 0.19477154969382948, 'fraction_numerical': 0.024258125294394726, 'mean_word_length': 3.0332383665717, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, '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': '7456621', 'n_tokens_mistral': 1680, 'n_tokens_neox': 1594, 'n_words': 398}