text:
import numpy as np

class Piece:
    '''
    Parent Class for all pieces
    Methods common to all pieces are defined here

    <var> name -> string : Stores the name of the piece in order to differentiate pieces

    <var> offsets -> numpy.ndarray : Stores the offsets which indiciate the very next position a piece can go to

    <var> is_sliding_piece -> bool : Sliding Pieces are : Queen, Rook and Bishop
    The non sliding pieces are Pawn, Knight, King

    <var> files -> dict : Maps the file value from string to an integer

    <var> ranks -> dict : Maps the rank value from string to an integer
    NOTE: Mapping is done in a manner such that when the board is printed on console, the board appears from white's perspective

    <func> possible_moves : Generates all possible moves for the pieces. This is overriden by the sub classes

    <func> move : Makes the move on the board if the move is legal

    <func> is_move_legal : Determines if the selected piece can actually move to the target square according to the rules of the game

    <func> is_path_blocked : Determines if the path from the current location of the SLIDING piece to the target square is blocked by a friendly or unfriendly piece.

    <func> is_target_friendly : Determines if the target square is occupied by a friendly piece and declares the move illegal

    <func> is_target_unfriendly : Determines if the target square is occupied by a unfriendly piece.
    '''

    name = None
    offsets = None
    is_sliding_piece = True

    files = {
        'a': 0,
        'b': 1,
        'c': 2,
        'd': 3,
        'e': 4,
        'f': 5,
        'g': 6,
        'h': 7,
    }

    ranks = {
        '1': 7,
        '2': 6,
        '3': 5,
        '4': 4,
        '5': 3,
        '6': 2,
        '7': 1,
        '8': 0,
    }

    def __init__(self, color, current_position):
        '''
        <param> color -> string : 'white' or 'black' indicating the color of the piece

        <param> current_position -> string : The location of the square where the piece must be placed eg. 'a1' or 'h8'

        <optional param> is_sliding_piece -> bool : True indicates that the piece is a sliding piece. Should be False for Knights only
        '''
        assert isinstance(color, str), 'Color values can be either \'white\' or \'black\''
        assert isinstance(current_position, str), 'Position of the square must be in the standard algebraic notation for Chess, eg. \'A1\' or \'h8\''
        assert all([len(current_position) == 2,
                    current_position[0] in self.files,
                    current_position[1] in self.ranks]), 'Invalid Square'
        assert isinstance(self.is_sliding_piece, bool), 'is_sliding_piece must be a True or False value'
        self.color = color
        self.current_position_as_string = current_position.lower()
        self.current_position = np.array([self.ranks[current_position[1]], self.files[current_position[0]]])

    def get_current_position_as_string(self):
        '''
        <return> -> string : The current location of the piece as a string, eg. 'a1' or 'h8'
        '''
        return self.current_position_as_string

    def get_current_position(self):
        '''
        <return> -> numpy.ndarray : The current location of the piece as an element of a numpy array, eg. array([1 0]) or array([3 4])
        NOTE : The first element is the rank, the second element in the file
        '''
        return self.current_position

    def set_current_position(self, target_position):
        '''
        <param> target_position -> string : The co-ordinates of the target square
        '''
        rank, file = self.ranks[target_position[1]], self.files[target_position[0]]
        self.current_position = np.array([rank, file])

    def possible_moves(self, board):
        '''
        <param> board -> np.ndarray : Shows the board

        <return> -> generator object : Returns a generator of all possible moves of a piece when clicked
        '''
        possible_moves = []
        for offset in self.offsets:
            position = self.get_current_position()
            rank, file = position[0], position[1]
            while True:
                rank_offset, file_offset = offset[0], offset[1]
                rank_possible = rank + rank_offset
                file_possible = file + file_offset
                try:
                    _ = board[rank_possible, file_possible]
                    assert not rank_possible < 0
                    assert not file_possible < 0

                except IndexError:
                    break
                except AssertionError:
                    break
                else:
                    possible_moves.append([rank_possible, file_possible])
                    rank = rank_possible
                    file = file_possible
                finally:
                    if self.name == 'KING' or self.name == 'PAWN' or self.name == 'KNIGHT':
                        break
        return possible_moves

    def move(self, board_state, where):
        '''
        <param> board_state -> numpy.ndarray : Defines the state of the board at the time a move is made

        <param> where -> string : The location of the target square eg. 'a1' or 'h8'

        <return> None

        Makes the move on the board.
        '''
        check_move = self.is_move_legal(board_state, where)
        if check_move[0]:
            target = [self.ranks[where[1]], self.files[where[0]]]
            current_position = self.get_current_position()
            board_state[current_position[0], current_position[1]] = Empty(where)
            board_state[target[0], target[1]] = self
            self.set_current_position(where)
        else:
            raise IllegalMoveError(check_move[1])

    def is_move_legal(self, board_state, where):
        '''
        <param> board_state -> numpy.ndarray : Defines the state of the board at the time a move is made

        <param> where -> string : The location of the target square eg. 'a1' or 'h8'

        Checks the following (in order):
        1. If the target square is already occupied by any friendly piece
        2. If the path to the target square is blocked by a friendly or unfriendly piece
        3. If the target square is covered by an unfriendly piece, and allow a capture if the piece on the target square isn't the king
        (To be added at the end, the special castling rules!)
        '''
        target = [self.ranks[where[1]], self.files[where[0]]]
        try:
            assert not isinstance(self, Empty), "There isn't any piece to move"
            assert target in self.possible_moves(board_state), f'{self.name} cannot move to {where}'
            assert not self.is_target_friendly(board_state, where), f"{where} is occupied by a friendly piece: {board_state[target[0],target[1]]}"
            assert not self.is_path_blocked(board_state, where), f"{self.name}'s path to {where} is blocked"
        except AssertionError as e:
            return (False, e)
        else:
            return (True,)

    def is_path_blocked(self, board_state, where):
        '''
        <param> board_state -> numpy.ndarray : Defines the state of the board at the time a move is made

        <param> where -> string : The location of the target square eg. 'a1' or 'h8'

        Determines if the direct path to the target square is blocked by friendly or unfriendly pieces
        '''
        if not self.is_sliding_piece:
            return False
        current_position = self.get_current_position()
        target = np.array([self.ranks[where[1]], self.files[where[0]]])
        diff = current_position - target
        try:
            m = int(diff[0]) / int(diff[1])
        except ZeroDivisionError:
            offsets = np.array([[-1, 0], [1, 0]])
        else:
            if m == 1:
                offsets = np.array([[1, 1], [-1, -1]])
            elif m == -1:
                offsets = np.array([[-1, 1], [1, -1]])
            elif m == 0:
                offsets = np.array([[0, -1], [0, 1]])

        while (current_position != target).any():
            new_position1 = current_position + offsets[0]
            new_position2 = current_position + offsets[1]
            dist1 = np.sqrt(np.sum(np.square(new_position1 - target)))
            dist2 = np.sqrt(np.sum(np.square(new_position2 - target)))
            new_position = new_position1 if dist1 < dist2 else new_position2
            if (target == new_position).all():
                return False
            if not isinstance(board_state[new_position[0], new_position[1]], Empty):
                return True
            current_position = new_position

    def is_target_friendly(self, board_state, where):
        '''
        <param> board_state -> numpy.ndarray : Defines the state of the board at the time a move is made

        <param> where -> string : The location of the target square eg. 'a1' or 'h8'

        Determines if the target square is already occupied by a friendly piece
        '''
        rank, file = self.ranks[where[1]], self.files[where[0]]
        if isinstance(board_state[rank, file], Empty):
            return False
        else:
            return board_state[rank, file].color == self.color

    def is_target_unfriendly(self, board_state, where):
        '''
        <param> board_state -> numpy.ndarray : Defines the state of the board at the time a move is made

        <param> where -> string : The location of the target square eg. 'a1' or 'h8'

        Determines if the target square is already occupied by an unfriendly piece
        '''
        rank, file = self.ranks[where[1]], self.files[where[0]]
        if self.name != 'KING':
            if isinstance(board_state[rank, file], Empty):
                return True
            else:
                return not board_state[rank, file].color == self.color

class Rook(Piece):
    '''
    The white rooks start on squares a1 and h1, while the black rooks start on a8 and h8.
    The rook moves horizontally or vertically, through any number of unoccupied squares.
    The rook cannot jump over pieces.
    The rook also participates, with the king, in a special move called castling.

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Rook can move in
    '''

    name = 'ROOK'

    offsets = np.array([[-1, 0],
                        [1, 0],
                        [0, -1],
                        [0, 1]])
    image = 'rook.png'

    def __init__(self, color, position):
        super().__init__(color, position)
        self.image = f'{color}-{self.image}'

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[0]}'

class Queen(Piece):
    '''
    The white queen start on squares d1, while the black queen start on d8.
    The queen can be moved any number of unoccupied squares in a straight line vertically,
    horizontally, or diagonally, thus combining the moves of the rook and bishop

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Queen can move in
    '''

    name = 'QUEEN'

    offsets = np.array([[-1, 0],
                        [1, 0],
                        [0, -1],
                        [0, 1],
                        [1, 1],
                        [1, -1],
                        [-1, 1],
                        [-1, -1]])
    image = 'queen.png'

    def __init__(self, color):
        position = 'd1' if color == 'white' else 'd8'
        super().__init__(color, position)
        self.image = f'{color}-{self.image}'

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[0]}'

class Pawn(Piece):
    '''
    The white pawns start on the second rank
    Unlike the other pieces, pawns cannot move backwards.
    Normally a pawn moves by advancing a single square, but the first time a pawn moves, it has the option of advancing two squares.
    Pawns may not use the initial two-square advance to jump over an occupied square, or to capture.
    Any piece immediately in front of a pawn, friend or foe, blocks its advance.

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Rook can move in
    '''

    name = 'PAWN'

    is_sliding_piece = False

    image = 'pawn.png'

    def __init__(self, color, current_position):
        super().__init__(color, current_position)
        self.has_moved = False
        self.image = f'{color}-{self.image}'
        self.set_offsets()

    def set_offsets(self):
        if self.color == 'white':
            if self.get_current_position_as_string()[1] != '2':
                self.offsets = np.array([[-1, 0],
                                         [-1, 1],
                                         [-1, -1]])
            else:
                self.offsets = np.array([[-1, 0],
                                         [-1, 1],
                                         [-1, -1],
                                         [-2, 0]])
        else:
            if self.get_current_position_as_string()[1] != '7':
                self.offsets = np.array([[1, 0],
                                         [1, 1],
                                         [1, -1]])
            else:
                self.offsets = np.array([[1, 0],
                                         [1, 1],
                                         [1, -1],
                                         [2, 0]])

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[0]}'

class Knight(Piece):
    '''
    The white knight start on squares b1 and g1, while the black knight start on b8 and g8.
    Compared to other chess pieces, the knight's movement is unique: it may move two squares vertically and one square horizontally,
    or two squares horizontally and one square vertically (with both forming the shape of an L).
    This way, a knight can have a maximum of 8 moves. While moving, the knight can jump over pieces to reach its destination.

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Rook can move in
    '''

    is_sliding_piece = False
    name = 'KNIGHT'

    image = 'knight.png'

    def __init__(self, color, position):
        super().__init__(color, position)
        self.image = f'{color}-{self.image}'

    offsets = np.array([[1, 2],
                        [2, 1],
                        [-1, 2],
                        [-2, -1],
                        [-1, -2],
                        [-2, 1],
                        [1, -2],
                        [2, -1]])

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[1]}'

class King(Piece):
    '''
    The white king start on squares e1, while the black king start on e8.
    A king can move one square in any direction (horizontally, vertically, or diagonally),
    unless the square is already occupied by a friendly piece, or the move would place the king in check.
    As a result, opposing kings may never occupy adjacent squares (see opposition) to give check,
    as that would put the moving king in check as well.
    However, the king can give discovered check by unblocking a bishop, rook, or queen.
    The king is also involved in the special move of castling.

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Rook can move in
    '''

    name = 'KING'

    is_sliding_piece = False

    image = 'king.png'

    offsets = np.array([[-1, 0],
                        [1, 0],
                        [0, -1],
                        [0, 1],
                        [1, 1],
                        [1, -1],
                        [-1, 1],
                        [-1, -1]])

    def __init__(self, color):
        position = 'e1' if color == 'white' else 'e8'
        super().__init__(color, position)
        self.image = f'{color}-{self.image}'

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[0]}'

class Bishop(Piece):
    '''
    The white bishops start on squares c1 and f1, while the black bishops start on c8 and f8.
    The bishop has no restrictions in distance for each move, but is limited to diagonal movement.
    Bishops, like all other pieces except the knight, cannot jump over other pieces.

    <var> name -> string : Name of the piece

    <var> offsets -> np.ndarray : Offsets indicate which direction the Rook can move in
    '''

    name = 'BISHOP'

    offsets = np.array([[1, 1],
                        [-1, 1],
                        [1, -1],
                        [-1, -1]])

    image = 'bishop.png'

    def __init__(self, color, position):
        super().__init__(color, position)
        self.image = f'{color}-{self.image}'

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'{self.color[0]}{self.name[0]}'

class Empty(Piece):
    '''
    We store empty positions using objects of this class
    '''

    name = 'EMPTY'

    offsets = None

    image = 'empty.png'

    def __init__(self, position):
        super().__init__('white', position)
        self.color = None

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return '  '

class IllegalMoveError(Exception):
    def __init__(self, *args):
        self.message = f'THIS MOVE IS ILLEGAL!!!\n{args[0]}'

    def __str__(self):
        return self.message

max_stars_repo_path: Python Programs/Chess/Piece.py

max_issues_repo_path: Python Programs/Chess/Piece.py

max_forks_repo_path: Python Programs/Chess/Piece.py


text:
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.6
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% [markdown]
# ---
# # [python-control](https://python-control.readthedocs.io/en/0.8.3/)
# ---
#
# + [GitHub](https://github.com/python-control/python-control)
# + [Murray](https://www.cds.caltech.edu/~murray/wiki/Control_Systems_Library_for_Python)
# + [Example](https://www.cds.caltech.edu/~murray/wiki/index.php?title=Python-control/Example:_Vertical_takeoff_and_landing_aircraft)
# + [pdf](http://robertobucher.dti.supsi.ch/wp-content/uploads/2017/03/BookPythonForControl.pdf)

# %% [markdown]
# # Standard import
#
# There are two different ways to use the package. For the default interface described in Function reference, simply import the control package as follows:

# %%
import control as co

# %% [markdown]
# ---
# #  Dynamical systems
#
# ---
#
# Recall the system studied in the *machine learning* problem.
# + It was *static*, in the sense that the outputs depended only on the inputs and not on an internal state.
# + It was *black box*, meaning that we represented the output as a probability distribution conditioned on the input.
#
# In the study of dynamical systems we make a few changes.
# + The system is **dynamic**. The outputs at time $t$ depend, not only on the inputs at time $t$, but also on the **state of the system** at time $t$.
# + We have a deterministic model of the system. We are not building the model based on data (as in ML), but rather based on physical principles.
#
# There is also a change of notation.
# + $u(t)$: system input (in ML it was $x$).
# + $x(t)$: system state.
#
# The figure below depicts a **linear** dynamical system.
#
#
#
# ---
# #  Systems representations
#
# ---
#
# Linear dynamical systems can be **continuous time** or **discrete time**. The theory for these two is similar (with caveats) and is taught in **ME 132**.
#
# **Note**: We will focus this tutorial on continuous-time systems.
#
# Also, there are two levels of *granularity* of models: **state-space** models and **input-output** models.
# + input-output models predict the **output** of the system given the input. Their mathematical form is the **transfer function**.
# + state-space models predict the **output** and the **state** of the system, given the input. Their mathematical form is the **state and output equations**.
#
#
#
# ---
#

# %%
import matplotlib.pyplot as plt
import numpy as np

def plotresponse(t,y):
    plt.figure(figsize=(10,4))
    plt.plot(t,y,linewidth=3)
    plt.xlabel('t',fontsize=14)
    plt.ylabel('y',fontsize=14)
    plt.grid()

def plotstepresponse(sys):

    t, y = co.step_response(sys)
    r = np.ones(t.shape)

    plt.figure(figsize=(10,4))
    plt.plot(t,y,t,r,linewidth=3)
    plt.xlabel('t',fontsize=14)
    plt.legend(['y','r'],fontsize=14)
    plt.grid()

# %% [markdown]
# ---
# #  Systems creation
#
# ---
#
#
#
# \begin{align}
# \frac{d}{dt}\begin{bmatrix}x\\ \dot{x}\end{bmatrix} &=
# \begin{bmatrix} 0&1\\-\frac{k}{m}&-\frac{b}{m} \end{bmatrix}\begin{bmatrix}x\\ \dot{x}\end{bmatrix} +
# \begin{bmatrix} 0 \\ 1/m \end{bmatrix} F & \text{state equation} \\
# x &= \begin{bmatrix} 1 & 0 \end{bmatrix} \begin{bmatrix} x \\ \dot{x} \end{bmatrix}  & \text{output equation}
# \end{align}
#
#
# ##  State-space systems `ss(A,B,C,D)`

# %%
m = 1
k = 1
b = 2
A = [[0,1],[-k/m,-b/m]]
B = [[0],[1/m]]
C = [1,0]
D = [0]
sys = co.ss(A,B,C,D)

print(sys)

# %% [markdown]
# ## Input-output systems `tf(num, den)`
#
# \begin{equation}
# G(s) = \frac{X(s)}{F(s)} =   \frac{1}{m s^2 + b s + k }
# \end{equation}

# %%
G = co.tf(1,[m,b,k])

print(G)

# %% [markdown]
# **Note**: It is a common technique to define a *unit transfer function* `tf('s')`, and use it to compose other transfer functions, as follows:

# %%
s = co.tf('s')
G  = (s + 1)/(s**2 + 5*s + 6)
print(G)

# %% [markdown]
# ---
# #  Systems connections
#
# ---

# %% [markdown]
# ## [`series(sys1,sys2)`](https://python-control.readthedocs.io/en/v0.6d/bdalg_strings.html#control.series)
#
#
#

# %%
Gstiff = co.tf(1,[1,2,100])
Gsoft  = co.tf(1,[1,2,3])

Gseries = co.series(Gstiff,Gsoft)

t, y = co.impulse_response(Gseries)
plotresponse(t,y)

# %% [markdown]
# ## [`parallel(sys1,sys2)`](https://python-control.readthedocs.io/en/v0.6d/bdalg_strings.html#control.parallel)
#
#
#

# %%
Gparallel = co.parallel(Gstiff,Gsoft)

t, y = co.impulse_response(Gparallel)
plotresponse(t,y)

# %% [markdown]
#
# ---
# #  Time domain simulation
# ---
#
#

# %% [markdown]
# ## [`initial_response(sys,T=None,X0=0.0)`](https://python-control.readthedocs.io/en/v0.6d/timeresp.html#control.initial_response) Given initial condition, zero input.
# + Returns T, y

# %%
t, y = co.initial_response(Gseries,T=None,X0=1.0)
plotresponse(t,y)

# %% [markdown]
# ## `impulse_response(sys,T=None,...)`: Impulse input, zero initial condition.

# %%
t, y = co.impulse_response(Gseries)
plotresponse(t,y)

# %% [markdown]
# ## [`step_response(sys,T=None,...)`](https://python-control.readthedocs.io/en/v0.6d/timeresp.html#control.step_response): Step input, zero initial condition.

# %%
t, y = co.step_response(Gseries)
plotresponse(t,y)

# %% [markdown]
# ## [ `forced_response(sys,T=None,U=0.0,X0=0.0,...)`](https://python-control.readthedocs.io/en/v0.6d/timeresp.html#control.forced_response): Given input, given initial condition.

# %%
import numpy as np
t = np.linspace(0,10)
u = np.sin(t)
t, y = co.forced_response(sys,t,u)

plt.figure(figsize=(10,4))
plt.plot(t,y,t,u,linewidth=3)
plt.xlabel('t')
plt.grid()
plt.legend(['y','u'])

# %% [markdown]
# ---
# #  System analysis
# ---
#
#
# # Stability
#
# An equilibrium point of a dynamical system is **asymptotically stable** if small perturbations away from that equilibrium point die out in time. For *linear systems*, an equivalent definition is that all unforced trajectories eventually die out (ie converge to the origin).
#
# Linear systems are:
# + **asymptotically stable**: if **all** system poles have negative real part.
# + **unstable**: if **any** system poles all have positive real part.
#
# **Note**: The remaining cases (ie some have zero real part, none have positive real part) is more complicated.
#
# ## [`pole(sys)`](https://python-control.readthedocs.io/en/v0.6d/analysis.html#control.pole)

# %%
import math

is_stable = all( co.pole(Gseries).real<0 )

print(is_stable)

# %%
Gundamped = co.tf(1,[1,0,100])

print(co.pole(Gundamped))

all( co.pole(Gundamped).real<0 )

# %% [markdown]
#
# ---
# #  Controller design
# ---

# %% [markdown]
# # Output regulation
#
# ### [`feedback(fwd,fdbck)`](https://python-control.readthedocs.io/en/v0.6d/bdalg_strings.html#control.feedback)
#
#
#
#
#
# This block diagram is achieved with,
#
# ```python
# cl_sys = co.feedback( co.series(cntrl,plant) , 1 )
# ```
#
# ## Proportional feedback

# %%
plant = co.tf(1,[1,1,-2])

cntrl = 4

RtoY = co.feedback( co.series(cntrl,plant) ,1)

