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

import os
from os.path import join, dirname, abspath
import sys

import argparse

import multiprocessing
import numpy as np

import torch
from torch import nn, optim, Tensor
import torch.nn.functional as F

from ctp.util import make_batches
from ctp.clutrr import Fact, Data, Instance, accuracy
from simple import DataParser

from ctp.clutrr.models import BatchNeuralKB
from model_simple import BatchHoppy

from ctp.reformulators import BaseReformulator
from ctp.reformulators import StaticReformulator
from ctp.reformulators import LinearReformulator
from ctp.reformulators import AttentiveReformulator
from ctp.reformulators import MemoryReformulator
from ctp.reformulators import NTPReformulator

from ctp.kernels import BaseKernel, GaussianKernel
from ctp.regularizers import N2, N3, Entropy

from typing import List, Tuple, Dict, Optional

import logging

logger = logging.getLogger(os.path.basename(sys.argv[0]))
np.set_printoptions(linewidth=256, precision=4, suppress=True, threshold=sys.maxsize)

torch.set_num_threads(multiprocessing.cpu_count())

def decode(vector: Tensor,
           kernel: BaseKernel,
           relation_embeddings: nn.Module) -> Tuple[int, float]:
    weight = relation_embeddings.weight
    k = kernel.pairwise(vector, weight)[0, :]
    top_idx = k.argmax(dim=0).item()
    top_score = k[top_idx].item()
    return top_idx, top_score

def show_rules(model: BatchHoppy,
               kernel: BaseKernel,
               relation_embeddings: nn.Embedding,
               relation_to_idx: Dict[str, int],
               device: Optional[torch.device] = None):
    idx_to_relation = {i: r for r, i in relation_to_idx.items()}

    rel_idx_pair_lst = sorted(relation_to_idx.items(), key=lambda kv: kv[1])

    for r, i in rel_idx_pair_lst:
        indices = torch.tensor([i], dtype=torch.long, device=device)

        r_emb = relation_embeddings(indices)

        hops_lst = [p for p in model.hops_lst]

        for reformulator, _ in hops_lst:
            def _to_pair(hop: Tensor) -> Tuple[str, float]:
                idx, score = decode(hop, kernel, relation_embeddings)
                rel = idx_to_relation[idx]
                return rel, score

            hop_tensor_lst = [hop for hop in reformulator(r_emb)]

            r_hops = [_to_pair(hop) for hop in hop_tensor_lst]
            print(r, ' ← ', ', '.join(f'({a} {b:.4f})' for a, b in r_hops))

            if isinstance(model.model, BatchHoppy):
                for _r_emb in hop_tensor_lst:
                    _hops_lst = [p for p in model.model.hops_lst]

                    for j, (_reformulator, _) in enumerate(_hops_lst):
                        _hop_tensor_lst = [_hop for _hop in _reformulator(_r_emb)]
                        _r_hops = [_to_pair(_hop) for _hop in _hop_tensor_lst]
                        print(j, ' ← ', ', '.join(f'({_a} {_b:.4f})' for _a, _b in _r_hops))

    return

class Batcher:
    def __init__(self,
                 batch_size: int,
                 nb_examples: int,
                 nb_epochs: int,
                 random_state: Optional[np.random.RandomState]):
        self.batch_size = batch_size
        self.nb_examples = nb_examples
        self.nb_epochs = nb_epochs
        self.random_state = random_state

        size = self.nb_epochs * self.nb_examples
        self.curriculum = np.zeros(size, dtype=np.int32)

        for epoch_no in range(nb_epochs):
            start, end = epoch_no * nb_examples, (epoch_no + 1) * nb_examples
            if self.random_state is not None:
                self.curriculum[start: end] = self.random_state.permutation(nb_examples)
            else:
                self.curriculum[start: end] = np.arange(nb_examples)

        self.batches = make_batches(self.curriculum.shape[0], self.batch_size)
        self.nb_batches = len(self.batches)

    def get_batch(self,
                  batch_start: int,
                  batch_end: int) -> np.ndarray:
        return self.curriculum[batch_start:batch_end]

# takes in the list of facts [F,3]
# gives back the list of the embeddings corresponding to the relations of each fact [F,E]
def encode_relation(facts: List[Fact], # [F,3]: [(s,r,o), (s,r,o), ..., (s,r,o)]
                    relation_embeddings: Tensor,
                    relation_to_idx: Dict[str, int],
                    device: Optional[torch.device] = None) -> Tensor:
    indices_np = np.array([relation_to_idx[r] for _, r, _ in facts], dtype=np.int64)
    indices = torch.tensor(indices_np, dtype=torch.long, device=device)
    res = F.embedding(indices, relation_embeddings)
    return res

# takes in the list of facts [F,3]
# gives back the list of the embeddings corresponding to the subject and object of each fact [F,E],[F,E]
def encode_arguments(facts: List[Fact],
                     entity_embeddings: Tensor,
                     entity_to_idx: Dict[str, int],
                     device: Optional[torch.device] = None) -> Tuple[Tensor, Tensor]:
    indices_np = np.array([[entity_to_idx[s], entity_to_idx[o]] for s, _, o in facts], dtype=np.int64)
    indices = torch.tensor(indices_np, dtype=torch.long, device=device)
    emb = F.embedding(indices, entity_embeddings)
    return emb[:, 0, :], emb[:, 1, :]

# takes in the list of facts [F,3]
# gives back the list of the embeddings corresponding to all entities that are in the facts as either
# subject or object
# [N,E], where N is the number of different entities in the facts
def encode_entities(facts: List[Fact],
                    entity_embeddings: Tensor,
                    entity_to_idx: Dict[str, int],
                    device: Optional[torch.device]) -> Tensor:
    indices_lst = sorted({entity_to_idx[e] for s, r, o in facts for e in {s, o}})
    indices = torch.tensor(indices_lst, dtype=torch.long, device=device)
    emb = F.embedding(indices, entity_embeddings)
    return emb

def main(argv):

    debug = True

    train_path = test_path = join(dirname(dirname(abspath(__file__))),'data', 'clutrr-emnlp', 'data_test', '64.csv') # "data/clutrr-emnlp/data_test/64.csv"
    test_paths = [test_path]

    # model params

    # the size of the embedding of each atom and relationship
    embedding_size = 20
    # when proving the body of a rule, we consider the k best substitutions for each variable
    k_max = 10
    # how many times to reformulate the goal(s) --> bigger for bigger graph: this is for training
    max_depth = 1
    # how many times to reformulate the goal(s): this is for testing --> this depth can be bigger than for training
    test_max_depth = None

    # the shape of the reformulation:
    # 2: goal(X,Z) -> p(X,Y), q(Y,Z) (2 elements in the body)
    # 1: goal(X,Z) -> r(X,Z) (1 element in the body)
    # 1R: goal(X,Z) -> s(Z,X) (variables in reversed order)
    # if we have multiple in the array, that means at each reformulation step we actually reformulate the same goal
    # multiple times according to the elements in the array (so we have more choice to get to a good proof)
    hops_str = ['2'] # ['2', '2', '1R']

    # training params

    nb_epochs = 1 # 100
    learning_rate = 0.1
    # training batch size
    batch_size = 8
    # testing batch size --> this can be smaller than for training
    test_batch_size = batch_size # could be other as well

    optimizer_name = 'adagrad' # choices = ['adagrad', 'adam', 'sgd']

    seed = 0 # int

    # how often you want to evaluate
    evaluate_every = 1 # int
    evaluate_every_batches = None # int

    # whether you want to regularize
    #argparser.add_argument('--N2', action='store', type=float, default=None)
    #argparser.add_argument('--N3', action='store', type=float, default=None)
    #argparser.add_argument('--entropy', '-E', action='store', type=float, default=None)

    scoring_type = 'concat' # choices = ['concat', 'min']

    # how to get the score combined from the conjunction of the parts of the body in a rule
    # e.g. goal(X,Y) :- p(X,Z), q(Z,Y) --> how do we combine the scores of p and q
    # I can just keep min, it works well, and later add the others potentially
    tnorm_name = 'min' # choices = ['min', 'prod', 'mean']
    # which function to use for reformulating a goal to subgoals
    # I can just keep linear for the initial version, and if it works, I can add more
    reformulator_name = 'linear' # choices = ['static', 'linear', 'attentive', 'memory', 'ntp'] --> deleted code part

    nb_gradient_accumulation_steps = 1 # int

