text:
import numpy as np
import matplotlib.pyplot as plt
import Objectes as ob
import Sistemes as sis
#Exemple Camp gravitcional en estatic
"""
terra = ob.cos(5.972*10E24,10E12,0,5519,3,escala=10E3,color="blue")
sol = ob.cos(1.989 *10E30, 0,0,9400,escala= 10E2,color="yellow")

sistema_solar = ob.static_sistem(limits=[-(1.1)*10E12,(1.1)*10E12,-(1.1)*10E12,(1.1)*10E12])
sistema_solar.add(terra)
sistema_solar.add(sol)
plt.show()
"""

#Exemple Particula amb una velocitat donada en un camp gravitacional
# Atencio valors irrealistes nomes per ensenyar el funcionament

Gxmin = -100
Gxmax = 100
Gymin = -100
Gymax = 100

sys= sis.static_sistem([Gxmin,Gxmax,Gymin,Gymax])

#sys.add(ob.cos(3*10E16,0,0,escala=10E-7))
#sys.add(ob.cos(2*10E17,4.5,-2,escala=10E-7))
sys.add(ob.cos(2*10E16,-15,-15,escala=10E-6))

sys.field_plot(np.meshgrid(np.linspace(Gxmin, Gxmax, 2000), np.linspace(Gymin, Gymax, 2000)))
p =ob.particula_accelerada(sys, -13, 15, [],1,vx_0=0.01,vy_0=-0.01,Tf=0.3,Te=0.0001,color="red")

k =ob.particula_accelerada(sys, -13, 15, [[0,0.1,10E3,-10E2,],[0.1,2.1,-10E2,-10E3,],[10.1,5,-10E20,-10E20,],[10222,5,-10E20,-10E20,]],1,vx_0=0.01,vy_0=-0.01,Tf=20,Te=0.0001)

p.trajectoria()
k.trajectoria()
plt.show()


text:
import numpy as np
import time
import matplotlib.pyplot as plt
import tensorflow as tf
import copy
from tqdm import tqdm
from helper_funcs import frames_to_mp4
from policy_random import Policy_Random

from gym.spaces import Box
from cpprb import ReplayBuffer
from tf2rl.experiments.utils import save_path
from tf2rl.envs.utils import is_discrete, get_act_dim
from tf2rl.misc.get_replay_buffer import get_replay_buffer, get_default_rb_dict
from tf2rl.misc.prepare_output_dir import prepare_output_dir
from tf2rl.misc.initialize_logger import initialize_logger
from tf2rl.envs.normalizer import EmpiricalNormalizer
from tf2rl.misc.discount_cumsum import discount_cumsum
import logging
import wandb

from tf2rl.algos.ppo import PPO
from tf2rl.algos.sac import SAC
from tf2rl.algos.td3 import TD3

class CollectSamples(object):

	def __init__(self, env, policy, visualize_rollouts, which_agent, dt_steps, dt_from_xml, follow_trajectories):
		self.env = env
		self.policy = policy
		self.visualize_at_all = visualize_rollouts
		self.which_agent = which_agent

		self.low = self.env.observation_space.low
		self.high = self.env.observation_space.high
		self.shape = self.env.observation_space.shape

		self.use_low = self.low + (self.high-self.low)/3.0  ##invalid value encountered in substract
		self.use_high = self.high - (self.high-self.low)/3.0

		self.dt_steps = dt_steps
		self.dt_from_xml = dt_from_xml

		self.follow_trajectories = follow_trajectories

	def collect_samples(self, num_rollouts, steps_per_rollout):
		observations_list = []
		actions_list = []
		starting_states_list=[]
		rewards_list = []
		visualization_frequency = 10
		pbar = tqdm(total=num_rollouts, ascii=True)
		for rollout_number in range(num_rollouts):
			org_observation = self.env.reset()
			min_len=0

			observation = org_observation
			min_len = 2
			if (steps_per_rollout==23333):
				steps_per_rollout = 100
				min_len = 50
			starting_state = np.copy(observation)
			observations = np.zeros((1,2))
			trails_counter = 0
			while not (np.shape(observations)[0] >= min_len):
				observations, actions, reward_for_rollout = self.perform_rollout(observation, steps_per_rollout,
																		rollout_number, visualization_frequency)
				trails_counter += 1
				if (trails_counter > 100):
					if (min_len == 5):
						min_len -= 1
					else:
						min_len -= 5
					trails_counter = 0

			rewards_list.append(reward_for_rollout)
			observations= np.array(observations)
			actions= np.array(actions)
			observations_list.append(observations)
			actions_list.append(actions)
			starting_states_list.append(starting_state)
			pbar.update(1)
		pbar.close()

		#return list of length = num rollouts
		#each entry of that list contains one rollout
		#each entry is [steps_per_rollout x statespace_dim] or [steps_per_rollout x actionspace_dim]
		print('observations_list, actions_list, starting_states_list, rewards_list')  #(60, xxx, 29) (60, xxx, 8) (60, xxx) (60,)

		# print(np.shape(observations_list), np.shape(actions_list), np.shape(starting_states_list),np.shape(rewards_list))
		# print(observations_list[0][0])
		# print(starting_states_list[0])
		return observations_list, actions_list, starting_states_list, rewards_list

	def perform_rollout(self, observation, steps_per_rollout, rollout_number, visualization_frequency):
		observations = []
		actions = []
		visualize = False
		reward_for_rollout = 0
		self.env.reset()
		if((rollout_number%visualization_frequency)==0):
			#print("currently performing rollout #", rollout_number)
			if(self.visualize_at_all):
				all_states=[]
				print ("---- visualizing a rollout ----")
				visualize=True

		for step_num in range(steps_per_rollout):
			action, _ = self.policy.get_action(observation)
			#print("[MESSAGE FROM COLLECT_SAMPLES]  The action just took: \n", action)

			observations.append(observation)
			actions.append(action)

			org_next_observation, reward, terminal, env_info = self.env.step(action)
			try:
				cost = env_info['cost']
			except (TypeError, KeyError, IndexError):
				cost = 0

			next_observation = org_next_observation
			reward_for_rollout += reward

			observation = np.copy(next_observation)

			if terminal:
				#print(f"Had to stop rollout after {step_num} steps because terminal state was reached.")
				break

			if(visualize):
				if(self.which_agent==0):
					curr_state = self.env.render()
					all_states.append(np.expand_dims(curr_state, axis=0))
				else:
					self.env.render()
					time.sleep(self.dt_steps*self.dt_from_xml)

		if(visualize and (self.which_agent==0)):
			all_states= np.concatenate(all_states, axis=0)
			plt.plot(all_states[:,0], all_states[:,1], 'r')
			plt.show()
		return observations, actions, reward_for_rollout

class CollectSamples_random(object):

	def __init__(self, env, max_steps, _output_dir=None, _episode_max_steps=100, gpu=0, wandb_on=False):

		self.env = env
		self._policy = Policy_Random(env)
		self.visualize_at_all = False

		self.max_steps = max_steps

		self.wandb_on = wandb_on

		self.low = self.env.observation_space.low
		self.high = self.env.observation_space.high
		self.shape = self.env.observation_space.shape

		self.use_low = self.low + (self.high-self.low)/3.0  ##invalid value encountered in substract
		self.use_high = self.high - (self.high-self.low)/3.0

		self.global_total_steps = 0
		self.global_num_episodes = 0
		self.global_cost = 0

	def collect_samples(self):
		observations_list = []
		actions_list = []
		starting_states_list=[]
		rewards_list = []
		visualization_frequency = 10
		pbar = tqdm(total=self.max_steps, ascii=True)
		step_counter = 0
		while(step_counter < self.max_steps):
			org_observation = self.env.reset()
			min_len=0

			observation = org_observation

			starting_state = np.copy(observation)

			observations, actions, reward_for_rollout, episode_cost = self.perform_rollout(observation)
			self.global_num_episodes += 1
			self.global_cost += episode_cost
			step_counter += len(actions)
			cost_rate = self.global_cost/self.global_total_steps

			if self.wandb_on:
				wandb.log({'Training_Return': reward_for_rollout, 'Training_Cost': episode_cost, 'Cost_Rate': cost_rate}, step=self.global_num_episodes)
				# print("DEBUG! LOGGING ...")

			tf.summary.scalar(name="Common/Training_Return", data=reward_for_rollout, step=self.global_total_steps)
			tf.summary.scalar(name="Common/Training_Cost", data=episode_cost, step=self.global_total_steps)
			tf.summary.scalar(name="Common/Cost_Rate", data=cost_rate, step=self.global_total_steps)

			if not (step_counter>self.max_steps): # otherwise ignore this rollout

				rewards_list.append(reward_for_rollout)
				observations= np.array(observations)
				actions= np.array(actions)
				observations_list.append(observations)
				actions_list.append(actions)
				starting_states_list.append(starting_state)

				pbar.update(len(actions))

		pbar.close()

		#return list of length = num rollouts
		#each entry of that list contains one rollout
		#each entry is [steps_per_rollout x statespace_dim] or [steps_per_rollout x actionspace_dim]
		#print('observations_list, actions_list, starting_states_list, rewards_list')  #(60, xxx, 29) (60, xxx, 8) (60, xxx) (60,)

		# print(np.shape(observations_list), np.shape(actions_list), np.shape(starting_states_list),np.shape(rewards_list))
		# print(observations_list[0][0])
		# print(starting_states_list[0])
		return observations_list, actions_list, starting_states_list, rewards_list

	def perform_rollout(self, observation, steps_per_rollout=1000, rollout_number=1, visualization_frequency=1):
		observations = []
		actions = []
		visualize = False
		reward_for_rollout = 0
		episode_cost = 0
		self.env.reset()
		if((rollout_number%visualization_frequency)==0):
			#print("currently performing rollout #", rollout_number)
			if(self.visualize_at_all):
				all_states=[]
				print ("---- visualizing a rollout ----")
				visualize=True

		for step_num in range(steps_per_rollout):
			action, _ = self._policy.get_action(observation)
			#print("[MESSAGE FROM COLLECT_SAMPLES]  The action just took: \n", action)

			observations.append(observation)
			actions.append(action)

			org_next_observation, reward, terminal, env_info = self.env.step(action)
			try:
				cost = env_info['cost']
			except (TypeError, KeyError, IndexError):
				cost = 0

			self.global_total_steps += 1

			next_observation = org_next_observation
			reward_for_rollout += reward
			episode_cost += cost

			observation = np.copy(next_observation)

			if terminal:
				#print(f"Had to stop rollout after {step_num} steps because terminal state was reached.")
				break

			if(visualize):
				if(self.which_agent==0):
					curr_state = self.env.render()
					all_states.append(np.expand_dims(curr_state, axis=0))
				else:
					self.env.render()
					time.sleep(self.dt_steps*self.dt_from_xml)

		if(visualize and (self.which_agent==0)):
			all_states= np.concatenate(all_states, axis=0)
			plt.plot(all_states[:,0], all_states[:,1], 'r')
			plt.show()
		return observations, actions, reward_for_rollout, episode_cost

