text:
# test reading of binary initial heads (float) and also binary icelltype (int).
# 1. Have binary data in a separate record for each layer
# 2. Have binary data in a single record for all layers

import os
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 testing_framework
from simulation import Simulation

ex = ["binary01", "binary02"]
exdirs = []
for s in ex:
    exdirs.append(os.path.join("temp", s))
ddir = "data"

def get_model(idx, dir):

    nlay, nrow, ncol = 5, 6, 7
    nper = 1
    perlen = 1.0
    nstp = 1
    tsmult = 1.0
    steady = [True]
    lenx = 300.0
    delr = delc = lenx / float(nrow)
    botm = np.linspace(-1.0, -5.0, 5)
    hnoflo = 1e30
    hdry = -1e30
    hk = 1.0

    nouter, ninner = 100, 300
    hclose, rclose, relax = 1e-6, 1e-3, 1.0

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

    name = ex[idx]

    # build MODFLOW 6 files
    ws = dir
    sim = flopy.mf6.MFSimulation(
        sim_name=name, version="mf6", exe_name="mf6", sim_ws=ws
    )
    # create tdis package
    tdis = flopy.mf6.ModflowTdis(
        sim, time_units="DAYS", nper=nper, perioddata=tdis_rc
    )

    # create gwf model
    gwf = flopy.mf6.MFModel(
        sim,
        model_type="gwf6",
        modelname=name,
        model_nam_file="{}.nam".format(name),
    )

    # create iterative model solution and register the gwf model with it
    ims = flopy.mf6.ModflowIms(
        sim,
        print_option="SUMMARY",
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration="CG",
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
    )
    sim.register_ims_package(ims, [gwf.name])

    # discretization
    # write top to a binary file
    text = "TOP"
    fname = "top.bin"
    pth = os.path.join(exdirs[idx], fname)
    f = open(pth, "wb")
    header = flopy.utils.BinaryHeader.create(
        bintype="HEAD",
        precision="double",
        text=text,
        nrow=nrow,
        ncol=ncol,
        ilay=1,
        pertim=1.0,
        totim=1.0,
        kstp=1,
        kper=1,
    )
    flopy.utils.Util2d.write_bin(
        (nrow, ncol),
        f,
        np.zeros((nrow, ncol), dtype=np.float64),
        header_data=header,
    )
    f.close()
    top = {
        "factor": 1.0,
        "filename": fname,
        "data": None,
        "binary": True,
        "iprn": 1,
    }

    # write botarr to binary file
    if idx == 0:
        botarr = []
        for k in range(nlay):
            text = "BOTM_L{}".format(k + 1)
            fname = "botm.l{:02d}.bin".format(k + 1)
            pth = os.path.join(exdirs[idx], fname)
            f = open(pth, "wb")
            header = flopy.utils.BinaryHeader.create(
                bintype="HEAD",
                precision="double",
                text=text,
                nrow=nrow,
                ncol=ncol,
                ilay=k + 1,
                pertim=1.0,
                totim=1.0,
                kstp=1,
                kper=1,
            )
            flopy.utils.Util2d.write_bin(
                (nrow, ncol),
                f,
                np.ones((nrow, ncol), dtype=np.float64) * botm[k],
                header_data=header,
            )
            f.close()
            botarr.append(
                {
                    "factor": 1.0,
                    "filename": fname,
                    "data": None,
                    "binary": True,
                    "iprn": 1,
                }
            )
    elif idx == 1:
        fname = "botm.bin"
        pth = os.path.join(exdirs[idx], fname)
        f = open(pth, "wb")
        tarr = np.ones((nlay, nrow, ncol), dtype=np.float64)
        for k in range(nlay):
            tarr[k, :, :] = botm[k]
        tarr = tarr.flatten()
        header = flopy.utils.BinaryHeader.create(
            bintype="HEAD",
            precision="double",
            text="BOTM",
            nrow=1,
            ncol=ncol * nrow * nlay,
            ilay=1,
            pertim=1.0,
            totim=1.0,
            kstp=1,
            kper=1,
        )
        flopy.utils.Util2d.write_bin((nrow, ncol), f, tarr, header_data=header)
        f.close()
        botarr = {
            "factor": 1.0,
            "filename": fname,
            "data": None,
            "binary": True,
            "iprn": 1,
        }

    # write idomain to binary file
    if idx == 0:
        idomain = []
        for k in range(nlay):
            text = "IDOMAIN_L{}".format(k + 1)
            fname = "idomain.l{:02d}.bin".format(k + 1)
            pth = os.path.join(exdirs[idx], fname)
            f = open(pth, "wb")
            header = flopy.utils.BinaryHeader.create(
                bintype="HEAD",
                precision="double",
                text=text,
                nrow=nrow,
                ncol=ncol,
                ilay=k + 1,
                pertim=1.0,
                totim=1.0,
                kstp=1,
                kper=1,
            )
            flopy.utils.Util2d.write_bin(
                (nrow, ncol),
                f,
                np.ones((nrow, ncol), dtype=np.int32),
                header_data=header,
            )
            f.close()
            idomain.append(
                {
                    "factor": 1.0,
                    "filename": fname,
                    "data": None,
                    "binary": True,
                    "iprn": 1,
                }
            )
    elif idx == 1:
        fname = "idomain.bin"
        pth = os.path.join(exdirs[idx], fname)
        f = open(pth, "wb")
        header = flopy.utils.BinaryHeader.create(
            bintype="HEAD",
            precision="double",
            text="IDOMAIN",
            nrow=1,
            ncol=ncol * nrow * nlay,
            ilay=1,
            pertim=1.0,
            totim=1.0,
            kstp=1,
            kper=1,
        )
        flopy.utils.Util2d.write_bin(
            (nrow, ncol),
            f,
            np.ones((nrow * ncol * nlay), dtype=np.int32),
            header_data=header,
        )
        f.close()
        idomain = {
            "factor": 1.0,
            "filename": fname,
            "data": None,
            "binary": True,
            "iprn": 1,
        }

    dis = flopy.mf6.ModflowGwfdis(
        gwf,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botarr,
        idomain=idomain,
        filename="{}.dis".format(name),
    )

    # initial conditions
    # write initial heads to binary file
    if idx == 0:
        strt = []
        for k in range(nlay):
            text = "IC_L{}".format(k + 1)
            fname = "ic.strt_l{:02d}.bin".format(k + 1)
            pth = os.path.join(exdirs[idx], fname)
            f = open(pth, "wb")
            header = flopy.utils.BinaryHeader.create(
                bintype="HEAD",
                precision="double",
                text=text,
                nrow=nrow,
                ncol=ncol,
                ilay=k + 1,
                pertim=1.0,
                totim=1.0,
                kstp=1,
                kper=1,
            )
            flopy.utils.Util2d.write_bin(
                (nrow, ncol),
                f,
                np.ones((nrow, ncol), dtype=np.float64),
                header_data=header,
            )
            f.close()
            strt.append(
                {
                    "factor": 1.0,
                    "filename": fname,
                    "data": None,
                    "binary": True,
                    "iprn": 1,
                }
            )
    elif idx == 1:
        fname = "ic.strt.bin"
        pth = os.path.join(exdirs[idx], fname)
        f = open(pth, "wb")
        header = flopy.utils.BinaryHeader.create(
            bintype="HEAD",
            precision="double",
            text="HEAD",
            nrow=1,
            ncol=ncol * nrow * nlay,
            ilay=1,
            pertim=1.0,
            totim=1.0,
            kstp=1,
            kper=1,
        )
        flopy.utils.Util2d.write_bin(
            (nrow, ncol),
            f,
            np.ones((nrow * ncol * nlay), dtype=np.float64),
            header_data=header,
        )
        f.close()
        strt = {
            "factor": 1.0,
            "filename": fname,
            "data": None,
            "binary": True,
            "iprn": 1,
        }

    ic = flopy.mf6.ModflowGwfic(gwf, strt=strt, filename="{}.ic".format(name))

    # node property flow
    # write icelltype to binary file
    if idx == 0:
        icelltype = []
        for k in range(nlay):
            fname = "npf.icelltype.l{}.bin".format(k + 1)
            pth = os.path.join(exdirs[idx], fname)
            f = open(pth, "wb")
            header = flopy.utils.BinaryHeader.create(
                bintype="head",
                text="ICELLTYPE",
                precision="double",
                nrow=nrow,
                ncol=ncol,
                ilay=k + 1,
                pertim=1.0,
                totim=1.0,
                kstp=1,
                kper=1,
            )
            header.tofile(f)
            flopy.utils.Util2d.write_bin(
                (nrow, ncol),
                f,
                np.ones((nrow, ncol), dtype=np.int32),
                header_data=header,
            )
            f.close()
            icelltype.append(
                {
                    "factor": 1.0,
                    "filename": fname,
                    "data": None,
                    "binary": True,
                    "iprn": 1,
                }
            )
    elif idx == 1:
        fname = "npf.icelltype.bin"
        pth = os.path.join(exdirs[idx], fname)
        f = open(pth, "wb")
        header = flopy.utils.BinaryHeader.create(
            bintype="head",
            text="ICELLTYPE",
            precision="double",
            nrow=1,
            ncol=ncol * nrow * nlay,
            ilay=1,
            pertim=1.0,
            totim=1.0,
            kstp=1,
            kper=1,
        )
        header.tofile(f)
        flopy.utils.Util2d.write_bin(
            (nrow, ncol),
            f,
            np.ones((nrow * ncol * nlay), dtype=np.int32),
            header_data=header,
        )
        icelltype = {
            "factor": 1.0,
            "filename": fname,
            "data": None,
            "binary": True,
            "iprn": 1,
        }

        f.close()

    npf = flopy.mf6.ModflowGwfnpf(
        gwf,
        save_flows=True,
        icelltype=icelltype,
        k=hk,
        k33=hk,
        filename="{}.npf".format(name),
    )

    # chd files
    chdlist0 = []
    chdlist0.append([(0, 0, 0), 1.0])
    chdlist0.append([(nlay - 1, nrow - 1, ncol - 1), 0.0])

    chdspdict = {0: chdlist0}
    chd = flopy.mf6.ModflowGwfchd(
        gwf,
        stress_period_data=chdspdict,
        save_flows=False,
        filename="{}.chd".format(name),
    )

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

    return sim

def build_models():
    for idx, dir in enumerate(exdirs):
        sim = get_model(idx, dir)
        sim.write_simulation()
    return

# - No need to change any code below
def test_mf6model():
    # initialize testing framework
    test = testing_framework()

    # build the models
    build_models()

    # run the test models
    for idx, dir in enumerate(exdirs):
        yield test.run_mf6, Simulation(dir)

    return

def main():
    # initialize testing framework
    test = testing_framework()

    # build the models
    build_models()

    # run the test models
    for idx, dir in enumerate(exdirs):
        sim = Simulation(dir)
        test.run_mf6(sim)

    return

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

    # run main routine
    main()


text:
from netket import legacy as nk
import flax
from flax import linen as nn
import jax
from jax import numpy as jnp
import numpy as np

# 1D Lattice
L = 20
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
hi = nk.hilbert.Spin(s=0.5) ** L

