# Enzyme Modules

An "enzyme module" is defined as a mechanistic description of a reaction consisting of mass action rate laws for all known reaction steps (Du et al., 2016). In **MASSpy**, enzyme modules are represented by the `EnzymeModule` object.

To demonstrate the utility of an `EnzymeModule` object and how it aids in constructing mechanistic models of enzyme behavior, an `EnzymeModule` of hexokinase$^{1, 2}$ is constructed and then merged with a model of glycolysis$^{3}$ for verification.

## Constructing Enzyme Modules

In order to construct the `EnzymeModule` of hexokinase, the following information is necessary:

1. The enzyme is a monomer.
2. The enzyme binding of substrates follows a random sequential mechanism.
3. The enzyme experiences product inhibtion and is competitively inhibited by 23DPG when complexed with D-glucose.

Total HEX1 Concentration$^2$: $\text{[HEX1]}_{total} = 24 nM = 0.000024 mM$.

```python
from operator import attrgetter

from mass import MassMetabolite
from mass.enzyme_modules import EnzymeModule
from mass.example_data import create_example_model

# Load the glycolysis and hemoglobin models, then merge them
glycolysis = create_example_model("Glycolysis")
hemoglobin = create_example_model("Hemoglobin")
glyc_hb = glycolysis.merge(hemoglobin, inplace=False)
```

    Set parameter Username

The `EnzymeModule` is a subclass of the `MassModel`, meaning that it inherits the methods and behaviors of the `MassModel` object.
Like a `MassModel`, an `EnzymeModule` object requires a unique identifier in order to be created. Optionally, the `name` and `subsystem` attributes are set during initialization.

```python
HEX1 = EnzymeModule("HEX1", name="Hexokinase (D-glucose:ATP)",
                    subsystem="Glycolysis")
```

### Defining the enzyme ligands

The ligands that interact with the enzyme (e.g. as the substrates, activators, and inhibitors) are created as `MassMetabolite` objects and added to the model.

```python
glc__D_c = MassMetabolite(
    "glc__D_c",
    name="D-Glucose",
    formula="C6H12O6",
    charge=0,
    compartment="c")
g6p_c = MassMetabolite(
    "g6p_c",
    name="D-Glucose 6-phosphate",
    formula="C6H11O9P",
    charge=-2,
    compartment="c")
atp_c = MassMetabolite(
    "atp_c",
    name="ATP",
    formula="C10H12N5O13P3",
    charge=-4,
    compartment="c")
adp_c = MassMetabolite(
    "adp_c",
    name="ADP",
    formula="C10H12N5O10P2",
    charge=-3,
    compartment="c")
_23dpg_c = MassMetabolite(
    "_23dpg_c",
    name="2,3-Disphospho-D-glycerate",
    formula="C3H3O10P2",
    charge=-5,
    compartment="c")
h_c = MassMetabolite(
    "h_c",
    name="H+",
    formula="H",
    charge=1,
    compartment="c")

HEX1.add_metabolites([glc__D_c, g6p_c, atp_c, adp_c, _23dpg_c, h_c])
```

Once added to the `EnzymeModule`, ligands can be accessed using the `enzyme_module_ligands` attribute.

```python
HEX1.enzyme_module_ligands
```

    [<MassMetabolite glc__D_c at 0x7fc20268cc40>,
     <MassMetabolite g6p_c at 0x7fc20268cbb0>,
     <MassMetabolite atp_c at 0x7fc20268cc70>,
     <MassMetabolite adp_c at 0x7fc20268ccd0>,
     <MassMetabolite _23dpg_c at 0x7fc20268cd30>,
     <MassMetabolite h_c at 0x7fc20268cd60>]

To keep track of the roles played by various ligands in the module, the `enzyme_module_ligands_categorized` attribute is set. The attribute takes a `dict`, with categories as keys and relevant `MassMetabolite` objects as values. Note that an object can be a part of multiple categories.

```python
HEX1.enzyme_module_ligands_categorized =  {
    "substrates": glc__D_c,
    "cofactors": atp_c,
    "inhibitors": _23dpg_c,
    "products": [adp_c, g6p_c, h_c]}
HEX1.enzyme_module_ligands_categorized
```

    [<Group substrates at 0x7fc1f00ae2e0>,
     <Group cofactors at 0x7fc1f00ae970>,
     <Group inhibitors at 0x7fc1f00ae8b0>,
     <Group products at 0x7fc1f00ae280>]

For each category, a `cobra.Group` is created containing the relevant objects. Once set, the attribute returns a `cobra.DictList` that contains the categorized groups. The groups and their members are printed as follows:

```python
for group in HEX1.enzyme_module_ligands_categorized:
    print("{0}: {1}".format(
        group.id, str(sorted([m.id for m in group.members]))))
```

    substrates: ['glc__D_c']
    cofactors: ['atp_c']
    inhibitors: ['_23dpg_c']
    products: ['adp_c', 'g6p_c', 'h_c']

### Defining the enzyme module forms

After adding `MassMetabolite` objects of ligands to the model, the various forms of the enzyme must be defined. These forms are represented by `EnzymeModuleForm` objects.

The `EnzymeModuleForm` object inherits from the `MassMetabolite` and is treated like any other metabolite in the model. However, the `EnzymeModuleForm` object contains the additional `bound_metabolites` attribute to assist in tracking metabolites bound to the enzyme form.

The `EnzymeModule.make_enzyme_module_form()` method allows for the creation of an `EnzymeModuleForm` object while assigning categories for the  `EnzymeModuleForm` in the process. Using `make_enzyme_module_form()` also adds the species to the module upon creation, accessible via the `EnzymeModule.enzyme_module_forms` attribute.

```python
hex1_c = HEX1.make_enzyme_module_form(
    "hex1_c",
    name="automatic",
    categories="Active",
    compartment="c")

hex1_A_c = HEX1.make_enzyme_module_form(
    "hex1_A_c",  # A stands complexted with ATP
    name="automatic",
    categories="Active",
    bound_metabolites={atp_c: 1},
    compartment="c")

hex1_G_c = HEX1.make_enzyme_module_form(
    "hex1_G_c",  # G stands for complexed with Glucose
    name="automatic",
    categories="Active",
    bound_metabolites={glc__D_c: 1},
    compartment="c")

hex1_AG_c = HEX1.make_enzyme_module_form(
    "hex1_AG_c",
    name="automatic",
    categories="Active",
    bound_metabolites={glc__D_c: 1, atp_c: 1},
    compartment="c")

hex1_G_CI_c = HEX1.make_enzyme_module_form(
    "hex1_G_CI_c",  # CI stands for competitive inhibition
    name="automatic",
    categories="Inhibited",
    bound_metabolites={glc__D_c: 1, _23dpg_c: 1},
    compartment="c")

hex1_A_PI_c = HEX1.make_enzyme_module_form(
    "hex1_A_PI_c",  # PI stands for competitive inhibition
    name="automatic",
    categories="Inhibited",
    bound_metabolites={adp_c: 1},
    compartment="c")

hex1_G_PI_c = HEX1.make_enzyme_module_form(
    "hex1_G_PI_c",  # PI stands for competitive inhibition
    name="automatic",
    categories="Inhibited",
    bound_metabolites={g6p_c: 1},
    compartment="c")

HEX1.enzyme_module_forms
```

    [<EnzymeModuleForm hex1_c at 0x7fc2208ccb50>,
     <EnzymeModuleForm hex1_A_c at 0x7fc2208ccbb0>,
     <EnzymeModuleForm hex1_G_c at 0x7fc2208ccf40>,
     <EnzymeModuleForm hex1_AG_c at 0x7fc2208ccd30>,
     <EnzymeModuleForm hex1_G_CI_c at 0x7fc2208ccee0>,
     <EnzymeModuleForm hex1_A_PI_c at 0x7fc2208ccaf0>,
     <EnzymeModuleForm hex1_G_PI_c at 0x7fc2208ccb80>]

The `bound_metabolites` attribute represents the ligands bound to the site(s) of enzyme.

```python
# Print automatically generated names
for enzyme_form in HEX1.enzyme_module_forms:
    print("Bound to sites of {0}:\n{1}\n".format(
        enzyme_form.id, {
            ligand.id: coeff
            for ligand, coeff in enzyme_form.bound_metabolites.items()}))
```

    Bound to sites of hex1_c:
    {}

    Bound to sites of hex1_A_c:
    {'atp_c': 1}

    Bound to sites of hex1_G_c:
    {'glc__D_c': 1}

    Bound to sites of hex1_AG_c:
    {'glc__D_c': 1, 'atp_c': 1}

    Bound to sites of hex1_G_CI_c:
    {'glc__D_c': 1, '_23dpg_c': 1}

    Bound to sites of hex1_A_PI_c:
    {'adp_c': 1}

    Bound to sites of hex1_G_PI_c:
    {'g6p_c': 1}

Setting the `bound_metabolites` attribute upon creation allow the `formula` and `charge` attributes of the various forms also to be set while ensuring mass and charge balancing is maintained. Note that the enzyme is represented as a moiety, and the ligands bound to the enzyme are represented in the chemical formula.

```python
# Get the elemental matrix for the enzyme
df = HEX1.get_elemental_matrix(array_type="DataFrame")
# Use iloc to only look at EnzymeModuleForms
df.iloc[:, 6:]
```

      hex1_c
      hex1_A_c
      hex1_G_c
      hex1_AG_c
      hex1_G_CI_c
      hex1_A_PI_c
      hex1_G_PI_c

      C
      0.0
      10.0
      6.0
      16.0
      9.0
      10.0
      6.0

      H
      0.0
      12.0
      12.0
      24.0
      15.0
      12.0
      11.0

      O
      0.0
      13.0
      6.0
      19.0
      16.0
      10.0
      9.0

      P
      0.0
      3.0
      0.0
      3.0
      2.0
      2.0
      1.0

      N
      0.0
      5.0
      0.0
      5.0
      0.0
      5.0
      0.0

      S
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0

      q
      0.0
      -4.0
      0.0
      -4.0
      -5.0
      -3.0
      -2.0

      [HEX]
      1.0
      1.0
      1.0
      1.0
      1.0
      1.0
      1.0

Setting the `name` argument as "automatic" in the `EnzymeModule.make_enzyme_module_form()` method causes a name for the `EnzymeModuleForm` to be generated based on the metabolites in the `bound_metabolites` attribute.

```python
# Print automatically generated names
for enzyme_form in HEX1.enzyme_module_forms:
    print(enzyme_form.name)
```

    HEX1
    HEX1-atp complex
    HEX1-glc__D complex
    HEX1-glc__D-atp complex
    HEX1-glc__D-_23dpg complex
    HEX1-adp complex
    HEX1-g6p complex

