text:
import copy
import math
import random

import numpy as np

from agent import BaseAgent
from utils import Directions, MapTiles, tile_cost

class Location(object):
    """
    Location class. It indicates (x, y) location.
    """

    def __init__(self, *args):
        if len(args) == 1:
            if isinstance(args[0], tuple):
                self._x = args[0][1]
                self._y = args[0][0]
            if isinstance(args[0], Location):
                self._x = args[0].x
                self._y = args[0].y
        elif len(args) == 2:
            self._x = args[0]
            self._y = args[1]

    @property
    def x(self):
        """
        Returns x-coordinate.

        :return: x-coordinate
        """

        return self._x

    @property
    def y(self):
        """
        Returns y-coordinate.

        :return: y-coordinate
        """

        return self._y

    def __hash__(self):
        return hash((self.x, self.y))

    def __eq__(self, other):
        return isinstance(self, type(other)) and self.x == other.x and self.y == other.y

    def north(self):
        return Location(self.x, self.y - 1)

    def south(self):
        return Location(self.x, self.y + 1)

    def west(self):
        return Location(self.x - 1, self.y)

    def east(self):
        return Location(self.x + 1, self.y)

    def __str__(self):
        return '(%d, %d)' % (self.x, self.y)

class GameMap(object):
    """
    Game map class.
    """

    def __init__(self, agent, game_map, map_objects):
        if isinstance(game_map, np.ndarray):
            self._game_map = game_map
        elif isinstance(game_map, GameMap):
            self._game_map = game_map._game_map

        self._agent = agent
        self._map_objects = {Location(x): obj for x, obj in map_objects.items()}

        shape = self._game_map.shape
        self._tmp_goal = Location(random.randint(0, shape[1] - 1), random.randint(0, shape[0] - 1))

    def __iter__(self):
        shape = self._game_map.shape
        for y in range(shape[0]):
            for x in range(shape[1]):
                yield Location(x, y), self._game_map[y][x]

    def __getitem__(self, item):
        if isinstance(item, Location):
            loc = item
            if loc.x < 0 or loc.y < 0 or loc.x >= self._game_map.shape[1] or loc.y >= self._game_map.shape[0]:
                return None
            return self._game_map[loc.y][loc.x]

    def update_goal(self, location, probs):
        """
        Updates the goal in priority order; 'medkit', 'skeleton', 'boss', and unexplored tile.
        """

        for location, obj in self._map_objects.items():
            if obj.label == 'medkit':
                self._tmp_goal = location
                return

            if obj.label == 'skeleton':
                self._tmp_goal = location
                return

            if obj.label == 'boss':
                self._tmp_goal = location
                return

        shape = self._game_map.shape
        while self._tmp_goal is None or location == self._tmp_goal or probs[self._tmp_goal.y][self._tmp_goal.x] == 0.0:
            self._tmp_goal = Location(random.randint(0, shape[1] - 1), random.randint(0, shape[0] - 1))

    def size(self):
        shape = self._game_map.shape
        return shape[0] * shape[1]

    def goal(self):
        return self._tmp_goal

    def available_directions(self, location):
        north = self._game_map[location.north()]
        south = self._game_map[location.south()]
        west = self._game_map[location.west()]
        east = self._game_map[location.east()]

        directions = [(Directions.NORTH, north), (Directions.SOUTH, south),
                      (Directions.WEST, west), (Directions.EAST, east)]

        return [x for x in directions if x[1] is not None and x[1] is not MapTiles.W]

class Path(object):
    """
    Path object.
    """

    def __init__(self, *args):
        if len(args) == 1:
            if isinstance(args[0], Path):
                self._start = copy.deepcopy(args[0]._start)
                self._directions = copy.deepcopy(args[0]._directions)
                self._locations = copy.deepcopy(args[0]._locations)
                self._tiles = copy.deepcopy(args[0]._tiles)
            else:
                self._start = args[0]
                self._directions = []
                self._locations = []
                self._tiles = []

    def append(self, direction, location, tile):
        copy = Path(self)

        copy._directions.append(direction)
        copy._locations.append(location)
        copy._tiles.append(tile)

        return copy

    def has(self, location):
        return location == self._start or location in self._locations

    def cost(self, game_map):
        _cost = 0
        for location, tile in zip(self._locations, self._tiles):
            if tile is not MapTiles.U:
                _cost = _cost + tile_cost[tile]
        return _cost

    def __iter__(self):
        for (direction, location) in zip(self._directions, self._locations):
            yield direction, location

class DepthLimitedPathSearcher(object):
    """
    Path Searcher using depth-limited search for the performance.
    """

    def __init__(self, agent, game_map, depth_limit=None):
        self._agent = agent
        self._game_map = game_map
        self._depth_limit = depth_limit if depth_limit else math.sqrt(game_map.size()) / 2

    def search(self, location, path=None, depth=0):
        if path is None:
            path = Path(location)

        paths = []
        goal = self._game_map.goal()
        for (direction, next_location, tile) in self.expand(location):
            if path.has(next_location):
                continue

            tile = self._game_map[next_location]
            p = path.append(direction, next_location, tile)
            if tile != MapTiles.U and depth < self._depth_limit:
                p = self.search(next_location, p, depth + 1)

            if p is not None:
                if next_location == goal:
                    return p

                paths.append(p)

        if len(paths) == 0:
            if depth == 0:
                raise ValueError('No available paths')
            return None

        lowest_cost_path = sorted(paths, key=lambda x: x.cost(self._game_map))[0] if len(paths) > 1 else paths[0]
        return lowest_cost_path

    def expand(self, location):
        directions = [(Directions.NORTH, location.north(), self._game_map[location.north()]),
                      (Directions.SOUTH, location.south(), self._game_map[location.south()]),
                      (Directions.WEST, location.west(), self._game_map[location.west()]),
                      (Directions.EAST, location.east(), self._game_map[location.east()])]

        return [x for x in directions if x[2] is not None and x[2] is not MapTiles.W]

class ProblemSolvingAgent(BaseAgent):

    def __init__(self, height, width, initial_strength, name='rogue_agent'):
        super().__init__(height=height, width=width,
                         initial_strength=initial_strength, name=name)

        self._probs = np.ones(shape=(width, height))
        self._frontiers = []
        self._map_objects_size = 0

    def step(self, location, strength, game_map, map_objects):
        location, game_map = Location(location), GameMap(self, game_map, map_objects)
        self.update(location, strength, game_map, map_objects)

        if len(self._frontiers) == 0:
            path = self.best_path(location, game_map)
            if path is None:
                raise ValueError('Reach to dead-end')

            for (direction, loc) in path:
                self._frontiers.append(direction)

        return self._frontiers.pop(0)

    def best_path(self, location, game_map):
        path_searcher = DepthLimitedPathSearcher(self, game_map)
        return path_searcher.search(location)

    def update(self, location, strength, game_map, map_objects):
        if self._map_objects_size != len(map_objects):
            self._map_objects_size = len(map_objects)

            self._probs[location.y][location.x] = 0.0

            for (loc, tile) in game_map:
                if tile == MapTiles.W:
                    self._probs[loc.y][loc.x] = 0.0

            p = 1.0 / np.count_nonzero(self._probs)
            self._probs[np.where(self._probs != 0.0)] = p

            game_map.update_goal(location, self._probs)
            self._frontiers.clear()

if __name__ == '__main__':
    from driver import GameDriver

    height, width = 10, 10
    initial_strength = 100
    num_powerups = 2
    num_monsters = 1
    num_dynamic_monsters = 1
    save_dir = 'map1/'
    verbose = False
    show_map = False
    map_type = 'emoji'
    map_file = None

    results = []
    count = 1
    for _ in range(count):
        agent = ProblemSolvingAgent(height, width, initial_strength)
        agents = [agent]

        game_driver = GameDriver(
            height=height, width=width,
            num_powerups=num_powerups,
            num_monsters=num_monsters,
            num_dynamic_monsters=num_dynamic_monsters,
            agents=agents,
            initial_strength=initial_strength,
            show_map=show_map, map_type=map_type,
            save_dir=save_dir, map_file=map_file)

        print('Starting game')
        try:
            res = game_driver.play(verbose=verbose)
        except:
            continue

        results.append(res)

    print(sum([1 for x in results if x]) / float(len(results)))


text:
# Repository: https://gitlab.com/quantify-os/quantify-scheduler
# Licensed according to the LICENCE file on the master branch
"""Module containing Zurich Instruments InstrumentCoordinator Components."""
# pylint: disable=useless-super-delegation
# pylint: disable=too-many-arguments
# pylint: disable=too-many-ancestors

from __future__ import annotations

import shutil
from pathlib import Path
import logging
from typing import Dict, TYPE_CHECKING, Any, Optional, Tuple

import numpy as np
from zhinst import qcodes
from quantify_core.data import handling
from quantify_scheduler.backends.zhinst import helpers as zi_helpers
from quantify_scheduler.instrument_coordinator.components import base
from quantify_scheduler.backends.zhinst.settings import ZISerializeSettings

if TYPE_CHECKING:
    from zhinst.qcodes.base import ZIBaseInstrument
    from quantify_scheduler.backends.zhinst_backend import ZIDeviceConfig
    from quantify_scheduler.backends.zhinst.settings import ZISettings

logger = logging.getLogger(__name__)

def convert_to_instrument_coordinator_format(acquisition_results):
    """
    Converts the acquisition results format of the UHFQA component to
    the format required by InstrumentCoordinator.
    Converts from `Dict[int, np.ndarray]` to `Dict[Tuple[int, int], Any]`.
    """
    reformatted_results: Dict[Tuple[int, int], Any] = dict()
    for acq_channel in acquisition_results:
        results_array = acquisition_results.get(acq_channel)
        for i, complex_value in enumerate(results_array):
            separated_value = (np.real(complex_value), np.imag(complex_value))
            reformatted_results[(acq_channel, i)] = separated_value
    return reformatted_results

class ZIInstrumentCoordinatorComponent(base.InstrumentCoordinatorComponentBase):
    """Zurich Instruments InstrumentCoordinator component base class."""

    def __init__(self, instrument: ZIBaseInstrument, **kwargs) -> None:
        """Create a new instance of ZIInstrumentCoordinatorComponent."""
        super().__init__(instrument, **kwargs)
        self.zi_device_config: Optional[ZIDeviceConfig] = None
        self.zi_settings: Optional[ZISettings] = None
        self._data_path: Path = Path(".")

    @property
    def is_running(self) -> bool:
        raise NotImplementedError()

    # pylint: disable=arguments-differ
    def prepare(self, zi_device_config: ZIDeviceConfig) -> bool:
        """
        Prepare the InstrumentCoordinator component with configuration
        required to arm the instrument.

        The preparation is skipped when the new zi_device_config is the same as that
        from the previous time prepare was called. This saves significant time overhead.

        Parameters
        ----------
        zi_device_config :
            The ZI instrument configuration. See the link for details of the
            configuration format.

        Returns
        -------
        :
            A boolean indicating if the ZI component was configured in this call.
        """
        self.zi_device_config = zi_device_config

        new_zi_settings = zi_device_config.settings_builder.build()
        old_zi_settings = self.zi_settings

        if new_zi_settings == old_zi_settings:
            logger.info(
                f"{self.name}: device config and settings "
                + "are identical! Compilation skipped."
            )
            return False

        logger.info(f"Configuring {self.name}.")
        # if the settings are not identical, update the attributes of the
        # ic component and apply the settings to the hardware.
        self.zi_settings = new_zi_settings

        # Writes settings to filestorage
        self._data_path = Path(handling.get_datadir())
        self.zi_settings.serialize(
            self._data_path,
            ZISerializeSettings(
                self.name, self.instrument._serial, self.instrument._type
            ),
        )

        # Upload settings, seqc and waveforms
        self.zi_settings.apply(self.instrument)

        return True

    def retrieve_acquisition(self) -> Any:
        return None

class HDAWGInstrumentCoordinatorComponent(ZIInstrumentCoordinatorComponent):
    """Zurich Instruments HDAWG InstrumentCoordinator Component class."""

    def __init__(self, instrument: qcodes.HDAWG, **kwargs) -> None:
        """Create a new instance of HDAWGInstrumentCoordinatorComponent."""
        assert isinstance(instrument, qcodes.HDAWG)
        super().__init__(instrument, **kwargs)

    @property
    def instrument(self) -> qcodes.HDAWG:
        return super().instrument

    @property
    def is_running(self) -> bool:
        return any(
            self.get_awg(awg_index).is_running
            for awg_index in self.zi_settings.awg_indexes
        )

    def get_awg(self, index: int) -> qcodes.hdawg.AWG:
        """
        Returns the AWG by index.

        Parameters
        ----------
        index :
            The awg index.

        Returns
        -------
        :
            The HDAWG AWG instance.
        """
        return self.instrument.awgs[index]

    def start(self) -> None:
        """Starts all HDAWG AWG(s) in reversed order by index."""
        for awg_index in reversed(self.zi_settings.awg_indexes):
            self.get_awg(awg_index).run()

    def stop(self) -> None:
        """Stops all HDAWG AWG(s) in order by index."""
        for awg_index in self.zi_settings.awg_indexes:
            self.get_awg(awg_index).stop()

    def retrieve_acquisition(self) -> Any:
        return None

    def wait_done(self, timeout_sec: int = 10) -> None:
        for awg_index in reversed(self.zi_settings.awg_indexes):
            self.get_awg(awg_index).wait_done(timeout_sec)

class UHFQAInstrumentCoordinatorComponent(ZIInstrumentCoordinatorComponent):
    """Zurich Instruments UHFQA InstrumentCoordinator Component class."""

    def __init__(self, instrument: qcodes.UHFQA, **kwargs) -> None:
        """Create a new instance of UHFQAInstrumentCoordinatorComponent."""
        assert isinstance(instrument, qcodes.UHFQA)
        super().__init__(instrument, **kwargs)

    @property
    def instrument(self) -> qcodes.UHFQA:
        return super().instrument

    @property
    def is_running(self) -> bool:
        return self.instrument.awg.is_running

    def start(self) -> None:
        self.instrument.awg.run()

    def stop(self) -> None:
        self.instrument.awg.stop()

    def prepare(self, zi_device_config: ZIDeviceConfig) -> bool:
        """
        Prepares the component with configurations
        required to arm the instrument.

        After this step is complete, the waveform file is uploaded
        to the LabOne WebServer.

        Parameters
        ----------
        zi_device_config :
            The ZI instrument configuration. See the link for details of the
            configuration format.

        Returns
        -------
        :
            A boolean indicating if the ZI component was configured in this call.
        """
        # always start by resetting the counters and stopping the AWG
        self.instrument.qas[0].result.enable(0)
        self.instrument.awg.stop()

        self.instrument.qas[0].result.reset(1)
        self.instrument.qas[0].result.enable(1)

        try:
            # if settings where identical, no configuration is needed.
            configure = super().prepare(zi_device_config)
            if configure is False:
                return False

        # pylint: disable=broad-except
        # the exception being raised is "Upload failed", but the ZI backend raises it
        # as a general exception.
        except Exception as e:
            # whenever a new UHF device is used for the first time,
            # certain waveform files will not exist. The lines below copy files so
            # that it is possible to read from that location.
            # this line of code should only be logging a warning the very first time
            # a new setup is used, and then resolve auto.
            logger.warning(e)
            configure = True

        self._data_path = Path(handling.get_datadir())
        # Copy the UHFQA waveforms to the waves directory
        # This is required before compilation.

        # N.B. note this copies waves that were written during compilation, but are not
        # contained in the zi_device_config that is passed as an argument here.
        waves_path: Path = zi_helpers.get_waves_directory(self.instrument.awg)
        wave_files = list(self._data_path.glob(f"{self.name}*.csv"))
        for file in wave_files:
            shutil.copy2(str(file), str(waves_path))

        # prepare twice to resolve issue with waveform memory not being updated
        # correctly. In practice, we see that integration weights update correctly, but
        # the waveforms in pulses do not. This problem is not fully understood, but this
        # resolves the issue at a minor overhead.

        if configure:
            # Upload settings, seqc and waveforms
            self.zi_settings.apply(self.instrument)
        return True

    def retrieve_acquisition(self) -> Dict[int, np.ndarray]:
        if self.zi_device_config is None:
            raise RuntimeError("Undefined device config, first prepare UHFQA!")

        acq_config = self.zi_device_config.acq_config

        acq_channel_results: Dict[int, np.ndarray] = dict()
        for acq_channel, resolve in acq_config.resolvers.items():
            acq_channel_results[acq_channel] = resolve(uhfqa=self.instrument)

        reformatted_results = convert_to_instrument_coordinator_format(
            acq_channel_results
        )

        return reformatted_results

    def wait_done(self, timeout_sec: int = 10) -> None:
        self.instrument.awg.wait_done(timeout_sec)

ext: py

max_stars_repo_name: quantify-os/quantify-scheduler

max_issues_repo_name: quantify-os/quantify-scheduler

max_forks_repo_name: quantify-os/quantify-scheduler

reason: import numpy


text:
from pathlib import Path
import sys
sys.path.append("c:\\Users\\kpdav\\machine_learning\\projects\\PGA-portfolio-optimizer\\config")
import config

import pandas as pd
import numpy as np
from sklearn.model_selection import TimeSeriesSplit, GroupKFold

def timeseries_split_dataset(df, t_gap=0, t_size=None):
    """Split dataset using based on TimeSeriesSplit framework

    Args:
        df (pd.DataFrame) : historical features dataset (base/non-modified)

        t_gap (int) : period of gaps between train and test set split

        t_size (int) : default of n_samples // (n_splits + 1), maximum allowed

    Returns:
        TimeSeriesSplit of training and test set based on given parameters
    """
    df.sort_values(by="date", inplace=True)
    df.reset_index(drop=True, inplace=True)

    tscv = TimeSeriesSplit(gap=t_gap, test_size=t_size)

    for train_index, test_index in tscv.split(df):
        ts_train_set = df.loc[train_index]
        ts_test_set = df.loc[test_index]

    return ts_train_set, ts_test_set

def group_kfold_split(df, group="player_id"):
    """Split dataset using GroupKFold

    Args:
        df (pd.Dataframe) : historical features dataset

        group (str) : non-overlapping group used. Default value player id

    Returns:
        GroupKFold split of dataset based on given parameters
    """
    df.sort_values(by="date", inplace=True)
    df.reset_index(drop=True, inplace=True)

    group_kfold = GroupKFold()

    grouping = df[group]

    for train_index, test_index in group_kfold.split(df, group=grouping):
        gkfold_train_set = df.loc[train_index]
        gkfold_test_set = df.loc[test_index]

    return gkfold_train_set, gkfold_test_set

if __name__ == "__main__":

    ts_features_datapath = str(Path(config.TIMESERIES_FRAMEWORK_DIR, "features_hpd_2017_2020.csv"))
    df = pd.read_csv(ts_features_datapath, parse_dates=["date"])

    train_df, test_df = timeseries_split_dataset(df)

    timeseries_train_dataset = str(Path(config.TIMESERIES_FRAMEWORK_DIR, "train_dataset_default.csv"))

    timeseries_test_dataset = str(Path(config.TIMESERIES_FRAMEWORK_DIR, "test_dataset_default.csv"))

    train_df.to_csv(timeseries_train_dataset, index=False)
    test_df.to_csv(timeseries_test_dataset, index=False)


text:
import os
import numpy as np
from joblib import Parallel, delayed, dump,load
from .utils import timefn,for_all_methods
njobs = 20 # os.cpu_count()//2

def compute_recall(neighbors, ground_truth):
    total = 0
    for gt_row, row in zip(ground_truth, neighbors):
        total += np.intersect1d(gt_row, row).shape[0]
    return total / ground_truth.size

def recall_atN_(neighbors_matrix, ground_truth):
    """
    Args:
        neighbors_matrix:(nq,topK=512)
        ground_truth: 1d or 2d
    """
    topk = neighbors_matrix.shape[1]
    if ground_truth.ndim == 1:
        ng = 1
    if ground_truth.ndim == 2:
        ng = ground_truth.shape[1]

    N = [1,2,4,8,10,16,20,32,64,100,128,256,512]

    recall_list=[]
    N_list = []
    for tn in N:
        if ng <= tn <= topk:
            neighbors_matrix_topN = neighbors_matrix[:,0:tn]
            recall = compute_recall(neighbors_matrix_topN, ground_truth)
            recall_list.append(recall)
            N_list.append(tn)
            print(f"recall {ng}@{tn} = {recall}")
    return recall_list,N_list