ha = nk.operator.Ising(h=1.0, hilbert=hi, graph=g)

alpha = 1

class RBM(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=alpha * x.shape[-1], dtype=jax.numpy.float32)(x)
        x = jnp.log(jnp.cosh(x))
        return jnp.sum(x, axis=-1)

ma = nk.machine.Flax(hi, RBM())
ma.init_random_parameters(seed=1232)

# Jax Sampler
sa = nk.sampler.MetropolisLocal(machine=ma, n_chains=2)

# Using Sgd
op = nk.optimizer.Sgd(ma, learning_rate=0.1)

# Create the optimization driver
gs = nk.Vmc(
    hamiltonian=ha, sampler=sa, optimizer=op, n_samples=1000, sr=None, n_discard=None
)

# The first iteration is slower because of start-up jit times
gs.run(out="test", n_iter=2)

gs.run(out="test", n_iter=300)

ext: py

max_stars_repo_name: gpescia/MyNetKet

max_issues_repo_name: gpescia/MyNetKet

max_forks_repo_name: gpescia/MyNetKet

reason: import numpy,import jax,from jax


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
The Hessians from two chemical models (QM and sGDML) are diagonalized
and the frequencies and normal modes are compared graphically.
"""
import matplotlib
matplotlib.rc('xtick', labelsize=12)
matplotlib.rc('ytick', labelsize=12)
matplotlib.rc('legend', fontsize=12)
matplotlib.rc('axes', labelsize=12)
import matplotlib.pyplot as plt

import numpy as np
import numpy.linalg as la
import torch

# # Local Imports
from semiclassical import units
from semiclassical.readers import FormattedCheckpointFile
from semiclassical.gdml_predictor import GDMLPredict

def vibrational_analysis(hess, masses, zero_threshold=1.0e-9):
    """
    compute vibrational frequencies and modes from Hessian

    Parameters:
    -----------
    hess:  ndarray (shape (3*nat,3*nat))
       Hessian in cartesian coordinates
    masses: ndarray (shape (3*nat))
       masses for each cartesian coordinate (in atomic units)
    zero_threshold: float
       modes with frequencies^2 below that threshold are treated as zero

    Returns:
    --------
    freq:  ndarray (shape (3*nat-6)))
      vibrational frequencies in atomic units
    modes: ndarray (shape (3*nat, 3*nat-6)))
      displacement vectors in mass-weighted coordinates
    """
    # convert Hessian to mass-weighted coordinates
    hess_mwc = hess / np.sqrt(np.outer(masses, masses))
    # mass weighted coordinates are now qi = sqrt(mi) dxi
    # compute eigen values of hess_mwc
    omega2,modes = la.eigh(hess_mwc)
    # modes that are zero within numerical accuracy
    zero_modes = np.where((omega2 < zero_threshold))[0]
    vib_modes = np.where((omega2 >= zero_threshold))[0]

    # number of zeros modes, #(zero-modes)
    Nzero = len(zero_modes)

    freqs = np.sqrt(omega2*(1.0+0.0j))
    # period T = 2*pi*hbar/E
    periods = 2.0*np.pi/freqs
    print("Frequencies")
    print("===========")
    print("- Zero modes (should be close to zero)")
    for fr in freqs[zero_modes]:
        print( "   %5.7f Hartree     %15.7f cm-1   " % (fr, fr*units.hartree_to_wavenumbers))
    print("- Vibrations")
    for fr,T in zip(freqs[vib_modes].real, periods[vib_modes].real):
        print("   %5.7f Hartree      %15.7f cm-1      %15.7f fs" % (fr, fr*units.hartree_to_wavenumbers, T*units.autime_to_fs))
    print("")
    en_zp = np.sum(freqs[6:].real)/2.0
    print("zero-point energy:   %5.7f Hartree      %5.7f cm-1" % (en_zp, en_zp*units.hartree_to_wavenumbers))
    # remove 6 translational and rotational modes
    return freqs[6:].real, modes[:,6:].real

def compare_frequencies(freqs_qm, freqs_sgdml):
    mode_nums = range(1, len(freqs_qm)+1)
    plt.bar(mode_nums, freqs_sgdml*units.hartree_to_wavenumbers, label="sGDML")
    plt.bar(mode_nums, freqs_qm*units.hartree_to_wavenumbers, label="QM", alpha=0.8)

    plt.ylabel("Frequency / cm$^{-1}$")
    plt.xlabel("Normal mode")

    plt.legend()
    plt.savefig("frequencies_sgdml_qm.svg")

    plt.show()

def compare_modes(modes_qm, modes_sgdml):
    mode_nums = range(1, modes_qm.shape[1]+1)
    # overlap between normal modes
    overlap = np.dot(modes_qm.T, modes_sgdml)
    #
    plt.bar(mode_nums, abs(np.diag(overlap)), label="mode overlap")

    plt.ylabel("Overlap")
    plt.xlabel("Normal mode")

    plt.legend()
    plt.savefig("modes_sgdml_qm.svg")

    plt.show()

if __name__ == "__main__":
    import sys
    import os.path

    usage = """
 Usage: {script}  freq.fchk  model.npz

      diagonalize mass-weighted Hessians of QM and sGDML model
      and compare frequencies and normal modes.

    Input Files:
      freq.fchk               :   fchk-file from frequency calculation
      model.npz               :   sGDML model

    """.format(script=os.path.basename(sys.argv[0]))
    args = sys.argv[1:]
    if len(args) < 2:
        print(usage)
        exit(-1)

    fchk_file = args[0]
    model_file = args[1]

    # QM frequencies
    with open(fchk_file) as f:
        fchk = FormattedCheckpointFile(f)

    pos, energy_qm, _, hessian_qm = fchk.harmonic_approximation()
    masses = fchk.masses()

    # sGDML frequencies
    model = np.load(model_file, allow_pickle=True)
    gdml = GDMLPredict(model)

    energy, _, hess = gdml.forward( torch.from_numpy(pos).unsqueeze(0) )
    energy_sgdml = energy[0].item()
    hessian_sgdml = hess[0,:,:]

    # compare energies
    print(f"ground state energy (Hartree) :  {energy_qm:12.7f} (QM)   vs.  {energy_sgdml:12.7f} (sGDML)")
    assert abs(energy_qm - energy_sgdml) < 0.05, "QM and predicted sGDML differ by more than 0.05 Hartree"

    freqs_qm, modes_qm = vibrational_analysis(hessian_qm, masses)
    freqs_sgdml, modes_sgdml = vibrational_analysis(hessian_sgdml, masses)

    compare_frequencies(freqs_qm, freqs_sgdml)
    compare_modes(modes_qm, modes_sgdml)

max_stars_repo_name: humeniuka/semiclassical

max_issues_repo_name: humeniuka/semiclassical

max_forks_repo_name: humeniuka/semiclassical

reason: import numpy


text:
import numpy as np
from typing import Optional
from dataclasses import dataclass

@dataclass
class _Representatives:
    rows_representatives: np.ndarray
    unique_inverse: np.ndarray
    len_added_rows: int

class FeaturesTransformer:
    """Transform array and generate 'just-enough' representatives for DBSCAN clustering.
    Use remap_labels method for remapping labels to original array.

    Parameters
    ----------
    features : array (n_samples, n_features)
    min_samples : int
        DBSCAN min_samples value.
    round_decimals : int, or None
        If a value is provided, the features will be rounded to the given number of decimals.
        Use to reduce false precision (i.e. meaningless decimals).

    Attributes
    ----------
    representatives : _Representatives
        Contains generated representative array, indices for unique inversing
        and the length difference between representatives and unique rows (required for remapping).
    """
    def __init__(self, features: np.ndarray, min_samples: int, round_decimals: Optional[int] = None):
        self.round_decimals = round_decimals
        self.min_samples = min_samples

        self.features = self._features_prep(features)
        self.representatives = self._representatives()

    def remap_labels(self, labels_of_stacked_rows: np.ndarray) -> np.ndarray:
        """
        Remap DBSCAN labels of representatives back to original feature set.

        Parameters
        -----------
        labels_of_stacked_rows : array (n_samples_representatives,)

        Returns
        --------
        labels_remapped : array (n_samples, )
        """
        labels_of_unique_rows = labels_of_stacked_rows[:-self.representatives.len_added_rows]

        labels_remapped = labels_of_unique_rows[self.representatives.unique_inverse]
        return labels_remapped

    def _representatives(self) -> _Representatives:
        """Generate a _Representatives object
        """
        unique_rows, unique_inverse, unique_counts = np.unique(self.features,
                                                               axis=0,
                                                               return_inverse=True,
                                                               return_counts=True)

        unique_counts_clipped_upto_min_samples = np.clip(unique_counts, 1, self.min_samples + 1) - 1
        rows_to_stack = unique_rows.repeat(unique_counts_clipped_upto_min_samples, axis=0)
        unique_rows_appended_with_repeats = np.append(unique_rows, rows_to_stack, axis=0)

        return _Representatives(unique_rows_appended_with_repeats, unique_inverse, len(rows_to_stack))

    def _round_(self, decimals: int) -> np.ndarray:
        rounded = np.round(self.features, decimals=decimals)
        return rounded

    def _features_prep(self, features: np.ndarray) -> np.ndarray:
        if isinstance(self.round_decimals, int):
            return np.round(features, self.round_decimals)
        return features

max_stars_repo_name: pazacharis/GridRep

max_issues_repo_name: pazacharis/GridRep

max_forks_repo_name: pazacharis/GridRep


text:
# (C) Copyright 2005-2021 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

"""
Tests for the Range trait with value type int.
"""

import unittest

from traits.api import (
    BaseRange,
    Either,
    HasTraits,
    Instance,
    Range,
    TraitError,
)
from traits.testing.optional_dependencies import numpy, requires_numpy

def ModelFactory(name, RangeFactory):
    """
    Helper function to create various similar model classes.

    Parameters
    ----------
    name : str
        Name to give the created class.
    RangeFactory : callable(*range_args, **range_kwargs) -> TraitType
        Callable with the same signature as Range; this will be used
        to create the model traits.

    Returns
    -------
    HasTraits subclass
        Subclass containing various Range-like traits, for testing.

    """

    class ModelWithRanges(HasTraits):
        """
        Model containing various Range-like traits.
        """

        # Simple integer range trait.
        percentage = RangeFactory(0, 100)

        # Traits that exercise the various possiblities for inclusion
        # or exclusion of the endpoints.
        open_closed = RangeFactory(0, 100, exclude_low=True)

        closed_open = RangeFactory(0, 100, exclude_high=True)

        open = RangeFactory(0, 100, exclude_low=True, exclude_high=True)

        closed = RangeFactory(0, 100)

        # Traits for one-sided intervals
        steam_temperature = RangeFactory(low=100)

        ice_temperature = RangeFactory(high=0)

    ModelWithRanges.__name__ = name

    return ModelWithRanges

class Impossible(object):
    """
    Type that never gets instantiated.
    """

    def __init__(self):
        raise TypeError("Cannot instantiate this class")

# A trait type that has a fast validator but doesn't accept any values.
impossible = Instance(Impossible, allow_none=False)

def RangeCompound(*args, **kwargs):
    """
    Compound trait including a Range.
    """
    return Either(impossible, Range(*args, **kwargs))

def BaseRangeCompound(*args, **kwargs):
    """
    Compound trait including a BaseRange.
    """
    return Either(impossible, BaseRange(*args, **kwargs))

ModelWithRange = ModelFactory("ModelWithRange", RangeFactory=Range)

ModelWithBaseRange = ModelFactory("ModelWithBaseRange", RangeFactory=BaseRange)

ModelWithRangeCompound = ModelFactory(
    "ModelWithRangeCompound", RangeFactory=RangeCompound
)

ModelWithBaseRangeCompound = ModelFactory(
    "ModelWithBaseRangeCompound", RangeFactory=BaseRangeCompound,
)

class InheritsFromInt(int):
    """
    Object that's integer-like by virtue of inheriting from int.
    """

    pass

class IntLike(object):
    """
    Object that's integer-like by virtue of providing an __index__ method.
    """

    def __init__(self, value):
        self._value = value

    def __index__(self):
        return self._value

class BadIntLike(object):
    """
    Object whose __index__ method raises something other than TypeError.
    """

    def __index__(self):
        raise ZeroDivisionError("bogus error")

class CommonRangeTests(object):
    def test_accepts_int(self):
        self.model.percentage = 35
        self.assertIs(type(self.model.percentage), int)
        self.assertEqual(self.model.percentage, 35)
        with self.assertRaises(TraitError):
            self.model.percentage = -1
        with self.assertRaises(TraitError):
            self.model.percentage = 101

    def test_accepts_bool(self):
        self.model.percentage = False
        self.assertIs(type(self.model.percentage), int)
        self.assertEqual(self.model.percentage, 0)

        self.model.percentage = True
        self.assertIs(type(self.model.percentage), int)
        self.assertEqual(self.model.percentage, 1)

    def test_rejects_bad_types(self):
        # A selection of things that don't count as integer-like.
        non_integers = [
            "not a number",
            "\N{GREEK CAPITAL LETTER SIGMA}",
            b"not a number",
            "3.5",
            "3",
            3 + 4j,
            0j,
            [1.2],
            (1.2,),
            0.0,
            -27.8,
            35.0,
            None,
        ]

        for non_integer in non_integers:
            self.model.percentage = 73
            with self.assertRaises(TraitError):
                self.model.percentage = non_integer
            self.assertEqual(self.model.percentage, 73)

    @requires_numpy
    def test_accepts_numpy_types(self):
        numpy_values = [
            numpy.uint8(25),
            numpy.uint16(25),
            numpy.uint32(25),
            numpy.uint64(25),
            numpy.int8(25),
            numpy.int16(25),
            numpy.int32(25),
            numpy.int64(25),
        ]
        for numpy_value in numpy_values:
            self.model.percentage = numpy_value
            self.assertIs(type(self.model.percentage), int)
            self.assertEqual(self.model.percentage, 25)

    @requires_numpy
    def test_rejects_numpy_types(self):
        numpy_values = [
            numpy.float16(25),
            numpy.float32(25),
            numpy.float64(25),
        ]
        for numpy_value in numpy_values:
            self.model.percentage = 88
            with self.assertRaises(TraitError):
                self.model.percentage = numpy_value
            self.assertEqual(self.model.percentage, 88)

    def test_accepts_int_subclass(self):
        self.model.percentage = InheritsFromInt(44)
        self.assertIs(type(self.model.percentage), int)
        self.assertEqual(self.model.percentage, 44)
        with self.assertRaises(TraitError):
            self.model.percentage = InheritsFromInt(-1)
        with self.assertRaises(TraitError):
            self.model.percentage = InheritsFromInt(101)

    def test_accepts_int_like(self):
        self.model.percentage = IntLike(35)
        self.assertIs(type(self.model.percentage), int)
        self.assertEqual(self.model.percentage, 35)
        with self.assertRaises(TraitError):
            self.model.percentage = IntLike(-1)
        with self.assertRaises(TraitError):
            self.model.percentage = IntLike(101)

    def test_bad_int_like(self):
        # Check that the exception is propagated as expected.
        with self.assertRaises(ZeroDivisionError):
            self.model.percentage = BadIntLike()

    def test_endpoints(self):
        # point within the interior of the range
        self.model.open = self.model.closed = 50
        self.model.open_closed = self.model.closed_open = 50
        self.assertEqual(self.model.open, 50)
        self.assertEqual(self.model.closed, 50)
        self.assertEqual(self.model.open_closed, 50)
        self.assertEqual(self.model.closed_open, 50)

        # low endpoint
        self.model.closed = self.model.closed_open = 0
        self.assertEqual(self.model.closed, 0)
        self.assertEqual(self.model.closed_open, 0)
        with self.assertRaises(TraitError):
            self.model.open = 0
        with self.assertRaises(TraitError):
            self.model.open_closed = 0

        # high endpoint
        self.model.closed = self.model.open_closed = 100
        self.assertEqual(self.model.closed, 100)
        self.assertEqual(self.model.open_closed, 100)
        with self.assertRaises(TraitError):
            self.model.open = 100
        with self.assertRaises(TraitError):
            self.model.closed_open = 100

    def test_half_infinite(self):
        ice_temperatures = [-273, -100, -1]
        water_temperatures = [1, 50, 99]
        steam_temperatures = [101, 1000, 10 ** 100, 10 ** 1000]

        for temperature in steam_temperatures:
            self.model.steam_temperature = temperature
            self.assertEqual(self.model.steam_temperature, temperature)

        for temperature in ice_temperatures + water_temperatures:
            self.model.steam_temperature = 1729
            with self.assertRaises(TraitError):
                self.model.steam_temperature = temperature
            self.assertEqual(self.model.steam_temperature, 1729)

        for temperature in ice_temperatures:
            self.model.ice_temperature = temperature
            self.assertEqual(self.model.ice_temperature, temperature)

        for temperature in water_temperatures + steam_temperatures:
            self.model.ice_temperature = -1729
            with self.assertRaises(TraitError):
                self.model.ice_temperature = temperature
            self.assertEqual(self.model.ice_temperature, -1729)

class TestIntRange(CommonRangeTests, unittest.TestCase):
    def setUp(self):
        self.model = ModelWithRange()

class TestIntBaseRange(CommonRangeTests, unittest.TestCase):
    def setUp(self):
        self.model = ModelWithBaseRange()

class TestIntRangeCompound(CommonRangeTests, unittest.TestCase):
    def setUp(self):
        self.model = ModelWithRangeCompound()

class TestIntBaseRangeCompound(CommonRangeTests, unittest.TestCase):
    def setUp(self):
        self.model = ModelWithBaseRangeCompound()


text:
from random_bril import gencode
import subprocess
import sys, re
import json
from scipy.stats import chisquare

spacere = r'\s+'
def nospace_str(stdout) :
    return re.sub(spacere, '', stdout.decode('utf-8'))

def getfreq(counts, key):
    if key.startswith("\'done\'"):
        return counts.get(key[7:],0)
    return 0

N  = 100
if __name__ == '__main__':
    n_errs = 0
    for i in range(10):
        program = gencode()
        progbytes = bytes(json.dumps(program, sort_keys=True), 'utf-8')

        xbrili = subprocess.run(['xbrili'], input=progbytes, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        erstr = nospace_str(xbrili.stderr)
        if  len(erstr) > 1:
            n_errs += 1
            # print("ERROR", erstr)
            continue

        paths = nospace_str(xbrili.stdout).split('[')

        dist = {}
        for p in paths:
            if len(p) < 2:
                continue
            key, val = p.split(']')
            dist[key] = float(val)

        # print(dist)
        counts = {}
        for  j in range(N):
            # sys.stdout.write('.')
            # sys.stdout.flush()

            try:
                brili_j = subprocess.run(['brili', '--envdump', '--noprint'], input =progbytes, stdout=subprocess.PIPE, timeout=2)

                env = nospace_str(brili_j.stdout)
                counts[env] = counts.get(env,0) + 1
            except subprocess.TimeoutExpired:
                continue
        # print(counts)
        # print(list(zip(*(round(dist[k] * N), getfreq(counts,k)) for k in dist.keys() )))
        try:
            f_exp, f_obs = zip(* [(round(dist[k] * N), getfreq(counts,k)) for k in dist.keys() if k.startswith("\'done\'")])
        except ValueError:
            f_exp = []
            f_obs = []
            continue
        chisq,p = chisquare(f_obs, f_exp)
        print('chi2', chisq, 'p', p,'\t', f_exp, f_obs)

    print('random programs that could not be executed by xbrili', n_errs)

max_stars_repo_name: orichardson/bril

max_issues_repo_name: orichardson/bril

max_forks_repo_name: orichardson/bril

reason: from scipy


text:
# -*- coding: utf-8 -*-
"""
Geometrical imaging equations, Paraxial optical calculations

