#!/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 2", we quickly repeat all the steps discussed at length in part1,
# but using a much-finer resolution.
# This time, we'll start with slightly more complex initial concentrations built from 2 superposed sine waves.
# 
# **We'll also explore the effects of:**  
# -Spatial resolution ("delta x")  
# -Temporal resolution ("delta t")    
# -Alternate methods of estimating numerical derivatives
# 
# LAST REVISED: June 23, 2024 (using v. 1.0 beta34.1)

# In[1]:


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


# In[2]:


from life123 import BioSim1D
from life123 import ChemData as chem
from life123 import MovieArray
from life123 import Numerical as num

import numpy as np

import plotly.express as px


# In[3]:


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

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


# # BASELINE
# This will be our initial system, whose adherence to the diffusion equation we'll test.
# Afterwards, we'll tweak individual parameters - and observed their effect on the closeness of the approximation

# In[4]:


# Parameters of the simulation run (the diffusion rate got set earlier, and will never vary)
delta_t = 0.01
n_bins = 300
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[5]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[6]:


fig = px.line(data_frame=bio.system_snapshot(), y=["A"], 
              title= "Initial System State",
              color_discrete_sequence = ['red'],
              labels={"value":"concentration", "variable":"Chemical", "index":"Bin number"})
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[7]:


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


# In[8]:


# 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[9]:


# Do the 4 rounds of single-step diffusion; 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)

    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[10]:


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


# In[11]:


# Compute time derivatives (for each bin)
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)
df_dt_all_bins.shape


# In[12]:


# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[13]:


# Computer the second spatial derivative (simple method)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# #### Now, let's use the computed values to see how the left- and right-hand side of the diffusion equation compare
# at all bins (except 2 at each edge), at the middle simulation time t2

# In[14]:


lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
lhs.shape


# In[15]:


rhs = diffusion_rate*second_gradient_x_at_t2
rhs.shape


# In[16]:


num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# The above number is a measure of the discrepancy from the perfect match (zero distance) that an ideal solution would provide. 
# 
# Notice how vastly closer to zero it is than the counterpart value we got with the very coarse simulation in experiment "validate_diffusion_1"

# # VARIATIONS on the BASELINE
# We'll now tweak some parameters, and observe the effect on the estimate of the discrepancy from the exact diffusion equation
# 
# The baseline distance was **0.023163289760024783**

# ## Variation 1 : reduce spatial resolution
# Expectation: greater discrepancy

# In[17]:


n_bins = 100          # Reducing the spatial resolution


# In[18]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[19]:


history = MovieArray()   # All the system state will get collected in this object
# 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[20]:


# Do the 4 rounds of single-step diffusion; 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)

    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[21]:


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


# In[22]:


# Compute time derivatives (for each bin) : simple method, using central differences
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[23]:


# Computer the second spatial derivative (simple method, using central differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# In[24]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### The discrepancy value has indeed gone up from the baseline 0.023163289760024783

# ## Variation 2 : increase spatial resolution
# Expectation: smaller discrepancy

# In[25]:


n_bins = 900          # Increasing the spatial resolution (from the baseline of 300)


# In[26]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[27]:


history = MovieArray()   # All the system state will get collected in this object
# 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[28]:


# Do the 4 rounds of single-step diffusion; 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)

    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[29]:


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


# In[30]:


# Compute time derivatives (for each bin) : simple method, using central differences
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[31]:


# Computer the second spatial derivative (simple method, using central differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# In[32]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### The discrepancy value has indeed gone down from the baseline 0.023163289760024783

# ## Variation 3 : reduce time resolution
# Expectation: greater discrepancy

# In[33]:


n_bins = 300          # Restoring the baseline number of bins
delta_t = 0.02        # Reducing the time resolution (from baselin 0.01)


# In[34]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[35]:


history = MovieArray()   # All the system state will get collected in this object
# 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[36]:


# Do the 4 rounds of single-step diffusion; 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)

    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[37]:


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


