# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: collapsed,code_folding
#     formats: ipynb,py
#     notebook_metadata_filter: all
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   language_info:
#     codemirror_mode:
#       name: ipython
#       version: 3
#     file_extension: .py
#     mimetype: text/x-python
#     name: python
#     nbconvert_exporter: python
#     pygments_lexer: ipython3
#     version: 3.7.6
#   latex_envs:
#     LaTeX_envs_menu_present: true
#     autoclose: false
#     autocomplete: true
#     bibliofile: biblio.bib
#     cite_by: apalike
#     current_citInitial: 1
#     eqLabelWithNumbers: true
#     eqNumInitial: 1
#     hotkeys:
#       equation: Ctrl-E
#       itemize: Ctrl-I
#     labels_anchors: false
#     latex_user_defs: false
#     report_style_numbering: false
#     user_envs_cfg: false
# ---

# # Alternative Combinations of Parameter Values
#
# Please write the names and email addresses of everyone who worked on this notebook on the line below.
#
# YOUR NAMES HERE
#
# ## Introduction
#
# The notebook "Micro-and-Macro-Implications-of-Very-Impatient-HHs" is an exercise that demonstrates the consequences of changing a key parameter of the [cstwMPC](http://econ.jhu.edu/people/ccarroll/papers/cstwMPC) model, the time preference factor $\beta$.
#
# The [REMARK](https://github.com/econ-ark/REMARK) `SolvingMicroDSOPs` reproduces the last figure in the [SolvingMicroDSOPs](http://econ.jhu.edu/people/ccarroll/SolvingMicroDSOPs) lecture notes, which shows that there are classes of alternate values of $\beta$ and $\rho$ that fit the data almost as well as the exact 'best fit' combination.
#
# Inspired by this comparison, this notebook asks you to examine the consequences for:
#
# * The consumption function
# * The distribution of wealth
#
# Of _joint_ changes in $\beta$ and $\rho$ together.
#
# One way you can do this is to construct a list of alternative values of $\rho$ (say, values that range upward from the default value of $\rho$, in increments of 0.2, all the way to $\rho=5$).  Then for each of these values of $\rho$ you will find the value of $\beta$ that leads the same value for target market resources, $\check{m}$.
#
# As a reminder, $\check{m}$ is defined as the value of $m$ at which the optimal value of ${c}$ is the value such that, at that value of ${c}$, the expected level of ${m}$ next period is the same as its current value:
#
# $\mathbb{E}_{t}[{m}_{t+1}] = {m}_{t}$
#
# Other notes:
# * The cstwMPC model solves and simulates the problems of consumers with 7 different values of $\beta$
#    * You should do your exercise using the middle value of $\beta$ from that exercise:
#       * `DiscFac_mean   = 0.9855583`
# * You are likely to run into the problem, as you experiment with parameter values, that you have asked HARK to solve a model that does not satisfy one of the impatience conditions required for the model to have a solution.  Those conditions are explained intuitively in the [TractableBufferStock](http://econ.jhu.edu/people/ccarroll/public/lecturenotes/consumption/TractableBufferStock/) model.  The versions of the impatience conditions that apply to the $\texttt{IndShockConsumerType}$ model can be found in the paper [BufferStockTheory](http://econ.jhu.edu/people/ccarroll/papers/BufferStockTheory), table 2.
#    * The conditions that need to be satisfied are:
#       * The Growth Impatience Condition (GIC)
#       * The Return Impatience Condition (RIC)
# * Please accumulate the list of solved consumers' problems in a list called `MyTypes`
#    * For compatibility with a further part of the assignment below

# + {"code_folding": []}
# This cell merely imports and sets up some basic functions and packages

# %matplotlib inline
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
from copy import deepcopy

import HARK # Prevents import error from Demos repo
from HARK.utilities import plotFuncs

# + {"code_folding": [0, 4]}
# Import IndShockConsumerType
from HARK.ConsumptionSaving.ConsIndShockModel import IndShockConsumerType

# Define a dictionary with calibrated parameters
cstwMPC_calibrated_parameters = {
    "CRRA":1.0,                    # Coefficient of relative risk aversion
    "Rfree":1.01/(1.0 - 1.0/160.0), # Survival probability,
    "PermGroFac":[1.000**0.25], # Permanent income growth factor (no perm growth),
    "PermGroFacAgg":1.0,
    "BoroCnstArt":0.0,
    "CubicBool":False,
    "vFuncBool":False,
    "PermShkStd":[(0.01*4/11)**0.5],  # Standard deviation of permanent shocks to income
    "PermShkCount":5,  # Number of points in permanent income shock grid
    "TranShkStd":[(0.01*4)**0.5],  # Standard deviation of transitory shocks to income,
    "TranShkCount":5,  # Number of points in transitory income shock grid
    "UnempPrb":0.07,  # Probability of unemployment while working
    "IncUnemp":0.15,  # Unemployment benefit replacement rate
    "UnempPrbRet":None,
    "IncUnempRet":None,
    "aXtraMin":0.00001,  # Minimum end-of-period assets in grid
    "aXtraMax":40,  # Maximum end-of-period assets in grid
    "aXtraCount":32,  # Number of points in assets grid
    "aXtraExtra":[None],
    "aXtraNestFac":3,  # Number of times to 'exponentially nest' when constructing assets grid
    "LivPrb":[1.0 - 1.0/160.0],  # Survival probability
    "DiscFac":0.97,             # Default intertemporal discount factor; dummy value, will be overwritten
    "cycles":0,
    "T_cycle":1,
    "T_retire":0,
    'T_sim':1200,  # Number of periods to simulate (idiosyncratic shocks model, perpetual youth)
    'T_age': 400,
    'IndL': 10.0/9.0,  # Labor supply per individual (constant),
    'aNrmInitMean':np.log(0.00001),
    'aNrmInitStd':0.0,
    'pLvlInitMean':0.0,
    'pLvlInitStd':0.0,
    'AgentCount':10000,
}
# -

# Construct a list of solved consumers' problems, IndShockConsumerType is just a place holder
MyTypes = [IndShockConsumerType(**cstwMPC_calibrated_parameters)]

# ## Simulating the Distribution of Wealth for Alternative Combinations
#
# You should now have constructed a list of consumer types all of whom have the same _target_ level of market resources $\check{m}$.
#
# But the fact that everyone has the same target ${m}$ does not mean that the _distribution_ of ${m}$ will be the same for all of these consumer types.
#
# In the code block below, fill in the contents of the loop to solve and simulate each agent type for many periods.  To do this, you should invoke the methods $\texttt{solve}$, $\texttt{initializeSim}$, and $\texttt{simulate}$ in that order.  Simulating for 1200 quarters (300 years) will approximate the long run distribution of wealth in the population.

for ThisType in tqdm(MyTypes):
    ThisType.solve()
    ThisType.initializeSim()
    ThisType.simulate()

# Now that you have solved and simulated these consumers, make a plot that shows the relationship between your alternative values of $\rho$ and the mean level of assets

# +
# To help you out, we have given you the command needed to construct a list of the levels of assets for all consumers
aLvl_all = np.concatenate([ThisType.aLvlNow for ThisType in MyTypes])

# You should take the mean of aLvl for each consumer in MyTypes, divide it by the mean across all simulations
# and then plot the ratio of the values of mean(aLvl) for each group against the value of $\rho$
# -

# # Interpret
# Here, you should attempt to give an intiutive explanation of the results you see in the figure you just constructed

# ## The Distribution of Wealth...
#
# Your next exercise is to show how the distribution of wealth differs for the different parameter  values

# +
from HARK.utilities import getLorenzShares, getPercentiles

# Finish filling in this function to calculate the Euclidean distance between the simulated and actual Lorenz curves.
def calcLorenzDistance(SomeTypes):
    '''
    Calculates the Euclidean distance between the simulated and actual (from SCF data) Lorenz curves at the
    20th, 40th, 60th, and 80th percentiles.

    Parameters
    ----------
    SomeTypes : [AgentType]
        List of AgentTypes that have been solved and simulated.  Current levels of individual assets should
        be stored in the attribute aLvlNow.

    Returns
    -------
    lorenz_distance : float
        Euclidean distance (square root of sum of squared differences) between simulated and actual Lorenz curves.
    '''
    # Define empirical Lorenz curve points
    lorenz_SCF = np.array([-0.00183091,  0.0104425 ,  0.0552605 ,  0.1751907 ])

    # Extract asset holdings from all consumer types
    aLvl_sim = np.concatenate([ThisType.aLvlNow for ThisType in MyTypes])

    # Calculate simulated Lorenz curve points
    lorenz_sim = getLorenzShares(aLvl_sim,percentiles=[0.2,0.4,0.6,0.8])

    # Calculate the Euclidean distance between the simulated and actual Lorenz curves
    lorenz_distance = np.sqrt(np.sum((lorenz_SCF - lorenz_sim)**2))

    # Return the Lorenz distance
    return lorenz_distance

# -

# ## ...and the Marginal Propensity to Consume
#
# Now let's look at the aggregate MPC.  In the code block below, write a function that produces text output of the following form:
#
# $\texttt{The 35th percentile of the MPC is 0.15623}$
#
# Your function should take two inputs: a list of types of consumers and an array of percentiles (numbers between 0 and 1). It should return no outputs, merely print to screen one line of text for each requested percentile.  The model is calibrated at a quarterly frequency, but Carroll et al report MPCs at an annual frequency. To convert, use the formula:
#
# $\kappa_{Y} \approx 1.0 - (1.0 - \kappa_{Q})^4$

# +
# Write a function to tell us about the distribution of the MPC in this code block, then test it!
# You will almost surely find it useful to use a for loop in this function.
def describeMPCdstn(SomeTypes,percentiles):
    MPC_sim = np.concatenate([ThisType.MPCnow for ThisType in SomeTypes])
    MPCpercentiles_quarterly = getPercentiles(MPC_sim,percentiles=percentiles)
    MPCpercentiles_annual = 1.0 - (1.0 - MPCpercentiles_quarterly)**4

    for j in range(len(percentiles)):
        print('The ' + str(100*percentiles[j]) + 'th percentile of the MPC is ' + str(MPCpercentiles_annual[j]))

describeMPCdstn(MyTypes,np.linspace(0.05,0.95,19))
# -

# # If You Get Here ...
#
# If you have finished the above exercises quickly and have more time to spend on this assignment, for extra credit you can do the same exercise where, instead of exploring the consequences of alternative values of relative risk aversion $\rho$, you should test the consequences of different values of the growth factor $\Gamma$ that lead to the same $\check{m}$.


"""
Sets

AUTHORS:

- William Stein (2005) - first version

- William Stein (2006-02-16) - large number of documentation and
  examples; improved code

- Mike Hansen (2007-3-25) - added differences and symmetric
  differences; fixed operators

- Florent Hivert (2010-06-17) - Adapted to categories

- Nicolas M. Thiery (2011-03-15) - Added subset and superset methods

- Julian Rueth (2013-04-09) - Collected common code in
  :class:`Set_object_binary`, fixed ``__hash__``.

"""

#*****************************************************************************
#       Copyright (C) 2005 William Stein <wstein@gmail.com>
#                     2013 Julian Rueth <julian.rueth@fsfe.org>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from __future__ import print_function
import six
from six import integer_types

from sage.misc.latex import latex
from sage.misc.prandom import choice

from sage.structure.category_object import CategoryObject
from sage.structure.element import Element
from sage.structure.parent import Parent, Set_generic
from sage.structure.richcmp import richcmp_method, richcmp, rich_to_bool

from sage.categories.sets_cat import Sets
from sage.categories.enumerated_sets import EnumeratedSets

import sage.rings.infinity

def has_finite_length(obj):
    """
    Return ``True`` if ``obj`` is known to have finite length.

    This is mainly meant for pure Python types, so we do not call any
    Sage-specific methods.

    EXAMPLES::

        sage: from sage.sets.set import has_finite_length
        sage: has_finite_length(tuple(range(10)))
        True
        sage: has_finite_length(list(range(10)))
        True
        sage: has_finite_length(set(range(10)))
        True
        sage: has_finite_length(iter(range(10)))
        False
        sage: has_finite_length(GF(17^127))
        True
        sage: has_finite_length(ZZ)
        False
    """
    try:
        len(obj)
    except OverflowError:
        return True
    except Exception:
        return False
    else:
        return True

def Set(X=[]):
    r"""
    Create the underlying set of ``X``.

    If ``X`` is a list, tuple, Python set, or ``X.is_finite()`` is
    ``True``, this returns a wrapper around Python's enumerated immutable
    ``frozenset`` type with extra functionality.  Otherwise it returns a
    more formal wrapper.

    If you need the functionality of mutable sets, use Python's
    builtin set type.

    EXAMPLES::

        sage: X = Set(GF(9,'a'))
        sage: X
        {0, 1, 2, a, a + 1, a + 2, 2*a, 2*a + 1, 2*a + 2}
        sage: type(X)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: Y = X.union(Set(QQ))
        sage: Y
        Set-theoretic union of {0, 1, 2, a, a + 1, a + 2, 2*a, 2*a + 1, 2*a + 2} and Set of elements of Rational Field
        sage: type(Y)
        <class 'sage.sets.set.Set_object_union_with_category'>

    Usually sets can be used as dictionary keys.

    ::

        sage: d={Set([2*I,1+I]):10}
        sage: d                  # key is randomly ordered
        {{I + 1, 2*I}: 10}
        sage: d[Set([1+I,2*I])]
        10
        sage: d[Set((1+I,2*I))]
        10

    The original object is often forgotten.

    ::

        sage: v = [1,2,3]
        sage: X = Set(v)
        sage: X
        {1, 2, 3}
        sage: v.append(5)
        sage: X
        {1, 2, 3}
        sage: 5 in X
        False

    Set also accepts iterators, but be careful to only give *finite*
    sets::

        sage: from six.moves import range
        sage: sorted(Set(range(1,6)))
        [1, 2, 3, 4, 5]
        sage: sorted(Set(list(range(1,6))))
        [1, 2, 3, 4, 5]
        sage: sorted(Set(iter(range(1,6))))
        [1, 2, 3, 4, 5]

    We can also create sets from different types::

        sage: sorted(Set([Sequence([3,1], immutable=True), 5, QQ, Partition([3,1,1])]), key=str)
        [5, Rational Field, [3, 1, 1], [3, 1]]

    Sets with unhashable objects work, but with less functionality::

        sage: A = Set([QQ, (3, 1), 5])  # hashable
        sage: sorted(A.list(), key=repr)
        [(3, 1), 5, Rational Field]
        sage: type(A)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: B = Set([QQ, [3, 1], 5])  # unhashable
        sage: sorted(B.list(), key=repr)
        Traceback (most recent call last):
        ...
        AttributeError: 'Set_object_with_category' object has no attribute 'list'
        sage: type(B)
        <class 'sage.sets.set.Set_object_with_category'>

    TESTS::

        sage: Set(Primes())
        Set of all prime numbers: 2, 3, 5, 7, ...
        sage: Set(Subsets([1,2,3])).cardinality()
        8
        sage: S = Set(iter([1,2,3])); S
        {1, 2, 3}
        sage: type(S)
        <class 'sage.sets.set.Set_object_enumerated_with_category'>
        sage: S = Set([])
        sage: TestSuite(S).run()

    Check that :trac:`16090` is fixed::

        sage: Set()
        {}
    """
    if isinstance(X, CategoryObject):
        if isinstance(X, Set_generic):
            return X
        elif X in Sets().Finite():
            return Set_object_enumerated(X)
        else:
            return Set_object(X)

    if isinstance(X, Element):
        raise TypeError("Element has no defined underlying set")

    try:
        X = frozenset(X)
    except TypeError:
        return Set_object(X)
    else:
        return Set_object_enumerated(X)