class CollectSamples_SAC(CollectSamples_random):

	def __init__(self, env, max_steps, _output_dir, _episode_max_steps=100, gpu=0, wandb_on=False, *args, **kwargs):
		super().__init__(env, max_steps, _output_dir, _episode_max_steps, gpu, wandb_on, *args, **kwargs)

		#self._set_from_args(args)
		# experiment settings

		self._max_steps = max_steps
		self._episode_max_steps = _episode_max_steps
		self._n_experiments = 1
		self._show_progress = False
		self._save_model_interval = int(400)
		self._save_summary_interval = int(1e3)
		self._normalize_obs = False
		self._logdir = _output_dir
		self._model_dir = _output_dir + '/../checkpoints/policy'
		# replay buffer
		self._use_prioritized_rb = False
		self._use_nstep_rb = False
		self._n_step = 4
		# test settings
		self._evaluate = False
		self._test_interval = int(1e4)  #TODO
		self._show_test_progress = False
		self._test_episodes = 3
		self._save_test_path = False
		self._save_test_movie = False
		self._show_test_images = False
		self.wandb_on = wandb_on
		self.evaluation_interval = 2#20 #TODO

		self._policy = SAC(
						state_shape=env.observation_space.shape,
						action_dim=env.action_space.high.size,
						gpu=gpu,
						max_action=env.action_space.high[0],
						batch_size=100,
						n_warmup=1000)
		self._env = env
		self._test_env = copy.copy(self._env)
		if self._normalize_obs:
			assert isinstance(env.observation_space, Box)
			self._obs_normalizer = EmpiricalNormalizer(
				shape=env.observation_space.shape)

		# prepare log directory
		self._output_dir = self._logdir
		self.logger = initialize_logger(
			logging_level=logging.getLevelName('INFO'),
			output_dir=self._output_dir)

		if self._evaluate:
			assert self._model_dir is not None
		self._set_check_point(self._model_dir)

		self.replay_buffer = get_replay_buffer(
			self._policy, self._env, self._use_prioritized_rb,
			self._use_nstep_rb, self._n_step)

		# prepare TensorBoard output
		# self.writer = tf.summary.create_file_writer(self._output_dir)
		# self.writer.set_as_default()

	def collect_samples(self):#observations_list, actions_list, starting_states_list, rewards_list (60, xxx, 29) (60, xxx, 8) (60, xxx) (60,)
		observations_list = []
		actions_list = []
		starting_states_list=[]
		rewards_list = []

		tf.summary.experimental.set_step(self.global_total_steps)
		episode_steps = 0
		episode_return = 0
		episode_cost = 0
		episode_start_time = time.perf_counter()
		n_episode = 0

		# self.replay_buffer = get_replay_buffer(
		# 	self._policy, self._env, self._use_prioritized_rb,
		# 	self._use_nstep_rb, self._n_step)

		obs = self._env.reset()
		observations = []
		actions = []
		starting_states_list.append(obs)

		while self.global_total_steps < self._max_steps:
			if self.global_total_steps < self._policy.n_warmup:
				action = self._env.action_space.sample()
			else:
				action = self._policy.get_action(obs)

			next_obs, reward, done, env_info = self._env.step(action)
			try:
				cost = env_info['cost']
			except (TypeError, KeyError, IndexError):
				cost = 0

			observations.append(obs)
			actions.append(action)
			rewards_list.append(reward)

			if self._show_progress:
				self._env.render()
			episode_steps += 1
			episode_return += reward
			episode_cost += cost
			self.global_total_steps += 1
			tf.summary.experimental.set_step(self.global_total_steps)

			done_flag = done
			if hasattr(self._env, "_max_episode_steps") and \
					episode_steps == self._env._max_episode_steps:
				done_flag = False
			self.replay_buffer.add(obs=obs, act=action,
							  next_obs=next_obs, rew=reward, done=done_flag)
			obs = next_obs

			if done or episode_steps == self._episode_max_steps:
				obs = self._env.reset()

				n_episode += 1
				self.global_num_episodes += 1
				self.global_cost += episode_cost
				cost_rate = self.global_cost/self.global_total_steps

				fps = episode_steps / (time.perf_counter() - episode_start_time)
				self.logger.info("Total Epi: {0: 5} Steps: {1: 7} Episode Steps: {2: 5} Return: {3: 7.4f} Cost: {4: 5.4f} FPS: {5:5.2f}".format(
					n_episode, self.global_total_steps, episode_steps, episode_return, episode_cost, fps))
				tf.summary.scalar(name="Common/Training_Return", data=episode_return, step=self.global_total_steps)
				tf.summary.scalar(name="Common/Training_Cost", data=episode_cost, step=self.global_total_steps)
				tf.summary.scalar(name="Common/Cost_Rate", data=cost_rate, step=self.global_total_steps)
				tf.summary.scalar(name="MFAgent/FPS", data=fps, step=self.global_total_steps)

				if self.wandb_on:
					wandb.log({'Training_Return': episode_return, 'Training_Cost': episode_cost, 'Cost_Rate': cost_rate}, step=self.global_num_episodes)

				rewards_list.append(episode_return)
				observations_list.append(np.array(observations))
				actions_list.append(np.array(actions))
				starting_states_list.append(obs)

				#print('=============DEBUG:actions_list',np.shape(actions_list),'=============')
				#print(actions_list)

				observations = []
				actions = []
				episode_steps = 0
				episode_return = 0
				episode_cost = 0
				episode_start_time = time.perf_counter()

			if self.global_total_steps < self._policy.n_warmup:
				continue

			if self.global_total_steps % self._policy.update_interval == 0:
				samples = self.replay_buffer.sample(self._policy.batch_size)
				with tf.summary.record_if(self.global_total_steps % self._save_summary_interval == 0):
					self._policy.train(
						samples["obs"], samples["act"], samples["next_obs"],
						samples["rew"], np.array(samples["done"], dtype=np.float32),
						None if not self._use_prioritized_rb else samples["weights"])
				if self._use_prioritized_rb:
					td_error = self._policy.compute_td_error(
						samples["obs"], samples["act"], samples["next_obs"],
						samples["rew"], np.array(samples["done"], dtype=np.float32))
					self.replay_buffer.update_priorities(
						samples["indexes"], np.abs(td_error) + 1e-6)

			if self.global_total_steps % self._test_interval == 0:
				# print('Evaluating the MFAgent ...')
				time1 = time.time()
				self.evaluate_policy()
			if False:
				avg_test_return = self.evaluate_policy(self.global_total_steps)
				self.logger.info("Evaluation Total Steps: {0: 7} Average Reward {1: 5.4f} over {2: 2} episodes".format(
					total_steps, avg_test_return, self._test_episodes))
				tf.summary.scalar(
					name="Common/average_test_return", data=avg_test_return)
				tf.summary.scalar(name="Common/fps", data=fps)
				self.writer.flush()

			if self.global_total_steps % self._save_model_interval == 0:
				self.checkpoint_manager.save()

		tf.summary.flush()

			#return list of length = num rollouts
			#each entry of that list contains one rollout
			#each entry is [steps_per_rollout x statespace_dim] or [steps_per_rollout x actionspace_dim]
		return observations_list, actions_list, starting_states_list, rewards_list

	def train(self):
		samples = self.replay_buffer.sample(self._policy.batch_size)
		with tf.summary.record_if(self.global_total_steps % self._save_summary_interval == 0):
			self._policy.train(
				samples["obs"], samples["act"], samples["next_obs"],
				samples["rew"], np.array(samples["done"], dtype=np.float32),
				None if not self._use_prioritized_rb else samples["weights"])
		if self._use_prioritized_rb:
			td_error = self._policy.compute_td_error(
				samples["obs"], samples["act"], samples["next_obs"],
				samples["rew"], np.array(samples["done"], dtype=np.float32))
			self.replay_buffer.update_priorities(
				samples["indexes"], np.abs(td_error) + 1e-6)

	def _set_check_point(self, model_dir):
		# Save and restore model
		self._checkpoint = tf.train.Checkpoint(policy=self._policy)
		self.checkpoint_manager = tf.train.CheckpointManager(
			self._checkpoint, directory=model_dir, max_to_keep=5)

		# if model_dir is not None:
		# 	assert os.path.isdir(model_dir)
		# 	self._latest_path_ckpt = tf.train.latest_checkpoint(model_dir)
		# 	self._checkpoint.restore(self._latest_path_ckpt)
		# 	self.logger.info("Restored {}".format(self._latest_path_ckpt))

	def evaluate_policy(self, total_steps=None):
		tf.summary.experimental.set_step(self.global_total_steps)
		if self._normalize_obs:
			self._test_env.normalizer.set_params(
				*self._env.normalizer.get_params())
		avg_test_return = 0.
		avg_test_cost = 0.
		if self._save_test_path:
			replay_buffer = get_replay_buffer(
				self._policy, self._test_env, size=self._episode_max_steps)
		for i in range(self._test_episodes):
			timer0 = time.time()
			episode_return = 0.
			episode_cost = 0.
			frames = []
			obs = self._test_env.reset()
			# print(f'LOOP{i}-PERIOD1: ', time.time() - timer0)
			timer0 = time.time()
			for _ in range(self._episode_max_steps):
				action = self._policy.get_action(obs, test=True)
				next_obs, reward, done, env_info = self._test_env.step(action)
				try:
					cost = env_info['cost']
				except (TypeError, KeyError, IndexError):
					cost = 0
				if self._save_test_path:
					replay_buffer.add(obs=obs, act=action,
									  next_obs=next_obs, rew=reward, done=done)

				if self._save_test_movie:
					frames.append(self._test_env.render(mode='rgb_array'))
				elif self._show_test_progress:
					self._test_env.render()
				episode_return += reward
				episode_cost += cost
				obs = next_obs
				if done:
					break
			# print(f'LOOP{i}-PERIOD2: ', time.time() - timer0)
			prefix = "step_{0:08d}_epi_{1:02d}_return_{2:010.4f}_cost_{2:010.4f}".format(
				self.global_total_steps, i, episode_return, episode_cost)
			if self._save_test_path:
				save_path(replay_buffer._encode_sample(np.arange(self._episode_max_steps)),
						  os.path.join(self._output_dir, prefix + ".pkl"))
				replay_buffer.clear()
			if self._save_test_movie:
				# frames_to_gif(frames, prefix, self._output_dir)
				frames_to_mp4(frames, prefix, self._output_dir)
			avg_test_return += episode_return
			avg_test_cost += episode_cost
		if self._show_test_images:
			images = tf.cast(
				tf.expand_dims(np.array(obs).transpose(2, 0, 1), axis=3),
				tf.uint8)
			tf.summary.image('train/input_img', images,)

		avg_test_return = avg_test_return / self._test_episodes
		if self.wandb_on:
			wandb.log({'Evaluation_Return': avg_test_return, 'Evaluation_Cost': avg_test_cost}, step=self.global_num_episodes)

		tf.summary.scalar(name="MFAgent/Evaluation_Return", data=avg_test_return, step=self.global_total_steps)
		tf.summary.scalar(name="MFAgent/Evaluation_Cost", data=avg_test_cost, step=self.global_total_steps)

		return avg_test_return, avg_test_cost

	def load_checkpoints(self, global_total_steps, global_num_episodes, global_cost):
		self.global_total_steps = global_total_steps
		self.global_num_episodes = global_num_episodes
		self.global_cost = global_cost

		self._latest_path_ckpt = tf.train.latest_checkpoint(self._model_dir)
		self._checkpoint.restore(self._latest_path_ckpt)
		self.logger.info("Restored {}".format(self._latest_path_ckpt))

		tf.summary.experimental.set_step(self.global_total_steps)