plotstepresponse(RtoY)

# %% [markdown]
# ## Proportional-integral feedback

# %%
plant = co.tf(1,[1,1,-2])
cntrl = 4 + 0.5* (1/co.tf('s'))

RtoY = co.feedback( co.series(cntrl,plant) ,1)

plotstepresponse(RtoY)

# %% [markdown]
# ---
# #  Control design example:
 Balancing the inverted pendulum with pole placement
# ---
#
#

# %% [markdown]
# ## Find the equations of motion ...
# \begin{align}
# \ddot{p} &= \frac{1}{Z}\left( -m\ell\cos\theta\:\xi_1 + \ell\xi_2 \right) \\
# \label{eq:tdd}
# \ddot\theta &= \frac{1}{Z}\left( -(M+m)\xi_1 + \cos\theta\:\xi_2 \right)
# \end{align}
# where
# \begin{align}
# Z &= M\ell + m\ell\sin^2\theta \\
# \label{xi1}
# \xi_1 &= \frac{\sigma}{m\ell}\dot{\theta} - g\sin\theta \\
# \label{xi2}
# \xi_2 &= F - c\dot{p} - m\ell\sin\theta\:\dot{\theta}^2 + \frac{\sigma}{\ell}\cos\theta\:\dot{\theta}
# \end{align}
#
# ## Find equilibria ...
#
# These are values of $p$ and $\theta$ such that $\dot{p}=\ddot{p}=\dot\theta=\ddot\theta=F=0$.
#
# Plugging these into the EOM we find that the sole condition on $p$ and $\theta$ for equilibrium is $\sin\theta=0$
#
# The eqilibrium set is all points with $\theta=n\pi$, where $n$ is an integer.
# \begin{equation}
# \{ \; (p,\theta) \;|\; p\in\mathbb{R} \; , \; \theta=n\pi, n\in \mathbb{Z} \;\}
# \end{equation}
#
#
#
# ## Linearize the EOM about the upward-pointing equilibrium state ...
#
# With some work we find the linearized dynamics for the unstable equilibria.
#
# \begin{equation}
# \frac{d}{dt}
# \begin{bmatrix} \dot{p} \\ p \\ \dot{\theta} \\ \theta \end{bmatrix} =
# \begin{bmatrix}
# -\frac{c}{M} & 0 & 0 & \frac{m}{M}g \\
# 1 & 0 & 0 & 0 \\
# \mp\frac{c}{M\ell} & 0 & -\frac{\sigma}{m\ell^2} & \pm\frac{M+m}{m\ell}g \\
# 0 & 0 & 1 & 0
# \end{bmatrix}
# \begin{bmatrix} \dot{p} \\ p \\ \dot{\theta} \\ \theta \end{bmatrix}
# +
# \begin{bmatrix} \frac{1}{M} \\ 0 \\ \pm\frac{1}{M\ell} \\ 0 \end{bmatrix}
#  F
# \end{equation}
#
#
# Introduce some notation to simplify the formaulas.
# \begin{equation}
# \bar{c} = \frac{c}{M}  \quad , \quad
# \bar{\sigma}=\frac{\sigma}{m\ell^2}   \quad , \quad
# \mu = \frac{M}{m} \quad , \quad
# \bar{F} = \frac{F}{M}
# \end{equation}
# Then,
# \begin{equation}
# \frac{d}{dt}
# \begin{bmatrix} \dot{p} \\ p \\ \dot{\theta} \\ \theta \end{bmatrix} =
# \begin{bmatrix}
# -\bar{c} & 0 & 0 & g/\mu \\
# 1 & 0 & 0 & 0 \\
# \mp\bar{c}/\ell & 0 & -\bar{\sigma} & \pm(1+\mu)g/\ell \\
# 0 & 0 & 1 & 0
# \end{bmatrix}
# \begin{bmatrix} \dot{p} \\ p \\ \dot{\theta} \\ \theta \end{bmatrix}
# +
# \begin{bmatrix} 1 \\ 0 \\ \pm1/\ell \\ 0 \end{bmatrix}
#  \bar{F}
# \end{equation}
#

# %%
M = 0.5;         # mass of the cart
m = 0.2;         # mass of the pendulum
ell = 1;         # length of the pendulum
g = 9.8;         # gravity
c = 1.0;         # translational friction
sigma = 0.01;    # rotational friction

cbar = c/M
sigmabar = sigma/m/ell/ell
mu = M/m

A = [[-cbar,0,0,g/mu],[1,0,0,0],[-cbar/ell,0,-sigmabar,(1+mu)*g/ell],[0,0,1,0]]
B = [[1],[0],[1/ell],[0]]

# %% [markdown]
# ## Check controllability

# %%
np.linalg.matrix_rank(co.ctrb(A,B))==4

# %% [markdown]
# ## Select the closed-loop poles and generate the controller gains

# %%
closed_loop_poles = np.array([-3,-2,-1,-2.5])

# controller
K = co.place(A,B,closed_loop_poles);

print(K)

# %% [markdown]
# ## Create the closed-loop system

# %%
sys_cl = co.feedback( co.ss(A,B,np.eye(4),0) , K )

print(sys_cl)

# %% [markdown]
# ## Test run the system and plot

# %%
t, y = co.initial_response(sys_cl,T=None,X0=[1,1,1,1])

plt.figure(figsize=(10,4))
plt.plot(t,y.T,linewidth=3)
plt.grid()
plt.xlabel('time',fontsize=16)
plt.legend(['$p$','$\dot{p}$',r'$\theta$',r'$\dot\theta$'],fontsize=12);

# %% [markdown]
# ### Comment
#
# This controller will clearly fail on the real nonlinear system (notice that deviations in $\theta$ are reaching over two radians). The real test is whether you can design a controller with this linearized method that will perform well on the nonlinear system as well.

# %%

ext: py

reason: import numpy


text:
'''
High-level functions and classes for solving a wide variety of economic models.
The "core" of HARK is a framework for "microeconomic" and "macroeconomic"
models.  A micro model concerns the dynamic optimization problem for some type
of agents, where agents take the inputs to their problem as exogenous.  A macro
model adds an additional layer, endogenizing some of the inputs to the micro
problem by finding a general equilibrium dynamic rule.
'''
from __future__ import print_function, division
from __future__ import absolute_import

from builtins import str
from builtins import range
from builtins import object
import sys
import os
from distutils.dir_util import copy_tree
from .utilities import getArgNames, NullFunc
from copy import copy, deepcopy
import numpy as np
from time import clock
from .parallel import multiThreadCommands, multiThreadCommandsFake

def distanceMetric(thing_A, thing_B):
    '''
    A "universal distance" metric that can be used as a default in many settings.

    Parameters
    ----------
    thing_A : object
        A generic object.
    thing_B : object
        Another generic object.

    Returns:
    ------------
    distance : float
        The "distance" between thing_A and thing_B.
    '''
    # Get the types of the two inputs
    typeA = type(thing_A)
    typeB = type(thing_B)

    if typeA is list and typeB is list:
        lenA = len(thing_A)  # If both inputs are lists, then the distance between
        lenB = len(thing_B)  # them is the maximum distance between corresponding
        if lenA == lenB:  # elements in the lists.  If they differ in length,
            distance_temp = []  # the distance is the difference in lengths.
            for n in range(lenA):
                distance_temp.append(distanceMetric(thing_A[n], thing_B[n]))
            distance = max(distance_temp)
        else:
            distance = float(abs(lenA - lenB))
    # If both inputs are numbers, return their difference
    elif isinstance(thing_A, (int, float)) and isinstance(thing_B, (int, float)):
        distance = float(abs(thing_A - thing_B))
    # If both inputs are array-like, return the maximum absolute difference b/w
    # corresponding elements (if same shape); return largest difference in dimensions
    # if shapes do not align.
    elif hasattr(thing_A, 'shape') and hasattr(thing_B, 'shape'):
        if thing_A.shape == thing_B.shape:
            distance = np.max(abs(thing_A - thing_B))
        else:
            # Flatten arrays so they have the same dimensions
            distance = np.max(abs(thing_A.flatten().shape[0] - thing_B.flatten().shape[0]))
    # If none of the above cases, but the objects are of the same class, call
    # the distance method of one on the other
    elif thing_A.__class__.__name__ == thing_B.__class__.__name__:
        if thing_A.__class__.__name__ == 'function':
            distance = 0.0
        else:
            distance = thing_A.distance(thing_B)
    else:  # Failsafe: the inputs are very far apart
        distance = 1000.0
    return distance

class HARKobject(object):
    '''
    A superclass for object classes in HARK.  Comes with two useful methods:
    a generic/universal distance method and an attribute assignment method.
    '''
    def distance(self, other):
        '''
        A generic distance method, which requires the existence of an attribute
        called distance_criteria, giving a list of strings naming the attributes
        to be considered by the distance metric.

        Parameters
        ----------
        other : object
            Another object to compare this instance to.

        Returns
        -------
        (unnamed) : float
            The distance between this object and another, using the "universal
            distance" metric.
        '''
        distance_list = [0.0]
        for attr_name in self.distance_criteria:
            try:
                obj_A = getattr(self, attr_name)
                obj_B = getattr(other, attr_name)
                distance_list.append(distanceMetric(obj_A, obj_B))
            except AttributeError:
                distance_list.append(1000.0)  # if either object lacks attribute, they are not the same
        return max(distance_list)

    def assignParameters(self, **kwds):
        '''
        Assign an arbitrary number of attributes to this agent.

        Parameters
        ----------
        **kwds : keyword arguments
            Any number of keyword arguments of the form key=value.  Each value
            will be assigned to the attribute named in self.

        Returns
        -------
        none
        '''
        for key in kwds:
            setattr(self, key, kwds[key])

    def __call__(self, **kwds):
        '''
        Assign an arbitrary number of attributes to this agent, as a convenience.
        See assignParameters.
        '''
        self.assignParameters(**kwds)

    def getAvg(self, varname, **kwds):
        '''
        Calculates the average of an attribute of this instance.  Returns NaN if no such attribute.

        Parameters
        ----------
        varname : string
            The name of the attribute whose average is to be calculated.  This attribute must be an
            np.array or other class compatible with np.mean.

        Returns
        -------
        avg : float or np.array
            The average of this attribute.  Might be an array if the axis keyword is passed.
        '''
        if hasattr(self, varname):
            return np.mean(getattr(self, varname), **kwds)
        else:
            return np.nan

class Solution(HARKobject):
    '''
    A superclass for representing the "solution" to a single period problem in a
    dynamic microeconomic model.

    NOTE: This can be deprecated now that HARKobject exists, but this requires
    replacing each instance of Solution with HARKobject in the other modules.
    '''

class AgentType(HARKobject):
    '''
    A superclass for economic agents in the HARK framework.  Each model should
    specify its own subclass of AgentType, inheriting its methods and overwriting
    as necessary.  Critically, every subclass of AgentType should define class-
    specific static values of the attributes time_vary and time_inv as lists of
    strings.  Each element of time_vary is the name of a field in AgentSubType
    that varies over time in the model.  Each element of time_inv is the name of
    a field in AgentSubType that is constant over time in the model. The string
    'solveOnePeriod' should appear in exactly one of these lists, depending on
    whether the same solution method is used in all periods of the model.
    '''
    def __init__(self, solution_terminal=None, cycles=1, time_flow=True, pseudo_terminal=True,
                 tolerance=0.000001, seed=0, **kwds):
        '''
        Initialize an instance of AgentType by setting attributes.

        Parameters
        ----------
        solution_terminal : Solution
            A representation of the solution to the terminal period problem of
            this AgentType instance, or an initial guess of the solution if this
            is an infinite horizon problem.
        cycles : int
            The number of times the sequence of periods is experienced by this
            AgentType in their "lifetime".  cycles=1 corresponds to a lifecycle
            model, with a certain sequence of one period problems experienced
            once before terminating.  cycles=0 corresponds to an infinite horizon
            model, with a sequence of one period problems repeating indefinitely.
        time_flow : boolean
            Whether time is currently "flowing" forward(True) or backward(False) for this
            instance.  Used to flip between solving (using backward iteration)
            and simulating (etc).
        pseudo_terminal : boolean
            Indicates whether solution_terminal isn't actually part of the
            solution to the problem (as a known solution to the terminal period
            problem), but instead represents a "scrap value"-style termination.
            When True, solution_terminal is not included in the solution; when
            false, solution_terminal is the last element of the solution.
        tolerance : float
            Maximum acceptable "distance" between successive solutions to the
            one period problem in an infinite horizon (cycles=0) model in order
            for the solution to be considered as having "converged".  Inoperative
            when cycles>0.
        seed : int
            A seed for this instance's random number generator.

        Returns
        -------
        None
        '''
        if solution_terminal is None:
            solution_terminal = NullFunc()
        self.solution_terminal  = solution_terminal # NOQA
        self.cycles             = cycles # NOQA
        self.time_flow          = time_flow # NOQA
        self.pseudo_terminal    = pseudo_terminal # NOQA
        self.solveOnePeriod     = NullFunc() # NOQA
        self.tolerance          = tolerance # NOQA
        self.seed               = seed # NOQA
        self.track_vars         = [] # NOQA
        self.poststate_vars     = [] # NOQA
        self.read_shocks        = False # NOQA
        self.assignParameters(**kwds) # NOQA
        self.resetRNG() # NOQA

    def timeReport(self):
        '''
        Report to the user the direction that time is currently "flowing" for
        this instance.  Only exists as a reminder of how time_flow works.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        if self.time_flow:
            print('Time varying objects are listed in ordinary chronological order.')
        else:
            print('Time varying objects are listed in reverse chronological order.')

    def timeFlip(self):
        '''
        Reverse the flow of time for this instance.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for name in self.time_vary:
            self.__dict__[name].reverse()
        self.time_flow = not self.time_flow

    def timeFwd(self):
        '''
        Make time flow forward for this instance.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        if not self.time_flow:
            self.timeFlip()

    def timeRev(self):
        '''
        Make time flow backward for this instance.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        if self.time_flow:
            self.timeFlip()

    def addToTimeVary(self, *params):
        '''
        Adds any number of parameters to time_vary for this instance.

        Parameters
        ----------
        params : string
            Any number of strings naming attributes to be added to time_vary

        Returns
        -------
        None
        '''
        for param in params:
            if param not in self.time_vary:
                self.time_vary.append(param)

    def addToTimeInv(self, *params):
        '''
        Adds any number of parameters to time_inv for this instance.

        Parameters
        ----------
        params : string
            Any number of strings naming attributes to be added to time_inv

        Returns
        -------
        None
        '''
        for param in params:
            if param not in self.time_inv:
                self.time_inv.append(param)

    def delFromTimeVary(self, *params):
        '''
        Removes any number of parameters from time_vary for this instance.

        Parameters
        ----------
        params : string
            Any number of strings naming attributes to be removed from time_vary

        Returns
        -------
        None
        '''
        for param in params:
            if param in self.time_vary:
                self.time_vary.remove(param)

    def delFromTimeInv(self, *params):
        '''
        Removes any number of parameters from time_inv for this instance.

        Parameters
        ----------
        params : string
            Any number of strings naming attributes to be removed from time_inv

        Returns
        -------
        None
        '''
        for param in params:
            if param in self.time_inv:
                self.time_inv.remove(param)

    def solve(self, verbose=False):
        '''
        Solve the model for this instance of an agent type by backward induction.
        Loops through the sequence of one period problems, passing the solution
        from period t+1 to the problem for period t.

        Parameters
        ----------
        verbose : boolean
            If True, solution progress is printed to screen.

        Returns
        -------
        none
        '''

        # Ignore floating point "errors". Numpy calls it "errors", but really it's excep-
        # tions with well-defined answers such as 1.0/0.0 that is np.inf, -1.0/0.0 that is
        # -np.inf, np.inf/np.inf is np.nan and so on.
        with np.errstate(divide='ignore', over='ignore', under='ignore', invalid='ignore'):
            self.preSolve()  # Do pre-solution stuff
            self.solution = solveAgent(self, verbose)  # Solve the model by backward induction
            if self.time_flow:  # Put the solution in chronological order if this instance's time flow runs that way
                self.solution.reverse()
            self.addToTimeVary('solution')  # Add solution to the list of time-varying attributes
            self.postSolve()  # Do post-solution stuff

    def resetRNG(self):
        '''
        Reset the random number generator for this type.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        self.RNG = np.random.RandomState(self.seed)

    def checkElementsOfTimeVaryAreLists(self):
        """
        A method to check that elements of time_vary are lists.
        """
        for param in self.time_vary:
            assert type(getattr(self, param)) == list, param + ' is not a list, but should be' + \
                                                   ' because it is in time_vary'

    def checkRestrictions(self):
        """
        A method to check that various restrictions are met for the model class.
        """
        return

    def preSolve(self):
        '''
        A method that is run immediately before the model is solved, to check inputs or to prepare
        the terminal solution, perhaps.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        self.checkRestrictions()
        self.checkElementsOfTimeVaryAreLists()
        return None

    def postSolve(self):
        '''
        A method that is run immediately after the model is solved, to finalize
        the solution in some way.  Does nothing here.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        return None

    def initializeSim(self):
        '''
        Prepares this AgentType for a new simulation.  Resets the internal random number generator,
        makes initial states for all agents (using simBirth), clears histories of tracked variables.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        self.resetRNG()
        self.t_sim = 0
        all_agents = np.ones(self.AgentCount, dtype=bool)
        blank_array = np.zeros(self.AgentCount)
        for var_name in self.poststate_vars:
            setattr(self, var_name, copy(blank_array))
            # exec('self.' + var_name + ' = copy(blank_array)')
        self.t_age = np.zeros(self.AgentCount, dtype=int)   # Number of periods since agent entry
        self.t_cycle = np.zeros(self.AgentCount, dtype=int)  # Which cycle period each agent is on
        self.simBirth(all_agents)
        self.clearHistory()
        return None

    def simOnePeriod(self):
        '''
        Simulates one period for this type.  Calls the methods getMortality(), getShocks() or
        readShocks, getStates(), getControls(), and getPostStates().  These should be defined for
        AgentType subclasses, except getMortality (define its components simDeath and simBirth
        instead) and readShocks.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        if not hasattr(self, 'solution'):
            raise Exception('Model instance does not have a solution stored. To simulate, it is necessary'
                            ' to run the `solve()` method of the class first.')

        self.getMortality()  # Replace some agents with "newborns"
        if self.read_shocks:  # If shock histories have been pre-specified, use those
            self.readShocks()
        else:                # Otherwise, draw shocks as usual according to subclass-specific method
            self.getShocks()
        self.getStates()  # Determine each agent's state at decision time
        self.getControls()   # Determine each agent's choice or control variables based on states
        self.getPostStates()  # Determine each agent's post-decision / end-of-period states using states and controls

        # Advance time for all agents
        self.t_age = self.t_age + 1  # Age all consumers by one period
        self.t_cycle = self.t_cycle + 1  # Age all consumers within their cycle
        self.t_cycle[self.t_cycle == self.T_cycle] = 0  # Resetting to zero for those who have reached the end

    def makeShockHistory(self):
        '''
        Makes a pre-specified history of shocks for the simulation.  Shock variables should be named
        in self.shock_vars, a list of strings that is subclass-specific.  This method runs a subset
        of the standard simulation loop by simulating only mortality and shocks; each variable named
        in shock_vars is stored in a T_sim x AgentCount array in an attribute of self named X_hist.
        Automatically sets self.read_shocks to True so that these pre-specified shocks are used for
        all subsequent calls to simulate().

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        # Make sure time is flowing forward and re-initialize the simulation
        orig_time = self.time_flow
        self.timeFwd()
        self.initializeSim()

        # Make blank history arrays for each shock variable
        for var_name in self.shock_vars:
            setattr(self, var_name+'_hist', np.zeros((self.T_sim, self.AgentCount)) + np.nan)

        # Make and store the history of shocks for each period
        for t in range(self.T_sim):
            self.getMortality()
            self.getShocks()
            for var_name in self.shock_vars:
                exec('self.' + var_name + '_hist[self.t_sim,:] = self.' + var_name)
            self.t_sim += 1
            self.t_age = self.t_age + 1  # Age all consumers by one period
            self.t_cycle = self.t_cycle + 1  # Age all consumers within their cycle
            self.t_cycle[self.t_cycle == self.T_cycle] = 0  # Resetting to zero for those who have reached the end

        # Restore the flow of time and flag that shocks can be read rather than simulated
        self.read_shocks = True
        if not orig_time:
            self.timeRev()

    def getMortality(self):
        '''
        Simulates mortality or agent turnover according to some model-specific rules named simDeath
        and simBirth (methods of an AgentType subclass).  simDeath takes no arguments and returns
        a Boolean array of size AgentCount, indicating which agents of this type have "died" and
        must be replaced.  simBirth takes such a Boolean array as an argument and generates initial
        post-decision states for those agent indices.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        who_dies = self.simDeath()
        self.simBirth(who_dies)
        return None

    def simDeath(self):
        '''
        Determines which agents in the current population "die" or should be replaced.  Takes no
        inputs, returns a Boolean array of size self.AgentCount, which has True for agents who die
        and False for those that survive. Returns all False by default, must be overwritten by a
        subclass to have replacement events.

        Parameters
        ----------
        None

        Returns
        -------
        who_dies : np.array
            Boolean array of size self.AgentCount indicating which agents die and are replaced.
        '''
        who_dies = np.zeros(self.AgentCount, dtype=bool)
        return who_dies

    def simBirth(self, which_agents):
        '''
        Makes new agents for the simulation.  Takes a boolean array as an input, indicating which
        agent indices are to be "born".  Does nothing by default, must be overwritten by a subclass.

        Parameters
        ----------
        which_agents : np.array(Bool)
            Boolean array of size self.AgentCount indicating which agents should be "born".

        Returns
        -------
        None
        '''
        print('AgentType subclass must define method simBirth!')
        return None

    def getShocks(self):
        '''
        Gets values of shock variables for the current period.  Does nothing by default, but can
        be overwritten by subclasses of AgentType.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        return None

    def readShocks(self):
        '''
        Reads values of shock variables for the current period from history arrays.  For each var-
        iable X named in self.shock_vars, this attribute of self is set to self.X_hist[self.t_sim,:].

        This method is only ever called if self.read_shocks is True.  This can be achieved by using
        the method makeShockHistory() (or manually after storing a "handcrafted" shock history).

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        for var_name in self.shock_vars:
            setattr(self, var_name, getattr(self, var_name + '_hist')[self.t_sim, :])

    def getStates(self):
        '''
        Gets values of state variables for the current period, probably by using post-decision states
        from last period, current period shocks, and maybe market-level events.  Does nothing by
        default, but can be overwritten by subclasses of AgentType.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        return None

    def getControls(self):
        '''
        Gets values of control variables for the current period, probably by using current states.
        Does nothing by default, but can be overwritten by subclasses of AgentType.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        return None

    def getPostStates(self):
        '''
        Gets values of post-decision state variables for the current period, probably by current
        states and controls and maybe market-level events or shock variables.  Does nothing by
        default, but can be overwritten by subclasses of AgentType.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        return None

    def simulate(self, sim_periods=None):
        '''
        Simulates this agent type for a given number of periods (defaults to self.T_sim if no input).
        Records histories of attributes named in self.track_vars in attributes named varname_hist.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        # Ignore floating point "errors". Numpy calls it "errors", but really it's excep-
        # tions with well-defined answers such as 1.0/0.0 that is np.inf, -1.0/0.0 that is
        # -np.inf, np.inf/np.inf is np.nan and so on.
        with np.errstate(divide='ignore', over='ignore', under='ignore', invalid='ignore'):
            orig_time = self.time_flow
            self.timeFwd()
            if sim_periods is None:
                sim_periods = self.T_sim

            for t in range(sim_periods):
                self.simOnePeriod()
                for var_name in self.track_vars:
                    exec('self.' + var_name + '_hist[self.t_sim,:] = self.' + var_name)
                self.t_sim += 1

            if not orig_time:
                self.timeRev()

    def clearHistory(self):
        '''
        Clears the histories of the attributes named in self.track_vars.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        for var_name in self.track_vars:
            exec('self.' + var_name + '_hist = np.zeros((self.T_sim,self.AgentCount)) + np.nan')

