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

"""
Some example's of good and bad plotting practices

Resources:
	For all the great color paletes for python, got to
	https://jiffyclub.github.io/palettable/

	For more infomation about which colormaps are safe, see
	http://colorbrewer2.org/
	Note: all colorbrewer colormaps are in the palettable package

	For general infomation about map making see:
	http://gisgeography.com/map-elements-how-to-guide-map-making/
	http://www.esri.com/news/arcuser/0911/making-a-map-meaningful.html

	For a detailed guide to mapmaking see the Mapbox guide pdf in the repo
"""
__title__ = "Plotting and map making"
__author__ = "Arden Burrell"
__version__ = "1.0 (19.01.2018)"
__email__ = "arden.burrell@unsw.edu.au"

#==============================================================================

# general modules
# import os, sys # check for files and so on
import numpy as np # array manipulations

# plotting modules
import matplotlib.pyplot as plt
#import matplotlib.cm as cm
import matplotlib.colors as mpc
from matplotlib import ticker
# import matplotlib as mpl
import pdb
# for mapping
from mpl_toolkits.basemap import Basemap

# for all the awesome color palets!!!!!!!!!!
import palettable

#==============================================================================
# =========== Main function ==========

def main():
	"""

	Call the demonstration functions in order

	"""
	# Show the mapping of years
	yearmapper()

	# Map a variable that diverges
	divergingmap()

	# Map a variable that is in groups
	groupmap()

	# Map a variable that is sequental
	sequentialmap()

#==============================================================================
# =========== plotting functions ==========
def mapper(
	image, cmap, vmin, vmax, region="global", title=None,
	tks=False, tckstr=None, norm=False, nancolour='dimgrey'
	):
	"""
	A general mapping function that can be used to generate many different
	map types
	"""

	# ===== set the spatial cordinates of the grid, mask the ocean and draw coastlines =====
	# Note:
	# 	because all of the imput datastes are numpy arrays, the lon and lat bounds
	# 	and the spatial projection (epsg) are hard coded. This infomation can be
	# 	can be pulled automitically from nectdf, geoTif's and other raster data
	if region == "global":
		# Set the map frame
		map = Basemap(llcrnrlon=-180.0,llcrnrlat=-90.0,urcrnrlon=180,urcrnrlat=90.0, resolution = 'h', epsg=4326)
		#adds national borders
		map.drawcountries()
	elif region == "australia":
		# Set the map frame
		map = Basemap(llcrnrlon=112.0,llcrnrlat=-44.5,urcrnrlon=156.25,urcrnrlat=-10, resolution = 'h', epsg=4326)
		map.drawlsmask(land_color='none', )
	elif region == "mongolia":
		# Set the map frame
		map = Basemap(llcrnrlon=85.0,llcrnrlat=40,urcrnrlon=120,urcrnrlat=52, resolution = 'h', epsg=4326)
		map.drawlsmask(land_color='none', )
		# Draw national borders with a thicker line
		map.drawcountries(linewidth = 1.25)
	else:
		sys.exit("Unknown Grid region")
	# Set the color or NaN's in the data
	map.drawmapboundary(fill_color=nancolour) #used to be grey
	# For land data, it is usefull to mask lakes etc
	map.fillcontinents(color= 'none', lake_color='white')
	# Add all of the coastlines to the map
	map.drawcoastlines()
	# Add the US and Australian state lines (Doesn't yet support other counteries)
	map.drawstates()

	# ====== Add mpc.BoundaryNorm made custom cmaps =====
	if norm == False:
		map.imshow(np.flipud(image), cmap=cmap, vmin=vmin, vmax=vmax)
	else:
		map.imshow(np.flipud(image), cmap=cmap, vmin=vmin, vmax=vmax, norm=norm)

	cb = map.colorbar()

	# ===== Move and adjust the number and location of colorbar ticks =====
	if tks == False:
		tick_locator = ticker.MaxNLocator(nbins=20)
	elif type(tks)==int:
		tick_locator = ticker.MaxNLocator(nbins=tks)
	else:
		tick_locator = ticker.MaxNLocator(nbins=vmax-vmin)
	# Add the changes to the colorbar
	cb.locator = tick_locator
	cb.update_ticks()

	# ===== Overwrite the tick labels and replace them with user defined =====
	if tckstr is not None:
		cb.ax.set_yticklabels(tckstr)

	# ===== add in the Graticules (lon lat grid lines) ======
	if region == "global":
		map.drawparallels(np.arange(-90, 90, 10),labels=[1,0,0,0], dashes=[1,2])
		map.drawmeridians(np.arange(-180, 180, 20),labels=[0,0,0,1], dashes=[1,2])
	elif region == "australia":
		map.drawparallels(np.arange(-50, -10, 10),labels=[1,0,0,0], dashes=[1,2])
		map.drawmeridians(np.arange(110, 156.25, 10),labels=[0,0,0,1], dashes=[1,2])
	elif region == "mongolia":
		map.drawparallels(np.arange(40, 60, 5),labels=[1,0,0,0], dashes=[1,2])
		map.drawmeridians(np.arange(80, 120, 5),labels=[0,0,0,1], dashes=[1,2])

	# ===== Show the plot =====
	plt.show()

def yearmapper():
	"""
	Demonstrate how to map a variable where there a many subdevisions,
	in this case the years of a breakpoint detected by the TSS.RESTREND method
	applied to GIMMSv3.1g NDVI data over mongolia.

	In this data:
		- Nan values mean the method failed
		- 0 value means that the method detected no breakpoints
		- else: the year in which a ecosystem breakpoint occured
	"""

	# Load the data from a numpy array
	mong_bp = np.load("./Data/Mongolia_breakpoint_years.npy")

	# ===== Make a naive plot =====
	# Comments: this plot is meant to be terrible,
	# 	All the values that are relevant are hidden
	# 	There is no colorbar

	# Set the figure number
	#	This isnt strictly nessary, but usefull if you
	# 	want to generate multiple plots then show them together
	plt.figure(1)
	# Make the plot
	plt.imshow(mong_bp)
	plt.show()

	# pdb.set_trace()

	# ===== Make a slightly less terrible one =====
	# Comments: this plot easier to read but not a map
	# 	Added bounds, and a much better colormap
	# 	The colormap is still continuous (hard to read)

	#Hard code the start and end years, (these can be determined automatically)
	vmin = 1987
	vmax = 2011
	# Get a colormap thats a bit better and tweak it
	cmap = plt.cm.jet
	cmap.set_under('w') #sets values <vmin to white in the plot
	cmap.set_bad('dimgrey')

	# Create the image

	plt.imshow(mong_bp, vmin = vmin, vmax = vmax, cmap=cmap)
	plt.colorbar()
	plt.show()

	# ===== Turn this into a proper map with a segmented colorbar =====
	plt.figure(2)
	# Comments: Now its an actual map
	# 	Segmented colorbars are easier to read
	# 	The colormap is still continuous (hard to read)

	# +++++ turn a continuous colormap into a segmented one ++++
	# The continuous color map
	cmap = plt.cm.jet
	# extract all colors from the cmap
	cmaplist = [cmap(i) for i in range(cmap.N)]
	cmap = cmap.from_list('Custom cmap', cmaplist, cmap.N)
	# Define the coverage and segments (max values, min values, number of segments)
	bounds = np.linspace(vmin, vmax, vmax-(vmin-1))
	# Select the colors for each bound using BoundaryNorm function
	norm = mpc.BoundaryNorm(bounds, cmap.N)
	# set pixels below the vmin to color to be white
	cmap.set_under('w')
	# cmap.set_above('w') # set colors above the vmax value to a color
	tks = (vmax-vmin)/2+1

	# +++++ Pass the data and infomation the mapper function +++++
	mapper(mong_bp, cmap, vmin, vmax, region="mongolia", tks=tks, norm=norm)

	# Note: There is almost no reason to use Jet in science,
	# this is one of very few exceptions. Use palettable ones instead!!!!!

def divergingmap():
	"""
	Demonstrate how to map a variable where the values diverge,
	in this case,  the total change in NDVI between 1982-2015 as
	detected by the TSS.RESTREND method	applied to GIMMSv3.1g NDVI
	data.

	In this data:
		- Nan values mean the method failed
		- -1 values have no statistically significant change in vegetation
		- else: total change in NDVI (note:NDVI is normalised between 0 and 1)
	"""

	# ===== Load the data from a numpy array =====
	TC = np.load("./Data/Global_total_change.npy")

	# ===== setup the mapping variables =====
	# set the min and max
	vmax = 0.3
	vmin = -0.3
	# Grab a diverging colormap from palettable
	cmap = mpc.ListedColormap(palettable.colorbrewer.diverging.PiYG_10.mpl_colors)
	# set the color of values of pixels with no significant change (have been set to -1)
	cmap.set_under('w')
	# This number needs to be adjusted to get the ticks tp line up
	tks=11

	# Generate the map
	mapper(TC, cmap, vmin, vmax, region="global", tks=tks)

def sequentialmap():
	"""
	Demonstrate how to map a variable where the values are sequential.
	in this case,  the p between NDVI and rainfall detected by the
	TSS.RESTREND method	applied to GIMMSv3.1g NDVI data.

	In this data:
		- NaN values mean the method failed
		- else: and Rsquared value between 0 and 1
	"""

	# ===== Load the data from a numpy array =====
	pval = np.load("./Data/Global_pval.npy")

	# ===== setup the mapping variables =====
	# set the min and max
	vmax = 1
	vmin = 0
	# Grab a sequential colormap from palettable
	cmap=mpc.ListedColormap(palettable.cmocean.sequential.Thermal_20_r.mpl_colors)
	# set the color of values of pixels with no significant change (have been set to -1)
	# cmap.set_under('w')
	# set the number of ticks
	tks=21

	# ===== Generate the map =====
	mapper(pval, cmap, vmin, vmax, region="global", tks=tks)

	# Note:
	# 	This is a terrible plot. Hopefully MS will
	# 	update the script with a function that allows
	#	colormas to be scalled.

def groupmap():
	"""
	Demonstrate how to map a variable where the values are in
	classified groups. in this case,  the comparison of TSS.RESTREND
	direction between GIMMSv3.1g NDVI and VOD data.

	In this data:
		- The numbers refer to groups (0 - 8)
	"""

	# ===== Load the data from a numpy array =====
	groups = np.load("./Data/Global_classes.npy")
	# show what the data looks like
	plt.imshow(groups)
	plt.colorbar()
	plt.show()

	# ===== setup the mapping variables =====
	# set the min and max
	vmin = -0.5
	vmax = 8.5
	# Create the labels
	label2=['No Change', 'Both Dec', 'Both Inc','NDVI Inc, VOD Dec','NDVI Inc, VOD stable', 'NDVI stable, VOD Dec','VOD Inc, NDVI Dec', 'VOD Inc, NDVI stable', 'VOD stable, NDVI Dec']
	# Create a custom discrete colourmap from hex values
	cmaphex = ['#d1cdca','#AF0627', '#ffffff','#1a9850', '#91cf60','#d9ef8b','#6a51a3','#9e9ac8', '#cbc9e2']
	cmap = mpc.ListedColormap(cmaphex)
	# Set the color under
	cmap.set_under('w')

	# ===== Generate the map =====
	mapper(groups, cmap, vmin, vmax, region="global", tks=True, tckstr = label2)

if __name__ == '__main__':
	main()


"""
`Introduction <introyt1_tutorial.html>`_ ||
**Tensors** ||
`Autograd <autogradyt_tutorial.html>`_ ||
`Building Models <modelsyt_tutorial.html>`_ ||
`TensorBoard Support <tensorboardyt_tutorial.html>`_ ||
`Training Models <trainingyt_tutorial.html>`_ ||
`Model Understanding <captumyt_tutorial.html>`_

Introduction to PyTorch Tensors
===============================

Follow along with the video below or on `youtube <https://www.youtube.com/watch?v=r7QDUPb2dCM>`__.

.. raw:: html

Tensors are the central data abstraction in PyTorch. This interactive
notebook provides an in-depth introduction to the ``torch.Tensor``
class.

First things first, let’s import the PyTorch module. We’ll also add
Python’s math module to facilitate some of the examples.

"""

import torch
import math

#########################################################################
# Creating Tensors
# ----------------
#
# The simplest way to create a tensor is with the ``torch.empty()`` call:
#

x = torch.empty(3, 4)
print(type(x))
print(x)

##########################################################################
# Let’s unpack what we just did:
#
# -  We created a tensor using one of the numerous factory methods
#    attached to the ``torch`` module.
# -  The tensor itself is 2-dimensional, having 3 rows and 4 columns.
# -  The type of the object returned is ``torch.Tensor``, which is an
#    alias for ``torch.FloatTensor``; by default, PyTorch tensors are
#    populated with 32-bit floating point numbers. (More on data types
#    below.)
# -  You will probably see some random-looking values when printing your
#    tensor. The ``torch.empty()`` call allocates memory for the tensor,
#    but does not initialize it with any values - so what you’re seeing is
#    whatever was in memory at the time of allocation.
#
# A brief note about tensors and their number of dimensions, and
# terminology:
#
# -  You will sometimes see a 1-dimensional tensor called a
#    *vector.*
# -  Likewise, a 2-dimensional tensor is often referred to as a
#    *matrix.*
# -  Anything with more than two dimensions is generally just
#    called a tensor.
#
# More often than not, you’ll want to initialize your tensor with some
# value. Common cases are all zeros, all ones, or random values, and the
# ``torch`` module provides factory methods for all of these:
#

zeros = torch.zeros(2, 3)
print(zeros)

ones = torch.ones(2, 3)
print(ones)

torch.manual_seed(1729)
random = torch.rand(2, 3)
print(random)

#########################################################################
# The factory methods all do just what you’d expect - we have a tensor
# full of zeros, another full of ones, and another with random values
# between 0 and 1.
#
# Random Tensors and Seeding
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Speaking of the random tensor, did you notice the call to
# ``torch.manual_seed()`` immediately preceding it? Initializing tensors,
# such as a model’s learning weights, with random values is common but
# there are times - especially in research settings - where you’ll want
# some assurance of the reproducibility of your results. Manually setting
# your random number generator’s seed is the way to do this. Let’s look
# more closely:
#

torch.manual_seed(1729)
random1 = torch.rand(2, 3)
print(random1)

random2 = torch.rand(2, 3)
print(random2)

torch.manual_seed(1729)
random3 = torch.rand(2, 3)
print(random3)

random4 = torch.rand(2, 3)
print(random4)

############################################################################
# What you should see above is that ``random1`` and ``random3`` carry
# identical values, as do ``random2`` and ``random4``. Manually setting
# the RNG’s seed resets it, so that identical computations depending on
# random number should, in most settings, provide identical results.
#
# For more information, see the `PyTorch documentation on
# reproducibility <https://pytorch.org/docs/stable/notes/randomness.html>`__.
#
# Tensor Shapes
# ~~~~~~~~~~~~~
#
# Often, when you’re performing operations on two or more tensors, they
# will need to be of the same *shape* - that is, having the same number of
# dimensions and the same number of cells in each dimension. For that, we
# have the ``torch.*_like()`` methods:
#

x = torch.empty(2, 2, 3)
print(x.shape)
print(x)

empty_like_x = torch.empty_like(x)
print(empty_like_x.shape)
print(empty_like_x)

zeros_like_x = torch.zeros_like(x)
print(zeros_like_x.shape)
print(zeros_like_x)

ones_like_x = torch.ones_like(x)
print(ones_like_x.shape)
print(ones_like_x)

rand_like_x = torch.rand_like(x)
print(rand_like_x.shape)
print(rand_like_x)

#########################################################################
# The first new thing in the code cell above is the use of the ``.shape``
# property on a tensor. This property contains a list of the extent of
# each dimension of a tensor - in our case, ``x`` is a three-dimensional
# tensor with shape 2 x 2 x 3.
#
# Below that, we call the ``.empty_like()``, ``.zeros_like()``,
# ``.ones_like()``, and ``.rand_like()`` methods. Using the ``.shape``
# property, we can verify that each of these methods returns a tensor of
# identical dimensionality and extent.
#
# The last way to create a tensor that will cover is to specify its data
# directly from a PyTorch collection:
#

some_constants = torch.tensor([[3.1415926, 2.71828], [1.61803, 0.0072897]])
print(some_constants)

some_integers = torch.tensor((2, 3, 5, 7, 11, 13, 17, 19))
print(some_integers)

more_integers = torch.tensor(((2, 4, 6), [3, 6, 9]))
print(more_integers)

######################################################################
# Using ``torch.tensor()`` is the most straightforward way to create a
# tensor if you already have data in a Python tuple or list. As shown
# above, nesting the collections will result in a multi-dimensional
# tensor.
#
# .. note::
#      ``torch.tensor()`` creates a copy of the data.
#
# Tensor Data Types
# ~~~~~~~~~~~~~~~~~
#
# Setting the datatype of a tensor is possible a couple of ways:
#

a = torch.ones((2, 3), dtype=torch.int16)
print(a)

b = torch.rand((2, 3), dtype=torch.float64) * 20.
print(b)

c = b.to(torch.int32)
print(c)

##########################################################################
# The simplest way to set the underlying data type of a tensor is with an
# optional argument at creation time. In the first line of the cell above,
# we set ``dtype=torch.int16`` for the tensor ``a``. When we print ``a``,
# we can see that it’s full of ``1`` rather than ``1.`` - Python’s subtle
# cue that this is an integer type rather than floating point.
#
# Another thing to notice about printing ``a`` is that, unlike when we
# left ``dtype`` as the default (32-bit floating point), printing the
# tensor also specifies its ``dtype``.
#
# You may have also spotted that we went from specifying the tensor’s
# shape as a series of integer arguments, to grouping those arguments in a
# tuple. This is not strictly necessary - PyTorch will take a series of
# initial, unlabeled integer arguments as a tensor shape - but when adding
# the optional arguments, it can make your intent more readable.
#
# The other way to set the datatype is with the ``.to()`` method. In the
# cell above, we create a random floating point tensor ``b`` in the usual
# way. Following that, we create ``c`` by converting ``b`` to a 32-bit
# integer with the ``.to()`` method. Note that ``c`` contains all the same
# values as ``b``, but truncated to integers.
#
# Available data types include:
#
# -  ``torch.bool``
# -  ``torch.int8``
# -  ``torch.uint8``
# -  ``torch.int16``
# -  ``torch.int32``
# -  ``torch.int64``
# -  ``torch.half``
# -  ``torch.float``
# -  ``torch.double``
# -  ``torch.bfloat``
#
# Math & Logic with PyTorch Tensors
# ---------------------------------
#
# Now that you know some of the ways to create a tensor… what can you do
# with them?
#
# Let’s look at basic arithmetic first, and how tensors interact with
# simple scalars:
#

ones = torch.zeros(2, 2) + 1
twos = torch.ones(2, 2) * 2
threes = (torch.ones(2, 2) * 7 - 1) / 2
fours = twos ** 2
sqrt2s = twos ** 0.5

print(ones)
print(twos)
print(threes)
print(fours)
print(sqrt2s)

##########################################################################
# As you can see above, arithmetic operations between tensors and scalars,
# such as addition, subtraction, multiplication, division, and
# exponentiation are distributed over every element of the tensor. Because
# the output of such an operation will be a tensor, you can chain them
# together with the usual operator precedence rules, as in the line where
# we create ``threes``.
#
# Similar operations between two tensors also behave like you’d
# intuitively expect:
#

powers2 = twos ** torch.tensor([[1, 2], [3, 4]])
print(powers2)

fives = ones + fours
print(fives)

dozens = threes * fours
print(dozens)

##########################################################################
# It’s important to note here that all of the tensors in the previous code
# cell were of identical shape. What happens when we try to perform a
# binary operation on tensors if dissimilar shape?
#
# .. note::
#      The following cell throws a run-time error. This is intentional.
#
# ::
#
#    a = torch.rand(2, 3)
#    b = torch.rand(3, 2)
#
#    print(a * b)
#

##########################################################################
# In the general case, you cannot operate on tensors of different shape
# this way, even in a case like the cell above, where the tensors have an
# identical number of elements.
#
# In Brief: Tensor Broadcasting
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# .. note::
#      If you are familiar with broadcasting semantics in NumPy
#      ndarrays, you’ll find the same rules apply here.
#
# The exception to the same-shapes rule is *tensor broadcasting.* Here’s
# an example:
#

rand = torch.rand(2, 4)
doubled = rand * (torch.ones(1, 4) * 2)

print(rand)
print(doubled)

#########################################################################
# What’s the trick here? How is it we got to multiply a 2x4 tensor by a
# 1x4 tensor?
#
# Broadcasting is a way to perform an operation between tensors that have
# similarities in their shapes. In the example above, the one-row,
# four-column tensor is multiplied by *both rows* of the two-row,
# four-column tensor.
#
# This is an important operation in Deep Learning. The common example is
# multiplying a tensor of learning weights by a *batch* of input tensors,
# applying the operation to each instance in the batch separately, and
# returning a tensor of identical shape - just like our (2, 4) \* (1, 4)
# example above returned a tensor of shape (2, 4).
#
# The rules for broadcasting are:
#
# -  Each tensor must have at least one dimension - no empty tensors.
#
# -  Comparing the dimension sizes of the two tensors, *going from last to
#    first:*
#
#    -  Each dimension must be equal, *or*
#
#    -  One of the dimensions must be of size 1, *or*
#
#    -  The dimension does not exist in one of the tensors
#
# Tensors of identical shape, of course, are trivially “broadcastable”, as
# you saw earlier.
#
# Here are some examples of situations that honor the above rules and
# allow broadcasting:
#

