text:
import numpy as np
import pandas as pd
import psp.MD_lib as MDlib
import time
import os
import psp.PSP_lib as bd
from openbabel import openbabel as ob
import glob
import psp.output_lib as lib
from tqdm import tqdm
from LigParGenPSP import Converter
import psp.MoleculeBuilder as mb
import random

obConversion = ob.OBConversion()

class Builder:
    def __init__(
        self,
        Dataframe,
        ID_col="ID",
        SMILES_col="smiles",
        NumMole="Num",
        Length="Len",
        NumConf="NumConf",
        NumModel=1,
        LeftCap="LeftCap",
        RightCap="RightCap",
        Loop="Loop",
        OutFile="amor_model",
        OutDir="amorphous_models",
        OutDir_xyz="molecules",
        density=0.65,
        tol_dis=2.0,
        box_type="c",
        box_size=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        incr_per=0.4,
        BondInfo=True,
    ):
        self.Dataframe = Dataframe
        self.ID_col = ID_col
        self.SMILES_col = SMILES_col
        self.NumMole = NumMole
        self.Length = Length
        self.NumConf = NumConf
        self.NumModel = NumModel
        self.LeftCap = LeftCap
        self.RightCap = RightCap
        self.Loop = Loop
        self.OutFile = OutFile
        self.OutDir = os.path.join(OutDir, "")
        self.OutDir_xyz = os.path.join(OutDir, OutDir_xyz, "")
        self.OutDir_packmol = os.path.join(OutDir, "packmol", "")
        self.OutDir_ligpargen = os.path.join(OutDir, "ligpargen", "")
        self.OutDir_pysimm = os.path.join(OutDir, "pysimm", "")
        self.density = density
        self.tol_dis = tol_dis
        self.box_type = box_type
        self.box_size = box_size
        self.incr_per = incr_per
        self.BondInfo = BondInfo

    def Build(self):
        start_1 = time.time()
        lib.print_psp_info()  # Print PSP info
        lib.print_input("AmorphousBuilder", self.Dataframe)
        if self.box_type == "c":
            box_type_ = "Cubic"
        else:
            box_type_ = "Rectangular"

        print(
            "\n",
            "Additional information: ",
            "\n",
            "Number of models: ",
            self.NumModel,
            "\n",
            "Density (g/cm3): ",
            self.density,
            "\n",
            "Tolerance distance (angstrom): ",
            self.tol_dis,
            "\n",
            "Box type: ",
            box_type_,
            "\n",
            "Output directory: ",
            self.OutDir,
            "\n",
        )

        # location of directory for VASP inputs (polymers) and build a directory
        bd.build_dir(self.OutDir)
        bd.build_dir(self.OutDir_xyz)

        # PACKMOL
        packmol_path = os.getenv("PACKMOL_EXEC")
        # packmol_path = '/home/hari/.soft/packmol/packmol'

        xyz_gen_pd = pd.DataFrame()
        for i in self.Dataframe.index:
            df = pd.DataFrame(self.Dataframe.loc[i]).T

            mol = mb.Builder(
                df,
                ID_col=self.ID_col,
                SMILES_col=self.SMILES_col,
                LeftCap=self.LeftCap[0],
                RightCap=self.RightCap[0],
                OutDir=self.OutDir_xyz,
                Length=[int(df[self.Length].values)],
                NumConf=int(df[self.NumConf].values) * self.NumModel,
                Loop=eval(str(df[self.Loop].values[0])),
                NCores=-1,
                Subscript=True,
            )
            results = mol.Build()
            xyz_gen_pd = pd.concat([xyz_gen_pd, results])

        if len(list(set(xyz_gen_pd["Result"].values))) != 1:
            xyz_gen_pd.to_csv("molecules.csv")
            print(
                "Couldn't generate XYZ coordinates of molecules, check 'molecules.csv'"
            )

        XYZ_list, smi_list, NMol_list, NumConf_list = [], [], [], []
        for index, row in self.Dataframe.iterrows():
            # Get number of molecules for each conformer of molecules
            NMol_list += [int(row[self.NumMole] / row[self.NumConf])] * row[
                self.NumConf
            ]

            # Get SMILES string for oligomers
            smiles_each = xyz_gen_pd[xyz_gen_pd['ID'] == row['ID']]['SMILES'].values[0]
            smi_list += smiles_each * row[self.NumConf]

            # Get a list of filenames for XYZ coordinates
            XYZ_list_ind = glob.glob(self.OutDir_xyz + str(row[self.ID_col]) + "*.pdb")
            XYZ_list.append(XYZ_list_ind)
            NumConf_list.append(int(row[self.NumConf]))

        # Define boundary conditions
        if max(self.box_size) == 0.0:  # Box size is not provided
            NMol_type = len(NMol_list)
            Total_NMol = sum(NMol_list)
            total_vol = 0
            for i in range(NMol_type):
                molar_mass = MDlib.get_molar_mass(smi_list[i])
                total_vol += MDlib.get_vol(self.density, NMol_list[i], molar_mass)
            self.box_size = MDlib.get_box_size(
                total_vol, box_type=self.box_type, incr_per=self.incr_per
            )

        xmin, xmax, ymin, ymax, zmin, zmax = (
            self.box_size[0],
            self.box_size[1],
            self.box_size[2],
            self.box_size[3],
            self.box_size[4],
            self.box_size[5],
        )

        fix_dis = self.tol_dis / 2

        ind_mol_count = [0] * len(NumConf_list)
        count_model = 0
        for model in tqdm(range(1, self.NumModel + 1), desc='Building models ...'):
            if self.NumModel > 1:
                print("MODEL ", model)
                packmol_outdir_model = self.OutDir_packmol[:-1] + '_' + str(model) + "/"
                bd.build_dir(packmol_outdir_model)

                XYZ_list_ind_model = []
                count_mol = 0
                for ind_list in XYZ_list:
                    if len(ind_list) >= (count_model + 1) * NumConf_list[count_mol]:
                        XYZ_list_ind_model.append(
                            ind_list[
                                count_model
                                * NumConf_list[count_mol]: (count_model + 1)
                                * NumConf_list[count_mol]
                            ]
                        )
                    else:
                        XYZ_list_ind_model.append(
                            random.sample(ind_list, NumConf_list[count_mol])
                        )

                    count_mol += 1

                XYZ_list_model = [
                    item for sublist in XYZ_list_ind_model for item in sublist
                ]
                count_model += 1
            else:
                bd.build_dir(self.OutDir_packmol)

                packmol_outdir_model = self.OutDir_packmol
                XYZ_list_model = [item for sublist in XYZ_list for item in sublist]
            # exit()
            # PACKMOL input file
            MDlib.gen_packmol_inp(
                packmol_outdir_model,
                self.tol_dis,
                XYZ_list_model,
                NMol_list,
                xmin + fix_dis,
                xmax - fix_dis,
                ymin + fix_dis,
                ymax - fix_dis,
                zmin + fix_dis,
                zmax - fix_dis,
            )
            # PACKMOL calculation
            command = (
                packmol_path + " < " + os.path.join(packmol_outdir_model, "packmol.inp")
            )
            errout = MDlib.run_packmol(
                command, os.path.join(packmol_outdir_model, "packmol.out")
            )

            if errout is not None:
                print(" Error in packmol calculation")
                exit()
            elif (
                os.path.exists(os.path.join(packmol_outdir_model, "packmol.pdb"))
                is False
            ):
                print(" Error in packmol calculation")
                exit()

            mol = ob.OBMol()
            obConversion = ob.OBConversion()
            obConversion.SetInAndOutFormats("pdb", "mol2")
            obConversion.ReadFile(
                mol, os.path.join(packmol_outdir_model, "packmol.pdb")
            )
            obConversion.WriteFile(
                mol, os.path.join(packmol_outdir_model, "packmol.mol2")
            )

            packmol_xyz = MDlib.read_mol2_xyz(
                os.path.join(packmol_outdir_model, "packmol.mol2")
            )
            packmol_bond = MDlib.read_mol2_bond(
                os.path.join(packmol_outdir_model, "packmol.mol2")
            )

            # Output filename
            if self.NumModel > 1:
                output_filename = self.OutFile + "_N" + str(count_model)
            else:
                output_filename = self.OutFile

            MDlib.gen_sys_vasp(
                os.path.join(self.OutDir, output_filename + ".vasp"),
                packmol_xyz,
                xmin,
                xmax,
                ymin,
                ymax,
                zmin,
                zmax,
            )
            MDlib.gen_sys_data(
                os.path.join(self.OutDir, output_filename + ".data"),
                packmol_xyz,
                packmol_bond,
                xmin,
                xmax,
                ymin,
                ymax,
                zmin,
                zmax,
                self.BondInfo,
            )
        end_1 = time.time()
        lib.print_out(
            pd.DataFrame(), "Amorphous model", np.round((end_1 - start_1) / 60, 2)
        )

    def get_opls(self, output_fname='amor_opls.lmps'):
        print("\nGenerating OPLS parameter file ...\n")
        system_pdb_fname = os.path.join(self.OutDir_packmol, "packmol.pdb")
        r = MDlib.get_coord_from_pdb(system_pdb_fname)

        bd.build_dir(self.OutDir_ligpargen)

        system_stats = {
            'total_atoms': 0,
            'total_bonds': 0,
            'total_angles': 0,
            'total_dihedrals': 0,
            'total_impropers': 0,
            'total_atom_types': 0,
            'total_bond_types': 0,
            'total_angle_types': 0,
            'total_dihedral_types': 0,
            'total_improper_types': 0,
        }
        dicts = []

        # run LigParGen for every pdb file in the OutDir_xyz directory
        for index, row in self.Dataframe.iterrows():
            _id = str(row[self.ID_col])
            _length = row[self.Length]
            _num = row[self.NumMole]
            _conf = 1  # read in only the first conformer
            output_prefix = "{}_N{}_C{}".format(_id, _length, _conf)
            lig_output_fname = "{}.lmp".format(output_prefix)
            data_fname = os.path.join(self.OutDir_ligpargen, lig_output_fname)

            try:
                print("LigParGen working on {}.pdb".format(output_prefix))
                Converter.convert(
                    pdb=os.path.join(self.OutDir_xyz, output_prefix + '.pdb'),
                    resname=output_prefix,
                    charge=0,
                    opt=0,
                    lbcc=True,
                    outdir='.',
                )
                os.rename(lig_output_fname, data_fname)
            except BaseException:
                print('problem running LigParGen for {}.pdb.'.format(output_prefix))

            # quickly read the headers of LigParGen generated LAMMPS
            # files to count total number of atoms/bonds/angles...etc
            (
                natoms,
                nbonds,
                nangles,
                ndihedrals,
                nimpropers,
                natom_types,
                nbond_types,
                nangle_types,
                ndihedral_types,
                nimproper_types,
            ) = MDlib.read_lmps_header(data_fname)

            system_stats['total_atom_types'] += natom_types
            system_stats['total_bond_types'] += nbond_types
            system_stats['total_angle_types'] += nangle_types
            system_stats['total_dihedral_types'] += ndihedral_types
            system_stats['total_improper_types'] += nimproper_types
            system_stats['total_atoms'] += natoms * _num
            system_stats['total_bonds'] += nbonds * _num
            system_stats['total_angles'] += nangles * _num
            system_stats['total_dihedrals'] += ndihedrals * _num
            system_stats['total_impropers'] += nimpropers * _num

            # this switcher dict is to navigate through and store info for each section of a LAMMPS file
            switcher = {
                'Masses': [],
                'Pair Coeffs': [],
                'Bond Coeffs': [],
                'Angle Coeffs': [],
                'Dihedral Coeffs': [],
                'Improper Coeffs': [],
                'Atoms': [],
                'Bonds': [],
                'Angles': [],
                'Dihedrals': [],
                'Impropers': [],
                'Num': _num,
            }
            current_section = None

            # read all the info in the LigParGen generated LAMMPS file for modification
            with open(data_fname, 'rt') as lines:
                for line in lines:
                    if any(x in line for x in switcher.keys()):
                        current_section = line.strip()
                    elif line == '\n' or not current_section:
                        continue
                    else:
                        section_list = switcher.get(
                            current_section, 'Invalid current section'
                        )
                        section_list.append(line.split())
            dicts.append(switcher)

        lammps_output = os.path.join(self.OutDir, output_fname)
        MDlib.write_lammps_ouput(lammps_output, r, self.box_size, system_stats, dicts)
        print("\nOPLS parameter file generated.")

    def get_gaff2(
        self, output_fname='amor_gaff2.lmps', atom_typing='pysimm', am1bcc_charges=False, swap_dict=None
    ):
        print("\nGenerating GAFF2 parameter file ...\n")
        system_pdb_fname = os.path.join(self.OutDir_packmol, "packmol.pdb")
        r = MDlib.get_coord_from_pdb(system_pdb_fname)

        bd.build_dir(self.OutDir_pysimm)

        system_stats = {
            'total_atoms': 0,
            'total_bonds': 0,
            'total_angles': 0,
            'total_dihedrals': 0,
            'total_impropers': 0,
            'total_atom_types': 0,
            'total_bond_types': 0,
            'total_angle_types': 0,
            'total_dihedral_types': 0,
            'total_improper_types': 0,
        }
        dicts = []

        from pysimm import system, forcefield

        # run Pysimm for every cml (converted from pdb with Babel) file in the OutDir_xyz directory
        for index, row in self.Dataframe.iterrows():
            _id = str(row[self.ID_col])
            _length = row[self.Length]
            _num = row[self.NumMole]
            _conf = 1  # read in only the first conformer
            output_prefix = "{}_N{}_C{}".format(_id, _length, _conf)
            pdb_file = os.path.join(self.OutDir_xyz, "{}.pdb".format(output_prefix))
            cml_file = os.path.join(self.OutDir_xyz, "{}.cml".format(output_prefix))

            obConversion.SetInAndOutFormats("pdb", "cml")
            mol = ob.OBMol()
            obConversion.ReadFile(mol, pdb_file)
            obConversion.WriteFile(mol, cml_file)

            data_fname = os.path.join(
                self.OutDir_pysimm, "{}.lmp".format(output_prefix)
            )

            try:
                print("Pysimm working on {}".format(cml_file))
                s = system.read_cml(cml_file)
            except BaseException:
                print('problem reading {} for Pysimm.'.format(cml_file))
                exit()

            f = forcefield.Gaff2()
            if atom_typing == 'pysimm':
                if am1bcc_charges:
                    print('AM1BCC method is not available with pysimm, using gasteiger method instead')
                for b in s.bonds:
                    if b.a.bonds.count == 3 and b.b.bonds.count == 3:
                        b.order = 4
                s.apply_forcefield(f, charges='gasteiger')
            elif atom_typing == 'antechamber':
                mol2_file = os.path.join(
                    self.OutDir_xyz, "{}.mol2".format(output_prefix)
                )
                obConversion.SetInAndOutFormats("pdb", "mol2")
                mol = ob.OBMol()
                obConversion.ReadFile(mol, pdb_file)
                obConversion.WriteFile(mol, mol2_file)

                print("Antechamber working on {}".format(mol2_file))
                MDlib.get_type_from_antechamber(s, mol2_file, 'gaff2', f, am1bcc_charges, swap_dict)
                s.pair_style = 'lj'
                s.apply_forcefield(f, charges=None if am1bcc_charges else 'gasteiger', skip_ptypes=True)
            else:
                print(
                    'Invalid atom typing option, please select pysimm or antechamber.'
                )
                exit()
            s.write_lammps(data_fname)

            # quickly read the headers of Pysimm generated LAMMPS
            # files to count total number of atoms/bonds/angles...etc
            (
                natoms,
                nbonds,
                nangles,
                ndihedrals,
                nimpropers,
                natom_types,
                nbond_types,
                nangle_types,
                ndihedral_types,
                nimproper_types,
            ) = MDlib.read_lmps_header(data_fname)

            system_stats['total_atom_types'] += natom_types
            system_stats['total_bond_types'] += nbond_types
            system_stats['total_angle_types'] += nangle_types
            system_stats['total_dihedral_types'] += ndihedral_types
            system_stats['total_improper_types'] += nimproper_types
            system_stats['total_atoms'] += natoms * _num
            system_stats['total_bonds'] += nbonds * _num
            system_stats['total_angles'] += nangles * _num
            system_stats['total_dihedrals'] += ndihedrals * _num
            system_stats['total_impropers'] += nimpropers * _num

            # this switcher dict is to navigate through and store info for each section of a LAMMPS file
            switcher = {
                'Masses': [],
                'Pair Coeffs': [],
                'Bond Coeffs': [],
                'Angle Coeffs': [],
                'Dihedral Coeffs': [],
                'Improper Coeffs': [],
                'Atoms': [],
                'Velocities': [],
                'Bonds': [],
                'Angles': [],
                'Dihedrals': [],
                'Impropers': [],
                'Num': _num,
            }
            current_section = None

            # read all the info in the Pysimm generated LAMMPS file for modification
            with open(data_fname, 'rt') as lines:
                for line in lines:
                    if any(x in line for x in switcher.keys()):
                        current_section = line.strip()
                    elif line == '\n' or not current_section:
                        continue
                    else:
                        section_list = switcher.get(
                            current_section, 'Invalid current section'
                        )
                        section_list.append(line.split())
            dicts.append(switcher)

        lammps_output = os.path.join(self.OutDir, output_fname)
        MDlib.write_lammps_ouput(lammps_output, r, self.box_size, system_stats, dicts)
        print("\nGAFF2 parameter file generated.")

