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

"""
Copyright (C) 2016-2020 Centre National d'Etudes Spatiales (CNES), CSSI, CESBIO  All Rights Reserved

Licensed 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/LICENSE2.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 os
import tempfile
from osgeo import gdal
import logging
from datetime import datetime
from xml.etree import ElementTree
import numpy as np
from scipy import ndimage
from scipy.ndimage import zoom
import math
from collections import namedtuple
from StartMaja.Common import FileSystem, ImageTools, ImageIO, XMLTools
from StartMaja.prepare_mnt.mnt.DEMInfo import DEMInfo

# GlobalSurfaceWater (GSW) base url for direct-download of the files:
surface_water_url = "https://storage.googleapis.com/global-surface-water/downloads2/occurrence/occurrence_%s_v1_1.tif"

class MNT(object):
    """
    Base class to set the DEM/MNT interfaces for a given site.
    DEM := Digital Elevation model)
    MNT := Modèle numérique de terrain; french for DEM
    """
    def __init__(self, site, **kwargs):
        if not int(gdal.VersionInfo()) >= 2000000:
            raise ImportError("MNT creation needs Gdal >2.0!")
        self.site = site
        self.dem_dir = kwargs.get("dem_dir", os.path.join(tempfile.gettempdir(), "maja_dem_files"))
        if not os.path.isdir(self.dem_dir):
            FileSystem.create_directory(self.dem_dir)
        self.wdir = kwargs.get("wdir", None)
        if not self.wdir:
            self.wdir = tempfile.mkdtemp(prefix="raw_dem_")
        else:
            FileSystem.create_directory(self.wdir)
        # This needs to be provided for EuDEM:
        self.raw_dem = kwargs.get("raw_dem", None)
        if not self.raw_dem:
            self.raw_dem = tempfile.mkdtemp(prefix="raw_dem_")
        else:
            FileSystem.create_directory(self.raw_dem)
        self.raw_gsw = kwargs.get("raw_gsw", None)
        if not self.raw_gsw:
            self.raw_gsw = tempfile.mkdtemp(prefix="raw_gsw_")
        else:
            FileSystem.create_directory(self.raw_gsw)
        self.gsw_codes = self.get_gsw_codes(self.site)
        self.dem_version = kwargs.get("dem_version", None)
        self.gsw_threshold = kwargs.get("gsw_threshold", 30.)
        self.gsw_dst = kwargs.get("gsw_dst", os.path.join(self.wdir, "surface_water_mask.tif"))
        self.quiet = not kwargs.get("verbose", False)

    def get_raw_data(self):
        """
        Get the DEM raw-data from a given directory. If not existing, an attempt will be made to download
        it automatically.
        :return: Write a single raster containing the water-data
        """
        raise NotImplementedError

    def prepare_mnt(self):
        """
        Prepare the ALT file from a given MNT-type (e.g. SRTM).
        :return: Write a single raster containing the altitude-data
        """
        raise NotImplementedError

    @staticmethod
    def calc_gradient(mnt_arr, res_x, res_y):
        """
        Calculate the gradient in x and y direction.
        :param mnt_arr: The ALT numpy array
        :param res_x: The resolution in x-direction. *Both positive and negative will work.*
        :param res_y: The resolution in y-direction. *Both positive and negative will work.*
        :return: The gradient/derivative in x- and y-direction.
        """
        # TODO Find a pure numpy 2D convolution.
        kernel_horizontal = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
        kernel_vertical = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
        dz_dc = ndimage.convolve(mnt_arr, kernel_horizontal) / 8. / np.abs(res_x)
        dz_dl = ndimage.convolve(mnt_arr, kernel_vertical) / 8. / np.abs(res_y)
        return dz_dc, dz_dl

    @staticmethod
    def calc_slope_aspect(dz_dc, dz_dl):
        """
        Calculate the slope and aspect.
        :param dz_dc: The derivative in x-direction
        :param dz_dl: The derivative in y-direction
        :return: The slope and aspect as numpy arrays.
        """
        norme = np.sqrt(dz_dc * dz_dc + dz_dl * dz_dl)
        slope = np.arctan(norme)
        aspect = np.where(dz_dc > 0, np.arccos(dz_dl / norme), 2 * np.pi - np.arccos(dz_dl / norme))
        aspect = np.where(slope < 0.0001, 0, aspect)
        slope = np.array(slope * 100., dtype=np.int16)
        aspect = np.array(aspect * 100., dtype=np.int16)
        return slope, aspect

    @staticmethod
    def resample_to_full_resolution(image, mnt_resolution, full_resolution, order=3):
        """
        Resample an image from the mnt resolution (usually a low one) to the
        max resolution of a given product (10m for S2).
        :param image: A numpy image
        :param mnt_resolution: The input resolution as tuple (res_x, res_y)
        :param full_resolution: The output resolution as tuple (res_x, res_y)
        :param order: The spline interpolation order. Default is 3.
        :return:
        """
        zoom_factor = (np.abs(mnt_resolution[0] / full_resolution[0]),
                       np.abs(mnt_resolution[1] / full_resolution[1]))

        # Order 3 = Cubic interpolation
        return zoom(image, zoom_factor, order=order)

    @staticmethod
    def get_gsw_codes(site, grid_step=10):
        """
        Get the list of GSW files for a given site.
        :param site: The site class
        :param grid_step: The step size of the grid
        :return: The list of filenames of format 'XX(E/W)_YY(N/S)' needed in order to cover to whole site.
        """
        if site.ul_lonlat[0] > 170 and site.lr_lonlat[0] < 160:
            raise ValueError("Cannot wrap around longitude change")

        point = namedtuple("point", ("x", "y"))
        pts = []
        for pt in [site.ul_lonlat, site.lr_lonlat]:
            lat_dec = (math.fabs(pt[1]) / grid_step)
            lon_dec = (math.fabs(pt[0]) / grid_step)
            if pt[1] > 0:
                lat_id = int(math.ceil(lat_dec) * grid_step)
            else:
                lat_id = -1 * int(math.floor(lat_dec) * grid_step)

            if pt[0] < 0:
                lon_id = int(math.ceil(lon_dec) * grid_step)
            else:
                lon_id = -1 * int(math.floor(lon_dec) * grid_step)
            pts.append(point(lon_id, lat_id))
        gsw_granules = []
        for x in range(pts[1].x, pts[0].x + grid_step, grid_step):
            for y in range(pts[1].y, pts[0].y + grid_step, grid_step):
                code_lat = "S" if y < 0 else "N"
                code_lon = "W" if x > 0 else "E"
                gsw_granules.append("%s%s_%s%s" % (int(math.fabs(x)), code_lon, int(math.fabs(y)), code_lat))
        return gsw_granules

    def get_raw_water_data(self):
        """
        Find the given gsw files or download them if not existing.

        :return: The list of filenames downloaded.
        """
        filenames = []
        for code in self.gsw_codes:
            current_url = surface_water_url % code
            filename = os.path.basename(current_url)
            output_path = os.path.join(self.raw_gsw, filename)
            if not os.path.isfile(output_path):
                # Download file:
                FileSystem.download_file(current_url, output_path, log_level=logging.INFO)
            filenames.append(output_path)
        return filenames

    def prepare_water_data(self):
        """
        Prepare the water mask constituing of a set of gsw files.

        :return: Writes the tiles water_mask to the self.gsw_dst path.
        """
        occ_files = self.get_raw_water_data()
        vrt_path = os.path.join(self.wdir, "vrt_%s.vrt" % self.site.nom)
        ImageTools.gdal_buildvrt(*occ_files, dst=vrt_path)
        # Overlay occurrence image with same extent as the given site.
        # Should the occurrence files not be complete, this sets all areas not covered by the occurrence to 0.
        ds_warped = ImageTools.gdal_warp(vrt_path,
                                         r="near",
                                         te=self.site.te_str,
                                         t_srs=self.site.epsg_str,
                                         tr=self.site.tr_str,
                                         dstnodata=0,
                                         multi=True)
        # Threshold the final image and write to destination:
        image_bin = ds_warped.array > self.gsw_threshold
        FileSystem.remove_file(vrt_path)
        ImageIO.write_geotiff_existing(image_bin, self.gsw_dst, ds_warped.get_ds())

    def to_maja_format(self, platform_id, mission_field, mnt_resolutions, coarse_res, full_res_only=False):
        """
        Writes an MNT in Maja (=EarthExplorer) format: A folder .DBL.DIR containing the rasters and an
        accompanying .HDR xml-file.
        The two files follow the maja syntax::
            *AUX_REFDE2*.(HDR|DBL.DIR)
        :param platform_id: The platform ID of two digits (e.g. S2_ for Sentinel2A/B; VS for Venus)
        :param mission_field: Similar to the platform ID, this is used in the <Mission>-field for the HDR file.
                              e.g. SENTINEL-2 for S2
        :param mnt_resolutions: A dict containing the resolutions for the given sensor. E.g.::
            {"XS": (10, -10)}
        :param coarse_res: A tuple of int describing the coarse resolution. E.g.::
            (240, -240).
        :param full_res_only:  If True, no coarse_res rasters will be created.
        :return: Writes the .DBL.DIR and .HDR into the specified self.dem_dir
        """
        assert len(mnt_resolutions) >= 1
        basename = str("%s_TEST_AUX_REFDE2_%s_%s" % (platform_id, self.site.nom, str(self.dem_version).zfill(4)))

        # Get mnt data
        mnt_max_res = self.prepare_mnt()

        # Water mask not needed with optional coarse_res writing:
        if coarse_res and not full_res_only:
            # Get water data
            self.prepare_water_data()
        mnt_res = (self.site.res_x, self.site.res_y)
        dbl_base = basename + ".DBL.DIR"
        dbl_dir = os.path.join(self.dem_dir, dbl_base)
        FileSystem.create_directory(dbl_dir)
        hdr = os.path.join(self.dem_dir, basename + ".HDR")

        # Calulate gradient mask at MNT resolution:
        mnt_in, drv = ImageIO.tiff_to_array(mnt_max_res, array_only=False)
        grad_y_mnt, grad_x_mnt = self.calc_gradient(mnt_in, self.site.res_x, self.site.res_y)

        full_res = (int(mnt_resolutions[0]["val"].split(" ")[0]),
                    int(mnt_resolutions[0]["val"].split(" ")[1]))

        grad_x = self.resample_to_full_resolution(grad_x_mnt, mnt_resolution=mnt_res, full_resolution=full_res, order=3)
        grad_y = self.resample_to_full_resolution(grad_y_mnt, mnt_resolution=mnt_res, full_resolution=full_res, order=3)

        slope, aspect = self.calc_slope_aspect(grad_y, grad_x)

        # Write full res slope and aspect
        geotransform = list(drv.GetGeoTransform())
        geotransform[1] = float(full_res[0])
        geotransform[-1] = float(full_res[1])
        projection = drv.GetProjection()
        tmp_asp = tempfile.mktemp(dir=self.wdir, suffix="_asp.tif")
        ImageIO.write_geotiff(aspect, tmp_asp, projection, tuple(geotransform))
        tmp_slp = tempfile.mktemp(dir=self.wdir, suffix="_slp.tif")
        ImageIO.write_geotiff(slope, tmp_slp, projection, tuple(geotransform))

        # Full resolution:
        write_resolution_name = True if len(mnt_resolutions) > 1 else False
        # Names for R1, R2 etc.
        rasters_written = []
        path_alt, path_asp, path_slp = "", "", ""
        all_paths_alt = []
        for res in mnt_resolutions:
            # ALT:
            bname_alt = basename + "_ALT"
            bname_alt += "_" + str(res["name"]) if write_resolution_name else ""
            bname_alt += ".TIF"
            rel_alt = os.path.join(dbl_base, bname_alt)
            path_alt = os.path.join(self.dem_dir, rel_alt)
            all_paths_alt.append(path_alt)
            ImageTools.gdal_warp(mnt_max_res, dst=path_alt, tr=res["val"], r="cubic", multi=True)
            rasters_written.append(rel_alt)
            # ASP:
            bname_asp = basename + "_ASP"
            bname_asp += "_" + res["name"] if write_resolution_name else ""
            bname_asp += ".TIF"
            rel_asp = os.path.join(dbl_base, bname_asp)
            path_asp = os.path.join(self.dem_dir, rel_asp)
            ImageTools.gdal_warp(tmp_asp, dst=path_asp, tr=res["val"], r="cubic", multi=True)
            rasters_written.append(rel_asp)
            # SLP:
            bname_slp = basename + "_SLP"
            bname_slp += "_" + res["name"] if write_resolution_name else ""
            bname_slp += ".TIF"
            rel_slp = os.path.join(dbl_base, bname_slp)
            path_slp = os.path.join(self.dem_dir, rel_slp)
            ImageTools.gdal_warp(tmp_slp, dst=path_slp, tr=res["val"], r="cubic", multi=True)
            rasters_written.append(rel_slp)

        # Optional coarse_res writing:
        if coarse_res and not full_res_only:
            # Resize all rasters for coarse res.
            coarse_res_str = str(coarse_res[0]) + " " + str(coarse_res[1])
            # ALC:
            bname_alc = basename + "_ALC.TIF"
            rel_alc = os.path.join(dbl_base, bname_alc)
            path_alc = os.path.join(self.dem_dir, rel_alc)
            ImageTools.gdal_warp(path_alt, dst=path_alc, tr=coarse_res_str, multi=True)
            rasters_written.append(rel_alc)
            # ALC:
            bname_asc = basename + "_ASC.TIF"
            rel_asc = os.path.join(dbl_base, bname_asc)
            path_asc = os.path.join(self.dem_dir, rel_asc)
            ImageTools.gdal_warp(path_asp, dst=path_asc, tr=coarse_res_str, multi=True)
            rasters_written.append(rel_asc)
            # ALC:
            bname_slc = basename + "_SLC.TIF"
            rel_slc = os.path.join(dbl_base, bname_slc)
            path_slc = os.path.join(self.dem_dir, rel_slc)
            ImageTools.gdal_warp(path_slp, dst=path_slc, tr=coarse_res_str, multi=True)
            rasters_written.append(rel_slc)
            # Water mask:
            bname_msk = basename + "_MSK.TIF"
            rel_msk = os.path.join(dbl_base, bname_msk)
            path_msk = os.path.join(self.dem_dir, rel_msk)
            ImageTools.gdal_warp(self.gsw_dst, dst=path_msk, tr=coarse_res_str, multi=True)
            rasters_written.append(rel_msk)

        # Write HDR Metadata:

        date_start = datetime(1970, 1, 1)
        date_end = datetime(2100, 1, 1)
        dem_info = DEMInfo(self.site, all_paths_alt[0])
        root = self._get_root()
        self._create_hdr(root, mission_field, basename, rasters_written,
                         dem_info, date_start, date_end, self.dem_version)
        XMLTools.write_xml(root, hdr)

        # Remove temp files:
        FileSystem.remove_file(tmp_asp)
        FileSystem.remove_file(tmp_slp)
        FileSystem.remove_file(mnt_max_res)
        return hdr, dbl_dir

    @staticmethod
    def _get_root():
        """
        Create the xml-root of a single dem file
        :return: The `ElementTree.Element` root node with all namespaces and schemas filled out.
        """
        xmlns = "http://eop-cfi.esa.int/CFI"
        xsi = "http://www.w3.org/2001/XMLSchema-instance"
        schema_location = "%s ./%s" % (xmlns, "AUX_REFDE2_ReferenceDemDataLevel2.xsd")
        type_xsi = "REFDE2_Header_Type"
        root = ElementTree.Element("Earth_Explorer_Header",
                                   attrib={"xmlns": xmlns,
                                           "schema_version": "1.00",
                                           "{" + xsi + "}schemaLocation": schema_location,
                                           "{" + xsi + "}type": type_xsi})
        return root

    @staticmethod
    def _create_hdr(root, mission, basename_out, rel_files, dem_info, date_start, date_end, version):
        """
        Create a DEM HDR file with the needed information.
        :param root: The HDR xml-root.
        :param mission: The full mission name, e.g. 'SENTINEL2_'
        :param basename_out: The basename of the HDR
        :param dem_info: A DEMInfo object
        :param date_start: The validity start date, usually 1970-01-01
        :param date_end: The validity end date, usually 2100-01-01
        :param rel_files: The list of files inside the .DBL.DIR. The filepaths are all relative to the HDR.
        :return: Writes the HDR in the same directory as the .DBL.DIR.
        """
        creation_date = datetime.now()

        a1 = ElementTree.SubElement(root, "Fixed_Header")
        toto = ElementTree.SubElement(a1, "File_Name")
        toto.text = basename_out
        toto = ElementTree.SubElement(a1, "File_Description")
        toto.text = "ReferenceDemDataLevel2"
        toto = ElementTree.SubElement(a1, "Notes")
        toto.text = "Created using Start_Maja"
        toto = ElementTree.SubElement(a1, "Mission")
        toto.text = mission
        toto = ElementTree.SubElement(a1, "File_Class")
        toto.text = "TEST"
        toto = ElementTree.SubElement(a1, "File_Type")
        toto.text = "AUX_REFDE2"
        b1 = ElementTree.SubElement(a1, "Validity_Period")
        toto = ElementTree.SubElement(b1, "Validity_Start")
        toto.text = "UTC=%s" % date_start.strftime("%Y-%m-%dT%H:%M:%S")
        toto = ElementTree.SubElement(b1, "Validity_Stop")
        toto.text = "UTC=%s" % date_end.strftime("%Y-%m-%dT%H:%M:%S")
        toto = ElementTree.SubElement(a1, "File_Version")
        toto.text = str(version).zfill(4)
        b2 = ElementTree.SubElement(a1, "Source")
        toto = ElementTree.SubElement(b2, "System")
        toto.text = "MAJA"
        toto = ElementTree.SubElement(b2, "Creator")
        toto.text = "Start_Maja"
        toto = ElementTree.SubElement(b2, "Creator_Version")
        toto.text = "1.0"
        toto = ElementTree.SubElement(b2, "Creation_Date")
        toto.text = "UTC=" + creation_date.strftime("%Y-%m-%dT%H:%M:%S")

        a = ElementTree.SubElement(root, "Variable_Header")
        b2 = ElementTree.SubElement(a, "Main_Product_Header")
        ElementTree.SubElement(b2, "List_of_Consumers", count="0")
        ElementTree.SubElement(b2, "List_of_Extensions", count="0")
        b3 = ElementTree.SubElement(a, "Specific_Product_Header")
        b4 = ElementTree.SubElement(b3, "Instance_Id")
        ElementTree.SubElement(b4, "Applicable_Site_Nick_Name").text = dem_info.name
        ElementTree.SubElement(b4, "File_Version").text = str(version).zfill(4)
        b5 = ElementTree.SubElement(b3, "List_of_Applicable_SiteDefinition_Ids", count="1")
        site_defintion_base = "_".join(basename_out.split("_")[:2])  # Get e.g. 'VE_TEST' or 'S2__TEST'
        site_definiton_name = site_defintion_base + "_MPL_SITDEF_S_" + dem_info.name
        ElementTree.SubElement(b5, "Applicable_SiteDefinition_Id", sn="1").text = site_definiton_name
        b6 = ElementTree.SubElement(b3, "DEM_Information")
        c1 = ElementTree.SubElement(b6, "Cartographic")
        d1 = ElementTree.SubElement(c1, "Coordinate_Reference_System")
        ElementTree.SubElement(d1, "Code").text = str(dem_info.epsg)
        ElementTree.SubElement(d1, "Short_Description").text = str(dem_info.short_description)
        d2 = ElementTree.SubElement(c1, "Upper_Left_Corner")
        ElementTree.SubElement(d2, "X", unit="m").text = str(dem_info.ulx)
        ElementTree.SubElement(d2, "Y", unit="m").text = str(dem_info.uly)
        d3 = ElementTree.SubElement(c1, "Sampling_Interval")
        ElementTree.SubElement(d3, "By_Line", unit="m").text = str(dem_info.resx)
        ElementTree.SubElement(d3, "By_Column", unit="m").text = str(dem_info.resy)
        d4 = ElementTree.SubElement(c1, "Size")
        ElementTree.SubElement(d4, "Lines").text = str(dem_info.ly)
        ElementTree.SubElement(d4, "Columns").text = str(dem_info.lx)
        ElementTree.SubElement(b6, "Mean_Altitude_Over_L2_Coverage", unit="m").text = str(dem_info.mean_alt)
        ElementTree.SubElement(b6, "Altitude_Standard_Deviation_Over_L2_Coverage",
                               unit="m").text = str(dem_info.std_dev_alt)
        ElementTree.SubElement(b6, "Nodata_Value").text = str(int(dem_info.nodata))
        ElementTree.SubElement(b6, "L2_To_DEM_Subsampling_Ratio").text = str(dem_info.dem_subsampling_ratio)
        ElementTree.SubElement(b6, "Comment").text = "No comment"

        b4 = ElementTree.SubElement(b3, "DBL_Organization")
        b5 = ElementTree.SubElement(b4, "List_of_Packaged_DBL_Files", count=str(len(rel_files)))
        for index, dem_file in enumerate(rel_files):
            b6 = ElementTree.SubElement(b5, "Packaged_DBL_File", sn=str(index+1))
            b7 = ElementTree.SubElement(b6, "Relative_File_Path")
            b7.text = dem_file

if __name__ == "__main__":
    pass

max_stars_repo_name: spacebel/MAJA

max_issues_repo_name: spacebel/MAJA

max_forks_repo_name: spacebel/MAJA


text:
from argparse import Namespace
import csv
from typing import List, Optional

import numpy as np
import torch
from tqdm import tqdm

from .predict import predict
from chemprop_solvation.data import MoleculeDataset
from chemprop_solvation.data.utils import get_data, get_data_from_smiles
from chemprop_solvation.utils import load_args, load_checkpoint, load_scalers

def make_predictions(args: Namespace, smiles: List[str] = None) -> List[Optional[List[float]]]:
    """
    Makes predictions. If smiles is provided, makes predictions on smiles. Otherwise makes predictions on args.test_data.

    :param args: Arguments.
    :param smiles: Smiles to make predictions on.
    :return: A list of lists of target predictions.
    """
    if args.gpu is not None:
        torch.cuda.set_device(args.gpu)

    print('Loading training args')
    scaler, features_scaler = load_scalers(args.checkpoint_paths[0])
    train_args = load_args(args.checkpoint_paths[0])

    # Update args with training arguments
    for key, value in vars(train_args).items():
        if not hasattr(args, key):
            setattr(args, key, value)

    print('Loading data')
    if smiles is not None:
        test_data = get_data_from_smiles(smiles=smiles, skip_invalid_smiles=False)
    else:
        test_data = get_data(path=args.test_path, args=args, use_compound_names=args.use_compound_names, skip_invalid_smiles=False)

    print('Validating SMILES')
    valid_indices = [i for i in range(len(test_data)) if test_data[i].mol is not None]
    full_data = test_data
    test_data = MoleculeDataset([test_data[i] for i in valid_indices])

    # Edge case if empty list of smiles is provided
    if len(test_data) == 0:
        return [None] * len(full_data)

    if args.use_compound_names:
        compound_names = test_data.compound_names()
    print(f'Test size = {len(test_data):,}')

    # Normalize features
    if train_args.features_scaling:
        test_data.normalize_features(features_scaler)

    # Predict with each model individually and sum predictions
    if args.dataset_type == 'multiclass':
        sum_preds = np.zeros((len(test_data), args.num_tasks, args.multiclass_num_classes))
    else:
        sum_preds = np.zeros((len(test_data), args.num_tasks))
    print(f'Predicting with an ensemble of {len(args.checkpoint_paths)} models')
    for checkpoint_path in tqdm(args.checkpoint_paths, total=len(args.checkpoint_paths)):
        # Load model
        model = load_checkpoint(checkpoint_path, cuda=args.cuda)
        model_preds = predict(
            model=model,
            data=test_data,
            batch_size=args.batch_size,
            scaler=scaler
        )
        sum_preds += np.array(model_preds)

    # Ensemble predictions
    avg_preds = sum_preds / len(args.checkpoint_paths)
    avg_preds = avg_preds.tolist()

    # Save predictions
    assert len(test_data) == len(avg_preds)
    print(f'Saving predictions to {args.preds_path}')

    # Put Nones for invalid smiles
    full_preds = [None] * len(full_data)
    for i, si in enumerate(valid_indices):
        full_preds[si] = avg_preds[i]
    avg_preds = full_preds
    if args.solvation:
        test_inchis = full_data.solvation_set_inchi()
    else:
        test_inchis = full_data.inchi()

    # Write predictions
    with open(args.preds_path, 'w') as f:
        writer = csv.writer(f)

        header = []

        if args.use_compound_names:
            header.append('compound_names')

        header.append('InChI')

        if args.dataset_type == 'multiclass':
            for name in args.task_names:
                for i in range(args.multiclass_num_classes):
                    header.append(name + '_class' + str(i))
        else:
            header.extend(args.task_names)
        writer.writerow(header)

        for i in range(len(avg_preds)):
            row = []

            if args.use_compound_names:
                row.append(compound_names[i])

            row.append(test_inchis[i])

            if avg_preds[i] is not None:
                if args.dataset_type == 'multiclass':
                    for task_probs in avg_preds[i]:
                        row.extend(task_probs)
                else:
                    row.extend(avg_preds[i])
            else:
                if args.dataset_type == 'multiclass':
                    row.extend([''] * args.num_tasks * args.multiclass_num_classes)
                else:
                    row.extend([''] * args.num_tasks)

            writer.writerow(row)

    return avg_preds

ext: py

reason: import numpy


text:
# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, division

# third-party imports
import numpy as np
import unittest
from nose.tools import eq_ as eq, assert_raises, assert_is_instance, \
    assert_not_is_instance
from allel.test.tools import assert_array_equal as aeq

# internal imports
from allel import GenotypeArray, HaplotypeArray, AlleleCountsArray, GenotypeVector, \
    GenotypeAlleleCountsArray, GenotypeAlleleCountsVector
from allel.test.model.test_api import GenotypeArrayInterface, HaplotypeArrayInterface, \
    diploid_genotype_data, triploid_genotype_data, haplotype_data, \
    AlleleCountsArrayInterface, allele_counts_data, GenotypeAlleleCountsArrayInterface, \
    diploid_genotype_ac_data, triploid_genotype_ac_data

# noinspection PyMethodMayBeStatic
class GenotypeArrayTests(GenotypeArrayInterface, unittest.TestCase):

    _class = GenotypeArray

    def setup_instance(self, data, dtype=None):
        return GenotypeArray(data, dtype=dtype)

    def test_constructor(self):

        # missing data arg
        with assert_raises(TypeError):
            # noinspection PyArgumentList
            GenotypeArray()

        # data has wrong dtype
        data = 'foo bar'
        with assert_raises(TypeError):
            GenotypeArray(data)

        # data has wrong dtype
        data = [4., 5., 3.7]
        with assert_raises(TypeError):
            GenotypeArray(data)

        # data has wrong dimensions
        data = [1, 2, 3]
        with assert_raises(TypeError):
            GenotypeArray(data)

        # data has wrong dimensions
        data = [[1, 2], [3, 4]]  # use HaplotypeArray instead
        with assert_raises(TypeError):
            GenotypeArray(data)

        # diploid data (typed)
        g = GenotypeArray(diploid_genotype_data, dtype='i1')
        aeq(diploid_genotype_data, g)
        eq(np.int8, g.dtype)

        # polyploid data (typed)
        g = GenotypeArray(triploid_genotype_data, dtype='i1')
        aeq(triploid_genotype_data, g)
        eq(np.int8, g.dtype)

    def test_slice_types(self):

        g = GenotypeArray(diploid_genotype_data, dtype='i1')

        # row slice
        s = g[1:]
        assert_is_instance(s, GenotypeArray)

        # col slice
        s = g[:, 1:]
        assert_is_instance(s, GenotypeArray)

        # row index
        s = g[0]
        assert_is_instance(s, GenotypeVector)
        assert_not_is_instance(s, GenotypeArray)

        # col index
        s = g[:, 0]
        assert_is_instance(s, GenotypeVector)
        assert_not_is_instance(s, GenotypeArray)

        # ploidy index
        s = g[:, :, 0]
        assert_is_instance(s, np.ndarray)
        assert_not_is_instance(s, GenotypeArray)

        # item
        s = g[0, 0, 0]
        assert_is_instance(s, np.int8)
        assert_not_is_instance(s, GenotypeArray)

    def test_haploidify_samples(self):

        # diploid
        g = GenotypeArray([[[0, 1], [2, 3]],
                           [[4, 5], [6, 7]],
                           [[8, 9], [10, 11]]], dtype='i1')
        h = g.haploidify_samples()
        eq(2, h.ndim)
        eq(3, h.n_variants)
        eq(2, h.n_haplotypes)
        eq(np.int8, h.dtype)
        for i in range(g.n_variants):
            for j in range(g.n_samples):
                self.assertIn(h[i, j], set(g[i, j]))

        # triploid
        g = GenotypeArray([[[0, 1, 2], [3, 4, 5]],
                           [[6, 7, 8], [9, 10, 11]],
                           [[12, 13, 14], [15, 16, 17]]], dtype='i1')
        h = g.haploidify_samples()
        eq(2, h.ndim)
        eq(3, h.n_variants)
        eq(2, h.n_haplotypes)
        eq(np.int8, h.dtype)
        for i in range(g.n_variants):
            for j in range(g.n_samples):
                self.assertIn(h[i, j], set(g[i, j]))

    def test_take(self):
        g = self.setup_instance(diploid_genotype_data)
        # take variants not in original order
        indices = [2, 0]
        t = g.take(indices, axis=0)
        eq(2, t.n_variants)
        eq(g.n_samples, t.n_samples)
        eq(g.ploidy, t.ploidy)
        expect = np.array(diploid_genotype_data).take(indices, axis=0)
        aeq(expect, t)

    def test_no_subclass(self):
        g = self.setup_instance(diploid_genotype_data)
        self.assertNotIsInstance(g.reshape((g.shape[0], -1)), GenotypeArray)
        self.assertNotIsInstance(g.flatten(), GenotypeArray)
        self.assertNotIsInstance(g.ravel(), GenotypeArray)
        self.assertNotIsInstance(g.transpose(), GenotypeArray)
        self.assertNotIsInstance(g.T, GenotypeArray)
        self.assertNotIsInstance(g.astype('f4'), GenotypeArray)

# noinspection PyMethodMayBeStatic
class HaplotypeArrayTests(HaplotypeArrayInterface, unittest.TestCase):

    _class = HaplotypeArray

    def setup_instance(self, data, dtype=None):
        return HaplotypeArray(data, dtype=dtype)

    def test_constructor(self):

        # missing data arg
        with self.assertRaises(TypeError):
            # noinspection PyArgumentList
            HaplotypeArray()

        # data has wrong dtype
        data = 'foo bar'
        with self.assertRaises(TypeError):
            HaplotypeArray(data)

        # data has wrong dtype
        data = [4., 5., 3.7]
        with self.assertRaises(TypeError):
            HaplotypeArray(data)

        # data has wrong dimensions
        data = [1, 2, 3]
        with self.assertRaises(TypeError):
            HaplotypeArray(data)

        # data has wrong dimensions
        data = diploid_genotype_data  # use GenotypeArray instead
        with self.assertRaises(TypeError):
            HaplotypeArray(data)

        # haploid data (typed)
        h = HaplotypeArray(haplotype_data, dtype='i1')
        aeq(haplotype_data, h)
        eq(np.int8, h.dtype)

    def test_slice_types(self):

        h = HaplotypeArray(haplotype_data, dtype='i1')

        # row slice
        s = h[1:]
        assert_is_instance(s, HaplotypeArray)

        # col slice
        s = h[:, 1:]
        assert_is_instance(s, HaplotypeArray)

        # row index
        s = h[0]
        assert_is_instance(s, np.ndarray)
        assert_not_is_instance(s, HaplotypeArray)

        # col index
        s = h[:, 0]
        assert_is_instance(s, np.ndarray)
        assert_not_is_instance(s, HaplotypeArray)

        # item
        s = h[0, 0]
        assert_is_instance(s, np.int8)
        assert_not_is_instance(s, HaplotypeArray)

# noinspection PyMethodMayBeStatic
class AlleleCountsArrayTests(AlleleCountsArrayInterface, unittest.TestCase):

    _class = AlleleCountsArray

    def setup_instance(self, data):
        return AlleleCountsArray(data)

    def test_constructor(self):

        # missing data arg
        with assert_raises(TypeError):
            # noinspection PyArgumentList
            AlleleCountsArray()

        # data has wrong dtype
        data = 'foo bar'
        with assert_raises(TypeError):
            AlleleCountsArray(data)

        # data has wrong dtype
        data = [4., 5., 3.7]
        with assert_raises(TypeError):
            AlleleCountsArray(data)

        # data has wrong dimensions
        data = [1, 2, 3]
        with assert_raises(TypeError):
            AlleleCountsArray(data)

        # data has wrong dimensions
        data = diploid_genotype_data
        with assert_raises(TypeError):
            AlleleCountsArray(data)

        # valid data (typed)
        ac = AlleleCountsArray(allele_counts_data, dtype='u1')
        aeq(allele_counts_data, ac)
        eq(np.uint8, ac.dtype)

    def test_slice_types(self):
        ac = AlleleCountsArray(allele_counts_data, dtype='u1')

        # total slice
        s = ac[:]
        assert_is_instance(s, AlleleCountsArray)

        # row slice
        s = ac[1:]
        assert_is_instance(s, AlleleCountsArray)

        # col slice
        s = ac[:, 1:]
        assert_is_instance(s, np.ndarray)

        # row index
        s = ac[0]
        assert_is_instance(s, np.ndarray)

        # col index
        s = ac[:, 0]
        assert_is_instance(s, np.ndarray)

        # item
        s = ac[0, 0]
        assert_is_instance(s, np.uint8)

    def test_reduce_types(self):
        ac = AlleleCountsArray(allele_counts_data, dtype='u1')

        for m in 'sum', 'max', 'argmax':
            x = getattr(ac, m)(axis=1)
            assert_is_instance(x, np.ndarray)

# noinspection PyMethodMayBeStatic
class GenotypeVectorTests(unittest.TestCase):

    def test_properties(self):

        # diploid row
        gv = GenotypeVector(diploid_genotype_data[0])
        eq(2, gv.ndim)
        eq((3, 2), gv.shape)
        eq(np.dtype(int), gv.dtype)
        eq(3, gv.n_calls)
        eq(2, gv.ploidy)
        eq(6, gv.n_allele_calls)

        # diploid column
        gv = GenotypeVector(np.array(diploid_genotype_data, dtype='i1')[:, 0])
        eq(2, gv.ndim)
        eq((5, 2), gv.shape)
        eq(np.dtype('i1'), gv.dtype)
        eq(5, gv.n_calls)
        eq(2, gv.ploidy)
        eq(10, gv.n_allele_calls)

        # triploid row
        gv = GenotypeVector(triploid_genotype_data[0])
        eq(2, gv.ndim)
        eq((3, 3), gv.shape)
        eq(np.dtype(int), gv.dtype)
        eq(3, gv.n_calls)
        eq(3, gv.ploidy)
        eq(9, gv.n_allele_calls)

        # triploid column
        gv = GenotypeVector(np.array(triploid_genotype_data, dtype='i1')[:, 0])
        eq(2, gv.ndim)
        eq((4, 3), gv.shape)
        eq(np.dtype('i1'), gv.dtype)
        eq(4, gv.n_calls)
        eq(3, gv.ploidy)
        eq(12, gv.n_allele_calls)

    def test_input_data(self):

        # ndim 3
        with assert_raises(TypeError):
            GenotypeVector(diploid_genotype_data)
        with assert_raises(TypeError):
            GenotypeVector(triploid_genotype_data)

        # ndim 1
        with assert_raises(TypeError):
            GenotypeVector([0, 1, 0, -1])

        # ndim 1
        with assert_raises(TypeError):
            GenotypeVector([0, 1, 0, -1])

        # wrong dtype
        with assert_raises(TypeError):
            GenotypeVector(np.array(diploid_genotype_data, dtype='f4')[:, 0])

    def test_getitem(self):
        gv = GenotypeVector(diploid_genotype_data[0])

        # these should return the same type
        gs = gv[:]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[:, :]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[...]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[0:2]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[0:2, :]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[np.array([True, False, True], dtype=bool)]
        assert_is_instance(gs, GenotypeVector)
        gs = gv[[0, 2]]
        assert_is_instance(gs, GenotypeVector)

        # these should return plain array
        gs = gv[:, 0]
        assert_not_is_instance(gs, GenotypeVector)
        gs = gv[np.newaxis, :2, 0]  # change dimension semantics
        assert_not_is_instance(gs, GenotypeVector)

    # __getitem__
    # to_html_str
    # _repr_html_
    # mask getter
    # mask setter
    # fill_masked
    # is_called
    # is_missing
    # is_hom
    # is_hom_ref
    # is_hom_alt
    # is_het
    # is_call
    # count_called
    # count_missing
    # count_hom
    # count_hom_ref
    # count_hom_alt
    # count_het
    # count_call
    # to_haplotypes
    # to_n_ref
    # to_n_alt
    # to_allele_counts
    # haploidify
    # to_gt
    # map_alleles

# noinspection PyMethodMayBeStatic
class GenotypeAlleleCountsArrayTests(GenotypeAlleleCountsArrayInterface, unittest.TestCase):

    _class = GenotypeAlleleCountsArray

    def setup_instance(self, data, dtype=None):
        return GenotypeAlleleCountsArray(data, dtype=dtype)

    def test_constructor(self):

        # missing data arg
        with assert_raises(TypeError):
            # noinspection PyArgumentList
            GenotypeAlleleCountsArray()

        # data has wrong dtype
        data = 'foo bar'
        with assert_raises(TypeError):
            GenotypeAlleleCountsArray(data)

        # data has wrong dtype
        data = [4., 5., 3.7]
        with assert_raises(TypeError):
            GenotypeAlleleCountsArray(data)

        # data has wrong dimensions
        data = [1, 2, 3]
        with assert_raises(TypeError):
            GenotypeAlleleCountsArray(data)

        # data has wrong dimensions
        data = [[1, 2], [3, 4]]  # use HaplotypeArray instead
        with assert_raises(TypeError):
            GenotypeAlleleCountsArray(data)

        # diploid data (typed)
        g = GenotypeAlleleCountsArray(diploid_genotype_ac_data, dtype='i1')
        aeq(diploid_genotype_ac_data, g)
        eq(np.int8, g.dtype)

        # polyploid data (typed)
        g = GenotypeAlleleCountsArray(triploid_genotype_ac_data, dtype='i1')
        aeq(triploid_genotype_ac_data, g)
        eq(np.int8, g.dtype)

    def test_slice_types(self):

        g = GenotypeAlleleCountsArray(diploid_genotype_ac_data, dtype='i1')

        # row slice
        s = g[1:]
        assert_is_instance(s, GenotypeAlleleCountsArray)

        # col slice
        s = g[:, 1:]
        assert_is_instance(s, GenotypeAlleleCountsArray)

        # row index
        s = g[0]
        assert_is_instance(s, GenotypeAlleleCountsVector)
        assert_not_is_instance(s, GenotypeAlleleCountsArray)

        # col index
        s = g[:, 0]
        assert_is_instance(s, GenotypeAlleleCountsVector)
        assert_not_is_instance(s, GenotypeAlleleCountsArray)

        # ploidy index
        s = g[:, :, 0]
        assert_is_instance(s, np.ndarray)
        assert_not_is_instance(s, GenotypeAlleleCountsArray)

        # item
        s = g[0, 0, 0]
        assert_is_instance(s, np.int8)
        assert_not_is_instance(s, GenotypeAlleleCountsArray)

max_stars_repo_name: petrelharp/scikit-allel

max_issues_repo_name: petrelharp/scikit-allel

max_forks_repo_name: petrelharp/scikit-allel


text:
import numpy as np
from sklearn.cluster import KMeans

class GapAnalysis(object):
    """Performs gap analysis for evaluating clustering results.

    This object implements the gap analysis method from Elements
    of Statistical Learning, Hastie Et. Al.

    Paramaters
    ----------
    cluster_method: sklearn KMeans class (arg b/c might add EM object support)
     KMeans object.
    max_clusters: int, optional
     Maximum number of clusters to consider. Default is 10.
    num_samples: int, optional
     Number of randomly sampled datasets to use for computing wk_null.
    random_state: int, optional
     Random seed to use.
    cov_type: string, optional
     Covariance type for a GaussianMixture model. Also indicates model type.

    Attributes
    ----------
    wk: 1d-array
     Within clusters inertia values on original dataset.
    wk_null: 1d-array
     Mean within clusters inertia values on null dataset samples.
    wk_null_stdev: 1d-array
     Stdev of within clusters inertia values on null dataset samples.
    """

    def __init__(self, cluster_method, max_clusters=10, num_samples=20,
                 random_state=0, cov_type=None):
        self.cluster_method = cluster_method
        self.max_clusters = max_clusters
        self.num_samples = num_samples
        self.random_state = random_state
        self.cov_type = cov_type

    def fit(self, X):
        """"Fit clustering_method to original data and to noise samples."""
        rgen = np.random.RandomState(self.random_state)
        self.wk_ = np.zeros(self.max_clusters)
        self.wk_null_ = np.zeros(self.max_clusters)
        self.wk_null_stdev_ = np.zeros(self.max_clusters)
        self.gap_ = np.zeros(self.max_clusters)
        self.gap_stdev_ = np.zeros(self.max_clusters)

        # Generate results for original dataset
        for idx, n in enumerate(range(1,self.max_clusters+1)):
            if self.cov_type is not None:
                cm = self.cluster_method(n_components=n,
                                         covariance_type=self.cov_type)
                cm.fit(X)
                self.wk_[idx] = cm.bic(X)
            else:
                cm = self.cluster_method(n_clusters=n)
                cm.fit(X)
                self.wk_[idx] = np.log(cm.inertia_)

        # Generate results for noise samples
        for idx, n in enumerate(range(1,self.max_clusters+1)):
            temp_wk = np.zeros(self.num_samples)
            for i in range(self.num_samples):
                if self.cov_type is not None:
                    cm = self.cluster_method(n_components=n,
                                             covariance_type=self.cov_type)
                    X_sample = self.generate_sample(X, rgen)
                    cm.fit(X_sample)
                    temp_wk[i] = cm.bic(X_sample)
                else:
                    cm = self.cluster_method(n_clusters=n)
                    cm.fit(self.generate_sample(X, rgen))
                    temp_wk[i] = cm.inertia_

            if self.cov_type is not None:
                self.wk_null_[idx] = np.mean(temp_wk)
                self.wk_null_stdev_[idx] = np.std(temp_wk)
                temp_gap = temp_wk - self.wk_[idx]
            else:
                self.wk_null_[idx] = np.mean(np.log(temp_wk))
                self.wk_null_stdev_[idx] = np.std(np.log(temp_wk))
                temp_gap = np.log(temp_wk) - self.wk_[idx]
            self.gap_[idx] = np.mean(temp_gap)
            self.gap_stdev_[idx] = (np.std(temp_gap) *
                                    np.sqrt(1+(1/self.num_samples)))

    def generate_sample(self, X, rgen):
        """Generate random sample for clustering."""
        sample_X = np.zeros(X.shape)
        for j in range(X.shape[1]):
            min_j = np.min(X[:,j])
            max_j = np.max(X[:,j])
            sample_X[:,j] = rgen.uniform(min_j, max_j, size=X.shape[0])
        return sample_X