a =     torch.ones(4, 3, 2)

b = a * torch.rand(   3, 2) # 3rd & 2nd dims identical to a, dim 1 absent
print(b)

c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
print(c)

d = a * torch.rand(   1, 2) # 3rd dim identical to a, 2nd dim = 1
print(d)

#############################################################################
# Look closely at the values of each tensor above:
#
# -  The multiplication operation that created ``b`` was
#    broadcast over every “layer” of ``a``.
# -  For ``c``, the operation was broadcast over ever layer and row of
#    ``a`` - every 3-element column is identical.
# -  For ``d``, we switched it around - now every *row* is identical,
#    across layers and columns.
#
# For more information on broadcasting, see the `PyTorch
# documentation <https://pytorch.org/docs/stable/notes/broadcasting.html>`__
# on the topic.
#
# Here are some examples of attempts at broadcasting that will fail:
#
# .. note::
#      The following cell throws a run-time error. This is intentional.
#
# ::
#
#    a =     torch.ones(4, 3, 2)
#
#    b = a * torch.rand(4, 3)    # dimensions must match last-to-first
#
#    c = a * torch.rand(   2, 3) # both 3rd & 2nd dims different
#
#    d = a * torch.rand((0, ))   # can't broadcast with an empty tensor
#

###########################################################################
# More Math with Tensors
# ~~~~~~~~~~~~~~~~~~~~~~
#
# PyTorch tensors have over three hundred operations that can be performed
# on them.
#
# Here is a small sample from some of the major categories of operations:
#

# common functions
a = torch.rand(2, 4) * 2 - 1
print('Common functions:')
print(torch.abs(a))
print(torch.ceil(a))
print(torch.floor(a))
print(torch.clamp(a, -0.5, 0.5))

# trigonometric functions and their inverses
angles = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
sines = torch.sin(angles)
inverses = torch.asin(sines)
print('\nSine and arcsine:')
print(angles)
print(sines)
print(inverses)

# bitwise operations
print('\nBitwise XOR:')
b = torch.tensor([1, 5, 11])
c = torch.tensor([2, 7, 10])
print(torch.bitwise_xor(b, c))

# comparisons:
print('\nBroadcasted, element-wise equality comparison:')
d = torch.tensor([[1., 2.], [3., 4.]])
e = torch.ones(1, 2)  # many comparison ops support broadcasting!
print(torch.eq(d, e)) # returns a tensor of type bool

# reductions:
print('\nReduction ops:')
print(torch.max(d))        # returns a single-element tensor
print(torch.max(d).item()) # extracts the value from the returned tensor
print(torch.mean(d))       # average
print(torch.std(d))        # standard deviation
print(torch.prod(d))       # product of all numbers
print(torch.unique(torch.tensor([1, 2, 1, 2, 1, 2]))) # filter unique elements

# vector and linear algebra operations
v1 = torch.tensor([1., 0., 0.])         # x unit vector
v2 = torch.tensor([0., 1., 0.])         # y unit vector
m1 = torch.rand(2, 2)                   # random matrix
m2 = torch.tensor([[3., 0.], [0., 3.]]) # three times identity matrix

print('\nVectors & Matrices:')
print(torch.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.svd(m3))       # singular value decomposition

##################################################################################
# This is a small sample of operations. For more details and the full inventory of
# math functions, have a look at the
# `documentation <https://pytorch.org/docs/stable/torch.html#math-operations>`__.
#
# Altering Tensors in Place
# ~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Most binary operations on tensors will return a third, new tensor. When
# we say ``c = a * b`` (where ``a`` and ``b`` are tensors), the new tensor
# ``c`` will occupy a region of memory distinct from the other tensors.
#
# There are times, though, that you may wish to alter a tensor in place -
# for example, if you’re doing an element-wise computation where you can
# discard intermediate values. For this, most of the math functions have a
# version with an appended underscore (``_``) that will alter a tensor in
# place.
#
# For example:
#

a = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('a:')
print(a)
print(torch.sin(a))   # this operation creates a new tensor in memory
print(a)              # a has not changed

b = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('\nb:')
print(b)
print(torch.sin_(b))  # note the underscore
print(b)              # b has changed

#######################################################################
# For arithmetic operations, there are functions that behave similarly:
#

a = torch.ones(2, 2)
b = torch.rand(2, 2)

print('Before:')
print(a)
print(b)
print('\nAfter adding:')
print(a.add_(b))
print(a)
print(b)
print('\nAfter multiplying')
print(b.mul_(b))
print(b)

##########################################################################
# Note that these in-place arithmetic functions are methods on the
# ``torch.Tensor`` object, not attached to the ``torch`` module like many
# other functions (e.g., ``torch.sin()``). As you can see from
# ``a.add_(b)``, *the calling tensor is the one that gets changed in
# place.*
#
# There is another option for placing the result of a computation in an
# existing, allocated tensor. Many of the methods and functions we’ve seen
# so far - including creation methods! - have an ``out`` argument that
# lets you specify a tensor to receive the output. If the ``out`` tensor
# is the correct shape and ``dtype``, this can happen without a new memory
# allocation:
#

a = torch.rand(2, 2)
b = torch.rand(2, 2)
c = torch.zeros(2, 2)
old_id = id(c)

print(c)
d = torch.matmul(a, b, out=c)
print(c)                # contents of c have changed

assert c is d           # test c & d are same object, not just containing equal values
assert id(c), old_id    # make sure that our new c is the same object as the old one

torch.rand(2, 2, out=c) # works for creation too!
print(c)                # c has changed again
assert id(c), old_id    # still the same object!

##########################################################################
# Copying Tensors
# ---------------
#
# As with any object in Python, assigning a tensor to a variable makes the
# variable a *label* of the tensor, and does not copy it. For example:
#

a = torch.ones(2, 2)
b = a

a[0][1] = 561  # we change a...
print(b)       # ...and b is also altered

######################################################################
# But what if you want a separate copy of the data to work on? The
# ``clone()`` method is there for you:
#

a = torch.ones(2, 2)
b = a.clone()

assert b is not a      # different objects in memory...
print(torch.eq(a, b))  # ...but still with the same contents!

a[0][1] = 561          # a changes...
print(b)               # ...but b is still all ones

#########################################################################
# **There is an important thing to be aware of when using ``clone()``.**
# If your source tensor has autograd, enabled then so will the clone.
# **This will be covered more deeply in the video on autograd,** but if
# you want the light version of the details, continue on.
#
# *In many cases, this will be what you want.* For example, if your model
# has multiple computation paths in its ``forward()`` method, and *both*
# the original tensor and its clone contribute to the model’s output, then
# to enable model learning you want autograd turned on for both tensors.
# If your source tensor has autograd enabled (which it generally will if
# it’s a set of learning weights or derived from a computation involving
# the weights), then you’ll get the result you want.
#
# On the other hand, if you’re doing a computation where *neither* the
# original tensor nor its clone need to track gradients, then as long as
# the source tensor has autograd turned off, you’re good to go.
#
# *There is a third case,* though: Imagine you’re performing a computation
# in your model’s ``forward()`` function, where gradients are turned on
# for everything by default, but you want to pull out some values
# mid-stream to generate some metrics. In this case, you *don’t* want the
# cloned copy of your source tensor to track gradients - performance is
# improved with autograd’s history tracking turned off. For this, you can
# use the ``.detach()`` method on the source tensor:
#

a = torch.rand(2, 2, requires_grad=True) # turn on autograd
print(a)

b = a.clone()
print(b)

c = a.detach().clone()
print(c)

print(a)

#########################################################################
# What’s happening here?
#
# -  We create ``a`` with ``requires_grad=True`` turned on. **We haven’t
#    covered this optional argument yet, but will during the unit on
#    autograd.**
# -  When we print ``a``, it informs us that the property
#    ``requires_grad=True`` - this means that autograd and computation
#    history tracking are turned on.
# -  We clone ``a`` and label it ``b``. When we print ``b``, we can see
#    that it’s tracking its computation history - it has inherited
#    ``a``\ ’s autograd settings, and added to the computation history.
# -  We clone ``a`` into ``c``, but we call ``detach()`` first.
# -  Printing ``c``, we see no computation history, and no
#    ``requires_grad=True``.
#
# The ``detach()`` method *detaches the tensor from its computation
# history.* It says, “do whatever comes next as if autograd was off.” It
# does this *without* changing ``a`` - you can see that when we print
# ``a`` again at the end, it retains its ``requires_grad=True`` property.
#
# Moving to GPU
# -------------
#
# One of the major advantages of PyTorch is its robust acceleration on
# CUDA-compatible Nvidia GPUs. (“CUDA” stands for *Compute Unified Device
# Architecture*, which is Nvidia’s platform for parallel computing.) So
# far, everything we’ve done has been on CPU. How do we move to the faster
# hardware?
#
# First, we should check whether a GPU is available, with the
# ``is_available()`` method.
#
# .. note::
#      If you do not have a CUDA-compatible GPU and CUDA drivers
#      installed, the executable cells in this section will not execute any
#      GPU-related code.
#

if torch.cuda.is_available():
    print('We have a GPU!')
else:
    print('Sorry, CPU only.')

##########################################################################
# Once we’ve determined that one or more GPUs is available, we need to put
# our data someplace where the GPU can see it. Your CPU does computation
# on data in your computer’s RAM. Your GPU has dedicated memory attached
# to it. Whenever you want to perform a computation on a device, you must
# move *all* the data needed for that computation to memory accessible by
# that device. (Colloquially, “moving the data to memory accessible by the
# GPU” is shorted to, “moving the data to the GPU”.)
#
# There are multiple ways to get your data onto your target device. You
# may do it at creation time:
#

if torch.cuda.is_available():
    gpu_rand = torch.rand(2, 2, device='cuda')
    print(gpu_rand)
else:
    print('Sorry, CPU only.')

##########################################################################
# By default, new tensors are created on the CPU, so we have to specify
# when we want to create our tensor on the GPU with the optional
# ``device`` argument. You can see when we print the new tensor, PyTorch
# informs us which device it’s on (if it’s not on CPU).
#
# You can query the number of GPUs with ``torch.cuda.device_count()``. If
# you have more than one GPU, you can specify them by index:
# ``device='cuda:0'``, ``device='cuda:1'``, etc.
#
# As a coding practice, specifying our devices everywhere with string
# constants is pretty fragile. In an ideal world, your code would perform
# robustly whether you’re on CPU or GPU hardware. You can do this by
# creating a device handle that can be passed to your tensors instead of a
# string:
#

if torch.cuda.is_available():
    my_device = torch.device('cuda')
else:
    my_device = torch.device('cpu')
print('Device: {}'.format(my_device))

