text:
"""
Tests the server collection compute capabilities.
"""

import numpy as np
import pytest

import qcfractal.interface as ptl
from qcfractal import testing
from qcfractal.testing import fractal_compute_server, mark_slow

def test_collection_query(fractal_compute_server):
    client = ptl.FractalClient(fractal_compute_server)

    ds = ptl.collections.Dataset("CAPITAL", client)
    ds.save()

    cols = client.list_collections()
    assert cols.index.contains(("Dataset", "CAPITAL"))

    ds = client.get_collection("dataset", "capital")
    assert ds.name == "CAPITAL"

    ds = client.get_collection("DATAset", "CAPital")
    assert ds.name == "CAPITAL"

@testing.using_psi4
def test_dataset_compute_gradient(fractal_compute_server):
    client = ptl.FractalClient(fractal_compute_server)

    # Build a dataset
    ds = ptl.collections.Dataset("ds_gradient",
                                 client,
                                 default_program="psi4",
                                 default_driver="gradient",
                                 default_units="hartree")

    ds.add_entry("He1", ptl.Molecule.from_data("He -1 0 0\n--\nHe 0 0 1"))
    ds.add_entry("He2", ptl.Molecule.from_data("He -1.1 0 0\n--\nHe 0 0 1.1"))

    contrib = {
        "name": "Gradient",
        "theory_level": "pseudo-random values",
        "values": {
            "He1": [0.03, 0, 0.02, -0.02, 0, -0.03],
            "He2": [0.03, 0, 0.02, -0.02, 0, -0.03]
        },
        "units": "hartree"
    }
    ds.add_contributed_values(contrib)
    ds.save()

    ds = client.get_collection("dataset", "ds_gradient")

    # Compute
    ds.compute("HF", "sto-3g")
    fractal_compute_server.await_results()

    ds.query("HF", "sto-3g", as_array=True)

    # Test out some statistics
    stats = ds.statistics("MUE", "HF/sto-3g", "Gradient")
    assert pytest.approx(stats.mean(), 1.e-5) == 0.00984176986312362

    stats = ds.statistics("UE", "HF/sto-3g", "Gradient")
    assert pytest.approx(stats.loc["He1"].mean(), 1.e-5) == 0.01635020639
    assert pytest.approx(stats.loc["He2"].mean(), 1.e-5) == 0.00333333333

    assert ds.list_history().shape[0] == 1
    assert ds.get_history().shape[0] == 1

def test_reactiondataset_check_state(fractal_compute_server):
    client = ptl.FractalClient(fractal_compute_server)
    ds = ptl.collections.ReactionDataset("check_state", client, ds_type="ie", default_program="rdkit")
    ds.add_ie_rxn("He1", ptl.Molecule.from_data("He -3 0 0\n--\nHe 0 0 2"))

    with pytest.raises(ValueError):
        ds.compute("SCF", "STO-3G")

    with pytest.raises(ValueError):
        ds.query("SCF", "STO-3G")

    ds.save()
    assert ds.query("SCF", "STO-3G")

    ds.add_keywords("default", "psi4", ptl.models.KeywordSet(values={"a": 5}))

    with pytest.raises(ValueError):
        ds.query("SCF", "STO-3G")

    ds.save()
    assert ds.query("SCF", "STO-3G")

    contrib = {
        "name": "Benchmark",
        "doi": None,
        "theory_level": "very high",
        "values": {
            "He1": 0.0009608501557,
            "He2": -0.00001098794749
        },
        "units": "hartree"
    }
    ds.add_contributed_values(contrib)
    with pytest.raises(ValueError):
        ds.query("SCF", "STO-3G")

    assert "benchmark" in ds.list_contributed_values()
    assert ds.get_contributed_values("benchmark").name == "Benchmark"

@testing.using_psi4
@testing.using_dftd3
def test_compute_reactiondataset_dftd3(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)
    ds_name = "He_DFTD3"
    ds = ptl.collections.ReactionDataset(ds_name, client, ds_type="ie")

    # Add two helium dimers to the DB at 4 and 8 bohr
    HeDimer = ptl.Molecule.from_data([[2, 0, 0, -4.123], [2, 0, 0, 4.123]], dtype="numpy", units="bohr", frags=[1])
    ds.add_ie_rxn("HeDimer", HeDimer, attributes={"r": 4})
    ds.set_default_program("psi4")

    ds.save()

    ncomp1 = ds.compute("B3LYP-D3", "6-31G")
    assert len(ncomp1.ids) == 4
    assert len(ncomp1.submitted) == 4

    ncomp2 = ds.compute("B3LYP-D3(BJ)", "6-31G")
    assert len(ncomp2.ids) == 4
    assert len(ncomp2.submitted) == 2

    fractal_compute_server.await_results()
    assert ds.query("B3LYP", "6-31G")
    assert ds.query("B3LYP-D3", "6-31G")
    assert ds.query("B3LYP-D3(BJ)", "6-31G")

    for key, value in {"B3LYP/6-31g": -0.002135, "B3LYP-D3/6-31g": -0.005818, "B3LYP-D3(BJ)/6-31g": -0.005636}.items():

        assert pytest.approx(value, 1.e-3) == ds.df.loc["HeDimer", key]

@testing.using_psi4
def test_compute_reactiondataset_regression(fractal_compute_server):
    """
    Tests an entire server and interaction energy dataset run
    """

    client = ptl.FractalClient(fractal_compute_server)
    ds_name = "He_PES"
    ds = ptl.collections.ReactionDataset(ds_name, client, ds_type="ie")

    # Add two helium dimers to the DB at 4 and 8 bohr
    He1 = ptl.Molecule.from_data([[2, 0, 0, -2], [2, 0, 0, 2]], dtype="numpy", units="bohr", frags=[1])
    ds.add_ie_rxn("He1", He1, attributes={"r": 4})

    # Save the DB and re-acquire via classmethod
    r = ds.save()
    ds = ptl.collections.ReactionDataset.from_server(client, ds_name)
    ds.set_default_program("psi4")
    assert "ReactionDataset(" in str(ds)

    # Test collection lists
    ret = client.list_collections(aslist=True)
    assert ds_name in ret["ReactionDataset"]

    ret = client.list_collections("reactiondataset", aslist=True)
    assert ds_name in ret

    He2 = ptl.Molecule.from_data([[2, 0, 0, -4], [2, 0, 0, 4]], dtype="numpy", units="bohr", frags=[1])
    ds.add_ie_rxn("He2", He2, attributes={"r": 4})

    contrib = {
        "name": "Benchmark",
        "doi": None,
        "theory_level": "very high",
        "values": {
            "He1": 0.0009608501557,
            "He2": -0.00001098794749
        },
        "units": "hartree"
    }
    ds.add_contributed_values(contrib)
    ds.data.default_benchmark = "Benchmark"

    # Save the DB and overwrite the result, reacquire via client
    r = ds.save()
    ds = client.get_collection("reactiondataset", ds_name)

    with pytest.raises(KeyError):
        ret = ds.compute("SCF", "STO-3G", stoich="nocp")  # Should be 'default' not 'nocp'

    # Compute SCF/sto-3g
    ret = ds.compute("SCF", "STO-3G")
    assert len(ret.submitted) == 3
    fractal_compute_server.await_results()

    # Query computed results
    assert ds.query("SCF", "STO-3G")
    assert pytest.approx(0.6024530476, 1.e-5) == ds.df.loc["He1", "SCF/sto-3g"]
    assert pytest.approx(-0.0068950359, 1.e-5) == ds.df.loc["He2", "SCF/sto-3g"]

    # Check results
    assert pytest.approx(0.00024477933196125805, 1.e-5) == ds.statistics("MUE", "SCF/sto-3g")

    assert pytest.approx([0.081193, 7.9533e-05], 1.e-1) == list(ds.statistics("URE", "SCF/sto-3g"))
    assert pytest.approx(0.0406367, 1.e-5) == ds.statistics("MURE", "SCF/sto-3g")
    assert pytest.approx(0.002447793, 1.e-5) == ds.statistics("MURE", "SCF/sto-3g", floor=10)

    assert isinstance(ds.to_json(), dict)
    assert ds.list_history(keywords=None).shape[0] == 1

    ds.units = "eV"
    assert pytest.approx(0.00010614635, 1.e-5) == ds.statistics("MURE", "SCF/sto-3g", floor=10)

@testing.using_psi4
def test_compute_reactiondataset_keywords(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)

    mol1 = ptl.Molecule.from_data("He 0 0 -1.1\n--\nHe 0 0 1.1")

    # Build a dataset
    ds = ptl.collections.ReactionDataset("dataset_options", client, ds_type="ie")
    ds.set_default_program("Psi4")

    ds.add_ie_rxn("He2", mol1)
    ds.add_keywords("direct", "psi4", ptl.models.KeywordSet(values={"scf_type": "direct"}), default=True)
    ds.add_keywords("df", "psi4", ptl.models.KeywordSet(values={"scf_type": "df"}))

    ds.save()
    ds = client.get_collection("reactiondataset", "dataset_options")

    # Compute, should default to direct options
    r = ds.compute("SCF", "STO-3G")
    fractal_compute_server.await_results()
    assert ds.query("SCF", "STO-3G")
    assert pytest.approx(0.39323818102293856, 1.e-5) == ds.df.loc["He2", "SCF/sto-3g"]

    r = ds.compute("SCF", "sto-3g", keywords="df")
    fractal_compute_server.await_results()
    assert ds.query("SCF", "sto-3g", keywords="df") == "SCF/sto-3g-df"
    assert pytest.approx(0.38748602675524185, 1.e-5) == ds.df.loc["He2", "SCF/sto-3g-df"]

    assert ds.list_history().shape[0] == 2
    assert ds.list_history(keywords="DF").shape[0] == 1
    assert ds.list_history(keywords="DIRECT").shape[0] == 1

    # Check saved history
    ds = client.get_collection("reactiondataset", "dataset_options")
    assert ds.list_history().shape[0] == 2
    assert {"df", "direct"} == set(ds.list_history().reset_index()["keywords"])

    # Check keywords
    kw = ds.get_keywords("df", "psi4")
    assert kw.values["scf_type"] == "df"

@mark_slow
@testing.using_torsiondrive
@testing.using_geometric
@testing.using_rdkit
def test_compute_openffworkflow(fractal_compute_server):
    """
    Tests the openffworkflow collection
    """

    # Obtain a client and build a BioFragment
    client = ptl.FractalClient(fractal_compute_server)

    openff_workflow_options = {
        # Blank Fragmenter options
        "enumerate_states": {},
        "enumerate_fragments": {},
        "torsiondrive_input": {},

        # TorsionDriveRecord, Geometric, and QC options
        "torsiondrive_static_options": {
            "keywords": {},
            "optimization_spec": {
                "program": "geometric",
                "keywords": {
                    "coordsys": "tric",
                }
            },
            "qc_spec": {
                "driver": "gradient",
                "method": "UFF",
                "basis": "",
                "keywords": None,
                "program": "rdkit",
            }
        },
        "optimization_static_options": {
            "program": "geometric",
            "keywords": {
                "coordsys": "tric",
            },
            "qc_spec": {
                "driver": "gradient",
                "method": "UFF",
                "basis": "",
                "keywords": None,
                "program": "rdkit",
            }
        }
    }
    wf = ptl.collections.OpenFFWorkflow("Workflow1", client=client, **openff_workflow_options)

    # # Add a fragment and wait for the compute
    hooh = ptl.data.get_molecule("hooh.json")
    fragment_input = {
        "label1": {
            "type": "torsiondrive_input",
            "initial_molecule": hooh.json_dict(),
            "grid_spacing": [90],
            "dihedrals": [[0, 1, 2, 3]],
        },
    }
    wf.add_fragment("HOOH", fragment_input)
    assert set(wf.list_fragments()) == {"HOOH"}
    fractal_compute_server.await_services(max_iter=5)

    final_energies = wf.list_final_energies()
    assert final_energies.keys() == {"HOOH"}
    assert final_energies["HOOH"].keys() == {"label1"}

    final_molecules = wf.list_final_molecules()
    assert final_molecules.keys() == {"HOOH"}
    assert final_molecules["HOOH"].keys() == {"label1"}

    optimization_input = {
        "label2": {
            "type": "optimization_input",
            "initial_molecule": hooh.json_dict(),
            "constraints": {
                'set': [{
                    "type": 'dihedral',
                    "indices": [0, 1, 2, 3],
                    "value": 0
                }]
            }
        }
    }

    wf.add_fragment("HOOH", optimization_input)
    fractal_compute_server.await_services(max_iter=5)

    final_energies = wf.list_final_energies()
    assert final_energies["HOOH"].keys() == {"label1", "label2"}
    assert pytest.approx(0.00259754, 1.e-4) == final_energies["HOOH"]["label2"]

    final_molecules = wf.list_final_molecules()
    assert final_molecules.keys() == {"HOOH"}
    assert final_molecules["HOOH"].keys() == {"label1", "label2"}

    # Add a second fragment
    butane = ptl.data.get_molecule("butane.json")
    butane_id = butane.identifiers.canonical_isomeric_explicit_hydrogen_mapped_smiles

    fragment_input = {
        "label1": {
            "type": "torsiondrive_input",
            "initial_molecule": butane.json_dict(),
            "grid_spacing": [90],
            "dihedrals": [[0, 2, 3, 1]],
        },
    }
    wf.add_fragment(butane_id, fragment_input)
    assert set(wf.list_fragments()) == {butane_id, "HOOH"}

    final_energies = wf.list_final_energies()
    assert final_energies.keys() == {butane_id, "HOOH"}
    assert final_energies[butane_id].keys() == {"label1"}
    assert final_energies[butane_id]["label1"] == {}

def test_generic_collection(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)
    g = ptl.collections.Generic("Generic1", client=client)

    # Check getters/savers
    g["hello"] = 5
    assert g["hello"] == 5

    # Copy the data over
    cp = g.get_data()
    assert cp.data["hello"] == 5
    cp.data["hello"] = 10

    # Make sure we did not change underlying data.
    assert g["hello"] == 5

    # Save semantics slightly wonky, double check this works
    g.save()
    g2 = ptl.collections.Generic.from_server(client, "Generic1")
    assert g2["hello"] == 5

def test_missing_collection(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)
    with pytest.raises(KeyError):
        client.get_collection("reactiondataset", "_waffles_")

@mark_slow
@testing.using_torsiondrive
@testing.using_geometric
@testing.using_rdkit
def test_torsiondrive_dataset(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)

    ds = ptl.collections.TorsionDriveDataset("testing", client=client)

    hooh1 = ptl.data.get_molecule("hooh.json")
    hooh2 = hooh1.copy(update={"geometry": hooh1.geometry + np.array([0, 0, 0.2])})

    ds.add_entry("hooh1", [hooh1], [[0, 1, 2, 3]], [90], attributes={"something": "hooh1"})
    ds.add_entry("hooh2", [hooh2], [[0, 1, 2, 3]], [90], attributes={"something": "hooh2"})

    optimization_spec = {
        "program": "geometric",
        "keywords": {
            "coordsys": "tric",
        }
    }
    qc_spec = {
        "driver": "gradient",
        "method": "UFF",
        "basis": "",
        "keywords": None,
        "program": "rdkit",
    }

    ds.add_specification("Spec1", optimization_spec, qc_spec, description="This is a really cool spec")

    ncompute = ds.compute("spec1")
    assert ncompute == 2
    assert ds.status("spec1")["Spec1"].sum() == 2  # Might have completed from previous run.

    ds.save()

    fractal_compute_server.await_services(max_iter=1)

    # Check status
    status_detail = ds.status("Spec1", detail=True)
    assert status_detail.loc["hooh2", "Complete"] == 1
    assert status_detail.loc["hooh2", "Total Points"] == 4

    fractal_compute_server.await_services(max_iter=5)

    ds = client.get_collection("torsiondrivedataset", "testing")
    ds.query("spec1")

    # Add another fake set, should instantly return
    ds.add_specification("spec2", optimization_spec, qc_spec, description="This is a really cool spec")

    # Test subsets
    ncompute = ds.compute("spec2", subset=set())
    assert ncompute == 0

    ncompute = ds.compute("spec2")
    assert ncompute == 2
    ds.query("spec2")

    # We effectively computed the same thing twice with two duplicate specs
    for row in ["hooh1", "hooh2"]:
        for spec in ["Spec1", "spec2"]:
            assert pytest.approx(ds.df.loc[row, spec].get_final_energies(90), 1.e-5) == 0.00015655375994799847

    assert ds.status().loc["COMPLETE", "Spec1"] == 2
    assert ds.status(collapse=False).loc["hooh1", "Spec1"] == "COMPLETE"

    assert ds.counts("hooh1").loc["hooh1", "Spec1"] > 5
    assert ds.counts("hooh1", specs="spec1", count_gradients=True).loc["hooh1", "Spec1"] > 30

@testing.using_geometric
@testing.using_rdkit
def test_optimization_dataset(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)

    ds = ptl.collections.OptimizationDataset("testing", client=client)

    opt_spec = {"program": "geometric"}
    qc_spec = {"driver": "gradient", "method": "UFF", "program": "rdkit"}
    ds.add_specification("test", opt_spec, qc_spec)

    hooh1 = ptl.data.get_molecule("hooh.json")
    hooh2 = hooh1.copy(update={"geometry": hooh1.geometry + np.array([0, 0, 0.2])})

    ds.add_entry("hooh1", hooh1)
    ds.add_entry("hooh1-2", hooh1)
    ds.add_entry("hooh2", hooh2)

    ds.compute("test")
    fractal_compute_server.await_results()

    ds.query("test")
    assert ds.status().loc["COMPLETE", "test"] == 3

    assert ds.counts().loc["hooh1", "test"] >= 4

    final_energy = 0.00011456853977485626
    for idx, row in ds.df["test"].items():
        assert pytest.approx(row.get_final_energy(), abs=1.e-5) == final_energy

    opt = ds.get_record("hooh1", "test")
    assert pytest.approx(opt.get_final_energy(), abs=1.e-5) == final_energy

@testing.using_geometric
@testing.using_rdkit
def test_grid_optimization_dataset(fractal_compute_server):

    client = ptl.FractalClient(fractal_compute_server)

    ds = ptl.collections.GridOptimizationDataset("testing", client=client)

    opt_spec = {"program": "geometric", "keywords": {}}
    qc_spec = {"driver": "gradient", "method": "UFF", "program": "rdkit"}
    ds.add_specification("test", opt_spec, qc_spec)

    hooh1 = ptl.data.get_molecule("hooh.json")

    scans = [{"type": "dihedral", "indices": [0, 1, 2, 3], "steps": [-10, 10], "step_type": "relative"}]
    ds.add_entry("hooh1", hooh1, scans=scans, preoptimization=False)

    ds.compute("test")
    fractal_compute_server.await_services()

    ds.query("test")
    assert ds.get_record("hooh1", "test").status == "COMPLETE"

max_stars_repo_name: yudongqiu/QCFractal

max_issues_repo_name: yudongqiu/QCFractal

max_forks_repo_name: yudongqiu/QCFractal


text:
from typing import Tuple

import numpy as np

def get_dataset() -> Tuple[np.ndarray, np.ndarray]:
    """OR dataset."""
    x = np.array([[0, 0], [1, 0], [0, 1], [1, 1]])
    y = np.array([[0], [1], [1], [1]])
    return x, y

def to_categorical(y: np.ndarray, num_classes: int) -> np.ndarray:
    y_categorical = np.zeros(shape=(len(y), num_classes))
    for i, yi in enumerate(y):
        y_categorical[i, yi] = 1
    return y_categorical

def softmax(y_pred: np.ndarray) -> np.ndarray:
    probabilities = np.zeros_like(y_pred)
    for i in range(len(y_pred)):
        exps = np.exp(y_pred[i])
        probabilities[i] = exps / np.sum(exps)
    return probabilities

def cross_entropy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    num_samples = y_true.shape[0]
    loss = float(-np.sum(y_true * np.log(y_pred)) / num_samples)
    return loss

if __name__ == "__main__":
    x, y = get_dataset()
    print(y)

    y_true = to_categorical(y, num_classes=2)
    print(y_true)

    y_logits = np.array([[10.8, -3.3], [12.2, 21.8], [1.1, 4.9], [7.05, 3.95]])

    y_pred = softmax(y_logits)
    print(y_pred)

    loss = cross_entropy(y_true, y_pred)
    print(loss)

ext: py

reason: import numpy


text:
"""Project: PhiK - correlation analyzer library

Created: 2018/12/28

Description:
    A set of functions to check for data quality issues in input data.

Authors:
    KPMG Advanced Analytics & Big Data team, Amstelveen, The Netherlands

Redistribution and use in source and binary forms, with or without
modification, are permitted according to the terms listed in the file
LICENSE.
"""

import warnings
import copy
from typing import Tuple

import pandas as pd
import numpy as np