# In[38]:


# Compute time derivatives (for each bin) : simple method, using central differences
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[39]:


# Computer the second spatial derivative (simple method, using central differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# In[40]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### The discrepancy value has indeed gone up from the baseline 0.023163289760024783

# ## Variation 4 : increase time resolution
# Expectation: smaller discrepancy

# In[41]:


delta_t = 0.005        # Increasing the time resolution (from baselin 0.01)


# In[42]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[43]:


history = MovieArray()   # All the system state will get collected in this object
# 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[44]:


# Do the 4 rounds of single-step diffusion; 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)

    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[45]:


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


# In[46]:


# Compute time derivatives (for each bin) : simple method, using central differences
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[47]:


# Computer the second spatial derivative (simple method, using central differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# In[48]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### The discrepancy value has indeed gone down from the baseline 0.023163289760024783

# ## Variation 5 : Use a better (higher-order) method to numerically estimate the SPATIAL derivatives
# Expectation: presumably smaller discrepancy (unless dwarfed by errors from approximations in the simulation)

# In[49]:


n_bins = 300          # Restoring the baseline number of bins
delta_t = 0.01        # Reducing the baseline time resolution


# In[50]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[51]:


history = MovieArray()   # All the system state will get collected in this object
# 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[52]:


# Do the 4 rounds of single-step diffusion; 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)

    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[53]:


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


# In[54]:


# Compute time derivatives (for each bin) : simple method, using central differences
df_dt_all_bins = np.apply_along_axis(np.gradient, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[55]:


# Computer the second spatial derivative (MORE SOPHISTICATED METHOD, using 5-point stencils)
gradient_x_at_t2 = num.gradient_order4_1d(arr=f_at_t2, dx=delta_x)
second_gradient_x_at_t2 = num.gradient_order4_1d(arr=gradient_x_at_t2, dx=delta_x)
second_gradient_x_at_t2.shape


# In[56]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### Not surprisingly, the discrepancy value (in part caused by poor numeric estimation of spatial derivatives) has indeed gone down from the baseline 0.023163289760024783

# ## Variation 6 : Use a better (higher-order) method to numerically estimate the TIME derivatives
# Expectation: presumably smaller discrepancy (unless dwarfed by errors from approximations in the simulation)
# 
# (Note: we revert to the original method for the *spatial* derivatives)

# In[57]:


# Initialize the system
bio = BioSim1D(n_bins=n_bins, chem_data=chem_data)

# Initialize the concentrations to 2 superposed sine waves
bio.inject_sine_conc(species_name="A", frequency=1, amplitude=10, bias=50)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=8)


# In[58]:


history = MovieArray()   # All the system state will get collected in this object
# 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[59]:


# Do the 4 rounds of single-step diffusion; 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)

    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[60]:


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


# In[61]:


# Compute time derivatives (for each bin) : MORE SOPHISTICATED METHOD, using 5-point stencils
df_dt_all_bins = np.apply_along_axis(num.gradient_order4_1d, 0, all_history, delta_t)

# Let's consider the state at the midpoint in time (t2)
f_at_t2 = all_history[2]     # The middle of the 5 time snapshots
f_at_t2.shape


# In[62]:


# Computer the second spatial derivative (BACK TO SIMPLE METHOD, using central differences)
gradient_x_at_t2 = np.gradient(f_at_t2, delta_x)
second_gradient_x_at_t2 = np.gradient(gradient_x_at_t2, delta_x)
second_gradient_x_at_t2.shape


# In[63]:


# Compare the left and right hand sides of the diffusion equation
lhs = df_dt_all_bins[2]   # t2 is the middle point of the 5
rhs = diffusion_rate*second_gradient_x_at_t2

num.compare_vectors(lhs, rhs, trim_edges=2)  # Euclidean distance, ignoring 2 edge points at each end


# #### Here, the discrepancy value has remained largely the same from the baseline 0.023163289760024783

# In[ ]: