#!/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 3", we perform all the steps done in part2,
# with an even finer resolution, and more complex initial concentrations,
# repeated for 2 different diffusion algorithms.
# 
# 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]:


# Parameters of the simulation run.  We'll be considering just 1 chemical species, "A"
diffusion_rate = 10.
delta_t = 0.01
n_bins = 5000
delta_x = 2       # Note that the number of bins also define the fraction of the sine wave cycle in each bin


# In[4]:


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


# In[ ]:





# # ALGORITHM 1

# In[5]:


algorithm = None   # "Explicit, with 3+1 stencil"


# In[6]:


# 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=12, bias=40)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=10)
bio.inject_sine_conc(species_name="A", frequency=16, amplitude=5)


# 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()


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


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[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), 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[12]:


# Computer the second spacial derivative, 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[13]:


# 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 above number is a measure of the discrepancy from the perfect match (zero distance) that an ideal solution would provide. 

# In[ ]:





# # ALGORITHM 2

# In[14]:


algorithm = "5_1_explicit"   # "Explicit, with 5+1 stencil"


# In[15]:


# 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=12, bias=40)
bio.inject_sine_conc(species_name="A", frequency=2, amplitude=10)
bio.inject_sine_conc(species_name="A", frequency=16, amplitude=5)


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


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


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


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


# In[19]:


# Compute time derivatives (for each bin), 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[20]:


# Computer the second spacial derivative, 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[21]:


# 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


# # Both algorithms show good measures of accuracy

# In[ ]: