text:
from rlkit.data_management.contextual_replay_buffer import (
    ContextualRelabelingReplayBuffer,
)
from rlkit.envs.contextual import ContextualEnv, delete_info

from rlkit.envs.contextual.goal_conditioned import (
    GoalDictDistributionFromMultitaskEnv,
    ContextualRewardFnFromMultitaskEnv,
    AddImageDistribution,
    GoalConditionedDiagnosticsToContextualDiagnostics,
    IndexIntoAchievedGoal,
)
from rlkit.envs.contextual.mask_conditioned import (
    MaskDictDistribution,
    MaskPathCollector,
    ContextualMaskingRewardFn,
)
from rlkit.launchers.sets.mask_inference import get_mask_params
from rlkit.launchers.sets.example_set_gen import gen_example_sets

from rlkit.envs.images import EnvRenderer, InsertImagesEnv
from rlkit.launchers.contextual.util import (
    get_save_video_function,
)
from rlkit.samplers.data_collector.contextual_path_collector import (
    ContextualPathCollector
)

from rlkit.launchers.rl_exp_launcher_util import (
    preprocess_rl_variant,
    get_envs,
    create_exploration_policy,
)

import copy
import torch
from functools import partial
import numpy as np

def rl_context_experiment(variant):
    import rlkit.torch.pytorch_util as ptu
    from rlkit.torch.td3.td3 import TD3 as TD3Trainer
    from rlkit.torch.sac.sac import SACTrainer
    from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
    from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
    from rlkit.torch.sac.policies import TanhGaussianPolicy
    from rlkit.torch.sac.policies import MakeDeterministic

    preprocess_rl_variant(variant)
    max_path_length = variant['max_path_length']
    observation_key = variant.get('observation_key', 'latent_observation')
    desired_goal_key = variant.get('desired_goal_key', 'latent_desired_goal')
    achieved_goal_key = variant.get('achieved_goal_key', 'latent_achieved_goal')

    contextual_mdp = variant.get('contextual_mdp', True)
    print("contextual_mdp:", contextual_mdp)

    mask_variant = variant.get('mask_variant', {})
    mask_conditioned = mask_variant.get('mask_conditioned', False)
    print("mask_conditioned:", mask_conditioned)

    if mask_conditioned:
        assert contextual_mdp

    if 'sac' in variant['algorithm'].lower():
        rl_algo = 'sac'
    elif 'td3' in variant['algorithm'].lower():
        rl_algo = 'td3'
    else:
        raise NotImplementedError
    print("RL algorithm:", rl_algo)

    ### load the example dataset, if running checkpoints ###
    if 'ckpt' in variant:
        import os.path as osp
        example_set_variant = variant.get('example_set_variant', dict())
        example_set_variant['use_cache'] = True
        example_set_variant['cache_path'] = osp.join(variant['ckpt'], 'example_dataset.npy')

    if mask_conditioned:
        env = get_envs(variant)
        mask_format = mask_variant['param_variant']['mask_format']
        assert mask_format in ['vector', 'matrix', 'distribution', 'cond_distribution']
        goal_dim = env.observation_space.spaces[desired_goal_key].low.size
        if mask_format in ['vector']:
            context_dim_for_networks = goal_dim + goal_dim
        elif mask_format in ['matrix', 'distribution', 'cond_distribution']:
            context_dim_for_networks = goal_dim + (goal_dim * goal_dim)
        else:
            raise TypeError

        if 'ckpt' in variant:
            from rlkit.util.io import local_path_from_s3_or_local_path
            import os.path as osp

            filename = local_path_from_s3_or_local_path(osp.join(variant['ckpt'], 'masks.npy'))
            masks = np.load(filename, allow_pickle=True)[()]
        else:
            masks = get_mask_params(
                env=env,
                example_set_variant=variant['example_set_variant'],
                param_variant=mask_variant['param_variant'],
            )

        mask_keys = list(masks.keys())
        context_keys = [desired_goal_key] + mask_keys
    else:
        context_keys = [desired_goal_key]

    def contextual_env_distrib_and_reward(mode='expl'):
        assert mode in ['expl', 'eval']
        env = get_envs(variant)

        if mode == 'expl':
            goal_sampling_mode = variant.get('expl_goal_sampling_mode', None)
        elif mode == 'eval':
            goal_sampling_mode = variant.get('eval_goal_sampling_mode', None)
        if goal_sampling_mode not in [None, 'example_set']:
            env.goal_sampling_mode = goal_sampling_mode

        mask_ids_for_training = mask_variant.get('mask_ids_for_training', None)

        if mask_conditioned:
            context_distrib = MaskDictDistribution(
                env,
                desired_goal_keys=[desired_goal_key],
                mask_format=mask_format,
                masks=masks,
                max_subtasks_to_focus_on=mask_variant.get('max_subtasks_to_focus_on', None),
                prev_subtask_weight=mask_variant.get('prev_subtask_weight', None),
                mask_distr=mask_variant.get('train_mask_distr', None),
                mask_ids=mask_ids_for_training,
            )
            reward_fn = ContextualMaskingRewardFn(
                achieved_goal_from_observation=IndexIntoAchievedGoal(achieved_goal_key),
                desired_goal_key=desired_goal_key,
                achieved_goal_key=achieved_goal_key,
                mask_keys=mask_keys,
                mask_format=mask_format,
                use_g_for_mean=mask_variant['use_g_for_mean'],
                use_squared_reward=mask_variant.get('use_squared_reward', False),
            )
        else:
            if goal_sampling_mode == 'example_set':
                example_dataset = gen_example_sets(get_envs(variant), variant['example_set_variant'])
                assert len(example_dataset['list_of_waypoints']) == 1
                from rlkit.envs.contextual.set_distributions import GoalDictDistributionFromSet
                context_distrib = GoalDictDistributionFromSet(
                    example_dataset['list_of_waypoints'][0],
                    desired_goal_keys=[desired_goal_key],
                )
            else:
                context_distrib = GoalDictDistributionFromMultitaskEnv(
                    env,
                    desired_goal_keys=[desired_goal_key],
                )
            reward_fn = ContextualRewardFnFromMultitaskEnv(
                env=env,
                achieved_goal_from_observation=IndexIntoAchievedGoal(achieved_goal_key),
                desired_goal_key=desired_goal_key,
                achieved_goal_key=achieved_goal_key,
                additional_obs_keys=variant['contextual_replay_buffer_kwargs'].get('observation_keys', None),
            )
        diag_fn = GoalConditionedDiagnosticsToContextualDiagnostics(
            env.goal_conditioned_diagnostics,
            desired_goal_key=desired_goal_key,
            observation_key=observation_key,
        )
        env = ContextualEnv(
            env,
            context_distribution=context_distrib,
            reward_fn=reward_fn,
            observation_key=observation_key,
            contextual_diagnostics_fns=[diag_fn],
            update_env_info_fn=delete_info if not variant.get('keep_env_infos', False) else None,
        )
        return env, context_distrib, reward_fn

    env, context_distrib, reward_fn = contextual_env_distrib_and_reward(mode='expl')
    eval_env, eval_context_distrib, _ = contextual_env_distrib_and_reward(mode='eval')

    if mask_conditioned:
        obs_dim = (
            env.observation_space.spaces[observation_key].low.size
            + context_dim_for_networks
        )
    elif contextual_mdp:
        obs_dim = (
            env.observation_space.spaces[observation_key].low.size
            + env.observation_space.spaces[desired_goal_key].low.size
        )
    else:
        obs_dim = env.observation_space.spaces[observation_key].low.size

    action_dim = env.action_space.low.size

    if 'ckpt' in variant and 'ckpt_epoch' in variant:
        from rlkit.util.io import local_path_from_s3_or_local_path
        import os.path as osp

        ckpt_epoch = variant['ckpt_epoch']
        if ckpt_epoch is not None:
            epoch = variant['ckpt_epoch']
            filename = local_path_from_s3_or_local_path(osp.join(variant['ckpt'], 'itr_%d.pkl' % epoch))
        else:
            filename = local_path_from_s3_or_local_path(osp.join(variant['ckpt'], 'params.pkl'))
        print("Loading ckpt from", filename)
        data = torch.load(filename)
        qf1 = data['trainer/qf1']
        qf2 = data['trainer/qf2']
        target_qf1 = data['trainer/target_qf1']
        target_qf2 = data['trainer/target_qf2']
        policy = data['trainer/policy']
        eval_policy = data['evaluation/policy']
        expl_policy = data['exploration/policy']
    else:
        qf1 = ConcatMlp(
            input_size=obs_dim + action_dim,
            output_size=1,
            **variant['qf_kwargs']
        )
        qf2 = ConcatMlp(
            input_size=obs_dim + action_dim,
            output_size=1,
            **variant['qf_kwargs']
        )
        target_qf1 = ConcatMlp(
            input_size=obs_dim + action_dim,
            output_size=1,
            **variant['qf_kwargs']
        )
        target_qf2 = ConcatMlp(
            input_size=obs_dim + action_dim,
            output_size=1,
            **variant['qf_kwargs']
        )
        if rl_algo == 'td3':
            policy = TanhMlpPolicy(
                input_size=obs_dim,
                output_size=action_dim,
                **variant['policy_kwargs']
            )
            target_policy = TanhMlpPolicy(
                input_size=obs_dim,
                output_size=action_dim,
                **variant['policy_kwargs']
            )
            expl_policy = create_exploration_policy(
                env, policy,
                exploration_version=variant['exploration_type'],
                exploration_noise=variant['exploration_noise'],
            )
            eval_policy = policy
        elif rl_algo == 'sac':
            policy = TanhGaussianPolicy(
                obs_dim=obs_dim,
                action_dim=action_dim,
                **variant['policy_kwargs']
            )
            expl_policy = policy
            eval_policy = MakeDeterministic(policy)

    post_process_mask_fn = partial(
        full_post_process_mask_fn,
        mask_conditioned=mask_conditioned,
        mask_variant=mask_variant,
        context_distrib=context_distrib,
        context_key=desired_goal_key,
        achieved_goal_key=achieved_goal_key,
    )

    def context_from_obs_dict_fn(obs_dict):
        context_dict = {
            desired_goal_key: obs_dict[achieved_goal_key]
        }

        if mask_conditioned:
            sample_masks_for_relabeling = mask_variant.get('sample_masks_for_relabeling', True)
            if sample_masks_for_relabeling:
                batch_size = next(iter(obs_dict.values())).shape[0]
                sampled_contexts = context_distrib.sample(batch_size)
                for mask_key in mask_keys:
                    context_dict[mask_key] = sampled_contexts[mask_key]
            else:
                for mask_key in mask_keys:
                    context_dict[mask_key] = obs_dict[mask_key]

        return context_dict

    def concat_context_to_obs(batch, replay_buffer=None, obs_dict=None, next_obs_dict=None, new_contexts=None):
        obs = batch['observations']
        next_obs = batch['next_observations']
        batch_size = obs.shape[0]
        if mask_conditioned:
            if obs_dict is not None and new_contexts is not None:
                if not mask_variant.get('relabel_masks', True):
                    for k in mask_keys:
                        new_contexts[k] = next_obs_dict[k][:]
                    batch.update(new_contexts)
                if not mask_variant.get('relabel_goals', True):
                    new_contexts[desired_goal_key] = next_obs_dict[desired_goal_key][:]
                    batch.update(new_contexts)

                new_contexts = post_process_mask_fn(obs_dict, new_contexts)
                batch.update(new_contexts)

            if mask_format in ['vector', 'matrix']:
                goal = batch[desired_goal_key]
                mask = batch['mask'].reshape((batch_size, -1))
                batch['observations'] = np.concatenate([obs, goal, mask], axis=1)
                batch['next_observations'] = np.concatenate([next_obs, goal, mask], axis=1)
            elif mask_format == 'distribution':
                goal = batch[desired_goal_key]
                sigma_inv = batch['mask_sigma_inv'].reshape((batch_size, -1))
                batch['observations'] = np.concatenate([obs, goal, sigma_inv], axis=1)
                batch['next_observations'] = np.concatenate([next_obs, goal, sigma_inv], axis=1)
            elif mask_format == 'cond_distribution':
                goal = batch[desired_goal_key]
                mu_w = batch['mask_mu_w']
                mu_g = batch['mask_mu_g']
                mu_A = batch['mask_mu_mat']
                sigma_inv = batch['mask_sigma_inv']
                if mask_variant['use_g_for_mean']:
                    mu_w_given_g = goal
                else:
                    mu_w_given_g = mu_w + np.squeeze(mu_A @ np.expand_dims(goal - mu_g, axis=-1), axis=-1)
                sigma_w_given_g_inv = sigma_inv.reshape((batch_size, -1))
                batch['observations'] = np.concatenate([obs, mu_w_given_g, sigma_w_given_g_inv], axis=1)
                batch['next_observations'] = np.concatenate([next_obs, mu_w_given_g, sigma_w_given_g_inv], axis=1)
            else:
                raise NotImplementedError
        elif contextual_mdp:
            goal = batch[desired_goal_key]
            batch['observations'] = np.concatenate([obs, goal], axis=1)
            batch['next_observations'] = np.concatenate([next_obs, goal], axis=1)
        else:
            batch['observations'] = obs
            batch['next_observations'] = next_obs

        return batch

    if 'observation_keys' not in variant['contextual_replay_buffer_kwargs']:
        variant['contextual_replay_buffer_kwargs']['observation_keys'] = []
    observation_keys = variant['contextual_replay_buffer_kwargs']['observation_keys']
    if observation_key not in observation_keys:
        observation_keys.append(observation_key)
    if achieved_goal_key not in observation_keys:
        observation_keys.append(achieved_goal_key)

    replay_buffer = ContextualRelabelingReplayBuffer(
        env=env,
        context_keys=context_keys,
        context_distribution=context_distrib,
        sample_context_from_obs_dict_fn=context_from_obs_dict_fn,
        reward_fn=reward_fn,
        post_process_batch_fn=concat_context_to_obs,
        **variant['contextual_replay_buffer_kwargs']
    )

    if rl_algo == 'td3':
        trainer = TD3Trainer(
            policy=policy,
            qf1=qf1,
            qf2=qf2,
            target_qf1=target_qf1,
            target_qf2=target_qf2,
            target_policy=target_policy,
            **variant['td3_trainer_kwargs']
        )
    elif rl_algo == 'sac':
        trainer = SACTrainer(
            env=env,
            policy=policy,
            qf1=qf1,
            qf2=qf2,
            target_qf1=target_qf1,
            target_qf2=target_qf2,
            **variant['sac_trainer_kwargs']
        )

    def create_path_collector(
            env,
            policy,
            mode='expl',
            mask_kwargs={},
    ):
        assert mode in ['expl', 'eval']

        save_env_in_snapshot = variant.get('save_env_in_snapshot', True)

        if mask_conditioned:
            if 'rollout_mask_order' in mask_kwargs:
                rollout_mask_order = mask_kwargs['rollout_mask_order']
            else:
                if mode == 'expl':
                    rollout_mask_order = mask_variant.get('rollout_mask_order_for_expl', 'fixed')
                elif mode == 'eval':
                    rollout_mask_order = mask_variant.get('rollout_mask_order_for_eval', 'fixed')
                else:
                    raise TypeError

            if 'mask_distr' in mask_kwargs:
                mask_distr = mask_kwargs['mask_distr']
            else:
                if mode == 'expl':
                    mask_distr = mask_variant['expl_mask_distr']
                elif mode == 'eval':
                    mask_distr = mask_variant['eval_mask_distr']
                else:
                    raise TypeError

            if 'mask_ids' in mask_kwargs:
                mask_ids = mask_kwargs['mask_ids']
            else:
                if mode == 'expl':
                    mask_ids = mask_variant.get('mask_ids_for_expl', None)
                elif mode == 'eval':
                    mask_ids = mask_variant.get('mask_ids_for_eval', None)
                else:
                    raise TypeError

            prev_subtask_weight = mask_variant.get('prev_subtask_weight', None)
            max_subtasks_to_focus_on = mask_variant.get('max_subtasks_to_focus_on', None)
            max_subtasks_per_rollout = mask_variant.get('max_subtasks_per_rollout', None)

            mode = mask_variant.get('context_post_process_mode', None)
            if mode in ['dilute_prev_subtasks_uniform', 'dilute_prev_subtasks_fixed']:
                prev_subtask_weight = 0.5

            return MaskPathCollector(
                env,
                policy,
                observation_key=observation_key,
                context_keys_for_policy=context_keys,
                concat_context_to_obs_fn=concat_context_to_obs,
                save_env_in_snapshot=save_env_in_snapshot,
                mask_sampler=(context_distrib if mode=='expl' else eval_context_distrib),
                mask_distr=mask_distr.copy(),
                mask_ids=mask_ids,
                max_path_length=max_path_length,
                rollout_mask_order=rollout_mask_order,
                prev_subtask_weight=prev_subtask_weight,
                max_subtasks_to_focus_on=max_subtasks_to_focus_on,
                max_subtasks_per_rollout=max_subtasks_per_rollout,
            )
        elif contextual_mdp:
            return ContextualPathCollector(
                env,
                policy,
                observation_key=observation_key,
                context_keys_for_policy=context_keys,
                save_env_in_snapshot=save_env_in_snapshot,
            )
        else:
            return ContextualPathCollector(
                env,
                policy,
                observation_key=observation_key,
                context_keys_for_policy=[],
                save_env_in_snapshot=save_env_in_snapshot,
            )

    expl_path_collector = create_path_collector(env, expl_policy, mode='expl')
    eval_path_collector = create_path_collector(eval_env, eval_policy, mode='eval')

    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        max_path_length=max_path_length,
        **variant['algo_kwargs']
    )

    algorithm.to(ptu.device)

    if variant.get("save_video", True):
        save_period = variant.get('save_video_period', 50)
        dump_video_kwargs = variant.get("dump_video_kwargs", dict())
        dump_video_kwargs['horizon'] = max_path_length

        renderer = EnvRenderer(**variant.get('renderer_kwargs', {}))

        def add_images(env, state_distribution):
            state_env = env.env
            image_goal_distribution = AddImageDistribution(
                env=state_env,
                base_distribution=state_distribution,
                image_goal_key='image_desired_goal',
                renderer=renderer,
            )
            img_env = InsertImagesEnv(state_env, renderers={
                'image_observation' : renderer,
            })
            context_env = ContextualEnv(
                img_env,
                context_distribution=image_goal_distribution,
                reward_fn=reward_fn,
                observation_key=observation_key,
                update_env_info_fn=delete_info,
            )
            return context_env

        img_eval_env = add_images(eval_env, eval_context_distrib)

        if variant.get('log_eval_video', True):
            video_path_collector = create_path_collector(img_eval_env, eval_policy, mode='eval')
            rollout_function = video_path_collector._rollout_fn
            eval_video_func = get_save_video_function(
                rollout_function,
                img_eval_env,
                eval_policy,
                tag="eval",
                imsize=variant['renderer_kwargs']['width'],
                image_format='CHW',
                save_video_period=save_period,
                **dump_video_kwargs
            )
            algorithm.post_train_funcs.append(eval_video_func)

        # additional eval videos for mask conditioned case
        if mask_conditioned:
            default_list = [
                'atomic',
                'atomic_seq',
                'cumul_seq',
                'full',
            ]
            eval_rollouts_for_videos = mask_variant.get('eval_rollouts_for_videos', default_list)
            for key in eval_rollouts_for_videos:
                assert key in default_list

            if 'cumul_seq' in eval_rollouts_for_videos:
                video_path_collector = create_path_collector(
                    img_eval_env,
                    eval_policy,
                    mode='eval',
                    mask_kwargs=dict(
                        mask_distr=dict(
                            cumul_seq=1.0
                        ),
                    ),
                )
                rollout_function = video_path_collector._rollout_fn
                eval_video_func = get_save_video_function(
                    rollout_function,
                    img_eval_env,
                    eval_policy,
                    tag="eval_cumul" if mask_conditioned else "eval",
                    imsize=variant['renderer_kwargs']['width'],
                    image_format='HWC',
                    save_video_period=save_period,
                    **dump_video_kwargs
                )
                algorithm.post_train_funcs.append(eval_video_func)

            if 'full' in eval_rollouts_for_videos:
                video_path_collector = create_path_collector(
                    img_eval_env,
                    eval_policy,
                    mode='eval',
                    mask_kwargs=dict(
                        mask_distr=dict(
                            full=1.0
                        ),
                    ),
                )
                rollout_function = video_path_collector._rollout_fn
                eval_video_func = get_save_video_function(
                    rollout_function,
                    img_eval_env,
                    eval_policy,
                    tag="eval_full",
                    imsize=variant['renderer_kwargs']['width'],
                    image_format='HWC',
                    save_video_period=save_period,
                    **dump_video_kwargs
                )
                algorithm.post_train_funcs.append(eval_video_func)

            if 'atomic_seq' in eval_rollouts_for_videos:
                video_path_collector = create_path_collector(
                    img_eval_env,
                    eval_policy,
                    mode='eval',
                    mask_kwargs=dict(
                        mask_distr=dict(
                            atomic_seq=1.0
                        ),
                    ),
                )
                rollout_function = video_path_collector._rollout_fn
                eval_video_func = get_save_video_function(
                    rollout_function,
                    img_eval_env,
                    eval_policy,
                    tag="eval_atomic",
                    imsize=variant['renderer_kwargs']['width'],
                    image_format='HWC',
                    save_video_period=save_period,
                    **dump_video_kwargs
                )
                algorithm.post_train_funcs.append(eval_video_func)

        if variant.get('log_expl_video', True) and not variant['algo_kwargs'].get('eval_only', False):
            img_expl_env = add_images(env, context_distrib)
            video_path_collector = create_path_collector(img_expl_env, expl_policy, mode='expl')
            rollout_function = video_path_collector._rollout_fn
            expl_video_func = get_save_video_function(
                rollout_function,
                img_expl_env,
                expl_policy,
                tag="expl",
                imsize=variant['renderer_kwargs']['width'],
                image_format='CHW',
                save_video_period=save_period,
                **dump_video_kwargs
            )
            algorithm.post_train_funcs.append(expl_video_func)

    addl_collectors = []
    addl_log_prefixes = []
    if mask_conditioned and mask_variant.get('log_mask_diagnostics', True):
        default_list = [
            'atomic',
            'atomic_seq',
            'cumul_seq',
            'full',
        ]
        eval_rollouts_to_log = mask_variant.get('eval_rollouts_to_log', default_list)
        for key in eval_rollouts_to_log:
            assert key in default_list

        # atomic masks
        if 'atomic' in eval_rollouts_to_log:
            for mask_id in eval_path_collector.mask_ids:
                mask_kwargs=dict(
                    mask_ids=[mask_id],
                    mask_distr=dict(
                        atomic=1.0,
                    ),
                )
                collector = create_path_collector(eval_env, eval_policy, mode='eval', mask_kwargs=mask_kwargs)
                addl_collectors.append(collector)
            addl_log_prefixes += [
                'mask_{}/'.format(''.join(str(mask_id)))
                for mask_id in eval_path_collector.mask_ids
            ]

        # full mask
        if 'full' in eval_rollouts_to_log:
            mask_kwargs=dict(
                mask_distr=dict(
                    full=1.0,
                ),
            )
            collector = create_path_collector(eval_env, eval_policy, mode='eval', mask_kwargs=mask_kwargs)
            addl_collectors.append(collector)
            addl_log_prefixes.append('mask_full/')

        # cumulative, sequential mask
        if 'cumul_seq' in eval_rollouts_to_log:
            mask_kwargs=dict(
                rollout_mask_order='fixed',
                mask_distr=dict(
                    cumul_seq=1.0,
                ),
            )
            collector = create_path_collector(eval_env, eval_policy, mode='eval', mask_kwargs=mask_kwargs)
            addl_collectors.append(collector)
            addl_log_prefixes.append('mask_cumul_seq/')

        # atomic, sequential mask
        if 'atomic_seq' in eval_rollouts_to_log:
            mask_kwargs=dict(
                rollout_mask_order='fixed',
                mask_distr=dict(
                    atomic_seq=1.0,
                ),
            )
            collector = create_path_collector(eval_env, eval_policy, mode='eval', mask_kwargs=mask_kwargs)
            addl_collectors.append(collector)
            addl_log_prefixes.append('mask_atomic_seq/')

        def get_mask_diagnostics(unused):
            from rlkit.core.logging import append_log, add_prefix, OrderedDict
            log = OrderedDict()
            for prefix, collector in zip(addl_log_prefixes, addl_collectors):
                paths = collector.collect_new_paths(
                    max_path_length,
                    variant['algo_kwargs']['num_eval_steps_per_epoch'],
                    discard_incomplete_paths=True,
                )
                old_path_info = eval_env.get_diagnostics(paths)

                keys_to_keep = []
                for key in old_path_info.keys():
                    if ('env_infos' in key) and ('final' in key) and ('Mean' in key):
                        keys_to_keep.append(key)
                path_info = OrderedDict()
                for key in keys_to_keep:
                    path_info[key] = old_path_info[key]

                generic_info = add_prefix(
                    path_info,
                    prefix,
                )
                append_log(log, generic_info)

            for collector in addl_collectors:
                collector.end_epoch(0)
            return log

        algorithm._eval_get_diag_fns.append(get_mask_diagnostics)

    if 'ckpt' in variant:
        from rlkit.util.io import local_path_from_s3_or_local_path
        import os.path as osp
        assert variant['algo_kwargs'].get('eval_only', False)

        def update_networks(algo, epoch):
            if 'ckpt_epoch' in variant:
                return

            if epoch % algo._eval_epoch_freq == 0:
                filename = local_path_from_s3_or_local_path(osp.join(variant['ckpt'], 'itr_%d.pkl' % epoch))
                print("Loading ckpt from", filename)
                data = torch.load(filename)#, map_location='cuda:1')
                eval_policy = data['evaluation/policy']
                eval_policy.to(ptu.device)
                algo.eval_data_collector._policy = eval_policy
                for collector in addl_collectors:
                    collector._policy = eval_policy

        algorithm.post_train_funcs.insert(0, update_networks)

    algorithm.train()

def full_post_process_mask_fn(
        obs_dict, context_dict,
        mask_conditioned,
        mask_variant,
        context_distrib,
        context_key,
        achieved_goal_key,
):
    assert mask_conditioned
    mode = mask_variant.get('context_post_process_mode', None)
    assert mode in [
        'prev_subtasks_solved',
        'dilute_prev_subtasks_uniform',
        'dilute_prev_subtasks_fixed',
        'atomic_to_corresp_cumul',
        None
    ]
    if mode is None:
        return context_dict

    if mode in [
        'prev_subtasks_solved',
        'dilute_prev_subtasks_uniform',
        'dilute_prev_subtasks_fixed',
        'atomic_to_corresp_cumul'
    ]:
        frac = mask_variant.get('context_post_process_frac', 0.50)
        cumul_mask_to_indices = context_distrib.get_cumul_mask_to_indices(context_dict['mask'])
        for k in cumul_mask_to_indices:
            indices = cumul_mask_to_indices[k]
            subset = np.random.choice(len(indices), int(len(indices)*frac), replace=False)
            cumul_mask_to_indices[k] = indices[subset]
    else:
        cumul_mask_to_indices = None
    pp_context_dict = copy.deepcopy(context_dict)

    if mode in ['prev_subtasks_solved', 'dilute_prev_subtasks_uniform', 'dilute_prev_subtasks_fixed']:
        cumul_masks = list(cumul_mask_to_indices.keys())
        for i in range(1, len(cumul_masks)):
            curr_mask = cumul_masks[i]
            prev_mask = cumul_masks[i-1]
            prev_obj_indices = np.where(np.array(prev_mask) > 0)[0]
            indices = cumul_mask_to_indices[curr_mask]
            if mode == 'prev_subtasks_solved':
                pp_context_dict[context_key][indices][:,prev_obj_indices] = \
                    obs_dict[achieved_goal_key][indices][:,prev_obj_indices]
            elif mode == 'dilute_prev_subtasks_uniform':
                pp_context_dict['mask'][indices][:, prev_obj_indices] = \
                    np.random.uniform(size=(len(indices), len(prev_obj_indices)))
            elif mode == 'dilute_prev_subtasks_fixed':
                pp_context_dict['mask'][indices][:, prev_obj_indices] = 0.5
        indices_to_relabel = np.concatenate(list(cumul_mask_to_indices.values()))
        orig_masks = obs_dict['mask'][indices_to_relabel]
        atomic_mask_to_subindices = context_distrib.get_atomic_mask_to_indices(orig_masks)
        atomic_masks = list(atomic_mask_to_subindices.keys())
        cumul_masks = list(cumul_mask_to_indices.keys())
        for i in range(1, len(atomic_masks)):
            orig_atomic_mask = atomic_masks[i]
            relabeled_cumul_mask = cumul_masks[i]
            subindices = atomic_mask_to_subindices[orig_atomic_mask]
            pp_context_dict['mask'][indices_to_relabel][subindices] = relabeled_cumul_mask

    return pp_context_dict