class CollectSamples_PPO(CollectSamples_SAC):

	def __init__(self, env, max_steps, _output_dir, _episode_max_steps=100, gpu=0, wandb_on=False, *args, **kwargs):
		super(CollectSamples_PPO, self).__init__(env, max_steps, _output_dir, _episode_max_steps, gpu, wandb_on, *args, **kwargs)

		# experiment settings
		self._max_steps = max_steps
		self._episode_max_steps = _episode_max_steps
		self._n_experiments = 1
		self._show_progress = False
		self._save_model_interval = int(1e4)
		self._save_summary_interval = int(1e3)
		self._normalize_obs = False
		self._logdir = _output_dir
		# self._model_dir = None
		# replay buffer
		self._use_prioritized_rb = False
		self._use_nstep_rb = False
		self._n_step = 4
		# test settings
		self._evaluate = False
		# self._test_interval = int(1e4)
		self._show_test_progress = False
		# self._test_episodes = 5
		self._save_test_path = False
		self._save_test_movie = False
		self._show_test_images = False
		# PPO settings
		self._enable_gae = False
		self._normalize_adv = False
		self._horizon = 2048

		self._output_dir = _output_dir

		self._policy = PPO(
						state_shape=env.observation_space.shape,
						action_dim=get_act_dim(env.action_space),
						is_discrete=is_discrete(env.action_space),
						max_action=None if is_discrete(
							env.action_space) else env.action_space.high[0],
						batch_size=64,
						actor_units=(64, 64),
						critic_units=(64, 64),
						n_epoch=10,
						lr_actor=3e-4,
						lr_critic=3e-4,
						hidden_activation_actor="tanh",
						hidden_activation_critic="tanh",
						discount=0.99,
						lam=0.95,
						entropy_coef=0.,
						horizon=self._horizon,
						normalize_adv=self._normalize_adv,
						enable_gae=self._enable_gae,
						gpu=gpu)

		# Prepare buffer
		self.replay_buffer = get_replay_buffer(
			self._policy, self._env)
		kwargs_local_buf = get_default_rb_dict(
			size=self._policy.horizon, env=self._env)
		kwargs_local_buf["env_dict"]["logp"] = {}
		kwargs_local_buf["env_dict"]["val"] = {}
		if is_discrete(self._env.action_space):
			kwargs_local_buf["env_dict"]["act"]["dtype"] = np.int32
		self.local_buffer = ReplayBuffer(**kwargs_local_buf)

	def collect_samples(self):

		observations_list = []
		actions_list = []
		starting_states_list=[]
		rewards_list = []

		# # Prepare buffer
		# self.replay_buffer = get_replay_buffer(
		# 	self._policy, self._env)
		# kwargs_local_buf = get_default_rb_dict(
		# 	size=self._policy.horizon, env=self._env)
		# kwargs_local_buf["env_dict"]["logp"] = {}
		# kwargs_local_buf["env_dict"]["val"] = {}
		# if is_discrete(self._env.action_space):
		# 	kwargs_local_buf["env_dict"]["act"]["dtype"] = np.int32
		# self.local_buffer = ReplayBuffer(**kwargs_local_buf)

		episode_steps = 0
		episode_return = 0
		episode_cost = 0
		episode_start_time = time.time()
		total_steps = np.array(0, dtype=np.int32)
		# self.global_total_steps = total_steps
		n_episode = 0
		obs = self._env.reset()
		observations = []
		actions = []
		starting_states_list.append(obs)

		tf.summary.experimental.set_step(self.global_total_steps)
		while self.global_total_steps < self._max_steps:
			# Collect samples
			for _ in range(self._policy.horizon):
				if self._normalize_obs:
					obs = self._obs_normalizer(obs, update=False)
				act, logp, val = self._policy.get_action_and_val(obs)
				next_obs, reward, done, env_info = self._env.step(act)
				try:
					cost = env_info['cost']
				except (TypeError, KeyError, IndexError):
					cost = 0

				observations.append(obs)
				actions.append(act)

				if self._show_progress:
					self._env.render()

				episode_steps += 1
				self.global_total_steps += 1
				# self.global_total_steps = total_steps
				episode_return += reward
				episode_cost += cost

				done_flag = done
				if hasattr(self._env, "_max_episode_steps") and \
						episode_steps == self._env._max_episode_steps:
					done_flag = False
				self.local_buffer.add(
					obs=obs, act=act, next_obs=next_obs,
					rew=reward, done=done_flag, logp=logp, val=val)
				obs = next_obs

				if done or episode_steps == self._episode_max_steps:
					tf.summary.experimental.set_step(self.global_total_steps)
					self.finish_horizon()
					obs = self._env.reset()
					n_episode += 1
					self.global_num_episodes += 1
					self.global_cost += episode_cost
					cost_rate = self.global_cost/self.global_total_steps
					fps = episode_steps / (time.time() - episode_start_time)
					self.logger.info(
						"Total Epi: {0: 5} Steps: {1: 7} Episode Steps: {2: 5} Return: {3: 7.4f} Cost: {4: 5.4f} FPS: {5:5.2f}".format(
							n_episode, int(self.global_total_steps), episode_steps, episode_return, episode_cost, fps))
					# tf.summary.scalar(name="Common/training_return", data=episode_return)

					tf.summary.scalar(name="Common/Training_Return", data=episode_return, step=self.global_total_steps)
					tf.summary.scalar(name="Common/Training_Cost", data=episode_cost, step=self.global_total_steps)
					tf.summary.scalar(name="Common/Cost_Rate", data=cost_rate, step=self.global_total_steps)
					tf.summary.scalar(name="MFAgent/FPS", data=fps, step=self.global_total_steps)

					if self.wandb_on:
						wandb.log({'Training_Return': episode_return, 'Training_Cost': episode_cost, 'Cost_Rate': cost_rate}, step=self.global_num_episodes)

					starting_states_list.append(obs)
					rewards_list.append(episode_return)
					observations_list.append(np.array(observations))
					actions_list.append(np.array(actions))

					#print('=============DEBUG:actions_list',np.shape(actions_list),'=============')
					#print(actions_list)

					observations = []
					actions = []
					episode_steps = 0
					episode_return = 0
					episode_start_time = time.time()

				if self.global_total_steps % self._test_interval == 0:
					# print('Evaluating the MFAgent ...')
					self.evaluate_policy()
				#if total_steps % self._test_interval == 0:
				if False:
					avg_test_return = self.evaluate_policy(self.global_total_steps)
					self.logger.info("Evaluation Total Steps: {0: 7} Average Reward {1: 5.4f} over {2: 2} episodes".format(
						self.global_total_steps, avg_test_return, self._test_episodes))
					tf.summary.scalar(
						name="Common/average_test_return", data=avg_test_return)
					self.writer.flush()

				if self.global_total_steps % self._save_model_interval == 0:
					self.checkpoint_manager.save()

			self.finish_horizon(last_val=val)

			tf.summary.experimental.set_step(self.global_total_steps)

			# Train actor critic
			if self._policy.normalize_adv:
				samples = self.replay_buffer._encode_sample(np.arange(self._policy.horizon))
				mean_adv = np.mean(samples["adv"])
				std_adv = np.std(samples["adv"])
				# Update normalizer
				if self._normalize_obs:
					self._obs_normalizer.experience(samples["obs"])
			with tf.summary.record_if(self.global_total_steps % self._save_summary_interval == 0):
				for _ in range(self._policy.n_epoch):
					samples = self.replay_buffer._encode_sample(
						np.random.permutation(self._policy.horizon))
					if self._normalize_obs:
						samples["obs"] = self._obs_normalizer(samples["obs"], update=False)
					if self._policy.normalize_adv:
						adv = (samples["adv"] - mean_adv) / (std_adv + 1e-8)
					else:
						adv = samples["adv"]
					for idx in range(int(self._policy.horizon / self._policy.batch_size)):
						target = slice(idx * self._policy.batch_size,
									   (idx + 1) * self._policy.batch_size)
						self._policy.train(
							states=samples["obs"][target],
							actions=samples["act"][target],
							advantages=adv[target],
							logp_olds=samples["logp"][target],
							returns=samples["ret"][target])

		tf.summary.flush()

		#print('====================DEGUG:observations_list',np.shape(observations_list),'====================')
		#print(observations_list)
		#print('====================DEGUG:actions_list',np.shape(actions_list),'====================')

		return observations_list, actions_list, starting_states_list, rewards_list

	def finish_horizon(self, last_val=0):
		samples = self.local_buffer._encode_sample(
			np.arange(self.local_buffer.get_stored_size()))
		rews = np.append(samples["rew"], last_val)
		vals = np.append(samples["val"], last_val)

		# GAE-Lambda advantage calculation
		deltas = rews[:-1] + self._policy.discount * vals[1:] - vals[:-1]
		if self._policy.enable_gae:
			advs = discount_cumsum(
				deltas, self._policy.discount * self._policy.lam)
		else:
			advs = deltas

		# Rewards-to-go, to be targets for the value function
		rets = discount_cumsum(rews, self._policy.discount)[:-1]
		self.replay_buffer.add(
			obs=samples["obs"], act=samples["act"], done=samples["done"],
			ret=rets, adv=advs, logp=np.squeeze(samples["logp"]))
		self.local_buffer.clear()

	def evaluate_policy(self, total_steps=None):
		avg_test_return = 0.
		avg_test_cost = 0.
		if self._save_test_path:
			replay_buffer = get_replay_buffer(
				self._policy, self._test_env, size=self._episode_max_steps)
		for i in range(self._test_episodes):
			episode_return = 0.
			episode_cost = 0.
			frames = []
			obs = self._test_env.reset()
			for _ in range(self._episode_max_steps):
				if self._normalize_obs:
					obs = self._obs_normalizer(obs, update=False)
				act, _ = self._policy.get_action(obs, test=True)
				act = act if not hasattr(self._env.action_space, "high") else \
					np.clip(act, self._env.action_space.low, self._env.action_space.high)
				next_obs, reward, done, env_info = self._test_env.step(act)
				try:
					cost = env_info['cost']
				except (TypeError, KeyError, IndexError):
					cost = 0
				if self._save_test_path:
					replay_buffer.add(
						obs=obs, act=act, next_obs=next_obs,
						rew=reward, done=done)

				if self._save_test_movie:
					frames.append(self._test_env.render(mode='rgb_array'))
				elif self._show_test_progress:
					self._test_env.render()
				episode_return += reward
				episode_cost += cost
				obs = next_obs
				if done:
					break
			prefix = "step_{0:08d}_epi_{1:02d}_return_{2:010.4f}_cost_{2:010.4f}".format(
				self.global_total_steps, i, episode_return, episode_cost)
			if self._save_test_path:
				save_path(replay_buffer.sample(self._episode_max_steps),
						  os.path.join(self._output_dir, prefix + ".pkl"))
				replay_buffer.clear()
			if self._save_test_movie:
				# frames_to_gif(frames, prefix, self._output_dir)
				frames_to_mp4(frames, prefix, self._output_dir)
			avg_test_return += episode_return
			avg_test_cost += episode_cost
		if self._show_test_images:
			images = tf.cast(
				tf.expand_dims(np.array(obs).transpose(2, 0, 1), axis=3),
				tf.uint8)
			tf.summary.image('train/input_img', images, )

		avg_test_return = avg_test_return / self._test_episodes
		avg_test_cost = avg_test_cost / self._test_episodes
		if self.wandb_on:
			wandb.log({'Evaluation_Return': avg_test_return, 'Evaluation_Cost': avg_test_cost}, step=self.global_num_episodes)
		tf.summary.scalar(name="MFAgent/Evaluation_Return", data=avg_test_return, step=self.global_total_steps)
		tf.summary.scalar(name="MFAgent/Evaluation_Cost", data=avg_test_cost, step=self.global_total_steps)

		return avg_test_return, avg_test_cost

	def train(self):
		if self._policy.normalize_adv:
			samples = self.replay_buffer._encode_sample(np.arange(self._policy.horizon))
			mean_adv = np.mean(samples["adv"])
			std_adv = np.std(samples["adv"])
			# Update normalizer
			if self._normalize_obs:
				self._obs_normalizer.experience(samples["obs"])
		with tf.summary.record_if(self.global_total_steps % self._save_summary_interval == 0):
			for _ in range(self._policy.n_epoch):
				samples = self.replay_buffer._encode_sample(
					np.random.permutation(self._policy.horizon))
				if self._normalize_obs:
					samples["obs"] = self._obs_normalizer(samples["obs"], update=False)
				if self._policy.normalize_adv:
					adv = (samples["adv"] - mean_adv) / (std_adv + 1e-8)
				else:
					adv = samples["adv"]
				for idx in range(int(self._policy.horizon / self._policy.batch_size)):
					target = slice(idx * self._policy.batch_size,
								   (idx + 1) * self._policy.batch_size)
					self._policy.train(
						states=samples["obs"][target],
						actions=samples["act"][target],
						advantages=adv[target],
						logp_olds=samples["logp"][target],
						returns=samples["ret"][target])

class CollectSamples_TD3(CollectSamples_SAC):

	def __init__(self, env, max_steps, _output_dir, _episode_max_steps=100, gpu=0, wandb_on=False, *args, **kwargs):
		super(CollectSamples_TD3, self).__init__(env, max_steps, _output_dir, _episode_max_steps, gpu, wandb_on, *args, **kwargs)

		self._policy = TD3(
        state_shape=env.observation_space.shape,
        action_dim=env.action_space.high.size,
        gpu=gpu,
        max_action=env.action_space.high[0],
        batch_size=100,
        n_warmup=1000)
        # trainer = Trainer(policy, env, args, test_env=test_env)

ext: py

max_stars_repo_name: Chenaah/RAMCO

max_issues_repo_name: Chenaah/RAMCO

max_forks_repo_name: Chenaah/RAMCO

reason: import numpy


text:
import sys,random,glob,argparse,uuid,math
import sensenet
import pybullet as p
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.distributions import Categorical

forward = 3
left = 0
right = 1
up = 4
down = 5
indexup = 6
indexdown = 7
action_choices = [forward,left,right,up,down] #dont move index for now

plans = [([left]*5),([right] * 5),([up] * 5),([down] * 5),([indexup] * 5),([indexdown] * 5),] + (["random"] * 6)
def random_plan():
  a = []
  for i in range(5):
    a.append(random.randint(0,7))
  return a
plan = None
def save_data(env,label,touches,actions):
  path = "touch_data/"+str(label)+"/"+str(uuid.uuid1())
  env.mkdir_p(path)
  np.save(path+"/touches",touches)
  np.save(path+"/actions",actions)

class Policy(nn.Module):

  def __init__(self,observation_space_n,action_space_n):
    super(Policy, self).__init__()
    self.affine1 = nn.Linear(observation_space_n, 128)
    self.affine2 = nn.Linear(128, action_space_n)
    self.saved_log_probs = []
    self.rewards = []
    #self.init_weights()

  def init_weights(self):
    self.affine1.weight.data.uniform_(-0.1, 0.1)
    self.affine2.weight.data.uniform_(-0.1, 0.1)

  def forward(self, x):
    #if args.gpu and torch.cuda.is_available():
      # TODO: clean way to access "args.gpu" from here.
    #  x = x.cuda()
    x = F.relu(self.affine1(x))
    action_scores = self.affine2(x)
    return F.softmax(action_scores,dim=1)

def select_action(state,n_actions,steps,epsilon=0.2):
  if False and steps < 100:
    return np.random.choice(n_actions)
  else:
    state = torch.from_numpy(state).float().unsqueeze(0)
    print(state)
    probs = policy(Variable(state))
    action = probs.multinomial()
    m = Categorical(probs)
    action = m.sample()
    policy.saved_log_probs.append(m.log_prob(action))
    return action.data[0]
def finish_episode():
  R = 0
  policy_loss = []
  rewards = []
  for r in policy.rewards[::-1]:
    R = r + args.gamma * R
    rewards.insert(0, R)
  rewards = torch.Tensor(rewards)
  rewards = (rewards - rewards.mean()) / (rewards.std() + np.finfo(np.float32).eps)
  for log_prob, reward in zip(policy.saved_log_probs, rewards):
    policy_loss.append(-log_prob * reward)
  optimizer.zero_grad()
  policy_loss = torch.cat(policy_loss).sum()
  policy_loss.backward()
  optimizer.step()
  del policy.rewards[:]
  del policy.saved_log_probs[:]
parser = argparse.ArgumentParser()
parser.add_argument('--render', action='store_true', help='render the environment')
parser.add_argument('--environment', type=str, default="HandEnv-v0")
parser.add_argument('--epochs', type=int, default=1)
parser.add_argument('--gamma', type=float, default=0.99, metavar='G', help='discount factor (default: 0.99)')
parser.add_argument('--folder', type=str)
parser.add_argument('--file', type=str)
parser.add_argument('--fast_exit', type=int, default=0)
args = parser.parse_args()
if args.folder:
  files = list(glob.iglob(args.folder+"/**/*.obj", recursive=True))