max_stars_repo_name: andrewdoss/seattle-food

max_issues_repo_name: andrewdoss/seattle-food

max_forks_repo_name: andrewdoss/seattle-food


text:
#reference: https://towardsdatascience.com/reinforcement-learning-rl-101-with-python-e1aa0d37d43b

def imps():
	global np, random
	import numpy as np
	import random

def swapColumn(arr, c0, c1):
	temp = np.copy(arr[:, c0])
	arr[:, c0] = arr[:, c1]
	my_array[:, c1] = temp

def copyColumn(arr, to, fr):
	arr[:, to] = arr[:, fr]

def copyRow(arr, to, fr):
	arr[to] = arr[fr]

def copyArrayContentToBorder(arr):
	copyColumn(arr, 0, 1)
	copyColumn(arr, gridSizeX+1, gridSizeX)
	copyRow(arr, 0, 1)
	copyRow(arr, gridSizeY+1, gridSizeY)

def fillArray(arr):		#for test
	cnt = arr.shape[0] * arr.shape[1]
	arr1 = arr.reshape(cnt)
	for i in range(cnt):
		arr1[i] = i

#reward is from initPos.
def getTargetPosAndReward(initPos, action):
	targetPos = (initPos[0], initPos[1])
	reward = rewards[targetPos]
	if reward == 0:
		return targetPos, reward

	targetPos = (initPos[0] + action[0], initPos[1] + action[1])
	#reward = rewards[targetPos]
	return targetPos, reward