x = torch.rand(2, 2, device=my_device)
print(x)

#########################################################################
# If you have an existing tensor living on one device, you can move it to
# another with the ``to()`` method. The following line of code creates a
# tensor on CPU, and moves it to whichever device handle you acquired in
# the previous cell.
#

y = torch.rand(2, 2)
y = y.to(my_device)

##########################################################################
# It is important to know that in order to do computation involving two or
# more tensors, *all of the tensors must be on the same device*. The
# following code will throw a runtime error, regardless of whether you
# have a GPU device available:
#
# ::
#
#    x = torch.rand(2, 2)
#    y = torch.rand(2, 2, device='gpu')
#    z = x + y  # exception will be thrown
#

###########################################################################
# Manipulating Tensor Shapes
# --------------------------
#
# Sometimes, you’ll need to change the shape of your tensor. Below, we’ll
# look at a few common cases, and how to handle them.
#
# Changing the Number of Dimensions
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# One case where you might need to change the number of dimensions is
# passing a single instance of input to your model. PyTorch models
# generally expect *batches* of input.
#
# For example, imagine having a model that works on 3 x 226 x 226 images -
# a 226-pixel square with 3 color channels. When you load and transform
# it, you’ll get a tensor of shape ``(3, 226, 226)``. Your model, though,
# is expecting input of shape ``(N, 3, 226, 226)``, where ``N`` is the
# number of images in the batch. So how do you make a batch of one?
#

a = torch.rand(3, 226, 226)
b = a.unsqueeze(0)

print(a.shape)
print(b.shape)

##########################################################################
# The ``unsqueeze()`` method adds a dimension of extent 1.
# ``unsqueeze(0)`` adds it as a new zeroth dimension - now you have a
# batch of one!
#
# So if that’s *un*\ squeezing? What do we mean by squeezing? We’re taking
# advantage of the fact that any dimension of extent 1 *does not* change
# the number of elements in the tensor.
#

c = torch.rand(1, 1, 1, 1, 1)
print(c)

##########################################################################
# Continuing the example above, let’s say the model’s output is a
# 20-element vector for each input. You would then expect the output to
# have shape ``(N, 20)``, where ``N`` is the number of instances in the
# input batch. That means that for our single-input batch, we’ll get an
# output of shape ``(1, 20)``.
#
# What if you want to do some *non-batched* computation with that output -
# something that’s just expecting a 20-element vector?
#

a = torch.rand(1, 20)
print(a.shape)
print(a)

b = a.squeeze(0)
print(b.shape)
print(b)

c = torch.rand(2, 2)
print(c.shape)

d = c.squeeze(0)
print(d.shape)

#########################################################################
# You can see from the shapes that our 2-dimensional tensor is now
# 1-dimensional, and if you look closely at the output of the cell above
# you’ll see that printing ``a`` shows an “extra” set of square brackets
# ``[]`` due to having an extra dimension.
#
# You may only ``squeeze()`` dimensions of extent 1. See above where we
# try to squeeze a dimension of size 2 in ``c``, and get back the same
# shape we started with. Calls to ``squeeze()`` and ``unsqueeze()`` can
# only act on dimensions of extent 1 because to do otherwise would change
# the number of elements in the tensor.
#
# Another place you might use ``unsqueeze()`` is to ease broadcasting.
# Recall the example above where we had the following code:
#
# ::
#
#    a =     torch.ones(4, 3, 2)
#
#    c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
#    print(c)
#
# The net effect of that was to broadcast the operation over dimensions 0
# and 2, causing the random, 3 x 1 tensor to be multiplied element-wise by
# every 3-element column in ``a``.
#
# What if the random vector had just been 3-element vector? We’d lose the
# ability to do the broadcast, because the final dimensions would not
# match up according to the broadcasting rules. ``unsqueeze()`` comes to
# the rescue:
#

a = torch.ones(4, 3, 2)
b = torch.rand(   3)     # trying to multiply a * b will give a runtime error
c = b.unsqueeze(1)       # change to a 2-dimensional tensor, adding new dim at the end
print(c.shape)
print(a * c)             # broadcasting works again!

######################################################################
# The ``squeeze()`` and ``unsqueeze()`` methods also have in-place
# versions, ``squeeze_()`` and ``unsqueeze_()``:
#

batch_me = torch.rand(3, 226, 226)
print(batch_me.shape)
batch_me.unsqueeze_(0)
print(batch_me.shape)

##########################################################################
# Sometimes you’ll want to change the shape of a tensor more radically,
# while still preserving the number of elements and their contents. One
# case where this happens is at the interface between a convolutional
# layer of a model and a linear layer of the model - this is common in
# image classification models. A convolution kernel will yield an output
# tensor of shape *features x width x height,* but the following linear
# layer expects a 1-dimensional input. ``reshape()`` will do this for you,
# provided that the dimensions you request yield the same number of
# elements as the input tensor has:
#

output3d = torch.rand(6, 20, 20)
print(output3d.shape)

input1d = output3d.reshape(6 * 20 * 20)
print(input1d.shape)

# can also call it as a method on the torch module:
print(torch.reshape(output3d, (6 * 20 * 20,)).shape)

###############################################################################
# .. note::
#      The ``(6 * 20 * 20,)`` argument in the final line of the cell
#      above is because PyTorch expects a **tuple** when specifying a
#      tensor shape - but when the shape is the first argument of a method, it
#      lets us cheat and just use a series of integers. Here, we had to add the
#      parentheses and comma to convince the method that this is really a
#      one-element tuple.
#
# When it can, ``reshape()`` will return a *view* on the tensor to be
# changed - that is, a separate tensor object looking at the same
# underlying region of memory. *This is important:* That means any change
# made to the source tensor will be reflected in the view on that tensor,
# unless you ``clone()`` it.
#
# There *are* conditions, beyond the scope of this introduction, where
# ``reshape()`` has to return a tensor carrying a copy of the data. For
# more information, see the
# `docs <https://pytorch.org/docs/stable/torch.html#torch.reshape>`__.
#

#######################################################################
# NumPy Bridge
# ------------
#
# In the section above on broadcasting, it was mentioned that PyTorch’s
# broadcast semantics are compatible with NumPy’s - but the kinship
# between PyTorch and NumPy goes even deeper than that.
#
# If you have existing ML or scientific code with data stored in NumPy
# ndarrays, you may wish to express that same data as PyTorch tensors,
# whether to take advantage of PyTorch’s GPU acceleration, or its
# efficient abstractions for building ML models. It’s easy to switch
# between ndarrays and PyTorch tensors:
#

import numpy as np

numpy_array = np.ones((2, 3))
print(numpy_array)

pytorch_tensor = torch.from_numpy(numpy_array)
print(pytorch_tensor)

##########################################################################
# PyTorch creates a tensor of the same shape and containing the same data
# as the NumPy array, going so far as to keep NumPy’s default 64-bit float
# data type.
#
# The conversion can just as easily go the other way:
#

pytorch_rand = torch.rand(2, 3)
print(pytorch_rand)

numpy_rand = pytorch_rand.numpy()
print(numpy_rand)

##########################################################################
# It is important to know that these converted objects are using *the same
# underlying memory* as their source objects, meaning that changes to one
# are reflected in the other:
#

numpy_array[1, 1] = 23
print(pytorch_tensor)

pytorch_rand[1, 1] = 17
print(numpy_rand)


# ---
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# jupyter:
#   jupytext:
#     formats: ipynb,md:myst,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.10.0
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# # Autodidax: JAX core from scratch
#
# Ever want to learn how JAX works, but the implementation seemed too
# impenetrable? Well, you're in luck! By reading this tutorial, you'll learn
# every big idea in JAX's core system. You'll even get clued into our weird
# jargon!

# ## Part 1: Transformations as interpreters: standard evaluation, `jvp`, and `vmap`
#
# We want to transform functions that look like this:
#
# ```python
# def f(x):
#   y = sin(x) * 2
#   z = - y + x
#   return z
# ```
#
# Think of functions like `sin` and the arithmetic operations underlying the
# infix operators (`mul`, `add`, and `neg`) as primitive operations, meaning
# atomic units of processing rather than compositions.
#
# "Transform" means "interpret differently." Instead of standard interpretation
# where we apply primitive functions to numerical inputs to produce numerical
# outputs, we want to override primitive application and let different values
# flow through our program. For example, we might want to replace the
# application of every primitive with an application of [its JVP
# rule](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html),
# and let primal-tangent pairs flow through our program. Moreover, we want to
# apply a composition of multiple transformations, leading to stacks of
# interpreters.

# ### JAX core machinery
#
# We can implement stacks of interpreters and even have them all discharge on
# the fly as we execute the Python function to be transformed. To start, let's
# define these primitives so that we can intercept their application:

# +
from typing import NamedTuple

class Primitive(NamedTuple):
  name: str

add_p = Primitive('add')
mul_p = Primitive('mul')
neg_p = Primitive("neg")
sin_p = Primitive("sin")
cos_p = Primitive("cos")
reduce_sum_p = Primitive("reduce_sum")
greater_p = Primitive("greater")

def add(x, y): return bind(add_p, x, y)
def mul(x, y): return bind(mul_p, x, y)
def neg(x): return bind(neg_p, x)
def sin(x): return bind(sin_p, x)
def cos(x): return bind(cos_p, x)
def reduce_sum(x, axis=None): return bind(reduce_sum_p, x, axis=axis)
def greater(x, y): return bind(greater_p, x, y)

# -

# We'll set up array data types and infix operator methods in a moment.
#
# A `Primitive` is just an object with a name, to which we attach our
# interpretation rules (one for each transformation). The `bind` function is our
# interception point: it'll figure out which transformation rule to apply, based
# on how the arguments are boxed in tracers and what interpreters are active.
#
# The functions that user code calls, like `add` and `sin`, are just wrappers
# around calls to `bind`. These wrappers let us control how arguments are passed
# to `bind`, and in particular we follow a handy internal convention: when we
# call `bind`, we pass values representing array data as positional arguments,
# and we pass metadata like the `axis` argument to `sum_p` via keyword. This
# calling convention simplifies some core logic (since e.g. instances of the
# `Tracer` class to be defined below can only occurr in positional arguments to
# `bind`). The wrappers can also provide docstrings!
#
# We represent active interpreters as a stack. The stack is just a simple
# `list`, and each element is a container with an integer level (corresponding
# to the element's height in the stack), an interpreter type (which we'll call a
# `trace_type`), and an optional field for any global data the interpreter
# needs. We call each element a `MainTrace`, though maybe "Interpreter" would be
# more descriptive.

# +
from contextlib import contextmanager
from typing import Type, List, Optional, Any

class MainTrace(NamedTuple):
  level: int
  trace_type: Type['Trace']
  global_data: Optional[Any]

trace_stack: List[MainTrace] = []

@contextmanager
def new_main(trace_type: Type['Trace'], global_data=None):
  level = len(trace_stack)
  main = MainTrace(level, trace_type, global_data)
  trace_stack.append(main)

  try:
    yield main
  finally:
    trace_stack.pop()

# -

# When we're about to apply a transformed function, we'll push another
# interpreter onto the stack using `new_main`. Then, as we apply primitives in
# the function, we can think of the `bind` first being interpreted by the trace
# at the top of the stack (i.e. with the highest level). If that first
# interpreter itself binds other primitives in its interpretation rule for the
# primitive, like how the JVP rule of `sin_p` might bind `cos_p` and `mul_p`,
# then those `bind` calls will be handled by the interpreter at the next level
# down.
#
# What goes at the bottom of the interpreter stack? At the bottom, we know all
# the transformation interpreters are finished, and we just want to do standard
# evaluation. So at the bottom we'll put an evaluation interpreter.
#
# Let's sketch out the interface for interpreters, which is based on the `Trace`
# and `Tracer` base classes. A `Tracer` represents a boxed-up value, perhaps
# carrying some extra context data used by the interpreter. A `Trace` handles
# boxing up vales into `Tracers` and also handles primitive application.

class Trace:
  main: MainTrace

  def __init__(self, main: MainTrace) -> None:
    self.main = main

  def pure(self, val): assert False  # must override
  def lift(self, val): assert False  # must override

  def process_primitive(self, primitive, tracers, params):
    assert False  # must override

# The first two methods are about boxing up values in `Tracer`s, which are the
# objects that flow through the Python programs we transform. The last method is
# the callback we'll use to interpret primitive application.
#
# The `Trace` itself doesn't contain any data, other than a reference to its
# corresponding `MainTrace` instance. In fact, multiple instances of a `Trace`
# might be created and discarded during an application of a transformation,
# whereas only a single `MainTrace` instance is created per application of a
# transformation.
#
# As for `Tracer`s themselves, each one carries an abstract value (and forwards
# infix operators to it), and the rest is up to the transformation. (The
# relationship between `Tracer`s and `AbstractValue`s is that there's one
# `Tracer` per transformation, and at least one `AbstractValue` per base type,
# like arrays.)

# +
import numpy as np
from typing import Tuple

class Tracer:
  _trace: Trace

  __array_priority__ = 1000

  @property
  def aval(self):
    assert False  # must override

  def full_lower(self):
    return self  # default implementation

  def __neg__(self): return self.aval._neg(self)
  def __add__(self, other): return self.aval._add(self, other)
  def __radd__(self, other): return self.aval._radd(self, other)
  def __mul__(self, other): return self.aval._mul(self, other)
  def __rmul__(self, other): return self.aval._rmul(self, other)
  def __gt__(self, other): return self.aval._gt(self, other)
  def __bool__(self): return self.aval._bool(self)
  def __nonzero__(self): return self.aval._nonzero(self)

  def __getattr__(self, name):
    try:
      return getattr(self.aval, name)
    except AttributeError:
      raise AttributeError(f"{self.__class__.__name__} has no attribute {name}")

class ShapedArray:
  array_abstraction_level = 1
  shape: Tuple[int]
  dtype: np.dtype

  def __init__(self, shape, dtype):
    self.shape = shape
    self.dtype = dtype

  @property
  def ndim(self):
    return len(self.shape)

  _neg = staticmethod(neg)
  _add = staticmethod(add)
  _radd = staticmethod(add)
  _mul = staticmethod(mul)
  _rmul = staticmethod(mul)
  _gt = staticmethod(greater)

  @staticmethod
  def _bool(tracer):
    raise Exception("ShapedArray can't be unambiguously converted to bool")

  @staticmethod
  def _nonzero(tracer):
    raise Exception("ShapedArray can't be unambiguously converted to bool")

  def str_short(self):
    return f'{self.dtype.name}[{",".join(str(d) for d in self.shape)}]'

class ConcreteArray(ShapedArray):
  array_abstraction_level = 2
  val: np.ndarray

  def __init__(self, val):
    self.val = val
    self.shape = val.shape
    self.dtype = val.dtype

  @staticmethod
  def _bool(tracer):
    return bool(tracer.aval.val)

  @staticmethod
  def _nonzero(tracer):
    return bool(tracer.aval.val)

def get_aval(x):
  if isinstance(x, Tracer):
    return x.aval
  else:
    return ConcreteArray(np.asarray(x))

# -

# Notice that we actually have two `AbstractValue`s for arrays, representing
# different levels of abstraction. A `ShapedArray` represents the set of all
# possible arrays with a given shape and dtype. A `ConcreteArray` represents a
# singleton set consisting of a single array value.
#
# Now that we've set up the trace stack, the Trace/Tracer API for interpreters,
# and abstract values, we can come back to implement `bind`:

def bind(prim, *args, **params):
  top_trace = find_top_trace(args)
  tracers = [full_raise(top_trace, arg) for arg in args]
  out = top_trace.process_primitive(prim, tracers, params)
  return full_lower(out)

# The main action is that we call `find_top_trace` to figure out which
# interpreter should handle this primitive application as a function of the
# arguments and the active traces on the trace stack. We then call that top
# trace's `process_primitive` so that the trace can apply its interpretation
# rule. The calls to `full_raise` just ensure that the inputs are boxed in the
# top trace's `Tracer` instances, and the call to `full_lower` is an optional
# optimization so that we unbox values out of `Tracer`s as much as possible.

# +
from operator import attrgetter

def find_top_trace(xs) -> Trace:
  top_main = max((x._trace.main for x in xs if isinstance(x, Tracer)),
                 default=trace_stack[0], key=attrgetter('level'))
  return top_main.trace_type(top_main)

# -

# In words, `find_top_trace` returns the highest-level interpreter associated
# with the `Tracer`s on its inputs, and otherwise returns the interpreter at the
# bottom of the stack (which is always an evaluation trace, at least for now).
# This corresponds to JAX transformations mostly working by data dependence
# _except_ for the special bottom-of-the-stack interpreter, which interprets
# everything.

# +
def full_lower(val):
  if isinstance(val, Tracer):
    return val.full_lower()
  else:
    return val

def full_raise(trace, val) -> Tracer:
  if not isinstance(val, Tracer):
    return trace.pure(val)
  level = trace.main.level
  if val._trace.main is trace.main:
    return val
  elif val._trace.main.level < level:
    return trace.lift(val)
  elif val._trace.main.level > level:
    raise Exception(f"Can't lift level {val._trace.main.level} to {level}.")
  else:  # val._trace.level == level
    raise Exception(f"Different traces at same level: {val._trace}, {trace}.")

# -

# The logic in `full_raise` serves to box values into `Tracer`s for a particular
# `Trace`, calling different methods on the `Trace` based on context:
# `Trace.pure` is called on non-`Tracer` constants, and `Trace.lift` is called
# for values that are already `Tracer`s from a lower-level interpreter. These
# two methods could share the same implementation, but by distinguishing them in
# the core logic we can provide more information to the `Trace` subclass.
#
# That's it for the JAX core! Now we can start adding interpreters.

# ### Evaluation interpreter
#
# We'll start with the simplest interpreter: the evaluation interpreter that
# will sit at the bottom of the interpreter stack.

# +
class EvalTrace(Trace):
  pure = lift = lambda self, x: x  # no boxing in Tracers needed

  def process_primitive(self, primitive, tracers, params):
    return impl_rules[primitive](*tracers, **params)

trace_stack.append(MainTrace(0, EvalTrace, None))  # special bottom of the stack

impl_rules = {}
impl_rules[add_p] = np.add
impl_rules[mul_p] = np.multiply
impl_rules[neg_p] = np.negative
impl_rules[sin_p] = np.sin
impl_rules[cos_p] = np.cos
impl_rules[reduce_sum_p] = np.sum
impl_rules[greater_p] = np.greater

# -

# With this interpreter, we can evaluate user functions:

# +
def f(x):
  y = sin(x) * 2
  z = - y + x
  return z

print(f(3.0))

# -

# Woo! Like going around in a big circle. But the point of this indirection is
# that now we can add some real transformations.

# ### Forward-mode autodiff with `jvp`
#
# First, a couple of helper functions:

# +
def zeros_like(val):
  return np.zeros_like(val)

def unzip2(pairs):
  lst1, lst2 = [], []
  for x1, x2 in pairs:
    lst1.append(x1)
    lst2.append(x2)
  return lst1, lst2

# -

# The `Tracer` for forward-mode autodiff carries a primal-tangent pair. The
# `Trace` applies JVP rules.

# +
class JVPTracer(Tracer):
  def __init__(self, trace, primal, tangent):
    self._trace = trace
    self.primal = primal
    self.tangent = tangent

  @property
  def aval(self):
    return get_aval(self.primal)

class JVPTrace(Trace):
  pure = lift = lambda self, val: JVPTracer(self, val, zeros_like(val))

  def process_primitive(self, primitive, tracers, params):
    primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
    jvp_rule = jvp_rules[primitive]
    primal_out, tangent_out = jvp_rule(primals_in, tangents_in, **params)
    return JVPTracer(self, primal_out, tangent_out)

jvp_rules = {}

# -

# Notice both `lift` and `sublift` package a value into a `JVPTracer` with the
# minimal amount of context, which is a zero tangent value.

# Let's add some JVP rules for primitives:

# +
def add_jvp(primals, tangents):
  (x, y), (x_dot, y_dot) = primals, tangents
  return x + y, x_dot + y_dot
jvp_rules[add_p] = add_jvp

def mul_jvp(primals, tangents):
  (x, y), (x_dot, y_dot) = primals, tangents
  return x * y, x_dot * y + x * y_dot
jvp_rules[mul_p] = mul_jvp

def sin_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return sin(x), cos(x) * x_dot
jvp_rules[sin_p] = sin_jvp

def cos_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return cos(x), -sin(x) * x_dot
jvp_rules[cos_p] = cos_jvp

def neg_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return neg(x), neg(x_dot)
jvp_rules[neg_p] = neg_jvp

def reduce_sum_jvp(primals, tangents, *, axis):
  (x,), (x_dot,) = primals, tangents
  return reduce_sum(x, axis), reduce_sum(x_dot, axis)
jvp_rules[reduce_sum_p] = reduce_sum_jvp

def greater_jvp(primals, tangents):
  (x, y), _ = primals, tangents
  out_primal = greater(x, y)
  return out_primal, zeros_like(out_primal)
jvp_rules[greater_p] = greater_jvp

# -

# Finally, we add a transformation API to kick off the trace:

def jvp(f, primals, tangents):
  with new_main(JVPTrace) as main:
    trace = JVPTrace(main)
    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]
    out = f(*tracers_in)
    tracer_out = full_raise(trace, out)
    primal_out, tangent_out = tracer_out.primal, tracer_out.tangent
  return primal_out, tangent_out

# And with that, we can differentiate!

x = 3.0
y, sin_deriv_at_3 = jvp(sin, (x,), (1.0,))
print(sin_deriv_at_3)
print(cos(3.0))

# +
def f(x):
  y = sin(x) * 2
  z = - y + x
  return z

x, xdot = 3., 1.
y, ydot = jvp(f, (x,), (xdot,))
print(y)
print(ydot)

# +
def deriv(f):
  return lambda x: jvp(f, (x,), (1.,))[1]

print(deriv(sin)(3.))
print(deriv(deriv(sin))(3.))
print(deriv(deriv(deriv(sin)))(3.))
print(deriv(deriv(deriv(deriv(sin))))(3.))

# +
def f(x):
  if x > 0.:  # Python control flow
    return 2. * x
  else:
    return x

print(deriv(f)(3.))
print(deriv(f)(-3.))

# -

# ### Vectorized batching with `vmap`
#
# First, a couple helper functions, one for producing mapped abstract values
# from unmapped ones (by removing an axis), and one for moving batch dimensions
# around:

# +
def mapped_aval(batch_dim, aval):
  shape = list(aval.shape)
  del shape[batch_dim]
  return ShapedArray(tuple(shape), aval.dtype)

def move_batch_axis(axis_size, src, dst, x):
  if src is not_mapped:
    target_shape = list(np.shape(x))
    target_shape.insert(dst, axis_size)
    return np.broadcast_to(np.expand_dims(x, dst), target_shape)
  else:
    return np.moveaxis(x, src, dst)

# -

# The `Tracer` for vectorized batching carries a batched value and an optional
# integer indicating which axis (if any) is the batch axis.

# +
from typing import Union

class NotMapped: pass
not_mapped = NotMapped()

class BatchTracer(Tracer):
  def __init__(self, trace, val, batch_dim: Union[NotMapped, int]):
    self._trace = trace
    self.val = val
    self.batch_dim = batch_dim

  @property
  def aval(self):
    if self.batch_dim is not_mapped:
      return get_aval(self.val)
    else:
      return mapped_aval(self.batch_dim, get_aval(self.val))

  def full_lower(self):
    if self.batch_dim is not_mapped:
      return full_lower(self.val)
    else:
      return self

class BatchTrace(Trace):
  pure = lift = lambda self, val: BatchTracer(self, val, not_mapped)

  def process_primitive(self, primitive, tracers, params):
    vals_in, bdims_in = unzip2((t.val, t.batch_dim) for t in tracers)
    vmap_rule = vmap_rules[primitive]
    val_out, bdim_out = vmap_rule(self.axis_size, vals_in, bdims_in, **params)
    return BatchTracer(self, val_out, bdim_out)

  @property
  def axis_size(self):
    return self.main.global_data

vmap_rules = {}
# -

# Here we've implemented the optional `Tracer.full_lower` method, which lets us
# peel off a batching tracer if it's not needed because it doesn't represent a
# batched value.
#
# For `BatchTrace`, analogous to `JVPTrace`, the methods `pure` and `lift` just
# box a value in a `BatchTracer` with the minimal amount of context, which in
# this case is a `batch_dim` taking the sentinel value `not_mapped`. Notice we
# use the `MainTrace`'s interpreter-global data field to store the batch axis
# size.
#
# Next we can define batching interpreter rules for each primitive:

# +
from functools import partial

def broadcasting_binop_batching_rule(op, axis_size, vals_in, dims_in):
  (x, y), (x_bdim, y_bdim) = vals_in, dims_in
  if x_bdim != y_bdim:
    y = move_batch_axis(axis_size, y_bdim, x_bdim, y)
  return op(x, y), x_bdim
vmap_rules[add_p] = partial(broadcasting_binop_batching_rule, add)
vmap_rules[mul_p] = partial(broadcasting_binop_batching_rule, mul)

def vectorized_unop_batching_rule(op, axis_size, vals_in, dims_in):
  (x,), (x_bdim,) = vals_in, dims_in
  return op(x), x_bdim
vmap_rules[sin_p] = partial(vectorized_unop_batching_rule, sin)
vmap_rules[cos_p] = partial(vectorized_unop_batching_rule, cos)
vmap_rules[neg_p] = partial(vectorized_unop_batching_rule, neg)

def reduce_sum_batching_rule(axis_size, vals_in, dims_in, *, axis):
  (x,), (x_bdim,) = vals_in, dims_in
  new_axis = axis + (x_bdim <= axis)
  out_bdim = x_bdim - (new_axis < x_bdim)
  return reduce_sum(x, new_axis), out_bdim
vmap_rules[reduce_sum_p] = reduce_sum_batching_rule

# -

# Finally, we add a transformation API to kick off the trace:

def vmap(f, in_axes, out_axis):
  def batched_f(*args):
    axis_size, = {x.shape[ax] for x, ax in zip(args, in_axes)
                  if ax is not None}
    with new_main(BatchTrace, axis_size) as main:
      trace = BatchTrace(main)
      tracers_in = [BatchTracer(trace, x, ax) if ax is not None else x
                    for x, ax in zip(args, in_axes)]
      out = f(*tracers_in)
      tracer_out = full_raise(trace, out)
      val_out, batch_dim_out = tracer_out.val, tracer_out.batch_dim
    return move_batch_axis(axis_size, batch_dim_out, out_axis, val_out)
  return batched_f

# +
def add_one_to_a_scalar(scalar):
  assert np.ndim(scalar) == 0
  return 1 + scalar

vector_in = np.arange(3.)
vector_out = vmap(add_one_to_a_scalar, (0,), 0)(vector_in)

print(vector_in)
print(vector_out)

# +
def jacfwd(f, x):
  pushfwd = lambda v: jvp(f, (x,), (v,))[1]
  vecs_in = np.eye(np.size(x)).reshape(np.shape(x) * 2)
  return vmap(pushfwd, (0,), 0)(vecs_in)

def f(x):
  return sin(x)

jacfwd(f, np.arange(3.))
# -

# That's it for `jvp` and `vmap`! Before moving on, let's highlight a few
# simplifications in what we've seen so far compared to the full JAX
# implementation:
# 1. **Fewer, simpler primitives.** More primitives means more interpretation
# rules, and for more complex primitives (like for convolution or advanced
# indexing) each rule is harder to write. But the overarching design is no
# different.
# 1. **Transformations expect arrays in, single array out.**
# 2. **No symbolic zeros in autodiff.**
# 3. **No special call primitives yet.** The core machinery needs to be
#     generalized to handle the most flexible kind of higher-order primitive,
#     used by `jax.custom_jvp` and `jax.custom_vjp`.

# ## Part 2: Jaxprs, for `jit` and `vjp`
#
# The next transformations are the horizon are `jit` for just-in-time
# compilation and `vjp` for reverse-mode autodiff.  (`grad` is just a small
# wrapper around `vjp`.) For `jvp` and `vmap` we only needed each `Tracer` to
# carry a little    bit of extra context, but for both `jit` and `vjp` we need
# much richer context: we need to represent _programs_. That is, we need jaxprs!
#
# Jaxprs are JAX's internal intermediate representation of programs. Jaxprs are
# an explicitly typed, functional, first-order language. We need a program
# representation for `jit` because the purpose of `jit` is to stage computation
# out of Python. For any computation we want to stage out, we need to be able to
# represent it as data, and build it up as we trace a Python function.
# Similarly, `vjp` needs a way to represent the computation for the backward
# pass of reverse-mode autodiff. We use the same jaxpr program representation
# for both needs.
#
# (Building a program representation is the most
# [free](https://en.wikipedia.org/wiki/Free_object) kind of
# trace- transformation, and so except for issues around handling native Python
# control flow, any transformation could be implemented by first tracing to a
# jaxpr and then interpreting the jaxpr.)
#
# The jaxpr term syntax is roughly:
#
# ```
# jaxpr ::=
#   { lambda <binder> , ... .
#     let <eqn>
#         ...
#     in <atom> }
#
# binder ::= <var>:<array_type>
# var ::= a | b | c | ...
# atom ::= <var> | <literal>
# literal ::= <int32> | <float32>
#
# eqn ::= <binder> = <primitive> [ <params> ] <atom> , ...
# ```
#
# The syntax of types is:
#
# ```
# jaxpr_type ::= [<array_type>, ...] -> [<array_type>, ...]
# array_type ::= <dtype>[<shape>]
# dtype ::= f32 | f64 | i32 | i64
# shape ::= <int> , ...
# ```
#
# How do we represent these as Python data structures? We reuse ShapedArrays to
# represent types, and we can represent the term syntax with a few Python
# structs:

# +
from typing import Dict, Set

class Var:
  aval: ShapedArray
  def __init__(self, aval): self.aval = aval

class Lit:
  val: Any
  aval: ShapedArray

  def __init__(self, val):
    self.val = val
    self.aval = raise_to_shaped(get_aval(self.val))

Atom = Union[Var, Lit]

class JaxprEqn(NamedTuple):
  primitive: Primitive
  inputs: List[Atom]
  params: Dict[str, Any]
  out_binder: Var

class Jaxpr(NamedTuple):
  in_binders: List[Var]
  eqns: List[JaxprEqn]
  out: Atom

def raise_to_shaped(aval):
  return ShapedArray(aval.shape, aval.dtype)

# +
class JaxprType:
  in_types: List[ShapedArray]
  out_type: ShapedArray

  def __init__(self, in_types, out_type):
    self.in_types = in_types
    self.out_type = out_type

  def __repr__(self):
    in_types = ', '.join(aval.str_short() for aval in self.in_types)
    out_type = self.out_type.str_short()
    return f'({in_types}) -> {out_type}'

def typecheck_jaxpr(jaxpr: Jaxpr) -> JaxprType:
  env: Set[Var] = set()

  for v in jaxpr.in_binders:
    env.add(v)

  for eqn in jaxpr.eqns:
    in_types = [typecheck_atom(env, x) for x in eqn.inputs]
    out_type = abstract_eval_rules[eqn.primitive](*in_types, **eqn.params)
    if not types_equal(out_type, eqn.out_binder.aval): raise TypeError
    env.add(eqn.out_binder)

  out_type = typecheck_atom(env, jaxpr.out)
  return JaxprType([v.aval for v in jaxpr.in_binders], out_type)

def typecheck_atom(env: Set[Var], x: Atom) -> ShapedArray:
  if isinstance(x, Var):
    if x not in env: raise TypeError("unbound variable")
    return x.aval
  elif isinstance(x, Lit):
    return raise_to_shaped(get_aval(x.val))
  else:
    assert False

def types_equal(a: ShapedArray, b: ShapedArray) -> bool:
  return a.shape == b.shape and a.dtype == b.dtype

# -

# Now that we have jaxprs as a data structure, we need ways to produce these
# from tracing Python code. In general there are two variants of how we trace to
# a jaxpr; `jit` uses one and `vjp` uses the other. We'll start with the one
# used   by `jit`, which is also used by control flow primitives like
# `lax.cond`, `lax.while_loop`, and `lax.scan`.

# +
# NB: the analogous class in JAX is called 'DynamicJaxprTracer'
class JaxprTracer(Tracer):
  __slots__ = ['aval']
  aval: ShapedArray

  def __init__(self, trace, aval):
    self._trace = trace
    self.aval = aval

# NB: the analogous class in JAX is called 'DynamicJaxprTrace'
class JaxprTrace(Trace):
  def new_arg(self, aval: ShapedArray) -> JaxprTracer:
    aval = raise_to_shaped(aval)
    tracer = JaxprTracer(self, aval)
    self.builder.tracer_to_var[id(tracer)] = Var(aval)
    return tracer

  def get_or_make_const_tracer(self, val: Any) -> JaxprTracer:
    tracer = self.builder.const_tracers.get(id(val))
    if tracer is None:
      tracer = JaxprTracer(self, raise_to_shaped(get_aval(val)))
      self.builder.add_const(tracer, val)
    return tracer
  pure = lift = get_or_make_const_tracer

  def process_primitive(self, primitive, tracers, params):
    avals_in = [t.aval for t in tracers]
    aval_out = abstract_eval_rules[primitive](*avals_in, **params)
    out_tracer = JaxprTracer(self, aval_out)
    inputs = [self.builder.getvar(t) for t in tracers]
    outvar = self.builder.add_var(out_tracer)
    self.builder.add_eqn(JaxprEqn(primitive, inputs, params, outvar))
    return out_tracer

  @property
  def builder(self):
    return self.main.global_data

# NB: in JAX, instead of a dict we attach impl rules to the Primitive instance
abstract_eval_rules = {}

# -

# Notice that we keep as interpreter-global data a builder object, which keeps
# track of variables, constants, and eqns  as we build up the jaxpr.

class JaxprBuilder:
  eqns: List[JaxprEqn]
  tracer_to_var: Dict[int, Var]
  const_tracers: Dict[int, JaxprTracer]
  constvals: Dict[Var, Any]

  def __init__(self):
    self.eqns = []
    self.tracer_to_var = {}
    self.const_tracers = {}
    self.constvals = {}

  def add_eqn(self, eqn: JaxprEqn) -> None:
    self.eqns.append(eqn)

  def add_var(self, tracer: JaxprTracer) -> Var:
    var = self.tracer_to_var.get(id(tracer))
    assert var is None
    var = self.tracer_to_var[id(tracer)] = Var(tracer.aval)
    return var

  def getvar(self, tracer: JaxprTracer) -> Var:
    var = self.tracer_to_var.get(id(tracer))
    assert var is not None
    return var

  def add_const(self, tracer: JaxprTracer, val: Any) -> Var:
    var = self.add_var(tracer)
    self.const_tracers[id(val)] = tracer
    self.constvals[var] = val
    return var

  def build(self, in_tracers: List[JaxprTracer], out_tracer: JaxprTracer
            ) -> Tuple[Jaxpr, List[Any]]:
    constvars, constvals = unzip2(self.constvals.items())
    t2v = lambda t: self.tracer_to_var[id(t)]
    in_binders = constvars + [t2v(t) for t in in_tracers]
    jaxpr = Jaxpr(in_binders, self.eqns, t2v(out_tracer))
    typecheck_jaxpr(jaxpr)
    return jaxpr, constvals

# The rules we need for `JaxprTrace.process_primitive` are essentially typing
# rules for primitive applications: given   the primitive, its parameters, and
# types for the inputs, the rule must produce a type for the output, which is
# then   packaged with the output `JaxprTracer`. We can use abstract evaluation
# rules for this same purpose, even though they  can be more general (since
# abstract evaluation rules need to work on ConcreteArray inputs as well). We'll
# reuse these abstract evaluation rules for the other jaxpr-producing trace
# machinery, where the potential extra generality is useful.

def broadcast_shapes(*shapes):
  assert len(shapes) > 1
  for sizes in zip(*shapes):
    sizes = [d for d in sizes if d != 1]
    if sizes[:-1] != sizes[1:]:
      raise Exception
  return tuple(next((d for d in sizes if d != 1), 1) for sizes in zip(*shapes))

# +
def broadcasting_binop_abstract_eval_rule(*avals_in):
  out_dtype = np.result_type(*map(np.result_type, avals_in))
  out_shape = broadcast_shapes(*map(np.shape, avals_in))
  return ShapedArray(out_shape, out_dtype)

abstract_eval_rules[add_p] = broadcasting_binop_abstract_eval_rule
abstract_eval_rules[mul_p] = broadcasting_binop_abstract_eval_rule

def vectorized_unop_abstract_eval_rule(aval_in):
  return ShapedArray(np.shape(aval_in), np.result_type(aval_in))

abstract_eval_rules[sin_p] = vectorized_unop_abstract_eval_rule
abstract_eval_rules[cos_p] = vectorized_unop_abstract_eval_rule
abstract_eval_rules[neg_p] = vectorized_unop_abstract_eval_rule

def reduce_sum_abstract_eval_rule(aval_in, *, axis):
  new_shape = [d for i, d in enumerate(aval_in.shape) if i != axis]
  return ShapedArray(tuple(new_shape), aval_in.dtype)
abstract_eval_rules[reduce_sum_p] = reduce_sum_abstract_eval_rule

# -

# To check our implementation, we can add a `make_jaxpr` transformation and
# first pretty-printer:

def make_jaxpr(f, avals_in):
  builder = JaxprBuilder()
  with new_main(JaxprTrace, builder) as main:
    trace = JaxprTrace(main)
    tracers_in = [trace.new_arg(aval) for aval in avals_in]
    out = f(*tracers_in)
    tracer_out = full_raise(trace, out)
    return builder.build(tracers_in, tracer_out)

# +
from collections import defaultdict
import itertools as it
import string

class PPrint:
  lines: List[Tuple[int, str]]

  def __init__(self, lines):
    self.lines = lines

  def indent(self, indent: int) -> 'PPrint':
    return PPrint([(indent + orig_indent, s) for orig_indent, s in self.lines])

  def __add__(self, rhs: 'PPrint') -> 'PPrint':
    return PPrint(self.lines + rhs.lines)

  def __rshift__(self, rhs: 'PPrint') -> 'PPrint':
    if not rhs.lines: return self
    if not self.lines: return rhs
    indent, s = self.lines[-1]
    indented_block = rhs.indent(indent + len(s))
    common_line = s + ' ' * rhs.lines[0][0] + rhs.lines[0][1]
    return PPrint(self.lines[:-1]
                  + [(indent, common_line)]
                  + indented_block.lines[1:])

  def __str__(self) -> str:
    return '\n'.join(' ' * indent + s for indent, s in self.lines)

def pp(s: Any) -> PPrint:
  return PPrint([(0, line) for line in str(s).splitlines()])

def vcat(ps: List[PPrint]) -> PPrint:
  return sum(ps, pp(''))

def pp_jaxpr(jaxpr: Jaxpr):
  namegen = (''.join(s) for r in it.count(1)
             for s in it.permutations(string.ascii_lowercase, r))
  names = defaultdict(lambda: next(namegen))
  in_binders = ', '.join(var_str(names, x) for x in jaxpr.in_binders)
  eqns = vcat([pp_eqn(names, e) for e in jaxpr.eqns])
  out = names[jaxpr.out] if isinstance(jaxpr.out, Var) else str(jaxpr.out.val)
  return (pp(f'{{ lambda {in_binders} .') +
          ((pp('let ') >> eqns) + pp(f'in {out} }}')).indent(2))

def var_str(names: Dict[Var, str], v: Var) -> str:
  return f'{names[v]}:{v.aval.str_short()}'

def pp_eqn(names: Dict[Var, str], eqn: JaxprEqn) -> PPrint:
  lhs = pp(var_str(names, eqn.out_binder))
  rhs = (pp(eqn.primitive.name) >> pp_params(eqn.params) >>
         pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)
                     for x in eqn.inputs)))
  return lhs >> pp(' = ') >> rhs

def pp_params(params: Dict[str, Any]) -> PPrint:
  items = sorted(params.items())
  if items:
    return pp(' [ ') >> vcat([pp(f'{k}={v}') for k, v in items]) >> pp(' ] ')
  else:
    return pp(' ')

# -

jaxpr, consts = make_jaxpr(lambda x: 2. * x, [raise_to_shaped(get_aval(3.))])
print(pp_jaxpr(jaxpr))
print(typecheck_jaxpr(jaxpr))


"""
.. _tut-fnirs-processing-bad:

Importance of Controls and Parameter Selection
==============================================

.. warning:: DO NOT USE THIS SCRIPT TO PROCESS YOUR DATA

.. sidebar:: Relevant literature

   Huppert TJ. Commentary on the statistical properties of noise and its
   implication on general linear models in functional near-infrared
   spectroscopy. Neurophotonics. 2016;3(1)

This tutorial demonstrates how **NOT** to process your
functional near-infrared spectroscopy (fNIRS)
data.
Here I demonstrate that with a simple process pipelines even random noise
can look like a canonical haemodynamic response.

This tutorial was written to motivate further interest in to what affect
signal processing has on our data. And to provide a concrete example
about why you should understand your analysis and not vary parameters to meet
preconceived ideas of what your data should look like.
I recommend specifying analysis parameters in advance along with your experimental
protocol. Hopefully this tutorial also highlights the importance of
including a control condition in your experimental design.

The
:ref:`MNE fNIRS waveform tutorial <mne:tut-fnirs-processing>`
is used as a basis for this analysis, so most of the document looks similar.
However, at the midpoint we replace the real data with noise and demonstrate
that without careful attention to the analysis parameter it would still
appear as if a fNIRS response is observed.

.. contents:: Page contents
   :local:
   :depth: 2

"""
# sphinx_gallery_thumbnail_number = 7

# Authors: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)

import os
import numpy as np

import mne
import mne_nirs
np.random.seed(1)

fnirs_data_folder = mne.datasets.fnirs_motor.data_path()
fnirs_cw_amplitude_dir = os.path.join(fnirs_data_folder, 'Participant-1')
raw_intensity = mne.io.read_raw_nirx(fnirs_cw_amplitude_dir, verbose=True)
raw_intensity.load_data()

# %%
# Selecting channels appropriate for detecting neural responses
# -------------------------------------------------------------
#
# First we remove channels that are too close together (short channels) to
# detect a neural response (less than 1 cm distance between optodes).
# These short channels can be seen in the figure above.
# To achieve this we pick all the channels that are not considered to be short.

picks = mne.pick_types(raw_intensity.info, meg=False, fnirs=True)
dists = mne.preprocessing.nirs.source_detector_distances(
    raw_intensity.info, picks=picks)