def recall_atN(neighbors_matrix, ground_truth):
    """
    Args:
        neighbors_matrix:(nq, topK=512)
        ground_truth:(nq,) or (nq,>=10)
    """
    if ground_truth.ndim == 2 :
        ground_truth_1 = ground_truth[:,0]
        ground_truth_10 = ground_truth[:,0:10]
        r1, N1 = recall_atN_(neighbors_matrix,ground_truth_1)
        print("\n")
        r10, N10 = recall_atN_(neighbors_matrix,ground_truth_10)
        print("\n")
        print(f"N={N1}")
        print(f"recall1@N:{r1}")
        print(f"N={N10}")
        print(f"recall10@N:{r10}")

    if ground_truth.ndim == 1 :
        r1, N1 = recall_atN_(neighbors_matrix, ground_truth)
        print("\n")
        print(f"N={N1}")
        print(f"recall1@N:{r1}")

class SearchNeighbors:
    def __init__(self, metric) -> None:
        self.metric = metric
        if metric != "dot_product" and metric != "l2_distance":
            raise Exception(f"not suport {metric},optional:l2_distance or dot_product")

    def compute_recall(self, neighbors, ground_truth):
        # total = 0
        # for gt_row, row in zip(ground_truth, neighbors):
        #     total += np.intersect1d(gt_row, row).shape[0]
        # return total / ground_truth.size
        return compute_recall(neighbors,ground_truth)

    def brute_force_search(self, target_set, test_set, metric="l2_distance"):
        if metric != "dot_product" and metric != "l2_distance":
            raise Exception(f"not suport {metric},optional:l2_distance or dot_product")

        if metric == "dot_product":
            inner_product = target_set @ test_set.T
            return np.argmax(inner_product, axis=0)

        if metric == "l2_distance":
            q_2 = np.linalg.norm(test_set, axis=1).reshape(-1, 1) ** 2
            x_2 = np.linalg.norm(target_set, axis=1).reshape(1, -1) ** 2

            q_x = test_set @ target_set.T
            gt_distance = q_2 - 2 * q_x + x_2

            gt = np.argmin(gt_distance, axis=1)
            return gt

    # @jit
    def _sort_topk_adc_score(self, adc_score, topk):
        metric = self.metric

        if metric == "dot_product":
            ind_t = np.argpartition(adc_score, -topk)
            ind = ind_t[-topk:]
            t1 = np.argsort(adc_score[ind])
            t2 = ind[t1]
            t3 = np.flip(t2)
            return t3

        if metric == "l2_distance":
            ind = np.argpartition(adc_score, topk)[0:topk]
            return ind[np.argsort(adc_score[ind])]

# @for_all_methods(profile)
class SearchNeighbors_PQ(SearchNeighbors):
    """
    For the PQ indexing phase of database vectors,
    a `D`-dim input vector is divided into `M` `D`/`M`-dim sub-vectors.
    Each sub-vector is quantized into a small integer via `Ks` codewords.
    For the querying phase, given a new `D`-dim query vector, the distance beween the query
    and the database PQ-codes are efficiently approximated via Asymmetric Distance.
    All vectors must be np.ndarray
    .. [Jegou11] H. Jegou et al., "Product Quantization for Nearest Neighbor Search", IEEE TPAMI 2011

    Args:
        M (int): The number of sub-space
        Ks (int): The number of codewords for each subspace
            (typically 256, so that each sub-vector is quantized
            into 256 bits = 1 byte = uint8)
        D (int): The dim of each vector
        pq_codebook (np.ndarray): shape=(M, Ks, Ds) with dtype=np.float32.
            codebook[m][ks] means ks-th codeword (Ds-dim) for m-th subspace
        pq_codes (np.ndarray): PQ codes with shape=(n, M) and dtype=np.int
        metric (str): dot_product or l2_distance

    Attributes:
        M (int): The number of sub-space
        Ks (int): The number of codewords for each subspace
        D (int): The dim of each vector
        Ds (int): The dim of each sub-vector, i.e., Ds=D/M
        pq_codebook (np.ndarray): shape=(M, Ks, Ds) with dtype=np.float32
        pq_codes (np.ndarray): PQ codes with shape=(n, M) and dtype=np.int
        metric (str): dot_product or l2_distance
    """

    def __init__(self, M, Ks, D, pq_codebook, pq_codes, metric="l2_distance") -> None:
        SearchNeighbors.__init__(self, metric)
        self.M = M
        self.Ks = Ks
        self.D = D
        assert D % M == 0, "D must be divisible by M"
        self.Ds = D // M

        self.pq_codebook = pq_codebook
        self.pq_codes = pq_codes
        assert pq_codebook.shape == (M, Ks, self.Ds), "pq_codebook.shape must equal to (M,Ks,Ds)"

        self.tem_value = np.zeros(shape=(pq_codes.shape))
        self.index = np.zeros_like(pq_codes)

        for i in range(M):
            self.index[:,i] = i

    # @jit
    def compute_distance(self, query):
        """
        The distances (the squared Euclidean distance or inner product) are computed by comparing each sub-vector of
        the query to the codewords for each sub-subspace. `lookup_table[m][ks]` contains the squared Euclidean
        distance or inner product between the `m`-th sub-vector of the query and the `ks`-th codeword for the `m`-th
        sub-space (`self.codewords[m][ks]`). By looking up table to compute distance.
        Args:
            query (np.ndarray): Input vector with shape=(D, )
        Returns:
            np.ndarray: Asymmetric Distances with shape=(n, )
        """
        pq_codebook = self.pq_codebook
        codes = self.pq_codes
        metric = self.metric

        M = self.M
        Ds = self.Ds

        q = query.reshape(M, Ds)
        if metric == "dot_product":
            lookup_table = np.matmul(pq_codebook, q[:, :, np.newaxis])[:, :, 0]
            # table_vaule = lookup_table[range(M), codes] # (n,M)
            # inner_prod = np.sum(table_vaule, axis=1)

            self.tem_value  = lookup_table[self.index, codes]
            inner_prod = np.sum(self.tem_value, axis=1)
            return inner_prod

        if metric == "l2_distance":
            lookup_table = np.linalg.norm(pq_codebook - q[:, np.newaxis, :], axis=2) ** 2
            dists = np.sum(lookup_table[range(M), codes], axis=1)
            return dists

    def search_neighbors(self, query, topk):
        """
        Args:
            pq_codebook, codes, query, metric are the same as method compute_distance
            topk (int): this method will return topk neighbors
        Returns:
            index (np.darray): query's topk neighbors
        """
        adc_score = self.compute_distance(query)
        return self._sort_topk_adc_score(adc_score, topk)

    @timefn
    def neighbors(self, queries, topk):
        """
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)

        """
        assert queries.ndim == 2
        n = queries.shape[0]
        neighbors_matrix = np.zeros((n, topk), dtype=int)
        for i in range(n):
            q = queries[i]
            neighbors_matrix[i] = self.search_neighbors(q, topk)

        self.neighbors_matrix = neighbors_matrix

        return neighbors_matrix

    @timefn
    def par_neighbors(self, queries, topk, njobs=njobs):
        """
        Parallel version of method neighbors
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query
            njobs: the process numbers use
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)

        """
        assert queries.ndim == 2
        n = queries.shape[0]
        neighbors_matrix = np.zeros((n, topk), dtype=int)

        result = Parallel(n_jobs = njobs,backend='multiprocessing')(delayed(self.search_neighbors)(q,topk) for q in queries)

        for i in range(n):
            neighbors_matrix[i] = result[i]

        self.neighbors_matrix = neighbors_matrix

        return neighbors_matrix

    def pq_recall(self, queries, topk, ground_truth):
        """
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query.
            ground_truth(np.darray): shape=(nq,) or (nq, num), queries real neighbors (top1 or top num)
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)
        """
        ground_truth = np.array(ground_truth)

        try:
            if topk > self.neighbors_matrix.shape[1]:
                neighbors_matrix = self.neighbors(queries, topk)
            else:
                neighbors_matrix = self.neighbors_matrix[:, 0:topk]
        except AttributeError:
            neighbors_matrix = self.neighbors(queries, topk)

        recall = self.compute_recall(neighbors_matrix, ground_truth)

        nr = neighbors_matrix.shape[1]
        if ground_truth.ndim == 1:
            ng = 1
        if ground_truth.ndim == 2:
            ng = ground_truth.shape[1]

        print(f"recall {ng}@{nr} = {recall}")