def swapValueMaps():
	global valueMap, valueMapNext
	tmpMap = valueMapNext
	valueMapNext = valueMap
	valueMap = tmpMap

def init():
	global gamma, gridSizeX, gridSizeY, actions, numIterations, weight, rewards, valueMap, valueMapNext, states
	gamma = 1 # discounting rate
	gridSizeX = 4
	gridSizeY = 4
	actions = np.array([[-1, 0], [1, 0], [0, 1], [0, -1]])
	numIterations = 1
	weight = 1./len(actions)

	rewards = np.full((gridSizeY+2, gridSizeX+2), -1, dtype=float)
	rewards[1,1] = 0
	rewards[4,4] = 0
	copyArrayContentToBorder(rewards)

	#make this 2 units bigger to avoid boundary check, every iteration, copy the content from appropriate row/columns
	valueMap = np.zeros((gridSizeY+2, gridSizeX+2))
	valueMapNext = np.zeros((gridSizeY+2, gridSizeX+2))
	states = [[y, x] for y in range(1,gridSizeY+1) for x in range(1,gridSizeX+1)]

def PolicyEvaluation():
	global valueMap, copyValueMap, weight
	for it in range(numIterations):
		for state in states:
			#print(state, rewards[state[0], state[1]])
			weightedRewards = 0
			for action in actions:
				targetPos, reward = getTargetPosAndReward(np.array(state), action)
				#print(targetPos, reward)
				weightedRewards += weight*(reward+(gamma*valueMap[targetPos[0], targetPos[1]]))
			valueMapNext[state[0], state[1]] = weightedRewards

		copyArrayContentToBorder(valueMapNext)
		swapValueMaps()

	print(valueMap)

max_stars_repo_path: dynamic programming.py

max_stars_repo_name: dong-zhan/reinforcement-learning

max_issues_repo_path: dynamic programming.py

max_issues_repo_name: dong-zhan/reinforcement-learning

max_forks_repo_path: dynamic programming.py

max_forks_repo_name: dong-zhan/reinforcement-learning


text:
"""
###################################################
#Class of scripts to testing Feudal Network policy#
###################################################
"""

import numpy as np
import unittest
from FuN.scripts.training.feudal_networks.policies.feudal_policy import \
    FeudalPolicy
import tensorflow as tf

np.set_printoptions(suppress=True, precision=6)

class TestFeudalPolicy(unittest.TestCase):
    """
    Class for testing feudal network policy

    """

    def setUp(self):
        """
        Function for initializing the test

        """

        # reset graph before each test case
        tf.reset_default_graph()

    def test_init(self):
        """
        Function for initializing the test

        """

        # global_step = tf.get_variable("global_step", [], tf.int32,
        #                               initializer=tf.constant_initializer(
        #                                   0, dtype=tf.int32),
        #                               trainable=False)
        # feudal = FeudalPolicy((80, 80, 3), 4, global_step)

    def test_fit_simple_dataset(self):
        """
        Function for initializing the test

        """

        with tf.Session() as session:
            global_step = tf.get_variable("global_step", [], tf.int32,
                                          initializer=tf.constant_initializer(
                                              0, dtype=tf.int32),
                                          trainable=False)
            obs_space = (80, 80, 3)
            act_space = 2
            lr = 1e-5
            g_dim = 256
            # worker_hid_dim = 32
            # manager_hid_dim = 256
            pi = FeudalPolicy(obs_space, act_space, global_step)

            grads = tf.gradients(pi.loss, pi.var_list)

            prints = []
            for g in grads:
                prints.append(g.op.name)
                prints.append(g)
            # grads[0] = tf.Print(grads[0],prints)
            grads, _ = tf.clip_by_global_norm(grads, 40)
            grads_and_vars = list(zip(grads, pi.var_list))
            opt = tf.train.AdamOptimizer(lr)
            train_op = opt.apply_gradients(grads_and_vars)

            # train_op = tf.train.AdamOptimizer(lr).minimize(
            # pi.loss,var_list=pi.var_list)
            session.run(tf.global_variables_initializer())

            obs = [np.zeros(obs_space), np.zeros(obs_space)]
            a = [[1, 0], [0, 1]]
            returns = [0, 1]
            s_diff = [np.ones(g_dim), np.ones(g_dim)]
            gsum = [np.zeros((1, g_dim)), np.ones((1, g_dim))]
            ri = [0, 0]

            _, features = pi.get_initial_features()
            worker_features = features[0:2]
            manager_features = features[2:]

            feed_dict = {
                pi.obs: obs,
                pi.ac: a,
                pi.r: returns,
                pi.s_diff: s_diff,
                pi.prev_g: gsum,
                pi.ri: ri,
                pi.state_in[0]: worker_features[0],
                pi.state_in[1]: worker_features[1],
                pi.state_in[2]: manager_features[0],
                pi.state_in[3]: manager_features[1]
            }

            n_updates = 1000
            verbose = True
            for i in range(n_updates):
                loss, vf, policy, _ = session.run([pi.loss,
                                                   pi.manager_vf,
                                                   pi.pi, train_op],
                                                  feed_dict=feed_dict)
                if verbose:
                    print('loss: {}\npolicy: {}\nvalue: {}\n-------'.format(
                        loss, policy, vf))

    def test_simple_manager_behavior(self):
        """
        Function for initializing the test

        """

        with tf.Session() as session:
            global_step = tf.get_variable("global_step", [], tf.int32,
                                          initializer=tf.constant_initializer(
                                              0, dtype=tf.int32),
                                          trainable=False)
            obs_space = (80, 80, 3)
            act_space = 2
            lr = 5e-4
            g_dim = 256
            # worker_hid_dim = 32
            # manager_hid_dim = 256
            pi = FeudalPolicy(obs_space, act_space, global_step)
            train_op = tf.train.AdamOptimizer(lr).minimize(pi.loss)

            worker_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
            worker_vars = [v for v in worker_vars if 'worker' in v.name]
            worker_assign = tf.group(*[tf.assign(v, tf.zeros_like(v))
                                       for v in worker_vars])

            session.run(tf.global_variables_initializer())

            obs = [np.zeros(obs_space), np.zeros(obs_space)]
            a = [[1, 0], [0, 1]]
            returns = [0, 1]
            s_diff = [np.ones(g_dim), np.ones(g_dim)]
            gsum = [np.zeros((1, g_dim)), np.ones((1, g_dim))]
            ri = [0, 0]

            _, features = pi.get_initial_features()
            worker_features = features[0:2]
            manager_features = features[2:]

            feed_dict = {
                pi.obs: obs,
                pi.ac: a,
                pi.r: returns,
                pi.s_diff: s_diff,
                pi.prev_g: gsum,
                pi.ri: ri,
                pi.state_in[0]: worker_features[0],
                pi.state_in[1]: worker_features[1],
                pi.state_in[2]: manager_features[0],
                pi.state_in[3]: manager_features[1]
            }

            n_updates = 1000
            verbose = True
            for i in range(n_updates):
                loss, vf, policy, _, _ = session.run(
                    [pi.loss, pi.manager_vf, pi.pi, train_op, worker_assign],
                    feed_dict=feed_dict)
                if verbose:
                    print('loss: {}\npolicy: {}\nvalue: {}\n-------'.format(
                        loss, policy, vf))

                worker_var_values = session.run(worker_vars)
                print(worker_var_values)
                U = session.run(pi.U, feed_dict=feed_dict)
                print(U)
                input()