else:
  files = [args.file]
random.shuffle(files)
env = sensenet.make(args.environment,{'render':args.render})
policy = Policy(3,env.action_space_n())
optimizer = optim.Adam(policy.parameters(), lr=1e-2)
total_step = 0
running_reward = 10
for filename in files:
  label = int(filename.split("/")[-3].split("_")[0])
  print(filename)
  print(label)
  path = "touch_data/"+str(label)+"/"
  env._reset({'obj_path':filename})
  env.mkdir_p(path)
  training_sets = []
  observations = []
  actions = []
  touch_count = 0
  step = 0
  episode = 0
  plan_step = 0
  tries = 0
  winners = 0
  prev_distance = 10000000
  for epoch in range(args.epochs):
    env.reset()
    while(1):
      points = p.getClosestPoints(env.obj_to_classify,env.agent,10000000,-1,22)
      al = points[0][7]
      ol = points[0][6]
      xd = (al[0]-ol[0])/2
      yd = (al[1]-ol[1])/2
      zd = (al[2]-ol[2])/2
      state = np.asarray([xd,td,zd])
      #print("xd",xd,"yd",yd,"zd",zd:wu)
      learning = False

      if env.is_touching() and plan_step == 0:
        plan = random.choice(plans)
        if plan == "random":
          plan = random_plan()
        action = plan[0]
        plan_step +=1
      elif plan_step >0:
        if len(plan) >= plan_step:
          #reset to a new plan
          plan_step = 0
          plan = random.choice(plans)
          if plan == "random":
            plan = random_plan()
        action = plan[plan_step]
        plan_step += 1
      elif not env.is_touching():
          action = select_action(np.asarray([xd,yd,zd]),env.action_space_n(),total_step)
          learning = True

      observation,reward,done,info = env.step(action)
      distance = math.sqrt(xd**2+yd**2+zd**2)
      if distance < prev_distance:
        reward += 1
        #reward += 1 * (max_steps - self.steps)
      policy.rewards.append(reward)
      prev_distance = distance
      print(len(policy.rewards))
      #if total_step % 20 == 0:
      #if total_step > 200 and total_step % 50 == 0:
      if learning == True and total_step % 50 == 0:
        print("episodeDDDDD")
        finish_episode()

      if env.is_touching():
        print("touch")

        observations.append(observation.reshape(200,200))
        actions.append(action)
        touch_count += 1
        #print("is touching")
      if touch_count >= 20:
        print("WINNER!!")
        winners += 1
        touch_count = 0
        #training_sets.append([observations,actions])
        save_data(env,label,observations,actions)
        observations = []
        actions = []
        step = 0
        episode +=1
        plan_step = 0
        env.reset()
      elif step >= 100:
        print("closing episode,touch_count",touch_count)
        touch_count = 0
        if len(observations) > 2:
          #training_sets.append([observations,actions])
          save_data(env,label,observations,actions)
        else:
          tries +=1
          if tries > 5:
            print("couldnt get trainig data for item",filename)
            break
        observations = []
        actions = []
        step = 0
        episode +=1
        plan_step = 0
        env.reset()
      total_step +=1
      if args.fast_exit != 0 and episode >= args.fast_exit:
          sys.exit()
      if step % 20 == 0:
        print("episode",episode,"label",label,"step", step,"plan", plan,"plan_step",plan_step,"touch_count",touch_count,"winners",winners)
      step +=1

max_stars_repo_name: jtoy/sensenet


text:
"""Bounding box visualization functions."""
from __future__ import absolute_import, division

import numpy as np
import mxnet as mx

from ...data.transforms.mask import fill, proto_fill

def expand_mask(masks, bboxes, im_shape, scores=None, thresh=0.5, scale=1.0, sortby=None):
    """Expand instance segmentation mask to full image size.

    Parameters
    ----------
    masks : numpy.ndarray or mxnet.nd.NDArray
        Binary images with shape `N, M, M`
    bboxes : numpy.ndarray or mxnet.nd.NDArray
        Bounding boxes with shape `N, 4`. Where `N` is the number of boxes
    im_shape : tuple
        Tuple of length 2: (width, height)
    scores : numpy.ndarray or mxnet.nd.NDArray, optional
        Confidence scores of the provided `bboxes` with shape `N`.
    thresh : float, optional, default 0.5
        Display threshold if `scores` is provided. Scores with less than `thresh`
        will be ignored in display, this is visually more elegant if you have
        a large number of bounding boxes with very small scores.
    sortby : str, optional, default None
        If not None, sort the color palette for masks by the given attributes of each bounding box.
        Valid inputs are 'area', 'xmin', 'ymin', 'xmax', 'ymax'.
    scale : float
        The scale of output image, which may affect the positions of boxes

    Returns
    -------
    numpy.ndarray
        Binary images with shape `N, height, width`
    numpy.ndarray
        Index array of sorted masks
    """
    if len(masks) != len(bboxes):
        raise ValueError('The length of bboxes and masks mismatch, {} vs {}'
                         .format(len(bboxes), len(masks)))
    if scores is not None and len(masks) != len(scores):
        raise ValueError('The length of scores and masks mismatch, {} vs {}'
                         .format(len(scores), len(masks)))

    if isinstance(masks, mx.nd.NDArray):
        masks = masks.asnumpy()
    if isinstance(bboxes, mx.nd.NDArray):
        bboxes = bboxes.asnumpy()
    if isinstance(scores, mx.nd.NDArray):
        scores = scores.asnumpy()

    if sortby is not None:
        if sortby == 'area':
            areas = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
            sorted_inds = np.argsort(-areas)
        elif sortby == 'xmin':
            sorted_inds = np.argsort(-bboxes[:, 0])
        elif sortby == 'ymin':
            sorted_inds = np.argsort(-bboxes[:, 1])
        elif sortby == 'xmax':
            sorted_inds = np.argsort(-bboxes[:, 2])
        elif sortby == 'ymax':
            sorted_inds = np.argsort(-bboxes[:, 3])
        else:
            raise ValueError('argument sortby cannot take value {}'
                             .format(sortby))
    else:
        sorted_inds = np.argsort(range(len(masks)))

    full_masks = []
    bboxes *= scale
    for i in sorted_inds:
        if scores is not None and scores[i] < thresh:
            continue
        mask = masks[i]
        bbox = bboxes[i]
        full_masks.append(proto_fill(mask, bbox, im_shape))
    full_masks = np.array(full_masks)
    return full_masks, sorted_inds

def plot_mask(img, masks, alpha=0.5):
    """Visualize segmentation mask.

    Parameters
    ----------
    img : numpy.ndarray or mxnet.nd.NDArray
        Image with shape `H, W, 3`.
    masks : numpy.ndarray or mxnet.nd.NDArray
        Binary images with shape `N, H, W`.
    alpha : float, optional, default 0.5
        Transparency of plotted mask

    Returns
    -------
    numpy.ndarray
        The image plotted with segmentation masks

    """
    if isinstance(img, mx.nd.NDArray):
        img = img.asnumpy()
    if isinstance(masks, mx.nd.NDArray):
        masks = masks.asnumpy()

    for mask in masks:
        color = np.random.random(3) * 255
        mask = np.repeat((mask > 0)[:, :, np.newaxis], repeats=3, axis=2)
        img = np.where(mask, img * (1 - alpha) + color * alpha, img)
    return img.astype('uint8')

def cv_merge_two_images(img1, img2, alpha=0.5, size=None):
    """Merge two images with OpoenCV.

    Parameters
    ----------
    img1 : numpy.ndarray or mxnet.nd.NDArray
        Image with shape `H, W, 3`.
    img2 : numpy.ndarray or mxnet.nd.NDArray
        Image with shape `H, W, 3`.
    alpha : float, optional, default 0.5
        Transparency of `img2`
    size : list, optional, default None
        The output size of the merged image

    Returns
    -------
    numpy.ndarray
        The merged image

    """
    from ..filesystem import try_import_cv2
    cv2 = try_import_cv2()

    if isinstance(img1, mx.nd.NDArray):
        img1 = img1.asnumpy()
    if isinstance(img2, mx.nd.NDArray):
        img2 = img2.asnumpy()
    img = cv2.addWeighted(img1, 1-alpha, img2, alpha, 0)
    if size is not None:
        img = cv2.resize(img, (int(size[1]), int(size[0])))
    return img

def expand_yolactmask(masks, bboxes, im_shape, scores=None, thresh=0.5):
    if len(masks) != len(bboxes):
        raise ValueError('The length of bboxes and masks mismatch, {} vs {}'
                         .format(len(bboxes), len(masks)))
    if scores is not None and len(masks) != len(scores):
        raise ValueError('The length of scores and masks mismatch, {} vs {}'
                         .format(len(scores), len(masks)))

    if isinstance(masks, mx.nd.NDArray):
        masks = masks.asnumpy()
    if isinstance(bboxes, mx.nd.NDArray):
        bboxes = bboxes.asnumpy()
    if isinstance(scores, mx.nd.NDArray):
        scores = scores.asnumpy()

    areas = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
    sorted_inds = np.argsort(-areas)

    full_masks = []
    for i in sorted_inds:
        if scores is not None and scores[i] < thresh:
            continue
        mask = masks[i]
        full_masks.append(proto_fill(mask, im_shape))
    full_masks = np.array(full_masks)
    return full_masks


text:
# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

import os

import numpy as np
import numpy.testing as nt

from .mpi import MPITestCase

from ..tod import sim_focalplane as sfp