class SearchNeighbors_AQ(SearchNeighbors):
    """
    Babenko, Artem, and Victor Lempitsky. "Additive quantization for extreme vector compression." CVPR 2014.

    Args:
        M (int): The number of codebooks
        K (int): The number of codewords for each codebook
        D (int): The dim of each vector
        aq_codebooks (np.ndarray): shape=(M*K, D) with dtype=np.float32.
            aq_codebooks[0:K,:] represents the K codewords in the first codebook
            aq_codebooks[(m-1)*K:mK,:] represents the K codewords in the m-th codebook
        aq_codes (np.ndarray): AQ codes with shape=(n, M) and dtype=np.int, where n is the number of encoded datapoints
            aq_codes[i,j] is in {0,1,...,K-1} for all i,j
        metric (str): dot_product or l2_distance

    Attributes:
        M (int): The number of codebooks
        K (int): The number of codewords for each codebook
        D (int): The dim of each vector
        aq_codebooks (np.ndarray): shape=(M*K, D) with dtype=np.float32.
        aq_codes (np.ndarray): AQ codes with shape=(n, M) and dtype=np.int, where n is the number of encoded datapoints
        metric (str): dot_product or l2_distance
        norm_sq (np.ndarray):shape=(n,), the square of the length of the quantized datapoints is recorded in norm_sq.
    """

    def __init__(self, M, K, D, aq_codebooks, aq_codes, metric="l2_distance") -> None:
        SearchNeighbors.__init__(self, metric)
        assert aq_codebooks.shape == (M * K, D), "aq_codebooks.shape must equal to (M*K, D)"
        assert aq_codes.shape[1] == M, "aq_codebooks.shape must equal to (n, M)"

        self.M = M
        self.K = K
        self.D = D
        self.aq_codebooks = aq_codebooks
        self.aq_codes = aq_codes

        self.aq_codes_2 = aq_codes + np.arange(M) * K

        if self.metric == "l2_distance":
            n = aq_codes.shape[0]
            self.norm_sq = np.zeros(n)
            for i in range(n):
                x = np.sum(aq_codebooks[self.aq_codes_2[i]], axis=0)
                self.norm_sq[i] = np.vdot(x, x)

    def compute_distance(self, query):
        """
        Compute the distances (the squared Euclidean distance or inner product).
        Args:
            query (np.ndarray): Input vector with shape=(D, )
        Returns:
            np.ndarray: Asymmetric Distances with shape=(n, )
        """
        aq_codebooks = self.aq_codebooks
        codes_2 = self.aq_codes_2
        metric = self.metric

        if metric == "dot_product":
            lookup_table = aq_codebooks @ query
            inner_prod = np.sum(lookup_table[codes_2], axis=1)

            return inner_prod

        if metric == "l2_distance":
            lookup_table = aq_codebooks @ query
            inner_prod = np.sum(lookup_table[codes_2], axis=1)

            dists = np.vdot(query, query) + self.norm_sq - 2 * inner_prod
            return dists

    def search_neighbors(self, query, topk):
        """
        Args:
            topk (int): this method will return topk neighbors
        Returns:
            index (np.darray): query's topk neighbors for a query.
        """
        adc_score = self.compute_distance(query)
        return self._sort_topk_adc_score(adc_score, topk)

    def neighbors(self, queries, topk):
        """
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query.
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)
        """
        assert queries.ndim == 2
        n = queries.shape[0]
        neighbors_matrix = np.zeros((n, topk), dtype=int)
        for i in range(n):
            q = queries[i]
            neighbors_matrix[i] = self.search_neighbors(q, topk)

        self.neighbors_matrix = neighbors_matrix

        return neighbors_matrix

    @timefn
    def par_neighbors(self, queries, topk, njobs=njobs):
        """
        Parallel version of method neighbors
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query
            njobs: the process numbers use
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)

        """
        assert queries.ndim == 2
        n = queries.shape[0]
        neighbors_matrix = np.zeros((n, topk), dtype=int)

        result = Parallel(n_jobs = njobs,backend='multiprocessing')(delayed(self.search_neighbors)(q,topk) for q in queries)

        for i in range(n):
            neighbors_matrix[i] = result[i]

        self.neighbors_matrix = neighbors_matrix

        return neighbors_matrix

    def aq_recall(self, queries, topk, ground_truth):
        """
        Args:
            queries (np.ndarray): Input matrix with shape=(nq, D), where nq is the number of queries.
            topk (int): this method will return topk neighbors for each query.
            ground_truth(np.darray): shape=(nq,) or (nq, num), queries real neighbors (top1 or top num)
        Returns:
            np.ndarray: topk neighbors for each query with shape=(nq, topk)
        """
        ground_truth = np.array(ground_truth)

        try:
            if topk > self.neighbors_matrix.shape[1]:
                neighbors_matrix = self.neighbors(queries, topk)
            else:
                neighbors_matrix = self.neighbors_matrix[:, 0:topk]
        except AttributeError:
            neighbors_matrix = self.neighbors(queries, topk)

        recall = self.compute_recall(neighbors_matrix, ground_truth)

        nr = neighbors_matrix.shape[1]
        if ground_truth.ndim == 1:
            ng = 1
        if ground_truth.ndim == 2:
            ng = ground_truth.shape[1]

        print(f"recall {ng}@{nr} = {recall}")


text:
from dask.distributed import Client
from multiprocessing import Pool, Array, Value

import numpy as np

import os, glob, sys

from partition import Partition

import psi4
psi4.core.be_quiet()

# def _scf(mol_string,
#                 potential=None, restricted="UKS"):

#     frag_info = pdft_fragment()

#     mol = psi4.geometry(mol_string)
#     wfn_base = psi4.core.Wavefunction.build(mol, "cc-pvdz")
#     wfn = psi4.proc.scf_wavefunction_factory("svwn", wfn_base, restricted)
#     wfn.initialize()

#     if potential is not None:
#         wfn.H().np[:] += potential

#     wfn.iterations()
#     wfn.finalize_energy()

#     #Paste results to pdf_fragment
#     energies = { "enuc" : wfn.get_energies('Nuclear'),
#                  "e1"   : wfn.get_energies('One-Electron'),
#                  "e2"   : wfn.get_energies('Two-Electron'),
#                  "exc"  : wfn.get_energies('XC'),
#                  "total": wfn.get_energies('Total Energy')}

#     frag_info.mol_str  = mol_string
#     frag_info.Da       = wfn.Da().np
#     frag_info.Db       = wfn.Db().np
#     frag_info.Ca       = wfn.Ca().np
#     frag_info.Cb       = wfn.Cb().np
#     frag_info.eig_a    = wfn.epsilon_a().np
#     frag_info.eig_b    = wfn.epsilon_b().np
#     frag_info.energies = energies
#     frag_info.energy   = wfn.get_energies('Total Energy')

#     return frag_info

# def scf(frags,  vext = None,
#                 client = None):

#     #Determine wheter or not a client has been initializied
#     if client is None:
#         client = Client()

#     #Check wether or not there is a vext to be added for scf cycle
#     if vext is None:
#         ret = client.map( _scf, frags )
#     else:
#         ret = client.map( _scf, frags, vext )

#     frag_data = [ i.result() for i in ret ]

#     return frag_data

if __name__ == '__main__':

    molecule = """
    He 0.0 0.0 0.0
    He 0.0 0.0 1.0
    symmetry c1
    """

    f1 = """
    He 0.0 0.0 0.0
    He 0.0 0.0 1.0
    symmetry c1
    """

    frags        = [molecule]

    part = Partition( frags, "svwn", "cc-pvdz" )
    frag_data = part.scf()

    ############### Process of Generting Partition Potential ########################

    #Generate a new external potential.
    dummy_v1 = np.zeros_like( frag_data[0].Da )
    np.fill_diagonal(dummy_v1, dummy_v1.diagonal() + 0.1)

    mols = [molecule]
    vps  = [dummy_v1]

    ############### Feed Partition Potential Back to Fragment ########################

    frag_data = part.scf( vps )

    for f in glob.glob("psi.*.clean"):
        os.remove(f)

    # for f in glob.glob("dask-worker-space"):
    #     os.rmdir(f)

    print("I'm leaving victorious")

    part.client.close()

max_stars_repo_name: VHchavez/Partition

max_issues_repo_name: VHchavez/Partition

max_forks_repo_name: VHchavez/Partition


text:
# Prepares data for simple chord classification.
# It computes the chromagram from the audio file and splits the labels to blocks.

import os

import numpy as np
import pandas as pd
from tfr.reassignment import chromagram
from tfr.signal import SignalFrames
from tfr.spectrogram import create_window

from time_intervals import block_labels

def prepare_chomagram_and_labels(
    album,
    song_title,
    block_size,
    hop_size,
    bin_range,
    bin_division):

    song = 'The_Beatles/'+album+'/'+song_title
    data_dir = '../data/beatles'
    audio_file = data_dir + '/audio-cd/' + song + '.wav'
    chord_file = data_dir  + '/chordlab/' + song + '.lab.pcs.tsv'
    audio_file, chord_file

    # ## Load audio & Split audio to blocks
    print('loading audio:', audio_file)

    signal_frames = SignalFrames(audio_file, frame_size=block_size, hop_size=hop_size)
    x_blocks = signal_frames.frames
    x_time = signal_frames.start_times
    fs = signal_frames.sample_rate

    print('sampling rate:', fs, 'Hz')
    print('number of samples:', signal_frames.length)
    print('duration in audio:', signal_frames.duration, 'sec')

    print('blocks shape:', x_blocks.shape)
    print('number of blocks:', len(x_blocks))
    # start times for each block
    print('last block starts at:', x_times[-1], 'sec')

    # ## Load chords
    print('loading chords:', chord_file)
    chords = pd.read_csv(chord_file, sep='\t')
    print('shape:', chords.shape)
    print('duration in chords:', chords['end'].iloc[-1])

    pcs_cols = ['C','Db','D','Eb','E','F','Gb','G','Ab','A','Bb','B']
    label_cols = ['label','root','bass'] + pcs_cols

    # ##

    # ## Mapping of chords to blocks

    def chords_to_blocks(chords, block_center_times):
        chord_ix = 0
        for t in block_center_times:
            yield chords.iloc[i][pcs_cols]

    def time_to_samples(time):
        return np.round(time * fs)

    chords['start_sample'] = time_to_samples(chords['start'])
    chords['end_sample'] = time_to_samples(chords['end'])
    df_blocks = pd.DataFrame({'start': time_to_samples(x_times).astype(np.int64)})
    df_blocks['end'] = df_blocks['start'] + block_size

    label_dict = chords[label_cols].drop_duplicates().set_index('label')

    df_labels = chords[['start_sample', 'end_sample', 'label']].copy()
    df_labels.rename(columns={'start_sample': 'start', 'end_sample': 'end'}, inplace=True)

    df_labelled_blocks = block_labels(df_blocks, df_labels)

    df_block_pcs = df_labelled_blocks[['label']].join(label_dict, on='label')[['label'] + pcs_cols]

    assert len(df_block_pcs) == len(df_blocks)

    block_labels_file = '{}/chord-pcs/{}_{}/{}.pcs'.format(data_dir, block_size, hop_size, song)
    print('block labels file:', block_labels_file)

    os.makedirs(os.path.dirname(block_labels_file), exist_ok=True)
    df_block_pcs.to_csv(block_labels_file, sep='\t', index=False)

    # ## Chromagram features

    w = create_window(block_size)
    X_chromagram = chromagram(x_blocks, w, fs, to_log=True, bin_range=bin_range, bin_division=bin_division)

    chromagram_file = '{}/chromagram/block={}_hop={}_bins={},{}_div={}/{}.npz'.format(
        data_dir, block_size, hop_size, bin_range[0], bin_range[1], bin_division, song)

    print('chomagram file:', chromagram_file)

    os.makedirs(os.path.dirname(chromagram_file), exist_ok=True)
    np.savez_compressed(chromagram_file, X=X_chromagram, times=x_times)

def transform_all():
    # takes an hour - could be optimized...
    import timeit

    block_size = 4096
    hop_size = 2048
    bin_range = (-48, 67)
    bin_division = 1

    with open('../data/beatles/isophonic-songs.txt') as file:
        for line in file.readlines():
            band, album, song_title = line.rstrip('\n').split('/')
            print(album, song_title)
            prepare_chomagram_and_labels(album, song_title, block_size, hop_size, bin_range, bin_division)