"""

import numpy as np
import scipy.constants as const

# Need to clarify some of these functions
#could break into newton imaging equations and gaussian imaging equations.
#newtons were derived in 1666 using similar triangles, referenced to focal planes.

def obj_dist_from_EFL_and_m(f,m):
    """Calculate object distance (z) from focal length (f) and magnification (m)
    """

    obj_dist = -1*((1-m)/m)*f
    return obj_dist
    # todo: clarify assumptions in terms of +/- distances.

def img_dist_from_EFL_and_m(f,m):
    """Calculate image distance (z') from focal length (f) and magnification (m)
    """

    img_dist = (1-m)*f
    return img_dist
    # todo: clarify assumptions in terms of +/- distances.

def thin_lens_image_dist(obj_dist,efl):
    """Calculate the image distance given the object distance and efl.  Uses
    the thin lens equation 1/img + 1/obj = 1/f.  Returns same units.

    Parameters
    ----------
    obj_dist : float
        Distance to object
    efl : float
        Focal length of the thin lens

    Returns
    -------
    img_dist : float
        The image distance
    """

    efl = float(efl)
    obj_dist = float(obj_dist)      #convert to float if int
    img_dist = 1/(1/efl - 1/obj_dist)
    return img_dist
    #todo: catch infinite case
    #todo: clarify parameters.  This is the distance in front of the focal pt.

def two_lens_EFL(f1,f2,d):
    """Calculate the focal length of two thin lenses sepatated by air.  The
    units must match, and will return the same units.

    Parameters
    ----------
    f1 : float
        Focal length of lens 1
    f2 : float
        Focal length of lens 2
    d : float
        Separation distance between the two lenses

    Returns
    -------
    f : float
        The focal length of the two lens system
    """

    phi1 = 1.0/f1
    phi2 = 1.0/f2
    phi = phi1 + phi2 -phi1*phi2*d

    return 1.0/phi

def thick_lens_EFL(R1,R2,t,n):
    """Calculate the focal length of a thick lens via geometrical method,
    given the two surface radii, the center thickenss, and the index.
    The units must match, and will return the same units.

    Parameters
    ----------
    R1 : float
        Radius of surface 1
    R2 : float
        Radius of surface 2
    t : float
        Center thickenss of the lens
    n : float
        Index of refraction

    Returns
    -------
    f : float
        The focal length of the thick lens
    """
    tau = t/n
    C1 = 1.0/R1
    C2 = 1.0/R2
    phi = (n-1.0)*(C1-C2 + (n-1)*C1*C2*tau)
    efl = 1.0/phi

    return efl
    #test1 50,-50,10,1.5 matches Zemax exactly: 51.741
    #todo:  better way to convert units besides writing several if's

def thin_prism_deviation(angle, n):
    """Calculate the ray deviation caused by a thin prism given prism angle
        and index.

    Parameters
    ----------
    angle : float
        Angle of the prism (degrees or radians)
    n: float
        Index of refraction of the prism material at the wavelength of interest

    Returns
    -------
    d : float
        The ray deviation due to the prism (in units of input angle)
    """
    d = -1*(n-1)*angle
    return d

max_stars_repo_name: jfblanchard/optical-tools

max_issues_repo_name: jfblanchard/optical-tools

max_forks_repo_name: jfblanchard/optical-tools

reason: import numpy,import scipy


text:
import ipywidgets as widgets
from ipywidgets import HBox, VBox, AppLayout, Layout,GridspecLayout
from threading import Thread
import threading
import time
import numpy as np
import logging
from IPython.display import display, HTML
import visa
import os
import sys

work_dir =  os.path.abspath(__file__ + '/..')
path = os.path.abspath(work_dir + '/../')
if not path in sys.path:
    sys.path.insert(0, path)
    print(path)

from pyPowerMeter import ThorlabsP1xx

# ------------------------------------
# -- Get the detectors
# ------------------------------------
rm = visa.ResourceManager()
rm = visa.ResourceManager()
connected_instr = {}
try:
    dev = rm.list_resources()
    names = []
    for dd in dev:
        try:
            instr = rm.open_resource(dd, timeout=0.5)
            # print('-'*60)
            n = instr.query('*IDN?').strip().replace('Thorlabs,','').replace('1.3.0','').replace(',', ', ')
            names += [n]
            connected_instr[n] = dd
        except Exception as err:
            pass
except Exception as err:
    pass

# ------------------------------------
# -- Create the UI
# ------------------------------------
class _dic2struct():
    def __init__(self, d, which='sim', do_tr=True):
        self._dic = d
        for a, b in d.items():
           setattr(self, a, _dic2struct(b) if isinstance(b, dict) else b)
    def __repr__(self):
        return str(list(self._dic.keys()))
dd = {'IN':{}, 'OUT':{}}
for ii in [dd['IN'], dd['OUT']]:
    ii['P'] = widgets.Text(value='0 mW', disabled=False)
    ii['P'].add_class("ui_power")
    ii['λ']= widgets.IntSlider(description = 'λ (nm)',value=1550,min=700, max=1800,step=50,
                                     orientation = 'vertical',readout=True, layout=Layout(height='180px'))
    ii['λ'].add_class("ui_wlgth")
    ii['fact']= widgets.IntSlider(description = 'Factor',value=9,min=1, max=100,step=1,
                                     orientation = 'vertical',readout=True, layout=Layout(height='180px'))
    ii['fact'].add_class("ui_wlgth")
    ii['Prgs'] = widgets.FloatProgress(value=0.0, min=0.0, max=1.0, bar_style='info')
    ii['Prgs'].add_class("ui_prgs")
    ii['Cnct'] = widgets.ToggleButton(value=False,description='Connect', layout=Layout(width='100%'))
    ii['Cnct'].add_class("ui_cnct")
    ii['Det'] = widgets.Select(options=connected_instr.keys(),
                                     disabled=False)
    _ = ii['Det'].add_class("ui_det")
dd['IN']['label'] = widgets.Label(value="Input:")
dd['IN']['label'].add_class("ui_label_title")

dd['OUT']['label'] = widgets.Label(value="Output:")
dd['OUT']['label'].add_class("ui_label_title")

dd['Losses'] = {}
dd['Losses']['loss'] = widgets.Text(value='-Inf dB', disabled=False)
dd['Losses']['loss'].add_class("ui_power")
dd['Losses']['label_top'] = widgets.Label(value="Losses:")
dd['Losses']['label_top'].add_class("ui_label_losses")
dd['Losses']['label_bottom'] = widgets.Label(value="")
dd['Losses']['label_bottom'].add_class("ui_label_losses")

dd['running'] = widgets.Checkbox(value = True, description = "Running")
dd['running'].add_class("ui_running")
ui = _dic2struct(dd)

IN = HBox([ui.IN.λ,ui.IN.fact,
            VBox([ui.IN.Det,ui.IN.Cnct]),
            VBox([ui.IN.P,
            ui.IN.Prgs]) ])

OUT = HBox([ui.OUT.λ,ui.OUT.fact,
            VBox([ui.OUT.Det,ui.OUT.Cnct]),
            VBox([ui.OUT.P,
            ui.OUT.Prgs])])
LOSS = VBox([ui.Losses.label_top,
            ui.Losses.loss,
            ui.Losses.label_bottom])

grid = VBox([ui.running,
             HBox([VBox([ui.IN.label,
                   IN,
                   ui.OUT.label,
                   OUT]), LOSS])
            ])

stop_threads = False
output = widgets.Output

# ------------------------------------
# -- Create the Worker
# ------------------------------------

Pmeter = {'IN':None, 'OUT':None}
def work(running, ui_slot, P, Loss, in_out):
    total = 100
    maxi = 0
    need_connect = True
    while True:

        if ui_slot.Cnct.value:
            if need_connect:
                instr = connected_instr[ui_slot.Det.value]
                Pmeter[in_out] = ThorlabsP1xx(address = instr)
                Pmeter[in_out].connected = True
                Pmeter[in_out].lbd = float(ui_slot.λ.value)
                need_connect = False

            val = Pmeter[in_out].power*1e3 * ui_slot.fact.value

            if val>10:
                ui_slot.P.value = '{:.1f}'.format(val) + ' mW'
            else:
                ui_slot.P.value = '{:.3f}'.format(val) + ' mW'

            ui_slot.Prgs.value = val
            if val>maxi:
                maxi = val
                ui_slot.Prgs.max= maxi
                ui_slot.Prgs.value = val
            if in_out == 'IN':
                outP = float(P.value.replace(' mW', ''))
                inP = float(ui_slot.P.value.replace(' mW', ''))

                if outP>0:
                    Loss.value = '{:.3f} dB'.format(10*np.log10(outP/inP))
                else:
                    Loss.value = '-Inf dB'
        else:
            if not need_connect:
                Pmeter[in_out].connected = False
                need_connect = True
        if not ui.running.value:
            logging.debug('thread stop')
            break
        time.sleep(0.1)

display(HTML("<style>div.balanced-text .widget-label {max-width:200px; width:200px}"))

thread_in = threading.Thread(target=work, args=(ui.running, ui.IN, ui.OUT.P,ui.Losses.loss,  'IN'))
thread_out = threading.Thread(target=work, args=(ui.running, ui.OUT,None ,None, 'OUT'))
display(grid)

thread_in.start()
thread_out.start()

max_stars_repo_name: gregmoille/InstrumentControl

max_issues_repo_name: gregmoille/InstrumentControl

max_forks_repo_name: gregmoille/InstrumentControl


text:
import pytest

import numpy

from gisutils import utils
from .helpers import raises

@pytest.mark.parametrize(('filename', 'expected'), [
    ('example.shp', 'example_test.shp'),
    ('example', 'example_test'),
])
def test_add_suffix_to_filename(filename, expected):
    result = utils.add_suffix_to_filename(filename, 'test')
    assert result == expected

@pytest.mark.parametrize(('column', 'value', 'expected', 'err'), [
    ('ID', 'Junk', None, None),
    ('ID', 'A1', ('A1', 'Ocean', 'A1_x', 'A1_y'), None),
    ('DS_ID', 'A1', None, ValueError),
])
def test_find_row_in_array(column, value, expected, err):
    with raises(err):
        input_array = numpy.array(
            [
                ('A1', 'Ocean', 'A1_x', 'A1_y'), ('A2', 'Ocean', 'A2_x', 'A2_y'),
                ('B1', 'A1', 'None', 'B1_y'), ('B2', 'A1', 'B2_x', 'None'),
                ('B3', 'A2', 'B3_x', 'B3_y'), ('C1', 'B2', 'C1_x', 'None'),
            ], dtype=[('ID', '<U5'), ('DS_ID', '<U5'), ('Cu', '<U5'), ('Pb', '<U5')]
        )
        result = utils.find_row_in_array(input_array, column, value)
        if expected is None:
            assert expected is result
        else:
            assert list(result) == list(expected)

max_stars_repo_name: phobson/gisutils

max_issues_repo_name: phobson/gisutils

max_forks_repo_name: phobson/gisutils


text:
import functools
import typing

import numpy as np

import pandas as pd
from pandas.api.types import is_numeric_dtype

import scipy.sparse

import sklearn.utils
from sklearn import preprocessing
from sklearn.base import BaseEstimator
from sklearn.compose import ColumnTransformer
from sklearn.exceptions import NotFittedError

from autoPyTorch.data.base_feature_validator import BaseFeatureValidator, SUPPORTED_FEAT_TYPES

class TabularFeatureValidator(BaseFeatureValidator):
    def _fit(
        self,
        X: SUPPORTED_FEAT_TYPES,
    ) -> BaseEstimator:
        """
        In case input data is a pandas DataFrame, this utility encodes the user provided
        features (from categorical for example) to a numerical value that further stages
        will be able to use

        Arguments:
            X (SUPPORTED_FEAT_TYPES):
                A set of features that are going to be validated (type and dimensionality
                checks) and a encoder fitted in the case the data needs encoding
        Returns:
            self:
                The fitted base estimator
        """

        # The final output of a validator is a numpy array. But pandas
        # gives us information about the column dtype
        if isinstance(X, np.ndarray):
            X = self.numpy_array_to_pandas(X)

        if hasattr(X, "iloc") and not scipy.sparse.issparse(X):
            X = typing.cast(pd.DataFrame, X)
            # Treat a column with all instances a NaN as numerical
            # This will prevent doing encoding to a categorical column made completely
            # out of nan values -- which will trigger a fail, as encoding is not supported
            # with nan values.
            # Columns that are completely made of NaN values are provided to the pipeline
            # so that later stages decide how to handle them
            if np.any(pd.isnull(X)):
                for column in X.columns:
                    if X[column].isna().all():
                        X[column] = pd.to_numeric(X[column])
                        # Also note this change in self.dtypes
                        if len(self.dtypes) != 0:
                            self.dtypes[list(X.columns).index(column)] = X[column].dtype

            if not X.select_dtypes(include='object').empty:
                X = self.infer_objects(X)

            self.enc_columns, self.feat_type = self._get_columns_to_encode(X)

            if len(self.enc_columns) > 0:
                X = self.impute_nan_in_categories(X)

                self.encoder = ColumnTransformer(
                    [
                        ("encoder",
                         preprocessing.OrdinalEncoder(
                             handle_unknown='use_encoded_value',
                             unknown_value=-1,
                         ), self.enc_columns)],
                    remainder="passthrough"
                )

                # Mypy redefinition
                assert self.encoder is not None
                self.encoder.fit(X)

                # The column transformer reoders the feature types - we therefore need to change
                # it as well
                # This means columns are shifted to the right
                def comparator(cmp1: str, cmp2: str) -> int:
                    if (
                        cmp1 == 'categorical' and cmp2 == 'categorical'
                        or cmp1 == 'numerical' and cmp2 == 'numerical'
                    ):
                        return 0
                    elif cmp1 == 'categorical' and cmp2 == 'numerical':
                        return -1
                    elif cmp1 == 'numerical' and cmp2 == 'categorical':
                        return 1
                    else:
                        raise ValueError((cmp1, cmp2))

                self.feat_type = sorted(
                    self.feat_type,
                    key=functools.cmp_to_key(comparator)
                )

                self.categories = [
                    # We fit an ordinal encoder, where all categorical
                    # columns are shifted to the left
                    list(range(len(cat)))
                    for cat in self.encoder.transformers_[0][1].categories_
                ]

            for i, type_ in enumerate(self.feat_type):
                if 'numerical' in type_:
                    self.numerical_columns.append(i)
                else:
                    self.categorical_columns.append(i)

        # Lastly, store the number of features
        self.num_features = np.shape(X)[1]
        return self

    def transform(
        self,
        X: SUPPORTED_FEAT_TYPES,
    ) -> np.ndarray:
        """
        Validates and fit a categorical encoder (if needed) to the features.
        The supported data types are List, numpy arrays and pandas DataFrames.

        Arguments:
            X_train (SUPPORTED_FEAT_TYPES):
                A set of features, whose categorical features are going to be
                transformed

        Return:
            np.ndarray:
                The transformed array
        """
        if not self._is_fitted:
            raise NotFittedError("Cannot call transform on a validator that is not fitted")

        # If a list was provided, it will be converted to pandas
        if isinstance(X, list):
            X, _ = self.list_to_dataframe(X)

        if isinstance(X, np.ndarray):
            X = self.numpy_array_to_pandas(X)

        if hasattr(X, "iloc") and not scipy.sparse.issparse(X):
            if np.any(pd.isnull(X)):
                for column in X.columns:
                    if X[column].isna().all():
                        X[column] = pd.to_numeric(X[column])

            # Also remove the object dtype for new data
            if not X.select_dtypes(include='object').empty:
                X = self.infer_objects(X)

        # Check the data here so we catch problems on new test data
        self._check_data(X)

        # Pandas related transformations
        if hasattr(X, "iloc") and self.encoder is not None:
            if np.any(pd.isnull(X)):
                # After above check it means that if there is a NaN
                # the whole column must be NaN
                # Make sure it is numerical and let the pipeline handle it
                for column in X.columns:
                    if X[column].isna().all():
                        X[column] = pd.to_numeric(X[column])

            # We also need to fillna on the transformation
            # in case test data is provided
            X = self.impute_nan_in_categories(X)

            X = self.encoder.transform(X)

        # Sparse related transformations
        # Not all sparse format support index sorting
        if scipy.sparse.issparse(X) and hasattr(X, 'sort_indices'):
            X.sort_indices()

        try:
            X = sklearn.utils.check_array(
                X,
                force_all_finite=False,
                accept_sparse='csr'
            )
        except Exception as e:
            self.logger.exception(f"Conversion failed for input {X.dtypes} {X}"
                                  "This means AutoPyTorch was not able to properly "
                                  "Extract the dtypes of the provided input features. "
                                  "Please try to manually cast it to a supported "
                                  "numerical or categorical values.")
            raise e
        return X

    def _check_data(
        self,
        X: SUPPORTED_FEAT_TYPES,
    ) -> None:
        """
        Feature dimensionality and data type checks

        Arguments:
            X (SUPPORTED_FEAT_TYPES):
                A set of features that are going to be validated (type and dimensionality
                checks) and a encoder fitted in the case the data needs encoding
        """

        if not isinstance(X, (np.ndarray, pd.DataFrame)) and not scipy.sparse.issparse(X):
            raise ValueError("AutoPyTorch only supports Numpy arrays, Pandas DataFrames,"
                             " scipy sparse and Python Lists, yet, the provided input is"
                             " of type {}".format(type(X))
                             )

        if self.data_type is None:
            self.data_type = type(X)
        if self.data_type != type(X):
            self.logger.warning("AutoPyTorch previously received features of type %s "
                                "yet the current features have type %s. Changing the dtype "
                                "of inputs to an estimator might cause problems" % (
                                    str(self.data_type),
                                    str(type(X)),
                                ),
                                )

        # Do not support category/string numpy data. Only numbers
        if hasattr(X, "dtype"):
            if not np.issubdtype(X.dtype.type, np.number):  # type: ignore[union-attr]
                raise ValueError(
                    "When providing a numpy array to AutoPyTorch, the only valid "
                    "dtypes are numerical ones. The provided data type {} is not supported."
                    "".format(
                        X.dtype.type,  # type: ignore[union-attr]
                    )
                )

        # Then for Pandas, we do not support Nan in categorical columns
        if hasattr(X, "iloc"):
            # If entered here, we have a pandas dataframe
            X = typing.cast(pd.DataFrame, X)

            # Handle objects if possible
            if not X.select_dtypes(include='object').empty:
                X = self.infer_objects(X)

            # Define the column to be encoded here as the feature validator is fitted once
            # per estimator
            enc_columns, _ = self._get_columns_to_encode(X)

            column_order = [column for column in X.columns]
            if len(self.column_order) > 0:
                if self.column_order != column_order:
                    raise ValueError("Changing the column order of the features after fit() is "
                                     "not supported. Fit() method was called with "
                                     "{} whereas the new features have {} as type".format(self.column_order,
                                                                                column_order,)
                                     )
            else:
                self.column_order = column_order

            dtypes = [dtype.name for dtype in X.dtypes]
            if len(self.dtypes) > 0:
                if self.dtypes != dtypes:
                    raise ValueError("Changing the dtype of the features after fit() is "
                                     "not supported. Fit() method was called with "
                                     "{} whereas the new features have {} as type".format(self.dtypes,
                                                                                dtypes,
                                                                                )
                                     )
            else:
                self.dtypes = dtypes

    def _get_columns_to_encode(
        self,
        X: pd.DataFrame,
    ) -> typing.Tuple[typing.List[str], typing.List[str]]:
        """
        Return the columns to be encoded from a pandas dataframe

        Arguments:
            X (pd.DataFrame)
                A set of features that are going to be validated (type and dimensionality
                checks) and a encoder fitted in the case the data needs encoding
        Returns:
            enc_columns (List[str]):
                Columns to encode, if any
            feat_type:
                Type of each column numerical/categorical
        """
        # Register if a column needs encoding
        enc_columns = []

        # Also, register the feature types for the estimator
        feat_type = []

        # Make sure each column is a valid type
        for i, column in enumerate(X.columns):
            if X[column].dtype.name in ['category', 'bool']:

                enc_columns.append(column)
                feat_type.append('categorical')
            # Move away from np.issubdtype as it causes
            # TypeError: data type not understood in certain pandas types
            elif not is_numeric_dtype(X[column]):
                if X[column].dtype.name == 'object':
                    raise ValueError(
                        "Input Column {} has invalid type object. "
                        "Cast it to a valid dtype before using it in AutoPyTorch. "
                        "Valid types are numerical, categorical or boolean. "
                        "You can cast it to a valid dtype using "
                        "pandas.Series.astype ."
                        "If working with string objects, the following "
                        "tutorial illustrates how to work with text data: "
                        "https://scikit-learn.org/stable/tutorial/text_analytics/working_with_text_data.html".format(
                            # noqa: E501
                            column,
                        )
                    )
                elif pd.core.dtypes.common.is_datetime_or_timedelta_dtype(
                    X[column].dtype
                ):
                    raise ValueError(
                        "AutoPyTorch does not support time and/or date datatype as given "
                        "in column {}. Please convert the time information to a numerical value "
                        "first. One example on how to do this can be found on "
                        "https://stats.stackexchange.com/questions/311494/".format(
                            column,
                        )
                    )
                else:
                    raise ValueError(
                        "Input Column {} has unsupported dtype {}. "
                        "Supported column types are categorical/bool/numerical dtypes. "
                        "Make sure your data is formatted in a correct way, "
                        "before feeding it to AutoPyTorch.".format(
                            column,
                            X[column].dtype.name,
                        )
                    )
            else:
                feat_type.append('numerical')
        return enc_columns, feat_type

    def list_to_dataframe(
        self,
        X_train: SUPPORTED_FEAT_TYPES,
        X_test: typing.Optional[SUPPORTED_FEAT_TYPES] = None,
    ) -> typing.Tuple[pd.DataFrame, typing.Optional[pd.DataFrame]]:
        """
        Converts a list to a pandas DataFrame. In this process, column types are inferred.

        If test data is provided, we proactively match it to train data

        Arguments:
            X_train (SUPPORTED_FEAT_TYPES):
                A set of features that are going to be validated (type and dimensionality
                checks) and a encoder fitted in the case the data needs encoding
            X_test (typing.Optional[SUPPORTED_FEAT_TYPES]):
                A hold out set of data used for checking
        Returns:
            pd.DataFrame:
                transformed train data from list to pandas DataFrame
            pd.DataFrame:
                transformed test data from list to pandas DataFrame
        """

        # If a list was provided, it will be converted to pandas
        X_train = pd.DataFrame(data=X_train).infer_objects()
        self.logger.warning("The provided feature types to AutoPyTorch are of type list."
                            "Features have been interpreted as: {}".format([(col, t) for col, t in
                                                                            zip(X_train.columns, X_train.dtypes)]))
        if X_test is not None:
            if not isinstance(X_test, list):
                self.logger.warning("Train features are a list while the provided test data"
                                    "is {}. X_test will be casted as DataFrame.".format(type(X_test))
                                    )
            X_test = pd.DataFrame(data=X_test).infer_objects()
        return X_train, X_test

    def numpy_array_to_pandas(
        self,
        X: np.ndarray,
    ) -> pd.DataFrame:
        """
        Converts a numpy array to pandas for type inference

        Arguments:
            X (np.ndarray):
                data to be interpreted.

        Returns:
            pd.DataFrame
        """
        return pd.DataFrame(X).infer_objects().convert_dtypes()

    def infer_objects(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        In case the input contains object columns, their type is inferred if possible

        This has to be done once, so the test and train data are treated equally

        Arguments:
            X (pd.DataFrame):
                data to be interpreted.

        Returns:
            pd.DataFrame
        """
        if hasattr(self, 'object_dtype_mapping'):
            # Mypy does not process the has attr. This dict is defined below
            for key, dtype in self.object_dtype_mapping.items():  # type: ignore[has-type]
                if 'int' in dtype.name:
                    # In the case train data was interpreted as int
                    # and test data was interpreted as float, because of 0.0
                    # for example, honor training data
                    X[key] = X[key].applymap(np.int64)
                else:
                    try:
                        X[key] = X[key].astype(dtype.name)
                    except Exception as e:
                        # Try inference if possible
                        self.logger.warning(f"Tried to cast column {key} to {dtype} caused {e}")
                        pass
        else:
            X = X.infer_objects()
            for column in X.columns:
                if not is_numeric_dtype(X[column]):
                    X[column] = X[column].astype('category')
            self.object_dtype_mapping = {column: X[column].dtype for column in X.columns}
        self.logger.debug(f"Infer Objects: {self.object_dtype_mapping}")
        return X

    def impute_nan_in_categories(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        impute missing values before encoding,
        remove once sklearn natively supports
        it in ordinal encoding. Sklearn issue:
        "https://github.com/scikit-learn/scikit-learn/issues/17123)"

        Arguments:
            X (pd.DataFrame):
                data to be interpreted.

        Returns:
            pd.DataFrame
        """

        # To be on the safe side, map always to the same missing
        # value per column
        if not hasattr(self, 'dict_nancol_to_missing'):
            self.dict_missing_value_per_col: typing.Dict[str, typing.Any] = {}

        # First make sure that we do not alter the type of the column which cause:
        # TypeError: '<' not supported between instances of 'int' and 'str'
        # in the encoding
        for column in self.enc_columns:
            if X[column].isna().any():
                if column not in self.dict_missing_value_per_col:
                    try:
                        float(X[column].dropna().values[0])
                        can_cast_as_number = True
                    except Exception:
                        can_cast_as_number = False
                    if can_cast_as_number:
                        # In this case, we expect to have a number as category
                        # it might be string, but its value represent a number
                        missing_value: typing.Union[str, int] = '-1' if isinstance(X[column].dropna().values[0],
                                                                                str) else -1
                    else:
                        missing_value = 'Missing!'

                    # Make sure this missing value is not seen before
                    # Do this check for categorical columns
                    # else modify the value
                    if hasattr(X[column], 'cat'):
                        while missing_value in X[column].cat.categories:
                            if isinstance(missing_value, str):
                                missing_value += '0'
                            else:
                                missing_value += missing_value
                    self.dict_missing_value_per_col[column] = missing_value

                # Convert the frame in place
                X[column].cat.add_categories([self.dict_missing_value_per_col[column]],
                                             inplace=True)
                X.fillna({column: self.dict_missing_value_per_col[column]}, inplace=True)
        return X

max_stars_repo_name: franchuterivera/Auto-PyTorch

max_issues_repo_name: franchuterivera/Auto-PyTorch

max_forks_repo_name: franchuterivera/Auto-PyTorch


text:
import torch
from game import Game
from constants import *
from joblib import Parallel, delayed
import numpy as np

evaluators = [Game(None, mctsEnable=True) for c in range(NUM_CORES)]

def getRes(evaluator, opf):
	wins = 0
	for i in range(int(EVAL_GAMES/NUM_CORES)):
		wins += evaluator.play(opFirst = opf)
	return wins

def evaluateAndSave(model):
	best_model = torch.load(BEST_PATH).to(DEVICE)
	wins = 0
	for evaluator in evaluators:
		evaluator.player = model
		evaluator.opponent = best_model
	results = Parallel(n_jobs=NUM_CORES)(delayed(getRes)(evaluators[i], i<NUM_CORES/2) for i in range(NUM_CORES))
	print(results)
	wins = np.sum(results)
	updateModel = wins >= EVAL_THRESH * EVAL_GAMES
	if updateModel:
		print("Current model = best")
		torch.save(model, BEST_PATH)
		return model
	print("Current model is not best")
	return best_model

max_stars_repo_name: shreshthtuli/AlphaGoZero

max_issues_repo_name: shreshthtuli/AlphaGoZero

max_forks_repo_name: shreshthtuli/AlphaGoZero


text:
# -*- coding: utf-8 -*-
"""
Created on Thu May 13 09:01:13 2021

@author: thka
"""

# %% Define Data Pool
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
pd.options.mode.chained_assignment = None

# Data base as pandas dataframe
class DataPool:
    """Class holding a DB (based on a Pandas DataFrame)"""

    def __init__(self, name, monthly=False):
        """ Constructor of a DB (based on a Pandas DataFrame) """
        self.name = name
        self.df = pd.DataFrame()
        self.ddiff = pd.DataFrame()
        self.dfeat = pd.DataFrame()
        self.monthly = monthly

    def put(self, data):
        """ Add a new entry to the DB """
        if self.df.empty:
            self.df = pd.DataFrame([data])
        else:
            self.df = self.df.append(data, ignore_index=True)

    def isEmpty(self):
        """ Check if DB is empty """
        return self.df.empty

    # Save DB as CSV
    def saveAsCSV(self, ver='00'):
        """ Save DB as a CSV file """
        self.df.to_csv(self.name + '_' + ver + '.csv')

    # Load DB from CSV
    def loadAsCSV(self, ver='00'):
        """ Load DB from a CSV file """
        # print('LOADING CSV FILE: ' + self.name)
        try:
            fname = self.name + '_' + ver + '.csv'
            self.df = pd.read_csv(fname, index_col=[0])
        except FileNotFoundError:
            print('No file: ' + fname)

    def clean_data(self):
        """ Create the silver buckets from RAW bronze data """
        # Sort by if and day
        self.df = self.df.sort_values(by=['id', 'power_cycle'])
        # Remove duplicates
        self.df.drop_duplicates(subset=['id', 'power_cycle'], inplace=True)

    def diff_data(self):
        """ Differentiate counters"""
        if self.ddiff.empty:
            # Clean data first
            self.clean_data()

            # Which columns to differentiate
            k_list = [x for x in list(self.df.keys())
                      if x not in ['id', 'power_cycle', 'time',
                                   'date', 'usage-at-time']]
            # Get list of ID's
            id_list = list(self.df['id'].unique())
            # Diff the relevant data
            self.ddiff = self.df.copy()
            self.ddiff.drop(
                self.ddiff.loc[self.ddiff['power_cycle'] == 0].index,
                inplace=True)

            # Loop over all ID's
            for id in id_list:
                self.ddiff.loc[self.ddiff['id'] == id, k_list] = \
                    self.df.loc[self.df['id'] == id, k_list]\
                        .diff().iloc[1:].astype(int)

    # Normalize dats
    def normalize_data(self):
        """ Normalize Data"""
        # Need diff'd data
        self.diff_data()
        # Check if already normalized
        if 'Usage' not in self.ddiff.keys():
            self.ddiff['Usage'] = self.ddiff['usage'] / 3600
            self.ddiff['Charge'] = self.ddiff['charge'] / 3600
            self.ddiff['Speech'] = 100.0 * self.ddiff['speech'] \
                / self.ddiff['usage']
            self.ddiff['OwnVoice'] = 100.0 * self.ddiff['ovd'] \
                / self.ddiff['usage']
            # OVD pr SNR?
            if 'ovd-snr-low' in self.ddiff.keys():
                self.ddiff['OVD-snr-low'] = 100.0 * self.ddiff['ovd-snr-low'] \
                    / self.ddiff['usage']
                self.ddiff['OVD-snr-med'] = 100.0 * self.ddiff['ovd-snr-med'] \
                    / self.ddiff['usage']
                self.ddiff['OVD-snr-high'] = 100.0 * self.ddiff['ovd-snr-high'] \
                    / self.ddiff['usage']

    # Prepare Plot data
    def prepare_plot_data(self):
        """ Prepare data for plotting """
        # Need normalized data
        self.normalize_data()
        # Prepare coloring of usage data based on threshold
        if 'Usage Low' not in self.ddiff.keys():
            self.ddiff['Usage Low'] = self.ddiff['Usage']
            self.ddiff['Usage OK'] = self.ddiff['Usage']
            self.ddiff['Threshold'] = np.ones(self.ddiff.shape[0]) * 5
            self.ddiff.loc[self.ddiff['Usage'] >= 5.0, 'Usage Low'] = 0.0
            self.ddiff.loc[self.ddiff['Usage'] < 5.0, 'Usage OK'] = 0.0

    # Plot daily data
    def plot_daily(self, user_id=0, days=31, last_day=-1):
        """ Plot most important Daily data """
        # Prepare coloring of usage data based on threshold
        self.prepare_plot_data()
        # Narrow scope for plot
        dp = self.ddiff[self.ddiff['id'] == user_id].iloc[last_day
                                                          - days:last_day]
        self.plot_data(dp, user_id)

    # Plot monthly data
    def plot_monthly(self, user_id=0, month=30):
        """ Plot most important Monthly data """
        # Prepare coloring of usage data based on threshold
        self.prepare_plot_data()
        # Narrow scope for plot
        dp = self.ddiff[self.ddiff['id'] == user_id].iloc[::month]

        self.plot_data(dp, user_id)

    # Plot data
    def plot_data(self, dp, user_id=0):
        # Check for date
        x_col = 'date'
        x_label = 'Date'
        if x_col not in dp.keys():
            x_col = 'power_cycle'
            x_label = 'Sessions'

        # Plot Voice Overview pr Day
        ax = dp.plot.bar(x=x_col, y=['OwnVoice', 'Speech'],
                         stacked=True,
                         color=['gold', 'darkgoldenrod'],
                         title='Speech Overview for HI #'
                         + str(user_id)
                         + ', Day view')
        ax.set_xlabel(x_label)
        ax.set_ylabel('Percentage (%)')

        # Plot OVD pr SNR pr Day
        if 'ovd-snr-low' in dp.keys():
            ax = dp.plot.bar(x=x_col,
                             y=['OVD-snr-low', 'OVD-snr-med', 'OVD-snr-high'],
                             stacked=True,
                             color=['orangered', 'orange', 'greenyellow'],
                             title='OVR pr SNR Overview for HI #'
                             + str(user_id) + ', Day view')
            ax.set_xlabel(x_label)
            ax.set_ylabel('Percentage (%)')

        # Plot Usage Overview pr Day
        ax = dp.plot.bar(x=x_col,
                         y=['Usage Low', 'Usage OK', 'Charge'],
                         stacked=True,
                         color=['red', 'limegreen', 'steelblue'],
                         title='Usage Overview for HI #' + str(user_id)
                         + ', Day view')
        dp.plot(x=x_col, y='Threshold', ax=ax, linestyle='dotted',
                color='black')
        ax.legend(loc="upper right")
        ax.set_xlabel(x_label)
        ax.set_ylabel('Hours')
        plt.xticks(rotation=90)

    # Create features for analytics
    def create_features(self):
        """ Create DataFrame for holding features"""
        self.dfeat = pd.DataFrame(data=sorted(self.df['id'].unique()),
                                  columns=['user_id'])

        # Prepare features
        features = ['usage-pr-day', 'ovd-inc', 'speech-inc', 'ovd-snr-low-inc']
        zeros = np.zeros((len(self.dfeat), len(features)))
        self.dfeat = self.dfeat.join(pd.DataFrame(data=zeros,
                                                  columns=features))
        self.dfeat.set_index('user_id', inplace=True)

        # Extract features for each client
        secPyear = 365 * 3600
        for i, row in self.dfeat.iterrows():
            di = self.df.loc[self.df['id'] == i]
            row['usage-pr-day'] = di.iloc[-1]['usage'] / secPyear
            # row['charge'] = di.iloc[-1]['charge'] / secPyear
            row['ovd-inc'] = \
                (di.iloc[-1]['ovd'] - di.iloc[0]['ovd']) / secPyear
            row['speech-inc'] = \
                (di.iloc[-1]['speech'] - di.iloc[0]['speech']) / secPyear
            row['ovd-snr-low-inc'] = \
                (di.iloc[-1]['ovd-snr-low'] - di.iloc[0]['ovd-snr-low'])\
                / secPyear

    # Plot features
    def plot_features(self):
        """ Plot features"""
        # Create features if not already done
        if self.dfeat.empty:
            self.create_features()
        # Create pairplot to analyze correlations
        sns.color_palette("tab10")
        sns.pairplot(self.dfeat, hue="ovd-inc")

# Class holding all databases
class CDP:
    """Class holding a number of DBs - Common Data Platform """

    def __init__(self):
        """ CDP Constructor """
        self.app_daily = DataPool('app_daily')
        self.app_hourly = DataPool('app_hourly')
        self.fsw_daily = DataPool('fsw_daily')
        self.fsw_monthly = DataPool('fsw_monthly')

    # Save all DB as CSV
    def saveAsCSV(self, ver='00'):
        """ Save all DB's as CSV files"""
        self.app_daily.saveAsCSV(ver)
        self.app_hourly.saveAsCSV(ver)
        self.fsw_daily.saveAsCSV(ver)
        self.fsw_monthly.saveAsCSV(ver)

    # Load all DB from CSV
    def loadAsCSV(self, ver='00'):
        """ Load all DB's from CSV files"""
        self.app_daily.loadAsCSV(ver)
        self.app_hourly.loadAsCSV(ver)
        self.fsw_daily.saveAsCSV(ver)
        self.fsw_monthly.loadAsCSV(ver)

    def getAppDaily(self):
        """ Get Daily App DB """
        return self.app_daily

    def getAppHourly(self):
        """ Get Hourly App DB """
        return self.app_hourly

    def getFswDaily(self):
        """ Get Daily FSW DB """
        return self.fsw_daily

    def getFswMonthly(self):
        """ Get Monthly FSW DB """
        return self.fsw_monthly

max_stars_repo_name: tomkaul/logsim

max_issues_repo_name: tomkaul/logsim

max_forks_repo_name: tomkaul/logsim


text:
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import uuid
from typing import Callable, Dict, List, NoReturn, Optional, Iterable, Union

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import pyspark.sql as sql
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import StructType
from pyspark.sql.pandas.types import from_arrow_type
import ray
from ray.data import Dataset, from_arrow_refs
from ray.types import ObjectRef
import ray.util.data as ml_dataset
import ray.util.iter as parallel_it
from ray.util.data import MLDataset
from ray.util.data.interface import _SourceShard
from ray._private.client_mode_hook import client_mode_wrap
from raydp.spark.parallel_iterator_worker import ParallelIteratorWorkerWithLen
from raydp.utils import divide_blocks

logger = logging.getLogger(__name__)

class RecordPiece:
    def __init__(self, row_ids, num_rows: int):
        self.row_ids = row_ids
        self.num_rows = num_rows

    def read(self, shuffle: bool) -> pd.DataFrame:
        raise NotImplementedError

    def with_row_ids(self, new_row_ids) -> "RecordPiece":
        raise NotImplementedError

    def __len__(self):
        """Return the number of rows"""
        return self.num_rows

class RayObjectPiece(RecordPiece):
    def __init__(self,
                 obj_id: ray.ObjectRef,
                 row_ids: Optional[List[int]],
                 num_rows: int):
        super().__init__(row_ids, num_rows)
        self.obj_id = obj_id

    def read(self, shuffle: bool) -> pd.DataFrame:
        data = ray.get(self.obj_id)
        reader = pa.ipc.open_stream(data)
        tb = reader.read_all()
        df: pd.DataFrame = tb.to_pandas()
        if self.row_ids:
            df = df.loc[self.row_ids]

        if shuffle:
            df = df.sample(frac=1.0)
        return df

    def with_row_ids(self, new_row_ids) -> "RayObjectPiece":
        """chang the num_rows to the length of new_row_ids. Keep the original size if
        the new_row_ids is None.
        """

        if new_row_ids:
            num_rows = len(new_row_ids)
        else:
            num_rows = self.num_rows

        return RayObjectPiece(self.obj_id, new_row_ids, num_rows)

class ParquetPiece(RecordPiece):
    def __init__(self,
                 piece: pq.ParquetDatasetPiece,
                 columns: List[str],
                 partitions: pq.ParquetPartitions,
                 row_ids: Optional[List[int]],
                 num_rows: int):
        super().__init__(row_ids, num_rows)
        self.piece = piece
        self.columns = columns
        self.partitions = partitions

    def read(self, shuffle: bool) -> pd.DataFrame:
        pdf = self.piece.read(columns=self.columns,
                              use_threads=False,
                              partitions=self.partitions).to_pandas()
        if self.row_ids:
            pdf = pdf.loc[self.row_ids]

        if shuffle:
            pdf = pdf.sample(frac=1.0)
        return pdf

    def with_row_ids(self, new_row_ids) -> "ParquetPiece":
        """
        chang the num_rows to the length of new_row_ids. Keep the original size if
        the new_row_ids is None.
        """
        if new_row_ids:
            num_rows = len(new_row_ids)
        else:
            num_rows = self.num_rows
        return ParquetPiece(self.piece, self.columns, self.partitions, new_row_ids, num_rows)

class RecordBatch(_SourceShard):
    def __init__(self,
                 shard_id: int,
                 prefix: str,
                 record_pieces: List[RecordPiece],
                 shuffle: bool,
                 shuffle_seed: int):
        self._shard_id = shard_id
        self._prefix = prefix
        self.record_pieces = record_pieces
        self.shuffle = shuffle
        self.shuffle_seed = shuffle_seed

    def prefix(self) -> str:
        return self._prefix

    @property
    def shard_id(self) -> int:
        return self._shard_id

    def __iter__(self) -> Iterable[pd.DataFrame]:
        if self.shuffle:
            np.random.seed(self.shuffle_seed)
            np.random.shuffle(self.record_pieces)

        for piece in self.record_pieces:
            yield piece.read(self.shuffle)

    def __len__(self):
        return sum([len(piece) for piece in self.record_pieces])

class RayRecordBatch(RecordBatch):
    def __init__(self,
                 shard_id: int,
                 prefix: str,
                 record_pieces: List[RecordPiece],
                 shuffle: bool,
                 shuffle_seed: int):
        super().__init__(shard_id, prefix, record_pieces, shuffle, shuffle_seed)
        self.resolved: bool = False

    def resolve(self, timeout: Optional[float] = None) -> NoReturn:
        """
        This is just fetch object from remote object store to local and without deserialization.
        :param timeout: The maximum amount of time in seconds to wait before returning.
        """
        if self.resolved:
            return

        worker = ray.worker.global_worker
        worker.check_connected()
        timeout_ms = int(timeout * 1000) if timeout else -1
        object_ids = [record.obj_id for record in self.record_pieces]
        worker.core_worker.get_objects(object_ids, worker.current_task_id, timeout_ms)
        self.resolved = True

@client_mode_wrap
def _register_objects(records):
    worker = ray.worker.global_worker
    blocks: List[ray.ObjectRef] = []
    block_sizes: List[int] = []
    for obj_id, owner, num_record in records:
        object_ref = ray.ObjectRef(obj_id)
        # Register the ownership of the ObjectRef
        worker.core_worker.deserialize_and_register_object_ref(
            object_ref.binary(), ray.ObjectRef.nil(), owner, "")
        blocks.append(object_ref)
        block_sizes.append(num_record)
    return blocks, block_sizes

def _save_spark_df_to_object_store(df: sql.DataFrame, use_batch: bool = True,
                                   _use_owner: bool = False):
    # call java function from python
    jvm = df.sql_ctx.sparkSession.sparkContext._jvm
    jdf = df._jdf
    object_store_writer = jvm.org.apache.spark.sql.raydp.ObjectStoreWriter(jdf)

    if _use_owner is True:
        records = object_store_writer.save(use_batch, RAYDP_OBJ_HOLDER_NAME)
    else:
        records = object_store_writer.save(use_batch, "")

    record_tuples = [(record.objectId(), record.ownerAddress(), record.numRecords())
                        for record in records]
    blocks, block_sizes = _register_objects(record_tuples)

    if _use_owner is True:
        holder = ray.get_actor(RAYDP_OBJ_HOLDER_NAME)
        df_id = uuid.uuid4()
        ray.get(holder.add_objects.remote(df_id, blocks))

    return blocks, block_sizes

def _create_ml_dataset(name: str,
                       record_pieces: List[RecordPiece],
                       record_sizes: List[int],
                       num_shards: int,
                       shuffle: bool,
                       shuffle_seed: int,
                       RecordBatchCls,
                       node_hints: List[str] = None) -> MLDataset:
    if node_hints is not None:
        assert num_shards % len(node_hints) == 0,\
            f"num_shards: {num_shards} should be a multiple of length of node_hints: {node_hints}"
    if shuffle_seed:
        np.random.seed(shuffle_seed)
    else:
        np.random.seed(0)

    # split the piece into num_shards partitions
    divided_blocks = divide_blocks(blocks=record_sizes,
                                   world_size=num_shards,
                                   shuffle=shuffle,
                                   shuffle_seed=shuffle_seed)

    record_batches = []

    for rank, blocks in divided_blocks.items():
        pieces = []
        for index, num_samples in blocks:
            record_size = record_sizes[index]
            piece = record_pieces[index]
            if num_samples != record_size:
                assert num_samples < record_size
                new_row_ids = np.random.choice(
                    record_size, size=num_samples).tolist()
                piece = piece.with_row_ids(new_row_ids)
            pieces.append(piece)

        if shuffle:
            np.random.shuffle(pieces)
        record_batches.append(RecordBatchCls(shard_id=rank,
                                             prefix=name,
                                             record_pieces=pieces,
                                             shuffle=shuffle,
                                             shuffle_seed=shuffle_seed))

    worker_cls = ray.remote(ParallelIteratorWorkerWithLen)
    if node_hints is not None:
        actors = []
        multiplier = num_shards // len(node_hints)
        resource_keys = [f"node:{node_hints[i // multiplier]}" for i in range(num_shards)]
        for g, resource_key in zip(record_batches, resource_keys):
            actor = worker_cls.options(resources={resource_key: 0.01}).remote(g, False, len(g))
            actors.append(actor)
    else:
        worker_cls = ray.remote(ParallelIteratorWorkerWithLen)
        actors = [worker_cls.remote(g, False, len(g)) for g in record_batches]

    it = parallel_it.from_actors(actors, name)
    ds = ml_dataset.from_parallel_iter(
        it, need_convert=False, batch_size=0, repeated=False)
    return ds

class RayMLDataset:
    @staticmethod
    def from_spark(df: sql.DataFrame,
                   num_shards: int,
                   shuffle: bool = True,
                   shuffle_seed: int = None,
                   fs_directory: Optional[str] = None,
                   compression: Optional[str] = None,
                   node_hints: List[str] = None) -> MLDataset:
        """ Create a MLDataset from Spark DataFrame

        This method will create a MLDataset from Spark DataFrame.

        :param df: the pyspark.sql.DataFrame
        :param num_shards: the number of shards will be created for the MLDataset
        :param shuffle: whether need to shuffle the blocks when create the MLDataset
        :param shuffle_seed: the shuffle seed, default is 0
        :param fs_directory: an optional distributed file system directory for cache the
            DataFrame. We will write the DataFrame to the given directory with parquet
            format if this is provided. Otherwise, we will write the DataFrame to ray
            object store.
        :param compression: the optional compression for write the DataFrame as parquet
            file. This is only useful when the fs_directory set.
        :param node_hints: the node hints to create MLDataset actors
        :return: a MLDataset
        """
        df = df.repartition(num_shards)
        if fs_directory is None:
            # fs_directory has not provided, we save the Spark DataFrame to ray object store
            blocks, block_sizes = _save_spark_df_to_object_store(df)
            record_pieces = [RayObjectPiece(obj, None, num_rows)
                             for obj, num_rows in zip(blocks, block_sizes)]

            return _create_ml_dataset("from_spark", record_pieces, block_sizes, num_shards,
                                      shuffle, shuffle_seed, RayRecordBatch,
                                      node_hints)
        else:
            # fs_directory has provided, we write the Spark DataFrame as Parquet files
            df.write.parquet(fs_directory, compression=compression)
            # create the MLDataset from the parquet file
            ds = RayMLDataset.from_parquet(
                fs_directory, num_shards, shuffle, shuffle_seed, node_hints)
            return ds

    @staticmethod
    def from_parquet(paths: Union[str, List[str]],
                     num_shards: int,
                     shuffle: bool = True,
                     shuffle_seed: int = None,
                     columns: Optional[List[str]] = None,
                     node_hints: List[str] = None,
                     extra_parquet_arguments: Dict = None) -> MLDataset:
        """ Create a MLDataset from Parquet files.

        :param paths: the parquet files path
        :param num_shards: the number of shards will be created for the MLDataset
        :param shuffle: whether need to shuffle the blocks when create the MLDataset
        :param shuffle_seed: the shuffle seed, default is 0
        :param columns: the columns that need to read
        :param node_hints: the node hints to create MLDataset actors
        :param extra_parquet_arguments: the extra arguments need to pass into the parquet file
            reading
        :return: a MLDataset
        """
        if not extra_parquet_arguments:
            extra_parquet_arguments = {}
        ds = pq.ParquetDataset(paths, **extra_parquet_arguments)
        pieces = ds.pieces
        record_pieces = []
        record_sizes = []

        for piece in pieces:
            meta_data = piece.get_metadata().to_dict()
            num_row_groups = meta_data["num_row_groups"]
            row_groups = meta_data["row_groups"]
            for i in range(num_row_groups):
                num_rows = row_groups[i]["num_rows"]
                parquet_ds_piece = pq.ParquetDatasetPiece(piece.path, piece.open_file_func,
                                                          piece.file_options, i,
                                                          piece.partition_keys)
                # row_ids will be set later
                record_pieces.append(ParquetPiece(piece=parquet_ds_piece,
                                                  columns=columns,
                                                  partitions=ds.partitions,
                                                  row_ids=None,
                                                  num_rows=num_rows))
                record_sizes.append(num_rows)

        return _create_ml_dataset("from_parquet", record_pieces, record_sizes, num_shards,
                                  shuffle, shuffle_seed, RecordBatch, node_hints)

    @staticmethod
    def to_torch(
            ds: MLDataset,
            world_size: int,
            world_rank: int,
            batch_size: int,
            collate_fn: Callable,
            shuffle: bool = False,
            shuffle_seed: int = None,
            local_rank: int = -1,
            prefer_node: str = None,
            prefetch: bool = False):
        """
        Create DataLoader from a MLDataset
        :param ds: the MLDataset
        :param world_size: the world_size of distributed model training
        :param world_rank: create the DataLoader for the given world_rank
        :param batch_size: the batch_size of the DtaLoader
        :param collate_fn: the collate_fn that create tensors from a pandas DataFrame
        :param shuffle: whether shuffle each batch of data
        :param shuffle_seed: the shuffle seed
        :param local_rank: the node local rank. It must be provided if prefer_node is
               not None.
        :param prefer_node: the prefer node for create the MLDataset actor
        :param prefetch: prefetch the data of DataLoader with one thread
        :return: a pytorch DataLoader
        """
        # pylint: disable=C0415
        import torch
        from raydp.torch.torch_ml_dataset import PrefetchedDataLoader, TorchMLDataset

        num_shards = ds.num_shards()
        assert num_shards % world_size == 0, \
            (f"The number shards of MLDataset({ds}) should be a multiple of "
             f"world_size({world_size})")
        multiplier = num_shards // world_size

        selected_ds = None
        if prefer_node is not None:
            assert 0 <= local_rank < world_size

            # get all actors
            # there should be only one actor_set because of select_shards() is not allowed
            # after union()

            def location_check(actor):
                address = ray.actors(actor._actor_id.hex())["Address"]["IPAddress"]
                return address == prefer_node

            actors = ds.actor_sets[0].actors
            actor_indexes = [i for i, actor in enumerate(actors) if location_check(actor)]
            if len(actor_indexes) % multiplier != 0:
                selected_ds = None
                logger.warning(f"We could not find enough shard actor in prefer "
                               f"node({prefer_node}), fail back to normal select_shards(). "
                               f"Found: ({actor_indexes}) which length is not multiple of "
                               f"num_shards({num_shards}) // world_size({world_size}).")
            else:
                shard_ids = actor_indexes[local_rank: local_rank + multiplier]
                selected_ds = ds.select_shards(shard_ids)

        if selected_ds is None:
            shard_ids = []
            i = world_rank
            step = world_size
            while i < num_shards:
                shard_ids.append(i)
                i += step
            selected_ds = ds.select_shards(shard_ids)

        selected_ds = selected_ds.batch(batch_size)
        torch_ds = TorchMLDataset(selected_ds, collate_fn, shuffle, shuffle_seed)
        data_loader = torch.utils.data.DataLoader(dataset=torch_ds,
                                                  batch_size=None,
                                                  batch_sampler=None,
                                                  shuffle=False,
                                                  num_workers=0,
                                                  collate_fn=None,
                                                  pin_memory=False,
                                                  drop_last=False,
                                                  sampler=None)
        if prefetch:
            data_loader = PrefetchedDataLoader(data_loader)
        return data_loader

def create_ml_dataset_from_spark(df: sql.DataFrame,
                                 num_shards: int,
                                 shuffle: bool,
                                 shuffle_seed: int,
                                 fs_directory: Optional[str] = None,
                                 compression: Optional[str] = None,
                                 node_hints: List[str] = None) -> MLDataset:
    return RayMLDataset.from_spark(
        df, num_shards, shuffle, shuffle_seed, fs_directory, compression, node_hints)

def spark_dataframe_to_ray_dataset(df: sql.DataFrame,
                                   parallelism: Optional[int] = None,
                                   _use_owner: bool = False):
    num_part = df.rdd.getNumPartitions()
    if parallelism is not None:
        if parallelism != num_part:
            df = df.repartition(parallelism)
    blocks, _ = _save_spark_df_to_object_store(df, False, _use_owner)
    return from_arrow_refs(blocks)

@ray.remote
class RayDPConversionHelper():
    def __init__(self):
        self.objects = {}
        self.this_actor_id = None

    def add_objects(self, timestamp, objects):
        self.objects[timestamp] = objects

    def get_object(self, timestamp, idx):
        return self.objects[timestamp][idx]

    def get_ray_address(self):
        return ray.worker.global_worker.node.address

    def terminate(self):
        ray.actor.exit_actor()

    def get_actor_id(self):
        self.this_actor_id = ray.get_runtime_context().actor_id
        return self.this_actor_id

RAYDP_OBJ_HOLDER_NAME = "raydp_obj_holder"

def _convert_by_udf(spark: sql.SparkSession,
                    blocks: List[ObjectRef],
                    locations: List[bytes],
                    schema: StructType) -> DataFrame:
    holder = ray.get_actor(RAYDP_OBJ_HOLDER_NAME)
    df_id = uuid.uuid4()
    ray.get(holder.add_objects.remote(df_id, blocks))
    jvm = spark.sparkContext._jvm
    object_store_reader = jvm.org.apache.spark.sql.raydp.ObjectStoreReader
    # create the rdd then dataframe to utilize locality
    jdf = object_store_reader.createRayObjectRefDF(spark._jsparkSession, locations)
    current_namespace = ray.get_runtime_context().namespace
    ray_address = ray.get(holder.get_ray_address.remote())
    blocks_df = DataFrame(jdf, spark._wrapped)
    def _convert_blocks_to_dataframe(blocks):
        # connect to ray
        if not ray.is_initialized():
            ray.init(address=ray_address,
                     namespace=current_namespace,
                     logging_level=logging.WARN)
        obj_holder = ray.get_actor(RAYDP_OBJ_HOLDER_NAME)
        for block in blocks:
            dfs = []
            for idx in block["idx"]:
                ref = ray.get(obj_holder.get_object.remote(df_id, idx))
                data = ray.get(ref)
                dfs.append(data.to_pandas())
            yield pd.concat(dfs)
    df = blocks_df.mapInPandas(_convert_blocks_to_dataframe, schema)
    return df

def _convert_by_rdd(spark: sql.SparkSession,
                    blocks: Dataset,
                    locations: List[bytes],
                    schema: StructType) -> DataFrame:
    object_ids = [block.binary() for block in blocks]
    schema_str = schema.json()
    jvm = spark.sparkContext._jvm
    # create rdd in java
    rdd = jvm.org.apache.spark.rdd.RayDatasetRDD(spark._jsc, object_ids, locations)
    # convert the rdd to dataframe
    object_store_reader = jvm.org.apache.spark.sql.raydp.ObjectStoreReader
    jdf = object_store_reader.RayDatasetToDataFrame(spark._jsparkSession, rdd, schema_str)
    return DataFrame(jdf, spark._wrapped)

@client_mode_wrap
def get_locations(blocks):
    core_worker = ray.worker.global_worker.core_worker
    return [
        core_worker.get_owner_address(block)
        for block in blocks
    ]

def ray_dataset_to_spark_dataframe(spark: sql.SparkSession,
                                   arrow_schema: "pa.lib.Schema",
                                   blocks: List[ObjectRef],
                                   locations = None) -> DataFrame:
    locations = get_locations(blocks)
    if not isinstance(arrow_schema, pa.lib.Schema):
        raise RuntimeError(f"Schema is {type(arrow_schema)}, required pyarrow.lib.Schema. \n" \
                           f"to_spark does not support converting non-arrow ray datasets.")
    schema = StructType()
    for field in arrow_schema:
        schema.add(field.name, from_arrow_type(field.type), nullable=field.nullable)
    #TODO how to branch on type of block?
    sample = ray.get(blocks[0])
    if isinstance(sample, bytes):
        return _convert_by_rdd(spark, blocks, locations, schema)
    elif isinstance(sample, pa.Table):
        return _convert_by_udf(spark, blocks, locations, schema)
    else:
        raise RuntimeError("ray.to_spark only supports arrow type blocks")

max_stars_repo_name: KepingYan/raydp

max_issues_repo_name: KepingYan/raydp

max_forks_repo_name: KepingYan/raydp