@richcmp_method
class Set_object(Set_generic):
    r"""
    A set attached to an almost arbitrary object.

    EXAMPLES::

        sage: K = GF(19)
        sage: Set(K)
        {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
        sage: S = Set(K)

        sage: latex(S)
        \left\{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18\right\}
        sage: TestSuite(S).run()

        sage: latex(Set(ZZ))
        \Bold{Z}

    TESTS:

    See trac ticket :trac:`14486`::

        sage: 0 == Set([1]), Set([1]) == 0
        (False, False)
        sage: 1 == Set([0]), Set([0]) == 1
        (False, False)
    """
    def __init__(self, X, category=None):
        """
        Create a Set_object

        This function is called by the Set function; users
        shouldn't call this directly.

        EXAMPLES::

            sage: type(Set(QQ))
            <class 'sage.sets.set.Set_object_with_category'>
            sage: Set(QQ).category()
            Category of sets

        TESTS::

            sage: _a, _b = get_coercion_model().canonical_coercion(Set([0]), 0)
            Traceback (most recent call last):
            ...
            TypeError: no common canonical parent for objects with parents:
            '<class 'sage.sets.set.Set_object_enumerated_with_category'>'
            and 'Integer Ring'
        """
        from sage.rings.integer import is_Integer
        if isinstance(X, integer_types) or is_Integer(X):
            # The coercion model will try to call Set_object(0)
            raise ValueError('underlying object cannot be an integer')

        if category is None:
            category = Sets()
        Parent.__init__(self, category=category)
        self.__object = X

    def __hash__(self):
        """
        Return the hash value of ``self``.

        EXAMPLES::

            sage: hash(Set(QQ)) == hash(QQ)
            True
        """
        return hash(self.__object)

    def _latex_(self):
        r"""
        Return latex representation of this set.

        This is often the same as the latex representation of this
        object when the object is infinite.

        EXAMPLES::

            sage: latex(Set(QQ))
            \Bold{Q}

        When the object is finite or a special set then the latex
        representation can be more interesting.

        ::

            sage: print(latex(Primes()))
            \text{\texttt{Set{ }of{ }all{ }prime{ }numbers:{ }2,{ }3,{ }5,{ }7,{ }...}}
            sage: print(latex(Set([1,1,1,5,6])))
            \left\{1, 5, 6\right\}
        """
        return latex(self.__object)

    def _repr_(self):
        """
        Print representation of this set.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: X
            Set of elements of Integer Ring
            sage: X.rename('{ integers }')
            sage: X
            { integers }
        """
        return "Set of elements of " + repr(self.__object)

    def __iter__(self):
        """
        Iterate over the elements of this set.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            2
        """
        return iter(self.__object)

    an_element = EnumeratedSets.ParentMethods.__dict__['_an_element_from_iterator']

    def __contains__(self, x):
        """
        Return ``True`` if `x` is in ``self``.

        EXAMPLES::

            sage: X = Set(ZZ)
            sage: 5 in X
            True
            sage: GF(7)(3) in X
            True
            sage: 2/1 in X
            True
            sage: 2/1 in ZZ
            True
            sage: 2/3 in X
            False

        Finite fields better illustrate the difference between
        ``__contains__`` for objects and their underlying sets.

            sage: X = Set(GF(7))
            sage: X
            {0, 1, 2, 3, 4, 5, 6}
            sage: 5/3 in X
            False
            sage: 5/3 in GF(7)
            False
            sage: sorted(Set(GF(7)).union(Set(GF(5))), key=int)
            [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 6]
            sage: Set(GF(7)).intersection(Set(GF(5)))
            {}
        """
        return x in self.__object

    def __richcmp__(self, right, op):
        r"""
        Compare ``self`` and ``right``.

        If ``right`` is not a :class:`Set_object`, return ``NotImplemented``.
        If ``right`` is also a :class:`Set_object`, returns comparison
        on the underlying objects.

        .. NOTE::

           If `X < Y` is true this does *not* necessarily mean
           that `X` is a subset of `Y`.  Also, any two sets can be
           compared still, but the result need not be meaningful
           if they are not equal.

        EXAMPLES::

            sage: Set(ZZ) == Set(QQ)
            False
            sage: Set(ZZ) < Set(QQ)
            True
            sage: Primes() == Set(QQ)
            False

        The following is random, illustrating that comparison of
        sets is not the subset relation, when they are not equal::

            sage: Primes() < Set(QQ)             # random  # py2
            True or False
        """
        if not isinstance(right, Set_object):
            return NotImplemented
        return richcmp(self.__object, right.__object, op)

    def union(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set(QQ).union(Set(ZZ))
            Set-theoretic union of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: Set(QQ) + Set(ZZ)
            Set-theoretic union of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: X = Set(QQ).union(Set(GF(3))); X
            Set-theoretic union of Set of elements of Rational Field and {0, 1, 2}
            sage: 2/3 in X
            True
            sage: GF(3)(2) in X
            True
            sage: GF(5)(2) in X
            False
            sage: sorted(Set(GF(7)) + Set(GF(3)), key=int)
            [0, 0, 1, 1, 2, 2, 3, 4, 5, 6]
        """
        if isinstance(X, Set_generic):
            if self is X:
                return self
            return Set_object_union(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def __add__(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set(RealField()) + Set(QQ^5)
             Set-theoretic union of Set of elements of Real Field with 53 bits of precision and Set of elements of Vector space of dimension 5 over Rational Field
            sage: Set(GF(3)) + Set(GF(2))
            {0, 1, 2, 0, 1}
            sage: Set(GF(2)) + Set(GF(4,'a'))
            {0, 1, a, a + 1}
            sage: sorted(Set(GF(8,'b')) + Set(GF(4,'a')), key=str)
            [0, 0, 1, 1, a, a + 1, b, b + 1, b^2, b^2 + 1, b^2 + b, b^2 + b + 1]
        """
        return self.union(X)

    def __or__(self, X):
        """
        Return the union of ``self`` and ``X``.

        EXAMPLES::

            sage: Set([2,3]) | Set([3,4])
            {2, 3, 4}
            sage: Set(ZZ) | Set(QQ)
            Set-theoretic union of Set of elements of Integer Ring and Set of elements of Rational Field
        """

        return self.union(X)

    def intersection(self, X):
        r"""
        Return the intersection of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set(ZZ).intersection(Primes())
            sage: 4 in X
            False
            sage: 3 in X
            True

            sage: 2/1 in X
            True

            sage: X = Set(GF(9,'b')).intersection(Set(GF(27,'c')))
            sage: X
            {}

            sage: X = Set(GF(9,'b')).intersection(Set(GF(27,'b')))
            sage: X
            {}
        """
        if isinstance(X, Set_generic):
            if self is X:
                return self
            return Set_object_intersection(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def difference(self, X):
        r"""
        Return the set difference ``self - X``.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: 4 in X
            True
            sage: 3 in X
            False

            sage: 4/1 in X
            True

            sage: X = Set(GF(9,'b')).difference(Set(GF(27,'c')))
            sage: X
            {0, 1, 2, b, b + 1, b + 2, 2*b, 2*b + 1, 2*b + 2}

            sage: X = Set(GF(9,'b')).difference(Set(GF(27,'b')))
            sage: X
            {0, 1, 2, b, b + 1, b + 2, 2*b, 2*b + 1, 2*b + 2}
        """
        if isinstance(X, Set_generic):
            if self is X:
                return Set([])
            return Set_object_difference(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def symmetric_difference(self, X):
        r"""
        Returns the symmetric difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set([1,2,3]).symmetric_difference(Set([3,4]))
            sage: X
            {1, 2, 4}
        """

        if isinstance(X, Set_generic):
            if self is X:
                return Set([])
            return Set_object_symmetric_difference(self, X)
        raise TypeError("X (=%s) must be a Set"%X)

    def __sub__(self, X):
        """
        Return the difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: Y = Set(ZZ) - Primes()
            sage: X == Y
            True
        """
        return self.difference(X)

    def __and__(self, X):
        """
        Returns the intersection of ``self`` and ``X``.

        EXAMPLES::

            sage: Set([2,3]) & Set([3,4])
            {3}
            sage: Set(ZZ) & Set(QQ)
            Set-theoretic intersection of Set of elements of Integer Ring and Set of elements of Rational Field
        """

        return self.intersection(X)

    def __xor__(self, X):
        """
        Returns the symmetric difference of ``self`` and ``X``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.symmetric_difference(Y)
            {3, 4}
            sage: X.__xor__(Y)
            {3, 4}
        """
        return self.symmetric_difference(X)

    def cardinality(self):
        """
        Return the cardinality of this set, which is either an integer or
        ``Infinity``.

        EXAMPLES::

            sage: Set(ZZ).cardinality()
            +Infinity
            sage: Primes().cardinality()
            +Infinity
            sage: Set(GF(5)).cardinality()
            5
            sage: Set(GF(5^2,'a')).cardinality()
            25
        """
        if not self.is_finite():
            return sage.rings.infinity.infinity

        if self is not self.__object:
            try:
                return self.__object.cardinality()
            except (AttributeError, NotImplementedError):
                pass
            from sage.rings.integer import Integer
            try:
                return Integer(len(self.__object))
            except TypeError:
                pass

        raise NotImplementedError("computation of cardinality of %s not yet implemented"%self.__object)

    def is_empty(self):
        """
        Return boolean representing emptiness of the set.

        OUTPUT:

        True if the set is empty, False if otherwise.

        EXAMPLES::

            sage: Set([]).is_empty()
            True
            sage: Set([0]).is_empty()
            False
            sage: Set([1..100]).is_empty()
            False
            sage: Set(SymmetricGroup(2).list()).is_empty()
            False
            sage: Set(ZZ).is_empty()
            False

        TESTS::

            sage: Set([]).is_empty()
            True
            sage: Set([1,2,3]).is_empty()
            False
            sage: Set([1..100]).is_empty()
            False
            sage: Set(DihedralGroup(4).list()).is_empty()
            False
            sage: Set(QQ).is_empty()
            False
        """
        return not self

    def is_finite(self):
        """
        Return ``True`` if ``self`` is finite.

        EXAMPLES::

            sage: Set(QQ).is_finite()
            False
            sage: Set(GF(250037)).is_finite()
            True
            sage: Set(Integers(2^1000000)).is_finite()
            True
            sage: Set([1,'a',ZZ]).is_finite()
            True
        """
        obj = self.__object
        try:
            is_finite = obj.is_finite
        except AttributeError:
            return has_finite_length(obj)
        else:
            return is_finite()

    def object(self):
        """
        Return underlying object.

        EXAMPLES::

            sage: X = Set(QQ)
            sage: X.object()
            Rational Field
            sage: X = Primes()
            sage: X.object()
            Set of all prime numbers: 2, 3, 5, 7, ...
        """
        return self.__object

    def subsets(self,size=None):
        """
        Return the :class:`Subsets` object representing the subsets of a set.
        If size is specified, return the subsets of that size.

        EXAMPLES::

            sage: X = Set([1,2,3])
            sage: list(X.subsets())
            [{}, {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}]
            sage: list(X.subsets(2))
            [{1, 2}, {1, 3}, {2, 3}]

        """
        from sage.combinat.subset import Subsets
        return Subsets(self,size)

class Set_object_enumerated(Set_object):
    """
    A finite enumerated set.
    """
    def __init__(self, X):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(GF(19)); S
            {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
            sage: S.category()
            Category of finite sets
            sage: print(latex(S))
            \left\{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18\right\}
            sage: TestSuite(S).run()
        """
        Set_object.__init__(self, X, category=Sets().Finite())

    def random_element(self):
        r"""
        Return a random element in this set.

        EXAMPLES::

            sage: Set([1,2,3]).random_element() # random
            2
        """
        try:
            return self.object().random_element()
        except AttributeError:
            # TODO: this very slow!
            return choice(self.list())

    def is_finite(self):
        r"""
        Return ``True`` as this is a finite set.

        EXAMPLES::

            sage: Set(GF(19)).is_finite()
            True
        """
        return True

    def cardinality(self):
        """
        Return the cardinality of ``self``.

        EXAMPLES::

            sage: Set([1,1]).cardinality()
            1
        """
        from sage.rings.integer import Integer
        return Integer(len(self.set()))

    def __len__(self):
        """
        EXAMPLES::

            sage: len(Set([1,1]))
            1
        """
        return len(self.set())

    def __iter__(self):
        r"""
        Iterating through the elements of ``self``.

        EXAMPLES::

            sage: S = Set(GF(19))
            sage: I = iter(S)
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            2
            sage: next(I)
            3
        """
        return iter(self.set())

    def _latex_(self):
        r"""
        Return the LaTeX representation of ``self``.

        EXAMPLES::

            sage: S = Set(GF(2))
            sage: latex(S)
            \left\{0, 1\right\}
        """
        return '\\left\\{' + ', '.join([latex(x) for x in self.set()])  + '\\right\\}'

    def _repr_(self):
        r"""
        Return the string representation of ``self``.

        EXAMPLES::

            sage: S = Set(GF(2))
            sage: S
            {0, 1}

        TESTS::

            sage: Set()
            {}
        """
        py_set = self.set()
        if six.PY3:
            if not py_set:
                return "{}"
            return repr(py_set)
        else:
            return "{" + repr(py_set)[5:-2] + "}"

    def list(self):
        """
        Return the elements of ``self``, as a list.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: X.list()
            [0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1]
            sage: type(X.list())
            <... 'list'>

        .. TODO::

            FIXME: What should be the order of the result?
            That of ``self.object()``? Or the order given by
            ``set(self.object())``? Note that :meth:`__getitem__` is
            currently implemented in term of this list method, which
            is really inefficient ...
        """
        return list(set(self.object()))

    def set(self):
        """
        Return the Python set object associated to this set.

        Python has a notion of finite set, and often Sage sets
        have an associated Python set.  This function returns
        that set.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: X.set()
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: type(X.set())
            <... 'set'>
            sage: type(X)
            <class 'sage.sets.set.Set_object_enumerated_with_category'>
        """
        return set(self.object())

    def frozenset(self):
        """
        Return the Python frozenset object associated to this set,
        which is an immutable set (hence hashable).

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: s = X.set(); s
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: hash(s)
            Traceback (most recent call last):
            ...
            TypeError: unhashable type: 'set'
            sage: s = X.frozenset(); s
            frozenset({0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1})

            sage: hash(s) != hash(tuple(X.set()))
            True

            sage: type(s)
            <... 'frozenset'>
        """
        return frozenset(self.object())

    def __hash__(self):
        """
        Return the hash of ``self`` (as a ``frozenset``).

        EXAMPLES::

            sage: s = Set(GF(8,'c'))
            sage: hash(s) == hash(s)
            True
        """
        return hash(self.frozenset())

    def __richcmp__(self, other, op):
        """
        Compare the sets ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: X == Set(GF(8,'c'))
            True
            sage: X == Set(GF(4,'a'))
            False
            sage: Set(QQ) == Set(ZZ)
            False
            sage: Set([1]) == set([1])
            True
        """
        if not isinstance(other, Set_object_enumerated):
            if isinstance(other, (set, frozenset)):
                return self.set() == other
            return NotImplemented
        if self.set() == other.set():
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def issubset(self, other):
        r"""
        Return whether ``self`` is a subset of ``other``.

        INPUT:

         - ``other`` -- a finite Set

        EXAMPLES::

            sage: X = Set([1,3,5])
            sage: Y = Set([0,1,2,3,5,7])
            sage: X.issubset(Y)
            True
            sage: Y.issubset(X)
            False
            sage: X.issubset(X)
            True

        TESTS::

            sage: len([Z for Z in Y.subsets() if Z.issubset(X)])
            8
        """
        if not isinstance(other, Set_object_enumerated):
            raise NotImplementedError
        return self.set().issubset(other.set())

    def issuperset(self, other):
        r"""
        Return whether ``self`` is a superset of ``other``.

        INPUT:

         - ``other`` -- a finite Set

        EXAMPLES::

            sage: X = Set([1,3,5])
            sage: Y = Set([0,1,2,3,5])
            sage: X.issuperset(Y)
            False
            sage: Y.issuperset(X)
            True
            sage: X.issuperset(X)
            True

        TESTS::

            sage: len([Z for Z in Y.subsets() if Z.issuperset(X)])
            4
        """
        if not isinstance(other, Set_object_enumerated):
            raise NotImplementedError
        return self.set().issuperset(other.set())

    def union(self, other):
        """
        Return the union of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: Y = Set([GF(8,'c').0, 1, 2, 3])
            sage: X
            {0, 1, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1}
            sage: sorted(Y)
            [1, 2, 3, c]
            sage: sorted(X.union(Y), key=str)
            [0, 1, 2, 3, c, c + 1, c^2, c^2 + 1, c^2 + c, c^2 + c + 1]
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.union(self, other)
        return Set_object_enumerated(self.set().union(other.set()))

    def intersection(self, other):
        """
        Return the intersection of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set(GF(8,'c'))
            sage: Y = Set([GF(8,'c').0, 1, 2, 3])
            sage: X.intersection(Y)
            {1, c}
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.intersection(self, other)
        return Set_object_enumerated(self.set().intersection(other.set()))

    def difference(self, other):
        """
        Return the set difference ``self - other``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.difference(Y)
            {3, 4}
            sage: Z = Set(ZZ)
            sage: W = Set([2.5, 4, 5, 6])
            sage: W.difference(Z)
            {2.50000000000000}
        """
        if not isinstance(other, Set_object_enumerated):
            return Set([x for x in self if x not in other])
        return Set_object_enumerated(self.set().difference(other.set()))

    def symmetric_difference(self, other):
        """
        Return the symmetric difference of ``self`` and ``other``.

        EXAMPLES::

            sage: X = Set([1,2,3,4])
            sage: Y = Set([1,2])
            sage: X.symmetric_difference(Y)
            {3, 4}
            sage: Z = Set(ZZ)
            sage: W = Set([2.5, 4, 5, 6])
            sage: U = W.symmetric_difference(Z)
            sage: 2.5 in U
            True
            sage: 4 in U
            False
            sage: V = Z.symmetric_difference(W)
            sage: V == U
            True
            sage: 2.5 in V
            True
            sage: 6 in V
            False
        """
        if not isinstance(other, Set_object_enumerated):
            return Set_object.symmetric_difference(self, other)
        return Set_object_enumerated(self.set().symmetric_difference(other.set()))

class Set_object_binary(Set_object):
    r"""
    An abstract common base class for sets defined by a binary operation (ex.
    :class:`Set_object_union`, :class:`Set_object_intersection`,
    :class:`Set_object_difference`, and
    :class:`Set_object_symmetric_difference`).

    INPUT:

    - ``X``, ``Y`` -- sets, the operands to ``op``

    - ``op`` -- a string describing the binary operation

    - ``latex_op`` -- a string used for rendering this object in LaTeX

    EXAMPLES::

        sage: X = Set(QQ^2)
        sage: Y = Set(ZZ)
        sage: from sage.sets.set import Set_object_binary
        sage: S = Set_object_binary(X, Y, "union", "\\cup"); S
        Set-theoretic union of Set of elements of Vector space of dimension 2
         over Rational Field and Set of elements of Integer Ring
    """
    def __init__(self, X, Y, op, latex_op):
        r"""
        Initialization.

        TESTS::

            sage: from sage.sets.set import Set_object_binary
            sage: X = Set(QQ^2)
            sage: Y = Set(ZZ)
            sage: S = Set_object_binary(X, Y, "union", "\\cup")
            sage: type(S)
            <class 'sage.sets.set.Set_object_binary_with_category'>
        """
        self._X = X
        self._Y = Y
        self._op = op
        self._latex_op = latex_op
        Set_object.__init__(self, self)

    def _repr_(self):
        r"""
        Return a string representation of this set.

        EXAMPLES::

            sage: Set(ZZ).union(Set(GF(5)))
            Set-theoretic union of Set of elements of Integer Ring and {0, 1, 2, 3, 4}
        """
        return "Set-theoretic {} of {} and {}".format(self._op, self._X, self._Y)

    def _latex_(self):
        r"""
        Return a latex representation of this set.

        EXAMPLES::

            sage: latex(Set(ZZ).union(Set(GF(5))))
            \Bold{Z} \cup \left\{0, 1, 2, 3, 4\right\}
        """
        return latex(self._X) + self._latex_op + latex(self._Y)

    def __hash__(self):
        """
        The hash value of this set.

        EXAMPLES:

        The hash values of equal sets are in general not equal since it is not
        decidable whether two sets are equal::

            sage: X = Set(GF(13)).intersection(Set(ZZ))
            sage: Y = Set(ZZ).intersection(Set(GF(13)))
            sage: hash(X) == hash(Y)
            False

        TESTS:

        Test that :trac:`14432` has been resolved::

            sage: S = Set(ZZ).union(Set([infinity]))
            sage: T = Set(ZZ).union(Set([infinity]))
            sage: hash(S) == hash(T)
            True
        """
        return hash((self._X, self._Y, self._op))

class Set_object_union(Set_object_binary):
    """
    A formal union of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ^2)
            sage: T = Set(ZZ)
            sage: X = S.union(T); X
            Set-theoretic union of Set of elements of Vector space of dimension 2 over Rational Field and Set of elements of Integer Ring

            sage: latex(X)
            \Bold{Q}^{2} \cup \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "union", "\\cup")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,0))
            sage: Z = Set(Primes())
            sage: X.union(Y).is_finite()
            True
            sage: X.union(Z).is_finite()
            False
        """
        return self._X.is_finite() and self._Y.is_finite()

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic union of sets in
           a meaningful manner.  So be careful when using this.

        EXAMPLES::

            sage: Y = Set(ZZ^2).union(Set(ZZ^3))
            sage: X = Set(ZZ^3).union(Set(ZZ^2))
            sage: X == Y
            True
            sage: Y == X
            True

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: Set(ZZ).union(Set(QQ)) == Set(QQ)
            False
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_union):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator over the elements of ``self``.

        EXAMPLES::

            sage: [x for x in Set(GF(3)).union(Set(GF(2)))]
            [0, 1, 2, 0, 1]
        """
        for x in self._X:
            yield x
        for y in self._Y:
            yield y

    def __contains__(self, x):
        """
        Return ``True`` if ``x`` is an element of ``self``.

        EXAMPLES::

            sage: X = Set(GF(3)).union(Set(GF(2)))
            sage: GF(5)(1) in X
            False
            sage: GF(3)(2) in X
            True
            sage: GF(2)(0) in X
            True
            sage: GF(5)(0) in X
            False
        """
        return x in self._X or x in self._Y

    def cardinality(self):
        """
        Return the cardinality of this set.

        EXAMPLES::

            sage: X = Set(GF(3)).union(Set(GF(2)))
            sage: X
            {0, 1, 2, 0, 1}
            sage: X.cardinality()
            5

            sage: X = Set(GF(3)).union(Set(ZZ))
            sage: X.cardinality()
            +Infinity
        """
        return self._X.cardinality() + self._Y.cardinality()

class Set_object_intersection(Set_object_binary):
    """
    Formal intersection of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ^2)
            sage: T = Set(ZZ)
            sage: X = S.intersection(T); X
            Set-theoretic intersection of Set of elements of Vector space of dimension 2 over Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q}^{2} \cap \Bold{Z}

            sage: X = Set(IntegerRange(100)).intersection(Primes())
            sage: X.is_finite()
            True
            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "intersection", "\\cap")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(IntegerRange(100))
            sage: Y = Set(ZZ)
            sage: X.intersection(Y).is_finite()
            True
            sage: Y.intersection(X).is_finite()
            True
            sage: Y.intersection(Set(QQ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return True
        elif self._Y.is_finite():
            return True
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic intersection of
           sets in a meaningful manner.  So be careful when using this.

        EXAMPLES::

            sage: Y = Set(ZZ).intersection(Set(QQ))
            sage: X = Set(QQ).intersection(Set(ZZ))
            sage: X == Y
            True
            sage: Y == X
            True

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: Set(ZZ).intersection(Set(QQ)) == Set(QQ)
            False
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_intersection):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        ``self`` is a formal intersection of `X` and `Y` and this function is
        implemented by iterating through the elements of `X` and for
        each checking if it is in `Y`, and if yielding it.

        EXAMPLES::

            sage: X = Set(ZZ).intersection(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            2

        Check that known finite intersections have finite iterators (see
        :trac:`18159`)::

            sage: P = Set(ZZ).intersection(Set(range(10,20)))
            sage: list(P)
            [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
        """
        X = self._X
        Y = self._Y
        if not self._X.is_finite() and self._Y.is_finite():
            X,Y = Y,X
        for x in X:
            if x in Y:
                yield x

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is a formal intersection of `X` and `Y` this function
        returns ``True`` if both `X` and `Y` contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).intersection(Set(RR))
            sage: 5 in X
            True
            sage: ComplexField().0 in X
            False

        Any specific floating-point number in Sage is to finite precision,
        hence it is rational::

            sage: RR(sqrt(2)) in X
            True

        Real constants are not rational::

            sage: pi in X
            False
        """
        return x in self._X and x in self._Y

class Set_object_difference(Set_object_binary):
    """
    Formal difference of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ)
            sage: T = Set(ZZ)
            sage: X = S.difference(T); X
            Set-theoretic difference of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q} - \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "difference", "-")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,5))
            sage: Z = Set(QQ)
            sage: X.difference(Y).is_finite()
            True
            sage: X.difference(Z).is_finite()
            True
            sage: Z.difference(X).is_finite()
            False
            sage: Z.difference(Set(ZZ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return True
        elif self._Y.is_finite():
            return False
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic intersection of
           sets in a meaningful manner.  So be careful when using
           this.

        EXAMPLES::

            sage: Y = Set(ZZ).difference(Set(QQ))
            sage: Y == Set([])
            False
            sage: X = Set(QQ).difference(Set(ZZ))
            sage: Y == X
            False
            sage: Z = X.difference(Set(ZZ))
            sage: Z == X
            False

        This illustrates that equality testing for formal unions
        can be misleading in general.

        ::

            sage: X == Set(QQ).difference(Set(ZZ))
            True
        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_difference):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        ``self`` is a formal difference of `X` and `Y` and this function
        is implemented by iterating through the elements of `X` and for
        each checking if it is not in `Y`, and if yielding it.

        EXAMPLES::

            sage: X = Set(ZZ).difference(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            -2
            sage: next(I)
            -3
        """
        for x in self._X:
            if x not in self._Y:
                yield x

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is a formal intersection of `X` and `Y` this function
        returns ``True`` if both `X` and `Y` contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).difference(Set(ZZ))
            sage: 5 in X
            False
            sage: ComplexField().0 in X
            False
            sage: sqrt(2) in X     # since sqrt(2) is not a numerical approx
            False
            sage: sqrt(RR(2)) in X # since sqrt(RR(2)) is a numerical approx
            True
            sage: 5/2 in X
            True
        """
        return x in self._X and x not in self._Y

class Set_object_symmetric_difference(Set_object_binary):
    """
    Formal symmetric difference of two sets.
    """
    def __init__(self, X, Y):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: S = Set(QQ)
            sage: T = Set(ZZ)
            sage: X = S.symmetric_difference(T); X
            Set-theoretic symmetric difference of Set of elements of Rational Field and Set of elements of Integer Ring
            sage: latex(X)
            \Bold{Q} \bigtriangleup \Bold{Z}

            sage: TestSuite(X).run()
        """
        Set_object_binary.__init__(self, X, Y, "symmetric difference", "\\bigtriangleup")

    def is_finite(self):
        r"""
        Return whether this set is finite.

        EXAMPLES::

            sage: X = Set(range(10))
            sage: Y = Set(range(-10,5))
            sage: Z = Set(QQ)
            sage: X.symmetric_difference(Y).is_finite()
            True
            sage: X.symmetric_difference(Z).is_finite()
            False
            sage: Z.symmetric_difference(X).is_finite()
            False
            sage: Z.symmetric_difference(Set(ZZ)).is_finite()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        if self._X.is_finite():
            return self._Y.is_finite()
        elif self._Y.is_finite():
            return False
        raise NotImplementedError

    def __richcmp__(self, right, op):
        r"""
        Try to compare ``self`` and ``right``.

        .. NOTE::

           Comparison is basically not implemented, or rather it could
           say sets are not equal even though they are.  I don't know
           how one could implement this for a generic symmetric
           difference of sets in a meaningful manner.  So be careful
           when using this.

        EXAMPLES::

            sage: Y = Set(ZZ).symmetric_difference(Set(QQ))
            sage: X = Set(QQ).symmetric_difference(Set(ZZ))
            sage: X == Y
            True
            sage: Y == X
            True

        """
        if not isinstance(right, Set_generic):
            return rich_to_bool(op, -1)
        if not isinstance(right, Set_object_symmetric_difference):
            return rich_to_bool(op, -1)
        if self._X == right._X and self._Y == right._Y or \
           self._X == right._Y and self._Y == right._X:
            return rich_to_bool(op, 0)
        return rich_to_bool(op, -1)

    def __iter__(self):
        """
        Return iterator through elements of ``self``.

        This function is implemented by first iterating through the elements
        of `X` and  yielding it if it is not in `Y`.
        Then it will iterate throw all the elements of `Y` and yielding it if
        it is not in `X`.

        EXAMPLES::

            sage: X = Set(ZZ).symmetric_difference(Primes())
            sage: I = X.__iter__()
            sage: next(I)
            0
            sage: next(I)
            1
            sage: next(I)
            -1
            sage: next(I)
            -2
            sage: next(I)
            -3
        """
        for x in self._X:
            if x not in self._Y:
                yield x

        for y in self._Y:
            if y not in self._X:
                yield y

    def __contains__(self, x):
        """
        Return ``True`` if ``self`` contains ``x``.

        Since ``self`` is the formal symmetric difference of `X` and `Y`
        this function returns ``True`` if either `X` or `Y` (but not both)
        contains ``x``.

        EXAMPLES::

            sage: X = Set(QQ).symmetric_difference(Primes())
            sage: 4 in X
            True
            sage: ComplexField().0 in X
            False
            sage: sqrt(2) in X      # since sqrt(2) is currently symbolic
            False
            sage: sqrt(RR(2)) in X # since sqrt(RR(2)) is currently approximated
            True
            sage: pi in X
            False
            sage: 5/2 in X
            True
            sage: 3 in X
            False
        """
        return (x in self._X and x not in self._Y) \
               or (x in self._Y and x not in self._X)

def is_Set(x):
    """
    Deprecated. Use ``isinstance(x, Set_generic)`` instead.

    TESTS::

        sage: from sage.sets.set import is_Set
        sage: is_Set(Primes())
        doctest:...: DeprecationWarning: Please use isinstance(x, Set_generic)
        See http://trac.sagemath.org/24443 for details.
        True
    """
    from sage.misc.superseded import deprecation
    deprecation(24443, "Please use isinstance(x, Set_generic)")
    return isinstance(x, Set_generic)


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

#

#

#

#
Applied Seismology

#
Lab: Instrument Response

#

#

#

# Seismo-Live: http://seismo-live.org
#
# ##### Authors:
# * Carl Tape ([@carltape](https://github.com/carltape))
# * Lion Krischer ([@krischer](https://github.com/krischer))
# ---

# based on *GEOS 626: Applied Seismology from Carl Tape*
#
# ---

# This cell does two things:
# (1) make plots appear in the notebook and (2) creates prettier plots.
# Make sure to execute it at least once!
# %matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('ggplot')
plt.rcParams['figure.figsize'] = 10, 5

# ### L1 Instructions
#
# This is a Python and ObsPy version of the above mentioned lab. We will perform the same calculations and create similar plots but we will be using different services, data formats, and tools.
#
# It aims to teach students how to calculate, visualize and use instruments responses.

# ### L2 Computing and plotting instrument response, $I(ω)$

# ##### Acquiring Data
#
# The first step when working with instrument data is to acquire it. Nowadays this is mostly done using the *station* web service defined by the [FDSN](http://www.fdsn.org/) (International Federation of Digital Seismograph Networks). Once you know how to use this you can use it to download data from data centers across the globe including IRIS, ORFEUS, ETH, RESIF, GEONET, and many more. Notable exceptions are most Asian countries and countries that don't freely share their data.

# With ObsPy this works by first importing the FDSN `Client` class and then initializing it with the data center you want to download from. For this lab we will use data from the Geoscope project, which can be downloaded from IRIS (its "home" data center is the IPGP in Paris which you can also use - just replace `"IRIS"` with `"IPGP"` in the following code box - but most of you will be more familiar with IRIS).
#
# See the [documentation of obspy.clients.fdsn](http://docs.obspy.org/packages/obspy.clients.fdsn.html) for more details.

from obspy.clients.fdsn import Client
# ObsPy knows which website to connect to for many data center.
# For others you can also pass the full URL.
c = Client("IRIS")

# Station information is time dependent due to for example changed sensors, new calibrations, or other changes. It is thus important to specify the time frame for which to request response information.
#
# ObsPy handles absolute times with the `UTCDateTime` class. See the [documentation](http://docs.obspy.org/packages/autogen/obspy.core.utcdatetime.UTCDateTime.html) and the [tutorial](http://docs.obspy.org/tutorial/code_snippets/utc_date_time.html) for more information.

import obspy
# Create two times which we will later use to download time series data.
starttime = obspy.UTCDateTime(2004, 12, 26, 0, 58, 50)
# Create a new time object by adding 5 days to the previous ones. Time
# differences are in seconds.
endtime = starttime + 86400 * 5

# We will use the `CAN` station which is located in Canberra, Australia, and featured in [*Park et al.* (2005, Figure 1)](http://dx.doi.org/10.1126/science.1112305). Please note that in many cases the station code is not unique and in practice the network code is also required. See the ObsPy tutorials for more details about these concepts. The `get_station()` method of the FDSN `Client` object returns station/instrument information. Once again see its [documentation](http://docs.obspy.org/packages/autogen/obspy.clients.fdsn.client.Client.get_stations.html) for more information.

# This will get all stations that satisfy all constraints simultaneosly.
inv = c.get_stations(network="G", station="CAN", channel="LHZ",
                     level="response",
                     starttime=starttime, endtime=endtime)
print(inv)

# The `station` services return `StationXML` files which are the most modern and future proof data format for station and instrument information. **If possible at all: Always try to use StationXML files!**

# ##### Plotting the Instrument Response
#
# We can now plot the instrument response with the `plot_response()` method. The `min_freq` argument determines the lowest plotted frequency - the maximum plotted frequency is always the Nyquist frequency.

inv.plot_response(min_freq=0.001);

# These are [Bode plots](https://en.wikipedia.org/wiki/Bode_plot) and show the amplitude $A(\omega)$ and the phase $\phi(\omega)$ response where
#
# $$
# I(\omega) = A(\omega) e^{i \phi(\omega)}
# $$
#
# Note that the formulas use $\omega$ but the plots show the frequency $f$ with $\omega = 2 \pi f$. The first variant usually simplifies notation whereas the latter one facilitates physical interpretation.
#
# This plot does not just show the instrument response - it shows the response of the whole recording chain from observed ground motion to actually stored data. This includes the analoge response of the actual seismometer, the analog-to-digitial converter (DAC), and possible digital filter stages.

# The sampling rate of this channel is 1 sample per second ($\Delta T$ = 1s), so the **Nyquist frequency**, which is defined as
#
# $$
# f_{Nyq} = 1 / (2 \Delta t),
# $$
#
# is $f_{Nyq}$ = 0.5 Hz. The above plot does not show frequencies higher than the Nyquist frequency (the vertical dashed line) as these frequencies do not exist in the data and thus have no physical meaning.

# Have a look at the details of the response information of this channel:

inv[0][0][0].response

# Here we see that it consists of multiple stages that convert from one unit to another. The final instrument response is the multipication (in the frequency domain) of all of these.
#
# **QUESTION:** What are the input and output units of the whole instrument? What does that mean physically?
#
# **QUESTION:** How does one conceptually remove the instrument response from recorded data?
#
# ##### Incomplete Responses
#
# You can also only plot the response of a range of stages. The plotted sensitivity in Obspy 1.0.1 is the overall sensitivity of all stages and not only the chosen ones - thus the plot looks a bit strange. This will be rectified [soon](https://github.com/obspy/obspy/issues/1368).

inv[0][0][0].plot(0.001, start_stage=1, end_stage=1);

# Note that the major differences (except a constant factor) are close to the Nyquist frequency - these are due to the here not plotted decimation and FIR filter stages. In many cases it is sufficient to only use Poles & Zeros information but there are stations where this is not true. As it is tedious to manually check everytime: **There is no reason to not use the full chain instrument response if it is available!**

# ##### Output Units

# Modern seismometers usually record the velocity of ground motion, strong motions sensor record acceleration.
#
# **QUESTION:** Why is that?
#
# The instrument response describes how incoming data (in whatever the instrument measures) it transformed into what is finally stored in a file. Restoring original ground motion requires the inversion of that process. The deconvolution of the instrument response as written for example in a StationXML file will convert data back to its original units.
#
# If your analysis requires other units it is best to perform the integration or differentiation during the instrument deconvolution. That deconvolution is usually performed in the frequency domain - once there integration/differentiation is almost free and the most accurate it gets.

# The following plots show the instrument response $I_d(\omega)$ from displacement, velocity $I_v(\omega)$, and acceleration $I_a(\omega)$.

print("Displacement")
inv.plot_response(0.001, output="DISP")
print("Velocity")
inv.plot_response(0.001, output="VEL")
print("Acceleration")
inv.plot_response(0.001, output="ACC");

#
# Basic properties of the Fourier transform result in
#
# $$
# \begin{align}
# X_v(\omega) &= (i \omega) X_d(\omega) \\
# X_a(\omega) &= (i \omega) X_v(\omega)
# \end{align}
# $$
#
# where $X_d$, $X_v$, and $X_a$ are the Fourier transforms of displacement $x_d(t)$, velocity $x_v(t)$, and acceleration $x_a(t)$.
#
# We also have the relationship between ground motion, instrument response, and the output form the seismometer:
#
# $$
# \begin{align}
# x_a(t) * i_a(t) &= c(t) \\
# X_a(\omega) I_a (\omega) &= C(\omega)
# \end{align}
# $$
#
# Here we have assumed that the instrument response is described with respect to acceleration. But we could alternatively consider the instrument response with respect to velocity or displacement; **the key point is that what comes out of the seismometer,** $C(\omega)$**, is fixed.** But we can descibe the input ground motion as displacement, velocity, or acceleration. Showing all three together and omitting explicit $\omega$ dependence, we have
#
# $$
# \begin{align}
# C &= X_a I_a = X_v I_v = X_d I_d\\
#   &= (i\omega) X_v I_a = (i \omega) X_d I_V = X_d I_d \\
#   &= (i\omega)^2 X_d I_a = (i \omega) X_d I_V = X_d I_d \\
# I_v &= I_d \ (i\omega)\\
# I_a&= I_v \ (i\omega)\\
# \end{align}
# $$
#
# It turns out that the effect of differentiation in the time domain leads to an *increase by a factor of one* in the slope of the amplitude spectrum, $|H(\omega)|$, in log-log space, for example, by changing from $X_d(\omega)$ to $X_v(\omega)$. But when we are looking at the *instrument response*, the slope will *decrease by a factor of one* when changing from, say, $I_d(\omega)$ to $I_v(\omega)$.
#
# We see this in the above figure. Consider the flat segment in the velocity correction spectrum: the displacement shows a corresponding slope increase whereas the acceleration show a slope decrease.
#
# **QUESTION:** What is meant by a "broadband" seismometer?

# ### L3 Deconvolve instrument response for a seismogram

# We will download some data and remove its instrument response.
#
# The first step is to get some data - we will once again use the FDSN web services but this time also to download waveform data.

# +
otime = obspy.UTCDateTime(2004, 12, 26, 0, 58, 50)
duration = 5 * 86400
starttime = otime
endtime = starttime + duration

st = c.get_waveforms(network="G", station="CAN", location="", channel="LH*",
                     starttime=starttime, endtime=endtime)

# Do it once again to also get the response information of the vertical channels.
inv = c.get_stations(network="G", station="CAN", location="", channel="LH*",
                     starttime=starttime, endtime=endtime, level="response")
# -

# We will now show what the main arrival looks like at bandpass 50-500s with and without the instrument deconvolution.

# Plot the first hour - slice() will leave the original
# data stream intact and just returns a new view on the
# data.
st.slice(endtime=otime + 3600).copy().taper(0.05).filter(
    "bandpass", freqmin=1.0 / 500.0, freqmax=1.0 / 50.0).plot()

# Now plot the same but with the instrument removed.
# Note the copy() here as remove_response() would otherwise
# remove the contents of the original data stream.
st.slice(endtime=otime + 3600).copy().remove_response(
    inventory=inv, output="VEL").taper(0.05).filter(
    "bandpass", freqmin=1.0 / 500.0, freqmax=1.0 / 50.0).plot()

# What are the differences? What is their cause?

# We will now plot the amplitude spectrum over 0.2 - 1.0 mHz to show the gravest normal mode peaks. We will only use the vertical component to simplify things.
#
# We will first do it without instrument correction

# +
import numpy as np

tr = st.select(component="Z")[0].copy().taper(0.05)

# rfft() is a varient of the FFT for real valued input.
# The negative frequencies of a DFT for purely real valued
# input are just the complex conjugates of the corresponding
# positive frequencies and thus redundant.
D = np.fft.rfft(tr.data)
# Return the corresponding frequencies.
freqs = np.fft.rfftfreq(tr.stats.npts, d=tr.stats.delta)

plt.plot(freqs, np.abs(D))
plt.xlim(0.2E-3, 1.0E-3)
plt.ylim(0, 1E7)
plt.xlabel("Frequency [mHz]")
plt.show()
# -

# Instrument removal is, in practice, a tricky and unstable operation. Have a look at the [documentation of the remove_response() method](http://docs.obspy.org/packages/autogen/obspy.core.trace.Trace.remove_response.html) and try to understand the meaning of the various parameters and when you would want to use each.

# +
# Remember the copy()!
tr = st.select(component="Z").copy().remove_response(inventory=inv,
                                                     output="DISP")[0]

D = np.fft.rfft(tr.data)
freqs = np.fft.rfftfreq(tr.stats.npts, d=tr.stats.delta)

plt.plot(freqs, np.abs(D))
plt.xlim(0.2E-3, 1.0E-3)
plt.ylim(0, 5.0)
plt.xlabel("Frequency [mHz]")
plt.show()
# -

# Are there any difference to the previous versions? If yes - what are they? Keep in mind that the previous plot is based on velocity data, whereas this one has been corrected to displacement.

# We are interested in the spectrum of the data so we don't really need to use instrument corrected data to calculate the spectrum - we can just calculate the spectrum and divide by the response.

# +
tr = st.select(component="Z")[0].copy().taper(0.05)

D = np.fft.rfft(tr.data)
freqs = np.fft.rfftfreq(tr.stats.npts, d=tr.stats.delta)
# Advanced ObsPy-foo to directly get the complex response.
freq_response, _ = \
    inv.select(channel="LHZ")[0][0][0].response.get_evalresp_response(
        tr.stats.delta, tr.stats.npts, output="DISP")

plt.plot(freqs, np.abs(D)/ np.abs(freq_response))
plt.xlim(0.2E-3, 1.0E-3)
plt.ylim(0, 35.0)
plt.xlabel("Frequency [mHz]")
plt.show()
# -

# Note that this is quite different then the previous spectrum. **Why?** Keep in mind that we are operating far from the passband of the instrument. The following two plots as well as the
# [documentation of the remove_response() method](http://docs.obspy.org/packages/autogen/obspy.core.trace.Trace.remove_response.html) should help you answer that question. Which alternative yields the more correct answer?

inv.select(channel="LHZ").plot_response(min_freq=0.0001);

tr = st.select(component="Z").copy().remove_response(
    inventory=inv, output="DISP", plot=True)[0]


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.3
#   kernelspec:
#     display_name: tcv-x21
#     language: python
#     name: tcv-x21
# ---

# %% [markdown] pycharm={"name": "#%% md\n"}
# # Setting up simulations of your own
#
# The ultimate aim of TCV-X21 is that the validation case can be replicated by
# other simulations, and that the experimental dataset serves as a benchmark
# for comparing and validating turbulence simulations.
#
# For that, we need *you* to be able to set up simulations with a turbulence
# code of your own.
#
# For this, you'll need
#
# 1. A turbulence code of your own
# 2. The reference equilibrium
# 3. Sources
# 4. Physical parameters
#
# You can find data files for the experimental reference scenario in
# `data/experimental_reference/reference_scenario`.

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
# Set up the analysis environment

# %load_ext autoreload
# %autoreload 2
# %matplotlib inline

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from pathlib import Path
from netCDF4 import Dataset

import tcvx21
from tcvx21 import Quantity, test_session

# Apply the custom style sheet, which makes the plots look the same
plt.style.use(tcvx21.style_sheet_inline)

# %% [markdown] pycharm={"name": "#%% md\n"}
# ## The reference equilibrium
#
# We develop a simple lower-single-null diverted magnetic geometry as the "reference
# scenario". We performed several discharges in the reference scenario to gather
# the experimental reference dataset. Next, we performed simulations in both
# field directions of a "reference scenario": TCV shot number 65402 at
# $t=1$s. The reference scenario represents a typical discharge from the
# experimental discharges, and it is the case modelled by the codes. Note that
# the simulations *do not use baffles* (unlike the shot number given here).
#
# The reference equilibrium is available in several file-formats: including
# `eqdsk`, `.mat` (MATLAB data file), and a labelled PARALLAX-type NetCDF.
# We demonstrate here the use of the labelled NetCDF. Note that in the figure
# shown here is in *normalised units*. All the spatial distances are
# divided by $R_0 = 0.90586$m.

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
# Load the equilibrium
equi = Dataset(
    tcvx21.experimental_reference_dir / "reference_scenario/reference_equilibrium.nc"
)

fig, ax = plt.subplots()

# Plot the flux surfaces
plt.contourf(
    equi["Magnetic_geometry/R"][:],
    equi["Magnetic_geometry/Z"][:],
    equi["Magnetic_geometry/psi"][:],
    levels=20,
)

cbar = plt.colorbar()

# Plot the separatrix
plt.contour(
    equi["Magnetic_geometry/R"][:],
    equi["Magnetic_geometry/Z"][:],
    equi["Magnetic_geometry/psi"][:],
    levels=[equi["Psi_limits"].psi_seperatrix],
    colors="red",
)

# Plot the divertor
plt.plot(
    equi["divertor_polygon/R_points"], equi["divertor_polygon/Z_points"], color="blue"
)

# Highlight the magnetic axis and X-point
R0 = equi["Magnetic_geometry"].magnetic_axis_R  # in metres
plt.scatter(
    equi["Magnetic_geometry"].magnetic_axis_R / R0,
    equi["Magnetic_geometry"].magnetic_axis_Z / R0,
    marker="x",
    color="blue",
)
plt.scatter(
    equi["Magnetic_geometry"].x_point_R / R0,
    equi["Magnetic_geometry"].x_point_Z / R0,
    marker="x",
    color="green",
)

# Make the plot look nice and add labels
ax.set_aspect("equal")
ax.set_title("TCV-X21 reference equilibrium")
ax.set_xlabel("$R/R_0$ normalised major radius")
ax.set_ylabel("$Z/R_0$ normalised vertical position")

cbar.ax.set_ylabel("$\Psi$ [Weber]", rotation=270, labelpad=20)

if test_session:
    plt.close("all")

# %% [markdown]
# The equilibrium file should give you enough information to exactly replicate
# the magnetic field structure of the TCV-X21 reference equilibrium. To
# construct the poloidal magnetic field from the equilibrium file, you can
# check the methods of `tcvx21.grillix_analysis.equi_m.Equi`.
#
# For the toroidal magnetic field strength, we assume
# $B_\phi = B_{\phi, axis}\frac{R_0}{R}$

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
from tcvx21.grillix_post.components.equi_m import Equi

grillix_equi = Equi(
    tcvx21.experimental_reference_dir / "reference_scenario/reference_equilibrium.nc"
)

fig, ax = plt.subplots()

r_test, z_test = np.linspace(0.6, 1.3), np.linspace(-0.9, 0.9)
plt.contourf(
    r_test,
    z_test,
    np.sqrt(
        grillix_equi.magnetic_field_r(r_test, z_test, grid=True) ** 2
        + grillix_equi.magnetic_field_z(r_test, z_test, grid=True) ** 2
    ),
    levels=20,
)
cbar = plt.colorbar()

plt.plot(
    equi["divertor_polygon/R_points"], equi["divertor_polygon/Z_points"], color="blue"
)

ax.set_aspect("equal")
ax.set_title("Poloidal magnetic field strength")
ax.set_xlabel("$R/R_0$ normalised major radius")
ax.set_ylabel("$Z/R_0$ normalised vertical position")

cbar.ax.set_ylabel("$B_{pol}$ [Tesla]", rotation=270, labelpad=20)

if test_session:
    plt.close("all")

# %% [markdown]
# ## Sources
#
# The source functions used by GRILLIX are given in `tcvx21/grillix_post/components/sources_m.py`
#
# ### Density source or neutrals
#
# The scenario is fuelled from above, away from the divertor, to try to
# reduce the amount of neutrals ionised in the divertor.
# We expect that the neutrals are mostly ionised just inside the confined region (this is
# called "low-recycling": this doesn't mean that there is no neutral recycling,
# just that the recycled neutrals should be primarily ionised in the confined
# region).
#
# There are two ways that this could be treated.
#
# 1. Approximate the neutral ionisation as a confined-region source
# 2. Self-consistently treat the neutral dynamics to determine the source rate
#
# For the initial evaluation we select option 1 for simplicity.
# Nevertheless, we note that using this approximation
# is likely our greatest source of disagreement, and are interested in
# comparing to simulations with neutrals.
#
# The total particle source rate isn't known exactly. We expect that
# most of the density source results from recycled neutrals rather than
# from neutrals injected from the gas fuelling port (and we don't have a
# good estimate of the gas fuelling rate anyway). As such, the
# density source rate (or neutral pressure) must be manually adjusted (see
# tuning the sources).
#
# The density source used by GRILLIX is
#
# $S_n = \hat{S}_n \exp\left[ -(\rho^2 - \rho_c^2)/\rho_w\right]$
#
# where $\hat{S}_n$ is a source rate, $\rho_c = 0.915141$ is the source centre and
# $\rho_w = 0.083797$ is the source width.
#
# ### Temperature or power source
#
# The shots are Ohmically heated (resistive heating due to transformer-driven
# current). We approximate this as a power source near the magnetic axis.
#
# The electron temperature source is modified to give a constant power injection. The
# power integral is
#
# $ P = \frac{3}{2}\iiint n S_{T_e} + (T_e + T_i) S_n \textrm{d}^3V$
#
# which has contributions both from the electron temperature source and from the density
# source (since adding plasma at a certain temperature requires energy). In GRILLIX,
# we rearrange this to find an expression for $S_{T_e}$ which gives a constant power
# rate (where $\hat{x} = x / x_0$ is a normalised form of $x$, with normalisation factor $x_0$)
#
# $S_{T_e} = T_{e0} / \tau_0 \left(\frac{1}{\hat{n}}\hat{S}_{T_{e,core}} - \frac{\hat{T}_e + \hat{T}_i}{\hat{n}} \hat{S}_n \right)$
#
# This allows us to write
#
# $P = \frac{3}{2}\iiint n_0 \hat{n} T_{e0} / \tau_0 \left(\frac{1}{\hat{n}}\hat{S}_{T_{e,core}} - \frac{\hat{T}_e + \hat{T}_i}{\hat{n}} \hat{S}_n \right) + T_{e0}(\hat{T}_e + \hat{T}_i) n_0 / \tau_0 \hat{S}_n \textrm{d}^3V$
#
# $P= \frac{3}{2}\frac{n_0 T_{e0}}{\tau_0} \iiint \hat{S}_{T_{e,core}} - (\hat{T}_e + \hat{T}_i) \hat{S}_n + (\hat{T}_e + \hat{T}_i) \hat{S}_n \textrm{d}^3V$
#
# $P= \frac{3}{2}\frac{n_0 T_{e0}}{\tau_0} \hat{S}_{T_{e,core}} \iiint f_{S,Te}(R,Z) \textrm{d}^3V$
#
# $P= \frac{3}{2}\frac{n_0 T_{e0}}{\tau_0} \hat{S}_{T_{e,core}} \mathcal{V}_w$
#
#
# Substituting this into the power integral gives
# $ P = \frac{3}{2} \int S_{T_{e,core}} d^3V$
#
# Then we set
# $S_{T_{e,core}} = 1 - \mathcal{S}_3(\rho, \rho_c, \rho_w)$
# for $\mathcal{S}_3$ a third-order [smoothstep](https://en.wikipedia.org/wiki/Smoothstep)
# function, centred at $\rho_c = 0.3$ and of width $\rho_w =0.3$.
#
# The term $\mathcal{V}_w$ is the effective weighted volume of the source, which we can
# use to work out the required source rate.
#
# Note that we are only taking energy out of the electron population to compensate for the density source.

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
from tcvx21.grillix_post import components
from tcvx21.grillix_post.components.sources_m import density_source_function, smoothstep

file_path = tcvx21.test_dir / "sample_data"

grid = components.Grid(file_path / "vgrid.nc")
norm = components.Normalisation.initialise_from_normalisation_file(
    file_path / "physical_parameters.nml"
)
snaps = components.read_snaps_from_file(
    file_path, norm, time_slice=slice(None), all_planes=True
)
equi = components.Equi(
    file_path / "TCV_ortho.nc", file_path / "pen_metainfo.nc", flip_z=True
)

parameter_filepath = tcvx21.grillix_post.filepath_resolver(file_path, "params.in")
params = components.convert_params_filepaths(
    parameter_filepath, components.read_fortran_namelist(parameter_filepath)
)

rho = equi.normalised_flux_surface_label(grid.r_s, grid.z_s)

density_source_rate = params["params_srcsnk"]["csrcn"]
etemp_source_rate = params["params_srcsnk"]["csrcte"]

density_source = density_source_function(rho, grid.districts, norm, source_strength=1.0)
core_source = xr.DataArray(1 - smoothstep(rho, step_centre=0.3, step_width=0.3)).where(
    grid.districts == "CLOSED", 0.0
)

if test_session:
    plt.close("all")

# %% [markdown] pycharm={"name": "#%% md\n"}
# ### Visualising the sources
#
#

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
from tcvx21.plotting.plot_array_as_transparency_m import plot_array_as_transparency

fig, ax = plt.subplots()

flip_z = -1.0 if equi.flipped_z else 1.0

# First, let's plot the magnetic data flux surfaces
ax.contour(grid.r_s * R0, grid.z_s * R0 * flip_z, rho, levels=20)

# Plot the sources, and bring them to in front of the flux-surfaces
im = plot_array_as_transparency(
    ax,
    grid.r_s * R0,
    grid.z_s * R0 * flip_z,
    alphas=np.flipud(density_source),
    cmap=plt.cm.Greens,
    intensity=0.7,
)
im.set_zorder(np.inf)
im = plot_array_as_transparency(
    ax,
    grid.r_s * R0,
    grid.z_s * R0 * flip_z,
    alphas=np.flipud(core_source),
    cmap=plt.cm.Blues,
    intensity=0.7,
)
im.set_zorder(np.inf)

divertor_ = tcvx21.read_from_json(
    tcvx21.experimental_reference_dir / "reference_scenario/divertor_polygon.json"
)
separatrix = tcvx21.analysis.find_contours(
    grid.r_s * R0, grid.z_s * R0, rho, level=1.0
)[0]

# Mark the separatrix
ax.plot(
    separatrix[:, 0],
    separatrix[:, 1] * flip_z,
    "C2",
    label="Separatrix",
    linestyle="--",
)
# Mark the divertor
ax.plot(
    divertor_["r_points"],
    divertor_["z_points"],
    color="C0",
    label="Vessel",
    linewidth=2.0 * plt.rcParams["lines.linewidth"],
)
# Set the plot limits in terms of the divertor
ax.set_xlim(divertor_["r_points"].min(), divertor_["r_points"].max())
ax.set_ylim(divertor_["z_points"].min(), divertor_["z_points"].max())
ax.set_aspect("equal")

plt.xlabel("R [m]")
plt.ylabel("Z [m]")
plt.title("Source positions")

if test_session:
    plt.close("all")

# %% [markdown] pycharm={"name": "#%% md\n"}
# ### Tuning the sources
#
# Depending on the type of source that you use, you have a number of free parameters. Either you
# will need to adjust
#
# * Density source rate and temperature source rate (electrons and possibly ions)
# * Density source rate and power source rate
#
# For this, you can tune the sources to match the following experimental constraints
#
# * Separatrix density of $\approx 6\times 10^{18} m^{-3}$
# * Separatrix electron temperature of $\approx 40 eV$
# * Total injected power of $150 kW$
# * Power crossing the separatrix of $\approx 120 kW$ (estimated from bolometry)
# * Total particle recycling of (very roughly) $\approx 3\times 10^{31} s^{-1}$ (estimated from total out-flux to
#   wall Langmuir probes, assuming perfect recycling).
#
# In GRILLIX, we set the total injected power to $150 kW$ and then tune the density source
# to match the separatrix density. We use the same sources in both field directions.
#
# We can work out exactly what the required core $T_e$ source rate is, by computing the
# $\mathcal{V}_w$ "effective weighted volume" of the core source, which allows us to then
# find the source rate as
#
# $\hat{S}_{T_e}= P / \left(\frac{3}{2}\frac{n_0 T_0}{\tau_0} \mathcal{V}_w\right)$

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
from tcvx21.grillix_post.observables.integrations_m import (
    axisymmetric_cylindrical_integration,
)
from tcvx21.grillix_post.components.sources_m import core_temperature_source_function
from tcvx21 import Dimensionless

core_source_volume = axisymmetric_cylindrical_integration(
    grid,
    norm,
    core_temperature_source_function(
        rho,
        grid.districts,
        norm,
        source_strength=1.0,
        source_centre=0.3,
        source_width=0.3,
    ).assign_attrs(norm=Dimensionless),
)

desired_power = Quantity(150, "kilowatt")
print(f"Effective weighted core source volume: {core_source_volume}")
print(f"Required source rate: {(desired_power / core_source_volume).to('MW m^-3')}")
print(
    f"Dimensionless source rate: {(desired_power / (1.5 * core_source_volume * norm.Te0 * norm.n0 / norm.tau_0)).to('')}"
)

# %% [markdown] pycharm={"name": "#%% md\n"}
# ## Physical parameters
#
# In turbulence simulations, usually you'll convert to dimensionless quantities, such that all
# modelled equations are in terms of unitless quantities. The normalisation factors can be
# computed from physical parameters of the system.
#
# The codes do not fix the profiles, so the exact choice of the normalisation factors isn't so
# important. The file `data/experimental_reference/reference_scenario/physical_parameters.nml` contains
# the following physical parameters (in Fortran namelist format)
#
# ```
# &physical_parameters
#     case_name = TCV_65402
#     # # # ! Magnetic field normalisation, usually taken on axis, in Tesla
#     B0 = 0.929
#     # # # ! Electron temperature normalisation, in electron-volts
#     Te0 = 41.3
#     # # # ! Ion temperature normalisation, in electron-volts
#     Ti0 = 41.3
#     # # # ! Density normalisation, in particles-per-cubic-metres
#     n0 = 1E19
#     # # # ! Major radius, in metres
#     R0 = 0.906
#     # # # ! Ion mass, in amu
#     Mi = 2
#     # # # ! Ion charge, in e
#     Z = 1
#     # # # ! Ion effective charge, in e
#     Z_eff = 1.5
# /
# ```
#
# In GRILLIX, we use these physical parameters to calculate normalisation factors (using the excellent `pint`
# unit handling library). You can see the calculation and definition of the terms in
# `tcvx21/grillix_post/components/normalisation_m.py`.
#
# You'll need to do this for your own code, but this might provide a helpful starting point.
#
# Of the normalisation parameters, the reference ion drift scale is important, since
# we set our grid resolution in terms of it.

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
print(norm)
print(f"Reference ion drift scale is {norm.rho_s0.to('mm'):3.2}")

# %% [markdown] pycharm={"name": "#%% md\n"}
# ### Core profiles
#
# You might also need core profiles, such as for setting up initial condition.
#
# Since these were not included in the analysis, these haven't been checked as
# thoroughly as the edge profiles, but they still might be helpful.
#
# Note that ion temperature in forward field is not available.

# %% jupyter={"outputs_hidden": false} pycharm={"name": "#%%\n"}
plt.style.use(tcvx21.style_sheet)
core_profiles = tcvx21.read_struct_from_file(
    tcvx21.experimental_reference_dir / "reference_scenario/dataset_TCV_CORE.mat"
)

def plot_from_matlab(data, label):
    r, r_units, val, err, units = [
        data[key] for key in ["r", "r_units", "val", "err", "val_units"]
    ]

    r = Quantity(r, r_units.lstrip("[").rstrip("]"))
    val = Quantity(val, units.lstrip("[").rstrip("]"))
    err = Quantity(err, units.lstrip("[").rstrip("]"))

    plt.errorbar(r, val, err, label=label)
    plt.xlabel(f"$R^u - R^u_{{sep}}$ [{r.units}]")
    plt.ylabel(f"{val.units}")

fig, ax = plt.subplots()
plot_from_matlab(core_profiles["Forw"]["ne"], "n forward")
plot_from_matlab(core_profiles["Rev"]["ne"], "n reversed")
plt.title("Core density")
plt.legend()

fig, ax = plt.subplots()
plot_from_matlab(core_profiles["Forw"]["Te"], "$T_e$ forward")
plot_from_matlab(core_profiles["Rev"]["Te"], "$T_e$ reversed")
plot_from_matlab(core_profiles["Rev"]["Ti"], "$T_i$ reversed")
plt.title("Core temperatures")
plt.legend()

if test_session:
    plt.close("all")

# %% [markdown] pycharm={"name": "#%% md\n"}
# ## Conclusion
#
# Hopefully, this is enough for you to start a simulation of
# your own. If not, let us know via issues, and we'll
# extend the documentation and data accordingly.
#
# To post-process and compare to TCV-x21, you'll need to extract the data at a
# series of measurement positions. We do this in `simulation_postprocessing.ipynb`.
#
#