diff --git a/README.md b/README.md
index fe548354..174a36a1 100644
--- a/README.md
+++ b/README.md
@@ -8,12 +8,68 @@ The code used in this exercise is based on [Chapter 7 of the book "Learning Scie
 
 ## Project description
 
+This an example python library, that solves diffusion equations in 2D.
+
+The code solves the equations over a square domain, which is at a fixed temperature.
+A higher temperature is applied to a circular area in the center.
+
+The method used for solving is the Finite Difference Method.
+
+Some parameters, like the thermal diffusivity, may be changed.
+
+The code creates four plots at some timesteps of the simulation.
+
+For the theoretical background refer to [Chapter 7 of the book "Learning Scientific Programming with Python"](https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/)
+
 ## Installing the package
 
 ### Using pip3 to install from PyPI
 
+Use 
+```shell
+ $ pip3 install \
+        --index-url https://test.pypi.org/simple/ \
+        --extra-index-url https://pypi.org/simple/ \
+        krampfkn_diffusion2d
+```
+
+to install the package from TestPyPI.
+
+
 ### Required dependencies
 
+The requirements for using this code are:
+ - `matplotlib>=3.9`
+ - `numpy>=2.1`
+
+All dependencies should be available in reasonably recent versions.
+In some cases older versions might work, but have not been tested.
+
 ## Running this package
 
+A minimal example to run this code is given below:
+
+```python
+from krampfkn_diffusion2d.diffusion2d import solve
+
+solve(
+    4., # D: Thermal diffusivity
+    .1, # dx: Discrete Interval for x-direction
+    .1  # dy: Discrete Interval for y-direction
+)
+```
+
 ## Citing
+
+When citing this repository use this reference:
+
+```bibtex
+@misc{diffusion2d_2024,
+    url = {https://github.com/Simulation-Software-Engineering/diffusion2D},
+    author = {Krampf, Kilian \and Desai, Ishaan},
+    title = {diffusion2D},
+    subtitle = {Example python code for packaging},
+    year = {2024},
+}
+```
+
diff --git a/diffusion2d.py b/diffusion2d.py
deleted file mode 100644
index c0c6083a..00000000
--- a/diffusion2d.py
+++ /dev/null
@@ -1,81 +0,0 @@
-"""
-Solving the two-dimensional diffusion equation
-
-Example acquired from https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/
-"""
-
-import numpy as np
-import matplotlib.pyplot as plt
-
-# plate size, mm
-w = h = 10.
-# intervals in x-, y- directions, mm
-dx = dy = 0.1
-# Thermal diffusivity of steel, mm^2/s
-D = 4.
-
-# Initial cold temperature of square domain
-T_cold = 300
-
-# Initial hot temperature of circular disc at the center
-T_hot = 700
-
-# Number of discrete mesh points in X and Y directions
-nx, ny = int(w / dx), int(h / dy)
-
-# Computing a stable time step
-dx2, dy2 = dx * dx, dy * dy
-dt = dx2 * dy2 / (2 * D * (dx2 + dy2))
-
-print("dt = {}".format(dt))
-
-u0 = T_cold * np.ones((nx, ny))
-u = u0.copy()
-
-# Initial conditions - circle of radius r centred at (cx,cy) (mm)
-r = min(h, w) / 4.0
-cx = w / 2.0
-cy = h / 2.0
-r2 = r ** 2
-for i in range(nx):
-    for j in range(ny):
-        p2 = (i * dx - cx) ** 2 + (j * dy - cy) ** 2
-        if p2 < r2:
-            u0[i, j] = T_hot
-
-
-def do_timestep(u_nm1, u, D, dt, dx2, dy2):
-    # Propagate with forward-difference in time, central-difference in space
-    u[1:-1, 1:-1] = u_nm1[1:-1, 1:-1] + D * dt * (
-            (u_nm1[2:, 1:-1] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[:-2, 1:-1]) / dx2
-            + (u_nm1[1:-1, 2:] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[1:-1, :-2]) / dy2)
-
-    u_nm1 = u.copy()
-    return u_nm1, u
-
-
-# Number of timesteps
-nsteps = 101
-# Output 4 figures at these timesteps
-n_output = [0, 10, 50, 100]
-fig_counter = 0
-fig = plt.figure()
-
-# Time loop
-for n in range(nsteps):
-    u0, u = do_timestep(u0, 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))
-
-# 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()
diff --git a/krampfkn_diffusion2d/__init__.py b/krampfkn_diffusion2d/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/krampfkn_diffusion2d/diffusion2d.py b/krampfkn_diffusion2d/diffusion2d.py
new file mode 100644
index 00000000..20db4e3f
--- /dev/null
+++ b/krampfkn_diffusion2d/diffusion2d.py
@@ -0,0 +1,69 @@
+"""
+Solving the two-dimensional diffusion equation
+
+Example acquired from https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/
+"""
+
+import numpy as np
+from krampfkn_diffusion2d.output import output_plots
+
+
+def solve(D = 4., dx = .1, dy = .1):
+    # plate size, mm
+    w = h = 10.
+
+    # Initial cold temperature of square domain
+    T_cold = 300
+
+    # Initial hot temperature of circular disc at the center
+    T_hot = 700
+
+    # Number of discrete mesh points in X and Y directions
+    nx, ny = int(w / dx), int(h / dy)
+
+    # Computing a stable time step
+    dx2, dy2 = dx * dx, dy * dy
+    dt = dx2 * dy2 / (2 * D * (dx2 + dy2))
+
+    print("dt = {}".format(dt))
+
+    u0 = T_cold * np.ones((nx, ny))
+    u = u0.copy()
+
+    # Initial conditions - circle of radius r centred at (cx,cy) (mm)
+    r = min(h, w) / 4.0
+    cx = w / 2.0
+    cy = h / 2.0
+    r2 = r ** 2
+    for i in range(nx):
+        for j in range(ny):
+            p2 = (i * dx - cx) ** 2 + (j * dy - cy) ** 2
+            if p2 < r2:
+                u0[i, j] = T_hot
+
+
+    def do_timestep(u_nm1, u, D, dt, dx2, dy2):
+        # Propagate with forward-difference in time, central-difference in space
+        u[1:-1, 1:-1] = u_nm1[1:-1, 1:-1] + D * dt * (
+                (u_nm1[2:, 1:-1] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[:-2, 1:-1]) / dx2
+                + (u_nm1[1:-1, 2:] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[1:-1, :-2]) / dy2)
+
+        u_nm1 = u.copy()
+        return u_nm1, u
+
+
+    # Number of timesteps
+    nsteps = 101
+    # Output 4 figures at these timesteps
+    n_output = [0, 10, 50, 100]
+    output_data = []
+
+    # Time loop
+    for n in range(nsteps):
+        u0, u = do_timestep(u0, u, D, dt, dx2, dy2)
+
+        # Create figure
+        if n in n_output:
+            output_data.append((n * dt * 1000, u.copy()))
+
+    output_plots(output_data, range(T_cold, T_hot))
\ No newline at end of file
diff --git a/krampfkn_diffusion2d/output.py b/krampfkn_diffusion2d/output.py
new file mode 100644
index 00000000..74ce1ead
--- /dev/null
+++ b/krampfkn_diffusion2d/output.py
@@ -0,0 +1,33 @@
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.image import AxesImage
+from numpy import ndarray
+
+def create_plot(data: ndarray, ax: Axes, title: str, boundaries: range) -> AxesImage:
+    im = ax.imshow(data.copy(), cmap=plt.get_cmap('hot'), vmin=boundaries.start, vmax=boundaries.stop)  # image for color bar axes
+    ax.set_axis_off()
+    ax.set_title(title)
+    return im
+
+def output_plots(data: [(float, ndarray)], boundaries: range) -> Figure:
+    fig_counter = 0
+    fig = plt.figure()
+
+    if len(data) == 0:
+        return fig
+
+    for (ts,u) in data:
+        print(ts, u)
+        fig_counter += 1
+        ax = fig.add_subplot(220 + fig_counter)
+        im = create_plot(u, ax, '{:.1f} ms'.format(ts), boundaries)
+
+    # 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()
+
+    return fig
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..0a10c9c4
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools"]
+
+[project]
+name = "krampfkn_diffusion2d"
+version = "0.0.4"
+description = "Example implementation for solving diffusion equations"
+readme = "README.md"
+license = { file = "LICENSE" }
+keywords = ["simulation", "diffusion"]
+classifiers = [
+    "Programming Language :: Python :: 3"
+]
+dependencies = [
+    "matplotlib>=3.9",
+    "numpy>=2.1",
+]
+
+[project.urls]
+Homepage = "https://github.com/Simulation-Software-Engineering/diffusion2D"
+Repository = "https://github.com/Simulation-Software-Engineering/diffusion2D"
+
+[project.entry-points."simulation"]
+solve = "krampfkn_diffusion2d:solve"