def generate_hex(npix, width, poltype, fwhm):
    if poltype == "qu":
        pol_a = sfp.hex_pol_angles_qu(npix)
        pol_b = sfp.hex_pol_angles_qu(npix, offset=90.0)
    elif poltype == "radial":
        pol_a = sfp.hex_pol_angles_radial(npix)
        pol_b = sfp.hex_pol_angles_radial(npix, offset=90.0)
    dets_a = sfp.hex_layout(npix, width, "", "A", pol_a)
    dets_b = sfp.hex_layout(npix, width, "", "B", pol_b)

    dets = dict()
    dets.update(dets_a)
    dets.update(dets_b)

    # Pol color different for A/B detectors
    detpolcolor = dict()
    detpolcolor.update({x: "red" for x in dets_a.keys()})
    detpolcolor.update({x: "blue" for x in dets_b.keys()})

    # set the label to just the detector name
    detlabels = {x: x for x in dets.keys()}

    # fwhm and face color the same
    detfwhm = {x: fwhm for x in dets.keys()}

    # cycle through some colors just for fun
    pclr = [
        (1.0, 0.0, 0.0, 0.1),
        (1.0, 0.5, 0.0, 0.1),
        (0.25, 0.5, 1.0, 0.1),
        (0.0, 0.75, 0.0, 0.1),
    ]
    detcolor = {y: pclr[(x // 2) % 4] for x, y in enumerate(sorted(dets.keys()))}

    # split out quaternions for plotting
    detquats = {x: dets[x]["quat"] for x in dets.keys()}

    return dets, detquats, detfwhm, detcolor, detpolcolor, detlabels

def generate_rhombus(npix, width, fwhm, prefix, center):
    pol_a = sfp.rhomb_pol_angles_qu(npix)
    pol_b = sfp.rhomb_pol_angles_qu(npix, offset=90.0)
    dets_a = sfp.rhombus_layout(npix, width, prefix, "A", pol_a, center=center)
    dets_b = sfp.rhombus_layout(npix, width, prefix, "B", pol_b, center=center)

    dets = dict()
    dets.update(dets_a)
    dets.update(dets_b)

    # Pol color different for A/B detectors
    detpolcolor = dict()
    detpolcolor.update({x: "red" for x in dets_a.keys()})
    detpolcolor.update({x: "blue" for x in dets_b.keys()})

    # set the label to just the detector name
    detlabels = {x: x for x in dets.keys()}

    # fwhm and face color the same
    detfwhm = {x: fwhm for x in dets.keys()}

    # cycle through some colors just for fun
    pclr = [
        (1.0, 0.0, 0.0, 0.1),
        (1.0, 0.5, 0.0, 0.1),
        (0.25, 0.5, 1.0, 0.1),
        (0.0, 0.75, 0.0, 0.1),
    ]
    detcolor = {y: pclr[(x // 2) % 4] for x, y in enumerate(sorted(dets.keys()))}

    # split out quaternions for plotting
    detquats = {x: dets[x]["quat"] for x in dets.keys()}

    return dets, detquats, detfwhm, detcolor, detpolcolor, detlabels

class SimFocalplaneTest(MPITestCase):
    def setUp(self):
        self.outdir = "toast_test_output"
        self.rank = 0
        if self.comm is not None:
            self.rank = self.comm.rank
        if self.rank == 0:
            if not os.path.isdir(self.outdir):
                os.mkdir(self.outdir)

    def test_cart_quat(self):
        xincr = np.linspace(-5.0, 5.0, num=10, endpoint=True)
        yincr = np.linspace(-5.0, 5.0, num=10, endpoint=True)
        offsets = list()
        for x in xincr:
            for y in yincr:
                ang = 3.6 * (x - xincr[0]) * (y - yincr[0])
                offsets.append([x, y, ang])
        quats = sfp.cartesian_to_quat(offsets)
        detquats = {"{}".format(x): y for x, y in enumerate(quats)}
        fwhm = {x: 30.0 for x in detquats.keys()}
        outfile = os.path.join(self.outdir, "out_test_cart2quat.png")
        if self.rank == 0:
            sfp.plot_focalplane(detquats, 12.0, 12.0, outfile, fwhm=fwhm)
        return

    def test_hex_nring(self):
        result = {
            1: 1,
            7: 2,
            19: 3,
            37: 4,
            61: 5,
            91: 6,
            127: 7,
            169: 8,
            217: 9,
            271: 10,
            331: 11,
            397: 12,
        }
        for npix, check in result.items():
            test = sfp.hex_nring(npix)
            nt.assert_equal(test, check)
        return

    def test_vis_hex_small(self):
        dets, detquats, detfwhm, detcolor, detpolcolor, detlabels = generate_hex(
            7, 5.0, "qu", 15.0
        )
        outfile = os.path.join(self.outdir, "out_test_vis_hex_small.png")
        if self.rank == 0:
            sfp.plot_focalplane(
                detquats,
                6.0,
                6.0,
                outfile,
                fwhm=detfwhm,
                facecolor=detcolor,
                polcolor=detpolcolor,
                labels=detlabels,
            )
        return

    def test_vis_hex_small_rad(self):
        dets, detquats, detfwhm, detcolor, detpolcolor, detlabels = generate_hex(
            7, 5.0, "radial", 15.0
        )
        outfile = os.path.join(self.outdir, "out_test_vis_hex_small_rad.png")
        if self.rank == 0:
            sfp.plot_focalplane(
                detquats,
                6.0,
                6.0,
                outfile,
                fwhm=detfwhm,
                facecolor=detcolor,
                polcolor=detpolcolor,
                labels=detlabels,
            )
        return

    def test_vis_hex_medium(self):
        dets, detquats, detfwhm, detcolor, detpolcolor, detlabels = generate_hex(
            91, 5.0, "qu", 10.0
        )
        outfile = os.path.join(self.outdir, "out_test_vis_hex_medium.png")
        if self.rank == 0:
            sfp.plot_focalplane(
                detquats,
                6.0,
                6.0,
                outfile,
                fwhm=detfwhm,
                facecolor=detcolor,
                polcolor=detpolcolor,
                labels=detlabels,
            )
        return

    def test_vis_hex_large(self):
        dets, detquats, detfwhm, detcolor, detpolcolor, detlabels = generate_hex(
            217, 5.0, "qu", 5.0
        )
        outfile = os.path.join(self.outdir, "out_test_vis_hex_large.png")
        if self.rank == 0:
            sfp.plot_focalplane(
                detquats,
                6.0,
                6.0,
                outfile,
                fwhm=detfwhm,
                facecolor=detcolor,
                polcolor=detpolcolor,
                labels=detlabels,
            )
        return

    def test_vis_rhombus(self):
        sixty = np.pi / 3.0
        thirty = np.pi / 6.0
        rtthree = np.sqrt(3.0)

        rdim = 8
        rpix = rdim ** 2

        hexwidth = 5.0
        rwidth = hexwidth / rtthree

        # angular separation of rhombi
        margin = 0.60 * hexwidth

        centers = [
            np.array([0.5 * margin, 0.0, 0.0]),
            np.array(
                [-0.5 * np.cos(sixty) * margin, 0.5 * np.sin(sixty) * margin, 120.0]
            ),
            np.array(
                [-0.5 * np.cos(sixty) * margin, -0.5 * np.sin(sixty) * margin, 240.0]
            ),
        ]

        cquats = sfp.cartesian_to_quat(centers)

        dets = dict()
        detquats = dict()
        detfwhm = dict()
        detcolor = dict()
        detpolcolor = dict()
        detlabels = dict()
        for w, c in enumerate(cquats):
            (
                wdets,
                wdetquats,
                wdetfwhm,
                wdetcolor,
                wdetpolcolor,
                wdetlabels,
            ) = generate_rhombus(rpix, rwidth, 7.0, "{}".format(w), c)
            dets.update(wdets)
            detquats.update(wdetquats)
            detfwhm.update(wdetfwhm)
            detcolor.update(wdetcolor)
            detpolcolor.update(wdetpolcolor)
            detlabels.update(wdetlabels)

        outfile = os.path.join(self.outdir, "out_test_vis_rhombus.png")
        if self.rank == 0:
            sfp.plot_focalplane(
                detquats,
                1.2 * hexwidth,
                1.2 * hexwidth,
                outfile,
                fwhm=detfwhm,
                facecolor=detcolor,
                polcolor=detpolcolor,
                labels=detlabels,
            )
        return


text:
import pandas as pd
import numpy as np

def read_covid19_data():
  url = 'https://covid.ourworldindata.org/data/owid-covid-data.csv'
  print(f'Reading COVID-19 data from : {url}')
  return pd.read_csv(url)


text:
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 28 12:20:39 2017

@author: L1566
"""

import numpy as np
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
import xarray as xr
import matplotlib.lines as mlines
import matplotlib.gridspec as gridspec
from datetime import datetime, timedelta

def plot_streamflow(ID,scen,q_fmi,q_clim,q_clim1,dailymean=True,timeseries=False,fmi_clim=False):
    ### daily average ###
    D_q_fmi = q_fmi.groupby('dtime.dayofyear').mean()
    D_q_clim = q_clim.groupby('dtime.dayofyear').mean()
    D_q_clim1 = q_clim1.groupby('dtime.dayofyear').mean()
    # mean daily value within each month of the year
    Mmean_q_fmi = q_fmi.groupby('dtime.month').mean()
    Mmean_q_clim = q_clim.groupby('dtime.month').mean()
    Mmean_q_clim1 = q_clim1.groupby('dtime.month').mean()
    ### ploting streamflow over 30 years ###
    if dailymean:
        sns.set_style('ticks')
        plt.figure(figsize = (12,8))
        #gs = gridspec.GridSpec(2, 2,width_ratios=[2, 1])
        plt.subplot(211)
        D_q_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
        D_q_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
        D_q_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
        plt.xlim([0,365])
        plt.xlabel('day of year')
        plt.ylabel('streamflow (mm/d)')
        plt.legend(loc='best')
        ###
        plt.subplot(212)
        Mmean_q_fmi.plot.line('b--',linewidth = 1,label='fmi(1980-2009)')
        Mmean_q_clim.plot.line('r',linewidth = 1,label='climate(1980-2009)')
        Mmean_q_clim1.plot.line('k',linewidth = 1,label='climate(2070-2099)')
        plt.xlabel('month of year')
        plt.ylabel('streamflow (mm/d)')
        plt.legend(loc='best')

        sns.despine()
        plt.tight_layout()
        plt.show()
        plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_daily_monthly_mean_flow.pdf',dpi=300)

    if fmi_clim:
        sns.set_style('ticks')
        plt.figure(figsize = (12,4))
        q_fmi.plot.line('bo',markersize = 4, alpha = 0.2,label='fmi(1980-2009)')
        q_clim.plot.line('r',linewidth = 0.5,alpha = 0.5,label='climate(1980-2009)')
        plt.xlabel('date')
        plt.ylabel('streamflow')
        plt.legend(loc='best')
        sns.despine()
        plt.show()
        plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\fmi_clim_flow_c_'+ID+'_'+scen+'.pdf',dpi=300)

    if timeseries:
        sns.set_style('ticks')
        plt.figure(figsize = (12,4))
        plt.plot(q_clim,'bo',markersize = 4, alpha = 0.2,label='climate(1980-2009)')
        plt.plot(q_clim1,'r',linewidth = 0.5,alpha = 0.5,label='climate(2070-2099)')
        plt.xlabel('number of days')
        plt.ylabel('streamflow')
        plt.legend(loc='best')
        sns.despine()
        plt.show()
        plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_clim_clim1_flow.pdf',dpi=300)

def plot_sat_deficit(ID,scen,S_fmi,S_clim,S_clim1):
    #daily mean
    D_S_fmi = S_fmi.groupby('dtime.dayofyear').mean()
    D_S_clim = S_clim.groupby('dtime.dayofyear').mean()
    D_S_clim1 = S_clim1.groupby('dtime.dayofyear').mean()
    # mean daily value within each month of the year
    Mmean_S_fmi = S_fmi.groupby('dtime.month').mean()
    Mmean_S_clim = S_clim.groupby('dtime.month').mean()
    Mmean_S_clim1 = S_clim1.groupby('dtime.month').mean()
    ### ploting daily sat. deficit over 30 years ###
    sns.set_style('ticks')
    plt.figure(figsize = (12,8))
    plt.subplot(211)
    D_S_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_S_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_S_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('sat. deficit (m)')
    plt.legend(loc='best')
    ###
    plt.subplot(212)
    Mmean_S_fmi.plot.line('b--',linewidth = 1,label='fmi(1980-2009)')
    Mmean_S_clim.plot.line('r',linewidth = 1,label='climate(1980-2009)')
    Mmean_S_clim1.plot.line('k',linewidth = 1,label='climate(2070-2099)')
    plt.xlabel('month of year')
    plt.ylabel('sat. deficit (m)')
    plt.legend(loc='best')
    sns.despine()
    plt.tight_layout()
    plt.show()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_daily_monthly_mean_sdeficit.pdf',dpi=300)

def mapping_changes(ID,scen,d_fmi,d_clim,d_clim1):
    ## mean ET components over 30 years
    TmET_fmi = (d_fmi.Transpi+d_fmi.Efloor+d_fmi.Evap).mean(dim=['dtime'])
    TmTrans_fmi = d_fmi.Transpi.mean(dim=['dtime'])
    TmEfloor_fmi = d_fmi.Efloor.mean(dim=['dtime'])
    TmEvap_fmi = d_fmi.Evap.mean(dim=['dtime'])
    TmSWE_fmi = d_fmi.SWE.mean(dim=['dtime'])
    ## mean ET components over 30 years
    TmET_clim = (d_clim.Transpi+d_clim.Efloor+d_clim.Evap).mean(dim=['dtime'])
    TmTrans_clim = d_clim.Transpi.mean(dim=['dtime'])
    TmEfloor_clim = d_clim.Efloor.mean(dim=['dtime'])
    TmEvap_clim = d_clim.Evap.mean(dim=['dtime'])
    TmSWE_clim = d_clim.SWE.mean(dim=['dtime'])
    ## mean ET components over 30 years
    TmET_clim1 = (d_clim1.Transpi+d_clim1.Efloor+d_clim1.Evap).mean(dim=['dtime'])
    TmTrans_clim1 = d_clim1.Transpi.mean(dim=['dtime'])
    TmEfloor_clim1 = d_clim1.Efloor.mean(dim=['dtime'])
    TmEvap_clim1 = d_clim1.Evap.mean(dim=['dtime'])
    TmSWE_clim1 = d_clim1.SWE.mean(dim=['dtime'])
    sns.set_style('white')
    plt.figure(figsize = (25,8))
    plt.subplot(251)
    TmET_clim.plot(cmap='jet',vmin=0.45,vmax=0.78)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(1980-2009)')
    plt.subplot(256)
    TmET_clim1.plot(cmap='jet',vmin=0.45,vmax=0.78)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(2070-2099)')
    plt.subplot(252)
    TmTrans_clim.plot(cmap='jet',vmin=0.0,vmax=0.72)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(1980-2009)')
    plt.subplot(257)
    TmTrans_clim1.plot(cmap='jet',vmin=0.0,vmax=0.72)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(2070-2099)')
    plt.subplot(253)
    TmEfloor_clim.plot(cmap='jet',vmin=0.0,vmax=0.66)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(1980-2009)')
    plt.subplot(258)
    TmEfloor_clim1.plot(cmap='jet',vmin=0.0,vmax=0.66)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(2070-2099)')
    plt.subplot(254)
    TmEvap_clim.plot(cmap='jet',vmin=0.0,vmax=0.43)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(1980-2009)')
    plt.subplot(259)
    TmEvap_clim1.plot(cmap='jet',vmin=0.0,vmax=0.43)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(2070-2099)')
    plt.subplot(255)
    TmSWE_clim.plot(cmap='jet',vmin=8.0,vmax=36)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(1980-2009)')
    plt.subplot(2,5,10)
    TmSWE_clim1.plot(cmap='jet',vmin=8.0,vmax=36)
    plt.xlabel('longitude')
    plt.ylabel('lattitude')
    plt.title('climate(2070-2099)')
    plt.tight_layout()
    plt.show()

    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_mapping_changes.pdf',dpi=300)

def plot_ET_components(ID,scen,mET_clim,mTrans_clim,mEfloor_clim,mEvap_clim,mSWE_clim,mET_clim1,mTrans_clim1,mEfloor_clim1,mEvap_clim1,mSWE_clim1):
    #1980-2009 daily mean
    D_mET_clim = mET_clim.groupby('dtime.dayofyear').mean()
    D_mTrans_clim = mTrans_clim.groupby('dtime.dayofyear').mean()
    D_mEfloor_clim = mEfloor_clim.groupby('dtime.dayofyear').mean()
    D_mEvap_clim = mEvap_clim.groupby('dtime.dayofyear').mean()
    D_mSWE_clim = mSWE_clim.groupby('dtime.dayofyear').max()
    #2070-2099
    D_mET_clim1 = mET_clim1.groupby('dtime.dayofyear').mean()
    D_mTrans_clim1 = mTrans_clim1.groupby('dtime.dayofyear').mean()
    D_mEfloor_clim1 = mEfloor_clim1.groupby('dtime.dayofyear').mean()
    D_mEvap_clim1 = mEvap_clim1.groupby('dtime.dayofyear').mean()
    D_mSWE_clim1 = mSWE_clim1.groupby('dtime.dayofyear').max()
    # 1980-2009 mean daily value within each month of the year
    Mmean_mET_clim = mET_clim.groupby('dtime.month').mean()
    Mmean_mTrans_clim = mTrans_clim.groupby('dtime.month').mean()
    Mmean_mEfloor_clim = mEfloor_clim.groupby('dtime.month').mean()
    Mmean_mEvap_clim = mEvap_clim.groupby('dtime.month').mean()
    Mmean_mSWE_clim = mSWE_clim.groupby('dtime.month').mean()
    Mmax_mSWE_clim = mSWE_clim.groupby('dtime.month').max()
    #2070-2099
    Mmean_mET_clim1 = mET_clim1.groupby('dtime.month').mean()
    Mmean_mTrans_clim1 = mTrans_clim1.groupby('dtime.month').mean()
    Mmean_mEfloor_clim1 = mEfloor_clim1.groupby('dtime.month').mean()
    Mmean_mEvap_clim1 = mEvap_clim1.groupby('dtime.month').mean()
    Mmean_mSWE_clim1 = mSWE_clim1.groupby('dtime.month').mean()
    Mmax_mSWE_clim1 = mSWE_clim1.groupby('dtime.month').max()
    ### ploting daily and monthly cat. av. ET, Trans, Efloor, Evap over 30 years ###
    sns.set_style('ticks')
    plt.figure(figsize = (18,15))
    gs = gridspec.GridSpec(5, 2,width_ratios=[2, 1])
    plt.subplot(gs[0])
    #D_mET_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mET_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_mET_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('ET (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[1])
    #plt.plot(Mmean_mET_fmi.index,Mmean_mET_fmi,'b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mET_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    Mmean_mET_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlabel('month')
    plt.ylabel('ET (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[2])
    #D_mTrans_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mTrans_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_mTrans_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Transpiration (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[3])
    #plt.plot(Mmean_mTrans_fmi.index,Mmean_mTrans_fmi,'b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mTrans_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    Mmean_mTrans_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlabel('month')
    plt.ylabel('Transpiration (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[4])
    #D_mEfloor_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mEfloor_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_mEfloor_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Efloor (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[5])
    #plt.plot(Mmean_mEfloor_fmi.index,Mmean_mEfloor_fmi,'b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mEfloor_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    Mmean_mEfloor_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlabel('month')
    plt.ylabel('Efloor (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[6])
    #D_mEvap_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mEvap_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_mEvap_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Evap (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[7])
    #plt.plot(Mmean_mEvap_fmi.index,Mmean_mEvap_fmi,'b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mEvap_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    Mmean_mEvap_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlabel('month')
    plt.ylabel('Evap (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[8])
    #D_mSWE_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mSWE_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    D_mSWE_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('SWE (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[9])
    #plt.plot(Mmean_mSWE_fmi.index,Mmean_mSWE_fmi,'b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmax_mSWE_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    Mmax_mSWE_clim1.plot.line('k',linewidth = 0.5,label='climate(2070-2099)')
    plt.xlabel('month')
    plt.ylabel('SWE (mm)')
    plt.legend(loc='best')
    plt.tight_layout()
    sns.despine()
    plt.show()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_daily_monthly_mean_ETs.pdf',dpi=300)

def plt_ETs_fmi_comparison(ID,scen,mET_clim,mTrans_clim,mEfloor_clim,mEvap_clim,mSWE_clim,mET_fmi,mTrans_fmi,mEfloor_fmi,mEvap_fmi,mSWE_fmi):
    #1980-2009 daily mean
    D_mET_clim = mET_clim.groupby('dtime.dayofyear').mean()
    D_mTrans_clim = mTrans_clim.groupby('dtime.dayofyear').mean()
    D_mEfloor_clim = mEfloor_clim.groupby('dtime.dayofyear').mean()
    D_mEvap_clim = mEvap_clim.groupby('dtime.dayofyear').mean()
    D_mSWE_clim = mSWE_clim.groupby('dtime.dayofyear').mean()
    #fmi data
    D_mET_fmi = mET_fmi.groupby('dtime.dayofyear').mean()
    D_mTrans_fmi = mTrans_fmi.groupby('dtime.dayofyear').mean()
    D_mEfloor_fmi = mEfloor_fmi.groupby('dtime.dayofyear').mean()
    D_mEvap_fmi = mEvap_fmi.groupby('dtime.dayofyear').mean()
    D_mSWE_fmi = mSWE_fmi.groupby('dtime.dayofyear').mean()
    # 1980-2009 mean daily value within each month of the year
    Mmean_mET_clim = mET_clim.groupby('dtime.month').mean()
    Mmean_mTrans_clim = mTrans_clim.groupby('dtime.month').mean()
    Mmean_mEfloor_clim = mEfloor_clim.groupby('dtime.month').mean()
    Mmean_mEvap_clim = mEvap_clim.groupby('dtime.month').mean()
    Mmean_mSWE_clim = mSWE_clim.groupby('dtime.month').mean()
    Mmax_mSWE_clim = mSWE_clim.groupby('dtime.month').max()
    #fmi data
    Mmean_mET_fmi = mET_fmi.groupby('dtime.month').mean()
    Mmean_mTrans_fmi = mTrans_fmi.groupby('dtime.month').mean()
    Mmean_mEfloor_fmi = mEfloor_fmi.groupby('dtime.month').mean()
    Mmean_mEvap_fmi = mEvap_fmi.groupby('dtime.month').mean()
    Mmean_mSWE_fmi = mSWE_fmi.groupby('dtime.month').mean()
    Mmax_mSWE_fmi = mSWE_fmi.groupby('dtime.month').max()
    ### ploting daily and monthly cat. av. ET, Trans, Efloor, Evap over 30 years ###
    sns.set_style('ticks')
    plt.figure(figsize = (18,15))
    gs = gridspec.GridSpec(5, 2,width_ratios=[2, 1])
    plt.subplot(gs[0])
    D_mET_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mET_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('ET (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[1])
    Mmean_mET_fmi.plot.line('b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mET_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlabel('month')
    plt.ylabel('ET (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[2])
    D_mTrans_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mTrans_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Transpiration (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[3])
    Mmean_mTrans_fmi.plot.line('b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mTrans_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlabel('month')
    plt.ylabel('Transpiration (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[4])
    D_mEfloor_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mEfloor_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Efloor (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[5])
    Mmean_mEfloor_fmi.plot.line('b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mEfloor_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlabel('month')
    plt.ylabel('Efloor (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[6])
    D_mEvap_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mEvap_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('Evap (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[7])
    Mmean_mEvap_fmi.plot.line('b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mEvap_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlabel('month')
    plt.ylabel('Evap (mm)')
    plt.legend(loc='best')

    plt.subplot(gs[8])
    D_mSWE_fmi.plot.line('bo',markersize = 4, alpha = 0.3,label='fmi(1980-2009)')
    D_mSWE_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlim([0,365])
    plt.xlabel('day of year')
    plt.ylabel('SWE (mm)')
    plt.legend(loc='best')
    ###
    plt.subplot(gs[9])
    Mmean_mSWE_fmi.plot.line('b-',linewidth = 0.5,label='fmi(1980-2009)')
    Mmean_mSWE_clim.plot.line('r',linewidth = 0.5,label='climate(1980-2009)')
    plt.xlabel('month')
    plt.ylabel('SWE (mm)')
    plt.legend(loc='best')
    plt.tight_layout()
    sns.despine()
    plt.show()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_daily_monthly_fmi_clim_ETs.pdf',dpi=300)
def Ann_Change(ID,scen,q_clim,mET_clim,mTrans_clim,mEfloor_clim,mEvap_clim,q_clim1,
               mET_clim1,mTrans_clim1,mEfloor_clim1,mEvap_clim1,A_p_clim,A_p_clim1):
    #############################################
    ## variables from climate data simulations (1981-2010)
    #############################################
    A_q_clim = q_clim.groupby('dtime.year').sum()
    A_mET_clim = mET_clim.groupby('dtime.year').sum()
    A_mTrans_clim = mTrans_clim.groupby('dtime.year').sum()
    A_mEfloor_clim = mEfloor_clim.groupby('dtime.year').sum()
    A_mEvap_clim = mEvap_clim.groupby('dtime.year').sum()
    #############################################
    ## variables from climate data simulations (2071-2100)
    #############################################
    A_q_clim1 = q_clim1.groupby('dtime.year').sum()
    A_mET_clim1 = mET_clim1.groupby('dtime.year').sum()
    A_mTrans_clim1 = mTrans_clim1.groupby('dtime.year').sum()
    A_mEfloor_clim1 = mEfloor_clim1.groupby('dtime.year').sum()
    A_mEvap_clim1 = mEvap_clim1.groupby('dtime.year').sum()

    ################################################
    ###  data in a Dataframe for boxplot ###########
    ##################################################
    mAQ = np.zeros(len(A_q_clim))
    mAET = np.zeros(len(A_q_clim))
    mATr = np.zeros(len(A_q_clim))
    mAEf = np.zeros(len(A_q_clim))
    mAEv = np.zeros(len(A_q_clim))
    for j in range(len(A_q_clim)):
        mAQ[j]=100.*(A_q_clim1[j]/A_p_clim1[j]-A_q_clim[j]/A_p_clim[j])/(A_q_clim[j]/A_p_clim[j])
        mAET[j]=100.*(A_mET_clim1[j]/A_p_clim1[j]-A_mET_clim[j]/A_p_clim[j])/(A_mET_clim[j]/A_p_clim[j])
        mATr[j]=100.*(A_mTrans_clim1[j]/A_p_clim1[j]-A_mTrans_clim[j]/A_p_clim[j])/(A_mTrans_clim[j]/A_p_clim[j])
        mAEf[j]=100.*(A_mEfloor_clim1[j]/A_p_clim1[j]-A_mEfloor_clim[j]/A_p_clim[j])/(A_mEfloor_clim[j]/A_p_clim[j])
        mAEv[j]=100.*(A_mEvap_clim1[j]/A_p_clim1[j]-A_mEvap_clim[j]/A_p_clim[j])/(A_mEvap_clim[j]/A_p_clim[j])

    mAM= pd.DataFrame()
    mAM['Q_P_ANN'] = mAQ
    mAM['ET_P_ANN'] = mAET
    mAM['Tr_P_ANN'] = mATr
    mAM['Ef_P_ANN'] = mAEf
    mAM['Evap_P_ANN'] = mAEv

    ##### plotting boxplot
    sns.set_style('ticks')
    plt.figure(figsize=(6,6))
    sns.boxplot(data=mAM,showmeans=True,showfliers=False,linewidth=0.5,meanprops=dict(markersize=2))
    plt.ylabel('relative change (%)')
    plt.xticks(rotation='vertical')
    plt.axhline(linewidth=0.5)
    sns.despine()
    plt.tight_layout()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_ann_change.pdf',dpi=300)

def Seasonal_Change(ID,scen,d_clim,d_clim1):
    #############################################
    ## variables from climate data simulations (1981-2010)
    #############################################
    q_clim = 1e3*d_clim.Qt
    #S_clim = d_clim.S
    mET_clim = (d_clim.Transpi+d_clim.Efloor+d_clim.Evap).mean(dim=['dlat','dlon'])
    mTrans_clim = d_clim.Transpi.mean(dim=['dlat','dlon'])
    mEfloor_clim = d_clim.Efloor.mean(dim=['dlat','dlon'])
    mEvap_clim = d_clim.Evap.mean(dim=['dlat','dlon'])
    mSWE_clim = d_clim.SWE.mean(dim=['dlat','dlon'])
    #############################################
    ## variables from climate data simulations (2071-2100)
    #############################################
    q_clim1 = 1e3*d_clim1.Qt
    #S_clim1 = d_clim1.S
    mET_clim1 = (d_clim1.Transpi+d_clim1.Efloor+d_clim1.Evap).mean(dim=['dlat','dlon'])
    mTrans_clim1 = d_clim1.Transpi.mean(dim=['dlat','dlon'])
    mEfloor_clim1 = d_clim1.Efloor.mean(dim=['dlat','dlon'])
    mEvap_clim1 = d_clim1.Evap.mean(dim=['dlat','dlon'])
    mSWE_clim1 = d_clim1.SWE.mean(dim=['dlat','dlon'])
    ################################################
    ###  seasonal mean daily ###########
    ##################################################
    month_clim = q_clim.groupby('dtime.month').apply(lambda x: x).month
    DJF = (month_clim <= 2) | (month_clim >= 12)
    MAM = (month_clim >= 3) & (month_clim <= 5)
    JJA = (month_clim >= 6) & (month_clim <= 8)
    SON = (month_clim >= 9) & (month_clim <= 11)
    # seasonal daily mean flow
    q_DJF = q_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    q_MAM = q_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    q_JJA = q_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    q_SON = q_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean ET
    mET_DJF = mET_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mET_MAM = mET_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mET_JJA = mET_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mET_SON = mET_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Transpi
    mTrans_DJF = mTrans_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mTrans_MAM = mTrans_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mTrans_JJA = mTrans_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mTrans_SON = mTrans_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Efloor
    mEfloor_DJF = mEfloor_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEfloor_MAM = mEfloor_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEfloor_JJA = mEfloor_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEfloor_SON = mEfloor_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Evap
    mEvap_DJF = mEvap_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEvap_MAM = mEvap_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEvap_JJA = mEvap_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEvap_SON = mEvap_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean SWE
    mSWE_DJF = mSWE_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mSWE_MAM = mSWE_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mSWE_JJA = mSWE_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mSWE_SON = mSWE_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    #####################################
    # seasonal daily mean flow for 2070-2099
    #####################################
    month_clim = q_clim1.groupby('dtime.month').apply(lambda x: x).month
    DJF = (month_clim <= 2) | (month_clim >= 12)
    MAM = (month_clim >= 3) & (month_clim <= 5)
    JJA = (month_clim >= 6) & (month_clim <= 8)
    SON = (month_clim >= 9) & (month_clim <= 11)
    q1_DJF = q_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    q1_MAM = q_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    q1_JJA = q_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    q1_SON = q_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean ET
    mET1_DJF = mET_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mET1_MAM = mET_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mET1_JJA = mET_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mET1_SON = mET_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Transpi
    mTrans1_DJF = mTrans_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mTrans1_MAM = mTrans_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mTrans1_JJA = mTrans_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mTrans1_SON = mTrans_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Efloor
    mEfloor1_DJF = mEfloor_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEfloor1_MAM = mEfloor_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEfloor1_JJA = mEfloor_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEfloor1_SON = mEfloor_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Evap
    mEvap1_DJF = mEvap_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEvap1_MAM = mEvap_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEvap1_JJA = mEvap_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEvap1_SON = mEvap_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean SWE
    mSWE1_DJF = mSWE_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mSWE1_MAM = mSWE_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mSWE1_JJA = mSWE_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mSWE1_SON = mSWE_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    ################################################
    ###  data in a Dataframe for boxplot ###########
    ##################################################
    ## seasonal mean daily flow
    mQ_DJF = np.zeros(len(q_DJF)-1)
    mQ_MAM = np.zeros(len(q_MAM)-1)
    mQ_JJA = np.zeros(len(q_JJA)-1)
    mQ_SON = np.zeros(len(q_SON)-1)
    for j in range(1,len(q_MAM)):
        if j==1: mQ_DJF[j-1]=np.NaN
        else: mQ_DJF[j-1]=100.*(q1_DJF[j-1]-q_DJF[j-1])/q_DJF[j-1]
        mQ_MAM[j-1]=100.*(q1_MAM[j]-q_MAM[j])/q_MAM[j]
        mQ_JJA[j-1]=100.*(q1_JJA[j]-q_JJA[j])/q_JJA[j]
        mQ_SON[j-1]=100.*(q1_SON[j]-q_SON[j])/q_SON[j]
    ## seasonal mean daily ET
    mmET_DJF = np.zeros(len(mET_DJF)-1)
    mmET_MAM = np.zeros(len(mET_MAM)-1)
    mmET_JJA = np.zeros(len(mET_JJA)-1)
    mmET_SON = np.zeros(len(mET_SON)-1)
    for j in range(1,len(mET_MAM)):
        if j==1: mmET_DJF[j-1]=np.NaN
        else: mmET_DJF[j-1]=100.*(mET1_DJF[j-1]-mET_DJF[j-1])/mET_DJF[j-1]
        mmET_MAM[j-1]=100.*(mET1_MAM[j]-mET_MAM[j])/mET_MAM[j]
        mmET_JJA[j-1]=100.*(mET1_JJA[j]-mET_JJA[j])/mET_JJA[j]
        mmET_SON[j-1]=100.*(mET1_SON[j]-mET_SON[j])/mET_SON[j]
    ## seasonal mean daily Trans
    mmTr_DJF = np.zeros(len(mTrans_DJF)-1)
    mmTr_MAM = np.zeros(len(mTrans_MAM)-1)
    mmTr_JJA = np.zeros(len(mTrans_JJA)-1)
    mmTr_SON = np.zeros(len(mTrans_SON)-1)
    for j in range(1,len(mTrans_MAM)):
        if j==1: mmTr_DJF[j-1]=np.NaN
        else: mmTr_DJF[j-1]=100.*(mTrans1_DJF[j-1]-mTrans_DJF[j-1])/mTrans_DJF[j-1]
        mmTr_MAM[j-1]=100.*(mTrans1_MAM[j]-mTrans_MAM[j])/mTrans_MAM[j]
        mmTr_JJA[j-1]=100.*(mTrans1_JJA[j]-mTrans_JJA[j])/mTrans_JJA[j]
        mmTr_SON[j-1]=100.*(mTrans1_SON[j]-mTrans_SON[j])/mTrans_SON[j]
    ## seasonal mean daily Efloor
    mmEf_DJF = np.zeros(len(mEfloor_DJF)-1)
    mmEf_MAM = np.zeros(len(mEfloor_MAM)-1)
    mmEf_JJA = np.zeros(len(mEfloor_JJA)-1)
    mmEf_SON = np.zeros(len(mEfloor_SON)-1)
    for j in range(1,len(mEfloor_MAM)):
        if j==1: mmEf_DJF[j-1]=np.NaN
        else: mmEf_DJF[j-1]=100.*(mEfloor1_DJF[j-1]-mEfloor_DJF[j-1])/mEfloor_DJF[j-1]
        mmEf_MAM[j-1]=100.*(mEfloor1_MAM[j]-mEfloor_MAM[j])/mEfloor_MAM[j]
        mmEf_JJA[j-1]=100.*(mEfloor1_JJA[j]-mEfloor_JJA[j])/mEfloor_JJA[j]
        mmEf_SON[j-1]=100.*(mEfloor1_SON[j]-mEfloor_SON[j])/mEfloor_SON[j]
    mmEf_DJF[mmEf_DJF>1e+03]=np.NaN
    mmEf_MAM[mmEf_MAM>1e+03]=np.NaN
    mmEf_JJA[mmEf_JJA>1e+03]=np.NaN
    mmEf_SON[mmEf_SON>1e+03]=np.NaN
    ## seasonal mean daily Evap
    mmEvap_DJF = np.zeros(len(mEvap_DJF)-1)
    mmEvap_MAM = np.zeros(len(mEvap_MAM)-1)
    mmEvap_JJA = np.zeros(len(mEvap_JJA)-1)
    mmEvap_SON = np.zeros(len(mEvap_SON)-1)
    for j in range(1,len(mEvap_MAM)):
        if j==1: mEvap_DJF[j-1]=np.NaN
        else: mmEvap_DJF[j-1]=100.*(mEvap1_DJF[j-1]-mEvap_DJF[j-1])/mEvap_DJF[j-1]
        mmEvap_MAM[j-1]=100.*(mEvap1_MAM[j]-mEvap_MAM[j])/mEvap_MAM[j]
        mmEvap_JJA[j-1]=100.*(mEvap1_JJA[j]-mEvap_JJA[j])/mEvap_JJA[j]
        mmEvap_SON[j-1]=100.*(mEvap1_SON[j]-mEvap_SON[j])/mEvap_SON[j]
    ## seasonal mean daily SWE
    mmSWE_DJF = np.zeros(len(mSWE_DJF)-1)
    mmSWE_MAM = np.zeros(len(mSWE_MAM)-1)
    mmSWE_JJA = np.zeros(len(mSWE_JJA)-1)
    mmSWE_SON = np.zeros(len(mSWE_SON)-1)
    for j in range(1,len(mSWE_MAM)):
        if j==1: mSWE_DJF[j-1]=np.NaN
        else: mmSWE_DJF[j-1]=100.*(mSWE1_DJF[j-1]-mSWE_DJF[j-1])/mSWE_DJF[j-1]
        mmSWE_MAM[j-1]=100.*(mSWE1_MAM[j]-mSWE_MAM[j])/mSWE_MAM[j]
        mmSWE_JJA[j-1]=100.*(mSWE1_JJA[j]-mSWE_JJA[j])/mSWE_JJA[j]
        mmSWE_SON[j-1]=100.*(mSWE1_SON[j]-mSWE_SON[j])/mSWE_SON[j]

    mAM= pd.DataFrame()
    mAM['Mean_Q_DJF'] = mQ_DJF
    mAM['Mean_Q_MAM'] = mQ_MAM
    mAM['Mean_Q_JJA'] = mQ_JJA
    mAM['Mean_Q_SON'] = mQ_SON
    mAM['Mean_ET_DJF'] = mmET_DJF
    mAM['Mean_ET_MAM'] = mmET_MAM
    mAM['Mean_ET_JJA'] = mmET_JJA
    mAM['Mean_ET_SON'] = mmET_SON
    mAM['Mean_Tr_DJF'] = mmTr_DJF
    mAM['Mean_Tr_MAM'] = mmTr_MAM
    mAM['Mean_Tr_JJA'] = mmTr_JJA
    mAM['Mean_Tr_SON'] = mmTr_SON
    mAM['Mean_Ef_DJF'] = mmEf_DJF
    mAM['Mean_Ef_MAM'] = mmEf_MAM
    mAM['Mean_Ef_JJA'] = mmEf_JJA
    mAM['Mean_Ef_SON'] = mmEf_SON
    mAM['Mean_Evap_DJF'] = mmEvap_DJF
    mAM['Mean_Evap_MAM'] = mmEvap_MAM
    mAM['Mean_Evap_JJA'] = mmEvap_JJA
    mAM['Mean_Evap_SON'] = mmEvap_SON
    mAM['Mean_SWE_DJF'] = mmSWE_DJF
    mAM['Mean_SWE_MAM'] = mmSWE_MAM
    mAM['Mean_SWE_JJA'] = mmSWE_JJA
    mAM['Mean_SWE_SON'] = mmSWE_SON

    ##### plotting boxplot
    sns.set_style('ticks')
    plt.figure(figsize=(8,4))
    sns.boxplot(data=mAM,showmeans=True,showfliers=False,linewidth=0.5,meanprops=dict(markersize=2))
    plt.ylabel('relative change (%)')
    plt.xticks(rotation='vertical')
    plt.axhline(linewidth=0.5)
    sns.despine()
    plt.tight_layout()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_seasonal.pdf',dpi=300)

def Ann_Seasonal_Change(ID,scen,d_clim,d_clim1):
    #############################################
    ## variables from climate data simulations (1981-2010)
    #############################################
    q_clim = 1e3*d_clim.Qt
    #S_clim = d_clim.S
    mET_clim = (d_clim.Transpi+d_clim.Efloor+d_clim.Evap).mean(dim=['dlat','dlon'])
    mTrans_clim = d_clim.Transpi.mean(dim=['dlat','dlon'])
    mEfloor_clim = d_clim.Efloor.mean(dim=['dlat','dlon'])
    mEvap_clim = d_clim.Evap.mean(dim=['dlat','dlon'])
    mSWE_clim = d_clim.SWE.mean(dim=['dlat','dlon'])
    # mean daily value within a year
    mA_q_clim = q_clim.groupby('dtime.year').mean()
    #mA_S_clim = S_clim.groupby('dtime.year').mean()
    mA_mET_clim = mET_clim.groupby('dtime.year').mean()
    mA_mTrans_clim = mTrans_clim.groupby('dtime.year').mean()
    mA_mEfloor_clim = mEfloor_clim.groupby('dtime.year').mean()
    mA_mEvap_clim = mEvap_clim.groupby('dtime.year').mean()
    mA_mSWE_clim = mSWE_clim.groupby('dtime.year').mean()
    #############################################
    ## variables from climate data simulations (2071-2100)
    #############################################
    q_clim1 = 1e3*d_clim1.Qt
    #S_clim1 = d_clim1.S
    mET_clim1 = (d_clim1.Transpi+d_clim1.Efloor+d_clim1.Evap).mean(dim=['dlat','dlon'])
    mTrans_clim1 = d_clim1.Transpi.mean(dim=['dlat','dlon'])
    mEfloor_clim1 = d_clim1.Efloor.mean(dim=['dlat','dlon'])
    mEvap_clim1 = d_clim1.Evap.mean(dim=['dlat','dlon'])
    mSWE_clim1 = d_clim1.SWE.mean(dim=['dlat','dlon'])
    ### mean daily value within a year
    mA_q_clim1 = q_clim1.groupby('dtime.year').mean()
    #mA_S_clim1 = S_clim1.groupby('dtime.year').mean()
    mA_mET_clim1 = mET_clim1.groupby('dtime.year').mean()
    mA_mTrans_clim1 = mTrans_clim1.groupby('dtime.year').mean()
    mA_mEfloor_clim1 = mEfloor_clim1.groupby('dtime.year').mean()
    mA_mEvap_clim1 = mEvap_clim1.groupby('dtime.year').mean()
    mA_mSWE_clim1 = mSWE_clim1.groupby('dtime.year').mean()
    ################################################
    ###  seasonal mean daily ###########
    ##################################################
    month_clim = q_clim.groupby('dtime.month').apply(lambda x: x).month
    DJF = (month_clim <= 2) | (month_clim >= 12)
    MAM = (month_clim >= 3) & (month_clim <= 5)
    JJA = (month_clim >= 6) & (month_clim <= 8)
    SON = (month_clim >= 9) & (month_clim <= 11)
    # seasonal daily mean flow
    q_DJF = q_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    q_MAM = q_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    q_JJA = q_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    q_SON = q_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean ET
    mET_DJF = mET_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mET_MAM = mET_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mET_JJA = mET_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mET_SON = mET_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Transpi
    mTrans_DJF = mTrans_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mTrans_MAM = mTrans_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mTrans_JJA = mTrans_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mTrans_SON = mTrans_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Efloor
    mEfloor_DJF = mEfloor_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEfloor_MAM = mEfloor_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEfloor_JJA = mEfloor_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEfloor_SON = mEfloor_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Evap
    mEvap_DJF = mEvap_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEvap_MAM = mEvap_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEvap_JJA = mEvap_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEvap_SON = mEvap_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean SWE
    mSWE_DJF = mSWE_clim.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mSWE_MAM = mSWE_clim.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mSWE_JJA = mSWE_clim.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mSWE_SON = mSWE_clim.where(SON).resample('AS-Sep', 'dtime', how='mean')
    #####################################
    # seasonal daily mean flow for 2070-2099
    #####################################
    month_clim = q_clim1.groupby('dtime.month').apply(lambda x: x).month
    DJF = (month_clim <= 2) | (month_clim >= 12)
    MAM = (month_clim >= 3) & (month_clim <= 5)
    JJA = (month_clim >= 6) & (month_clim <= 8)
    SON = (month_clim >= 9) & (month_clim <= 11)
    q1_DJF = q_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    q1_MAM = q_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    q1_JJA = q_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    q1_SON = q_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean ET
    mET1_DJF = mET_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mET1_MAM = mET_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mET1_JJA = mET_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mET1_SON = mET_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Transpi
    mTrans1_DJF = mTrans_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mTrans1_MAM = mTrans_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mTrans1_JJA = mTrans_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mTrans1_SON = mTrans_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Efloor
    mEfloor1_DJF = mEfloor_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEfloor1_MAM = mEfloor_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEfloor1_JJA = mEfloor_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEfloor1_SON = mEfloor_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean Evap
    mEvap1_DJF = mEvap_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mEvap1_MAM = mEvap_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mEvap1_JJA = mEvap_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mEvap1_SON = mEvap_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    # seasonal daily mean SWE
    mSWE1_DJF = mSWE_clim1.where(DJF).resample('AS-Dec', 'dtime', how='mean')
    mSWE1_MAM = mSWE_clim1.where(MAM).resample('AS-Mar', 'dtime', how='mean')
    mSWE1_JJA = mSWE_clim1.where(JJA).resample('AS-Jun', 'dtime', how='mean')
    mSWE1_SON = mSWE_clim1.where(SON).resample('AS-Sep', 'dtime', how='mean')
    ################################################
    ###  data in a Dataframe for boxplot ###########
    ##################################################
    mAQ = np.zeros(len(mA_q_clim))
    mAET = np.zeros(len(mA_mET_clim))
    mATr = np.zeros(len(mA_mTrans_clim))
    mAEf = np.zeros(len(mA_mEfloor_clim))
    mAEv = np.zeros(len(mA_mEvap_clim))
    mASWE = np.zeros(len(mA_mSWE_clim))
    for j in range(len(mA_q_clim)):
        mAQ[j]=100.*(mA_q_clim1[j]-mA_q_clim[j])/mA_q_clim[j]
        mAET[j]=100.*(mA_mET_clim1[j]-mA_mET_clim[j])/mA_mET_clim[j]
        mATr[j]=100.*(mA_mTrans_clim1[j]-mA_mTrans_clim[j])/mA_mTrans_clim[j]
        mAEf[j]=100.*(mA_mEfloor_clim1[j]-mA_mEfloor_clim[j])/mA_mEfloor_clim[j]
        mAEv[j]=100.*(mA_mEvap_clim1[j]-mA_mEvap_clim[j])/mA_mEvap_clim[j]
        mASWE[j]=100.*(mA_mSWE_clim1[j]-mA_mSWE_clim[j])/mA_mSWE_clim[j]
    ## seasonal mean daily flow
    mQ_DJF = np.zeros(len(q_DJF)-1)
    mQ_MAM = np.zeros(len(q_MAM)-1)
    mQ_JJA = np.zeros(len(q_JJA)-1)
    mQ_SON = np.zeros(len(q_SON)-1)
    for j in range(1,len(q_MAM)):
        if j==1: mQ_DJF[j-1]=np.NaN
        else: mQ_DJF[j-1]=100.*(q1_DJF[j-1]-q_DJF[j-1])/q_DJF[j-1]
        mQ_MAM[j-1]=100.*(q1_MAM[j]-q_MAM[j])/q_MAM[j]
        mQ_JJA[j-1]=100.*(q1_JJA[j]-q_JJA[j])/q_JJA[j]
        mQ_SON[j-1]=100.*(q1_SON[j]-q_SON[j])/q_SON[j]
    ## seasonal mean daily ET
    mmET_DJF = np.zeros(len(mET_DJF)-1)
    mmET_MAM = np.zeros(len(mET_MAM)-1)
    mmET_JJA = np.zeros(len(mET_JJA)-1)
    mmET_SON = np.zeros(len(mET_SON)-1)
    for j in range(1,len(mET_MAM)):
        if j==1: mmET_DJF[j-1]=np.NaN
        else: mmET_DJF[j-1]=100.*(mET1_DJF[j-1]-mET_DJF[j-1])/mET_DJF[j-1]
        mmET_MAM[j-1]=100.*(mET1_MAM[j]-mET_MAM[j])/mET_MAM[j]
        mmET_JJA[j-1]=100.*(mET1_JJA[j]-mET_JJA[j])/mET_JJA[j]
        mmET_SON[j-1]=100.*(mET1_SON[j]-mET_SON[j])/mET_SON[j]
    ## seasonal mean daily Trans
    mmTr_DJF = np.zeros(len(mTrans_DJF)-1)
    mmTr_MAM = np.zeros(len(mTrans_MAM)-1)
    mmTr_JJA = np.zeros(len(mTrans_JJA)-1)
    mmTr_SON = np.zeros(len(mTrans_SON)-1)
    for j in range(1,len(mTrans_MAM)):
        if j==1: mmTr_DJF[j-1]=np.NaN
        else: mmTr_DJF[j-1]=100.*(mTrans1_DJF[j-1]-mTrans_DJF[j-1])/mTrans_DJF[j-1]
        mmTr_MAM[j-1]=100.*(mTrans1_MAM[j]-mTrans_MAM[j])/mTrans_MAM[j]
        mmTr_JJA[j-1]=100.*(mTrans1_JJA[j]-mTrans_JJA[j])/mTrans_JJA[j]
        mmTr_SON[j-1]=100.*(mTrans1_SON[j]-mTrans_SON[j])/mTrans_SON[j]
    ## seasonal mean daily Efloor
    mmEf_DJF = np.zeros(len(mEfloor_DJF)-1)
    mmEf_MAM = np.zeros(len(mEfloor_MAM)-1)
    mmEf_JJA = np.zeros(len(mEfloor_JJA)-1)
    mmEf_SON = np.zeros(len(mEfloor_SON)-1)
    for j in range(1,len(mEfloor_MAM)):
        if j==1: mmEf_DJF[j-1]=np.NaN
        else: mmEf_DJF[j-1]=100.*(mEfloor1_DJF[j-1]-mEfloor_DJF[j-1])/mEfloor_DJF[j-1]
        mmEf_MAM[j-1]=100.*(mEfloor1_MAM[j]-mEfloor_MAM[j])/mEfloor_MAM[j]
        mmEf_JJA[j-1]=100.*(mEfloor1_JJA[j]-mEfloor_JJA[j])/mEfloor_JJA[j]
        mmEf_SON[j-1]=100.*(mEfloor1_SON[j]-mEfloor_SON[j])/mEfloor_SON[j]
    mmEf_DJF[mmEf_DJF>1e+03]=np.NaN
    mmEf_MAM[mmEf_MAM>1e+03]=np.NaN
    mmEf_JJA[mmEf_JJA>1e+03]=np.NaN
    mmEf_SON[mmEf_SON>1e+03]=np.NaN
    ## seasonal mean daily Evap
    mmEvap_DJF = np.zeros(len(mEvap_DJF)-1)
    mmEvap_MAM = np.zeros(len(mEvap_MAM)-1)
    mmEvap_JJA = np.zeros(len(mEvap_JJA)-1)
    mmEvap_SON = np.zeros(len(mEvap_SON)-1)
    for j in range(1,len(mEvap_MAM)):
        if j==1: mEvap_DJF[j-1]=np.NaN
        else: mmEvap_DJF[j-1]=100.*(mEvap1_DJF[j-1]-mEvap_DJF[j-1])/mEvap_DJF[j-1]
        mmEvap_MAM[j-1]=100.*(mEvap1_MAM[j]-mEvap_MAM[j])/mEvap_MAM[j]
        mmEvap_JJA[j-1]=100.*(mEvap1_JJA[j]-mEvap_JJA[j])/mEvap_JJA[j]
        mmEvap_SON[j-1]=100.*(mEvap1_SON[j]-mEvap_SON[j])/mEvap_SON[j]
    ## seasonal mean daily SWE
    mmSWE_DJF = np.zeros(len(mSWE_DJF)-1)
    mmSWE_MAM = np.zeros(len(mSWE_MAM)-1)
    mmSWE_JJA = np.zeros(len(mSWE_JJA)-1)
    mmSWE_SON = np.zeros(len(mSWE_SON)-1)
    for j in range(1,len(mSWE_MAM)):
        if j==1: mSWE_DJF[j-1]=np.NaN
        else: mmSWE_DJF[j-1]=100.*(mSWE1_DJF[j-1]-mSWE_DJF[j-1])/mSWE_DJF[j-1]
        mmSWE_MAM[j-1]=100.*(mSWE1_MAM[j]-mSWE_MAM[j])/mSWE_MAM[j]
        mmSWE_JJA[j-1]=100.*(mSWE1_JJA[j]-mSWE_JJA[j])/mSWE_JJA[j]
        mmSWE_SON[j-1]=100.*(mSWE1_SON[j]-mSWE_SON[j])/mSWE_SON[j]

    mAM= pd.DataFrame()
    mAM['Mean_Q_ANN'] = mAQ
    mAM['Mean_Q_DJF'] = mQ_DJF
    mAM['Mean_Q_MAM'] = mQ_MAM
    mAM['Mean_Q_JJA'] = mQ_JJA
    mAM['Mean_Q_SON'] = mQ_SON
    mAM['Mean_ET_ANN'] = mAET
    mAM['Mean_ET_DJF'] = mmET_DJF
    mAM['Mean_ET_MAM'] = mmET_MAM
    mAM['Mean_ET_JJA'] = mmET_JJA
    mAM['Mean_ET_SON'] = mmET_SON
    mAM['Mean_Tr_ANN'] = mATr
    mAM['Mean_Tr_DJF'] = mmTr_DJF
    mAM['Mean_Tr_MAM'] = mmTr_MAM
    mAM['Mean_Tr_JJA'] = mmTr_JJA
    mAM['Mean_Tr_SON'] = mmTr_SON
    mAM['Mean_Ef_ANN'] = mAEf
    mAM['Mean_Ef_DJF'] = mmEf_DJF
    mAM['Mean_Ef_MAM'] = mmEf_MAM
    mAM['Mean_Ef_JJA'] = mmEf_JJA
    mAM['Mean_Ef_SON'] = mmEf_SON
    mAM['Mean_Evap_ANN'] = mAEv
    mAM['Mean_Evap_DJF'] = mmEvap_DJF
    mAM['Mean_Evap_MAM'] = mmEvap_MAM
    mAM['Mean_Evap_JJA'] = mmEvap_JJA
    mAM['Mean_Evap_SON'] = mmEvap_SON
    mAM['Mean_SWE_ANN'] = mASWE
    mAM['Mean_SWE_DJF'] = mmSWE_DJF
    mAM['Mean_SWE_MAM'] = mmSWE_MAM
    mAM['Mean_SWE_JJA'] = mmSWE_JJA
    mAM['Mean_SWE_SON'] = mmSWE_SON

    ##### plotting boxplot
    sns.set_style('ticks')
    plt.figure(figsize=(8,4))
    sns.boxplot(data=mAM,showmeans=True,showfliers=False,linewidth=0.5,meanprops=dict(markersize=2))
    plt.ylabel('relative change (%)')
    plt.xticks(rotation='vertical')
    plt.axhline(linewidth=0.5)
    sns.despine()
    plt.tight_layout()
    plt.savefig('C:\Apps\WinPython-64bit-2.7.10.3\SAMULI_ARI\spathy\\results\\climate\\c_'+ID+'_'+scen+'_ann_seasonal.pdf',dpi=300)