diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..8ebdc17
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,4 @@
+__pycache__/
+dist/
+build/
+*.egg-info/
diff --git a/README.md b/README.md
index fe54835..3b401b3 100644
--- a/README.md
+++ b/README.md
@@ -7,13 +7,30 @@ Please follow the instructions in [pypi_exercise.md](https://github.com/Simulati
 The code used in this exercise is based on [Chapter 7 of the book "Learning Scientific Programming with Python"](https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/).
 
 ## Project description
+This project simulates the diffusion equation in 2D domain using finite difference methods. The simulation demonstrates how heat diffuses in a square region with an initial high-temperature circle at the center. The resulting temperature distrubution is visualized using Matplotliib.
 
 ## Installing the package
+You can install this package from TestPyPI using the following command:
+```bash
+pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple <your_username>_diffusion2D
+```
 
 ### Using pip3 to install from PyPI
+Once your package is live on PyPI, you can install it using:
+```bash
+pip install <your_username>_diffusion2D
+```
 
 ### Required dependencies
+This package requires numpy and matplotlib libraries.
+These dependencies are automatically installed when you install the package via pip.
 
 ## Running this package
+After installation, you can use following command to run the simulation:
+```bash
+python3 -m diffusion2d
+```
+This will generate a 2D plot showing the diffusion process at different time steps.
 
 ## Citing
+This project is based on materials from book [Learning Scientific Programming with Python](https://scipython.com/book/chapter-7-matplotlib/examples/two-dimensional-diffusion-equation/)
diff --git a/diffusion_plots.png b/diffusion_plots.png
new file mode 100644
index 0000000..42b9491
Binary files /dev/null and b/diffusion_plots.png differ
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..a1c8bc3
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,20 @@
+[build-system]
+requires = ["setuptools","wheel","build"]
+
+[project]
+name = "wanivi_diffusion2d"
+description = "A 2D diffusion equation solver with visualization output"
+version = "0.0.2"
+authors = [
+	{name="Vaishnavi", email="vaishnaviwani671@gmail.com"}
+]
+readme = "README.md"
+keywords = ["diffusion2D", "solver"]
+classifiers = [
+	"Programming Language :: Python :: 3"
+]
+dependencies = [
+	"matplotlib",
+	"numpy"
+]
+requires-python = ">=3.0"
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..83b3f95
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,4 @@
+from setuptools import setup
+
+if __name__ == "__main__":
+	setup()
diff --git a/src/diffusion2d_pkg/__init__.py b/src/diffusion2d_pkg/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/diffusion2d.py b/src/diffusion2d_pkg/diffusion2d.py
similarity index 80%
rename from diffusion2d.py
rename to src/diffusion2d_pkg/diffusion2d.py
index c0c6083..84af1d3 100644
--- a/diffusion2d.py
+++ b/src/diffusion2d_pkg/diffusion2d.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
+from .output import create_plot, output_plots
 
 # plate size, mm
 w = h = 10.
@@ -68,14 +69,7 @@ def do_timestep(u_nm1, u, D, dt, dx2, dy2):
     # Create figure
     if n in n_output:
         fig_counter += 1
-        ax = fig.add_subplot(220 + fig_counter)
-        im = ax.imshow(u.copy(), cmap=plt.get_cmap('hot'), vmin=T_cold, vmax=T_hot)  # image for color bar axes
-        ax.set_axis_off()
-        ax.set_title('{:.1f} ms'.format(n * dt * 1000))
+        ax, im = create_plot(u,n,dt,T_cold,T_hot,fig,fig_counter)
 
 # Plot output figures
-fig.subplots_adjust(right=0.85)
-cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
-cbar_ax.set_xlabel('$T$ / K', labelpad=20)
-fig.colorbar(im, cax=cbar_ax)
-plt.show()
+output_plots(fig, im)
diff --git a/src/diffusion2d_pkg/output.py b/src/diffusion2d_pkg/output.py
new file mode 100644
index 0000000..689f6fd
--- /dev/null
+++ b/src/diffusion2d_pkg/output.py
@@ -0,0 +1,15 @@
+import matplotlib.pyplot as plt
+
+def create_plot(u, n, dt, T_cold, T_hot, fig, fig_counter):
+	ax=fig.add_subplot(220+fig_counter)
+	im=ax.imshow(u.copy(),cmap=plt.get_cmap('hot'),vmin=T_cold,vmax=T_hot)
+	ax.set_axis_off()
+	ax.set_title('{:.1f} ms'.format(n*dt*1000))
+	return ax, im
+
+def output_plots(fig, im):
+	fig.subplots_adjust(right=0.85)
+	cbar_ax=fig.add_axes([0.9,0.15,0.03,0.7])
+	cbar_ax.set_xlabel('$T$/K',labelpad=20)
+	fig.colorbar(im,cax=cbar_ax)
+	plt.show()