max_stars_repo_name: Ramprasad-Group/PSP

max_issues_repo_name: Ramprasad-Group/PSP

max_forks_repo_name: Ramprasad-Group/PSP


text:
import multiprocessing as mp
from multiprocessing import Process
import numpy as np
import pybullet as p
from pybullet_utils import bullet_client
import random
from threading import Thread
import time

import robodesk

import opentamp.main as main
from opentamp.core.parsing import parse_domain_config, parse_problem_config
import opentamp.core.util_classes.common_constants as const
from opentamp.core.util_classes.robots import Baxter
from opentamp.core.util_classes.openrave_body import *
from opentamp.core.util_classes.transform_utils import *
from opentamp.pma.hl_solver import *
from opentamp.pma.pr_graph import *
from opentamp.pma import backtrack_ll_solver_OSQP as bt_ll
from opentamp.pma.robot_solver import RobotSolverOSQP
import opentamp.core.util_classes.transform_utils as T

from opentamp.policy_hooks.utils.file_utils import load_config, setup_dirs, LOG_DIR
from opentamp.policy_hooks.run_training import argsparser
from opentamp.policy_hooks.utils.load_agent import *
import opentamp.policy_hooks.robodesk.hyp as hyp
import opentamp.policy_hooks.robodesk.desk_prob as prob

# def stepSimulation(t_limit=None, delay=0.):
#     import pybullet as op1
#     start_t = time.time()

#     p1 = bullet_client.BulletClient(op1.SHARED_MEMORY)
#     if p1 < 0:
#         print('Subprocess could not find PyBullet server, exiting...')
#         return

#     time.sleep(0.1)
#     print('-----> Subprocess Connected to physics server')
#     while (t_limit is None or time.time() - start_t < t_limit):
#         p1.stepSimulation()
#         if delay > 0: time.sleep(delay)