The `categories` argument allows for `EnzymeModuleForm` objects to be placed into `cobra.Group` objects representing those categories. As with the ligands, the categorized enzyme module forms are returned in a `DictList` of `Group` objects by the `enzyme_module_forms_categorized` attribute.

```python
for group in HEX1.enzyme_module_forms_categorized:
    print("{0}: {1}".format(
        group.id, str(sorted([m.id for m in group.members]))))
```

    Active: ['hex1_AG_c', 'hex1_A_c', 'hex1_G_c', 'hex1_c']
    Inhibited: ['hex1_A_PI_c', 'hex1_G_CI_c', 'hex1_G_PI_c']

Alternatively, the `enzyme_module_forms_categorized` attribute can be set using a `dict`:

```python
HEX1.enzyme_module_forms_categorized =  {
    "competitively_inhibited": hex1_G_CI_c}

for group in HEX1.enzyme_module_forms_categorized:
    print("{0}: {1}".format(
        group.id, str(sorted([m.id for m in group.members]))))
```

    Active: ['hex1_AG_c', 'hex1_A_c', 'hex1_G_c', 'hex1_c']
    Inhibited: ['hex1_A_PI_c', 'hex1_G_CI_c', 'hex1_G_PI_c']
    competitively_inhibited: ['hex1_G_CI_c']

### Defining enzyme module reactions

The next step is to define all of the reaction steps that represent the catalytic mechanism and regulation of the enzyme module. These reactions are represented as `EnzymeModuleReaction` objects.

The `EnzymeModuleReaction` object inherits from the `MassReaction` and is treated like any other reaction in the model. Like the `make_enzyme_module_form()` method, the `make_enzyme_module_reaction()` method allows for the creation of an `EnzymeModuleReaction` object while assigning categories for the `EnzymeModuleReaction` in the process.

Species that exist in the model can also be added to the reaction by providing a dictionary of metabolites and their stoichiometric coefficients to the `metabolites_to_add` argument.

```python
HEX1_1 = HEX1.make_enzyme_module_reaction(
    "HEX1_1",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="product_inhibition",
    metabolites_to_add={
        "hex1_c": -1,
        "adp_c": -1,
        "hex1_A_PI_c": 1})

HEX1_2 = HEX1.make_enzyme_module_reaction(
    "HEX1_2",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="product_inhibition",
    metabolites_to_add={
        "hex1_c": -1,
        "g6p_c": -1,
        "hex1_G_PI_c": 1})

HEX1_3 = HEX1.make_enzyme_module_reaction(
    "HEX1_3",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="glc__D_c_binding",
    metabolites_to_add={
        "hex1_c": -1,
        "glc__D_c": -1,
        "hex1_G_c": 1})

HEX1_4 = HEX1.make_enzyme_module_reaction(
    "HEX1_4",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="atp_c_binding",
    metabolites_to_add={
        "hex1_c": -1,
        "atp_c": -1,
        "hex1_A_c": 1})

HEX1_5 = HEX1.make_enzyme_module_reaction(
    "HEX1_5",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="competitive_inhibition",
    metabolites_to_add={
        "hex1_G_c": -1,
        "_23dpg_c": -1,
        "hex1_G_CI_c": 1})

HEX1_6 = HEX1.make_enzyme_module_reaction(
    "HEX1_6",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="atp_c_binding",
    metabolites_to_add={
        "hex1_G_c": -1,
        "atp_c": -1,
        "hex1_AG_c": 1})

HEX1_7 = HEX1.make_enzyme_module_reaction(
    "HEX1_7",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="glc__D_c_binding",
    metabolites_to_add={
        "hex1_A_c": -1,
        "glc__D_c": -1,
        "hex1_AG_c": 1})

HEX1_8 = HEX1.make_enzyme_module_reaction(
    "HEX1_8",
    name="Automatic",
    subsystem="Glycolysis",
    reversible=True,
    categories="catalyzation",
    metabolites_to_add={
        "hex1_AG_c": -1,
        "hex1_c": 1,
        "adp_c": 1,
        "g6p_c": 1,
        "h_c": 1})

for reaction in HEX1.enzyme_module_reactions:
    print(reaction)
```

    HEX1_1: adp_c + hex1_c <=> hex1_A_PI_c
    HEX1_2: g6p_c + hex1_c <=> hex1_G_PI_c
    HEX1_3: glc__D_c + hex1_c <=> hex1_G_c
    HEX1_4: atp_c + hex1_c <=> hex1_A_c
    HEX1_5: _23dpg_c + hex1_G_c <=> hex1_G_CI_c
    HEX1_6: atp_c + hex1_G_c <=> hex1_AG_c
    HEX1_7: glc__D_c + hex1_A_c <=> hex1_AG_c
    HEX1_8: hex1_AG_c <=> adp_c + g6p_c + h_c + hex1_c

The `categories` argument allows for `EnzymeModuleReactions` objects to be placed into `cobra.Group` objects representing those categories. As with the ligands and enzyme forms, a `DictList` of the relevant groups are returned with the `enzyme_module_reactions_categorized` attribute.

```python
HEX1.enzyme_module_reactions_categorized
```

    [<Group product_inhibition at 0x7fc1e10e3f40>,
     <Group glc__D_c_binding at 0x7fc2026703a0>,
     <Group atp_c_binding at 0x7fc1e10babe0>,
     <Group competitive_inhibition at 0x7fc1e10baf40>,
     <Group catalyzation at 0x7fc1e10c8fa0>]

#### Unifying rate parameters

For this `EnzymeModule`, the reactions representing glucose binding to the enzyme and ATP binding to the enzyme have the same forward rate and equilibrium constants. Instead of defining the parameter values for each individual reaction, the `unify_rate_parameters()` method can be used to create custom rate laws for the given reactions that all depend on the same rate parameters.

The `unify_rate_parameters()` method takes a list of reactions and an identifier to use for the unified parameter. The `enzyme_prefix` flag can be set to `True` to prefix the new parameter identifier with the identifier of the `EnzymeModule`, ensuring that any existing custom parameters are not overwritten.

```python
for ligand, pid in zip([glc__D_c, atp_c],["G", "A"]):
    # Get the group of reactions corresponding to the ligand
    category = "_".join((ligand.id, "binding"))
    group = HEX1.enzyme_module_reactions_categorized.get_by_id(category)

    # Unify the parameters
    HEX1.unify_rate_parameters(
        group.members, new_parameter_id=pid, enzyme_prefix=True)

    # Print the new reaction rates
    print("\n" + category + "\n" + "-" * len(category))
    for reaction in sorted(group.members, key=attrgetter("id")):
        print(reaction.id + ": " + str(reaction.rate))
```

    glc__D_c_binding
    ----------------
    HEX1_3: kf_HEX1_G*(glc__D_c(t)*hex1_c(t) - hex1_G_c(t)/Keq_HEX1_G)
    HEX1_7: kf_HEX1_G*(glc__D_c(t)*hex1_A_c(t) - hex1_AG_c(t)/Keq_HEX1_G)

    atp_c_binding
    -------------
    HEX1_4: kf_HEX1_A*(atp_c(t)*hex1_c(t) - hex1_A_c(t)/Keq_HEX1_A)
    HEX1_6: kf_HEX1_A*(atp_c(t)*hex1_G_c(t) - hex1_AG_c(t)/Keq_HEX1_A)

## Determining Enzyme Form Concentrations and Rate Constants

The next step is to solve for the steady state concentrations for the various forms of the enzyme symbolically using **SymPy**. Because the numerical values for the dissociation constants have been defined, these equations are solved in terms of the rate constants. The rate constants can be approximated using the total enzyme concentration as a constraint and substituted back into the equations to calculate the numerical values of the steady state concentrations.

```python
from sympy import Eq, Symbol, lambdify, simplify, solveset

from mass import strip_time
from mass.util.matrix import matrix_rank
```

### Solving steady state concentrations symbolically

To get the symbolic solutions for the individual enzyme forms, the ODEs are first collected in a `dict`. Keys are the enzyme forms, and values are their ODEs with the time dependency stripped via the `strip_time` function.

```python
ode_dict = {
    enzyme_form.id: Eq(strip_time(enzyme_form.ode), 0)
    for enzyme_form in HEX1.enzyme_module_forms}
# Matrix rank of enzyme stoichiometric matrix without substrates
rank = matrix_rank(HEX1.S[6:])
print("Rank Deficiency: {0}".format(len(ode_dict) - rank))
```

    Rank Deficiency: 1

Because the stoichiometric matrix (without ligands) has a rank deficiency of one, there is a dependent variable in the system unless another equation is added. Therefore, the completely free enzyme form is treated as the dependent variable, and all of the enzyme forms are solved in terms of the free enzyme form.

```python
enzyme_solutions = {}
for enzyme_form in HEX1.enzyme_module_forms:
    # Skip dependent variable
    if enzyme_form.id == "hex1_c":
        continue
    # Get the ODE for the enzyme form from the ODE dict
    equation = ode_dict[enzyme_form.id]
    # Solve the equation for the enzyme form, substituting
    # previously found enzyme form solutions into the equation
    solution = solveset(equation.subs(enzyme_solutions),
                        enzyme_form.id)
    # Store the solution
    enzyme_solutions[enzyme_form.id] = list(solution)[0]
    # Substitute the new solution into existing solutions
    enzyme_solutions.update({
        enzyme_form: sol.subs(enzyme_solutions)
        for enzyme_form, sol in enzyme_solutions.items()})

args = set()
for solution in enzyme_solutions.values():
    args.update(solution.atoms(Symbol))
```

#### Defining the Rate Equation

To make up for the rank deficiency, an additional equation is needed. Typically, the rate of the enzyme is the summation of the rates for the catalyzation reaction step(s) of the enzyme. The `make_enzyme_rate_equation()` method can be used to create the rate equation from a list of reactions. If `use_rates=True`, the rate expressions of the reactions are added together. If `update_enzyme=True`, the rate equation is set as a symbolic expression for the `enzyme_rate_equation` attribute.

```python
# Get the catalyzation reactions
catalyzation_group = HEX1.enzyme_module_reactions_categorized.get_by_id(
    "catalyzation")

HEX1.make_enzyme_rate_equation(catalyzation_group.members,
                               use_rates=True,
                               update_enzyme=True)

print(HEX1.enzyme_rate_equation)
```

    kf_HEX1_8*(Keq_HEX1_8*hex1_AG_c(t) - adp_c(t)*g6p_c(t)*hex1_c(t))/Keq_HEX1_8

With the rate equation defined, the `enzyme_rate_error()` method is used to get the equation as the difference between the flux value and the rate equation.

```python
enzyme_rate_equation = strip_time(HEX1.enzyme_rate_error(use_values=False))
print(enzyme_rate_equation)
```

    v_HEX1 - kf_HEX1_8*(Keq_HEX1_8*hex1_AG_c - adp_c*g6p_c*hex1_c)/Keq_HEX1_8

The solutions for the enzyme forms are substituted into the rate equation, and the equation is solved for the free enzyme form. The solutions are subsequently updated, resulting in symbolic equations that do not depend on any enzyme form.

```python
# Solve for last unknown concentration symbolically
solution = solveset(enzyme_rate_equation.subs(enzyme_solutions),
                    "hex1_c")

# Update solution dictionary with the new solution
enzyme_solutions["hex1_c"] = list(solution)[0]

# Update solutions with free variable solutions
enzyme_solutions = {
    enzyme_form: simplify(solution.subs(enzyme_solutions))
    for enzyme_form, solution in enzyme_solutions.items()}

args = set()
for solution in enzyme_solutions.values():
    args.update(solution.atoms(Symbol))
print(args)
```

    {Keq_HEX1_5, adp_c, kf_HEX1_8, _23dpg_c, g6p_c, Keq_HEX1_2, Keq_HEX1_8, kf_HEX1_G, v_HEX1, Keq_HEX1_A, Keq_HEX1_1, Keq_HEX1_G, glc__D_c, atp_c, kf_HEX1_A}

Numerical values for known quantities are substituted into the equations. For this `EnzymeModule` of Hexokinase, the following dissociation constants are used:

$$\begin{align}
K_{d,\ \text{GLC-D}} &= 0.038\ \text{mM} \\
K_{d,\ \text{ATP}} &= 2.06\ \text{mM} \\
K_{i,\ \text{23DPG}} &= 5.5\ \text{mM} \\
K_{i,\ \text{ADP}} &= 1\ \text{mM} \\
K_{i,\ \text{G6P}} &= 66.67\ \text{mM} \\
\end{align}$$

A value of $K_{\text{HEX1}}= 313.12$ is used for the catalyzation step. Note that the inverse of the dissociation constant is used for reactions that form complexes.

```python
numerical_values = {
    "Keq_HEX1_1": 1,
    "Keq_HEX1_2": 1 / 66.67,
    "Keq_HEX1_G": 1 / 0.038,
    "Keq_HEX1_A": 1 / 2.06,
    "Keq_HEX1_5": 1 / 5.5,
    "Keq_HEX1_8": 313.12}
# Update the model with the parameters
HEX1.update_parameters(numerical_values)
```

The ligand concentrations and the rate for the enzyme are extracted from the merged glycolysis and hemoglobin model.

```python
# Get steady state flux for EnzymeModule
HEX1.enzyme_rate = glyc_hb.reactions.get_by_id("HEX1").steady_state_flux
numerical_values[HEX1.enzyme_flux_symbol_str] = HEX1.enzyme_rate

# Get the ligand concentrations
for met in HEX1.enzyme_module_ligands:
    concentration = glyc_hb.metabolites.get_by_id(met.id).initial_condition
    # Set the ligand initial condition and add to numercal values dictionary
    met.initial_condition = concentration
    numerical_values[met.id] = concentration
```

The numerical values are substituted into the symbolic equations, resulting in the steady state concentrations that depend only on the rate constants.

```python
enzyme_solutions = {
    enzyme_form: simplify(sol.subs(numerical_values))
    for enzyme_form, sol in enzyme_solutions.items()}

args = set()
for solution in enzyme_solutions.values():
    args.update(solution.atoms(Symbol))
print(args)
```

    {kf_HEX1_G, kf_HEX1_8, kf_HEX1_A}

### Approximating Rate Constants

To determine the set of rate constants for the enzyme module, the absolute error between the total hexokinase concentration value (found in literature) and the computed hexokinase concentration is minimized. For this example, the `minimize()` function of the **SciPy** package is utilized to find a feasible set of rate constants.

```python
from scipy.optimize import minimize
```

The objective function for the minimization is first made symbolically. The `enzyme_total_symbol_str` property can be used to represent the total enzyme concentration, while the `enzyme_concentration_total_equation` property creates a symbolic expression for the sum of all enzyme forms.

```python
enzyme_total_error = abs(
    Symbol(HEX1.enzyme_total_symbol_str)
    - strip_time(HEX1.enzyme_concentration_total_equation))
print(enzyme_total_error)
```

    Abs(-HEX1_Total + hex1_AG_c + hex1_A_PI_c + hex1_A_c + hex1_G_CI_c + hex1_G_PI_c + hex1_G_c + hex1_c)

The `enzyme_concentration_total` attribute stores the total amount of enzyme in the model and substituted into the expression. The total HEX1 concentration is $24 * 10^{-6} \text{mM}$.

```python
HEX1.enzyme_concentration_total = 24e-6
enzyme_total_error = enzyme_total_error.subs({
    HEX1.enzyme_total_symbol_str: HEX1.enzyme_concentration_total})
print(enzyme_total_error)
```

    Abs(hex1_AG_c + hex1_A_PI_c + hex1_A_c + hex1_G_CI_c + hex1_G_PI_c + hex1_G_c + hex1_c - 2.4e-5)

Finally, the symbolic equations for the enzyme forms are substituted into the enzyme total error equation, resulting in an expression that represents the objective function with the only unknown variables being rate constants. The `lambdify()` function of the **SymPy** package converts the symbolic objective into a lambda function that can be used with the `minimize()` function of **SciPy**.

```python
enzyme_total_error = simplify(enzyme_total_error.subs(enzyme_solutions))

# Sort the arguments to ensure input format remains consistent
args = sorted(list(map(str, args)))
# Use lambdify to make objective function as a lambda function
obj_fun = lambda x: lambdify(args, enzyme_total_error)(*x)
```

The `minimize()` function is now used to approximate the rate constants. The optimization problems for enzyme rate constants are typically nonlinear, and require nonlinear optimization routines to find feasible solutions.

```python
# Minimize the objective function, initial guess based on publication values
initial_guess = [1e8, 9376585, 52001]
variable_bounds = ((0, 1e9), (0, 1e9), (0, 1e9))
solution = minimize(obj_fun, x0=initial_guess,
                    method="trust-constr",
                    bounds=variable_bounds)
# Map solution array to variables
rate_constants = dict(zip(args, solution.x))
print(rate_constants)
```

    {'kf_HEX1_8': 100000000.0025878, 'kf_HEX1_A': 9376585.030755484, 'kf_HEX1_G': 52006.59981223971}

Because the rate constants associated with the inhibition of the enzyme forms are not necessary for computing the concentrations, a rapid binding assumption is made for the inhibition reactions. Therefore, a large number is set for the rate constants. The parameters are set using the `update_parameters()` method.

```python
rate_constants["kf_HEX1_1"] = 1e6
rate_constants["kf_HEX1_2"] = 1e6
rate_constants["kf_HEX1_5"] = 1e6
HEX1.update_parameters(rate_constants)
```

### Calculating numerical values for concentrations

Once the rate constants have been estimated, they are substituted back into the symbolic concentration equations in order to obtain their numerical values.

```python
for enzyme_form, solution in enzyme_solutions.items():
    # Get the enzyme form object, determine the steady state concentration
    enzyme_form = HEX1.enzyme_module_forms.get_by_id(enzyme_form)
    enzyme_form.initial_condition = float(solution.subs(rate_constants))
    print("{0}: {1:e}".format(enzyme_form.id,
                              enzyme_form.initial_condition))
```

    hex1_A_c: 9.401421e-06
    hex1_G_c: 5.718872e-08
    hex1_AG_c: 1.174630e-08
    hex1_G_CI_c: 3.223364e-08
    hex1_A_PI_c: 3.519706e-06
    hex1_G_PI_c: 8.847367e-09
    hex1_c: 1.213692e-05

#### Error values

As a quality assurance check, the `enzyme_concentration_total_error()` method can be used to get the error between the `enzyme_concentration_total` attribute and the sum of the enzyme form concentrations. A positive value indicates the `enzyme_concentration_total` attribute is greater than the sum of the individual enzyme form concentrations that were computed.

```python
print("Total Enzyme Concentration Error: {0}".format(
    HEX1.enzyme_concentration_total_error(use_values=True)))
```

    Total Enzyme Concentration Error: -1.1680622689093244e-06

Similarly, the error between the `enzyme_rate` attribute and the computed value from the `enzyme_rate_equation` can be also checked using the `enzyme_rate_error()` method, in which a positive value indicates that the `enzyme_rate` attribute is greater than the value computed when using the rate equation.

```python
print("Enzyme Rate Error: {0}".format(
    HEX1.enzyme_rate_error(use_values=True)))
```

    Enzyme Rate Error: 4.440892098500626e-16

## Adding EnzymeModules to Models

With the `EnzymeModule` built, it can be integrated into a larger network and simulated. To add an `EnzymeModule` to an existing `MassModel`, the `merge()` method is used. After merging, the `remove_reactions()` method is used to remove the reaction replaced with the enzyme module. The `EnzymeModule` should always be merged into the `MassModel` as demonstrated below:

```python
glyc_hb_HEX1 = glyc_hb.merge(HEX1, inplace=False)
glyc_hb_HEX1.remove_reactions([
    glyc_hb_HEX1.reactions.get_by_id("HEX1")])
```

All objects, numerical values, and certain attributes of the `EnzymeModule` are transferred into the `MassModel` upon merging. This includes all enzyme forms, reactions steps, initial conditions, rate parameters, and category groups.

```python
glyc_hb_HEX1
```

        Name	Glycolysis_Hemoglobin_HEX1

        Memory address	0x07fc1e114ce50

        Stoichiometric Matrix
        35x37

        Matrix Rank
        32

        Number of metabolites
        35

        Initial conditions defined
        35/35

        Number of reactions
        37

        Number of genes
        0

        Number of enzyme modules
        1

        Number of groups
        12

        Objective expression
        0

        Compartments
        Cytosol

### The EnzymeModuleDict object

During the merge process, an `EnzymeModuleDict` is created from the `EnzymeModule` and added to the `MassModel.enzyme_modules` attribute.

```python
print(glyc_hb_HEX1.enzyme_modules)
HEX1_dict = glyc_hb_HEX1.enzyme_modules.get_by_id("HEX1")
HEX1_dict
```

    [<EnzymeModuleDict HEX1 at 0x7fc1f08e3dc0>]

        Name	HEX1

        Memory address	0x07fc1f08e3dc0

        Stoichiometric Matrix
        13x8

        Matrix Rank
        7

        Subsystem
        Glycolysis

        Number of Ligands
        6

        Number of EnzymeForms
        7

        Number of EnzymeModuleReactions
        8

        Enzyme Concentration Total
        2.4e-05

        Enzyme Net Flux
        1.12

The `EnzymeModuleDict` inherits from an `OrderedDict`, thereby inheriting ordered dictionary methods such as `keys()`:

```python
print("\n".join(HEX1_dict.keys()))
```

    id
    name
    subsystem
    enzyme_module_ligands
    enzyme_module_forms
    enzyme_module_reactions
    enzyme_module_ligands_categorized
    enzyme_module_forms_categorized
    enzyme_module_reactions_categorized
    enzyme_concentration_total
    enzyme_rate
    enzyme_concentration_total_equation
    enzyme_rate_equation
    S
    model

The `EnzymeModuleDict` stores several of the enzyme-specific attributes so that they are still accessible after integrating the enzyme module into a larger network. The keys of the `EnzymeModuleDict`  also can be treated as attribute accessors:

```python
print("Enzyme Rate:\n{0} = {1}".format(
    HEX1_dict["enzyme_rate"],       # Returned using dict key
    HEX1_dict.enzyme_rate_equation  # Returned using attribute accessor
))
```

    Enzyme Rate:
    1.12 = kf_HEX1_8*(Keq_HEX1_8*hex1_AG_c(t) - adp_c(t)*g6p_c(t)*hex1_c(t))/Keq_HEX1_8

### Steady State Validation
The last step is to ensure that a steady state is reached with the completed enzyme module within a larger network context.

```python
import matplotlib.pyplot as plt

from mass import Simulation
from mass.visualization import plot_time_profile
```

Here, the model is simulated, and the enzyme's ability to reach steady state is graphically verified:

```python
# Setup simulation object
sim = Simulation(glyc_hb_HEX1, verbose=True)
# Simulate from 0 to 1000 with 10001 points in the output
conc_sol, flux_sol = sim.simulate(
    glyc_hb_HEX1, time=(0, 1e3))

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 4))
plot_time_profile(
    conc_sol, observable=HEX1_dict.enzyme_module_forms, ax=ax,
    xlim=(1e-3, 1e3),
    legend="right outside", plot_function="loglog",
    xlabel="Time [hr]", ylabel="Concentration [mM]",
    title="TIme profile of Concentrations for Enzyme Forms");
```

The plot shows that the enzyme can reach a steady state when integrated into a larger network, meaning the enzyme module that represents hexokinase in this system is complete!

## Additional Examples
For additional examples of analyzing and visualizing systems with enzyme modules, see the following:

* [Visualizing Catalytic Potentials of Glycolytic Regulatory Kinases](../gallery/visualization/catalytic_potential_visualizations.ipynb)

$^1$ Procedure outlined in Du et al., 2016

$^2$ Hexokinase based on Yurkovich et al., 2018, Du et al., 2016, and Mulquiney and Kuchel, 1999.

$^3$ Glycolysis model based on Yurkovich et al., 2018 and Chapter 10 of Systems Biology: Simulation of Dynamic Network States


[](https://kaggle.com/kernels/welcome?src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W2D1_DeepLearning/W2D1_Tutorial3.ipynb)

# Tutorial 3: Building and Evaluating Normative Encoding Models
**Week 2, Day 1: Deep Learning**

**By Neuromatch Academy**

**Content creators**: Jorge A. Menendez, Yalda Mohsenzadeh, Carsen Stringer

**Conent reviewers**: Roozbeh Farhoodi, Madineh Sarvestani, Kshitij Dwivedi, Spiros Chavlis, Ella Batty, Michael Waskom

**Our 2021 Sponsors, including Presenting Sponsor Facebook Reality Labs**

---
# Tutorial Objectives

*Estimated timing of tutorial: 1 hour, 10 minutes*

In this tutorial, we'll be using deep learning to build an encoding model of the visual system, and then compare its internal representations to those observed in neural data.

Its parameters won't be directly optimized to fit the neural data. Instead, we will optimize its parameters to solve a particular visual task that we know the brain can solve. We therefore refer to it as a **"normative"** encoding model, since it is optimized for a specific behavioral task. It is the optimal model to solve the problem (optimal for the specified architecture).  See Section 3 of the bonus tutorial to fit a convolutional neural network directly to neural data (an alternative approach for investigating encoding with deep neural networks).

To then evaluate whether this normative encoding model is actually a good model of the brain, we'll analyze its internal representations and compare them to the representations observed in mouse primary visual cortex. Since we understand exactly what the encoding model's representations are optimized to do, any similarities will hopefully shed light on why the representations in the brain look the way they do.

More concretely, our goal will be learn how to:
* Visualize and analyze the internal representations of a deep network
* Quantify the similarity between distributed representations in a model and neural representations observed in recordings, using Representational Similarity Analysis (RSA)

```python
# @title Tutorial slides

# @markdown These are the slides for all videos in this tutorial.
from IPython.display import IFrame
IFrame(src=f"https://mfr.ca-1.osf.io/render?url=https://osf.io/kwyvp/?direct%26mode=render%26action=download%26mode=render", width=854, height=480)
```

---
# Setup

```python
# Imports

import numpy as np
from scipy.stats import zscore
import matplotlib as mpl
from matplotlib import pyplot as plt
import torch
from torch import nn, optim
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
```

```python
#@title Data retrieval and loading
import os
import hashlib
import requests

fname = "W3D4_stringer_oribinned1.npz"
url = "https://osf.io/683xc/download"
expected_md5 = "436599dfd8ebe6019f066c38aed20580"

if not os.path.isfile(fname):
  try:
    r = requests.get(url)
  except requests.ConnectionError:
    print("!!! Failed to download data !!!")
  else:
    if r.status_code != requests.codes.ok:
      print("!!! Failed to download data !!!")
    elif hashlib.md5(r.content).hexdigest() != expected_md5:
      print("!!! Data download appears corrupted !!!")
    else:
      with open(fname, "wb") as fid:
        fid.write(r.content)
```

```python
#@title Figure Settings

%matplotlib inline
%config InlineBackend.figure_format='retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Plotting Functions

def show_stimulus(img, ax=None):
  """Visualize a stimulus"""
  if ax is None:
    ax = plt.gca()
  ax.imshow(img, cmap=mpl.cm.binary)
  ax.set_aspect('auto')
  ax.set_xticks([])
  ax.set_yticks([])
  ax.spines['left'].set_visible(False)
  ax.spines['bottom'].set_visible(False)

def plot_corr_matrix(rdm, ax=None):
  """Plot dissimilarity matrix

  Args:
    rdm (numpy array): n_stimuli x n_stimuli representational dissimilarity
      matrix
    ax (matplotlib axes): axes onto which to plot

  Returns:
    nothing

  """
  if ax is None:
    ax = plt.gca()
  image = ax.imshow(rdm, vmin=0.0, vmax=2.0)
  ax.set_xticks([])
  ax.set_yticks([])
  cbar = plt.colorbar(image, ax=ax, label='dissimilarity')

def plot_multiple_rdm(rdm_dict):
  """Draw multiple subplots for each RDM in rdm_dict."""
  fig, axs = plt.subplots(1, len(rdm_dict),
                          figsize=(4 * len(resp_dict), 3.5))

  # Compute RDM's for each set of responses and plot
  for i, (label, rdm) in enumerate(rdm_dict.items()):

    image = plot_corr_matrix(rdm, axs[i])
    axs[i].set_title(label)

def plot_rdm_rdm_correlations(rdm_sim):
  """Draw a bar plot showing between-RDM correlations."""
  f, ax = plt.subplots()
  ax.bar(rdm_sim.keys(), rdm_sim.values())
  ax.set_xlabel('Deep network model layer')
  ax.set_ylabel('Correlation of model layer RDM\nwith mouse V1 RDM')

def plot_rdm_rows(ori_list, rdm_dict, rdm_oris):
  """Plot the dissimilarity of response to each stimulus with response to one
  specific stimulus

  Args:
    ori_list (list of float): plot dissimilarity with response to stimulus with
      orientations closest to each value in this list
    rdm_dict (dict): RDM's from which to extract dissimilarities
    rdm_oris (np.ndarray): orientations corresponding to each row/column of RDMs
    in rdm_dict

  """
  n_col = len(ori_list)
  f, axs = plt.subplots(1, n_col, figsize=(4 * n_col, 4), sharey=True)

  # Get index of orientation closest to ori_plot
  for ax, ori_plot in zip(axs, ori_list):
    iori = np.argmin(np.abs(rdm_oris - ori_plot))

    # Plot dissimilarity curves in each RDM
    for label, rdm in rdm_dict.items():
      ax.plot(rdm_oris, rdm[iori, :], label=label)

    # Draw vertical line at stimulus we are plotting dissimilarity w.r.t.
    ax.axvline(rdm_oris[iori], color=".7", zorder=-1)

    # Label axes
    ax.set_title(f'Dissimilarity with response\nto {ori_plot: .0f}$^o$ stimulus')
    ax.set_xlabel('Stimulus orientation ($^o$)')

  axs[0].set_ylabel('Dissimilarity')
  axs[-1].legend(loc="upper left", bbox_to_anchor=(1, 1))
```

```python
# @title Helper Functions

def load_data(data_name=fname, bin_width=1):
  """Load mouse V1 data from Stringer et al. (2019)

  Data from study reported in this preprint:
  https://www.biorxiv.org/content/10.1101/679324v2.abstract

  These data comprise time-averaged responses of ~20,000 neurons
  to ~4,000 stimulus gratings of different orientations, recorded
  through Calcium imaginge. The responses have been normalized by
  spontanous levels of activity and then z-scored over stimuli, so
  expect negative numbers. They have also been binned and averaged
  to each degree of orientation.

  This function returns the relevant data (neural responses and
  stimulus orientations) in a torch.Tensor of data type torch.float32
  in order to match the default data type for nn.Parameters in
  Google Colab.

  This function will actually average responses to stimuli with orientations
  falling within bins specified by the bin_width argument. This helps
  produce individual neural "responses" with smoother and more
  interpretable tuning curves.

  Args:
    bin_width (float): size of stimulus bins over which to average neural
      responses

  Returns:
    resp (torch.Tensor): n_stimuli x n_neurons matrix of neural responses,
        each row contains the responses of each neuron to a given stimulus.
        As mentioned above, neural "response" is actually an average over
        responses to stimuli with similar angles falling within specified bins.
    stimuli: (torch.Tensor): n_stimuli x 1 column vector with orientation
        of each stimulus, in degrees. This is actually the mean orientation
        of all stimuli in each bin.

  """
  with np.load(data_name) as dobj:
    data = dict(**dobj)
  resp = data['resp']
  stimuli = data['stimuli']

  if bin_width > 1:
    # Bin neural responses and stimuli
    bins = np.digitize(stimuli, np.arange(0, 360 + bin_width, bin_width))
    stimuli_binned = np.array([stimuli[bins == i].mean() for i in np.unique(bins)])
    resp_binned = np.array([resp[bins == i, :].mean(0) for i in np.unique(bins)])
  else:
    resp_binned = resp
    stimuli_binned = stimuli

  # only use stimuli <= 180
  resp_binned = resp_binned[stimuli_binned <= 180]
  stimuli_binned = stimuli_binned[stimuli_binned <= 180]

  stimuli_binned -= 90  # 0 means vertical, -ve means tilted left, +ve means tilted right

  # Return as torch.Tensor
  resp_tensor = torch.tensor(resp_binned, dtype=torch.float32)
  stimuli_tensor = torch.tensor(stimuli_binned, dtype=torch.float32).unsqueeze(1)  # add singleton dimension to make a column vector

  return resp_tensor, stimuli_tensor

def grating(angle, sf=1 / 28, res=0.1, patch=False):
  """Generate oriented grating stimulus

  Args:
    angle (float): orientation of grating (angle from vertical), in degrees
    sf (float): controls spatial frequency of the grating
    res (float): resolution of image. Smaller values will make the image
      smaller in terms of pixels. res=1.0 corresponds to 640 x 480 pixels.
    patch (boolean): set to True to make the grating a localized
      patch on the left side of the image. If False, then the
      grating occupies the full image.

  Returns:
    torch.Tensor: (res * 480) x (res * 640) pixel oriented grating image

  """

  angle = np.deg2rad(angle)  # transform to radians

  wpix, hpix = 640, 480  # width and height of image in pixels for res=1.0

  xx, yy = np.meshgrid(sf * np.arange(0, wpix * res) / res, sf * np.arange(0, hpix * res) / res)

  if patch:
    gratings = np.cos(xx * np.cos(angle + .1) + yy * np.sin(angle + .1))  # phase shift to make it better fit within patch
    gratings[gratings < 0] = 0
    gratings[gratings > 0] = 1
    xcent = gratings.shape[1] * .75
    ycent = gratings.shape[0] / 2
    xxc, yyc = np.meshgrid(np.arange(0, gratings.shape[1]), np.arange(0, gratings.shape[0]))
    icirc = ((xxc - xcent) ** 2 + (yyc - ycent) ** 2) ** 0.5 < wpix / 3 / 2 * res
    gratings[~icirc] = 0.5

  else:
    gratings = np.cos(xx * np.cos(angle) + yy * np.sin(angle))
    gratings[gratings < 0] = 0
    gratings[gratings > 0] = 1

  # Return torch tensor
  return torch.tensor(gratings, dtype=torch.float32)

def filters(out_channels=6, K=7):
  """ make example filters, some center-surround and gabors
  Returns:
      filters: out_channels x K x K
  """
  grid = np.linspace(-K/2, K/2, K).astype(np.float32)
  xx,yy = np.meshgrid(grid, grid, indexing='ij')

  # create center-surround filters
  sigma = 1.1
  gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*sigma**2))
  wide_gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*(sigma*2)**2))
  center_surround = gaussian - 0.5 * wide_gaussian

  # create gabor filters
  thetas = np.linspace(0, 180, out_channels-2+1)[:-1] * np.pi/180
  gabors = np.zeros((len(thetas), K, K), np.float32)
  lam = 10
  phi = np.pi/2
  gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*(sigma*0.4)**2))
  for i,theta in enumerate(thetas):
    x = xx*np.cos(theta) + yy*np.sin(theta)
    gabors[i] = gaussian * np.cos(2*np.pi*x/lam + phi)

  filters = np.concatenate((center_surround[np.newaxis,:,:],
                            -1*center_surround[np.newaxis,:,:],
                            gabors),
                           axis=0)
  filters /= np.abs(filters).max(axis=(1,2))[:,np.newaxis,np.newaxis]
  # convert to torch
  filters = torch.from_numpy(filters)
  # add channel axis
  filters = filters.unsqueeze(1)

  return filters

class CNN(nn.Module):
  """Deep convolutional network with one convolutional + pooling layer followed
  by one fully connected layer

  Args:
    h_in (int): height of input image, in pixels (i.e. number of rows)
    w_in (int): width of input image, in pixels (i.e. number of columns)

  Attributes:
    conv (nn.Conv2d): filter weights of convolutional layer
    pool (nn.MaxPool2d): max pooling layer
    dims (tuple of ints): dimensions of output from pool layer
    fc (nn.Linear): weights and biases of fully connected layer
    out (nn.Linear): weights and biases of output layer

  """

  def __init__(self, h_in, w_in):
    super().__init__()
    C_in = 1  # input stimuli have only 1 input channel
    C_out = 6  # number of output channels (i.e. of convolutional kernels to convolve the input with)
    K = 7  # size of each convolutional kernel
    Kpool = 8  # size of patches over which to pool
    self.conv = nn.Conv2d(C_in, C_out, kernel_size=K, padding=K//2)  # add padding to ensure that each channel has same dimensionality as input
    self.pool = nn.MaxPool2d(Kpool)
    self.dims = (C_out, h_in // Kpool, w_in // Kpool)  # dimensions of pool layer output
    self.fc = nn.Linear(np.prod(self.dims), 10)  # flattened pool output --> 10D representation
    self.out = nn.Linear(10, 1)  # 10D representation --> scalar
    self.conv.weight = nn.Parameter(filters(C_out, K))
    self.conv.bias = nn.Parameter(torch.zeros((C_out,), dtype=torch.float32))

  def forward(self, x):
    """Classify grating stimulus as tilted right or left

    Args:
      x (torch.Tensor): p x 48 x 64 tensor with pixel grayscale values for
          each of p stimulus images.

    Returns:
      torch.Tensor: p x 1 tensor with network outputs for each input provided
          in x. Each output should be interpreted as the probability of the
          corresponding stimulus being tilted right.

    """
    x = x.unsqueeze(1)  # p x 1 x 48 x 64, add a singleton dimension for the single stimulus channel
    x = torch.relu(self.conv(x))  # output of convolutional layer
    x = self.pool(x)  # output of pooling layer
    x = x.view(-1, np.prod(self.dims))  # flatten pooling layer outputs into a vector
    x = torch.relu(self.fc(x))  # output of fully connected layer
    x = torch.sigmoid(self.out(x))  # network output
    return x

def train(net, train_data, train_labels,
          n_epochs=25, learning_rate=0.0005,
          batch_size=100, momentum=.99):
  """Run stochastic gradient descent on binary cross-entropy loss for a given
  deep network (cf. appendix for details)

  Args:
    net (nn.Module): deep network whose parameters to optimize with SGD
    train_data (torch.Tensor): n_train x h x w tensor with stimulus gratings
    train_labels (torch.Tensor): n_train x 1 tensor with true tilt of each
      stimulus grating in train_data, i.e. 1. for right, 0. for left
    n_epochs (int): number of times to run SGD through whole training data set
    batch_size (int): number of training data samples in each mini-batch
    learning_rate (float): learning rate to use for SGD updates
    momentum (float): momentum parameter for SGD updates

  """

  # Initialize binary cross-entropy loss function
  loss_fn = nn.BCELoss()

  # Initialize SGD optimizer with momentum
  optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=momentum)

  # Placeholder to save loss at each iteration
  track_loss = []

  # Loop over epochs
  for i in range(n_epochs):

    # Split up training data into random non-overlapping mini-batches
    ishuffle = torch.randperm(train_data.shape[0])  # random ordering of training data
    minibatch_data = torch.split(train_data[ishuffle], batch_size)  # split train_data into minibatches
    minibatch_labels = torch.split(train_labels[ishuffle], batch_size)  # split train_labels into minibatches

    # Loop over mini-batches
    for stimuli, tilt in zip(minibatch_data, minibatch_labels):

      # Evaluate loss and update network weights
      out = net(stimuli)  # predicted probability of tilt right
      loss = loss_fn(out, tilt)  # evaluate loss
      optimizer.zero_grad()  # clear gradients
      loss.backward()  # compute gradients
      optimizer.step()  # update weights

      # Keep track of loss at each iteration
      track_loss.append(loss.item())

    # Track progress
    if (i + 1) % (n_epochs // 5) == 0:
      print(f'epoch {i + 1} | loss on last mini-batch: {loss.item(): .2e}')

  print('training done!')

def get_hidden_activity(net, stimuli, layer_labels):
  """Retrieve internal representations of network

  Args:
    net (nn.Module): deep network
    stimuli (torch.Tensor): p x 48 x 64 tensor with stimuli for which to
      compute and retrieve internal representations
    layer_labels (list): list of strings with labels of each layer for which
      to return its internal representations

  Returns:
    dict: internal representations at each layer of the network, in
      numpy arrays. The keys of this dict are the strings in layer_labels.

  """

  # Placeholder
  hidden_activity = {}

  # Attach 'hooks' to each layer of the network to store hidden
  # representations in hidden_activity
  def hook(module, input, output):
    module_label = list(net._modules.keys())[np.argwhere([module == m for m in net._modules.values()])[0, 0]]
    if module_label in layer_labels:  # ignore output layer
      hidden_activity[module_label] = output.view(stimuli.shape[0], -1).detach().numpy()
  hooks = [layer.register_forward_hook(hook) for layer in net.children()]

  # Run stimuli through the network
  pred = net(stimuli)

  # Remove the hooks
  [h.remove() for h in hooks]

  return hidden_activity
```

---
# Section 1: Setting up deep network and neural data

In the future sections, we will compare the activity in a deep network, specifically in a CNN, with neural activity. First, we need to understand the task we are using (Section 1.1), train our deep network (Section 1.2), and load in neural data (Section 1.3).

```python
# @title Video 1: Deep convolutional network for orientation discrimination
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV1ip4y1i7Yo", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="KlXtKJCpV4I", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

## Section 1.1: Orientation discrimination task

We will build our normative encoding model by optimizing its parameters to solve an orientation discrimination task.

The task is to tell whether a given grating stimulus is tilted to the "right" or "left"; that is, whether its angle relative to the vertical is positive or negative, respectively. We show example stimuli below, which were constructed using the helper function `grating()`.

Note that this is a task that we know many mammalian visual systems are capable of solving. It is therefore conceivable that the representations in a deep network model optimized for this task might resemble those in the brain. To test this hypothesis, we will compare the representations of our optimized encoding model to neural activity recorded in response to these very same stimuli, courtesy of [Stringer et al 2019](https://www.biorxiv.org/content/10.1101/679324v2.abstract).

```python
#@title
#@markdown Execute this cell to plot example stimuli

orientations = np.linspace(-90, 90, 5)

h = 3
n_col = len(orientations)
fig, axs = plt.subplots(1, n_col, figsize=(h * n_col, h))

h, w  = grating(0).shape  # height and width of stimulus
print('stimulus size: %i x %i' % (h, w))

for i, ori in enumerate(orientations):
  stimulus = grating(ori)
  axs[i].set_title(f'{ori: .0f}$^o$')
  show_stimulus(stimulus, axs[i])
```

## Section 1.2: A deep network model of orientation discrimination

*Estimated timing to here from start of tutorial: 10 min*

Our goal is to build a model that solves the orientation discrimination task outlined above. The model should take as input a stimulus image and output the probability of that stimulus being tilted right.

To do this, we will use a **convolutional neural network (CNN)**, which is the type of network we saw in Tutorial 2. Here, we will use a CNN that performs *two-dimensional* convolutions on the raw stimulus image (which is a 2D matrix of pixels), rather than *one-dimensional* convolutions on a categorical 1D vector representation of the stimulus. CNNs are commonly used for image processing.

The particular CNN we will use here has two layers:
1. a *convolutional layer*, which convolves the images with a set of filters
2. a *fully connected layer*, which transforms the output of this convolution into a 10-dimensional representation

Finally, a set of output weights transforms this 10-dimensional representation into a single scalar $p$, denoting the predicted probability of the input stimulus being tilted right.

See [bonus section 1](#b1) for in-depth instructions for how to code up such a network in PyTorch. For now, however, we'll leave these details aside and focus on training this network, called `CNN`, and analyzing its internal representations.

Run the next cell to train such a network to solve this task. After initializing our CNN model, it builds a dataset of oriented grating stimuli to use for training it. These are then passed into a function called `train()` that uses SGD to optimize the model's parameters, taking similar arguments as the `train()` function we wrote in Tutorial 1.

Note that it may take ~30 seconds for the training to complete.

```python
help(train)
```

```python
# Set random seeds for reproducibility
np.random.seed(12)
torch.manual_seed(12)

# Initialize CNN model
net = CNN(h, w)

# Build training set to train it on
n_train = 1000  # size of training set

# sample n_train random orientations between -90 and +90 degrees
ori = (np.random.rand(n_train) - 0.5) * 180

# build orientated grating stimuli
stimuli = torch.stack([grating(i) for i in ori])

# stimulus tilt: 1. if tilted right, 0. if tilted left, as a column vector
tilt = torch.tensor(ori > 0).type(torch.float).unsqueeze(-1)

# Train model
train(net, stimuli, tilt)
```

## Section 1.3: Load data

*Estimated timing to here from start of tutorial: 15 min*

In the next cell, we provide code for loading in some data from [this paper](https://www.biorxiv.org/content/10.1101/679324v2.abstract), which contains the responses of about ~20,000 neurons in mouse primary visual cortex to grating stimuli like those used to train our network (this is the same data used in Tutorial 1). These data are stored in two variables:
* `resp_v1` is a matrix where each row contains the responses of all neurons to a single stimulus.
* `ori` is a vector with the orientations of each stimulus, in degrees. As in the above convention, negative angles denote stimuli tilted to the left and positive angles denote stimuli tilted to the right.

We will then extract our deep CNN model's representations of these same stimuli (i.e. oriented gratings with the orientations in `ori`). We will run the same stimuli through our CNN model and use the helper function `get_hidden_activity()` to store the model's internal representations. The output of this function is a Python `dict`, which contains a matrix of population responses (just like `resp_v1`) for each layer of the network specified by the `layer_labels` argument. We'll focus on looking at the representations in
* the output of the first convolutional layer, stored in the model as `'pool'` (see Bonus Section 1 for the details of the CNN architecture to understand why it's called this way)
* the 10-dimensional output of the fully connected layer, stored in the model as `'fc'`

```python
# Load mouse V1 data
resp_v1, ori = load_data()

# Extract model internal representations of each stimulus in the V1 data
# construct grating stimuli for each orientation presented in the V1 data
stimuli = torch.stack([grating(a.item()) for a in ori])
layer_labels = ['pool', 'fc']
resp_model = get_hidden_activity(net, stimuli, layer_labels)

# Aggregate all responses into one dict
resp_dict = {}
resp_dict['V1 data'] = resp_v1
for k, v in resp_model.items():
  label = f"model\n'{k}' layer"
  resp_dict[label] = v
```

---
# Section 2: Quantitative comparisons of CNNs and neural activity

*Estimated timing to here from start of tutorial: 20 min*

Let's now analyze the internal representations of our deep CNN model of orientation discrimination and compare them to population responses in mouse primary visual cortex.

In this section, we'll try to quantitatively compare CNN and primary visual cortex representations. In the next section, we will visualize their representations and get some intuition for their structure.

```python
# @title Video 2: Quantitative comparisons of CNNs and neural activity
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV1KT4y1j7nn", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="2Jbk7jFBvbU", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

We noticed above some similarities and differences between the population responses in mouse primary visual cortex and in different layers in our model. Let's now try to quantify this.

To do this, we'll use a technique called [**Representational Similarity Analysis**](https://www.frontiersin.org/articles/10.3389/neuro.06.004.2008/full?utm_source=FWEB&utm_medium=NBLOG&utm_campaign=ECO_10YA_top-research). The idea is to look at the similarity structure between representations of different stimuli. We can say that a brain area and a model use a similar representational scheme if stimuli that are represented (dis)similarly in the brain are represented (dis)similarly in the model as well.

## Section 2.1: Representational dissimilarity matrix (RDM)

To quantify this, we begin by computing the **representational dissimilarity matrix (RDM)** for the mouse V1 data and each model layer. This matrix, which we'll call $\mathbf{M}$, is computed as one minus the correlation coefficients between population responses to each stimulus. We can efficiently compute this by using the $z$-scored responses.

The $z$-scored response of all neurons $\mathbf{r}$ to stimulus $s$ is the response mean-subtracted across neurons $i$ and normalized to standard deviation 1 across neurons $i$ where $N$ is the total number of neurons:
\begin{equation}
  \mathbf{z}^{(s)} = \frac{\mathbf{r}^{(s)} - \mu^{(s)}}
  {\sigma^{(s)}}
\end{equation}
where $\mu^{(s)} = \frac{1}{N}\sum_{i=1}^N r_i^{(s)}$ and
$\sigma^{(s)} = \sqrt{\frac{1}{N}\sum_{i=1}^N \left( r_i^{(s)} - \mu^{(s)} \right)^2}$.

Then the full matrix can be computed as:
\begin{gather}
  \mathbf{M} = 1 - \frac{1}{N} \mathbf{ZZ}^T \\
\end{gather}
where $\mathbf{Z}$ is the z-scored response matrix with rows $\mathbf{r}^{(s)}$ and N is the number of neurons (or units). See [bonus section 3](#b3) for full explanation.

### Coding Exercise 2.1: Compute RDMs

Complete the function `RDM()` for computing the RDM for a given set of population responses to each stimulus. Use the above formula in terms of $z$-scored population responses. We will use the helper function `zscore()` to compute the matrix of $z$-scored responses.

The subsequent cell uses this function to plot the RDM of the population responses in the V1 data and in each layer of our model CNN.

```python
def RDM(resp):
  """Compute the representational dissimilarity matrix (RDM)

  Args:
    resp (ndarray): S x N matrix with population responses to
      each stimulus in each row

  Returns:
    ndarray: S x S representational dissimilarity matrix
  """
  #########################################################
  ## TO DO for students: compute representational dissimilarity matrix
  # Fill out function and remove
  raise NotImplementedError("Student exercise: complete function RDM")
  #########################################################

  # z-score responses to each stimulus
  zresp = zscore(resp, axis=1)

  # Compute RDM
  RDM = ...

  return RDM

# Compute RDMs for each layer
rdm_dict = {label: RDM(resp) for label, resp in resp_dict.items()}

# Plot RDMs
plot_multiple_rdm(rdm_dict)
```

```python
# to_remove solution
def RDM(resp):
  """Compute the representational dissimilarity matrix (RDM)

  Args:
    resp (ndarray): S x N matrix with population responses to
      each stimulus in each row

  Returns:
    ndarray: S x S representational dissimilarity matrix
  """

  # z-score responses to each stimulus
  zresp = zscore(resp, axis=1)

  # Compute RDM
  RDM = 1 - (zresp @ zresp.T) / zresp.shape[1]

  return RDM

# Compute RDMs for each layer
rdm_dict = {label: RDM(resp) for label, resp in resp_dict.items()}

# Plot RDMs
with plt.xkcd():
  plot_multiple_rdm(rdm_dict)
```

```python
# @title Video 3: Coding Exercise 2.1 solution discussion
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV16a4y1a7nc", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="otzR-KXDjus", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

## Section 2.2: Determing representation similarity

*Estimated timing to here from start of tutorial: 35 min*

To quantify how similar the representations are, we can simply correlate their dissimilarity matrices. For this, we'll again use the correlation coefficient. Note that dissimilarity matrices are symmetric ($M_{ss'} = M_{s's}$), so we should only use the off-diagonal terms on one side of the diagonal when computing this correlation to avoid overcounting. Moreover, we should leave out the diagonal terms, which are always equal to 0, so will always be perfectly correlated across any pair of RDM's.

### Coding Exercise 2.2: Correlate RDMs

Complete the function `correlate_rdms()` below that computes this correlation. The code for extracting the off-diagonal terms is provided.

We will then use function to compute the correlation between the RDM's for each layer of our model CNN and that of the V1 data.

```python
def correlate_rdms(rdm1, rdm2):
  """Correlate off-diagonal elements of two RDM's

  Args:
    rdm1 (np.ndarray): S x S representational dissimilarity matrix
    rdm2 (np.ndarray): S x S representational dissimilarity matrix to
      correlate with rdm1

  Returns:
    float: correlation coefficient between the off-diagonal elements
      of rdm1 and rdm2

  """

  # Extract off-diagonal elements of each RDM
  ioffdiag = np.triu_indices(rdm1.shape[0], k=1)  # indices of off-diagonal elements
  rdm1_offdiag = rdm1[ioffdiag]
  rdm2_offdiag = rdm2[ioffdiag]

  #########################################################
  ## TO DO for students: compute correlation coefficient
  # Fill out function and remove
  raise NotImplementedError("Student exercise: complete correlate rdms")
  #########################################################
  corr_coef = np.corrcoef(..., ...)[0,1]

  return corr_coef

# Split RDMs into V1 responses and model responses
rdm_model = rdm_dict.copy()
rdm_v1 = rdm_model.pop('V1 data')

# Correlate off-diagonal terms of dissimilarity matrices
rdm_sim = {label: correlate_rdms(rdm_v1, rdm) for label, rdm in rdm_model.items()}

# Visualize
plot_rdm_rdm_correlations(rdm_sim)
```

```python
# to_remove solution
def correlate_rdms(rdm1, rdm2):
  """Correlate off-diagonal elements of two RDM's

  Args:
    rdm1 (np.ndarray): S x S representational dissimilarity matrix
    rdm2 (np.ndarray): S x S representational dissimilarity matrix to
      correlate with rdm1

  Returns:
    float: correlation coefficient between the off-diagonal elements
      of rdm1 and rdm2

  """

  # Extract off-diagonal elements of each RDM
  ioffdiag = np.triu_indices(rdm1.shape[0], k=1)  # indices of off-diagonal elements
  rdm1_offdiag = rdm1[ioffdiag]
  rdm2_offdiag = rdm2[ioffdiag]

  corr_coef = np.corrcoef(rdm1_offdiag, rdm2_offdiag)[0,1]

  return corr_coef

# Split RDMs into V1 responses and model responses
rdm_model = rdm_dict.copy()
rdm_v1 = rdm_model.pop('V1 data')

# Correlate off-diagonal terms of dissimilarity matrices
rdm_sim = {label: correlate_rdms(rdm_v1, rdm) for label, rdm in rdm_model.items()}

# Visualize
with plt.xkcd():
  plot_rdm_rdm_correlations(rdm_sim)
```

According to this metric, which layer's representations most resemble those in the data? Does this agree with your intuitions from coding exercise 2.1?

## Section 2.3: Further understanding RDMs

*Estimated timing to here from start of tutorial: 45 min*

To better understand how these correlations in RDM's arise, we can try plotting individual rows of the RDM matrix. The resulting curves show the similarity of the responses to each stimulus with that to one specific stimulus.

```python
ori_list = [-75, -25, 25, 75]
plot_rdm_rows(ori_list, rdm_dict, ori)
```

---
# Section 3: Qualitative comparisons of CNNs and neural activity

To visualize the representations in the data and in each of these model layers, we'll use two classic techniques from systems neuroscience:

1. **tuning curves**: plotting the response of single neurons (or units, in the case of the deep network) as a function of the stimulus orientation

2. **dimensionality reduction**: plotting full population responses to each stimulus in two dimensions via dimensionality reduction. We'll use the non-linear dimensionality reduction technique t-SNE for this. We use dimensionality reduction because there are many units and it's difficult to visualize all of them at once. We use a non-linear dimensionality reduction technique because it can capture complex relationships between stimuli (see W1D5 for more details).

## Section 3.1: Tuning Curves

*Estimated timing to here from start of tutorial: 50 min*

Below, we show some example tuning curves for different neurons and units in the CNN we trained above. How are the single neuron responses similar/different between the model and the data? Try running this cell multiple times to get an idea of shared properties in the tuning curves of the neurons within each population.

```python
#@title
#@markdown Execute this cell to visualize tuning curves

fig, axs = plt.subplots(1, len(resp_dict), figsize=(len(resp_dict) * 6, 6))

for i, (label, resp) in enumerate(resp_dict.items()):

  ax = axs[i]
  ax.set_title('%s responses' % label)

  # Pick three random neurons whose tuning curves to plot
  ineurons = np.random.choice(resp.shape[1], 3, replace=False)

  # Plot tuning curves of ineurons
  ax.plot(ori, resp[:, ineurons])

  ax.set_xticks(np.linspace(-90, 90, 5))
  ax.set_xlabel('stimulus orientation')
  ax.set_ylabel('neural response')

plt.tight_layout()
plt.show()
```

## Section 3.2: Dimensionality reduction of representations

*Estimated timing to here from start of tutorial: 55 min*

We can visualize a dimensionality-reduced version of the internal representations of the mouse primary visual cortex or CNN internal representations in order to potentially uncover informative structure. Here, we use PCA to reduce the dimensionality to 20 dimensions, and then use tSNE to further reduce dimensionality to 2 dimensions. We use the first step of PCA so that tSNE runs faster (this is standard practice in the field).

```python
# @markdown Execute this cell to visualize low-d representations
def plot_resp_lowd(resp_dict):
    """Plot a low-dimensional representation of each dataset in resp_dict."""
    n_col = len(resp_dict)
    fig, axs = plt.subplots(1, n_col, figsize=(4.5 * len(resp_dict), 4.5))
    for i, (label, resp) in enumerate(resp_dict.items()):

      ax = axs[i]
      ax.set_title('%s responses' % label)

      # First do PCA to reduce dimensionality to 20 dimensions so that tSNE is faster
      resp_lowd = PCA(n_components=min(20, resp.shape[1]), random_state=0).fit_transform(resp)

      # Then do tSNE to reduce dimensionality to 2 dimensions
      resp_lowd = TSNE(n_components=2, random_state=0).fit_transform(resp_lowd)

      # Plot dimensionality-reduced population responses 'resp_lowd'
      # on 2D axes, with each point colored by stimulus orientation
      x, y = resp_lowd[:, 0], resp_lowd[:, 1]
      pts = ax.scatter(x, y, c=ori, cmap='twilight', vmin=-90, vmax=90)
      fig.colorbar(pts, ax=ax, ticks=np.linspace(-90, 90, 5), label='Stimulus orientation')

      ax.set_xlabel('Dimension 1')
      ax.set_ylabel('Dimension 2')
      ax.set_xticks([])
      ax.set_yticks([])

with plt.xkcd():
  plot_resp_lowd(resp_dict)
```

### Think! 3.2: Visualizing reduced dimensionality representations

Interpret the figure above. Why do these representations look the way they do? Here are a few specific questions to think about:
  * How are the population responses similar/different between the model and the data? Can you explain these population-level responses from the single neuron responses seen in the previous exercise, or vice-versa?
  * How do the representations in the different layers of the model differ, and how does this relate to the orientation discrimination task the model was optimized for?
  * Which layer of our deep network encoding model most closely resembles the V1 data?

```python
# to_remove explanation

"""
The single unit activations of the 'pool' layer in the model have peaks at various
orientations, but they appear to have more peaks than the tuning curves from the
original data do. When we look at the population level responses we see that they
are not quite as smooth across orientations in the t-SNE embedding, which is likely
due to the fact that the 'pool' layer activations do not have localized responses
to orientations like the neural data.

The representations in the fully-connected 'fc' layer appear to be much more
clustered than the 'pool' layer. Stimuli which correspond to the same choice are
clustered together. It seems like the 'pool' layer is still working hard to
represent information about the orientation of the stimulus (much like the V1
population), whereas the 'fc' layer only cares about tilt category, representing
all the stimuli with the same category in a similar way regardless of their different
orientations.

From this analysis, it appears that the 'pool' layer is more similar to the neural
data at the population level.

""";
```

---
# Summary

*Estimated timing of tutorial: 1 hour, 10 minutes*

In this notebook, we learned
* how to use deep learning to build a normative encoding model of the visual system
* how to use RSA to evaluate how the model's representations match to those in the brain

Our approach was to optimize a deep convolutional network to solve an orientation discrimination task. But note that many other approaches could have been taken.

Firstly, there are many other "normative" ways to solve this orientation discrimination task. We could have used different neural network architectures, or even used a completely different algorithm that didn't involve a neural network at all, but instead used other kinds of image transformations (e.g. Fourier transforms). Neural network approaches, however, are special in that they explicitly uses abstract distributed representations to compute, which feels a lot closer to the kinds of algorithms the brain uses. *Convolutional* neural networks in particular are well-suited for building normative models of the visual system.

Secondly, our choice of visual task was mostly arbitrary. For example, we could have trained our network to directly estimate the orientation of the stimulus, rather than just discriminating between two classes of tilt. Or, we could have trained the network to perform a more naturalistic task, such as recognizing the rotation of an arbitrary image. Or we could try a task like object recognition. Is this something that mice compute in their visual cortex?

Training on different tasks could lead to different representations of the oriented grating stimuli, which might match the observed V1 representations better or worse.

**See Section 3 in the Bonus Tutorial to fit a convolutional neural network directly to neural activity to build an encoding model.**

---
# Bonus

## Bonus Section 1: Building CNN's with PyTorch

Here we walk through building the different types of layers in a CNN using PyTorch, culminating in the CNN model used above.

### Bonus Section 1.1: Fully connected layers

In a fully connected layer, each unit computes a weighted sum over all the input units and applies a non-linear function to this weighted sum. You have used such layers many times already in parts 1 and 2. As you have already seen, these are implemented in PyTorch using the `nn.Linear` class.

  See the next cell for code for constructing a deep network with one fully connected layer that will classify an input image as being tilted left or right. Specifically, its output is the predicted probability of the input image being tilted right. To ensure that its output is a probability (i.e. a number between 0 and 1), we use a sigmoid activation function to squash the output into this range (implemented with `torch.sigmoid()`).

```python
class FC(nn.Module):
  """Deep network with one fully connected layer

    Args:
      h_in (int): height of input image, in pixels (i.e. number of rows)
      w_in (int): width of input image, in pixels (i.e. number of columns)

    Attributes:
      fc (nn.Linear): weights and biases of fully connected layer
      out (nn.Linear): weights and biases of output layer

    """

  def __init__(self, h_in, w_in):
    super().__init__()
    self.dims = h_in * w_in  # dimensions of flattened input
    self.fc = nn.Linear(self.dims, 10)  # flattened input image --> 10D representation
    self.out = nn.Linear(10, 1)  # 10D representation --> scalar

  def forward(self, x):
    """Classify grating stimulus as tilted right or left

    Args:
      x (torch.Tensor): p x 48 x 64 tensor with pixel grayscale values for
          each of p stimulus images.

    Returns:
      torch.Tensor: p x 1 tensor with network outputs for each input provided
          in x. Each output should be interpreted as the probability of the
          corresponding stimulus being tilted right.

    """
    x = x.view(-1, self.dims)  # flatten each input image into a vector
    x = torch.relu(self.fc(x))  # output of fully connected layer
    x = torch.sigmoid(self.out(x))  # network output
    return x
```

### Bonus Section 1.2: Convolutional layers

In a convolutional layer, each unit computes a weighted sum over a two-dimensional $K \times K$ patch of inputs. As we saw in part 2, the units are arranged in **channels** (see figure below), whereby units in the same channel compute the same weighted sum over different parts of the input, using the weights of that channel's **convolutional filter (or kernel)**. The output of a convolutional layer is thus a three-dimensional tensor of shape $C^{out} \times H \times W$, where $C^{out}$ is the number of channels (i.e. the number of convolutional filters/kernels), and $H$ and $W$ are the height and width of the input.

Such layers can be implemented in Python using the PyTorch class `nn.Conv2d as we saw in tutorial 2 (documentation [here](https://pytorch.org/docs/master/generated/torch.nn.Conv2d.html)).

See the next cell for code incorporating a convolutional layer with 8 convolutional filters of size 5 $\times$ 5 into our above fully connected network. Note that we have to flatten the multi-channel output in order to pass it on to the fully connected layer.

```python
class ConvFC(nn.Module):
  """Deep network with one convolutional layer and one fully connected layer

  Args:
    h_in (int): height of input image, in pixels (i.e. number of rows)
    w_in (int): width of input image, in pixels (i.e. number of columns)

  Attributes:
    conv (nn.Conv2d): filter weights of convolutional layer
    dims (tuple of ints): dimensions of output from conv layer
    fc (nn.Linear): weights and biases of fully connected layer
    out (nn.Linear): weights and biases of output layer

  """

  def __init__(self, h_in, w_in):
    super().__init__()
    C_in = 1  # input stimuli have only 1 input channel
    C_out = 6  # number of output channels (i.e. of convolutional kernels to convolve the input with)
    K = 7  # size of each convolutional kernel (should be odd number for the padding to work as expected)
    self.conv = nn.Conv2d(C_in, C_out, kernel_size=K, padding=K//2)  # add padding to ensure that each channel has same dimensionality as input
    self.dims = (C_out, h_in, C_out)  # dimensions of conv layer output
    self.fc = nn.Linear(np.prod(self.dims), 10)  # flattened conv output --> 10D representation
    self.out = nn.Linear(10, 1)  # 10D representation --> scalar

  def forward(self, x):
    """Classify grating stimulus as tilted right or left

    Args:
      x (torch.Tensor): p x 48 x 64 tensor with pixel grayscale values for
          each of p stimulus images.

    Returns:
      torch.Tensor: p x 1 tensor with network outputs for each input provided
          in x. Each output should be interpreted as the probability of the
          corresponding stimulus being tilted right.

    """
    x = x.unsqueeze(1)  # p x 1 x 48 x 64, add a singleton dimension for the single stimulus channel
    x = torch.relu(self.conv(x))  # output of convolutional layer
    x = x.view(-1, np.prod(self.dims))  # flatten convolutional layer outputs into a vector
    x = torch.relu(self.fc(x))  # output of fully connected layer
    x = torch.sigmoid(self.out(x))  # network output
    return x
```

### Bonus Section 1.3: Max pooling layers

In a max pooling layer, each unit computes the maximum over a small two-dimensional $K^{pool} \times K^{pool}$ patch of inputs. Given a multi-channel input of dimensions $C \times H \times W$, the output of a max pooling layer has dimensions $C \times H^{out} \times W^{out}$, where:
\begin{align}
  H^{out} &= \left\lfloor \frac{H}{K^{pool}} \right\rfloor\\
  W^{out} &= \left\lfloor \frac{W}{K^{pool}} \right\rfloor
\end{align}
where $\lfloor\cdot\rfloor$ denotes rounding down to the nearest integer below (i.e. floor division `//` in Python).

  Max pooling layers can be implemented with the PyTorch `nn.MaxPool2d` class, which takes as a single argument the size $K^{pool}$ of the pooling patch. See the next cell for an example, which builds upon the previous example by adding in a max pooling layer just after the convolutional layer. Note again that we need to calculate the dimensions of its output in order to set the dimensions of the subsequent fully connected layer.

```python
class ConvPoolFC(nn.Module):
  """Deep network with one convolutional layer followed by a max pooling layer
  and one fully connected layer

  Args:
    h_in (int): height of input image, in pixels (i.e. number of rows)
    w_in (int): width of input image, in pixels (i.e. number of columns)

  Attributes:
    conv (nn.Conv2d): filter weights of convolutional layer
    pool (nn.MaxPool2d): max pooling layer
    dims (tuple of ints): dimensions of output from pool layer
    fc (nn.Linear): weights and biases of fully connected layer
    out (nn.Linear): weights and biases of output layer

  """

  def __init__(self, h_in, w_in):
    super().__init__()
    C_in = 1  # input stimuli have only 1 input channel
    C_out = 6  # number of output channels (i.e. of convolutional kernels to convolve the input with)
    K = 7  # size of each convolutional kernel
    Kpool = 8  # size of patches over which to pool
    self.conv = nn.Conv2d(C_in, C_out, kernel_size=K, padding=K//2)  # add padding to ensure that each channel has same dimensionality as input
    self.pool = nn.MaxPool2d(Kpool)
    self.dims = (C_out, h_in // Kpool, w_in // Kpool)  # dimensions of pool layer output
    self.fc = nn.Linear(np.prod(self.dims), 10)  # flattened pool output --> 10D representation
    self.out = nn.Linear(10, 1)  # 10D representation --> scalar

  def forward(self, x):
    """Classify grating stimulus as tilted right or left

    Args:
      x (torch.Tensor): p x 48 x 64 tensor with pixel grayscale values for
          each of p stimulus images.

    Returns:
      torch.Tensor: p x 1 tensor with network outputs for each input provided
          in x. Each output should be interpreted as the probability of the
          corresponding stimulus being tilted right.

    """
    x = x.unsqueeze(1)  # p x 1 x 48 x 64, add a singleton dimension for the single stimulus channel
    x = torch.relu(self.conv(x))  # output of convolutional layer
    x = self.pool(x)  # output of pooling layer
    x = x.view(-1, np.prod(self.dims))  # flatten pooling layer outputs into a vector
    x = torch.relu(self.fc(x))  # output of fully connected layer
    x = torch.sigmoid(self.out(x))  # network output
    return x
```

This pooling layer completes the CNN model trained above to perform orientation discrimination. We can think of this architecture as having two primary layers:
1. a convolutional + pooling layer
2. a fully connected layer

We group together the convolution and pooling layers into one, as they really form one full unit of convolutional processing, where each patch of the image is passed through a convolutional filter and pooled with neighboring patches. It is standar practice to follow up any convolutional layer with a pooling layer, so they are generally treated as a single block of processing.

## Bonus Section 2: Orientation discrimination as a binary classification problem

What loss function should we minimize to optimize orientation discrimination performance? We first note that the orientation discrimination task is a **binary classification problem**, where the goal is to classify a given stimulus into one of two classes: being tilted left or being tilted right.

Our goal is thus to output a high probability of the stimulus being tilted right (i.e. large $p$) whenever the stimulus is tilted right, and a high probability of the stimulus being tilted left (i.e. large $1-p \Leftrightarrow$ small $p$) whenever the stimulus is tilted left.

Let $\tilde{y}^{(n)}$ be the label of the $n$th stimulus in the mini-batch, indicating its true tilt:
\begin{equation}
  \tilde{y}^{(n)} =
  \begin{cases}
    1 &\text{if stimulus }n\text{ is tilted right} \\
    0 &\text{if stimulus }n\text{ is tilted left}
  \end{cases}
\end{equation}
Let $p^{(n)}$ be the predicted probability of that stimulus being tilted right assigned by our network. Note that that $1-p^{(n)}$ is the predicted probability of that stimulus being tilted left. We'd now like to modify the parameters so as to maximize the predicted probability of the true class $\tilde{y}^{(n)}$. One way to formalize this is as maximizing the *log* probability
\begin{align}
  \log \left( \text{predicted probability of stimulus } n \text{ being of class } \tilde{y}^{(n)}\right) &=
  \begin{cases}
    \log p^{(n)} &\text{if }\tilde{y}^{(n)} = 1 \\
    \log (1 - p^{(n)}) &\text{if }\tilde{y}^{(n)} = 0
  \end{cases}
  \\
  &= \tilde{y}^{(n)} \log p^{(n)} + (1 - \tilde{y}^{(n)})\log(1 - p^{(n)})
\end{align}
You should recognize this expression as the log likelihood of the Bernoulli distribution under the predicted probability $p^{(n)}$. This is the same quantity that is maximized in logistic regression, where the predicted probability $p^{(n)}$ is just a simple linear sum of its inputs (rather than a complicated non-linear operation, like in the deep networks used here).

To turn this into a loss function, we simply multiply it by -1, resulting in the so-called **binary cross-entropy**, or **negative log likelihood**. Summing over $P$ samples in a batch, the binary cross entropy loss is given by
\begin{equation}
  L = -\sum_{n=1}^P \tilde{y}^{(n)} \log p^{(n)} + (1 - \tilde{y}^{(n)})\log(1 - p^{(n)})
\end{equation}
The binary cross-entropy loss can be implemented in PyTorch using the `nn.BCELoss()` loss function (cf. [documentation](https://pytorch.org/docs/master/generated/torch.nn.BCELoss.html)).

Feel free to check out the code used to optimize the CNN in the `train()` function defined in the hidden cell of helper functions at the top of the notebook. Because the CNN's used here have lots of parameters, we have to use two tricks that we didn't use in the previous parts of this tutorial:
1. We have to use *stochastic* gradient descent (SGD), rather than just gradient descent (GD).
2. We have to use [momentum](https://distill.pub/2017/momentum/) in our SGD updates. This is easily incorporated into our PyTorch implementation by just setting the `momentum` argument of the built-in `optim.SGD` optimizer.

## Bonus Section 3: RDM Z-Score Explanation

If $r^{(s)}_i$ is the response of the $i$th neuron to the $s$th stimulus, then
\begin{gather}
  M_{ss'} = 1 - \frac{\text{Cov}\left[ r_i^{(s)}, r_i^{(s')} \right]}{\sqrt{\text{Var}\left[ r_i^{(s)} \right] \text{Var}\left[ r_i^{(s')} \right]}} = 1 - \frac{\sum_{i=1}^N (r_i^{(s)} - \bar{r}^{(s)})(r_i^{(s')} - \bar{r}^{(s')}) }{\sqrt{\sum_{i=1}^N \left( r_i^{(s)} - \bar{r}^{(s)} \right)^2 \sum_{i=1}^N \left( r_i^{(s')} - \bar{r}^{(s')} \right)^2 }} \\
  \bar{r}^{(s)} = \frac{1}{N} \sum_{i=1}^N r_i^{(s)}
\end{gather}
This can be computed efficiently by using the $z$-scored responses
\begin{equation}
  z_i^{(s)} = \frac{r_i^{(s)} - \bar{r}^{(s)}}{\sqrt{\frac{1}{N}\sum_{i=1}^N \left( r_i^{(s)} - \bar{r}^{(s)} \right)^2}} \Rightarrow M_{ss'} = 1 - \frac{1}{N}\sum_{i=1}^N z_i^{(s)}z_i^{(s')}
\end{equation}
such that the full matrix can be computed through the matrix multiplication
\begin{gather}
  \mathbf{M} = 1 - \frac{1}{N} \mathbf{ZZ}^T \\
  \mathbf{Z} =
  \begin{bmatrix}
    z_1^{(1)} & z_2^{(1)} & \ldots & z_N^{(1)} \\
    z_1^{(2)} & z_2^{(2)} & \ldots & z_N^{(2)} \\
    \vdots & \vdots & \ddots & \vdots \\
    z_1^{(S)} & z_2^{(S)} & \ldots & z_N^{(S)}
  \end{bmatrix}
\end{gather}

where $S$ is the total number of stimuli. Note that $\mathbf{Z}$ is an $S \times N$ matrix, and $\mathbf{M}$ is an $S \times S$ matrix.