# Pricing with Rational Inattnetion: Mackowiak and Wiederholt (2009)

This notebook replicates [Mackowiak and Wiederholt (2009)](https://www.aeaweb.org/articles?id=10.1257/aer.99.3.769) (henceforth MW) using the methods and the [solver](https://github.com/choongryulyang/dynamic_multivariate_RI) from [Afrouzi and Yang (2019)](http://www.afrouzi.com/dynamic_inattention.pdf) that was adapted for Julia by [Miguel Acosta](https://www.acostamiguel.com/home).

## Contents
* [Setup](#setup)
* [Mapping the Problem to a D.R.I.P.](#map_drip)
* [Initialization](#param)
* [Functions](#funcs)
    * [Solving for the fixed point given $\omega$](#fpointomega)
    * [Solving for the optimal $\omega$](#solveomega)
* [Figures](#figures)
    * [Benchmark Economy](#fig_benchmark)
    * [Other values of real rigidity](#robust_alpha)
    * [Other values of information capacity](#robust_kappa)
* [Measuring Performance/Speed](#performance)

## Setup

The problem in MW is
\begin{align}
	&\min_{\{\hat{\Delta}_{i,t},\hat{z}_{i,t}\}} \left\{ E\left[(\Delta_t-\hat{\Delta}_{i,t})^2\right]+\underset{\equiv \xi}{\underbrace{\left(\frac{\hat{\pi}_{14}}{\hat{\pi}_{11}}\right)}}^2E\left[(z_{i,t}-\hat{z}_{i,t})^2\right]\right\}, \\
	s.t.\quad & \mathcal{I}(\{\Delta_t\};\{\hat{\Delta}_{i,t}\})+\mathcal{I}(\{z_{i,t}\};\{\hat{z}_{i,t}\})\leq \kappa, \tag{capacity constraint}\\
	& \{\Delta_t,\hat{\Delta}_{i,t}\} \perp \{z_{i,t},\hat{z}_{i,t}\} \tag{independence assumption}
\end{align}
where
\begin{align}
    \Delta_t&\equiv p_t + \underset{\equiv 1-\alpha}{\underbrace{\left(\frac{|\hat{\pi}_{13}|}{|\hat{\pi}_{11}|}\right)}}(q_t-p_t) \\
    p_t &= \int_0^1 \hat{\Delta}_{i,t}di \\
    q_t &= \rho q_{t-1} + \nu_t, \nu_{q,t}\sim \mathcal{N}(0,\sigma_q^2) \\
    z_{i,t}&= \rho z_{i,t-1} + \nu_{z,t}, \nu_{z,t}\sim \mathcal{N}(0,\sigma_z^2)
\end{align}

## Mapping the Problem to a D.R.I.P.

There are a few ways of translating the problem above to a D.R.I.P. structure; however, the most efficient way, due to the independence assumption, is to write it as the sum of two D.R.I.P.'s: one that solves the attention problem for the idiosyncratic shock, and one that solves the attention problem for the aggregate shock which also has endogenous feedback.

Moreover, since the problem above has a fixed capacity, instead of a fixed cost of attention ($\omega$) as in D.R.I.P. pacakge, we need to iterate over $\omega$'s to find the one that corresponds with $\kappa$.

## Initialization
Include the solver and import packages:

```julia
include("DRIP.jl");
using Printf, Plots, LaTeXStrings, BenchmarkTools; pyplot();
```

Set parameters:

```julia
## Parameters
ρ  = 0.95;
σq = 0.01;
σz = 11.8*σq;
κ  = 3;
ξ  = 1;
α  = 1 - 0.15;

## Primitives of drip
L  = 21; # length of trunction
A  = [zeros(1,L);[Matrix(I,L-1,L-1);zeros(1,L-1)]']; # MW truncate the state space with linear irfs of length 20
Qq = zeros(L,1); Qq[1]=σq;
Qz = zeros(L,1); Qz[1]=σz;
H  = zeros(L,1); H[1:21] = Array(1:-1/20:0);
```

## Functions
We start with a function that solves the aggregate problem with feedback for a given $\omega$.
### Solving for the fixed point given $\omega$

```julia
function agg_drip(ω,A,Qq,        #primitives of drip except for H because H is endogenous
                  α,             #strategic complementarity
                  H;             #state space rep. of q
                  β     = 1,     #optional: discount factor, MW's parameterization implies β = 1
                  H0    = H,     #optional: initial guess for HΔ (H is the true solution when α=0)
                  maxit = 10000, #optional: max number of iterations for GE code
                  tol   = 1e-4,  #optional: tolerance for iterations
                  w     = 1)     #optional: update weight for RI
    # set primitives
    errmin= 1;
    err   = 1;
    iter  = 0;
    L     = length(H);
    while (err > tol) & (iter < maxit)
            if iter == 0
                global agg  = solve_drip(ω,β,A,Qq,H0;w = w);
            else
                global agg  = solve_drip(ω,β,A,Qq,H0;Ω0 = agg.Ω , Σ0 = agg.Σ_1,w = w);
            end

            XFUN(jj) = ((I-agg.K*agg.Y')*agg.A)^jj * (agg.K*agg.Y') * (agg.A')^jj
            X = infinitesum(XFUN; maxit=200, start = 0);  #E[x⃗]=X×x⃗

            XpFUN(jj) = α^jj * X^(jj)
            Xp = infinitesum(XpFUN; maxit=200, start = 0);

            H1 = (1-α)*Xp'*H;

            err= 0.5*norm(H1-H0,2)/norm(H0)+0.5*err;
            if capacity(agg) < 1e-2
                H0 = 0.9*H0+0.1*rand(L);
            else
                H0 = H1;
                if err < errmin
                    global aggmin = agg;
                    errmin = err;
                end
            end
            iter += 1;
    end
    return(aggmin)
end;
```

Now we need a function that iterates over $\omega$'s to find the one that corresponds to a given capacity for the MW problem.
### Solving for the optimal $\omega$

```julia
function MW(κ,α,A,Qq,Qz,Hq,Hz; #primitives of MW problem
            ω     = σq^2,      #optional: initial guess for ω
            tol   = 1e-3,      #optional: tolerance for κ
            maxit = 10000)     #optional: max iterations
    ωs    = [ω; 2*ω];
    caps  = [];
    iter  = 0;
    err   = 1;
    it    = 0;
    while (err > tol) & (iter < maxit)
        agg = agg_drip(ω,A,Qq,α,H; H0 = rand(L),maxit=20,w=0.95);
        idi = solve_drip(ω,1,A,Qz,H,w = 0.9) ;
        cap = capacity(agg, unit = "bit") + capacity(idi, unit = "bit");
        x = ω/σq^2;
        @printf("ω = %.2fσq² for κ = %.2f \n",x,cap)
        push!(caps,cap);
        if it == 0
            ω = ωs[end];
        else
            slope = (caps[end]-caps[end-1])/(ωs[end]-ωs[end-1]);
            ω     = ω + (κ-caps[end])/slope;
            push!(ωs,ω);
        end
        err = abs(caps[end] - κ)/κ;
        it  += 1;
    end
    return(ω);
end;
```

## Figures
Start with the benchmark calibration:
### Benchmark Economy

```julia
ω    = MW(3,α,A,Qq,Qz,H,H);
agg  = agg_drip(ω,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95);
idi  = solve_drip(ω,1,A,Qz,H,w = 0.9);
@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",capacity(agg),capacity(idi));
```

    ω = 1.00σq² for κ = 3.79
    ω = 2.00σq² for κ = 3.13
    ω = 2.20σq² for κ = 3.07
    ω = 2.43σq² for κ = 2.98
    ω = 2.38σq² for κ = 3.01
    ω = 2.40σq² for κ = 2.98
    ω = 2.39σq² for κ = 2.98
    ω = 2.32σq² for κ = 3.02
    ω = 2.36σq² for κ = 3.00
    Agg. Capacity = 0.08 bits, Idio. Capacity = 2.93 bits

```julia
iirfs  = dripirfs(idi,L)
airfs  = dripirfs(agg,L)

p1 = plot([iirfs.a[1,1,:],σz*H],
     label             = ["Perfect information" "Rational inattention"],
     marker            = [:square :circle],
     color             = [:gray25 :black],
     markercolor       = [false :black],
     markerstrokecolor = [:gray25 :black],
     markersize        = [7 3],
     xlabel            = "Periods",
     ylabel            = "Impulse responses to shocks \n of one standard deviation");
p2 = plot([σq*H,airfs.a[1,1,1:end],σq*agg.H],
     label             = ["Perfect information" "Rational inattention" L"$\Delta$ at the fixed point w/ RI"],
     marker            = [:circle :square :utriangle],
     color             = [:black :gray25 :gray50],
     markercolor       = [:black false false] ,
     markerstrokecolor = [:black :gray25 :gray50],
     markersize        = [3 7 7],
     xlabel            = "Periods",
     ylabel            = "Impulse responses to shocks \n of one standard deviation")

plot(p1,p2,
    layout     = (2,1),
    xlim       = (0,L+1),
    lw         = 1,
    legend     = :topright,
    legendfont = font(12),
    tickfont   = font(12),
    size       = (1000,550),
    xticks     = 1:2:21,
    framestyle = :box)
```

### Other values of real rigidity ($\alpha$)
$\alpha=0.7$:

```julia
ω_α7   = MW(3,0.7,A,Qq,Qz,H,H);
agg_α7 = agg_drip(ω_α7,A,Qq,0.7,H; H0 = rand(L), maxit = 100, w = 0.95);
idi_α7 = solve_drip(ω_α7,1,A,Qz,H,w = 0.9);
@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",capacity(agg_α7),capacity(idi_α7));
```

    ω = 1.00σq² for κ = 3.94
    ω = 2.00σq² for κ = 3.31
    ω = 2.48σq² for κ = 3.11
    ω = 2.77σq² for κ = 3.01
    ω = 2.81σq² for κ = 3.01
    ω = 2.84σq² for κ = 3.00
    Agg. Capacity = 0.19 bits, Idio. Capacity = 2.81 bits

$\alpha$ = 0:

```julia
ω_α0   = MW(3,0,A,Qq,Qz,H,H);
agg_α0 = agg_drip(ω_α0,A,Qq,0,H; H0 = rand(L), maxit = 100, w = 0.95);
idi_α0 = solve_drip(ω_α0,1,A,Qz,H,w = 0.9);
@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",capacity(agg_α0),capacity(idi_α0));
```

    ω = 1.00σq² for κ = 4.16
    ω = 2.00σq² for κ = 3.49
    ω = 2.73σq² for κ = 3.21
    ω = 3.27σq² for κ = 3.05
    ω = 3.43σq² for κ = 3.00
    Agg. Capacity = 0.32 bits, Idio. Capacity = 2.68 bits

```julia
airfs_α7  = dripirfs(agg_α7,L);
airfs_α0  = dripirfs(agg_α0,L);

plot(1:L,0.75*[σq*H,airfs.a[1,1,:],airfs_α7.a[1,1,:],airfs_α0.a[1,1,:]],
    label             = ["Perfect information" L"Rational inattention, benchmark economy, $\alpha = 0.85$" L"Lower degree of real rigidity, $\alpha = 0.7$" L"Lower degree of real rigidity, $\alpha = 0$"],
    marker            = [:circle :square :x :circle],
    color             = [:black :gray20 :gray40 :gray60],
    markercolor       = [:black false :gray40 false],
    markerstrokecolor = [:black :gray20 :gray40 :gray60],
    markersize        = [3 7 7 7],
    xlim              = (0,L+1),
    lw                = 1,
    xticks            = 1:2:21,
    legend            = :topright,
    legendfont        = font(11),
    tickfont          = font(11),
    size              = (1000,275),
    ylim              = (0,σq),
    framestyle        = :box,
    xlabel            = "Periods",
    ylabel            = "Impulse responses to shocks \n of one standard deviation")
```

### Other values of information capacity ($\kappa$)
$\kappa=4$:

```julia
ω_κ4   = MW(4,α,A,Qq,Qz,H,H);
agg_κ4 = agg_drip(ω_κ4,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
idi_κ4 = solve_drip(ω_κ4,1,A,Qz,H,w = 0.9)
@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",capacity(agg_κ4),capacity(idi_κ4));
```

    ω = 1.00σq² for κ = 3.80
    ω = 2.00σq² for κ = 3.14
    ω = 0.69σq² for κ = 4.14
    ω = 0.87σq² for κ = 3.93
    ω = 0.81σq² for κ = 4.00
    Agg. Capacity = 0.35 bits, Idio. Capacity = 3.64 bits

$\kappa=5$:

```julia
ω_κ5   = MW(5,α,A,Qq,Qz,H,H; ω = 0.1*σq^2);
agg_κ5 = agg_drip(ω_κ5,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
idi_κ5 = solve_drip(ω_κ5,1,A,Qz,H,w = 0.9)
@printf("Agg. Capacity = %.2f bits, Idio. Capacity = %.2f bits",capacity(agg_κ5),capacity(idi_κ5));
```

    ω = 0.10σq² for κ = 6.25
    ω = 0.20σq² for κ = 5.41
    ω = 0.25σq² for κ = 5.17
    ω = 0.28σq² for κ = 5.03
    ω = 0.29σq² for κ = 5.00
    Agg. Capacity = 0.70 bits, Idio. Capacity = 4.30 bits

Plot:

```julia
airfs_κ4  = dripirfs(agg_κ4,L);
airfs_κ5  = dripirfs(agg_κ5,L);

plot(1:L,0.75*[σq*H,airfs.a[1,1,:],airfs_κ4.a[1,1,:],airfs_κ5.a[1,1,:]],
    label             = ["Perfect information" L"Rational inattention, benchmark economy, $\kappa = 3$" L"Rational inattention $\kappa = 4$" L"Rational inattention $\kappa = 5$"],
    marker            = [:circle :square :x :circle],
    color             = [:black :gray20 :gray40 :gray60],
    markercolor       = [:black false :gray40 false],
    markerstrokecolor = [:black :gray20 :gray40 :gray60],
    markersize        = [3 7 7 7],
    xlim              = (0,L+1),
    lw                = 1,
    xticks            = 1:2:21,
    legend            = :topright,
    legendfont        = font(11),
    tickfont          = font(11),
    size              = (1000,275),
    ylim              = (0,σq),
    framestyle        = :box,
    xlabel            = "Periods",
    ylabel            = "Impulse responses to shocks \n of one standard deviation")
```

### Measure Performance/Speed
#### Performance of the code for aggregate problem with feedback

For random values of $\omega$ and benchmark values of other parameters:

```julia
@benchmark agg_drip(ω,A,Qq,α,H; H0 = rand(L), maxit = 100, w = 0.95) setup = (ω = σq^2*5*rand())
```

    BenchmarkTools.Trial:
      memory estimate:  486.46 MiB
      allocs estimate:  191569
      --------------
      minimum time:     638.661 ms (13.03% GC)
      median time:      866.816 ms (12.90% GC)
      mean time:        859.116 ms (12.88% GC)
      maximum time:     1.060 s (13.66% GC)
      --------------
      samples:          6
      evals/sample:     1

For $\alpha=0.7$ and random values of $\omega$:

```julia
@benchmark agg_drip(ω_α7,A,Qq,0.7,H; H0 = rand(L), maxit = 100, w = 0.95) setup = (ω = σq^2*5*rand())
```

    BenchmarkTools.Trial:
      memory estimate:  93.13 MiB
      allocs estimate:  35347
      --------------
      minimum time:     107.963 ms (11.21% GC)
      median time:      156.021 ms (14.04% GC)
      mean time:        157.944 ms (13.58% GC)
      maximum time:     229.831 ms (11.66% GC)
      --------------
      samples:          32
      evals/sample:     1

#### Performance of the code for idiosyncratic problem

```julia
@benchmark solve_drip(ω,1,A,Qz,H,w = 0.9) setup = (ω = σq^2*5*rand())
```

    RI Code hit maxit

    BenchmarkTools.Trial:
      memory estimate:  22.45 MiB
      allocs estimate:  8812
      --------------
      minimum time:     29.107 ms (0.00% GC)
      median time:      35.026 ms (14.28% GC)
      mean time:        177.988 ms (11.59% GC)
      maximum time:     4.329 s (12.01% GC)
      --------------
      samples:          30
      evals/sample:     1


```python
# Header starts here.
from sympy.physics.units import *
from sympy import *

# Rounding:
import decimal
from decimal import Decimal as DX
from copy import deepcopy
def iso_round(obj, pv, rounding=decimal.ROUND_HALF_EVEN):
    import sympy
    """
    Rounding acc. to DIN EN ISO 80000-1:2013-08
    place value = Rundestellenwert
    """
    assert pv in set([
        # place value   #  round to:
        1,              #  1
        0.1,            #  1st digit after decimal
        0.01,           #  2nd
        0.001,          #  3rd
        0.0001,         #  4th
        0.00001,        #  5th
        0.000001,       #  6th
        0.0000001,      #  7th
        0.00000001,     #  8th
        0.000000001,    #  9th
        0.0000000001,   # 10th
        ])
    objc = deepcopy(obj)
    try:
        tmp = DX(str(float(objc)))
        objc = tmp.quantize(DX(str(pv)), rounding=rounding)
    except:
        for i in range(len(objc)):
            tmp = DX(str(float(objc[i])))
            objc[i] = tmp.quantize(DX(str(pv)), rounding=rounding)
    return objc

# LateX:
kwargs = {}
kwargs["mat_str"] = "bmatrix"
kwargs["mat_delim"] = ""
# kwargs["symbol_names"] = {FB: "F^{\mathsf B}", }

# Units:
(k, M, G ) = ( 10**3, 10**6, 10**9 )
(mm, cm) = ( m/1000, m/100 )
Newton = kg*m/s**2
Pa     = Newton/m**2
MPa    = M*Pa
GPa    = G*Pa
kN     = k*Newton
deg    = pi/180

half = S(1)/2

# Header ends here.
#
# https://colab.research.google.com/github/kassbohm/tm-snippets/blob/master/ipynb/TM_3/3_ST/3.3.G/3.3.G_cc.ipynb

a = var("a", positive=True)
w1 = var("omega1")

w2z, w3z = var("w2z, w3z")
a1z, a2z, a3z = var("a1z, alpha2, alpha3")
vBx, vBy = var("vBx, vBy")
aBx, aBy = var("aBx, aBy")
vCx, vCy = var("vCx, vCy")
aCx, aCy = var("aCx, aCy")

sub_list=[
    (a, 1 *m),
    (w1, 2 / s),
]

w1 = Matrix([0,0,-w1])
w2 = Matrix([0,0, w2z])
w3 = Matrix([0,0, w3z])

a1 = Matrix([0, 0, 0])
a2 = Matrix([0,0,a2z])
a3 = Matrix([0,0,a3z])

vA = Matrix([0, 0, 0])
aA = Matrix([0, 0, 0])
vB = Matrix([vBx, vBy, 0])
aB = Matrix([aBx, aBy, 0])
vC = Matrix([vCx, vCy, 0])
aC = Matrix([aCx, aCy, 0])
vD = Matrix([0, 0, 0])
aD = Matrix([0, 0, 0])

dAB = Matrix([a, 0, 0])
h = sqrt( (2*a)**2 - (a/2)**2 )
dBC = Matrix([ a/2, h, 0])
dDC = Matrix([-a/2, h, 0])

pprint("\nvB and aB:")
eq1a = Eq( vB, vA + w1.cross(dAB) )
eq1b = Eq( aB, aA + a1.cross(dAB)  + w1.cross(w1.cross(dAB)) )
sol = solve([eq1a, eq1b], [vBx, vBy, aBx, aBy])
vBx, vBy = sol[vBx], sol[vBy]
aBx, aBy = sol[aBx], sol[aBy]
vB = Matrix([vBx, vBy, 0])
aB = Matrix([aBx, aBy, 0])
pprint(vB)
pprint(aB)

pprint("\nvC aC, w2, w3:")
eq2a = Eq( vC, vB + w2.cross(dBC) )
eq3a = Eq( vC, vD + w3.cross(dDC) )
a
sol = solve([eq2a, eq3a], [vCx, vCy, w2z, w3z], dict=True)
sol = sol[0]
pprint(sol)
vCx, vCy = sol[vCx], sol[vCy]
w2z, w3z = sol[w2z], sol[w3z]
vC = Matrix([vCx, vCy, 0])
w2 = Matrix([0,0, w2z])
w3 = Matrix([0,0, w3z])

eq2b = Eq( aC, aB + a2.cross(dBC)  + w2.cross(w2.cross(dBC)) )
eq3b = Eq( aC, aD + a3.cross(dDC)  + w3.cross(w3.cross(dDC)) )

pprint("\n(2b):")
pprint(eq2b)
# pprint(latex(eq2b,**kwargs))
pprint("\n(3b):")
pprint(eq3b)
# pprint(latex(eq3b,**kwargs))

sol = solve([eq2b, eq3b], [aCx, aCy, a2z, a3z], dict=True)
sol = sol[0]
pprint(sol)
aCx, aCy = sol[aCx], sol[aCy]
a2z, a3z = sol[a2z], sol[a3z]

aC = Matrix([aCx, aCy, 0])
a2 = Matrix([0,0,a2z])
a3 = Matrix([0,0,a3z])
pprint(vC)
pprint(aC)

pprint(w2)
pprint(w3)
pprint("\na2 and a3:")
pprint(a2)
pprint(a3)

pprint("\nvC / (m/s):")
tmp = vC.subs(sub_list)
tmp /= m/s
tmp = iso_round(tmp,0.01)
pprint(tmp)

# vB and aB:
# ⎡  0  ⎤
# ⎢     ⎥
# ⎢-a⋅ω₁⎥
# ⎢     ⎥
# ⎣  0  ⎦
# ⎡     2⎤
# ⎢-a⋅ω₁ ⎥
# ⎢      ⎥
# ⎢  0   ⎥
# ⎢      ⎥
# ⎣  0   ⎦
#
# vC and aC:
# ⎡-√15⋅a⋅ω₁ ⎤
# ⎢──────────⎥
# ⎢    2     ⎥
# ⎢          ⎥
# ⎢  -a⋅ω₁   ⎥
# ⎢  ──────  ⎥
# ⎢    2     ⎥
# ⎢          ⎥
# ⎣    0     ⎦
# ⎡        2     ⎤
# ⎢   -a⋅ω₁      ⎥
# ⎢   ───────    ⎥
# ⎢      2       ⎥
# ⎢              ⎥
# ⎢            2 ⎥
# ⎢-17⋅√15⋅a⋅ω₁  ⎥
# ⎢──────────────⎥
# ⎢      30      ⎥
# ⎢              ⎥
# ⎣      0       ⎦
#
# w2 and w3:
# ⎡0 ⎤
# ⎢  ⎥
# ⎢0 ⎥
# ⎢  ⎥
# ⎣ω₁⎦
# ⎡0 ⎤
# ⎢  ⎥
# ⎢0 ⎥
# ⎢  ⎥
# ⎣ω₁⎦
#
# a2 and a3:
# ⎡     0     ⎤
# ⎢           ⎥
# ⎢     0     ⎥
# ⎢           ⎥
# ⎢         2 ⎥
# ⎢-2⋅√15⋅ω₁  ⎥
# ⎢───────────⎥
# ⎣     15    ⎦
# ⎡    0    ⎤
# ⎢         ⎥
# ⎢    0    ⎥
# ⎢         ⎥
# ⎢        2⎥
# ⎢2⋅√15⋅ω₁ ⎥
# ⎢─────────⎥
# ⎣    15   ⎦
#
# vC / (m/s):
# ⎡-3.87⎤
# ⎢     ⎥
# ⎢-1.0 ⎥
# ⎢     ⎥
# ⎣ 0.0 ⎦

```


```python
import sys
sys.path.append('..')
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from sympy import simplify_logic
import time
from sklearn.metrics import accuracy_score
import pandas as pd
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.tree import _tree, export_text
from torchvision import datasets, transforms
from torch.utils.data import TensorDataset, DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
import torch.nn.functional as F
from sklearn.preprocessing import OneHotEncoder, MinMaxScaler, LabelEncoder
from sklearn.model_selection import StratifiedKFold

from lens.utils.base import validate_network, set_seed, tree_to_formula
from lens.utils.relu_nn import get_reduced_model, prune_features
from lens.utils.psi_nn import prune_equal_fanin
from lens import logic

results_dir = 'results/omalizumab'
if not os.path.isdir(results_dir):
    os.makedirs(results_dir)

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
seed = 42
n_rep = 10
tot_epochs = 6001
prune_epochs = 3001
```

```python
gene_expression_matrix = pd.read_csv('data/omalizumab/reduced_w_1/data.csv', index_col=None, header=None)
labels = pd.read_csv('data/omalizumab/reduced_w_1/tempLabels_W-1.csv', index_col=None, header=None)
genes = pd.read_csv('data/omalizumab/reduced_w_1/features.csv', index_col=None, header=None)
gene_expression_matrix
```

      0
      1
      2
      3
      4

      0
      3.320000
      3.320000
      3.32000
      6.941536
      6.590419

      1
      4.232978
      3.320000
      3.32000
      7.279548
      6.476784

      2
      3.320000
      4.200609
      3.32000
      7.741600
      4.643134

      3
      3.320000
      3.320000
      3.32000
      7.276600
      5.953452

      4
      3.320000
      3.320000
      3.32000
      7.224628
      6.555227

      ...
      ...
      ...
      ...
      ...
      ...

      56
      3.320000
      3.320000
      3.32000
      7.660182
      6.128603

      57
      3.320000
      3.700430
      3.45131
      7.809826
      6.153968

      58
      3.320000
      3.320000
      3.32000
      7.580588
      6.134398

      59
      4.174319
      3.320000
      3.32000
      7.016004
      7.124143

      60
      3.699251
      3.320000
      3.32000
      7.568044
      6.236546

61 rows × 5 columns

```python
encoder = LabelEncoder()
labels_encoded = encoder.fit_transform(labels.values)
labels_encoded_noncontrols = labels_encoded[labels_encoded!=0] - 1

data_controls = gene_expression_matrix[labels_encoded==0]
data = gene_expression_matrix[labels_encoded!=0]

gene_signature = data_controls.mean(axis=0)
data_scaled = data - gene_signature

scaler = MinMaxScaler((0, 1))
scaler.fit(data_scaled)
data_normalized = scaler.transform(data_scaled)

x = torch.FloatTensor(data_normalized)
y = torch.FloatTensor(labels_encoded_noncontrols)
print(x.shape)
print(y.shape)
```

    torch.Size([40, 5])
    torch.Size([40])

    c:\users\pietr\anaconda3\envs\deep-logic\lib\site-packages\sklearn\utils\validation.py:63: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples, ), for example using ravel().
      return f(*args, **kwargs)

```python
concepts = list(genes.values.squeeze())
concepts
```

    ['ILMN_3286286',
     'ILMN_1775520',
     'ILMN_1656849',
     'ILMN_1781198',
     'ILMN_1665457']

```python
n_splits = 10
skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=0)
```

```python
def train_nn(x_train, y_train, need_pruning, seed, device, relu=False, verbose=False):
    set_seed(seed)
    x_train = x_train.to(device)
    y_train = y_train.to(device)
    layers = [
        torch.nn.Linear(x_train.size(1), 10),
        torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
        torch.nn.Linear(10, 3),
        torch.nn.ReLU() if relu else torch.nn.LeakyReLU(),
        torch.nn.Linear(3, 1),
        torch.nn.Sigmoid(),
    ]
    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
    loss_form = torch.nn.BCELoss()
    model.train()
    for epoch in range(tot_epochs):
        # forward pass
        optimizer.zero_grad()
        y_pred = model(x_train).squeeze()
        # Compute Loss
        loss = loss_form(y_pred, y_train)

        for module in model.children():
            if isinstance(module, torch.nn.Linear):
                loss += 0.001 * torch.norm(module.weight, 1)
                loss += 0.001 * torch.norm(module.bias, 1)
                break

        # backward pass
        loss.backward()
        optimizer.step()

        if epoch > prune_epochs and need_pruning:
            prune_features(model, n_classes=1, device=device)
            need_pruning = False

        # compute accuracy
        if epoch % 500 == 0 and verbose:
            y_pred_d = y_pred > 0.5
            accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
            print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')

    return model
```

```python
def train_psi_nn(x_train, y_train, need_pruning, seed, device, verbose=False):
    set_seed(seed)
    x_train = x_train.to(device)
    y_train = y_train.to(device).to(torch.float)
    layers = [
        torch.nn.Linear(x_train.size(1), 10),
        torch.nn.Sigmoid(),
        torch.nn.Linear(10, 4),
        torch.nn.Sigmoid(),
        torch.nn.Linear(4, 1),
        torch.nn.Sigmoid(),
    ]
    model = torch.nn.Sequential(*layers).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    loss_form = torch.nn.BCELoss()
    model.train()
    for epoch in range(tot_epochs):
        # forward pass
        optimizer.zero_grad()
        y_pred = model(x_train).squeeze()
        # Compute Loss
        loss = loss_form(y_pred, y_train)

        for module in model.children():
            if isinstance(module, torch.nn.Linear):
                loss += 0.00001 * torch.norm(module.weight, 1)

        # backward pass
        loss.backward()
        optimizer.step()

        if epoch > prune_epochs and need_pruning:
            model = prune_equal_fanin(model, 2, validate=True, device=device)
            need_pruning = False

        # compute accuracy
        if epoch % 500 == 0 and verbose:
            y_pred_d = y_pred > 0.5
            accuracy = y_pred_d.eq(y_train).sum().item() / y_train.size(0)
            print(f'\t Epoch {epoch}: train accuracy: {accuracy:.4f}')

    return model
```

```python
def c_to_y(method, need_pruning, relu, verbose=False):
    methods = []
    splits = []
    explanations = []
    explanations_inv = []
    model_accuracies = []
    explanation_accuracies = []
    explanation_accuracies_inv = []
    elapsed_times = []
    elapsed_times_inv = []
    for split, (train_index, test_index) in enumerate(skf.split(x.cpu().detach().numpy(), y.cpu().detach().numpy())):
        print(f'Split [{split+1}/{n_splits}]')
        x_train, x_test = torch.FloatTensor(x[train_index]), torch.FloatTensor(x[test_index])
        y_train, y_test = torch.FloatTensor(y[train_index]), torch.FloatTensor(y[test_index])

        explanation, explanation_inv = '', ''
        explanation_accuracy, explanation_accuracy_inv = 0, 0

        if method == 'tree':
            classifier = DecisionTreeClassifier(random_state=seed)
            classifier.fit(x_train.detach().numpy(), y_train.detach().numpy())
            y_preds = classifier.predict(x_test.detach().numpy())
            model_accuracy = accuracy_score(y_test.detach().numpy(), y_preds)

            target_class = 1
            start = time.time()
            explanation = tree_to_formula(classifier, concepts, target_class)
            elapsed_time = time.time() - start

            target_class_inv = 0
            start = time.time()
            explanation_inv = tree_to_formula(classifier, concepts, target_class_inv)
            elapsed_time = time.time() - start

        else:
            if method == 'psi':
                # positive class
                target_class = 1
                model = train_psi_nn(x_train, y_train.eq(target_class), need_pruning, seed, device, verbose)
                y_preds = model(x_test.to(device)).cpu().detach().numpy()
                model_accuracy = accuracy_score(y_test.cpu().detach().numpy(), y_preds>0.5)

            else:
                model = train_nn(x_train, y_train, need_pruning, seed, device, relu, verbose)
                y_preds = model(x_test.to(device)).cpu().detach().numpy()
                model_accuracy = accuracy_score(y_test.cpu().detach().numpy(), y_preds>0.5)

            # positive class
            target_class = 1
            start = time.time()
            if method == 'psi':
                global_explanation = logic.generate_fol_explanations(model, device)[0]
            else:
                global_explanation, _, _ = logic.relu_nn.combine_local_explanations(model,
                                                                                x_train.to(device), y_train.to(device),
                                                                                target_class=target_class,
                                                                                method=method, device=device)
            elapsed_time = time.time() - start

            if global_explanation:
                explanation_accuracy, _ = logic.base.test_explanation(global_explanation, target_class, x_test, y_test)
                explanation = logic.base.replace_names(global_explanation, concepts)

            # negative class
            target_class_inv = 0
            if method == 'psi':
                model = train_psi_nn(x_train, y_train.eq(target_class_inv), need_pruning, seed, device, verbose)

            start = time.time()
            if method == 'psi':
                global_explanation_inv = logic.generate_fol_explanations(model, device)[0]
            else:
                global_explanation_inv, _, _ = logic.relu_nn.combine_local_explanations(model,
                                                                                x_train.to(device),
                                                                                y_train.to(device),
                                                                                target_class=target_class_inv,
                                                                                method=method, device=device)
            elapsed_time_inv = time.time() - start
            if global_explanation_inv:
                explanation_accuracy_inv, _ = logic.base.test_explanation(global_explanation_inv,
                                                                          target_class_inv, x_test, y_test)
                explanation_inv = logic.base.replace_names(global_explanation_inv, concepts)

        if verbose:
            print(f'\t Model\'s accuracy: {model_accuracy:.4f}')
            print(f'\t Class {target_class} - Global explanation: "{explanation}" - Accuracy: {explanation_accuracy:.4f}')
            print(f'\t Elapsed time {elapsed_time}')
            print(f'\t Class {target_class_inv} - Global explanation: "{explanation_inv}" - Accuracy: {explanation_accuracy_inv:.4f}')
            print(f'\t Elapsed time {elapsed_time_inv}')

        methods.append(method)
        splits.append(split)
        explanations.append(explanation)
        explanations_inv.append(explanation_inv)
        model_accuracies.append(model_accuracy)
        explanation_accuracies.append(explanation_accuracy)
        explanation_accuracies_inv.append(explanation_accuracy_inv)
        elapsed_times.append(elapsed_time)
        elapsed_times_inv.append(elapsed_time_inv)

    results = pd.DataFrame({
        'method': methods,
        'split': splits,
        'explanation': explanations,
        'explanation_inv': explanations_inv,
        'model_accuracy': model_accuracies,
        'explanation_accuracy': explanation_accuracies,
        'explanation_accuracy_inv': explanation_accuracies_inv,
        'elapsed_time': elapsed_times,
        'elapsed_time_inv': elapsed_times_inv,
    })
    results.to_csv(os.path.join(results_dir, f'results_{method}.csv'))

    return results
```

```python
model = train_nn(x, y, need_pruning=True, seed=seed, device=device, relu=False, verbose=True)
y_preds = model(x.to(device)).cpu().detach().numpy()
model_accuracy = accuracy_score(y.cpu().detach().numpy(), y>0.5)

# positive class
target_class = 1
start = time.time()
global_explanation, _, _ = logic.relu_nn.combine_local_explanations(model,
                                                                    x.to(device), y.to(device),
                                                                    target_class=target_class,
                                                                    topk_explanations=10,
                                                                    method='pruning', device=device)
elapsed_time = time.time() - start

if global_explanation:
    explanation_accuracy, _ = logic.base.test_explanation(global_explanation, target_class, x, y)
    explanation = logic.base.replace_names(global_explanation, concepts)

print(f'\t Model\'s accuracy: {model_accuracy:.4f}')
print(f'\t Class {target_class} - Global explanation: "{explanation}" - Accuracy: {explanation_accuracy:.4f}')
print(f'\t Elapsed time {elapsed_time}')
```

    	 Epoch 0: train accuracy: 0.7500
    	 Epoch 500: train accuracy: 1.0000
    	 Epoch 1000: train accuracy: 1.0000
    	 Epoch 1500: train accuracy: 1.0000
    	 Epoch 2000: train accuracy: 1.0000
    	 Epoch 2500: train accuracy: 1.0000
    	 Epoch 3000: train accuracy: 1.0000
    	 Epoch 3500: train accuracy: 1.0000
    	 Epoch 4000: train accuracy: 1.0000
    	 Epoch 4500: train accuracy: 1.0000
    	 Epoch 5000: train accuracy: 1.0000
    	 Epoch 5500: train accuracy: 1.0000
    	 Epoch 6000: train accuracy: 1.0000
    	 Model's accuracy: 1.0000
    	 Class 1 - Global explanation: "(~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849) | (~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198) | (~ILMN_1775520 & ~ILMN_1656849 & ~ILMN_1781198)" - Accuracy: 97.5000
    	 Elapsed time 0.04890179634094238

# General pruning

```python
method = 'pruning'
need_pruning = True
relu = False
results_pruning = c_to_y(method, need_pruning, relu, verbose=False)
results_pruning
```

    Split [1/10]
    Split [2/10]
    Split [3/10]
    Split [4/10]
    Split [5/10]
    Split [6/10]
    Split [7/10]
    Split [8/10]
    Split [9/10]
    Split [10/10]

      method
      split
      explanation
      explanation_inv
      model_accuracy
      explanation_accuracy
      explanation_accuracy_inv
      elapsed_time
      elapsed_time_inv

      0
      pruning
      0
      ~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198
      (ILMN_1656849 & ILMN_1781198 & ILMN_1665457 & ...
      1.0
      50.0
      75.0
      0.055851
      0.048870

      1
      pruning
      1
      ~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198
      (ILMN_3286286 & ILMN_1781198 & ~ILMN_1656849 &...
      1.0
      50.0
      75.0
      0.052821
      0.042885

      2
      pruning
      2
      ~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849
      (ILMN_1775520 & ILMN_1781198 & ~ILMN_3286286 &...
      1.0
      100.0
      75.0
      0.048869
      0.064827

      3
      pruning
      3
      ~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198
      (ILMN_3286286 & ILMN_1781198 & ~ILMN_1656849 &...
      1.0
      50.0
      75.0
      0.059353
      0.053855

      4
      pruning
      4
      ~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198
      (ILMN_3286286 & ILMN_1781198 & ~ILMN_1656849 &...
      1.0
      50.0
      100.0
      0.051861
      0.081782

      5
      pruning
      5
      ~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849 ...
      ILMN_1656849 & ILMN_1781198 & ~ILMN_3286286 & ...
      1.0
      25.0
      75.0
      0.070779
      0.072300

      6
      pruning
      6
      ~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849
      (ILMN_1775520 & ILMN_1781198 & ~ILMN_3286286 &...
      1.0
      100.0
      75.0
      0.083778
      0.070807

      7
      pruning
      7
      ~ILMN_3286286 & ~ILMN_1656849 & ~ILMN_1781198
      (ILMN_3286286 & ILMN_1781198 & ~ILMN_1656849 &...
      1.0
      50.0
      75.0
      0.063799
      0.051859

      8
      pruning
      8
      ~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849
      (ILMN_1775520 & ILMN_1781198 & ~ILMN_3286286 &...
      1.0
      100.0
      75.0
      0.051926
      0.054788

      9
      pruning
      9
      ~ILMN_3286286 & ~ILMN_1775520 & ~ILMN_1656849 ...
      ILMN_1656849 & ILMN_1781198 & ~ILMN_3286286 & ...
      1.0
      75.0
      75.0
      0.065828
      0.075335

# LIME

```python
method = 'lime'
need_pruning = False
relu = False
results_lime = c_to_y(method, need_pruning, relu, False)
results_lime
```

# ReLU

```python
method = 'weights'
need_pruning = False
relu = True
results_weights = c_to_y(method, need_pruning, relu, False)
results_weights
```

# Psi network

```python
method = 'psi'
need_pruning = True
relu = False
results_psi = c_to_y(method, need_pruning, relu, verbose=False)
results_psi
```

# Decision tree

```python
method = 'tree'
need_pruning = False
relu = False
results_tree = c_to_y(method, need_pruning, relu, verbose=False)
results_tree
```

# Summary

```python
cols = ['model_accuracy', 'explanation_accuracy', 'explanation_accuracy_inv', 'elapsed_time', 'elapsed_time_inv']
mean_cols = [f'{c}_mean' for c in cols]
sem_cols = [f'{c}_sem' for c in cols]

# pruning
df_mean = results_pruning[cols].mean()
df_sem = results_pruning[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_pruning = pd.concat([df_mean, df_sem])
summary_pruning.name = 'pruning'

# lime
df_mean = results_lime[cols].mean()
df_sem = results_lime[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_lime = pd.concat([df_mean, df_sem])
summary_lime.name = 'lime'

# weights
df_mean = results_weights[cols].mean()
df_sem = results_weights[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_weights = pd.concat([df_mean, df_sem])
summary_weights.name = 'weights'

# psi
df_mean = results_psi[cols].mean()
df_sem = results_psi[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_psi = pd.concat([df_mean, df_sem])
summary_psi.name = 'psi'

# tree
df_mean = results_tree[cols].mean()
df_sem = results_tree[cols].sem()
df_mean.columns = mean_cols
df_sem.columns = sem_cols
summary_tree = pd.concat([df_mean, df_sem])
summary_tree.name = 'tree'

summary = pd.concat([summary_pruning,
                     summary_lime,
                     summary_weights,
                     summary_psi,
                     summary_tree], axis=1).T
summary.columns = mean_cols + sem_cols
summary
```

      model_accuracy_mean
      explanation_accuracy_mean
      explanation_accuracy_inv_mean
      elapsed_time_mean
      elapsed_time_inv_mean
      model_accuracy_sem
      explanation_accuracy_sem
      explanation_accuracy_inv_sem
      elapsed_time_sem
      elapsed_time_inv_sem

      pruning
      0.975
      0.675
      0.800
      0.068861
      0.041138
      0.025000
      0.098953
      0.050000
      0.004212
      0.003383

      lime
      0.850
      1.000
      0.250
      24.077748
      10.371542
      0.040825
      0.000000
      0.000000
      0.115518
      0.075365

      weights
      1.000
      0.750
      0.775
      0.109159
      0.050979
      0.000000
      0.091287
      0.078617
      0.003653
      0.003413

      psi
      0.800
      0.900
      0.800
      0.035007
      0.029819
      0.033333
      0.040825
      0.033333
      0.002218
      0.003012

      tree
      0.875
      0.875
      0.875
      0.000000
      0.000000
      0.055902
      0.055902
      0.055902
      0.000000
      0.000000

```python
summary.to_csv(os.path.join(results_dir, 'summary.csv'))
```