if __name__ == '__main__':
    args = argsparser()
    args.config = 'opentamp.policy_hooks.robodesk.hyp'
    args.render = True

    base_config = hyp.refresh_config()
    base_config['id'] = 0
    base_config.update(vars(args))
    base_config['args'] = args
    config, config_module = load_config(args, base_config)
    config.update(base_config)
    config['load_render'] = True
    agent_config = load_agent(config)
    agent = build_agent(agent_config)
    env = agent.base_env
    mjc_env = agent.mjc_env

    '''
    agent.add_viewer()
    agent.mjc_env.set_attr('panda', 'right', np.zeros(7), forward=True)
    agent.render_viewer(env.render(resize=True))
    agent.render_viewer(env.render(resize=True))
    agent.render_viewer(env.render(resize=True))
    import ipdb; ipdb.set_trace()
    '''

    try:
        p.disconnect()
    except Exception as e:
        print(e)

    #const.NEAR_GRIP_COEFF = 5e-2
    #const.GRASP_DIST = 0.2
    #const.APPROACH_DIST = 0.025
    #const.EEREACHABLE_ROT_COEFF = 8e-3
    bt_ll.DEBUG = True
    openrave_bodies = None
    domain_fname = opentamp.__path__._path[0] + "/domains/robot_manipulation_domain/right_desk.domain"
    prob = opentamp.__path__._path[0] + "/domains/robot_manipulation_domain/probs/robodesk_prob.prob"
    d_c = main.parse_file_to_dict(domain_fname)
    domain = parse_domain_config.ParseDomainConfig.parse(d_c)
    hls = FDSolver(d_c, cleanup_files=False)
    p_c = main.parse_file_to_dict(prob)
    visual = True # len(os.environ.get('DISPLAY', '')) > 0
    #visual = False
    problem = parse_problem_config.ParseProblemConfig.parse(p_c, domain, None, use_tf=True, sess=None, visual=visual)
    params = problem.init_state.params

    cam_pos = (2.5, 111., -43., (0., 0., 0.)) # dist, yaw, pitch, target pos
    p.resetDebugVisualizerCamera(*cam_pos)
    # proc1 = Process(target=stepSimulation, args=(600, 0.))
    # proc1.name = 'PyBullet Sim Step'
    # proc1.start()

    sucs = []
    N_RUNS = 50
    for run_num in range(N_RUNS):
        agent.mjc_env.reset()

        for param in ['ball', 'upright_block', 'flat_block', \
                      'drawer_handle', 'shelf_handle', 'shelf', 'drawer']:
            pose = agent.mjc_env.get_item_pose(param, euler=True)
            targ = '{}_init_target'.format(param)
            params[param].pose[:,0] = pose[0]
            params[param].rotation[:,0] = pose[1]
            if targ in params:
                params[targ].value[:,0] = pose[0]
                params[targ].rotation[:,0] = pose[1]
            #params[param].pose[:,0] = env.physics.named.data.qpos[param][:3]
            #quat = env.physics.named.data.qpos[param][3:7]
            #quat = [quat[1], quat[2], quat[3], quat[0]]
            #euler = T.quaternion_to_euler(quat)
            #params[param].rotation[:,0] = euler
        #params['ball'].rotation[:,0] = [0., -0.4, 1.57]
        params['drawer'].hinge[:,0] = agent.mjc_env.get_attr('drawer', 'hinge')
        params['shelf'].hinge[:,0] = agent.mjc_env.get_attr('shelf', 'hinge')

        right_jnts = agent.mjc_env.get_attr('panda', 'right')
        lb, ub = params['panda'].geom.get_joint_limits('right')
        lb = np.array(lb) + 2e-3
        ub = np.array(ub) - 2e-3
        right_jnts = np.clip(right_jnts, lb, ub)
        params['panda'].right[:,0] = right_jnts

        params['robot_init_pose'].right[:,0] = right_jnts
        params['robot_init_pose'].right_gripper[:,0] = agent.mjc_env.get_attr('panda', 'right_gripper')

        for param in params:
            targ = '{}_init_target'.format(param)
            if targ in params:
                params[targ].value[:,0] = params[param].pose[:,0]
                params[targ].rotation[:,0] = params[param].rotation[:,0]

        #goal = '(and (InSlideDoor ball drawer) (Stacked upright_block flat_block) (NearGripperRight panda green_button))'
        goal = '(and (SlideDoorClose shelf_handle shelf) (Stacked upright_block flat_block) (NearGripperRight panda green_button))'
        #goal = '(and (SlideDoorOpen drawer_handle drawer) (NearApproachRight panda upright_block))'
        #goal = '(and (InSlideDoor upright_block shelf) (NearApproachRight panda ball))'
        #goal = '(and (SlideDoorClose drawer_handle drawer) (InSlideDoor ball drawer))'
        #goal = '(and (InSlideDoor ball drawer) (InSlideDoor upright_block shelf) (SlideDoorClose drawer_handle drawer) (SlideDoorClose shelf_handle shelf))'
        #goal = '(Stacked upright_block flat_block)'
        #goal = '(and (SlideDoorClose shelf_handle shelf) (InSlideDoor upright_block shelf))'
        #goal = '(Lifted flat_block panda)'
        #goal = '(Lifted upright_block panda)'
        goal = '(Lifted ball panda)'
        #goal = '(SlideDoorClose shelf_handle shelf)'
        #goal = '(SlideDoorOpen drawer_handle drawer)'
        #goal = '(InSlideDoor ball drawer)'
        goal = '(At flat_block bin_target)'
        #goal = '(At ball bin_target)'
        #goal = '(InSlideDoor upright_block shelf)'
        #goal = '(InSlideDoor flat_block shelf)'
        #goal = '(and (At flat_block bin_target) (At upright_block off_desk_target))'
        #goal = '(and (InSlideDoor ball drawer) (NearGripperRight panda green_button))'
        #goal = '(NearGripperRight panda green_button)'
        #goal = '(Lifted upright_block panda)'
        #goal = "(NearGripperRight panda green_button)"

        goals = [
             #'(Lifted upright_block panda)',
             #'(Lifted ball panda)',
             #'(Stacked upright_block flat_block)',
             #'(SlideDoorClose shelf_handle shelf)',
             #'(SlideDoorOpen drawer_handle drawer)',
             #'(Near flat_block bin_target)',
             #'(Near upright_block off_desk_target)',
             #'(InSlideDoor flat_block shelf)',
             #'(InGripperRight panda green_button)',
             '(and (InSlideDoor upright_block shelf) (SlideDoorClose shelf_handle shelf))',
                ]

        goal = random.choice(goals)
        goal_info = []
        goal_str = goal.strip()[1:-1]
        if goal_str.startswith('and'):
            goals = goal_str.split('(')
            for g in goals:
                if g.find(')') >= 0:
                    goal_info.append(g.strip()[:-1].split(' '))
        else:
            goal_info = [goal_str.split(' ')]
        print('SOLVING:', goal, goal_info)

        print('CONSISTENT?', problem.init_state.is_consistent())
        solver = RobotSolverOSQP()

        plan, descr = p_mod_abs(hls, solver, domain, problem, goal=goal, debug=True, n_resamples=5, max_iter=2)

        if type(plan) is str or plan is None:
            sucs.append([goal_info[0][0], 'OPT FAIL'])
            continue
        #import ipdb; ipdb.set_trace()

        if visual:
            agent.add_viewer()

        panda = plan.params['panda']
        for act in plan.actions:
            st, et = act.active_timesteps
            x = agent.get_state()
            for (pname, aname), inds in agent.state_inds.items():
                getattr(plan.params[pname], aname)[:, st] = x[inds]
            print('FAILED', act, plan.get_failed_preds((st,st)))
            for t in range(st, et):
                grip = panda.right_gripper[:, min(t+1, plan.horizon-1)]
                grip = -0.005 * np.ones(2) if grip[0] < 0.01 else 0.07 * np.ones(2)
                ctrl = np.r_[panda.right[:,t], grip]
                obs, rew, done, info = agent.mjc_env.step(ctrl)
                if 'hand_image' in obs:
                    agent.render_viewer(np.r_[obs['image'], obs['hand_image']])
                else:
                    agent.render_viewer(obs['image'])

        x = agent.get_state()
        goal_suc = [agent.parse_goal(x, g[0], g[1:]) for g in goal_info]
        print('Goal?', goal, goal_suc)
        #import ipdb; ipdb.set_trace()
        sucs.append([goal_info[0], goal_suc])

        print('\n\n-----------\nALL GOALS:')
        for item in sucs:
            print(item)
        print('------------\n\n')
        import ipdb; ipdb.set_trace()

ext: py

max_stars_repo_name: Algorithmic-Alignment-Lab/openTAMP

max_issues_repo_name: Algorithmic-Alignment-Lab/openTAMP

max_forks_repo_name: Algorithmic-Alignment-Lab/openTAMP

reason: import numpy


text:
# Visualization of a MOS instance using pygame
#
# Note to run this file, you need to run the following
# in the parent directory of multi_object_search:
#
#   python -m multi_object_search.env.visual
#

import pygame
import cv2
import math
import numpy as np
import random
from pomdp_py import util
from pomdp_problems.multi_object_search.env.env import *
from pomdp_problems.multi_object_search.domain.observation import *
from pomdp_problems.multi_object_search.domain.action import *
from pomdp_problems.multi_object_search.domain.state import *
from pomdp_problems.multi_object_search.example_worlds import *

# Deterministic way to get object color
def object_color(objid, count):
    color = [107, 107, 107]
    if count % 3 == 0:
        color[0] += 100 + (3 * (objid*5 % 11))
        color[0] = max(12, min(222, color[0]))
    elif count % 3 == 1:
        color[1] += 100 + (3 * (objid*5 % 11))
        color[1] = max(12, min(222, color[1]))
    else:
        color[2] += 100 + (3 * (objid*5 % 11))
        color[2] = max(12, min(222, color[2]))
    return tuple(color)