def dq_check_nunique_values(df: pd.DataFrame, interval_cols: list, dropna:bool=True) -> Tuple[pd.DataFrame, list]:
    """
    Basic data quality checks per column in a DataFrame.

    The following checks are done:

    1. For all non-interval variables, if the number of unique values per variable is larger than 100 a warning is printed.
    When the number of unique values is large, the variable is likely to be an interval variable. Calculation of phik
    will be slow(ish) for pairs of variables where one (or two) have many different values (i.e. many bins).

    2. For all interval variables, the number of unique values must be at least two. If the number of unique values is
    zero (i.e. all NaN) the column is removed. If the number of unique values is one, it is not possible to
    automatically create a binning for this variable (as min and max are the same). The variable is therefore dropped,
    irrespective of whether dropna is True or False.

    3. For all non-interval variables, the number of unique values must be at least either
    a) 1 if dropna=False (NaN is now also considered a valid category), or
    b) 2 if dropna=True

    The function returns a DataFrame where all columns with invalid data are removed. Also the list of interval_cols
    is updated and returned.

    :param pd.DataFrame df: input data
    :param list interval_cols: column names of columns with interval variables.
    :param bool dropna: remove NaN values when True
    :returns: cleaned data, updated list of interval columns
    """
    # check for existing columns
    interval_cols = [col for col in interval_cols if col in df.columns]

    # check non-interval variable for number of unique values
    for col in sorted(list(set(df.columns)-set(interval_cols))):
        if df[col].nunique() > 1000:
            warnings.warn(
                'The number of unique values of variable {0:s} is large: {1:d}. Are you sure this is '
                'not an interval variable? Analysis for pairs of variables including {0:s} can be slow.'
                .format(col, df[col].nunique())
            )

    drop_cols = []

    # check for interval values whether there are at least two unique values (otherwise I cannot bin automatically)
    for col in interval_cols:
        if df[col].nunique() < 2:
            drop_cols.append(col)
            warnings.warn(
                'Not enough unique value for variable {0:s} for analysis {1:d}. Dropping this column'
                .format(col, df[col].nunique())
            )

    # check non-interval values whether there are at least two different values OR 1 value and NaN if dropna==False
    for col in sorted(list(set(df.columns) - set(interval_cols))):
        if df[col].nunique() == 0 or (df[col].nunique() == 1 and dropna):
            drop_cols.append(col)
            warnings.warn(
                'Not enough unique value for variable {0:s} for analysis {1:d}. Dropping this column'
                .format(col, df[col].nunique())
            )

    df_clean = df.copy()
    interval_cols_clean = copy.copy(interval_cols)
    if len(drop_cols) > 0:
        # preserves column order: https://github.com/KaveIO/PhiK/issues/1
        df_clean.drop(columns=drop_cols, inplace=True)
        interval_cols_clean = [col for col in interval_cols if col not in drop_cols]

    return df_clean, interval_cols_clean

def dq_check_hist2d(hist2d: np.ndarray) -> bool:
    """Basic data quality checks for a contingency table

    The Following checks are done:

    1. There must be at least two bins in both the x and y direction.

    2. If the number of bins in the x and/or y direction is larger than 100 a warning is printed.

    :param hist2d: contingency table
    :return: bool passed_check
    """

    if 0 in hist2d.shape or 1 in hist2d.shape:
        warnings.warn(
            'Too few unique values for variable x ({0:d}) or y ({1:d})'.format(hist2d.shape[0], hist2d.shape[1])
        )
        return False
    if hist2d.shape[0] > 1000:
        warnings.warn(
            'The number of unique values of variable x is large: {0:d}. '
            'Are you sure this is not an interval variable? Analysis might be slow.'
            .format(hist2d.shape[0])
        )
    if hist2d.shape[1] > 1000:
        warnings.warn(
            'The number of unique values of variable y is large: {0:d}. '
            'Are you sure this is not an interval variable? Analysis might be slow.'
            .format(hist2d.shape[0])
        )

    return True


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    Functions for calling estimator predictions in batches
"""

__author__ = "John Hoff"
__email__ = "john.hoff@braindonor.net"
__copyright__ = "Copyright 2019, John Hoff"
__license__ = "Creative Commons Attribution-ShareAlike 4.0 International License"
__version__ = "1.0"

import numpy as np

def batch_fit_classifier(estimator, x, y, transformer=None, increment=500, verbose=True):
    if len(x) > increment:
        step_size = len(x) // increment
    else:
        step_size = 1
    steps = 0
    x_chunks = np.array_split(x, step_size)
    y_chunks = np.array_split(y, step_size)
    classes = np.unique(y)
    for i in range(len(x_chunks)):
        if verbose:
            print('.', end='')
            if i > 0 and i % 100 == 0:
                print(i*increment)
        x_chunk = x_chunks[i]
        y_chunk = y_chunks[i]
        if transformer is not None:
            x_chunk = transformer.transform(x_chunk)
        estimator.partial_fit(x_chunk, y_chunk, classes=classes)
    if verbose:
        print('')

max_stars_repo_name: theBraindonor/chicago-crime-arrests

max_issues_repo_name: theBraindonor/chicago-crime-arrests

max_forks_repo_name: theBraindonor/chicago-crime-arrests


text:
from dopamine.gym import make_env
from keras.models import load_model
import numpy as np
from mathtools.utils import Vessel
from ipdb import set_trace as debug
from gym import wrappers
import pylab as plt

if __name__ == '__main__':

    # Create an environment.
    env = make_env('wiggle')

    # Create a policy model.
    name = 'wiggles'
    location = '../weights/policy_{:s}.h5'.format(name)
    policy = load_model(location)
    obsfilt = Vessel('../weights/filters_{:s}.dat'.format(name))

    done = False
    obs = env.reset()
    while not done:
        obs = obsfilt.filt(obs)
        action = policy.predict(np.atleast_2d(obs))
        obs, reward, done, info = env.step(action[0])

        if done:
            print('REWARD: {:2f}'.format(reward))
            break

    plt.ion()
    plt.close('all')
    plt.figure(100)
    plt.plot(env.T, env.X)
    plt.ylim([0.5, 2.5])
    plt.xlabel('Time (seconds)')
    plt.ylabel('Mass Position (meters)')
    plt.savefig('rl-snap.png')

    plt.figure(200)
    plt.plot(env.T, env.A)
    # plt.ylim([0.5, 2.5])
    plt.xlabel('Time (seconds)')
    plt.ylabel('Mass Position (meters)')
    plt.savefig('control-signal.png')

max_stars_repo_name: lightscalar/dopamine

max_issues_repo_name: lightscalar/dopamine

max_forks_repo_name: lightscalar/dopamine


text:
#################################################################################
# Copyright (c) 2018-2021, Texas Instruments Incorporated - http://www.ti.com
# All Rights Reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#################################################################################

import os
import numpy as np
import torch.utils
import cv2

def split2list(images, split):
    if isinstance(split, str):
        with open(split) as f:
            split_values = [x.strip() == '1' for x in f.readlines()]
        assert(len(images) == len(split_values))
    elif isinstance(split, float):
        split_values = np.random.uniform(0,1,len(images)) < split
    else:
        assert False, 'split could not be understood'
    #
    train_samples = [sample for sample, sval in zip(images, split_values) if sval]
    test_samples = [sample for sample, sval in zip(images, split_values) if not sval]
    return train_samples, test_samples

def load_flo(path):
    with open(path, 'rb') as f:
        magic = np.fromfile(f, np.float32, count=1)
        assert(202021.25 == magic),'Magic number incorrect. Invalid .flo file'
        h = np.fromfile(f, np.int32, count=1)[0]
        w = np.fromfile(f, np.int32, count=1)[0]
        data = np.fromfile(f, np.float32, count=2*w*h)
    # Reshape data into 3D array (columns, rows, bands)
    data2D = np.resize(data, (w, h, 2))
    return data2D

def default_loader(root, path_imgs, path_flows):
    imgs = [os.path.join(root,path) for path in path_imgs]
    flows = [os.path.join(root,path_flo) for path_flo in path_flows]
    imgs = [cv2.imread(img)[:,:,::-1] for img in imgs]
    imgs = [img.astype(np.float32) for img in imgs]
    flows = [load_flo(flo) for flo in flows]
    return imgs,flows

class ListDataset(torch.utils.data.Dataset):
    def __init__(self, root, path_list, transform=None, loader=default_loader):
        self.root = root
        self.path_list = path_list
        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        inputs, targets = self.path_list[index]
        inputs, targets = self.loader(self.root, inputs, targets)
        if self.transform is not None:
            inputs, targets = self.transform(inputs, targets)
        #
        return inputs, targets

    def __len__(self):
        return len(self.path_list)

class ListDatasetWithAdditionalInfo(torch.utils.data.Dataset):
    def __init__(self, root, path_list, transform=None, loader=default_loader):
        self.root = root
        self.path_list = path_list
        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        inputs, targets = self.path_list[index]
        inputs, targets, input_paths, target_paths = self.loader(self.root, inputs, targets, additional_info=True)
        if self.transform is not None:
            inputs, targets = self.transform(inputs, targets)
        return inputs, targets, input_paths, target_paths

    def __len__(self):
        return len(self.path_list)

max_stars_repo_name: TexasInstruments/vision


text:
import os
import argparse
import shutil
import tempfile
import csv
import pickle
import pydevd_pycharm

from typing import List

from supervised.automl import AutoML

import pandas as pd
import numpy as np

from azureml.core import Run

parser = argparse.ArgumentParser()
parser.add_argument("--input_name", required=True)
parser.add_argument("--model_directory", required=True)
parser.add_argument("--binary_model_directory", required=True)
parser.add_argument("--regression_model_directory", required=True)
parser.add_argument("--output", required=True)
args = parser.parse_args()

# # Attach PyCharm debugger
# print("Starting debugger...")
# pydevd_pycharm.settrace(
#     os.environ.get("PYCHARM_DEBUG_HOST"),
#     port=int(os.environ.get("PYCHARM_DEBUG_PORT")),
#     stdoutToServer=False,
#     stderrToServer=False)
# print('PyCharm debugger attached to prepare.py.')
#

current_directory = os.getcwd()
temporary_prefix = f"{current_directory}/"
temporary_path = tempfile.mkdtemp(prefix=temporary_prefix)

input_dataset_filepaths = Run.get_context() \
    .input_datasets[args.input_name] \
    .download(target_path=temporary_path)

if not isinstance(input_dataset_filepaths, list):
    input_dataset_filepaths = [input_dataset_filepaths]

model_directory_filepaths = Run.get_context() \
    .input_datasets[args.model_directory] \
    .download(target_path=temporary_path)

def get_report_dir(
    model_directory_filepaths: List[str],
    model_specific_directory: str
) -> str:

    model_specific_filepath = [
        f for f in model_directory_filepaths
        if model_specific_directory in f
    ][0]

    model_directory_end_idx = \
        model_specific_filepath.find(model_specific_directory) \
        + len(model_specific_directory)

    return model_specific_filepath[:model_directory_end_idx]

binary_model_directory = get_report_dir(
    model_directory_filepaths, args.binary_model_directory)

regression_model_directory = get_report_dir(
    model_directory_filepaths, args.regression_model_directory)

print(f"binary_model_directory: {binary_model_directory}")
print(f"regression_model_directory: {regression_model_directory}")

assert isinstance(binary_model_directory, str)
assert isinstance(regression_model_directory, str)

os.makedirs(args.output, exist_ok=True)

# Load models
automl_binary = AutoML(results_path=binary_model_directory)
automl_regression = AutoML(results_path=regression_model_directory)

# Load transformer
transformer_filepath_list = [f for f in model_directory_filepaths if f.split(".")[-1] == "pkl"]
if transformer_filepath_list:
    assert len(transformer_filepath_list) == 1
    transformer_filepath = transformer_filepath_list[0]
    transformer = pickle.load(open(transformer_filepath, "rb"))

# Predict
res_binary = []
res_regression = []
res_ground_truth_binary = []
res_ground_truth_regression = []

# Dataset filepath is csv (there can be other stuff like w2v model)
input_dataset_filepaths = [
    f for f in input_dataset_filepaths
    if f.split(".")[-1] == "csv"
]

for filepath in input_dataset_filepaths:

    X = pd.read_csv(filepath)

    if 'label_0' in X.columns and 'label_1' in X.columns:

        # Ground truth
        # res_ground_truth_binary.append(X.label_0.values)
        # res_ground_truth_regression.append(X.label_1.values)

        X.drop(columns=['label_0', 'label_1'], inplace=True)

    predictions_binary = automl_binary.predict_proba(X)[:, 1]

    if transformer_filepath_list:
        predictions_regression = transformer.inverse_transform(
            automl_regression.predict(X).reshape(-1, 1)
        ).flatten()
    else:
        predictions_regression = np.exp(
            automl_regression.predict(X)
        ) * 1.7

    res_binary.append(predictions_binary)
    res_regression.append(predictions_regression)

# Flatten to 1 array per model
submission_array_binary = np.array(res_binary).flatten()
submission_array_regression = np.array(res_regression).flatten()

binary_threshold = 0.485
submission_array_regression[submission_array_binary < binary_threshold] = 0.0

# if 'label_0' in X_all.columns and 'label_1' in X_all.columns:
#     ground_truth_array_binary = np.array(res_ground_truth_binary).flatten()
#     ground_truth_array_regression = np.array(res_ground_truth_regression).flatten()

# Save submission to csv
submission_filepath = os.path.join(args.output, "submission.csv")
submission = csv.writer(open(submission_filepath, 'w', encoding='UTF8'))
for x, y in zip(submission_array_binary, submission_array_regression):
    submission.writerow([x, y])

print(f"Submission saved to '{submission_filepath}'")

# # Save ground truth as well
# if 'label_0' in X_all.columns and 'label_1' in X_all.columns:
#     try:
#         ground_truth_filepath = os.path.join(args.output, "ground_truth.csv")
#         df_ground_truth = pd.DataFrame({
#             'label_0': ground_truth_array_binary,
#             'label_1': ground_truth_array_regression
#         })
#
#         df_ground_truth.to_csv(ground_truth_filepath, index=False)
#         print(f"Ground truth saved to '{ground_truth_filepath}'")
#     except:
#         pass

if os.path.exists(temporary_path):
    shutil.rmtree(temporary_path)

max_stars_repo_name: mananabot/alphamev-winning-submission

max_issues_repo_name: mananabot/alphamev-winning-submission

max_forks_repo_name: mananabot/alphamev-winning-submission


text:
# -*- coding: utf-8 -*-
"""
Modules to support data reduction in Python.

The main purpose of the base module ``Data_Reduction`` is to provide a
suplerclass with a good set of attributes and methods to cover all common needs.

The base module is also able to read data from a text file as a ``numpy``
structured array.  This is done with a class called ``DataGetterMixin`` which
must be invoked after the base class has been initiated.

The module function ``examine_text_data_file()`` reveals the structure of the
file(s) that provide the data..

Examples
========

Here we initiate a base class after mixing in the data getter. The first line o
the file has column names but the first three columns are all under one
name ``UTC`` so we specify column widths to consider the first three columns
to be one column.  We use the names from the first line of the file, which
could have been done with an ``open()``, ``readline()``, and ``close()``::

  mixIn(Observation, DataGetterMixin)
  obs = Observation(dss=28, date="2012/127", project="SolarPatrol")
  obs.open_datafile('t12127.10',
                          delimiter=[17,16,3,11,7,9,8,2,6],
                          skip_header=1,
                      names="UTC Epoch Chan Tsys Int Az El Diode Level".split())

Now the data getter is already mixed in to Observation so we don't need to do
it again. In this case we specify the names of the columns, changing ``Int`` to
``Integr``::

  obs2 = Observation(dss=28, date="2012/127", project="SolarPatrol")
  obs2.open_datafile('t12127.10', skip_header=1,
          names="Year DOY UTC Epoch Chan Tsys Integr Az El Diode Level".split())

The class Map inherits from DataGetterMixin, so no explicit mixin required::

  obsmap = Map(dss=84, date="2020/163", project="SolarPatrol")
  obsmap.initialize('sim-venus.dat', source="Venus")

Let's examine ``obsmap``.  We have only one signal column::

  In [3]: obsmap.channel.keys()
  Out[3]: dict_keys(['xl'])
  In [4]: obsmap.channel['xl'].keys()
  Out[4]: dict_keys(['freq', 'bw', 'pol', 'ifmode', 'atten', 'power'])
"""
# standard Python modules
import datetime
import glob
import h5py
import logging
import math
import matplotlib.dates as MPLd
import numpy as NP
import os
import re
import readline
import scipy.interpolate
import scipy.fftpack

import Astronomy as A
import Astronomy.DSN_coordinates as coords
import Astronomy.Ephem as AE
import DatesTimes as DT
import local_dirs
import Math.clusters as VQ # vector quantization
import support

# enable raw_input Tab completion
readline.parse_and_bind("tab: complete")

logger = logging.getLogger(__name__) # module logger

class Observation(object):
  """
  superclass for a data structure and methods

  Attributes
  ==========
    aliases      - (dict) data keys to replace those in original data
    channel      - (dict) signal paths, e.g., different freqs and pols
    data         - (dict) original data, e.g., read from file or database

    DOY          - (int) day of year of observation
    end          - (float) UNIX time at the end
    latitude     - (float) from obs
    logger       - (logging.Logger)
    longitude    - (float) from obs
    name         - (str) user assigned, defaults to YEAR/DOY
    numdata      - (int) number of data samples
    obs          - (AE.DSS) observatory
    session      - (Session) set of observations, parent to Observation
    session_path - (str) directory for session files
    start        - (float) UNIX time at the beginning
    year         - (int) year of observation

  **Reserved Column Names**

  These column names are recognized.  They are also the keys for attribute
  ``data``.

  These quantities must be present in some form::

    unixtime   (float) UNIX time in sec
    chan_name  (str)   channel name
    integr     (float) integration (exposure) in sec
    azel       (float,float) azimuth and elevation in decimal deg
    power      (float) power level if only a single channel

  Optional::

    diode      (float) 0 or power in K (integers OK)
    level      (float) (unidentified -- in ``tlog`` table)
    cryotemp   (float) cryostat temp in K
    windspeed  (float) km/hr
    winddir    (float) deg
    ambtemp    (float) deg C
    pressure   (float) mbar

  Columns to be computed::

    mpldatenum (float) matplotlib ``datenum``

  Alternative for ``power``::

    tsys       (float) system temperature (calibrated power)
    top        (float) alternative for ``tsys`` (used in DSN)
    vfc_counts (int)   VFC counts (rate times ``integr``)

  Any column with a name which is not a reserved name is assumed to be
  power-like data from the channel with that name, unless that name is in a
  list provided to the argument ``ignore`` in the method ``get_data_channels``
  of the class ``DataGetterMixin``.

  Alternative for ``unixtime``::

    year       (int) year of observation
    doy        (int) day of year
    utc        (str) HH:MM:SS
    timestr    (str) something like 2020/06/14/14:22:21.00

  Alternative for ``chan_name``::

    chan       (int) index in receiver channel names

  Alternative for ``azel``::

    radec      (float,float) precessed right ascension in decimal hours and
               precessed declination in decimal deg
    radec1950  (float,float) mean right ascension in decimal hours and
               mean declination in decimal deg at epoch
    radec2000  (float,float) mean right ascension in decimal hours and
               mean declination at epoch in decimal deg
    az         (float) azimuth in decimal deg
    el         (float) elevation in decimal deg
    ra         (float) precessed right ascension in decimal hours
    dec        (float) precessed declination in decimal deg
    ra1950     (float) mean right ascension in decimal hours at epoch
    dec1950    (float) mean declination in decimal deg at epoch
    ra2000     (float) mean right ascension in decimal hours at epoch
    dec2000    (float) mean declination in decimal deg at epoch

  Notes
  =====
  * The ``data`` structure is a dict.
  * The value of a ``data`` item is either a numpy array or a object
    like ``float``, ``int``, or ``str``.
  * The keys have reserved words defined above and will be lowercase.
  * Items with other keys may be added, typically by a child class.
  * Coordinates shall be in pairs, `e.g. ``azel``, ``radec``. (This way you
    never get one without the other.)
  """
  reserved = ['unixtime','chan_name','integr','az','el','year','doy','utc',
              'timestr','chan','tsys','top','diode','level','cryotemp',
              'windspeed','winddir','ambtemp','pressure',
              'ra','dec','ra1950','dec1950','ra2000','dec2000']
  power_keys = ['tsys', 'top', 'vfc_counts', 'power']

  def __init__(self, parent=None, name=None, dss=None,
                     date=None, project=None):
    """
    Create a base Observation object.

    This is not meant to be initialized by itself.  A subclass generally
    determines how data are read in.  However, method ``initialize()``
    provides a basic data read capability using ``numpy.genfromtxt()``
    and creates the object's data structure.

    Args:
      parent (Session):   session to which this observation belongs
      name (str):         an identifier; default is station ID + "obs"
      dss (int):          station number
      date (str):         "YEAR/DOY"
      project (str):      directory under /usr/local/projects
    """
    self.logger = logging.getLogger(logger.name+".Observation")
    self.session = parent
    # observatory must be specified
    if dss:
      self.obs = coords.DSS(dss)
      self.longitude = self.obs.long*180/math.pi # deg
      self.latitude = self.obs.lat*180/math.pi   # deg
    else:
      self.logger.error("__init__: requires observatory location")
      raise Exception("Where were the data taken?")
    # give the object a name
    if name:
      self.name = name
    else:
      self.name = "DSS"+str(dss)+"obs"
    self.logger = logging.getLogger(logger.name+".Observation")
    # the observation was part of some project
    if project:
      self.project = project
    else:
      self.logger.error("__init__: requires a project")
      raise Exception("Where are the session's working files?")
    # the observation was done on some date
    if date:
      y,d = date.split('/')
      self.year = int(y);
      self.DOY = int(d)
      projdatapath, self.sessionpath, rawdatapath = \
                              get_obs_dirs(project, dss, self.year, self.DOY,
                                           datafmt=None)
      self.logger.debug("__init__: session path: %s", self.sessionpath)
    else:
      self.logger.error("__init__: requires a date")
      raise Exception("When were the date taken?")
    # accomodate subclass arguments
    self.aliases = {}

    # what I really want to do here is see if this was called by a subclass,
    # in which case I do not try to get the channel info until this
    # initialization has finished.
    #
    #if hasattr(self, "get_data_channels"):
    #  channels = self, get_data_channels()
    #  self.make_channels(channels)
    #else:
    #  self.logger.info("__init__: initialize() may now be called")

  def splitkey(self, longlat):
    """
    Checks for presence of coordinates in pairs or singles

    @param longlat : "azel", or "radec", or "radecEPOC"
    @type  longlat : str
    """
    longitude = longlat[:2] # 'az' or 'ra'
    if len(longlat) > 5: # has epoch
        epoch = longlat[-4:]
        longitude += epoch
        latitude = longlat[2:-4]+epoch
    else: # date of observation
        latitude = longlat[2:]
        epoch = None
    return longitude, latitude, epoch

  def check_for(self, data, longlat):
    """
    Checks for separate coordinates and splits if coord pairs

    Args:
      data (dict):    attribute ``data``
      longlat (str): "azel", or "radec", or "radecEPOC"
    """
    longitude, latitude, epoch = self.splitkey(longlat)
    if longitude in data.dtype.names and \
       latitude  in data.dtype.names:
      self.logger.debug("check_for: data has %s and %s", longitude, latitude)
      self.data[longitude] = data[longitude]
      self.data[latitude]  = data[latitude]
      return True
    elif longlat in data.dtype.names:
      self.logger.debug("check_for: data has %s", longlat)
      self.data[longitude],self.data[latitude] = map(None, *data[longlat])
      self.logger.debug("check_for: added %s and %s to data",
                        longitude, latitude)
      return True
    else:
      # coords need to be computed from other coords

      return False

  def unpack_to_complex(self, rawdata):
    """
    Converts a sequence of alternating real/imag samples to complex

    @param rawdata : alternating real and imaginary bytes
    @type  rawdata : numpy array of signed int8

    @return: numpy array of complex
    """
    datalen = len(rawdata)
    real = rawdata[0:datalen:2]
    imag = rawdata[1:datalen:2]
    data = real + 1j*imag
    return data

  def sideband_separate(self, data):
    """
    Converts a complex spectrum array and returns two reals with USB and LSB

    This applies a Hilbert transform to the complex data.
    """
    usb = (data.real + scipy.fftpack.hilbert(data).imag)
    lsb = (scipy.fftpack.hilbert(data).real + data.imag)
    return lsb,usb

  class Channel(support.PropertiedClass):
    """
    Class for a signal path
    """
    def __init__(self, parent, name, freq=None, bw=None, pol=None, IFtype=None,
                       atten=None):
      """
      Notes
      =====
      The properties can be accessed as if the class were a dict.

      Arguments
      =========
      freq:float or int: center frequency in MHz
      bw:float or int:   bandwidth in MHz
      pol:str:           polarization code
      """
      support.PropertiedClass.__init__(self)
      self.parent = parent
      self.logger = logging.getLogger(self.parent.name+".Channel")
      self.logger.debug("__init__: created %s", self.logger.name)
      self.logger.debug("__init__: parent is %s", self.parent)
      self.name = name
      self.data['freq'] = freq
      self.data['bw'] = bw
      self.data['pol'] = pol
      self.data['ifmode'] = IFtype
      self.data['atten'] = atten

class DataGetterMixin(object):
  """
  Class for getting data from a CSV file.
  """
  def initialize(self, filename, delimiter=" ", names=True, skip_header=0,
                       source=None):
    """
    Get the data and make a data structure for the observations.

    This is not included by default in ``__init__()`` to keep it simple for
    subclasses.

    Args:
      filename (str):     name only, required; the path is provided
      delimiter (str):    what separates the columns
      names (bool):       the first line has column names
      skip_header (int) : number of rows to skip
    """
    # get the data
    data = self.open_datafile(filename, delimiter=delimiter, names=names,
                              skip_header=skip_header)
    # get the signal columns and names
    metadata, signals = self.get_data_channels(data)
    # create Channel objects for the signal properties
    self.make_channels(signals)
    # create the data structure
    self.make_data_struct(data, metadata, signals)
    # compute the offsets from the source center for each data point
    if source:
      self.get_offsets(source=source)
    else:
      self.logger.warning("initialize: no source specified; no offsets")

  def open_datafile(self, filename, delimiter=" ", names=True, skip_header=0):
    """
    Opens and reads a data file

    This is used by ``Malargue`` (one data files) and ``GAVRT`` (one data file
    for each signal).

    Args:
      filename (str):    text data file name
      delimiter (str):   separator between columns (default: whitespace)
      names (bool):      file row has column names (default: True)
      skip_header (int): number of rows to skip at beginning of file

    Returns:
      ndarray:
    """
    data = NP.genfromtxt(self.sessionpath+filename,
                         delimiter=delimiter,
                         dtype=None,
                         names=names,
                         case_sensitive='lower',
                         skip_header=skip_header,
                         encoding=None)
    return data

  def get_data_channels(self, data, ignore=None):
    """
    Gets or sets the names of the signal columns

    Column names are separated into metadata and signals. Names in
    ``ignore`` re ignored. Names in ``aliases`` are replaced.
    Args:
      data (ndarray): data read from text file
      ignore (list of str): columns to ignore; default None

    Returns:
      (list of str, list of str): metadata, signals
    """
    names = data.dtype.names
    metadata = []
    signals = []
    for name in names:
      if ignore:
        if name in ignore:
          pass
      if name.casefold() in map(str.casefold, self.aliases):
        key = self.aliases[name].lower() # we use only lower case names
      else:
        key = name.lower()
      self.logger.debug("get_data_channels: doing %s for %s", key, name)
      if key in map(str.casefold, Observation.reserved):
        if key.casefold() in ['top', 'tsys']:
          signals.append(key)
        else:
          metadata.append(key)
      else:
        signals.append(key)
    self.logger.debug("get_data_channels: signals: %s", signals)
    self.logger.debug("get_data_channels: metadata: %s", metadata)
    return metadata, signals

  def make_data_struct(self, data, metadata, signals):
    """
    Takes a text table with headers and converts it into a numpy ``ndarray``.

    That means that a column can be extracted using `data[label]`.

    Args
    ====
      data:     (ndarray) the data from the text file
      metadata: (list of str) the column names for metadata
      signals:  (list of str) the column names for power-like data
    """
    # get the known columns:
    self.data = {}
    self.numdata = len(data)
    #self.logger.debug("make_data_struct: using aliases: %s", self.aliases)
    # get columns that are not metadata; each has power for a channel
    for signal in signals:
      #self.logger.debug("make_data_struct: for signal: %s", signal)
      #if signal in self.aliases.items():
        # get the key in 'data' which matches 'value' in 'aliases'
      #  power = data[next(key for key, value in self.aliases.items()
      #                                                if value == signal)][idx]
      #else:
      #  power = data[signal]
      #self.channel[signal]['power'] = power
      self.channel[signal]['power'] = data[signal]
    # get UNIX time
    if 'unixtime' in metadata:
      if 'unixtime' in data.dtype.names:
        self.data['unixtime'] = data['unixtime']
      else:
        # look up the equivalent of UNIX time in the data table
        self.data['unixtime'] = data[next(key
                                          for key, value in self.aliases.items()
                                          if value == 'unixtime')]
      # compute other convenient forms of time
      self.data['datetime'] = [] # Python datetime.date
      self.data['date_num'] = [] # matplotlib.dates date number
      for idx in list(range(self.numdata)):
        if 'unixtime' in data.dtype.names:
          tm = data['unixtime'][idx]
        else:
          tm = data[next(key for key, value in self.aliases.items()
                                                   if value == 'unixtime')][idx]
        dt = datetime.datetime.utcfromtimestamp(tm)
        self.data['datetime'].append(dt)
        self.data['date_num'].append(MPLd.date2num(dt))
      self.start = self.data['unixtime'][0]
      self.end = self.data['unixtime'][-1]
    else:
      # figure out how to process the time data columns
      pass
    # compute alternate coordinates
    if self.check_for(data, 'azel'):
      # azel exists; compute radec if needed; then radec2000 if needed
      if self.check_for(data, 'radec'):
        pass
      else:
        self.radec_from_azel()
        if self.check_for(data, 'radec2000'):
          # ra2000 and dec2000 already exist
          pass
        else:
          self.radec2000_from_radec()
    elif self.check_for(data, 'radec2000'):
      # coordinates exist; compute back to azimuth and elevation
      if self.check_for(data, 'radec'):
        pass
      else:
        # compute observed RA and dec
        self.radec_from_radec2000()
        if self.check_for(data, 'azel'):
          pass
        else:
          self.azel_from_radec()
    # in here check for 'radec'
    else:
      self.logger.error("no coordinates found in data")
      raise Exception("check INFO logging for columns found")
    self.start = self.data['unixtime'].min()
    self.end   = self.data['unixtime'].max()

  def make_channels(self, signals, props=None):
    """
    Assign properties to the channels.

    The prop keys are "freq", "pol", and "IFtype".

    Args:
      props (dict of dicts): signal channel properties.
    """
    self.channel = {}
    for ch in signals:
      chindex = signals.index(ch)
      if props:
        self.channel[ch] = self.Channel(self, ch,
                                        freq  =props[ch]['freq'],
                                        bw    =props[ch]['bw'],
                                        pol   =props[ch]['pol'],
                                        IFtype=props[ch]['IFtype'],
                                        atten =props[ch]['atten'])
      else:
        self.channel[ch] = self.Channel(self, ch)

class GriddingMixin(object):
  """
  Class for all the data and methods associated with a raster scan map

  It is expected that the parent class is a subclass of ``Observation`` already
  by virtue of it being a superclass of subclass which inherits these methods.

  Attrs:
    cfg (dict):
    data (numpy array):      from ``Observation``
    logger (logging.Logger): replaces ``Observation`` logger
    name (str):              replaces ``Observation`` name
    session (Session):
    source (str):
    step (float):            map step size

  """
  def get_grid_stepsize(self, xy=None):
    """
    Determine the stepsize of gridded data

    This assumes xdec and dec data increase incrementally by 'stepsize'.
    The sequences may repeat in a sawtooth-like series. The number of
    'xdec' and 'dec' points is multiple times the gridsize.

    Arguments:
      xy (tuple or list) - X-array and Y-array (default Map.data)
    """
    # get the absolute value of coordinate intervals
    if xy:
      dxdecs = abs(xy[0][1:] - xy[0][:-1])
      ddecs  = abs(xy[1][1:] - xy[1][:-1])
    else:
      dxdecs = abs(self.data['xdec_offset'][1:]-self.data['xdec_offset'][:-1])
      ddecs  = abs(self.data['dec_offset'][1:] -self.data['dec_offset'][:-1])
    # form array of X,Y pairs
    coords =  NP.array(list(zip(dxdecs,ddecs)))
    # expect two clusters (default)
    cluster_pos = VQ.find_clusters(coords).round(4) # tenths of mdeg
    # return the non-zero intervals
    return cluster_pos[0].max(), cluster_pos[1].max()

  def regrid(self, width=1.0, height=1.0, step=None, power_key=None):
    """
    converts a map from observed coordinates to map coordinates

    If ``step`` is not given then the step size will be the average step size
    in X and the average step in Y.  In this case, the effect is to make a
    regular grid if the original positions were not exact, i.e., pointing error.

    @param width : map width in deg
    @type  width : float

    @param height : map height in deg
    @type  height : float

    @param step : map step size in X and Y in deg
    @type  step : (float, float)

    @param power_key : dict key of Z-value
    @type  power_key : str
    """
    # what is the power-like quantity?
    if power_key:
      pass
    else:
      # take the first that matches
      for key in Observation.power_keys:
        if key in self.data:
          power_key = key
          self.logger.info("regrid: using '%s'", power_key)
          break
        else:
          continue
    if power_key:
      pass
    else:
      self.logger.error("regrid: no power data key found")
      return None
    if step == None:
      # use the original stepsize
      self.xstep, self.ystep = self.get_grid_stepsize()
    else:
      self.xstep, self.ystep = step
    self.data['grid_x'] = NP.arange(
                                  -width/2, width/2+self.xstep/2, self.xstep/2)
    self.data['grid_y'] = NP.arange(
                                 -height/2,height/2+self.ystep/2, self.ystep/2)
    self.logger.debug("regrid: grid shape is %dx%d", len(self.data['grid_x']),
                                                     len(self.data['grid_y']))
    self.data['grid_z'] = {}
    for chnl in self.channel:
      self.logger.debug("regrid: processing %s", chnl)
      points = list(zip(self.data['xdec_offset'],self.data['dec_offset']))
      self.logger.debug("regrid: %d positions", len(points))
      values = self.data[power_key][chnl]
      self.logger.debug("regrid: %d values", len(values))
      xi, yi = NP.meshgrid(self.data['grid_x'], self.data['grid_y'])
      try:
        self.data['grid_z'][chnl] = scipy.interpolate.griddata(points, values,
                                                    (xi, yi), method='nearest')
      except ValueError as details:
        self.logger.error("regrid: gridding failed: %s", str(details))
        self.logger.debug("regrid: channel %s length of points is %d",
                                                              chnl, len(points))
        self.logger.debug("regrid: channel %s length of values is %d", chnl,
                                                                    len(values))
        continue

  def radec_from_azel(self):
    """
    compute RA and dec from az and el
    """
    RA = []; decs = []; RAdecs = []
    for idx in list(range(self.numdata)):
      # setup
      dt = self.data['datetime'][idx]
      # format time as (YEAR, DOY.fff)
      time_tuple = (dt.year,
                    DT.day_of_year(dt.year,dt.month,dt.day)
                      + (  dt.hour
                         + dt.minute/60.
                         + dt.second/3600.
                         + dt.microsecond/3600./1e6)/24.)
      azimuth = self.data['az'][idx]
      elevation = self.data['el'][idx]
      # compute
      ra,dec = A.AzEl_to_RaDec(azimuth, elevation,
                               self.latitude,
                               -self.longitude,
                               time_tuple)
      RA.append(ra)
      decs.append(dec)
      RAdecs.append((RA,decs))
    self.data['ra'] = RA
    self.data['dec'] = decs
    self.data['radec'] = RAdecs

  def radec2000_from_radec(self):
    """
    compute RA2000 and dec2000 from observed RA and dec
    """
    RA2000 = []; decs2000 = []; RAdec2000 = []
    for idx in list(range(self.numdata)):
      # setup
      tm = self.data['unixtime'][idx]
      mjd = DT.UnixTime_to_MJD(tm)
      MJD = int(mjd)
      UT = 24*(mjd-MJD)
      ra = self.data['ra']
      dec = self.data['dec']
      # compute
      ra2000,dec2000 = A.apparent_to_J2000(MJD,UT,
                                           ra, dec,
                                           self.longitude, self.latitude)
      RA2000.append(ra2000)
      decs2000.append(dec2000)
      RAdec2000.append((ra2000,dec2000))
    self.data['ra2000'] = RA2000
    self.data['dec2000'] = dec2000
    self.data['radec2000'] = RAdec2000

  def radec_from_radec2000(self):
    """
    compute apparent RA and dec. from J2000 RA and dec
    """
    RA = []; decs = []; RAdecs = []
    for idx in list(range(self.numdata)):
      # setup
      tm = self.data['unixtime'][idx]
      mjd = DT.UnixTime_to_MJD(tm)
      MJD = int(mjd)
      UT = 24*(mjd-MJD)
      ra2000 = self.data['ra2000'][idx]
      dec2000 = self.data['dec2000'][idx]
      # compute
      ra, dec = A.J2000_to_apparent(MJD, UT,
                                    ra2000*math.pi/12, dec2000*math.pi/180)
      RA.append(ra)
      decs.append(dec)
      RAdecs.append((ra,dec))
    self.data['ra'] = RA
    self.data['dec'] = decs
    self.data['radec'] = RAdecs

  def azel_from_radec(self):
    """
    compute azimuth and elevation from apparent right ascension and declination
    """
    azs = []; els = []; azels = []
    for idx in list(range(self.numdata)):
      # setup
      ra = self.data['ra'][idx]
      dec = self.data['dec'][idx]
      timetuple = self.data['datetime'][idx].timetuple()
      year = timetuple.tm_year
      doy = timetuple.tm_yday + (timetuple.tm_hour
                              +(timetuple.tm_min+timetuple.tm_sec/60)/60)/24
      # compute
      az, el = A.RaDec_to_AzEl(ra, dec,
                                     self.latitude, self.longitude, (year,doy))
      azs.append(az)
      els.append(el)
      azels.append((az,el))
    self.data['az'] = azs
    self.data['el'] = els
    self.data['azel'] = azels

  def get_offsets(self, source="Sun", xdec_ofst=0., dec_ofst=0.):
    """
    Generates a map in coordinates relative to a source

    If the source is the default, the position of the Sun will be computed for
    the time of each sample. IT SEEMS LIKE A GOOD IDEA TO DO THIS FOR PLANETS
    ALSO.

    This adds elements with keys ``xdec_offset`` and ``dec_offset`` to the
    attribute ``data``.

    @param source : source at map center
    @type  source : ephem source instance

    @param xdec_ofst : relative X-dec position of sample
    @type  xdec_ofst : float

    @param dec_ofst : relative dec position of sample
    @type  dec_ofst : float

    @return: (dxdecs,ddecs) in degrees
    """
    if source.lower() == "sun":
      src = AE.ephem.Sun()
    else:
      src = AE.calibrator(source)
    self.data['dec_offset'] = []
    self.data['xdec_offset'] = []
    for count in range(len(self.data['unixtime'])):
      dt = datetime.datetime.utcfromtimestamp(
                                     self.data['unixtime'][count])
      if type(src) == AE.Quasar:
        pass
      else:
        src.compute(dt)
      ra_center = src.ra*12/math.pi    # hours
      dec_center = src.dec*180/math.pi # degrees
      decrad = src.dec
      # right ascension increases to the left, cross-dec to the right
      self.data['xdec_offset'].append(xdec_ofst -
                      (self.data['ra'][count]  - ra_center)*15*math.cos(decrad) )
      self.data['dec_offset'].append(  dec_ofst +
                       self.data['dec'][count] - dec_center)
    # change list to NP.array
    self.data['xdec_offset'] = NP.array(self.data['xdec_offset'])
    self.data['dec_offset'] = NP.array(self.data['dec_offset'])

class Map(Observation, GriddingMixin):
  """
  Map class without special features for GAVRT and Malargue

  Most of the methods are mixed in to avoid conflicting with subclasses
  """
  def __init__(self, parent=None, name=None, dss=None, date=None, project=None):
    """
    Create a Map object

    Args:
      parent (Session): an observing session to which this belongs
      name (str):       an identifier, like a scan number
      dss (int):        station where the data were taken
      date (str):       date of observation as "YEAR/DOY"
      project (str):    project for which this observation was made
    """
    Observation.__init__(self, parent=parent, name=name, dss=dss, date=date,
                               project=project)

class Recording(h5py.File):
  """
  Class for raw data

  This is typically the contents of a data file transcribed into a standard
  format.  It may be the data of one Observation object, or data for multiple
  Observation objects, or contain part of the data for an Observation object.

  If the data being curated are not in a standard project, and they are not
  in a standard place,
  """
  def __init__(self, session=None, path=None, date=None, dss=None, name=None):
    """
    Initialize a metadata container and data directory

    Args
    ====
      session (Session): required, unless:
      path (str)       : location of raw data files
      date
    """
    self.logger = logging.getLogger(logger.name+".Recording")
    if session:
      self.session = session
      if not name:
        name = session.project + "-" + str(session.year) + "-" + \
                    ('%03d' % session.doy) + "-dss" + str(session.dss)+".info"
      self.year = session.year
      self.doy = session.doy
      self.dss = session.dss
      self.project = session.project
      self.session_dir = session.session_dir
    elif path and name:
      self.session = Session() # for its methods and attributes
      self.session_dir = path
      self.name = name
    else:
      raise RuntimeError("either a session or a path and filename required")
    h5py.File.__init__(self, name, 'w')
    self.attrs['project'] = self.project
    self.attrs['dss'] = self.dss
    self.attrs['year'] = self.year
    self.attrs['doy'] = self.doy

class Session(object):
  """
  Base class for an observing session on a given year and DOY

  Public Attributes::

    doy (int)               - day of year for session
    logger (logging.Logger) - logging.Logger object
    parent (object)         - a data reduction session (mult. observ. sessions)
    year (int)              -
    doy (int)               -
    project (str)           -
    session_dir (str)       - path to results from this session

  A session usually refers to a telescope, date and project.  This will
  normally define a path to the session directory.
  """
  def __init__(self, parent=None, date=None, project=None, dss=None,
                     path=None):
    """
    initialize data reduction for one observing session

    Args
    ====
      parent:  (object) optional class for a data reduction tool
      date:    (str) required, format YEAR/DOY
      project: (str) required
      dss      (int) required
      path     (str) optional

    If `path` is given for a non-standard observing files location, and it does
    not exist, it will be created.  Then the Recording and Observation instances
    must be directed to where the files are.
    """
    self.logger = logging.getLogger(logger.name+".Session")
    if parent:
      self.session = parent
    if date and project and dss:
      y,d = date.split('/')
      self.year = int(y);
      self.doy = int(d)
      self.project = project
      self.dss = dss
      self.name = "'%s %4d/%03d'" % (self.project, self.year, self.doy)
    else:
      self.logger.error("__init__: missing DSS or year or DOY or project")
      raise Exception("Where and when and for what project were the data taken?")
    self.find_session_dir(path=path)

  def find_session_dir(self, path=None):
    """
    find or make the sessions directory

    Args:
      path (str) - explicit path to files
    """
    self.logger.debug("find_session_dir: entered for path=%s", path)
    if path:
      self.session_dir = path
    else:
      obs_dir = local_dirs.projects_dir + self.project \
                                          +"/Observations/dss"+str(self.dss)+"/"
      self.session_dir = obs_dir+ "%4d" % self.year +"/"+ "%03d" % self.doy +"/"
    if not os.path.exists(self.session_dir):
      os.makedirs(self.session_dir, mode=0o775)

  def select_data_files(self, datapath=None, name_pattern="", auto=True,
                              load_hdf=False):
    """
    Provide the user with menu to select data files.

    Finding the right data store is complicated as there are many kinds of data
    files

    * If datapath is ...RA_data/HDF5/... then the files could be .h5 (Ashish)
      or .hdf5 (Dean).
    * If datapath is ...RA_data/FITS/... then the extent is .fits.
    * If datapath is ...project_data/... then the extent is .pkl
    * If datapath is ...projects/...     (default) then the extent is probably
      .csv or .dat or .prd.

    @param datapath : path to top of the tree where the DSS subdirectories are
    @type  datapath : str

    @param name_pattern : pattern for selecting file names, e.g. source
    @type  name_pattern : str

    @param load_hdf : use RA_data/HDF5 directory if True
    @type  load_hdf : bool

    @para auto : take all files found
    @type auto : bool

    @return: list of str
    """
    # Get the data files to be processed
    self.logger.debug("select_data_files: looking in %s", datapath)
    if name_pattern:
      name,extent = os.path.splitext(name_pattern)
      if extent.isalpha(): # a proper extent with no wildcards
        # take name pattern as is
        pass
      else:
        # only one * at front and back of pattern
        name_pattern = "*"+name_pattern.rstrip('*')+"*"
    else:
      # no pattern specified.  All files.
      name_pattern = "*"
    self.logger.debug("select_data_files: for pattern %s", name_pattern)
    if datapath:
      if re.search('HDF5', datapath):
        load_hdf = True
      elif re.search('project_data', datapath):
        load_hdf = False
        datafiles = support.text.select_files(datapath+name_pattern+"[0-9].pkl")
      elif re.search('FITS', datapath):
        datafiles = support.text.select_files(datapath+name_pattern+".fits")
      if load_hdf:
        full = datapath+name_pattern+".h*5"
      else:
        full = datapath+name_pattern
    else:
      full = self.session_dir + name_pattern
    self.logger.debug("select_data_files: from: %s", full)
    if auto:
      datafiles = glob.glob(full)
    else:
      datafiles = support.text.select_files(full)

    self.logger.debug("select_data_files: found %s", datafiles)
    if datafiles == []:
      self.logger.error(
                "select_data_files: None found. Is the data directory mounted?")
      raise RuntimeError('No data files found.')
    if type(datafiles) == str:
      datafiles = [datafiles]
    self.logger.info("select_data_files: to be processed: %s", datafiles)
    return datafiles

class Spectrum(Observation):
  """
  Class for spectra
  """
  def __init__(self):
    """
    needs a spectrum attribute
    """
    self.logger = logging.getLogger(logger.name+".Spectrum")

  def get_num_chans(self, linefreq, bandwidth, max_vel_width):
    """
    compute the base 2 number of output channels for the specified resolution
    """
    kmpspMHz = 300000./linefreq
    BW_kmps = bandwidth*kmpspMHz
    est_num_chan_out = BW_kmps/max_vel_width
    self.logger.debug("get_num_chans: estimated num chans out = %d",
                      est_num_chan_out)
    return 2**int(math.ceil(math.log(est_num_chan_out,2)))

  def reduce_spectrum_channels(self, refval, refpix, delta,
                             num_chan=1024, axis=0):
    """
    Reduce the number of channels in the spectrum.

    The default option is to reduce the spectrum to a specified number of
    channels with a default of 1024. The input spectrum is presumed to have
    2**N channels so that num_chan/num_chan_in is an integer.

    If 'spectrum' is an N-D array, then the spectrum axis is given by 'axis'
    which defaults to 0.

    'delta' is negative for lower sideband or reversed double sideband spectra.

    @param spectrum : spectrum values
    @type  spectrum : list or nparray

    @param refval : X-axis value at the reference pixel of 'spectrum'
    @type  refval : float

    @param refpix : reference pixel for 'spectrum'
    @type  refpix : int

    @param delta : interval between pixels on the X-axis
    @type  delta : float

    @param num_chan : optional number of channels to be returned (default: 2^10)
    @type  num_chan : int

    @return: numpy.array
    """
    if math.log(num_chan,2) % 1:
      raise RuntimeError("num_chan = %d is not a power of 2", num_chan)
    if type(self.spectrum) == NP.ndarray:
      num_chans_in = self.spectrum.shape[axis]
    else:
      num_chans_in = len(self.spectrum)
    if math.log(num_chans_in,2) % 1:
      raise RuntimeError("input spectrum length = %d is not a power of 2",
                                                                  num_chans_in)
    self.logger.debug("reduce_spectrum_channels: %d channels in", num_chans_in)

    num_chan_avg = num_chans_in/num_chan
    newrefpix = refpix/num_chan_avg
    self.logger.debug("reduce_spectrum_channels: refpix from %d to %d",
               refpix, newrefpix)

    newdelta = delta*num_chan_avg
    self.logger.debug("reduce_spectrum_channels: delta from %.3f to %.3f",
               delta, newdelta)
    newrefval = refval + delta*(num_chan_avg/2 - 1)
    self.logger.debug("reduce_spectrum_channels: refval from %.3f to %.3f",
               refval, newrefval)
    self.logger.debug("reduce_spectrum_channels: averaging %d channels", num_chan_avg)

    specout = NP.array([spectrum[index*num_chan_avg:(index+1)*num_chan_avg].mean()
                                                 for index in range(num_chan)])
    self.logger.debug("reduce_spectrum_channels: %d channels out", num_chan)
    return specout, newrefval, newrefpix, newdelta

  def get_freq_array(self, bandwidth, n_chans):
    """
    Create an array of frequencies for the channels of a backend

    @param bandwidth : bandwidth
    @type  bandwidth : float

    @param n_chans : number of channels
    @type  n_chans : int

    @return: frequency of each channel in same units as bandwidth
    """
    return NP.arange(n_chans)*float(bandwidth)/n_chans

  def freq_to_chan(frequency,bandwidth,n_chans):
    """
    Returns the channel number where a given frequency is to be found.

    @param frequency : frequency of channel in sane units as bandwidth.
    @type  frequency : float

    @param bandwidth : upper limit of spectrometer passband
    @type  bandwidth : float

    @param n_chans : number of channels in the spectrometer
    @type  n_chans : int

    @return: channel number (int)
    """
    if frequency < 0:
      frequency = bandwidth + frequency
    if frequency > bandwidth:
      raise RuntimeError("that frequency is too high.")
    return round(float(frequency)/bandwidth*n_chans) % n_chans

  def get_smoothed_bandshape(self, degree = None, poly_order=15):
    """
    Do a Gaussian smoothing of the spectrum and then fit a polynomial.
    Optionally, the raw and smoothed data and the fitted polynomial can be
    plotted.

    Note
    ====
    ``numpy.polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False)``
    Least squares polynomial fit.
    Fit a polynomial::

      p(x) = p[0] * x**deg + ... + p[deg]

    of degree deg to points (x, y).
    Returns a vector of coefficients p that minimises the squared error.

    @param spectrum : input data
    @type  spectrum : list of float

    @param degree : number of samples to smoothed (Gaussian FWHM)
    @type  degree : int

    @param poly_order : order of the polynomial
    @type  poly_order : int

    @param plot : plotting option
    @type  plot : boolean

    @return: (polynomial_coefficient, smoothed_spectrum)
    """
    if degree == None:
      degree = len(self.spectrum)/100
    # normalize the spectrum so max is 1 and convert to dB.
    max_lev = NP.max(self.spectrum)
    norm_spec = NP.array(self.spectrum)/float(max_lev)
    norm_spec_db = 10*NP.log10(norm_spec)
    # do a Gaussian smoothing
    norm_spec_db_smoothed = smoothListGaussian(norm_spec_db, degree=degree)
    # deal with the edges by making them equal to the smoothed end points
    norm_spec_db_smoothed_resized = NP.ones(len(self.spectrum))
    # left end
    norm_spec_db_smoothed_resized[0:degree] = norm_spec_db_smoothed[0]
    # middle
    norm_spec_db_smoothed_resized[degree:degree+len(norm_spec_db_smoothed)] = \
      norm_spec_db_smoothed
    # right end
    norm_spec_db_smoothed_resized[degree+len(norm_spec_db_smoothed):] = \
      norm_spec_db_smoothed[-1]
    return poly, norm_spec_db_smoothed_resized

# ------------------------ module functions -------------------------------

def examine_text_data_file(filename):
  """
  Examine a file to guide ``genfromtxt()``

  Things to look for::

    * Is there a header line with column names? If not, use argument ``names``.
    * Is the number of names equal to the number of columns? If not::
      - use argument ``names`` and ``skip_header=1``, or
      - use argument ``delimiter`` with a list of column widths
        and ``skip_header=1``.
  """
  print(examine_text_data_file.__doc__)
  fd = open(filename, "r")
  lines = fd.readlines()
  fd.close()
  topline = lines[0].strip().split()
  print("          1         2         3         4         5         6         7")
  print("01234567890123456789012345678901234567890123456789012345678901234567890123456789")
  print(lines[0].strip())
  print(lines[1].strip())
  print(" ...")
  print(lines[-1].strip())
  data = NP.genfromtxt(filename, dtype=None, names=None, skip_header=1, encoding=None)
  print("%d datatypes:" % len(data.dtype.fields))
  for item in data.dtype.fields:
    print(item, data.dtype.fields[item])

def get_obs_dirs(project, station, year, DOY, datafmt=None):
  """
  Returns the directories where data and working files are kept

  @param project : project code string, e.g., RRL
  @type  project : str

  @param station : DSN station number
  @type  station : int

  @param year : year of observation
  @type  year : int

  @param DOY : day of year of observations
  @type  DOY : int

  @param datafmt : raw data format
  @type  datafmt : str
  """
  #logger.debug("get_obs_dirs: type %s for %s, DSS%d, %4d/%03d",
  #             datafmt, project, station, year, DOY)
  obspath = "dss%2d/%4d/%03d/" %  (station,year,DOY)
  if project:
    projdatapath = "/usr/local/project_data/"+project+"/"+obspath
    projworkpath = "/usr/local/projects/"+project+"/Observations/"+obspath
  else:
    projdatapath = ""
    projworkpath = ""
  if datafmt:
    rawdatapath = "/usr/local/RA_data/"+datafmt+"/"+obspath
  else:
    rawdatapath = ""
  return projdatapath, projworkpath, rawdatapath

# --------- old stuff to be discarded still needed for now ---------------

def old_get_obs_session(project=None, dss=None, date=None, path='proj'):
  """
  Provides project, station, year and DOY, asking as needed.

  It follows one of several possible paths to get to the session::

    proj - path through /usr/local/projects/<project>
    hdf5 - path through /usr/local/RA_data/HDF5
    fits - path through /usr/local/RA_data/FITS
    wvsr - path through /data

  @param project : optional name as defined in /usr/local/projects
  @type  project : str

  @param dss : optional station number
  @type  dss : int

  @param date : optional YYYY/DDD
  @type  date : str

  @return: project, DSS, year, DOY.
  """
  def get_directory(path):
    """
    """
    # only one trailing /
    path = path.rstrip('/')+"/*"
    logger.debug("get_obs_session:get_directory: from %s", path)
    names = glob.glob(path)
    if names:
      dirs = []
      for name in names:
        if os.path.isdir(name):
          dirs.append(os.path.basename(name))
      dirs.sort()
      for name in dirs:
        print((name), end=' ')
      return input('\n>')
    else:
      return []

  def from_wvsr_dir():
    """
    this needs to be completed and tested on crab14 or an auto host
    """
    session = get_directory(local_dirs.wvsr_dir)
    return session

  cwd = os.getcwd()
  # get the project
  if project:
    pass
  else:
    os.chdir(local_dirs.projects_dir)
    project = get_directory(local_dirs.projects_dir)
  logger.debug("from_wvsr_dir: project is %s", project)
  projectpath = local_dirs.projects_dir+project
  # get the station
  if path[:4].lower() == 'wvsr':
    # special call
    print("from_wvsr_dir()")
  if path[:4].lower() == 'proj':
    os.chdir(projectpath+"/Observations/")
  elif path[:4].lower() == 'hdf5':
    os.chdir(local_dirs.hdf5_dir)
  elif path[:4].lower() == 'fits':
    os.chdir(local_dirs.fits_dir)

  # get the station
  if dss:
    pass
  else:
    # This seems odd but get_directory() needs '/' and int does not
    station = get_directory(os.getcwd()+"/").rstrip('/')
    dss = int(station[-2:])
  stationpath = os.getcwd()+"/dss"+str(dss)
  # get the date
  if date:
    items = date.split('/')
    year = int(items[0])
    DOY = int(items[1])
  else:
    year = int(get_directory(stationpath))
    yearpath = stationpath+"/"+str(year)
    DOY = int(get_directory(yearpath))
  os.chdir(cwd)
  return project, dss, year, DOY

reason: import numpy,import scipy


text:
import os
import gzip
import math
import time
import requests
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import torch
from torch.distributions.uniform import Uniform
from torch.nn.functional import affine_grid, grid_sample

"""
==========
simulate bouncing mnist using the training dataset in mnist
==========
"""
class Sim_BMNIST():
    def __init__(self, timesteps, num_digits, frame_size, delta_t, chunk_size):
        '''
        X : coordinates
        V : velocity
        '''
        super(Sim_BMNIST, self).__init__()
        self.timesteps = timesteps
        self.num_digits = num_digits
        self.frame_size = frame_size
        self.mnist_size = 28 ## by default
        self.delta_t = delta_t
        self.chunk_size = chunk_size ## datasets are dividied into pieces with this number and saved separately

    def load_mnist(self, MNIST_DIR):
        MNIST_PATH = os.path.join(MNIST_DIR, 'train-images-idx3-ubyte.gz')
        if not os.path.exists(MNIST_PATH):
            print('===Downloading MNIST train dataset from \'http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\'===')
            if not os.path.exists(MNIST_DIR):
                os.makedirs(MNIST_DIR)
            r = requests.get('http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz')
            open(MNIST_PATH,'wb').write(r.content)
            print('===Saved to \'%s\'===' % MNIST_PATH)

        with gzip.open(MNIST_PATH, 'rb') as f:
            mnist = np.frombuffer(f.read(), np.uint8, offset=16)
            mnist = mnist.reshape(-1, 28, 28)
        return mnist

    def sim_trajectory(self, init_xs):
        ''' Generate a random sequence of a MNIST digit '''
        v_norm = Uniform(0, 1).sample() * 2 * math.pi
        #v_norm = torch.ones(1) * 2 * math.pi
        v_y = torch.sin(v_norm).item()
        v_x = torch.cos(v_norm).item()
        V0 = torch.Tensor([v_x, v_y])
        X = torch.zeros((self.timesteps, 2))
        V = torch.zeros((self.timesteps, 2))
        X[0] = init_xs
        V[0] = V0
        for t in range(0, self.timesteps -1):
            X_new = X[t] + V[t] * self.delta_t
            V_new = V[t]

            if X_new[0] < -1.0:
                X_new[0] = -1.0 + torch.abs(-1.0 - X_new[0])
                V_new[0] = - V_new[0]
            if X_new[0] > 1.0:
                X_new[0] = 1.0 - torch.abs(X_new[0] - 1.0)
                V_new[0] = - V_new[0]
            if X_new[1] < -1.0:
                X_new[1] = -1.0 + torch.abs(-1.0 - X_new[1])
                V_new[1] = - V_new[1]
            if X_new[1] > 1.0:
                X_new[1] = 1.0 - torch.abs(X_new[1] - 1.0)
                V_new[1] = - V_new[1]
            V[t+1] = V_new
            X[t+1] = X_new
        return X, V

    def sim_trajectories(self, num_tjs, save_flag=False):
        Xs = []
        Vs = []
        x0 = Uniform(-1, 1).sample((num_tjs, 2))
        # a2 = 0.5**2
        # while(True):
        #     if ((x0[0] - x0[1])**2).sum() > a2:
        #         break
        #     x0 = Uniform(-1, 1).sample((num_tjs, 2))
        for i in range(num_tjs):
            x, v = self.sim_trajectory(init_xs=x0[i])
            Xs.append(x.unsqueeze(0))
            Vs.append(v.unsqueeze(0))
        if save_flag:
            np.save('pos', torch.cat(Xs, 0).data.numpy())
            np.save('disp', torch.cat(Vs, 0).data.numpy())
        return torch.cat(Xs, 0), torch.cat(Vs, 0)

    def sim_one_bmnist(self, mnist, mnist_index):
        '''
        Get random trajectories for the digits and generate a video.
        '''
        s_factor = self.frame_size / self.mnist_size
        t_factor = (self.frame_size - self.mnist_size) / self.mnist_size
        bmnist = []
        Xs, Vs = self.sim_trajectories(num_tjs=self.num_digits)
        for k in range(self.num_digits):
            digit_image = torch.from_numpy(mnist[mnist_index[k]] / 255.0).float()
            S = torch.Tensor([[s_factor, 0], [0, s_factor]]).repeat(self.timesteps, 1, 1)
            Thetas = torch.cat((S, Xs[k].unsqueeze(-1) * t_factor), -1)
            grid = affine_grid(Thetas, torch.Size((self.timesteps, 1, self.frame_size, self.frame_size)), align_corners=True)
            bmnist.append(grid_sample(digit_image.repeat(self.timesteps, 1, 1).unsqueeze(1), grid, mode='nearest', align_corners=True))
            # TJ.append(Xs[n].unsqueeze(0))
            # Init_V.append(V[0.unsqueeze()])
        bmnist = torch.cat(bmnist, 1).sum(1).clamp(min=0.0, max=1.0)
        return bmnist

    def sim_save_data(self, num_seqs, PATH):
        """
        ==========
        way it saves data:
        if num_seqs <= N, then one round of indexing is enough
        if num_seqs > N, then more than one round is needed
        ==========
        """
        if not os.path.exists(PATH):
            os.makedirs(PATH)
        mnist = self.load_mnist(MNIST_DIR=PATH)
        N = mnist.shape[0]
        assert num_seqs > 0, 'number of sequences must be a positive number'
        assert isinstance(num_seqs, int)
        consecutive = torch.arange(N).repeat(1, self.num_digits).squeeze(0)
        mnist_indices = consecutive[torch.randperm(N*self.num_digits)].view(N, self.num_digits) ## generate random indices
        num_seqs_left = num_seqs
        print('Start to simulate bouncing mnist sequences...')
        counter = 1
        while(num_seqs_left > 0):
            time_start = time.time()
            num_this_round = min(self.chunk_size, num_seqs_left)
            bmnists = []
            if mnist_indices.shape[0] < num_this_round: ## more indices are needed for this round
                new_indices = consecutive[torch.randperm(N*self.num_digits)].view(N, self.num_digits)
                mnist_indices = torch.cat((mnist_indices, new_indices), 0)
            for i in range(num_this_round):
                bmnist = self.sim_one_bmnist(mnist=mnist, mnist_index=mnist_indices[i])
                bmnists.append(bmnist.unsqueeze(0))
            mnist_indices = mnist_indices[num_this_round:]
            bmnists = torch.cat(bmnists, 0)
            assert bmnists.shape == (num_this_round, self.timesteps, self.frame_size, self.frame_size), "ERROR! unexpected chunk shape."
            incremental_PATH = PATH + 'ob-%d' % counter
            np.save(incremental_PATH, bmnists)
            counter += 1
            num_seqs_left = max(num_seqs_left - num_this_round, 0)
            time_end = time.time()
            print('(%ds) Simulated %d sequences, saved to \'%s\', %d sequences left.' % ((time_end - time_start), num_this_round, incremental_PATH, num_seqs_left))

    def viz_data(self, MNIST_DIR, num_seqs=5, fs=2):
        mnist = self.load_mnist(MNIST_DIR=MNIST_DIR)
        N = mnist.shape[0]
        mnist_indices = torch.arange(N).repeat(1, self.num_digits).squeeze(0)
        mnist_indices = mnist_indices[torch.randperm(N*self.num_digits)].view(N, self.num_digits)
        num_cols = self.timesteps
        num_rows = num_seqs
        gs = gridspec.GridSpec(num_rows, num_cols)
        gs.update(left=0.0 , bottom=0.0, right=1.0, top=1.0, wspace=0.05, hspace=0.05)
        fig = plt.figure(figsize=(fs * num_cols, fs * num_rows))
        for i in range(num_rows):
            bmnist = self.sim_one_bmnist(mnist, mnist_indices[i])
            for j in range(num_cols):
                ax = fig.add_subplot(gs[i, j])
                ax.set_xticks([])
                ax.set_yticks([])
                ax.imshow(bmnist[j], cmap='gray', vmin=0.0, vmax=1.0)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser('Bouncing MNIST DATA')
    parser.add_argument('--num_instances', default=60000, type=int)
    parser.add_argument('--data_path', default='../../data/bmnist/')
    parser.add_argument('--timesteps', default=10, help='number of video frames in one video')
    parser.add_argument('--num_digits', default=3, help='number of digitis in one video')
    parser.add_argument('--delta_t', default=0.3, help='constant velocity of the digits')
    parser.add_argument('--frame_size', default=96, help='squared size of the canvas')
    parser.add_argument('--chunk_size', default=1000, help='number of sqeuences that are stored in one single file (for the purpose of memory saving)')
    args = parser.parse_args()
    simulator = Sim_BMNIST(args.timesteps, args.num_digits, args.frame_size, args.delta_t, args.chunk_size)
    simulator.sim_save_data(args.num_instances, args.data_path)

max_stars_repo_name: Hao-W/svae

max_issues_repo_name: hao-w/apgs

max_forks_repo_name: hao-w/apgs