if __name__ == '__main__':
    unittest.main()


text:
#!/usr/bin/env python3

from pyhiir.hiir import hiir
from pyhiir.allpass import LowPass
from scipy import signal
import matplotlib.pyplot as plt
import numpy as np
import unittest

class MainTest(unittest.TestCase):
    def test_create(self):
        h = hiir()
        self.assertIsNotNone(h)
        #x = h.compute_coefs_order_tbw(2, 0.1)
# f = LowPass(x)

#exit(0)

# tf = f.get_transfer_function()

# w, h = signal.freqz(tf.b, tf.a)

# plt.plot(w / np.pi, 20 * np.log10(abs(h)), 'b')
# plt.show()

if __name__=='__main__':
    unittest.main()

max_stars_repo_name: asanza/pyhiir

max_issues_repo_name: asanza/pyhiir

max_forks_repo_name: asanza/pyhiir

reason: import numpy,from scipy


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

class CNN(nn.Module):
    def __init__(self, config, input_dim):
        super().__init__()

        self.convs = nn.ModuleList()
        d_prev = input_dim
        d = config.encoder_conv_dim
        w = config.encoder_image_size
        for k, s in zip(config.encoder_kernel_size, config.encoder_stride):
            self.convs.append(nn.Conv2d(d_prev, d, int(k), int(s)))
            w = int(np.floor((w - (int(k) - 1) - 1) / int(s) + 1))
            d_prev = d

        print("Output of CNN (%d) = %d x %d x %d" % (w * w * d, w, w, d))
        self.output_dim = config.encoder_conv_output_dim

        self.fc = nn.Linear(w * w * d, self.output_dim)
        self.ln = nn.LayerNorm(self.output_dim)

        self.apply(weight_init)

    def forward(self, ob, detach_conv=False):
        out = ob
        for conv in self.convs:
            out = F.relu(conv(out))
        out = out.flatten(start_dim=1)

        if detach_conv:
            out = out.detach()

        out = self.fc(out)
        out = self.ln(out)
        out = F.tanh(out)

        return out

    # from https://github.com/MishaLaskin/rad/blob/master/encoder.py
    def copy_conv_weights_from(self, source):
        """Tie convolutional layers"""
        # only tie conv layers
        for i, conv in enumerate(self.convs):
            assert type(source.convs[i]) == type(conv)
            conv.weight = source.convs[i].weight
            conv.bias = source.convs[i].bias

# from https://github.com/denisyarats/drq/blob/master/utils.py#L62
def weight_init(tensor):
    if isinstance(tensor, nn.Linear):
        nn.init.orthogonal_(tensor.weight.data)
        tensor.bias.data.fill_(0.0)
    elif isinstance(tensor, nn.Conv2d) or isinstance(tensor, nn.ConvTranspose2d):
        tensor.weight.data.fill_(0.0)
        tensor.bias.data.fill_(0.0)
        mid = tensor.weight.size(2) // 2
        gain = nn.init.calculate_gain("relu")
        nn.init.orthogonal_(tensor.weight.data[:, :, mid, mid], gain)
        # nn.init.orthogonal_(tensor.weight.data, gain)

class MLP(nn.Module):
    def __init__(
        self, config, input_dim, output_dim, hid_dims=[], activation_fn=None,
    ):
        super().__init__()
        self.activation_fn = activation_fn
        if activation_fn is None:
            self.activation_fn = getattr(F, config.policy_activation)

        self.fcs = nn.ModuleList()
        prev_dim = input_dim
        for d in hid_dims + [output_dim]:
            self.fcs.append(nn.Linear(prev_dim, d))
            prev_dim = d

        self.output_dim = output_dim
        self.apply(weight_init)

    def forward(self, ob):
        out = ob
        for fc in self.fcs[:-1]:
            out = self.activation_fn(fc(out))
        out = self.fcs[-1](out)
        return out

def flatten_ob(ob: dict, ac=None):
    """
    Flattens the observation dictionary. The observation dictionary
    can either contain a single ob, or a batch of obs.
    Any images must be flattened to 1D tensors, but
    we must be careful to check if we are doing a single instance
    or batch before we flatten.

    Returns a list of dim [N x D] where N is batch size and D is sum of flattened
    dims of observations
    """
    inp = []
    images = []
    single_ob = False
    for k, v in ob.items():
        if k in ["camera_ob", "depth_ob", "segmentation_ob"]:
            images.append(v)
        else:
            if len(v.shape) == 1:
                single_ob = True
            inp.append(v)
    # concatenate images into 1D
    for image in images:
        if single_ob:
            img = torch.flatten(image)
        else:  # batch of obs, flatten after bs dim
            img = torch.flatten(image, start_dim=1)
        inp.append(img)
    # now flatten into Nx1 tensors
    if single_ob:
        inp = [x.unsqueeze(0) for x in inp]

    if ac is not None:
        ac = list(ac.values())
        if len(ac[0].shape) == 1:
            ac = [x.unsqueeze(0) for x in ac]
        inp.extend(ac)
    inp = torch.cat(inp, dim=-1)
    return inp

def flatten_ac(ac: dict):
    ac = list(ac.values())
    if len(ac[0].shape) == 1:
        ac = [x.unsqueeze(0) for x in ac]
    ac = torch.cat(ac, dim=-1)
    return ac

max_stars_repo_name: youngwoon/robot-learning

max_issues_repo_name: haoyu-x/robot-learning

max_forks_repo_name: haoyu-x/robot-learning


text:
from abc import abstractmethod
from copy import copy
from .graph_vector import GraphVector, GraphModule
from .graph_vector_dict import GraphVector_dict, GraphModule_dict
from .graph_vector_vector import GraphVector_vector, GraphModule_vector
from .formality_graph import FormalityGraph
from .formality_graph_basis import FormalityGraphBasis
# for insertion:
from .undirected_graph_vector import UndirectedGraphVector_dict, UndirectedGraphVector_vector

class FormalityGraphVector(GraphVector):
    """
    Vector representing a linear combination of formality graphs.
    """
    @abstractmethod
    def nground(self):
        """
        Return the number of ground vertices in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of ground vertices.
        """
        pass

    @abstractmethod
    def is_aerial(self):
        """
        Return True if this graph vector is aerial, and False otherwise.
        """
        pass

    @abstractmethod
    def set_aerial(self, is_aerial=True):
        """
        Set this graph vector to be aerial if ``is_aerial`` is True, respectively not aerial if ``is_aerial`` is False.
        """
        pass

    def ground_symmetrization(self):
        """
        Return the symmetrization of this graph vector with respect to the ground vertices.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of ground vertices.
        """
        from itertools import permutations
        result = self.parent().zero()
        for sigma in permutations(range(self.nground())):
            result += self.map_graphs(lambda g: g.ground_relabeled(sigma))
        return result

    def ground_skew_symmetrization(self):
        """
        Return the skew-symmetrization (or anti-symmetrization) of this graph vector with respect to the ground vertices.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of ground vertices.
        """
        from itertools import permutations
        from util.permutation import selection_sort
        result = self.parent().zero()
        for sigma in permutations(range(self.nground())):
            sign = selection_sort(list(sigma))
            result += sign * self.map_graphs(lambda g: g.ground_relabeled(sigma))
        return result

    def differential_orders(self):
        """
        Return an iterator over the tuples of in-degrees of ground vertices of graphs in this graph vector.
        """
        seen = set([])
        for (_, g) in self:
            diff_order = g.differential_orders()
            if diff_order not in seen:
                yield diff_order
                seen.add(diff_order)

    def part_of_differential_order(self, diff_order):
        """
        Return the graph vector which is the summand of this graph vector containing only graphs such that the in-degrees of the ground vertices are ``diff_order``.
        """
        terms = []
        for (c,g) in self:
            if g.differential_orders() == diff_order:
                terms.append((c,g))
        return self.parent()(terms)