#### Visualization through pygame ####
class MosViz:

    def __init__(self, env,
                 res=30, fps=30, controllable=False):
        self._env = env

        self._res = res
        self._img = self._make_gridworld_image(res)
        self._last_observation = {}  # map from robot id to MosOOObservation
        self._last_viz_observation = {}  # map from robot id to MosOOObservation
        self._last_action = {}  # map from robot id to Action
        self._last_belief = {}  # map from robot id to OOBelief

        self._controllable = controllable
        self._running = False
        self._fps = fps
        self._playtime = 0.0

        # Generate some colors, one per target object
        colors = {}
        for i, objid in enumerate(env.target_objects):
            colors[objid] = object_color(objid, i)
        self._target_colors = colors

    def _make_gridworld_image(self, r):
        # Preparing 2d array
        w, l = self._env.width, self._env.length
        arr2d = np.full((self._env.width,
                         self._env.length), 0)  # free grids
        state = self._env.state
        for objid in state.object_states:
            pose = state.object_states[objid]["pose"]
            if state.object_states[objid].objclass == "robot":
                arr2d[pose[0], pose[1]] = 0  # free grid
            elif state.object_states[objid].objclass == "obstacle":
                arr2d[pose[0], pose[1]] = 1  # obstacle
            elif state.object_states[objid].objclass == "target":
                arr2d[pose[0], pose[1]] = 2  # target

        # Creating image
        img = np.full((w*r,l*r,3), 255, dtype=np.int32)
        for x in range(w):
            for y in range(l):
                if arr2d[x,y] == 0:    # free
                    cv2.rectangle(img, (y*r, x*r), (y*r+r, x*r+r),
                                  (255, 255, 255), -1)
                elif arr2d[x,y] == 1:  # obstacle
                    cv2.rectangle(img, (y*r, x*r), (y*r+r, x*r+r),
                                  (40, 31, 3), -1)
                elif arr2d[x,y] == 2:  # target
                    cv2.rectangle(img, (y*r, x*r), (y*r+r, x*r+r),
                                  (255, 165, 0), -1)
                cv2.rectangle(img, (y*r, x*r), (y*r+r, x*r+r),
                              (0, 0, 0), 1, 8)
        return img

    @property
    def img_width(self):
        return self._img.shape[0]

    @property
    def img_height(self):
        return self._img.shape[1]

    @property
    def last_observation(self):
        return self._last_observation

    def update(self, robot_id, action, observation, viz_observation, belief):
        """
        Update the visualization after there is new real action and observation
        and updated belief.

        Args:
            observation (MosOOObservation): Real observation
            viz_observation (MosOOObservation): An observation used to visualize
                                                the sensing region.
        """
        self._last_action[robot_id] = action
        self._last_observation[robot_id] = observation
        self._last_viz_observation[robot_id] = viz_observation
        self._last_belief[robot_id] = belief

    @staticmethod
    def draw_robot(img, x, y, th, size, color=(255,12,12)):
        radius = int(round(size / 2))
        cv2.circle(img, (y+radius, x+radius), radius, color, thickness=2)

        endpoint = (y+radius + int(round(radius*math.sin(th))),
                    x+radius + int(round(radius*math.cos(th))))
        cv2.line(img, (y+radius,x+radius), endpoint, color, 2)

    @staticmethod
    def draw_observation(img, z, rx, ry, rth, r, size, color=(12,12,255)):
        assert type(z) == MosOOObservation, "%s != MosOOObservation" % (str(type(z)))
        radius = int(round(r / 2))
        for objid in z.objposes:
            if z.for_obj(objid).pose != ObjectObservation.NULL:
                lx, ly = z.for_obj(objid).pose
                cv2.circle(img, (ly*r+radius,
                                 lx*r+radius), size, color, thickness=-1)

    @staticmethod
    def draw_belief(img, belief, r, size, target_colors):
        """belief (OOBelief)"""
        radius = int(round(r / 2))

        circle_drawn = {}  # map from pose to number of times drawn

        for objid in belief.object_beliefs:
            if isinstance(belief.object_belief(objid).random(), RobotState):
                continue
            hist = belief.object_belief(objid).get_histogram()
            color = target_colors[objid]

            last_val = -1
            count = 0
            for state in reversed(sorted(hist, key=hist.get)):
                if state.objclass == 'target':
                    if last_val != -1:
                        color = util.lighter(color, 1-hist[state]/last_val)
                    if np.mean(np.array(color) / np.array([255, 255, 255])) < 0.99:
                        tx, ty = state['pose']
                        if (tx,ty) not in circle_drawn:
                            circle_drawn[(tx,ty)] = 0
                        circle_drawn[(tx,ty)] += 1

                        cv2.circle(img, (ty*r+radius,
                                         tx*r+radius), size//circle_drawn[(tx,ty)], color, thickness=-1)
                        last_val = hist[state]

                        count +=1
                        if last_val <= 0:
                            break

    # PyGame interface functions
    def on_init(self):
        """pygame init"""
        pygame.init()  # calls pygame.font.init()
        # init main screen and background
        self._display_surf = pygame.display.set_mode((self.img_width,
                                                      self.img_height),
                                                     pygame.HWSURFACE)
        self._background = pygame.Surface(self._display_surf.get_size()).convert()
        self._clock = pygame.time.Clock()

        # Font
        self._myfont = pygame.font.SysFont('Comic Sans MS', 30)
        self._running = True

    def on_event(self, event):
        # TODO: Keyboard control multiple robots
        robot_id = list(self._env.robot_ids)[0]  # Just pick the first one.

        if event.type == pygame.QUIT:
            self._running = False
        elif event.type == pygame.KEYDOWN:
            u = None  # control signal according to motion model
            action = None  # control input by user

            # odometry model
            if event.key == pygame.K_LEFT:
                action = MoveLeft
            elif event.key == pygame.K_RIGHT:
                action = MoveRight
            elif event.key == pygame.K_UP:
                action = MoveForward
            elif event.key == pygame.K_DOWN:
                action = MoveBackward
            # euclidean axis model
            elif event.key == pygame.K_a:
                action = MoveWest
            elif event.key == pygame.K_d:
                action = MoveEast
            elif event.key == pygame.K_s:
                action = MoveSouth
            elif event.key == pygame.K_w:
                action = MoveNorth
            elif event.key == pygame.K_SPACE:
                action = Look
            elif event.key == pygame.K_RETURN:
                action = Find

            if action is None:
                return

            if self._controllable:
                if isinstance(action, MotionAction):
                    reward = self._env.state_transition(action, execute=True, robot_id=robot_id)
                    z = None
                elif isinstance(action, LookAction) or isinstance(action, FindAction):
                    robot_pose = self._env.state.pose(robot_id)
                    z = self._env.sensors[robot_id].observe(robot_pose,
                                                            self._env.state)
                    self._last_observation[robot_id] = z
                    self._last_viz_observation[robot_id] = z
                    reward = self._env.state_transition(action, execute=True, robot_id=robot_id)
                print("robot state: %s" % str(self._env.state.object_states[robot_id]))
                print("     action: %s" % str(action.name))
                print("     observation: %s" % str(z))
                print("     reward: %s" % str(reward))
                print("------------")
            return action

    def on_loop(self):
        self._playtime += self._clock.tick(self._fps) / 1000.0

    def on_render(self):
        # self._display_surf.blit(self._background, (0, 0))
        self.render_env(self._display_surf)
        robot_id = list(self._env.robot_ids)[0]  # Just pick the first one.
        rx, ry, rth = self._env.state.pose(robot_id)
        fps_text = "FPS: {0:.2f}".format(self._clock.get_fps())
        last_action = self._last_action.get(robot_id, None)
        last_action_str = "no_action" if last_action is None else str(last_action)
        pygame.display.set_caption("%s | Robot%d(%.2f,%.2f,%.2f) | %s | %s" %
                                   (last_action_str, robot_id, rx, ry, rth*180/math.pi,
                                    str(self._env.state.object_states[robot_id]["objects_found"]),
                                    fps_text))
        pygame.display.flip()

    def on_cleanup(self):
        pygame.quit()

    def on_execute(self):
        if self.on_init() == False:
            self._running = False

        while( self._running ):
            for event in pygame.event.get():
                self.on_event(event)
            self.on_loop()
            self.on_render()
        self.on_cleanup()

    def render_env(self, display_surf):
        # draw robot, a circle and a vector
        img = np.copy(self._img)
        for i, robot_id in enumerate(self._env.robot_ids):
            rx, ry, rth = self._env.state.pose(robot_id)
            r = self._res  # Not radius!
            last_observation = self._last_observation.get(robot_id, None)
            last_viz_observation = self._last_viz_observation.get(robot_id, None)
            last_belief = self._last_belief.get(robot_id, None)
            if last_belief is not None:
                MosViz.draw_belief(img, last_belief, r, r//3, self._target_colors)
            if last_viz_observation is not None:
                MosViz.draw_observation(img, last_viz_observation,
                                        rx, ry, rth, r, r//4, color=(200, 200, 12))
            if last_observation is not None:
                MosViz.draw_observation(img, last_observation,
                                        rx, ry, rth, r, r//8, color=(20, 20, 180))

            MosViz.draw_robot(img, rx*r, ry*r, rth, r, color=(12, 255*(0.8*(i+1)), 12))
        pygame.surfarray.blit_array(display_surf, img)

def unittest():
    # If you don't want occlusion, use this:
    laserstr = make_laser_sensor(90, (1, 8), 0.5, False)
    # If you want occlusion, use this
    # (the difference is mainly in angle_increment; this
    #  is due to the discretization - discretization may
    #  cause "strange" behavior when checking occlusion
    #  but the model is actually doing the right thing.)
    laserstr_occ = make_laser_sensor(360, (1, 8), 0.5, True)
    # Proximity sensor
    proxstr = make_proximity_sensor(1.5, False)
    proxstr_occ = make_proximity_sensor(1.5, True)

    worldmap, robot = world1
    worldstr = equip_sensors(worldmap, {robot: laserstr})

    dim, robots, objects, obstacles, sensors = interpret(worldstr)
    init_state = MosOOState({**objects, **robots})
    env = MosEnvironment(dim,
                         init_state, sensors,
                         obstacles=obstacles)
    viz = MosViz(env, controllable=True)
    viz.on_execute()

if __name__ == '__main__':
    unittest()


text:
# Copyright (C) 2018 Zhixian MA <zx@mazhixian.me>

"""
A class to define the convolutional autoencoder (CAE) by tensorflow.

Reference
=========
[1] tensorflow_tutorials
    https://github.com/pkmital/tensorflow_tutorials/blob/master/
    python/07_autoencoder.py

Update
======
[2017-09-04] Add methods to fine-tuning the network.
[2017-09-04] Change padding and stride strtegy to custom providing.
[2017-09-04] Remove pooling staffs

[2017-09-07] Add validation
[2017-09-07] Add accuracy evalutation of CNN, i.e., fine-tuning
[2017-09-07] Add draw_training_curve()

[2017-09-14] Variable space to fix the checkpoint, and solve restore problem
[2018-05-03] Add adjustable learning rate
[2018-08-24] Fixed and adjusted parameters
"""

import tensorflow as tf
import numpy as np
import time

class ConvAE():
    """
    A convolutional autoencoder (CAE) constructor

    inputs
    ======
    input_shape: np.ndarray
        Shape of the input of the network = (d,r,c).
        d: dimensions,
        r: rows,
        c: cols
    kernel_size: list
        Shape of the kernels in each ConvLayer
    Kernel_num: list
        Number of kernels in each ConvLayer
    fc_nodes: list
        The dense layers after the fully connected layer
        of the last ConvLayer or pooling layer
    encode_nodes: int
        Number of nodes in the encode layer
    padding: tuple (pad_en, pad_de)
        Padding modes for encode and decode blocks,
        default as ('SAME','SAME')
    stride: 2D tuple
        Stride step, default as (2, 2)
    numclass: int
        Number of classes at the sofrmax layer
    valrate: float
        Validation rate, range from 0.0 to 0.1, default as 0.3
    sess: tf.Session
        The restored session of pretrained, default as None
    name: str
        Name of the saved pretrained network

    Methods
    =======
    gen_BatchIterator: Generate a iterator for batch training.
    cae_build: Build the CAE network.
    gen_cost: Generate the CAE cost function.
    gen_optimizer: Generate the ADAM optimizer.
    gen_validation: Get training and validation subsets with provided rate.
    cae_train: Train the CAE network.
    cae_test: Test the CAE network.
    cae_encode: Get the feature vector of the input image.
    cae_decode: Reconstruct from a feature vector.
    cae_save: Save the CAE network.
    gen_cnn_BatchIterator: Generate a iterator for cnn batch training.
    cnn_build: Build the CNN network.
    gen_cnn_cost: Generate the CNN cost function.
    gen_cnn_optimizer: Generate the CNN optimizer.
    cnn_train: Train the CNN network.
    cnn_predict: Predict or classify a RG image.
    gen_accuracy: Generate the accuracy tensor.
    cnn_accuracy: Test The CNN network.
    cnn_save: Save the CNN network.
    """

    def __init__(self, input_shape,
                 kernel_size=[3,3,3], kernel_num=[10,10,10],
                 fc_nodes=[64],
                 encode_nodes=10,
                 padding = ('SAME','SAME'),
                 stride = (2, 2),
                 numclass = 2,
                 valrate = 0.3,
                 sess = None,
                 name = None,
                 ):
        """
        The initializer
        """
        self.input_shape = input_shape
        self.kernel_size = kernel_size
        self.kernel_num = kernel_num
        self.fc_nodes = fc_nodes
        self.encode_nodes = encode_nodes
        self.pad_en = padding[0]
        self.pad_de = padding[1]
        self.stride = stride
        self.numclass = numclass # used for fine tuning, maybe need reconsidered.
        self.y_ = tf.placeholder(tf.float32, shape=[None, numclass], name="cnn-labels") # placeholder of y
        self.l_cnn = tf.placeholder(tf.float32, shape=[None, numclass], name="l_cnn") # cnn
        self.l_en = tf.placeholder(tf.float32, shape=[None, self.encode_nodes], name="l_en")
        # Input layer
        in_depth = self.input_shape[3]
        in_row = self.input_shape[1]
        in_col = self.input_shape[2]
        self.l_in = tf.placeholder(
            tf.float32,
            [None,in_row,in_col,in_depth],
            name='l_in')
        # decoder layer
        self.l_de = tf.placeholder(tf.float32,
                                   [None,in_row,in_col,in_depth], name='l_de')
        self.droprate = tf.placeholder(tf.float32, name="droprate")
        self.valrate = 0.2
        self.get_learning_rate()
        # sess for reload
        self.sess = sess
        if self.sess is not None:
            try:
                self.l_in = sess.graph.get_tensor_by_name(name['l_in'])
                self.l_en = sess.graph.get_tensor_by_name(name['l_en'])
                self.l_de = sess.graph.get_tensor_by_name(name['l_de'])
            except:
                print("Some thing wrong when load the pretrained network...")

    def gen_BatchIterator(self, data, batch_size=100, shuffle=True):
        """
        Return the next 'batch_size' examples
        from the X_in dataset

        Reference
        =========
        [1] tensorflow.examples.tutorial.mnist.input_data.next_batch
        Input
        =====
        data: 4d np.ndarray
            The samples to be batched
        batch_size: int
            Size of a single batch.
        shuffle: bool
            Whether shuffle the indices.

        Output
        ======
        Yield a batch generator
        """
        if shuffle:
            indices = np.arange(len(data))
            np.random.shuffle(indices)
        for start_idx in range(0, len(data) - batch_size + 1, batch_size):
            if shuffle:
                excerpt = indices[start_idx: start_idx + batch_size]
            else:
                excerpt = slice(start_idx, start_idx + batch_size)
            yield data[excerpt]

    def get_learning_rate(self):
        """Get the exponentially decreased learning rate."""
        self.init_lr = tf.placeholder(tf.float32, name="init_lr")
        self.global_step = tf.placeholder(tf.float32, name="global_step")
        self.decay_step = tf.placeholder(tf.float32, name="decay_step")
        self.decay_rate = tf.placeholder(tf.float32, name="decay_rate")
        self.learning_rate = tf.train.exponential_decay(
            learning_rate=self.init_lr ,
            global_step=self.global_step,
            decay_steps=self.decay_step,
            decay_rate=self.decay_rate,
            staircase=False,
            name="learning_rate")

    # Pre-training part
    def cae_build(self):
        """Construct the network, including ConvLayers and FC layers."""

        # Useful methods
        def weight_variable(shape, name):
            initial = tf.truncated_normal(shape, stddev=0.1)
            return tf.Variable(initial, name=name)

        def bias_variable(shape, name):
            initial = tf.constant(0.1, shape=shape)
            return tf.Variable(initial, name=name)

        # Init
        self.encoder = []
        self.shapes_en = []

        # Encoder layers
        current_input = self.l_in
        for i, depth_output in enumerate(self.kernel_num):
            depth_input = current_input.get_shape().as_list()[3]
            self.shapes_en.append(current_input.get_shape().as_list())
            W_name = "Conv_En_W{0}".format(i)
            b_name = "Conv_En_b{0}".format(i)
            W = weight_variable(
                shape=[self.kernel_size[i],
                       self.kernel_size[i],
                       depth_input,
                       depth_output],
                name=W_name)
            b = bias_variable(shape=[depth_output], name=b_name)
            self.encoder.append(W) # share
            output = tf.add(
                tf.nn.conv2d(current_input,
                             W, strides=[1,self.stride[0],self.stride[1],1],
                             padding=self.pad_en), b)
            output = tf.nn.relu(output)
            output = tf.nn.dropout(output, self.droprate)
            current_input = output

        if self.encode_nodes is not None:
            # Dense layer
            shape_conv = current_input.get_shape().as_list()
            depth_dense = shape_conv[1] * shape_conv[2] * shape_conv[3]
            l_en_dense = tf.reshape(current_input, [-1, depth_dense])

            # fully connected
            depth_input = depth_dense
            current_input = l_en_dense
            self.en_fc = [] # save and share weights of the fc layers
            self.fc_depth=[] # depth_input of the encoder
            for i, depth_output in enumerate(self.fc_nodes):
                self.fc_depth.append(depth_input)
                W_name = "FC_En_W{0}".format(i)
                b_name = "FC_En_b{0}".format(i)
                W = weight_variable(shape=[depth_input, depth_output], name=W_name)
                b = bias_variable(shape=[depth_output], name=b_name)
                self.en_fc.append(W) # share weight
                output = tf.nn.relu(tf.matmul(current_input, W) + b)
                # dropout
                output = tf.nn.dropout(output, self.droprate)
                current_input = output
                depth_input = depth_output

            # encode layer
            W_en = weight_variable(shape=[depth_input, self.encode_nodes], name="En_W")
            b_en = bias_variable(shape=[self.encode_nodes], name="En_b")
            output = tf.nn.relu(tf.matmul(current_input, W_en) + b_en)
            current_input = output
            self.l_en = current_input

            # decoder layers
            W_de = tf.transpose(W_en)
            if len(self.fc_nodes):
                depth_output = self.fc_nodes[-1]
            else:
                depth_output = depth_dense
            b_de = bias_variable(shape=[depth_output], name="De_b")
            output = tf.nn.relu(tf.matmul(current_input, W_de) + b_de)
            current_input = output

            # fc layers
            for i in range(len(self.fc_nodes)-1, -1, -1):
                depth_output = self.fc_depth[i]
                b_name = "FC_De_b{0}".format(i)
                W = tf.transpose(self.en_fc[i])
                b = bias_variable(shape=[depth_output], name=b_name)
                output = tf.nn.relu(tf.matmul(current_input, W) + b)
                # dropout
                output = tf.nn.dropout(output, self.droprate)
                current_input = output

            # Dense layer
            # reshape
            current_input = tf.reshape(current_input,
                                    [-1,
                                     shape_conv[1],
                                     shape_conv[2],
                                     shape_conv[3]])
        else:
            self.l_en = current_input

        self.decoder = self.encoder.copy()
        self.decoder.reverse()
        self.shapes_de = self.shapes_en.copy()
        self.shapes_de.reverse()

        for i, shape in enumerate(self.shapes_de):
            W = self.decoder[i]
            b_name = "Conv_De_b{0}".format(i)
            b = bias_variable(shape=(W.get_shape().as_list()[2],), name=b_name)
            output = tf.add(tf.nn.conv2d_transpose(
                current_input, W,
                tf.stack([tf.shape(self.l_in)[0], shape[1], shape[2], shape[3]]),
                strides = [1,self.stride[0],self.stride[1],1], padding=self.pad_de),
                b)
            output = tf.nn.relu(output)
            output = tf.nn.dropout(output, self.droprate)
            current_input = output

        # Decoder output
        self.l_de = current_input

    def gen_cost(self):
        """Generate the cost, usually be the mean square error."""
        # cost function
        self.cost = tf.reduce_mean(tf.square(self.l_de-self.l_in))

    def gen_optimizer(self, learning_rate=0.01):
        """Generate the optimizer to minimize the cose"""
        self.gen_cost()
        self.optimizer = tf.train.AdamOptimizer(self.learning_rate, name="cae-optimizer").minimize(self.cost)

    def gen_validation(self, data, label=None):
        """Separate the dataset into training and validation subsets.

        inputs
        ======
        data: np.ndarray
            The input data, 4D matrix
        label: np.ndarray or list, opt
            The labels w.r.t. input data, optional

        outputs
        =======
        data_train: {"data": , "label": }
        data_val: {"data":, "label":}
        """
        if label is None:
            label_train = None
            label_val = None
            idx = np.random.permutation(len(data))
            num_val = int(len(data)*self.valrate)
            data_val = {"data": data[idx[0:num_val],:,:,:],
                        "label": label_val}
            # train
            data_train = {"data": data[idx[num_val:],:,:,:],
                          "label": label_train}
        else:
            # Training and validation
            idx = np.random.permutation(len(data))
            num_val = int(len(data)*self.valrate)
            data_val = {"data": data[idx[0:num_val],:,:,:],
                        "label": label[idx[0:num_val],:]}
            # train
            data_train = {"data": data[idx[num_val:],:,:,:],
                          "label": label[idx[num_val:],:]}

        return data_train,data_val

    def cae_train(self, data, num_epochs=20,
                  learning_rate=0.01, derate=1/16, deepoch=20,
                  batch_size=100, droprate=0.5):
        """Train the net"""
        timestamp = time.strftime('%Y-%m-%d: %H:%M:%S', time.localtime(time.time()))
        print("[%s] Training parameters\n" % (timestamp))
        print("[%s] Epochs: %d\tLearning rate: %.2f\n" % (timestamp, num_epochs, learning_rate))
        print("[%s] Batch size: %d\tDrop rate: %.2f\n" % (timestamp, batch_size, droprate))
        self.gen_optimizer(learning_rate = learning_rate)
        init_op = tf.global_variables_initializer()
        self.sess = tf.InteractiveSession()
        self.sess.run(init_op)
        # save loss for drawing
        self.cae_trainloss = []
        self.cae_valloss = []

        decay_rate = 0.95
        # generate training and validation samples
        data_train,data_val = self.gen_validation(data=data, label=None)
        numbatch_trn = len(data_train["data"]) // batch_size
        numbatch_val = len(data_val["data"]) // batch_size
        for epoch in range(num_epochs):
            # update learning rate
            lr_dict = {
                self.init_lr: learning_rate,
                self.global_step: epoch,
                self.decay_step: batch_size,
                self.decay_rate: decay_rate}
            cost_trn = 0
            cost_val = 0
            for batch in self.gen_BatchIterator(data=data_train["data"], batch_size=batch_size):
                trn_dict = {self.l_in: batch, self.droprate: droprate}
                trn_dict.update(lr_dict)
                self.sess.run(self.optimizer, feed_dict=trn_dict)
                cost_trn += self.sess.run(self.cost,
                               feed_dict={self.l_in: batch, self.droprate: 0.0})

            for batch in self.gen_BatchIterator(data=data_val["data"], batch_size=batch_size):
                cost_val += self.sess.run(self.cost, feed_dict={self.l_in: batch, self.droprate: 0.0})

            timestamp = time.strftime('%Y-%m-%d: %H:%M:%S', time.localtime(time.time()))
            print("[%s] Epoch: %03d\tTraining loss: %.6f\tValidation loss: %.6f" %
                (timestamp, epoch+1, cost_trn/numbatch_trn, cost_val/numbatch_val))
            self.cae_trainloss.append(cost_trn/numbatch_trn)
            self.cae_valloss.append(cost_val/numbatch_val)

    def cae_test(self, img):
        """Test the trained network,

        Input
        =====
        img: np.ndarray
            The images to be test.

        Output
        ======
        img_pre: np.ndarray
            The predicted images.
        """
        # params
        depth = self.input_shape[3]
        rows = self.input_shape[1]
        cols = self.input_shape[2]
        # Reshape the images
        shapes = img.shape
        if len(shapes) == 2:
            if shapes[0] != rows or shapes[1] != cols:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 3:
            if shapes[0] != rows or shapes[1] != cols or shapes[2] != depth:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 4:
            if shapes[1] != rows or shapes[2] != cols or shapes[3] != depth :
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(shapes[0],rows,cols,depth)

        # generate predicted images
        if hasattr(self, 'sess'):
            img_pre = self.sess.run(self.l_de, feed_dict={self.l_in: img_te, self.droprate: 0.0})
        else:
            print("You should train the network.")
            return None

        return img_pre

    def cae_encode(self, img):
        """Test the trained network,

        Input
        =====
        img: np.ndarray
            The images to be test.

        Output
        ======
        code: np.ndarray
            The encoded codes.
        """
        # params
        depth = self.input_shape[3]
        rows = self.input_shape[1]
        cols = self.input_shape[2]
        # Reshape the images
        shapes = img.shape
        if len(shapes) == 2:
            if shapes[0] != rows or shapes[1] != cols:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 3:
            if shapes[0] != rows or shapes[1] != cols or shapes[2] != depth:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 4:
            if shapes[1] != rows or shapes[2] != cols or shapes[3] != depth :
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(shapes[0],rows,cols,depth)

        # generate predicted images
        if hasattr(self, 'sess'):
            code = self.sess.run(self.l_en, feed_dict={self.l_in: img_te, self.droprate: 0.0})
        else:
            print("You should train the network.")
            return None

        return code

    def cae_decode(self, code):
        """Generate agn images with provided codes.

        Input
        =====
        code: np.ndarray
            The code to be decoded.

        Output
        ======
        img_de: np.ndarray
            The decoded and reconstructed images.
        """
        # Compare code length
        code_len = self.l_en.get_shape()[1]
        if code.shape[1] != code_len:
            print("The length of provided codes should be equal to the network's")
            return None
        else:
            # decoding
            if hasattr(self, 'sess'):
                l_in_shape = self.l_in.get_shape().as_list()
                l_in_shape[0] = code.shape[0]
                p_in = np.zeros(l_in_shape) # pseudo input
                img_de = self.sess.run(self.l_de,
                                       feed_dict={self.l_in: p_in,
                                                  self.l_en: code,
                                                  self.droprate: 0.0})
            else:
                print("You should train the network firstly.")
                return None

        return img_de

    def cae_save(self, namepath, netpath):
        """Save the net"""
        import pickle
        import sys
        sys.setrecursionlimit(1000000)

        # save the names
        savedict = {'l_in': self.l_in.name,
                    'l_en': self.l_en.name,
                    'l_de': self.l_de.name,
                    'l_cnn': self.l_cnn.name,
                    'y_':self.y_.name,
                    'cnn_cost':self.cnn_cost.name,
                    'droprate': self.droprate.name,
                    'netpath': netpath}
        with open(namepath, 'wb') as fp:
            pickle.dump(savedict, fp)

        # save the net
        saver = tf.train.Saver()
        saver.save(self.sess, netpath)

    ##################
    #Fine-tuning part#
    ##################
    def gen_cnn_BatchIterator(self, data, label, batch_size=100, shuffle=True):
        """
        Return the next 'batch_size' examples
        from the X_in dataset

        Reference
        =========
        [1] tensorflow.examples.tutorial.mnist.input_data.next_batch
        Input
        =====
        data: 4d np.ndarray
            The samples to be batched
        label: np.ndarray
            The labels to be batched w.r.t. data
        batch_size: int
            Size of a single batch.
        shuffle: bool
            Whether shuffle the indices.

        Output
        ======
        Yield a batch generator
        """
        if shuffle:
            indices = np.arange(len(data))
            np.random.shuffle(indices)
        else:
            indices = np.arange(len(data))
        data = data[indices]
        label = label[indices]
        numbatch = len(indices) // batch_size
        return data,label,numbatch

    def cnn_build(self):
        """Build the cnn for fine-tuning, just add a softmax layer after the encoder layer.
           Since the weights are shared between encoder and decoder, the fine-tuned weights
           can be reused.
        """
        def weight_variable(shape, name):
            initial = tf.truncated_normal(shape, stddev=0.1)
            return tf.Variable(initial, name=name)

        def bias_variable(shape, name):
            initial = tf.constant(0.1, shape=shape)
            return tf.Variable(initial, name=name)
        # add a softmax layer
        W_soft = weight_variable([self.encode_nodes, self.numclass], name="cnn_W")
        b_soft = bias_variable([self.numclass], name="cnn_b")

        self.l_cnn = tf.nn.softmax(tf.matmul(self.l_en, W_soft) + b_soft)
        # generate the optimizer
        self.gen_cnn_optimizer()
        # generate the accuracy estimator
        self.gen_accuracy()

    def gen_cnn_cost(self):
        """Generate the cost function for the added softmax layer at the end of encoder."""
        # cost function
        self.cnn_cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=self.y_, logits=self.l_cnn))

    def gen_cnn_optimizer(self):
        """Generate the optimizer for the fine-tuning cnn.

        Reference
        =========
        [1] Tensorflow: No gradients provided for any variable
            https://stackoverflow.com/questions/38778760/tensorflow-no-gradients-provided-for-any-variable
        """
        self.gen_cnn_cost()
        self.cnn_optimizer = tf.train.AdamOptimizer(self.learning_rate,name='cnn-optimizer').minimize(self.cnn_cost)

    def cnn_train(self, data, label, num_epochs=20, learning_rate=0.0005, batch_size=100, droprate=0.5):
        """Train the net"""
        timestamp = time.strftime('%Y-%m-%d: %H:%M:%S', time.localtime(time.time()))
        print("[%s] Training parameters\n" % (timestamp))
        print("[%s] Epochs: %d\tLearning rate: %.2f\n" % (timestamp, num_epochs, learning_rate))
        print("[%s] Batch size: %d\tDrop rate: %.2f\n" % (timestamp, batch_size, droprate))
        if self.sess is None:
            init_op = tf.global_variables_initializer()
            self.sess = tf.InteractiveSession()
            self.sess.run(init_op)

        # save loss for drawing
        self.cnn_trainloss = []
        self.cnn_valloss = []
        self.cnn_trainacc = []
        self.cnn_valacc = []

        decay_rate = 0.95
        # generate training and validation samples
        data_train,data_val = self.gen_validation(data=data, label=label)
        for epoch in range(num_epochs):
            cost = 0
            acc = 0
            data_batch,label_batch,numbatch = self.gen_cnn_BatchIterator(
                data=data_train["data"], label=data_train["label"], batch_size=batch_size)
            # update learning rate
            lr_dict = {
                self.init_lr: learning_rate,
                self.global_step: epoch,
                self.decay_step: batch_size,
                self.decay_rate: decay_rate}
            for i in range(numbatch):
                batch = data_batch[i*batch_size:(i+1)*batch_size,:,:,:]
                batch_label = label_batch[i*batch_size:(i+1)*batch_size,:]
                trn_dict = {self.l_in: batch, self.y_: batch_label, self.droprate: droprate}
                trn_dict.update(lr_dict)
                # print(batch_label.shape)
                self.sess.run(self.cnn_optimizer,
                              feed_dict=trn_dict)
                cost += self.sess.run(
                    self.cnn_cost,
                    feed_dict={self.l_in: batch, self.y_: batch_label, self.droprate: 0.0})
                acc += self.cnn_accuracy(img=batch, label=batch_label)
            # validation
            valloss = self.sess.run(
                self.cnn_cost,
                feed_dict={self.l_in: data_val["data"],
                           self.y_: data_val["label"],
                           self.droprate: 0.0})
            valacc = self.cnn_accuracy(img=data_val["data"], label=data_val["label"])
            timestamp = time.strftime('%Y-%m-%d: %H:%M:%S', time.localtime(time.time()))
            print("[%s] Epoch: %03d Trn loss: %.6f Trn acc: %.3f Val loss: %.6f Val acc: %.3f" %
                (timestamp, epoch+1, cost/numbatch, acc/numbatch, valloss, valacc))
            self.cnn_trainloss.append(cost/numbatch)
            self.cnn_valloss.append(valloss)
            self.cnn_trainacc.append(acc/numbatch)
            self.cnn_valacc.append(valacc)

    def cnn_predict(self, img):
        """Predict types of the sample with the
           trained fine-tuned network.

        Input
        =====
        img: np.ndarray
            The images to be test.

        Output
        ======
        label_pre: np.ndarray
            The predicted label.
        """
        # params
        depth = self.input_shape[3]
        rows = self.input_shape[1]
        cols = self.input_shape[2]
        # Reshape the images
        shapes = img.shape
        if len(shapes) == 2:
            if shapes[0] != rows or shapes[1] != cols:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 3:
            if shapes[0] != rows or shapes[1] != cols or shapes[2] != depth:
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(1,rows,cols,depth)
        elif len(shapes) == 4:
            if shapes[1] != rows or shapes[2] != cols or shapes[3] != depth :
                print('The shape of the test images do not match the network.')
                return None
            img_te = img.reshape(shapes[0],rows,cols,depth)

        # generate predicted images
        if hasattr(self, 'sess'):
            label_pre = tf.argmax(self.sess.run(self.l_cnn, feed_dict={self.l_in: img_te, self.droprate: 0.0}))
        else:
            print("You should train the network.")
            return None

        return label_pre.eval()

    def gen_accuracy(self):
        """Generate the accuracy tensor"""
        self.correction_prediction = tf.equal(tf.argmax(self.l_cnn, 1), tf.argmax(self.y_,1), name="corr_pre")
        self.accuracy = tf.reduce_mean(tf.cast(self.correction_prediction, tf.float32), name="acc")

    def cnn_accuracy(self, img, label):
        """Calculate the classification accuracy of the cnn network.

        inputs
        ======
        img: np.ndarray
            The images to be classified
        label: np.ndarray, one hot
            The labels w.r.t. the images

        output
        ======
        acc: float
            Classification accuracy
        """
        return self.accuracy.eval(feed_dict={self.l_in: img, self.y_: label, self.droprate: 0.0})

    def cnn_save(self, namepath, netpath):
        """Save the net"""
        import pickle
        import sys
        sys.setrecursionlimit(1000000)

        # save the names
        savedict = {'l_in': self.l_in.name,
                    'l_cnn': self.l_cnn.name,
                    'y_':self.y_.name,
                    'droprate': self.droprate.name,
                    'netpath': netpath}
        with open(namepath, 'wb') as fp:
            pickle.dump(savedict, fp)

        # save the net
        saver = tf.train.Saver()
        saver.save(self.sess, netpath)

max_stars_repo_name: myinxd/MCRGNet

max_issues_repo_name: myinxd/MCRGNet

max_forks_repo_name: myinxd/MCRGNet


text:
import os.path
import numpy as np
import scipy.sparse as sp
import torch
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from sklearn.metrics import *
from sklearn.model_selection import *

from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier

def encode_onehot(labels):
    classes = set(labels)
    classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
                    enumerate(classes)}
    labels_onehot = np.array(list(map(classes_dict.get, labels)),
                             dtype=np.int32)
    return labels_onehot

def process_data(graph_type='corr_coef_pval', group='', rand_cut=0.1, pval_cut=0.05, coef_cut=0.7, dataset='adni', feats='dma'):

    if(dataset == 'adni'):
        # 506 samples
        # <id> <74 features> <class label>
        clinical = pd.read_excel('/data1/GCN_AB/adni/CTX_Demo_n506.xlsx')
        if(group == '.mci'):
            clinical = clinical[clinical.Dx == 1]
        elif(group == '.nc'):
            clinical = clinical[clinical.Dx == 0]

        id = clinical.loc[:, "ID"]
        demo = clinical.loc[:, ["Age","Sex", "Edu", "APOEe4"]]
        mri = clinical.loc[:, "SupraTentorial":"RightMiddleTemporal"]

        label = clinical.loc[:, "111CUTOFF"]

        # prs = pd.read_csv('{}prs_output_IGAP.profile'.format(path), delimiter=r'\s+')

    elif(dataset == 'smc'):
        # 521 samples
        # <id> <25 features> <class label>
        clinical = pd.read_excel('/data1/GCN_AB/smc/CTX_Demo_n521.xlsx')
        if(group == '.mci'):
            clinical = clinical[clinical.ABPET_Dx == 1]
        elif(group == '.nc'):
            clinical = clinical[clinical.ABPET_Dx == 0]

        id = clinical.loc[:, "IID"]
        demo = clinical.loc[:, ["test_age", "sex", "education", "APOE4"]]
        mri = clinical.loc[:, "ICV.mm3.":"HV_native_R"]

        label = clinical.loc[:, "amyloid_beta"]

    feat = pd.concat([demo, mri], axis=1)

    if(feats == 'dm'):
        feat = feat.drop(['APOEe4'], axis=1)

    feat = (feat - feat.min()) / (feat.max() - feat.min())
    print(feat.shape)

    df = pd.concat([id, feat, label], axis=1)
    df.to_csv("/data1/GCN_AB/{}/{}_{}{}.content".format(dataset, dataset, feats, group), sep=' ', index=False, header=False)

    if(graph_type == 'random'):
        edge_cut = np.around(rand_cut/100*len(id)*len(id)).astype(int)

        o = np.ones(edge_cut, dtype=np.int)
        z = np.zeros(np.product((len(id), len(id))) - edge_cut, dtype=np.int)
        board = np.concatenate([o, z])
        np.random.shuffle(board)
        board = board.reshape((len(id), len(id)))

        graph = nx.from_pandas_adjacency(pd.DataFrame(board, index=id, columns=id))

        # graph = nx.binomial_graph(n=506, p=edge_cut/((506*506)/2), seed=22)
        # graph = nx.relabel_nodes(graph, {new:old for new,old in enumerate(df.ID)}, copy=False)

        graph.remove_edges_from(nx.selfloop_edges(graph))
        print(nx.info(graph))
        # plot_degree_dist(graph)

        nx.write_edgelist(graph, "/data1/GCN_AB/{}/{}_{}_{}_{}{}.edges".format(dataset, dataset, feats, graph_type, rand_cut, group), data=False)

    elif(graph_type == 'corr_coef_pval'):
        from scipy.stats import pearsonr
        corr_df = feat.T
        corr_df.columns = id

        corr_coef_df = corr_df.corr()
        corr_pval_df = corr_df.corr(method=lambda x, y: pearsonr(x, y)[1]) - np.eye(len(corr_df.columns))

        df1 = corr_coef_df.stack().reset_index()
        df1.columns = ['s', 't', 'w']
        print(df1.head())

        df2 = corr_pval_df.stack().reset_index()
        df2.columns = ['s', 't', 'p']
        print(df2.head())

        merged_df = pd.concat([df1, df2], axis=1, join='inner')
        print(merged_df.head())
        print('df1 {}, df2 {}, merged {}'.format(df1.shape[0], df2.shape[0], merged_df.shape[0]))

        # graph = nx.from_pandas_adjacency((np.abs(corr_coef_df) > 0.9).astype(int))
        # graph = nx.from_pandas_adjacency((corr_pval_df < edge_cut).astype(int))
        # graph = nx.from_pandas_adjacency(((np.abs(corr_coef_df) > coef_cut) & (corr_pval_df < pval_cut)).astype(int))
        # graph = nx.from_pandas_adjacency(corr_coef_df[corr_pval_df < pval_cut].fillna(0))

        # graph.remove_edges_from(nx.selfloop_edges(graph))
        # print(nx.info(graph))
        # print('sparsity {:.4f}%'.format(graph.number_of_edges() / (graph.number_of_nodes() * graph.number_of_nodes()) * 100))

        # plot_degree_dist(graph)

        # edge_cut = "{:.1f}_{}".format(coef_cut, pval_cut)
        edge_cut = 'weighted'
        fname = "/data1/GCN_AB/{}/{}_{}_{}_{}{}.edges".format(dataset, dataset, feats, graph_type, edge_cut, group)

        # nx.write_edgelist(graph, fname, data=['weight'])
        merged_df.to_csv(fname, sep=' ', header=False, index=False)
        sort_edgelist(fname)

    # nx.draw_networkx(graph_pval, with_labels=False, node_size=30)

def sort_edgelist(fname):
    edgelist = pd.read_table(fname, header=None, sep=' ')
    edgelist.columns = ['s', 't', 'w', 'p']
    edgelist = edgelist.sort_values(by=['w', 'p'], ascending=[False, True])
    edgelist.to_csv(fname, sep=' ', header=False, index=False)

def plot_degree_dist(G):
    degrees = [G.degree(n) for n in G.nodes()]
    plt.hist(degrees)
    plt.show()

def load_data(dataset="adni", feats="dma", graph="corr_coef_pval", group="", rand_cut=0.1, pval_cut=0.05, coef_cut=0.7, top_edges_p=100):

    print('Loading {} dataset...'.format(dataset))

    if(graph == 'corr_coef_pval'):
        # cut = "{:.1f}_{}".format(coef_cut, pval_cut)
        cut = 'weighted'
    else:
        cut = rand_cut

    content_file = "/data1/GCN_AB/{}/{}_{}{}.content".format(dataset, dataset, feats, group)
    edge_file = "/data1/GCN_AB/{}/{}_{}_{}_{}{}.edges".format(dataset, dataset, feats, graph, cut, group)

    if((os.path.isfile(content_file) == False) | (os.path.isfile(edge_file) == False)):
        print('Processing {} dataset...'.format(dataset))
        process_data(graph_type=graph, group=group, rand_cut=rand_cut, pval_cut=pval_cut, coef_cut=coef_cut, dataset=dataset, feats=feats)

    idx_features_labels = np.genfromtxt(content_file, dtype=np.dtype(str))
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    labels = encode_onehot(idx_features_labels[:, -1])

    features = normalize(features)
    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[1])

    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)

    # build graph
    adj=""

    top_edges = labels.shape[0] * labels.shape[0] * top_edges_p / 100

    idx_map = {j: i for i, j in enumerate(idx)}

    edges_unordered = np.genfromtxt(edge_file, dtype=np.int32, usecols=[0,1], max_rows=top_edges)
    weights = np.genfromtxt(edge_file, dtype=np.float64, usecols=[2], max_rows=top_edges)

    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())),
                     dtype=np.int32).reshape(edges_unordered.shape)
    print(edges.shape)

    #adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
    #                    shape=(labels.shape[0], labels.shape[0]),
    #                    dtype=np.float32)
    adj = sp.coo_matrix((weights, (edges[:,0], edges[:,1])), shape=(labels.shape[0], labels.shape[0]), dtype=np.float64)

    # build symmetric adjacency matrix
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

    adj = normalize(adj + sp.eye(adj.shape[0]))

    adj = sparse_mx_to_torch_sparse_tensor(adj)

    return adj, features, labels, idx

def normalize(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1))
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    mx = r_mat_inv.dot(mx)
    return mx

def accuracy(output, labels):
    preds = output.max(1)[1].type_as(labels)
    correct = preds.eq(labels).double()
    correct = correct.sum()
    accuracy = correct / len(labels)

    preds = output.max(1)[1]
    precision = precision_score(labels, preds)
    recall = recall_score(labels, preds)
    fone = f1_score(labels, preds)
    ap = average_precision_score(labels, preds)
    auc = roc_auc_score(labels, preds)

    return accuracy, precision, recall, fone, ap, auc

def sparse_mx_to_torch_sparse_tensor(sparse_mx):
    """Convert a scipy sparse matrix to a torch sparse tensor."""
    sparse_mx = sparse_mx.tocoo().astype(np.float32)
    indices = torch.from_numpy(
        np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
    values = torch.from_numpy(sparse_mx.data)
    shape = torch.Size(sparse_mx.shape)
    return torch.sparse.FloatTensor(indices, values, shape)

def plot_embedding(emb, labels, ppl, fname, cmap_code='jet'):
    from sklearn.manifold import TSNE

    trans = TSNE(n_components=2, perplexity=ppl)
    emb = trans.fit_transform(emb)

    import matplotlib.pyplot as plt
    plt.scatter(
        emb[:, 0],
        emb[:, 1],
        c=labels,
        cmap=cmap_code,
        alpha=0.7,
    )
    plt.savefig(fname)

def classification_test(X, y, CFs, show_plot=False, k=10, rep=5, dict={}):

    if(show_plot):
        f, axes = plt.subplots(1, 2, figsize=(10, 5))
        axes[0].plot([0, 1], [0, 1], color='darkblue', linestyle='--')

    for CF in CFs:
        aucs = []

        cv = RepeatedStratifiedKFold(n_splits=k, n_repeats=rep)

        y_real = []
        y_pred = []
        y_prob = []
        for train_id, test_id in cv.split(X, y):
            X_train, X_test = X[train_id], X[test_id]
            y_train, y_test = y[train_id], y[test_id]

            if CF == "SVC":
                model = SVC(kernel='linear', C=1.0, probability=True, class_weight='balanced', random_state=42)
                clf = model.fit(X_train, y_train)
            elif CF == "RFC":
                model = RandomForestClassifier(n_estimators=100, max_depth=5, class_weight='balanced', random_state=42)
                clf = model.fit(X_train, y_train)
            elif CF == "LR-lasso":
                model = LogisticRegression(penalty="l1", solver="liblinear", class_weight='balanced', random_state=42)
                clf = model.fit(X_train, y_train)
            elif CF == "LR-ridge":
                model = LogisticRegression(penalty="l2", solver="liblinear", class_weight='balanced', random_state=42)
                clf = model.fit(X_train, y_train)
            elif CF == "MLP":
                model = MLPClassifier(hidden_layer_sizes=(32,), max_iter=550, learning_rate_init=0.005, alpha=5e-10,
                                      early_stopping=True, validation_fraction=0.2, random_state=42)
                clf = model.fit(X_train, y_train, class_weight='balanced')

            pred = clf.predict(X_test)
            prob = clf.predict_proba(X_test)[:, 1]

            y_real.append(y_test)
            y_pred.append(pred)
            y_prob.append(prob)

            fpr, tpr, _ = roc_curve(y_test, prob)
            aucs.append(auc(fpr, tpr))

        dict[CF] = aucs

        y_real = np.concatenate(y_real)
        y_pred = np.concatenate(y_pred)
        y_prob = np.concatenate(y_prob)

        if (show_plot):
            fpr, tpr, _ = roc_curve(y_real, y_prob)
            lab = '%s (%.4f)' % (CF, auc(fpr, tpr))
            axes[0].step(fpr, tpr, label=lab, lw=1)

            precision, recall, _ = precision_recall_curve(y_real, y_prob)
            lab = '%s (%.4f)' % (CF, auc(recall, precision))
            axes[1].step(recall, precision, label=lab, lw=1)

            axes[0].set_xlabel('FPR')
            axes[0].set_ylabel('TPR')
            axes[0].set_ylim(0, 1.05)
            axes[0].legend(loc='lower right', fontsize='small')

            axes[1].set_xlabel('Recall')
            axes[1].set_ylabel('Precision')
            axes[1].set_ylim(0, 1.05)
            axes[1].legend(loc='lower right', fontsize='small')

    return dict

reason: import numpy,import scipy,from scipy,import networkx


text:
from time import time
import numpy as np

class LineSearch:

    def __init__(self, params):
        self.name = type(self).__name__
        self.params = params

        for (key, value) in params.items():
            setattr(self, key, value)

    def __call__(self, theta, *args, **kwargs):
        start_time = time()
        min_point, steps, iters = self.run(theta, **self.params)
        if hasattr(self, 'max_iterations'):
            if iters >= self.max_iterations:
                print(
                    f'Warning: {self.name} reached max iters ({iters}), ' +
                    f'lambda: {min_point}, tolerance: {self.tolerance}.'
                )
        return {
            'steps': steps,
            'min_point': np.float64(min_point),
            'min_value': theta(min_point),
            'performance': {
                'iterations': iters,
                'duration': round((time() - start_time) * 1000, 1),
            }
        }

    def run(self, **kwargs):
        raise NotImplementedError('Method must be overriden in subclass')

class ConstantSearch(LineSearch):
    param_symbols = {
        'delta': '\\lambda',
    }

    def run(self, theta, delta, gamma=None, **kwargs):
        delta = theta.domain(delta, gamma)
        return [delta], delta, 0

class GoldenSectionSearch(LineSearch):
    param_symbols = {
        'a': 'a',
        'b': 'b',
        'tolerance': 'l',
    }

    def run(self, theta, a, b, tolerance, max_iterations,
            gamma=None, **kwargs):
        a = theta.domain(a, gamma)
        b = theta.domain(b, gamma)
        steps = [[a, b]]
        k = 0

        alpha = 0.618
        delta = a + (1 - alpha) * (b - a)
        mu = a + alpha * (b - a)

        while b - a > tolerance and k < max_iterations:
            if theta(delta) > theta(mu):
                a = delta
                delta = mu
                mu = a + alpha * (b - a)
            else:
                b = mu
                mu = delta
                delta = a + (1 - alpha) * (b - a)
            steps.append([a, b])
            k += 1

        return (a + b) / 2, steps, k

class BisectionSearch(LineSearch):
    param_symbols = {
        'a': 'a',
        'b': 'b',
        'tolerance': 'l',
    }

    def run(self, theta, a, b, tolerance, max_iterations,
            gamma=None, **kwargs):
        a = theta.domain(a, gamma)
        b = theta.domain(b, gamma)
        steps = [[a, b]]
        k = 0

        while abs(b - a) > tolerance and k < max_iterations:
            delta = (b + a) / 2
            # delta = theta.domain(a, gamma)
            dtheta = theta.derivative(delta)
            if dtheta == 0:
                break
            elif dtheta > 0:
                b = delta
            else:
                a = delta
            steps.append([a, b])
            k += 1

        return (a + b) / 2, steps, k

class DichotomousSearch(LineSearch):
    param_symbols = {
        'a': 'a',
        'b': 'b',
        'tolerance': 'l',
        'epsilon': '\\epsilon',
    }

    def run(self, theta, a, b, tolerance, epsilon,
            max_iterations, gamma=None, **kwargs):
        a = theta.domain(a, gamma)
        b = theta.domain(b, gamma)
        steps = [[a, b]]
        k = 0

        while b - a > tolerance and k < max_iterations:
            delta = (b + a) / 2 - epsilon
            mu = (b + a) / 2 + epsilon

            delta = theta.domain(delta, gamma)
            mu = theta.domain(mu, gamma)

            if theta(delta) < theta(mu):
                b = mu
            else:
                a = delta
            steps.append([a, b])
            k += 1

        return (a + b) / 2, steps, k

class FibonacciSearch(LineSearch):
    param_symbols = {
        'a': 'a',
        'b': 'b',
        'tolerance': 'l',
        'epsilon': '\\epsilon',
    }

    def fib(self, n):
        a, b = 0, 1
        for i in range(n):
            a, b = b, a + b
        return a

    def run(self, theta, a, b, tolerance, epsilon,
            max_iterations, gamma=None, **kwargs):
        a = theta.domain(a, gamma)
        b = theta.domain(b, gamma)
        steps = [[a, b]]
        k = 0

        n = 1
        while (b - a) / self.fib(n) > tolerance:
            n += 1

        delta = a + self.fib(n - 2) / self.fib(n) * (b - a)
        mu = a + self.fib(n - 1) / self.fib(n) * (b - a)

        while k <= n - 1 and k < max_iterations:
            if theta(delta) > theta(mu):
                a = delta
                delta = mu
                mu = a + self.fib(n - k - 1) / self.fib(n - k) * (b - a)
            else:
                b = mu
                mu = delta
                delta = a + self.fib(n - k - 2) / self.fib(n - k) * (b - a)
            steps.append([a, b])
            k += 1

        delta = theta.domain(delta, gamma)
        _mu = mu + epsilon
        _mu = theta.domain(_mu, gamma)
        if theta(delta) > theta(_mu):
            a = delta
        else:
            b = _mu

        return (a + b) / 2, steps, k

class UniformSearch(LineSearch):
    param_symbols = {
        'a': 'a',
        'b': 'b',
        'tolerance': 'l',
        'interval_count': 'n',
        'interval_multiplier': 'm',
    }

    def run(self, theta, a, b, tolerance, interval_count, interval_multiplier,
            max_iterations, gamma=None, **kwargs):
        a = theta.domain(a, gamma)
        # b = theta.domain(b, gamma)
        steps = [a]
        k = 0

        interval_size = (b - a) / interval_count
        min_point = a
        min_value = theta(a)

        while interval_size > tolerance and k < max_iterations:
            for i in range(interval_count + 1):
                k += 1
                x = a + i * interval_size
                x = theta.domain(x, gamma)
                y = theta(x)

                if y < min_value:
                    min_point = x
                    min_value = y
                    steps.append(x)

            a = min_point - interval_size
            b = min_point + interval_size
            interval_count = int(interval_count * interval_multiplier)
            interval_size = (b - a) / interval_count
        return min_point, steps, k

class NewtonsSearch(LineSearch):
    param_symbols = {
        'delta': '\\lambda',
        'tolerance': 'l',
    }

    def run(self, theta, delta, tolerance, max_iterations,
            gamma=None, **kwargs):
        delta = theta.domain(delta, gamma)
        steps = [delta]
        k = 0
        while abs(theta.derivative(delta)) > tolerance and k < max_iterations:
            delta = delta - theta.derivative(delta) / theta.derivative2(delta)
            delta = theta.domain(delta, gamma)
            steps.append(delta)
            k += 1
        return delta, steps, k

class ArmijoSearch(LineSearch):
    param_symbols = {
        'delta': '\\lambda',
        'alpha': '\\alpha',
        'beta': '\\beta',
    }

    def run(self, theta, delta, alpha, beta, max_iterations,
            gamma=None, **kwargs):
        delta = theta.domain(delta, gamma)
        steps = [delta]
        k = 0

        theta0 = theta(0)
        Dtheta0 = theta.derivative(0)

        condition = theta0 + alpha * delta * Dtheta0

        while theta(delta) > condition and k < max_iterations:
            delta *= beta
            condition = theta0 + alpha * delta * Dtheta0
            steps.append(delta)
            k += 1

        return delta, steps, k

line_search_dict = {c.__name__: c for c in LineSearch.__subclasses__()}

max_stars_repo_name: EinariTuukkanen/line-search-comparison

max_issues_repo_name: EinariTuukkanen/line-search-comparison

max_forks_repo_name: EinariTuukkanen/line-search-comparison


text:
import numpy as np
import matplotlib.pyplot as plt
import sympy as sm

x = "hallo"

print(3+4+5)
# stupid

def optimized_tax_revenue(x):
    t_0 = x[0]
    t_1 = x[1]
    k   = x[2]
    return -tax_revenue(e,v,m,t0,t1,k,N)

Guess = [0.2,0.3,0.4]

solution = optimize.minimize(optimized_tax_revenue,Guess,method='SLSQP',bounds=((0,1),(0,1),(0,1))))
t0  = solution.x[0]
t1  = solution.x[1]
k   = solution.x[2]
T   = tax_revenue(e,v,m,t0,t1,k,N)

print(t0)

max_stars_repo_path: Patient 0.py

max_issues_repo_path: Patient 0.py

max_forks_repo_path: Patient 0.py

reason: import numpy,import sympy