def solveAgent(agent, verbose):
    '''
    Solve the dynamic model for one agent type.  This function iterates on "cycles"
    of an agent's model either a given number of times or until solution convergence
    if an infinite horizon model is used (with agent.cycles = 0).

    Parameters
    ----------
    agent : AgentType
        The microeconomic AgentType whose dynamic problem is to be solved.
    verbose : boolean
        If True, solution progress is printed to screen (when cycles != 1).

    Returns
    -------
    solution : [Solution]
        A list of solutions to the one period problems that the agent will
        encounter in his "lifetime".  Returns in reverse chronological order.
    '''
    # Record the flow of time when the Agent began the process, and make sure time is flowing backwards
    original_time_flow = agent.time_flow
    agent.timeRev()

    # Check to see whether this is an (in)finite horizon problem
    cycles_left      = agent.cycles # NOQA
    infinite_horizon = cycles_left == 0 # NOQA
    # Initialize the solution, which includes the terminal solution if it's not a pseudo-terminal period
    solution = []
    if not agent.pseudo_terminal:
        solution.append(deepcopy(agent.solution_terminal))

    # Initialize the process, then loop over cycles
    solution_last    = agent.solution_terminal # NOQA
    go               = True # NOQA
    completed_cycles = 0 # NOQA
    max_cycles       = 5000 # NOQA  - escape clause
    if verbose:
        t_last = clock()
    while go:
        # Solve a cycle of the model, recording it if horizon is finite
        solution_cycle = solveOneCycle(agent, solution_last)
        if not infinite_horizon:
            solution += solution_cycle

        # Check for termination: identical solutions across cycle iterations or run out of cycles
        solution_now = solution_cycle[-1]
        if infinite_horizon:
            if completed_cycles > 0:
                solution_distance = solution_now.distance(solution_last)
                go = (solution_distance > agent.tolerance and completed_cycles < max_cycles)
            else:  # Assume solution does not converge after only one cycle
                solution_distance = 100.0
                go = True
        else:
            cycles_left += -1
            go = cycles_left > 0

        # Update the "last period solution"
        solution_last = solution_now
        completed_cycles += 1

        # Display progress if requested
        if verbose:
            t_now = clock()
            if infinite_horizon:
                print('Finished cycle #' + str(completed_cycles) + ' in ' + str(t_now-t_last) +
                      ' seconds, solution distance = ' + str(solution_distance))
            else:
                print('Finished cycle #' + str(completed_cycles) + ' of ' + str(agent.cycles) +
                      ' in ' + str(t_now-t_last) + ' seconds.')
            t_last = t_now

    # Record the last cycle if horizon is infinite (solution is still empty!)
    if infinite_horizon:
        solution = solution_cycle  # PseudoTerminal=False impossible for infinite horizon

    # Restore the direction of time to its original orientation, then return the solution
    if original_time_flow:
        agent.timeFwd()
    return solution

def solveOneCycle(agent, solution_last):
    '''
    Solve one "cycle" of the dynamic model for one agent type.  This function
    iterates over the periods within an agent's cycle, updating the time-varying
    parameters and passing them to the single period solver(s).

    Parameters
    ----------
    agent : AgentType
        The microeconomic AgentType whose dynamic problem is to be solved.
    solution_last : Solution
        A representation of the solution of the period that comes after the
        end of the sequence of one period problems.  This might be the term-
        inal period solution, a "pseudo terminal" solution, or simply the
        solution to the earliest period from the succeeding cycle.

    Returns
    -------
    solution_cycle : [Solution]
        A list of one period solutions for one "cycle" of the AgentType's
        microeconomic model.  Returns in reverse chronological order.
    '''
    # Calculate number of periods per cycle, defaults to 1 if all variables are time invariant
    if len(agent.time_vary) > 0:
        # name = agent.time_vary[0]
        # T = len(eval('agent.' + name))
        T = len(agent.__dict__[agent.time_vary[0]])
    else:
        T = 1

    # Check whether the same solution method is used in all periods
    always_same_solver = 'solveOnePeriod' not in agent.time_vary
    if always_same_solver:
        solveOnePeriod = agent.solveOnePeriod
        these_args = getArgNames(solveOnePeriod)

    # Construct a dictionary to be passed to the solver
    # time_inv_string = ''
    # for name in agent.time_inv:
    #     time_inv_string += ' \'' + name + '\' : agent.' + name + ','
    # time_vary_string = ''
    # for name in agent.time_vary:
    #     time_vary_string += ' \'' + name + '\' : None,'
    # solve_dict = eval('{' + time_inv_string + time_vary_string + '}')
    solve_dict = {parameter: agent.__dict__[parameter] for parameter in agent.time_inv}
    solve_dict.update({parameter: None for parameter in agent.time_vary})

    # Initialize the solution for this cycle, then iterate on periods
    solution_cycle = []
    solution_next = solution_last
    for t in range(T):
        # Update which single period solver to use (if it depends on time)
        if not always_same_solver:
            solveOnePeriod = agent.solveOnePeriod[t]
            these_args = getArgNames(solveOnePeriod)

        # Update time-varying single period inputs
        for name in agent.time_vary:
            if name in these_args:
                # solve_dict[name] = eval('agent.' + name + '[t]')
                solve_dict[name] = agent.__dict__[name][t]
        solve_dict['solution_next'] = solution_next

        # Make a temporary dictionary for this period
        temp_dict = {name: solve_dict[name] for name in these_args}

        # Solve one period, add it to the solution, and move to the next period
        solution_t = solveOnePeriod(**temp_dict)
        solution_cycle.append(solution_t)
        solution_next = solution_t

    # Return the list of per-period solutions
    return solution_cycle

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