class FormalityGraphModule(GraphModule):
    """
    Module spanned by formality graphs.
    """
    def element_from_kgs_encoding(self, kgs_encoding, hbar):
        """
        Return the linear combination of Kontsevich graphs specified by an encoding, as an element of this module.

        INPUT:

        - ``kgs_encoding`` -- a string, containing an encoding of a graph series expansion as used in the ``kontsevich_graph_series-cpp`` program

        - ``hbar`` -- an element of the base ring, to be used as the series expansion parameter
        """
        from sage.misc.sage_eval import sage_eval
        terms = []
        exponent = 0
        for line in kgs_encoding.split('\n'):
            line = line.strip()
            if line[0] == 'h':
                exponent = int(line[2:-1])
                continue
            elif line[0] == '#':
                continue
            encoding_str, coeff_str = line.rsplit('    ')
            coeff = self.base_ring()(sage_eval(coeff_str)) * hbar**exponent
            encoding_parts = encoding_str.split('   ')
            if encoding_parts[1] == '':
                # no edges
                prefix = [int(v) for v in encoding_parts[0].split(' ')]
                terms.append((coeff, FormalityGraph(prefix[0], prefix[1], [])))
                continue
            prefix_str, targets_str = encoding_parts
            prefix = [int(v) for v in prefix_str.split(' ')]
            targets_flat = [int(v) for v in targets_str.split(' ')]
            edges = []
            for k in range(prefix[1]):
                edges.append((prefix[0] + k, targets_flat[2*k]))
                edges.append((prefix[0] + k, targets_flat[2*k+1]))
            terms.append((coeff, FormalityGraph(prefix[0], prefix[1], edges)))
        return self(terms)

class FormalityGraphVector_dict(FormalityGraphVector, GraphVector_dict):
    """
    Vector representing a linear combination of formality graphs (stored as a dictionary).
    """
    def __init__(self, parent, vector, is_aerial=False):
        """
        Initialize this formality graph vector.

        INPUT:

        - ``parent`` -- a FormalityGraphModule

        - ``vector`` -- a dictionary, representing a sparse vector of coefficients with respect to the basis of ``parent``

        - ``is_aerial`` -- (default: False) a boolean, if True then this graph vector will be aerial
        """
        if not isinstance(parent, FormalityGraphModule_dict):
            raise ValueError("parent must be a FormalityGraphModule_dict")
        super().__init__(parent, vector)
        self._is_aerial = is_aerial

    def nvertices(self):
        """
        Return the number of vertices in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of vertices.
        """
        for key in self._vector:
            gv, av, e = key[:3]
            if not self._vector[key].is_zero():
                return gv + av

    def nedges(self):
        """
        Return the number of edges in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of edges.
        """
        for key in self._vector:
            gv, av, e = key[:3]
            if not self._vector[key].is_zero():
                return e

    def nground(self):
        """
        Return the number of ground vertices in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of ground vertices.
        """
        for key in self._vector:
            gv, av, e = key[:3]
            if not self._vector[key].is_zero():
                return gv

    def is_aerial(self):
        """
        Return True if this graph vector is aerial, and False otherwise.
        """
        return self._is_aerial

    def set_aerial(self, is_aerial=True):
        """
        Set this graph vector to be aerial if ``is_aerial`` is True, respectively not aerial if ``is_aerial`` is False.
        """
        if is_aerial and not self.nground() in [0, None]:
            raise ValueError("graph vector containing graphs with more than zero ground vertices can't be aerial")
        self._is_aerial = is_aerial

    def insertion(self, position, other, **kwargs):
        """
        Return the insertion of ``other`` into this graph vector at the vertex ``position``.
        """
        if self.is_aerial() and not other.is_aerial():
            raise ValueError("can't insert non-aerial graph vector into aerial graph vector")
        result = UndirectedGraphVector_dict.insertion(self, position, other, **kwargs)
        result.set_aerial(self.is_aerial())
        return result

class FormalityGraphModule_dict(FormalityGraphModule, GraphModule_dict):
    """
    Module spanned by formality graphs (with elements stored as dictionaries).
    """
    def __init__(self, base_ring, graph_basis):
        """
        Initialize this formality graph module.

        INPUT:

        - ``base_ring`` -- a ring, to be used as the ring of coefficients

        - ``graph_basis`` -- a FormalityGraphBasis
        """
        if not isinstance(graph_basis, FormalityGraphBasis):
            raise ValueError('graph_basis must be a FormalityGraphBasis')
        super().__init__(base_ring, graph_basis)
        self.element_class = FormalityGraphVector_dict

class FormalityGraphVector_vector(FormalityGraphVector, GraphVector_vector):
    """
    Vector representing a linear combination of formality graphs (stored as a dictionary of vectors).
    """
    def __init__(self, parent, vectors, is_aerial=False):
        """
        Initialize this graph vector.

        INPUT:

        - ``parent`` -- a FormalityGraphModule

        - ``vectors`` -- a dictionary, mapping tri-gradings to sparse vectors of coefficients with respect to the basis of ``parent``

        - ``is_aerial`` -- (default: False) a boolean, if True then this graph vector will be aerial
        """
        if not isinstance(parent, FormalityGraphModule_vector):
            raise ValueError("parent must be a FormalityGraphModule_vector")
        super().__init__(parent, vectors)
        self._is_aerial = is_aerial

    def nvertices(self):
        """
        Return the number of vertices in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of vertices.
        """
        for grading in self._vectors:
            if not self._vectors[grading].is_zero():
                return grading[0] + grading[1]

    def nedges(self):
        """
        Return the number of edges in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of edges.
        """
        for grading in self._vectors:
            if not self._vectors[grading].is_zero():
                return grading[2]

    def nground(self):
        """
        Return the number of ground vertices in each graph in this graph vector.

        ASSUMPTIONS:

        Assumes all graphs in this graph vector have the same number of ground vertices.
        """
        for grading in self._vectors:
            if not self._vectors[grading].is_zero():
                return grading[0]

    def is_aerial(self):
        """
        Return True if this graph vector is aerial, and False otherwise.
        """
        return self._is_aerial

    def set_aerial(self, is_aerial=True):
        """
        Set this graph vector to be aerial if ``is_aerial`` is True, respectively not aerial if ``is_aerial`` is False.
        """
        if is_aerial and not self.nground() in [0, None]:
            raise ValueError("graph vector containing graphs with more than zero ground vertices can't be aerial")
        self._is_aerial = is_aerial

    def insertion(self, position, other, **kwargs):
        """
        Return the insertion of ``other`` into this graph vector at the vertex ``position``.
        """
        if self.is_aerial() and not other.is_aerial():
            raise ValueError("can't insert non-aerial graph vector into aerial graph vector")
        result = UndirectedGraphVector_vector.insertion(self, position, other, **kwargs)
        result.set_aerial(self.is_aerial())
        return result

class FormalityGraphModule_vector(FormalityGraphModule, GraphModule_vector):
    """
    Module spanned by formality graphs (with elements stored as dictionaries of vectors).
    """
    def __init__(self, base_ring, graph_basis, vector_constructor, matrix_constructor):
        """
        Initialize this formality graph module.

        INPUT:

        - ``base_ring`` -- a ring, to be used as the ring of coefficients

        - ``graph_basis`` -- a FormalityGraphBasis

        - ``vector_constructor`` -- constructor of (sparse) vectors

        - ``matrix_constructor`` -- constructor of (sparse) matrices
        """
        if not isinstance(graph_basis, FormalityGraphBasis):
            raise ValueError('graph_basis must be a FormalityGraphBasis')
        super().__init__(base_ring, graph_basis, vector_constructor, matrix_constructor)
        self.element_class = FormalityGraphVector_vector

max_stars_repo_name: rburing/gcaops

max_issues_repo_name: rburing/gcaops

max_forks_repo_name: rburing/gcaops

reason: import sage,from sage


text:
import numpy as np

def clip(x, min_x, max_x):
    """
    clips x between min_x and max_x
    """
    return max(min(x, max_x), min_x)

def get_angle_sequence_rad(angle):
    """
    cuts an angle bigger than pi into a sequence of angles <= pi/2
    """
    aa = abs(angle)
    if aa >= np.pi:  # equality is important because the direction is undefined
        chunks = int(aa // (np.pi / 2.0))
        last = aa % (np.pi / 2.0)
        res = [np.pi / 2.0, ] * chunks
        if last != 0:
            res.append(last)
    else:
        res = [aa, ]
    res = np.array(res) * np.sign(angle)
    return res.tolist()

def smallest_angle_diff_rad(theta_new, theta_old):
    """
    computes theta_new - theta_old (radian)
    """
    diff = theta_new - theta_old
    mod_p = diff % (2 * np.pi)
    mod_n = diff % (-2 * np.pi)
    return mod_p if abs(mod_p) < abs(mod_n) else mod_n

max_stars_repo_name: thecognifly/cognifly-python

max_issues_repo_name: thecognifly/cognifly-python

max_forks_repo_name: thecognifly/cognifly-python


text:
# vortex.py
# -------------------------------------------------------------------------
# Create a quiver plot.
# -------------------------------------------------------------------------
import numpy as np
import matplotlib.pyplot as plt

coords = np.linspace(-1, 1, 11)
X, Y = np.meshgrid(coords, coords)
Vx, Vy = Y, -X

plt.figure()
# pivot='mid' places centers of arrows on grid points
plt.quiver(X, Y, Vx, Vy, pivot='mid')
plt.axis('square')
plt.show()

max_stars_repo_name: KinderPy/code-samples

max_issues_repo_name: KinderPy/code-samples

max_forks_repo_name: KinderPy/code-samples


text:
# The content of this file is based on: OpenAI Spinning Up https://spinningup.openai.com/.
from typing import Optional, List

import numpy as np

from molgym.spaces import ObservationType
from molgym.tools import util
from molgym.tools.mpi import mpi_mean_std

class PPOBuffer:
    """
    A buffer for storing trajectories experienced by a PPO agent interacting
    with the environment, and using Generalized Advantage Estimation (GAE-Lambda)
    for calculating the advantages of state-action pairs.
    """
    def __init__(self, int_act_dim: int, size: int, gamma=0.99, lam=0.95) -> None:
        self.obs_buf: List[Optional[ObservationType]] = [None] * size
        self.act_buf = np.empty((size, int_act_dim), dtype=np.float32)
        self.rew_buf = np.zeros(size, dtype=np.float32)
        self.next_obs_buf: List[Optional[ObservationType]] = [None] * size
        self.term_buf = np.zeros(size, dtype=np.bool)

        self.val_buf = np.zeros(size, dtype=np.float32)
        self.logp_buf = np.zeros(size, dtype=np.float32)

        # Filled when path is finished
        self.adv_buf = np.zeros(size, dtype=np.float32)
        self.ret_buf = np.zeros(size, dtype=np.float32)

        self.gamma = gamma
        self.lam = lam

        self.ptr = 0
        self.path_start_idx = 0
        self.max_size = size

    def store(self, obs: ObservationType, act: np.ndarray, reward: float, next_obs: ObservationType, terminal: bool,
              value: float, logp: float):
        """Append one time step of agent-environment interaction to the buffer."""
        assert self.ptr < self.max_size  # buffer has to have room so you can store

        self.obs_buf[self.ptr] = obs
        self.act_buf[self.ptr] = act
        self.rew_buf[self.ptr] = reward
        self.next_obs_buf[self.ptr] = next_obs
        self.term_buf[self.ptr] = terminal

        self.val_buf[self.ptr] = value
        self.logp_buf[self.ptr] = logp

        self.ptr += 1

    def finish_path(self, last_val: float) -> float:
        """
        Call this at the end of a trajectory, or when one gets cut off
        by an epoch ending. This looks back in the buffer to where the
        trajectory started, and uses rewards and value estimates from
        the whole trajectory to compute advantage estimates with GAE-Lambda,
        as well as compute the rewards-to-go for each state, to use as
        the targets for the value function.

        The "last_val" argument should be 0 if the trajectory ended
        because the agent reached a terminal state (died), and otherwise
        should be V(s_T), the value function estimated for the last state.
        This allows us to bootstrap the reward-to-go calculation to account
        for timesteps beyond the arbitrary episode horizon (or epoch cutoff).
        """

        path_slice = slice(self.path_start_idx, self.ptr)
        rews = np.append(self.rew_buf[path_slice], last_val)
        vals = np.append(self.val_buf[path_slice], last_val)

        # the next two lines implement GAE-Lambda advantage calculation
        deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1]
        self.adv_buf[path_slice] = util.discount_cumsum(deltas, self.gamma * self.lam)

        # the next line computes rewards-to-go, to be targets for the value function
        self.ret_buf[path_slice] = util.discount_cumsum(rews, self.gamma)[:-1]
        episodic_return = self.ret_buf[self.path_start_idx]

        self.path_start_idx = self.ptr

        return episodic_return

    def get(self):
        """
        Call this at the end of an epoch to get all of the data from
        the buffer, with advantages appropriately normalized (shifted to have
        mean zero and std one). Also, resets some pointers in the buffer.
        """
        assert self.is_full()  # buffer has to be full before you can get
        self.ptr, self.path_start_idx = 0, 0

        # the next two lines implement the advantage normalization trick
        adv_mean, adv_std = mpi_mean_std(self.adv_buf, axis=-1)
        self.adv_buf = (self.adv_buf - adv_mean) / adv_std

        return dict(obs=self.obs_buf, act=self.act_buf, ret=self.ret_buf, adv=self.adv_buf, logp=self.logp_buf)

    def is_full(self) -> bool:
        return self.ptr == self.max_size

