<commit_before>import matplotlib as mpl
mpl.use('Agg')
import os
import time
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from xmitgcm import open_mdsdataset
# import argparse
from dask.diagnostics import ProgressBar
from dask.array import from_array
from dask.multiprocessing import get

def mitgcm_Movie(ddir,
                prefix=['tracer_snapshots'],
                maskname='hFacC',
                varname='TRAC01',
                clim =[-1,35]):
    ds   = open_mdsdataset(ddir,prefix=prefix,
                            swap_dims=False)
    ds   = ds[varname].where(ds[maskname]==1)
    run  = os.path.basename(os.path.normpath(ddir))
    odir = ddir+'/movie'
    Movie(ds,odir,clim=clim,moviename=run)

def Movie(da,odir,
            varname     = None,
            framedim    = 'time',
            moviename   = 'movie',
            plotstyle   = 'simple',
            clim        = None,
            cmap        = None,
            bgcolor     = np.array([1,1,1])*0.3,
            framewidth  = 1280,
            frameheight = 720,
            dpi         = 100,
            dask        = True,
            ):
    # Set defaults:

    if not isinstance(da,xr.DataArray):
        raise RuntimeError('input has to be an xarray DataStructure, instead\
        is '+str(type(da)))

    if not os.path.exists(odir):
        os.makedirs(odir)

    # Infer defaults from data
    if not clim:
        print('clim will be inferred from data, this can take very long...')
        clim = [da.min(),da.max()]
    if not cmap:
        cmap = plt.cm.RdYlBu_r

    # Annnd here we go
    print('+++ Execute plot function +++')
    if dask:
        data   = da.data
        frame_axis = da.get_axis_num(framedim)
        drop_axis = [da.get_axis_num(a) for a in da.dims if not a == framedim]
        chunks = list(data.shape)
        chunks[frame_axis] = 1
        data = data.rechunk(chunks)
        # with ProgressBar():
        dummy = data.map_blocks(FramePrint,chunks = [1],
                                drop_axis = drop_axis,
                                dtype=np.float64,
                                dask=dask,
                                frame_axis = frame_axis,
                                odir        = odir,
                                cmap        = cmap,
                                clim        = clim,
                                framewidth  = framewidth,
                                frameheight = frameheight,
                                dpi         = dpi,
                                bgcolor     = bgcolor
                                # this line is some dask 'magic': it parallelizes
                                # the print function with processes and not threads,
                                # which is a lot faster for custom functions
                                # apparently!
                                ).compute(get=get)

    else:
        # do it with a simple for loop...can this really be quicker?
        for ii in range(0,len(da.time)):
            start_time = time.time()
            da_slice = da[{framedim:ii}]
            # fig,ax,h = FramePrint(da_slice,
            dummy = FramePrint(da_slice,
                                    frame       = ii,
                                    odir        = odir,
                                    cmap        = cmap,
                                    clim        = clim,
                                    framewidth  = framewidth,
                                    frameheight = frameheight,
                                    dpi         = dpi,
                                    bgcolor     = bgcolor
                                    )
            if ii % 100 == 0:
                remaining_time = (len(da.time)-ii)*(time.time() - start_time)/60
                print('FRAME---%04d---' %ii)
                print('Estimated time left : %d minutes' %remaining_time)<commit_msg>started to pep 8 this babay<commit_after>import matplotlib as mpl
import os
import time
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from xmitgcm import open_mdsdataset
# import argparse
# from dask.diagnostics import ProgressBar
from dask.array import from_array
from dask.multiprocessing import get
mpl.use('Agg')


def mitgcm_Movie(ddir,
                 prefix=['tracer_snapshots'],
                 maskname='hFacC',
                 varname='TRAC01',
                 clim=[-1,  35]):
    ds = open_mdsdataset(ddir, prefix=prefix,
                         swap_dims=False)
    ds = ds[varname].where(ds[maskname] == 1)
    run = os.path.basename(os.path.normpath(ddir))
    odir = ddir+'/movie'
    Movie(ds, odir, clim=clim, moviename=run)


def Movie(da, odir,
          varname=None,
          framedim='time',
          moviename='movie',
          plotstyle='simple',
          clim=None,
          cmap=None,
          bgcolor=np.array([1, 1, 1])*0.3,
          framewidth=1280,
          frameheight=720,
          dpi=100,
          dask=True,
          ):
    # Set defaults:

    if not isinstance(da, xr.DataArray):
        raise RuntimeError('input has to be an xarray DataStructure, instead\
        is '+str(type(da)))

    if not os.path.exists(odir):
        os.makedirs(odir)

    # Infer defaults from data
    if not clim:
        print('clim will be inferred from data, this can take very long...')
        clim = [da.min(),da.max()]
    if not cmap:
        cmap = plt.cm.RdYlBu_r

    # Annnd here we go
    print('+++ Execute plot function +++')
    if dask:
        data   = da.data
        frame_axis = da.get_axis_num(framedim)
        drop_axis = [da.get_axis_num(a) for a in da.dims if not a == framedim]
        chunks = list(data.shape)
        chunks[frame_axis] = 1
        data = data.rechunk(chunks)
        # with ProgressBar():
        dummy = data.map_blocks(FramePrint,chunks = [1],
                                drop_axis = drop_axis,
                                dtype=np.float64,
                                dask=dask,
                                frame_axis = frame_axis,
                                odir        = odir,
                                cmap        = cmap,
                                clim        = clim,
                                framewidth  = framewidth,
                                frameheight = frameheight,
                                dpi         = dpi,
                                bgcolor     = bgcolor
                                # this line is some dask 'magic': it parallelizes
                                # the print function with processes and not threads,
                                # which is a lot faster for custom functions
                                # apparently!
                                ).compute(get=get)

    else:
        # do it with a simple for loop...can this really be quicker?
        for ii in range(0,len(da.time)):
            start_time = time.time()
            da_slice = da[{framedim:ii}]
            # fig,ax,h = FramePrint(da_slice,
            dummy = FramePrint(da_slice,
                                    frame=ii,
                                    odir        = odir,
                                    cmap        = cmap,
                                    clim        = clim,
                                    framewidth  = framewidth,
                                    frameheight = frameheight,
                                    dpi         = dpi,
                                    bgcolor     = bgcolor
                                    )
            if ii % 100 == 0:
                remaining_time = (len(da.time)-ii)*(time.time() - start_time)/60
                print('FRAME---%04d---' %ii)
                print('Estimated time left : %d minutes' %remaining_time)
{'fraction_non_alphanumeric': 0.08637407501913753, 'fraction_numerical': 0.008803266139321255, 'mean_word_length': 1.1342227062346855, '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': 46, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '334308', 'n_tokens_mistral': 2220, 'n_tokens_neox': 2103, 'n_words': 523}