class Market(HARKobject):
    '''
    A superclass to represent a central clearinghouse of information.  Used for
    dynamic general equilibrium models to solve the "macroeconomic" model as a
    layer on top of the "microeconomic" models of one or more AgentTypes.
    '''
    def __init__(self, agents=[], sow_vars=[], reap_vars=[], const_vars=[], track_vars=[], dyn_vars=[],
                 millRule=None, calcDynamics=None, act_T=1000, tolerance=0.000001):
        '''
        Make a new instance of the Market class.

        Parameters
        ----------
        agents : [AgentType]
            A list of all the AgentTypes in this market.
        sow_vars : [string]
            Names of variables generated by the "aggregate market process" that should
            be "sown" to the agents in the market.  Aggregate state, etc.
        reap_vars : [string]
            Names of variables to be collected ("reaped") from agents in the market
            to be used in the "aggregate market process".
        const_vars : [string]
            Names of attributes of the Market instance that are used in the "aggregate
            market process" but do not come from agents-- they are constant or simply
            parameters inherent to the process.
        track_vars : [string]
            Names of variables generated by the "aggregate market process" that should
            be tracked as a "history" so that a new dynamic rule can be calculated.
            This is often a subset of sow_vars.
        dyn_vars : [string]
            Names of variables that constitute a "dynamic rule".
        millRule : function
            A function that takes inputs named in reap_vars and returns an object
            with attributes named in sow_vars.  The "aggregate market process" that
            transforms individual agent actions/states/data into aggregate data to
            be sent back to agents.
        calcDynamics : function
            A function that takes inputs named in track_vars and returns an object
            with attributes named in dyn_vars.  Looks at histories of aggregate
            variables and generates a new "dynamic rule" for agents to believe and
            act on.
        act_T : int
            The number of times that the "aggregate market process" should be run
            in order to generate a history of aggregate variables.
        tolerance: float
            Minimum acceptable distance between "dynamic rules" to consider the
            Market solution process converged.  Distance is a user-defined metric.

        Returns
        -------
        None
    '''
        self.agents     = agents # NOQA
        self.reap_vars  = reap_vars # NOQA
        self.sow_vars   = sow_vars # NOQA
        self.const_vars = const_vars # NOQA
        self.track_vars = track_vars # NOQA
        self.dyn_vars   = dyn_vars # NOQA
        if millRule is not None:  # To prevent overwriting of method-based millRules
            self.millRule = millRule
        if calcDynamics is not None:  # Ditto for calcDynamics
            self.calcDynamics = calcDynamics
        self.act_T     = act_T # NOQA
        self.tolerance = tolerance # NOQA
        self.max_loops = 1000 # NOQA

        self.print_parallel_error_once = True
        # Print the error associated with calling the parallel method
        # "solveAgents" one time. If set to false, the error will never
        # print. See "solveAgents" for why this prints once or never.

    def solveAgents(self):
        '''
        Solves the microeconomic problem for all AgentTypes in this market.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        # for this_type in self.agents:
        # this_type.solve()
        try:
            multiThreadCommands(self.agents, ['solve()'])
        except Exception as err:
            if self.print_parallel_error_once:
                # Set flag to False so this is only printed once.
                self.print_parallel_error_once = False
                print("**** WARNING: could not execute multiThreadCommands in HARK.core.Market.solveAgents() ",
                      "so using the serial version instead. This will likely be slower. "
                      "The multiTreadCommands() functions failed with the following error:", '\n',
                      sys.exc_info()[0], ':', err)  # sys.exc_info()[0])
            multiThreadCommandsFake(self.agents, ['solve()'])

    def solve(self):
        '''
        "Solves" the market by finding a "dynamic rule" that governs the aggregate
        market state such that when agents believe in these dynamics, their actions
        collectively generate the same dynamic rule.

        Parameters
        ----------
        None

        Returns
        -------
        None
        '''
        go = True
        max_loops = self.max_loops  # Failsafe against infinite solution loop
        completed_loops = 0
        old_dynamics = None

        while go:  # Loop until the dynamic process converges or we hit the loop cap
            self.solveAgents()  # Solve each AgentType's micro problem
            self.makeHistory()  # "Run" the model while tracking aggregate variables
            new_dynamics = self.updateDynamics()  # Find a new aggregate dynamic rule

            # Check to see if the dynamic rule has converged (if this is not the first loop)
            if completed_loops > 0:
                distance = new_dynamics.distance(old_dynamics)
            else:
                distance = 1000000.0

            # Move to the next loop if the terminal conditions are not met
            old_dynamics = new_dynamics
            completed_loops += 1
            go = distance >= self.tolerance and completed_loops < max_loops

        self.dynamics = new_dynamics  # Store the final dynamic rule in self

    def reap(self):
        '''
        Collects attributes named in reap_vars from each AgentType in the market,
        storing them in respectively named attributes of self.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for var_name in self.reap_vars:
            harvest = []
            for this_type in self.agents:
                harvest.append(getattr(this_type, var_name))
            setattr(self, var_name, harvest)

    def sow(self):
        '''
        Distributes attrributes named in sow_vars from self to each AgentType
        in the market, storing them in respectively named attributes.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for var_name in self.sow_vars:
            this_seed = getattr(self, var_name)
            for this_type in self.agents:
                setattr(this_type, var_name, this_seed)

    def mill(self):
        '''
        Processes the variables collected from agents using the function millRule,
        storing the results in attributes named in aggr_sow.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        # Make a dictionary of inputs for the millRule
        reap_vars_string = ''
        for name in self.reap_vars:
            reap_vars_string += ' \'' + name + '\' : self.' + name + ','
        const_vars_string = ''
        for name in self.const_vars:
            const_vars_string += ' \'' + name + '\' : self.' + name + ','
        mill_dict = eval('{' + reap_vars_string + const_vars_string + '}')

        # Run the millRule and store its output in self
        product = self.millRule(**mill_dict)
        for j in range(len(self.sow_vars)):
            this_var = self.sow_vars[j]
            this_product = getattr(product, this_var)
            setattr(self, this_var, this_product)

    def cultivate(self):
        '''
        Has each AgentType in agents perform their marketAction method, using
        variables sown from the market (and maybe also "private" variables).
        The marketAction method should store new results in attributes named in
        reap_vars to be reaped later.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for this_type in self.agents:
            this_type.marketAction()

    def reset(self):
        '''
        Reset the state of the market (attributes in sow_vars, etc) to some
        user-defined initial state, and erase the histories of tracked variables.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for var_name in self.track_vars:  # Reset the history of tracked variables
            setattr(self, var_name + '_hist', [])
        for var_name in self.sow_vars:  # Set the sow variables to their initial levels
            initial_val = getattr(self, var_name + '_init')
            setattr(self, var_name, initial_val)
        for this_type in self.agents:  # Reset each AgentType in the market
            this_type.reset()

    def store(self):
        '''
        Record the current value of each variable X named in track_vars in an
        attribute named X_hist.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        for var_name in self.track_vars:
            value_now = getattr(self, var_name)
            getattr(self, var_name + '_hist').append(value_now)

    def makeHistory(self):
        '''
        Runs a loop of sow-->cultivate-->reap-->mill act_T times, tracking the
        evolution of variables X named in track_vars in attributes named X_hist.

        Parameters
        ----------
        none

        Returns
        -------
        none
        '''
        self.reset()  # Initialize the state of the market
        for t in range(self.act_T):
            self.sow()  # Distribute aggregated information/state to agents
            self.cultivate()  # Agents take action
            self.reap()  # Collect individual data from agents
            self.mill()  # Process individual data into aggregate data
            self.store()  # Record variables of interest

    def updateDynamics(self):
        '''
        Calculates a new "aggregate dynamic rule" using the history of variables
        named in track_vars, and distributes this rule to AgentTypes in agents.

        Parameters
        ----------
        none

        Returns
        -------
        dynamics : instance
            The new "aggregate dynamic rule" that agents believe in and act on.
            Should have attributes named in dyn_vars.
        '''
        # Make a dictionary of inputs for the dynamics calculator
        history_vars_string = ''
        arg_names = list(getArgNames(self.calcDynamics))
        if 'self' in arg_names:
            arg_names.remove('self')
        for name in arg_names:
            history_vars_string += ' \'' + name + '\' : self.' + name + '_hist,'
        update_dict = eval('{' + history_vars_string + '}')

        # Calculate a new dynamic rule and distribute it to the agents in agent_list
        dynamics = self.calcDynamics(**update_dict)  # User-defined dynamics calculator
        for var_name in self.dyn_vars:
            this_obj = getattr(dynamics, var_name)
            for this_type in self.agents:
                setattr(this_type, var_name, this_obj)
        return dynamics

# ------------------------------------------------------------------------------
# Code to copy entire modules to a local directory
# ------------------------------------------------------------------------------

#  Define a function to run the copying:
def copy_module(target_path, my_directory_full_path, my_module):
    '''
    Helper function for copy_module_to_local(). Provides the actual copy
    functionality, with highly cautious safeguards against copying over
    important things.

    Parameters
    ----------
    target_path : string
        String, file path to target location

    my_directory_full_path: string
        String, full pathname to this file's directory

    my_module : string
        String, name of the module to copy

    Returns
    -------
    none
    '''

    if target_path == 'q' or target_path == 'Q':
        print("Goodbye!")
        return
    elif target_path == os.path.expanduser("~") or os.path.normpath(target_path) == os.path.expanduser("~"):
        print("You have indicated that the target location is " + target_path +
              " -- that is, you want to wipe out your home directory with the contents of " + my_module +
              ". My programming does not allow me to do that.\n\nGoodbye!")
        return
    elif os.path.exists(target_path):
        print("There is already a file or directory at the location " + target_path +
              ". For safety reasons this code does not overwrite existing files.\n Please remove the file at "
              + target_path +
              " and try again.")
        return
    else:
        user_input = input("""You have indicated you want to copy module:\n    """ + my_module
                           + """\nto:\n    """ + target_path + """\nIs that correct? Please indicate: y / [n]\n\n""")
        if user_input == 'y' or user_input == 'Y':
            # print("copy_tree(",my_directory_full_path,",", target_path,")")
            copy_tree(my_directory_full_path, target_path)
        else:
            print("Goodbye!")
            return

def print_helper():

    my_directory_full_path = os.path.dirname(os.path.realpath(__file__))

    print(my_directory_full_path)

def copy_module_to_local(full_module_name):
    '''
    This function contains simple code to copy a submodule to a location on
    your hard drive, as specified by you. The purpose of this code is to provide
    users with a simple way to access a *copy* of code that usually sits deep in
    the Econ-ARK package structure, for purposes of tinkering and experimenting
    directly. This is meant to be a simple way to explore HARK code. To interact
    with the codebase under active development, please refer to the documentation
    under github.com/econ-ark/HARK/

    To execute, do the following on the Python command line:

        from HARK.core import copy_module_to_local
        copy_module_to_local("FULL-HARK-MODULE-NAME-HERE")

    For example, if you want SolvingMicroDSOPs you would enter

        from HARK.core import copy_module_to_local
        copy_module_to_local("HARK.SolvingMicroDSOPs")

    '''

    # Find a default directory -- user home directory:
    home_directory_RAW = os.path.expanduser("~")
    # Thanks to https://stackoverflow.com/a/4028943

    # Find the directory of the HARK.core module:
    # my_directory_full_path = os.path.dirname(os.path.realpath(__file__))
    hark_core_directory_full_path = os.path.dirname(os.path.realpath(__file__))
    # From https://stackoverflow.com/a/5137509
    # Important note from that answer:
    # (Note that the incantation above won't work if you've already used os.chdir()
    # to change your current working directory,
    # since the value of the __file__ constant is relative to the current working directory and is not changed by an
    #  os.chdir() call.)
    #
    # NOTE: for this specific file that I am testing, the path should be:
    # '/home/npalmer/anaconda3/envs/py3fresh/lib/python3.6/site-packages/HARK/SolvingMicroDSOPs/---example-file---

    # Split out the name of the module. Break if proper format is not followed:
    all_module_names_list = full_module_name.split('.')  # Assume put in at correct format
    if all_module_names_list[0] != "HARK":
        print("\nWarning: the module name does not start with 'HARK'. Instead it is: '"
              + all_module_names_list[0]+"' --please format the full namespace of the module you want. \n"
              "For example, 'HARK.SolvingMicroDSOPs'")
        print("\nGoodbye!")
        return

    # Construct the pathname to the module to copy:
    my_directory_full_path = hark_core_directory_full_path
    for a_directory_name in all_module_names_list[1:]:
        my_directory_full_path = os.path.join(my_directory_full_path, a_directory_name)

    head_path, my_module = os.path.split(my_directory_full_path)

    home_directory_with_module = os.path.join(home_directory_RAW, my_module)

    print("\n\n\nmy_directory_full_path:", my_directory_full_path, '\n\n\n')

    # Interact with the user:
    #     - Ask the user for the target place to copy the directory
    #         - Offer use "q/y/other" option
    #     - Check if there is something there already
    #     - If so, ask if should replace
    #     - If not, just copy there
    #     - Quit

    target_path = input("""You have invoked the 'replicate' process for the current module:\n    """ +
                        my_module + """\nThe default copy location is your home directory:\n    """ +
                        home_directory_with_module + """\nPlease enter one of the three options in single quotes below, excluding the quotes:

        'q' or return/enter to quit the process
        'y' to accept the default home directory: """+home_directory_with_module+"""
        'n' to specify your own pathname\n\n""")

    if target_path == 'n' or target_path == 'N':
        target_path = input("""Please enter the full pathname to your target directory location: """)

        # Clean up:
        target_path = os.path.expanduser(target_path)
        target_path = os.path.expandvars(target_path)
        target_path = os.path.normpath(target_path)

        # Check to see if they included the module name; if not add it here:
        temp_head, temp_tail = os.path.split(target_path)
        if temp_tail != my_module:
            target_path = os.path.join(target_path, my_module)

    elif target_path == 'y' or target_path == 'Y':
        # Just using the default path:
        target_path = home_directory_with_module
    else:
        # Assume "quit"
        return

    if target_path != 'q' and target_path != 'Q' or target_path == '':
        # Run the copy command:
        copy_module(target_path, my_directory_full_path, my_module)

    return

    if target_path != 'q' and target_path != 'Q' or target_path == '':
        # Run the copy command:
        copy_module(target_path, my_directory_full_path, my_module)

    return

def main():
    print("Sorry, HARK.core doesn't actually do anything on its own.")
    print("To see some examples of its frameworks in action, try running a model module.")
    print("Several interesting model modules can be found in /ConsumptionSavingModel.")
    print('For an extraordinarily simple model that demonstrates the "microeconomic" and')
    print('"macroeconomic" frameworks, see /FashionVictim/FashionVictimModel.')

if __name__ == '__main__':
    main()


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
from threading import Lock

import cv2
import numpy as np

import cv_bridge
from cv_bridge.boost.cv_bridge_boost import getCvType
import dynamic_reconfigure.server
import rospy

from jsk_perception.cfg import ImagePublisherConfig
from sensor_msgs.msg import CameraInfo
from sensor_msgs.msg import Image

class ImagePublisher(object):

    def __init__(self):
        self.lock = Lock()
        self.imgmsg = None
        self.encoding = rospy.get_param('~encoding', 'bgr8')
        self.frame_id = rospy.get_param('~frame_id', 'camera')
        self.fovx = rospy.get_param('~fovx', None)
        self.fovy = rospy.get_param('~fovy', None)
        if (self.fovx is None) != (self.fovy is None):
            rospy.logwarn('fovx and fovy should be specified, but '
                          'specified only {}'
                          .format('fovx' if self.fovx else 'fovy'))
        dynamic_reconfigure.server.Server(
            ImagePublisherConfig, self._cb_dyn_reconfig)
        self.pub = rospy.Publisher('~output', Image, queue_size=1)
        self.publish_info = rospy.get_param('~publish_info', True)
        if self.publish_info:
            self.pub_info = rospy.Publisher(
                '~output/camera_info', CameraInfo, queue_size=1)
        rate = rospy.get_param('~rate', 1.)
        rospy.Timer(rospy.Duration(1. / rate), self.publish)

    def _cb_dyn_reconfig(self, config, level):
        file_name = config['file_name']
        config['file_name'] = os.path.abspath(file_name)
        img = cv2.imread(file_name, cv2.IMREAD_UNCHANGED)
        # when file is gray scale but encoding is not grayscale,
        # load the image again in color
        if (len(img.shape) == 2 and
                getCvType(self.encoding) not in [
                    cv2.CV_8UC1, cv2.CV_16UC1, cv2.CV_32FC1]):
            img = cv2.imread(file_name, cv2.IMREAD_COLOR)
        if img is None:
            rospy.logwarn('Could not read image file: {}'.format(file_name))
            with self.lock:
                self.imgmsg = None
        else:
            rospy.loginfo('Read the image file: {}'.format(file_name))
            with self.lock:
                self.imgmsg = self.cv2_to_imgmsg(img, self.encoding)
        return config

    def publish(self, event):
        if self.imgmsg is None:
            return
        now = rospy.Time.now()
        # setup ros message and publish
        with self.lock:
            self.imgmsg.header.stamp = now
            self.imgmsg.header.frame_id = self.frame_id
        self.pub.publish(self.imgmsg)
        if self.publish_info:
            info = CameraInfo()
            info.header.stamp = now
            info.header.frame_id = self.frame_id
            info.width = self.imgmsg.width
            info.height = self.imgmsg.height
            if self.fovx is not None and self.fovy is not None:
                fx = self.imgmsg.width / 2.0 / \
                    np.tan(np.deg2rad(self.fovx / 2.0))
                fy = self.imgmsg.height / 2.0 / \
                    np.tan(np.deg2rad(self.fovy / 2.0))
                cx = self.imgmsg.width / 2.0
                cy = self.imgmsg.height / 2.0
                info.K = np.array([fx, 0, cx,
                                   0, fy, cy,
                                   0, 0, 1.0])
                info.P = np.array([fx, 0, cx, 0,
                                   0, fy, cy, 0,
                                   0, 0, 1, 0])
                info.R = [1, 0, 0, 0, 1, 0, 0, 0, 1]
            self.pub_info.publish(info)

    def cv2_to_imgmsg(self, img, encoding):
        bridge = cv_bridge.CvBridge()
        # resolve encoding
        if getCvType(encoding) in [cv2.CV_8UC1, cv2.CV_16UC1, cv2.CV_32FC1]:
            # mono8
            if len(img.shape) == 3:
                if img.shape[2] == 4:
                    code = cv2.COLOR_BGRA2GRAY
                else:
                    code = cv2.COLOR_BGR2GRAY
                img = cv2.cvtColor(img, code)
            if getCvType(encoding) == cv2.CV_16UC1:
                # 16UC1
                img = img.astype(np.float32)
                img = img / 255 * (2 ** 16)
                img = img.astype(np.uint16)
            elif getCvType(encoding) == cv2.CV_32FC1:
                # 32FC1
                img = img.astype(np.float32)
                img /= 255
        elif getCvType(encoding) == cv2.CV_8UC3 and len(img.shape) == 3:
            # 8UC3
            # BGRA, BGR -> BGR
            img = img[:, :, :3]
            # BGR -> RGB
            if encoding in ('rgb8', 'rgb16'):
                img = img[:, :, ::-1]
        elif (getCvType(encoding) == cv2.CV_8UC4 and
                len(img.shape) == 3 and img.shape[2] == 4):
            # 8UC4
            if encoding in ('rgba8', 'rgba16'):
                # BGRA -> RGBA
                img = img[:, :, [2, 1, 0, 3]]
        else:
            rospy.logerr('unsupported encoding: {0}'.format(encoding))
            return
        return bridge.cv2_to_imgmsg(img, encoding=encoding)

if __name__ == '__main__':
    rospy.init_node('image_publisher')
    ImagePublisher()
    rospy.spin()


text:
import torch
import torchvision
import numpy as np
import os
import csv
import logging
import json

#from fcos_core.data.datasets.load_ct_img import load_prep_img as load_prep_img_dl
from maskrcnn_benchmark.data.datasets.load_ct_img_kits import load_prep_img as load_prep_img_kits
from maskrcnn_benchmark.data.datasets.load_ct_img_lits import load_prep_img as load_prep_img_lits
from maskrcnn_benchmark.data.datasets.load_ct_img_ircadb import load_prep_img as load_prep_img_ircadb

from maskrcnn_benchmark.structures.bounding_box import BoxList
from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.data.datasets.DeepLesion_utils import load_tag_dict_from_xlsfile, gen_mask_polygon_from_recist, load_lesion_tags
from maskrcnn_benchmark.data.datasets.DeepLesion_utils import gen_parent_list, gen_exclusive_list, gen_children_list
from maskrcnn_benchmark.data.transforms.build import build_d2transforms

class combinedDataset_fewshot(object):

    def __init__(
        self, split, data_dir, ann_file, transforms=None
    ):
        self.transforms = transforms
        self.split = split
        self.data_path = data_dir
        self.classes = ['__background__',  # always index 0
                        'lesion']
        self.num_classes = len(self.classes)
        ann_file = "/data/nihcc/FCOS/csv_files_datasets/target_combined_fewshot_train.csv"

        self.loadinfo(ann_file)
        self.image_fn_list, self.lesion_idx_grouped = self.load_split_index()
        self.num_images = len(self.image_fn_list)
        self.logger = logging.getLogger(__name__)
        #self.id_to_img_map = {k: v for k, v in enumerate(self.ids)}

        # for classification
#         if cfg.MODEL.TAG_ON:
#             self._process_tags()
#             if split == 'test':
#                 self.logger.info('loading 500 hand-labeled test tags')
#                 self._process_manual_annot_test_tags()

        self.logger.info('Combined target %s num_images: %d' % (split, self.num_images))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target, info).
        """
        #image_fn = 'Copy of ' + self.image_fn_list[index]
        image_fn =  self.image_fn_list[index]
        lesion_idx_grouped = self.lesion_idx_grouped[index]

        boxes0 = self.boxes[lesion_idx_grouped]
        #print('lesion_idx_grouped',lesion_idx_grouped)

        # slice_no = self.slice_idx[lesion_idx_grouped][0]
        slice_intv = self.slice_intv[lesion_idx_grouped][0]
        spacing = self.spacing[lesion_idx_grouped][0]

        #print(image_fn)
        num_slice = cfg.INPUT.NUM_SLICES * cfg.INPUT.NUM_IMAGES_3DCE
        is_train = self.split=='train'

        #print('image_fn',image_fn)
        if 'case' in image_fn:   ########### major change
          im, im_scale, crop = load_prep_img_kits(self.data_path, image_fn, spacing, slice_intv,cfg.INPUT.IMG_DO_CLIP, num_slice=num_slice, is_train=is_train)

        elif 'volume' in image_fn:   ########### major change
          im, im_scale, crop = load_prep_img_lits(self.data_path, image_fn, spacing, slice_intv,cfg.INPUT.IMG_DO_CLIP, num_slice=num_slice, is_train=is_train)

        else:                ########### major change
          im, im_scale, crop = load_prep_img_ircadb(self.data_path, image_fn, spacing, slice_intv,cfg.INPUT.IMG_DO_CLIP, num_slice=num_slice, is_train=is_train)

        current_im_shape = im.shape[0:2]
        boxes_new = boxes0.copy()
        if cfg.INPUT.IMG_DO_CLIP:
            offset = [crop[2], crop[0]]
            boxes_new -= offset*2
        boxes_new *= im_scale

        # for i,box in enumerate(boxes_new):
        #   x1,y1,x2,y2 = box
        #   boxes_new[i,:] = x1,y1,x2+x1,y2+y1

        #print('boxes_new',boxes_new)
        im = im.astype(np.uint8)
        im,boxes_new =  build_d2transforms(im,boxes_new)
        im = im.copy() - cfg.INPUT.PIXEL_MEAN
        im = torch.from_numpy(im.transpose((2, 0, 1))).to(dtype=torch.float)

        boxes = torch.as_tensor(boxes_new).reshape(-1, 4)  # guard against no boxes
        target = BoxList(boxes, (im.shape[2], im.shape[1]), mode="xyxy")

        num_boxes = boxes.shape[0]
        classes = torch.ones(num_boxes, dtype=torch.int)  # lesion/nonlesion
        target.add_field("labels", classes)

        infos = {'im_index': index, 'lesion_idxs': lesion_idx_grouped, 'image_fn': image_fn,
                 'crop': crop,  'slice_intv': slice_intv, 'im_scale': im_scale, 'spacing' : spacing, 'current_im_shape' : current_im_shape}

        return im, target, infos

    def __len__(self):
        return len(self.image_fn_list)

    def load_split_index(self):
        """
        need to group lesion indices to image indices, since one image can have multiple lesions
        :return:
        """

        split_list = ['train', 'val', 'test', 'small']
        index = split_list.index(self.split)
        if self.split != 'small':
            lesion_idx_list = np.where((self.train_val_test == index + 1) )[0]
        else:
            lesion_idx_list = np.arange(30)
        fn_list = self.filenames[lesion_idx_list]
        fn_list_unique, inv_ind = np.unique(fn_list, return_inverse=True)
        lesion_idx_grouped = [lesion_idx_list[inv_ind==i] for i in range(len(fn_list_unique))]
        return fn_list_unique, lesion_idx_grouped

    def loadinfo(self, path):
        """load annotations and meta-info from DL_info.csv"""
        info = []
        with open(path) as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                filename = row[0]  # replace the last _ in filename with / or \
                #print(filename)
                #idx = filename.rindex('_')
                row[0] = filename
                info.append(row)
        info = info[1:]

        # the information not used in this project are commented
        self.filenames = np.array([row[0] for row in info])

        self.slice_idx = np.array([int(row[1]) for row in info])

        self.boxes = np.array([[float(x) for x in row[2].split('[')[-1].split(']')[0].split(',')] for row in info])
        #self.boxes -= 1  # coordinates in info file start from 1

        self.spacing = np.array([row[4] for row in info])
        self.slice_intv = np.array([row[3] for row in info])

        self.cr_image_size = np.array([[int(x) for x in row[5].split('(')[-1].split(')')[0].split(',')] for row in info])

        self.train_val_test = np.array([int(row[6]) for row in info])

max_stars_repo_name: manusheoran/UaDAN

max_issues_repo_name: manusheoran/UaDAN

max_forks_repo_name: manusheoran/UaDAN


text:
# Copyright 2018, The Johns Hopkins University Applied Physics Laboratory LLC
# All rights reserved.
# Distributed under the terms of the Modified BSD License.

from abc import ABC, abstractmethod
import logging
import numpy as np

class Environment(ABC):
    """Environment for the simulation

    Args:
        size (int): length of the state vector

    Attributes:
        state (np.array): environment state vector of floats
    """
    logger = logging.getLogger(__name__)

    def __init__(self, size):
        if size > 0:
            self.state = np.zeros(size, dtype='f')
        else:
            self.state = None

    def init(self, now):
        """Initialize the environment's state

        Implement method for any initialization that requires knowledge of the start time.

        Args:
            now (int, float): Start time of the environment
        """
        pass

    @abstractmethod
    def step(self, now, agents):
        """Update the environment's state

        Args:
            now (int, float): Current time of the simulation
            agents (list): list of Agent objects
        """
        pass

    def complete(self, now, agents):
        """Complete a time step

        Implement this if you need to perform any operations at end of time step.

        Args:
            now (int, float): Current time of the simulation
            agents (list): list of Agent objects
        """
        pass

class NullEnvironment(Environment):
    """Empty environment"""
    def __init__(self):
        super().__init__(0)

    def step(self, now, agents):
        pass

class NetworkEnvironment(Environment):
    """Environment with a network over the agents

    Attributes:
        state (np.array): environment state vector of floats
        network (obj): network object
    """
    def __init__(self, size, network):
        super().__init__(size)
        self.network = network


text:
#!/usr/bin/env python

import numpy as np
import sklearn
from sklearn.preprocessing import LabelEncoder

import pickle

from sensor_stick.srv import GetNormals
from sensor_stick.features import compute_color_histograms
from sensor_stick.features import compute_normal_histograms
from visualization_msgs.msg import Marker

from sensor_stick.marker_tools import *
from sensor_stick.msg import DetectedObjectsArray
from sensor_stick.msg import DetectedObject
from sensor_stick.pcl_helper import *

g_model = None
g_clf = None
g_encoder = None
g_scaler = None

def get_normals(cloud):
    get_normals_prox = rospy.ServiceProxy('/feature_extractor/get_normals', GetNormals)
    return get_normals_prox(cloud).cluster

# Callback function for your Point Cloud Subscriber
def pcl_callback(pcl_msg):

    # Convert ROS msg to PCL data
    _pcl_cloud = ros_to_pcl( pcl_msg )

    # Voxel Grid Downsampling
    _voxel_filter = _pcl_cloud.make_voxel_grid_filter()
    _voxel_filter.set_leaf_size( 0.01, 0.01, 0.01 )
    _pcl_cloud = _voxel_filter.filter()

    # PassThrough Filter
    _pass_filter = _pcl_cloud.make_passthrough_filter()
    _pass_filter.set_filter_field_name( 'z' )
    _pass_filter.set_filter_limits( 0.6, 1.1 )
    _pcl_cloud = _pass_filter.filter()

    # RANSAC Plane Segmentation
    _ransac_seg = _pcl_cloud.make_segmenter()
    _ransac_seg.set_model_type( pcl.SACMODEL_PLANE )
    _ransac_seg.set_method_type( pcl.SAC_RANSAC )
    _ransac_seg.set_distance_threshold( 0.02 )
    _plane_indices, _ = _ransac_seg.segment()

    # Extract inliers and outliers
    _table_cloud = _pcl_cloud.extract( _plane_indices, negative = False )
    _objects_cloud = _pcl_cloud.extract( _plane_indices, negative = True )

    # Euclidean Clustering

    _xyz_objects_cloud = XYZRGB_to_XYZ( _objects_cloud )
    _kdtree = _xyz_objects_cloud.make_kdtree()

    _cluster_extractor = _xyz_objects_cloud.make_EuclideanClusterExtraction()
    _cluster_extractor.set_ClusterTolerance( 0.05 )
    _cluster_extractor.set_MinClusterSize( 30 )
    _cluster_extractor.set_MaxClusterSize( 2000 )
    _cluster_extractor.set_SearchMethod( _kdtree )
    _cluster_indices = _cluster_extractor.Extract()

    # Create Cluster-Mask Point Cloud to visualize each cluster separately
    _cluster_color = get_color_list( len( _cluster_indices ) )
    _color_cluster_point_list = []

    for j, indices in enumerate( _cluster_indices ) :
        for i, index in enumerate( indices ) :
            _color_cluster_point_list.append( [ _xyz_objects_cloud[ index ][0],
                                                _xyz_objects_cloud[ index ][1],
                                                _xyz_objects_cloud[ index ][2],
                                                rgb_to_float( _cluster_color[j] ) ] )

    _cluster_cloud = pcl.PointCloud_PointXYZRGB()
    _cluster_cloud.from_list( _color_cluster_point_list )
    _ros_cloud_cluster = pcl_to_ros( _cluster_cloud )

    # Convert PCL data to ROS messages
    _ros_cloud_objects = pcl_to_ros( _objects_cloud )
    _ros_cloud_table = pcl_to_ros( _table_cloud )

    # Publish ROS messages
    pcl_objects_pub.publish( _ros_cloud_objects )
    pcl_table_pub.publish( _ros_cloud_table )
    pcl_cluster_pub.publish( _ros_cloud_cluster )

    # Exercise-3 *****************************************

    # Classify the clusters! (loop through each detected cluster one at a time)
    _detected_objects_labels = []
    _detected_objects = []

    for _index, _pts_list in enumerate( _cluster_indices ) :

        # Grab the points for the cluster
        _pcl_cluster_cloud = _objects_cloud.extract( _pts_list )

        # Compute the associated feature vector
        _ros_cluster_cloud = pcl_to_ros( _pcl_cluster_cloud )

        _chists = compute_color_histograms( _ros_cluster_cloud )
        _normals = get_normals( _ros_cluster_cloud )
        _nhists = compute_normal_histograms( _normals )
        _feature_vector = np.concatenate( ( _chists, _nhists ) )

        # Make the prediction
        _prediction = g_clf.predict( g_scaler.transform( _feature_vector.reshape( 1, -1 ) ) )
        _label = g_encoder.inverse_transform( _prediction )[0]
        _detected_objects_labels.append( _label )

        # Publish a label into RViz
        _label_pos = list( _xyz_objects_cloud[ _pts_list[0] ] )
        _label_pos[2] += 0.4
        object_markers_pub.publish( make_label( _label, _label_pos, _index ) )

        # Add the detected object to the list of detected objects.
        _do = DetectedObject()
        _do.label = _label
        _do.cloud = _ros_cluster_cloud
        _detected_objects.append( _do )

    rospy.loginfo( 'Detected {} objects: {}'.format( len( _detected_objects_labels ), _detected_objects_labels ) )

    # Publish the list of detected objects
    detected_objects_pub.publish( _detected_objects )

if __name__ == '__main__':

    # ROS node initialization
    rospy.init_node( 'clustering', anonymous = True )
    # Create Subscribers
    pcl_sub = rospy.Subscriber( "/sensor_stick/point_cloud", pc2.PointCloud2, pcl_callback, queue_size = 1 )
    # Create Publishers
    pcl_objects_pub = rospy.Publisher( "/pcl_objects", PointCloud2, queue_size = 1 )
    pcl_table_pub = rospy.Publisher( "/pcl_table", PointCloud2, queue_size = 1 )
    pcl_cluster_pub = rospy.Publisher( "/pcl_cluster", PointCloud2, queue_size = 1 )
    object_markers_pub = rospy.Publisher( "/object_markers", Marker, queue_size = 1 )
    detected_objects_pub = rospy.Publisher( "/detected_objects", DetectedObjectsArray, queue_size = 1 )

    g_model = pickle.load( open( 'model.sav', 'rb' ) )
    g_clf = g_model['classifier']
    g_encoder = LabelEncoder()
    g_encoder.classes_ = g_model['classes']
    g_scaler = g_model['scaler']

    # Initialize color_list
    get_color_list.color_list = []

    # Spin while node is not shutdown
    while not rospy.is_shutdown() :
        rospy.spin()

max_stars_repo_name: wpumacay/RoboND-Perception-Project

max_issues_repo_name: wpumacay/RoboND-Perception-Project

max_forks_repo_name: wpumacay/RoboND-Perception-Project