max_stars_repo_name: polyzer/molgym

max_issues_repo_name: polyzer/molgym

max_forks_repo_name: polyzer/molgym


text:
import os
import shutil
import xml.etree.ElementTree as ET
import numpy as np
import cv2
import matplotlib.pyplot as plt
import csv

######################
## HELPER FUNCTIONS ##
######################

# Get files's name that match the given the file type containing in the given root directory
def filelist(root, f_type):
    # INPUT:
    #   root (root directory),
    #   f_type (type of file such as '.png', '.jpg', or '.xml')
    # OUTPUT: list of files' name under the root directory
    return [os.path.join(d_path, f) for d_path, d_name, files in os.walk(root) for f in files if f.endswith(f_type)]

# Retrieve images' info (file name, width, height, bounding box's location, and index) sorted by filename from .xml files
def get_image_info_xml(annotations_path):
    # INPUT: annotations_path (path to the directory that contains .xml annotation files)
    # OUTPUT: 2D NumPy array containing images' info extracted from .xml files
    annotations = filelist(annotations_path, '.xml')
    annotations_list = []
    for path in annotations:
        root = ET.parse(path).getroot()
        filename = int((root.find("./filename").text)[4:-4])
        width = int(root.find("./size/width").text)
        height = int(root.find("./size/height").text)
        x_min = int(root.find("./object/bndbox/xmin").text)
        y_min = int(root.find("./object/bndbox/ymin").text)
        x_max = int(root.find("./object/bndbox/xmax").text)
        y_max = int(root.find("./object/bndbox/ymax").text)
        annotation = [filename, width, height, x_min, y_min, x_max, y_max]
        annotations_list.append(annotation)
    annotations_list = np.array(annotations_list)
    sorted = annotations_list[np.argsort(annotations_list[:, 0])]
    sorted = np.hstack((sorted, np.arange(sorted.shape[0]).reshape(sorted.shape[0], 1)))
    return sorted

# Retrieve images' info (row number, image name, and bounding box's location) from .csv file
def get_image_info_csv(annotations_path):
    # INPUT: annotations_path (path to .csv annotation file)
    # OUTPUT: 2D NumPy array containing images' info extracted from .csv file
    data = []
    with open(annotations_path) as csvfile:
        reader = csv.reader(csvfile, delimiter=",")
        # Skip the first row (which contains columns' names)
        next(reader)
        for row in reader:
            # r = []
            # for item in row:
            # r.append(item)
            data.append(row)
    return np.array(data)

# Read an image from the given file's location and convert into RGB values
def read_image(path):
    # INPUT: path (path to image file's location)
    # OUTPUT: 3D NumPy array of image data in RGB values
    img = cv2.imread(path)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Write an image to the given path and convert back into BGR values
def write_image(path, img):
    # INPUT:
    #   path (path to image file's location),
    #   img (3D NumPy array of image data in RGB values)
    # OUTPUT: -
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(path, img)

# Create a mask from the given bounding box array of the same shape as image
def boundingbox_to_mask(bounding_box, img):
    # INPUT:
    #   bounding_box (1D NumPy array with size = 4 containing bounding box's location [x_min y_min x_max y_max]),
    #   img (3D NumPy array of image data)
    # OUTPUT: 2D NumPy array containing 1's at the corresponding location of bounding box and 0's otherwise
    mask = np.zeros((img.shape[0], img.shape[1]))
    mask[bounding_box[1]:bounding_box[3]+1, bounding_box[0]:bounding_box[2]+1] = 1
    return mask

# Convert a mask to a bounding box array
def mask_to_boundingbox(mask):
    # INPUT: mask (2D NumPy array containing 1's if it is a bounding box's location, otherwise 0's)
    # OUTPUT: 1D NumPy array with size = 4 containing bounding box's location [y_min x_min y_max x_max]
    rows, cols = np.nonzero(mask)
    # If bounding box does not exist, return 1D NumPy array of 4 zeros
    if len(rows) == 0 or len(cols) == 0:
        return np.zeros(4)
    return np.array([np.min(rows), np.min(cols), np.max(rows), np.max(cols)])

# Create a bounding box array from each image's info in 2D NumPy array images' info extracted from .xml files
def create_boundingbox_array_xml(row):
    # INPUT: row (1D NumPy array containing single image's info [filename width height x_min y_min x_max y_max index])
    # OUTPUT: 1D NumPy array with size = 4 containing bounding box's location [x_min y_min x_max y_max]
    return row[3:7]

# Create a bounding box array from each image's info in 2D NumPy array images' info extracted from .csv file(s)
def create_boundingbox_array_csv(row):
    # INPUT: row (1D NumPy array containing single image's info [row_num x_min y_min x_max y_max])
    # OUTPUT: 1D NumPy array with size = 4 containing bounding box's location [x_min y_min x_max y_max]
    return row[1:]

# Preprocess an image
def preprocess_image(img, write_directory, bounding_box, size, index):
    # INPUT:
    #   img (3D NumPy array of image data),
    #   write_directory (path to write a preprocessed image),
    #   bounding_box (1D NumPy array with size = 4 containing bounding box's location [x_min y_min x_max y_max]),
    #   size (size of preprocessed square image),
    #   index (image's index that helps writing the image)
    # OUTPUT: 1D NumPy array with size = 4 containing relocated bounding box's location [y_min x_min y_max x_max]
    # Resize the image to be a square image with the given size
    resize_img = cv2.resize(img, (size, size))
    # Relocate the bounding box's location corresponding to the image resizing
    resized_mask = cv2.resize(boundingbox_to_mask(bounding_box, img), (size, size))
    relocated_boundingbox = mask_to_boundingbox(resized_mask)
    # Write the resized image to the given path
    write_path = write_directory + "/Cars" + str(index) + ".png"
    write_image(write_path, resize_img)
    return relocated_boundingbox

# Delete the directory if it is exist; otherwise, create a blank directory at the given path
def create_directory(path):
    # INPUT: path (path to directory)
    # OUTPUT: -
    if os.path.exists(path):
        shutil.rmtree(path)
    os.makedirs(path)

# Convert the bounding box's location from ratio of height/width into integers
def boundingbox_location_float_to_int(img, img_info):
    # INPUT:
    #   img (3D NumPy array of image data),
    #   img_info (1D NumPy array containing single image's info [row_num x_min y_min x_max y_max]) in float values
    # OUTPUT: 1D NumPy array containing single image's info [row_num x_min y_min x_max y_max] in integer values
    height, width, _ = img.shape
    new_img_info = [0, 0, 0, 0, 0]
    new_img_info[0] = int(img_info[0])
    new_img_info[1] = round(float(img_info[1]) * width)
    new_img_info[2] = round(float(img_info[2]) * height)
    new_img_info[3] = round(float(img_info[3]) * width)
    new_img_info[4] = round(float(img_info[4]) * height)
    return new_img_info

# Find whether the image annotated with .csv file exists
def find_csv_read_path(image_file, full_paths):
    for path in full_paths:
        if image_file in path:
            return True
    return False

########################
## DATA VISUALIZATION ##
########################

# Visualizer
def show_corner_bb(im, bb):
    def create_corner_rect(bb, color="red"):
        bb = np.array(bb, dtype=np.float32)
        return plt.Rectangle((bb[1], bb[0]), bb[3]-bb[1], bb[2]-bb[0], color=color, fill=False, lw=4)
    plt.imshow(im)
    plt.gca().add_patch(create_corner_rect(bb))
    plt.show()

#######################
## DATA AUGMENTATION ##
#######################

def rotate_image(img, deg, isMask, mode=cv2.BORDER_REFLECT, interpolation=cv2.INTER_AREA):
    height, width, _ = img.shape
    rotationMatrix = cv2.getRotationMatrix2D((height/2, width/2), deg, 1)
    if isMask:
        return cv2.warpAffine(img, rotationMatrix, (height, width), borderMode=cv2.BORDER_CONSTANT)
    else:
        return cv2.warpAffine(img, rotationMatrix, (height, width), borderMode=mode, flags=cv2.WARP_FILL_OUTLIERS+interpolation)

def transformsXY(path, bounding_box):
    img = cv2.imread(str(path)).astype(np.float32)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)/255
    mask = boundingbox_to_mask(bounding_box, img)
    rotation_deg = (np.random.random() - 0.5) * 20
    img = rotate_image(img, rotation_deg, False)
    mask = rotate_image(mask, rotation_deg, True)
    if np.random.random() > 0.5:
        img = np.fliplr(img).copy()
        mask = np.fliplr(mask).copy()
    return img, mask_to_boundingbox(mask)

###################
## MAIN FUNCTION ##
###################

def preprocessing(output_image_size=224, training_size=0.85):
    print("Begin preprocessing ...")

    # Input paths
    xml_images_read_path = "raw_data/dataset1/images"
    xml_annotations_read_path = "raw_data/dataset1/annotations"
    csv_images_read_path = "raw_data/dataset5/Cars"
    csv_annotations_read_path = "raw_data/dataset5/license.csv"

    # Output paths
    train_images_directory_path = "train_data"
    train_images_write_path = train_images_directory_path + "/data"
    create_directory(train_images_write_path)
    train_boundingboxes_write_path = train_images_directory_path + "/boundingbox.csv"
    test_images_directory_path = "test_data"
    test_images_write_path = test_images_directory_path + "/data"
    create_directory(test_images_write_path)
    test_boundingboxes_write_path = test_images_directory_path + "/boundingbox.csv"

    csv_images_info = get_image_info_csv(csv_annotations_read_path)
    csv_images_info = np.delete(csv_images_info, 1, 1)
    csv_images_filelist = filelist(csv_images_read_path, ".jpg")
    csv_cutoff_index = int(training_size * len(csv_images_info))
    xml_images_info = get_image_info_xml(xml_annotations_read_path)
    xml_cutoff_index = int(training_size * len(xml_images_info))

    # Training Data
    csv_train_images_info = csv_images_info[:csv_cutoff_index]
    index = 0
    train_boundingboxes_info = []
    for img_info in csv_train_images_info:
        isExist = find_csv_read_path("car" + str(img_info[0]) + ".jpg", csv_images_filelist)
        if isExist:
            read_path = csv_images_read_path + "/car" + str(img_info[0]) + ".jpg"
            img = read_image(read_path)
            new_img_info = boundingbox_location_float_to_int(img, img_info)
            bounding_box = create_boundingbox_array_csv(new_img_info)
            relocated_boundingbox = preprocess_image(img, train_images_write_path, bounding_box, output_image_size, index)
            boundingbox_info = [index] + list(relocated_boundingbox)
            train_boundingboxes_info.append(boundingbox_info)
            index += 1
    xml_train_images_info = xml_images_info[:xml_cutoff_index]
    for img_info in xml_train_images_info:
        read_path = xml_images_read_path + "/Cars" + str(img_info[0]) + ".png"
        img = read_image(read_path)
        bounding_box = create_boundingbox_array_xml(img_info)
        relocated_boundingbox = preprocess_image(img, train_images_write_path, bounding_box, output_image_size, index)
        boundingbox_info = [index] + list(relocated_boundingbox)
        train_boundingboxes_info.append(boundingbox_info)
        index += 1
    np.savetxt(train_boundingboxes_write_path, train_boundingboxes_info, delimiter=',', fmt='%d')

    # Testing Data
    csv_test_images_info = csv_images_info[csv_cutoff_index:]
    index = 0
    test_boundingboxes_info = []
    for img_info in csv_test_images_info:
        isExist = find_csv_read_path("car" + str(img_info[0]) + ".jpg", csv_images_filelist)
        if isExist:
            read_path = csv_images_read_path + "/car" + str(img_info[0]) + ".jpg"
            img = read_image(read_path)
            new_img_info = boundingbox_location_float_to_int(img, img_info)
            bounding_box = create_boundingbox_array_csv(new_img_info)
            relocated_boundingbox = preprocess_image(img, test_images_write_path, bounding_box, output_image_size, index)
            boundingbox_info = [index] + list(relocated_boundingbox)
            test_boundingboxes_info.append(boundingbox_info)
            index += 1
    xml_test_images_info = xml_images_info[xml_cutoff_index:]
    for img_info in xml_test_images_info:
        read_path = xml_images_read_path + "/Cars" + str(img_info[0]) + ".png"
        img = read_image(read_path)
        bounding_box = create_boundingbox_array_xml(img_info)
        relocated_boundingbox = preprocess_image(img, test_images_write_path, bounding_box, output_image_size, index)
        boundingbox_info = [index] + list(relocated_boundingbox)
        test_boundingboxes_info.append(boundingbox_info)
        index += 1
    np.savetxt(test_boundingboxes_write_path, test_boundingboxes_info, delimiter=',', fmt='%d')

    print("Preprocessing Done!")

