<commit_before>"""
Input/output routines for tracpy.

Contains:
    setupROMSfiles
    readgrid
    readfields
    savetracks
    loadtracks
    loadtransport
    save_ll2grid
"""

import netCDF4 as netCDF
import glob
import numpy as np
from datetime import datetime, timedelta
import pdb
from matplotlib import delaunay
import octant
import time
from matplotlib.pyplot import *
import op
import os
import tracpy
from matplotlib.mlab import find

def setupROMSfiles(loc,date,ff,tout, time_units, tstride=1):
    '''
    setupROMSfiles()
    <NAME>, March 2013

    Figures out necessary files to read in for track times and what
    model output indices within those files to use.

    Input:
     loc        File location. loc can be:
                * a thredds server web address
                * a single string of a file location
                * a list of strings of multiple file locations to be searched
                through
     date       datetime format start date
     ff         Time direction. ff=1 forward, ff=-1 backward
     tout       Number of model outputs to use
     time_units To convert to datetime
     tstride    Stride in time, in case want to use less model output
                 than is available. Default is 1, using all output.

    Output:
     nc         NetCDF object for relevant files
     tinds      Indices of outputs to use from fname files
    '''
    # This addresses an issue in netCDF4 that was then fixed, but
    # this line makes updating unnecessary. Issue described here: 
    # http://code.google.com/p/netcdf4-python/issues/detail?id=170
    netCDF._set_default_format(format='NETCDF3_64BIT')

    # For thredds server where all information is available in one place
    # or for a single file
    if 'http' in loc or type(loc)==str:
        nc = netCDF.Dataset(loc)

    # This is for the case when we have a bunch of files to sort through
    else:
        # the globbing should happen ahead of time so this case looks different than
        # the single file case
        nc = netCDF.MFDataset(loc) # files in fname are in chronological order

    # Convert date to number
    dates = netCDF.num2date(nc.variables['ocean_time'][:], time_units)
    # time index with time value just below date (relative to file ifile)
    istart = find(dates<=date)[-1]

    # Select indices 
    if ff==1:
        tinds = range(istart,istart+tout, tstride) # indices of model outputs desired
    else: # backward in time
        # have to shift istart since there are now new indices behind since going backward
        tinds = range(istart,istart-tout, -tstride)

    return nc, tinds

def readgrid(grid_filename, vert_filename=None, llcrnrlon=-98.5, llcrnrlat=22.5, 
            urcrnrlon=-87.5, urcrnrlat=31.0, lat_0=30, lon_0=-94, res='i', 
            usebasemap=False, usespherical=True):
    '''
    readgrid(loc)
    <NAME>, March 2013
    This function should be read in at the beginnind of a run.py call.
    It reads in all necessary grid information that won't change in time
    and stores it in a dictionary called grid.
    All arrays are changed to Fortran ordering (from Python ordering)
    and to tracmass variables ordering from ROMS ordering 
    i.e. from [t,k,j,i] to [i,j,k,t]
    right away after reading in.

    Input:
     grid_filename    File name (with extension) where grid information is stored
     vert_filename    (optional) File name (with extension) where vertical grid information
                    is stored, if not in grid_loc. Can also skip this if don't need 
                    vertical grid info.
     also optional basemap box parameters. Default is for full shelf model.
     usebasemap          (False) Whether to use load basemap into grid (True) or pyproj (False).
                    Basemap is slower but can be used for plotting, and pyproj is the opposite.

    Output:
     grid           Dictionary containing all necessary time-independent grid fields

    grid dictionary contains: (array sizing is for tracmass ordering)
     imt,jmt,km     Grid index sizing constants in (x,y,z), are for horizontal rho grid [scalar]
     dxv            Horizontal grid cell walls areas in x direction [imt,jmt-1]
     dyu            Horizontal grid cell walls areas in y direction [imt-1,jmt]
     dxdy           Horizontal area of cells defined at cell centers [imt,jmt]
     mask           Land/sea mask [imt,jmt] 
     pm,pn          Difference in horizontal grid spacing in x and y [imt,jmt]
     kmt            Number of vertical levels in horizontal space [imt,jmt]
     dzt0           Thickness in meters of grid at each k-level with time-independent free surface. 
                    Surface is at km [imt,jmt,km].
     zrt0           Depth in meters of grid at each k-level on vertical rho grid with time-independent 
                    free surface. Surface is at km [imt,jmt,km]
     zwt0           Depth in meters of grid at each k-level on vertical w grid with time-independent 
                    free surface. Surface is at km [imt,jmt,km]
     xr,yr          Rho grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xu,yu          U grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xv,yv          V grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xpsi,ypsi      Psi grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     X,Y            Grid index arrays
     tri,trir       Delaunay triangulations
     Cs_r,sc_r      Vertical grid streching paramters [km-1]
     hc             Critical depth [scalar]
     h              Depths [imt,jmt]
     theta_s        Vertical stretching parameter [scalar]. A parameter 
                    (typically 0.0 <= theta_s < 5.0) that defines the amount 
                    of grid focusing. A higher value for theta_s will focus the grid more.
     theta_b        Vertical stretching parameter [scalar]. A parameter (0.0 < theta_b < 1.0) 
                    that says whether the coordinate will be focused at the surface 
                    (theta_b -> 1.0) or split evenly between surface and bottom (theta_b -> 0)
     basemap        Basemap object

    Note: all are in fortran ordering and tracmass ordering except for X, Y, tri, and tric
    To test: [array].flags['F_CONTIGUOUS'] will return true if it is fortran ordering
    '''
    keeptime = 0 # do timing for readgrid
    if keeptime: starttime = time.time()

    # Read in grid parameters and find x and y in domain on different grids
    # use full dataset to get grid information
    # This addresses an issue in netCDF4 that was then fixed, but
    # this line makes updating unnecessary. Issue described here:
    # http://code.google.com/p/netcdf4-python/issues/detail?id=170
    netCDF._set_default_format(format='NETCDF3_64BIT')
    gridfile = netCDF.Dataset(grid_filename)

    # # Read in whether grid is spherical or not
    # try:
    #     usesphericaltemp = gridfile.variables['spherical'][:]
    #     if usesphericaltemp == 'T':
    #         usespherical = True
    #     else:
    #         usespherical = False
    # except KeyError: # Assume not lon/lat if spherical flag is not in grid file
    #     usespherical = False

    # Basemap parameters.
    if usespherical:
        llcrnrlon=llcrnrlon; llcrnrlat=llcrnrlat; 
        urcrnrlon=urcrnrlon; urcrnrlat=urcrnrlat; projection='lcc'
        lat_0=lat_0; lon_0=lon_0; resolution=res; area_thresh=0.
        # pdb.set_trace()
        if usebasemap:
            from mpl_toolkits.basemap import Basemap
            basemap = Basemap(llcrnrlon=llcrnrlon,
                         llcrnrlat=llcrnrlat,
                         urcrnrlon=urcrnrlon,
                         urcrnrlat=urcrnrlat,
                         projection=projection,
                         lat_0=lat_0,
                         lon_0=lon_0,
                         resolution=resolution,
                         area_thresh=area_thresh)
        else:
            from pyproj import Proj
            # this gives somewhat different differences between projected coordinates as compared with previous basemap
            # definition for the default values.
            basemap = Proj(proj='lcc', lat_1=llcrnrlat, lat_2=urcrnrlat, lat_0=lat_0, lon_0=lon_0, x_0=0, y_0=0,ellps='clrk66',datum='NAD27')
            # basemap = (proj='lcc',lat_1=44.33333333333334,lat_2=46,lat_0=43.66666666666666, lon_0=-120.5,x_0=609601.2192024384, y_0=0,ellps='clrk66',datum='NAD27')
            # basemap = Proj("+proj=lcc +lat_0=lat_0 +lon_0=lon_0")
                            # +x_0=1700000 \<commit_msg>commented out netcdf file format lines since newer version seems to die from them<commit_after>"""
Input/output routines for tracpy.

Contains:
    setupROMSfiles
    readgrid
    readfields
    savetracks
    loadtracks
    loadtransport
    save_ll2grid
"""

import netCDF4 as netCDF
import glob
import numpy as np
from datetime import datetime, timedelta
import pdb
from matplotlib import delaunay
import octant
import time
from matplotlib.pyplot import *
import op
import os
import tracpy
from matplotlib.mlab import find

def setupROMSfiles(loc,date,ff,tout, time_units, tstride=1):
    '''
    setupROMSfiles()
    <NAME>, March 2013

    Figures out necessary files to read in for track times and what
    model output indices within those files to use.

    Input:
     loc        File location. loc can be:
                * a thredds server web address
                * a single string of a file location
                * a list of strings of multiple file locations to be searched
                through
     date       datetime format start date
     ff         Time direction. ff=1 forward, ff=-1 backward
     tout       Number of model outputs to use
     time_units To convert to datetime
     tstride    Stride in time, in case want to use less model output
                 than is available. Default is 1, using all output.

    Output:
     nc         NetCDF object for relevant files
     tinds      Indices of outputs to use from fname files
    '''
    # # This addresses an issue in netCDF4 that was then fixed, but
    # # this line makes updating unnecessary. Issue described here: 
    # # http://code.google.com/p/netcdf4-python/issues/detail?id=170
    # netCDF._set_default_format(format='NETCDF3_64BIT')

    # For thredds server where all information is available in one place
    # or for a single file
    if 'http' in loc or type(loc)==str:
        nc = netCDF.Dataset(loc)

    # This is for the case when we have a bunch of files to sort through
    else:
        # the globbing should happen ahead of time so this case looks different than
        # the single file case
        nc = netCDF.MFDataset(loc) # files in fname are in chronological order

    # Convert date to number
    dates = netCDF.num2date(nc.variables['ocean_time'][:], time_units)
    # time index with time value just below date (relative to file ifile)
    istart = find(dates<=date)[-1]

    # Select indices 
    if ff==1:
        tinds = range(istart,istart+tout, tstride) # indices of model outputs desired
    else: # backward in time
        # have to shift istart since there are now new indices behind since going backward
        tinds = range(istart,istart-tout, -tstride)

    return nc, tinds

def readgrid(grid_filename, vert_filename=None, proj='lcc', llcrnrlon=-98.5, llcrnrlat=22.5, 
            urcrnrlon=-87.5, urcrnrlat=31.0, lat_0=30, lon_0=-94, res='i', 
            usebasemap=False, usespherical=True):
    '''
    readgrid(loc)
    <NAME>, March 2013
    This function should be read in at the beginnind of a run.py call.
    It reads in all necessary grid information that won't change in time
    and stores it in a dictionary called grid.
    All arrays are changed to Fortran ordering (from Python ordering)
    and to tracmass variables ordering from ROMS ordering 
    i.e. from [t,k,j,i] to [i,j,k,t]
    right away after reading in.

    Input:
     grid_filename    File name (with extension) where grid information is stored
     vert_filename    (optional) File name (with extension) where vertical grid information
                    is stored, if not in grid_loc. Can also skip this if don't need 
                    vertical grid info.
     also optional basemap box parameters. Default is for full shelf model.
     proj           Basemap projection.
     usebasemap          (False) Whether to use load basemap into grid (True) or pyproj (False).
                    Basemap is slower but can be used for plotting, and pyproj is the opposite.

    Output:
     grid           Dictionary containing all necessary time-independent grid fields

    grid dictionary contains: (array sizing is for tracmass ordering)
     imt,jmt,km     Grid index sizing constants in (x,y,z), are for horizontal rho grid [scalar]
     dxv            Horizontal grid cell walls areas in x direction [imt,jmt-1]
     dyu            Horizontal grid cell walls areas in y direction [imt-1,jmt]
     dxdy           Horizontal area of cells defined at cell centers [imt,jmt]
     mask           Land/sea mask [imt,jmt] 
     pm,pn          Difference in horizontal grid spacing in x and y [imt,jmt]
     kmt            Number of vertical levels in horizontal space [imt,jmt]
     dzt0           Thickness in meters of grid at each k-level with time-independent free surface. 
                    Surface is at km [imt,jmt,km].
     zrt0           Depth in meters of grid at each k-level on vertical rho grid with time-independent 
                    free surface. Surface is at km [imt,jmt,km]
     zwt0           Depth in meters of grid at each k-level on vertical w grid with time-independent 
                    free surface. Surface is at km [imt,jmt,km]
     xr,yr          Rho grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xu,yu          U grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xv,yv          V grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     xpsi,ypsi      Psi grid zonal (x) and meriodional (y) coordinates [imt,jmt]
     X,Y            Grid index arrays
     tri,trir       Delaunay triangulations
     Cs_r,sc_r      Vertical grid streching paramters [km-1]
     hc             Critical depth [scalar]
     h              Depths [imt,jmt]
     theta_s        Vertical stretching parameter [scalar]. A parameter 
                    (typically 0.0 <= theta_s < 5.0) that defines the amount 
                    of grid focusing. A higher value for theta_s will focus the grid more.
     theta_b        Vertical stretching parameter [scalar]. A parameter (0.0 < theta_b < 1.0) 
                    that says whether the coordinate will be focused at the surface 
                    (theta_b -> 1.0) or split evenly between surface and bottom (theta_b -> 0)
     basemap        Basemap object

    Note: all are in fortran ordering and tracmass ordering except for X, Y, tri, and tric
    To test: [array].flags['F_CONTIGUOUS'] will return true if it is fortran ordering
    '''
    keeptime = 0 # do timing for readgrid
    if keeptime: starttime = time.time()

    # Read in grid parameters and find x and y in domain on different grids
    # use full dataset to get grid information
    # This addresses an issue in netCDF4 that was then fixed, but
    # this line makes updating unnecessary. Issue described here:
    # http://code.google.com/p/netcdf4-python/issues/detail?id=170
    # Newer netCDF is complaining about this line so let's try without it.
    # netCDF._set_default_format(format='NETCDF3_64BIT')
    gridfile = netCDF.Dataset(grid_filename)

    # # Read in whether grid is spherical or not
    # try:
    #     usesphericaltemp = gridfile.variables['spherical'][:]
    #     if usesphericaltemp == 'T':
    #         usespherical = True
    #     else:
    #         usespherical = False
    # except KeyError: # Assume not lon/lat if spherical flag is not in grid file
    #     usespherical = False

    # Basemap parameters.
    if usespherical:
        llcrnrlon=llcrnrlon; llcrnrlat=llcrnrlat; 
        urcrnrlon=urcrnrlon; urcrnrlat=urcrnrlat; projection=proj
        lat_0=lat_0; lon_0=lon_0; resolution=res; area_thresh=0.
        # pdb.set_trace()
        if usebasemap:
            from mpl_toolkits.basemap import Basemap
            basemap = Basemap(llcrnrlon=llcrnrlon,
                         llcrnrlat=llcrnrlat,
                         urcrnrlon=urcrnrlon,
                         urcrnrlat=urcrnrlat,
                         projection=projection,
                         lat_0=lat_0,
                         lon_0=lon_0,
                         resolution=resolution,
                         area_thresh=area_thresh)
        else:
            from pyproj import Proj
            # this gives somewhat different differences between projected coordinates as compared with previous basemap
            # definition for the default values.
            basemap = Proj(proj='lcc', lat_1=llcrnrlat, lat_2=urcrnrlat, lat_0=lat_0, lon_0=lon_0, x_0=0, y_0=0,ellps='clrk66',datum='NAD27')
            # basemap = (proj='lcc',lat_1=44.33333333333334,lat_2=46,lat_0=43.66666666666666, lon_0=-120.5,x_0=609601.2192024384, y_0=0,ellps='clrk66',datum='NAD27')
            # basemap = Proj("+proj=lcc +lat_0=lat_0 +lon_0=lon_0")
                            # +x_0=1700000 \
{'fraction_non_alphanumeric': 0.06595744680851064, 'fraction_numerical': 0.019436457734330073, 'mean_word_length': 2.303133903133903, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, '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': '3776857', 'n_tokens_mistral': 5412, 'n_tokens_neox': 5039, 'n_words': 1873}