raw_intensity.pick(picks[dists > 0.01])
raw_intensity.plot(n_channels=len(raw_intensity.ch_names),
                   duration=500, show_scrollbars=False)

# %%
# Converting from raw intensity to optical density
# ------------------------------------------------
#
# The raw intensity values are then converted to optical density.

raw_od = mne.preprocessing.nirs.optical_density(raw_intensity)
raw_od.plot(n_channels=len(raw_od.ch_names),
            duration=500, show_scrollbars=False)

# %%
# Converting from optical density to haemoglobin
# ----------------------------------------------
#
# Next we convert the optical density data to haemoglobin concentration using
# the modified Beer-Lambert law.

raw_haemo = mne.preprocessing.nirs.beer_lambert_law(raw_od)
raw_haemo.plot(n_channels=len(raw_haemo.ch_names),
               duration=500, show_scrollbars=False)

# %%
# !!!!!!Replace real data with white noise!!!!!
# ----------------------------------------------
#
# Here we replace the signals with white noise.
# We make the HbR white noise 3 time smaller than HbO as is commonly observed.

raw_haemo._data = np.random.randn(40, 23239) / 1.0e6 * 1
raw_haemo._data[::2, :]= np.random.randn(20, 23239) / 1.0e6 * 3

# %%
# Removing heart rate from signal (BAD DONT COPY)
# -----------------------------------------------
#
# This analysis is BAD and is provided just as a demonstration.
# Do not do this!!
#

fig = raw_haemo.plot_psd(average=True)
fig.suptitle('Before filtering', weight='bold', size='x-large')
fig.subplots_adjust(top=0.88)
raw_haemo = raw_haemo.filter(0.05, 0.1, h_trans_bandwidth=0.2,
                             l_trans_bandwidth=0.02)
fig = raw_haemo.plot_psd(average=True)
fig.suptitle('After filtering', weight='bold', size='x-large')
fig.subplots_adjust(top=0.88)

# %%
# Apply Cui negative correlation method
# -------------------------------------
#
# Here we apply the Cui signal enhancement technique.

raw_haemo = mne_nirs.signal_enhancement.enhance_negative_correlation(raw_haemo)

# %%
# Extract epochs
# --------------
#
# Now that the signal has been converted to relative haemoglobin concentration,
# and the unwanted heart rate component has been removed, we can extract epochs
# related to each of the experimental conditions.
#
# First we extract the events of interest and visualise them to ensure they are
# correct.

events, _ = mne.events_from_annotations(raw_haemo, event_id={'1.0': 1,
                                                             '2.0': 2,
                                                             '3.0': 3})
event_dict = {'Control': 1, 'Tapping/Left': 2, 'Tapping/Right': 3}
fig = mne.viz.plot_events(events, event_id=event_dict,
                          sfreq=raw_haemo.info['sfreq'])
fig.subplots_adjust(right=0.7)  # make room for the legend

# %%
# Next we define the range of our epochs, the rejection criteria,
# baseline correction, and extract the epochs. We visualise the log of which
# epochs were dropped.

reject_criteria = dict(hbo=10e-6)
tmin, tmax = -1.5, 10.5

epochs = mne.Epochs(raw_haemo, events, event_id=event_dict,
                    tmin=tmin, tmax=tmax,
                    reject=reject_criteria, reject_by_annotation=True,
                    proj=True, baseline=(None, 0), preload=True,
                    detrend=1, verbose=True)

# %%
# Plot standard fNIRS response image
# ----------------------------------
#
# Next we generate the most common visualisation of fNIRS data: plotting
# both the HbO and HbR on the same figure to illustrate the relation between
# the two signals.

evoked_dict = {'Tapping/HbO': epochs['Tapping'].average(picks='hbo'),
               'Tapping/HbR': epochs['Tapping'].average(picks='hbr')}

# Rename channels until the encoding of frequency in ch_name is fixed
for condition in evoked_dict:
    evoked_dict[condition].rename_channels(lambda x: x[:-4])

color_dict = dict(HbO='#AA3377', HbR='b')

mne.viz.plot_compare_evokeds(evoked_dict, combine="mean", ci=0.95,
                             colors=color_dict)

# %%
# Summary
# --------------
#
# Above is a waveform that has many of the characteristics that are expected
# from a haeomodynamic response.
# We see a baseline that is around zero, the oxyhaemoglobin is positive
# and the deoxyhaemoglobin is negative,
# the peak occurs around 5 seconds then returns to
# baseline.
# However, it was generated from random noise!
#
# .. sidebar:: PLEASE READ THIS PAPER!!
#
#    Huppert TJ. Commentary on the statistical properties of noise and its
#    implication on general linear models in functional near-infrared
#    spectroscopy. Neurophotonics. 2016;3(1)
#
# As this is an introduction tutorial I do not go in to details about the
# signal process theory behind why the settings above are wrong.
# But if you are interested you should start with a DSP textbook
# on signals and systems.
#
# There are many signal processing parameter choices that interact to produce
# your end result. It is important to have a control condition to ensure
# that you aren't forcing your data to have these fake responses.
# Its also important to understand the effect that each parameter choice
# has on your output. Using a receiver operating characteristic is a good
# approach to understand your analysis.
# This also highlights some of the concerns with the time domain analysis
# approach, and the benefits of the GLM style of analysis
# :ref:`(see MNE-NIRS fNIRS GLM tutorial) <tut-fnirs-hrf>` (and Huppert 2016).
#
# .. warning:: DO NOT USE THIS SCRIPT TO PROCESS YOUR DATA


# -*- coding: utf-8 -*-
"""
Examples of Sandpile

AUTHORS:

- David Perkinson (2015-05) [Using `examples.py` from homology as
  template.]

This file constructs some examples of Sandpiles.

The examples are accessible by typing ``sandpiles.NAME``, where
``NAME`` is the name of the example.  You can get a list by typing
``sandpiles.`` and hitting the TAB key::

   sandpiles.Complete
   sandpiles.Cycle
   sandpiles.Diamond
   sandpiles.Grid
   sandpiles.House

See the documentation for each particular type of example for full details.
"""

from sage.sandpiles.sandpile import Sandpile
from sage.graphs.graph_generators import graphs

class SandpileExamples(object):
    """
    Some examples of sandpiles.

    Here are the available examples; you can also type
    ``sandpiles.``  and hit tab to get a list:

    - :meth:`Complete`
    - :meth:`Cycle`
    - :meth:`Diamond`
    - :meth:`Grid`
    - :meth:`House`

    EXAMPLES::

        sage: s = sandpiles.Complete(4)
        sage: s.invariant_factors()
        [1, 4, 4]
        sage: s.laplacian()
        [ 3 -1 -1 -1]
        [-1  3 -1 -1]
        [-1 -1  3 -1]
        [-1 -1 -1  3]
    """
    def __call__(self):
        r"""
        If sandpiles() is executed, return a helpful message.

        INPUT:

        None

        OUTPUT:

        None

        EXAMPLES::

            sage: sandpiles()
            Try sandpile.FOO() where FOO is in the list:
            <BLANKLINE>
                Complete, Cycle, Diamond, Fan, Grid, House, Wheel
        """
        print 'Try sandpile.FOO() where FOO is in the list:\n'
        print "    " + ", ".join([str(i) for i in dir(sandpiles) if i[0]!='_'])

    def Complete(self, n):
        """
        The complete sandpile graph with `n` vertices.

        INPUT:

        -  ``n`` -- positive integer

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: s = sandpiles.Complete(4)
            sage: s.group_order()
            16
            sage: sandpiles.Complete(3) == sandpiles.Cycle(3)
            True
        """
        return Sandpile(graphs.CompleteGraph(n),0)

    def Cycle(self, n):
        """
        Sandpile on the cycle graph with `n` vertices.

        INPUT:

        -  ``n`` -- a non-negative integer

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: s = sandpiles.Cycle(4)
            sage: s.edges()
            [(0, 1, 1),
             (0, 3, 1),
             (1, 0, 1),
             (1, 2, 1),
             (2, 1, 1),
             (2, 3, 1),
             (3, 0, 1),
             (3, 2, 1)]
        """
        return Sandpile(graphs.CycleGraph(n),0)

    def Diamond(self):
        """
        Sandpile on the diamond graph.

        INPUT:

        None

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: s = sandpiles.Diamond()
            sage: s.invariant_factors()
            [1, 1, 8]
        """
        return Sandpile(graphs.DiamondGraph(),0)

    def Fan(self, n, deg_three_verts=False):
        """
        Sandpile on the Fan graph with a total of `n` vertices.

        INPUT:

        -  ``n`` -- a non-negative integer

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: f = sandpiles.Fan(10)
            sage: f.group_order() == fibonacci(18)
            True
            sage: f = sandpiles.Fan(10,True)  # all nonsink vertices have deg 3
            sage: f.group_order() == fibonacci(20)
            True
        """
        f = graphs.WheelGraph(n)
        if n>2:
            f.delete_edge(1,n-1)
            if deg_three_verts:
                f.allow_multiple_edges(True)
                f.add_edges([(0,1),(0,n-1)])
            return Sandpile(f,0)
        elif n==1:
            return Sandpile(f,0)
        elif n==2:
            if deg_three_verts:
                return Sandpile({0:{1:3}, 1:{0:3}})
            else:
                return Sandpile(f,0)

    def Grid(self, m, n):
        """
        Sandpile on the diamond graph.

        INPUT:

        -  ``m``, ``n`` -- negative integers

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: s = sandpiles.Grid(2,3)
            sage: s.vertices()
            [(0, 0), (1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3)]
            sage: s.invariant_factors()
            [1, 1, 1, 1, 1, 2415]
            sage: s = sandpiles.Grid(1,1)
            sage: s.dict()
            {(0, 0): {(1, 1): 4}, (1, 1): {(0, 0): 4}}
        """
        G = graphs.Grid2dGraph(m+2,n+2)
        G.allow_multiple_edges(True)  # to ensure each vertex ends up with degree 4
        V = [(i,j) for i in [0,m+1] for j in range(n+2)] + [(i,j) for j in [0,n+1] for i in range(m+2)]
        G.merge_vertices(V)
        return Sandpile(G, (0,0))

    def House(self):
        """
        Sandpile on the House graph.

        INPUT:

        None

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: s = sandpiles.House()
            sage: s.invariant_factors()
            [1, 1, 1, 11]
        """
        return Sandpile(graphs.HouseGraph(),0)

    def Wheel(self, n):
        """
        Sandpile on the wheel graph with a total of `n` vertices.

        INPUT:

        -  ``n`` -- a non-negative integer

        OUTPUT:

        - Sandpile

        EXAMPLES::

            sage: w = sandpiles.Wheel(6)
            sage: w.invariant_factors()
            [1, 1, 1, 11, 11]
        """
        return Sandpile(graphs.WheelGraph(n),0)

sandpiles = SandpileExamples()