#!/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.

# ### TAGS :  "diffusion 1D", "under-the-hood"

# In[1]:


LAST_REVISED = "May 3, 2025"
LIFE123_VERSION = "1.0.0rc3"       # Library version this experiment is based on


# In[2]:


#import set_path              # Using MyBinder?  Uncomment this before running the next cell!


# In[3]:


#import sys
#sys.path.append("C:/some_path/my_env_or_install")   # CHANGE to the folder containing your venv or libraries installation!
# NOTE: If any of the imports below can't find a module, uncomment the lines above, or try:  import set_path   

from life123 import BioSim1D, ChemData, CollectionArray, Numerical, check_version

import numpy as np


# In[4]:


check_version(LIFE123_VERSION)


# In[ ]:





# In[5]:


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


chem_data = ChemData(diffusion_rates=diffusion_rate, names="A")


# In[ ]:





# # ALGORITHM 1

# In[7]:


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


# In[8]:


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


# In[9]:


# Visualize the initial system state
bio.visualize_system(title_prefix="Initial System State (for the tiny system)")


# In[10]:


# Show as heatmap
bio.system_heatmaps(title_prefix="Initial System State (for the tiny system)")


# In[ ]:





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


history = CollectionArray()   # All the system state will get collected in this object
# Store the initial state
arr = bio.lookup_species(chem_index=0, copy=True)
history.store(par=bio.system_time, data_snapshot=arr, caption=f"State at time {bio.system_time}")


# In[12]:


# 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(chem_index=0, copy=True)
    history.store(par=bio.system_time, data_snapshot=arr, caption=f"State at time {bio.system_time}")


# In[13]:


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


# In[14]:


# Compute time derivatives (for each bin), using 5-point stencils
df_dt_all_bins = np.apply_along_axis(Numerical.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[15]:


# Computer the second spacial derivative, using 5-point stencils
gradient_x_at_t2 = Numerical.gradient_order4_1d(arr=f_at_t2, dx=delta_x)
second_gradient_x_at_t2 = Numerical.gradient_order4_1d(arr=gradient_x_at_t2, dx=delta_x)
second_gradient_x_at_t2.shape


# In[16]:


# 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

Numerical.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[ ]:





# In[ ]:





# # ALGORITHM 2

# In[17]:


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


# 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(chem_label="A", number_cycles=1, amplitude=12, bias=40)
bio.inject_sine_conc(chem_label="A", number_cycles=2, amplitude=10)
bio.inject_sine_conc(chem_label="A", number_cycles=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[19]:


history = CollectionArray()   # All the system state will get collected in this object
# Store the initial state
arr = bio.lookup_species(chem_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(chem_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), using 5-point stencils
df_dt_all_bins = np.apply_along_axis(Numerical.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[23]:


# Computer the second spacial derivative, using 5-point stencils
gradient_x_at_t2 = Numerical.gradient_order4_1d(arr=f_at_t2, dx=delta_x)
second_gradient_x_at_t2 = Numerical.gradient_order4_1d(arr=gradient_x_at_t2, dx=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

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


# # Both algorithms show good measures of accuracy

# In[ ]: