text:
# Simple one-layer model with a drn and sfr network on top.  Purpose is to
# test movement of solute between stress and advanced packages.  In this case
# water from a drain is moved into the first sfr reach.  The test confirms
# that the solute from the drain is moved into the sfr reach.
# There is no flow between the stream and the aquifer.

import os
import shutil
import numpy as np

try:
    import flopy
except:
    msg = "Error. FloPy package is not available.\n"
    msg += "Try installing using the following command:\n"
    msg += " pip install flopy"
    raise Exception(msg)

from framework import set_teardown_test

import targets
exe_name_mf6 = targets.target_dict["mf6"]
exe_name_mf6 = os.path.abspath(exe_name_mf6)

testdir = "./temp"
testgroup = "mvt02fmi"
d = os.path.join(testdir, testgroup)
if os.path.isdir(d):
    shutil.rmtree(d)

ex = ["mvt02fmi"]
exdirs = []
for s in ex:
    exdirs.append(os.path.join("temp", s))

# parameters
lx = 7.0
lz = 1.0
nlay = 1
nrow = 1
ncol = 7
nper = 1
delc = 1.0
delr = lx / ncol
delz = lz / nlay
top = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
botm = list(top - np.arange(delz, nlay * delz + delz, delz))
botm[2] = -1.0
idomain = np.full((nlay, nrow, ncol), 1)

perlen = [10.0]
nstp = [10]
tsmult = [1.0]
tdis_rc = []
for i in range(nper):
    tdis_rc.append((perlen[i], nstp[i], tsmult[i]))

nouter, ninner = 20, 10
hclose, rclose, relax = 1e-8, 1e-6, 0.97

def run_flow_model():

    name = "flow"
    gwfname = name
    wsf = os.path.join(testdir, testgroup, name)
    sim = flopy.mf6.MFSimulation(
        sim_name=name, sim_ws=wsf, exe_name=exe_name_mf6
    )

    # create tdis package
    tdis = flopy.mf6.ModflowTdis(
        sim, time_units="DAYS", nper=nper, perioddata=tdis_rc
    )

    # create gwf model
    gwf = flopy.mf6.ModflowGwf(
        sim,
        modelname=gwfname,
        save_flows=True,
    )

    imsgwf = flopy.mf6.ModflowIms(
        sim,
        print_option="ALL",
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration="BICGSTAB",
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
        filename="{}.ims".format(gwfname),
    )

    dis = flopy.mf6.ModflowGwfdis(
        gwf,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botm,
        idomain=idomain,
    )

    # initial conditions
    ic = flopy.mf6.ModflowGwfic(gwf, strt=np.arange(ncol))

    # node property flow
    Kh = 20.0
    Kv = 20.0
    npf = flopy.mf6.ModflowGwfnpf(
        gwf,
        xt3doptions=False,
        save_flows=True,
        save_specific_discharge=True,
        save_saturation=True,
        icelltype=0,
        k=Kh,
        k33=Kv,
    )

    # chd files
    drnlist1 = [
        [(0, 0, ncol - 1), 0.0, Kh],
    ]
    drn1 = flopy.mf6.ModflowGwfdrn(
        gwf,
        stress_period_data=drnlist1,
        print_input=True,
        print_flows=True,
        save_flows=False,
        mover=True,
        pname="DRN-1",
    )

    # wel files
    wellist1 = [
        [(0, 0, 0), 1.0, 100.0],
    ]
    wel1 = flopy.mf6.ModflowGwfwel(
        gwf,
        stress_period_data=wellist1,
        print_input=True,
        print_flows=True,
        save_flows=False,
        pname="WEL-1",
        auxiliary="CONCENTRATION",
    )

    # pak_data = [<rno> <cellid(ncelldim)> <rlen> <rwid> <rgrd> <rtp> <rbth> <rhk> <man> <ncon> <ustrf> <ndv> [<aux(naux)>] [<boundname>]]
    rlen = delr
    rwid = delc
    rgrd = 1.0
    rtp = 0.0
    rbth = 0.1
    rhk = 0.0  # 0.01
    rman = 1.0
    ncon = 2
    ustrf = 1.0
    ndv = 0
    pak_data = []
    for irno in range(ncol):
        ncon = 2
        if irno in [0, ncol - 1]:
            ncon = 1
        cellid = (0, 0, irno)
        t = (
            irno,
            cellid,
            rlen,
            rwid,
            rgrd,
            rtp,
            rbth,
            rhk,
            rman,
            ncon,
            ustrf,
            ndv,
        )
        pak_data.append(t)

    con_data = []
    for irno in range(ncol):
        if irno == 0:
            t = (irno, -(irno + 1))
        elif irno == ncol - 1:
            t = (irno, irno - 1)
        else:
            t = (irno, irno - 1, -(irno + 1))
        con_data.append(t)

    p_data = [
        (0, "INFLOW", 0.0),
    ]

    # note: for specifying sfr number, use fortran indexing!
    # sfr_obs = {('sfr_obs.csv'): [('lakestage', 'stage', 1),
    #                             ('lakevolume', 'volume', 1),
    #                             ('lak1', 'lak', 1, 1),
    #                             ('lak2', 'lak', 1, 2),
    #                             ('lak3', 'lak', 1, 3)]}

    sfr_on = True
    if sfr_on:
        sfr = flopy.mf6.modflow.ModflowGwfsfr(
            gwf,
            save_flows=True,
            print_input=True,
            print_flows=True,
            print_stage=True,
            mover=True,
            stage_filerecord=f"{gwfname}.sfr.stg",
            budget_filerecord=f"{gwfname}.sfr.bud",
            nreaches=ncol,
            packagedata=pak_data,
            pname="SFR-1",
            connectiondata=con_data,
            perioddata=p_data,
            # observations=lak_obs,
            # auxiliary=['CONCENTRATION',
            #           'DENSITY'],
        )

    # output control
    oc = flopy.mf6.ModflowGwfoc(
        gwf,
        budget_filerecord=f"{gwfname}.bud",
        head_filerecord=f"{gwfname}.hds",
        headprintrecord=[("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")],
        saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
        printrecord=[("HEAD", "LAST"), ("BUDGET", "LAST")],
    )

    packages = [
        ("drn-1",),
        ("sfr-1",),
    ]
    perioddata = [
        ("drn-1", 0, "sfr-1", 0, "factor", 1.0),
    ]
    mvr = flopy.mf6.ModflowGwfmvr(
        gwf,
        maxmvr=len(perioddata),
        budget_filerecord=f"{gwfname}.mvr.bud",
        maxpackages=len(packages),
        print_flows=True,
        packages=packages,
        perioddata=perioddata,
    )

    sim.write_simulation()
    success, buff = sim.run_simulation(silent=False)
    errmsg = "flow model did not terminate successfully\n{}".format(buff)
    assert success, errmsg

    return

def run_transport_model():

    name = "transport"
    gwtname = name
    wst = os.path.join(testdir, testgroup, name)
    sim = flopy.mf6.MFSimulation(
        sim_name=name,
        version="mf6",
        exe_name=exe_name_mf6,
        sim_ws=wst,
        continue_=False,
        memory_print_option=["ALL"],
    )

    # create tdis package
    tdis = flopy.mf6.ModflowTdis(
        sim, time_units="DAYS", nper=nper, perioddata=tdis_rc
    )

    gwt = flopy.mf6.ModflowGwt(
        sim,
        modelname=gwtname,
        model_nam_file="{}.nam".format(gwtname),
    )

    imsgwt = flopy.mf6.ModflowIms(
        sim,
        print_option="ALL",
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration="BICGSTAB",
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
        filename="{}.ims".format(gwtname),
    )

    dis = flopy.mf6.ModflowGwtdis(
        gwt,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botm,
        idomain=idomain,
    )

    # initial conditions
    ic = flopy.mf6.ModflowGwtic(
        gwt, strt=10.0, filename="{}.ic".format(gwtname)
    )

    # advection
    adv = flopy.mf6.ModflowGwtadv(
        gwt, scheme="UPSTREAM", filename="{}.adv".format(gwtname)
    )

    # storage
    porosity = 1.0
    sto = flopy.mf6.ModflowGwtmst(
        gwt, porosity=porosity, filename="{}.sto".format(gwtname)
    )
    # sources
    sourcerecarray = [
        ("WEL-1", "AUX", "CONCENTRATION"),
    ]
    ssm = flopy.mf6.ModflowGwtssm(
        gwt, sources=sourcerecarray, filename="{}.ssm".format(gwtname)
    )

    # sft
    sftpackagedata = []
    for irno in range(ncol):
        t = (irno, 0.0, 99.0, 999.0, "myreach{}".format(irno + 1))
        sftpackagedata.append(t)

    sftperioddata = [
        (0, "STATUS", "ACTIVE"),
    ]

    sft_obs = {
        (gwtname + ".sft.obs.csv",): [
            ("sft-{}-conc".format(i + 1), "CONCENTRATION", i + 1)
            for i in range(7)
        ]
        + [
            ("sft-1-extinflow", "EXT-INFLOW", 1),
            ("sft-1-rain", "RAINFALL", 1),
            ("sft-1-roff", "RUNOFF", 1),
            ("sft-1-evap", "EVAPORATION", 1),
            ("sft-1-const", "CONSTANT", 1),
            ("sft-1-gwt1", "SFT", 1, 1),
            ("sft-1-gwt2", "SFT", 2, 1),
            ("sft-1-gwt3", "SFT", 3, 1),
            ("sft-1-myreach1", "SFT", "MYREACH1"),
        ],
    }
    # append additional obs attributes to obs dictionary
    sft_obs["digits"] = 12
    sft_obs["print_input"] = True
    sft_obs["filename"] = gwtname + ".sft.obs"

    sft = flopy.mf6.modflow.ModflowGwtsft(
        gwt,
        boundnames=True,
        save_flows=True,
        print_input=True,
        print_flows=True,
        print_concentration=True,
        concentration_filerecord=gwtname + ".sft.bin",
        budget_filerecord=gwtname + ".sft.bud",
        packagedata=sftpackagedata,
        reachperioddata=sftperioddata,
        observations=sft_obs,
        pname="SFR-1",
        auxiliary=["aux1", "aux2"],
    )

    # mover transport package
    fname = "{}.mvt.bud".format(gwtname)
    mvt = flopy.mf6.modflow.ModflowGwtmvt(
        gwt, print_flows=True, budget_filerecord=fname
    )

    pd = [
        ("GWFHEAD", "../flow/flow.hds", None),
        ("GWFBUDGET", "../flow/flow.bud", None),
        ("GWFMOVER", "../flow/flow.mvr.bud", None),
        ("SFR-1", "../flow/flow.sfr.bud", None),
    ]
    fmi = flopy.mf6.ModflowGwtfmi(gwt, packagedata=pd)

    # output control
    oc = flopy.mf6.ModflowGwtoc(
        gwt,
        budget_filerecord=f"{gwtname}.bud",
        budgetcsv_filerecord=f"{gwtname}.bud.csv",
        concentration_filerecord=f"{gwtname}.ucn",
        concentrationprintrecord=[
            ("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")
        ],
        saverecord=[("CONCENTRATION", "ALL"), ("BUDGET", "ALL")],
        printrecord=[("CONCENTRATION", "ALL"), ("BUDGET", "ALL")],
    )

    sim.write_simulation()
    success, buff = sim.run_simulation(silent=False)
    errmsg = "transport model did not terminate successfully\n{}".format(buff)
    assert success, errmsg

    print("evaluating results...")

    # Load csv budget and make sure names are correct
    fname = f"{gwtname}.bud.csv"
    fname = os.path.join(gwt.model_ws, fname)
    budcsv = np.genfromtxt(fname, names=True, delimiter=",", deletechars="")
    answer = ['time', 'STORAGE-AQUEOUS(MST)_IN', 'WEL(SSM_WEL-1)_IN', 'DRN(SSM_DRN-1)_IN', 'DRN-TO-MVR(SSM_DRN-1)_IN',
              'SFT(SFR-1)_IN', 'STORAGE-AQUEOUS(MST)_OUT', 'WEL(SSM_WEL-1)_OUT', 'DRN(SSM_DRN-1)_OUT',
              'DRN-TO-MVR(SSM_DRN-1)_OUT', 'SFT(SFR-1)_OUT', 'TOTAL_IN', 'TOTAL_OUT', 'PERCENT_DIFFERENCE']
    for i, name in enumerate(budcsv.dtype.names[:len(answer)]):
        assert answer[i] == name

    # ensure sfr concentrations were saved
    csft = sft.output.concentration().get_data().flatten()

    # load the aquifer concentrations
    caq = gwt.output.concentration().get_data().flatten()

    # sft observation results
    fpth = os.path.join(gwt.model_ws, gwtname + ".sft.obs.csv")
    try:
        tc = np.genfromtxt(fpth, names=True, delimiter=",")
    except:
        assert False, 'could not load data from "{}"'.format(fpth)

    # compare observation concs with binary file concs
    for i in range(7):
        oname = f"SFT{i+1}CONC"
        assert np.allclose(tc[oname][-1], csft[i]), "{} {}".format(
            tc[oname][-1], csft[i]
        )

    simres = tc["SFT1CONC"]
    answer = [
        5.35156250000,
        9.25781250000,
        13.6718750000,
        19.5703125000,
        27.1337890625,
        35.9912109375,
        45.4956054688,
        54.9609375000,
        63.8175964355,
        71.6825866699,
    ]

    assert np.allclose(simres, answer), "{} {}".format(simres, answer)

    # load the mvt budget file
    mobj = mvt.output.budget()

    # load the sft budget file
    bobj = sft.output.budget()

    # check the flow-ja-face terms
    res = bobj.get_data(text="flow-ja-face")[-1]
    # print(res)

    # check the storage terms, which include the total mass in the reach as an aux variable
    res = bobj.get_data(text="storage")[-1]
    # print(res)

    # uncomment when testing so files aren't deleted
    # assert False

    return

def test_mvt02fmi():
    run_flow_model()
    run_transport_model()
    d = os.path.join(testdir, testgroup)

    teardowntest = set_teardown_test()
    if teardowntest:
        if os.path.isdir(d):
            shutil.rmtree(d)
    return

if __name__ == "__main__":
    # print message
    print("standalone run of {}".format(os.path.basename(__file__)))

    # run tests
    test_mvt02fmi()


text:
import torch
import torch.nn.functional as F
import numpy as np

from advbias.augmentor.adv_bias import AdvBias

"""
    Create a 3d bspline kernel matrix
"""

def bspline_kernel_3d(sigma=[1, 1, 1], order=2, asTensor=False, dtype=torch.float32, device='gpu'):
    kernel_ones = torch.ones(1, 1, *sigma)
    kernel = kernel_ones
    padding = np.array(sigma) - 1

    for i in range(1, order + 1):
        # change 2d to 3d
        kernel = F.conv3d(kernel, kernel_ones, padding=(
            padding).tolist())/(sigma[0]*sigma[1]*sigma[2])
    if asTensor:
        return kernel[0, 0, ...].to(dtype=dtype, device=device)
    else:
        return kernel[0, 0, ...].numpy()

class AdvBias3D(AdvBias):
    """
     Adv Bias field
    """

    def __init__(self,
                 config_dict={
                     'epsilon': 0.3,  # bias field magnitude
                     # spacings between two points along z, y and x direction.
                     'control_point_spacing': [64, 64, 64],
                     # we downsample images to reduce computational costs, especially with large spacing along images. To get true spacings, one should multiply control point spacing with downscale.
                     'downscale': 2,
                     'data_size': [2, 1, 128, 128, 128],
                     'interpolation_order': 3,
                     'init_mode': 'random',
                     'space': 'log'},
                 power_iteration=False,
                 use_gpu=True, debug=False):
        super(AdvBias, self).__init__(
            config_dict=config_dict, use_gpu=use_gpu, debug=debug)
        self.power_iteration = power_iteration

    def compute_smoothed_bias(self, cpoint=None, interpolation_kernel=None, padding=None, stride=None):
        '''
        generate bias field given the cpoints N*1*k*l
        :return: bias field bs*1*H*W
        '''
        if interpolation_kernel is None:
            interpolation_kernel = self.interp_kernel
        if padding is None:
            padding = self._padding
        if stride is None:
            stride = self._stride
        if cpoint is None:
            cpoint = self.param
        bias_field = F.conv_transpose3d(cpoint, interpolation_kernel,
                                        padding=padding, stride=stride, groups=1)
        # crop bias
        bias_field_tmp = bias_field[:, :,
                                    stride[0] + self._crop_start[0]:-stride[0] - self._crop_end[0],
                                    stride[1] + self._crop_start[1]:-stride[1] - self._crop_end[1],
                                    stride[2] + self._crop_start[2]:-stride[2] - self._crop_end[2],
                                    ]

        # recover bias field to original image resolution for efficiency.
        if self.debug:
            print('after bspline intep, size:', bias_field_tmp.size())
        scale_factor_d = self._image_size[0] / bias_field_tmp.size(2)
        scale_factor_h = self._image_size[1] / bias_field_tmp.size(3)
        scale_factor_w = self._image_size[2] / bias_field_tmp.size(4)

        if scale_factor_h > 1 or scale_factor_w > 1:
            upsampler = torch.nn.Upsample(scale_factor=(scale_factor_d, scale_factor_h, scale_factor_w), mode='trilinear',
                                          align_corners=False)
            diff_bias = upsampler(bias_field_tmp)
            print('recover resolution, size of bias field:', diff_bias.size())

        else:
            diff_bias = bias_field_tmp

        bias_field = torch.exp(diff_bias)

        return bias_field

    def get_bspline_kernel(self, spacing, order=3):
        '''

        :param order init: bspline order, default to 3
        :param spacing tuple of int: spacing between control points along h and w.
        :return:  kernel matrix
        '''
        self._kernel = bspline_kernel_3d(
            spacing, order=order, asTensor=True, dtype=self._dtype, device=self._device)
        self._padding = (np.array(self._kernel.size()) - 1) / 2
        self._padding = self._padding.astype(dtype=int).tolist()
        self._kernel.unsqueeze_(0).unsqueeze_(0)
        self._kernel = self._kernel.to(dtype=self._dtype, device=self._device)
        return self._kernel

    def get_name(self):
        return 'bias'

    def is_geometric(self):
        return 0

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    images = torch.zeros(2, 1, 128, 128, 128).cuda()
    images[:, :, 0:80, 60:80, 60:80] = 0.2
    images[:, :, 0:80, 20:80, 20:80] = 0.8

    images = images.float()
    images.requires_grad = False
    print('input:', images.size())
    augmentor = AdvBias3D(
        config_dict={'epsilon': 0.3,
                     'control_point_spacing': [64, 64, 64],
                     'downscale': 1,  # increase the downscale factor to save interpolation time
                     'data_size': [2, 1, 128, 128, 128],
                     'interpolation_order': 3,
                     'init_mode': 'gaussian',
                     'space': 'log'},
        power_iteration=False,
        debug=True, use_gpu=True)

    # perform random bias field
    augmentor.init_parameters()
    transformed = augmentor.forward(images)
    error = transformed-images
    print('sum error', torch.sum(error))

    plt.subplot(231)
    plt.imshow(images.detach().cpu().numpy()[0, 0, 0])

    plt.subplot(232)
    plt.imshow(transformed.detach().cpu().numpy()[0, 0, 0])

    plt.subplot(233)
    plt.imshow((augmentor.bias_field.detach()).detach().cpu().numpy()[0, 0, 0])
    plt.subplot(234)
    plt.imshow(images.detach().cpu().numpy()[0, 0, 28])
    plt.subplot(235)
    plt.imshow(transformed.detach().cpu().numpy()[0, 0, 28])

    plt.subplot(236)
    plt.imshow(augmentor.bias_field.detach().cpu().numpy()[0, 0, 28])
    plt.savefig('./result/log/test_bias_3D.png')

ext: py

reason: import numpy


text:
import os
import numpy as np
import scipy.io as sio

DEFAULT_FOLDER = 'DataBase/HadamardMatrix/'

def load_hadamard_matrix(n, folder_name=DEFAULT_FOLDER):
    """
        Loads or creates Hadamard matrix with given input dimention n
        and normalizes it such that H'H=I
    """

    fname = os.path.join(folder_name, 'H%d.mat' % n)
    if not os.path.exists(fname):
        if n == 1:
            H = np.array([1], dtype=np.float32)
        elif n % 2 != 0:
            H = None
        else:
            H = load_hadamard_matrix(n//2, folder_name)
            if H is not None:
                H = np.kron(np.array([[1, 1], [1, -1]]), H) / np.sqrt(2)
    else:
        data = sio.loadmat(fname)
        H = data['H'].astype(np.float32) / np.sqrt(n)

    return H

max_stars_repo_name: afshinabdi/DistributedTraining-QCS

max_issues_repo_name: afshinabdi/DistributedTraining-QCS

max_forks_repo_name: afshinabdi/DistributedTraining-QCS

reason: import numpy,import scipy


text:
""" Market Watch View """
__docformat__ = "numpy"

import argparse
from typing import List
import re
import requests
import pandas as pd
from bs4 import BeautifulSoup
import numpy as np
from gamestonk_terminal.helper_funcs import (
    check_positive,
    get_user_agent,
    clean_data_values_to_float,
    int_or_round_float,
    parse_known_args_and_warn,
)

# pylint: disable=too-many-branches

def sec_fillings(other_args: List[str], ticker: str):
    """Display SEC filings for a given stock ticker

    Parameters
    ----------
    other_args : List[str]
        argparse other args - ["-n", "10"]
    ticker : str
        Stock ticker
    """

    parser = argparse.ArgumentParser(
        add_help=False,
        prog="sec",
        description="""
            Prints SEC filings of the company. The following fields are expected: Filing Date,
            Document Date, Type, Category, Amended, and Link. [Source: Market Watch]
        """,
    )

    parser.add_argument(
        "-n",
        "--num",
        action="store",
        dest="n_num",
        type=check_positive,
        default=5,
        help="number of latest SEC filings.",
    )

    try:
        ns_parser = parse_known_args_and_warn(parser, other_args)
        if not ns_parser:
            return

        pd.set_option("display.max_colwidth", None)

        url_financials = f"https://www.marketwatch.com/investing/stock/{ticker}/financials/secfilings"

        text_soup_financials = BeautifulSoup(
            requests.get(url_financials, headers={"User-Agent": get_user_agent()}).text,
            "lxml",
        )

        # a_financials_header = list()
        df_financials = None
        b_ready_to_process_info = False
        soup_financials = text_soup_financials.findAll("tr", {"class": "table__row"})
        for financials_info in soup_financials:
            a_financials = financials_info.text.split("\n")

            # If header has been processed and dataframe created ready to populate the SEC information
            if b_ready_to_process_info:
                l_financials_info = [a_financials[2]]
                l_financials_info.extend(a_financials[5:-1])
                l_financials_info.append(financials_info.a["href"])
                # Append data values to financials
                df_financials.loc[len(df_financials.index)] = l_financials_info  # type: ignore

            if "Filing Date" in a_financials:
                l_financials_header = [a_financials[2]]
                l_financials_header.extend(a_financials[5:-1])
                l_financials_header.append("Link")

                df_financials = pd.DataFrame(columns=l_financials_header)
                df_financials.set_index("Filing Date")
                b_ready_to_process_info = True

        # Set Filing Date as index
        df_financials = df_financials.set_index("Filing Date")  # type: ignore

        print(df_financials.head(n=ns_parser.n_num).to_string())
        print("")

    except Exception as e:
        print(e)
        print("")
        return

def sean_seah_warnings(other_args: List[str], ticker: str):
    """Display Sean Seah warnings

    Parameters
    ----------
    other_args : List[str]
        argparse other args
    ticker : str
        Stock ticker
    """

    parser = argparse.ArgumentParser(
        add_help=False,
        prog="warnings",
        description="""
            Sean Seah warnings. Check: Consistent historical earnings per share;
            Consistently high return on equity; Consistently high return on assets; 5x Net
            Income > Long-Term Debt; and Interest coverage ratio more than 3. See
            https://www.drwealth.com/gone-fishing-with-buffett-by-sean-seah/comment-page-1/
            [Source: Market Watch]
        """,
    )

    parser.add_argument(
        "-i",
        "--info",
        action="store_true",
        default=False,
        dest="b_info",
        help="provide more information about Sean Seah warning rules.",
    )
    parser.add_argument(
        "-d",
        "--debug",
        action="store_true",
        default=False,
        dest="b_debug",
        help="print insights into warnings calculation.",
    )

    try:
        ns_parser = parse_known_args_and_warn(parser, other_args)
        if not ns_parser:
            return

        if ns_parser.b_info:
            filepath = "fundamental_analysis/info_sean_seah.txt"
            with open(filepath) as fp:
                line = fp.readline()
                while line:
                    print(f"{line.strip()}")
                    line = fp.readline()
                print("")

        # From INCOME STATEMENT, get: 'EPS (Basic)', 'Net Income', 'Interest Expense', 'EBITDA'
        url_financials = (
            f"https://www.marketwatch.com/investing/stock/{ticker}/financials/income"
        )
        text_soup_financials = BeautifulSoup(
            requests.get(url_financials, headers={"User-Agent": get_user_agent()}).text,
            "lxml",
        )

        # Define financials columns
        a_financials_header = list()
        for financials_header in text_soup_financials.findAll(
            "th", {"class": "overflow__heading"}
        ):
            a_financials_header.append(
                financials_header.text.strip("\n").split("\n")[0]
            )
        df_financials = pd.DataFrame(columns=a_financials_header[0:-1])

        # Add financials values
        soup_financials = text_soup_financials.findAll(
            lambda tag: tag.name == "tr" and tag.get("class") == ["table__row"]
        )
        soup_financials += text_soup_financials.findAll(
            "tr", {"class": "table__row is-highlighted"}
        )
        for financials_info in soup_financials:
            financials_row = financials_info.text.split("\n")
            if len(financials_row) > 5:
                for item in financials_row:
                    if bool(re.search(r"\d", item)):
                        a_financials_info = financials_info.text.split("\n")
                        l_financials = [a_financials_info[2]]
                        l_financials.extend(a_financials_info[5:-2])
                        # Append data values to financials
                        df_financials.loc[len(df_financials.index)] = l_financials
                        break

        l_fin = ["EPS (Basic)", "Net Income", "Interest Expense", "EBITDA"]

        if not all(elem in df_financials["Item"].values for elem in l_fin):
            print("The source doesn't contain all necessary financial data")
            print(url_financials, "\n")
            return

        # Set item name as index
        df_financials = df_financials.set_index("Item")

        df_sean_seah = df_financials.loc[l_fin]

        # From BALANCE SHEET, get: 'Liabilities & Shareholders\' Equity', 'Long-Term Debt'
        url_financials = f"https://www.marketwatch.com/investing/stock/{ticker}/financials/balance-sheet"
        text_soup_financials = BeautifulSoup(
            requests.get(url_financials, headers={"User-Agent": get_user_agent()}).text,
            "lxml",
        )

        # Define financials columns
        a_financials_header = list()
        for financials_header in text_soup_financials.findAll(
            "th", {"class": "overflow__heading"}
        ):
            a_financials_header.append(
                financials_header.text.strip("\n").split("\n")[0]
            )

        s_header_end_trend = "5-year trend"
        df_financials = pd.DataFrame(
            columns=a_financials_header[
                0 : a_financials_header.index(s_header_end_trend)
            ]
        )

        # Add financials values
        soup_financials = text_soup_financials.findAll(
            lambda tag: tag.name == "tr" and tag.get("class") == ["table__row"]
        )
        soup_financials += text_soup_financials.findAll(
            "tr", {"class": "table__row is-highlighted"}
        )
        for financials_info in soup_financials:
            financials_row = financials_info.text.split("\n")
            if len(financials_row) > 5:
                for item in financials_row:
                    if bool(re.search(r"\d", item)):
                        a_financials_info = financials_info.text.split("\n")
                        l_financials = [a_financials_info[2]]
                        l_financials.extend(a_financials_info[5:-2])
                        # Append data values to financials
                        df_financials.loc[len(df_financials.index)] = l_financials
                        break

        # Set item name as index
        df_financials = df_financials.set_index("Item")

        # Create dataframe to compute meaningful metrics from sean seah book
        df_sean_seah = df_sean_seah.append(
            df_financials.loc[
                [
                    "Total Shareholders' Equity",
                    "Liabilities & Shareholders' Equity",
                    "Long-Term Debt",
                ]
            ]
        )

        # Clean these metrics by parsing their values to float
        df_sean_seah = df_sean_seah.applymap(lambda x: clean_data_values_to_float(x))

        # Add additional necessary metrics
        series = (
            df_sean_seah.loc["Net Income"]
            / df_sean_seah.loc["Total Shareholders' Equity"]
        )
        series.name = "ROE"
        df_sean_seah = df_sean_seah.append(series)

        series = df_sean_seah.loc["EBITDA"] / df_sean_seah.loc["Interest Expense"]
        series.name = "Interest Coverage Ratio"
        df_sean_seah = df_sean_seah.append(series)

        series = (
            df_sean_seah.loc["Net Income"]
            / df_sean_seah.loc["Liabilities & Shareholders' Equity"]
        )
        series.name = "ROA"
        df_sean_seah = df_sean_seah.append(series)

        print(df_sean_seah.applymap(lambda x: int_or_round_float(x)).to_string())

        n_warnings = 0
        print("\nWARNINGS:")

        if np.any(df_sean_seah.loc["EPS (Basic)"].diff().dropna().values < 0):
            print("NO consistent historical earnings per share")
            n_warnings += 1
            if ns_parser.b_debug:
                sa_eps = np.array2string(
                    df_sean_seah.loc["EPS (Basic)"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   EPS: {sa_eps}")
                sa_growth = np.array2string(
                    df_sean_seah.loc["EPS (Basic)"].diff().dropna().values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   Growth: {sa_growth} < 0")

        if np.any(df_sean_seah.loc["ROE"].values < 0.15):
            print("NOT consistently high return on equity")
            n_warnings += 1
            if ns_parser.b_debug:
                sa_roe = np.array2string(
                    df_sean_seah.loc["ROE"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   ROE: {sa_roe} < 0.15")

        if np.any(df_sean_seah.loc["ROA"].values < 0.07):
            print("NOT consistently high return on assets")
            n_warnings += 1
            if ns_parser.b_debug:
                sa_roa = np.array2string(
                    df_sean_seah.loc["ROA"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   ROA: {sa_roa} < 0.07")

        if np.any(
            df_sean_seah.loc["Long-Term Debt"].values
            > 5 * df_sean_seah.loc["Net Income"].values
        ):
            print("5x Net Income < Long-Term Debt")
            n_warnings += 1
            if ns_parser.b_debug:
                sa_5_net_income = np.array2string(
                    5 * df_sean_seah.loc["Net Income"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   5x NET Income: {sa_5_net_income}")
                sa_long_term_debt = np.array2string(
                    df_sean_seah.loc["Long-Term Debt"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   lower than Long-Term Debt: {sa_long_term_debt}")

        if np.any(df_sean_seah.loc["Interest Coverage Ratio"].values < 3):
            print("Interest coverage ratio less than 3")
            n_warnings += 1
            if ns_parser.b_debug:
                sa_interest_coverage_ratio = np.array2string(
                    100 * df_sean_seah.loc["Interest Coverage Ratio"].values,
                    formatter={"float_kind": lambda x: int_or_round_float(x)},
                )
                print(f"   Interest Coverage Ratio: {sa_interest_coverage_ratio} < 3")

        if n_warnings == 0:
            print("None. Good stonk")

        print("")

    except Exception as e:
        print(e)
        print("")
        return


text:
from functools import partial

import hypothesis.extra.numpy as hnp
import hypothesis.strategies as st
import numpy as np
import pytest
from hypothesis import given

import mygrad as mg
from mygrad.math.arithmetic.ops import Positive, Square
from tests.custom_strategies import tensors
from tests.utils.checkers import expected_constant
from tests.utils.wrappers import adds_constant_arg

from ..wrappers.uber import backprop_test_factory, fwdprop_test_factory

@adds_constant_arg
def custom_pow(x, p):
    return x ** p

@adds_constant_arg
def in_place_custom_pow(x, p):
    out = +x
    out **= p
    return out

def any_scalar(*args, p):
    return st.sampled_from([int(p), float(p), np.array(p)])

@given(x=tensors(elements=st.floats(-10, 10)), p=st.sampled_from([2, 3]))
def test_special_pow_propagate_constant(x, p):
    y = x ** p
    assert y.constant is x.constant

@pytest.mark.parametrize("power, op", [(1, Positive), (2, Square)])
def test_pow_uses_special_case(power, op):
    @given(exp=st.sampled_from([int(power), float(power), np.array(power)]))
    def wrapped_func(exp):
        out = mg.arange(2) ** exp
        assert isinstance(out.creator, op)

    wrapped_func()

@fwdprop_test_factory(
    mygrad_func=custom_pow,
    true_func=custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=1)},
    permit_0d_array_as_float=False,
)
def test_pow_1_fwd():
    pass

@backprop_test_factory(
    mygrad_func=custom_pow,
    true_func=custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=1)},
)
def test_pow_1_bkwd():
    pass

@fwdprop_test_factory(
    mygrad_func=in_place_custom_pow,
    true_func=in_place_custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=1)},
    permit_0d_array_as_float=False,
)
def test_inplace_pow_1_fwd():
    pass

@backprop_test_factory(
    mygrad_func=in_place_custom_pow,
    true_func=in_place_custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=1)},
)
def test_inplace_pow_1_bkwd():
    pass

@fwdprop_test_factory(
    mygrad_func=custom_pow,
    true_func=custom_pow,
    num_arrays=1,
    permit_0d_array_as_float=False,
    kwargs={"p": partial(any_scalar, p=2)},
)
def test_pow_2_fwd():
    pass

@backprop_test_factory(
    mygrad_func=custom_pow,
    true_func=custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=2)},
)
def test_pow_2_bkwd():
    pass

@fwdprop_test_factory(
    mygrad_func=in_place_custom_pow,
    true_func=in_place_custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=2)},
    permit_0d_array_as_float=False,
)
def test_inplace_pow_2_fwd():
    pass

@backprop_test_factory(
    mygrad_func=in_place_custom_pow,
    true_func=in_place_custom_pow,
    num_arrays=1,
    kwargs={"p": partial(any_scalar, p=2)},
)
def test_inplace_pow_2_bkwd():
    pass


text:
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd

import nose.tools as nt
import numpy.testing as npt
import pandas.util.testing as pdt
from numpy.testing.decorators import skipif
from nose import SkipTest

try:
    import statsmodels.regression.linear_model as smlm
    _no_statsmodels = False
except ImportError:
    _no_statsmodels = True

from .. import linearmodels as lm
from .. import algorithms as algo
from .. import utils
from ..palettes import color_palette

rs = np.random.RandomState(0)

class TestLinearPlotter(object):

    rs = np.random.RandomState(77)
    df = pd.DataFrame(dict(x=rs.normal(size=60),
                           d=rs.randint(-2, 3, 60),
                           y=rs.gamma(4, size=60),
                           s=np.tile(list("abcdefghij"), 6)))
    df["z"] = df.y + rs.randn(60)
    df["y_na"] = df.y.copy()
    df.y_na.ix[[10, 20, 30]] = np.nan

    def test_establish_variables_from_frame(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y="y")
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_frame_equal(p.data, self.df)

    def test_establish_variables_from_series(self):

        p = lm._LinearPlotter()
        p.establish_variables(None, x=self.df.x, y=self.df.y)
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        nt.assert_is(p.data, None)

    def test_establish_variables_from_array(self):

        p = lm._LinearPlotter()
        p.establish_variables(None,
                              x=self.df.x.values,
                              y=self.df.y.values)
        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        nt.assert_is(p.data, None)

    def test_establish_variables_from_mix(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y=self.df.y)
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_frame_equal(p.data, self.df)

    def test_establish_variables_from_bad(self):

        p = lm._LinearPlotter()
        with nt.assert_raises(ValueError):
            p.establish_variables(None, x="x", y=self.df.y)

    def test_dropna(self):

        p = lm._LinearPlotter()
        p.establish_variables(self.df, x="x", y_na="y_na")
        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y_na, self.df.y_na)

        p.dropna("x", "y_na")
        mask = self.df.y_na.notnull()
        pdt.assert_series_equal(p.x, self.df.x[mask])
        pdt.assert_series_equal(p.y_na, self.df.y_na[mask])

class TestRegressionPlotter(object):

    rs = np.random.RandomState(49)

    grid = np.linspace(-3, 3, 30)
    n_boot = 100
    bins_numeric = 3
    bins_given = [-1, 0, 1]

    df = pd.DataFrame(dict(x=rs.normal(size=60),
                           d=rs.randint(-2, 3, 60),
                           y=rs.gamma(4, size=60),
                           s=np.tile(list(range(6)), 10)))
    df["z"] = df.y + rs.randn(60)
    df["y_na"] = df.y.copy()

    bw_err = rs.randn(6)[df.s.values] * 2
    df.y += bw_err

    p = 1 / (1 + np.exp(-(df.x * 2 + rs.randn(60))))
    df["c"] = [rs.binomial(1, p_i) for p_i in p]
    df.y_na.ix[[10, 20, 30]] = np.nan

    def test_variables_from_frame(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, units="s")

        pdt.assert_series_equal(p.x, self.df.x)
        pdt.assert_series_equal(p.y, self.df.y)
        pdt.assert_series_equal(p.units, self.df.s)
        pdt.assert_frame_equal(p.data, self.df)

    def test_variables_from_series(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y, units=self.df.s)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        npt.assert_array_equal(p.units, self.df.s)
        nt.assert_is(p.data, None)

    def test_variables_from_mix(self):

        p = lm._RegressionPlotter("x", self.df.y + 1, data=self.df)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y + 1)
        pdt.assert_frame_equal(p.data, self.df)

    def test_dropna(self):

        p = lm._RegressionPlotter("x", "y_na", data=self.df)
        nt.assert_equal(len(p.x), pd.notnull(self.df.y_na).sum())

        p = lm._RegressionPlotter("x", "y_na", data=self.df, dropna=False)
        nt.assert_equal(len(p.x), len(self.df.y_na))

    def test_ci(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, ci=95)
        nt.assert_equal(p.ci, 95)
        nt.assert_equal(p.x_ci, 95)

        p = lm._RegressionPlotter("x", "y", data=self.df, ci=95, x_ci=68)
        nt.assert_equal(p.ci, 95)
        nt.assert_equal(p.x_ci, 68)

    @skipif(_no_statsmodels)
    def test_fast_regression(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fit with the "fast" function, which just does linear algebra
        yhat_fast, _ = p.fit_fast(self.grid)

        # Fit using the statsmodels function with an OLS model
        yhat_smod, _ = p.fit_statsmodels(self.grid, smlm.OLS)

        # Compare the vector of y_hat values
        npt.assert_array_almost_equal(yhat_fast, yhat_smod)

    @skipif(_no_statsmodels)
    def test_regress_poly(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fit an first-order polynomial
        yhat_poly, _ = p.fit_poly(self.grid, 1)

        # Fit using the statsmodels function with an OLS model
        yhat_smod, _ = p.fit_statsmodels(self.grid, smlm.OLS)

        # Compare the vector of y_hat values
        npt.assert_array_almost_equal(yhat_poly, yhat_smod)

    def test_regress_logx(self):

        x = np.arange(1, 10)
        y = np.arange(1, 10)
        grid = np.linspace(1, 10, 100)
        p = lm._RegressionPlotter(x, y, n_boot=self.n_boot)

        yhat_lin, _ = p.fit_fast(grid)
        yhat_log, _ = p.fit_logx(grid)

        nt.assert_greater(yhat_lin[0], yhat_log[0])
        nt.assert_greater(yhat_log[20], yhat_lin[20])
        nt.assert_greater(yhat_lin[90], yhat_log[90])

    @skipif(_no_statsmodels)
    def test_regress_n_boot(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, n_boot=self.n_boot)

        # Fast (linear algebra) version
        _, boots_fast = p.fit_fast(self.grid)
        npt.assert_equal(boots_fast.shape, (self.n_boot, self.grid.size))

        # Slower (np.polyfit) version
        _, boots_poly = p.fit_poly(self.grid, 1)
        npt.assert_equal(boots_poly.shape, (self.n_boot, self.grid.size))

        # Slowest (statsmodels) version
        _, boots_smod = p.fit_statsmodels(self.grid, smlm.OLS)
        npt.assert_equal(boots_smod.shape, (self.n_boot, self.grid.size))

    @skipif(_no_statsmodels)
    def test_regress_without_bootstrap(self):

        p = lm._RegressionPlotter("x", "y", data=self.df,
                                  n_boot=self.n_boot, ci=None)

        # Fast (linear algebra) version
        _, boots_fast = p.fit_fast(self.grid)
        nt.assert_is(boots_fast, None)

        # Slower (np.polyfit) version
        _, boots_poly = p.fit_poly(self.grid, 1)
        nt.assert_is(boots_poly, None)

        # Slowest (statsmodels) version
        _, boots_smod = p.fit_statsmodels(self.grid, smlm.OLS)
        nt.assert_is(boots_smod, None)

    def test_numeric_bins(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_numeric)
        npt.assert_equal(len(bins), self.bins_numeric)
        npt.assert_array_equal(np.unique(x_binned), bins)

    def test_provided_bins(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_given)
        npt.assert_array_equal(np.unique(x_binned), self.bins_given)

    def test_bin_results(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x_binned, bins = p.bin_predictor(self.bins_given)
        nt.assert_greater(self.df.x[x_binned == 0].min(),
                          self.df.x[x_binned == -1].max())
        nt.assert_greater(self.df.x[x_binned == 1].min(),
                          self.df.x[x_binned == 0].max())

    def test_scatter_data(self):

        p = lm._RegressionPlotter(self.df.x, self.df.y)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.d)
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y, x_jitter=.1)
        x, y = p.scatter_data
        nt.assert_true((x != self.df.d).any())
        npt.assert_array_less(np.abs(self.df.d - x), np.repeat(.1, len(x)))
        npt.assert_array_equal(y, self.df.y)

        p = lm._RegressionPlotter(self.df.d, self.df.y, y_jitter=.05)
        x, y = p.scatter_data
        npt.assert_array_equal(x, self.df.d)
        npt.assert_array_less(np.abs(self.df.y - y), np.repeat(.1, len(y)))

    def test_estimate_data(self):

        p = lm._RegressionPlotter(self.df.d, self.df.y, x_estimator=np.mean)

        x, y, ci = p.estimate_data

        npt.assert_array_equal(x, np.sort(np.unique(self.df.d)))
        npt.assert_array_almost_equal(y, self.df.groupby("d").y.mean())
        npt.assert_array_less(np.array(ci)[:, 0], y)
        npt.assert_array_less(y, np.array(ci)[:, 1])

    def test_estimate_cis(self):

        # set known good seed to avoid the test stochastically failing
        np.random.seed(123)

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=95)
        _, _, ci_big = p.estimate_data

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=50)
        _, _, ci_wee = p.estimate_data
        npt.assert_array_less(np.diff(ci_wee), np.diff(ci_big))

        p = lm._RegressionPlotter(self.df.d, self.df.y,
                                  x_estimator=np.mean, ci=None)
        _, _, ci_nil = p.estimate_data
        npt.assert_array_equal(ci_nil, [None] * len(ci_nil))

    def test_estimate_units(self):

        # Seed the RNG locally
        np.random.seed(345)

        p = lm._RegressionPlotter("x", "y", data=self.df,
                                  units="s", x_bins=3)
        _, _, ci_big = p.estimate_data
        ci_big = np.diff(ci_big, axis=1)

        p = lm._RegressionPlotter("x", "y", data=self.df, x_bins=3)
        _, _, ci_wee = p.estimate_data
        ci_wee = np.diff(ci_wee, axis=1)

        npt.assert_array_less(ci_wee, ci_big)

    def test_partial(self):

        x = self.rs.randn(100)
        y = x + self.rs.randn(100)
        z = x + self.rs.randn(100)

        p = lm._RegressionPlotter(y, z)
        _, r_orig = np.corrcoef(p.x, p.y)[0]

        p = lm._RegressionPlotter(y, z, y_partial=x)
        _, r_semipartial = np.corrcoef(p.x, p.y)[0]
        nt.assert_less(r_semipartial, r_orig)

        p = lm._RegressionPlotter(y, z, x_partial=x, y_partial=x)
        _, r_partial = np.corrcoef(p.x, p.y)[0]
        nt.assert_less(r_partial, r_orig)

    @skipif(_no_statsmodels)
    def test_logistic_regression(self):

        p = lm._RegressionPlotter("x", "c", data=self.df,
                                  logistic=True, n_boot=self.n_boot)
        _, yhat, _ = p.fit_regression(x_range=(-3, 3))
        npt.assert_array_less(yhat, 1)
        npt.assert_array_less(0, yhat)

    @skipif(_no_statsmodels)
    def test_robust_regression(self):

        p_ols = lm._RegressionPlotter("x", "y", data=self.df,
                                      n_boot=self.n_boot)
        _, ols_yhat, _ = p_ols.fit_regression(x_range=(-3, 3))

        p_robust = lm._RegressionPlotter("x", "y", data=self.df,
                                         robust=True, n_boot=self.n_boot)
        _, robust_yhat, _ = p_robust.fit_regression(x_range=(-3, 3))

        nt.assert_equal(len(ols_yhat), len(robust_yhat))

    @skipif(_no_statsmodels)
    def test_lowess_regression(self):

        p = lm._RegressionPlotter("x", "y", data=self.df, lowess=True)
        grid, yhat, err_bands = p.fit_regression(x_range=(-3, 3))

        nt.assert_equal(len(grid), len(yhat))
        nt.assert_is(err_bands, None)

    def test_regression_options(self):

        with nt.assert_raises(ValueError):
            lm._RegressionPlotter("x", "y", data=self.df,
                                  lowess=True, order=2)

        with nt.assert_raises(ValueError):
            lm._RegressionPlotter("x", "y", data=self.df,
                                  lowess=True, logistic=True)

    def test_regression_limits(self):

        f, ax = plt.subplots()
        ax.scatter(self.df.x, self.df.y)
        p = lm._RegressionPlotter("x", "y", data=self.df)
        grid, _, _ = p.fit_regression(ax)
        xlim = ax.get_xlim()
        nt.assert_equal(grid.min(), xlim[0])
        nt.assert_equal(grid.max(), xlim[1])

        p = lm._RegressionPlotter("x", "y", data=self.df, truncate=True)
        grid, _, _ = p.fit_regression()
        nt.assert_equal(grid.min(), self.df.x.min())
        nt.assert_equal(grid.max(), self.df.x.max())

        plt.close("all")

class TestDiscretePlotter(object):

    rs = np.random.RandomState(341)
    df = pd.DataFrame(dict(x=np.repeat(list("abc"), 30),
                           y=rs.randn(90),
                           g=np.tile(list("xy"), 45),
                           u=np.tile(np.arange(6), 15)))
    bw_err = rs.randn(6)[df.u.values]
    df.y += bw_err
    df["y_na"] = df.y.copy()
    df.y_na.ix[[10, 20, 30]] = np.nan

    def test_variables_from_frame(self):

        p = lm._DiscretePlotter("x", "y", "g", data=self.df, units="u")

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        npt.assert_array_equal(p.hue, self.df.g)
        npt.assert_array_equal(p.units, self.df.u)
        pdt.assert_frame_equal(p.data, self.df)

    def test_variables_from_series(self):

        p = lm._DiscretePlotter(self.df.x, self.df.y, self.df.g,
                                units=self.df.u)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y)
        npt.assert_array_equal(p.hue, self.df.g)
        npt.assert_array_equal(p.units, self.df.u)
        nt.assert_is(p.data, None)

    def test_variables_from_mix(self):

        p = lm._DiscretePlotter("x", self.df.y + 1, data=self.df)

        npt.assert_array_equal(p.x, self.df.x)
        npt.assert_array_equal(p.y, self.df.y + 1)
        pdt.assert_frame_equal(p.data, self.df)

    def test_variables_var_order(self):

        p = lm._DiscretePlotter("x", "y", "g", data=self.df)

        npt.assert_array_equal(p.x_order, list("abc"))
        npt.assert_array_equal(p.hue_order, list("xy"))

        x_order = list("bca")
        hue_order = list("yx")
        p = lm._DiscretePlotter("x", "y", "g", data=self.df,
                                x_order=x_order, hue_order=hue_order)

        npt.assert_array_equal(p.x_order, x_order)
        npt.assert_array_equal(p.hue_order, hue_order)

    def test_count_x(self):

        p = lm._DiscretePlotter("x", hue="g", data=self.df)
        nt.assert_true(p.y_count)
        npt.assert_array_equal(p.x, p.y)
        nt.assert_is(p.estimator, len)

    def test_dropna(self):

        p = lm._DiscretePlotter("x", "y_na", data=self.df)
        nt.assert_equal(len(p.x), pd.notnull(self.df.y_na).sum())

        p = lm._DiscretePlotter("x", "y_na", data=self.df, dropna=False)
        nt.assert_equal(len(p.x), len(self.df.y_na))

    def test_palette(self):

        p = lm._DiscretePlotter("x", "y", data=self.df)
        nt.assert_equal(p.palette, [color_palette()[0]] * 3)

        p = lm._DiscretePlotter("x", "y", data=self.df, color="green")
        nt.assert_equal(p.palette, ["green"] * 3)

        p = lm._DiscretePlotter("x", "y", data=self.df, palette="husl")
        nt.assert_equal(p.palette, color_palette("husl", 3))
        nt.assert_true(p.x_palette)

        p = lm._DiscretePlotter("x", "y", "g", data=self.df)
        nt.assert_equal(p.palette, color_palette(n_colors=2))

        pal = {"x": "pink", "y": "green"}
        p = lm._DiscretePlotter("x", "y", "g", data=self.df, palette=pal)
        nt.assert_equal(p.palette, color_palette(["pink", "green"], 2))

        p = lm._DiscretePlotter("x", "y", "g", data=self.df,
                                palette=pal, hue_order=list("yx"))
        nt.assert_equal(p.palette, color_palette(["green", "pink"], 2))

    def test_markers(self):

        p = lm._DiscretePlotter("x", "y", hue="g", data=self.df)
        nt.assert_equal(p.markers, ["o", "o"])

        markers = ["o", "s"]
        p = lm._DiscretePlotter("x", "y", hue="g", data=self.df,
                                markers=markers)
        nt.assert_equal(p.markers, markers)

        with nt.assert_raises(ValueError):
            p = lm._DiscretePlotter("x", "y", hue="g", data=self.df,
                                    markers=["o", "s", "d"])

    def test_linestyles(self):

        p = lm._DiscretePlotter("x", "y", hue="g", data=self.df)
        nt.assert_equal(p.linestyles, ["-", "-"])

        linestyles = ["-", "--"]
        p = lm._DiscretePlotter("x", "y", hue="g", data=self.df,
                                linestyles=linestyles)
        nt.assert_equal(p.linestyles, linestyles)

        with nt.assert_raises(ValueError):
            p = lm._DiscretePlotter("x", "y", hue="g", data=self.df,
                                    linestyles=["-", "--", ":"])

    def test_plot_kind(self):

        p = lm._DiscretePlotter("x", "y", data=self.df, kind="bar")
        nt.assert_equal(p.kind, "bar")

        p = lm._DiscretePlotter("x", "y", data=self.df, kind="point")
        nt.assert_equal(p.kind, "point")

        p = lm._DiscretePlotter("x", data=self.df, kind="auto")
        nt.assert_equal(p.kind, "bar")

        p = lm._DiscretePlotter("x", np.ones(len(self.df)),
                                data=self.df, kind="auto")
        nt.assert_equal(p.kind, "point")

        with nt.assert_raises(ValueError):
            p = lm._DiscretePlotter("x", "y", data=self.df, kind="dino")

    def test_positions(self):

        p = lm._DiscretePlotter("x", "y", data=self.df, kind="bar")
        npt.assert_array_equal(p.positions, [0, 1, 2])
        npt.assert_array_equal(p.offset, [0])

        p = lm._DiscretePlotter("x", "y", "g", data=self.df, kind="bar")
        npt.assert_array_equal(p.positions, [0, 1, 2])
        npt.assert_array_equal(p.offset, [-.2, .2])

        p = lm._DiscretePlotter("x", "y", "g", data=self.df, kind="point")
        npt.assert_array_equal(p.positions, [0, 1, 2])
        npt.assert_array_equal(p.offset, [0, 0])

        p = lm._DiscretePlotter("x", "y", "g", data=self.df,
                                kind="point", dodge=.4)
        npt.assert_array_equal(p.positions, [0, 1, 2])
        npt.assert_array_equal(p.offset, [-.2, .2])

    def test_estimate_data(self):

        p = lm._DiscretePlotter("x", "y", data=self.df)
        nt.assert_equal(len(list(p.estimate_data)), 1)
        pos, height, ci = next(p.estimate_data)

        npt.assert_array_equal(pos, [0, 1, 2])

        height_want = self.df.groupby("x").y.mean()
        npt.assert_array_almost_equal(height, height_want)

        get_cis = lambda x: utils.ci(algo.bootstrap(x, random_seed=0), 95)
        ci_want = np.array(self.df.groupby("x").y.apply(get_cis).tolist())
        npt.assert_array_almost_equal(np.squeeze(ci), ci_want, 1)

        p = lm._DiscretePlotter("x", "y", "g", data=self.df)
        nt.assert_equal(len(list(p.estimate_data)), 2)
        data_gen = p.estimate_data

        first_hue = self.df[self.df.g == "x"]
        pos, height, ci = next(data_gen)

        npt.assert_array_equal(pos, [-.2, .8, 1.8])

        height_want = first_hue.groupby("x").y.mean()
        npt.assert_array_almost_equal(height, height_want)

        ci_want = np.array(first_hue.groupby("x").y.apply(get_cis).tolist())
        npt.assert_array_almost_equal(np.squeeze(ci), ci_want, 1)

        second_hue = self.df[self.df.g == "y"]
        pos, height, ci = next(data_gen)

        npt.assert_array_equal(pos, [.2, 1.2, 2.2])

        height_want = second_hue.groupby("x").y.mean()
        npt.assert_array_almost_equal(height, height_want)

        ci_want = np.array(second_hue.groupby("x").y.apply(get_cis).tolist())
        npt.assert_array_almost_equal(np.squeeze(ci), ci_want, 1)

    def test_plot_cis(self):

        p = lm._DiscretePlotter("x", "y", data=self.df, ci=95)
        _, _, ci_big = next(p.estimate_data)
        ci_big = np.diff(ci_big, axis=1)

        p = lm._DiscretePlotter("x", "y", data=self.df, ci=68)
        _, _, ci_wee = next(p.estimate_data)
        ci_wee = np.diff(ci_wee, axis=1)

        npt.assert_array_less(ci_wee, ci_big)

    def test_plot_units(self):

        p = lm._DiscretePlotter("x", "y", data=self.df, units="u")
        _, _, ci_big = next(p.estimate_data)
        ci_big = np.diff(ci_big, axis=1)

        p = lm._DiscretePlotter("x", "y", data=self.df)
        _, _, ci_wee = next(p.estimate_data)
        ci_wee = np.diff(ci_wee, axis=1)

        npt.assert_array_less(ci_wee, ci_big)

    def test_annotations(self):

        f, ax = plt.subplots()
        p = lm._DiscretePlotter("x", "y", "g", data=self.df)
        p.plot(ax)
        nt.assert_equal(ax.get_xlabel(), "x")
        nt.assert_equal(ax.get_ylabel(), "y")
        nt.assert_equal(ax.legend_.get_title().get_text(), "g")

        plt.close("all")

class TestDiscretePlots(object):

    rs = np.random.RandomState(341)
    df = pd.DataFrame(dict(x=np.repeat(list("abc"), 30),
                           v=rs.randn(90),
                           y=rs.randn(90) + 5,
                           z=rs.uniform(90),
                           g=np.tile(list("xy"), 45),
                           h=np.repeat(list("abc"), 30),
                           u=np.tile(np.arange(6), 15)))
    bw_err = rs.randn(6)[df.u.values]
    df.y += bw_err

    def test_barplot(self):

        f, ax = plt.subplots()
        lm.barplot("x", "y", data=self.df, hline=None, ax=ax)

        nt.assert_equal(len(ax.patches), 3)
        nt.assert_equal(len(ax.lines), 3)

        f, ax = plt.subplots()
        lm.barplot("x", "y", data=self.df, palette="husl",
                   hline=None, ax=ax)

        nt.assert_equal(len(ax.patches), 3)
        nt.assert_equal(len(ax.lines), 3)
        bar_colors = np.array([el.get_facecolor() for el in ax.patches])
        npt.assert_array_equal(color_palette("husl", 3), bar_colors[:, :3])

        plt.close("all")

    def test_bar_data(self):

        f, ax = plt.subplots()
        lm.barplot("x", "y", data=self.df, ax=ax)

        nt.assert_equal(len(ax.patches), 3)
        nt.assert_equal(len(ax.lines), 3)

        f, ax = plt.subplots()
        lm.barplot("x", "y", data=self.df, palette="husl", ax=ax)

        nt.assert_equal(len(ax.patches), 3)
        nt.assert_equal(len(ax.lines), 3)
        bar_colors = np.array([el.get_facecolor() for el in ax.patches])
        npt.assert_array_equal(color_palette("husl", 3), bar_colors[:, :3])

        plt.close("all")

    def test_pointplot(self):

        f, ax = plt.subplots()
        lm.pointplot("x", "y", data=self.df, join=False, ax=ax)

        nt.assert_equal(len(ax.collections), 1)
        nt.assert_equal(len(ax.lines), 3)

        f, ax = plt.subplots()
        lm.pointplot("x", "y", "g", data=self.df, palette="husl", ax=ax)

        nt.assert_equal(len(ax.collections), 2)
        nt.assert_equal(len(ax.lines), 8)
        point_colors = [c.get_facecolor()[:, :3] for c in ax.collections]
        expected_palette = color_palette("husl", 2)
        npt.assert_array_equal(np.squeeze(point_colors), expected_palette)

        plt.close("all")

    def test_point_data(self):

        f, ax = plt.subplots()
        lm.pointplot("x", "y", data=self.df, join=False, ax=ax)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, np.arange(3))
        npt.assert_array_almost_equal(y, self.df.groupby("x").y.mean())

        f, ax = plt.subplots()
        lm.pointplot("x", "y", "g", data=self.df, join=False, ax=ax)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, np.arange(3))
        expected_y = self.df[self.df.g == "x"].groupby("x").y.mean()
        npt.assert_array_almost_equal(y, expected_y)

        plt.close("all")

    def test_factorplot_bar(self):

        g = lm.factorplot("x", "y", data=self.df, kind="bar")
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.patches), 3)
        nt.assert_equal(len(ax.lines), 3)

        g = lm.factorplot("x", "y", "g", data=self.df, kind="bar")
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.patches), 6)
        nt.assert_equal(len(ax.lines), 6)

        plt.close("all")

    def test_factorplot_point(self):

        g = lm.factorplot("x", "y", data=self.df, kind="point")
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.collections), 1)
        nt.assert_equal(len(ax.lines), 4)

        g = lm.factorplot("x", "y", "g", data=self.df, kind="point")
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.collections), 2)
        nt.assert_equal(len(ax.lines), 8)

        plt.close("all")

    def test_factorplot_auto(self):

        g = lm.factorplot("x", "y", data=self.df)
        nt.assert_equal(len(g.axes[0, 0].collections), 1)

        g = lm.factorplot("x", data=self.df)
        nt.assert_equal(len(g.axes[0, 0].patches), 3)

        g = lm.factorplot("x", "v", data=self.df)
        nt.assert_equal(len(g.axes[0, 0].patches), 3)

        g = lm.factorplot("x", "z", data=self.df)
        nt.assert_equal(len(g.axes[0, 0].collections), 1)

        plt.close("all")

    def test_factorplot_facets(self):

        g = lm.factorplot("x", "y", data=self.df, row="g", col="h")
        nt.assert_equal(g.axes.shape, (2, 3))

        g = lm.factorplot("x", "y", data=self.df, col="u", col_wrap=4)
        nt.assert_equal(g.axes.shape, (6,))

        g = lm.factorplot("x", "y", "u", data=self.df, col="u")
        nt.assert_equal(g.axes.shape, (1, 6))
        nt.assert_is(g._legend, None)

        plt.close("all")

    def test_factorplot_missing(self):

        d = pd.DataFrame(dict(a=["a", "a", "b", "c", "c"],
                              b=["x", "y", "x", "x", "y"],
                              c=[1, 2, 3, 4, 5]))

        g = lm.factorplot("a", "c", data=d, col="b", kind="point")

        ax = g.axes[0, 0]
        x1, y1 = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x1, [0, 1, 2])
        npt.assert_array_equal(y1, [1, 3, 4])

        ax = g.axes[0, 1]
        x1, y1 = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x1, [0, 2])
        npt.assert_array_equal(y1, [2, 5])

    def test_factorplot_hline(self):

        g = lm.factorplot("x", "v", data=self.df, kind="bar", hline=0)
        ax = g.axes[0, 0]
        hline = ax.lines[-1]
        npt.assert_array_equal(hline.get_data(), [(0, 1), (0, 0)])

        plt.close("all")

class TestRegressionPlots(object):

    rs = np.random.RandomState(56)
    df = pd.DataFrame(dict(x=rs.randn(90),
                           y=rs.randn(90) + 5,
                           z=rs.randint(0, 1, 90),
                           g=np.repeat(list("abc"), 30),
                           h=np.tile(list("xy"), 45),
                           u=np.tile(np.arange(6), 15)))
    bw_err = rs.randn(6)[df.u.values]
    df.y += bw_err

    def test_regplot_basic(self):

        f, ax = plt.subplots()
        lm.regplot("x", "y", self.df)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 2)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

        plt.close("all")

    def test_regplot_selective(self):

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, scatter=False, ax=ax)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, fit_reg=False)
        nt.assert_equal(len(ax.lines), 0)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

        f, ax = plt.subplots()
        ax = lm.regplot("x", "y", self.df, ci=None)
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 1)
        ax.clear()

        plt.close("all")

    def test_regplot_scatter_kws_alpha(self):

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_is(ax.collections[0]._alpha, None)
        nt.assert_equal(ax.collections[0]._facecolors[0, 3], 0.5)

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_equal(ax.collections[0]._alpha, 0.8)

        f, ax = plt.subplots()
        color = np.array([[0.3, 0.8, 0.5]])
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color,
                                                        'alpha': 0.4})
        nt.assert_equal(ax.collections[0]._alpha, 0.4)

        f, ax = plt.subplots()
        color = 'r'
        ax = lm.regplot("x", "y", self.df, scatter_kws={'color': color})
        nt.assert_equal(ax.collections[0]._alpha, 0.8)

        plt.close("all")

    def test_regplot_binned(self):

        ax = lm.regplot("x", "y", self.df, x_bins=5)
        nt.assert_equal(len(ax.lines), 6)
        nt.assert_equal(len(ax.collections), 2)

        plt.close("all")

    def test_lmplot_basic(self):

        g = lm.lmplot("x", "y", self.df)
        ax = g.axes[0, 0]
        nt.assert_equal(len(ax.lines), 1)
        nt.assert_equal(len(ax.collections), 2)

        x, y = ax.collections[0].get_offsets().T
        npt.assert_array_equal(x, self.df.x)
        npt.assert_array_equal(y, self.df.y)

        plt.close("all")

    def test_lmplot_hue(self):

        g = lm.lmplot("x", "y", data=self.df, hue="h")
        ax = g.axes[0, 0]

        nt.assert_equal(len(ax.lines), 2)
        nt.assert_equal(len(ax.collections), 4)
        plt.close("all")

    def test_lmplot_markers(self):

        g1 = lm.lmplot("x", "y", data=self.df, hue="h", markers="s")
        nt.assert_equal(g1.hue_kws, {"marker": ["s", "s"]})

        g2 = lm.lmplot("x", "y", data=self.df, hue="h", markers=["o", "s"])
        nt.assert_equal(g2.hue_kws, {"marker": ["o", "s"]})

        with nt.assert_raises(ValueError):
            lm.lmplot("x", "y", data=self.df, hue="h", markers=["o", "s", "d"])

        plt.close("all")

    def test_lmplot_marker_linewidths(self):

        if mpl.__version__ == "1.4.2":
            raise SkipTest

        g = lm.lmplot("x", "y", data=self.df, hue="h",
                      fit_reg=False, markers=["o", "+"])
        c = g.axes[0, 0].collections
        nt.assert_equal(c[0].get_linewidths()[0], 0)
        rclw = mpl.rcParams["lines.linewidth"]
        nt.assert_equal(c[1].get_linewidths()[0], rclw)
        plt.close("all")

    def test_lmplot_facets(self):

        g = lm.lmplot("x", "y", data=self.df, row="g", col="h")
        nt.assert_equal(g.axes.shape, (3, 2))

        g = lm.lmplot("x", "y", data=self.df, col="u", col_wrap=4)
        nt.assert_equal(g.axes.shape, (6,))

        g = lm.lmplot("x", "y", data=self.df, hue="h", col="u")
        nt.assert_equal(g.axes.shape, (1, 6))
        plt.close("all")

    def test_lmplot_hue_col_nolegend(self):

        g = lm.lmplot("x", "y", data=self.df, col="h", hue="h")
        nt.assert_is(g._legend, None)
        plt.close("all")

    def test_lmplot_scatter_kws(self):

        g = lm.lmplot("x", "y", hue="h", data=self.df, ci=None)
        red_scatter, blue_scatter = g.axes[0, 0].collections

        red, blue = color_palette(n_colors=2)
        npt.assert_array_equal(red, red_scatter.get_facecolors()[0, :3])
        npt.assert_array_equal(blue, blue_scatter.get_facecolors()[0, :3])

        plt.close("all")

    def test_residplot(self):

        x, y = self.df.x, self.df.y
        ax = lm.residplot(x, y)

        resid = y - np.polyval(np.polyfit(x, y, 1), x)
        x_plot, y_plot = ax.collections[0].get_offsets().T

        npt.assert_array_equal(x, x_plot)
        npt.assert_array_almost_equal(resid, y_plot)
        plt.close("all")

    @skipif(_no_statsmodels)
    def test_residplot_lowess(self):

        ax = lm.residplot("x", "y", self.df, lowess=True)
        nt.assert_equal(len(ax.lines), 2)

        x, y = ax.lines[1].get_xydata().T
        npt.assert_array_equal(x, np.sort(self.df.x))

        plt.close("all")

max_stars_repo_name: kjemmett/seaborn

max_issues_repo_name: kjemmett/seaborn

max_forks_repo_name: kjemmett/seaborn

reason: import numpy,from numpy,import statsmodels


text:
"""Some utils for celebA dataset"""

import png
import dcgan_utils
import numpy as np
import utils

def display_transform(image):
    image = dcgan_utils.inverse_transform(image)
    return image

def view_image(image, hparams, mask=None):
    """Process and show the image"""
    image = display_transform(image)
    if len(image) == hparams.n_input:
        image = image.reshape(hparams.image_shape)
        if mask is not None:
            mask = mask.reshape(hparams.image_shape)
            image = np.maximum(np.minimum(1.0, image - 1.0*image*(1-mask)), -1.0)
    utils.plot_image(image)

def save_image(image, path):
    """Save an image as a png file"""
    image = dcgan_utils.inverse_transform(image)
    png_writer = png.Writer(64, 64)
    with open(path, 'wb') as outfile:
        png_writer.write(outfile, 255*image.reshape([64,-1]))


text:
import numpy as np
import sys
from taylor_approximation import compute_derivatives
import json

z = float(sys.argv[1])
rmin, rmax, dr = 50, 160, 0.5

# Remake the data grid:
order = 4
Npoints = 2*order + 1
# these are OmegaM, h, sigma8
x0s = [0.31, 0.68, 0.73]; Nparams = len(x0s)
dxs = [0.01, 0.01, 0.05]

rr = np.arange(rmin, rmax, dr)
output_shape = (len(rr),6)

center_ii = (order,)*Nparams
X0grid = np.zeros( (Npoints,)*Nparams+ output_shape)
X2grid = np.zeros( (Npoints,)*Nparams+ output_shape)

# Load data
for ii in range(Npoints):
    for jj in range(Npoints):
        for kk in range(Npoints):
            #print(ii,jj,kk)
            X0grid[ii,jj,kk] = np.loadtxt('data/boss_z_%.2f/boss_xi0_%d_%d_%d.txt'%(z,ii,jj,kk))
            X2grid[ii,jj,kk] = np.loadtxt('data/boss_z_%.2f/boss_xi2_%d_%d_%d.txt'%(z,ii,jj,kk))

# Now compute the derivatives
derivs0 = compute_derivatives(X0grid, dxs, center_ii, 5)
derivs2 = compute_derivatives(X2grid, dxs, center_ii, 5)

# Now save:
outfile = 'emu/boss_z_%.2f_xiells.json'%(z)

list0 = [ dd.tolist() for dd in derivs0 ]
list2 = [ dd.tolist() for dd in derivs2 ]

outdict = {'params': ['omegam', 'h', 'sigma8'],\
           'x0': x0s,\
           'rvec': rr.tolist(),\
           'derivs0': list0,\
           'derivs2': list2,}

json_file = open(outfile, 'w')
json.dump(outdict, json_file)
json_file.close()

max_stars_repo_name: LBJ-Wade/CobayaLSS

max_issues_repo_name: LBJ-Wade/CobayaLSS

max_forks_repo_name: LBJ-Wade/CobayaLSS