text:
import cv2
import numpy as np

#set the below variable to True if using the tracking accessory
useTrackingCam = True
if useTrackingCam:
    import tracker_servo as  tracker

# LINES_BODY is used when drawing the skeleton onto the source image.
# Each variable is a list of continuous lines.
# Each line is a list of keypoints as defined at https://google.github.io/mediapipe/solutions/pose.html#pose-landmark-model-blazepose-ghum-3d
LINES_BODY = [[28,30,32,28,26,24,12,11,23,25,27,29,31,27],
                [23,24],
                [22,16,18,20,16,14,12],
                [21,15,17,19,15,13,11],
                [8,6,5,4,0,1,2,3,7],
                [10,9],
                ]

class BlazeposeRenderer:
    def __init__(self,
                pose,
                show_3d=False,
                output=None):
        self.pose = pose
        self.show_3d = show_3d

        # Rendering flags
        self.show_rot_rect = False
        self.show_landmarks = True
        self.show_score = False
        self.show_fps = True

        if self.show_3d:

            self.vis3d = o3d.visualization.Visualizer()
            self.vis3d.create_window()
            opt = self.vis3d.get_render_option()
            opt.background_color = np.asarray([0, 0, 0])
            z = min(pose.img_h, pose.img_w)/3
            self.grid_floor = create_grid([0,pose.img_h,-z],[pose.img_w,pose.img_h,-z],[pose.img_w,pose.img_h,z],[0,pose.img_h,z],5,2, color=(1,1,1))
            self.grid_wall = create_grid([0,0,z],[pose.img_w,0,z],[pose.img_w,pose.img_h,z],[0,pose.img_h,z],5,2, color=(1,1,1))
            self.vis3d.add_geometry(self.grid_floor)
            self.vis3d.add_geometry(self.grid_wall)
            view_control = self.vis3d.get_view_control()
            view_control.set_up(np.array([0,-1,0]))
            view_control.set_front(np.array([0,0,-1]))

        if output is None:
            self.output = None
        else:
            fourcc = cv2.VideoWriter_fourcc(*"MJPG")
            self.output = cv2.VideoWriter(output,fourcc,pose.video_fps,(pose.img_w, pose.img_h))

    def draw_landmarks(self, body):
        if self.show_rot_rect:
            cv2.polylines(self.frame, [np.array(body.rect_points)], True, (0,255,255), 2, cv2.LINE_AA)
        if self.show_landmarks:
            list_connections = LINES_BODY
            lines = [np.array([body.landmarks[point,:2] for point in line]) for line in list_connections]
            # lines = [np.array([body.landmarks_padded[point,:2] for point in line]) for line in list_connections]
            cv2.polylines(self.frame, lines, False, (255, 180, 90), 2, cv2.LINE_AA)

############our code for tracker cam#####################
            #get the far left and far right X position output of the landmarks
            xMin = 1000000
            xMax = -100000
            yMin = 1000000
            yMax = -100000
            for i,x_y in enumerate(body.landmarks[:self.pose.nb_kps,:2]):
                x = x_y[0]
                y = x_y[1]
                if x > xMax:
                    xMax = x
                if x < xMin:
                    xMin = x
                if y > yMax:
                    yMax = y
                if y < yMin:
                    yMin = y
# this colors the circles for the rendered lines (not part of tracking cam)
                if i > 10:
                    color = (0,255,0) if i%2==0 else (0,0,255)
                elif i == 0:
                    color = (0,255,255)
                elif i in [4,5,6,8,10]:
                    color = (0,255,0)
                else:
                    color = (0,0,255)
                cv2.circle(self.frame, (x_y[0], x_y[1]), 4, color, -11)

############our code for tracker cam#####################
            if useTrackingCam:
                xMiddle = ((xMax-xMin)/2)+xMin
                yMiddle = ((yMax-yMin)/2)+yMin
                frameXmiddle = int(self.frame.shape[1]/2)
                frameYmiddle = int(self.frame.shape[0]/2)
                #draw a circle in the middle of the frame
                cv2.circle(self.frame, (int(frameXmiddle),int(frameYmiddle)), 35, (220,80,149), -1)
                #draw a circle in the middle of our target we want to track
                cv2.circle(self.frame, (int(xMiddle),int(yMiddle)), 20, (100,150,200), -1)

                #set the threshold on the movement bounding box
                #adjust the values (based on with of input frame)
                left_wall = .45 * self.frame.shape[1]
                right_wall = .55 *self.frame.shape[1]
                ceiling = .7 * self.frame.shape[0] #values measured from TOP of image
                ground = .8 * self.frame.shape[0]

                #get distance from tracking object to center of frame. Adjust movement angle of servo the further the object is from center of frame
                targetDistance = abs(xMiddle - (self.frame.shape[1]/2))
                yawMovementSpeed = 0
                if targetDistance > self.frame.shape[1] * .05:
                    yawMovementSpeed = 1
                if targetDistance > self.frame.shape[1] * .2:
                    yawMovementSpeed = 2
                if targetDistance > self.frame.shape[1] * .4:
                    yawMovementSpeed = 3
                tracker.set_yaw_speed(yawMovementSpeed)

                currentlyMoving = 'X'
                if xMiddle < left_wall:
                    tracker.left()
                    currentlyMoving = '<--'
                if xMiddle > right_wall:
                    tracker.right()
                    currentlyMoving = '-->'
                if yMiddle > ceiling:
                    tracker.down()
                if yMiddle < ground:
                    tracker.up()
                cv2.putText(self.frame, 'Track speed: ' + str(yawMovementSpeed) + ' direction: ' + currentlyMoving,(20,50),
                    cv2.FONT_HERSHEY_SIMPLEX, 2,(50,100,255),2,cv2.LINE_AA)

    def draw(self, frame, body):
        self.frame = frame
        if body:
            self.draw_landmarks(body)
            self.body = body
        else:
############ our tracker cam code
            # tracker cam will scan if no target object detected
            if useTrackingCam:
                tracker.scan()
        if self.frame is None:
            self.frame = frame
            self.body = None

        return self.frame

    def exit(self):
        if self.output:
            self.output.release()

    def waitKey(self, delay=1):
        key = cv2.waitKey(delay)
        return key

ext: py

reason: import numpy


text:
import itertools

import numpy as np

MACHINE_EPSILON = (
    7.0 / 3 - 4.0 / 3 - 1
)  # the difference between 1 and the smallest floating point number greater than 1

def clamp(x: float, left: float, right: float) -> float:
    return max(left, min(right, x))

def transform(vectors, matrix, directions=False):
    """Applies an affine transform to a list of vectors, entry-wise."""
    shape = np.array(vectors.shape)
    shape[-1] += 1
    vectors_4d = np.empty(shape, dtype=vectors.dtype)
    vectors_4d[..., :-1] = vectors
    # if directions=True translation components of transforms will be ignored
    vectors_4d[..., -1] = 0 if directions else 1
    # ensure type conditions are met so we can use np.dot(out=) kwarg for performance
    matrix = matrix.astype(vectors_4d.dtype, copy=False).T
    return np.dot(vectors_4d, matrix, out=vectors_4d)[..., :-1]

def aabb_to_sphere(aabb):
    """Returns a sphere that fits exactly around an axis-aligned bounding box."""
    diagonal = aabb[1] - aabb[0]
    center = aabb[0] + diagonal * 0.5
    radius = np.linalg.norm(diagonal) * 0.5
    return np.array([*center, radius])

def transform_aabb(aabb, matrix):
    """Applies an affine transform to an axis-aligned bounding box."""
    corners = np.array(list(itertools.product(*aabb.T)))
    corners_world = transform(corners, matrix)
    return np.array([corners_world.min(axis=0), corners_world.max(axis=0)])


text:
#
# Circular convolution of two 1-dim vectors
# 06_fou/convolution.py
#
# (c) Dr. Yves J. Hilpisch
# Derivatives Analytics with Python
#
import numpy as np

#
# Function Definitions
#

def revpy(a):
    ''' Reversing the order of the vector's numbers (for loop). '''
    a = np.array(a)
    n = len(a)
    c = np.zeros(n, dtype=np.float)
    c[0] = a[0]
    for j in range(1, n):
        c[j] = a[n - j]
    return c

def revnp(a):
    ''' Reversing the order of the vector's numbers (NumPy version). '''
    b = a.copy()
    b[1:] = b[1:][::-1]
    return b

def convolution(a, b):
    ''' Convolution of two vectors. '''
    if len(a) != len(b):
        raise ValueError("Lengths of vectors do not match.")
    n = len(a)
    c = np.zeros(n, dtype=np.float)
    for j in range(n):
        s = 0
        for k in range(n):
            if j - k >= 0:
                s += a[j - k] * b[k]
            else:
                s += a[j - k + n] * b[k]
        c[j] = s
    return c

max_stars_repo_name: PyCN/dawp

max_issues_repo_name: joshuatu/dawp

max_forks_repo_name: joshuatu/dawp


text:
import numpy as np
from pathlib import Path
import json
import os.path
import sys
import argparse
import time
import json, pickle
import torch
from math import ceil
from scipy import signal

THIS_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.abspath(os.path.join(THIS_DIR, os.pardir))
DATA_DIR = os.path.join(ROOT_DIR, 'data')
sys.path.append(ROOT_DIR)

PAD_STATE = 0
START_STATE = 1
END_STATE = 2
EMPTY_STATE = 3
NUM_SPECIAL_STATES = 4

HUMAN_DELTA = 0.125 # minimum block separation that we are happy with

#from models import create_model
from models.ddc_model import DDCModel
from feature_extraction.utils import distribute_tasks,ResampleLinear1D

parser = argparse.ArgumentParser(description='Get DDC features from song features')
parser.add_argument("data_path", type=str, help="Directory contining Beat Saber level folders")
parser.add_argument('--checkpoints_dir', type=str)
parser.add_argument('--audio_format', type=str, default="wav")
parser.add_argument('--experiment_name', type=str)
parser.add_argument('--peak_threshold', type=float, default=0.0148)
parser.add_argument('--fps', type=float, default=60)
parser.add_argument('--checkpoint', type=str, default="latest")
parser.add_argument('--temperature', type=float, default=1.00)
parser.add_argument('--cuda', action="store_true")
parser.add_argument("--step_size", metavar='', type=float, default=0.01666666666)
parser.add_argument("--replace_existing", action="store_true")

args = parser.parse_args()

# makes arugments into global variables of the same name, used later in the code
globals().update(vars(args))
data_path = Path(data_path)

## distributing tasks accross nodes ##
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
print(rank)
# print("creating {} of size {}".format(feature_name, feature_size))

experiment_name = args.experiment_name+"/"
checkpoint = args.checkpoint
temperature=args.temperature

from pathlib import Path

''' LOAD MODEL, OPTS, AND WEIGHTS'''
#%%

##loading opt object from experiment
opt = json.loads(open(ROOT_DIR.__str__()+"/feature_extraction/"+experiment_name+"opt.json","r").read())
# we assume we have 1 GPU in generating machine :P
if args.cuda:
    opt["gpu_ids"] = [0]
else:
    opt["gpu_ids"] = []
opt["checkpoints_dir"] = args.checkpoints_dir
# print(opt["checkpoints_dir"])
opt["load_iter"] = int(checkpoint)
if args.cuda:
    opt["cuda"] = True
else:
    opt["cuda"] = False
opt["experiment_name"] = args.experiment_name.split("/")[0]
if "dropout" not in opt: #for older experiments
    opt["dropout"] = 0.0
# construct opt Struct object
class Struct:
    def __init__(self, **entries):
        self.__dict__.update(entries)
opt = Struct(**opt)

assert opt.binarized

#model = create_model(opt)
model = DDCModel(opt)
#model.setup()
receptive_field = 1

checkpoint = "iter_"+checkpoint
model.load_networks(checkpoint)

#assuming mp3 for now. TODO: generalize
candidate_feature_files = sorted(data_path.glob('**/*'+audio_format+'_multi_mel_80.npy'), key=lambda path: path.parent.__str__())
tasks = distribute_tasks(candidate_feature_files,rank,size)