#### IGNORE --> code as well
    gntp_R = None # int

    slope = None # float
    init_size = 1.0 # float

    init_type = 'uniform'
    ref_init_type = 'random'

#### IGNORE
    #argparser.add_argument('--fix-relations', '--FR', action='store_true', default=False)
    is_fixed_relations = False
    # whether you want to train on the smallest graph first (it's easier to train on them)
    start_simple = None # int

    is_debug = False

    load_path = None # str
    save_path = None # str

    # took out predicates from the code, no need for them (they take only 1 argument, but an edge in PyG always takes 2)
    #is_predicate = False

# IGNORED variables:
    #N2_weight = args.N2
    #N3_weight = args.N3
    #entropy_weight = args.entropy

    #is_fixed_relations = args.fix_relations

    np.random.seed(seed)
    random_state = np.random.RandomState(seed)
    torch.manual_seed(seed)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f'Device: {device}')

    if torch.cuda.is_available():
        torch.set_default_tensor_type(torch.cuda.FloatTensor)

# Initializing data and embeddings

    data = DataParser(train_path=train_path, test_paths=test_paths)
    entity_lst, relation_lst = data.entity_lst, data.relation_lst

    test_relation_lst = ["aunt", "brother", "daughter", "daughter-in-law", "father", "father-in-law", "granddaughter",
                         "grandfather", "grandmother", "grandson", "mother", "mother-in-law", "nephew", "niece",
                         "sister", "son", "son-in-law", "uncle"]

    nb_entities = len(entity_lst)
    nb_relations = len(relation_lst)

    entity_to_idx = {e: i for i, e in enumerate(entity_lst)}
    relation_to_idx = {r: i for i, r in enumerate(relation_lst)}

    kernel = GaussianKernel(slope=slope)

    entity_embeddings = nn.Embedding(nb_entities, embedding_size, sparse=True).to(device)
    nn.init.uniform_(entity_embeddings.weight, -1.0, 1.0)
    entity_embeddings.requires_grad = False

    relation_embeddings = nn.Embedding(nb_relations, embedding_size, sparse=True).to(device)

    if is_fixed_relations is True:
        relation_embeddings.requires_grad = False

    if init_type in {'uniform'}:
        nn.init.uniform_(relation_embeddings.weight, -1.0, 1.0)

    relation_embeddings.weight.data *= init_size

    # the model that lets you look up in the KB
    model = BatchNeuralKB(kernel=kernel, scoring_type=scoring_type).to(device)
    memory: Dict[int, MemoryReformulator.Memory] = {}

    # generates a reformulator of type given as an argument
    def make_hop(s: str) -> Tuple[BaseReformulator, bool]:
        nonlocal memory
        if s.isdigit():
            nb_hops, is_reversed = int(s), False
        else:
            nb_hops, is_reversed = int(s[:-1]), True
        res = None

        if reformulator_name in {'linear'}:
            res = LinearReformulator(nb_hops, embedding_size, init_name=ref_init_type)

        assert res is not None
        return res.to(device), is_reversed

    hops_lst = [make_hop(s) for s in hops_str] # hops_str = [2,2,1R]

    # "model" is a neural KB for checking whether the facts are true or not
    # hoppy is the model that does the reasoning, using the neural KB
    hoppy = BatchHoppy(model=model, k=k_max, depth=max_depth, tnorm_name=tnorm_name,
                       hops_lst=hops_lst, R=gntp_R).to(device)

    def scoring_function(instances_batch: List[Instance],
                         relation_lst: List[str],
                         is_train: bool = False,
                         _depth: Optional[int] = None) -> Tuple[Tensor, List[Tensor]]:

        rel_emb_lst: List[Tensor] = []
        arg1_emb_lst: List[Tensor] = []
        arg2_emb_lst: List[Tensor] = []

        story_rel_lst: List[Tensor] = []
        story_arg1_lst: List[Tensor] = []
        story_arg2_lst: List[Tensor] = []

        embeddings_lst: List[Tensor] = []

        label_lst: List[int] = []

        for i, instance in enumerate(instances_batch):

            story, target = instance.story, instance.target
            s, r, o = target # 1 target relation

            # the relation embeddings from all facts (==story)
            # [F,E], where F is the number of facts
            story_rel = encode_relation(story, relation_embeddings.weight, relation_to_idx, device)
            # the subject,object embeddings from all facts (==story)
            # [F,E],[F,E]
            story_arg1, story_arg2 = encode_arguments(story, entity_embeddings.weight, entity_to_idx, device)

            # the subject,object embeddings from all facts (==story) as a list of all different entity embeddings
            # [N,E], where N is the number of different entities in the facts
            embeddings = encode_entities(story, entity_embeddings.weight, entity_to_idx, device)

            # target subject (s), object(o) paired with all possible relations from the relation_lst (x)
            # [R,3] where R is the number of relations
            target_lst: List[Tuple[str, str, str]] = [(s, x, o) for x in relation_lst]

            assert len(target_lst) == len(test_relation_lst)

            # for each instance in the batch we are iterating through: [0,0,...,0,1,0,0,...,0] list
            # where 1 is at the relation that is the target relation
            label_lst += [int(tr == r) for tr in relation_lst]

            # relation embeddings of the target subject,object paired with all possible relations - R relations
            # [R,E]
            rel_emb = encode_relation(target_lst, relation_embeddings.weight, relation_to_idx, device)
            # target subject,object embeddings
            # [R,E],[R,E] (actually all will be the same as the subject, object doesn't change just the relations)
            arg1_emb, arg2_emb = encode_arguments(target_lst, entity_embeddings.weight, entity_to_idx, device)

            batch_size = rel_emb.shape[0] # R - the number of relations --> B from now on
            fact_size = story_rel.shape[0] # F - the number of facts in "story"
            entity_size = embeddings.shape[0] # N - the number of different entities in the facts (=="story")

            # [B, E]
            # at each instance in the batch we add [B,E] to it (B==R)
            # --> rel_emb_lst will become [batch,B,E], where batch is number of instances in batch
            rel_emb_lst += [rel_emb]
            arg1_emb_lst += [arg1_emb]
            arg2_emb_lst += [arg2_emb]

            # [B, F, E]
            # repeat the same facts for each rel subst in instance of current batch
            # (1 batch will be all possible relation substitutions for the target relation)
            # --> story_rel_lst will become [batch,B,F,E], where batch is number of instances in batch
            story_rel_lst += [story_rel.view(1, fact_size, -1).repeat(batch_size, 1, 1)]
            story_arg1_lst += [story_arg1.view(1, fact_size, -1).repeat(batch_size, 1, 1)]
            story_arg2_lst += [story_arg2.view(1, fact_size, -1).repeat(batch_size, 1, 1)]

            # [B, N, E]
            # repeat the same entity embeddings for each instance in batch (== for each relation)
            # (1 batch will be all possible relation substitutions for the target relation)
            # --> embeddings_lst will become [batch,B,N,E], where batch is number of instances in batch
            embeddings_lst += [embeddings.view(1, entity_size, -1).repeat(batch_size, 1, 1)]

        def cat_pad(t_lst: List[Tensor]) -> Tuple[Tensor, Tensor]:
            # t: [B,F,E] --> t.shape[1] == F --> F possibly different each time
            # OR
            # t: [B,N,E] --> t.shape[1] == N --> N possibly different each time
            lengths: List[int] = [t.shape[1] for t in t_lst]
            max_len: int = max(lengths)

            # pad the _t tensor: [B,F,E] (which is all t in t_lst) with embeddings of 0 where F is smaller than max_len
            # OR
            # pad the _t tensor: [B,N,E] (which is all t in t_lst) with embeddings of 0 where N is smaller than max_len
            def my_pad(_t: Tensor, pad: List[int]) -> Tensor:
                return torch.transpose(F.pad(torch.transpose(_t, 1, 2), pad=pad), 1, 2)

            # [batch*B,F,E], where each F is the same (sometimes padded with extra 0 embeddings to reach that)
            # OR
            # [batch*B,N,E], where each N is the same (sometimes padded with extra 0 embeddings to reach that)
            res_t: Tensor = torch.cat([my_pad(t, pad=[0, max_len - lengths[idx]]) for idx, t in enumerate(t_lst)], dim=0)
            # [batch*B], the number of facts (before padding) in an instance
            # OR
            # [batch*B], the number of entities (before padding) in an instance
            res_l: Tensor = torch.tensor([t.shape[1] for t in t_lst for _ in range(t.shape[0])],
                                         dtype=torch.long, device=device)
            return res_t, res_l # [batch*B,F,E],[batch*B] OR [batch*B,N,E],[batch*B]

        # [batch*B,E], where batch is number of batches & B is batch_size that is number of relations
        rel_emb = torch.cat(rel_emb_lst, dim=0)
        arg1_emb = torch.cat(arg1_emb_lst, dim=0)
        arg2_emb = torch.cat(arg2_emb_lst, dim=0)

        # story_rel,story_arg1,story_arg2: [batch*B,F,E], where each F is the same (sometimes padded with extra 0 embeddings to reach that)
        # nb_facts: [batch*B], the number of facts (before padding) in each instance
        story_rel, nb_facts = cat_pad(story_rel_lst) # story_rel_lst: [batch,B,F,E]
        story_arg1, _ = cat_pad(story_arg1_lst)
        story_arg2, _ = cat_pad(story_arg2_lst)
        facts = [story_rel, story_arg1, story_arg2] # [3,batch*B,F,E]

        # _embeddings: [batch*B,N,E], where each N is the same (sometimes padded with extra 0 embeddings to reach that)
        # nb_embeddings: [batch*B], the number of entities (before padding) in each instance
        _embeddings, nb_embeddings = cat_pad(embeddings_lst)

        max_depth_ = hoppy.depth
        if not is_train and test_max_depth is not None:
            hoppy.depth = test_max_depth

        if _depth is not None:
            hoppy.depth = _depth

        scores = hoppy.score(rel_emb, arg1_emb, arg2_emb, facts, nb_facts, _embeddings, nb_embeddings)

        if not is_train and test_max_depth is not None:
            hoppy.depth = max_depth_

        if _depth is not None:
            hoppy.depth = max_depth_ #??

        return scores, [rel_emb, arg1_emb, arg2_emb]

    def evaluate(instances: List[Instance],
                 path: str,
                 sample_size: Optional[int] = None) -> float:
        res = 0.0
        if len(instances) > 0:
            res = accuracy(scoring_function=scoring_function,
                           instances=instances,
                           sample_size=sample_size,
                           relation_lst=test_relation_lst,
                           batch_size=test_batch_size,
                           relation_to_predicate=None,
                           is_debug=is_debug)
            logger.info(f'Test Accuracy on {path}: {res:.6f}')
        return res

    loss_function = nn.BCELoss()

    params_lst = [p for p in hoppy.parameters() if not torch.equal(p, entity_embeddings.weight)]

    if is_fixed_relations is False:
        params_lst += relation_embeddings.parameters()

    params = nn.ParameterList(params_lst).to(device)

    if load_path is not None:
        model.load_state_dict(torch.load(load_path))

    for tensor in params_lst:
        logger.info(f'\t{tensor.size()}\t{tensor.device}')

    #TODO: set just adam for example, no need for all these choices
    optimizer_factory = {
        'adagrad': lambda arg: optim.Adagrad(arg, lr=learning_rate),
        'adam': lambda arg: optim.Adam(arg, lr=learning_rate),
        'sgd': lambda arg: optim.SGD(arg, lr=learning_rate)
    }

    assert optimizer_name in optimizer_factory
    optimizer = optimizer_factory[optimizer_name](params)

    global_step = 0

    for epoch_no in range(1, nb_epochs + 1):

        training_set, is_simple = data.train_graph, False # HeteroData