def parse_args():
    import argparse

    parser = argparse.ArgumentParser(description='Prepare data for chord classification - chromagram and labels.')
    parser.add_argument('ALBUM')
    parser.add_argument('SONG_TITLE')
    parser.add_argument('-b', '--block-size', type=int, default=4096, help='block size')
    parser.add_argument('-p', '--hop-size', type=int, default=2048, help='hop size')
    parser.add_argument('-r', '--bin-range', type=int, nargs=2, default=[-48, 67], help='chromagram bin range')
    parser.add_argument('-d', '--bin-division', type=int, default=1, help='bins per semitone in chromagram')

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()

    prepare_chomagram_and_labels(
        album=args.ALBUM,
        song_title=args.SONG_TITLE,
        block_size=args.block_size,
        hop_size=args.hop_size,
        bin_range=args.bin_range,
        bin_division=args.bin_division)

# example usage:
# $ python prepare_data_for_chord_classification.py \
#   "03_-_A_Hard_Day's_Night" "05_-_And_I_Love_Her" \
#   -b 4096 -p 2048 -r -48 67 -d 1

max_stars_repo_name: bzamecnik/ml


text:
import sys
sys.path.append(r'/wynton/home/degradolab/lonelu/GitHub_Design/Combs2')
import combs2
import prody as pr
import pandas as pd
import os
import numpy as np
from sklearn.neighbors import NearestNeighbors
from scipy.sparse import lil_matrix

class CgCombInfo:
    def __init__(self):
        self.ligand_id = None
        self.vdm_cgs = {}

def get_ligand_coords(filtered_ligands, lgd_sel):
    # Get all ligand coords.
    ligand_coords = []
    for i in range(len(filtered_ligands)):
        lgd = filtered_ligands[i]
        coords = []
        for lsa in lgd_sel:
            coords.extend(lgd.select('name ' + lsa).getCoords().flatten())
        ligand_coords.append(coords)
    ligand_coords = np.array(ligand_coords)

    return ligand_coords

def get_vdm_labels_coords(dfa, represent_name, correspond_resname, correspond_names):
    '''
    Example:
    correspond_resname = 'ASP'
    represent_name = 'OD2'
    correspond_names = ['CG', 'OD1', 'OD2']
    '''
    labels = dfa[(dfa['chain'] == 'Y') & (dfa['resname'] == correspond_resname) & (dfa['name'] == represent_name)][['CG', 'rota', 'probe_name', 'seg_chain_resnum']]

    vdm_coords =[]
    for k in correspond_names:
        df_contacts = dfa[
            (dfa['resname'] == correspond_resname)
            & (dfa['chain'] == 'Y')
            & (dfa['name'] == k)
        ]
        vdm_coords.append(df_contacts[['c_x', 'c_y', 'c_z']].values.T)
    vdm_coords = np.concatenate(vdm_coords).T

    return labels, vdm_coords

def get_nearest_vdms(vdm_coords, ligand_coords, radius = 2):
    '''
    #to be predecated.
    '''
    # Nearest Neighbor
    nbr = NearestNeighbors(radius=radius).fit(vdm_coords)
    adj_matrix = nbr.radius_neighbors_graph(ligand_coords).astype(bool)
    print(adj_matrix.sum())
    print(adj_matrix.shape)

    m_adj_matrix = adj_matrix.tolil()

    all_inds = {}
    for r in range(m_adj_matrix.shape[0]):
        inds = m_adj_matrix.rows[r]
        if len(inds) <= 0:
            continue
        for ind in inds:
            all_inds.add(ind)

    return all_inds

def write_vdms(outdir, all_inds, labels, dfa, prefix):

    for ind in all_inds:
        x = labels.iloc[ind]
        # if x['seg_chain_resnum'][2] != 45:
        #     #print(x['seg_chain_resnum'][2])
        #     continue
        v = dfa[(dfa['CG'] == x['CG']) & (dfa['rota'] == x['rota']) & (dfa['probe_name'] == x['probe_name']) & (dfa['seg_chain_resnum'] == x['seg_chain_resnum'])]
        # if v['resname'].iloc[-1] != 'ARG':
        #     continue
        print(ind)
        combs2.design.functions.print_dataframe(v, outpath=outdir, tag = '_' + str(ind), prefix = prefix)
    return

def get_nearest_vdms_rmsd(vdm_coords, ligand_coords, radius = 1):
    # Nearest Neighbor
    nbr = NearestNeighbors(radius=radius).fit(vdm_coords)
    dists, inds = nbr.radius_neighbors(ligand_coords)

    return dists, inds

def vdm_ligand_clash(vdm, ligand, clash_radius = 2.7):
    '''
    The vdm X sidechain may clash with the ligand.
    return True if clash.
    '''
    vdm_coords =[]
    for i in range(vdm[(vdm['chain'] == 'X')].shape[0]):
        k = vdm[(vdm['chain'] == 'X')].iloc[i]
        if k['name'][0] == 'H':
            continue
        vdm_coords.append(k[['c_x', 'c_y', 'c_z']].values)

    ligand_coords = ligand.select('heavy').getCoords()

    nbr = NearestNeighbors(radius=clash_radius).fit(vdm_coords)
    adj_matrix = nbr.radius_neighbors_graph(ligand_coords).astype(bool)

    if adj_matrix.sum() > 0:
        return True
    return False

def search_vdm(cg_dict, ligands, cg_id, input_dict, labels_cgs, df_cgs, dist_ind_cgs, rmsd = 0.5):
    '''
    s: combs2.design._sample.Sample

    cg_id: (0, 0)  #The first is to record the cg in ligand, the second is to record the cg in vdms.

    input_dict[cg_id] = {

        cg : 'coo'  #dfa = s.cg_dict['coo']

        lgd_sel : ['C9', 'O3', 'O4']  # correpond to coo: ASP ['CG', 'OD1', 'OD2']

        represent_name : 'OD2'

        correspond_resname : 'ASP'

        correspond_names : ['CG', 'OD1', 'OD2']
    }
    '''

    dfa = cg_dict[input_dict[cg_id]['cg']]

    ligand_coords = get_ligand_coords(ligands, input_dict[cg_id]['lgd_sel'])

    labels, vdm_coords = get_vdm_labels_coords(dfa, input_dict[cg_id]['represent_name'], input_dict[cg_id]['correspond_resname'], input_dict[cg_id]['correspond_names'])

    if vdm_coords.shape[0] <=0:
        return

    num_cg_atoms = len(input_dict[cg_id]['lgd_sel'])
    radius = np.sqrt(num_cg_atoms) * rmsd

    dists, inds = get_nearest_vdms_rmsd(vdm_coords, ligand_coords, radius = radius)

    labels_cgs[cg_id] = labels
    df_cgs[cg_id] = dfa
    dist_ind_cgs[cg_id] = (dists, inds)

    return

def construct_vdm_write(outdir, ligands, labels_cgs, input_dict, df_cgs, dist_ind_cgs, clash_radius = 2.7, benchmark_filters = []):
    '''
    dist_ind_cgs: dict. {cg: (dists, inds)}, where dists in shape: (len(ligands), )
    df_cgs: {cg: df}
    '''
    CgCombInfoDict = {}

    for i in range(len(ligands)):
        cgCombInfo = CgCombInfo()
        cgCombInfo.ligand_id = i

        for cg in dist_ind_cgs.keys():
            info = []

            num_cg_atoms = len(input_dict[cg]['lgd_sel'])
            labels = labels_cgs[cg]
            dfa = df_cgs[cg]

            dists = dist_ind_cgs[cg][0][i]
            inds = dist_ind_cgs[cg][1][i]

            for j in range(len(inds)):
                ind = inds[j]
                rmsd = dists[j]/np.sqrt(num_cg_atoms)
                if ind > labels.shape[0]:
                    print(cg)
                    print(labels.shape)
                    print(inds)
                x = labels.iloc[ind]
                v = dfa[(dfa['CG'] == x['CG']) & (dfa['rota'] == x['rota']) & (dfa['probe_name'] == x['probe_name']) & (dfa['seg_chain_resnum'] == x['seg_chain_resnum'])]
                if vdm_ligand_clash(v, ligands[i], clash_radius):
                    continue

                seg_chain_resnum = v['seg_chain_resnum'].iloc[-1]
                #print(seg_chain_resnum)
                resname = v['resname'].iloc[-1]

                if len(benchmark_filters) > 0:
                    if seg_chain_resnum in benchmark_filters.keys():
                        if benchmark_filters[seg_chain_resnum] != resname:
                            continue

                info.append((rmsd, v))
                prefix = 'Lig-'+ str(i) + '_'  + '_'.join(str(z) for z in seg_chain_resnum) + '_key_' + '-'.join(str(z) for z in cg) + '_rmsd_' + str(round(rmsd, 2))  + '_' + resname + '_' # + str(x['C_score_bb_ind']) + '_'
                combs2.design.functions.print_dataframe(v, outpath=outdir, tag = '_' + str(ind), prefix = prefix)
            cgCombInfo.vdm_cgs[cg] = info
        CgCombInfoDict[i] = cgCombInfo

        if any([len(cgCombInfo.vdm_cgs[k]) > 0 for k in cgCombInfo.vdm_cgs.keys()]):
            pr.writePDB(outdir + 'Lig-' + str(i) + '_' + ligands[i].getTitle(), ligands[i])

    return CgCombInfoDict

def write_summary(outdir, CgCombInfoDict, name = '_summary.tsv'):
    with open(outdir + name, 'w') as f:
        f.write('ligand_id\tcg_id\trmsd\n')
        for k in CgCombInfoDict.keys():
            cgInfo = CgCombInfoDict[k]
            for cg in cgInfo.vdm_cgs.keys():
                info = cgInfo.vdm_cgs[cg]
                for o in info:
                    f.write(str(cgInfo.ligand_id) + '\t')
                    f.write(str(cg) + '\t')
                    f.write(str(o[0]) + '\n')
    return

max_stars_repo_name: lonelu/Metalprot

max_issues_repo_name: lonelu/Metalprot

max_forks_repo_name: lonelu/Metalprot

reason: import numpy,from scipy


text:
# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. 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. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.
from __future__ import absolute_import

from sagemaker import Model, RealTimePredictor, Session
from sagemaker.content_types import CONTENT_TYPE_CSV
from sagemaker.fw_registry import registry
from sagemaker.predictor import csv_serializer

framework_name = 'sparkml-serving'
repo_name = 'sagemaker-sparkml-serving'

class SparkMLPredictor(RealTimePredictor):
    """
    Performs predictions against an MLeap serialized SparkML model.

    The implementation of :meth:`~sagemaker.predictor.RealTimePredictor.predict` in this
    `RealTimePredictor` requires a json as input. The input should follow the json format
    as documented.

    ``predict()`` returns a csv output, comma separated if the output is a list.
    """

    def __init__(self, endpoint, sagemaker_session=None):
        """
        Initializes a SparkMLPredictor which should be used with SparkMLModel to perform predictions against SparkML
        models serialized via MLeap. The response is returned in text/csv format which is the default response format
        for SparkML Serving container.
        Args:
            endpoint (str): The name of the endpoint to perform inference on.
            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
                using the default AWS configuration chain.
        """
        sagemaker_session = sagemaker_session or Session()
        super(SparkMLPredictor, self).__init__(endpoint=endpoint, sagemaker_session=sagemaker_session,
                                               serializer=csv_serializer, content_type=CONTENT_TYPE_CSV)

class SparkMLModel(Model):
    """
    Model data and S3 location holder for MLeap serialized SparkML model. Calling
    :meth:`~sagemaker.model.Model.deploy` creates an Endpoint and return
    a Predictor to performs predictions against an MLeap serialized SparkML model .
    """

    def __init__(self, model_data, role=None, spark_version=2.2, sagemaker_session=None, **kwargs):
        """Initialize a SparkMLModel..
        Args:
            model_data (str): The S3 location of a SageMaker model data ``.tar.gz`` file. For SparkML, this will be the
                output that has been produced by the Spark job after serializing the Model via MLeap.
            role (str): An AWS IAM role (either name or full ARN). The Amazon SageMaker training jobs and APIs
                that create Amazon SageMaker endpoints use this role to access training data and model artifacts.
                After the endpoint is created, the inference code might use the IAM role,
                if it needs to access an AWS resource.
            spark_version (str): Spark version you want to use for executing the inference (default: '2.2').
            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
                using the default AWS configuration chain. For local mode, please do not pass this variable.
        """
        # for local mode, sagemaker_session should be passed as None but we need a session to get boto_region_name
        region_name = (sagemaker_session or Session()).boto_region_name
        image = "{}/{}:{}".format(registry(region_name, framework_name), repo_name, spark_version)
        super(SparkMLModel, self).__init__(model_data, image, role, predictor_cls=SparkMLPredictor,
                                           sagemaker_session=sagemaker_session, **kwargs)

max_stars_repo_name: yifeim/sagemaker-python-sdk

max_issues_repo_name: yifeim/sagemaker-python-sdk

max_forks_repo_name: yifeim/sagemaker-python-sdk

reason: from sage


text:
import nltk
from functools import reduce
import random
from os import listdir
import numpy as np
import matplotlib.pyplot as plt

stemmer = nltk.stem.SnowballStemmer('italian')
my_stop_words = ['può','sempre','modo','quando','stesso','senza','ogni','tempo','dopo','fatto','così','cosa','ora','poi']
argomenti = ['conflitti','delusioni','fallimento','inquietudine','odio','sofferenza']
stopwords = nltk.corpus.stopwords.words("italian")

def get_index(s):
    nome = s.split('.')[0]
    indice = int(nome.split('-')[1])
    return indice

def new_img():
    nome = sorted(listdir('img'),key=get_index)[-1]
    indice = get_index(nome)
    indice += 1
    return f'img/pie-{indice}.png'

def distribuzione(lista_parole):
    fd = nltk.FreqDist(lista_parole)
    piu_comuni = fd.most_common(10)
    nome = new_img()
    valori = [coppia[1] for coppia in piu_comuni]
    label = [coppia[0] for coppia in piu_comuni]
    x = np.array(valori)
    plt.switch_backend('agg')
    plt.pie(x,labels=label)
    plt.savefig(nome,transparent=True)
    plt.clf()

    return nome

def estrai_frasi(testo):
    return nltk.sent_tokenize(testo)

def stem_parole(lista_parole):
    """
    Mappo ogni parola della lista nel suo stem
    """
    return [stemmer.stem(parola) for parola in lista_parole]

def pipe(lista_fn):
    return lambda input: reduce(lambda input,fn: fn(input),lista_fn,input)

def estrai_parole(testo):
    """
    Dato un testo, produco una lista delle sole parole utili all' analisi
    """
    tokens = nltk.word_tokenize(testo)
    # rimuovo punteggiatura
    parole = [w for w in tokens if w.isalpha()]
    parole_utili = [w.lower() for w in parole if w.lower() not in stopwords]
    parole_utili = [parola for parola in parole_utili if parola not in my_stop_words]
    parole_utili = [w for w in parole_utili if len(w) > 2]
    return parole_utili

def appendi_frasi(nome_file,label,lista):
    with open(nome_file) as f:
            for line in f.readlines():
                lista.append((line,label))

def mischia_lista(li):
    random.shuffle(li)
    return li

def get_frasi():
    frasi = []
    for argomento in argomenti:
        appendi_frasi(f'frasi_{argomento}.txt','SI',frasi)
    #appendo le frasi positive dal file felicita
    appendi_frasi('frasi_felicita.txt','NO',frasi)
    return frasi

def database():
    return mischia_lista(get_frasi())

if __name__ == '__main__':
    pass
    #test di alcune funzioni
    ##print(len(database()))
    ##print(estrai_parole('Cerca il testo ad un determinato link, e vediamo che umore ne esce'))
    # print(new_img())


text:
import numpy as np
import matplotlib.pyplot as plt
from random import random

Tmax = 1
Tmin = 1e-3
tau = 1e4
x_0 = 2
y_0 = 2

#def f(x,y):
#	return x**2 - np.cos(4*np.pi*x) + (y-1)**2

def f(x,y):
	if 0<x<50 and -20<y<20:
		return np.cos(x) + np.cos(np.sqrt(2)*x + np.cos(np.sqrt(3)*x)) + (y-1)**2
	else:
		return 1e10
t = 0
T = Tmax
fxy = f(x_0,y_0)
x = x_0
y = y_0

x_results = [x]
y_results = [y]
t_results = [t]

while T>Tmin:

    # Cooling
    t += 1
    T = Tmax*np.exp(-t/tau)

    # Choose two cities to swap and make sure they are distinct
    rx = np.random.normal(0,1)
    ry = np.random.normal(0,1)
    oldx = x
    oldy = y
    oldfxy = fxy
    x += rx
    y += ry
    fxy = f(x,y)
    diff = fxy - oldfxy

    t_results.append(t)
    x_results.append(x)
    y_results.append(y)

    if random()>np.exp(-diff/T):
        x = oldx
        y = oldy
        fxy = oldfxy

# print("x =", x, "y =", y,"f(x,y) =", fxy)
fig, graph = plt.subplots()
graph.plot(t_results, x_results, ".")
graph.set(xlabel='time', ylabel='x value',
       title="plot of x vs t")
graph.grid()
fig.savefig("q1b_x.png")
plt.show()

fig, graph = plt.subplots()
graph.plot(t_results, y_results, ".")
graph.set(xlabel='time', ylabel='y value',
       title="plot of y vs t")
graph.grid()
fig.savefig("q1b_y.png")
plt.show()


text:
#
# Parabolic error measure.
#
# This file is part of PINTS.
#  Copyright (c) 2017-2019, University of Oxford.
#  For licensing information, see the LICENSE file distributed with the PINTS
#  software package.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import pints
import numpy as np

class ParabolicError(pints.ErrorMeasure):
    """
    Error measure based ona simple parabola centered around a user specified
    point.

    .. math::
        f(x) = (x - c)^s

    *Extends:* :class:`pints.ErrorMeasure`.
    """
    def __init__(self, c=[0, 0]):
        self._c = pints.vector(c)
        self._n = len(self._c)

    def __call__(self, x):
        return np.sum((self._c - x)**2)

    def n_parameters(self):
        """ See :meth:`pints.ErrorMeasure.n_parameters()`. """
        return self._n

    def optimum(self):
        """
        Returns the global optimum for this function.
        """
        return np.array(self._c, copy=True)

max_stars_repo_name: iamleeg/pints

max_issues_repo_name: iamleeg/pints

max_forks_repo_name: iamleeg/pints


text:
import numpy as np
import logging
from matplotlib.colors import Normalize
import matplotlib.pyplot as plt
import plotly.express as px
from mpl_toolkits.axes_grid1 import make_axes_locatable
import xarray as xr

def plot_ndi(image, factor=1.0, clip_range=None, is_bar=True, **kwargs):
    """
    Utility function for plotting Normalized Difference Index images.

    Args:
        param factor: factor to rescale image.
        type factor: float.

        param clip_range: range to clip image.
        type clip_range: float.

        param is_bar: whether to plot colorbar.
        type is_bar: bool.
    """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(15, 15))
    if clip_range is not None:
        im = ax.imshow(np.clip(image * factor, *clip_range), **kwargs)
    else:
        im = ax.imshow(image * factor, **kwargs)
    if is_bar:
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax)
    ax.set_xticks([])
    ax.set_yticks([])

def plot_truecolor(image, factor=2.5, clip_range=[0,1],is_bar=False):
    """
    Utility function for plotting Normalized Difference Index images.

    Args:
        param factor: factor to rescale image.
        type factor: float.

        param clip_range: range to clip image.
        type clip_range: float.

        param is_bar: whether to plot colorbar.
        type is_bar: bool.
    """

    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(15, 15))
    if clip_range != None:
        im = ax.imshow(np.clip(image[...,[3,2,1]] * factor, *clip_range), vmin=0, vmax=1)
    else:
        im = ax.imshow(image[...,[3,2,1]])
    if is_bar:
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax)
    ax.set_xticks([])
    ax.set_yticks([])

# testing ideas:
# 1.) check input arguments data type
# 2.) check if any inputs arguments are None and if yes return value error
def plot_LST_true(True_Image,LST,date,cmp=None):
    """
    Function which plots satellite image and its corresponding land surface temperature as subplots.
    Args:
        param True_Image: numpy array(RGB) of satellite image of shape (h,w,3)
        type factor: np.array.

        param LST: LST matrix of the corresponding satellite image of shape (h,w).
        type LST: np.array.

        param date: timestamp info of the satellite image so can be added in the plot.
        type is_bar: datetime.datetime .

        param cmap: colormap for the plot. If None looks at the max and min value of the LST matrix
        type cmap: matplotlib cmap object

    """

    # to remove the clipping warning/logging.
    if True_Image is None or LST is None or date is None:
        raise TypeError("NoneType value in one of the arguments")
    if not isinstance(True_Image,np.ndarray):
        raise TypeError('Please provide a np array for true image.')

    if not isinstance(LST,np.ndarray):
        raise TypeError('Please provide a np array for lsr.')
    logger = logging.getLogger()
    old_level = logger.level
    logger.setLevel(100)
    ## the above part is only to remove the clipping warning and has no other relevance to the plot code.
    if cmp is not None:
        norm = Normalize(vmin=cmp[0],vmax=cmp[1])
    else:
        norm = Normalize(vmin=np.min(LST),vmax=np.max(LST))

    plt.rcParams['figure.figsize'] = [12, 8]
    f, (ax1,ax2) = plt.subplots(1,2)
    im = ax2.imshow(LST,cmap=plt.cm.jet,norm=norm)
    f.colorbar(im,orientation="horizontal",fraction=0.07)
    ax1.imshow(True_Image)
    ax1.set_title(date)
    plt.show()

def plot_all_LST(True_Image,LST,date):
    """
    Function which plots all satellite image and there corresponding land surface temperature.
    This function will only work for Jupyter Notebooks for .py u will have to write subplots.
    Args:
        param True_Image: numpy array(RGB) of satellite image of shape (n,h,w,3) where n is number of images
        type factor: np.array.

        param LST: LST matrix of the corresponding satellite image of shape (n,h,w). where n is number of images
        type LST: np.array.

        param date: timestamp info of the satellite image so can be added in the plot.
        type is_bar: datetime.datetime .
    """
    if True_Image is None or LST is None or date is None:
        raise TypeError("NoneType value in one of the arguments")
    if not isinstance(True_Image,np.ndarray):
        raise TypeError('Please provide a np array for true image.')

    if not isinstance(LST,np.ndarray):
        raise TypeError('Please provide a np array for lsr.')

    vmax = np.max(LST)
    vmin = np.min(LST)
    cmp = (vmax,vmin)
    assert True_Image.shape[:-1] == LST.shape
    vSze = True_Image.shape[0]
    for i in range(vSze):
        plot_LST_true(True_Image[i],LST[i],date[i],cmp = None )

def plot_dash_line(df,select_index,height=260,color = 'green'):
    """Function used to plot dash line.

    Args:
        param df : pd.DataFrame contains all plot data. df has the shape (Number of blobs, Timestamps )
        type df: pd.DataFrame

        param select_index: the index of blobs
        type select_index: int

        height (int, optional): This is the default value of height of line map. Defaults to 260.

    Returns:
        fig[px.figure]: the corresponding line plot will be return.
    """

    fig = px.line(df.iloc[select_index],labels={'index':'Time'})
    fig.update_yaxes(autorange=True,automargin=True)
    fig.update_layout(
        showlegend=False,
        height=height,
    margin=dict(r=0, l=0, t=0.3, b=0.1),)

    return fig

def plot_dash_index(df,select_index,year_value,years_timestamp,coord_data,height=300,color_continuous_scale='Greens'):
    """Function used to plot dash index image.

    Args:
        param df : pd.DataFrame contains all plot data. df has the shape (Number of blobs, Timestamps )
        type df: pd.DataFrame

        param select_index: the index of blobs
        type select_index: int

        param year_value: the index of blobs
        type year_value: str or int

        param years_timestamp: the index of blobs
        type years_timestamp: str

        height (int, optional): This is the default value of height of line map. Defaults to 260.

    Returns:
        fig[px.figure]: the corresponding index plot will be return.
    """
    xr_data = df.loc[select_index,str(year_value)]
    xr_time = [sub_year for sub_year in years_timestamp if sub_year.year==year_value]
    space_lati = coord_data[select_index]['space_lati']
    space_long = coord_data[select_index]['space_long']
    xr_array = xr.DataArray(xr_data,coords=[xr_time,space_lati[::-1],space_long],dims=['time','lati','long'])
    fig=px.imshow(xr_array,animation_frame='time',color_continuous_scale=color_continuous_scale)
    fig.update_yaxes(autorange=True,automargin=True)
    fig.update_layout(
        height=300,
    margin=dict(r=0, l=0, t=0.5, b=0.1),)

    return fig

def plot_islands(original,lst,detected_island,limit=None,enhance_radius=False):
    """
    Function which plots all satellite image and there corresponding land surface temperature.
    This function will only work for Jupyter Notebooks for .py u will have to write subplots.
    Args:
        param original: numpy array(RGB) of satellite image of shape (n,h,w,3) where n is number of images
        type factor: np.array.

        param LST: LST matrix of the corresponding satellite image of shape (n,h,w). where n is number of images
        type LST: np.array.

        param detected_island: list of the coordinates of detected island which are to be plotted in the original
        and its corresponding image.
        type detected_island:  list of tuple.

        param enhance_radius: whether to plot circle with bigger radius instead of the one computed by hessian algorithm.
        type enhance_radius: boolean  .

    """
    if original is None or lst is None or detected_island is None:
        raise TypeError("NoneType value in one of the arguments")
    if not isinstance(original,np.ndarray):
        raise TypeError('Please provide a numpy array for original.')

    if not isinstance(lst,np.ndarray):
        raise TypeError('Please provide a numpy array for lst.')
    if not isinstance(detected_island,np.ndarray):
        raise TypeError('Please provide a numpy array for detected_island.')

    fig, axes = plt.subplots(1, 2, figsize=(40, 40), sharex=True, sharey=True)
    ax = axes.ravel()
    #plt.imshow(vLST[0],cmap='gray')
    ax[0].imshow(original)
    ax[1].imshow(lst,cmap=plt.cm.jet)
    for blob in detected_island:
        y, x, r = blob
        if limit is None:
            if enhance_radius:
                c = plt.Circle((x, y), r*3, color='green', linewidth=2, fill=False)
            else:
                c =plt.Circle((x, y), r, color='green', linewidth=2, fill=False)
            c2 = plt.Circle((x, y), r, color='green', linewidth=2, fill=False)
            ax[0].add_patch(c)
            ax[1].add_patch(c2)
        else:
            if r>limit:
                if enhance_radius:
                    print("here")
                    c = plt.Circle((x, y), r*3, color='green', linewidth=2, fill=False)
                else:
                    c =plt.Circle((x, y), r, color='green', linewidth=2, fill=False)
                c2 = plt.Circle((x, y), r, color='green', linewidth=2, fill=False)
                ax[0].add_patch(c)
                ax[1].add_patch(c2)
    plt.show()


text:
# Modified based on the DEQ repo.

import torch
from torch import nn
import torch.nn.functional as functional
from torch.autograd import Function
import numpy as np
import pickle
import sys
import os
from scipy.optimize import root
import time
from termcolor import colored

from mdeq_lib.modules.broyden import _safe_norm, scalar_search_armijo, matvec, rmatvec, line_search

def adj_broyden(
        g,
        x0,
        threshold,
        eps,
        ls=False,
        name="unknown",
        adj_type='C',
        inverse_direction_fun=None,
        inverse_direction_freq=None,
):
    """Adjoint Broyden method

    Parameters:
        - g (fun): the root defining function
        - x0 (torch.Tensor): the initial estimate for the root
        - threshold (int): the maximum number of vectors to store for the Broyden matrix
        - eps (float): the tolerance
        - ls (bool): whether to perform line search. False by default and in practice.
        - name (str): tag to check whether you are in forward or backward mode.
        - adj_type (str): whether to use the B or C type update from the paper
            "On the local convergence of adjoint Broyden methods" Schlenkrich et al. 2010
            Definition (3). For now only 'C' is implemented, the adjoint Broyden residual update.
        - inverse_direction_fun (callable): when called with the current solution, it should return the
            inverse direction to be used for OPA. OPA not used when None.
        - inverse_direction_freq (int): the frequency at which to use the OPA update. Defaults
            to None, which means no OPA is used.
    """
    bsz, total_hsize, n_elem = x0.size()
    dev = x0.device

    x_est = x0           # (bsz, 2d, L')
    gx = g(x_est)        # (bsz, 2d, L')
    nstep = 0
    n_updates = 0
    n_opa_updates = 0
    tnstep = 0
    LBFGS_thres = min(threshold, 27)

    # For fast calculation of inv_jacobian (approximately)
    Us = torch.zeros(bsz, total_hsize, n_elem, LBFGS_thres).to(dev)
    VTs = torch.zeros(bsz, LBFGS_thres, total_hsize, n_elem).to(dev)
    update = gx
    new_objective = init_objective = torch.norm(gx).item()
    prot_break = False
    trace = [init_objective]
    new_trace = [-1]

    # To be used in protective breaks
    protect_thres = 1e6 * n_elem
    lowest = new_objective
    lowest_xest, lowest_gx, lowest_step = x_est, gx, nstep
    opa = inverse_direction_fun is not None and inverse_direction_freq is not None
    while new_objective >= eps and nstep < threshold:
        if opa and nstep > 1 and nstep % inverse_direction_freq == 0:
            n_opa_updates += 1
            inverse_direction = inverse_direction_fun(x_est).detach()
            e = matvec(Us[:,:,:,:n_updates], VTs[:,:n_updates], inverse_direction)
            e = e / torch.norm(e) * torch.norm(update)
            gx.backward(e, retain_graph=False)
            b = x_temp.grad
            a = matvec(Us[:,:,:,:n_updates], VTs[:,:n_updates], e)
            b = rmatvec(Us[:,:,:,:n_updates], VTs[:,:n_updates], b)
            c = (e - b) / torch.einsum('bij, bij -> b', b, e)[:, None, None]
            x_temp.grad.zero_()
            u = a
            vT = c
            if torch.isnan(vT).any() or torch.isnan(u).any():
                msg = colored(f"WARNING: OPA resulted in some nan for the qN matrices.", 'red')
                print(msg)
            vT[vT != vT] = 0
            u[u != u] = 0
            n_updates += 1
            VTs[:,(n_updates-1) % LBFGS_thres] = vT
            Us[:,:,:,(n_updates-1) % LBFGS_thres] = u
        x_est, delta_x, ite = line_search(update, x_est, gx, g, nstep=nstep, on=ls, compute_g=False)
        x_temp = x_est.clone().detach().requires_grad_()
        with torch.enable_grad():
            gx = g(x_temp)
        nstep += 1
        tnstep += (ite+1)
        new_objective = torch.norm(gx).item()
        trace.append(new_objective)
        try:
            new2_objective = torch.norm(delta_x).item() / (torch.norm(x_est - delta_x).item())   # Relative residual
        except:
            new2_objective = torch.norm(delta_x).item() / (torch.norm(x_est - delta_x).item() + 1e-9)
        new_trace.append(new2_objective)
        if new_objective < lowest:
            lowest_xest, lowest_gx = x_est.clone().detach(), gx.clone().detach()
            lowest = new_objective
            lowest_step = nstep
        if new_objective < eps:
            break
        if new_objective < 3*eps and nstep > 30 and np.max(trace[-30:]) / np.min(trace[-30:]) < 1.3:
            # if there's hardly been any progress in the last 30 steps
            break
        if new_objective > init_objective * protect_thres:
            prot_break = True
            break

        # this is the part that changes between Broyden and Adjoint Broyden
        if adj_type != 'C':
            raise NotImplementedError('Use adj_type C for now')
        else:
            sigma = gx
        n_updates += 1
        part_Us, part_VTs = Us[..., :n_updates - 1], VTs[:, :n_updates - 1]
        # a = An^{-1} sigma
        a = matvec(part_Us, part_VTs, sigma)
        # b = sigma^T g'(xn) An^{-1}
        #######
        # Backprop on g
        #######
        gx.backward(sigma, retain_graph=(opa and nstep>1 and nstep%inverse_direction_freq==0))
        b = x_temp.grad
        #######
        b = rmatvec(part_Us, part_VTs, b)
        # c = (sigma - b) / (b^T sigma)
        c = (sigma - b) / torch.einsum('bij, bij -> b', b, sigma)[:, None, None]
        x_temp.grad.zero_()
        # these next 2 assignments allow us to get back to the original writing of
        # broyden by Shaojie
        u = a
        vT = c
        if torch.isnan(vT).any() or torch.isnan(u).any():
            msg = colored(f"WARNING: Adjoint Broyden resulted in some nan for the qN matrices.", 'red')
            print(msg)
        vT[vT != vT] = 0
        u[u != u] = 0
        VTs[:,(n_updates-1) % LBFGS_thres] = vT
        Us[:,:,:,(n_updates-1) % LBFGS_thres] = u
        update = -matvec(Us[:,:,:,:n_updates], VTs[:,:n_updates], gx)

    # NOTE: why was this present originally? is it a question of memory?
    # Us, VTs = None, None
    return {"result": lowest_xest,
            "nstep": nstep,
            "tnstep": tnstep,
            "lowest_step": lowest_step,
            'n_opa_updates': n_opa_updates,
            'opa_freq': inverse_direction_freq,
            "diff": torch.norm(lowest_gx).item(),
            "diff_detail": torch.norm(lowest_gx, dim=1),
            "prot_break": prot_break,
            "trace": trace,
            "new_trace": new_trace,
            "eps": eps,
            "threshold": threshold,
            "Us": Us,
            "VTs": VTs}

max_stars_repo_name: zaccharieramzi/shine

max_issues_repo_name: zaccharieramzi/shine

max_forks_repo_name: zaccharieramzi/shine


text:
import sys
from graph_stuffie import Node, Graph

import networkx as nx
import matplotlib.pyplot as plt

import pydot
from networkx.drawing.nx_pydot import write_dot

print("Loading Graph...")
g = Graph()
g.load(sys.argv[1])

print("Creatig nx graph...")
G = nx.DiGraph()

print("Adding nodes to nx graph...")
G.add_nodes_from(list(g.nodes.keys()))
# print(g.nodes.keys())

print("Adding edges to nx graph...")
for node in g.nodes:
    edges = g.nodes[node].edges
    edges = [(node, other_node) for other_node in edges]
    G.add_edges_from(edges)

print("Drawing graph")
layout = nx.spring_layout(G)
nx.draw_networkx_labels(G, layout)
nx.draw_networkx_nodes(G, layout, cmap=plt.get_cmap('jet'), node_size=100)
nx.draw_networkx_edges(G, layout, edge_list=G.edges(), arrows=True)

write_dot(G, 'plot.dot')

max_stars_repo_name: hthuwal/ARC-Challenge

max_issues_repo_name: hthuwal/ARC-Challenge

max_forks_repo_name: hthuwal/ARC-Challenge

reason: import networkx,from networkx


text:
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
import tempfile

import numpy as np

from dense_deep_event_stereo import _events
from dense_deep_event_stereo import events

def create_mockup_event_sequence():
    """Creates mockup events sequence."""
    np.random.seed(0)
    timestamp = np.array([0.5, 0.6, 0.2, 0.3, 0.1])
    polarity = np.array([-1, 1, 1, 1, -1])
    x = np.array([50, 200, 200, 300, 150])
    y = np.array([100, 50, 50, 0, 150])
    camera = np.array([1, 1, 0, 0, 0])
    disparity = np.full((5), float('inf'))

    features = np.stack([timestamp, x, y, polarity, camera, disparity], axis=1)
    features_names = [
        'timestamp', 'x', 'y', 'polarity', 'camera_id', 'disparity'
    ]
    event_sequence = events.EventSequence(features,
                                          features_names,
                                          image_height=260,
                                          image_width=346)
    event_sequence.sort_by_timestamp()
    return event_sequence

def test_compute_spatial_binning():
    sequence = np.array(
        [[0, 0, 1], [0, 0, 2], [0, 0, 3], [1, 0, 4], [1, 0, 5], [0, 0, 0]],
        dtype=np.float).reshape((1, 6, 3))
    sequence_length = np.array([5], dtype=np.long)
    features_indices = np.array([0, 2], dtype=np.long)
    height = 3
    width = 2
    maximum_local_sequence_length = 2
    y_column = 1
    x_column = 0

    local_sequences, local_sequences_lenghts = _events.compute_spatial_binning(
        sequence, sequence_length, features_indices, height, width,
        maximum_local_sequence_length, y_column, x_column)

    expected_local_sequences_lenghts = np.zeros((height, width), dtype=np.long)
    expected_local_sequences_lenghts[0, 0] = 2
    expected_local_sequences_lenghts[0, 1] = 2
    assert np.all(expected_local_sequences_lenghts == local_sequences_lenghts)
    expected_local_sequences = np.zeros(
        (height, width, maximum_local_sequence_length, 2))
    expected_local_sequences[0, 0, 0, :] = np.array([0, 1])
    expected_local_sequences[0, 0, 1, :] = np.array([0, 2])
    expected_local_sequences[0, 1, 0, :] = np.array([1, 4])
    expected_local_sequences[0, 1, 1, :] = np.array([1, 5])
    assert np.all(expected_local_sequences == local_sequences)

def test_event_sequence():
    file_with_events = tempfile.mkstemp(suffix='.npy')[1]
    np.save(file_with_events, np.array([
        [0.1, 10, 10, -1],
        [0.2, 30, 100, 1],
    ]))
    sequence_from_file = events.EventSequence.from_npy(file_with_events,
                                                       camera_name='left',
                                                       image_height=260,
                                                       image_width=346)
    assert len(sequence_from_file) == 2
    assert {
        'timestamp': 0.1,
        'x': 10.0,
        'y': 10.0,
        'polarity': -1.0,
        'camera_id': 0.0,
        'disparity': float('inf')
    } == sequence_from_file[0]

    sequence = create_mockup_event_sequence()

    assert np.all(
        sequence._features[:, 0] == np.array([0.1, 0.2, 0.3, 0.5, 0.6]))
    assert len(sequence) == 5
    assert sequence.start_timestamp() == 0.1
    assert sequence.end_timestamp() == 0.6
    assert {
        'timestamp': 0.1,
        'x': 150.0,
        'y': 150.0,
        'polarity': -1.0,
        'camera_id': 0.0,
        'disparity': float('inf')
    } == sequence[0]

    sequence.set_disparity(0, 10)
    assert {
        'timestamp': 0.1,
        'x': 150.0,
        'y': 150.0,
        'polarity': -1.0,
        'camera_id': 0.0,
        'disparity': 10.0
    } == sequence[0]

    image = sequence.to_image()
    assert np.all(image[:, 0, 0] == np.array([[255, 255, 255]]))
    assert np.all(image[:, 100, 50] == np.array([[0, 0, 255]]))
    assert np.all(image[:, 50, 200] == np.array([[255, 0, 0]]))

    disparity_image = sequence.to_disparity_image()
    assert disparity_image[0, 0] == 255
    assert disparity_image[150, 150] == 7 * 10

    negative_event_sequence = sequence.filter_by_polarity(-1)
    assert not (
        negative_event_sequence.
        _features[:, negative_event_sequence.get_column_index('polarity')] == 1
    ).any()

    sequence0 = sequence.filter_by_camera_name('left')
    sequence1 = sequence.filter_by_camera_name('right')
    assert len(sequence0) == 3
    assert len(sequence1) == 2
    assert len(sequence) == len(sequence0 + sequence1)

    sequence_half = sequence.filter_by_timestamp(0.0, 0.4)
    assert len(sequence_half) == 3

    sequence_clone = sequence.clone()
    assert np.all(sequence_clone._features == sequence._features)

    sequence.spatial_hashing()
    assert sequence._spatial_hash_table is not None
    assert sequence._events_in_location is not None

    # Check if the copy is deep.
    sequence_copy = sequence.clone()
    sequence_copy._features_names += ['embedding']
    sequence_copy._features[0, 0] = 1000
    assert len(sequence._features_names) == 6
    assert len(sequence._features) != 1000

def test_compute_maximum_events_in_location():
    event_sequence = np.array([[0.1, 0, 0], [0.2, 0, 0], [0.3, 1, 0]])
    y_column = 2
    x_column = 1
    image_height = 2
    image_width = 2

    assert 2 == _events.compute_maximum_events_in_location(
        event_sequence, y_column, x_column, image_height, image_width)

def test_compute_spatial_hash_table():
    event_sequence = np.array([[0.1, 0, 0], [0.2, 0, 0], [0.3, 1, 0]],
                              dtype=np.float)
    y_column = 2
    x_column = 1
    image_height = 2
    image_width = 2

    expected_spatial_hash_table = np.zeros((image_height, image_height, 2),
                                           dtype=np.int)
    expected_spatial_hash_table[0, 0, 0] = 0
    expected_spatial_hash_table[0, 0, 1] = 1
    expected_spatial_hash_table[0, 1, 0] = 2

    expected_events_in_location = np.array([[2, 1], [0, 0]], dtype=np.int)

    actual_spatial_hash_table, actual_events_in_location = \
        _events.compute_spatial_hash_table(event_sequence, y_column,
                                           x_column, image_height, image_width)
    assert np.all(actual_spatial_hash_table == expected_spatial_hash_table)
    assert np.all(actual_events_in_location == expected_events_in_location)