for i in tasks:
    path = candidate_feature_files[i]
    features_file = str(path)+"_"+"ddc_hidden"+".npy"
    if replace_existing or not os.path.isfile(features_file):
        print(path)

        sr = opt.sampling_rate
        hop = int(opt.step_size*sr)
        features = np.load(path)

        #generate level
        # first_samples basically works as a padding, for the first few outputs, which don't have any "past part" of the song to look at.
        first_samples = torch.full((1,opt.output_channels,receptive_field//2),START_STATE,dtype=torch.float)
        print(features.shape)
        features, peak_probs = model.generate_features(features)
        peak_probs = peak_probs[0,:,-1].cpu().detach().numpy()
        features = features.cpu().detach().numpy()
        features = features[0]
        features = ResampleLinear1D(features,int(np.floor(features.shape[0]*0.01*fps)))
        # features = downsample_signal(features[0], 0.01666666666667/0.01)
        print(features.shape)
        np.save(features_file,features)
        window = signal.hamming(ceil(HUMAN_DELTA/opt.step_size))
        smoothed_peaks = np.convolve(peak_probs,window,mode='same')

        thresholded_peaks = smoothed_peaks*(smoothed_peaks>args.peak_threshold)
        peaks = signal.find_peaks(thresholded_peaks)[0]
        print("number of peaks", len(peaks))

max_stars_repo_name: guillefix/mt-lightning

max_issues_repo_name: guillefix/mt-lightning

max_forks_repo_name: guillefix/transflower-lightning

reason: import numpy,from scipy


text:
from __future__ import print_function
import os, sys, re, json
from collections import defaultdict
import numpy as np
import pandas as pd
try:
    from word2keypress.weight_matrix import WEIGHT_MATRIX
    from word2keypress.weighted_edist import (
        STARTSTR, ENDSTR, KB, BLANK, SHIFT_KEY, CAPS_KEY, all_edits, _editdist)
except ImportError:
    from weight_matrix import WEIGHT_MATRIX
    from weighted_edist import (
        STARTSTR, ENDSTR, KB, BLANK, SHIFT_KEY, CAPS_KEY, all_edits, _editdist
    )

EDIT_DIST_CUTOFF = 1
WEIGHT_MATRIX = [
    (e,w) for e, w in WEIGHT_MATRIX
    if _editdist(e[0], e[1], limit=EDIT_DIST_CUTOFF)[1]
]

# giant_regex = re.complie('|'.join(
#     re.escape(l) for ((l,r),w) in WEIGHT_MATRIX))

def allowed_edits(pw_key_str):
    """
    Returns all the edits that are allowed for @pw.
    An edit=(l -> r) is allowed if @l is in @pw
    returns the filtered WEIGHT_MATRIX
    """
    if not pw_key_str.startswith(STARTSTR):
        pw_key_str = STARTSTR + pw_key_str + ENDSTR
        # print(pw_key_str)
    return sorted(
        [((l,r), w) for ((l,r),w) in WEIGHT_MATRIX
         if l.replace(BLANK, '') in pw_key_str],
        key=lambda x: x[1], reverse=True
    )

def edit_locs(pw_key_str, l):
    matched_indexes = [
        (m.start(), m.end())
        for m in re.finditer('({})'.format(re.escape(l.replace(BLANK, ''))),
                             pw_key_str)
        if m.start()<len(pw_key_str) and m.end()>0
    ]
    return matched_indexes

def apply_edit(pw_key_str, e):
    """
    Applies edits on the pw_key_str, whereever the edit e is possible
    If there are multiple location, then apply only at one location.
    """
    l, r = e
    matched_indexes = edit_locs(pw_key_str, l)
    assert matched_indexes, "Wow!! matched index is empty for pw={}, e={}"\
        .format(pw, e)
    # Choose one index at random from the possible options
    # i = np.random.randint(0, len(matched_indexes))
    # pos_s, pos_e = matched_indexes[i]
    # if BLANK in l:
    #     typo_key_str = _insert_edit(pw_key_str, l, r, pos_s, pos_e)
    # else:
    for m in matched_indexes:
        pos_s, pos_e = m
        typo_key_str = pw_key_str[:pos_s] + r + pw_key_str[pos_e:]
        yield typo_key_str.replace(BLANK, ''), 1.0/len(matched_indexes)

def num_typos(n, ed):
    # type: (int, int) -> int
    assert ed>=0, "edit distance should be no less than 0. Got = {}".format(ed)
    t = (2*96)**ed
    a = n+1
    for i in range(2, ed+1):
        a *= (n+i)
    return a*t

def get_prob(rpw, tpw):
    """
    Probability that rpw is mistyped to tpw,
    get all the edit locations. sum their probabiliries
    """
    edits = set(all_edits(rpw, tpw, N=1, edit_cutoff=1))
    pw_key_str = STARTSTR + KB.word_to_keyseq(rpw) + ENDSTR
    E = allowed_edits(pw_key_str)
    s = float(sum(x[1] for x in E))
    if(s==0): return 0.0
    # print("s = {} (len(E)={})".format(s, len(E)))
    # print(edits)
    total_ed1_typos_estimate = 2*96*(len(rpw) + 1)
    f = 1.0/num_typos(len(rpw), 1 if edits else 2)
    for e, w in E:
        if e not in edits: continue
        for typo_key_str, w_frac in apply_edit(pw_key_str, e):
            typo_key_str = typo_key_str.strip(STARTSTR).strip(ENDSTR)
            typo = KB.keyseq_to_word(typo_key_str)
            if typo == tpw:
                f += w*w_frac
    return f/s

def get_topk_typos(pw, k=10):
    """
    Returns top k typos of the word pw
    """
    pw_key_str = STARTSTR + KB.word_to_keyseq(pw) + ENDSTR
    E = sorted(allowed_edits(pw_key_str), key=lambda x: x[1]*len(x[0][0]),
               reverse=True)
    tt = defaultdict(float)
    s = float(sum(x[1] for x in E))
    # print("s = {} (len(E)={})".format(s, len(E)))
    i = 0
    debug_pw = {pw.swapcase()}
    while len(tt)<k*len(pw)*10 and i <len(E):
        e, w = E[i]
        for typo_key_str, w_frac in apply_edit(pw_key_str, e):
            typo_key_str = typo_key_str.strip(STARTSTR).strip(ENDSTR)
            typo = KB.keyseq_to_word(typo_key_str)
            tt[typo] += w * w_frac/s
            # if typo in debug_pw:
            #     print("{!r} ->> {} ({}, {})".format(typo_key_str, e, w, w*w_frac/s))
        i += 1
    return sorted(tt.items(), key=lambda x: x[1], reverse=True)[:k]

def read_typos(f_name):
    d = pd.read_csv(f_name, skipinitialspace=False)\
          .astype(str)
    d_ts = d[d.rpw != d.tpw].sample(int(0.03*len(d.index)), random_state=435)
    return d_ts

def test_model_rank(train_f):
    from pyxdameraulevenshtein import damerau_levenshtein_distance as dldist
    d_ts = read_typos(train_f)
    a = np.array([get_prob(rpw, tpw)
            for rpw, tpw in zip(d_ts.rpw, d_ts.tpw)
            if dldist(rpw.lower(), tpw.lower())<=1])
    a = a[a>0]
    rank = []
    for rpw, tpw in  zip(d_ts.rpw, d_ts.tpw):
        if dldist(rpw.lower(), tpw.lower())>1: continue
        k = 20
        typos = [tp for tp, w in get_topk_typos(rpw, k)]
        if tpw in typos:
            rank.append(typos.index(tpw)+1)
        else:
            rank.append(k)
    print("Avg_rank: ", sum(rank)/float(len(rank)*k))
    print(d_ts.shape, a.shape, a.mean(), a.std())
    return  a

def test_model_likelihood(train_f):
    from pyxdameraulevenshtein import damerau_levenshtein_distance as dldist
    d_ts = read_typos(train_f)
    ed = d_ts.apply(lambda r: dldist(r.rpw, r.tpw), axis=1)
    probs = d_ts[ed<=1].apply(lambda r: get_prob(r.rpw, r.tpw), axis=1)
    likelihood = np.log(probs[probs>0]).mean()
    return likelihood

if __name__ == '__main__':
    USAGE = """Usage:
$ {} [options] [arguments]
-allowed-edits <rpw> : returns the allowed edits of rpw
-sample <password>: samples typos for the password from the model
-prob <rpw> <tpw>: probability of rpw -> tpw
-topktypos <rpw> [<n>]  : returns n (default 10) typos of rpw
-test <typo-fname> : Tests the efficacy of the model, ~/pwtypos-code/typodata/typos.csv
-keypress  <rpw> : Return the keypress representation of
        """.format(sys.argv[0])
    if len(sys.argv)<=1:
        print(USAGE)
        exit(1)
    if sys.argv[1] == '-allowed-edits':
        pw = KB.word_to_keyseq(sys.argv[2])
        (l,r),w  = WEIGHT_MATRIX[0]
        assert l.replace(BLANK, '') in pw, "{!r}, {} {}"\
                .format(l.replace(BLANK, ''), pw, w)
        print(allowed_edits(pw))
    elif sys.argv[1] == '-sample':
        pw = sys.argv[2]
        print("{} --> {}".format(pw, len(set((sample_typos(pw, 100))))))
    elif sys.argv[1] == '-topktypos':
        pw = sys.argv[2]
        n = int(sys.argv[3]) if len(sys.argv)>3 else 10
        typos = get_topk_typos(pw, n)
        print('\n'.join("{}: {:.5f}".format(x,y) for x, y in typos))
        print("{} --> {}".format(pw, len(typos)))
    elif sys.argv[1] == '-prob':
        print(get_prob(sys.argv[2], sys.argv[3]))
    elif sys.argv[1] == '-test':
        # test_model_rank(sys.argv[2])
        print("Log-Likelihood: ", test_model_likelihood(sys.argv[2]))
    elif sys.argv[1] == '-keypress':
        print(repr(KB.word_to_keyseq(sys.argv[2])))
    else:
        print(USAGE)


text:
import numpy as np
import matplotlib.pyplot as plt
import math

size = 6
x = np.arange(size)

Learn_180 = [76.69, 74.43, 69.17, 61.65, 53.39, 53.39]

Learn_360 = [81.95, 75.19, 70.68, 62.41, 54.14, 54.14]

Person_720 = [84.96, 81.95, 72.18, 63.91, 61.65, 60.9]

total_width, n = 0.8, 3
width = total_width / n
x = x - (total_width - width) / 2

plt.xlabel('Prediction delay (s)', fontsize=20)
plt.ylabel('Average accuracy (%)', fontsize=20)
plt.bar(x-0.9*width, Learn_180, fc='#faa755', width=0.75*width, label='180s training model')
plt.bar(x, Learn_360, fc='#6b473c', width=0.75*width, label='360s training model')
plt.bar(x+0.9*width, Person_720, fc='#8a5d19', width=0.75*width, label='720s training model')
plt.xticks(x, ('1/3', 1, 3, 5, 10, 15), fontsize=18)
plt.yticks(fontsize=18)
plt.legend(loc='lower right', fontsize=15)
plt.show()


text:
import numpy as np
import matplotlib.pyplot as plt
from fbm import FBM

def generate_sin(x, T, A):
    """Generate a mixed sinusoidal sequence"""
    y = np.zeros(len(x))
    for i in range(len(T)):
        y += A[i] * np.sin(2 * np.pi / T[i] * x)

    return y

def gen_covariates(x, index):
    """Generate covariates"""
    covariates = np.zeros((x.shape[0], 4))
    covariates[:, 0] = (x // 24) % 7
    covariates[:, 1] = x % 24
    covariates[:, 2] = (x // (24 * 30)) % 12
    covariates[:, 0] = covariates[:, 0] / 6
    covariates[:, 1] = covariates[:, 1] / 23
    covariates[:, 2] = covariates[:, 2] / 11

    covariates[:, -1] = np.zeros(x.shape[0]) + index
    return covariates

def fractional_brownian_noise(length, hurst, step):
    """Genereate fractional brownian noise"""
    f = FBM(length, hurst, step)
    noise = f.fbm()
    return noise

def synthesis_data():
    """synthesis a mixed sinusoidal dataset"""
    T = [24, 168, 720]
    seq_num = 60
    seq_len = T[-1] * 20
    data = []
    covariates = []
    for i in range(seq_num):
        start = int(np.random.uniform(0, T[-1]))
        x = start + np.arange(seq_len)
        A = np.random.uniform(5, 10, 3)
        y = generate_sin(x, T, A)
        data.append(y)
        covariates.append(gen_covariates(x, i))
        # plt.plot(x[:T[-1]], y[:T[-1]])
        # plt.show()

    data = np.array(data)
    mean, cov = polynomial_decay_cov(seq_len)
    noise = multivariate_normal(mean, cov, seq_num)
    data = data + noise
    covariates = np.array(covariates)
    data = np.concatenate([data[:, :, None], covariates], axis=2)
    np.save('data/synthetic.npy', data)

def covariance(data):
    """compute the covariance of the data"""
    data_mean = data.mean(0)
    data = data - data_mean
    length = data.shape[1]
    data_covariance = np.zeros((length, length))

    for i in range(length):
        for j in range(length):
            data_covariance[i, j] = (data[:, i] * data[:, j]).mean()

    return data_covariance

def test_fbm():
    """Plot the covariance of the generated fractional brownian noise"""
    f = FBM(300, 0.3, 1)
    fbm_data = []
    for i in range(100):
        sample = f.fbm()
        fbm_data.append(sample[1:])
    fbm_data = np.array(fbm_data)
    cov = covariance(fbm_data)
    plt.imshow(cov)
    plt.savefig('fbm_cov.jpg')

def polynomial_decay_cov(length):
    """Define the function of covariance decay with distance"""
    mean = np.zeros(length)

    x_axis = np.arange(length)
    distance = x_axis[:, None] - x_axis[None, :]
    distance = np.abs(distance)
    cov = 1 / (distance + 1)
    return mean, cov

def multivariate_normal(mean, cov, seq_num):
    """Generate multivariate normal distribution"""
    noise = np.random.multivariate_normal(mean, cov, (seq_num,), 'raise')
    return noise

if __name__ == '__main__':
    synthesis_data()

max_stars_repo_name: alipay/Pyraformer

max_issues_repo_name: alipay/Pyraformer

max_forks_repo_name: alipay/Pyraformer


text:
""" Statistical tests for evaluating distribution of 1 or 2 categorical
variables with at least one having more than two levels

functions:

1. one_categorical_hypothesis
2. two_categorical_hypothesis
"""

import numpy as np
import scipy.stats

def one_categorical_hypothesis(counts: np.ndarray, nobs: np.ndarray) -> tuple:
    """Applying chi square test goodness of fit

    Ho: the observed counts of the input groups follow population distribution
    HA: the observed counts of groups do not follow population distribution
        (not random pick form population)

    Args:
        counts (np.ndarray): input group  observed counts
        nobs (np.ndarray): input group total count

    Returns:
        tuple: chi square value, p value
    """
    p_expected = sum(counts) / sum(nobs)
    expected_counts = nobs * p_expected
    chi_square = sum((counts-expected_counts)**2/expected_counts)
    df = len(nobs) - 1
    p_value = scipy.stats.chi2.sf(chi_square, df)
    return chi_square, p_value

def two_categorical_hypothesis(observed: np.ndarray) -> tuple:
    """Applying chi square independence test to compare two variables

    Ho: two variables are independent
    Ha: two variables are dependent

    Args:
        observed (np.ndarray): 2d array the rows represent first variable
                                        the columns represent second variable

    Returns:
        tuple: chi square value, p value
    """
    nrow, ncol = observed.shape
    row_totals = np.sum(observed, axis=1).reshape(nrow, 1)
    column_totals = np.sum(observed, axis=0).reshape(ncol, 1)
    total = np.sum(observed)

    expected = column_totals.T * (row_totals/total)
    print(expected.shape)
    chi_square = np.sum((observed-expected)**2/expected)
    df = (nrow - 1) * (ncol - 1)
    p_value = scipy.stats.chi2.sf(chi_square, df)
    return chi_square, p_value

reason: import numpy,import scipy


text:
#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import PointCloud2, PointField
from pyquaternion import Quaternion
import sensor_msgs.point_cloud2 as pc2
from std_msgs.msg import Header
from jsk_recognition_msgs.msg import BoundingBox, BoundingBoxArray

import sys
import time
sys.path.append("/home/indra/rospy3_melodic/src/second_ros/second.pytorch")

import numpy as np
import math
import pickle
from pathlib import Path

import torch
from google.protobuf import text_format
from second.pytorch.train import build_network
from second.protos import pipeline_pb2
from second.utils import config_tool

class Second_ROS:
    def __init__(self):
        config_path, ckpt_path = self.init_ros()
        self.init_second(config_path, ckpt_path)

    def init_second(self, config_path, ckpt_path):
        """ Initialize second model """

        config = pipeline_pb2.TrainEvalPipelineConfig()
        with open(config_path, "r") as f:
            proto_str = f.read()
            text_format.Merge(proto_str, config)
        input_cfg = config.eval_input_reader
        model_cfg = config.model.second
        #config_tool.change_detection_range_v2(model_cfg, [-50, -50, 50, 50])
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # device = torch.device("cpu")

        self.net = build_network(model_cfg).to(self.device).eval()
        self.net.load_state_dict(torch.load(ckpt_path))
        target_assigner = self.net.target_assigner
        self.voxel_generator = self.net.voxel_generator

        grid_size = self.voxel_generator.grid_size
        feature_map_size = grid_size[:2] // config_tool.get_downsample_factor(model_cfg)
        feature_map_size = [*feature_map_size, 1][::-1]

        self.anchors = target_assigner.generate_anchors(feature_map_size)["anchors"]
        self.anchors = torch.tensor(self.anchors, dtype=torch.float32, device=self.device)
        self.anchors = self.anchors.view(1, -1, 7)

    def init_ros(self):
        """ Initialize ros parameters """

        self.sub_velo = rospy.Subscriber("/kitti/velo/pointcloud", PointCloud2, self.lidar_callback, queue_size=1, buff_size=2**24)
        self.pub_bbox = rospy.Publisher("/detections", BoundingBoxArray, queue_size=1)

        config_path = rospy.get_param("/config_path", "/home/indra/fyp_perception/traveller/second.pytorch/second/configs/car.fhd.config")
        ckpt_path = rospy.get_param("/ckpt_path", "/home/indra/Documents/pretrained_models_v1.5/car_onestage/voxelnet-27855.tckpt")

        #config_path = rospy.get_param("/config_path", "/home/indra/rospy3_melodic/src/second_ros/config/all.fhd.config")
        #ckpt_path = rospy.get_param("/ckpt_path", "/home/indra/rospy3_melodic/src/second_ros/trained_models/voxelnet-99040.tckpt")

        return config_path, ckpt_path

    def inference(self, points):
        num_features = 4
        points = points.reshape([-1, num_features])
        print(points.shape)
        dic = self.voxel_generator.generate(points)
        voxels, coords, num_points = dic['voxels'], dic['coordinates'], dic['num_points_per_voxel']
        coords = np.pad(coords, ((0, 0), (1, 0)), mode='constant', constant_values=0)
        voxels = torch.tensor(voxels, dtype=torch.float32, device=self.device)
        coords = torch.tensor(coords, dtype=torch.int32, device=self.device)
        num_points = torch.tensor(num_points, dtype=torch.int32, device=self.device)

        input_points = {
        "anchors": self.anchors,
        "voxels": voxels,
        "num_points": num_points,
        "coordinates": coords,
        }

        pred = self.net(input_points)[0]
        boxes_lidar = pred["box3d_lidar"].detach().cpu().numpy()
        scores = pred["scores"].detach().cpu().numpy()
        label = pred["label_preds"].detach().cpu().numpy()

        return boxes_lidar, scores, label

    def lidar_callback(self, msg):
        """ Captures pointcloud data and feed into second model for inference """

        #arr_box = BoundingBoxArray()

        pcl_msg = pc2.read_points(msg, skip_nans=False, field_names=("x", "y", "z","intensity","ring"))
        np_p = np.array(list(pcl_msg), dtype=np.float32)
        #print("np_p shape: ", np_p.shape)
        #np_p = np.delete(np_p, -1, 1)  #  delete "ring" field

        # convert to xyzi point cloud
        x = np_p[:, 0].reshape(-1)
        y = np_p[:, 1].reshape(-1)
        z = np_p[:, 2].reshape(-1)
        if np_p.shape[1] == 4: # if intensity field exists
            i = np_p[:, 3].reshape(-1)
        else:
            i = np.zeros((np_p.shape[0], 1)).reshape(-1)
        cloud = np.stack((x, y, z, i)).T

        # start processing
        tic = time.time()
        boxes_lidar, scores, label = self.inference(cloud)
        toc = time.time()
        fps = 1/(toc-tic)
        rospy.loginfo("FPS: %f", fps)

        num_detections = len(boxes_lidar)
        arr_bbox = BoundingBoxArray()

        for i in range(num_detections):
            bbox = BoundingBox()

            bbox.header.frame_id = msg.header.frame_id
            bbox.header.stamp = rospy.Time.now()

            bbox.pose.position.x = float(boxes_lidar[i][0])
            bbox.pose.position.y = float(boxes_lidar[i][1])
            bbox.pose.position.z = float(boxes_lidar[i][2]) + float(boxes_lidar[i][5]) / 2
            bbox.dimensions.x = float(boxes_lidar[i][3])  # width
            bbox.dimensions.y = float(boxes_lidar[i][4])  # length
            bbox.dimensions.z = float(boxes_lidar[i][5])  # height

            q = Quaternion(axis=(0, 0, 1), radians=float(boxes_lidar[i][6]))
            bbox.pose.orientation.x = q.x
            bbox.pose.orientation.y = q.y
            bbox.pose.orientation.z = q.z
            bbox.pose.orientation.w = q.w
            bbox.value = scores[i]
            bbox.label = label[i]

            arr_bbox.boxes.append(bbox)

        arr_bbox.header.frame_id = msg.header.frame_id
        arr_bbox.header.stamp = rospy.Time.now()
        #print("Number of detections: {}".format(num_detections))

        self.pub_bbox.publish(arr_bbox)

if __name__ == '__main__':
    sec = Second_ROS()
    rospy.init_node('second_ros_node', anonymous=True)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        del sec
        print("Shutting down")


text:
from typing import Any, Callable, Sequence

import flax.linen as nn
import jax.numpy as jnp
from flax import linen as nn

ModuleDef = Any

class MLPResNetV2Block(nn.Module):
    """MLPResNet block."""
    features: int
    act: Callable

    @nn.compact
    def __call__(self, x):
        residual = x
        y = nn.LayerNorm()(x)
        y = self.act(y)
        y = nn.Dense(self.features)(y)
        y = nn.LayerNorm()(y)
        y = self.act(y)
        y = nn.Dense(self.features)(y)

        if residual.shape != y.shape:
            residual = nn.Dense(self.features)(residual)

        return residual + y

class MLPResNetV2(nn.Module):
    """MLPResNetV2."""
    num_blocks: int
    num_classes: int
    features: int = 256
    dtype: Any = jnp.float32
    act: Callable = nn.relu

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.features)(x)
        for _ in range(self.num_blocks):
            x = MLPResNetV2Block(self.features, act=self.act)(x)
        x = nn.LayerNorm()(x)
        x = self.act(x)
        x = nn.Dense(self.num_classes, dtype=self.dtype)(x)
        x = jnp.asarray(x, self.dtype)
        return x

max_stars_repo_name: ikostrikov/linenplus

max_issues_repo_name: ikostrikov/linenplus

max_forks_repo_name: ikostrikov/linenplus

reason: import jax


text:
# The MIT License (MIT)
# Copyright (c) 2020 by Brockmann Consult GmbH and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import unittest
from typing import List

import numpy as np
import xarray as xr

from nc2zarr.error import ConverterError
from nc2zarr.preprocessor import DatasetPreProcessor
from tests.helpers import new_test_dataset

class DatasetPreProcessorTest(unittest.TestCase):

    def test_time_dim_is_added_and_time_bounds_variable_is_converted_to_var(self):
        ds = new_test_dataset(day=5,
                              add_time_bnds=True,
                              add_time_as_dim_to_vars=False,
                              time_bounds_name='date_bnds',
                              add_time_bounds_as_var=True)
        pre_processor = DatasetPreProcessor(input_concat_dim='time')
        ds = pre_processor.preprocess_dataset(ds)
        self.assertIn('time', ds)
        self.assertEqual(np.array(['2020-12-05T10:00:00.000000000'], dtype='datetime64[ns]'),
                         ds.time.data)
        self.assertIn('bounds', ds.time.attrs)
        self.assertEqual('date_bnds', ds.time.bounds)
        self.assertIn('date_bnds', ds)
        self.assertIn('date_bnds', ds.coords)
        self.assertEqual(('time', 'bnds'), ds.date_bnds.dims)
        self.assertIn('r_ui16', ds)
        self.assertEqual(('time', 'lat', 'lon'), ds.r_ui16.dims)

    def test_select_variables(self):
        ds = new_test_dataset(day=1)
        self.assertIn('time', ds)
        pre_processor = DatasetPreProcessor(input_variables=['r_i32', 'lon', 'lat', 'time'], input_concat_dim='time')
        new_ds = pre_processor.preprocess_dataset(ds)
        self.assertIsInstance(new_ds, xr.Dataset)
        self.assertAllInDataset(['r_i32', 'lon', 'lat', 'time'], new_ds)
        self.assertNoneInDataset(['r_ui16', 'r_f32'], new_ds)

    def test_leaves_time_coord_untouched(self):
        ds = new_test_dataset(day=1)
        self.assertIn('time', ds)
        pre_processor = DatasetPreProcessor(input_variables=None, input_concat_dim='time')
        new_ds = pre_processor.preprocess_dataset(ds)
        self.assertIsInstance(new_ds, xr.Dataset)
        self.assertAllInDataset(['r_ui16', 'r_ui16', 'r_i32', 'lon', 'lat', 'time'], new_ds)
        self.assertIn('time', new_ds)
        self.assertEqual(ds.time, new_ds.time)

    def test_time_var_is_scalar(self):
        ds = new_test_dataset(day=2)
        time = ds.time
        ds = ds.rename_vars({'time': 't'})
        ds = ds.swap_dims({'time': 't'})
        ds = ds.assign_coords({'time': xr.DataArray(time[0].values.item(),
                                                    dims=())})
        ds.time.attrs.update(time.attrs)
        ds.time.encoding.update(time.encoding)
        ds = ds.drop_vars('t')
        pre_processor = DatasetPreProcessor(input_variables=None, input_concat_dim='time')
        new_ds = pre_processor.preprocess_dataset(ds)
        self.assertIsInstance(new_ds, xr.Dataset)
        self.assertAllInDataset(['r_ui16', 'r_ui16', 'r_i32', 'lon', 'lat', 'time'],
                                new_ds)
        self.assertIn('time', new_ds)
        self.assertEqual(('time',), new_ds.time.dims)
        self.assertEqual((1,), new_ds.time.shape)
        self.assertEqual(ds.time, new_ds.time)
        self.assertEqual({"long_name": "time"}, new_ds.time.attrs)
        self.assertEqual({'calendar': 'proleptic_gregorian',
                          'units': 'seconds since 1970-01-01 00:00:00'},
                         new_ds.time.encoding)

    def test_adds_time_dim_from_iso_format_attrs(self):
        ds = new_test_dataset(day=None)
        ds.attrs.update(time_coverage_start='2020-09-08 10:30:00',
                        time_coverage_end='2020-09-08 12:30:00')
        self._test_adds_time_dim(ds)

    def test_adds_time_dim_from_iso_format_attrs_2(self):
        ds = new_test_dataset(day=None)
        ds.attrs.update(time_coverage_start='2020-09-08T10:30:00Z',
                        time_coverage_end='2020-09-08T12:30:00Z')
        self._test_adds_time_dim(ds)

    def test_adds_time_dim_from_non_iso_format_attrs(self):
        ds = new_test_dataset(day=None)
        ds.attrs.update(time_coverage_start='20200908103000',
                        time_coverage_end='20200908123000')
        self._test_adds_time_dim(ds)

    def test_illegal_time_coverage(self):
        ds = new_test_dataset(day=None)
        ds.attrs.update(time_coverage_start='yesterday',
                        time_coverage_end='20200908123000')
        self._test_raises(ds, 'Cannot parse timestamp from "yesterday".')

    def test_missing_time_coverage(self):
        ds = new_test_dataset(day=None)
        self._test_raises(ds, 'Missing time_coverage_start and/or time_coverage_end in dataset attributes.')

        ds = new_test_dataset(day=None)
        ds.attrs.update(start_time='2020-09-08T10:30:00Z',
                        end_time='2020-09-08T12:30:00Z')
        self._test_raises(ds, 'Missing time_coverage_start and/or time_coverage_end in dataset attributes.')

    def test_illegal_concat_dim(self):
        ds = new_test_dataset(day=None)
        self._test_raises(ds, 'Missing (coordinate) variable "t" for dimension "t".',
                          input_concat_dim='t')

    def _test_adds_time_dim(self, ds: xr.Dataset):
        self.assertNotIn('time', ds)
        pre_processor = DatasetPreProcessor(input_variables=None, input_concat_dim='time')
        new_ds = pre_processor.preprocess_dataset(ds)
        self.assertIsInstance(new_ds, xr.Dataset)
        self.assertAllInDataset(['r_ui16', 'r_ui16', 'r_i32', 'lon', 'lat', 'time', 'time_bnds'], new_ds)
        self.assertEqual(1, len(new_ds.time))
        self.assertEqual(np.array(['2020-09-08T11:30:00'], dtype='datetime64[ns]'),
                         np.array(new_ds.time[0], dtype='datetime64[ns]'))
        self.assertEqual({'bounds': 'time_bnds'}, new_ds.time.attrs)
        self.assertEqual(1, len(new_ds.time_bnds))
        self.assertEqual(np.array(['2020-09-08T10:30:00'], dtype='datetime64[ns]'),
                         np.array(new_ds.time_bnds[0][0], dtype='datetime64[ns]'))
        self.assertEqual(np.array(['2020-09-08T12:30:00'], dtype='datetime64[ns]'),
                         np.array(new_ds.time_bnds[0][1], dtype='datetime64[ns]'))

    def _test_raises(self, ds, expected_message: str, input_concat_dim='time'):
        pre_processor = DatasetPreProcessor(input_variables=None, input_concat_dim=input_concat_dim)
        with self.assertRaises(ConverterError) as cm:
            pre_processor.preprocess_dataset(ds)
        self.assertEqual(expected_message, f'{cm.exception}')

    def assertAllInDataset(self, var_names: List[str], ds: xr.Dataset):
        for var_name in var_names:
            self.assertIn(var_name, ds)

    def assertNoneInDataset(self, var_names: List[str], ds: xr.Dataset):
        for var_name in var_names:
            self.assertNotIn(var_name, ds)

class CustomPreprocessorTest(unittest.TestCase):
    @classmethod
    def my_preprocessor(cls, ds: xr.Dataset) -> xr.Dataset:
        return ds.swap_dims({"sounding_dim": "time"})

    @classmethod
    def new_input_dataset(cls, /, offset, size):
        return xr.Dataset(dict(time=xr.DataArray(np.arange(offset, offset + size),
                                                 dims='sounding_dim'),
                               pressure=xr.DataArray(np.random.random(size * 20).reshape((size, 20)),
                                                     dims=['sounding_dim', 'levels_dim'])))

    def test_swap_dims(self):
        pre_processor = DatasetPreProcessor(input_variables=None,
                                            input_custom_preprocessor="tests.test_preprocessor:"
                                                                      "CustomPreprocessorTest.my_preprocessor",
                                            input_concat_dim='time')
        ds = self.new_input_dataset(0, size=100)
        self.assertEqual({'sounding_dim': 100, 'levels_dim': 20}, ds.dims)
        self.assertEqual(('sounding_dim',), ds.time.dims)
        self.assertEqual(('sounding_dim', 'levels_dim'), ds.pressure.dims)

        ds = pre_processor.preprocess_dataset(ds)
        self.assertEqual({'time': 100, 'levels_dim': 20}, ds.dims)
        self.assertEqual(('time',), ds.time.dims)
        self.assertEqual(('time', 'levels_dim'), ds.pressure.dims)


text:
import processor2000
import fitter2000
from cv2 import cv2
import jigsaw2000
import os
import queue
import threading
import time
import numpy as np

datadir = r'C:\jigsaw\data\2000'

font = cv2.FONT_HERSHEY_SIMPLEX

class calibrator():
	def __init__(self):
		self.factor_r = self.factor_g = self.factor_b = 0

	def calibrate(self, video, seconds):
		cal_time_start = time.time()
		calibrated = False
		while not calibrated:
			check,img = video.read()
			d = time.time() - cal_time_start
			hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
			text = ''
			hok = hisogramok(hsv)
			if hok:
				cal_time_start = time.time()
				cv2.putText(img, 'Remove peice', (10,450), font, 3, (0, 255, 0), 2, cv2.LINE_AA)
				cv2.imshow('Webcam view', img)
				cv2.waitKey(1)
			elif (d < seconds): # allow 3 seconds for alignment following histogram match confirmation
				cv2.putText(img, 'Calibrating '+str(seconds-int(d)), (10,450), font, 3, (0, 255, 0), 2, cv2.LINE_AA)
				cv2.imshow('Webcam view', img)
				cv2.waitKey(1)
			else:
				check,img = video.read()
				b,g,r = cv2.split(img)
				pixels = (img.shape[0]*img.shape[1])
				avg_r =  r.sum() / pixels
				avg_g =  g.sum() / pixels
				avg_b =  b.sum() / pixels
				blur = cv2.GaussianBlur(img,(65,65),0)
				blur_b,blur_g,blur_r = cv2.split(blur)
				self.factor_r = blur_r / avg_r
				self.factor_g = blur_g / avg_g
				self.factor_b = blur_b / avg_b
				cv2.waitKey(1)
				calibrated = True

	def capture(self, video):
		check,img = video.read()
		b,g,r = cv2.split(img)
		r_ = r / self.factor_r
		g_ = g / self.factor_g
		b_ = b / self.factor_b

		r = np.uint8(np.clip(r_, 0, 255))
		g = np.uint8(np.clip(g_, 0, 255))
		b = np.uint8(np.clip(b_, 0, 255))

		return cv2.merge((b,g,r))

class database():
	foldernames = ['Yellow','White','Orange','Brown','Red','Pink','Purple','Blue','LghtGrn','DarkGrn','Other1','Other2','Other3']
	sx = 10
	sy = 20

	def getNextId(self):
		while (os.path.exists(database.id2path(self.id))):
			self.id += 1
		return self.id

	def nextPage(self):
		page,x,y = database.id2crd(self.id)
		page = page + 1
		page = page % len(database.foldernames)
		x = y = 0
		self.id = page * self.sx * self.sy
		self.getNextId()
		page,x,y = database.id2crd(self.id)
		print('Skipped to page:', '[',page,']',database.foldernames[page],'\t','\tx:',x,'\ty:',y,'\tid:',self.id,'name:',database.id2shortname(self.id))

	def greet(self):
		page,x,y = database.id2crd(self.id)
		print('Current:', '[',page,']',database.foldernames[page],'\t','\tx:',x,'\ty:',y,'\tid:',self.id,'name:',database.id2shortname(self.id))

	@staticmethod
	def id2path(id):
		return datadir + '\\png' + '\\' + database.id2filename(id) + '.png'

	@staticmethod
	def alum(i):
		return chr(ord('a')+i)

	@staticmethod
	def getfilename(page,x,y):
		name =  database.foldernames[page] + '/' + database.alum(x) + '_' + str(y+1)
		return name

	@staticmethod
	def id2filename(id):
		page,x,y = database.id2crd(id)
		path = database.getfilename(page,x,y)
		return path

	@staticmethod
	def id2shortname(id):
		page,x,y = database.id2crd(id)
		return database.alum(x) + '_' + str(1+y)

	@staticmethod
	def id2crd(id):
		per_page = database.sx*database.sy
		page = id // per_page
		y = (id % per_page) // database.sx
		x = (id % database.sx)
		return (page,x,y)

	def __init__(self):
		for f in database.foldernames:
			path = datadir + '\\png' + '\\' + f
			if (not os.path.isdir(path)):
				os.mkdir(path)
			path = datadir + '\\npz' + '\\' + f
			if (not os.path.isdir(path)):
				os.mkdir(path)
		self.id = 0
		self.getNextId()

def hisogramok(hsv):
	h,s,v = cv2.split(hsv)
	blacksPixels = cv2.inRange(s, 0, 128)
	tot =  blacksPixels.sum() / 255
	pct = tot / (h.shape[0]*h.shape[1])
	return (pct >= 0.25)

def wait4histogram(xvid, video, seconds, db, name=''):
	nohist = time.time()
	if (seconds > 0):
		while(True):
			img = xvid.capture(video)
			hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
			text = ''
			hok = hisogramok(hsv)
			d = time.time() - nohist
			if not (hok):
				nohist = time.time()
				text = 'Not detected'
			else:
				text = 'Capture in '+str(seconds-int(d))
			if (d >= seconds):
				ok = processor2000.check_before_processing(img)
				if (ok):
					cv2.putText(img, 'Analysing', (10,250), font, 3, (255, 0, 0), 2, cv2.LINE_AA)
					cv2.putText(img, name, (10,450), font, 3, (255, 0, 0), 2, cv2.LINE_AA)
					cv2.imshow('Webcam view', img)
					cv2.waitKey(1)
					return
				else:
					nohist = time.time()
			cv2.putText(img, text, (10,250), font, 3, (0, 255, 0), 2, cv2.LINE_AA)
			cv2.imshow('Webcam view', img)
			key = cv2.waitKey(1)
			if (key == ord('j')):
				db.nextPage()
			if (key == ord('?')):
				db.greet()
			if (key == ord('q')):
				db.greet()
				exit(1)
	else:
		while(True):
			img = xvid.capture(video)
			hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
			text = ''
			hok = hisogramok(hsv)
			if not (hok): return
			cv2.putText(img, 'Captured', (10,250), font, 3, (255, 0, 0), 2, cv2.LINE_AA)
			cv2.putText(img, name, (10,450), font, 3, (255, 0, 0), 2, cv2.LINE_AA)
			cv2.imshow('Webcam view', img)
			cv2.waitKey(1)

def main():
	d = database()
	video = cv2.VideoCapture(0)

	# calibrate:
	xvid = calibrator()
	xvid.calibrate(video, 2)

	while True:
		id = d.getNextId()
		pngfile = database.id2path(id)
		shortname = database.id2shortname(id)
		wait4histogram(xvid,video, 3, d, shortname)
		# page might change...
		id = d.getNextId()
		shortname = database.id2shortname(id)
		pngfile = database.id2path(id)
		##########
		img = xvid.capture(video)
		p = processor2000.process_cam(img)
		p.id = database.id2filename(id)
		#p.show(True)
		p.save()
		cv2.imwrite(pngfile, img)
		wait4histogram(xvid,video,-1, d, shortname)

main()

max_stars_repo_name: EitanTal/PyJigsaw

max_issues_repo_name: EitanTal/PyJigsaw

max_forks_repo_name: EitanTal/PyJigsaw


text:
"""Functionality to distill a complex model knowledge to a simpler model"""

from log_calls import record_history
from ...automl.base import AutoML
from ...automl.presets.tabular_presets import TabularAutoML
from ...ml_algo.base import TabularMLAlgo
from ...ml_algo.boost_cb import BoostCB
from ...ml_algo.boost_lgbm import BoostLGBM
from ...pipelines.features.lgb_pipeline import LGBSimpleFeatures
from ...reader.base import PandasToPandasReader
from ...automl.base import MLPipeline
from ...utils.logging import get_logger
from ...tasks import Task
from typing import Sequence
from sklearn.exceptions import NotFittedError
from sklearn.metrics import roc_auc_score, accuracy_score, mean_squared_error, log_loss
from sklearn.feature_selection import RFECV
from pandas import DataFrame
import numpy as np

logger = get_logger(__name__)

class AlgoDistillError(Exception):
    def __init__(self, error):
        self.text = error

@record_history(enabled=False)
class Distiller:
    """This class contains utilities to perform knowledge distillation from complex models to simpler models.

    A model to perform distillation from is called Teacher.
    A model to learn knowledge from the teacher is called Student.
    """
    @property
    def is_fitted(self):
        return self._fitted

    def __init__(self, teacher: TabularAutoML, students = None):
        """Takes a teacher and a list of students as input.

        Args:
            teacher: is used to learn knowledge from given data.
            students: are used to learn knowledge from the teacher.
        """
        self.teacher = teacher
        self.students = list()
        self._fitted = False
        self.models_scores = None
        self.metric = None

        if not students:
            students = [BoostCB, BoostLGBM]
        elif type(students).__name__ == 'list':
            for algo in students:
                if type(algo).__name__ != 'TabularAutoML':
                    raise AlgoDistillError('Algorithm TypeError')
        else:
            if type(students).__name__ != 'TabularAutoML':
                raise AlgoDistillError('Algorithm TypeError')
            student = list()
            student.append(students)
            students = student

        for algo in students:
            reader = PandasToPandasReader(Task(self.teacher.task.name), samples=None, max_nan_rate=1, max_constant_rate=1,
                                          advanced_roles=True, drop_score_co=-1, n_jobs=1)
            pipeline_lvl1 = MLPipeline(ml_algos=[algo(default_params={'verbose': 0})],
                                       pre_selection=None,
                                       features_pipeline=LGBSimpleFeatures(),
                                       post_selection=None)
            self.students.append(AutoML(reader, [[pipeline_lvl1]], skip_conn=False, verbose=0))

    def fit(self, data, **kwargs):
        """Fits the teacher to given data.

        Args:
            data: a dataset to fit the teacher to.
            **kwargs: optional params for fitting the teacher.
        """
        assert not self._fitted, 'The distiller is already fitted'
        self.teacher.fit_predict(data, **kwargs)
        self._fitted = True

    def fit_predict(self, data, **kwargs):
        """Fits the teacher to given data and returns labels for fitting the students.

        Args:
            data: a dataset to fit the teacher to.
            **kwargs: optional params for fitting the teacher.

        Returns:
            labels for fitting the students.
        """
        assert not self._fitted, 'The distiller is already fitted'
        self.teacher.fit_predict(data, **kwargs)
        self._fitted = True
        return self.teacher.predict(data)

    def predict(self, data):
        """Returns labels for fitting the students.

        Args:
            data: a dataset to compute labels for.

        Returns:
            labels for fitting the students.
        """
        if not self._fitted:
            raise NotFittedError
        return self.teacher.predict(data)

    def distill(self, data, labels=None, metric=None):
        """Fits students to given data and finds the best one.

        Args:
            data: dataset to fit students to.
            labels: labels for fitting students. If not provided, predict labels using the teacher.
            metric: metric to maximize to find the best student.

        Returns:
            the best student.
        """
        if not self._fitted and not labels:
            raise NotFittedError
        _data = data.drop(columns=self.teacher.reader.used_array_attrs['target'])
        preds = dict()

        if metric:
            self.metric = metric

        else:
            if self.teacher.task.name == 'binary':
                self.metric = 'AUC'

            if self.teacher.task.name == 'reg':
                self.metric = 'mse'

            if self.teacher.task.name == 'multiclass':
                self.metric = 'crossentropy'

        if self.teacher.task.name == 'binary':
            if labels is not None:
                _data['__target__'] = labels
            else:
                _data['__target__'] = [1 if i > 0.5 else 0 for i in self.teacher.predict(_data).data]
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.fit_predict(_data, roles={'target': '__target__'})
            models_scores = DataFrame(columns=[self.metric], index=preds.keys())
            for name, pred in preds.items():
                models_scores.loc[name, self.metric] = roc_auc_score(data[self.teacher.reader.used_array_attrs['target']],
                                                                     pred.data[:, 0])
        if self.teacher.task.name == 'reg':
            if labels is not None:
                _data['__target__'] = labels
            else:
                _data['__target__'] = self.teacher.predict(_data).data[:, 0]
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.fit_predict(_data, roles={'target': '__target__'})
            models_scores = DataFrame(columns=[self.metric], index=preds.keys())
            for name, pred in preds.items():
                models_scores.loc[name, self.metric] = mean_squared_error(data[self.teacher.reader.used_array_attrs['target']],
                                                                          pred.data)
        if self.teacher.task.name == 'multiclass':
            if labels is not None:
                _data['__target__'] = labels
            else:
                _data['__target__'] = np.argmax(self.teacher.predict(_data).data, axis=1)
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.fit_predict(_data, roles={'target': '__target__'})
            models_scores = DataFrame(columns=[self.metric], index=preds.keys())
            for name, pred in preds.items():
                models_scores.loc[name, self.metric] = log_loss(data[self.teacher.reader.used_array_attrs['target']].values,
                                                                pred.data)

        self.models_scores = models_scores

        return self.students[np.argmax(models_scores[self.metric])]

    def eval_metrics(self, data, metrics=None):
        """Compute students' quality by calculating provided metrics.

        Args:
            data: dataset to compute metrics at.
            metrics: metrics to compute.

        Returns:
            pandas.DataFrame with metrics' values.
        """
        preds = dict()
        if not self._fitted:
            raise NotFittedError
        if not metrics:
            if self.teacher.task.name == 'binary':
                metrics = [roc_auc_score]
            if self.teacher.task.name == 'reg':
                metrics = [mean_squared_error]
            if self.teacher.task.name == 'multiclass':
                metrics = [log_loss]

        if self.teacher.task.name == 'binary':
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.predict(data).data[:, 0]
        if self.teacher.task.name == 'reg':
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.predict(data).data
        if self.teacher.task.name == 'multiclass':
            for estimator in self.students:
                preds[estimator.levels[0][0].ml_algos[0].name] = estimator.predict(data).data

        models_scores = DataFrame(columns=[metric.__name__ for metric in metrics],
                                  index=[estimator.levels[0][0].ml_algos[0].name for estimator in self.students])
        for metric in metrics:
            for name, pred in preds.items():
                if metric == accuracy_score:
                    pred = pred > 0.5
                models_scores.loc[name, metric.__name__] = metric(data[self.teacher.reader.used_array_attrs['target']], pred.data)

        return models_scores

    def distill_data(self, estimator, features, target):
        if self.task == 'reg':
            scoring = "neg_mean_squared_error"
        else:
            scoring = 'f1'
        rfeat = RFECV(estimator=estimator, step=1, scoring=scoring)
        rfeat.fit(features, target)
        return rfeat.transform(features)

max_stars_repo_name: Andrey-Nikitin/LightAutoML

max_issues_repo_name: Andrey-Nikitin/LightAutoML

max_forks_repo_name: Andrey-Nikitin/LightAutoML


text:
"""
Run the greedy grid search for BNMTF with the Exp priors on the Sanger dataset.
"""

import sys, os
project_location = os.path.dirname(__file__)+"/../../../../"
sys.path.append(project_location)

from BNMTF.code.models.bnmtf_vb_optimised import bnmtf_vb_optimised
from BNMTF.code.cross_validation.mask import compute_Ms, compute_folds
from BNMTF.code.cross_validation.greedy_search_bnmtf import GreedySearch
from BNMTF.data_drug_sensitivity.gdsc.load_data import load_gdsc

import numpy, matplotlib.pyplot as plt
import scipy.interpolate

##########

standardised = False #standardised Sanger or unstandardised
no_folds = 5

restarts = 5
iterations = 1000
I, J = 622,138
values_K = range(1,30+1)
values_L = range(1,30+1)

alpha, beta = 1., 1.
lambdaF = 1./10.
lambdaS = 1./10.
lambdaG = 1./10.
priors = { 'alpha':alpha, 'beta':beta, 'lambdaF':lambdaF, 'lambdaS':lambdaS, 'lambdaG':lambdaG }

initFG = 'kmeans'
initS = 'random'

classifier = bnmtf_vb_optimised

search_metric = 'AIC'

# Load in data
(_,X_min,M,_,_,_,_) = load_gdsc(standardised=standardised)

folds_test = compute_folds(I,J,no_folds,M)
folds_training = compute_Ms(folds_test)
(M_train,M_test) = (folds_training[0],folds_test[0])

# Run the line search
priors = { 'alpha':alpha, 'beta':beta, 'lambdaF':lambdaF, 'lambdaS':lambdaS, 'lambdaG':lambdaG }
greedy_search = GreedySearch(classifier,values_K,values_L,X_min,M,priors,initS,initFG,iterations,restarts=restarts)
greedy_search.search(search_metric)

# Plot the performances of all metrics
metrics = ['loglikelihood', 'BIC', 'AIC', 'MSE']
for metric in metrics:
    # Make three lists of indices X,Y,Z (K,L,metric)
    KLvalues = numpy.array(greedy_search.all_values(metric))
    (list_values_K,list_values_L,values) = zip(*KLvalues)

    # Set up a regular grid of interpolation points
    Ki, Li = (numpy.linspace(min(list_values_K), max(list_values_K), 100),
              numpy.linspace(min(list_values_L), max(list_values_L), 100))
    Ki, Li = numpy.meshgrid(Ki, Li)

    # Interpolate
    rbf = scipy.interpolate.Rbf(list_values_K, list_values_L, values, function='linear')
    values_i = rbf(Ki, Li)

    # Plot
    plt.figure()
    plt.imshow(values_i, cmap='jet_r',
               vmin=min(values), vmax=max(values), origin='lower',
               extent=[min(list_values_K)-1, max(list_values_K)+1, min(list_values_L)-1, max(list_values_L)+1])
    plt.scatter(list_values_K, list_values_L, c=values, cmap='jet_r')
    plt.colorbar()
    plt.title("Metric: %s." % metric)
    plt.xlabel("K")
    plt.ylabel("L")
    plt.show()

    # Print the best value
    best_K,best_L = greedy_search.best_value(metric)
    print "Best K,L for metric %s: %s,%s." % (metric,best_K,best_L)

# Also print out all values in a dictionary
all_values = {}
for metric in metrics:
    (_,_,values) = zip(*numpy.array(greedy_search.all_values(metric)))
    all_values[metric] = list(values)

print "all_values = %s \nlist_values_K=%s \nlist_values_L=%s" % \
    (all_values,list(list_values_K),list(list_values_L))

'''
all_values = {'MSE': [3.0272042551947203, 3.027204256305112, 3.0272042923576148, 2.5914654932112464, 2.5918836849320201, 2.5914602381010914, 2.3493739958858635, 2.3511225674996381, 2.3584324978814539, 2.1868222893761833, 2.1911559705091568, 2.2016668628098452, 2.0510257720785683, 2.0546897432717603, 2.0586496735360251, 2.0826309185454925], 'loglikelihood': [-138379.73430838491, -138379.74014614287, -138380.57362950334, -132935.31284836732, -132949.35927074254, -132936.87960196467, -129506.33023969264, -129543.6258747291, -129641.90589727147, -127005.00858615834, -127072.80234078577, -127243.67779180428, -124768.25065830135, -124835.58530247367, -124903.988439383, -125310.14633691005], 'AIC': [278283.46861676982, 279529.48029228573, 278565.14725900668, 268922.62569673464, 270198.71854148508, 269207.75920392934, 263596.66047938529, 264921.2517494582, 264151.81179454294, 260130.01717231667, 261517.60468157154, 260893.35558360856, 257196.50131660269, 258585.17060494734, 257755.976878766, 259824.29267382011], 'BIC': [285262.09744653094, 292213.73348023742, 286825.93886588927, 282898.19996735867, 289889.07547585049, 284474.6545572257, 284587.4968019739, 291636.02904133906, 286443.12750535476, 288154.43215797161, 295275.11894206965, 290227.40826303756, 292272.81157642574, 299403.7384451644, 294151.08313791396, 301970.81481891294]}
list_values_K=[1.0, 2.0, 1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0, 4.0, 5.0, 6.0, 5.0, 6.0]
list_values_L=[1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 6.0, 6.0]
'''

max_stars_repo_name: ThomasBrouwer/BNMTF

max_issues_repo_name: ThomasBrouwer/BNMTF

max_forks_repo_name: ThomasBrouwer/BNMTF


text:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn
from sklearn.preprocessing import Imputer
from keras.models import Sequential
from keras.layers import Dense
import numpy
from keras.models import Sequential
from keras.layers import Dense

#loading the model
df=pd.read_csv('maternal_mortality.csv')

df.fillna(df.mean(), inplace=True)

print(df.head)

df.to_csv('new.csv')

#print shape of the data
#print(df.shape)

numpy.random.seed(5)

X = df[:,0:7]
Y = df[:,7]

#model layers
model = Sequential()
model.add(Dense(12, input_dim=7, activation='relu'))
model.add(Dense(8, activation='relu'))
model.add(Dense(1, activation='softmax'))

#compiling the model

model.compile(Adam(lr=0.0001), loss='binary_crossentropy', metrics=['accuracy'])

#fitting the model
model.fit(X, Y, epochs=1000, batch_size=10)

#evaluate the model

scores = model.evaluate(X,Y)
print("\n%s: %.2f%%" % (model.metrics_names[1], scores[1]*100))