### IGNORE for now
        #if start_simple is not None and epoch_no <= start_simple:
        #    training_set = [ins for ins in training_set if len(ins.story) == 2]
        #    is_simple = True
        #    logger.info(f'{len(data.train)} → {len(training_set)}')

        batcher = Batcher(batch_size=batch_size, nb_examples=len(training_set), nb_epochs=1, random_state=random_state)

        nb_batches = len(batcher.batches)
        epoch_loss_values = []

        counter = 0
        for batch_no, (batch_start, batch_end) in enumerate(batcher.batches, start=1):
            counter += 1
            if debug:
                if counter > 1:
                    break
            global_step += 1

            # getting current batch from the training set
            indices_batch = batcher.get_batch(batch_start, batch_end)
            instances_batch = [training_set[i] for i in indices_batch]

            # label_lst: list of 1s and 0s indicating which query (==target) relation is where in the test_relation_lst
            label_lst: List[int] = [int(ins.target[1] == tr) for ins in instances_batch for tr in test_relation_lst]

            labels = torch.tensor(label_lst, dtype=torch.float32, device=device)

            #returns scores of what??
            scores, query_emb_lst = scoring_function(instances_batch,
                                                     test_relation_lst,
                                                     is_train=True,
                                                     _depth=1 if is_simple else None)

            loss = loss_function(scores, labels)

            factors = [hoppy.factor(e) for e in query_emb_lst]

            loss_value = loss.item()
            epoch_loss_values += [loss_value]

            if nb_gradient_accumulation_steps > 1:
                loss = loss / nb_gradient_accumulation_steps

            loss.backward()

            if nb_gradient_accumulation_steps == 1 or global_step % nb_gradient_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()

            logger.info(f'Epoch {epoch_no}/{nb_epochs}\tBatch {batch_no}/{nb_batches}\tLoss {loss_value:.4f}')

            if evaluate_every_batches is not None:
                if global_step % evaluate_every_batches == 0:
                    for test_path in test_paths:
                        evaluate(instances=data.test[test_path], path=test_path)

        if epoch_no % evaluate_every == 0:
            for test_path in test_paths:
                evaluate(instances=data.test[test_path], path=test_path)

            # TODO: don't need for now, but we may look at it later
            if is_debug is True:
                with torch.no_grad():
                    show_rules(model=hoppy, kernel=kernel, relation_embeddings=relation_embeddings,
                               relation_to_idx=relation_to_idx, device=device)

        loss_mean, loss_std = np.mean(epoch_loss_values), np.std(epoch_loss_values)

        slope = kernel.slope.item() if isinstance(kernel.slope, Tensor) else kernel.slope
        logger.info(f'Epoch {epoch_no}/{nb_epochs}\tLoss {loss_mean:.4f} ± {loss_std:.4f}\tSlope {slope:.4f}')

    import time
    start = time.time()

    for test_path in test_paths:
        evaluate(instances=data.test[test_path], path=test_path)

    end = time.time()
    logger.info(f'Evaluation took {end - start} seconds.')

    if save_path is not None:
        torch.save(model.state_dict(), save_path)

    logger.info("Training finished")

if __name__ == '__main__':
    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
    print(' '.join(sys.argv))
    main(sys.argv[1:])

max_stars_repo_name: Vikicsizmadia/ctp

max_issues_repo_name: Vikicsizmadia/ctp

max_forks_repo_name: Vikicsizmadia/ctp


text:
from gym import spaces
from pybullet_envs.env_bases import MJCFBaseBulletEnv
from pybullet_envs.gym_pendulum_envs import InvertedPendulumBulletEnv, InvertedPendulumSwingupBulletEnv
from pybullet_envs.scene_abstract import SingleRobotEmptyScene

from rlberry.envs.bullet3.pybullet_envs.robot_pendula import Pendulum, PendulumSwingup
import numpy as np

class PendulumBulletEnv(InvertedPendulumBulletEnv):
    """Simple pendulum"""

    def __init__(self):
        self.robot = Pendulum()
        MJCFBaseBulletEnv.__init__(self, self.robot)
        self.stateId = -1

    def create_single_player_scene(self, bullet_client):
        return SingleRobotEmptyScene(bullet_client, gravity=9.81, timestep=0.02, frame_skip=1)

    def step(self, a):
        self.robot.apply_action(a)
        self.scene.global_step()
        state = self.robot.calc_state()  # sets self.pos_x self.pos_y
        vel_penalty = 0
        if self.robot.swingup:
            reward = np.cos(self.robot.theta)
            done = False
        else:
            reward = 1.0
            done = np.abs(self.robot.theta) > .2
        self.rewards = [float(reward)]
        self.HUD(state, a, done)
        return state, sum(self.rewards), done, {}

class PendulumSwingupBulletEnv(PendulumBulletEnv):
    def __init__(self):
        self.robot = PendulumSwingup()
        MJCFBaseBulletEnv.__init__(self, self.robot)
        self.stateId = -1

class DiscretePendulumBulletEnv(PendulumBulletEnv):
    """pybullet's InvertedPendulum with discrete actions"""

    def __init__(self):
        super().__init__()
        self.continuous_action_space = self.action_space
        self.action_space = spaces.Discrete(3)

    def step(self, a):
        if a == 0:
            return super().step(self.continuous_action_space.low)
        elif a == 1:
            return super().step(self.continuous_action_space.high)
        elif a == 2:
            return super().step(np.zeros(self.continuous_action_space.shape))
        else:
            raise IndexError

class DiscretePendulumSwingupBulletEnv(PendulumSwingupBulletEnv):
    """pybullet's InvertedPendulumSwingup with discrete actions"""

    def __init__(self):
        super().__init__()
        self.continuous_action_space = self.action_space
        self.action_space = spaces.Discrete(3)

    def step(self, a):
        if a == 0:
            return super().step(self.continuous_action_space.low)
        elif a == 1:
            return super().step(self.continuous_action_space.high)
        elif a == 2:
            return super().step(np.zeros(self.continuous_action_space.shape))
        else:
            raise IndexError

ext: py

max_stars_repo_name: rlberry-py/rlberry

max_issues_repo_name: rlberry-py/rlberry

max_forks_repo_name: rlberry-py/rlberry

reason: import numpy


text:
import dynamic_graph.sot_talos_balance.talos.ft_calibration_conf as conf
import numpy as np
from dynamic_graph.sot_talos_balance.ft_calibration import FtCalibration
from numpy.testing import assert_almost_equal

robot_name = 'robot'
ftc = FtCalibration('ftc')
ftc.init(robot_name)
rfw = conf.rfw
lfw = conf.lfw
ftc.setLeftFootWeight(lfw)
ftc.setRightFootWeight(rfw)
print('Weights of both feet set to {0}'.format(rfw))

print('The robot should be in the air.')

ftc.right_foot_force_in.value = np.array([1, 1, 8, 1, 1, 1])
ftc.left_foot_force_in.value = np.array([1, 1, 6, 1, 1, 1])

print("Let's calibrate the ft sensors...")
ftc.calibrateFeetSensor()
for i in range(2, 1003):
    ftc.right_foot_force_in.value = np.random.randn(6) * 0.0000001 + [1, 1, 8, 1, 1, 1]
    ftc.left_foot_force_in.value = np.random.randn(6) * 0.0000001 + [1, 1, 6, 1, 1, 1]
    ftc.right_foot_force_out.recompute(i)
    ftc.left_foot_force_out.recompute(i)

assert_almost_equal(ftc.right_foot_force_out.value, np.array((0, 0, -rfw, 0, 0, 0)), 5)
assert_almost_equal(ftc.left_foot_force_out.value, np.array((0, 0, -lfw, 0, 0, 0)), 5)

print("Ft sensors calibrated!")
print("Value outputed after calibration:")
print(ftc.right_foot_force_out.value)
print(ftc.left_foot_force_out.value)

reason: import numpy,from numpy


text:
import os
import random
from math import pi
from typing import Any, List, Callable, Tuple, Iterator, Reversible, Union, Optional

import numpy as np
import pandas as pd
import torch
from torch import nn

def atanh(x: torch.Tensor) -> torch.Tensor:
    return (torch.log(1 + x) - torch.log(1 - x)) / 2

def clamp_min_parameter(parameter: nn.Parameter, min_value: float) -> None:
    parameter.data = parameter.data.clamp_min(min_value)

def clamp_min_tensor(tensor: torch.Tensor, lo: torch.Tensor) -> None:
    tensor[tensor < lo] = lo[tensor < lo]

def clamp_max_tensor(tensor: torch.Tensor, hi: torch.Tensor) -> None:
    tensor[tensor > hi] = hi[tensor > hi]

def covariance(x: torch.Tensor) -> torch.Tensor:
    x = x.numpy()
    cov = np.cov(x, rowvar=False, bias=True)
    return torch.from_numpy(cov)

def dot_products(tensor: torch.Tensor) -> torch.Tensor:
    """
    Computes the dot product between each pair of rows in the given Tensor.
    :param tensor: the Tensor to compute the dot products between the rows of. Size [N, D].
    :return: the computed dot products, where result[i, j] = tensor[i, :] * tensor[j, :]. Size [N, N].
    """
    products = tensor[:, None, :] * tensor[None, :, :]   # [N, N, D]
    return torch.sum(products, dim=2)   # [N, N]

def euclidean_distances(tensor: torch.Tensor) -> torch.Tensor:
    differences = tensor[:, None, :] - tensor[None, :, :]   # [N, N, D]
    return torch.sqrt(torch.sum(differences ** 2, dim=2))   # [N, N]

def model_device(model: nn.Module) -> torch.device:
    return next(model.parameters()).device

def radians_to_degrees(tensor: torch.Tensor) -> torch.Tensor:
    return tensor * 180 / pi

def random_bool(true_prob: float) -> torch.Tensor:
    return random.random() < true_prob

def random_bool_tensor(size: Tuple[int, ...], true_prob: float, device: Optional[torch.device] = None) -> torch.Tensor:
    return torch.rand(size, device=device) < true_prob

def random_bool_like(tensor: torch.Tensor, true_prob: float) -> torch.Tensor:
    return torch.rand_like(tensor) < true_prob

def random_choice_tensor(
    size: Tuple[int, ...], choice0: Any, choice1: Any, prob0: float, device: Optional[torch.device] = None
) -> torch.Tensor:
    mask = random_bool_tensor(size, prob0, device=device)
    choice0 = torch.tensor(choice0, device=device)
    choice1 = torch.tensor(choice1, device=device)
    return torch.where(mask, choice0, choice1)

def random_choice_like(
    tensor: torch.Tensor, choice0: Union[int, float], choice1: Union[int, float], prob0: float
) -> torch.Tensor:
    mask = random_bool_like(tensor, prob0)
    choice0 = torch.tensor(choice0, device=tensor.device)
    choice1 = torch.tensor(choice1, device=tensor.device)
    return torch.where(mask, choice0, choice1)

def random_float(lo: float, hi: float) -> torch.Tensor:
    return lo + (hi - lo) * random.random()

def random_float_like(tensor: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
    return lo + (hi - lo) * torch.rand_like(tensor)

def max_tensor(value: float, tensor: torch.Tensor) -> torch.Tensor:
    return torch.max(torch.tensor(value, dtype=torch.float, device=tensor.device), tensor)

def min_tensor(value: float, tensor: torch.Tensor) -> torch.Tensor:
    return torch.min(torch.tensor(value, dtype=torch.float, device=tensor.device), tensor)

def normalize_tensor(tensor: torch.Tensor) -> torch.Tensor:
    min_value = tensor.min()
    max_value = tensor.max()

    if min_value == max_value:
        return torch.full_like(tensor, 0)   # prevent division by zero

    return (tensor - min_value) / (max_value - min_value)

def read_lines(filepath: str, constructor: Callable[[str], Any]) -> List[Any]:
    with open(filepath, "r") as file:
        lines = file.readlines()

    return [constructor(line[:-1]) for line in lines]

def remove_diagonal(tensor: torch.Tensor) -> torch.Tensor:
    mask = ~torch.eye(tensor.size(0), dtype=torch.bool)
    return tensor.masked_select(mask)

def remove_outliers(df: pd.DataFrame, data_col: str, group_col: str, min_q: float = 0, max_q: float = 100) \
        -> pd.DataFrame:
    """
    TODO
    :param df:
    :param data_col:
    :param group_col:
    :param min_q:
    :param max_q:
    :return:
    """
    builder = []
    for group, subset in df.groupby(group_col):
        min_p = np.percentile(subset[data_col], min_q)
        max_p = np.percentile(subset[data_col], max_q)

        mask = (min_p <= subset[data_col]) & (subset[data_col] <= max_p)
        builder.append(subset[mask])
    return pd.concat(builder, axis=0)

def remove_row_items(tensor: torch.Tensor, column_indices: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = tensor.size()

    mask = torch.full((n_rows, n_cols), True, dtype=torch.bool, device=tensor.device)
    row_indices = torch.arange(n_rows)
    mask[row_indices, column_indices] = False

    result = tensor.masked_select(mask)
    return result.reshape(n_rows, n_cols - 1)

def reversed_zip(*lists: Union[Reversible[Any], nn.ModuleList]) -> Iterator[Tuple[Any]]:
    return zip(reversed(a_list) for a_list in lists)

def select_row_items(tensor: torch.Tensor, column_indices: torch.Tensor) -> torch.Tensor:
    row_indices = torch.arange(tensor.size(0))
    return tensor[row_indices, column_indices]

def set_deterministic_seed() -> None:
    random.seed(1731)
    np.random.seed(1731)
    torch.manual_seed(1731)
    # noinspection PyUnresolvedReferences
    torch.cuda.manual_seed(1731)

    # noinspection PyUnresolvedReferences
    torch.backends.cudnn.deterministic = True
    # noinspection PyUnresolvedReferences
    torch.backends.cudnn.benchmark = False

def spatial_norm(x: torch.Tensor) -> torch.Tensor:
    return x.flatten(start_dim=1).norm(p=2, dim=1)

def write_lines(filepath: str, lines: List[Any]) -> None:
    lines = [str(line) + "\n" for line in lines]

    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as file:
        file.writelines(lines)

max_stars_repo_name: googleinterns/out-of-distribution

max_issues_repo_name: googleinterns/out-of-distribution

max_forks_repo_name: googleinterns/out-of-distribution


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Aug 26 14:05:26 2018

@author: Kazuki
"""

import gc, os
from tqdm import tqdm
import pandas as pd
import numpy as np
import sys
sys.path.append(f'/home/{os.environ.get("USER")}/PythonLibrary')
import lgbextension as ex
import lightgbm as lgb
from multiprocessing import cpu_count
from sklearn.model_selection import GroupKFold

from glob import glob
#import count
import utils, utils_best
#utils.start(__file__)
#==============================================================================

SEED = 71

NFOLD = 6

new_features = ['f022']

COMMENT = new_features

param = {
         'objective': 'binary',
         'metric': 'auc',
         'learning_rate': 0.01,

         'max_depth': 6,
         'num_leaves': 63,
         'max_bin': 255,

         'min_child_weight': 10,
         'min_data_in_leaf': 150,
         'reg_lambda': 0.5,  # L2 regularization term on weights.
         'reg_alpha': 0.5,  # L1 regularization term on weights.

         'colsample_bytree': 0.9,
         'subsample': 0.9,
#         'nthread': 32,
         'nthread': cpu_count(),
         'bagging_freq': 1,
         'verbose':-1,
         'seed': SEED
         }

loader = utils_best.Loader('LB804')

# =============================================================================
# load
# =============================================================================
X = loader.train()
y = utils.read_pickles('../data/label').TARGET

col = []
files = []
for new_feature in new_features:
    col += [c for c in X.columns if new_feature in c]
    files += glob(f'../feature/train_{new_feature}*')

print('files:', len(files))

X.drop(col, axis=1, inplace=True)
X_ = pd.concat([pd.read_feather(f) for f in tqdm(files, mininterval=60)
                ], axis=1)
X = pd.concat([X, X_], axis=1)
del X_

if X.columns.duplicated().sum()>0:
    raise Exception(f'duplicated!: { X.columns[X.columns.duplicated()] }')
print('no dup :) ')
print(f'X.shape {X.shape}')

gc.collect()

CAT = list( set(X.columns) & set(loader.category()) )
CAT += [new_features[0]+'_user_id']

tmp = pd.read_csv('../data/user_id_v3.csv.gz').set_index('SK_ID_CURR') # TODO: change
sub_train = pd.read_csv('../input/application_train.csv.zip', usecols=['SK_ID_CURR']).set_index('SK_ID_CURR')
sub_train['user_id'] = tmp.user_id
sub_train['g'] = sub_train.user_id % NFOLD

sub_train['y'] = y.values
#sub_train['cnt'] = sub_train.index.value_counts()
#sub_train['w'] = 1 / sub_train.cnt.values

group_kfold = GroupKFold(n_splits=NFOLD)
folds = group_kfold.split(X, sub_train['y'], sub_train['g'])

# =============================================================================
# cv
# =============================================================================
dtrain = lgb.Dataset(X, y, categorical_feature=CAT )
gc.collect()

ret, models = lgb.cv(param, dtrain, 9999, folds=folds,
                     early_stopping_rounds=100, verbose_eval=50,
                     seed=SEED)

y_pred = ex.eval_oob(X, y, models, SEED)

result = f"CV auc-mean({COMMENT}): {ret['auc-mean'][-1]} + {ret['auc-stdv'][-1]}"
print(result)
utils.send_line(result)

imp = ex.getImp(models)

imp['split'] /= imp['split'].max()
imp['gain'] /= imp['gain'].max()
imp['total'] = imp['split'] + imp['gain']

imp.sort_values('total', ascending=False, inplace=True)
imp.reset_index(drop=True, inplace=True)

imp.to_csv('LOG/imp_f022.csv', index=False)
# =============================================================================
#
# =============================================================================

X = loader.train()
y = utils.read_pickles('../data/label').TARGET

files = []
for new_feature in imp[imp.split!=0][imp.feature.str.startswith(new_features[0])].feature:
    files += glob(f'../feature/train_{new_feature}*')

X_ = pd.concat([pd.read_feather(f) for f in tqdm(files, mininterval=60)
                ], axis=1)

for i in range(30,100,10):
    X_new = pd.concat([X, X_.iloc[:,:i] ], axis=1)
    CAT = list( set(X_new.columns) & set(loader.category()+[new_features[0]+'_user_id']) )

    dtrain = lgb.Dataset(X_new, y, categorical_feature=CAT )
    gc.collect()

    folds = group_kfold.split(X, sub_train['y'], sub_train['g'])
    ret, models = lgb.cv(param, dtrain, 9999, folds=folds,
                         early_stopping_rounds=100, verbose_eval=50,
                         seed=SEED)

    y_pred = ex.eval_oob(X, y, models, SEED)

    result = f"CV auc-mean({COMMENT, i}): {ret['auc-mean'][-1]} + {ret['auc-stdv'][-1]}"
    print(result)
    utils.send_line(result)


text:
""" The following examples illustrate the SAR models with different possible set of configurations.
Please refer to
* html docs -> Instrument Models -> Synthetic Aperture Radar Model
* html docs -> API Reference -> Synthetic Aperture Radar Model

"""
import json
import numpy
import sys, os
from flask import Flask, request

from instrupy.radiometer_model import RadiometerModel

Re = 6378.137e3 # [m]
c = 299792458 # m/s

def evaluate_rad(daz_m, delv_m, alt):
    rad_dict = {"@type": "Radiometer",
                "name": "Test radiometer",
                "orientation": {
                   "referenceFrame": "NADIR_POINTING",
                   "convention": "REF_FRAME_ALIGNED"
                },
                "antenna": {"shape": "RECTANGULAR", "height": str(daz_m), "width": str(delv_m), "radiationEfficiency": 0.6,
                           "apertureEfficiency": 0.6, "apertureExcitationProfile": "UNIFORM", "phyTemp": 50},
                "operatingFrequency": 1275.7e6,
                "system": {
                    "@type": "TOTAL_POWER",
                    "integrationTime": 0.1,
                    "bandwidth": 40e6,
                    "predetectionGain": 83,
                    "predetectionInpNoiseTemp": 200,
                    "predetectionGainVariation": 2,
                    "integratorVoltageGain": 1
                },
                "scan":{
                    "@type": "FIXED"
                },
                "targetBrightnessTemp": 290
                }
    test_rad = RadiometerModel.from_dict(rad_dict)
    h = float(alt)*1e3
    orb_speed = numpy.sqrt(3.986004418e14/(Re + h)) # [m/s] orbital speed
    gnd_spd = orb_speed*(Re/(Re+h))
    obsv_metrics = test_rad.calc_data_metrics(alt_km = h*1e-3, gnd_spd = orb_speed*(Re/(Re+h)))
    return obsv_metrics

app = Flask(__name__)

@app.route('/', methods=['GET','POST'])
def index():
	if request.method == 'POST':
		print(request.values)
		h = request.form.get("height")
		w = request.form.get("width")
		alt = request.form.get("altitude")
		metrics = evaluate_rad(h,w,alt)
		print(metrics)
		return metrics
	else:
		return "hello world"

if __name__ == '__main__':
	app.run(port='5000')

max_stars_repo_name: bgorr/instrupy

max_issues_repo_name: bgorr/instrupy

max_forks_repo_name: bgorr/instrupy


text:
import numpy as np

from mrubis_controller.marl.agent import Agent
from mrubis_controller.marl.mrubis_data_helper import build_replay_buffer
from mrubis_controller.marl.mrubis_data_helper import has_shop_remaining_issues

class RobustnessComponent:
    def __init__(self, number_of_agents, config=None):
        if config is None:
            self.config = {
                'actor_lower_threshold': 0,  # value against which the specified statistic is compared
                'actor_upper_threshold': 0.75,  # value against which the specified statistic is compared
                'critic_lower_threshold': 5,  # value against which the specified statistic is compared
                'critic_upper_threshold': 15,  # value against which the specified statistic is compared
                'evaluation_periods': 5,  # number of periods over which data is compared to the specified threshold
                'datapoints_to_alarm': 3,
                # number of data points within the Evaluation Periods that must be breaching to cause the alarm
                'retrain_episodes': 100,
            }

        self.agents_status = {index: 'INSUFFICIENT_DATA' for index in range(number_of_agents)}
        # only copy agents that have been OK
        self.previously_stable_agents = {index: False for index in range(number_of_agents)}
        # OK –                  The metric or expression is within the defined threshold.
        # ALARM –               The metric or expression is outside of the defined threshold.
        # INSUFFICIENT_DATA –   The alarm has just started, the metric is not available,
        #                       or not enough data is available for the metric to determine the alarm state.
        # CALIBRATION           The agent is currently trying to recalibrate due to a prior ALARM status
        # RETRAIN               The agent is trying to improve
        # RETIRED               The agent has no shops left and has no tasks to do and is waiting for challenges

        self.calibration_distribution = {}  # stores the ranked shop distribution during calibration mode
        self.retrain_count = {}  # counter for retraining episodes
        self.metrics = []  # store metrics of agent learning
        self.replay_buffer = []  # store history

    def validate_execution(self, agents, agent_index, observations_):
        """
        recheck if calibration was correct or wrong
        if calibration was correct: move shops to other agents by updating shop_distribution
        and learn action for new agent. If calibration was wrong don't learn this and skip
        update calibration_distribution and test next agent
        """

        if self._is_agent_in_charge_for_other_calibrating_agent(agent_index):
            # check if other agent was doing a better job
            failing_agent_shops = self._get_shops_of_agent_in_charge_for_other_calibrating_agent(agent_index)
            for n in range(len(failing_agent_shops['failing_agent_index'])):
                challenged_shops = failing_agent_shops['shops'][n].difference(agents[agent_index].shops)
                for shop in challenged_shops:
                    if has_shop_remaining_issues(observations_, shop):
                        # remove current challenge and wait for next agent to solve this challenge
                        self.calibration_distribution[failing_agent_shops['failing_agent_index'][n]].popitem()

                        # check if all agents were not able to help the agent
                        if len(self.calibration_distribution[
                                   failing_agent_shops['failing_agent_index'][n]].values()) == 0:
                            if self.previously_stable_agents[failing_agent_shops['failing_agent_index'][n]]:
                                self._copy_agent(agents, failing_agent_shops['failing_agent_index'][n])
                            self.agents_status[failing_agent_shops['failing_agent_index'][n]] = 'RETRAIN'
                            self.retrain_count[failing_agent_shops['failing_agent_index'][n]] = 0
                    else:
                        self._move_shops(agents, failing_agent_shops['failing_agent_index'][n], agent_index,
                                         challenged_shops)

    def get_execution_plan(self, agent_index):
        """ this is the EXECUTE stage for the challenge of the agents
            If the agent is next in line to test his strength for helping a shop, this method returns the shops of that
            other agent and its own shops. If there is no challenge available, None is returned
        """
        if self.agents_status[agent_index] == 'CALIBRATION':
            # don't use an agent that is currently in calibration mode
            return
        elif self._is_agent_in_charge_for_other_calibrating_agent(agent_index):
            # get next challenge
            failing_agent_shops = self._get_shops_of_agent_in_charge_for_other_calibrating_agent(agent_index)
            return set.union(*failing_agent_shops['shops'])
        else:
            return None

    def skip_agent(self, index):
        """ skip agents that are in CALIBRATION or ALARM and skip retired agents only if there are not challenged """
        return self.agents_status[index] in ['CALIBRATION', 'ALARM'] or (
                self.agents_status[index] == 'RETIRED' and self._is_agent_in_charge_for_other_calibrating_agent(index))

    def monitor(self, metrics, history):
        """ this is the MONITOR stage that saves metrics and history """
        self.replay_buffer.append(history)
        self.metrics.append(metrics)

    def analyze(self, agents):
        """ this is the ANALYSIS state that checks whether metrics have breached a threshold and trigger an alarm
        """
        eval_periods = self.config['evaluation_periods']
        for index, agent in enumerate(agents):
            if self.agents_status[index] in ['ALARM', 'CALIBRATION', 'RETIRED']:
                continue
            if self.agents_status[index] in ['RETRAIN']:
                # after 'retrain_episodes' steps, the agent should move to OK or ALARM to enable validation again
                self.retrain_count[index] += 1
                if self.retrain_count[index] < self.config['retrain_episodes']:
                    continue
                else:
                    # reset retrain count
                    # status will be set to ALARM or OK following the code below
                    self.retrain_count[index] = 0
            agent_metric = [agent_metric for metric in self.metrics for agent_metric in metric[index]][-eval_periods:]
            if len(agent_metric) < eval_periods:
                self.agents_status[index] = 'INSUFFICIENT_DATA'
            else:
                # check if any NN is breaching the defined threshold
                # at least datapoints_to_alarm steps must breach the threshold to set an ALARM
                # otherwise the status is OK
                # TODO: why does the critic loss change by changing the data generator?
                for network in ['actor']:
                    breached_metrics = [item[network] for item in agent_metric if
                                        self.config[f"{network}_lower_threshold"] > abs(item[network]) or
                                        abs(item[network]) > self.config[f"{network}_upper_threshold"]]
                    if len(breached_metrics) >= self.config['datapoints_to_alarm']:
                        self.agents_status[index] = 'ALARM'
                        break
                    else:
                        self.agents_status[index] = 'OK'
                        # agent that was OK at some point can be copied before retrained
                        self.previously_stable_agents[index] = True

    def plan(self, agents):
        """ before selecting an action all agents have to be checked if they are in ALARM and init the calibration """
        for index, agent in enumerate(agents):
            if self.agents_status[index] == 'ALARM':
                self._init_calibration_distribution(index, agent, agents)

    def _init_calibration_distribution(self, failing_agent_index, failing_agent, agents):
        """ this is the PLANNING stage
            as soon as an agent is in alarm all available agents (in state OK) that do not fail a pre-check
            (don't make the same last mistake), are ranked by a challenge (last observation).
            The challenge is used to rank the agents by the lowest critic loss which should help to find the best agent
            faster. The execution plan is stored inside the calibration_distribution if any agent is left that did pass
            the pre-checks. Otherwise, the agent has to be retrained as no agent is available to help that agent
        """
        relevant_history = build_replay_buffer(agents, failing_agent_index, self.replay_buffer)
        agents_predictions = {}
        for index, agent in enumerate(agents):
            if agent != failing_agent:
                # ignore agents in alarm or calibration state
                if self.agents_status[index] not in ['OK']:
                    continue

                # ignore agents similar to failing agent
                if not self._is_agent_check_plausible(agent, relevant_history):
                    continue

                # calc predictions for last step by all agents
                agent_predictions = {}
                for shop_name, components in relevant_history['observations'].items():
                    probabilities = agent.get_probabilities_for_observations(components)
                    max_value = np.amax(probabilities)
                    agent_predictions[shop_name] = max_value
                agents_predictions[index] = agent_predictions

        # rank possible agents bei certainty of action probability
        ranked_agents = self._rank_agents_predictions(relevant_history['observations'], agents_predictions)

        # init calibration_distribution
        ranked_test_distributions = {
            agent_index: {
                'shops': agents[agent_index].shops.union(agents[failing_agent_index].shops)
            }
            for agent_index in ranked_agents.keys()
        }

        if not ranked_test_distributions:
            # no agent is available to start a calibration so retrain that agent
            self.agents_status[failing_agent_index] = 'RETRAIN'
            self.retrain_count[failing_agent_index] = 0
        else:
            self.calibration_distribution[failing_agent_index] = ranked_test_distributions
            self.agents_status[failing_agent_index] = 'CALIBRATION'

    def _is_agent_in_charge_for_other_calibrating_agent(self, index):
        """
        checks if index_agent is next in line and has currently a failing component so the agent can have a challenge
        """
        return any(len(distribution.keys()) > 0 and index == list(distribution.keys())[0] for distribution in
                   self.calibration_distribution.values())

    def _get_shops_of_agent_in_charge_for_other_calibrating_agent(self, agent_index):
        """
        get all distributions if agent_index is next in line and the failing agent has shops with failures
        so we can have challenge for that agent, otherwise take the next agent
        """
        # TODO: return only shops that have remaining issues
        calibration_distribution = {'shops': [], 'failing_agent_index': []}
        for failing_agent_index, distribution in self.calibration_distribution.items():
            if len(distribution.keys()) > 0 and agent_index == list(distribution.keys())[0]:
                calibration_distribution['shops'].append(list(distribution.values())[0]['shops'])
                calibration_distribution['failing_agent_index'].append(failing_agent_index)
        return calibration_distribution or None

    def _rank_agents_predictions(self, observations, agents_predictions):
        """ ranks the agent for a given observation by checking how certain they are to have the correct action """
        ranking_per_shop = {}
        # each agent gets points per ranking per shop
        # the less points the better
        score_per_agent = {agent_index: 0 for agent_index, _ in agents_predictions.items()}  # init scores
        for shop_name, components in observations.items():
            # prepare data for sorting
            max_prediction_per_agent = [{'agent_index': agent_index, 'amax_prediction': predictions[shop_name]} for
                                        agent_index, predictions in agents_predictions.items()]
            # sort predictions
            ranking_per_shop[shop_name] = sorted(max_prediction_per_agent, key=lambda d: d['amax_prediction'],
                                                 reverse=True)
        for shop, ranked_predictions in ranking_per_shop.items():
            for index, prediction in enumerate(ranked_predictions):
                score_per_agent[prediction['agent_index']] += index
        return dict(sorted(score_per_agent.items(), key=lambda item: item[1]))

    def _is_agent_check_plausible(self, agent, relevant_history):
        """ filter agents which would have chosen the exact same action that was wrong """
        actions = agent.choose_action(relevant_history['observations'])
        actions = {action['shop']: action['component'] for action in actions}
        if any(actions[historic_action['shop']] == historic_action['component'] for historic_action in
               relevant_history['actions'].values()):
            return False
        return True

    def _move_shops(self, agents, source_index, destination_index, shops):
        """ moves shops from source agent to destination
            if no shop is left in source agent, that agent retires
        """
        agents[source_index].remove_shops(shops)
        agents[destination_index].add_shops(shops)

        # check if agent has any shops left
        if len(agents[source_index].shops) == 0:
            self.agents_status[source_index] = 'RETIRED'

    def _copy_agent(self, agents, agent_index):
        """ copy agent and retire copy of that agent
            retired agent is a snapshot of a shop configuration that could eventually be relevant in the future
            e.g. a rollback of a shop
        """
        agents[agent_index].save('copy')
        load_model = {'start_time': agents[agent_index].start_time, 'episode': 'copy', 'index': agent_index}
        copy_index = len(agents)
        copy_agent = Agent(
            copy_index,
            {},
            agents[agent_index].action_space_inverted,
            load_model,
            agents[agent_index].ridge_regression_train_data_path,
        )
        copy_agent.start_time = agents[agent_index].start_time

        # append new agent
        agents.append(copy_agent)

        # update status
        self.agents_status[copy_index] = 'RETIRED'

        # setting this to false prevents to copy agent each retrain loop
        self.previously_stable_agents[agent_index] = False
        self.previously_stable_agents[copy_index] = False


text:
# -*- coding: utf-8 -*-
#
#  This file is part of Sequana software
#
#  Copyright (c) 2016 - Sequana Development Team
#
#  File author(s):
#      Thomas Cokelaer <thomas.cokelaer@pasteur.fr>
#
#  Distributed under the terms of the 3-clause BSD license.
#  The full license is in the LICENSE file, distributed with this software.
#
#  website: https://github.com/sequana/sequana
#  documentation: http://sequana.readthedocs.io
#
##############################################################################

"""Heatmap and dendograms"""
import warnings
import scipy.cluster.hierarchy as hierarchy
import scipy.spatial.distance as distance
import easydev

__all__ = ['Linkage']

class Linkage(object):
    """Linkage used in other tools such as Heatmap"""

    methods = ["single", "complete", "average", "weighted", "centroid",
        "median", "ward"]
    metrics = ['braycurtis', 'canberra', 'chebyshev', 'cityblock',
        'correlation', 'cosine', 'dice', 'euclidean', 'hamming',
        'jaccard', 'jensenshannon', 'kulsinski', 'mahalanobis', 'matching',
        'minkowski', 'rogerstanimoto', 'russellrao', 'seuclidean',
        'sokalmichener', 'sokalsneath', 'sqeuclidean', 'yule']

    def __init__(self):
        """.. rubric:: constructor

        :param data: a dataframe or possibly a numpy matrix.

        """
        pass

    def check_metric(self, value):
        easydev.check_param_in_list(value, self.metrics)

    def check_method(self, value):
        # None is possible
        # in R, in addition to single, complete, average, centroid,
        # median and ward
        # there are  ward.D, wardD2 and mcquitty
        # default is complete
        easydev.check_param_in_list(str(value), self.methods)

    def linkage(self, df, method, metric):
        self.check_metric(metric)
        self.check_method(method)
        d = distance.pdist(df)
        D = distance.squareform(d)

        # hierarchy.ClusterWarning
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore')
            Y = hierarchy.linkage(D, method=method, metric=metric)
            return Y

max_stars_repo_name: khourhin/sequana

max_issues_repo_name: khourhin/sequana

max_forks_repo_name: khourhin/sequana

reason: import scipy


text:
# encoding=utf-8

from sys import version_info
import os

IS_PY3 = version_info.major >= 3
from ._dicomsdl import *
# If you have `ImportError: DLL load failed in Microsoft Windows` on fresh installed Windows system, install Microsoft Visual C++ Redistributable. https://docs.microsoft.com/en-us/cpp/windows/latest-supported-vc-redist?view=msvc-170
from . import util

def __dataset__getitem__(self, tagstr):
  de = self.getDataElement(tagstr)
  vr = de.vr()
  if vr == VR.SQ:
    return de.toSequence()
  elif vr == VR.PIXSEQ:
    return de.toPixelSequence()
  elif vr == VR.NONE:
    return None
  else:
    return de.value()
DataSet.__getitem__ = __dataset__getitem__

def __dataset__setitem__(self, tag_vr, value):
  if isinstance(tag_vr, tuple):
    tag, vr = tag_vr
    if not isinstance(vr, VR.type):
      vr = VR.from_string(vr)
  else:
    tag = tag_vr
    vr = VR.NONE
  if vr in (VR.PIXSEQ, VR.SQ):
    raise AttributeError("Cannot set value of DataSet or Pixel Sequence")
  self.addDataElement(tag, vr).setValue(value)
DataSet.__setitem__ = __dataset__setitem__

def __dataset__getattr__(self, key):
    tag = TAG.from_keyword(key)
    if tag == 0xffffffff:
        raise AttributeError("'DataSet' object has no attribute '%s'"%(key))
    else:
        return __dataset__getitem__(self, key)
DataSet.__getattr__ = __dataset__getattr__

def __dataset__setattr__(self, key, value):
    tag = TAG.from_keyword(key)
    if tag == 0xffffffff:
        raise AttributeError("'DataSet' object has no attribute '%s'"%(key))
    else:
        return __dataset__setitem__(self, key, value)
DataSet.__setattr__ = __dataset__setattr__

def __dataset____dir__(self):
    """Return attributes of this DataSet with keywords of all DataElements.

    This enables word completions when user use 'tab' on an DataSet object
    under interactive console environment.
    """
    return self.____dir__() +\
        [TAG.keyword(de.tag())
         for de in self
         if de.tag() & 0x10000 == 0 and de.tag() & 0xffff != 0]
DataSet.____dir__ = DataSet.__dir__ if IS_PY3 else lambda _: dir(DataSet)
DataSet.__dir__ = __dataset____dir__

def __dataset__pixelData__(self, index=0, storedvalue=False):
  """Returns pixel values in this DataSet.

  Args:
    index (int): pixelData return index'th image if dataset holds multiframe
                 data.
    storedvalue (bool): True for get stored values; pixel values before LUT
                        transformation using RescaleSlope and RescaleIntercept.

  Returns:
    Numpy array containing pixel values of `index`'th image if dataset holds
    multiframe data. If `storedvalue` is False, RescaleSlope and
    RescaleIntercept are applied to pixel values.
    ( pixel values = stored values * RescsaleSlope + RescaleIntercept )

  Example:
    >>> dset=dicom.open_file('some_CT_image_file')
    >>> dset.RescaleIntercept
    -1024.0
    >>> dset.pixelData(storedvalue=True).min()  # before LUT transformation
    -2000
    >>> dset.pixelData().min()  # after LUT transformation
    -3024.0

  """
  import numpy as np
  # https://stackoverflow.com/questions/44659924/returning-numpy-arrays-via-pybind11
  info = self.getPixelDataInfo()

  dtype = info['dtype']
  if info['PlanarConfiguration'] == 'RRRGGGBBB':
    shape = [3, info['Rows'], info['Cols']]
  elif info['PlanarConfiguration'] == 'RGBRGBRGB':
    shape = [info['Rows'], info['Cols'], 3]
  else:
    shape = [info['Rows'], info['Cols']]

  outarr = np.empty(shape, dtype=dtype)
  self.copyFrameData(index, outarr)

  check = lambda x: x[index] if isinstance(x, list) else x

  if storedvalue is False:
    intercept = check(info['RescaleIntercept'])
    intercept = intercept if intercept is not None else 0.0
    slope = check(info['RescaleSlope'])
    slope = slope if slope is not None else 1.0
    outarr = np.float32(outarr)
    outarr *= slope
    outarr += intercept

  return outarr
DataSet.pixelData = __dataset__pixelData__

def __dataset__to_pil_image(self, index=0):
  """Returns a pillow image.

  Args:
    index (int): pixelData return index'th image if dataset holds multiframe
                 image data.

  Example:
    >>> dset.to_pil_image().show()  # pop up system's default image viewer

  """
  import numpy as np
  from PIL import Image
  # https://stackoverflow.com/questions/44659924/returning-numpy-arrays-via-pybind11
  info = self.getPixelDataInfo()

  dtype = info['dtype']
  if info['PlanarConfiguration'] == 'RRRGGGBBB':
    shape = [3, info['Rows'], info['Cols']]
    assert False # Need Implementation
  elif info['PlanarConfiguration'] == 'RGBRGBRGB':
    shape = [info['Rows'], info['Cols'], 3]

  else:
    shape = [info['Rows'], info['Cols']]

  outarr = np.empty(shape, dtype=dtype)
  self.copyFrameData(index, outarr)

  if len(shape) == 3 and shape[-1] == 3:
    return Image.fromarray(outarr)

  check = lambda x: x[index] if isinstance(x, list) else x

  intercept = check(info['RescaleIntercept'])
  intercept = intercept if intercept is not None else 0.0
  slope = check(info['RescaleSlope'])
  slope = slope if slope is not None else 1.0
  c = check(info['WindowCenter'])
  w = check(info['WindowWidth'])

  if c is not None and w is not None:
    # xmin, xmax after applying rescale intensity
    xmin = c - 0.5 - (w - 1) * 0.5
    xmax = c - 0.5 + (w - 1) * 0.5

    # xmin, xmax before applying rescale intensity
    xmin = (xmin - intercept) / slope
    xmax = (xmax - intercept) / slope
  else:
    xmin = outarr.min()
    xmax = outarr.max()

  data8 = np.empty_like(outarr, dtype=np.uint8)
  util.convert_to_uint8(outarr, data8, xmin, xmax)

  if info['PhotometricInterpretation'] == 'MONOCHROME1':
    data8 = 255 - data8

  return Image.fromarray(data8)
DataSet.to_pil_image = __dataset__to_pil_image
DataSet.toPilImage = __dataset__to_pil_image

max_stars_repo_name: tsangel/dicomsdl

max_issues_repo_name: tsangel/dicomsdl

max_forks_repo_name: tsangel/dicomsdl