if __name__ == "__main__":
    preprocessing()

max_stars_repo_name: jaysarva/Automatic-License-Plate-Search

max_issues_repo_name: jaysarva/Automatic-License-Plate-Search

max_forks_repo_name: jaysarva/Automatic-License-Plate-Search


text:
#!/usr/bin/python3

from __future__ import absolute_import, division, print_function
from . import config
import logging
import os
import fnmatch
import subprocess
import numpy as np
import skimage.transform as st
from . import helpers

logger = logging.getLogger(__name__)

UM_PER_CM = 1e4

def sart_n(im, n=1, imr=None, theta=None, **kwargs):
    """Perform n iterations of sart reconstruction

    Parameters
    ----------
    im : 2d ndarray of floats
        input sinogram
    n : int
        number of sart iterations to perform
    imr : 2d ndarray of floats
        initial image estimate
    theta : array-like iterable of projection angles
        If None, infer n angles are from 0 to 180 (non-inclusive), where n is the number of
        sinogram columns in im

    Returns
    -------
    2d ndarray of floats

    """
    for i in range(n):
        imr = st.iradon_sart(im, theta=theta, image=imr)
    return imr

def iradon_absorption(sino, angles):
    """Wrapper for absorption (as opposed to emission) iradon call for getting the matrix
    estimate

    Parameters
    ----------
    sino :  2d ndarray of floats
        sinogram to backproject
    angles : 1d array-like
        sinogram projection angles

    Returns
    -------
    2d ndarray of float32's

    """
    outdir = config.xlict_recon_mpi_fbp_PATH_OUTPUT_DATA  # output directory
    indir = config.xlict_recon_mpi_fbp_PATH_INPUT_TIFFS  # input directory
    base_filename = 'temp_sino.tif'

    def xtract_exe_string(filename, indir, outdir, angles, ps,
                          xlict_recon_mpi_fbp_filter='ramp'):
        assert xlict_recon_mpi_fbp_filter in {'ramp', 'shepp-logan', 'cosine',
                                              'hamming', 'hann', 'none'}
        # lookup dict for reconstruction filter type and filter enable flag
        filter_control = {
            'ramp': (0, 1),
            'shepp-logan': (1, 1),
            'cosine': (2, 1),
            'hamming': (3, 1),
            'hann': (4, 1),
            'none': (0, 0),
        }

        recon_filter, filter_enable = filter_control[xlict_recon_mpi_fbp_filter]
        implementation = 'xlict_recon_mpi_fbp'
        rm = {'xlict_recon_mpi_fbp': 0, 'xlict_recon_gridrec': 1}[implementation]
        xtract_options = \
            r'-id {indir} --sino {prg} -od {outdir} -pfct r_.tif -rmu 1 ' \
            r'-e {e_keV} --recon_method {rm} -fr {recon_filter} ' \
            r'-as {astep} --recon_filter_enabled {filter_enable} ' \
            r'-ps {ps} -corm {corm} --force_cpu {force_cpu}'.format(
                indir=indir,
                prg=filename,  # read file name
                outdir=outdir,
                e_keV=config.energy_keV,
                rm=rm,  # 0 (FBP), 1 (Gridrec)
                recon_filter=recon_filter,
                filter_enable=filter_enable,
                astep=angles[1] - angles[0],  # Angle step (deg)
                ps=ps,  # pixel size (um)
                force_cpu=config.xlict_force_cpu,
                corm=config.xlict_recon_mpi_corm,  # centre-of-rot'n offset
            )
        return xtract_options

    sino = sino.T  # orient sinogram axes to what X-TRACT expects
    # write to disk for XTRACT
    filename = os.path.join(indir, base_filename)
    helpers.write_tiff32(filename, sino)

    # reconstruct (XLI/XTRACT writes the result to disk), then read result from disk
    xes = xtract_exe_string(base_filename, indir, outdir, angles, ps=1)
    PATH_XTRACT_EXE = config.xlict_recon_mpi_exe
    with open(os.devnull, 'w') as fnull:
        _ = subprocess.call('{} {}'.format(PATH_XTRACT_EXE, xes),
                            stdout=fnull, stderr=subprocess.STDOUT)

    im = helpers.read_tiff32(os.path.join(outdir, 'r_0.tif'))
    # Reorient the reconstruction to match the convention established by other
    # reconstructors here.
    im = im.T

    # Apply a hard circular mask to the result, to match the behaviour of skimage
    im = helpers.zero_outside_circle(im)

    # write im to the output path with name r_<nnn+1>.tif,
    # where nnn is the highest number in the existing files
    files = os.listdir(outdir)
    if 'r_0.tif' in files:
        matches = sorted(fnmatch.filter(files, 'r_[0-9][0-9][0-9][0-9].tif'))
        if not matches:
            dest = 'r_0000.tif'
        else:
            # files exist matching r_nnn.tif; get the highest and add 1
            nnn = int(matches[-1][2:5])
            dest = 'r_%04d.tif' % (nnn + 1)
        helpers.write_tiff32(os.path.join(outdir, dest), im)

    # im *= 1e-4  # This factor was determined by direct comparison with iradon above

    return im

def iradon(sino, angles, filter='none'):
    """Wrapper for potentially a bunch of different backprojector
    implementations, wrapped with a common call signature. Selection of
    the implementation and specific parameters for individual implementations
    are pulled from the config file.

    Parameters
    ----------
    sino : 2d ndarray of floats
        sinogram to backproject
    angles : 1d array-like
        sinogram projection angles
    filter : str
        A string that selects the iradon filter type based on those supported by xlict.
        One of 'ramp', 'shepp-logan', 'cosine', 'hamming', 'hann', 'none'

    Returns
    -------
    2d ndarray of float32's

    """
    implementation = config.iradon_implementation

    if implementation == 'skimage_iradon':
        im = st.iradon(sino, theta=angles,
                       circle=config.skimage_iradon_circle,
                       filter=config.skimage_iradon_filter)

    elif implementation == 'skimage_iradon_sart':
        # Clip solution to positive values. MLEM requires this.
        im = st.iradon_sart(sino, theta=angles,
                            clip=(1e-6, sino.max() * sino.shape[0] * sino.shape[1]))

    elif implementation in ('xlict_recon_mpi_fbp', 'xlict_recon_gridrec'):
        outdir = config.xlict_recon_mpi_fbp_PATH_OUTPUT_DATA  # output directory
        indir = config.xlict_recon_mpi_fbp_PATH_INPUT_TIFFS  # input directory
        base_filename = 'temp_sino.tif'

        def xtract_exe_string(filename, indir, outdir, angles, ps):
            assert filter in {'ramp', 'shepp-logan', 'cosine',
                              'hamming', 'hann', 'none'}
            # lookup dict for reconstruction filter type and filter enable flag
            filter_control = {
                'ramp': (0, 1),
                'shepp-logan': (1, 1),
                'cosine': (2, 1),
                'hamming': (3, 1),
                'hann': (4, 1),
                'none': (0, 0),
            }

            recon_filter, filter_enable = filter_control[filter]
            rm = {'xlict_recon_mpi_fbp': 0, 'xlict_recon_gridrec': 1}[implementation]
            xtract_options = \
                r'-id {indir} --sino {prg} -od {outdir} -pfct r_.tif -rmu 1 ' \
                r'-e {e_keV} --recon_method {rm} -fr {recon_filter} ' \
                r'-as {astep} --recon_filter_enabled {filter_enable} ' \
                r'-ps {ps} -corm {corm} --force_cpu {force_cpu}'.format(
                    indir=indir,
                    prg=filename,  # read file name
                    outdir=outdir,
                    e_keV=config.energy_keV,
                    rm=rm,  # 0 (FBP), 1 (Gridrec)
                    recon_filter=recon_filter,
                    filter_enable=filter_enable,
                    astep=angles[1] - angles[0],  # Angle step (deg)
                    ps=ps,  # pixel size (um)
                    force_cpu=config.xlict_force_cpu,
                    corm=config.xlict_recon_mpi_corm,  # centre-of-rot'n offset
                )
            return xtract_options

        sino = sino.T  # orient sinogram axes to what X-TRACT expects
        # write to disk for XTRACT
        filename = os.path.join(indir, base_filename)
        helpers.write_tiff32(filename, sino)

        # reconstruct (XLI/XTRACT writes the result to disk), then read result from disk
        xes = xtract_exe_string(base_filename, indir, outdir, angles, ps=1)
        PATH_XTRACT_EXE = config.xlict_recon_mpi_exe
        with open(os.devnull, 'w') as fnull:
            _ = subprocess.call('{} {}'.format(PATH_XTRACT_EXE, xes),
                                stdout=fnull, stderr=subprocess.STDOUT)

        im = helpers.read_tiff32(os.path.join(outdir, 'r_0.tif'))
        # Reorient the reconstruction to match the convention established by other
        # reconstructors here.
        im = im.T

        # Apply a hard circular mask to the result, to match the behaviour of skimage
        im = helpers.zero_outside_circle(im)

        # write im to the output path with name r_<nnn+1>.tif,
        # where nnn is the highest number in the existing files
        files = os.listdir(outdir)
        if 'r_0.tif' in files:
            matches = sorted(fnmatch.filter(files, 'r_[0-9][0-9][0-9][0-9].tif'))
            if not matches:
                dest = 'r_0000.tif'
            else:
                # files exist matching r_nnn.tif; get the highest and add 1
                nnn = int(matches[-1][2:5])
                dest = 'r_%04d.tif' % (nnn + 1)
            helpers.write_tiff32(os.path.join(outdir, dest), im)

        im *= 5e-5  # This factor was determined by direct comparison with iradon above

    return im

if __name__ == "__main__":
    angles = np.linspace(0, 360, 180, endpoint=False)
    side = 5
    m = np.zeros((5, 5))
    m[2, 2] = 1.0
    p = st.radon(m, angles, circle=True)
    config.iradon_implementation = 'skimage_iradon'
    im1 = iradon(p, angles)
    config.iradon_implementation = 'xlict_recon_mpi_fbp'
    # config.iradon_implementation = 'xlict_recon_gridrec'
    im2 = iradon(p, angles)
    print(im1)
    print(im2)
    print('scale_factor =', im1 / im2)