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

# ## Carry out a few single indivual steps of diffusion, 
# ### and directly verify that the values satisfy the diffusion equation
# 
# In this "PART 1", we'll be looking at a tiny system with just 10 bins, to easily inspect the numbers directly.
# We'll use concentrations initialized to an upward-shifted sine wave with 1 cycle over the length system
# 
# LAST REVISED: May 27, 2023

# In[1]:


import set_path      # Importing this module will add the project's home directory to sys.path


# In[2]:


from experiments.get_notebook_info import get_notebook_basename

from src.life_1D.bio_sim_1d import BioSim1D
from src.modules.reactions.reaction_data import ReactionData as chem
from src.modules.movies.movies import MovieArray
from src.modules.numerical.numerical import Numerical as num

import numpy as np

import plotly.express as px
import plotly.graph_objects as go


# In[3]:


# We'll be considering just 1 chemical species, "A"
diffusion_rate = 10.

chem_data = chem(diffusion_rates=[diffusion_rate], names=["A"])


# In[4]:


# Initialize the system with just a few bins
bio = BioSim1D(n_bins=10, chem_data=chem_data)


# In[5]:


# Initialize the concentrations to a sine wave with 1 cycle over the system
# (with a bias to always keep it > 0)
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)


# In[6]:


bio.show_system_snapshot()


# In[7]:


fig = px.line(data_frame=bio.system_snapshot(), y=["A"], 
              title= "Initial System State (for the tiny system)",
              color_discrete_sequence = ['red'],
              labels={"value":"concentration", "variable":"Chemical", "index":"Bin number"})
fig.show()


# In[8]:


# Show as heatmap
fig = px.imshow(bio.system_snapshot().T, 
                title= "Initial System State (for the tiny system)", 
                labels=dict(x="Bin number", y="Chem. species", color="Concentration"),
                text_auto=False, color_continuous_scale="gray_r") 

fig.data[0].xgap=1
fig.data[0].ygap=1

fig.show()


# ### Now do 4 rounds of single-step diffusion, to collect the system state at a total of 5 time points: 
# #### t0 (the initial state), plus t1, t2, t3 and t4

# In[9]:


# All the system states will get collected in this object
history = MovieArray()


# In[10]:


# Store the initial state
arr = bio.lookup_species(species_index=0, copy=True)
history.store(par=bio.system_time, data_snapshot=arr, caption=f"State at time {bio.system_time}")


# In[11]:


# Take a look at what got stored so far (a matrix whose only row is the initial state)
history.get_array()


# In[12]:


# Additional parameters of the simulation run (the diffusion rate got set earlier)
delta_t = 0.01
delta_x = 2       # Note that the number of bins also define the fraction of the sine wave cycle in each bin
algorithm = None  # "Explicit, with 3+1 stencil"


# In[13]:


# Do the 4 rounds of single-step diffusion; show the system state after each step, and accumulate all data
# in the history object
for _ in range(4):
    bio.diffuse(time_step=delta_t, n_steps=1, delta_x=delta_x , algorithm=algorithm)
    bio.describe_state(concise=True)

    arr = bio.lookup_species(species_index=0, copy=True)
    history.store(par=bio.system_time, data_snapshot=arr, caption=f"State at time {bio.system_time}")


# In[14]:


# Now, let's examine the data collected at the 5 time points
all_history = history.get_array()
all_history


# #### Each row in the above matrix is a state snapshot at a different time:
# first row is the initial state at time 0; the successive rows are at t1, t2, t3, t4

# In[15]:


# Let's consider the state, i.e. the concentration across the bins, at the midpoint in time (t2)
f_at_t2 = all_history[2]
f_at_t2


# If one compares the above state (at t2) with the initial state (the top row in the matrix), 
# one can see, for example, that the leftmost bin's concentration is increasing ("pulled up" by its neighbor to the right):
# 50. has now become 50.28881576
# 
# The rightmost bin's concentration is decreasing ("pulled down" by its neighbor to the left):
# 44.12214748 has now become 43.94505274

# ## The diffusion equation states that the partial derivative of the concentration values with respect to time  must equal the (diffusion rate) times (the 2nd partial derivative  with respect to space).
# Let's see if that is the case for the values at time t2!  We are picking t2 because we need at least a value at the earlier time, and a value at the later time

# ## A. The 2nd partial derivative  with respect to space

# In[16]:


# A simple-minded way of computing the 2nd spacial derivative is to use the Numpy gradient function TWICE across the x value
# (that function approximates the derivative using differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)   # Start with taking the first derivative
gradient_x_at_t2    # This will be the partial derivative of the concentration with respect to x, at time t2


# For example, the 2nd entry in the above array of estimated derivatives, can be manually checked from the 1st and 3rd value in the function f_at_t2(x), as follows:

# In[17]:


(59.32979676 - 50.28881576) / (2*delta_x)    # This way of numerically estimating derivatives is called "Central Differences"


# #### Now take the derivative again, with the Numpy gradient function, to arrive at a coarse estimate of the 2nd derivative with respect to x:

# In[18]:


second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2


# Note how the 2nd derivative is 0 at bin 5 (bins are numbered 0 thru 9): if you look at the earlier sine plot, x5 is the inflection point.

# ### B. The partial derivative with respect to time
# Now, let's look at how concentrations change with time.  Let's first revisit the full history:

# In[19]:


all_history


# #### For simplicity, let's start by just inspecting how the values change over time at the *3rd bin* from the left, 
# i.e. the 3rd *column* (index 2 because counting starts at 0) of the above matrix:

# In[20]:


f_of_t_at_x2 = all_history[ : , 2]  # The column with index 2
f_of_t_at_x2


# ### The above a function of time 
# (we took an entry from each row - the rows being the system states at t0, t1, t2, t3, t4); let's look at its time derivative:

# In[21]:


gradient_t_at_x2 = np.gradient(f_of_t_at_x2, delta_t)
gradient_t_at_x2


# ### The above is the rate of change of the concentration, as the diffusion proceeds, at the position x2
# At time t2, the midpoint in the simulation, the value is:

# In[22]:


gradient_t_at_x2[2]


# ### C. All said and done, we have collected the time derivative and the 2nd spacial derivative, at the point (x2, t2).
# Do those values satisfy the diffusion equation??  Does the time derivative indeed equal the (diffusion rate) x (the 
# 2nd spacial derivative)?  Let's see:

# In[23]:


gradient_t_at_x2[2]


# In[24]:


diffusion_rate * second_gradient_x_at_t2[2]


# ## D. The 2 value indeed roughly match - considering the coarseness of the large spacial grid, and the coarseness of estimating the derivatives numerically. 
# We have evidence that the diffusion equation may indeed be satisfied by the values we obtained from our diffusion algorithm, in the proximity of the point (x2, t2)

# ### E. Finally, instead of just scrutining the match at the point x2, let's do that for all the points in space at time t2,
# WITH THE EXCEPTION of the outmost points (because the numeric estimation of the derivatives gets very crummy at the boundary).
#  
# Let's first re-visit all the data once again:

# In[25]:


all_history


# In[26]:


# The following is just an expanded version of what we did before for the time derivative; 
# instead of just considering 1 column,
# like we did before, we're now repeating the computations along all columns
# (the computations are applied vertically, "along axis 0" in Numpy-speak)
gradient_t = np.apply_along_axis(np.gradient, 0, all_history, delta_t)
gradient_t


# In[27]:


# Again, we focus on time t2 (the 3rd row), to stay away from the edges
gradient_t_at_t2 = gradient_t[2]
gradient_t_at_t2


# Note the value -8.94751865, 3rd from left : that's the single value we looked at before (at point x2)

# #### Time to again check the match of the two sides of the diffusion equation

# In[28]:


lhs = gradient_t_at_t2   # The LEFT-hand side of the diffusion equation, as a vector for all the spacial points at time t2
lhs


# In[29]:


rhs = diffusion_rate * second_gradient_x_at_t2  # The RIGHT-hand side of the diffusion equation, again as a vector
rhs


# ## The left-hand side and the right-hand side of the diffusion equation appear to generally agree, except at the boundary points, where our approximations are just too crummy

# In[30]:


lhs - rhs


# In[31]:


# Here we use a handy function to compare two equal-sized vectors,
# while opting to disregarding a specified number of entries at each edge.
# It returns the Euclidean distance ("L2 norm") of the shortened vectors
num.compare_vectors(lhs, rhs, trim_edges=1)


# #### IMPORTANT: all values in this experiment are VERY coarse, because of the large effective delta_x (the tiny number of bins.)
# In part2 (notebook "validate_diffusion_2"), much-better approximations will get looked at!

# In[ ]: